from __future__ import division
import numpy as np
from scipy.stats import binom
import matplotlib.pyplot as plt
from mc_tools import mc_compute_stationary, mc_sample_path
from quantecon.mc_tools import MarkovChain
import quantecon as qe

%matplotlib inline

def KMR_2x2_P_simultaneous(N, p, epsilon):
          P = np.empty((N+1, N+1), dtype=float)
          for x in range(N+1):
                                                           
            if x/N < p: 
                P[x,:]  = binom.pmf(range(N+1), N, epsilon/2)
            elif x/N == p:
                P[x,:]  = binom.pmf(range(N+1), N, 1/2)           
            else:
                P[x,:]  = binom.pmf(range(N+1), N, 1-epsilon/2) 
          return P
      

#Parameter
p = 1/3
N = 6
epsilon = 0.1

kmr_simul=KMR_2x2_P_simultaneous(N, p, epsilon)

print(kmr_simul)

def KMR_2x2_P_sequential(N, p, epsilon):
          P = np.zeros((N+1, N+1), dtype=float)
          
          P[0, 0], P[0, 1] = 1 - epsilon * (1/2), epsilon * (1/2)    
          P[N, N-1], P[N, N] = epsilon * (1/2), 1 - epsilon * (1/2) 
     
          
          for x in range(1, N):
             
                if x/N >p:
                    P[x, x-1] = (x/N)*epsilon * (1/2)
                    P[x, x+1] = (1-(x/N))*(1 - epsilon*(1/2))
                    P[x, x] = 1 - P[x, x-1] - P[x, x+1] 
                elif x/N<p:
                    P[x, x-1] =(x/N)*(1 - epsilon*(1/2))
                    P[x, x+1] =(1-(x/N))*epsilon * (1/2)
                    P[x, x] = 1 - P[x, x-1] - P[x, x+1]     
                else:
                    P[x, x-1] = (x/N)*(1/2)
                    P[x, x+1] = (1-(x/N))*(1/2)
                    P[x, x] = 1 - P[x, x-1] - P[x, x+1]           
          return P
               
               

kmr_seq=KMR_2x2_P_sequential(N, p, epsilon)

print(kmr_seq)

mc_kmr_simul = MarkovChain(kmr_simul)
path_kmr_simul = mc_kmr_simul.simulate(init=0, sample_size=10000)

mc_kmr_seq = MarkovChain(kmr_seq)
path_kmr_seq = mc_kmr_seq.simulate(init=0, sample_size=10000)

fig, ax = plt.subplots(figsize=(9, 6))
ax.grid()
ax.plot(path_kmr_simul, label=r'Transition of state $x$')
ax.legend(loc='lower right')

fig, ax = plt.subplots(figsize=(9, 6))
ax.grid()
ax.plot(path_kmr_seq, label=r'Transition of state $x$')
ax.legend(loc='lower right')

kmr_simul_m=qe.MarkovChain(kmr_simul)
kmr_simul_s=kmr_simul_m.stationary_distributions
print(kmr_simul_s)


kmr_seq_m=qe.MarkovChain(kmr_seq)
kmr_seq_s=kmr_seq_m.stationary_distributions
print(kmr_seq_s)

"""
Average dist of the states.

"""

def kmr_simulate_dist(matrix, init, T, num_reps):
    
    mc = MarkovChain(matrix)
    x = []
    for i in range(num_reps):
        x.append(mc.simulate(init=init, sample_size=T+1))
    x_array = np.array(x)
    y = np.empty(matrix.shape[0])
    for i in range(matrix.shape[0]):
        y[i] = (x_array == i).sum() / ((T + 1) * num_reps) 
    return y

"""
Compute a fraction of states.

"""


def kmr_simul_simulate_state(N, p, init, T, eps_min, eps_max, eps_length, num_reps, state):
   
     epsilons = np.linspace(eps_min, eps_max, eps_length)
   
     fraction_sim_simul = np.zeros(eps_length)
     for i, epsilon in enumerate(epsilons):
         fraction_sim_simul[i] = kmr_simulate_dist(KMR_2x2_P_simultaneous(N, p, epsilon), init, T, num_reps)[state]
     return fraction_sim_simul

def kmr_seq_simulate_state(N, p, init, T, eps_min, eps_max, eps_length, num_reps, state):
   
     epsilons = np.linspace(eps_min, eps_max, eps_length)

     fraction_sim_seq = np.zeros(eps_length)
     for i, epsilon in enumerate(epsilons):
         fraction_sim_seq[i] = kmr_simulate_dist(KMR_2x2_P_sequential(N, p, epsilon), init, T, num_reps)[state]
     return fraction_sim_seq

frac_simul_6=kmr_simul_simulate_state(N, p, init=0, T=1000, eps_min=0.00001, eps_max=0.1, eps_length=100, num_reps=100,state=6)

fig, ax = plt.subplots(figsize=(9, 6))
ax.grid()
ax.set_xlabel(r'Value of $\varepsilon$')
ax.set_ylabel('Fraction of state 6')
ax.set_xlim(eps_min, eps_max)
ax.plot(epsilons, frac_simul_6)

frac_seq_6=kmr_seq_simulate_state(N, p, init=0, T=1000, eps_min=0.00001, eps_max=0.1, eps_length=100, num_reps=100,state=6)

fig, ax = plt.subplots(figsize=(9, 6))
ax.grid()
ax.set_xlabel(r'Value of $\varepsilon$')
ax.set_ylabel('Fraction of state 6')
ax.set_xlim(eps_min, eps_max)
ax.plot(epsilons, frac_seq_6)