Cleaned up and documented the code, take 2