"""Unit tests for generic score/LM tests and conditional moment tests

Created on Mon Nov 17 08:44:06 2014

Author: Josef Perktold
License: BSD-3

"""

import numpy as np
from numpy.testing import assert_allclose

from statsmodels.regression.linear_model import OLS
import statsmodels.stats._diagnostic_other as diao
from statsmodels.stats._diagnostic_other import CMTNewey, CMTTauchen


class CheckCMT:

    def test_score(self):
        expected = self.results_score
        for msg, actual in self.res_score():
            # not all cases provide all 3 elements,
            # TODO: fix API, returns of functions
            assert_allclose(
                actual, expected[: np.size(actual)], rtol=1e-13, err_msg=msg
            )

    def test_scorehc0(self):
        expected = self.results_hc0
        for msg, actual in self.res_hc0():
            assert_allclose(
                actual, expected[: np.size(actual)], rtol=1e-13, err_msg=msg
            )

    def test_scoreopg(self):
        expected = self.results_opg
        for msg, actual in self.res_opg():
            assert_allclose(
                actual, expected[: np.size(actual)], rtol=1e-13, err_msg=msg
            )


class TestCMTOLS(CheckCMT):

    @classmethod
    def setup_class(cls):

        np.random.seed(864259)

        nobs, k_vars = 100, 4
        sig_e = 1
        x0 = np.random.randn(nobs, k_vars)
        x0[:, 0] = 1
        y_true = x0.sum(1)
        y = y_true + sig_e * np.random.randn(nobs)

        x1 = np.random.randn(nobs, 2)
        x = np.column_stack((x0, x1))

        cls.exog_full = x
        cls.exog_add = x1
        cls.res_ols = OLS(y, x0).fit()

        cls.attach_moment_conditions()

        # results from initial run, not reference package (drooped 2 digits)
        cls.results_score = (1.6857659627548, 0.43046770240535, 2)
        cls.results_hc0 = (1.6385932313952, 0.4407415561953, 2)
        cls.results_opg = (1.72226002418488, 0.422684174119544, 2.0)

    @classmethod
    def attach_moment_conditions(cls):
        # TODO: a better structure ?
        res_ols = cls.res_ols
        # assumes x = column_stack([x0, x1])
        x = cls.exog_full
        # x0 = cls.res_ols.model.exog  # not used here
        x1 = cls.exog_add

        nobs, k_constraints = x1.shape

        # TODO: cleanup after initial copy past
        moms_obs = res_ols.resid[:, None] * x
        moms = moms_obs.sum(0)
        cov_moms = res_ols.mse_resid * x.T.dot(x)  # np.linalg.inv(x.T.dot(x))
        cov_moms *= res_ols.df_resid / nobs

        # weights used for GMM to replicate OLS
        weights = np.linalg.inv(cov_moms)
        # we do not use last two variables
        weights[:, -k_constraints:] = 0
        weights[-k_constraints:, :] = 0

        k_moms = moms.shape[0]
        # TODO: Newey has different versions that all produce the same result
        #       in example
        L = np.eye(k_moms)[-k_constraints:]  # .dot(np.linalg.inv(cov_moms))

        moms_deriv = cov_moms[:, :-k_constraints]

        covm = moms_obs.T.dot(moms_obs)

        # attach
        cls.nobs = nobs
        cls.moms = moms
        cls.moms_obs = moms_obs
        cls.cov_moms = cov_moms
        cls.covm = covm
        cls.moms_deriv = moms_deriv
        cls.weights = weights
        cls.L = L

    def res_score(self):
        res_ols = self.res_ols
        nobs = self.nobs
        moms = self.moms
        cov_moms = self.cov_moms
        weights = self.weights
        L = self.L
        x = self.exog_full  # for auxiliary regression only

        res_all = []

        # auxiliary regression
        stat = nobs * OLS(res_ols.resid, x).fit().rsquared
        res_all.append(("ols R2", stat))

        stat = moms.dot(np.linalg.solve(cov_moms, moms))
        res_all.append(("score simple", stat))

        tres = diao.lm_robust(
            moms, np.eye(moms.shape[0])[-2:], np.linalg.inv(cov_moms), cov_moms
        )
        res_all.append(("score mle", tres))

        tres = CMTNewey(moms, cov_moms, cov_moms[:, :-2], weights, L).chisquare
        res_all.append(("Newey", tres))

        tres = CMTTauchen(
            moms[:-2], cov_moms[:-2, :-2], moms[-2:], cov_moms[-2:, :-2], cov_moms
        ).chisquare
        res_all.append(("Tauchen", tres))

        return res_all

    def res_opg(self):
        res_ols = self.res_ols
        nobs = self.nobs
        moms = self.moms
        moms_obs = self.moms_obs
        covm = self.covm
        weights = self.weights
        L = self.L
        x = self.exog_full

        res_ols2_hc0 = OLS(res_ols.model.endog, x).fit(cov_type="HC0")

        res_all = []

        # auxiliary regression
        ones = np.ones(nobs)
        stat = nobs * OLS(ones, moms_obs).fit().rsquared
        res_all.append(("ols R2", stat))

        tres = res_ols2_hc0.compare_lm_test(res_ols, demean=False)
        res_all.append(("comp_lm uc", tres))

        tres = CMTNewey(moms, covm, covm[:, :-2], weights, L).chisquare
        res_all.append(("Newey", tres))

        tres = CMTTauchen(
            moms[:-2], covm[:-2, :-2], moms[-2:], covm[-2:, :-2], covm
        ).chisquare
        res_all.append(("Tauchen", tres))

        tres = diao.lm_robust_subset(moms[-2:], 2, covm, covm)
        res_all.append(("score subset QMLE", tres))

        tres = diao.lm_robust(
            moms, np.eye(moms.shape[0])[-2:], np.linalg.inv(covm), covm, cov_params=None
        )
        res_all.append(("scoreB QMLE", tres))

        tres = diao.lm_robust(
            moms,
            np.eye(moms.shape[0])[-2:],
            np.linalg.inv(covm),
            None,
            cov_params=np.linalg.inv(covm),
        )
        res_all.append(("scoreV QMLE", tres))

        return res_all

    def res_hc0(self):
        res_ols = self.res_ols
        moms = self.moms
        moms_obs = self.moms_obs
        cov_moms = self.cov_moms  # Hessian with scale
        covm = self.covm
        moms_deriv = self.moms_deriv
        weights = self.weights
        L = self.L

        x0 = res_ols.model.exog
        x1 = self.exog_add

        res_all = []

        tres = diao.cm_test_robust(
            resid=res_ols.resid, resid_deriv=x0, instruments=x1, weights=1
        )
        # TODO: extra return and no df in cm_test_robust Wooldridge
        res_all.append(("Wooldridge", tres[:2]))

        tres = CMTNewey(moms, covm, moms_deriv, weights, L).chisquare
        res_all.append(("Newey", tres))

        tres = CMTTauchen(
            moms[:-2], cov_moms[:-2, :-2], moms[-2:], cov_moms[-2:, :-2], covm
        ).chisquare
        res_all.append(("Tauchen", tres))

        tres = diao.lm_robust_subset(moms[-2:], 2, cov_moms, covm)
        res_all.append(("score subset QMLE", tres))

        tres = diao.lm_robust(
            moms, np.eye(moms.shape[0])[-2:], np.linalg.inv(cov_moms), covm
        )
        res_all.append(("scoreB QMLE", tres))

        # need sandwich cov_params V
        Ainv = np.linalg.inv(cov_moms)
        vv = Ainv.dot(covm).dot(Ainv)
        tres = diao.lm_robust(
            moms,
            np.eye(moms.shape[0])[-2:],
            np.linalg.inv(cov_moms),
            None,
            cov_params=vv,
        )
        res_all.append(("scoreV QMLE", tres))

        tres = diao.conditional_moment_test_generic(
            moms_obs[:, -2:], cov_moms[-2:, :-2], moms_obs[:, :-2], cov_moms[:-2, :-2]
        )
        tres_ = (tres.stat_cmt, tres.pval_cmt)
        res_all.append(("cmt", tres_))

        # using unscaled hessian instead of scaled
        x = self.exog_full
        hess_unscaled = x.T.dot(x)
        tres = diao.conditional_moment_test_generic(
            moms_obs[:, -2:],
            hess_unscaled[-2:, :-2],
            moms_obs[:, :-2],
            hess_unscaled[:-2, :-2],
        )  # , covm)
        tres_ = (tres.stat_cmt, tres.pval_cmt)
        res_all.append(("cmt", tres_))

        score_deriv_uu = cov_moms[:-2, :-2]
        score_deriv_cu = cov_moms[-2:, :-2]
        cov_score_cc = covm[-2:, -2:]
        cov_score_cu = covm[-2:, :-2]
        cov_score_uu = covm[:-2, :-2]
        tres = diao.lm_robust_subset_parts(
            moms[-2:],
            2,
            score_deriv_uu,
            score_deriv_cu,
            cov_score_cc,
            cov_score_cu,
            cov_score_uu,
        )

        res_all.append(("score subset_parts QMLE", tres))

        params_deriv = np.eye(x.shape[1], x.shape[1] - 2)
        # params_deriv[[-2, -1], [-2, -1]] = 0
        score = moms
        score_deriv = cov_moms
        cov_score = covm

        tres = diao.lm_robust_reparameterized(
            score, params_deriv, score_deriv, cov_score
        )

        res_all.append(("score reparam QMLE", tres))

        return res_all
