from scipy.stats import norm
from numpy import array, polyval, inf, asarray

__all__ = ['mackinnonp', 'mackinnoncrit']

# These are the cut-off values for the left-tail vs. the rest of the
# tau distribution, for getting the p-values

tau_star_nc = [-1.04, -1.53, -2.68, -3.09, -3.07, -3.77]
tau_min_nc = [-19.04, -19.62, -21.21, -23.25, -21.63, -25.74]
tau_max_nc = [inf, 1.51, 0.86, 0.88, 1.05, 1.24]
tau_star_c = [-1.61, -2.62, -3.13, -3.47, -3.78, -3.93]
tau_min_c = [-18.83, -18.86, -23.48, -28.07, -25.96, -23.27]
tau_max_c = [2.74, 0.92, 0.55, 0.61, 0.79, 1]
tau_star_ct = [-2.89, -3.19, -3.50, -3.65, -3.80, -4.36]
tau_min_ct = [-16.18, -21.15, -25.37, -26.63, -26.53, -26.18]
tau_max_ct = [0.7, 0.63, 0.71, 0.93, 1.19, 1.42]
tau_star_ctt = [-3.21, -3.51, -3.81, -3.83, -4.12, -4.63]
tau_min_ctt = [-17.17, -21.1, -24.33, -24.03, -24.33, -28.22]
tau_max_ctt = [0.54, 0.79, 1.08, 1.43, 3.49, 1.92]

_tau_maxs = {
    "nc": tau_max_nc,
    "c": tau_max_c,
    "ct": tau_max_ct,
    "ctt": tau_max_ctt,
}
_tau_mins = {
    "nc": tau_min_nc,
    "c": tau_min_c,
    "ct": tau_min_ct,
    "ctt": tau_min_ctt,
}
_tau_stars = {
    "nc": tau_star_nc,
    "c": tau_star_c,
    "ct": tau_star_ct,
    "ctt": tau_star_ctt,
}


small_scaling = array([1, 1, 1e-2])
tau_nc_smallp = [
    [0.6344, 1.2378, 3.2496],
    [1.9129, 1.3857, 3.5322],
    [2.7648, 1.4502, 3.4186],
    [3.4336, 1.4835, 3.19],
    [4.0999, 1.5533, 3.59],
    [4.5388, 1.5344, 2.9807]]
tau_nc_smallp = asarray(tau_nc_smallp)*small_scaling

tau_c_smallp = [
    [2.1659, 1.4412, 3.8269],
    [2.92, 1.5012, 3.9796],
    [3.4699, 1.4856, 3.164],
    [3.9673, 1.4777, 2.6315],
    [4.5509, 1.5338, 2.9545],
    [5.1399, 1.6036, 3.4445]]
tau_c_smallp = asarray(tau_c_smallp)*small_scaling

tau_ct_smallp = [
    [3.2512, 1.6047, 4.9588],
    [3.6646, 1.5419, 3.6448],
    [4.0983, 1.5173, 2.9898],
    [4.5844, 1.5338, 2.8796],
    [5.0722, 1.5634, 2.9472],
    [5.53, 1.5914, 3.0392]]
tau_ct_smallp = asarray(tau_ct_smallp)*small_scaling

tau_ctt_smallp = [
    [4.0003, 1.658, 4.8288],
    [4.3534, 1.6016, 3.7947],
    [4.7343, 1.5768, 3.2396],
    [5.214, 1.6077, 3.3449],
    [5.6481, 1.6274, 3.3455],
    [5.9296, 1.5929, 2.8223]]
tau_ctt_smallp = asarray(tau_ctt_smallp)*small_scaling

_tau_smallps = {
    "nc": tau_nc_smallp,
    "c": tau_c_smallp,
    "ct": tau_ct_smallp,
    "ctt": tau_ctt_smallp,
}


large_scaling = array([1, 1e-1, 1e-1, 1e-2])
tau_nc_largep = [
    [0.4797, 9.3557, -0.6999, 3.3066],
    [1.5578, 8.558, -2.083, -3.3549],
    [2.2268, 6.8093, -3.2362, -5.4448],
    [2.7654, 6.4502, -3.0811, -4.4946],
    [3.2684, 6.8051, -2.6778, -3.4972],
    [3.7268, 7.167, -2.3648, -2.8288]]
tau_nc_largep = asarray(tau_nc_largep)*large_scaling

tau_c_largep = [
    [1.7339, 9.3202, -1.2745, -1.0368],
    [2.1945, 6.4695, -2.9198, -4.2377],
    [2.5893, 4.5168, -3.6529, -5.0074],
    [3.0387, 4.5452, -3.3666, -4.1921],
    [3.5049, 5.2098, -2.9158, -3.3468],
    [3.9489, 5.8933, -2.5359, -2.721]]
tau_c_largep = asarray(tau_c_largep)*large_scaling

tau_ct_largep = [
    [2.5261, 6.1654, -3.7956, -6.0285],
    [2.85, 5.272, -3.6622, -5.1695],
    [3.221, 5.255, -3.2685, -4.1501],
    [3.652, 5.9758, -2.7483, -3.2081],
    [4.0712, 6.6428, -2.3464, -2.546],
    [4.4735, 7.1757, -2.0681, -2.1196]]
tau_ct_largep = asarray(tau_ct_largep)*large_scaling

tau_ctt_largep = [
    [3.0778, 4.9529, -4.1477, -5.9359],
    [3.4713, 5.967, -3.2507, -4.2286],
    [3.8637, 6.7852, -2.6286, -3.1381],
    [4.2736, 7.6199, -2.1534, -2.4026],
    [4.6679, 8.2618, -1.822, -1.9147],
    [5.0009, 8.3735, -1.6994, -1.6928]]
tau_ctt_largep = asarray(tau_ctt_largep)*large_scaling

_tau_largeps = {
    "nc": tau_nc_largep,
    "c": tau_c_largep,
    "ct": tau_ct_largep,
    "ctt": tau_ctt_largep,
}


# NOTE: The Z-statistic is used when lags are included to account for
#  serial correlation in the error term

z_star_nc = [-2.9, -8.7, -14.8, -20.9, -25.7, -30.5]
z_star_c = [-8.9, -14.3, -19.5, -25.1, -29.6, -34.4]
z_star_ct = [-15.0, -19.6, -25.3, -29.6, -31.8, -38.4]
z_star_ctt = [-20.7, -25.3, -29.9, -34.4, -38.5, -44.2]


# These are Table 5 from MacKinnon (1994)
# small p is defined as p in .005 to .150 ie p = .005 up to z_star
# Z* is the largest value for which it is appropriate to use these
# approximations
# the left tail approximation is
# p = norm.cdf(d_0 + d_1*log(abs(z)) + d_2*log(abs(z))**2 + d_3*log(abs(z))**3)
# there is no Z-min, ie., it is well-behaved in the left tail

z_nc_smallp = array([
    [.0342, -.6376, 0, -.03872],
    [1.3426, -.7680, 0, -.04104],
    [3.8607, -2.4159, .51293, -.09835],
    [6.1072, -3.7250, .85887, -.13102],
    [7.7800, -4.4579, 1.00056, -.14014],
    [4.0253, -.8815, 0, -.04887]])

z_c_smallp = array([
    [2.2142, -1.7863, .32828, -.07727],
    [1.1662, .1814, -.36707, 0],
    [6.6584, -4.3486, 1.04705, -.15011],
    [3.3249, -.8456, 0, -.04818],
    [4.0356, -.9306, 0, -.04776],
    [13.9959, -8.4314, 1.97411, -.22234]])

z_ct_smallp = array([
    [4.6476, -2.8932, 0.5832, -0.0999],
    [7.2453, -4.7021, 1.127, -.15665],
    [3.4893, -0.8914, 0, -.04755],
    [1.6604, 1.0375, -0.53377, 0],
    [2.006, 1.1197, -0.55315, 0],
    [11.1626, -5.6858, 1.21479, -.15428]])

z_ctt_smallp = array([
    [3.6739, -1.1549, 0, -0.03947],
    [3.9783, -1.0619, 0, -0.04394],
    [2.0062, 0.8907, -0.51708, 0],
    [4.9218, -1.0663, 0, -0.04691],
    [5.1433, -0.9877, 0, -0.04993],
    [23.6812, -14.6485, 3.42909, -.33794]])
# These are Table 6 from MacKinnon (1994).
# These are well-behaved in the right tail.
# the approximation function is
# p = norm.cdf(d_0 + d_1 * z + d_2*z**2 + d_3*z**3 + d_4*z**4)
z_large_scaling = array([1, 1e-1, 1e-2, 1e-3, 1e-5])
z_nc_largep = array([
    [0.4927, 6.906, 13.2331, 12.099, 0],
    [1.5167, 4.6859, 4.2401, 2.7939, 7.9601],
    [2.2347, 3.9465, 2.2406, 0.8746, 1.4239],
    [2.8239, 3.6265, 1.6738, 0.5408, 0.7449],
    [3.3174, 3.3492, 1.2792, 0.3416, 0.3894],
    [3.729, 3.0611, 0.9579, 0.2087, 0.1943]])
z_nc_largep *= z_large_scaling

z_c_largep = array([
    [1.717, 5.5243, 4.3463, 1.6671, 0],
    [2.2394, 4.2377, 2.432, 0.9241, 0.4364],
    [2.743, 3.626, 1.5703, 0.4612, 0.567],
    [3.228, 3.3399, 1.2319, 0.3162, 0.3482],
    [3.6583, 3.0934, 0.9681, 0.2111, 0.1979],
    [4.0379, 2.8735, 0.7694, 0.1433, 0.1146]])
z_c_largep *= z_large_scaling

z_ct_largep = array([
    [2.7117, 4.5731, 2.2868, 0.6362, 0.5],
    [3.0972, 4.0873, 1.8982, 0.5796, 0.7384],
    [3.4594, 3.6326, 1.4284, 0.3813, 0.4325],
    [3.806, 3.2634, 1.0689, 0.2402, 0.2304],
    [4.1402, 2.9867, 0.8323, 0.16, 0.1315],
    [4.4497, 2.7534, 0.6582, 0.1089, 0.0773]])
z_ct_largep *= z_large_scaling

z_ctt_largep = array([
    [3.4671, 4.3476, 1.9231, 0.5381, 0.6216],
    [3.7827, 3.9421, 1.5699, 0.4093, 0.4485],
    [4.052, 3.4947, 1.1772, 0.2642, 0.2502],
    [4.3311, 3.1625, 0.9126, 0.1775, 0.1462],
    [4.594, 2.8739, 0.707, 0.1181, 0.0838],
    [4.8479, 2.6447, 0.5647, 0.0827, 0.0518]])
z_ctt_largep *= z_large_scaling


# TODO: finish this and then integrate them into adf function
def mackinnonp(teststat, regression="c", N=1, lags=None):
    """
    Returns MacKinnon's approximate p-value for teststat.

    Parameters
    ----------
    teststat : float
        "T-value" from an Augmented Dickey-Fuller regression.
    regression : str {"c", "nc", "ct", "ctt"}
        This is the method of regression that was used.  Following MacKinnon's
        notation, this can be "c" for constant, "nc" for no constant, "ct" for
        constant and trend, and "ctt" for constant, trend, and trend-squared.
    N : int
        The number of series believed to be I(1).  For (Augmented) Dickey-
        Fuller N = 1.

    Returns
    -------
    p-value : float
        The p-value for the ADF statistic estimated using MacKinnon 1994.

    References
    ----------
    .. [*] MacKinnon, J.G. 1994  "Approximate Asymptotic Distribution Functions
        for Unit-Root and Cointegration Tests." Journal of Business & Economics
        Statistics, 12.2, 167-76.

    Notes
    -----
    For (A)DF
    H_0: AR coefficient = 1
    H_a: AR coefficient < 1
    """
    maxstat = _tau_maxs[regression]
    minstat = _tau_mins[regression]
    starstat = _tau_stars[regression]
    if teststat > maxstat[N-1]:
        return 1.0
    elif teststat < minstat[N-1]:
        return 0.0
    if teststat <= starstat[N-1]:
        tau_coef = _tau_smallps[regression][N-1]
    else:
        # Note: above is only for z stats
        tau_coef = _tau_largeps[regression][N-1]
    return norm.cdf(polyval(tau_coef[::-1], teststat))


# These are the new estimates from MacKinnon 2010
# the first axis is N -1
# the second axis is 1 %, 5 %, 10 %
# the last axis is the coefficients

tau_nc_2010 = [[
    [-2.56574, -2.2358, -3.627, 0],  # N = 1
    [-1.94100, -0.2686, -3.365, 31.223],
    [-1.61682, 0.2656, -2.714, 25.364]]]
tau_nc_2010 = asarray(tau_nc_2010)

tau_c_2010 = [
    [[-3.43035, -6.5393, -16.786, -79.433],  # N = 1, 1%
     [-2.86154, -2.8903, -4.234, -40.040],   # 5 %
     [-2.56677, -1.5384, -2.809, 0]],        # 10 %
    [[-3.89644, -10.9519, -33.527, 0],       # N = 2
     [-3.33613, -6.1101, -6.823, 0],
     [-3.04445, -4.2412, -2.720, 0]],
    [[-4.29374, -14.4354, -33.195, 47.433],  # N = 3
     [-3.74066, -8.5632, -10.852, 27.982],
     [-3.45218, -6.2143, -3.718, 0]],
    [[-4.64332, -18.1031, -37.972, 0],       # N = 4
     [-4.09600, -11.2349, -11.175, 0],
     [-3.81020, -8.3931, -4.137, 0]],
    [[-4.95756, -21.8883, -45.142, 0],       # N = 5
     [-4.41519, -14.0405, -12.575, 0],
     [-4.13157, -10.7417, -3.784, 0]],
    [[-5.24568, -25.6688, -57.737, 88.639],  # N = 6
     [-4.70693, -16.9178, -17.492, 60.007],
     [-4.42501, -13.1875, -5.104, 27.877]],
    [[-5.51233, -29.5760, -69.398, 164.295],  # N = 7
     [-4.97684, -19.9021, -22.045, 110.761],
     [-4.69648, -15.7315, -5.104, 27.877]],
    [[-5.76202, -33.5258, -82.189, 256.289],  # N = 8
     [-5.22924, -23.0023, -24.646, 144.479],
     [-4.95007, -18.3959, -7.344, 94.872]],
    [[-5.99742, -37.6572, -87.365, 248.316],  # N = 9
     [-5.46697, -26.2057, -26.627, 176.382],
     [-5.18897, -21.1377, -9.484, 172.704]],
    [[-6.22103, -41.7154, -102.680, 389.33],  # N = 10
     [-5.69244, -29.4521, -30.994, 251.016],
     [-5.41533, -24.0006, -7.514, 163.049]],
    [[-6.43377, -46.0084, -106.809, 352.752],  # N = 11
     [-5.90714, -32.8336, -30.275, 249.994],
     [-5.63086, -26.9693, -4.083, 151.427]],
    [[-6.63790, -50.2095, -124.156, 579.622],  # N = 12
     [-6.11279, -36.2681, -32.505, 314.802],
     [-5.83724, -29.9864, -2.686, 184.116]]]
tau_c_2010 = asarray(tau_c_2010)

tau_ct_2010 = [
    [[-3.95877, -9.0531, -28.428, -134.155],   # N = 1
     [-3.41049, -4.3904, -9.036, -45.374],
     [-3.12705, -2.5856, -3.925, -22.380]],
    [[-4.32762, -15.4387, -35.679, 0],         # N = 2
     [-3.78057, -9.5106, -12.074, 0],
     [-3.49631, -7.0815, -7.538, 21.892]],
    [[-4.66305, -18.7688, -49.793, 104.244],   # N = 3
     [-4.11890, -11.8922, -19.031, 77.332],
     [-3.83511, -9.0723, -8.504, 35.403]],
    [[-4.96940, -22.4694, -52.599, 51.314],    # N = 4
     [-4.42871, -14.5876, -18.228, 39.647],
     [-4.14633, -11.2500, -9.873, 54.109]],
    [[-5.25276, -26.2183, -59.631, 50.646],    # N = 5
     [-4.71537, -17.3569, -22.660, 91.359],
     [-4.43422, -13.6078, -10.238, 76.781]],
    [[-5.51727, -29.9760, -75.222, 202.253],   # N = 6
     [-4.98228, -20.3050, -25.224, 132.03],
     [-4.70233, -16.1253, -9.836, 94.272]],
    [[-5.76537, -33.9165, -84.312, 245.394],   # N = 7
     [-5.23299, -23.3328, -28.955, 182.342],
     [-4.95405, -18.7352, -10.168, 120.575]],
    [[-6.00003, -37.8892, -96.428, 335.92],    # N = 8
     [-5.46971, -26.4771, -31.034, 220.165],
     [-5.19183, -21.4328, -10.726, 157.955]],