from scipy import *
from scipy import linalg
from numpy import random 
import pylab

# Draws one realisation from the discrete distribution given by probs
def discrete_sample(probs):
    return sum(cumsum(probs)<random.random_sample())

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

# Defintion of a class MarkovChain
class MarkovChain:
    # Constructor: creates a new MarkovChain object
    def __init__(self, K, initial):
        self.K = K
        self.initial = initial       
 
    # Compute m-step transition kernel
    def K_m_step(self, m): 
        Km = eye(self.K.shape[0])       # Create identity matrix of correct size
        for i in xrange(m):
            Km = dot(Km, self.K)        # Repeatedly multiply it by K
        return Km

    # Compute distribution of states at time m
    def distribution_at_m(self, m):
        return dot(self.initial, self.K_m_step(m))
                                        # Compute lambda'K(m)

    # Compute the invariant distribution
    def invariant_distribution(self):
        eigen = linalg.eig(self.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

    # Compute the invariant distribution (alternative implementation)
    def invariant_distribution_alternative(self):
        n = self.K.shape[0]             # Get number of rows
        A = concatenate((self.K.T - eye(n), [ones(n)]))
                                        # Compute K'-I and add row of ones
        b = zeros(n+1)                  # Set b = [0 ... 0 1]
        b[n] = 1                  
        mu = linalg.lstsq(A, b)[0]      # Solve this system of equations
        return mu

   # Draw path of length n from the Markov chain
    def sample_chain(self, n):
        sample_path = zeros(n)          # Set up vector to hold results
        sample_path[0] = discrete_sample(self.initial)
                                        # Draw X[0] from the initial distribution
        for t in xrange(1,n):           # Draw remaining X[t] from conditional
                                        # distribution if X[t+1] | X[t]=x[t]
           sample_path[t] = discrete_sample(K[sample_path[t-1],:])
                                       
        return sample_path



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

# Task 1

K = array([[0.9, 0.1],                  # Define exaple transition kernel
           [0.3, 0.7]]) 
lam = array([0.2, 0.8])                 # Set (arbitrary) initial distribution

myMC = MarkovChain(K, lam)              # Create MarkovChain object

print "2-step Transition Kernel:"
print myMC.K_m_step(2)                  # Print 2-step transition kernel

# Task 2

print "Distribution at Time t=2:"
print myMC.distribution_at_m(2)         # Print distribution of states at t=2

# Task 3

print "Invariant Distribution"
print myMC.invariant_distribution()     # Print invariant distribution

# Task 4
n = 20
distn = array([myMC.distribution_at_m(t) for t in xrange(n+1)])
                                        # Compute distribution of states
                                        # for t=0..n
pylab.figure()
pylab.plot(distn[:,0], 'bs-', distn[:,1], 'rs-')
pylab.show()


# Task 5

n = 100
mySample= myMC.sample_chain(n)          # Draw sample path of length n
pylab.figure()
pylab.plot(mySample, 'gs-')             # Plot the path
pylab.show()

# Task 6

print "Proportion of 0's:"
print float(sum(mySample==0))/n
print "Proportion of 1's:"
print float(sum(mySample==1))/n