Wrote some comments and modified code a bit