lots of refactoring
Fixed test cases