"""
Discretization of continuous processes as markov chain
"""


import scipy as sp
import scipy.stats
import numpy as np
import math

# univariate tauchen code is taken from dpsolve ( https://bitbucket.org/stevejb/dpsolve )

def tauchen(N, mu, rho, sigma, m=2):
    """
    Approximate an AR1 process by a finite markov chain using Tauchen's method.

    :param N: scalar, number of nodes for Z
    :param mu: scalar, unconditional mean of process
    :param rho: scalar
    :param sigma: scalar, std. dev. of epsilons
    :param m: max +- std. devs.
    :returns: Z, N*1 vector, nodes for Z. Zprob, N*N matrix, transition probabilities

    SJB: This is a port of Martin Floden's 1996 Matlab code to implement Tauchen 1986 Economic Letters method The following comments are Floden's. Finds a Markov chain whose sample paths approximate those of the AR(1) process z(t+1) = (1-rho)*mu + rho * z(t) + eps(t+1) where eps are normal with stddev sigma.
    """
    Z     = np.zeros((N,1))
    Zprob = np.zeros((N,N))
    a     = (1-rho)*mu

    Z[-1]  = m * math.sqrt(sigma**2 / (1 - (rho**2)))
    Z[0]  = -1 * Z[-1]
    zstep = (Z[-1] - Z[0]) / (N - 1)

    for i in range(1,N):
        Z[i] = Z[0] + zstep * (i)

    Z = Z + a / (1-rho)

    for j in range(0,N):
        for k in range(0,N):
            if k == 0:
                Zprob[j,k] = sp.stats.norm.cdf((Z[0] - a - rho * Z[j] + zstep / 2) / sigma)
            elif k == (N-1):
                Zprob[j,k] = 1 - sp.stats.norm.cdf((Z[-1] - a - rho * Z[j] - zstep / 2) / sigma)
            else:
                up = sp.stats.norm.cdf((Z[k] - a - rho * Z[j] + zstep / 2) / sigma)
                down = sp.stats.norm.cdf( (Z[k] - a - rho * Z[j] - zstep / 2) / sigma)
                Zprob[j,k] = up - down

    return( (Z, Zprob) )


def rouwenhorst(rho, sigma, N):
    """
    Approximate an AR1 process by a finite markov chain using Rouwenhorst's method.

    :param rho: autocorrelation of the AR1 process
    :param sigma: conditional standard deviation of the AR1 process
    :param N: number of states
    :return [nodes, P]: equally spaced nodes and transition matrix
    """

    from numpy import sqrt, linspace, array,zeros

    sigma = float(sigma)

    if N == 1:
      nodes = array([0.0])
      transitions = array([[1.0]])
      return [nodes, transitions]

    p = (rho+1)/2
    q = p
    nu = sqrt( (N-1)/(1-rho**2) )*sigma

    nodes = linspace( -nu, nu, N)
    sig_a = sigma
    n = 1
    #    mat0 = array( [[1]] )
    mat0 = array([[p,1-p],[1-q,q]])
    if N == 2:
        return [nodes,mat0]
    for n in range(3,N+1):
        mat = zeros( (n,n) )
        mat_A = mat.copy()
        mat_B = mat.copy()
        mat_C = mat.copy()
        mat_D = mat.copy()
        mat_A[:-1,:-1] = mat0
        mat_B[:-1,1:] = mat0
        mat_C[1:,:-1] = mat0
        mat_D[1:,1:] = mat0

        mat0 = p*mat_A + (1-p)*mat_B + (1-q)*mat_C + q*mat_D
        mat0[1:-1,:] = mat0[1:-1,:]/2
    P = mat0
    return [nodes, P]



def multidimensional_discretization(rho, sigma, N=3, method='rouwenhorst', m=2):
    """
    Discretize an VAR(1) into a markov chain. The autoregression matrix is supposed to be a scalar.

    :param rho:
    :param sigma:
    :param N:
    :param method:
    :param m:
    :return:
    """

    # rho is assumed to be a scalar
    # sigma is a positive symmetric matrix
    # N number of points in each non-degenerate dimension
    # m : standard deviations to approximate

    import scipy.linalg
    from itertools import product

    d = sigma.shape[1]

    sigma = sigma.copy()

    zero_columns = np.where(sigma.sum(axis=0)==0)[0]
    for i in zero_columns:
        sigma[i,i] = 1

    L = scipy.linalg.cholesky(sigma)

    N = int(N)

    if method=='tauchen':
        [nodes_1d, probas_1d] = tauchen(N, 0, rho, 1, m=m)
    elif method=='rouwenhorst':
        [nodes_1d, probas_1d] = rouwenhorst(rho, 1, N)

    markov_nodes = np.array( list( product( *([nodes_1d]*d))) ).T
    markov_indices = np.array( list( product( *([range(N)]*d)  )  ) ).T

    markov_nodes = np.dot(L, markov_nodes)


    transition_matrix = 1
    for i in range(d):
        transition_matrix = np.kron(transition_matrix, probas_1d)

    markov_nodes = np.ascontiguousarray(markov_nodes.T)

    for i in zero_columns:
        markov_nodes[:,i] = 0

    return [markov_nodes, transition_matrix]

def tensor_markov( *args ):
    """Computes the product of two independent markov chains.
    
    :param m1: a tuple containing the nodes and the transition matrix of the first chain
    :param m2: a tuple containing the nodes and the transition matrix of the second chain
    :return: a tuple containing the nodes and the transition matrix of the product chain
    """
    
    if len(args) > 2:

        m1 = args[0]
        m2 = args[1]
        tail = args[2:]
                
        prod = tensor_markov(m1,m2)
        return tensor_markov( prod , *tail )

    elif len(args) == 2:

        m1,m2 = args
        
        n1, t1 = m1
        n2, t2 = m2
        
        n1 = np.array(n1, dtype=float)
        n2 = np.array(n2, dtype=float)
        t1 = np.array(t1, dtype=float)
        t2 = np.array(t2, dtype=float)

        assert(n1.shape[0] == t1.shape[0] == t1.shape[1])
        assert(n2.shape[0] == t2.shape[0] == t2.shape[1])

        t = np.kron(t1, t2)

        p = t1.shape[0]
        q = t2.shape[0]

        np.tile( n2, (1,p))
        # n = np.row_stack([
        #     np.repeat(n1, q, axis=1),
        #     np.tile( n2, (1,p))
        # ])
        n = np.column_stack([
            np.repeat(n1, q, axis=0),
            np.tile( n2, (p,1))
        ])
        return n,t
    
    else:
        raise Exception("Incorrect number of arguments. Expected at least 2. Found {}.".format(len(args)))

#
# quantization_data = '/home/pablo/quantization_grids/'
#
#
# def quantization_nodes(N,sigma):
#     import numpy
#     import numpy.linalg
#     assert( len(sigma.shape) == 2 )
#     var = numpy.diag(sigma)
#     d = sigma.shape[0]
#     [w, x] = standard_quantization_weights(N,d )
#     A = numpy.linalg.cholesky(sigma)
#     x = numpy.dot(A, x)
#     return [x,w]
#
#
# def quantization_weights(N,sigma):
#     [x,w] = quantization_nodes(N,sigma)
#     return [w,x]
#
# def standard_quantization_weights(N,d):
#     filename = quantization_data + '{0}_{1}_nopti'.format(N,d)
#     import numpy
#
#     try:
#         G = numpy.loadtxt(filename)
#     except Exception as e:
#         raise e
#
#     w = G[:N,0]
#     x = G[:N,1:d+1]
#
#     s = numpy.dot(w, x)
#     x = x - numpy.outer(w,s)
#
#     return [w,x.T]

if __name__ == '__main__':

    [Z,Zprob] = tauchen( 5,0,0.8,0.1,1.5 )

    import numpy
    sigma = numpy.diag([0.1, 0.1])**2
    sigma[0,1]  = numpy.sqrt( 0.5*sigma[0,0]*sigma[1,1] )
    sigma[1,0]  = numpy.sqrt( 0.5*sigma[0,0]*sigma[1,1] )
    rho = 0.9

    [nodes, transition] = multidimensional_discretization(rho, sigma, 2, method='rouwenhorst')

    transition0 = transition.copy()*0
    transition0[1,1] = 1.0

    print(nodes.shape)

    print(transition.shape)

    [nodes, transition] = tensor_markov( (nodes, transition0), (nodes, transition) )

    print(nodes.shape)
    print(transition.shape)
