added some test and cleaned up a bit more the code. Need more tests still