Update main.py

replace big for loop with a function that does the same thing