Code cleanup

Cleaning up the code and refactoring.