"""
Test functions for GEE

External comparisons are to R and Stata.  The statsmodels GEE
implementation should generally agree with the R GEE implementation
for the independence and exchangeable correlation structures.  For
other correlation structures, the details of the correlation
estimation differ among implementations and the results will not agree
exactly.
"""

import os
import warnings

import numpy as np
from numpy.testing import (
    assert_,
    assert_allclose,
    assert_almost_equal,
    assert_array_less,
    assert_equal,
)
import pandas as pd
import pytest
from scipy.stats.distributions import norm

from statsmodels import tools
from statsmodels.compat import lrange
import statsmodels.discrete.discrete_model as discrete
from statsmodels.genmod import cov_struct, families
import statsmodels.genmod.generalized_estimating_equations as gee
import statsmodels.regression.linear_model as lm
from statsmodels.tools.sm_exceptions import SpecificationWarning

try:
    import matplotlib.pyplot as plt
except ImportError:
    pass

pdf_output = False

if pdf_output:
    from matplotlib.backends.backend_pdf import PdfPages

    pdf = PdfPages("test_glm.pdf")
else:
    pdf = None


def close_or_save(pdf, fig):
    if pdf_output:
        pdf.savefig(fig)


def load_data(fname, icept=True):
    """
    Load a data set from the results directory.  The data set should
    be a CSV file with the following format:

    Column 0: Group indicator
    Column 1: endog variable
    Columns 2-end: exog variables

    If `icept` is True, an intercept is prepended to the exog
    variables.
    """

    cur_dir = os.path.dirname(os.path.abspath(__file__))
    Z = np.genfromtxt(os.path.join(cur_dir, "results", fname), delimiter=",")

    group = Z[:, 0]
    endog = Z[:, 1]
    exog = Z[:, 2:]

    if icept:
        exog = np.concatenate((np.ones((exog.shape[0], 1)), exog), axis=1)

    return endog, exog, group


def check_wrapper(results):
    # check wrapper
    assert_(isinstance(results.params, pd.Series))
    assert_(isinstance(results.fittedvalues, pd.Series))
    assert_(isinstance(results.resid, pd.Series))
    assert_(isinstance(results.centered_resid, pd.Series))

    assert_(isinstance(results._results.params, np.ndarray))
    assert_(isinstance(results._results.fittedvalues, np.ndarray))
    assert_(isinstance(results._results.resid, np.ndarray))
    assert_(isinstance(results._results.centered_resid, np.ndarray))


class TestGEE:

    def test_margins_gaussian(self):
        # Check marginal effects for a Gaussian GEE fit.  Marginal
        # effects and ordinary effects should be equal.

        n = 40
        np.random.seed(34234)
        exog = np.random.normal(size=(n, 3))
        exog[:, 0] = 1

        groups = np.kron(np.arange(n / 4), np.r_[1, 1, 1, 1])
        endog = exog[:, 1] + np.random.normal(size=n)

        model = gee.GEE(endog, exog, groups)
        result = model.fit(start_params=[-4.88085602e-04, 1.18501903, 4.78820100e-02])

        marg = result.get_margeff()

        assert_allclose(marg.margeff, result.params[1:])
        assert_allclose(marg.margeff_se, result.bse[1:])

        # smoke test
        marg.summary()

    def test_margins_gaussian_lists_tuples(self):
        # Check marginal effects for a Gaussian GEE fit using lists and
        # tuples. Marginal effects and ordinary effects should be equal.

        n = 40
        np.random.seed(34234)
        exog_arr = np.random.normal(size=(n, 3))
        exog_arr[:, 0] = 1

        groups_arr = np.kron(np.arange(n / 4), np.r_[1, 1, 1, 1])
        endog_arr = exog_arr[:, 1] + np.random.normal(size=n)

        # check that GEE accepts lists
        exog_list = [list(row) for row in exog_arr]
        groups_list = list(groups_arr)
        endog_list = list(endog_arr)

        model = gee.GEE(endog_list, exog_list, groups_list)
        result = model.fit(start_params=[-4.88085602e-04, 1.18501903, 4.78820100e-02])

        marg = result.get_margeff()

        assert_allclose(marg.margeff, result.params[1:])
        assert_allclose(marg.margeff_se, result.bse[1:])

        # check that GEE accepts tuples
        exog_tuple = tuple(tuple(row) for row in exog_arr)
        groups_tuple = tuple(groups_arr)
        endog_tuple = tuple(endog_arr)

        model = gee.GEE(endog_tuple, exog_tuple, groups_tuple)
        result = model.fit(start_params=[-4.88085602e-04, 1.18501903, 4.78820100e-02])

        marg = result.get_margeff()

        assert_allclose(marg.margeff, result.params[1:])
        assert_allclose(marg.margeff_se, result.bse[1:])

    def test_margins_logistic(self):
        # Check marginal effects for a binomial GEE fit.  Comparison
        # comes from Stata.

        np.random.seed(34234)
        endog = np.r_[0, 0, 0, 0, 1, 1, 1, 1]
        exog = np.ones((8, 2))
        exog[:, 1] = np.r_[1, 2, 1, 1, 2, 1, 2, 2]

        groups = np.arange(8)

        model = gee.GEE(endog, exog, groups, family=families.Binomial())
        result = model.fit(cov_type="naive", start_params=[-3.29583687, 2.19722458])

        marg = result.get_margeff()

        assert_allclose(marg.margeff, np.r_[0.4119796])
        assert_allclose(marg.margeff_se, np.r_[0.1379962], rtol=1e-6)

    def test_margins_multinomial(self):
        # Check marginal effects for a 2-class multinomial GEE fit,
        # which should be equivalent to logistic regression.  Comparison
        # comes from Stata.

        np.random.seed(34234)
        endog = np.r_[0, 0, 0, 0, 1, 1, 1, 1]
        exog = np.ones((8, 2))
        exog[:, 1] = np.r_[1, 2, 1, 1, 2, 1, 2, 2]

        groups = np.arange(8)

        model = gee.NominalGEE(endog, exog, groups)
        result = model.fit(cov_type="naive", start_params=[3.295837, -2.197225])

        marg = result.get_margeff()

        assert_allclose(marg.margeff, np.r_[-0.41197961], rtol=1e-5)
        assert_allclose(marg.margeff_se, np.r_[0.1379962], rtol=1e-6)

    @pytest.mark.smoke
    @pytest.mark.matplotlib
    def test_nominal_plot(self, close_figures):
        np.random.seed(34234)
        endog = np.r_[0, 0, 0, 0, 1, 1, 1, 1]
        exog = np.ones((8, 2))
        exog[:, 1] = np.r_[1, 2, 1, 1, 2, 1, 2, 2]

        groups = np.arange(8)

        model = gee.NominalGEE(endog, exog, groups)
        result = model.fit(cov_type="naive", start_params=[3.295837, -2.197225])

        fig = result.plot_distribution()
        assert_equal(isinstance(fig, plt.Figure), True)

    def test_margins_poisson(self):
        # Check marginal effects for a Poisson GEE fit.

        np.random.seed(34234)
        endog = np.r_[10, 15, 12, 13, 20, 18, 26, 29]
        exog = np.ones((8, 2))
        exog[:, 1] = np.r_[0, 0, 0, 0, 1, 1, 1, 1]

        groups = np.arange(8)

        model = gee.GEE(endog, exog, groups, family=families.Poisson())
        result = model.fit(cov_type="naive", start_params=[2.52572864, 0.62057649])

        marg = result.get_margeff()

        assert_allclose(marg.margeff, np.r_[11.0928], rtol=1e-6)
        assert_allclose(marg.margeff_se, np.r_[3.269015], rtol=1e-6)

    def test_multinomial(self):
        """
        Check the 2-class multinomial (nominal) GEE fit against
        logistic regression.
        """

        np.random.seed(34234)
        endog = np.r_[0, 0, 0, 0, 1, 1, 1, 1]
        exog = np.ones((8, 2))
        exog[:, 1] = np.r_[1, 2, 1, 1, 2, 1, 2, 2]

        groups = np.arange(8)

        model = gee.NominalGEE(endog, exog, groups)
        results = model.fit(cov_type="naive", start_params=[3.295837, -2.197225])

        logit_model = gee.GEE(endog, exog, groups, family=families.Binomial())
        logit_results = logit_model.fit(cov_type="naive")

        assert_allclose(results.params, -logit_results.params, rtol=1e-5)
        assert_allclose(results.bse, logit_results.bse, rtol=1e-5)

    def test_weighted(self):

        # Simple check where the answer can be computed by hand.
        exog = np.ones(20)
        weights = np.ones(20)
        weights[0:10] = 2
        endog = np.zeros(20)
        endog[0:10] += 1
        groups = np.kron(np.arange(10), np.r_[1, 1])
        model = gee.GEE(endog, exog, groups, weights=weights)
        result = model.fit()
        assert_allclose(result.params, np.r_[2 / 3.0])

        # Comparison against stata using groups with different sizes.
        weights = np.ones(20)
        weights[10:] = 2
        endog = np.r_[1, 2, 3, 2, 3, 4, 3, 4, 5, 4, 5, 6, 5, 6, 7, 6, 7, 8, 7, 8]
        exog1 = np.r_[1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 3, 3, 3, 3]
        groups = np.r_[1, 1, 2, 2, 2, 2, 4, 4, 5, 5, 6, 6, 6, 6, 8, 8, 9, 9, 10, 10]
        exog = np.column_stack((np.ones(20), exog1))

        # Comparison using independence model
        model = gee.GEE(
            endog, exog, groups, weights=weights, cov_struct=cov_struct.Independence()
        )
        g = np.mean([2, 4, 2, 2, 4, 2, 2, 2])
        fac = 20 / float(20 - g)
        result = model.fit(ddof_scale=0, scaling_factor=fac)

        assert_allclose(result.params, np.r_[1.247573, 1.436893], atol=1e-6)
        assert_allclose(result.scale, 1.808576)

        # Stata multiples robust SE by sqrt(N / (N - g)), where N is
        # the total sample size and g is the average group size.
        assert_allclose(result.bse, np.r_[0.895366, 0.3425498], atol=1e-5)

        # Comparison using exchangeable model
        # Smoke test for now
        model = gee.GEE(
            endog, exog, groups, weights=weights, cov_struct=cov_struct.Exchangeable()
        )
        model.fit(ddof_scale=0)

    # This is in the release announcement for version 0.6.
    def test_poisson_epil(self):

        cur_dir = os.path.dirname(os.path.abspath(__file__))
        fname = os.path.join(cur_dir, "results", "epil.csv")
        data = pd.read_csv(fname)

        fam = families.Poisson()
        ind = cov_struct.Independence()
        mod1 = gee.GEE.from_formula(
            "y ~ age + trt + base", data["subject"], data, cov_struct=ind, family=fam
        )
        rslt1 = mod1.fit(cov_type="naive")

        # Coefficients should agree with GLM
        from statsmodels.genmod.generalized_linear_model import GLM

        mod2 = GLM.from_formula("y ~ age + trt + base", data, family=families.Poisson())
        rslt2 = mod2.fit()

        # do not use wrapper, asserts_xxx do not work
        rslt1 = rslt1._results
        rslt2 = rslt2._results

        assert_allclose(rslt1.params, rslt2.params, rtol=1e-6, atol=1e-6)
        assert_allclose(rslt1.bse, rslt2.bse, rtol=1e-6, atol=1e-6)

    def test_missing(self):
        # Test missing data handling for calling from the api.  Missing
        # data handling does not currently work for formulas.

        np.random.seed(34234)
        endog = np.random.normal(size=100)
        exog = np.random.normal(size=(100, 3))
        exog[:, 0] = 1
        groups = np.kron(lrange(20), np.ones(5))

        endog[0] = np.nan
        endog[5:7] = np.nan
        exog[10:12, 1] = np.nan

        mod1 = gee.GEE(endog, exog, groups, missing="drop")
        rslt1 = mod1.fit()

        assert_almost_equal(len(mod1.endog), 95)
        assert_almost_equal(np.asarray(mod1.exog.shape), np.r_[95, 3])

        ii = np.isfinite(endog) & np.isfinite(exog).all(1)

        mod2 = gee.GEE(endog[ii], exog[ii, :], groups[ii], missing="none")
        rslt2 = mod2.fit()

        assert_almost_equal(rslt1.params, rslt2.params)
        assert_almost_equal(rslt1.bse, rslt2.bse)

    def test_missing_formula(self):
        # Test missing data handling for formulas.

        np.random.seed(34234)
        endog = np.random.normal(size=100)
        exog1 = np.random.normal(size=100)
        exog2 = np.random.normal(size=100)
        exog3 = np.random.normal(size=100)
        groups = np.kron(lrange(20), np.ones(5))

        endog[0] = np.nan
        endog[5:7] = np.nan
        exog2[10:12] = np.nan

        data0 = pd.DataFrame(
            {
                "endog": endog,
                "exog1": exog1,
                "exog2": exog2,
                "exog3": exog3,
                "groups": groups,
            }
        )

        for k in 0, 1:
            data = data0.copy()
            kwargs = {}
            if k == 1:
                data["offset"] = 0
                data["time"] = 0
                kwargs["offset"] = "offset"
                kwargs["time"] = "time"

            mod1 = gee.GEE.from_formula(
                "endog ~ exog1 + exog2 + exog3",
                groups="groups",
                data=data,
                missing="drop",
                **kwargs,
            )
            rslt1 = mod1.fit()

            assert_almost_equal(len(mod1.endog), 95)
            assert_almost_equal(np.asarray(mod1.exog.shape), np.r_[95, 4])

            data = data.dropna()

            kwargs = {}
            if k == 1:
                kwargs["offset"] = data["offset"]
                kwargs["time"] = data["time"]

            mod2 = gee.GEE.from_formula(
                "endog ~ exog1 + exog2 + exog3",
                groups=data["groups"],
                data=data,
                missing="none",
                **kwargs,
            )
            rslt2 = mod2.fit()

            assert_almost_equal(rslt1.params.values, rslt2.params.values)
            assert_almost_equal(rslt1.bse.values, rslt2.bse.values)

    @pytest.mark.parametrize("k1", [False, True])
    @pytest.mark.parametrize("k2", [False, True])
    def test_invalid_args(self, k1, k2):

        for j in range(3):

            p = [20, 20, 20]
            p[j] = 18

            endog = np.zeros(p[0])
            exog = np.zeros((p[1], 2))

            kwargs = {}
            kwargs["groups"] = np.zeros(p[2])
            if k1:
                kwargs["exposure"] = np.zeros(18)
            if k2:
                kwargs["time"] = np.zeros(18)
            with pytest.raises(ValueError):
                gee.GEE(endog, exog, **kwargs)

    def test_default_time(self):
        # Check that the time defaults work correctly.

        endog, exog, group = load_data("gee_logistic_1.csv")

        # Time values for the autoregressive model
        T = np.zeros(len(endog))
        idx = set(group)
        for ii in idx:
            jj = np.flatnonzero(group == ii)
            T[jj] = lrange(len(jj))

        family = families.Binomial()
        va = cov_struct.Autoregressive(grid=False)

        md1 = gee.GEE(endog, exog, group, family=family, cov_struct=va)
        mdf1 = md1.fit()

        md2 = gee.GEE(endog, exog, group, time=T, family=family, cov_struct=va)
        mdf2 = md2.fit()

        assert_almost_equal(mdf1.params, mdf2.params, decimal=6)
        assert_almost_equal(mdf1.standard_errors(), mdf2.standard_errors(), decimal=6)

    def test_logistic(self):
        # R code for comparing results:

        # library(gee)
        # Z = read.csv("results/gee_logistic_1.csv", header=FALSE)
        # Y = Z[,2]
        # Id = Z[,1]
        # X1 = Z[,3]
        # X2 = Z[,4]
        # X3 = Z[,5]

        # mi = gee(Y ~ X1 + X2 + X3, id=Id, family=binomial,
        #         corstr="independence")
        # smi = summary(mi)
        # u = coefficients(smi)
        # cfi = paste(u[,1], collapse=",")
        # sei = paste(u[,4], collapse=",")

        # me = gee(Y ~ X1 + X2 + X3, id=Id, family=binomial,
        #         corstr="exchangeable")
        # sme = summary(me)
        # u = coefficients(sme)
        # cfe = paste(u[,1], collapse=",")
        # see = paste(u[,4], collapse=",")

        # ma = gee(Y ~ X1 + X2 + X3, id=Id, family=binomial,
        #         corstr="AR-M")
        # sma = summary(ma)
        # u = coefficients(sma)
        # cfa = paste(u[,1], collapse=",")
        # sea = paste(u[,4], collapse=",")

        # sprintf("cf = [[%s],[%s],[%s]]", cfi, cfe, cfa)
        # sprintf("se = [[%s],[%s],[%s]]", sei, see, sea)

        endog, exog, group = load_data("gee_logistic_1.csv")

        # Time values for the autoregressive model
        T = np.zeros(len(endog))
        idx = set(group)
        for ii in idx:
            jj = np.flatnonzero(group == ii)
            T[jj] = lrange(len(jj))

        family = families.Binomial()
        ve = cov_struct.Exchangeable()
        vi = cov_struct.Independence()
        va = cov_struct.Autoregressive(grid=False)

        # From R gee
        cf = [
            [0.0167272965285882, 1.13038654425893, -1.86896345082962, 1.09397608331333],
            [0.0178982283915449, 1.13118798191788, -1.86133518416017, 1.08944256230299],
            [0.0109621937947958, 1.13226505028438, -1.88278757333046, 1.09954623769449],
        ]
        se = [
            [
                0.127291720283049,
                0.166725808326067,
                0.192430061340865,
                0.173141068839597,
            ],
            [
                0.127045031730155,
                0.165470678232842,
                0.192052750030501,
                0.173174779369249,
            ],
            [
                0.127240302296444,
                0.170554083928117,
                0.191045527104503,
                0.169776150974586,
            ],
        ]

        for j, v in enumerate((vi, ve, va)):
            md = gee.GEE(endog, exog, group, T, family, v)
            mdf = md.fit()
            if id(v) != id(va):
                assert_almost_equal(mdf.params, cf[j], decimal=6)
                assert_almost_equal(mdf.standard_errors(), se[j], decimal=6)

        # Test with formulas
        D = np.concatenate((endog[:, None], group[:, None], exog[:, 1:]), axis=1)
        D = pd.DataFrame(D)
        D.columns = [
            "Y",
            "Id",
        ] + ["X%d" % (k + 1) for k in range(exog.shape[1] - 1)]
        for j, v in enumerate((vi, ve)):
            md = gee.GEE.from_formula(
                "Y ~ X1 + X2 + X3", "Id", D, family=family, cov_struct=v
            )
            mdf = md.fit()
            assert_almost_equal(mdf.params, cf[j], decimal=6)
            assert_almost_equal(mdf.standard_errors(), se[j], decimal=6)

        # FIXME: do not leave commented-out
        # Check for run-time exceptions in summary
        # print(mdf.summary())

    def test_autoregressive(self):

        dep_params_true = [0, 0.589208623896, 0.559823804948]

        params_true = [
            [1.08043787, 1.12709319, 0.90133927],
            [0.9613677, 1.05826987, 0.90832055],
            [1.05370439, 0.96084864, 0.93923374],
        ]

        np.random.seed(342837482)

        num_group = 100
        ar_param = 0.5
        k = 3

        ga = families.Gaussian()

        for gsize in 1, 2, 3:

            ix = np.arange(gsize)[:, None] - np.arange(gsize)[None, :]
            ix = np.abs(ix)
            cmat = ar_param**ix
            cmat_r = np.linalg.cholesky(cmat)

            endog = []
            exog = []
            groups = []
            for i in range(num_group):
                x = np.random.normal(size=(gsize, k))
                exog.append(x)
                expval = x.sum(1)
                errors = np.dot(cmat_r, np.random.normal(size=gsize))
                endog.append(expval + errors)
                groups.append(i * np.ones(gsize))

            endog = np.concatenate(endog)
            groups = np.concatenate(groups)
            exog = np.concatenate(exog, axis=0)

            ar = cov_struct.Autoregressive(grid=False)
            md = gee.GEE(endog, exog, groups, family=ga, cov_struct=ar)
            mdf = md.fit()
            assert_almost_equal(ar.dep_params, dep_params_true[gsize - 1])
            assert_almost_equal(mdf.params, params_true[gsize - 1])

    def test_post_estimation(self):

        family = families.Gaussian()
        endog, exog, group = load_data("gee_linear_1.csv")

        ve = cov_struct.Exchangeable()

        md = gee.GEE(endog, exog, group, None, family, ve)
        mdf = md.fit()

        assert_almost_equal(np.dot(exog, mdf.params), mdf.fittedvalues)
        assert_almost_equal(endog - np.dot(exog, mdf.params), mdf.resid)

    def test_scoretest(self):
        # Regression tests

        np.random.seed(6432)
        n = 200  # Must be divisible by 4
        exog = np.random.normal(size=(n, 4))
        endog = exog[:, 0] + exog[:, 1] + exog[:, 2]
        endog += 3 * np.random.normal(size=n)
        group = np.kron(np.arange(n / 4), np.ones(4))

        # Test under the null.
        L = np.array([[1.0, -1, 0, 0]])
        R = np.array(
            [
                0.0,
            ]
        )
        family = families.Gaussian()
        va = cov_struct.Independence()
        mod1 = gee.GEE(
            endog, exog, group, family=family, cov_struct=va, constraint=(L, R)
        )
        res1 = mod1.fit()
        assert_almost_equal(res1.score_test()["statistic"], 1.08126334)
        assert_almost_equal(res1.score_test()["p-value"], 0.2984151086)

        # Test under the alternative.
        L = np.array([[1.0, -1, 0, 0]])
        R = np.array(
            [
                1.0,
            ]
        )
        family = families.Gaussian()
        va = cov_struct.Independence()
        mod2 = gee.GEE(
            endog, exog, group, family=family, cov_struct=va, constraint=(L, R)
        )
        res2 = mod2.fit()
        assert_almost_equal(res2.score_test()["statistic"], 3.491110965)
        assert_almost_equal(res2.score_test()["p-value"], 0.0616991659)

        # Compare to Wald tests
        exog = np.random.normal(size=(n, 2))
        L = np.array([[1, -1]])
        R = np.array([0.0])
        f = np.r_[1, -1]
        for i in range(10):
            endog = (
                exog[:, 0] + (0.5 + i / 10.0) * exog[:, 1] + np.random.normal(size=n)
            )
            family = families.Gaussian()
            va = cov_struct.Independence()
            mod0 = gee.GEE(endog, exog, group, family=family, cov_struct=va)
            rslt0 = mod0.fit()
            family = families.Gaussian()
            va = cov_struct.Independence()
            mod1 = gee.GEE(
                endog, exog, group, family=family, cov_struct=va, constraint=(L, R)
            )
            res1 = mod1.fit()
            se = np.sqrt(np.dot(f, np.dot(rslt0.cov_params(), f)))
            wald_z = np.dot(f, rslt0.params) / se
            wald_p = 2 * norm.cdf(-np.abs(wald_z))
            score_p = res1.score_test()["p-value"]
            assert_array_less(np.abs(wald_p - score_p), 0.02)

    @pytest.mark.parametrize(
        "cov_struct", [cov_struct.Independence, cov_struct.Exchangeable]
    )
    def test_compare_score_test(self, cov_struct):

        np.random.seed(6432)
        n = 200  # Must be divisible by 4
        exog = np.random.normal(size=(n, 4))
        group = np.kron(np.arange(n / 4), np.ones(4))

        exog_sub = exog[:, [0, 3]]
        endog = exog_sub.sum(1) + 3 * np.random.normal(size=n)

        L = np.asarray([[0, 1, 0, 0], [0, 0, 1, 0]])
        R = np.zeros(2)
        mod_lr = gee.GEE(endog, exog, group, constraint=(L, R), cov_struct=cov_struct())
        mod_lr.fit()

        mod_sub = gee.GEE(endog, exog_sub, group, cov_struct=cov_struct())
        res_sub = mod_sub.fit()

        for call_fit in [False, True]:
            mod = gee.GEE(endog, exog, group, cov_struct=cov_struct())
            if call_fit:
                # Should work with or without fitting the parent model
                mod.fit()
            score_results = mod.compare_score_test(res_sub)
            assert_almost_equal(
                score_results["statistic"], mod_lr.score_test_results["statistic"]
            )
            assert_almost_equal(
                score_results["p-value"], mod_lr.score_test_results["p-value"]
            )
            assert_almost_equal(score_results["df"], mod_lr.score_test_results["df"])

    def test_compare_score_test_warnings(self):

        np.random.seed(6432)
        n = 200  # Must be divisible by 4
        exog = np.random.normal(size=(n, 4))
        group = np.kron(np.arange(n / 4), np.ones(4))
        exog_sub = exog[:, [0, 3]]
        endog = exog_sub.sum(1) + 3 * np.random.normal(size=n)

        # Mismatched cov_struct
        mod_sub = gee.GEE(
            endog, exog_sub, group, cov_struct=cov_struct.Exchangeable()
        )
        res_sub = mod_sub.fit()
        mod = gee.GEE(endog, exog, group, cov_struct=cov_struct.Independence())
        with pytest.warns(UserWarning, match="Model and submodel have"):
            mod.compare_score_test(res_sub)  # smoketest

        # Mismatched family
        mod_sub = gee.GEE(endog, exog_sub, group, family=families.Gaussian())
        res_sub = mod_sub.fit()
        mod = gee.GEE(endog, exog, group, family=families.Poisson())
        with pytest.warns(SpecificationWarning, match="Model and submodel have"):
            mod.compare_score_test(res_sub)  # smoketest

        # Mismatched size
        with pytest.raises(ValueError):
            mod_sub = gee.GEE(endog, exog_sub, group)
            res_sub = mod_sub.fit()
            mod = gee.GEE(endog[0:100], exog[:100, :], group[0:100])
            mod.compare_score_test(res_sub)  # smoketest

        # Mismatched weights
        w = np.random.uniform(size=n)
        mod_sub = gee.GEE(endog, exog_sub, group, weights=w)
        res_sub = mod_sub.fit()
        mod = gee.GEE(endog, exog, group)
        with pytest.warns(SpecificationWarning, match="Model and submodel should"):
            mod.compare_score_test(res_sub)  # smoketest

        # Parent and submodel are the same dimension
        w = np.random.uniform(size=n)
        mod_sub = gee.GEE(endog, exog, group)
        res_sub = mod_sub.fit()
        mod = gee.GEE(endog, exog, group)
        with pytest.warns(SpecificationWarning, match="Model and submodel have"):
            mod.compare_score_test(res_sub)  # smoketest

    def test_constraint_covtype(self):
        # Test constraints with different cov types
        np.random.seed(6432)
        n = 200
        exog = np.random.normal(size=(n, 4))
        endog = exog[:, 0] + exog[:, 1] + exog[:, 2]
        endog += 3 * np.random.normal(size=n)
        group = np.kron(np.arange(n / 4), np.ones(4))
        L = np.array([[1.0, -1, 0, 0]])
        R = np.array(
            [
                0.0,
            ]
        )
        family = families.Gaussian()
        va = cov_struct.Independence()
        for cov_type in "robust", "naive", "bias_reduced":
            model = gee.GEE(
                endog, exog, group, family=family, cov_struct=va, constraint=(L, R)
            )
            result = model.fit(cov_type=cov_type)
            result.standard_errors(cov_type=cov_type)
            assert_allclose(result.cov_robust.shape, np.r_[4, 4])
            assert_allclose(result.cov_naive.shape, np.r_[4, 4])
            if cov_type == "bias_reduced":
                assert_allclose(result.cov_robust_bc.shape, np.r_[4, 4])

    def test_linear(self):
        # library(gee)

        # Z = read.csv("results/gee_linear_1.csv", header=FALSE)
        # Y = Z[,2]
        # Id = Z[,1]
        # X1 = Z[,3]
        # X2 = Z[,4]
        # X3 = Z[,5]
        # mi = gee(Y ~ X1 + X2 + X3, id=Id, family=gaussian,
        #         corstr="independence", tol=1e-8, maxit=100)
        # smi = summary(mi)
        # u = coefficients(smi)

        # cfi = paste(u[,1], collapse=",")
        # sei = paste(u[,4], collapse=",")

        # me = gee(Y ~ X1 + X2 + X3, id=Id, family=gaussian,
        #         corstr="exchangeable", tol=1e-8, maxit=100)
        # sme = summary(me)
        # u = coefficients(sme)

        # cfe = paste(u[,1], collapse=",")
        # see = paste(u[,4], collapse=",")

        # sprintf("cf = [[%s],[%s]]", cfi, cfe)
        # sprintf("se = [[%s],[%s]]", sei, see)

        family = families.Gaussian()

        endog, exog, group = load_data("gee_linear_1.csv")

        vi = cov_struct.Independence()
        ve = cov_struct.Exchangeable()

        # From R gee
        cf = [
            [-0.01850226507491, 0.81436304278962, -1.56167635393184, 0.794239361055003],
            [
                -0.0182920577154767,
                0.814898414022467,
                -1.56194040106201,
                0.793499517527478,
            ],
        ]
        se = [
            [
                0.0440733554189401,
                0.0479993639119261,
                0.0496045952071308,
                0.0479467597161284,
            ],
            [
                0.0440369906460754,
                0.0480069787567662,
                0.049519758758187,
                0.0479760443027526,
            ],
        ]

        for j, v in enumerate((vi, ve)):
            md = gee.GEE(endog, exog, group, None, family, v)
            mdf = md.fit()
            assert_almost_equal(mdf.params, cf[j], decimal=10)
            assert_almost_equal(mdf.standard_errors(), se[j], decimal=10)

        # Test with formulas
        D = np.concatenate((endog[:, None], group[:, None], exog[:, 1:]), axis=1)
        D = pd.DataFrame(D)
        D.columns = [
            "Y",
            "Id",
        ] + ["X%d" % (k + 1) for k in range(exog.shape[1] - 1)]
        for j, v in enumerate((vi, ve)):
            md = gee.GEE.from_formula(
                "Y ~ X1 + X2 + X3", "Id", D, family=family, cov_struct=v
            )
            mdf = md.fit()
            assert_almost_equal(mdf.params, cf[j], decimal=10)
            assert_almost_equal(mdf.standard_errors(), se[j], decimal=10)

    def test_linear_constrained(self):

        family = families.Gaussian()

        np.random.seed(34234)
        exog = np.random.normal(size=(300, 4))
        exog[:, 0] = 1
        endog = np.dot(exog, np.r_[1, 1, 0, 0.2]) + np.random.normal(size=300)
        group = np.kron(np.arange(100), np.r_[1, 1, 1])

        vi = cov_struct.Independence()
        ve = cov_struct.Exchangeable()

        L = np.r_[[[0, 0, 0, 1]]]
        R = np.r_[0,]

        for _, v in enumerate((vi, ve)):
            md = gee.GEE(endog, exog, group, None, family, v, constraint=(L, R))
            mdf = md.fit()
            assert_almost_equal(mdf.params[3], 0, decimal=10)

    def test_nested_linear(self):

        family = families.Gaussian()

        endog, exog, group = load_data("gee_nested_linear_1.csv")

        group_n = []
        for _ in range(endog.shape[0] // 10):
            group_n.extend(
                [
                    0,
                ]
                * 5
            )
            group_n.extend(
                [
                    1,
                ]
                * 5
            )
        group_n = np.array(group_n)[:, None]

        dp = cov_struct.Independence()
        md = gee.GEE(endog, exog, group, None, family, dp)
        mdf1 = md.fit()

        # From statsmodels.GEE (not an independent test)
        cf = np.r_[-0.1671073, 1.00467426, -2.01723004, 0.97297106]
        se = np.r_[0.08629606, 0.04058653, 0.04067038, 0.03777989]
        assert_almost_equal(mdf1.params, cf, decimal=6)
        assert_almost_equal(mdf1.standard_errors(), se, decimal=6)

        ne = cov_struct.Nested()
        md = gee.GEE(endog, exog, group, None, family, ne, dep_data=group_n)
        mdf2 = md.fit(start_params=mdf1.params)

        # From statsmodels.GEE (not an independent test)
        cf = np.r_[-0.16655319, 1.02183688, -2.00858719, 1.00101969]
        se = np.r_[0.08632616, 0.02913582, 0.03114428, 0.02893991]
        assert_almost_equal(mdf2.params, cf, decimal=6)
        assert_almost_equal(mdf2.standard_errors(), se, decimal=6)

        smry = mdf2.cov_struct.summary()
        assert_allclose(
            smry.Variance, np.r_[1.043878, 0.611656, 1.421205], atol=1e-5, rtol=1e-5
        )

    def test_nested_pandas(self):

        np.random.seed(4234)
        n = 10000

        # Outer groups
        groups = np.kron(np.arange(n // 100), np.ones(100)).astype(int)

        # Inner groups
        groups1 = np.kron(np.arange(n // 50), np.ones(50)).astype(int)
        groups2 = np.kron(np.arange(n // 10), np.ones(10)).astype(int)

        # Group effects
        groups_e = np.random.normal(size=n // 100)
        groups1_e = 2 * np.random.normal(size=n // 50)
        groups2_e = 3 * np.random.normal(size=n // 10)

        y = groups_e[groups] + groups1_e[groups1] + groups2_e[groups2]
        y += 0.5 * np.random.normal(size=n)

        df = pd.DataFrame(
            {"y": y, "TheGroups": groups, "groups1": groups1, "groups2": groups2}
        )

        model = gee.GEE.from_formula(
            "y ~ 1",
            groups="TheGroups",
            dep_data="0 + groups1 + groups2",
            cov_struct=cov_struct.Nested(),
            data=df,
        )
        result = model.fit()

