#!/usr/bin/env python
# coding: utf-8

# # Foundations of Computational Economics #21
# 
# by Fedor Iskhakov, ANU
# 
# <img src="_static/img/dag3logo.png" style="width:256px;">

# ## Computing a stationary distribution of a Markov chain
# 
# <img src="_static/img/lab.png" style="width:64px;">

# <img src="_static/img/youtube.png" style="width:65px;">
# 
# [https://youtu.be/eEbDbM17soU](https://youtu.be/eEbDbM17soU)
# 
# Description: Successive approximations and direct linear solver.

# Let stochastic matrix $ P $ denote the transition probability matrix of a Markov chain which takes values from $ S=\{0,\dots,n-1\} $
# 
# Assume that $ P $ is aperiodic and irreducible, so there exists a unique stationary distribution $ \psi^\star $ such that
# 
# $$
# \psi^\star = \psi^\star \cdot P
# $$
# 
# Our task is to compute $ \psi^\star $

# ### Method 1: successive approximations
# 
# Due to convergence results, we can use the following algorithm (see previous video on Markov chains)
# 
# 1. Start from arbitrary distribution $ \psi_0 $  
# 1. Compute the updated distribution $ \psi_t = \psi_{t-1} P $ until $ \psi_t $ and $ \psi_{t-1} $ are indistinguishable  

# In[1]:


import numpy as np
P = np.array([[.5,.3,.2],[.5,.4,.1],[.1,.1,.8]])
ψ = np.array([1,0,0])


# In[2]:


def stationary_sa(P,psi0,tol=1e-6,maxiter=100,callback=None):
    '''Computes stationary distribution for the Markov chain given by transition probability
    matrix P, with given maximum number of iterations, and convergence tolerance.
    callback function is called at each iteration if given.
    Method: successive approximations
    '''
    pass

stationary_sa(P,ψ)


# In[3]:


def stationary_sa(P,psi0=[None,],tol=1e-6,maxiter=100,callback=None):
    '''Computes stationary distribution for the Markov chain given by transition probability
    matrix P, with given maximum number of iterations, and convergence tolerance.
    callback function is called at each iteration if given.
    Method: successive approximations
    '''
    if psi0[0]==None:
        # degenrate initial distribution
        psi0 = [0,]*P.shape[0]
        psi0[0]=1.0
    P,psi0 = np.asarray(P),np.asarray(psi0)  # convert to np arrays (in case lists were passed)
    assert np.all(np.abs(P.sum(axis=1)-1)<1e-10), 'Passed P is not stochastic matrix'
    assert np.abs(psi0.sum()-1)<1e-10, 'Passed probabilities do not sum up to one'
    for i in range(maxiter):  # main loop
        psi1 = psi0 @ P  # update approximation of psi^star
        err = np.amax(abs(psi0-psi1))  # error is the max absolute deviation element-wise
        if callback != None: callback(err=err,iter=i,psi=psi1)
        if err<tol:
            break  # break out if converged
        psi0 = psi1  # get ready to the next iteration
    else:
        raise RuntimeError('Failed to converge in %d iterations'%maxiter)
    return psi1

def printme(**kwargs):
    print('iter %d, psi = %r'%(kwargs['iter'],kwargs['psi']))
stationary_sa(P,callback=printme)


# ### Method 2: direct solution with linear solver
# 
# $ \psi^\star \in \mathbb{R}^n $ is a row vector of probabilities, $ \sum \psi^\star_i = 1 $
# 
# $$
# \psi^\star = \psi^\star \cdot P
# $$
# 
# Write a linear system of equations:
# 
# $$
# \psi^\star (I - P) = 0, \quad \psi^\star
# $$
# 
# $$
# (I - P') \tilde{\psi}^\star = 0, \quad \tilde{\psi}^\star \text{ is a column vector!}
# $$
# 
# We can use the second form with the standard `np.linalg.solver`

# In[4]:


# have to be careful with the zero solution!
m = P.shape[0]
A = np.eye(m) - P.T
print(np.linalg.solve(A,np.zeros(m)))
print(np.linalg.matrix_rank(A))


# ### Include the linear constraint into the system itself
# 
# Let $ e $ be the column vector of appropriate length with all elements equal 1
# 
# The system
# 
# $$
# \left(
# \begin{array}{cc}
# I - P' & e \\
# e' & 1
# \end{array}
# \right) \cdot
# \left(
# \begin{array}{c}
# \psi \\
# \epsilon
# \end{array}
# \right)
# =
# \left(
# \begin{array}{c}
# e \\
# 2
# \end{array}
# \right)
# $$
# 
# has full rank and therefore can be solved directly

# In[5]:


def stationary_linalg(P,psi0=None):
    '''Computes stationary distribution for the Markov chain given by transition probability
    matrix P. Method: linear solver.
    '''
    pass


# In[6]:


def stationary_linalg(P,psi0=[None,]):
    '''Computes stationary distribution for the Markov chain given by transition probability
    matrix P. Method: linear solver.
    '''
    if psi0[0]==None:
        # degenrate initial distribution
        psi0 = [0,]*P.shape[0]
        psi0[0]=1.0
    P,psi0 = np.asarray(P),np.asarray(psi0)  # convert to np arrays (in case lists were passed)
    assert np.all(np.abs(P.sum(axis=1)-1)<1e-10), 'Passed P is not stochastic matrix'
    assert np.abs(psi0.sum()-1)<1e-10, 'Passed probabilities do not sum up to one'
    m = P.shape[0]  # dimension of the problem
    A = np.ones((m+1,m+1))  # square matrix of ones
    A[:-1,:-1] = np.eye(m)-P.T
    b = np.ones(m+1)
    b[-1]=2
    res = np.linalg.solve(A,b)
    return res[:-1]

print(stationary_linalg(P))


# #### Further learning resources
# 
# - Post on Brilliant.org
#   [https://brilliant.org/wiki/stationary-distributions/](https://brilliant.org/wiki/stationary-distributions/)