Added some comments and code cleanup