'''
    DS2000
    Fall 2024
    Sample code from class - iterating over lists when we need to modify or
        do two lists at the asame time

    Reminder:
    - for num in lst: (iteration by value)
    - for i in range(len(lst)): (iteration by position)

    Goals for today:
    - read in each line from the csv file (years, khoury, dmsb)
    - create three lists (probably of integers)
    - Compare khoury vs dmsb for each year
'''

import matplotlib.pyplot as plt

FILENAME = "neu_enrollments.csv"

def main():
    # step one: gather data by reading from the CSV file
    with open(FILENAME, "r") as infile:
        line = infile.readline()
        years = line.split(",")

        line = infile.readline()
        khoury = line.split(",")

        line = infile.readline()
        dmsb = line.split(",")

    # Sanity check print, what is in lists?
    # yes! I have three lists of strings
    print("Just after reading in from csv file...")
    print(years)
    print(khoury)
    print(dmsb)

    # Turn all our lists into lists of ints
    # I need to modify the lists, iteration by position
    for i in range(len(years)):
        years[i] = int(years[i])
        khoury[i] = int(khoury[i])
        dmsb[i] = int(dmsb[i])

    # Sanity check print: do I have lists of ints?
    print("After converting to int, here's what's in my lists...")
    print(years)
    print(khoury)
    print(dmsb)

    # How many times did Khoury have more students than DMSB?
    # Human confirmation from looking at the file... this should be zero
    count_khoury = 0
    for i in range(len(khoury)):
        if khoury[i] > dmsb[i]:
            count_khoury = count_khoury + 1

    # Step three: communication - print the khoury vs dmsb results
    for i in range(len(khoury)):
        print("In", years[i], "Khoury had enrollment of", khoury[i],
              "while dmsb had enrollment of", dmsb[i])
    print("Overall there were", count_khoury, "years when Khoury had more students!")

    # Plot the enrollments over time
    # PLot a list: x values are positions, y values are values from the list
    # Add a label and plt.legend() to make a legend pop up!
    plt.plot(dmsb, color = "dodgerblue", label = "DMSB")
    plt.plot(khoury, color = "magenta", label = "Khoury")
    plt.legend()
    plt.title("Total undergrad enrollments at DMSB vs Khoury, 2018-2023")

    # Here's an example of not using the default xticks.
    # Give matplotlib the positions and the labels for each position
    year_positions = [0, 1, 2, 3, 4, 5]
    plt.xticks(year_positions, years)

    # label the x and y axes
    plt.xlabel("Enrollment Year")
    plt.ylabel("Total Undergrad Enrollment")
    plt.show()



main()