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

# Packages
import numpy as np
from ogusa import utils

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


def replacement_rate_vals(nssmat, wss, factor_ss, j, p):
    '''
    Calculates replacement rate values for the social security system.

    Args:
        nssmat (Numpy array): initial guess at labor supply, size = SxJ
        new_w (scalar): steady state real wage rate
        factor_ss (scalar): scaling factor converting model units to
            dollars
        j (int): index of lifetime income group
        p (OG-USA Specifications object): model parameters

    Returns:
        theta (Numpy array): social security replacement rate value for
            lifetime income group j

    '''
    if j is not None:
        e = p.e[:, j]
    else:
        e = p.e
    # adjust number of calendar years AIME computed from int model periods
    equiv_periods = int(round((p.S / 80.0) * p.AIME_num_years)) - 1
    if e.ndim == 2:
        dim2 = e.shape[1]
    else:
        dim2 = 1
    earnings = (e * (wss * nssmat * factor_ss)).reshape(p.S, dim2)
    # get highest earning years for number of years AIME computed from
    highest_earn =\
        (-1.0 * np.sort(-1.0 * earnings[:p.retire[-1], :],
                        axis=0))[:equiv_periods]
    AIME = highest_earn.sum(0) / ((12.0 * (p.S / 80.0)) * equiv_periods)
    PIA = np.zeros(dim2)
    # Compute level of replacement using AIME brackets and PIA rates
    for j in range(dim2):
        if AIME[j] < p.AIME_bkt_1:
            PIA[j] = p.PIA_rate_bkt_1 * AIME[j]
        elif AIME[j] < p.AIME_bkt_2:
            PIA[j] = (p.PIA_rate_bkt_1 * p.AIME_bkt_1 +
                      p.PIA_rate_bkt_2 * (AIME[j] - p.AIME_bkt_1))
        else:
            PIA[j] = (p.PIA_rate_bkt_1 * p.AIME_bkt_1 +
                      p.PIA_rate_bkt_2 * (p.AIME_bkt_2 - p.AIME_bkt_1) +
                      p.PIA_rate_bkt_3 * (AIME[j] - p.AIME_bkt_2))
    # Set the maximum monthly replacment rate from SS benefits tables
    PIA[PIA > p.PIA_maxpayment] = p.PIA_maxpayment
    if p.PIA_minpayment != 0.0:
        PIA[PIA < p.PIA_minpayment] = p.PIA_minpayment
    theta = (PIA * (12.0 * p.S / 80.0)) / (factor_ss * wss)
    return theta


def ETR_wealth(b, h_wealth, m_wealth, p_wealth):
    r'''
    Calculates the effective tax rate on wealth.

    .. math::
        T_{j,s,t}^{w} = \frac{h^{w}p_{w}b_{j,s,t}}{h^{w}b_{j,s,t} + m^{w}}

    Args:
        b (Numpy array): savings
        h_wealth (scalar): parameter of wealth tax function
        p_wealth (scalar): parameter of wealth tax function
        m_wealth (scalar): parameter of wealth tax function

    Returns:
        tau_w (Numpy array): effective tax rate on wealth, size = SxJ

    '''
    tau_w = (p_wealth * h_wealth * b) / (h_wealth * b + m_wealth)
    return tau_w


def MTR_wealth(b, h_wealth, m_wealth, p_wealth):
    r'''
    Calculates the marginal tax rate on wealth from the wealth tax.

    .. math::
        \frac{\partial T_{j,s,t}^{w}}{\partial b_{j,s,t}} = \frac{h^{w}m^{w}p_{w}}{(b_{j,s,t}h^{w}m^{w})^{2}}

    Args:
        b (Numpy array): savings
        h_wealth (scalar): parameter of wealth tax function
        p_wealth (scalar): parameter of wealth tax function
        m_wealth (scalar): parameter of wealth tax function

    Returns:
        tau_prime (Numpy array): marginal tax rate on wealth, size = SxJ

    '''
    tau_prime = ((b * h_wealth * m_wealth * p_wealth) /
                 ((b * h_wealth + m_wealth) ** 2) +
                 ETR_wealth(b, h_wealth, m_wealth, p_wealth))
    return tau_prime


def ETR_income(r, w, b, n, factor, e, etr_params, p):
    '''
    Calculates effective personal income tax rate.

    Args:
        r (array_like): real interest rate
        w (array_like): real wage rate
        b (Numpy array): savings
        n (Numpy array): labor supply
        factor (scalar): scaling factor converting model units to
            dollars
        e (Numpy array): effective labor units
        etr_params (Numpy array): effective tax rate function parameters
        p (OG-USA Specifications object): model parameters

    Returns:
        tau (Numpy array): effective tax rate on total income

    '''
    X = (w * e * n) * factor
    Y = (r * b) * factor
    X2 = X ** 2
    Y2 = Y ** 2
    income = X + Y
    income2 = income ** 2

    if p.tax_func_type == 'GS':
        phi0 = np.squeeze(etr_params[..., 0])
        phi1 = np.squeeze(etr_params[..., 1])
        phi2 = np.squeeze(etr_params[..., 2])
        tau = ((phi0 * (income - ((income ** -phi1) + phi2) **
                        (-1 / phi1))) / income)
    elif p.tax_func_type == 'DEP_totalinc':
        A = np.squeeze(etr_params[..., 0])
        B = np.squeeze(etr_params[..., 1])
        max_income = np.squeeze(etr_params[..., 4])
        min_income = np.squeeze(etr_params[..., 5])
        shift_income = np.squeeze(etr_params[..., 8])
        shift = np.squeeze(etr_params[..., 10])
        tau_income = (((max_income - min_income) *
                       (A * income2 + B * income) /
                       (A * income2 + B * income + 1)) + min_income)
        tau = tau_income + shift_income + shift
    else:  # DEP or linear
        A = np.squeeze(etr_params[..., 0])
        B = np.squeeze(etr_params[..., 1])
        C = np.squeeze(etr_params[..., 2])
        D = np.squeeze(etr_params[..., 3])
        max_x = np.squeeze(etr_params[..., 4])
        min_x = np.squeeze(etr_params[..., 5])
        max_y = np.squeeze(etr_params[..., 6])
        min_y = np.squeeze(etr_params[..., 7])
        shift_x = np.squeeze(etr_params[..., 8])
        shift_y = np.squeeze(etr_params[..., 9])
        shift = np.squeeze(etr_params[..., 10])
        share = np.squeeze(etr_params[..., 11])

        tau_x = ((max_x - min_x) * (A * X2 + B * X) /
                 (A * X2 + B * X + 1) + min_x)
        tau_y = ((max_y - min_y) * (C * Y2 + D * Y) /
                 (C * Y2 + D * Y + 1) + min_y)
        tau = (((tau_x + shift_x) ** share) *
               ((tau_y + shift_y) ** (1 - share))) + shift

    return tau


def MTR_income(r, w, b, n, factor, mtr_capital, e, etr_params,
               mtr_params, p):
    r'''
    Generates the marginal tax rate on labor income for households.

    Args:
        r (array_like): real interest rate
        w (array_like): real wage rate
        b (Numpy array): savings
        n (Numpy array): labor supply
        factor (scalar): scaling factor converting model units to
            dollars
        mtr_capital (bool): whether to compute the marginal tax rate on
            capital income or labor income
        e (Numpy array): effective labor units
        etr_params (Numpy array): effective tax rate function parameters
        p (OG-USA Specifications object): model parameters

    Returns:
        tau (Numpy array): marginal tax rate on income source

    '''
    X = (w * e * n) * factor
    Y = (r * b) * factor
    X2 = X ** 2
    Y2 = Y ** 2
    income = X + Y
    income2 = income ** 2

    if p.tax_func_type == 'GS':
        if p.analytical_mtrs:
            phi0 = np.squeeze(etr_params[..., 0])
            phi1 = np.squeeze(etr_params[..., 1])
            phi2 = np.squeeze(etr_params[..., 2])
        else:
            phi0 = np.squeeze(mtr_params[..., 0])
            phi1 = np.squeeze(mtr_params[..., 1])
            phi2 = np.squeeze(mtr_params[..., 2])
        tau = (phi0*(1 - (income ** (-phi1 - 1) *
                          ((income ** -phi1) + phi2) **
                          ((-1 - phi1) / phi1))))
    elif p.tax_func_type == 'DEP_totalinc':
        if p.analytical_mtrs:
            A = np.squeeze(etr_params[..., 0])
            B = np.squeeze(etr_params[..., 1])
            max_income = np.squeeze(etr_params[..., 4])
            min_income = np.squeeze(etr_params[..., 5])
            shift_income = np.squeeze(etr_params[..., 8])
            shift = np.squeeze(etr_params[..., 10])
            d_etr = ((max_income - min_income) * ((2 * A * income + B) /
                     ((A * income2 + B * income + 1) ** 2)))
            etr = (((max_income - min_income) *
                    ((A * income2 + B * income) /
                     (A * income2 + B * income + 1)) + min_income) +
                   shift_income + shift)
            tau = (d_etr * income) + (etr)
        else:
            A = np.squeeze(mtr_params[..., 0])
            B = np.squeeze(mtr_params[..., 1])
            max_income = np.squeeze(mtr_params[..., 4])
            min_income = np.squeeze(mtr_params[..., 5])
            shift_income = np.squeeze(mtr_params[..., 8])
            shift = np.squeeze(mtr_params[..., 10])
            tau_income = (((max_income - min_income) *
                           (A * income2 + B * income) /
                           (A * income2 + B * income + 1)) + min_income)
            tau = tau_income + shift_income + shift
    else:  # DEP or linear
        if p.analytical_mtrs:
            A = np.squeeze(etr_params[..., 0])
            B = np.squeeze(etr_params[..., 1])
            C = np.squeeze(etr_params[..., 2])
            D = np.squeeze(etr_params[..., 3])
            max_x = np.squeeze(etr_params[..., 4])
            min_x = np.squeeze(etr_params[..., 5])
            max_y = np.squeeze(etr_params[..., 6])
            min_y = np.squeeze(etr_params[..., 7])
            shift_x = np.squeeze(etr_params[..., 8])
            shift_y = np.squeeze(etr_params[..., 9])
            shift = np.squeeze(etr_params[..., 10])
            share = np.squeeze(etr_params[..., 11])

            tau_x = ((max_x - min_x) * (A * X2 + B * X) /
                     (A * X2 + B * X + 1) + min_x)
            tau_y = ((max_y - min_y) * (C * Y2 + D * Y) /
                     (C * Y2 + D * Y + 1) + min_y)
            etr = (((tau_x + shift_x) ** share) *
                   ((tau_y + shift_y) ** (1 - share))) + shift
            if mtr_capital:
                d_etr = ((1-share) * ((tau_y + shift_y) ** (-share)) *
                         (max_y - min_y) * ((2 * C * Y + D) /
                                            ((C * Y2 + D * Y + 1)
                                             ** 2)) *
                         ((tau_x + shift_x) ** share))
                tau = d_etr * income + etr
            else:
                d_etr = (share * ((tau_x + shift_x) ** (share - 1)) *
                         (max_x - min_x) * ((2 * A * X + B) /
                                            ((A * X2 + B * X + 1)
                                             ** 2)) *
                         ((tau_y + shift_y) ** (1 - share)))
                tau = d_etr * income + etr
        else:
            A = np.squeeze(mtr_params[..., 0])
            B = np.squeeze(mtr_params[..., 1])
            C = np.squeeze(mtr_params[..., 2])
            D = np.squeeze(mtr_params[..., 3])
            max_x = np.squeeze(mtr_params[..., 4])
            min_x = np.squeeze(mtr_params[..., 5])
            max_y = np.squeeze(mtr_params[..., 6])
            min_y = np.squeeze(mtr_params[..., 7])
            shift_x = np.squeeze(mtr_params[..., 8])
            shift_y = np.squeeze(mtr_params[..., 9])
            shift = np.squeeze(mtr_params[..., 10])
            share = np.squeeze(mtr_params[..., 11])

            tau_x = ((max_x - min_x) * (A * X2 + B * X) /
                     (A * X2 + B * X + 1) + min_x)
            tau_y = ((max_y - min_y) * (C * Y2 + D * Y) /
                     (C * Y2 + D * Y + 1) + min_y)
            tau = (((tau_x + shift_x) ** share) *
                   ((tau_y + shift_y) ** (1 - share))) + shift

    return tau


def get_biz_tax(w, Y, L, K, p, method):
    r'''
    Finds total business income tax revenue.

    .. math::
        R_{t}^{b} = \tau_{t}^{b}(Y_{t} - w_{t}L_{t}) - \tau_{t}^{b}\delta_{t}^{\tau}K_{t}^{\tau}
    Args:
        r (array_like): real interest rate
        Y (array_like): aggregate output
        L (array_like): aggregate labor demand
        K (array_like): aggregate capital demand

    Returns:
        business_revenue (array_like): aggregate business tax revenue

    '''
    if method == 'SS':
        delta_tau = p.delta_tau[-1]
        tau_b = p.tau_b[-1]
    else:
        delta_tau = p.delta_tau[:p.T]
        tau_b = p.tau_b[:p.T]
    business_revenue = tau_b * (Y - w * L) - tau_b * delta_tau * K
    return business_revenue


def net_taxes(r, w, b, n, bq, factor, tr, theta, t, j, shift, method,
              e, etr_params, p):
    '''
    Calculate net taxes paid for each household.

    Args:
        r (array_like): real interest rate
        w (array_like): real wage rate
        b (Numpy array): savings
        n (Numpy array): labor supply
        bq (Numpy array): bequests received
        factor (scalar): scaling factor converting model units to
            dollars
        tr (Numpy array): government transfers to the household
        theta (Numpy array): social security replacement rate value for
            lifetime income group j
        t (int): time period
        j (int): index of lifetime income group
        shift (bool): whether computing for periods 0--s or 1--(s+1),
            =True for 1--(s+1)
        method (str): adjusts calculation dimensions based on 'SS' or
            'TPI'
        e (Numpy array): effective labor units
        etr_params (Numpy array): effective tax rate function parameters
        p (OG-USA Specifications object): model parameters

    Returns:
        net_tax (Numpy array): net taxes paid for each household

    '''
    T_I = income_tax_liab(r, w, b, n, factor, t, j, method, e, etr_params, p)
    pension = pension_amount(w, n, theta, t, j, shift, method, e, p)
    T_BQ = bequest_tax_liab(r, b, bq, t, j, method, p)
    T_W = wealth_tax_liab(r, b, t, j, method, p)

    net_tax = T_I - pension + T_BQ + T_W - tr

    return net_tax


def income_tax_liab(r, w, b, n, factor, t, j, method, e, etr_params, p):
    '''
    Calculate income and payroll tax liability for each household

    Args:
        r (array_like): real interest rate
        w (array_like): real wage rate
        b (Numpy array): savings
        n (Numpy array): labor supply
        factor (scalar): scaling factor converting model units to
            dollars
        t (int): time period
        j (int): index of lifetime income group
        method (str): adjusts calculation dimensions based on 'SS' or
            'TPI'
        e (Numpy array): effective labor units
        etr_params (Numpy array): effective tax rate function parameters
        p (OG-USA Specifications object): model parameters

    Returns:
        T_I (Numpy array): total income and payroll taxes paid for each
            household

    '''
    if j is not None:
        if method == 'TPI':
            if b.ndim == 2:
                r = r.reshape(r.shape[0], 1)
                w = w.reshape(w.shape[0], 1)
    else:
        if method == 'TPI':
            r = utils.to_timepath_shape(r)
            w = utils.to_timepath_shape(w)

    income = r * b + w * e * n
    labor_income = w * e * n
    T_I = ETR_income(r, w, b, n, factor, e, etr_params, p) * income
    if method == 'SS':
        T_P = p.tau_payroll[-1] * labor_income
    elif method == 'TPI':
        length = w.shape[0]
        if len(b.shape) == 1:
            T_P = p.tau_payroll[t: t + length] * labor_income
        elif len(b.shape) == 2:
            T_P = (p.tau_payroll[t: t + length].reshape(length, 1) *
                   labor_income)
        else:
            T_P = (p.tau_payroll[t:t + length].reshape(length, 1, 1) *
                   labor_income)
    elif method == 'TPI_scalar':
        T_P = p.tau_payroll[0] * labor_income

    income_payroll_tax_liab = T_I + T_P

    return income_payroll_tax_liab


def pension_amount(w, n, theta, t, j, shift, method, e, p):
    '''
    Calculate public pension benefit amounts for each household.

    Args:
        w (array_like): real wage rate
        n (Numpy array): labor supply
        theta (Numpy array): social security replacement rate value for
            lifetime income group j
        t (int): time period
        j (int): index of lifetime income group
        shift (bool): whether computing for periods 0--s or 1--(s+1),
            =True for 1--(s+1)
        method (str): adjusts calculation dimensions based on 'SS' or
            'TPI'
        e (Numpy array): effective labor units
        p (OG-USA Specifications object): model parameters

    Returns:
        pension (Numpy array): pension amount for each household

    '''
    if j is not None:
        if method == 'TPI':
            if n.ndim == 2:
                w = w.reshape(w.shape[0], 1)
    else:
        if method == 'TPI':
            w = utils.to_timepath_shape(w)

    pension = np.zeros_like(n)
    if method == 'SS':
        # Depending on if we are looking at b_s or b_s+1, the
        # entry for retirement will change (it shifts back one).
        # The shift boolean makes sure we start replacement rates
        # at the correct age.
        if shift is False:
            pension[p.retire[-1]:] = theta * w
        else:
            pension[p.retire[-1] - 1:] = theta * w
    elif method == 'TPI':
        length = w.shape[0]
        if not shift:
            # retireTPI is different from retire, because in TP income
            # we are counting backwards with different length lists.
            # This will always be the correct location of retirement,
            # depending on the shape of the lists.
            retireTPI = (p.retire[t: t + length] - p.S)
        else:
            retireTPI = (p.retire[t: t + length] - 1 - p.S)
        if len(n.shape) == 1:
            if not shift:
                retireTPI = p.retire[t] - p.S
            else:
                retireTPI = p.retire[t] - 1 - p.S
            pension[retireTPI:] = (
                theta[j] * p.replacement_rate_adjust[t] * w[retireTPI:])
        elif len(n.shape) == 2:
            for tt in range(pension.shape[0]):
                pension[tt, retireTPI[tt]:] = (
                    theta * p.replacement_rate_adjust[t + tt] * w[tt])
        else:
            for tt in range(pension.shape[0]):
                pension[tt, retireTPI[tt]:, :] = (
                    theta.reshape(1, p.J) *
                    p.replacement_rate_adjust[t + tt] * w[tt])
    elif method == 'TPI_scalar':
        # The above methods won't work if scalars are used.  This option
        # is only called by the SS_TPI_firstdoughnutring function in TPI.
        pension = theta * p.replacement_rate_adjust[0] * w

    return pension


def wealth_tax_liab(r, b, t, j, method, p):
    '''
    Calculate wealth tax liability for each household.

    Args:
        r (array_like): real interest rate
        b (Numpy array): savings
        t (int): time period
        j (int): index of lifetime income group
        method (str): adjusts calculation dimensions based on 'SS' or
            'TPI'
        p (OG-USA Specifications object): model parameters

    Returns:
        T_W (Numpy array): wealth tax liability for each household

    '''
    if j is not None:
        if method == 'TPI':
            if b.ndim == 2:
                r = r.reshape(r.shape[0], 1)
    else:
        if method == 'TPI':
            r = utils.to_timepath_shape(r)

    if method == 'SS':
        T_W = (ETR_wealth(b, p.h_wealth[-1], p.m_wealth[-1],
                          p.p_wealth[-1]) * b)
    elif method == 'TPI':
        length = r.shape[0]
        if len(b.shape) == 1:
            T_W = (ETR_wealth(b, p.h_wealth[t:t + length],
                              p.m_wealth[t:t + length],
                              p.p_wealth[t:t + length]) * b)
        elif len(b.shape) == 2:
            T_W = (ETR_wealth(b, p.h_wealth[t:t + length],
                              p.m_wealth[t:t + length],
                              p.p_wealth[t:t + length]) * b)
        else:
            T_W = (ETR_wealth(
                b, p.h_wealth[t:t + length].reshape(length, 1, 1),
                p.m_wealth[t:t + length].reshape(length, 1, 1),
                p.p_wealth[t:t + length].reshape(length, 1, 1)) * b)
    elif method == 'TPI_scalar':
        T_W = (ETR_wealth(b, p.h_wealth[0], p.m_wealth[0],
                          p.p_wealth[0]) * b)

    return T_W


def bequest_tax_liab(r, b, bq, t, j, method, p):
    '''
    Calculate liability due from taxes on bequests for each household.

    Args:
        r (array_like): real interest rate
        b (Numpy array): savings
        bq (Numpy array): bequests received
        t (int): time period
        j (int): index of lifetime income group
        method (str): adjusts calculation dimensions based on 'SS' or
            'TPI'
        p (OG-USA Specifications object): model parameters

    Returns:
        T_BQ (Numpy array): bequest tax liability for each household

    '''
    if j is not None:
        lambdas = p.lambdas[j]
        if method == 'TPI':
            if b.ndim == 2:
                r = r.reshape(r.shape[0], 1)
    else:
        lambdas = np.transpose(p.lambdas)
        if method == 'TPI':
            r = utils.to_timepath_shape(r)

    if method == 'SS':
        T_BQ = p.tau_bq[-1] * bq
    elif method == 'TPI':
        length = r.shape[0]
        if len(b.shape) == 1:
            T_BQ = p.tau_bq[t:t + length] * bq
        elif len(b.shape) == 2:
            T_BQ = p.tau_bq[t:t + length].reshape(length, 1) * bq / lambdas
        else:
            T_BQ = p.tau_bq[t:t + length].reshape(length, 1, 1) * bq
    elif method == 'TPI_scalar':
        # The above methods won't work if scalars are used.  This option
        # is only called by the SS_TPI_firstdoughnutring function in TPI.
        T_BQ = p.tau_bq[0] * bq

    return T_BQ

[((51, 10, 51, 24), 'numpy.zeros', 'np.zeros', ({(51, 19, 51, 23): 'dim2'}, {}), '(dim2)', True, 'import numpy as np\n'), ((461, 14, 461, 30), 'numpy.zeros_like', 'np.zeros_like', ({(461, 28, 461, 29): 'n'}, {}), '(n)', True, 'import numpy as np\n'), ((142, 15, 142, 45), 'numpy.squeeze', 'np.squeeze', ({(142, 26, 142, 44): 'etr_params[..., 0]'}, {}), '(etr_params[..., 0])', True, 'import numpy as np\n'), ((143, 15, 143, 45), 'numpy.squeeze', 'np.squeeze', ({(143, 26, 143, 44): 'etr_params[..., 1]'}, {}), '(etr_params[..., 1])', True, 'import numpy as np\n'), ((144, 15, 144, 45), 'numpy.squeeze', 'np.squeeze', ({(144, 26, 144, 44): 'etr_params[..., 2]'}, {}), '(etr_params[..., 2])', True, 'import numpy as np\n'), ((579, 18, 579, 41), 'numpy.transpose', 'np.transpose', ({(579, 31, 579, 40): 'p.lambdas'}, {}), '(p.lambdas)', True, 'import numpy as np\n'), ((48, 16, 49, 31), 'numpy.sort', 'np.sort', (), '', True, 'import numpy as np\n'), ((148, 12, 148, 42), 'numpy.squeeze', 'np.squeeze', ({(148, 23, 148, 41): 'etr_params[..., 0]'}, {}), '(etr_params[..., 0])', True, 'import numpy as np\n'), ((149, 12, 149, 42), 'numpy.squeeze', 'np.squeeze', ({(149, 23, 149, 41): 'etr_params[..., 1]'}, {}), '(etr_params[..., 1])', True, 'import numpy as np\n'), ((150, 21, 150, 51), 'numpy.squeeze', 'np.squeeze', ({(150, 32, 150, 50): 'etr_params[..., 4]'}, {}), '(etr_params[..., 4])', True, 'import numpy as np\n'), ((151, 21, 151, 51), 'numpy.squeeze', 'np.squeeze', ({(151, 32, 151, 50): 'etr_params[..., 5]'}, {}), '(etr_params[..., 5])', True, 'import numpy as np\n'), ((152, 23, 152, 53), 'numpy.squeeze', 'np.squeeze', ({(152, 34, 152, 52): 'etr_params[..., 8]'}, {}), '(etr_params[..., 8])', True, 'import numpy as np\n'), ((153, 16, 153, 47), 'numpy.squeeze', 'np.squeeze', ({(153, 27, 153, 46): 'etr_params[..., 10]'}, {}), '(etr_params[..., 10])', True, 'import numpy as np\n'), ((159, 12, 159, 42), 'numpy.squeeze', 'np.squeeze', ({(159, 23, 159, 41): 'etr_params[..., 0]'}, {}), '(etr_params[..., 0])', True, 'import numpy as np\n'), ((160, 12, 160, 42), 'numpy.squeeze', 'np.squeeze', ({(160, 23, 160, 41): 'etr_params[..., 1]'}, {}), '(etr_params[..., 1])', True, 'import numpy as np\n'), ((161, 12, 161, 42), 'numpy.squeeze', 'np.squeeze', ({(161, 23, 161, 41): 'etr_params[..., 2]'}, {}), '(etr_params[..., 2])', True, 'import numpy as np\n'), ((162, 12, 162, 42), 'numpy.squeeze', 'np.squeeze', ({(162, 23, 162, 41): 'etr_params[..., 3]'}, {}), '(etr_params[..., 3])', True, 'import numpy as np\n'), ((163, 16, 163, 46), 'numpy.squeeze', 'np.squeeze', ({(163, 27, 163, 45): 'etr_params[..., 4]'}, {}), '(etr_params[..., 4])', True, 'import numpy as np\n'), ((164, 16, 164, 46), 'numpy.squeeze', 'np.squeeze', ({(164, 27, 164, 45): 'etr_params[..., 5]'}, {}), '(etr_params[..., 5])', True, 'import numpy as np\n'), ((165, 16, 165, 46), 'numpy.squeeze', 'np.squeeze', ({(165, 27, 165, 45): 'etr_params[..., 6]'}, {}), '(etr_params[..., 6])', True, 'import numpy as np\n'), ((166, 16, 166, 46), 'numpy.squeeze', 'np.squeeze', ({(166, 27, 166, 45): 'etr_params[..., 7]'}, {}), '(etr_params[..., 7])', True, 'import numpy as np\n'), ((167, 18, 167, 48), 'numpy.squeeze', 'np.squeeze', ({(167, 29, 167, 47): 'etr_params[..., 8]'}, {}), '(etr_params[..., 8])', True, 'import numpy as np\n'), ((168, 18, 168, 48), 'numpy.squeeze', 'np.squeeze', ({(168, 29, 168, 47): 'etr_params[..., 9]'}, {}), '(etr_params[..., 9])', True, 'import numpy as np\n'), ((169, 16, 169, 47), 'numpy.squeeze', 'np.squeeze', ({(169, 27, 169, 46): 'etr_params[..., 10]'}, {}), '(etr_params[..., 10])', True, 'import numpy as np\n'), ((170, 16, 170, 47), 'numpy.squeeze', 'np.squeeze', ({(170, 27, 170, 46): 'etr_params[..., 11]'}, {}), '(etr_params[..., 11])', True, 'import numpy as np\n'), ((213, 19, 213, 49), 'numpy.squeeze', 'np.squeeze', ({(213, 30, 213, 48): 'etr_params[..., 0]'}, {}), '(etr_params[..., 0])', True, 'import numpy as np\n'), ((214, 19, 214, 49), 'numpy.squeeze', 'np.squeeze', ({(214, 30, 214, 48): 'etr_params[..., 1]'}, {}), '(etr_params[..., 1])', True, 'import numpy as np\n'), ((215, 19, 215, 49), 'numpy.squeeze', 'np.squeeze', ({(215, 30, 215, 48): 'etr_params[..., 2]'}, {}), '(etr_params[..., 2])', True, 'import numpy as np\n'), ((217, 19, 217, 49), 'numpy.squeeze', 'np.squeeze', ({(217, 30, 217, 48): 'mtr_params[..., 0]'}, {}), '(mtr_params[..., 0])', True, 'import numpy as np\n'), ((218, 19, 218, 49), 'numpy.squeeze', 'np.squeeze', ({(218, 30, 218, 48): 'mtr_params[..., 1]'}, {}), '(mtr_params[..., 1])', True, 'import numpy as np\n'), ((219, 19, 219, 49), 'numpy.squeeze', 'np.squeeze', ({(219, 30, 219, 48): 'mtr_params[..., 2]'}, {}), '(mtr_params[..., 2])', True, 'import numpy as np\n'), ((405, 16, 405, 42), 'ogusa.utils.to_timepath_shape', 'utils.to_timepath_shape', ({(405, 40, 405, 41): 'r'}, {}), '(r)', False, 'from ogusa import utils\n'), ((406, 16, 406, 42), 'ogusa.utils.to_timepath_shape', 'utils.to_timepath_shape', ({(406, 40, 406, 41): 'w'}, {}), '(w)', False, 'from ogusa import utils\n'), ((459, 16, 459, 42), 'ogusa.utils.to_timepath_shape', 'utils.to_timepath_shape', ({(459, 40, 459, 41): 'w'}, {}), '(w)', False, 'from ogusa import utils\n'), ((528, 16, 528, 42), 'ogusa.utils.to_timepath_shape', 'utils.to_timepath_shape', ({(528, 40, 528, 41): 'r'}, {}), '(r)', False, 'from ogusa import utils\n'), ((581, 16, 581, 42), 'ogusa.utils.to_timepath_shape', 'utils.to_timepath_shape', ({(581, 40, 581, 41): 'r'}, {}), '(r)', False, 'from ogusa import utils\n'), ((225, 16, 225, 46), 'numpy.squeeze', 'np.squeeze', ({(225, 27, 225, 45): 'etr_params[..., 0]'}, {}), '(etr_params[..., 0])', True, 'import numpy as np\n'), ((226, 16, 226, 46), 'numpy.squeeze', 'np.squeeze', ({(226, 27, 226, 45): 'etr_params[..., 1]'}, {}), '(etr_params[..., 1])', True, 'import numpy as np\n'), ((227, 25, 227, 55), 'numpy.squeeze', 'np.squeeze', ({(227, 36, 227, 54): 'etr_params[..., 4]'}, {}), '(etr_params[..., 4])', True, 'import numpy as np\n'), ((228, 25, 228, 55), 'numpy.squeeze', 'np.squeeze', ({(228, 36, 228, 54): 'etr_params[..., 5]'}, {}), '(etr_params[..., 5])', True, 'import numpy as np\n'), ((229, 27, 229, 57), 'numpy.squeeze', 'np.squeeze', ({(229, 38, 229, 56): 'etr_params[..., 8]'}, {}), '(etr_params[..., 8])', True, 'import numpy as np\n'), ((230, 20, 230, 51), 'numpy.squeeze', 'np.squeeze', ({(230, 31, 230, 50): 'etr_params[..., 10]'}, {}), '(etr_params[..., 10])', True, 'import numpy as np\n'), ((239, 16, 239, 46), 'numpy.squeeze', 'np.squeeze', ({(239, 27, 239, 45): 'mtr_params[..., 0]'}, {}), '(mtr_params[..., 0])', True, 'import numpy as np\n'), ((240, 16, 240, 46), 'numpy.squeeze', 'np.squeeze', ({(240, 27, 240, 45): 'mtr_params[..., 1]'}, {}), '(mtr_params[..., 1])', True, 'import numpy as np\n'), ((241, 25, 241, 55), 'numpy.squeeze', 'np.squeeze', ({(241, 36, 241, 54): 'mtr_params[..., 4]'}, {}), '(mtr_params[..., 4])', True, 'import numpy as np\n'), ((242, 25, 242, 55), 'numpy.squeeze', 'np.squeeze', ({(242, 36, 242, 54): 'mtr_params[..., 5]'}, {}), '(mtr_params[..., 5])', True, 'import numpy as np\n'), ((243, 27, 243, 57), 'numpy.squeeze', 'np.squeeze', ({(243, 38, 243, 56): 'mtr_params[..., 8]'}, {}), '(mtr_params[..., 8])', True, 'import numpy as np\n'), ((244, 20, 244, 51), 'numpy.squeeze', 'np.squeeze', ({(244, 31, 244, 50): 'mtr_params[..., 10]'}, {}), '(mtr_params[..., 10])', True, 'import numpy as np\n'), ((251, 16, 251, 46), 'numpy.squeeze', 'np.squeeze', ({(251, 27, 251, 45): 'etr_params[..., 0]'}, {}), '(etr_params[..., 0])', True, 'import numpy as np\n'), ((252, 16, 252, 46), 'numpy.squeeze', 'np.squeeze', ({(252, 27, 252, 45): 'etr_params[..., 1]'}, {}), '(etr_params[..., 1])', True, 'import numpy as np\n'), ((253, 16, 253, 46), 'numpy.squeeze', 'np.squeeze', ({(253, 27, 253, 45): 'etr_params[..., 2]'}, {}), '(etr_params[..., 2])', True, 'import numpy as np\n'), ((254, 16, 254, 46), 'numpy.squeeze', 'np.squeeze', ({(254, 27, 254, 45): 'etr_params[..., 3]'}, {}), '(etr_params[..., 3])', True, 'import numpy as np\n'), ((255, 20, 255, 50), 'numpy.squeeze', 'np.squeeze', ({(255, 31, 255, 49): 'etr_params[..., 4]'}, {}), '(etr_params[..., 4])', True, 'import numpy as np\n'), ((256, 20, 256, 50), 'numpy.squeeze', 'np.squeeze', ({(256, 31, 256, 49): 'etr_params[..., 5]'}, {}), '(etr_params[..., 5])', True, 'import numpy as np\n'), ((257, 20, 257, 50), 'numpy.squeeze', 'np.squeeze', ({(257, 31, 257, 49): 'etr_params[..., 6]'}, {}), '(etr_params[..., 6])', True, 'import numpy as np\n'), ((258, 20, 258, 50), 'numpy.squeeze', 'np.squeeze', ({(258, 31, 258, 49): 'etr_params[..., 7]'}, {}), '(etr_params[..., 7])', True, 'import numpy as np\n'), ((259, 22, 259, 52), 'numpy.squeeze', 'np.squeeze', ({(259, 33, 259, 51): 'etr_params[..., 8]'}, {}), '(etr_params[..., 8])', True, 'import numpy as np\n'), ((260, 22, 260, 52), 'numpy.squeeze', 'np.squeeze', ({(260, 33, 260, 51): 'etr_params[..., 9]'}, {}), '(etr_params[..., 9])', True, 'import numpy as np\n'), ((261, 20, 261, 51), 'numpy.squeeze', 'np.squeeze', ({(261, 31, 261, 50): 'etr_params[..., 10]'}, {}), '(etr_params[..., 10])', True, 'import numpy as np\n'), ((262, 20, 262, 51), 'numpy.squeeze', 'np.squeeze', ({(262, 31, 262, 50): 'etr_params[..., 11]'}, {}), '(etr_params[..., 11])', True, 'import numpy as np\n'), ((285, 16, 285, 46), 'numpy.squeeze', 'np.squeeze', ({(285, 27, 285, 45): 'mtr_params[..., 0]'}, {}), '(mtr_params[..., 0])', True, 'import numpy as np\n'), ((286, 16, 286, 46), 'numpy.squeeze', 'np.squeeze', ({(286, 27, 286, 45): 'mtr_params[..., 1]'}, {}), '(mtr_params[..., 1])', True, 'import numpy as np\n'), ((287, 16, 287, 46), 'numpy.squeeze', 'np.squeeze', ({(287, 27, 287, 45): 'mtr_params[..., 2]'}, {}), '(mtr_params[..., 2])', True, 'import numpy as np\n'), ((288, 16, 288, 46), 'numpy.squeeze', 'np.squeeze', ({(288, 27, 288, 45): 'mtr_params[..., 3]'}, {}), '(mtr_params[..., 3])', True, 'import numpy as np\n'), ((289, 20, 289, 50), 'numpy.squeeze', 'np.squeeze', ({(289, 31, 289, 49): 'mtr_params[..., 4]'}, {}), '(mtr_params[..., 4])', True, 'import numpy as np\n'), ((290, 20, 290, 50), 'numpy.squeeze', 'np.squeeze', ({(290, 31, 290, 49): 'mtr_params[..., 5]'}, {}), '(mtr_params[..., 5])', True, 'import numpy as np\n'), ((291, 20, 291, 50), 'numpy.squeeze', 'np.squeeze', ({(291, 31, 291, 49): 'mtr_params[..., 6]'}, {}), '(mtr_params[..., 6])', True, 'import numpy as np\n'), ((292, 20, 292, 50), 'numpy.squeeze', 'np.squeeze', ({(292, 31, 292, 49): 'mtr_params[..., 7]'}, {}), '(mtr_params[..., 7])', True, 'import numpy as np\n'), ((293, 22, 293, 52), 'numpy.squeeze', 'np.squeeze', ({(293, 33, 293, 51): 'mtr_params[..., 8]'}, {}), '(mtr_params[..., 8])', True, 'import numpy as np\n'), ((294, 22, 294, 52), 'numpy.squeeze', 'np.squeeze', ({(294, 33, 294, 51): 'mtr_params[..., 9]'}, {}), '(mtr_params[..., 9])', True, 'import numpy as np\n'), ((295, 20, 295, 51), 'numpy.squeeze', 'np.squeeze', ({(295, 31, 295, 50): 'mtr_params[..., 10]'}, {}), '(mtr_params[..., 10])', True, 'import numpy as np\n'), ((296, 20, 296, 51), 'numpy.squeeze', 'np.squeeze', ({(296, 31, 296, 50): 'mtr_params[..., 11]'}, {}), '(mtr_params[..., 11])', True, 'import numpy as np\n')]