Refactor main.py

For Issue #87