Cleaned up code from work on #44 and fixed #45