'''
    DS2000
    Spring 2023
    Sample code from class - Jan 20th
    
    Plotting to understand the relationship
    between overtime and total income
    (boston firefighter salaries)

'''

import matplotlib.pyplot as plt

SALARY_FILE = "firefighter_salaries.txt"

def main():
    # Gather data - read in each employee's overtime,
    # regular pay, and color
    with open(SALARY_FILE, "r") as infile:
        overtime1 = float(infile.readline())
        regular1 = float(infile.readline())
        color1 = infile.readline().strip()
        
        overtime2 = float(infile.readline())
        regular2 = float(infile.readline())
        color2 = infile.readline().strip()
        
        overtime3 = float(infile.readline())
        regular3 = float(infile.readline())
        color3 = infile.readline().strip()
    
    # printing sanity check
    print("Data so far:", overtime1, regular1, overtime2, regular2)
    
    # Computation - figure out everyone's total salary and the avg
    total1 = overtime1 + regular1
    total2 = overtime2 + regular2
    total3 = overtime3 + regular3
    total_avg = (total1 + total2 + total3) / 3
    
    # Communication - plot the overtime vs total pay
    # overtime = independenet variable (x axis)
    # total pay = dependent variable (y axis)
    plt.plot(overtime1, total1, "o", color = color1)
    plt.plot(overtime2, total2, "o", color = color2)
    plt.plot(overtime3, total3, "o", color = color3)
    
    # Change the xlimit and ylimit
    plt.xlim(5000, 30000)
    plt.ylim(100000, 175000)
    
    # Add labels and a title
    plt.xlabel("Overtime ($)")
    plt.ylabel("Total income ($)")
    plt.title("Overtime vs Total Income for Boston Firefighters (2021)")
    
    # PLot a horizontal line for the average
    plt.axhline(total_avg, color = "deeppink", label = "average total income")
    
    # Add a legend (uses labels from above)
    plt.legend()
    
      
    
main()