Code cleanup

Removed unused imports and simple cleaning of the code