Added some inline docs and cleaned up code