"""
Heckman correction for sample selection bias (the Heckman procedure).

This is a modified version of the code available at
https://github.com/statsmodels/statsmodels/blob/92ea62232fd63c7b60c60bee4517ab3711d906e3/statsmodels/regression/Heckit.py

Created by Tetsu Haruyama
Last Modified 2 July 2019
"""

import numpy as np
import statsmodels.api as sm
import statsmodels.base.model as base
from statsmodels.iolib import summary
from statsmodels.tools.numdiff import approx_fprime
from scipy.stats import norm
import warnings
import pandas as pd
from statsmodels.tools.decorators import cache_readonly, cache_writable    # Tetsu
from statsmodels.compat.python import lrange

class Heckit(base.LikelihoodModel):
    """
    Class for Heckit correction for sample selection bias model.

    Parameters
    ----------
    endog : 1darray
        Data for the dependent variable. Should be set to np.nan for
        censored observations.
    exog : 2darray
        Data for the regression (response) equation. If a constant
        term is desired, the user should directly add the constant
        column to the data before using it as an argument here.
    exog_select : 2darray
        Data for the selection equation. If a constant
        term is desired, the user should directly add the constant
        column to the data before using it as an argument here.
    **kwargs:
        missing=, which can be 'none', 'drop', or 'raise'

    Notes
    -----
    The selection equation should contain at least one variable that
    is not in the regression (response) equation, i.e. the selection
    equation should contain at least one instrument. However, this
    module should still work if the user chooses not to do this.
    """

    def __init__(self, endog, exog, exog_select, **kwargs):

        # check that Z has same index as X (and consequently Y through super().__init__)
        if pd.__name__ in type(endog).__module__ and pd.__name__ in type(exog).__module__:
            if not all(endog.index==exog_select.index):
                raise ValueError("Z indices need to be the same as X and Y indices")


        # shape checks
        if (len(endog) == len(exog)) and (len(endog) == len(exog_select)):
            pass
        else:
            raise ValueError("Y, X, and Z data shapes do not conform with each other.")

        try:
            if (endog.ndim == 1) and (exog.ndim <= 2) and (exog_select.ndim <= 2):
                pass
            else:
                raise ValueError("Y, X, and Z data shapes do not conform with each other.")
        except AttributeError:
            if (np.asarray(endog).ndim == 1) and (np.asarray(exog).ndim <= 2) and (np.asarray(exog_select).ndim <= 2):
                pass
            else:
                raise ValueError("Y, X, and Z data shapes do not conform with each other.")

        # give missing (treated) values in endog variable finite values so that super().__init__
        # does not strip them out -- they will be put back after the call to super().__init__
        treated = np.asarray(~np.isnan(endog))

        try:
            endog_nomissing = endog.copy()
            endog_nomissing[~treated] = -99999
        except (TypeError, AttributeError):
            endog_nomissing = [endog[i] if treated[i] else -99999 for i in range(len(treated))]

        # create 1-D array that will be np.nan for every row of exog_select that has any missing
        # values and a finite value otherwise for the call to super().__init__ so that it can
        # strip out rows where exog_select has missing data if missing option is set

        if np.asarray(exog_select).ndim==2:
            exog_select_1dnan_placeholder = \
                [np.nan if any(np.isnan(row)) else 1 for row in np.asarray(exog_select)]
        else:  # assume ==1
            exog_select_1dnan_placeholder = [np.nan if np.isnan(row) else 1 for row in np.asarray(exog_select)]

        if pd.__name__ in type(endog).__module__:
            exog_select_1dnan_placeholder = pd.Series(exog_select_1dnan_placeholder, index=endog.index)
        else:
            exog_select_1dnan_placeholder = np.array(exog_select_1dnan_placeholder)

        # create array of sequential row positions so that rows of exog_select that have missing
        # data can be identified after call to super().__init__
        obsno = np.array(list(range(len(endog))))

        # call super().__init__
        super(Heckit, self).__init__(
            endog_nomissing, exog=exog,
            exog_select_1dnan_placeholder=exog_select_1dnan_placeholder, obsno=obsno,
            treated=treated,
            **kwargs)

        # put np.nan back into endog for treated rows
        self.endog = self.data.endog = \
            np.asarray(
                [self.endog[i] if self.treated[i] else np.nan for i in range(len(self.treated))]
                )

        # strip out rows stripped out by call to super().__init__ in Z variable
        self.exog_select = np.asarray([np.asarray(exog_select)[obs] for obs in self.obsno])

        # store variable names of exog_select
        try:
            self.exog_select_names = exog_select.columns.tolist()
        except AttributeError:
            self.exog_select_names = None

        # delete attributes created by the call to super().__init__ that are no longer needed
        del self.exog_select_1dnan_placeholder
        del self.obsno


        # store observation counts
        self.nobs_total = len(endog)
        self.nobs_uncensored = self.nobs = np.sum(self.treated)
        self.nobs_censored = self.nobs_total - self.nobs_uncensored


    def initialize(self):
        self.wendog = self.endog
        self.wexog = self.exog
        self._df_model = None   # Tetsu
        self.rank = None    # Tetsu
        self._df_resid = None # Tetsu


    def whiten(self, data):
        """
        Model whitener for Heckit correction model does nothing.
        """
        return data


    def get_datamats(self):
        Y = np.asarray(self.endog)
        Y = Y[self.treated]

        X = np.asarray(self.exog)
        X = X[self.treated,:]
        if X.ndim==1:
            X = np.atleast_2d(X).T

        Z = np.asarray(self.exog_select)
        if Z.ndim==1:
            Z = np.atleast_2d(Z).T

        return Y, X, Z


    # Tetsu ------ starts -----
    @property
    def df_model(self):
        """
        The model degree of freedom, defined as the rank of the regressor
        matrix minus 1 if a constant is included.
        """
        X = np.asarray(self.exog)
        X = X[self.treated,:]
        if self._df_model is None:
            if self.rank is None:
                self.rank = np.linalg.matrix_rank(X)
                self._df_model = float(self.rank - self.k_constant)
        return self._df_model


    @df_model.setter
    def df_model(self, value):
        self._df_model = value

    @property
    def df_resid(self):
        """
        The residual degree of freedom, defined as the number of observations
        minus the rank of the regressor matrix.
        """
        X = np.asarray(self.exog)
        X = X[self.treated,:]

        if self._df_resid is None:
            if self.rank is None:
                self.rank = np.linalg.matrix_rank(X)
            self._df_resid = self.nobs - self.rank
        return self._df_resid

    @df_resid.setter
    def df_resid(self, value):
        self._df_resid = value

    # Tetsu ----- ends -----





    def fit(self, method='twostep',
            cov_type_1='nonrobust', cov_type_2='nonrobust',  #  Tetsu ################
             start_params_mle=None, method_mle=None, maxiter_mle=None, **kwargs_mle):
        """
        Fit the Heckit selection model.

        Parameters
        ----------
        method : str
            Can only be "2step", which uses Heckit's two-step method.
        start_params_mle: 1darray
            If using MLE, starting parameters.
        method_mle: str
            If using MLE, the MLE estimation method.
        maxiter_mle: scalar
            If using MLE, the maximum number of iterations for MLE estimation.
        **kwargs_mle:
            Other arguments to pass to optimizer for MLE estimation.

        Returns
        -------
        A HeckitResults class instance.

        See Also
        ---------
        HeckitResults


        """

        ## Show warning to user if estimation is by two-step but MLE arguments were also provided
        if method=='twostep':
            if start_params_mle is not None or method_mle is not None or maxiter_mle is not None or \
                len(kwargs_mle.keys())>0:
                warnings.warn('The user chose to estimate the Heckit model by Two-Step,' + \
                    ' but MLE arguments were provided. Extraneous MLE arguments will be ignored.')

        results = self._fit_twostep(cov_type_1=cov_type_1,cov_type_2=cov_type_2)  # Tetsu : cov_type_1=cov_type_1,cov_type_2=cov_type_2
        return results


    def _fit_twostep(self,cov_type_1='nonrobust',cov_type_2='nonrobust'):  # Tetsu : cov_type_1=cov_type_1,cov_type_2=cov_type_2
        ########################################################################
        # PRIVATE METHOD
        # Fits using Heckit two-step from Heckit (1979).
        ########################################################################

        ## prep data
        Y, X, Z = self.get_datamats()

        ## Step 1
        step1model = sm.Probit(self.treated, Z)
        step1res = step1model.fit(disp=False,cov_type=cov_type_1)  # Tets cov_type=cov_type_1
        step1_fitted = np.atleast_2d(step1res.fittedvalues).T
        step1_varcov = step1res.cov_params()

        inverse_mills = norm.pdf(step1_fitted)/norm.cdf(step1_fitted)

        ## Step 2
        W = np.hstack((X, inverse_mills[self.treated] ) )
        step2model = sm.OLS(Y, W)
        step2res = step2model.fit(cov_type=cov_type_2)    # Tetsu : cov_type_2=cov_type

        params = step2res.params[:-1]
        betaHat_inverse_mills = step2res.params[-1]


        ## Compute standard errors
        # Compute estimated error variance of censored regression
        delta = np.multiply(inverse_mills, inverse_mills + step1_fitted)[self.treated]

        sigma2Hat = step2res.resid.dot(step2res.resid) / self.nobs_uncensored + \
            (betaHat_inverse_mills**2 * sum(delta)) / self.nobs_uncensored
        sigma2Hat = sigma2Hat[0]
        sigmaHat = np.sqrt(sigma2Hat)
        rhoHat = betaHat_inverse_mills / sigmaHat

        # compute standard errors of beta estimates of censored regression
        delta_1d = delta.T[0]

        Q = rhoHat**2 * ((W.T*delta_1d).dot(Z[self.treated])).dot(step1_varcov).dot((Z[self.treated].T*delta_1d).dot(W))

        WT_W_inv = np.linalg.inv(W.T.dot(W))
        WT_R = W.T*(1 - rhoHat**2 * delta_1d)
        normalized_varcov_all = WT_W_inv.dot(WT_R.dot(W)+Q).dot(WT_W_inv)
        del WT_W_inv
        del WT_R

        del delta_1d

        normalized_varcov = normalized_varcov_all[:-1,:-1]

        varcov_all = sigma2Hat * normalized_varcov_all
        varcov = varcov_all[:-1,:-1]

        stderr_all = np.sqrt(np.diag(varcov_all))
        stderr = stderr_all[:-1]
        stderr_betaHat_inverse_mills = stderr_all[-1]


        ## store results
        results = HeckitResults(self, params, normalized_varcov, sigma2Hat,
            select_res=step1res,
            params_inverse_mills=betaHat_inverse_mills, stderr_inverse_mills=stderr_betaHat_inverse_mills,
            var_reg_error=sigma2Hat, corr_eqnerrors=rhoHat,
            method='twostep', cov_type_1=cov_type_1, cov_type_2=cov_type_2)  #  cov_type_1=cov_type_1, cov_type_2=cov_type_2

        return results


    def predict(self, params, exog=None):
        """
        Return linear predicted values from a design matrix.

        Parameters
        ----------
        exog : array-like
            Design / exogenous data
        params : array-like, optional after fit has been called
            Parameters of a linear model

        Returns
        -------
        An array of fitted values

        Notes
        -----
        If the model has not yet been fit, params is not optional.
        """
        if exog is None:
            exog = self.exog
        return np.dot(exog, params)

        if self._results is None and params is None:
            raise ValueError("If the model has not been fit, then you must specify the params argument.")
        if self._results is not None:
            return np.dot(exog, self._results.params)
        else:
            return np.dot(exog, params)


class HeckitResults(base.LikelihoodModelResults):
    """
    Class to represent results/fits for Heckit model.

    Attributes
    ----------
    select_res : ProbitResult object
        The ProbitResult object created when estimating the selection equation.
    params_inverse_mills : scalar
        Parameter estimate of the coef on the inverse Mills term in the second step.
    stderr_inverse_mills : scalar
        Standard error of the parameter estimate of the coef on the inverse Mills
        term in the second step.
    var_reg_error : scalar
        Estimate of the "sigma" term, i.e. the error variance estimate of the
        regression (response) equation
    corr_eqnerrors : scalar
        Estimate of the "rho" term, i.e. the correlation estimate of the errors between the
        regression (response) equation and the selection equation
    method : string
        The method used to produce the estimates, i.e. 'twostep', 'mle'
    """

    def __init__(self, model, params, normalized_cov_params=None, scale=1.,
        select_res=None,
        params_inverse_mills=None, stderr_inverse_mills=None,
        var_reg_error=None, corr_eqnerrors=None,
        method=None, cov_type_1=None, cov_type_2=None):  # Tetsu : cov_type_1=None, cov_type_2=None

        super(HeckitResults, self).__init__(model, params,
                                                normalized_cov_params,
                                                scale)

        self.select_res = select_res
        self.params_inverse_mills = params_inverse_mills
        self.stderr_inverse_mills = stderr_inverse_mills
        self.var_reg_error = var_reg_error
        self.corr_eqnerrors = corr_eqnerrors
        self.method = method


    # Tetsu ----- starts ------
        self.cov_type_1 = cov_type_1
        self.cov_type_2 = cov_type_2
        self.df_model = model.df_model
        self.df_resid = model.df_resid
        dropnan = np.asarray(~np.isnan(self.model.endog))
        self.endog = self.model.endog[dropnan]
        self.exog = self.model.exog[dropnan]
        self.wendog = self.endog
        self.wexog = self.exog
        self.pinv_wexog = np.linalg.pinv(self.exog)

        if not hasattr(self, 'use_t'):
            self.use_t = False
            # self.use_t = True

        if not hasattr(self.select_res, 'use_t'):
            self.select_res.use_t = False

    @cache_readonly
    def nobs(self):
        """Number of observations n."""
        # return float(self.model.nobs_total)
        return float(len(self.endog))

    @cache_readonly
    def fittedvalues(self):
        """The predicted values for the original (unwhitened) design."""
        # return self.model.predict(self.params, self.model.exog)
        return self.model.predict(self.params, self.exog)

    @cache_readonly
    def resid(self):
        """The residuals of the model."""
        return self.endog - self.model.predict(self.params, self.exog)

    @cache_readonly
    def wresid(self):
        """
        The residuals of the transformed/whitened regressand and regressor(s)
        """
        return self.wendog - self.model.predict(self.params, self.wexog)

    @cache_readonly
    def ssr(self):
        """Sum of squared (whitened) residuals."""
        wresid = self.wresid
        return np.dot(wresid, wresid)

    @cache_readonly
    def centered_tss(self):
        """The total (weighted) sum of squares centered about the mean."""
        model = self.model
        weights = getattr(model, 'weights', None)
        sigma = getattr(model, 'sigma', None)
        if weights is not None:
            # mean = np.average(model.endog, weights=weights)
            mean = np.average(self.endog, weights=weights)
            return np.sum(weights * (self.endog - mean)**2)

        elif sigma is not None:
            # Exactly matches WLS when sigma is diagonal
            iota = np.ones_like(self.endog)
            iota = model.whiten(iota)
            mean = self.wendog.dot(iota) / iota.dot(iota)
            err = self.endog - mean
            err = model.whiten(err)
            return np.sum(err**2)
        else:
            centered_endog = self.wendog - self.wendog.mean()
            return np.dot(centered_endog, centered_endog)

    @cache_readonly
    def uncentered_tss(self):
        """
        Uncentered sum of squares.  Sum of the squared values of the
        (whitened) endogenous response variable.
        """
        wendog = self.wendog
        return np.dot(wendog, wendog)

    @cache_readonly
    def ess(self):
        """Explained sum of squares. If a constant is present, the centered
        total sum of squares minus the sum of squared residuals. If there is no
        constant, the uncentered total sum of squares is used."""

        if self.k_constant:
            return self.centered_tss - self.ssr
        else:
            return self.uncentered_tss - self.ssr

    @cache_readonly
    def rsquared(self):
        """
        R-squared of a model with an intercept.  This is defined here
        as 1 - `ssr`/`centered_tss` if the constant is included in the
        model and 1 - `ssr`/`uncentered_tss` if the constant is
        omitted.
        """
        if self.k_constant:
            return 1 - self.ssr/self.centered_tss
        else:
            return 1 - self.ssr/self.uncentered_tss

    @cache_readonly
    def rsquared_adj(self):
        """
        Adjusted R-squared.  This is defined here as 1 -
        (`nobs`-1)/`df_resid` * (1-`rsquared`) if a constant is
        included and 1 - `nobs`/`df_resid` * (1-`rsquared`) if no
        constant is included.
        """
        return 1 - (np.divide(self.nobs - self.k_constant, self.df_resid)
                    * (1 - self.rsquared))

    @cache_readonly
    def mse_model(self):
        """
        Mean squared error the model. This is the explained sum of
        squares divided by the model degrees of freedom.
        """
        return self.ess/self.df_model


    @cache_readonly
    def mse_resid(self):
        """
        Mean squared error of the residuals.  The sum of squared
        residuals divided by the residual degrees of freedom.
        """
        return self.ssr/self.df_resid


    @cache_readonly
    def mse_total(self):
        """
        Total mean squared error.  Defined as the uncentered total sum
        of squares divided by n the number of observations.
        """
        if self.k_constant:
            return self.centered_tss / (self.df_resid + self.df_model)
        else:
            return self.uncentered_tss / (self.df_resid + self.df_model)


    @cache_readonly
    def fvalue(self):
        """F-statistic of the fully specified model.  Calculated as the mean
        squared error of the model divided by the mean squared error of the
        residuals."""
        if hasattr(self, 'cov_type_1') and self.cov_type_1 != 'nonrobust':
            # with heteroscedasticity or correlation robustness
            k_params = self.normalized_cov_params.shape[0]
            mat = np.eye(k_params)
            const_idx = self.model.data.const_idx
            # TODO: What if model includes implicit constant, e.g. all
            #       dummies but no constant regressor?
            # TODO: Restats as LM test by projecting orthogonalizing
            #       to constant?
            if self.model.data.k_constant == 1:
                # if constant is implicit, return nan see #2444
                if const_idx is None:
                    return np.nan

                idx = lrange(k_params)
                idx.pop(const_idx)
                mat = mat[idx]  # remove constant
                if mat.size == 0:  # see  #3642
                    return np.nan
            ft = self.f_test(mat)
            # using backdoor to set another attribute that we already have
            self._cache['f_pvalue'] = ft.pvalue
            return ft.fvalue
        else:
            # for standard homoscedastic case
            return self.mse_model/self.mse_resid


    @cache_readonly
    def f_pvalue(self):
        """p-value of the F-statistic"""
        from scipy import stats
        return stats.f.sf(self.fvalue, self.df_model, self.df_resid)


    def _HCCM(self, scale):
        H = np.dot(self.pinv_wexog,
                   scale[:, None] * self.pinv_wexog.T)
        return H

    @cache_readonly
    def cov_HC1(self):
        """
        See statsmodels.RegressionResults
        """

        self.het_scale = self.nobs/(self.df_resid)*(self.wresid**2)
        cov_HC1 = self._HCCM(self.het_scale)
        return cov_HC1

    @cache_readonly
    def HC1_se(self):
        """
        See statsmodels.RegressionResults
        """
        return np.sqrt(np.diag(self.cov_HC1))

    # Tetsu ----- ends -----



    def summary(self, yname=None, xname=None, zname=None, title=None, alpha=.05):
        """Summarize the Heckit model Results

        Parameters
        -----------
        yname : string, optional
            Default is `y`
        xname : list of strings, optional
            Default is `x_##` for ## in p the number of regressors
            in the regression (response) equation.
        zname : list of strings, optional
            Default is `z_##` for ## in p the number of regressors
            in the selection equation.
        title : string, optional
            Title for the top table. If not None, then this replaces the
            default title
        alpha : float
            significance level for the confidence intervals

        Returns
        -------
        smry : Summary instance
            this holds the summary tables and text, which can be printed or
            converted to various output formats.

        See Also
        --------
        statsmodels.iolib.summary.Summary : class to hold summary
            results

        """

        ## Put in Z name detected from data if none supplied, unless that too could not be
        ## inferred from data, then put in generic names
        if zname is None and self.model.exog_select_names is not None:
            zname=self.model.exog_select_names
        elif zname is None and self.model.exog_select_names is None:
            try:
                zname = ['z' + str(i) for i in range(len(self.model.exog_select[0]))]
                zname[0]  = 'z0_or_zconst'
            except TypeError:
                zname = 'z0_or_zconst'

        try:  # for Python 3
            if isinstance(zname, str):
                zname = [zname]
        except NameError:  # for Python 2
            if isinstance(zname, basestring):
                zname = [zname]


        ## create summary object
        # instantiate the object
        smry = summary.Summary()

        top_left = [('Dep. Variable:', None),
                    ('Model:', None),
                    ('Method:', ['Heckman Two-Step']),
                    ('Date:', None),
                    ('Time:', None),
                    ('No. Total Obs.:', ["%#i" % self.model.nobs_total]),
                    ('No. Censored Obs.:', ["%#i" % self.model.nobs_censored]),
                    ('No. Uncensored Obs.:', ["%#i" % self.model.nobs_uncensored]),
                    ]


        top_right = [('R-squared:', ["%#8.3f" % self.rsquared]),
                     ('Adj. R-squared:', ["%#8.3f" % self.rsquared_adj]),
                     ('F-statistics:', ["%#8.3f" % self.fvalue]),
                     ('Prob (F-statistic):', ["%#8.3f" % self.f_pvalue]),  # Tetsu added
                     ('Cov in 1st Stage:', [self.cov_type_1]),  # Tetsu added
                     ('Cov in 2nd Stage:', [self.cov_type_2]),  # Tetsu added
                     ]

        if title is None:
            title = self.model.__class__.__name__ + ' ' + "Regression Results"

        smry.add_table_2cols(self, gleft=top_left, gright=top_right,
                          yname=yname, xname=xname, title=title)

        # add the Heckit-corrected regression table
        smry.add_table_params(self, yname=yname, xname=xname, alpha=alpha,
                             use_t=self.use_t)

        # add the selection equation estimates table
        smry.add_table_params(self.select_res, yname=yname, xname=zname, alpha=alpha,
                             use_t=self.select_res.use_t)

        # add the estimate to the inverse Mills estimate (z-score)
        smry.add_table_params(
            base.LikelihoodModelResults(None, np.atleast_1d(self.params_inverse_mills),
            normalized_cov_params=np.atleast_1d(self.stderr_inverse_mills**2), scale=1.),
            yname=None, xname=['IMR (Lambda)'], alpha=alpha,
            use_t=False)

        # add point estimates for rho and sigma
        diagn_left = [('rho:', ["%#6.3f" % self.corr_eqnerrors]),
                      ('sigma:', ["%#6.3f" % np.sqrt(self.var_reg_error)]),
                      ]

        diagn_right = [
                       ]

        smry.add_table_2cols(self, gleft=diagn_left, gright=diagn_right,
                          yname=yname, xname=xname,
                          title="")

        # add text at end
        smry.add_extra_txt(['First table are the estimates for the regression (response) equation.',
            'Second table are the estimates for the selection equation.',
            'Third table is the estimate for the coef of the inverse Mills ratio (Heckman\'s Lambda).'])

        ## return
        return smry
