Cleaned up code and added a few missing comments