# GamSurface.py
'''
  Gambler's Ruin Problem
  Specify N (the win pos)

  A range of start position and probabilities
  are used. For a given (p,i,N), the game is
  run NUM_TRIALS times, and the
  fraction of wins is returned. i is always in
  the range [0..N] in integer steps.

  A surface 3D plot is created with the XY plane
  being the p and i, and the z-axis is the 
  fraction of win.

  input: 10
'''

import random
import numpy as np
import matplotlib.pyplot as plt

NUM_TRIALS = 20

@np.vectorize
def calcWins(p, start, N):
  # print(p, start, N)
  numWins = 0
  for i in range(NUM_TRIALS):
    time, isWin = gamble(p, start, N)
    if isWin:
      numWins += 1
  return numWins/NUM_TRIALS

def gamble(p, start, N):
  t = 0
  pos = start
  while (pos > 0) and (pos < N):
    r = random.choices(['win', 'lose'], [p, 1-p])
    t += 1
    if r[0] == 'win':
      pos += 1
    else:
      pos -= 1
  return t, (pos == N)

N = int( input("N? "))
# probability ranges
pStart = 0.4; pEnd = 0.6   # close to 0.5
print("Prob range:", pStart, "--", pEnd)
ps = np.linspace(pStart, pEnd, 41)  
# start position ranges
starts = np.linspace(0, N, N+1)
X, Y = np.meshgrid(ps, starts)    # numpy output required
Z = calcWins(X, Y, N)  
ax = plt.axes(projection='3d')
ax.plot_surface(X, Y, Z, 
    cmap="Blues_r", edgecolor="Black", lw=0.2)
plt.xlabel('Prob')
plt.ylabel('Start Pos')
ax.set_zlabel('Win Frac')  # no plt.zlabel()
plt.title(f"Gambler's Ruin Wins (N={N})")
plt.show()