Code refactoring

Cleanup the code and refactor it to make it more readable