Refactored code and added comments