Minor refactoring. Added some comments as well.