''''
    DS2000
    Fall 2024
    Sample code from class - 10/15/24

    Building on dunks code from Friday, with a few things to fix...
    - split is not perfect (separate city, state name)
    - print when we find a city, but we don't count
    - plot individual points instead of lists (not *awful* but could be better!)
'''

import matplotlib.pyplot as plt

FILENAME = "dunks.csv"
LONG_POS = 0
LAT_POS = 1
CITY_POS = 2

def process_str(whole_str, long, lat):
    '''
        parameters: a string, two ints (positions)
        returns: a list of strings, except at the two given positions
        does: splits the given string into a list of strings, converts
              the two positions to floats
    '''
    lst = whole_str.split(",")
    lst[lat] = float(lst[lat])
    lst[long] = float(lst[long])
    return lst

def find_city(lst, city, city_pos):
    '''
        parameters: lst of floats/strings, string for city, int for city_pos
        returns: boolean
        does: finds if the given city exists in the list at the given position
    '''
    if city in lst[city_pos]:
        return True
    else:
        return False


def main():
    # prompt the user for a city
    city = input("What city to look for Dunks in?\n")

    # start a counter varialbe at 0
    count = 0

    # step one - gather data by reading from the file
    with open(FILENAME, "r") as infile:
        infile.readline()

        for line in infile:
            # turn line into a list, because line is a string right now
            lst = process_str(line, LONG_POS, LAT_POS)

            # Is this dunks in the city the user gave us?
            if find_city(lst, city, CITY_POS):
                print(lst)
                count += 1

            # communication - plot the lat/long
            plt.plot(lst[LONG_POS], lst[LAT_POS], "x", color = "orange")

    # How many did I find in a given city?
    print(count, "Dunks in", city)

    # more communication, final touches and show the plot
    plt.title("Dunkin Donuts Locations in the US")
    plt.xlabel("Longitude")
    plt.ylabel("Latitude")
    plt.show()

main()