added a few more tests
slightly different training
beautify the code