import numpy as np

np.random.seed(1234)

n_states = 5
n_steps = 50
tolerance = 1e-5

# Random transition matrix and state vector
transitionMatrix = np.array([[0.6, 0.1, 0.3],
                             [0.1, 0.7, 0.2],
                             [0.2, 0.2, 0.6]])

# initial dist = [L, w, W]
intital_dist = np.array([1, 0, 0])

# Normalize rows in transitionMatrix
# transitionMatrix /= transitionMatrix.sum(axis=1)[:, np.newaxis]

# # Normalize p
# p /= p.sum()

# Take steps
for k in range(n_steps):
    intital_dist = transitionMatrix.T.dot(intital_dist)

p_50 = intital_dist
print(f'p {p_50}')

# Compute stationary state
w, v = np.linalg.eig(transitionMatrix.T)

j_stationary = np.argmin(abs(w - 1.0))
p_stationary = v[:, j_stationary].real
p_stationary /= p_stationary.sum()
print(p_stationary)