#   gambler's ruin program
import sys, random
import matplotlib.pyplot as plt

def gamble(p1, g1_amt, g2_amt):
    g1_seq = [g1_amt]
    g2_seq = [g2_amt]
    cur_g1 = g1_amt
    cur_g2 = g2_amt
    steps  = 0
    while cur_g1 > 0 and cur_g2 > 0 :
#
#   YOUR CODE GOES HERE to execute random step
#    Hint : use random.uniform(0,1) and if-else
#     statement. Check if rand number is less than
#     p1. If so, g1 wins, else g2 wins.
#
#   
        g1_seq.append(cur_g1)          #  Add to g1 and g2 sequence lists
        g2_seq.append(cur_g2)  
        steps += 1 
    return g1_seq, g2_seq, steps

   
p  = float(input('Enter probability of gambler1 winning:  '))
if p<= 0.0 or p>=1 :
    print 'Probability entered =', p, ' prob must be between 0 and 1'
    sys.exit(0)
g1 = int(input('Enter bankroll for gambler1: '))
g2 = int(input('Enter bankroll for gambler2: '))
#   Execute gambler's game
#
g1_walk,g2_walk, num_steps = gamble(p, g1, g2) 
#
#   Plot Gambler 1 and 2 step sequences for game
#
x = range(num_steps+1)
plt.plot(x, g1_walk,'r-', label = 'Gambler1')
plt.plot(x, g2_walk,'b-', label = 'Gambler2')
plt.legend(loc= 'upper left')
plt.title(' Gambler\'s ruin: p='+ str(p)+ ', g1 = $'+ str(g1)+ ', g2 = \$= '+ str(g2))
plt.grid(True)  
plt.show()