Tidy up a little and add some tests

- Move docstrings inside the functions (where they should be)
- Add some high level doctests
- Replace print statements with logging
- Tidy up some of the code