Update main.py

added some test code for distance