import numpy as np
import pandas as pd

import pytest
from numpy.testing import assert_equal, assert_allclose, assert_raises

from statsmodels.tsa.statespace.tools import (
    constrain_stationary_univariate as constrain,
    unconstrain_stationary_univariate as unconstrain)

from statsmodels.tsa.arima import specification


def check_attributes(spec, order, seasonal_order, enforce_stationarity,
                     enforce_invertibility, concentrate_scale):
    p, d, q = order
    P, D, Q, s = seasonal_order

    assert_equal(spec.order, (p, d, q))
    assert_equal(spec.seasonal_order, (P, D, Q, s))

    assert_equal(spec.ar_order, p)
    assert_equal(spec.diff, d)
    assert_equal(spec.ma_order, q)

    assert_equal(spec.seasonal_ar_order, P)
    assert_equal(spec.seasonal_diff, D)
    assert_equal(spec.seasonal_ma_order, Q)
    assert_equal(spec.seasonal_periods, s)

    assert_equal(spec.ar_lags,
                 (p if isinstance(p, list) else np.arange(1, p + 1)))
    assert_equal(spec.ma_lags,
                 (q if isinstance(q, list) else np.arange(1, q + 1)))

    assert_equal(spec.seasonal_ar_lags,
                 (P if isinstance(P, list) else np.arange(1, P + 1)))
    assert_equal(spec.seasonal_ma_lags,
                 (Q if isinstance(Q, list) else np.arange(1, Q + 1)))

    max_ar_order = p[-1] if isinstance(p, list) else p
    max_ma_order = q[-1] if isinstance(q, list) else q
    max_seasonal_ar_order = P[-1] if isinstance(P, list) else P
    max_seasonal_ma_order = Q[-1] if isinstance(Q, list) else Q
    assert_equal(spec.max_ar_order, max_ar_order)
    assert_equal(spec.max_ma_order, max_ma_order)
    assert_equal(spec.max_seasonal_ar_order, max_seasonal_ar_order)
    assert_equal(spec.max_seasonal_ma_order, max_seasonal_ma_order)
    assert_equal(spec.max_reduced_ar_order,
                 max_ar_order + max_seasonal_ar_order * s)
    assert_equal(spec.max_reduced_ma_order,
                 max_ma_order + max_seasonal_ma_order * s)

    assert_equal(spec.enforce_stationarity, enforce_stationarity)
    assert_equal(spec.enforce_invertibility, enforce_invertibility)
    assert_equal(spec.concentrate_scale, concentrate_scale)


def check_properties(spec, order, seasonal_order, enforce_stationarity,
                     enforce_invertibility, concentrate_scale,
                     is_ar_consecutive, is_ma_consecutive, exog_names,
                     ar_names, ma_names, seasonal_ar_names, seasonal_ma_names):
    p, d, q = order
    P, D, Q, s = seasonal_order

    k_exog_params = len(exog_names)
    k_ar_params = len(p) if isinstance(p, list) else p
    k_ma_params = len(q) if isinstance(q, list) else q
    k_seasonal_ar_params = len(P) if isinstance(P, list) else P
    k_seasonal_ma_params = len(Q) if isinstance(Q, list) else Q
    k_variance_params = int(not concentrate_scale)

    param_names = (exog_names + ar_names + ma_names + seasonal_ar_names +
                   seasonal_ma_names)
    if not concentrate_scale:
        param_names.append('sigma2')

    assert_equal(spec.is_ar_consecutive, is_ar_consecutive)
    assert_equal(spec.is_ma_consecutive, is_ma_consecutive)
    assert_equal(spec.is_integrated, d + D > 0)
    assert_equal(spec.is_seasonal, s > 0)

    assert_equal(spec.k_exog_params, k_exog_params)
    assert_equal(spec.k_ar_params, k_ar_params)
    assert_equal(spec.k_ma_params, k_ma_params)
    assert_equal(spec.k_seasonal_ar_params, k_seasonal_ar_params)
    assert_equal(spec.k_seasonal_ma_params, k_seasonal_ma_params)
    assert_equal(spec.k_params,
                 k_exog_params + k_ar_params + k_ma_params +
                 k_seasonal_ar_params + k_seasonal_ma_params +
                 k_variance_params)

    assert_equal(spec.exog_names, exog_names)
    assert_equal(spec.ar_names, ar_names)
    assert_equal(spec.ma_names, ma_names)
    assert_equal(spec.seasonal_ar_names, seasonal_ar_names)
    assert_equal(spec.seasonal_ma_names, seasonal_ma_names)
    assert_equal(spec.param_names, param_names)


def check_methods(spec, order, seasonal_order, enforce_stationarity,
                  enforce_invertibility, concentrate_scale,
                  exog_params, ar_params, ma_params, seasonal_ar_params,
                  seasonal_ma_params, sigma2):
    params = np.r_[exog_params, ar_params, ma_params, seasonal_ar_params,
                   seasonal_ma_params, sigma2]

    # Test methods
    desired = {
        'exog_params': exog_params,
        'ar_params': ar_params,
        'ma_params': ma_params,
        'seasonal_ar_params': seasonal_ar_params,
        'seasonal_ma_params': seasonal_ma_params}
    if not concentrate_scale:
        desired['sigma2'] = sigma2
    assert_equal(spec.split_params(params), desired)

    assert_equal(spec.join_params(**desired), params)

    assert_equal(spec.validate_params(params), None)

    # Wrong shape
    assert_raises(ValueError, spec.validate_params, [])

    # Wrong dtype
    assert_raises(ValueError, spec.validate_params,
                  ['a'] + params[1:].tolist())

    # NaN / Infinity
    assert_raises(ValueError, spec.validate_params,
                  np.r_[np.inf, params[1:]])
    assert_raises(ValueError, spec.validate_params,
                  np.r_[np.nan, params[1:]])

    # Non-stationary / non-invertible
    if spec.max_ar_order > 0:
        params = np.r_[exog_params, np.ones_like(ar_params), ma_params,
                       np.zeros_like(seasonal_ar_params),
                       seasonal_ma_params, sigma2]
        if enforce_stationarity:
            assert_raises(ValueError, spec.validate_params, params)
        else:
            assert_equal(spec.validate_params(params), None)
    if spec.max_ma_order > 0:
        params = np.r_[exog_params, ar_params, np.ones_like(ma_params),
                       seasonal_ar_params, np.zeros_like(seasonal_ma_params),
                       sigma2]
        if enforce_invertibility:
            assert_raises(ValueError, spec.validate_params, params)
        else:
            assert_equal(spec.validate_params(params), None)
    if spec.max_seasonal_ar_order > 0:
        params = np.r_[exog_params, np.zeros_like(ar_params), ma_params,
                       np.ones_like(seasonal_ar_params), seasonal_ma_params,
                       sigma2]
        if enforce_stationarity:
            assert_raises(ValueError, spec.validate_params, params)
        else:
            assert_equal(spec.validate_params(params), None)
    if spec.max_seasonal_ma_order > 0:
        params = np.r_[exog_params, ar_params, np.zeros_like(ma_params),
                       seasonal_ar_params, np.ones_like(seasonal_ma_params),
                       sigma2]
        if enforce_invertibility:
            assert_raises(ValueError, spec.validate_params, params)
        else:
            assert_equal(spec.validate_params(params), None)

    # Invalid variances
    if not concentrate_scale:
        params = np.r_[exog_params, ar_params, ma_params, seasonal_ar_params,
                       seasonal_ma_params, 0.]
        assert_raises(ValueError, spec.validate_params, params)
        params = np.r_[exog_params, ar_params, ma_params, seasonal_ar_params,
                       seasonal_ma_params, -1]
        assert_raises(ValueError, spec.validate_params, params)

    # Constrain / unconstrain
    unconstrained_ar_params = ar_params
    unconstrained_ma_params = ma_params
    unconstrained_seasonal_ar_params = seasonal_ar_params
    unconstrained_seasonal_ma_params = seasonal_ma_params
    unconstrained_sigma2 = sigma2

    if spec.max_ar_order > 0 and enforce_stationarity:
        unconstrained_ar_params = unconstrain(np.array(ar_params))
    if spec.max_ma_order > 0 and enforce_invertibility:
        unconstrained_ma_params = unconstrain(-np.array(ma_params))
    if spec.max_seasonal_ar_order > 0 and enforce_stationarity:
        unconstrained_seasonal_ar_params = (
            unconstrain(np.array(seasonal_ar_params)))
    if spec.max_seasonal_ma_order > 0 and enforce_invertibility:
        unconstrained_seasonal_ma_params = (
            unconstrain(-np.array(unconstrained_seasonal_ma_params)))
    if not concentrate_scale:
        unconstrained_sigma2 = unconstrained_sigma2**0.5

    unconstrained_params = np.r_[
        exog_params, unconstrained_ar_params, unconstrained_ma_params,
        unconstrained_seasonal_ar_params, unconstrained_seasonal_ma_params,
        unconstrained_sigma2]
    params = np.r_[exog_params, ar_params, ma_params, seasonal_ar_params,
                   seasonal_ma_params, sigma2]

    assert_allclose(spec.unconstrain_params(params), unconstrained_params)

    assert_allclose(spec.constrain_params(unconstrained_params), params)

    assert_allclose(
        spec.constrain_params(spec.unconstrain_params(params)), params)


@pytest.mark.parametrize("n,d,D,s,params,which", [
    # AR models
    (0, 0, 0, 0, np.array([1.]), 'p'),
    (1, 0, 0, 0, np.array([0.5, 1.]), 'p'),
    (1, 0, 0, 0, np.array([-0.2, 100.]), 'p'),
    (2, 0, 0, 0, np.array([-0.2, 0.5, 100.]), 'p'),
    (20, 0, 0, 0, np.array([0] * 20 + [100.]), 'p'),
    # ARI models
    (0, 1, 0, 0, np.array([1.]), 'p'),
    (0, 1, 1, 4, np.array([1.]), 'p'),
    (1, 1, 0, 0, np.array([0.5, 1.]), 'p'),
    (1, 1, 1, 4, np.array([0.5, 1.]), 'p'),
    # MA models
    (0, 0, 0, 0, np.array([1.]), 'q'),
    (1, 0, 0, 0, np.array([0.5, 1.]), 'q'),
    (1, 0, 0, 0, np.array([-0.2, 100.]), 'q'),
    (2, 0, 0, 0, np.array([-0.2, 0.5, 100.]), 'q'),
    (20, 0, 0, 0, np.array([0] * 20 + [100.]), 'q'),
    # IMA models
    (0, 1, 0, 0, np.array([1.]), 'q'),
    (0, 1, 1, 4, np.array([1.]), 'q'),
    (1, 1, 0, 0, np.array([0.5, 1.]), 'q'),
    (1, 1, 1, 4, np.array([0.5, 1.]), 'q'),
])
def test_specification_ar_or_ma(n, d, D, s, params, which):
    if which == 'p':
        p, d, q = n, d, 0
        ar_names = ['ar.L%d' % i for i in range(1, p + 1)]
        ma_names = []
    else:
        p, d, q = 0, d, n
        ar_names = []
        ma_names = ['ma.L%d' % i for i in range(1, q + 1)]
    ar_params = params[:p]
    ma_params = params[p:-1]
    sigma2 = params[-1]
    P, D, Q, s = 0, D, 0, s

    args = ((p, d, q), (P, D, Q, s))
    kwargs = {
        'enforce_stationarity': None,
        'enforce_invertibility': None,
        'concentrate_scale': None
    }

    properties_kwargs = kwargs.copy()
    properties_kwargs.update({
        'is_ar_consecutive': True,
        'is_ma_consecutive': True,
        'exog_names': [],
        'ar_names': ar_names,
        'ma_names': ma_names,
        'seasonal_ar_names': [],
        'seasonal_ma_names': []})

    methods_kwargs = kwargs.copy()
    methods_kwargs.update({
        'exog_params': [],
        'ar_params': ar_params,
        'ma_params': ma_params,
        'seasonal_ar_params': [],
        'seasonal_ma_params': [],
        'sigma2': sigma2})

    # Test the spec created with order, seasonal_order
    spec = specification.SARIMAXSpecification(
            order=(p, d, q), seasonal_order=(P, D, Q, s))

    check_attributes(spec, *args, **kwargs)
    check_properties(spec, *args, **properties_kwargs)
    check_methods(spec, *args, **methods_kwargs)

    # Test the spec created with ar_order, etc.
    spec = specification.SARIMAXSpecification(
            ar_order=p, diff=d, ma_order=q, seasonal_ar_order=P,
            seasonal_diff=D, seasonal_ma_order=Q, seasonal_periods=s)

    check_attributes(spec, *args, **kwargs)
    check_properties(spec, *args, **properties_kwargs)
    check_methods(spec, *args, **methods_kwargs)


@pytest.mark.parametrize(("endog,exog,p,d,q,P,D,Q,s,"
                          "enforce_stationarity,enforce_invertibility,"
                          "concentrate_scale"), [
    (None, None, 0, 0, 0, 0, 0, 0, 0, True, True, False),
    (None, None, 1, 0, 1, 0, 0, 0, 0, True, True, False),
    (None, None, 1, 1, 1, 0, 0, 0, 0, True, True, False),
    (None, None, 1, 0, 0, 0, 0, 0, 4, True, True, False),
    (None, None, 0, 0, 0, 1, 1, 1, 4, True, True, False),
    (None, None, 1, 0, 0, 1, 0, 0, 4, True, True, False),
    (None, None, 1, 0, 0, 1, 1, 1, 4, True, True, False),
    (None, None, 2, 1, 3, 4, 1, 3, 12, True, True, False),

    # Non-consecutive lag orders
    (None, None, [1, 3], 0, 0, 1, 0, 0, 4, True, True, False),
    (None, None, 0, 0, 0, 0, 0, [1, 3], 4, True, True, False),
    (None, None, [2], 0, [1, 3], [1, 3], 0, [1, 4], 4, True, True, False),

    # Modify enforce / concentrate
    (None, None, 2, 1, 3, 4, 1, 3, 12, False, False, True),
    (None, None, 2, 1, 3, 4, 1, 3, 12, True, False, True),
    (None, None, 2, 1, 3, 4, 1, 3, 12, False, True, True),

    # Endog / exog
    (True, None, 2, 1, 3, 4, 1, 3, 12, False, True, True),
    (None, 2, 2, 1, 3, 4, 1, 3, 12, False, True, True),
    (True, 2, 2, 1, 3, 4, 1, 3, 12, False, True, True),
    ('y', None, 2, 1, 3, 4, 1, 3, 12, False, True, True),
    (None, ['x1'], 2, 1, 3, 4, 1, 3, 12, False, True, True),
    ('y', ['x1'], 2, 1, 3, 4, 1, 3, 12, False, True, True),
    ('y', ['x1', 'x2'], 2, 1, 3, 4, 1, 3, 12, False, True, True),
    (True, ['x1', 'x2'], 2, 1, 3, 4, 1, 3, 12, False, True, True),
    ('y', 2, 2, 1, 3, 4, 1, 3, 12, False, True, True),
])
def test_specification(endog, exog, p, d, q, P, D, Q, s,
                       enforce_stationarity, enforce_invertibility,
                       concentrate_scale):
    # Assumptions:
    # - p, q, P, Q are either integers or lists of non-consecutive integers
    #   (i.e. we are not testing boolean lists or consecutive lists here, which
    #   should be tested in the `standardize_lag_order` tests)

    # Construct the specification
    if isinstance(p, list):
        k_ar_params = len(p)
        max_ar_order = p[-1]
    else:
        k_ar_params = max_ar_order = p

    if isinstance(q, list):
        k_ma_params = len(q)