# headsLeads.py
'''
  simulating NUM_TRIALS games of Heads or Tails, with each game consisting of NUM_FLIPS tosses. The distribution of the number of times that Peter is in the lead is given plotted, together with the arc sine pdf.

  Player A is said to be in the lead at time n if the random walk is above the t-axis at that time, or if the random walk is on the t-axis at time n but above the t-axis at time n−1. (At time 0, neither player is in the lead.)
'''

import random
import math
import collections
import itertools
import matplotlib.pyplot as plt

NUM_TRIALS = 10000
NUM_FLIPS = 80

def arcsinePDF(x):
  return 1/(math.pi*math.sqrt(x*(1-x)))

leads = []      # stores lead lengths
numEquals = 1   # when heads == tails
for i in range(NUM_TRIALS):
  hDiff = 0
  leadLen = 0
  for j in range(NUM_FLIPS):
    if random.choice(['H', 'T']) == 'H':
      hDiff += 1
    else:
      hDiff -= 1 
    if hDiff > 0:
      leadLen += 1
    elif hDiff == 0:
      numEquals += 1
  leads.append(leadLen)
print(f"Avg. No. Equals: {(numEquals/NUM_TRIALS):.1f}")

leadCounts = collections.Counter(leads)
leadCounts = dict(sorted(leadCounts.items()))
print(leadCounts)
xs = list(leadCounts.keys())
ys = list(leadCounts.values())
print(f"Positive % = {(sum(ys[1:])/NUM_TRIALS):.2%}")
maxYs = max(ys[1:])

plt.bar(xs[1:], ys[1:])
plt.xlabel("Lead Lengths")
plt.ylabel('Frequency')
plt.title("Frequency of Lead Lengths for " + str(NUM_FLIPS) + " flips")
ars = [arcsinePDF(i/(NUM_FLIPS+1)) for i in range(1, NUM_FLIPS+1)]
maxArs = max(ars)
plotArs = [a*maxYs/maxArs for a in ars]
plt.plot(xs[1:], plotArs, color='g', lw=2)
plt.show()