Code clean up in main.py