Cleaned the code / Refactoring