"""Defines functions for calibration of the main model and for solving steady state"""
from model_blocks import (hh_ha, asset_clearing, rsimple, compute_iMPCs)
import numpy as np
from scipy.optimize import brentq, fsolve
import matplotlib.pyplot as plt
from sequence_jacobian import create_model

# This initializes a dictionary with parameters and (r,nfa,B,G,Transfer) etc
def init_calibration():

    calibration = {'r' : 0.0/4, 'eis' : 1, 'gamma' : 1, 'eta' : 1, 'alpha' : 0.157, 'omega':0.15, 'omega_US': 0.15,
                     'frisch' : 0.5, 'varphi' : 1, 'prop': 1/3, 'betashock':1, 'x':1,
                    'epsilon_w' : 2, 'lbda':0.181, 'mpctarget': 0.25}

    gamma_w = calibration['epsilon_w']/(calibration['epsilon_w'] - 1) * (1-calibration['lbda'])
    chi     = (1-calibration['alpha'])*calibration['eta'] + calibration['gamma']
    calibration.update({'gamma_w': gamma_w, 'chi': chi})

    # Add household side parameters
    calibration.update({'rho_e': 0.92, 'sd_e_uncscale': 0.92, 'nE': 7,
                 'amin': 0.0, 'amax': 150, 'nA': 500,
                 'n_beta' : 3, 'beta_spread': 0.06, 'beta' : 0.98
                })
    sd_e = calibration['sd_e_uncscale']*(1-calibration['lbda'])
    calibration.update({'sd_e': sd_e})

    # Other parameters for transition
    calibration.update({'kappa_w': 0.1, 'phi_pi': 1.5,  'phi_piH':0.0, 'realrate': 1
                })

    # Variables
    calibration.update({'Theta': 1,'P': 1, 'G_over_Y' : 0.141, 'B_over_Y' : 0.819, 'nfa' : 0.0*4, 'Transfer': 0.0,
                         'Q': 1, 'Y': 0.5219  })
    calibration['G'] = calibration['G_over_Y'] * calibration['Y']
    calibration['B'] = calibration['B_over_Y'] * calibration['Y'] * 4

    # Other variables for transition
    # calibration.update({'piw': 0,'piFH': 0, 'piHHres':0, 'PHHres':0, 'piFHres':0, 'PFHres':0, 'Pres':0,
    #                    'pioutres':0.0, 'piwres':0, 'fisher':0, 'uip':0, 'uip_E':0,
    #                     'piw_target': 0, 'piwres_eq': 0,
    #                    'phh':1,'pfh':1})

    calibration.update({'piw': 0, 'piHres':0, 'PHres':0,
                       'pioutres':0.0, 'piwres':0, 'fisher':0, 'uip':0,
                        'piw_target': 0, 'piwres_eq': 0,
                       'php':1,'pwp':1, 'PH':1, 'phw':1})

    return calibration


# ## Step 1 - Not used
# - Solve for $(Y,C,Q)$
# Equations for SS to hit Y=Q=1
def step1_eqns(calibration):

    r, nfa, B, G, Transfer, Y = calibration['r'], calibration['nfa'], calibration['B'], calibration['G'],  calibration['Transfer'], calibration['Y']
    alpha, gamma_w, eis = calibration['alpha'], calibration['gamma_w'], calibration['eis']

    C = r*nfa + Y - G
    Cstar = -(1-alpha)/alpha*(r*nfa) + Y - G
    varphi = gamma_w * (Y - G - r*B - Transfer) * C**(-1/eis)

    return C, Cstar, varphi


# ## Step 2
# - Solve for $(\beta, \Delta)$
# Calibrates beta given beta_spread, nfa, B and (r,Z)
def solve_beta(model, calibration, print_beta = 0):

    calibration['Z'] = calibration['Y'] - calibration['G'] - calibration['r']*calibration['B'] - calibration['Transfer']

    unknowns_ss = {'beta': calibration["beta"] }
    targets_ss = {'asset_mkt': 0 }
    ss = model.solve_steady_state(calibration, unknowns_ss, targets_ss, solver="brentq",
                                   solver_kwargs={"bracket": (0.98, -0.002 + 1/(1+calibration['r'])  )})
    if print_beta == 1:
        print("beta: ", ss['beta'])
    return ss

# Calibrates beta_spread and beta
def solve_beta_spread(beta_spread, model, calibration, opt = 0, print_spread = 0, print_beta = 1):

    calibration.update({ "beta_spread": beta_spread })
    if print_spread == 1:
        print("Delta: ", beta_spread)

    # Solve beta for a given beta_spread
    ss = solve_beta(model, calibration, print_beta = print_beta)
    calibration.update({'beta':ss['beta'], "iMPC":compute_iMPCs(ss, 0)})
    if opt == 1:
        return calibration['iMPC'][0,0] - calibration['mpctarget']   # If solving for beta_spread using optimizer
    else:
        return calibration


# - Plot MPC vs $\Delta$
def step2_mpc(model, calibration,  bmin=0, bmax=0.2, nB = 2):

    # Plot MPC for different Delta (calibrating beta to clear A=nfa+B)
    beta_spread_list = np.linspace(bmin,bmax,nB)
    MPCs = np.zeros((len(beta_spread_list),2))
    for b, beta_spread in enumerate(beta_spread_list):
        calibration = solve_beta_spread(beta_spread, model, calibration, opt=0)
        MPCs[b,:] = [calibration['iMPC'][0,0], calibration['iMPC'][1,0]]
    plt.plot(beta_spread_list, MPCs[:,0]  , label = r'MPC0')
    plt.plot(beta_spread_list, MPCs[:,1]  , label = r'MPC1')
    plt.axhline(y=calibration['mpctarget'], color='#808080', linestyle=':', label = 'MPC Target')
    plt.xlabel(r"beta spread $\Delta$")
    plt.legend()
    plt.show()


# ## Combine Step 1 + Step 2
# Calibrate SS given parameters and (r,nfa,B,G,Transfer)
# If opt = 1, then calibrate (beta,delta) to hit A = nfa+B and mpc0
def calibrate_ss(calibration, use_data = 1, plot = 0, opt = 0, bounds = [0.05, 0.1], bmin=0, bmax=0.15 ):

    ############################################################
    # Step 1 - Get (C,Cstar,varphi)
    if use_data == 0:
        C, Cstar, varphi = step1_eqns(calibration)
        calibration.update( {'C': C, 'Cstar': Cstar, 'varphi': varphi})

    ############################################################
    # Step 2 to solve for (beta,delta)
    blocks_step2 = [hh_ha, asset_clearing, rsimple]
    hank_step2   = create_model(blocks_step2, name = "Baseline SS Step 2")

    # 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)
    calibration.update({'xi': xi })

    # Plot MPC vs beta_spread for Step 2 (calibrating beta to A=nfa+b)
    if plot == 1:
        step2_mpc(hank_step2, calibration, bmin=bmin, bmax=bmax, nB = 2) # this is to get bounds

    if opt == 1:
        print("Solving for beta and beta_spread...")
        # Solve for \Delta
        beta_spread, sol = brentq(solve_beta_spread, bounds[0], bounds[1], args = (hank_step2, calibration, opt, 0, 0), full_output=True, xtol=1e-12, rtol = 1e-12)
#         beta_spread = fsolve(solve_beta_spread, [calibration['beta_spread']], args = (hank_step2, calibration,  opt, 0, 0))
        calibration.update({ "beta_spread": beta_spread })

        # Get Base SS and calibration
        print("Done! Calibrated parameters are:")
        calibration_base = solve_beta_spread(beta_spread, hank_step2, calibration, opt = 0, print_spread = 1, print_beta = 1)
        return calibration_base
    else:
        return calibration
