cleaned up code and added more tests.