Refactors to remove most of the code. Yey!