# flipLoc.py
# input: 500
''' 
  NUM _TRIALS tests of of n steps
  where the final landing position (loc) is
  recorded in the fin list for each trial.
'''

import matplotlib.pyplot as plt
import random, math

NUM_TRIALS = 5000

def probN(n, x):
 # after n steps, what is prob of landing on x?
 # x is counted from the left side
 return math.comb(n, x) / (2**n)


nf = int(input("no. steps? "))

fin = [0]*(2*nf+1)
offset = nf
minPos = 0; maxPos = 0

for _ in range(NUM_TRIALS):
  loc = 0
  for i in range(nf):
    flip = random.choice(['H', 'T'])
    if flip == 'H':
      loc += 1
    else:
      loc -= 1

  fin[offset+loc] += 1
  if loc > maxPos:
    maxPos = loc
  if loc < minPos:
    minPos = loc

print(f"min: {minPos}; max: {maxPos}")

print("Prob of being at origin:", probN(nf, nf//2))

tot = sum(fin)
probs = [f/tot for f in fin]

plt.figure(figsize=(10, 5))
plt.bar(range(-nf, nf+1), probs, width=1)
plt.xlim([minPos-1, maxPos+2])
plt.xlabel("Position")
plt.ylabel('Prob')
plt.title("Prob of being at pos after " + str(nf) + " steps")
plt.show()