Adding Comments
Refactoring a little
Added some test cases