"""Defines all the model blocks"""
import numpy as np

# SSJ imports
from sequence_jacobian.blocks.het_block import het
from sequence_jacobian import grids
from sequence_jacobian import simple, create_model, solved, combine
from sequence_jacobian import utilities as utils


"""
Main Model
"""

# This is just for initialization
def household_init(e_grid, a_grid, Z, Transfer, rpost, eis):
    coh = (1 + rpost) * a_grid + Z * e_grid[:, np.newaxis] + Transfer
    Va = (1 + rpost) * (0.1 * coh) ** (-1 / eis)
    return Va

@het(exogenous='Pi', policy='a', backward='Va', backward_init=household_init)
def household_ha(Va_p, Pi, a_grid, e_grid, coh, rpost, beta, beta_int, eis, betashock):

    """
    Single backward iteration step using endogenous gridpoint method for households with CRRA utility.
    """

    # Solve HH problem
    uc_nextgrid = ( betashock * (beta+beta_int[:, np.newaxis]) * Pi) @ Va_p
    c_nextgrid = uc_nextgrid ** (-eis)
    a = utils.interpolate.interpolate_y(c_nextgrid + a_grid, coh, a_grid)
    iconst = np.nonzero(a < a_grid[0])
    a[iconst] = a_grid[0]
    c = coh - a
    Va = (1 + rpost) * c ** (-1 / eis)

    # Compute cstar
    uce = c ** (-1 / eis) * e_grid[:,np.newaxis]

    return Va, a, c, uce

def make_grid(rho_e, sd_e, nE, amin, amax, nA, beta_spread, n_beta, prop ):
    e_grid, pi_e, Pi = grids.markov_rouwenhorst(rho=rho_e, sigma=sd_e, N=nE)
    a_grid = grids.asset_grid(amin=amin, amax=amax, n=nA)

    # Now make everything bigger (prop is proportion of agents with beta)
    beta_int = np.linspace(-beta_spread, 0, n_beta)
    e_grid = np.kron(e_grid, np.ones(n_beta))

    # Create vector for share of beta hetegoreneity
    vec = np.zeros(n_beta)
    vec[-1] = prop                      # Ones with beta
    vec[0:-1] = (1 -  prop)/(n_beta-1)      # Others have equal mass
    vec = vec * n_beta
    pi_e = np.kron(pi_e, vec)
    Pi = np.kron(Pi, np.eye(n_beta))
    beta_int = np.kron(np.ones(nE), beta_int)

    return e_grid, pi_e, Pi, a_grid, beta_int


# This changes whenever Z or Transfer changes
def hetinput_ge(e_grid, a_grid, Z, Transfer, rpost):
    """ Define inputs that go into household's problem """
    coh = (1 + rpost) * a_grid + Z * e_grid[:, np.newaxis] + Transfer
    return coh

hh_ha0 = household_ha.add_hetinputs([make_grid, hetinput_ge])
hh_ha0 = hh_ha0.remap({'A':'ASCALED', 'C':'CSCALED', 'Z':'ZSCALED'})

@simple
def scale_Z(Z, omega, omega_US):
    ZSCALED = Z * omega_US/omega
    return ZSCALED

@simple
def scale_A_C(ASCALED, CSCALED, omega, omega_US):
    scale = omega_US / omega
    A = ASCALED/scale
    C = CSCALED/scale
    return A, C

hh_ha1 = combine([hh_ha0, scale_Z, scale_A_C])

"""Now the preference shock stuff"""
hh_ha1 = hh_ha1.remap({'Z': 'Zmod', 'rpost': 'rpostmod', 'A': 'AMOD', 'C': 'CMOD'})

def ces_price(alpha, eta, x1, x2):
    if eta != 1:
        return ((1-alpha)*x1**(1-eta) + alpha*x2**(1-eta))**(1/(1-eta))
    else:
        return x1**(1-alpha)*x2**alpha

@simple
def modified_P(alpha, eta, php, pwp, x):          # php = P_H / P; pwp = P_W / P
    pw_ph = pwp / php    # ratio of P_W / P_H
    num = ces_price(alpha, eta, 1/x, pw_ph)
    den = ces_price(alpha, eta, 1, pw_ph)
    pmod = num / den
    return pmod

@simple
def modified_inputs(rpost, Z, pmod):
    rpostmod = (1+rpost)*(pmod(-1)/pmod) - 1
    Zmod = Z / pmod
    return rpostmod, Zmod

@simple
def modified_outputs(AMOD, CMOD, pmod, x, alpha, eta, php, pwp):
    A = AMOD * pmod
    C = CMOD * pmod
    cH = (1-alpha)*(php/pmod/x)**(-eta)*CMOD/x
    cW = alpha*(pwp/pmod)**(-eta)*CMOD
    Ccheck = php*cH + pwp*cW
    return A, C, Ccheck, cH, cW

hh_ha = combine([hh_ha1, modified_P, modified_inputs, modified_outputs])

# MPC
def compute_iMPCs(ss, scaled = 0, T=1000):
    if scaled == 0:
        J_household = hh_ha.jacobian(ss, inputs=["Z"],  T=T)
        M = J_household['C']['Z']
    else:
        J_household = hh_ha.jacobian(ss, inputs=["ZSCALED"],  T=T)
        M = J_household['CSCALED']['ZSCALED']
    return M


def compute_iMPCs_beh(M):
    MPC_beh = M.copy()
    T = len(M)
    for t in range(1,T):
        MPC_beh[t:,t] = MPC_beh[0:-t,0]
        MPC_beh[:t,t] = 0

    return MPC_beh

# foreign demand
@simple
def foreign_c(phw,  omega, gamma, Cstar):
    cHstar = omega*phw**(-gamma) * Cstar     # phw = P_H / P_W = (P_H/P) / Q since P_W = E = Q * P
    return cHstar


# Real exchange rate and domestic price
@simple
def domestic_prices(Q, alpha, eta):     # Q_t -> p_ht, p_wt
    if alpha < 1:
        if eta == 1:
            php = Q**(-alpha/(1-alpha))
        else:
            php = ( (1 - alpha*Q**(1-eta)) / (1-alpha) )**(1/(1-eta))
        pwp = Q
    else:
        php = 1 # undefined
        pwp = 1
    return php, pwp

# foreign prices (with P_t^* = P_{Ft}^* = 1)
@simple
def foreign_prices(php, Q):            # Q_t -> p_hw
    phw = php/Q                        # phw = P_H / P_W = (P_H/P) / Q since P_W = E = Q * P
    return phw

# production
@simple
def labor(Y, Theta):
    N = Y / Theta
    return N

# wages and income
@simple
def wages(php, Theta, Y, T, N):
    w = php*Theta               # real wage
    atw = php*Y
    Z = php * (Y - T)
    return w, atw, Z

# Steady State wage PC with piw=0
@simple
def unions_ss(N, C, Z, frisch, eis, varphi, gamma_w):
    den = gamma_w * Z/N
    piwres_ss = varphi*N**(1/frisch) * C**(1/eis) / den - 1
    return piwres_ss

# Steady State wage PC without distortion with piw=0
@simple
def unions_wage_ss(N, C, Z, frisch, eis, varphi, epsilon_w, w):
    den = epsilon_w/(epsilon_w-1) * w        # real wage instead of (1-lambda)*Z/N
    piwres_ss = varphi*N**(1/frisch) * C**(1/eis) / den - 1
    return piwres_ss

# wage PC
@solved(unknowns={'piw': (-0.2, 0.2)}, targets={'piwres'}, solver="brentq")
def unions(N, C, Z, piw, kappa_w, beta, frisch, eis, varphi, gamma_w):
    den = gamma_w * Z/N
    piwres = piw - beta*piw(1) - kappa_w*(varphi*N**(1/frisch) * C**(1/eis) / den - 1)
    return piwres

# wage PC without distortion
@solved(unknowns={'piw': (-0.2, 0.2)}, targets={'piwres'}, solver="brentq")
def unions_wage(N, C, Z, piw, kappa_w, beta, frisch, eis, varphi, gamma_w, epsilon_w, w):
    den = epsilon_w/(epsilon_w-1) * w        # real wage instead of (1-lambda)*Z/N
    piwres = piw - beta*piw(1) - kappa_w*(varphi*N**(1/frisch) * C**(1/eis) / den - 1)
    return piwres

# producer price inflation
@solved(unknowns={'piH':0,'PH':1}, targets=['piHres','PHres'], solver="broyden_custom" )
def nkpc(piH, PH, piw, Theta):
    PHres = 1 + piH - PH/PH(-1)
    piHres = 1 + piH - (1+piw)*Theta(-1)/Theta
    return piHres, PHres

# CPI
@simple
def cpi(piH,Q,alpha, pi):
    piout = piH + alpha/(1-alpha)*(Q-Q(-1))
    pioutres = piout - pi
    return piout, pioutres

@solved(unknowns={'Q':(0.1,2.0)}, targets=['uip'], solver="brentq") # r_t -> Q_t
def UIP(r, Q, istar):
    uip = 1 + r - (1+istar) * Q(+1)/Q
    return uip

@simple
def UIP_E(PH,phw):
    E = PH/phw
    return E

@simple
def rsimple(r):
    rpost = r(-1)
    return rpost

## Other prices useful for transition
@simple
def prices_ss(r, php, pwp, phw, PH, piw):
    piH = piw
    pi = piH
    E = PH/phw
    P = PH / php
    i = r + pi
    piout = pi
    rstar = r
    istar = r

    return piH, pi, E, P, i, piout, rstar, istar

@simple
def taylor(pi,piH,rstar,phi_pi,phi_piH,realrate): # pi_t -> i_t
    if realrate == 0:                                      # taylor
        i = rstar + phi_pi*pi + phi_piH*piH
    elif realrate == 1:                                    # constant-r
        i = rstar + pi(+1)
    elif realrate == 2:                                    # ZLB
        i = rstar + pi - pi
    return i

# other monetary policies

# natural rate piw = 0
@simple
def piw_natural(piw,piw_target):
    piwres_eq = piw - piw_target
    return piwres_eq

@simple
def eq_quant_ZLB(r, rstar):
    pi = rstar - r(-1)    # i = rstar; from fisher, i = r + pi(+1)
    return pi

@simple
def taylor_world(pi, rstar, phi_Cstar, Cstar, Cstar_ss):     # Cstar -> i_t
    i = pi(+1) + rstar + phi_Cstar*(Cstar/Cstar_ss - 1)
    return i

@simple
def fiscal_tax(r, G, B, php, Transfer):
    T = ( (1+r(-1))*B(-1) - B + Transfer ) / php + G
    FD = B - B(-1)
    return T, FD


@simple
def fiscal_transfer(r, G, B, php, T):
    Transfer = B - (1+r(-1))*B(-1) - php*(G - T)
    FD = B - B(-1)
    return FD, Transfer

@simple
def eq_quant(pi, r, i):
    fisher = i - r - pi(+1)
    return fisher
#
# @simple
# def eq_quant_world(pi, i):
#     r = i - pi(+1)
#     return r

@simple
def CA(nfa, php, pwp, Y, G, C, cW, cH, r, rpost, A, Transfer):
    nfares = r(-1)*nfa(-1) + php*(Y-G) - C - Transfer + (rpost - r(-1))*A(-1) + nfa(-1) - nfa
    netexports = php*(Y-G) - C
    imports = pwp*cW
    exports = netexports + imports
    CAD = nfa(-1) - nfa
    return nfares, netexports, imports, exports, CAD

@simple
def asset_clearing(A, B, nfa):
    asset_mkt = A - nfa - B
    return asset_mkt

@simple
def mkt_clearing(A, B, nfa, cH, cHstar, G, Y):
    asset_mkt = A - nfa - B
    goods_mkt = cH + cHstar + G - Y
    return asset_mkt, goods_mkt

# Other aggregate equations for SS
@simple
def to_solve(C, Z, r, A):
    C_tosolve = C - r*A - Z
    return C_tosolve


"""
Blocks for World Economy
"""

# wage PC
@solved(unknowns={'pi': (-1.0, 1.0)}, targets={'pires'}, solver="brentq")
def unions_world(N, C, Z, pi, kappa_w, beta, frisch, eis, varphi, gamma_w):
    den = gamma_w * Z/N
    pires = pi - beta*pi(1) - kappa_w*(varphi*N**(1/frisch) * C**(1/eis) / den - 1)
    return pires

# wage PC SS where pi = constant
@simple
def unions_world_ss(N, C, Z, kappa_w, beta, frisch, eis, varphi, gamma_w, pi):
    den = gamma_w * Z/N
    piout = 1/(1-beta) * kappa_w*(varphi*N**(1/frisch) * C**(1/eis) / den - 1)
    pires_ss = piout - pi
    return pires_ss

@simple
def natural_allocation_Y_world(Y, Y0):
    Yres = Y - Y0
    return Yres

@simple
def mkt_clearing_world(A, B, C, G, Y):
    asset_mkt = A -  B
    goods_mkt = C + G - Y
    return asset_mkt, goods_mkt

# domestic demand for home and foreign goods for Blanchard Model
@simple
def hh_outputs_blanchard(C, php, pwp, eta, alpha):
    cH = (1-alpha) * php**(-eta) * C
    cW = alpha * pwp**(-eta) * C
    return cH, cW
