"""
Test  for ordinal models
"""

import warnings

import numpy as np
from numpy.testing import assert_allclose, assert_equal
import pandas as pd
import pytest
import scipy.stats as stats

from statsmodels.discrete.discrete_model import Logit
from statsmodels.miscmodels.ordinal_model import OrderedModel
from statsmodels.tools.sm_exceptions import HessianInversionWarning
from statsmodels.tools.tools import add_constant

from .results.results_ordinal_model import data_store as ds


class CheckOrdinalModelMixin:

    def test_basic(self):
        # checks basic results againt R MASS package
        n_cat = ds.n_ordinal_cat
        res1 = self.res1
        res2 = self.res2
        # coefficients values, standard errors, t & p values
        assert_allclose(res1.params[: -n_cat + 1], res2.coefficients_val, atol=2e-4)
        assert_allclose(
            res1.bse[: -n_cat + 1], res2.coefficients_stdE, rtol=0.003, atol=1e-5
        )
        assert_allclose(
            res1.tvalues[: -n_cat + 1], res2.coefficients_tval, rtol=0.003, atol=7e-4
        )
        assert_allclose(
            res1.pvalues[: -n_cat + 1], res2.coefficients_pval, rtol=0.009, atol=1e-5
        )
        # thresholds are given with exponentiated increments
        # from the first threshold
        assert_allclose(
            res1.model.transform_threshold_params(res1.params)[1:-1],
            res2.thresholds,
            atol=4e-4,
        )

        # probabilities
        assert_allclose(res1.predict()[:7, :], res2.prob_pred, atol=5e-5)

    def test_pandas(self):
        # makes sure that the Pandas ecosystem is supported
        res1 = self.res1
        resp = self.resp
        # converges slightly differently why?
        assert_allclose(res1.params, resp.params, atol=1e-10)
        assert_allclose(res1.bse, resp.bse, atol=1e-10)

        assert_allclose(res1.model.endog, resp.model.endog, rtol=1e-10)
        assert_allclose(res1.model.exog, resp.model.exog, rtol=1e-10)

    def test_formula(self):
        # makes sure the "R-way" of writing models is supported
        res1 = self.res1
        resf = self.resf
        # converges slightly differently why? yet e-5 is ok
        assert_allclose(res1.params, resf.params, atol=5e-5)
        assert_allclose(res1.bse, resf.bse, atol=5e-5)

        assert_allclose(res1.model.endog, resf.model.endog, rtol=1e-10)
        assert_allclose(res1.model.exog, resf.model.exog, rtol=1e-10)

    def test_unordered(self):
        # makes sure that ordered = True is optional for the endog Serie
        # et categories have to be set in the right order
        res1 = self.res1
        resf = self.resu
        # converges slightly differently why?
        assert_allclose(res1.params, resf.params, atol=1e-10)
        assert_allclose(res1.bse, resf.bse, atol=1e-10)

        assert_allclose(res1.model.endog, resf.model.endog, rtol=1e-10)
        assert_allclose(res1.model.exog, resf.model.exog, rtol=1e-10)

    def test_results_other(self):

        res1 = self.res1  # numpy
        resp = self.resp  # pandas

        param_names_np = ["x1", "x2", "x3", "0/1", "1/2"]
        param_names_pd = [
            "pared",
            "public",
            "gpa",
            "unlikely/somewhat likely",
            "somewhat likely/very likely",
        ]

        assert res1.model.data.param_names == param_names_np
        assert self.resp.model.data.param_names == param_names_pd
        assert self.resp.model.endog_names == "apply"

        # results
        if hasattr(self, "pred_table"):
            table = res1.pred_table()
            assert_equal(table.values, self.pred_table)

        # smoke test
        res1.summary()

        # inherited
        tt = res1.t_test(np.eye(len(res1.params)))
        assert_allclose(tt.pvalue, res1.pvalues, rtol=1e-13)

        tt = resp.t_test(["pared", "public", "gpa"])  # pandas names
        assert_allclose(tt.pvalue, res1.pvalues[:3], rtol=1e-13)

        pred = res1.predict(exog=res1.model.exog[-5:])
        fitted = res1.predict()
        assert_allclose(pred, fitted[-5:], rtol=1e-13)

        pred = resp.predict(exog=resp.model.data.orig_exog.iloc[-5:])
        fitted = resp.predict()
        assert_allclose(pred, fitted[-5:], rtol=1e-13)

        dataf = self.resf.model.data.frame  # is a dict
        dataf_df = pd.DataFrame.from_dict(dataf)
        pred = self.resf.predict(exog=dataf_df.iloc[-5:])
        fitted = self.resf.predict()
        assert_allclose(pred, fitted[-5:], rtol=1e-13)

        n, k = res1.model.exog.shape
        assert_equal(self.resf.df_resid, n - (k + 2))

        # check wrapper
        assert resp.params.index.tolist() == resp.model.exog_names
        assert resp.bse.index.tolist() == resp.model.exog_names


class TestLogitModel(CheckOrdinalModelMixin):

    @classmethod
    def setup_class(cls):
        data = ds.df
        data_unordered = ds.df_unordered

        # standard fit
        mod = OrderedModel(
            data["apply"].values.codes,
            np.asarray(data[["pared", "public", "gpa"]], float),
            distr="logit",
        )
        res = mod.fit(method="bfgs", disp=False)
        # standard fit with pandas input
        modp = OrderedModel(
            data["apply"], data[["pared", "public", "gpa"]], distr="logit"
        )
        resp = modp.fit(method="bfgs", disp=False)
        # fit with formula
        modf = OrderedModel.from_formula(
            "apply ~ pared + public + gpa - 1",
            data={
                "apply": data["apply"].values.codes,
                "pared": data["pared"],
                "public": data["public"],
                "gpa": data["gpa"],
            },
            distr="logit",
        )
        resf = modf.fit(method="bfgs", disp=False)
        # fit on data with ordered=False
        modu = OrderedModel(
            data_unordered["apply"].values.codes,
            np.asarray(data_unordered[["pared", "public", "gpa"]], float),
            distr="logit",
        )
        resu = modu.fit(method="bfgs", disp=False)

        from .results.results_ordinal_model import res_ord_logit as res2

        cls.res2 = res2
        cls.res1 = res
        cls.resp = resp
        cls.resf = resf
        cls.resu = resu

    def test_postestimation(self):
        res1 = self.res1
        res2 = self.res2
        resid_prob = res1.resid_prob

        assert_allclose(resid_prob[: len(res2.resid_prob)], res2.resid_prob, atol=1e-4)
        stats_prob = [
            resid_prob.mean(),
            resid_prob.min(),
            resid_prob.max(),
            resid_prob.var(ddof=1),
        ]
        assert_allclose(stats_prob, res2.resid_prob_stats, atol=1e-5)

        # from R generalhoslem
        # > logitgof(ologit_ucla$apply2, fitted(r_logit), g = 10, ord = TRUE)
        chi2 = 20.958760713111
        df = 17
        p_value = 0.2281403796588
        # values in Stata using ologitgof are a bit different,
        # I guess different sort algorithm and because of ties, see #7095

        import statsmodels.stats.diagnostic_gen as dia

        # TODO: add more properties or methods to Results class
        fitted = res1.predict()
        y_dummy = (res1.model.endog[:, None] == np.arange(3)).astype(int)
        sv = (fitted * np.arange(1, 3 + 1)).sum(1)
        dt = dia.test_chisquare_binning(
            y_dummy,
            fitted,
            sort_var=sv,
            bins=10,
            df=None,
            ordered=True,
            sort_method="stable",
        )
        assert_allclose(dt.statistic, chi2, rtol=5e-5)
        assert_allclose(dt.pvalue, p_value, rtol=1e-4)
        assert_equal(dt.df, df)


class TestProbitModel(CheckOrdinalModelMixin):

    @classmethod
    def setup_class(cls):
        data = ds.df
        data_unordered = ds.df_unordered

        mod = OrderedModel(
            data["apply"].values.codes,
            np.asarray(data[["pared", "public", "gpa"]], float),
            distr="probit",
        )
        res = mod.fit(method="bfgs", disp=False)

        modp = OrderedModel(
            data["apply"], data[["pared", "public", "gpa"]], distr="probit"
        )
        resp = modp.fit(method="bfgs", disp=False)

        modf = OrderedModel.from_formula(
            "apply ~ pared + public + gpa - 1",
            data={
                "apply": data["apply"].values.codes,
                "pared": data["pared"],
                "public": data["public"],
                "gpa": data["gpa"],
            },
            distr="probit",
        )
        resf = modf.fit(method="bfgs", disp=False)

        modu = OrderedModel(
            data_unordered["apply"].values.codes,
            np.asarray(data_unordered[["pared", "public", "gpa"]], float),
            distr="probit",
        )
        resu = modu.fit(method="bfgs", disp=False)

        from .results.results_ordinal_model import res_ord_probit as res2

        cls.res2 = res2
        cls.res1 = res
        cls.resp = resp
        cls.resf = resf
        cls.resu = resu

        # regression numbers
        cls.pred_table = np.array(
            [
                [202, 18, 0, 220],
                [112, 28, 0, 140],
                [27, 13, 0, 40],  # noqa
                [341, 59, 0, 400],
            ],
            dtype=np.int64,
        )

    def test_loglikerelated(self):

        res1 = self.res1
        # res2 = self.res2

        mod = res1.model
        fact = 1.1  # evaluate away from optimum
        score1 = mod.score(res1.params * fact)
        score_obs_numdiff = mod.score_obs(res1.params * fact)
        score_obs_exog = mod.score_obs_(res1.params * fact)
        # Relax atol due to small failures on OSX
        assert_allclose(score_obs_numdiff.sum(0), score1, atol=1e-6)
        assert_allclose(score_obs_exog.sum(0), score1[: mod.k_vars], atol=1e-6)

        # null model
        mod_null = OrderedModel(
            mod.endog, None, offset=np.zeros(mod.nobs), distr=mod.distr
        )
        null_params = mod.start_params
        res_null = mod_null.fit(method="bfgs", disp=False)
        assert_allclose(res_null.params, null_params[mod.k_vars :], rtol=1e-8)
        assert_allclose(res1.llnull, res_null.llf, rtol=1e-8)

    def test_formula_categorical(self):

        resp = self.resp
        data = ds.df

        formula = "apply ~ pared + public + gpa - 1"
        modf2 = OrderedModel.from_formula(formula, data, distr="probit")
        resf2 = modf2.fit(method="bfgs", disp=False)
        assert_allclose(resf2.params, resp.params, atol=1e-8)
        assert modf2.exog_names == resp.model.exog_names
        assert modf2.data.ynames == resp.model.data.ynames
        assert hasattr(modf2.data, "frame")
        assert not hasattr(modf2, "frame")

        msg = "Only ordered pandas Categorical"
        with pytest.raises(ValueError, match=msg):
            # only ordered categorical or numerical endog are allowed
            # string endog raises ValueError
            OrderedModel.from_formula(
                "apply ~ pared + public + gpa - 1",
                data={
                    "apply": np.asarray(data["apply"]),
                    "pared": data["pared"],
                    "public": data["public"],
                    "gpa": data["gpa"],
                },
                distr="probit",
            )

    def test_offset(self):

        resp = self.resp
        data = ds.df
        offset = np.ones(len(data))

        formula = "apply ~ pared + public + gpa - 1"
        modf2 = OrderedModel.from_formula(formula, data, offset=offset, distr="probit")
        resf2 = modf2.fit(method="bfgs", disp=False)

        resf2_params = np.asarray(resf2.params)
        resp_params = np.asarray(resp.params)
        assert_allclose(resf2_params[:3], resp_params[:3], atol=2e-4)
        assert_allclose(resf2_params[3], resp_params[3] + 1, atol=2e-4)

        fitted = resp.predict()
        fitted2 = resf2.predict()
        assert_allclose(fitted2, fitted, atol=2e-4)

        pred_ones = resf2.predict(data[:6], offset=np.ones(6))
        assert_allclose(pred_ones, fitted[:6], atol=2e-4)

        # check default is 0. if exog provided
        pred_zero1 = resf2.predict(data[:6])
        pred_zero2 = resf2.predict(data[:6], offset=0)
        assert_allclose(pred_zero1, pred_zero2, atol=2e-4)

        # compare with equivalent results frp, no-offset model
        pred_zero = resp.predict(
            data[["pared", "public", "gpa"]].iloc[:6], offset=-np.ones(6)
        )
        assert_allclose(pred_zero1, pred_zero, atol=2e-4)

        params_adj = resp.params.copy()
        params_adj.iloc[3] += 1
        fitted_zero = resp.model.predict(params_adj)
        assert_allclose(pred_zero1, fitted_zero[:6], atol=2e-4)


class TestLogitModelFormula:

    @classmethod
    def setup_class(cls):
        data = ds.df
        nobs = len(data)
        data["dummy"] = (np.arange(nobs) < (nobs / 2)).astype(float)
        # alias to correspond to patsy name
        data["C(dummy)[T.1.0]"] = data["dummy"]
        cls.data = data

        columns = ["C(dummy)[T.1.0]", "pared", "public", "gpa"]
        # standard fit
        mod = OrderedModel(
            data["apply"].values.codes, np.asarray(data[columns], float), distr="logit"
        )
        cls.res = mod.fit(method="bfgs", disp=False)
        # standard fit with pandas input
        modp = OrderedModel(data["apply"], data[columns], distr="logit")
        cls.resp = modp.fit(method="bfgs", disp=False)

    def test_setup(self):
        data = self.data
        resp = self.resp
        fittedvalues = resp.predict()

        formulas = [
            "apply ~ 1 + pared + public + gpa + C(dummy)",
            "apply ~ pared + public + gpa + C(dummy)",
        ]
        for formula in formulas:
            modf1 = OrderedModel.from_formula(formula, data, distr="logit")
            resf1 = modf1.fit(method="bfgs")
            summf1 = resf1.summary()
            summf1_str = str(summf1)
            assert resf1.model.exog_names == resp.model.exog_names
            assert resf1.model.data.param_names == resp.model.exog_names
            assert all(name in summf1_str for name in resp.model.data.param_names)
            assert_allclose(resf1.predict(data[:5]), fittedvalues[:5])

        # test over parameterized model with implicit constant
        formula = "apply ~ 0 + pared + public + gpa + C(dummy)"

        with pytest.raises(ValueError, match="not be a constant"):
            OrderedModel.from_formula(formula, data, distr="logit")

        # ignore constant, so we get results without exception
        modf2 = OrderedModel.from_formula(formula, data, distr="logit", hasconst=False)
        # we get a warning in some environments
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", HessianInversionWarning)
            resf2 = modf2.fit(method="bfgs")

        assert_allclose(resf2.predict(data[:5]), fittedvalues[:5], rtol=1e-4)


class TestCLogLogModel(CheckOrdinalModelMixin):

    @classmethod
    def setup_class(cls):
        data = ds.df
        data_unordered = ds.df_unordered

        # a Scipy distribution defined minimally
        class CLogLog(stats.rv_continuous):
            def _ppf(self, q):
                return np.log(-np.log(1 - q))

            def _cdf(self, x):
                return 1 - np.exp(-np.exp(x))

        cloglog = CLogLog()

        mod = OrderedModel(
            data["apply"].values.codes,
            np.asarray(data[["pared", "public", "gpa"]], float),
            distr=cloglog,
        )
        res = mod.fit(method="bfgs", disp=False)

        modp = OrderedModel(
            data["apply"], data[["pared", "public", "gpa"]], distr=cloglog
        )
        resp = modp.fit(method="bfgs", disp=False)

        # with pytest.warns(UserWarning):
        modf = OrderedModel.from_formula(
            "apply ~ pared + public + gpa - 1",
            data={
                "apply": data["apply"].values.codes,
                "pared": data["pared"],
                "public": data["public"],
                "gpa": data["gpa"],
            },
            distr=cloglog,
        )
        resf = modf.fit(method="bfgs", disp=False)

        modu = OrderedModel(
            data_unordered["apply"].values.codes,
            np.asarray(data_unordered[["pared", "public", "gpa"]], float),
            distr=cloglog,
        )
        resu = modu.fit(method="bfgs", disp=False)

        from .results.results_ordinal_model import res_ord_cloglog as res2

        cls.res2 = res2
        cls.res1 = res
        cls.resp = resp
        cls.resf = resf
        cls.resu = resu


class TestLogitBinary:
    # compare OrderedModel with discrete Logit for binary case
    def test_attributes(self):
        data = ds.df

        mask_drop = data["apply"] == "somewhat likely"
        data2 = data.loc[~mask_drop, :].copy()
        # we need to remove the category also from the Categorical Index
        data2["apply"] = data2["apply"].cat.remove_categories("somewhat likely")

        # standard fit with pandas input
        modp = OrderedModel(
            data2["apply"], data2[["pared", "public", "gpa"]], distr="logit"
        )
        resp = modp.fit(method="bfgs", disp=False)

        exog = add_constant(data2[["pared", "public", "gpa"]], prepend=False)
        mod_logit = Logit(data2["apply"].cat.codes, exog)
        res_logit = mod_logit.fit()

        attributes = "bse df_resid llf aic bic llnull".split()
        attributes += "llnull llr llr_pvalue prsquared".split()
        params = np.asarray(resp.params)
        logit_params = np.asarray(res_logit.params)
        assert_allclose(params[:3], logit_params[:3], rtol=1e-5)
        assert_allclose(params[3], -logit_params[3], rtol=1e-5)
        for attr in attributes:
            assert_allclose(getattr(resp, attr), getattr(res_logit, attr), rtol=1e-4)

        resp = modp.fit(
            method="bfgs", disp=False, cov_type="hac", cov_kwds={"maxlags": 2}
        )
        res_logit = mod_logit.fit(
            method="bfgs", disp=False, cov_type="hac", cov_kwds={"maxlags": 2}
        )
        for attr in attributes:
            assert_allclose(getattr(resp, attr), getattr(res_logit, attr), rtol=1e-4)

        resp = modp.fit(method="bfgs", disp=False, cov_type="hc1")
        res_logit = mod_logit.fit(method="bfgs", disp=False, cov_type="hc1")
        for attr in attributes:
            assert_allclose(getattr(resp, attr), getattr(res_logit, attr), rtol=1e-4)


def test_nan_endog_exceptions():
    nobs = 15
    y = np.repeat(np.arange(3), nobs // 3)
    x = np.column_stack((np.ones(nobs), np.arange(nobs)))
    with pytest.raises(ValueError, match="not be a constant"):
        OrderedModel(y, x, distr="logit")

    y_nan = y.astype(float)
    y_nan[0] = np.nan
    with pytest.raises(ValueError, match="NaN in dependent variable"):
        OrderedModel(y_nan, x[:, 1:], distr="logit")

    if hasattr(pd, "CategoricalDtype"):
        cats = pd.Series(y, dtype=pd.CategoricalDtype([0, 1, 2], ordered=True))
        cats.loc[cats == 0] = np.nan
        df = pd.DataFrame({"endog": cats, "exog": x[:, 1]})

        msg = "missing values in categorical endog"
        with pytest.raises(ValueError, match=msg):
            OrderedModel(df["endog"], df[["exog"]])