Tidy up a little and add some tests

[robyoung]

• 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

[alexrudnick]

Looks good! Thank you! :)