import random
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
#Assume you always pick the first door

def run():
    iterations = 100 #The bigger number it is, the more the numbers fall to its true average

    stayCorrect = 0 #raw score for staying
    stayPercentage = [] #rate of success if you were to stay

    switchCorrect = 0 #Raw score for the amount of wins when you switch
    switchPercentage = [] #rate of success if you were to switch

    numDoors = 10
    for cycle in range(1, iterations + 1):
        print("----------------------------------------------------------------------")
        print("iteration: " + str(cycle))

        #assign superCar, it is 1 in the array
        doors = [0 for _ in range(numDoors)] #create flexible number of doors
        doors[random.randrange(0, numDoors)] = 1 #Assign one of the doors as car

        
###########################################################################################################################
        #DEPCRECATED AFTER MAKING NUMBER OF DOORS FLEXIBLE
        # first = doors[0]
        #second = doors[1]
        #third = doors[2]
        # 
        # Create 3 different squares representing doors
                                             #1    #2    #3
        #print(f" ___   ___   ___\n|   | |   | |   |\n| {first} | | {second} | | {third} |\n|___| |___| |___|\n")
###########################################################################################################################


        
        
        #Remove all the doors except for 3
        newDoors = [doors[0]]
        extraZero = False
        for index in range(1, numDoors): #you go 2nd door -> the rest of the doors


            # In the game, Monty removes all the doors that are zeros leaving the one supercar. This is analagous to adding a single door
            # that has a supercar
            if (doors[index] == 1):
                newDoors.append(doors[index])
                break

        ###### You Stay
        if (newDoors[0] == 1):
            stayCorrect = stayCorrect + 1
            print("By staying, you were correct")
        ###### You switch
        elif (newDoors[1] == 1):
            switchCorrect = switchCorrect + 1
            print("By switching, you were correct")

        
        
        stayPercentage.append(stayCorrect / cycle)
        switchPercentage.append(switchCorrect / cycle)

        print("Percentage of wins when staying: " + str(stayPercentage[-1]))
        print("Percentage of wins when switching: " + str(switchPercentage[-1]))
    

    ####################################################################
    # visual representation 
    ####################################################################

    listIterations = np.arange(1, iterations + 1)
    plt.plot(listIterations, stayPercentage, label = "Stay")
    plt.plot(listIterations, switchPercentage, label = "Switch")


    stayDataframe = pd.DataFrame(stayPercentage)
    switchDataframe = pd.DataFrame(switchPercentage)
    sum = stayDataframe + switchDataframe
    
    plt.plot(listIterations, sum, label = "Sum")


    plt.xlabel('Trial')
    plt.ylabel('Percentage')
    plt.title('Rate of growth of switching and staying')

    plt.legend()
    plt.show()


run()