Refactored code so it sucks less