Refactor layout to be more pythonic

Use proper main() function and function out some of the messy logic in
main().