Refactor code into a more object oriented style