"""
Procedures for fitting marginal regression models to dependent data
using Generalized Estimating Equations.

References
----------
KY Liang and S Zeger. "Longitudinal data analysis using
generalized linear models". Biometrika (1986) 73 (1): 13-22.

S Zeger and KY Liang. "Longitudinal Data Analysis for Discrete and
Continuous Outcomes". Biometrics Vol. 42, No. 1 (Mar., 1986),
pp. 121-130

A Rotnitzky and NP Jewell (1990). "Hypothesis testing of regression
parameters in semiparametric generalized linear models for cluster
correlated data", Biometrika, 77, 485-497.

Xu Guo and Wei Pan (2002). "Small sample performance of the score
test in GEE".
http://www.sph.umn.edu/faculty1/wp-content/uploads/2012/11/rr2002-013.pdf

LA Mancl LA, TA DeRouen (2001). A covariance estimator for GEE with
improved small-sample properties.  Biometrics. 2001 Mar;57(1):126-34.
"""

from statsmodels.compat.pandas import Appender
from statsmodels.compat.python import lzip

from collections import defaultdict
import warnings

import numpy as np
import pandas as pd
from scipy import stats

import statsmodels.base.model as base
import statsmodels.base.wrapper as wrap
from statsmodels.discrete.discrete_margins import (
    _check_at_is_all,
    _check_discrete_args,
    _check_margeff_args,
    _effects_at,
    _get_count_index,
    _get_dummy_index,
    _get_margeff_exog,
    _transform_names,
    margeff_cov_with_se,
)
from statsmodels.formula._manager import FormulaManager
from statsmodels.formula.formulatools import advance_eval_env
from statsmodels.genmod import cov_struct as cov_structs, families
from statsmodels.genmod.families import varfuncs
from statsmodels.genmod.families.links import Link
from statsmodels.genmod.generalized_linear_model import GLM, GLMResults
from statsmodels.graphics._regressionplots_doc import (
    _plot_added_variable_doc,
    _plot_ceres_residuals_doc,
    _plot_partial_residuals_doc,
)

# used for wrapper:
import statsmodels.regression.linear_model as lm
from statsmodels.tools.decorators import cache_readonly
from statsmodels.tools.sm_exceptions import (
    ConvergenceWarning,
    DomainWarning,
    IterationLimitWarning,
    ModelWarning,
    SingularMatrixWarning,
    SpecificationWarning,
    ValueWarning,
)


class ParameterConstraint:
    """
    A class for managing linear equality constraints for a parameter
    vector.
    """

    def __init__(self, lhs, rhs, exog):
        """
        Parameters
        ----------
        lhs : ndarray
           A q x p matrix which is the left hand side of the
           constraint lhs * param = rhs.  The number of constraints is
           q >= 1 and p is the dimension of the parameter vector.
        rhs : ndarray
          A 1-dimensional vector of length q which is the right hand
          side of the constraint equation.
        exog : ndarray
          The n x p exognenous data for the full model.
        """

        # In case a row or column vector is passed (patsy linear
        # constraints passes a column vector).
        rhs = np.atleast_1d(rhs.squeeze())

        if rhs.ndim > 1:
            raise ValueError(
                "The right hand side of the constraint must be a vector."
            )

        if len(rhs) != lhs.shape[0]:
            raise ValueError(
                "The number of rows of the left hand "
                "side constraint matrix L must equal "
                "the length of the right hand side "
                "constraint vector R."
            )

        self.lhs = lhs
        self.rhs = rhs

        # The columns of lhs0 are an orthogonal basis for the
        # orthogonal complement to row(lhs), the columns of lhs1 are
        # an orthogonal basis for row(lhs).  The columns of lhsf =
        # [lhs0, lhs1] are mutually orthogonal.
        lhs_u, lhs_s, lhs_vt = np.linalg.svd(lhs.T, full_matrices=1)
        self.lhs0 = lhs_u[:, len(lhs_s) :]
        self.lhs1 = lhs_u[:, 0 : len(lhs_s)]
        self.lhsf = np.hstack((self.lhs0, self.lhs1))

        # param0 is one solution to the underdetermined system
        # L * param = R.
        self.param0 = np.dot(self.lhs1, np.dot(lhs_vt, self.rhs) / lhs_s)

        self._offset_increment = np.dot(exog, self.param0)

        self.orig_exog = exog
        self.exog_fulltrans = np.dot(exog, self.lhsf)

    def offset_increment(self):
        """
        Returns a vector that should be added to the offset vector to
        accommodate the constraint.

        Parameters
        ----------
        exog : array_like
           The exogeneous data for the model.
        """

        return self._offset_increment

    def reduced_exog(self):
        """
        Returns a linearly transformed exog matrix whose columns span
        the constrained model space.

        Parameters
        ----------
        exog : array_like
           The exogeneous data for the model.
        """
        return self.exog_fulltrans[:, 0 : self.lhs0.shape[1]]

    def restore_exog(self):
        """
        Returns the full exog matrix before it was reduced to
        satisfy the constraint.
        """
        return self.orig_exog

    def unpack_param(self, params):
        """
        Converts the parameter vector `params` from reduced to full
        coordinates.
        """

        return self.param0 + np.dot(self.lhs0, params)

    def unpack_cov(self, bcov):
        """
        Converts the covariance matrix `bcov` from reduced to full
        coordinates.
        """

        return np.dot(self.lhs0, np.dot(bcov, self.lhs0.T))


_gee_init_doc = """
    Marginal regression model fit using Generalized Estimating Equations.

    GEE can be used to fit Generalized Linear Models (GLMs) when the
    data have a grouped structure, and the observations are possibly
    correlated within groups but not between groups.

    Parameters
    ----------
    endog : array_like
        1d array of endogenous values (i.e. responses, outcomes,
        dependent variables, or 'Y' values).
    exog : array_like
        2d array of exogeneous values (i.e. covariates, predictors,
        independent variables, regressors, or 'X' values). A `nobs x
        k` array where `nobs` is the number of observations and `k` is
        the number of regressors. An intercept is not included by
        default and should be added by the user. See
        `statsmodels.tools.add_constant`.
    groups : array_like
        A 1d array of length `nobs` containing the group labels.
    time : array_like
        A 2d array of time (or other index) values, used by some
        dependence structures to define similarity relationships among
        observations within a cluster.
    family : family class instance
%(family_doc)s
    cov_struct : CovStruct class instance
        The default is Independence.  To specify an exchangeable
        structure use cov_struct = Exchangeable().  See
        statsmodels.genmod.cov_struct.CovStruct for more
        information.
    offset : array_like
        An offset to be included in the fit.  If provided, must be
        an array whose length is the number of rows in exog.
    dep_data : array_like
        Additional data passed to the dependence structure.
    constraint : (ndarray, ndarray)
        If provided, the constraint is a tuple (L, R) such that the
        model parameters are estimated under the constraint L *
        param = R, where L is a q x p matrix and R is a
        q-dimensional vector.  If constraint is provided, a score
        test is performed to compare the constrained model to the
        unconstrained model.
    update_dep : bool
        If true, the dependence parameters are optimized, otherwise
        they are held fixed at their starting values.
    weights : array_like
        An array of case weights to use in the analysis.
    %(extra_params)s

    See Also
    --------
    statsmodels.genmod.families.family
    :ref:`families`
    :ref:`links`

    Notes
    -----
    Only the following combinations make sense for family and link ::

                   + ident log logit probit cloglog pow opow nbinom loglog logc
      Gaussian     |   x    x                        x
      inv Gaussian |   x    x                        x
      binomial     |   x    x    x     x       x     x    x           x      x
      Poisson      |   x    x                        x
      neg binomial |   x    x                        x          x
      gamma        |   x    x                        x

    Not all of these link functions are currently available.

    Endog and exog are references so that if the data they refer
    to are already arrays and these arrays are changed, endog and
    exog will change.

    The "robust" covariance type is the standard "sandwich estimator"
    (e.g. Liang and Zeger (1986)).  It is the default here and in most
    other packages.  The "naive" estimator gives smaller standard
    errors, but is only correct if the working correlation structure
    is correctly specified.  The "bias reduced" estimator of Mancl and
    DeRouen (Biometrics, 2001) reduces the downward bias of the robust
    estimator.

    The robust covariance provided here follows Liang and Zeger (1986)
    and agrees with R's gee implementation.  To obtain the robust
    standard errors reported in Stata, multiply by sqrt(N / (N - g)),
    where N is the total sample size, and g is the average group size.
    %(notes)s
    Examples
    --------
    %(example)s
"""

_gee_nointercept = """
    The nominal and ordinal GEE models should not have an intercept
    (either implicit or explicit).  Use "0 + " in a formula to
    suppress the intercept.
"""

_gee_family_doc = """\
        The default is Gaussian.  To specify the binomial
        distribution use `family=sm.families.Binomial()`. Each family
        can take a link instance as an argument.  See
        statsmodels.genmod.families.family for more information."""

_gee_ordinal_family_doc = """\
        The only family supported is `Binomial`.  The default `Logit`
        link may be replaced with `probit` if desired."""

_gee_nominal_family_doc = """\
        The default value `None` uses a multinomial logit family
        specifically designed for use with GEE.  Setting this
        argument to a non-default value is not currently supported."""

_gee_fit_doc = """
    Fits a marginal regression model using generalized estimating
    equations (GEE).

    Parameters
    ----------
    maxiter : int
        The maximum number of iterations
    ctol : float
        The convergence criterion for stopping the Gauss-Seidel
        iterations
    start_params : array_like
        A vector of starting values for the regression
        coefficients.  If None, a default is chosen.
    params_niter : int
        The number of Gauss-Seidel updates of the mean structure
        parameters that take place prior to each update of the
        dependence structure.
    first_dep_update : int
        No dependence structure updates occur before this
        iteration number.
    cov_type : str
        One of "robust", "naive", or "bias_reduced".
    ddof_scale : scalar or None
        The scale parameter is estimated as the sum of squared
        Pearson residuals divided by `N - ddof_scale`, where N
        is the total sample size.  If `ddof_scale` is None, the
        number of covariates (including an intercept if present)
        is used.
    scaling_factor : scalar
        The estimated covariance of the parameter estimates is
        scaled by this value.  Default is 1, Stata uses N / (N - g),
        where N is the total sample size and g is the average group
        size.
    scale : str or float, optional
        `scale` can be None, 'X2', or a float
        If a float, its value is used as the scale parameter.
        The default value is None, which uses `X2` (Pearson's
        chi-square) for Gamma, Gaussian, and Inverse Gaussian.
        The default is 1 for the Binomial and Poisson families.

    Returns
    -------
    An instance of the GEEResults class or subclass

    Notes
    -----
    If convergence difficulties occur, increase the values of
    `first_dep_update` and/or `params_niter`.  Setting
    `first_dep_update` to a greater value (e.g. ~10-20) causes the
    algorithm to move close to the GLM solution before attempting
    to identify the dependence structure.

    For the Gaussian family, there is no benefit to setting
    `params_niter` to a value greater than 1, since the mean
    structure parameters converge in one step.
"""

_gee_results_doc = """
    Attributes
    ----------

    cov_params_default : ndarray
        default covariance of the parameter estimates. Is chosen among one
        of the following three based on `cov_type`
    cov_robust : ndarray
        covariance of the parameter estimates that is robust
    cov_naive : ndarray
        covariance of the parameter estimates that is not robust to
        correlation or variance misspecification
    cov_robust_bc : ndarray
        covariance of the parameter estimates that is robust and bias
        reduced
    converged : bool
        indicator for convergence of the optimization.
        True if the norm of the score is smaller than a threshold
    cov_type : str
        string indicating whether a "robust", "naive" or "bias_reduced"
        covariance is used as default
    fit_history : dict
        Contains information about the iterations.
    fittedvalues : ndarray
        Linear predicted values for the fitted model.
        dot(exog, params)
    model : class instance
        Pointer to GEE model instance that called `fit`.
    normalized_cov_params : ndarray
        See GEE docstring
    params : ndarray
        The coefficients of the fitted model.  Note that
        interpretation of the coefficients often depends on the
        distribution family and the data.
    scale : float
        The estimate of the scale / dispersion for the model fit.
        See GEE.fit for more information.
    score_norm : float
        norm of the score at the end of the iterative estimation.
    bse : ndarray
        The standard errors of the fitted GEE parameters.
"""

_gee_example = """
    Logistic regression with autoregressive working dependence:

    >>> import statsmodels.api as sm
    >>> family = sm.families.Binomial()
    >>> va = sm.cov_struct.Autoregressive()
    >>> model = sm.GEE(endog, exog, group, family=family, cov_struct=va)
    >>> result = model.fit()
    >>> print(result.summary())

    Use formulas to fit a Poisson GLM with independent working
    dependence:

    >>> import statsmodels.api as sm
    >>> fam = sm.families.Poisson()
    >>> ind = sm.cov_struct.Independence()
    >>> model = sm.GEE.from_formula("y ~ age + trt + base", "subject",
                                 data, cov_struct=ind, family=fam)
    >>> result = model.fit()
    >>> print(result.summary())

    Equivalent, using the formula API:

    >>> import statsmodels.api as sm
    >>> import statsmodels.formula.api as smf
    >>> fam = sm.families.Poisson()
    >>> ind = sm.cov_struct.Independence()
    >>> model = smf.gee("y ~ age + trt + base", "subject",
                    data, cov_struct=ind, family=fam)
    >>> result = model.fit()
    >>> print(result.summary())
"""

_gee_ordinal_example = """
    Fit an ordinal regression model using GEE, with "global
    odds ratio" dependence:

    >>> import statsmodels.api as sm
    >>> gor = sm.cov_struct.GlobalOddsRatio("ordinal")
    >>> model = sm.OrdinalGEE(endog, exog, groups, cov_struct=gor)
    >>> result = model.fit()
    >>> print(result.summary())

    Using formulas:

    >>> import statsmodels.formula.api as smf
    >>> model = smf.ordinal_gee("y ~ 0 + x1 + x2", groups, data,
                                    cov_struct=gor)
    >>> result = model.fit()
    >>> print(result.summary())
"""

_gee_nominal_example = """
    Fit a nominal regression model using GEE:

    >>> import statsmodels.api as sm
    >>> import statsmodels.formula.api as smf
    >>> gor = sm.cov_struct.GlobalOddsRatio("nominal")
    >>> model = sm.NominalGEE(endog, exog, groups, cov_struct=gor)
    >>> result = model.fit()
    >>> print(result.summary())

    Using formulas:

    >>> import statsmodels.api as sm
    >>> model = sm.NominalGEE.from_formula("y ~ 0 + x1 + x2", groups,
                     data, cov_struct=gor)
    >>> result = model.fit()
    >>> print(result.summary())

    Using the formula API:

    >>> import statsmodels.formula.api as smf
    >>> model = smf.nominal_gee("y ~ 0 + x1 + x2", groups, data,
                                cov_struct=gor)
    >>> result = model.fit()
    >>> print(result.summary())
"""


def _check_args(endog, exog, groups, time, offset, exposure):

    if endog.size != exog.shape[0]:
        raise ValueError(
            "Leading dimension of 'exog' should match length of 'endog'"
        )

    if groups.size != endog.size:
        raise ValueError("'groups' and 'endog' should have the same size")

    if time is not None and (time.size != endog.size):
        raise ValueError("'time' and 'endog' should have the same size")

    if offset is not None and (offset.size != endog.size):
        raise ValueError("'offset and 'endog' should have the same size")

    if exposure is not None and (exposure.size != endog.size):
        raise ValueError("'exposure' and 'endog' should have the same size")


class GEE(GLM):

    __doc__ = (
        "    Marginal Regression Model using Generalized Estimating "
        "Equations.\n"
        + _gee_init_doc
        % {
            "extra_params": base._missing_param_doc,
            "family_doc": _gee_family_doc,
            "example": _gee_example,
            "notes": "",
        }
    )

    cached_means = None

    def __init__(
        self,
        endog,
        exog,
        groups,
        time=None,
        family=None,
        cov_struct=None,
        missing="none",
        offset=None,
        exposure=None,
        dep_data=None,
        constraint=None,
        update_dep=True,
        weights=None,
        **kwargs,
    ):

        if type(self) is GEE:
            self._check_kwargs(kwargs)
        if family is not None:
            if not isinstance(family.link, tuple(family.safe_links)):
                msg = (
                    "The {0} link function does not respect the "
                    "domain of the {1} family."
                )
                warnings.warn(
                    msg.format(
                        family.link.__class__.__name__, family.__class__.__name__
                    ),
                    DomainWarning,
                    stacklevel=2,
                )

        groups = np.asarray(groups)  # in case groups is pandas

        if "missing_idx" in kwargs and kwargs["missing_idx"] is not None:
            # If here, we are entering from super.from_formula; missing
            # has already been dropped from endog and exog, but not from
            # the other variables.
            ii = ~kwargs["missing_idx"]
            groups = groups[ii]
            if time is not None:
                time = time[ii]
            if offset is not None:
                offset = offset[ii]
            if exposure is not None:
                exposure = exposure[ii]
            del kwargs["missing_idx"]

        self.missing = missing
        self.dep_data = dep_data
        self.constraint = constraint
        self.update_dep = update_dep

        self._fit_history = defaultdict(list)

        # Pass groups, time, offset, and dep_data so they are
        # processed for missing data along with endog and exog.
        # Calling super creates self.exog, self.endog, etc. as
        # ndarrays and the original exog, endog, etc. are
        # self.data.endog, etc.
        super().__init__(
            endog,
            exog,
            groups=groups,
            time=time,
            offset=offset,
            exposure=exposure,
            weights=weights,
            dep_data=dep_data,
            missing=missing,
            family=family,
            **kwargs,
        )

        _check_args(
            self.endog,
            self.exog,
            self.groups,
            self.time,
            getattr(self, "offset", None),
            getattr(self, "exposure", None),
        )

        self._init_keys.extend(["update_dep", "constraint", "family", "cov_struct"])
        # remove keys added by super that are not supported
        try:
            self._init_keys.remove("freq_weights")
            self._init_keys.remove("var_weights")
        except ValueError:
            pass

        # Handle the family argument
        if family is None:
            family = families.Gaussian()
        elif not issubclass(family.__class__, families.Family):
            raise ValueError("GEE: `family` must be a genmod family instance")
        self.family = family

        # Handle the cov_struct argument
        if cov_struct is None:
            cov_struct = cov_structs.Independence()
        elif not issubclass(cov_struct.__class__, cov_structs.CovStruct):
            raise ValueError(
                "GEE: `cov_struct` must be a genmod cov_struct instance"
            )

        self.cov_struct = cov_struct

        # Handle the constraint
        self.constraint = None
        if constraint is not None:
            if len(constraint) != 2:
                raise ValueError("GEE: `constraint` must be a 2-tuple.")
            if constraint[0].shape[1] != self.exog.shape[1]:
                raise ValueError(
                    "GEE: the left hand side of the constraint must have "
                    "the same number of columns as the exog matrix."
                )
            self.constraint = ParameterConstraint(
                constraint[0], constraint[1], self.exog
            )

            if self._offset_exposure is not None:
                self._offset_exposure += self.constraint.offset_increment()
            else:
                self._offset_exposure = self.constraint.offset_increment().copy()
            self.exog = self.constraint.reduced_exog()

        # Create list of row indices for each group
        group_labels, ix = np.unique(self.groups, return_inverse=True)
        se = pd.Series(index=np.arange(len(ix)), dtype="int")
        gb = se.groupby(ix).groups
        dk = [(lb, np.asarray(gb[k])) for k, lb in enumerate(group_labels)]
        self.group_indices = dict(dk)
        self.group_labels = group_labels

        # Convert the data to the internal representation, which is a
        # list of arrays, corresponding to the groups.
        self.endog_li = self.cluster_list(self.endog)
        self.exog_li = self.cluster_list(self.exog)

        if self.weights is not None:
            self.weights_li = self.cluster_list(self.weights)

        self.num_group = len(self.endog_li)

        # Time defaults to a 1d grid with equal spacing
        if self.time is not None:
            if self.time.ndim == 1:
                self.time = self.time[:, None]
            self.time_li = self.cluster_list(self.time)
        else:
            self.time_li = [
                np.arange(len(y), dtype=np.float64)[:, None] for y in self.endog_li
            ]
            self.time = np.concatenate(self.time_li)

        if self._offset_exposure is None or (
            np.isscalar(self._offset_exposure) and self._offset_exposure == 0.0
        ):
            self.offset_li = None
        else:
            self.offset_li = self.cluster_list(self._offset_exposure)
        if constraint is not None:
            self.constraint.exog_fulltrans_li = self.cluster_list(
                self.constraint.exog_fulltrans
            )

        self.family = family

        self.cov_struct.initialize(self)

        # Total sample size
        group_ns = [len(y) for y in self.endog_li]
        self.nobs = sum(group_ns)
        # The following are column based, not on rank see #1928
        self.df_model = self.exog.shape[1] - 1  # assumes constant
        self.df_resid = self.nobs - self.exog.shape[1]

        # Skip the covariance updates if all groups have a single
        # observation (reduces to fitting a GLM).
        maxgroup = max([len(x) for x in self.endog_li])
        if maxgroup == 1:
            self.update_dep = False

    # Override to allow groups and time to be passed as variable
    # names.
    @classmethod
    def from_formula(
        cls,
        formula,
        groups,
        data,
        subset=None,
        time=None,
        offset=None,
        exposure=None,
        *args,
        **kwargs,
    ):
        """
        Create a GEE model instance from a formula and dataframe.

        Parameters
        ----------
        formula : str or generic Formula object
            The formula specifying the model
        groups : array_like or string
            Array of grouping labels.  If a string, this is the name
            of a variable in `data` that contains the grouping labels.
        data : array_like
            The data for the model.
        subset : array_like
            An array-like object of booleans, integers, or index
            values that indicate the subset of the data to used when
            fitting the model.
        time : array_like or string
            The time values, used for dependence structures involving
            distances between observations.  If a string, this is the
            name of a variable in `data` that contains the time
            values.
        offset : array_like or string
            The offset values, added to the linear predictor.  If a
            string, this is the name of a variable in `data` that
            contains the offset values.
        exposure : array_like or string
            The exposure values, only used if the link function is the
            logarithm function, in which case the log of `exposure`
            is added to the offset (if any).  If a string, this is the
            name of a variable in `data` that contains the offset
            values.
        {missing_param_doc}
        args : extra arguments
            These are passed to the model
        kwargs : extra keyword arguments
            These are passed to the model with two exceptions. `dep_data`
            is processed as described below.  The ``eval_env`` keyword is
            passed to patsy. It can be either a
            :class:`patsy:patsy.EvalEnvironment` object or an integer
            indicating the depth of the namespace to use. For example, the
            default ``eval_env=0`` uses the calling namespace.
            If you wish to use a "clean" environment set ``eval_env=-1``.

        Optional arguments
        ------------------
        dep_data : str or array_like
            Data used for estimating the dependence structure.  See
            specific dependence structure classes (e.g. Nested) for
            details.  If `dep_data` is a string, it is interpreted as
            a formula that is applied to `data`. If it is an array, it
            must be an array of strings corresponding to column names in
            `data`.  Otherwise it must be an array-like with the same
            number of rows as data.

        Returns
        -------
        model : GEE model instance

        Notes
        -----
        `data` must define __getitem__ with the keys in the formula
        terms args and kwargs are passed on to the model
        instantiation. E.g., a numpy structured or rec array, a
        dictionary, or a pandas DataFrame.
        """.format(
            missing_param_doc=base._missing_param_doc
        )

        groups_name = "Groups"
        if isinstance(groups, str):
            groups_name = groups
            groups = data[groups]

        if isinstance(time, str):
            time = data[time]

        if isinstance(offset, str):
            offset = data[offset]

        if isinstance(exposure, str):
            exposure = data[exposure]

        dep_data = kwargs.get("dep_data")
        dep_data_names = None
        if dep_data is not None:
            if isinstance(dep_data, str):
                dep_data = FormulaManager().get_matrices(dep_data, data, pandas=True)
                dep_data_names = dep_data.columns.tolist()
            else:
                dep_data_names = list(dep_data)
                dep_data = data[dep_data]
            kwargs["dep_data"] = np.asarray(dep_data)

        family = None
        if "family" in kwargs:
            family = kwargs["family"]
            del kwargs["family"]

        advance_eval_env(kwargs)
        model = super().from_formula(
            formula,
            data,
            *args,
            subset=subset,
            groups=groups,
            time=time,
            offset=offset,
            exposure=exposure,
            family=family,
            **kwargs,
        )

        if dep_data_names is not None:
            model._dep_data_names = dep_data_names
        model._groups_name = groups_name

        return model

    def cluster_list(self, array):
        """
        Returns `array` split into subarrays corresponding to the
        cluster structure.
        """

        if array.ndim == 1:
            return [np.array(array[self.group_indices[k]]) for k in self.group_labels]
        else:
            return [
                np.array(array[self.group_indices[k], :]) for k in self.group_labels
            ]

    def compare_score_test(self, submodel):
        """
        Perform a score test for the given submodel against this model.

        Parameters
        ----------
        submodel : GEEResults instance
            A fitted GEE model that is a submodel of this model.

        Returns
        -------
        A dictionary with keys "statistic", "p-value", and "df",
        containing the score test statistic, its chi^2 p-value,
        and the degrees of freedom used to compute the p-value.

        Notes
        -----
        The score test can be performed without calling 'fit' on the
        larger model.  The provided submodel must be obtained from a
        fitted GEE.

        This method performs the same score test as can be obtained by
        fitting the GEE with a linear constraint and calling `score_test`
        on the results.

        References
        ----------
        Xu Guo and Wei Pan (2002). "Small sample performance of the score
        test in GEE".
        http://www.sph.umn.edu/faculty1/wp-content/uploads/2012/11/rr2002-013.pdf
        """

        # Since the model has not been fit, its scaletype has not been
        # set.  So give it the scaletype of the submodel.
        self.scaletype = submodel.model.scaletype

        # Check consistency between model and submodel (not a comprehensive
        # check)
        submod = submodel.model
        if self.exog.shape[0] != submod.exog.shape[0]:
            msg = "Model and submodel have different numbers of cases."
            raise ValueError(msg)
        if self.exog.shape[1] == submod.exog.shape[1]:
            msg = "Model and submodel have the same number of variables"
            warnings.warn(msg, SpecificationWarning, stacklevel=2)
        if not isinstance(self.family, type(submod.family)):
            msg = "Model and submodel have different GLM families."
            warnings.warn(msg, SpecificationWarning, stacklevel=2)
        if not isinstance(self.cov_struct, type(submod.cov_struct)):
            warnings.warn(
                "Model and submodel have different GEE covariance structures.",
                SpecificationWarning,
                stacklevel=2,
            )
        if not np.equal(self.weights, submod.weights).all():
            msg = "Model and submodel should have the same weights."
            warnings.warn(msg, SpecificationWarning, stacklevel=2)

        # Get the positions of the submodel variables in the
        # parent model
        qm, qc = _score_test_submodel(self, submodel.model)
        if qm is None:
            msg = "The provided model is not a submodel."
            raise ValueError(msg)

        # Embed the submodel params into a params vector for the
        # parent model
        params_ex = np.dot(qm, submodel.params)

        # Attempt to preserve the state of the parent model
        cov_struct_save = self.cov_struct
        import copy

        cached_means_save = copy.deepcopy(self.cached_means)

        # Get the score vector of the submodel params in
        # the parent model
        self.cov_struct = submodel.cov_struct
        self.update_cached_means(params_ex)
        _, score = self._update_mean_params()
        if score is None:
