# -*- coding: utf-8 -*-
'''
Shared function for HANK
'''
from __future__ import print_function

import numpy as np
import scipy as sc
from scipy.stats import norm
from scipy.interpolate import interp1d, interp2d
from scipy import sparse as sp


def Transition(N,rho,sigma_e,bounds):
    '''
    Calculate transition probability matrix for a given grid for a Markov chain
    with long-run variance equal to 1 and mean 0
    
     Parameters
    ----------
    N : float
        number of states
    rho : float
    sigma_e : float
    bounds : np.array (1,N+1)

    Returns
    ----------
    P : np.array    
        transition matrix
    '''
    pijfunc = lambda x, bound1, bound2 : norm.pdf(x)*(norm.cdf((bound2-rho*x)/sigma_e)-norm.cdf((bound1-rho*x)/sigma_e))
    
    
    P=np.zeros((N,N))
    for i in range(0, int(np.floor((N-1)/2)+1)):
        for j in range(0, N):
        
            pijvalue, err = sc.integrate.quad(pijfunc, bounds[i], bounds[i+1], args=(bounds[j], bounds[j+1]))
            P[i,j]=pijvalue/(norm.cdf(bounds[i+1])-norm.cdf(bounds[i]))
            
             
    P[int(np.floor((N-1)/2)+1):N,:] = P[int(np.ceil((N-1)/2)-1)::-1, ::-1]
    ps=np.sum(P, axis=1)
    
    P=P.copy()/np.transpose(np.tile(ps,(N,1)))
    
    return P


def ExTransitions(S, grid, mpar, par):
    '''
    Generate transition probabilities and grid
    
    Parameters
    ----------
    S : float
        Aggregate exogenous state
    grid : dict
        grid['m']=grid.m : np.array
        grid['h']=grid.h : np.array
        grid['boundsH']=grid.boundsH : np.array (1,mpar['nh'])
    par : dict
        par['xi]=par.xi : float
        par['rhoS']=par.rhoS : float
        par['rhoH']=par.rhoH : float
    mpar : dict
        mpar['nm']=mpar.nm : int
        mpar['nh']=mpar.nh : int   
        mpar['in']=mpar.in : float
        mpar['out']=mpar.out : float   
    
     Returns
     -------
     P_H : np.array
         Transition probabilities
     grid : dict
         Grid
     par : dict
         Parameters
    '''
    
    aux = np.sqrt(S) * np.sqrt(1-par['rhoH']**2)
    
    P = Transition(mpar['nh']-1, par['rhoH'], aux, grid['boundsH'].copy())
    
    P_H = np.concatenate((P, np.tile(mpar['in'],(int(mpar['nh']-1),1))), axis=1)
    lastrow = np.concatenate((np.zeros((1,mpar['nh']-1)), [[1-mpar['out']]]), axis=1)
    lastrow[0,int(np.ceil(mpar['nh']/2))-1] = mpar['out']
    P_H = np.concatenate((P_H.copy(),lastrow.copy()),axis=0)
    P_H = P_H.copy()/np.transpose(np.tile(np.sum(P_H, axis=1),(mpar['nh'],1)))
    
    return {'P_H': P_H, 'grid': grid, 'par': par}


def GenWeight(x,xgrid):
    '''
    Generate weights and indexes used for linear interpolation
    (no extrapolation allowed)
    
    Parameters
    ----------
    x: np.array
        Points at which function is to be interpolated 
    xgrid: np.array
        grid points at which function is measured
        
    Returns
    -------
    weight : np.array
        weight for each index
    index : np.array
        index for integration    
    '''
    
    index = np.digitize(x, xgrid)-1
    index[x <= xgrid[0]] = 0
    index[x >= xgrid[-1]] = len(xgrid)-2
    
    weight = (x-xgrid[index])/(xgrid[index+1]-xgrid[index]) # weight xm of higher gridpoint
    weight[weight.copy()<=0] = 10**(-16) # no extrapolation
    weight[weight.copy()>=1] = 1-10**(-16)
    
    return {'weight': weight, 'index': index}

def MakeGrid(mpar, grid):
    '''
    Make a quadruble log grid
    
    Parameters
    ----------
    mpar : dict
        mpar['nm']=mpar.nm : int
    grid : dict
        grid['m']=grid.m : np.array
        
    Returns
    -------
    grid : np.array
        new grid
    '''
    m_min = -10.87 # natural borrowing limit
    m_max = 150
    grid['m'] = np.exp(np.exp(np.linspace(0., np.log(np.log(m_max - m_min +1)+1), mpar['nm']))-1)-1+m_min
    
    grid['m'][np.abs(grid['m'])==np.min(grid['m'])]=0.
    
    return grid

def Tauchen(rho, N, sigma, mue, types):
    '''
    Generates a discrete approximation to an AR 1 process following Tauchen(1987)
    
    Parameters
    ----------
    rho : float
        coefficient for AR1
    N : int
        number of gridpoints
    sigma : float
        long-run variance
    mue :  float   
        mean of AR1 process
    types : string
        grid transition generation alogrithm
        'importance' : importance sampling (Each bin has probability 1/N to realize)
        'equi' : bin-centers are equi-spaced between +-3 std
        'simple' : like equi + Transition Probabilities are calculated without using integrals
        'simple importance' : like simple but with grid from importance
        
    return
    -----------
    grid : np.array
        grid 
    P : np.array
        Markov probability
    bounds : np.array
        bounds
    
    '''
    pijfunc = lambda x, bound1, bound2 : norm.pdf(x)*(norm.cdf((bound2-rho*x)/sigma_e)-norm.cdf((bound1-rho*x)/sigma_e))
    
    if types in {'importance','equi','simple','simple importance'}:
       types = types
    else:
       types = 'importance'
       print('Warning: TAUCHEN:NoOpt','No valid type set. Importance sampling used instead')
       
    if types == 'importance': # Importance sampling
           
       grid_probs = np.linspace(0,1,N+1) 
       bounds = norm.ppf(grid_probs)
       
       # replace (-)Inf bounds by finite numbers
       bounds[0] = bounds[1].copy()-99
       bounds[-1] = bounds[-2].copy()+99
        
       # Calculate grid - centers
       grid = 1*N*( norm.pdf(bounds[:-1]) - norm.pdf(bounds[1:]))
      
       sigma_e = np.sqrt(1-rho**2) # Calculate short run variance
       P=np.zeros((N,N))
    
       for i in range( int(np.floor((N-1)/2+1)) ): # Exploit symmetrie
          for j in range(N):
              pijvalue, err = sc.integrate.quad(pijfunc, bounds[i], bounds[i+1], args=(bounds[j], bounds[j+1]),epsabs=10**(-6))
              P[i,j] = N*pijvalue
              
       #P=np.array([[0.9106870252,0.0893094991,0.0000037601],[0.0893075628,0.8213812539,0.0893075628],[0.0000037601,0.0893094991,0.9106899258]])
       #print bounds 
       #print P
       P[int(np.floor((N-1)/2)+1):N,:] = P[int(np.ceil((N-1)/2))-1::-1,::-1].copy()
       
    elif types == 'equi': # use +-3 std equi-spaced grid
        # Equi-spaced
        step = 6/(N-1)
        grid = np.range(-3.,3+step,step)
        
        bounds = np.concatenate(([-99],grid[:-1].copy()+step/2,[99]),axis=1)
        sigma_e = np.sqrt(1-rho**2) # calculate short run variance
        P=np.zeros((N,N))
        
        pijfunc = lambda x, bound1, bound2 : norm.pdf(x)*(norm.cdf((bound2-rho*x)/sigma_e)-norm.cdf((bound1-rho*x)/sigma_e))
        
        for i in range( int(np.floor((N-1)/2+1)) ): # Exploit symmetrie
          for j in range(N):
              pijvalue, err = sc.integrate.quad(pijfunc, bounds[i], bounds[i+1], args=(bounds[j], bounds[j+1]))
              P[i,j] = pijvalue/(norm.cdf(bounds[i])-norm.cdf(bounds[i-1]))
       
        P[int(np.floor((N-1)/2)+1):N,:] = P[int(np.ceil((N-1)/2))-1::-1,::-1].copy()
        
    elif types == 'simple': # use simple transition probabilities
        
        step = 12/(N-1)
        grid = np.range(-6.,6+step, step)
        bounds=[]
        sigma_e = np.sqrt(1-rho**2)
        P=np.zeros((N,N))
        
        pijfunc = lambda x, bound1, bound2 : norm.pdf(x)*(norm.cdf((bound2-rho*x)/sigma_e)-norm.cdf((bound1-rho*x)/sigma_e))
        
        for i in range(N):
            P[i,0] = norm.cdf((grid[0]+step/2-rho*grid[i])/sigma_e)
            P[i,-1] = 1- norm.cdf((grid[-1]+step/2-rho*grid[i])/sigma_e)
            for j in range(1,N-1):
                P[i,j] = norm.cdf((grid[j]+step/2-rho*grid[i])/sigma_e) - norm.cdf((grid[j]-step/2-rho*grid[i])/sigma_e)
                
    elif types == 'simple importance': # use simple transition probabilities
        
        grid_probs = np.linspace(0.,1.,N+1)            
        bounds = norm.ppd(grid_probs.copy())
        
        # calculate grid - centers
        grid = N*(norm.pdf(bounds[:-1])-norm.pdf(bounds[1:]))
        
        #replace -Inf bounds by finite numbers
        bounds[0] = bounds[1] - 99
        bounds[-1] = bounds[-2] + 99
        
        sigma_e = np.sqrt(1-rho**2)
        P=np.zeros((N,N))
        pijfunc = lambda x, bound1, bound2 : norm.pdf(x)*(norm.cdf((bound2-rho*x)/sigma_e)-norm.cdf((bound1-rho*x)/sigma_e))
        
        for i in range(int(np.floor((N-1)/2))+1):
            P[i,0] = norm.cdf((bounds[1]-rho*grid[i])/sigma_e)
            P[i,-1] = 1- norm.cdf((bounds[-2]-rho*grid[i])/sigma_e)
            for j in range(int(np.floor((N-1)/2))+1):
                P[i,j] = norm.cdf((bounds[j+1]-rho*grid[i])/sigma_e) -norm.cdf((bounds[j]-rho*grid[i])/sigma_e)
            
        P[int(np.floor((N-1)/2))+1:,:] = P[int(np.ceil((N-1)/2))-1::-1,::-1].copy()
    
    
    ps = np.sum(P,axis=1)
    P=P.copy()/np.transpose(np.tile(ps.copy(),(N,1)))
    
    grid = grid.copy()*np.sqrt(sigma) + mue
        
    return {'grid': grid, 'P':P, 'bounds':bounds}

