Fixed some small problems
Cleaned up code.