import copy
import numpy as np
import sequence_jacobian as sj
from model_solution_funcs import solve_hank_ss

# Scale the debt change (model_input contains path of debt change)
def get_dB(model_input, which_countries, scale):

    dB_c = {}    # path of dB_c
    dB_end = {}  # final value

    for c, cou in enumerate(which_countries):

        temp2 = model_input.loc[model_input['Code'] == cou].copy()
        temp2['Date'] = temp2['Date'].dt.strftime('%Y-%m')
        temp2 = temp2.set_index('Date')

        dB_c[cou]  = (temp2[['value']]/100*4) / scale # 'value' is annual percent increase in B relative to initial Y
        dB_end[cou]  = dB_c[cou].iloc[-1].iloc[0]

    return dB_c, dB_end


## Computes Jacobians w.r.t inputs; ss_c is either the dict of country-level SS
def compute_jacobians(hank, ss_c, unknowns, targets, inputs, outputs=['Y','C', 'A','Q'], T=500):

    print("Computing Jacobian of", outputs, "w.r.t.", inputs, "...")
    G_c = {}
    for c in ss_c:
        if outputs==[]:
            G_c[c] = hank.solve_jacobian(ss_c[c], inputs=inputs, unknowns = unknowns, targets=targets, T=T)
        else:
            G_c[c] = hank.solve_jacobian(ss_c[c], inputs=inputs, outputs=outputs, unknowns = unknowns, targets=targets, T=T)
    print("Done!")

    return G_c


## World SS function
# New World SS - takes in dB for each country and calibration that includes initial B
# Here dB is percent of SS output
def compute_Cstar_and_r_world(x, dB, country_calibrations, opt=1, tax_distortion = 1):

    Cstar, r = x
    country_ss_new = {}
    which_countries = list(country_calibrations.keys())

    eq1 = 0 # World Cstar
    eq2 = 0 # World P
    agg_B = 0
    agg_A = 0
    agg_Y = 0
    agg_G = 0
    agg_C = 0
    # compute steady state for export market clearing
    for c, cou in enumerate(which_countries):

        calibration_c = copy.deepcopy(country_calibrations[cou])
        calibration_c.update({'r' : r, 'Cstar' : Cstar, 'B': calibration_c['B'] + dB[cou]*calibration_c['Y'],
                              'istar': r, 'rstar':r  }) # istar is world rate, rstar is Taylor rule intercept
        ss_c = solve_hank_ss(calibration_c, tax_distortion = tax_distortion)     # Solve for (Y,nfa,Q,C) as a fn of (r,G,B,C^*,r^*)

        # export market and price consistency condition
        eq1 = eq1 + ss_c['alpha']*ss_c['Q']**(-ss_c['eta'])*ss_c['C']

        ss_c['PHF'] = ss_c['PH']/ss_c['E']

        if ss_c['gamma']==1:
            eq2 = eq2 + ss_c['omega']*np.log(ss_c['PHF'])
        else:
            eq2 = eq2 + ss_c['omega']*ss_c['PHF']**(1-ss_c['gamma'])

        # other aggregates
        agg_B = agg_B + ss_c['B']/ss_c['Q']
        agg_A = agg_A + ss_c['A']/ss_c['Q']
        agg_Y = agg_Y + ss_c['Y']*ss_c['PHF']
        agg_G = agg_G + ss_c['G']*ss_c['PHF']
        agg_C = agg_C + ss_c['C']/ss_c['Q']

        country_ss_new[cou] = ss_c

    if ss_c['gamma']==1:
        eq2 = np.exp(eq2)

    if opt == 1:
        return [agg_A - agg_B, eq2-1] # Looks like at r=0, not equivalent to clear goods and asset market!!
    else:
        return country_ss_new, agg_B, agg_A, agg_Y, agg_G, agg_C, eq1, eq2

## Jacobian approach to compute (Cstar,istar) path
def solve_Cstar_istar(hank, ss_c, G_c, T, B_path_c, rstar_path, betashock_path_c, x_path_c):

    # initialize everything
    imp = {'A': np.zeros(T), 'Q': np.zeros(T)}
    J = {'A': {'Cstar': np.zeros((T,T)), 'istar': np.zeros((T,T))},
         'Q': {'Cstar': np.zeros((T,T)), 'istar': np.zeros((T,T))}}

    for c in ss_c:

        ss = ss_c[c].copy()
        ss.update({'realrate': 0, 'phi_pi':0, 'phi_piH': 1.5})

        Jc = G_c[c]

        Q_weight = ss['omega']/(1-ss['alpha'])

        J['A']['Cstar'] += Jc['A', 'Cstar']
        J['A']['istar'] += Jc['A', 'istar']
        J['Q']['Cstar'] += Q_weight*Jc['Q', 'Cstar']
        J['Q']['istar'] += Q_weight*Jc['Q', 'istar']

        direct_shock = sj.ImpulseDict({'B': B_path_c[c] - ss['B'],
                                       'rstar': rstar_path - ss['r'],
                                       'betashock': betashock_path_c[c] - ss['betashock'],
                                       'x': x_path_c[c] - ss['x']})

        direct_results = Jc @ direct_shock

        imp['A'] += direct_shock['B']
        imp['A'] -= direct_results['A']
        imp['Q'] -= Q_weight * direct_results['Q']

    J = sj.JacobianDict(J, T=T)
    imp = sj.ImpulseDict(imp, T=T)

    return -(J.factored() @ imp), J


# Solve World transition (non-linearly)
def solve_world_transition_nl(x, hank, unknowns, targets, ss_by_country_initial, ss_by_country_final,
                           B_path_c, rstar_path, betashock_path_c, x_path_c, T):

    Cstar_path = x[0:T]
    istar_path = x[T:]

    irf_C_R = {}
    ss0_C = {}
    ss_new_C = {}

    asset_path = np.zeros(T)
    bond_path = np.zeros(T)
    nfa_path = np.zeros(T)
    asset_clear = np.zeros(T)
    price_clear = np.zeros(T)

    which_countries = list(ss_by_country_initial.keys())

    for c, cou in enumerate(which_countries) :

        # Initial steady state (realrate does not matter for SS but need to update for transition)
        ss0 = copy.deepcopy(ss_by_country_initial[cou])
        ss0.update({'realrate': 0, 'phi_pi':0, 'phi_piH': 1.5}) # taylor rule
        ss0_C[cou] = ss0

        # Final steady state
        ss_new = ss_by_country_final[cou]
        ss_new.update({'realrate': 0, 'phi_pi':0, 'phi_piH': 1.5}) # taylor rule
        ss_new_C[cou] = ss_new

        # exogenous path to feed (in difference from steady state values)
        print("Running:", cou)
        exogenous = {"B": B_path_c[cou] - ss_new['B'],
                     'betashock': betashock_path_c[cou] - ss_new['betashock'],
                     'x': x_path_c[cou] - ss_new['x'],
                     "Cstar": Cstar_path - ss_new['Cstar'],
                     "istar": istar_path - ss_new['istar'],
                     "rstar": rstar_path - ss_new['rstar']}  # rstar is the Taylor rule intercept
            # nonlinear transition for this country (in deviation from final steady state) at these inputs

        irf_r =  hank.solve_impulse_nonlinear(ss_new, unknowns, targets, exogenous, ss_initial = ss0, internals=['household_ha'])
        irf_C_R[cou] = irf_r.copy()

        # clearing conditions ss_new['A']+irf_r['A']
        asset_path = asset_path + (ss_new['A']+irf_r['A'])/(ss_new['Q']+irf_r['Q'])
        bond_path = bond_path +  (ss_new['B']+irf_r['B'])/(ss_new['Q']+irf_r['Q'])
        nfa_path = nfa_path + (ss_new['nfa']+irf_r['nfa'])/(ss_new['Q']+irf_r['Q'])

        if ss_new['gamma']==1:
            price_clear = price_clear +  ss_new['omega']*np.log(ss_new['phw']+ irf_r['phw'])
        else:
            price_clear = price_clear +  ss_new['omega']*(ss_new['phw']+irf_r['phw'])**(1-ss_new['gamma'])

    if ss_new['gamma']==1:
        price_clear = np.exp(price_clear) -1
    else:
        price_clear = price_clear -1

    asset_clear = asset_path - bond_path

    return (irf_C_R, Cstar_path, istar_path, asset_path, bond_path, price_clear, nfa_path, ss_new_C)
