Updated main.py with some code.