import numpy as np
import random

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

#PART A
def nextStepEhr(X,M):
    U = random.randint(0,1)
    V = random.uniform(0,1)
    Y=0

    if U == 0:
        Y = X
        return Y
    elif U==1 and V <= (X/M):
        Y = X-1
        return Y
    else:
        Y = X+1
        return Y
    return

#PART C
def matrixEhr(M):
    T = []
    for r in range(0,(M+1)):
        if r==0:
            P = np.zeros((M+1))
            P[r] = 0.5
            P[r+1]= 0.5
            T.append(P)
        elif r==(M):
            P = np.zeros((M+1))
            P[r] = 0.5
            P[r-1] = 0.5
            T.append(P)
        else:
            P = np.zeros((M+1))
            P[r] = 0.5
            P[r-1] = r/(2*(M))
            P[r+1] = ((M)-r)/(2*(M))
            T.append(P)
    return T

#PART B
X=0
states = np.arange(0,1001)
statecount = np.zeros(1001)
for a in range(0,10000):
    X = nextStepEhr(X,1000)


for a in range(0,60000):
    X = nextStepEhr(X,1000)
    statecount[X] += 1

plt.bar(states, statecount,align='center', alpha=0.5)
plt.ylabel('Time in state i')
plt.xlabel('State')
plt.title('Lazy Ehrenfest Coin Chain')


#PART C
mu = np.zeros(21)
mu[0] = 1

v = np.zeros(21)
v[0] = 0.5
v[20] = 0.5
k = np.arange(1,61)

F = np.zeros((60,21))
F_2 = np.zeros((60,21))
A = matrixEhr(20)
for x in k:
    F[x-1,:] = mu@(np.linalg.matrix_power(A, x))
    F_2[x-1,:] = v@(np.linalg.matrix_power(A, x))

x= np.arange(1,22)
X,Y = np.meshgrid(x,k)

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
fig_2 = plt.figure()
ax_2 = fig_2.add_subplot(111, projection='3d')
ax.plot_wireframe(X,Y,F)
ax.view_init(30,100)
ax.set_title('F for μ')
ax.set_xlabel('States')
ax.set_ylabel('Time')
ax.set_zlabel('F')

ax_2.plot_wireframe(X,Y,F_2)
ax_2.view_init(30,100)
ax_2.set_title('F for v')
ax_2.set_xlabel('States')
ax_2.set_ylabel('Time')
ax_2.set_zlabel('F')