# montyhall.py

import matplotlib.pyplot as plt
import random 

NUM_ITERS = 100000
NUM_DOORS = 3

trials = range(NUM_ITERS)
round1 = range(NUM_DOORS)  # door choices in round 1
stayCount, chgCount = 0, 0
stayProbs = []
chgProbs = []
for i in trials:
  carDoor = random.choice(round1)
  playerPick1 = random.choice(round1) 
  montyChoices = [n for n in round1 
                     if n!= carDoor and n!= playerPick1]
  goatDoor = random.choice(montyChoices) # what Monty opens
  round2 = [n for n in round1 
                     if n!= playerPick1 and n!=goatDoor]
  playerPick2 = random.choice(round2)
  if playerPick1 == carDoor: 
    stayCount += 1    # would win by staying
  if playerPick2 == carDoor: 
    chgCount += 1     # would win by switching
  p1 = stayCount/(i+1)
  p2 = chgCount/(i+1)
  stayProbs.append(p1)
  chgProbs.append(p2)
  
plt.plot(trials, chgProbs, label="Switch")
plt.plot(trials, stayProbs, label="Stay")
plt.xscale('log')
plt.xlabel('No. trials')
plt.ylabel('Winning Prob.')
plt.legend()
plt.title("Monty Hall with " + str(NUM_DOORS) + " doors")
plt.grid('on')
plt.show()