from statsmodels.compat.pandas import MONTH_END

import os
import pickle

import numpy as np
from numpy.testing import assert_allclose
import pandas as pd
import pytest

from statsmodels.datasets import co2
from statsmodels.tsa.seasonal import STL, DecomposeResult

cur_dir = os.path.dirname(os.path.abspath(__file__))
file_path = os.path.join(cur_dir, "results", "stl_test_results.csv")
results = pd.read_csv(file_path)
results.columns = [c.strip() for c in results.columns]
results.scenario = results.scenario.apply(str.strip)
results = results.set_index(["scenario", "idx"])


@pytest.fixture(scope="module", params=[True, False])
def robust(request):
    return request.param


def default_kwargs_base():
    file_path = os.path.join(cur_dir, "results", "stl_co2.csv")
    co2 = np.asarray(pd.read_csv(file_path, header=None).iloc[:, 0])
    y = co2
    nobs = y.shape[0]
    nperiod = 12
    work = np.zeros((nobs + 2 * nperiod, 7))
    rw = np.ones(nobs)
    trend = np.zeros(nobs)
    season = np.zeros(nobs)
    return dict(
        y=y,
        n=y.shape[0],
        np=nperiod,
        ns=35,
        nt=19,
        nl=13,
        no=2,
        ni=1,
        nsjump=4,
        ntjump=2,
        nljump=2,
        isdeg=1,
        itdeg=1,
        ildeg=1,
        rw=rw,
        trend=trend,
        season=season,
        work=work,
    )


@pytest.fixture(scope="function")
def default_kwargs():
    return default_kwargs_base()


@pytest.fixture(scope="function")
def default_kwargs_short():
    kwargs = default_kwargs_base()
    y = kwargs["y"][:-1]
    nobs = y.shape[0]
    work = np.zeros((nobs + 2 * kwargs["np"], 7))
    rw = np.ones(nobs)
    trend = np.zeros(nobs)
    season = np.zeros(nobs)
    kwargs.update(
        dict(y=y, n=nobs, rw=rw, trend=trend, season=season, work=work)
    )
    return kwargs


def _to_class_kwargs(kwargs, robust=False):
    endog = kwargs["y"]
    np = kwargs["np"]
    ns = kwargs["ns"]
    nt = kwargs["nt"]
    nl = kwargs["nl"]
    isdeg = kwargs["isdeg"]
    itdeg = kwargs["itdeg"]
    ildeg = kwargs["ildeg"]
    nsjump = kwargs["nsjump"]
    ntjump = kwargs["ntjump"]
    nljump = kwargs["nljump"]
    outer_iter = kwargs["no"]
    inner_iter = kwargs["ni"]
    class_kwargs = dict(
        endog=endog,
        period=np,
        seasonal=ns,
        trend=nt,
        low_pass=nl,
        seasonal_deg=isdeg,
        trend_deg=itdeg,
        low_pass_deg=ildeg,
        robust=robust,
        seasonal_jump=nsjump,
        trend_jump=ntjump,
        low_pass_jump=nljump,
    )
    return class_kwargs, outer_iter, inner_iter


def test_baseline_class(default_kwargs):
    class_kwargs, outer, inner = _to_class_kwargs(default_kwargs)
    mod = STL(**class_kwargs)
    res = mod.fit(outer_iter=outer, inner_iter=inner)

    expected = results.loc["baseline"].sort_index()
    assert_allclose(res.trend, expected.trend)
    assert_allclose(res.seasonal, expected.season)
    assert_allclose(res.weights, expected.rw)
    resid = class_kwargs["endog"] - expected.trend - expected.season
    assert_allclose(res.resid, resid)


def test_short_class(default_kwargs_short):
    class_kwargs, outer, inner = _to_class_kwargs(default_kwargs_short)
    mod = STL(**class_kwargs)
    res = mod.fit(outer_iter=outer, inner_iter=inner)

    expected = results.loc["short"].sort_index()
    assert_allclose(res.seasonal, expected.season)
    assert_allclose(res.trend, expected.trend)
    assert_allclose(res.weights, expected.rw)


def test_nljump_1_class(default_kwargs):
    default_kwargs["nljump"] = 1
    class_kwargs, outer, inner = _to_class_kwargs(default_kwargs)
    mod = STL(**class_kwargs)
    res = mod.fit(outer_iter=outer, inner_iter=inner)

    expected = results.loc["nljump-1"].sort_index()
    assert_allclose(res.seasonal, expected.season)
    assert_allclose(res.trend, expected.trend)
    assert_allclose(res.weights, expected.rw)


def test_ntjump_1_class(default_kwargs):
    default_kwargs["ntjump"] = 1
    class_kwargs, outer, inner = _to_class_kwargs(default_kwargs)
    mod = STL(**class_kwargs)
    res = mod.fit(outer_iter=outer, inner_iter=inner)

    expected = results.loc["ntjump-1"].sort_index()
    assert_allclose(res.seasonal, expected.season)
    assert_allclose(res.trend, expected.trend)
    assert_allclose(res.weights, expected.rw)


def test_nljump_1_ntjump_1_class(default_kwargs):
    default_kwargs["nljump"] = 1
    default_kwargs["ntjump"] = 1
    class_kwargs, outer, inner = _to_class_kwargs(default_kwargs)
    mod = STL(**class_kwargs)
    res = mod.fit(outer_iter=outer, inner_iter=inner)

    expected = results.loc["nljump-1-ntjump-1"].sort_index()
    assert_allclose(res.seasonal, expected.season)
    assert_allclose(res.trend, expected.trend)
    assert_allclose(res.weights, expected.rw)


def test_parameter_checks_period(default_kwargs):
    class_kwargs, _, _ = _to_class_kwargs(default_kwargs)
    endog = class_kwargs["endog"]
    endog2 = np.hstack((endog[:, None], endog[:, None]))
    period = class_kwargs["period"]
    with pytest.raises(ValueError, match="endog is required to have ndim 1"):
        STL(endog=endog2, period=period)
    match = "period must be a positive integer >= 2"
    with pytest.raises(ValueError, match=match):
        STL(endog=endog, period=1)
    with pytest.raises(ValueError, match=match):
        STL(endog=endog, period=-12)
    with pytest.raises(ValueError, match=match):
        STL(endog=endog, period=4.0)


def test_parameter_checks_seasonal(default_kwargs):
    class_kwargs, _, _ = _to_class_kwargs(default_kwargs)
    endog = class_kwargs["endog"]
    period = class_kwargs["period"]
    match = "seasonal must be an odd positive integer >= 3"
    with pytest.raises(ValueError, match=match):
        STL(endog=endog, period=period, seasonal=2)
    with pytest.raises(ValueError, match=match):
        STL(endog=endog, period=period, seasonal=-7)
    with pytest.raises(ValueError, match=match):
        STL(endog=endog, period=period, seasonal=13.0)


def test_parameter_checks_trend(default_kwargs):
    class_kwargs, _, _ = _to_class_kwargs(default_kwargs)
    endog = class_kwargs["endog"]
    period = class_kwargs["period"]
    match = "trend must be an odd positive integer >= 3 where trend > period"
    with pytest.raises(ValueError, match=match):
        STL(endog=endog, period=period, trend=14)
    with pytest.raises(ValueError, match=match):
        STL(endog=endog, period=period, trend=11)
    with pytest.raises(ValueError, match=match):
        STL(endog=endog, period=period, trend=-19)
    with pytest.raises(ValueError, match=match):
        STL(endog=endog, period=period, trend=19.0)


def test_parameter_checks_low_pass(default_kwargs):
    class_kwargs, _, _ = _to_class_kwargs(default_kwargs)
    endog = class_kwargs["endog"]
    period = class_kwargs["period"]

    match = (
        "low_pass must be an odd positive integer >= 3 where"
        " low_pass > period"
    )
    with pytest.raises(ValueError, match=match):
        STL(endog=endog, period=period, low_pass=14)
    with pytest.raises(ValueError, match=match):
        STL(endog=endog, period=period, low_pass=7)
    with pytest.raises(ValueError, match=match):
        STL(endog=endog, period=period, low_pass=-19)
    with pytest.raises(ValueError, match=match):
        STL(endog=endog, period=period, low_pass=19.0)


def test_jump_errors(default_kwargs):
    class_kwargs, _, _ = _to_class_kwargs(default_kwargs)
    endog = class_kwargs["endog"]
    period = class_kwargs["period"]
    with pytest.raises(ValueError, match="low_pass_jump must be a positive"):
        STL(endog=endog, period=period, low_pass_jump=0)
    with pytest.raises(ValueError, match="low_pass_jump must be a positive"):
        STL(endog=endog, period=period, low_pass_jump=1.0)
    with pytest.raises(ValueError, match="seasonal_jump must be a positive"):
        STL(endog=endog, period=period, seasonal_jump=0)
    with pytest.raises(ValueError, match="seasonal_jump must be a positive"):
        STL(endog=endog, period=period, seasonal_jump=1.0)
    with pytest.raises(ValueError, match="trend_jump must be a positive"):
        STL(endog=endog, period=period, trend_jump=0)
    with pytest.raises(ValueError, match="trend_jump must be a positive"):
        STL(endog=endog, period=period, trend_jump=1.0)


def test_defaults_smoke(default_kwargs, robust):
    class_kwargs, _, _ = _to_class_kwargs(default_kwargs, robust)
    endog = class_kwargs["endog"]
    period = class_kwargs["period"]
    mod = STL(endog=endog, period=period)
    mod.fit()


def test_pandas(default_kwargs, robust):
    class_kwargs, _, _ = _to_class_kwargs(default_kwargs, robust)
    endog = pd.Series(class_kwargs["endog"], name="y")
    period = class_kwargs["period"]
    mod = STL(endog=endog, period=period)
    res = mod.fit()
    assert isinstance(res.trend, pd.Series)
    assert isinstance(res.seasonal, pd.Series)
    assert isinstance(res.resid, pd.Series)
    assert isinstance(res.weights, pd.Series)


def test_period_detection(default_kwargs):
    class_kwargs, _, _ = _to_class_kwargs(default_kwargs)
    mod = STL(**class_kwargs)
    res = mod.fit()

    del class_kwargs["period"]
    endog = class_kwargs["endog"]
    index = pd.date_range("1-1-1959", periods=348, freq=MONTH_END)
    class_kwargs["endog"] = pd.Series(endog, index=index)
    mod = STL(**class_kwargs)

    res_implicit_period = mod.fit()
    assert_allclose(res.seasonal, res_implicit_period.seasonal)


def test_no_period(default_kwargs):
    class_kwargs, _, _ = _to_class_kwargs(default_kwargs)
    del class_kwargs["period"]
    class_kwargs["endog"] = pd.Series(class_kwargs["endog"])
    with pytest.raises(ValueError, match="Unable to determine period from"):
        STL(**class_kwargs)


@pytest.mark.matplotlib
def test_plot(default_kwargs, close_figures):
    class_kwargs, outer, inner = _to_class_kwargs(default_kwargs)
    res = STL(**class_kwargs).fit(outer_iter=outer, inner_iter=inner)
    res.plot()

    class_kwargs["endog"] = pd.Series(class_kwargs["endog"], name="CO2")
    res = STL(**class_kwargs).fit()
    res.plot()


def test_default_trend(default_kwargs):
    # GH 6686
    class_kwargs, _, _ = _to_class_kwargs(default_kwargs)
    class_kwargs["seasonal"] = 17
    class_kwargs["trend"] = None
    mod = STL(**class_kwargs)
    period = class_kwargs["period"]
    seasonal = class_kwargs["seasonal"]
    expected = int(np.ceil(1.5 * period / (1 - 1.5 / seasonal)))
    expected += 1 if expected % 2 == 0 else 0
    assert mod.config["trend"] == expected

    class_kwargs["seasonal"] = 7
    mod = STL(**class_kwargs)
    period = class_kwargs["period"]
    seasonal = class_kwargs["seasonal"]
    expected = int(np.ceil(1.5 * period / (1 - 1.5 / seasonal)))
    expected += 1 if expected % 2 == 0 else 0
    assert mod.config["trend"] == expected


def test_pickle(default_kwargs):
    class_kwargs, outer, inner = _to_class_kwargs(default_kwargs)
    mod = STL(**class_kwargs)
    res = mod.fit()
    pkl = pickle.dumps(mod)
    reloaded = pickle.loads(pkl)
    res2 = reloaded.fit()
    assert_allclose(res.trend, res2.trend)
    assert_allclose(res.seasonal, res2.seasonal)
    assert mod.config == reloaded.config


def test_squezable_to_1d():
    data = co2.load().data
    data = data.resample(MONTH_END).mean().ffill()
    res = STL(data).fit()
    assert isinstance(res, DecomposeResult)