Cleaned up code and fixed last few bugs