Clean code up (its still shitty)