changes to allow import as submodule

 - Add returns to functions
 - Move print statement to main