Refactor code
Add some comments