Clean up and added comments to code