"""
Implementation of proportional hazards regression models for duration
data that may be censored ("Cox models").

References
----------
T Therneau (1996).  Extending the Cox model.  Technical report.
http://www.mayo.edu/research/documents/biostat-58pdf/DOC-10027288

G Rodriguez (2005).  Non-parametric estimation in survival models.
http://data.princeton.edu/pop509/NonParametricSurvival.pdf

B Gillespie (2006).  Checking the assumptions in the Cox proportional
hazards model.
http://www.mwsug.org/proceedings/2006/stats/MWSUG-2006-SD08.pdf
"""

from statsmodels.compat.pandas import Appender

import numpy as np

from statsmodels.base import model
import statsmodels.base.model as base
from statsmodels.formula.formulatools import advance_eval_env
from statsmodels.tools.decorators import cache_readonly
from statsmodels.tools.sm_exceptions import SpecificationWarning

_predict_docstring = """
    Returns predicted values from the proportional hazards
    regression model.

    Parameters
    ----------%(params_doc)s
    exog : array_like
        Data to use as `exog` in forming predictions.  If not
        provided, the `exog` values from the model used to fit the
        data are used.%(cov_params_doc)s
    endog : array_like
        Duration (time) values at which the predictions are made.
        Only used if pred_type is either 'cumhaz' or 'surv'.  If
        using model `exog`, defaults to model `endog` (time), but
        may be provided explicitly to make predictions at
        alternative times.
    strata : array_like
        A vector of stratum values used to form the predictions.
        Not used (may be 'None') if pred_type is 'lhr' or 'hr'.
        If `exog` is None, the model stratum values are used.  If
        `exog` is not None and pred_type is 'surv' or 'cumhaz',
        stratum values must be provided (unless there is only one
        stratum).
    offset : array_like
        Offset values used to create the predicted values.
    pred_type : str
        If 'lhr', returns log hazard ratios, if 'hr' returns
        hazard ratios, if 'surv' returns the survival function, if
        'cumhaz' returns the cumulative hazard function.
    pred_only : bool
        If True, returns only an array of predicted values.  Otherwise
        returns a bunch containing the predicted values and standard
        errors.

    Returns
    -------
    A bunch containing two fields: `predicted_values` and
    `standard_errors`.

    Notes
    -----
    Standard errors are only returned when predicting the log
    hazard ratio (pred_type is 'lhr').

    Types `surv` and `cumhaz` require estimation of the cumulative
    hazard function.
"""

_predict_params_doc = """
    params : array_like
        The proportional hazards model parameters."""

_predict_cov_params_docstring = """
    cov_params : array_like
        The covariance matrix of the estimated `params` vector,
        used to obtain prediction errors if pred_type='lhr',
        otherwise optional."""


class PHSurvivalTime:

    def __init__(self, time, status, exog, strata=None, entry=None, offset=None):
        """
        Represent a collection of survival times with possible
        stratification and left truncation.

        Parameters
        ----------
        time : array_like
            The times at which either the event (failure) occurs or
            the observation is censored.
        status : array_like
            Indicates whether the event (failure) occurs at `time`
            (`status` is 1), or if `time` is a censoring time (`status`
            is 0).
        exog : array_like
            The exogeneous (covariate) data matrix, cases are rows and
            variables are columns.
        strata : array_like
            Grouping variable defining the strata.  If None, all
            observations are in a single stratum.
        entry : array_like
            Entry (left truncation) times.  The observation is not
            part of the risk set for times before the entry time.  If
            None, the entry time is treated as being zero, which
            gives no left truncation.  The entry time must be less
            than or equal to `time`.
        offset : array_like
            An optional array of offsets
        """

        # Default strata
        if strata is None:
            strata = np.zeros(len(time), dtype=np.int32)

        # Default entry times
        if entry is None:
            entry = np.zeros(len(time))

        # Parameter validity checks.
        self._check(time, status, strata, entry)

        # Get the row indices for the cases in each stratum
        stu = np.unique(strata)
        sth = {x: [] for x in stu}
        for i, k in enumerate(strata):
            sth[k].append(i)
        stratum_rows = [np.asarray(sth[k], dtype=np.int32) for k in stu]
        stratum_names = stu

        # Remove strata with no events
        ix = [i for i, ix in enumerate(stratum_rows) if status[ix].sum() > 0]
        self.nstrat_orig = len(stratum_rows)
        stratum_rows = [stratum_rows[i] for i in ix]
        stratum_names = [stratum_names[i] for i in ix]

        # The number of strata
        nstrat = len(stratum_rows)
        self.nstrat = nstrat

        # Remove subjects whose entry time occurs after the last event
        # in their stratum.
        for stx, ix in enumerate(stratum_rows):
            last_failure = max(time[ix][status[ix] == 1])

            # Stata uses < here, R uses <=
            ii = [i for i, t in enumerate(entry[ix]) if t <= last_failure]
            stratum_rows[stx] = stratum_rows[stx][ii]

        # Remove subjects who are censored before the first event in
        # their stratum.
        for stx, ix in enumerate(stratum_rows):
            first_failure = min(time[ix][status[ix] == 1])

            ii = [i for i, t in enumerate(time[ix]) if t >= first_failure]
            stratum_rows[stx] = stratum_rows[stx][ii]

        # Order by time within each stratum
        for stx, ix in enumerate(stratum_rows):
            ii = np.argsort(time[ix])
            stratum_rows[stx] = stratum_rows[stx][ii]

        if offset is not None:
            self.offset_s = []
            for stx in range(nstrat):
                self.offset_s.append(offset[stratum_rows[stx]])
        else:
            self.offset_s = None

        # Number of informative subjects
        self.n_obs = sum([len(ix) for ix in stratum_rows])

        self.stratum_rows = stratum_rows
        self.stratum_names = stratum_names

        # Split everything by stratum
        self.time_s = self._split(time)
        self.exog_s = self._split(exog)
        self.status_s = self._split(status)
        self.entry_s = self._split(entry)

        # Precalculate some indices needed to fit Cox models.
        # Distinct failure times within a stratum are always taken to
        # be sorted in ascending order.
        #
        # ufailt_ix[stx][k] is a list of indices for subjects who fail
        # at the k^th sorted unique failure time in stratum stx
        #
        # risk_enter[stx][k] is a list of indices for subjects who
        # enter the risk set at the k^th sorted unique failure time in
        # stratum stx
        #
        # risk_exit[stx][k] is a list of indices for subjects who exit
        # the risk set at the k^th sorted unique failure time in
        # stratum stx
        self.ufailt_ix, self.risk_enter, self.risk_exit, self.ufailt = [], [], [], []

        for stx in range(self.nstrat):

            # All failure times
            ift = np.flatnonzero(self.status_s[stx] == 1)
            ft = self.time_s[stx][ift]

            # Unique failure times
            uft = np.unique(ft)
            nuft = len(uft)

            # Indices of cases that fail at each unique failure time
            # uft_map = {x:i for i,x in enumerate(uft)} # requires >=2.7
            uft_map = {x: i for i, x in enumerate(uft)}  # 2.6
            uft_ix = [[] for k in range(nuft)]
            for ix, ti in zip(ift, ft):
                uft_ix[uft_map[ti]].append(ix)

            # Indices of cases (failed or censored) that enter the
            # risk set at each unique failure time.
            risk_enter1 = [[] for k in range(nuft)]
            for i, t in enumerate(self.time_s[stx]):
                ix = np.searchsorted(uft, t, "right") - 1
                if ix >= 0:
                    risk_enter1[ix].append(i)

            # Indices of cases (failed or censored) that exit the
            # risk set at each unique failure time.
            risk_exit1 = [[] for k in range(nuft)]
            for i, t in enumerate(self.entry_s[stx]):
                ix = np.searchsorted(uft, t)
                risk_exit1[ix].append(i)

            self.ufailt.append(uft)
            self.ufailt_ix.append([np.asarray(x, dtype=np.int32) for x in uft_ix])
            self.risk_enter.append([np.asarray(x, dtype=np.int32) for x in risk_enter1])
            self.risk_exit.append([np.asarray(x, dtype=np.int32) for x in risk_exit1])

    def _split(self, x):
        v = []
        if x.ndim == 1:
            for ix in self.stratum_rows:
                v.append(x[ix])
        else:
            for ix in self.stratum_rows:
                v.append(x[ix, :])
        return v

    def _check(self, time, status, strata, entry):
        n1, n2, n3, n4 = len(time), len(status), len(strata), len(entry)
        nv = [n1, n2, n3, n4]
        if max(nv) != min(nv):
            raise ValueError(
                "endog, status, strata, and " + "entry must all have the same length"
            )
        if min(time) < 0:
            raise ValueError("endog must be non-negative")
        if min(entry) < 0:
            raise ValueError("entry time must be non-negative")

        # In Stata, this is entry >= time, in R it is >.
        if np.any(entry > time):
            raise ValueError(
                "entry times may not occur " + "after event or censoring times"
            )


class PHReg(model.LikelihoodModel):
    """
    Cox Proportional Hazards Regression Model

    The Cox PH Model is for right censored data.

    Parameters
    ----------
    endog : array_like
        The observed times (event or censoring)
    exog : 2D array_like
        The covariates or exogeneous variables
    status : array_like
        The censoring status values; status=1 indicates that an
        event occurred (e.g. failure or death), status=0 indicates
        that the observation was right censored. If None, defaults
        to status=1 for all cases.
    entry : array_like
        The entry times, if left truncation occurs
    strata : array_like
        Stratum labels.  If None, all observations are taken to be
        in a single stratum.
    ties : str
        The method used to handle tied times, must be either 'breslow'
        or 'efron'.
    offset : array_like
        Array of offset values
    missing : str
        The method used to handle missing data

    Notes
    -----
    Proportional hazards regression models should not include an
    explicit or implicit intercept.  The effect of an intercept is
    not identified using the partial likelihood approach.

    `endog`, `event`, `strata`, `entry`, and the first dimension
    of `exog` all must have the same length
    """

    def __init__(
        self,
        endog,
        exog,
        status=None,
        entry=None,
        strata=None,
        offset=None,
        ties="breslow",
        missing="drop",
        **kwargs,
    ):

        # Default is no censoring
        if status is None:
            status = np.ones(len(endog))

        super().__init__(
            endog,
            exog,
            status=status,
            entry=entry,
            strata=strata,
            offset=offset,
            missing=missing,
            **kwargs,
        )

        # endog and exog are automatically converted, but these are
        # not
        if self.status is not None:
            self.status = np.asarray(self.status)
        if self.entry is not None:
            self.entry = np.asarray(self.entry)
        if self.strata is not None:
            self.strata = np.asarray(self.strata)
        if self.offset is not None:
            self.offset = np.asarray(self.offset)

        self.surv = PHSurvivalTime(
            self.endog, self.status, self.exog, self.strata, self.entry, self.offset
        )
        self.nobs = len(self.endog)
        self.groups = None

        # TODO: not used?
        self.missing = missing

        self.df_resid = float(self.exog.shape[0] - np.linalg.matrix_rank(self.exog))
        self.df_model = float(np.linalg.matrix_rank(self.exog))

        ties = ties.lower()
        if ties not in ("efron", "breslow"):
            raise ValueError("`ties` must be either `efron` or " + "`breslow`")

        self.ties = ties

    @classmethod
    def from_formula(
        cls,
        formula,
        data,
        status=None,
        entry=None,
        strata=None,
        offset=None,
        subset=None,
        ties="breslow",
        missing="drop",
        *args,
        **kwargs,
    ):
        """
        Create a proportional hazards regression model from a formula
        and dataframe.

        Parameters
        ----------
        formula : str or generic Formula object
            The formula specifying the model
        data : array_like
            The data for the model. See Notes.
        status : array_like
            The censoring status values; status=1 indicates that an
            event occurred (e.g. failure or death), status=0 indicates
            that the observation was right censored. If None, defaults
            to status=1 for all cases.
        entry : array_like
            The entry times, if left truncation occurs
        strata : array_like
            Stratum labels.  If None, all observations are taken to be
            in a single stratum.
        offset : array_like
            Array of offset values
        subset : array_like
            An array-like object of booleans, integers, or index
            values that indicate the subset of df to use in the
            model. Assumes df is a `pandas.DataFrame`
        ties : str
            The method used to handle tied times, must be either 'breslow'
            or 'efron'.
        missing : str
            The method used to handle missing data
        args : extra arguments
            These are passed to the model
        kwargs : extra keyword arguments
            These are passed to the model with one exception. 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``.

        Returns
        -------
        model : PHReg model instance
        """

        # Allow array arguments to be passed by column name.
        if isinstance(status, str):
            status = data[status]
        if isinstance(entry, str):
            entry = data[entry]
        if isinstance(strata, str):
            strata = data[strata]
        if isinstance(offset, str):
            offset = data[offset]

        import re

        terms = re.split(r"[+\-~]", formula)
        for term in terms:
            term = term.strip()
            if term in ("0", "1"):
                import warnings

                warnings.warn(
                    "PHReg formulas should not include any '0' or '1' terms",
                    SpecificationWarning,
                    stacklevel=2,
                )
        advance_eval_env(kwargs)
        mod = super().from_formula(
            formula,
            data,
            *args,
            status=status,
            entry=entry,
            strata=strata,
            offset=offset,
            subset=subset,
            ties=ties,
            missing=missing,
            drop_cols=["Intercept"],
            **kwargs,
        )

        return mod

    def fit(self, groups=None, **args):
        """
        Fit a proportional hazards regression model.

        Parameters
        ----------
        groups : array_like
            Labels indicating groups of observations that may be
            dependent.  If present, the standard errors account for
            this dependence. Does not affect fitted values.

        Returns
        -------
        PHRegResults
            Returns a results instance.
        """

        # TODO process for missing values
        if groups is not None:
            if len(groups) != len(self.endog):
                msg = "len(groups) = %d and len(endog) = %d differ" % (
                    len(groups),
                    len(self.endog),
                )
                raise ValueError(msg)
            self.groups = np.asarray(groups)
        else:
            self.groups = None

        if "disp" not in args:
            args["disp"] = False

        fit_rslts = super().fit(**args)

        if self.groups is None:
            cov_params = fit_rslts.cov_params()
        else:
            cov_params = self.robust_covariance(fit_rslts.params)

        results = PHRegResults(self, fit_rslts.params, cov_params)

        return results

    def fit_regularized(
        self, method="elastic_net", alpha=0.0, start_params=None, refit=False, **kwargs
    ):
        r"""
        Return a regularized fit to a linear regression model.

        Parameters
        ----------
        method : {'elastic_net'}
            Only the `elastic_net` approach is currently implemented.
        alpha : scalar or array_like
            The penalty weight.  If a scalar, the same penalty weight
            applies to all variables in the model.  If a vector, it
            must have the same length as `params`, and contains a
            penalty weight for each coefficient.
        start_params : array_like
            Starting values for `params`.
        refit : bool
            If True, the model is refit using only the variables that
            have non-zero coefficients in the regularized fit.  The
            refitted model is not regularized.
        **kwargs
            Additional keyword arguments used to fit the model.

        Returns
        -------
        PHRegResults
            Returns a results instance.

        Notes
        -----
        The penalty is the ``elastic net`` penalty, which is a
        combination of L1 and L2 penalties.

        The function that is minimized is:

        .. math::

            -loglike/n + alpha*((1-L1\_wt)*|params|_2^2/2 + L1\_wt*|params|_1)

        where :math:`|*|_1` and :math:`|*|_2` are the L1 and L2 norms.

        Post-estimation results are based on the same data used to
        select variables, hence may be subject to overfitting biases.

        The elastic_net method uses the following keyword arguments:

        maxiter : int
            Maximum number of iterations
        L1_wt  : float
            Must be in [0, 1].  The L1 penalty has weight L1_wt and the
            L2 penalty has weight 1 - L1_wt.
        cnvrg_tol : float
            Convergence threshold for line searches
        zero_tol : float
            Coefficients below this threshold are treated as zero.
        """

        from statsmodels.base.elastic_net import fit_elasticnet

        if method != "elastic_net":
            raise ValueError("method for fit_regularized must be elastic_net")

        defaults = {"maxiter": 50, "L1_wt": 1, "cnvrg_tol": 1e-10, "zero_tol": 1e-10}
        defaults.update(kwargs)

        return fit_elasticnet(
            self,
            method=method,
            alpha=alpha,
            start_params=start_params,
            refit=refit,
            **defaults,
        )

    def loglike(self, params):
        """
        Returns the log partial likelihood function evaluated at
        `params`.
        """

        if self.ties == "breslow":
            return self.breslow_loglike(params)
        elif self.ties == "efron":
            return self.efron_loglike(params)

    def score(self, params):
        """
        Returns the score function evaluated at `params`.
        """

        if self.ties == "breslow":
            return self.breslow_gradient(params)
        elif self.ties == "efron":
            return self.efron_gradient(params)

    def hessian(self, params):
        """
        Returns the Hessian matrix of the log partial likelihood
        function evaluated at `params`.
        """

        if self.ties == "breslow":
            return self.breslow_hessian(params)
        else:
            return self.efron_hessian(params)

    def breslow_loglike(self, params):
        """
        Returns the value of the log partial likelihood function
        evaluated at `params`, using the Breslow method to handle tied
        times.
        """

        surv = self.surv

        like = 0.0

        # Loop over strata
        for stx in range(surv.nstrat):

            uft_ix = surv.ufailt_ix[stx]
            exog_s = surv.exog_s[stx]
            nuft = len(uft_ix)

            linpred = np.dot(exog_s, params)
            if surv.offset_s is not None:
                linpred += surv.offset_s[stx]
            linpred -= linpred.max()
            e_linpred = np.exp(linpred)

            xp0 = 0.0

            # Iterate backward through the unique failure times.
            for i in range(nuft)[::-1]:

                # Update for new cases entering the risk set.
                ix = surv.risk_enter[stx][i]
                xp0 += e_linpred[ix].sum()

                # Account for all cases that fail at this point.
                ix = uft_ix[i]
                like += (linpred[ix] - np.log(xp0)).sum()

                # Update for cases leaving the risk set.
                ix = surv.risk_exit[stx][i]
                xp0 -= e_linpred[ix].sum()

        return like

    def efron_loglike(self, params):
        """
        Returns the value of the log partial likelihood function
        evaluated at `params`, using the Efron method to handle tied
        times.
        """

        surv = self.surv

        like = 0.0

        # Loop over strata
        for stx in range(surv.nstrat):

            # exog and linear predictor for this stratum
            exog_s = surv.exog_s[stx]
            linpred = np.dot(exog_s, params)
            if surv.offset_s is not None:
                linpred += surv.offset_s[stx]
            linpred -= linpred.max()
            e_linpred = np.exp(linpred)

            xp0 = 0.0

            # Iterate backward through the unique failure times.
            uft_ix = surv.ufailt_ix[stx]
            nuft = len(uft_ix)
            for i in range(nuft)[::-1]:

                # Update for new cases entering the risk set.
                ix = surv.risk_enter[stx][i]
                xp0 += e_linpred[ix].sum()
                xp0f = e_linpred[uft_ix[i]].sum()

                # Account for all cases that fail at this point.
                ix = uft_ix[i]
                like += linpred[ix].sum()

                m = len(ix)
                J = np.arange(m, dtype=np.float64) / m
                like -= np.log(xp0 - J * xp0f).sum()

                # Update for cases leaving the risk set.
                ix = surv.risk_exit[stx][i]
                xp0 -= e_linpred[ix].sum()

        return like

    def breslow_gradient(self, params):
        """
        Returns the gradient of the log partial likelihood, using the
        Breslow method to handle tied times.
        """

        surv = self.surv

        grad = 0.0

        # Loop over strata
        for stx in range(surv.nstrat):

            # Indices of subjects in the stratum
            # strat_ix = surv.stratum_rows[stx]

            # Unique failure times in the stratum
            uft_ix = surv.ufailt_ix[stx]
            nuft = len(uft_ix)

            # exog and linear predictor for the stratum
            exog_s = surv.exog_s[stx]
            linpred = np.dot(exog_s, params)
            if surv.offset_s is not None:
                linpred += surv.offset_s[stx]
            linpred -= linpred.max()
            e_linpred = np.exp(linpred)

            xp0, xp1 = 0.0, 0.0

            # Iterate backward through the unique failure times.
            for i in range(nuft)[::-1]:

                # Update for new cases entering the risk set.
                ix = surv.risk_enter[stx][i]
                if len(ix) > 0:
                    v = exog_s[ix, :]
                    xp0 += e_linpred[ix].sum()
                    xp1 += (e_linpred[ix][:, None] * v).sum(0)

                # Account for all cases that fail at this point.
                ix = uft_ix[i]
                grad += (exog_s[ix, :] - xp1 / xp0).sum(0)

                # Update for cases leaving the risk set.
                ix = surv.risk_exit[stx][i]
                if len(ix) > 0:
                    v = exog_s[ix, :]
                    xp0 -= e_linpred[ix].sum()
                    xp1 -= (e_linpred[ix][:, None] * v).sum(0)

        return grad

    def efron_gradient(self, params):
        """
        Returns the gradient of the log partial likelihood evaluated
        at `params`, using the Efron method to handle tied times.
        """

        surv = self.surv

        grad = 0.0

        # Loop over strata
        for stx in range(surv.nstrat):

            # Indices of cases in the stratum
            # strat_ix = surv.stratum_rows[stx]

            # exog and linear predictor of the stratum
            exog_s = surv.exog_s[stx]
            linpred = np.dot(exog_s, params)
            if surv.offset_s is not None:
                linpred += surv.offset_s[stx]
            linpred -= linpred.max()
            e_linpred = np.exp(linpred)

            xp0, xp1 = 0.0, 0.0

            # Iterate backward through the unique failure times.
            uft_ix = surv.ufailt_ix[stx]
            nuft = len(uft_ix)
            for i in range(nuft)[::-1]:

                # Update for new cases entering the risk set.
                ix = surv.risk_enter[stx][i]
                if len(ix) > 0:
                    v = exog_s[ix, :]
                    xp0 += e_linpred[ix].sum()
                    xp1 += (e_linpred[ix][:, None] * v).sum(0)
                ixf = uft_ix[i]
                if len(ixf) > 0:
                    v = exog_s[ixf, :]
                    xp0f = e_linpred[ixf].sum()
                    xp1f = (e_linpred[ixf][:, None] * v).sum(0)

                    # Consider all cases that fail at this point.
                    grad += v.sum(0)

                    m = len(ixf)
                    J = np.arange(m, dtype=np.float64) / m
                    numer = xp1 - np.outer(J, xp1f)
                    denom = xp0 - np.outer(J, xp0f)
                    ratio = numer / denom
                    rsum = ratio.sum(0)
                    grad -= rsum

                # Update for cases leaving the risk set.
                ix = surv.risk_exit[stx][i]
                if len(ix) > 0:
                    v = exog_s[ix, :]
                    xp0 -= e_linpred[ix].sum()
                    xp1 -= (e_linpred[ix][:, None] * v).sum(0)

        return grad

    def breslow_hessian(self, params):
        """
        Returns the Hessian of the log partial likelihood evaluated at
        `params`, using the Breslow method to handle tied times.
        """

        surv = self.surv

        hess = 0.0

        # Loop over strata
        for stx in range(surv.nstrat):

            uft_ix = surv.ufailt_ix[stx]
            nuft = len(uft_ix)

            exog_s = surv.exog_s[stx]

            linpred = np.dot(exog_s, params)
            if surv.offset_s is not None:
                linpred += surv.offset_s[stx]
            linpred -= linpred.max()
            e_linpred = np.exp(linpred)

            xp0, xp1, xp2 = 0.0, 0.0, 0.0

            # Iterate backward through the unique failure times.
            for i in range(nuft)[::-1]:

                # Update for new cases entering the risk set.
                ix = surv.risk_enter[stx][i]
                if len(ix) > 0:
                    xp0 += e_linpred[ix].sum()
                    v = exog_s[ix, :]
                    xp1 += (e_linpred[ix][:, None] * v).sum(0)
                    elx = e_linpred[ix]
                    xp2 += np.einsum("ij,ik,i->jk", v, v, elx)

                # Account for all cases that fail at this point.
                m = len(uft_ix[i])
                hess += m * (xp2 / xp0 - np.outer(xp1, xp1) / xp0**2)

                # Update for new cases entering the risk set.
                ix = surv.risk_exit[stx][i]
                if len(ix) > 0:
                    xp0 -= e_linpred[ix].sum()
                    v = exog_s[ix, :]
                    xp1 -= (e_linpred[ix][:, None] * v).sum(0)
                    elx = e_linpred[ix]
                    xp2 -= np.einsum("ij,ik,i->jk", v, v, elx)
        return -hess

    def efron_hessian(self, params):
        """
        Returns the Hessian matrix of the partial log-likelihood
        evaluated at `params`, using the Efron method to handle tied
        times.
        """

        surv = self.surv

        hess = 0.0

        # Loop over strata
        for stx in range(surv.nstrat):

            exog_s = surv.exog_s[stx]

            linpred = np.dot(exog_s, params)
            if surv.offset_s is not None:
                linpred += surv.offset_s[stx]
            linpred -= linpred.max()
            e_linpred = np.exp(linpred)

            xp0, xp1, xp2 = 0.0, 0.0, 0.0

            # Iterate backward through the unique failure times.
            uft_ix = surv.ufailt_ix[stx]
            nuft = len(uft_ix)
            for i in range(nuft)[::-1]:

                # Update for new cases entering the risk set.
                ix = surv.risk_enter[stx][i]
                if len(ix) > 0:
                    xp0 += e_linpred[ix].sum()
                    v = exog_s[ix, :]
                    xp1 += (e_linpred[ix][:, None] * v).sum(0)
                    elx = e_linpred[ix]
                    xp2 += np.einsum("ij,ik,i->jk", v, v, elx)

                ixf = uft_ix[i]
                if len(ixf) > 0:
                    v = exog_s[ixf, :]
                    xp0f = e_linpred[ixf].sum()
                    xp1f = (e_linpred[ixf][:, None] * v).sum(0)
                    elx = e_linpred[ixf]
                    xp2f = np.einsum("ij,ik,i->jk", v, v, elx)

                # Account for all cases that fail at this point.
                m = len(uft_ix[i])
                J = np.arange(m, dtype=np.float64) / m
                c0 = xp0 - J * xp0f
                hess += xp2 * np.sum(1 / c0)
                hess -= xp2f * np.sum(J / c0)
                mat = (xp1[None, :] - np.outer(J, xp1f)) / c0[:, None]
                hess -= np.einsum("ij,ik->jk", mat, mat)

                # Update for new cases entering the risk set.
                ix = surv.risk_exit[stx][i]
                if len(ix) > 0:
                    xp0 -= e_linpred[ix].sum()
                    v = exog_s[ix, :]
                    xp1 -= (e_linpred[ix][:, None] * v).sum(0)
                    elx = e_linpred[ix]
                    xp2 -= np.einsum("ij,ik,i->jk", v, v, elx)

        return -hess

    def robust_covariance(self, params):
        """
        Returns a covariance matrix for the proportional hazards model
        regresion coefficient estimates that is robust to certain
        forms of model misspecification.

        Parameters
        ----------
        params : ndarray
            The parameter vector at which the covariance matrix is
            calculated.

        Returns
        -------
        The robust covariance matrix as a square ndarray.

        Notes
        -----
        This function uses the `groups` argument to determine groups
        within which observations may be dependent.  The covariance
        matrix is calculated using the Huber-White "sandwich" approach.
        """

        if self.groups is None:
            raise ValueError(
                "`groups` must be specified to calculate the robust covariance matrix"
            )

        hess = self.hessian(params)

        score_obs = self.score_residuals(params)

        # Collapse
        grads = {}
        for i, g in enumerate(self.groups):
            if g not in grads:
                grads[g] = 0.0
            grads[g] += score_obs[i, :]
        grads = np.asarray(list(grads.values()))

        mat = grads[None, :, :]
        mat = mat.T * mat
