Cleanup
fixes
Fixes and improving the organization of the code