Cleaned up a little bit... WIP