Code refactoring

Much improved design
More tests