Update main.py

Added comments and removed test code