On going refactoring
Cleanup code