##################################
#   Python code for question 7   #
##################################

from scipy import *
from scipy import linalg
import pylab

# We will be using the code from the computer practicals to compute
# the invariant distribution 
def invariant_distribution(K):
    eigen = linalg.eig(K.T)         # Find eigenvals and vecs
    idx = eigen[0].argmax()         # Find which eigenval is largest (i.e. 1)
    mu = eigen[1][:,idx]            # Extract corresponding eigenvec
    mu = mu / sum(mu)	            # Normalise distribution
    return mu

# Define cross-tabulation matrix
C = array([ [50., 19., 26., 8., 7., 11., 6., 2],
            [16., 40., 34., 18., 11., 20., 8., 3],
            [12., 35., 65., 66., 35., 88., 23., 21],
            [11., 20., 58., 110., 40., 183., 64., 32],
            [2., 8., 12., 23., 25., 46., 28., 12],
            [12., 28., 102., 162., 90., 554., 230., 177],
            [0., 6., 19., 40., 21., 158., 143., 71],
            [0., 3., 14., 32., 15., 126., 91., 106] ])

# Absolute frequencies of occupations of fathers is row-wise sum
fathers = C.sum(1)
# Absolute frequencies of occupations of sons is column-wise sum
sons = C.sum(0)

# Compute the relative frequencies
freq_fathers = fathers / sum(fathers)
freq_sons = sons / sum(sons)
print(freq_fathers)
print(freq_sons)

# To turn C into a transition kernel we need to divide each row by its sum
K = (C.T / fathers).T

# Compute the invariant distribution
mu = invariant_distribution(K)
print(mu)

# Visualise the result using a barplot
pylab.bar(linspace(1,8,8)-.4,freq_fathers,0.25,color="r", 
          label="Actual distribution of fathers")
pylab.bar(linspace(1,8,8)-.1,freq_sons,0.25,color="g",
          label="Actual distribution of sons")
pylab.bar(linspace(1,8,8)+.2,mu,0.25,color="b",
          label="Invariant distribution")
pylab.legend(loc=2)
pylab.show()

####################################################

##################################
#   Python code for question 8   #
##################################

from scipy import *
from scipy import random
import pylab

# Simulates a path of length n+1 from the Wright-Fisher model
# Initial distribution is given by the array initial (length 2)
def simulate_wright_fisher(initial, n):
    result = zeros((n+1,2))                      # Create array to store result
    result[0,:] = initial                        # Set first line to initial dist'n
    two_N = sum(initial)                         # Compute 2*N
    for t in xrange(n):
        result[t+1,0] = random.binomial(two_N, result[t,0]/two_N)
                                                 # Draw X[t+1]|X[t]=x[t]
        result[t+1,1] = two_N - result[t+1,0]    # Compute 2*N - X[t+1] 
    return result


wf_path = simulate_wright_fisher([300,100], 1000)
print wf_path

pylab.figure()
pylab.plot(wf_path[:,0], 'bs-', wf_path[:,1], 'rs-')
pylab.show()