# GamTime.py
'''
  Gambler's Ruin Problem
  Specify p, i, N (prob, start pos, win pos)
  The game is run only once, and the progress
  of the gambler (i.e. his position) is plotted
  against time. The game stops when the gambler
  reaches N (a win) or 0 (ruin).

  inputs: 0.6
          2 10  or 20 100

  inputs: 0.49
          2 10  or  20 100

  inputs: 0.49
          75 100
'''

import random
import matplotlib.pyplot as plt


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

p = float( input("p? "))
i, N = map(int, input("i, N? ").split())
ts, locs = gambleTime(p, i, N)
print("Finish time:", ts[-1])
plt.axhline(y=N, color='g', linestyle='--')
plt.axhline(y=i, color='k', linestyle='--')
plt.axhline(y=0, color='r', linestyle='--')
plt.plot(ts, locs)
plt.title(f"Gambler's Ruin Time (p={p},i={i},N={N})")
plt.xlabel('Time')
# plt.xticks(range(0, ts[-1]+1))
plt.ylabel('Pos')
plt.show()