Cleaned up the code and added documentation.