Refactored a lot. Basically imports and moving functions around.
Fixed bad tests.