from model_blocks import *
import numpy as np
from sequence_jacobian import create_model


def solve_hank_ss(calibration, print_args = 0, tax_distortion=1):
    """
    Given calibration dict, solves for (Y,nfa,Q) as a function of (r,B,G,Cstar,rstar) which are already updated in calibration.
    Then, C comes out of HH problem
    """
    # These blocks combine Step 1 and Step 2
    if tax_distortion == 1:
        blocks_ss = [hh_ha, foreign_c, domestic_prices, foreign_prices, labor, wages, unions_ss, fiscal_tax, rsimple, CA, mkt_clearing, prices_ss]
    else:
        blocks_ss = [hh_ha, foreign_c, domestic_prices, foreign_prices, labor, wages, unions_wage_ss, fiscal_tax, rsimple, CA, mkt_clearing, prices_ss]
    hank_ss   = create_model(blocks_ss, name = "Baseline SS")

    if print_args == 1:
        print(f"Inputs: {hank_ss.inputs} \n")
        print(f"Outputs: {hank_ss.outputs}")

    # C comes out of the HH problem
    unknowns_ss = {'Y': calibration['Y'], 'nfa': calibration['nfa'], 'Q': calibration['Q'] }
    targets_ss = [ 'goods_mkt', 'piwres_ss', 'asset_mkt' ]
    ss = hank_ss.solve_steady_state(calibration, unknowns_ss, targets_ss, solver="broyden_custom" )
    calibration.update(ss)

    # auxiliary parameter xi for Prop 1
    t1 = (calibration['eis']**(-1) - 1)/(1 + calibration['frisch']**(-1))
    t2 = (1 - calibration['G']/calibration['Y']) * (calibration['chi'] - 1 + calibration['alpha'])/(calibration['chi'] - 1)
    xi = (t1+t2)**(-1)
    ss.update({'xi': xi })

    return ss



def aggregate_split(D, a, pctl):
    """Given a distribution 'D' and the array of assets 'a' held at each point in that
    distribution, return the amount of total assets by people below the 'pctl' percentile of
    assets, and by people above that percentile.

    Note that 'D' and 'a' must have the same dimensions."""
    # make D and a one-dimensional vectors sorted in ascending order of a
    D = D.ravel()
    a = a.ravel()
    ii = a.argsort()
    a = a[ii]
    D = D[ii]

    # determine the index i that takes us above pctl percentile
    Dcum = D.cumsum()
    i = np.argmax(Dcum > pctl)

    # what percent of people at gridpoint i are "under" pctl percentle?
    pi = (pctl - Dcum[i-1])/D[i]

    # split total assets into below and above pctl
    a_below = np.vdot(a[:i], D[:i]) + pi*a[i]*D[i]
    a_above = np.vdot(a[i+1:], D[i+1:]) + (1-pi)*a[i]*D[i]

    return a_below, a_above


def compute_irfs(irf_r, ss0, ss_new, scale, pctl, which_shocks, irf_vars):

    # responses
    for shock_var in which_shocks: # for all shocks
        irf_r[shock_var] = irf_r[shock_var] + ss_new[shock_var] - ss0[shock_var]

    for var in irf_vars:
        if var not in irf_r:  # assign a zero time series
            irf_r[var] = irf_r[which_shocks[0]] - irf_r[which_shocks[0]]
        elif var not in which_shocks:  # variables other than the shocks
            irf_r[var] = irf_r[var] + ss_new[var] - ss0[var]

    # Compute IRFs for Tax Revenue = T*PH/P
    irf_r_t = (irf_r['T'] + ss0['T']) * (irf_r['php'] + ss0['php']) - ss0['T']*ss0['php']
    irf_r['T'] = irf_r_t

    irf_r['CA'] = -1*irf_r['CAD']
    irf_r['FD'] = np.append([irf_r['B'][0]], np.diff(irf_r['B']))

    # Rescale
    irf_r = scale * irf_r

    # compute asset distribution
    a_below0, a_above0 = aggregate_split(ss0.internals['household_ha']['D'], ss0.internals['household_ha']['a'], pctl)
    a_below_new, a_above_new  = aggregate_split(ss_new.internals['household_ha']['D'], ss_new.internals['household_ha']['a'], pctl)
    scale_omega = ss0['omega']/ss0['omega_US']
    a_below0, a_above0 = a_below0*scale_omega, a_above0*scale_omega
    a_below_new, a_above_new = a_below_new*scale_omega, a_above_new*scale_omega

    T = len(irf_r[which_shocks[0]])
    a_below = np.zeros(T)
    a_above = np.zeros(T)
    for t in range(T):

        D_temp = irf_r.internals['household_ha']['D'][t] / scale
        a_temp = irf_r.internals['household_ha']['a'][t] / scale

        D_temp = D_temp + ss_new.internals['household_ha']['D']
        a_temp = a_temp + ss_new.internals['household_ha']['a']
        a_below[t], a_above[t] = aggregate_split(D_temp, a_temp, pctl)

        a_below[t] = a_below[t] * (irf_r['pmod'][t]/scale  + ss0['pmod']) * scale_omega - a_below0 * ss0['pmod']
        a_above[t] = a_above[t] * (irf_r['pmod'][t]/scale  + ss0['pmod']) * scale_omega - a_above0 * ss0['pmod']

    irf_r['Atop20'] = a_above * scale
    irf_r['Abottom80'] = a_below * scale

    # Get the iMPC matrix from the initial SS
    M = compute_iMPCs(ss0)
    M_new = compute_iMPCs(ss_new)

    # Compute Behavorial MPC
    MPC_beh = compute_iMPCs_beh(M)
    MPC_beh_new = compute_iMPCs_beh(M_new)


    return irf_r, a_below0, a_above0, a_below_new, a_above_new, M, M_new, MPC_beh, MPC_beh_new
