code refactoring and adding comments