Update main.py

Refactored the code a bit so that it is a bit more DRY. Probably more can be done though for now it seems a bit more readable.

Additionally the comments have been modified slightly so that they respect the margin at character 80.