'''
    buffy.py
    DS2000
    Spring 2022
    Code from class -- working with conditionals, random numbers,
                       and a little data viz to communicate how good or
                       bad the episodes of Buffy are                  
'''

import matplotlib.pyplot as plt
import random

def main():

    count = 0
    while count < 10:
        
        # Step one: gather data by generating random numbers
        # for season and episode
        season = random.randint(1, 7)
        if season == 1:
            ep = random.randint(1, 12)
        else:
            ep = random.randint(1, 22)
        
        # Step two: computations
        # Decide what height and color to plot this episode
        height = 5
        color = "darkorange"
        
        # If it's a season finale, it's good!
        if ep == 22 and season > 1:
            height += 3
            color = "greenyellow"
        elif ep == 12 and season == 1:
            height += 3
            color = "greenyellow"
        
        # If it's season 4, thren it's good
        # If it's season 4, ep 9 or 10, it's great
        if season == 4:
            height += 3
            color = "greenyellow"
            if ep == 9 or ep == 10:
                height += 2
                color = "lime"
                
        # If it's after saeson 5, it's not as good
        # EXCEPT FOR THE MUSICAL!
        if season == 6 and ep == 5:
            height += 10
            color = "magenta"
        elif season > 5:
            height -= 4
            color = "red"
        
        # Step three: communication -- visualize our plot
        print("Season:", season, "Episode:", ep)
        plt.plot(season * 22 + ep, height, 'o', color = color,
                 markersize = 15)
        plt.ylim(0, 12)
        plt.xlim(20, 180)
        
        # increase count for the loop
        count += 1
    
main()