"""
Tests for VAR models (via VARMAX)

These are primarily tests of VAR results from the VARMAX state space approach
compared to the output of estimation via CSS

loglikelihood_burn = k_ar is required, since the loglikelihood reported by the
CSS approach is the conditional loglikelihood.

Author: Chad Fulton
License: Simplified-BSD
"""
import os

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

from statsmodels.tsa.statespace import varmax

from .results import results_var_R

current_path = os.path.dirname(os.path.abspath(__file__))
results_var_R_output = pd.read_csv(
    os.path.join(current_path, "results", "results_var_R_output.csv"))

up2 = os.path.split(os.path.split(current_path)[0])[0]
dta = pd.read_stata(os.path.join(up2, "tests", "results", "lutkepohl2.dta"))
dta.index = pd.PeriodIndex(dta.qtr, freq="Q")
endog = dta[["dln_inv", "dln_inc", "dln_consump"]].loc["1960Q2":"1978"]

# Note that we start the time-trend at 0 here, to match
time_trend0 = np.arange(len(endog))
exog0 = np.c_[np.ones(len(endog)), time_trend0, time_trend0**2]
time_trend0_fcast = np.arange(len(endog), len(endog) + 10)
exog0_fcast = np.c_[np.ones(10), time_trend0_fcast, time_trend0_fcast**2]
# And here we start the time-trend at 1
time_trend1 = np.arange(1, len(endog) + 1)
time_trend1_fcast = np.arange(len(endog) + 1, len(endog) + 1 + 10)
exog1 = np.c_[np.ones(len(endog)), time_trend1, time_trend1**2]
exog1_fcast = np.c_[np.ones(10), time_trend1_fcast, time_trend1_fcast**2]


def check_irf(test, mod, results, params=None):
    # Note that 'vars' uses an estimator of Sigma_u with k_params fewer degrees
    # of freedom to compute impulse responses, so we need to reset the
    # state_cov matrix for this part of the test
    Sigma_u_mle = mod["state_cov"]
    nobs_effective = mod.nobs - mod.k_ar
    df_resid = (nobs_effective -
                (mod.k_ar * mod.k_endog + mod.k_trend + mod.k_exog))
    Sigma_u = Sigma_u_mle * nobs_effective / df_resid

    L = np.linalg.cholesky(Sigma_u)
    if params is None:
        params = np.copy(results["params"])
    params[-6:] = L[np.tril_indices_from(L)]
    res = mod.smooth(params)

    for i in range(3):
        impulse_to = endog.columns[i]

        # Non-orthogonalized
        columns = [f"{test}.irf.{impulse_to}.{name}"
                   for name in endog.columns]
        assert_allclose(res.impulse_responses(10, i),
                        results_var_R_output[columns])

        # Orthogonalized
        columns = [f"{test}.irf.ortho.{impulse_to}.{name}"
                   for name in endog.columns]
        assert_allclose(res.impulse_responses(10, i, orthogonalized=True),
                        results_var_R_output[columns])

        # Orthogonalized, cumulated
        columns = [f"{test}.irf.cumu.{impulse_to}.{name}"
                   for name in endog.columns]
        result = res.impulse_responses(10, i,
                                       orthogonalized=True, cumulative=True)
        assert_allclose(result,
                        results_var_R_output[columns])


def test_var_basic():
    test = "basic"

    # VAR(2), no trend or exog
    results = results_var_R.res_basic
    mod = varmax.VARMAX(endog, order=(2, 0), trend="n", loglikelihood_burn=2)
    res = mod.smooth(results["params"])

    assert_allclose(res.llf, results["llf"])

    # Forecast
    columns = [f"{test}.fcast.{name}.fcst" for name in endog.columns]
    assert_allclose(res.forecast(10), results_var_R_output[columns].iloc[:10])

    # IRF
    check_irf(test, mod, results)

    # FEVD: TODO


def test_var_c():
    test = "c"

    # VAR(2), constant trend, no exog
    results = results_var_R.res_c
    mod = varmax.VARMAX(endog, order=(2, 0), trend="c", loglikelihood_burn=2)
    res = mod.smooth(results["params"])

    assert_allclose(res.llf, results["llf"])

    # Forecast
    columns = [f"{test}.fcast.{name}.fcst" for name in endog.columns]
    assert_allclose(res.forecast(10), results_var_R_output[columns].iloc[:10])

    # IRF
    check_irf(test, mod, results)

    # FEVD: TODO


def test_var_ct():
    test = "ct"

    # VAR(2), constant and time trend, no exog
    results = results_var_R.res_ct
    mod = varmax.VARMAX(endog, order=(2, 0), trend="ct", loglikelihood_burn=2)
    res = mod.smooth(results["params"])

    assert_allclose(res.llf, results["llf"])

    # Forecast
    columns = [f"{test}.fcast.{name}.fcst" for name in endog.columns]
    assert_allclose(res.forecast(10), results_var_R_output[columns].iloc[:10])

    # IRF
    check_irf(test, mod, results)

    # FEVD: TODO


def test_var_ct_as_exog0():
    test = "ct_as_exog0"

    # VAR(2), no built-in trend, constant and time trend as exog
    # Here we start the time-trend at 0
    results = results_var_R.res_ct_as_exog0
    mod = varmax.VARMAX(endog, order=(2, 0), exog=exog0[:, :2], trend="n",
                        loglikelihood_burn=2)
    res = mod.smooth(results["params"])

    assert_allclose(res.llf, results["llf"])

    # Forecast
    columns = [f"{test}.fcast.{name}.fcst" for name in endog.columns]
    assert_allclose(res.forecast(10, exog=exog0_fcast[:, :2]),
                    results_var_R_output[columns].iloc[:10])

    # IRF
    check_irf(test, mod, results)

    # FEVD: TODO


def test_var_ct_as_exog1():
    test = "ct"

    # VAR(2), no built-in trend, constant and time trend as exog
    # Here we start the time-trend at 1 and so we can compare to the built-in
    # trend results "res_ct"
    results = results_var_R.res_ct
    mod = varmax.VARMAX(endog, order=(2, 0), exog=exog1[:, :2], trend="n",
                        loglikelihood_burn=2)
    # Since the params were given for the built-in trend case, we need to
    # re-order them
    params = results["params"]
    params = np.r_[params[6:-6], params[:6], params[-6:]]
    res = mod.smooth(params)

    assert_allclose(res.llf, results["llf"])

    # Forecast
    columns = [f"{test}.fcast.{name}.fcst" for name in endog.columns]
    assert_allclose(res.forecast(10, exog=exog1_fcast[:, :2]),
                    results_var_R_output[columns].iloc[:10])

    # IRF
    check_irf(test, mod, results, params)

    # FEVD: TODO


def test_var_ctt():
    test = "ctt_as_exog1"

    # VAR(2), constant, trend, and trend**2, no exog
    # Note that this is comparing against trend as exog in the R package,
    # since it does not have a built-in option for trend**2
    results = results_var_R.res_ctt_as_exog1
    mod = varmax.VARMAX(endog, order=(2, 0), trend="ctt",
                        loglikelihood_burn=2)
    params = results["params"]
    params = np.r_[params[-(6+9):-6], params[:-(6+9)], params[-6:]]
    res = mod.smooth(params)

    assert_allclose(res.llf, results["llf"])

    # Forecast
    columns = [f"{test}.fcast.{name}.fcst" for name in endog.columns]
    assert_allclose(res.forecast(10), results_var_R_output[columns].iloc[:10])

    # IRF
    check_irf(test, mod, results, params)

    # FEVD: TODO


def test_var_ct_exog():
    test = "ct_exog"

    # VAR(2), constant and trend, 'inc' variable as exog
    results = results_var_R.res_ct_exog
    exog = dta["inc"].loc["1960Q2":"1978"]
    exog_fcast = dta[["inc"]].loc["1979Q1":"1981Q2"]
    mod = varmax.VARMAX(endog, order=(2, 0), exog=exog, trend="ct",
                        loglikelihood_burn=2)
    res = mod.smooth(results["params"])

    assert_allclose(res.llf, results["llf"])

    # Forecast
    columns = [f"{test}.fcast.{name}.fcst" for name in endog.columns]
    assert_allclose(res.forecast(10, exog=exog_fcast),
                    results_var_R_output[columns].iloc[:10])

    # IRF
    check_irf(test, mod, results)

    # FEVD: TODO


def test_var_c_2exog():
    test = "c_2exog"

    # VAR(2), constant, 'inc', 'inv' variables as exog
    results = results_var_R.res_c_2exog
    exog = dta[["inc", "inv"]].loc["1960Q2":"1978"]
    exog_fcast = dta[["inc", "inv"]].loc["1979Q1":"1981Q2"]
    mod = varmax.VARMAX(endog, order=(2, 0), exog=exog, trend="c",
                        loglikelihood_burn=2)
    res = mod.smooth(results["params"])

    assert_allclose(res.llf, results["llf"])

    # Forecast
    columns = [f"{test}.fcast.{name}.fcst" for name in endog.columns]
    assert_allclose(res.forecast(10, exog=exog_fcast),
                    results_var_R_output[columns].iloc[:10])

    # IRF
    check_irf(test, mod, results)

    # FEVD: TODO
