Cleanup
Fixing main class