Reorganized main.py and broke it into functions