'''
    DS2000
    Spring 2023
    Sample code from class - coffee info!
    
    Tweaking code from Friday...
    - use the csv library to read a csv file (so we can search "Boston,MA")
    - use a function to turn a column into a list
'''

import csv
import matplotlib.pyplot as plt

DUNKS_FILE = "DD_US.csv"
CITY_COL = 2
LONG_COL = 0
LAT_COL = 1
LANEY_LAT = 42.30537
LANEY_LONG = -71.061

def read_csv(filename):
    ''' Function: read_csv
        Parameters: filename
        Returns: 2d list of strings
        Does: reads each line of the file to create a 2d list.
    '''
    lst = []
    with open(filename, "r") as infile:
        csvfile = csv.reader(infile)
        next(csvfile)
        for line in csvfile:
            lst.append(line)
    return lst
    
def col_to_lst(lst, col):
    ''' Function: col_to_lst
        Parameters: 2d list, int
        Returns: 1d list
        Does: Create a list made up of every row at the given column
    '''
    new_lst = []
    for row in lst:
        new_lst.append(row[col])
    return new_lst

def get_local(lst, city, col):
    ''' Function: get_local
        Parameters: 2d list of strings, string for city, int
        Returns: 2d list of strings
        Does: Look for given city at the position given by int
              in every row, returns a filtered 2d list
    '''
    smaller = []
    for row in lst:
        if city in row[col]:
            smaller.append(row)
    return smaller

def main():
    # Gather data - read from the file
    data = read_csv(DUNKS_FILE)
    
    # Computation - isolate lat/long into separate lists
    lats = col_to_lst(data, LAT_COL)
    longs = col_to_lst(data, LONG_COL)
    
    # Turn lists of strings into lists of floats
    lats = [float(lat) for lat in lats]
    longs = [float(long) for long in longs]
    
    # communicate - plot all the DD loatoins
    plt.scatter(longs, lats, color = "orange")
    plt.plot(LANEY_LONG, LANEY_LAT, "*", color = "fuchsia", markersize = 15)
    
    # Gather data - what city to look at?
    city = input("What city?\n")
    
    # Computations - find all the DD in that city
    city_results = get_local(data, city, CITY_COL)
    
    # Communication - print out all the locations
    print(city_results)
    print("How many Dunks in that city?", len(city_results))
    print("How many Dunks in US?", len(data))
    
main()