%pylab inline

p = 0.6
q = 1-p
n1 = 3
n2 = 3
P1 = (1-(q/p)**n1) / (1.-(q/p)**(n1+n2))
P2 = ((q/p)**n1 - (q/p)**(n1+n2)) / (1.-(q/p)**(n1+n2))
print P1,P2, P1+P2

from numpy import random
random.seed(12345)


max_steps = 500  # maximum number of steps for each random walk

def walk(n1,n2,p,verbose=True):
    n1path = [n1]   # keep track of n1 each step of the walk
    for nstep in range(1,max_steps+1):
        r = random.random(1)
        if r <= p:
            n1 = n1+1
            n2 = n2-1
        else:
            n1 = n1-1
            n2 = n2+1
        n1path.append(n1)
        if verbose:
            print "In step %i, r = %7.4f and n1 = %i, n2 = %i" % (nstep,r,n1,n2)
        if min(n1,n2) == 0:
            break
    if verbose:
        print "Stopped after %i steps with n1 = %i, n2 = %i" % (nstep,n1,n2)
    if n1 == 0:
        winner = 2
    elif n2 == 0:
        winner = 1
    else:
        winner = 0
    return nstep, winner, n1path
    

n1 = 4
n2 = 4
nstep, winner, n1path = walk(n1, n2,.6)
plot(n1path,'o-')
xlabel('Step in walk')
ylabel('Number of pennies player 1 has')
yt = yticks(range(n1+n2+1))

random.seed(1234)
for k in range(7):
    nstep, winner, n1path = walk(5,5,.6,verbose=False)
    plot(n1path,'o-')
xlabel('Step in walk')
ylabel('n1')


p = 0.6
n1 = 3
n2 = 3

wins = [0,0,0]
kwalks = 1000
for k in range(kwalks):
    nsteps, winner, n1path = walk(n1, n2, p, verbose=False)
    wins[winner] += 1
frac1wins = wins[1] / float(kwalks)
frac2wins = wins[2] / float(kwalks)

if wins[0] > 0:
    print "Warning: %i walks out of %i did not result in a win by either player" \
        % (wins[0],kwalks)
    
print "Player 1 won %s fraction of the time, Player 2 won %s fraction of the time" \
        % (frac1wins,frac2wins)

q = 1-p
P1 = (1-(q/p)**n1) / (1.-(q/p)**(n1+n2))
P2 = 1 - P1
print "True probabilities are P1 = %7.4f, P2 = %7.4f" % (P1,P2)

    

p = 0.8
q = 1 - p
n1 = 50
n2 = 50

wins = [0,0,0]
kwalks = 1000
nsteps_total = 0
for k in range(kwalks):
    nsteps, winner, n1path = walk(n1, n2, p, verbose=False)
    nsteps_total += nsteps
    wins[winner] += 1
    
if wins[0] > 0:
    print "Warning: %i walks out of %i did not result in a win by either player" \
        % (wins[0],kwalks)
        
print "The average path length is %7.4f" % (nsteps_total/float(kwalks))

mean_length = (n1/(q-p)) - ((n1+n2)/(q-p))*(1-(q/p)**n1)/(1-(q/p)**(n1+n2))
print "True mean path length is %7.4f" % mean_length