Finished refactoring to more object oriented design