Cleaned up some code and split things into classes