'''
------------------------------------------------------------------------
Household functions for taxes in the steady state and along the
transition path..

This file calls the following files:
    tax.py
------------------------------------------------------------------------
'''

# Packages
import numpy as np
import tax

'''
------------------------------------------------------------------------
    Functions
------------------------------------------------------------------------
'''


def marg_ut_cons(c, sigma):
    '''
    Computation of marginal utility of consumption.
    Inputs:
        c     = [T,S,J] array, household consumption
        sigma = scalar, coefficient of relative risk aversion
    Functions called: None
    Objects in function:
        output = [T,S,J] array, marginal utility of consumption
    Returns: output
    '''
    if np.ndim(c) == 0:
        c = np.array([c])
    # epsilon = 0.0001
    epsilon = 0.002
    cvec_cnstr = c < epsilon
    MU_c = np.zeros(c.shape)
    MU_c[~cvec_cnstr] = c[~cvec_cnstr] ** (-sigma)
    b2 = (-sigma * (epsilon ** (-sigma - 1))) / 2
    b1 = (epsilon ** (-sigma)) - 2 * b2 * epsilon
    MU_c[cvec_cnstr] = 2 * b2 * c[cvec_cnstr] + b1
    output = MU_c
    output = np.squeeze(output)

    return output


def marg_ut_labor(n, params):
    '''
    Computation of marginal disutility of labor.
    Inputs:
        n         = [T,S,J] array, household labor supply
        params    = length 4 tuple (b_ellipse, upsilon, ltilde, chi_n)
        b_ellipse = scalar, scaling parameter in elliptical utility function
        upsilon   = curvature parameter in elliptical utility function
        ltilde    = scalar, upper bound of household labor supply
        chi_n     = [S,] vector, utility weights on disutility of labor
    Functions called: None
    Objects in function:
        output = [T,S,J] array, marginal disutility of labor supply
    Returns: output
    '''
    b_ellipse, upsilon, ltilde, chi_n = params

    b_ellip = b_ellipse
    nvec = n
    if np.ndim(nvec) == 0:
        nvec = np.array([nvec])
    eps_low = 0.0001#0.000001
    eps_high = ltilde - 0.0001 #0.000001
    nvec_low = nvec < eps_low
    nvec_high = nvec > eps_high
    nvec_uncstr = np.logical_and(~nvec_low, ~nvec_high)
    MDU_n = np.zeros(nvec.shape)
    MDU_n[nvec_uncstr] = (
        (b_ellip / ltilde) *
        ((nvec[nvec_uncstr] / ltilde) ** (upsilon - 1)) *
        ((1 - ((nvec[nvec_uncstr] / ltilde) ** upsilon)) **
         ((1 - upsilon) / upsilon)))
    b2 = (0.5 * b_ellip * (ltilde ** (-upsilon)) * (upsilon - 1) *
          (eps_low ** (upsilon - 2)) *
          ((1 - ((eps_low / ltilde) ** upsilon)) **
          ((1 - upsilon) / upsilon)) *
          (1 + ((eps_low / ltilde) ** upsilon) *
          ((1 - ((eps_low / ltilde) ** upsilon)) ** (-1))))
    b1 = ((b_ellip / ltilde) * ((eps_low / ltilde) ** (upsilon - 1)) *
          ((1 - ((eps_low / ltilde) ** upsilon)) **
          ((1 - upsilon) / upsilon)) - (2 * b2 * eps_low))
    MDU_n[nvec_low] = 2 * b2 * nvec[nvec_low] + b1
    d2 = (0.5 * b_ellip * (ltilde ** (-upsilon)) * (upsilon - 1) *
          (eps_high ** (upsilon - 2)) *
          ((1 - ((eps_high / ltilde) ** upsilon)) **
          ((1 - upsilon) / upsilon)) *
          (1 + ((eps_high / ltilde) ** upsilon) *
          ((1 - ((eps_high / ltilde) ** upsilon)) ** (-1))))
    d1 = ((b_ellip / ltilde) * ((eps_high / ltilde) **
          (upsilon - 1)) * ((1 - ((eps_high / ltilde) ** upsilon)) **
          ((1 - upsilon) / upsilon)) - (2 * d2 * eps_high))
    MDU_n[nvec_high] = 2 * d2 * nvec[nvec_high] + d1
    output = MDU_n*chi_n
    output = np.squeeze(output)

    return output


def get_cons(r, w, b, b_splus1, n, BQ, net_tax, params):
    '''
    Calculation of househld consumption.

    Inputs:
        r        = [T,] vector, interest rates
        w        = [T,] vector, wage rates
        b        = [T,S,J] array, distribution of wealth/capital
        b_splus1 = [T,S,J] array, distribution of wealth/capital,
                    one period ahead
        n        = [T,S,J] array, distribution of labor supply
        BQ       = [T,J] array, bequests by lifetime income group
        net_tax  = [T,S,J] array, distribution of net taxes
        params    = length 3 tuple (e, lambdas, g_y)
        e        = [S,J] array, effective labor units by age and
                    lifetime income group
        lambdas  = [S,] vector, fraction of population in each lifetime
                    income group
        g_y      = scalar, exogenous labor augmenting technological growth

    Functions called: None

    Objects in function:
        cons = [T,S,J] array, household consumption

    Returns: cons
    '''
    e, lambdas, g_y = params

    cons = ((1 + r) * b + w * e * n + BQ / lambdas - b_splus1 *
            np.exp(g_y) - net_tax)
    return cons


def FOC_savings(r, w, b, b_splus1, b_splus2, n, BQ, factor, T_H,
                params):
    '''
    Computes Euler errors for the FOC for savings in the steady state.
    This function is usually looped through over J, so it does one
    lifetime income group at a time.

    Inputs:
        r           = scalar, interest rate
        w           = scalar, wage rate
        b           = [S,J] array, distribution of wealth/capital
        b_splus1    = [S,J] array, distribution of wealth/capital,
                        one period ahead
        b_splus2    = [S,J] array, distribution of wealth/capital, two
                        periods ahead
        n           = [S,J] array, distribution of labor supply
        BQ          = [J,] vector, aggregate bequests by lifetime income
                        group
        factor      = scalar, scaling factor to convert model income to
                        dollars
        T_H         = scalar, lump sum transfer
        params      = length 18 tuple (e, sigma, beta, g_y, chi_b,
                                       theta, tau_bq, rho, lambdas, J,
                                       S, etr_params, mtry_params,
                                       h_wealth, p_wealth, m_wealth,
                                       tau_payroll, tau_bq)
        e           = [S,J] array, effective labor units
        sigma       = scalar, coefficient of relative risk aversion
        beta        = scalar, discount factor
        g_y         = scalar, exogenous labor augmenting technological
                        growth
        chi_b       = [J,] vector, utility weight on bequests for each
                        lifetime income group
        theta       = [J,] vector, replacement rate for each lifetime
                        income group
        tau_bq      = scalar, bequest tax rate (scalar)
        rho         = [S,] vector, mortality rates
        lambdas     = [J,] vector, ability weights
        J           = integer, number of lifetime income groups
        S           = integer, number of economically active periods in
                        lifetime
        etr_params  = [S,12] array, parameters of effective income tax
                        rate function
        mtry_params = [S,12] array, parameters of marginal tax rate on
                        capital income function
        h_wealth    = scalar, parameter in wealth tax function
        p_wealth    = scalar, parameter in wealth tax function
        m_wealth    = scalar, parameter in wealth tax function
        tau_payroll = scalar, payroll tax rate
        tau_bq      = scalar, bequest tax rate

    Functions called:
        get_cons
        marg_ut_cons
        tax.total_taxes
        tax.MTR_capital

    Objects in function:
        tax1 = [S,J] array, net taxes in the current period
        tax2 = [S,J] array, net taxes one period ahead
        cons1 = [S,J] array, consumption in the current period
        cons2 = [S,J] array, consumption one period ahead
        deriv = [S,J] array, after-tax return on capital
        savings_ut = [S,J] array, marginal utility from savings
        euler = [S,J] array, Euler error from FOC for savings

    Returns: euler
    '''
    (e, sigma, beta, g_y, chi_b, theta, tau_bq, rho, lambdas, j, J, S,
     analytical_mtrs, etr_params, mtry_params, h_wealth, p_wealth,
     m_wealth, tau_payroll, retire, method) = params

    # In order to not have 2 savings euler equations (one that solves
    # the first S-1 equations, and one that solves the last one), we
    # combine them.  In order to do this, we have to compute a
    # consumption term in period t+1, which requires us to have a shifted
    # e and n matrix.  We append a zero on the end of both of these so
    # they will be the right size.  We could append any value to them,
    # since in the euler equation, the coefficient on the marginal
    # utility of consumption for this term will be zero (since rho is
    # one).
    e_extended = np.array(list(e) + [0])
    n_extended = np.array(list(n) + [0])
    etr_params_extended = np.append(etr_params,
                                    np.reshape(etr_params[-1, :],
                                               (1, etr_params.shape[1])),
                                    axis=0)[1:, :]
    mtry_params_extended = np.append(mtry_params,
                                     np.reshape(mtry_params[-1, :],
                                                (1, mtry_params.shape[1])),
                                     axis=0)[1:, :]
    if method == 'TPI':
        r_extended = np.append(r, r[-1])
        w_extended = np.append(w, w[-1])
        BQ_extended = np.append(BQ, BQ[-1])
        T_H_extended = np.append(T_H, T_H[-1])
    elif method == 'SS':
        r_extended = np.array([r, r])
        w_extended = np.array([w, w])
        BQ_extended = np.array([BQ, BQ])
        T_H_extended = np.array([T_H, T_H])

    tax1_params = (e, lambdas, method, retire, etr_params, h_wealth,
                   p_wealth, m_wealth, tau_payroll, theta, tau_bq, J, S)
    tax1 = tax.total_taxes(r, w, b, n, BQ, factor, T_H, j, False,
                           tax1_params)
    tax2_params = (e_extended[1:], lambdas, method, retire,
                   etr_params_extended, h_wealth, p_wealth, m_wealth,
                   tau_payroll, theta, tau_bq, J, S)
    tax2 = tax.total_taxes(r_extended[1:], w_extended[1:], b_splus1,
                           n_extended[1:], BQ_extended[1:], factor,
                           T_H_extended[1:], j, True, tax2_params)
    cons1_params = (e, lambdas, g_y)
    cons1 = get_cons(r, w, b, b_splus1, n, BQ, tax1, cons1_params)
    cons2_params = (e_extended[1:], lambdas, g_y)
    cons2 = get_cons(r_extended[1:], w_extended[1:], b_splus1, b_splus2,
                     n_extended[1:], BQ_extended[1:], tax2,
                     cons2_params)
    cons2[-1] = 0.01  # set to small positive number to avoid exception
    # errors when negative b/c this period value doesn't matter -
    # it's consumption after the last period of life
    mtr_cap_params = (e_extended[1:], etr_params_extended,
                      mtry_params_extended, analytical_mtrs)
    deriv = ((1 + r_extended[1:]) - r_extended[1:] *
             (tax.MTR_capital(r_extended[1:], w_extended[1:], b_splus1,
                              n_extended[1:], factor, mtr_cap_params)) -
             (tax.tau_w_prime(b_splus1, (h_wealth, p_wealth, m_wealth)) *
             b_splus1) - tax.tau_wealth(b_splus1, (h_wealth, p_wealth,
                                                   m_wealth)))

    savings_ut = (rho * np.exp(-sigma * g_y) * chi_b * b_splus1 **
                  (-sigma))

    euler_error = (marg_ut_cons(cons1, sigma) - beta * (1 - rho) *
                   deriv * marg_ut_cons(cons2, sigma) *
                   np.exp(-sigma * g_y) - savings_ut)

    return euler_error


def FOC_labor(r, w, b, b_splus1, n, BQ, factor, T_H, params):
    '''
    Computes Euler errors for the FOC for labor supply in the steady
    state.  This function is usually looped through over J, so it does
    one lifetime income group at a time.

    Inputs:
        r           = scalar, interest rate
        w           = scalar, wage rate
        b           = [S,J] array, distribution of wealth/capital
                        holdings
        b_splus1    = [S,J] array, distribution of wealth/capital
                        holdings one period ahead
        n           = [S,J] array, distribution of labor supply
        BQ          = [J,] vector, aggregate bequests by lifetime
                        income group
        factor      = scalar, scaling factor to convert model income to
                        dollars
        T_H         = scalar, lump sum transfer
        params      = length 19 tuple (e, sigma, g_y, theta, b_ellipse,
                                       upsilon, ltilde, chi_n, tau_bq,
                                       lambdas, J, S, etr_params,
                                       mtrx_params, h_wealth, p_wealth,
                                       m_wealth, tau_payroll, tau_bq)
        e           = [S,J] array, effective labor units
        sigma       = scalar, coefficient of relative risk aversion
        g_y         = scalar, exogenous labor augmenting technological
                        growth
        theta       = [J,] vector, replacement rate for each lifetime
                        income group
        b_ellipse   = scalar, scaling parameter in elliptical utility
                        function
        upsilon     = curvature parameter in elliptical utility function
        chi_n       = [S,] vector, utility weights on disutility of labor
        ltilde      = scalar, upper bound of household labor supply
        tau_bq      = scalar, bequest tax rate (scalar)
        lambdas     = [J,] vector, ability weights
        J           = integer, number of lifetime income groups
        S           = integer, number of economically active periods in
                        lifetime
        etr_params  = [S,10] array, parameters of effective income tax
                        rate function
        mtrx_params = [S,10] array, parameters of marginal tax rate on
                        labor income function
        h_wealth    = scalar, parameter in wealth tax function
        p_wealth    = scalar, parameter in wealth tax function
        m_wealth    = scalar, parameter in wealth tax function
        tau_payroll = scalar, payroll tax rate
        tau_bq      = scalar, bequest tax rate

    Functions called:
        get_cons
        marg_ut_cons
        marg_ut_labor
        tax.total_taxes
        tax.MTR_labor

    Objects in function:
        tax = [S,J] array, net taxes in the current period
        cons = [S,J] array, consumption in the current period
        deriv = [S,J] array, net of tax share of labor income
        euler = [S,J] array, Euler error from FOC for labor supply

    Returns: euler
    '''
    (e, sigma, g_y, theta, b_ellipse, upsilon, chi_n, ltilde, tau_bq,
     lambdas, j, J, S, analytical_mtrs, etr_params, mtrx_params,
     h_wealth, p_wealth, m_wealth, tau_payroll, retire, method) = params

    tax1_params = (e, lambdas, method, retire, etr_params, h_wealth,
                   p_wealth, m_wealth, tau_payroll, theta, tau_bq, J, S)
    tax1 = tax.total_taxes(r, w, b, n, BQ, factor, T_H, j, False,
                           tax1_params)
    cons_params = (e, lambdas, g_y)
    cons = get_cons(r, w, b, b_splus1, n, BQ, tax1, cons_params)
    mtr_lab_params = (e, etr_params, mtrx_params, analytical_mtrs)
    deriv = (1 - tau_payroll - tax.MTR_labor(r, w, b, n, factor,
                                             mtr_lab_params))

    lab_params = (b_ellipse, upsilon, ltilde, chi_n)
    FOC_error = (marg_ut_cons(cons, sigma) * w * deriv * e -
                 marg_ut_labor(n, lab_params))

    return FOC_error


def get_K(b, params):
    '''
    Calculates aggregate capital supplied.

    Inputs:
        b           = [T,S,J] array, distribution of wealth/capital holdings
        params      = length 4 tuple, (omega, lambdas, g_n, method)
        omega       = [S,T] array, population weights
        lambdas     = [J,] vector, fraction in each lifetime income group
        g_n         = [T,] vector, population growth rate
        method      = string, 'SS' or 'TPI'

    Functions called: None

    Objects in function:
        K_presum = [T,S,J] array, weighted distribution of wealth/capital holdings
        K        = [T,] vector, aggregate capital supply

    Returns: K
    '''

    omega, lambdas, imm_rates, g_n, method = params

    if method == 'SS':
        part1 = b* omega * lambdas
        omega_extended = np.append(omega[1:],[0.0])
        imm_extended = np.append(imm_rates[1:],[0.0])
        part2 = b*(omega_extended*imm_extended).reshape(omega.shape[0],1)*lambdas
        K_presum = part1+part2
        K = K_presum.sum()
    elif method == 'TPI':
        part1 = b* omega * lambdas
        #omega_extended = np.append(omega[1:,:,:],np.zeros((1,omega.shape[1],omega.shape[2])),axis=0)
        omega_shift = np.append(omega[:,1:,:],np.zeros((omega.shape[0],1,omega.shape[2])),axis=1)
        #imm_extended = np.append(imm_rates[1:,:,:],np.zeros((1,imm_rates.shape[1],imm_rates.shape[2])),axis=0)
        imm_shift = np.append(imm_rates[:,1:,:],np.zeros((imm_rates.shape[0],1,imm_rates.shape[2])),axis=1)
        #part2 = b*(omega_extended*imm_extended)*lambdas
        part2 = b*imm_shift*omega_shift*lambdas
        K_presum = part1+part2
        K = K_presum.sum(1).sum(1)
    K /= (1.0 + g_n)
    return K


def get_BQ(r, b_splus1, params):
    '''
    Calculation of bequests to each lifetime income group.

    Inputs:
        r           = [T,] vector, interest rates
        b_splus1    = [T,S,J] array, distribution of wealth/capital holdings one period ahead
        params      = length 5 tuple, (omega, lambdas, rho, g_n, method)
        omega       = [S,T] array, population weights
        lambdas     = [J,] vector, fraction in each lifetime income group
        rho         = [S,] vector, mortality rates
        g_n         = scalar, population growth rate
        method      = string, 'SS' or 'TPI'

    Functions called: None

    Objects in function:
        BQ_presum = [T,S,J] array, weighted distribution of wealth/capital holdings one period ahead
        BQ        = [T,J] array, aggregate bequests by lifetime income group

    Returns: BQ
    '''
    omega, lambdas, rho, g_n, method = params

    BQ_presum = b_splus1 * omega * rho * lambdas
    if method == 'SS':
        BQ = BQ_presum.sum(0)
    elif method == 'TPI':
        BQ = BQ_presum.sum(1)
    BQ *= (1.0 + r) / (1.0 + g_n)
    return BQ


def get_u(c, n, b_splus1, params):
    '''
    Computes flow utility for the household.

    Inputs:
        b_splus1 = [S,J] array, steady state distribution of capital
        n = [S,J] array, steady state distribution of labor
        c = [S,J] array, steady state distribution of consumption
        sigma = scalar, coefficient of relative risk aversion
        chi_n  = [S,] vector of utility weights for disulity of labor
        b_ellipse = scalar, scale parameter on elliptical utility
        ltilde = scalar, upper bound of household labor supply
        upsilon = scalar, curvature parameter on elliptical utility
        k_ellipse = scalar, shift parameter on elliptical utility
        rho_s = [S,] vector, mortality rates by age
        chi_b = [J,] vector, utility weights on bequests
        g_y = scalar, economic growth rate

    Functions called: None

    Objects in function:
        utility = [S,J] array, utility for all agents

    Returns:
        utility
    '''
    sigma, chi_n, b_ellipse, ltilde, upsilon, rho_s, chi_b = params

    utility = (((c ** (1-sigma) - 1) / (1 - sigma)) +
               (chi_n * ((b_ellipse * (1 - (n / ltilde) ** upsilon)
                          ** (1 / upsilon)))) +
               (rho_s * chi_b * ((b_splus1 ** (1-sigma) - 1)
                                 / (1 - sigma))))

    return utility


def get_C(c, params):
    '''
    Calculation of aggregate consumption.

    Inputs:
        cons        = [T,S,J] array, household consumption
        params      = length 3 tuple (omega, lambdas, method)
        omega       = [S,T] array, population weights by age (Sx1 array)
        lambdas     = [J,1] vector, lifetime income group weights
        method      = string, 'SS' or 'TPI'

    Functions called: None

    Objects in function:
        aggC_presum = [T,S,J] array, weighted consumption by household
        aggC        = [T,] vector, aggregate consumption

    Returns: aggC
    '''

    omega, lambdas, method = params

    aggC_presum = c * omega * lambdas
    if method == 'SS':
        aggC = aggC_presum.sum()
    elif method == 'TPI':
        aggC = aggC_presum.sum(1).sum(1)
    return aggC


def constraint_checker_SS(bssmat, nssmat, cssmat, ltilde):
    '''
    Checks constraints on consumption, savings, and labor supply in the
    steady state.

    Inputs:
        bssmat = [S,J] array, steady state distribution of capital
        nssmat = [S,J] array, steady state distribution of labor
        cssmat = [S,J] array, steady state distribution of consumption
        ltilde = scalar, upper bound of household labor supply

    Functions called: None

    Objects in function:
        flag2 = boolean, indicates if labor supply constraints violated
                (=False if not)

    Returns:
        # Prints warnings for violations of capital, labor, and
            consumption constraints.
    '''
    print 'Checking constraints on capital, labor, and consumption.'

    if (bssmat < 0).any():
        print '\tWARNING: There is negative capital stock'
    flag2 = False
    if (nssmat < 0).any():
        print '\tWARNING: Labor supply violates nonnegativity ' +\
            'constraints.'
        flag2 = True
    if (nssmat > ltilde).any():
        print '\tWARNING: Labor suppy violates the ltilde constraint.'
        flag2 = True
    if flag2 is False:
        print '\tThere were no violations of the constraints on labor'\
         + ' supply.'
    if (cssmat < 0).any():
        print '\tWARNING: Consumption violates nonnegativity' +\
            ' constraints.'
    else:
        print '\tThere were no violations of the constraints on' +\
            ' consumption.'


def constraint_checker_TPI(b_dist, n_dist, c_dist, t, ltilde):
    '''
    Checks constraints on consumption, savings, and labor supply along
    the transition path. Does this for each period t separately.

    Inputs:
        b_dist = [S,J] array, distribution of capital
        n_dist = [S,J] array, distribution of labor
        c_dist = [S,J] array, distribution of consumption
        t      = integer, time period
        ltilde = scalar, upper bound of household labor supply

    Functions called: None

    Objects in function: None

    Returns:
        # Prints warnings for violations of capital, labor, and
            consumption constraints.
    '''
    if (b_dist <= 0).any():
        print '\tWARNING: Aggregate capital is less than or equal to '\
            'zero in period %.f.' % t
    if (n_dist < 0).any():
        print '\tWARNING: Labor supply violates nonnegativity' +\
            ' constraints in period %.f.' % t
    if (n_dist > ltilde).any():
        print '\tWARNING: Labor suppy violates the ltilde constraint' +\
            ' in period %.f.' % t
    if (c_dist < 0).any():
        print '\tWARNING: Consumption violates nonnegativity' +\
            ' constraints in period %.f.' % t
