import math
import warnings
import patsy
import numpy as np
import statsmodels.api as sm
import statsmodels.formula.api as smf
from scipy.stats import logistic, norm

from zepid.causal.ipw import propensity_score
from zepid.calc import probability_to_odds


class TMLE:
    def __init__(self, df, exposure, outcome, measure='risk_difference', alpha = 0.05):
        """Implementation of a single time-point target maximum likelihood estimator. It uses standard logistic
        regression models to calculate Psi. The TMLE estimator allows a standard logistic regression to be used.
        Alternatively, users are able to directly input predicted outcomes from other methods (like machine learning
        algorithms).

        TMLE is a double robust substitution estimator. For further details, see the listed references. The following
        is a general description of the estimation process for TMLE

        1. Initial estimates for Y are predicted from a regression model. Expected values for each individual are
        generated under the scenarios of all treated vs all untreated

        .. math::

            E(Y|A, L)

        2. Predicted probabilities are generated from a regression model

        .. math::

            \Pr(A=1|L)

        3. The predicted Y is merged together with the IPW using

        .. math::

            H_a(A=a,L) = \frac{I(A=1)}{\pi_1} - \frac{I(A=0)}{\pi_0}

        for each individual. Afterwards, the predicted Y is set as an offset in the following logit model and used to
        predict values under each treatment strategy after fitted

        .. math::

            logit(E(Y|A,L)) = logit(Y_a) + \sigma * H_a

        4. The targeted Psi is estimated via one of the following formulas

        .. math::

            RD = \frac{1}{n} \sum_{i=1}^{n} (\hat{Y}_1 - \hat{Y}_0)

            RR = \frac{1}{n} \sum_{i=1}^{n} (\hat{Y}_1) / \frac{1}{n} \sum_{i=1}^{n} (\hat{Y}_0)

        Confidence intervals are constructed using influence curve theory

        Parameters
        ----------
        df : DataFrame
            Pandas dataframe containing the variables of interest
        exposure : str
            Column label for the exposure of interest
        outcome : str
            Column label for the outcome of interest
        measure : str, optional
            What the TMLE psi estimates. Current options include; risk difference comparing treated to untreated
            (F(A=1) - F(A=0)), risk ratio (F(A=1) / F(A=0)), and odds ratio. The following keywords are available
            * 'risk_difference'  :   F(A=1) - F(A=0)
            * 'risk_ratio'       :   F(A=1) / F(A=0)
            * 'odds_ratio'       :   (F(A=1) / (1 - F(A=1)) / (F(A=0) / (1 - F(A=0))
        alpha : int, optional
            Alpha for confidence interval level. Default is 0.05

        Examples
        --------
        Setting up environment
        >>>from zepid import load_sample_data, spline
        >>>from zepid.causal.doublyrobust import TMLE
        >>>df = load_sample_data(False).dropna()
        >>>df[['cd4_rs1', 'cd4_rs2']] = spline(df, 'cd40', n_knots=3, term=2, restricted=True)

        Estimating TMLE using logistic regression
        >>>tmle = TMLE(df, exposure='art', outcome='dead')
        >>># Specifying exposure/treatment model
        >>>tmle.exposure_model('male + age0 + cd40 + cd4_rs1 + cd4_rs2 + dvl0')
        >>># Specifying outcome model
        >>>tmle.exposure_model('art + male + age0 + cd40 + cd4_rs1 + cd4_rs2 + dvl0')
        >>># TMLE estimation procedure
        >>>tmle.fit()
        >>># Printing main results
        >>>tmle.summary()
        >>># Extracting point estimate and confidence intervals, respectively
        >>>tmle.psi
        >>>tmle.confint

        Estimating TMLE with machine learning algorithm from sklearn
        >>>from sklearn.linear_model import LogisticRegression
        >>>log1 = LogisticRegression(penalty='l1', random_state=201)
        >>>tmle = TMLE(df, 'art', 'dead')
        >>># custom_model allows specification of machine learning algorithms
        >>>tmle.exposure_model('male + age0 + cd40 + cd4_rs1 + cd4_rs2 + dvl0', custom_model=log1)
        >>>tmle.outcome_model('male + age0 + cd40 + cd4_rs1 + cd4_rs2 + dvl0', custom_model=log1)
        >>>tmle.fit()

        Demonstration of estimating g-model with symmetric bounds
        >>>tmle.exposure_model('male + age0 + cd40 + cd4_rs1 + cd4_rs2 + dvl0', bound=0.05)

        Demonstration of estimating g-model with asymmetric bounds
        >>>tmle.exposure_model('male + age0 + cd40 + cd4_rs1 + cd4_rs2 + dvl0', bound=[0.05, 0.9])

        References
        ----------
        Schuler, Megan S., and Sherri Rose. "Targeted maximum likelihood estimation for causal inference in
        observational studies." American journal of epidemiology 185.1 (2017): 65-73.

        Van der Laan, Mark J., and Sherri Rose. Targeted learning: causal inference for observational and experimental
        data. Springer Science & Business Media, 2011.

        Van Der Laan, Mark J., and Daniel Rubin. "Targeted maximum likelihood learning." The International Journal of
        Biostatistics 2.1 (2006).
        """
        if df.dropna().shape[0] != df.shape[0]:
            warnings.warn("There is missing data in the dataset. By default, TMLE will drop all missing data. TMLE will"
                          "fit "+str(df.dropna().shape[0])+' of '+str(df.shape[0])+' observations', UserWarning)

        # Detailed steps follow "Targeted Learning" chapter 4, figure 4.2 by van der Laan, Rose
        self._psi_correspond = measure
        self.df = df.copy().dropna().reset_index()
        self.alpha = alpha
        self._exposure = exposure
        self._outcome = outcome
        self._out_model = None
        self._exp_model = None
        self._fit_exposure_model = False
        self._fit_outcome_model = False
        self.QA0W = None
        self.QA1W = None
        self.QAW = None
        self.g1W = None
        self.g0W = None
        self._epsilon = None
        self.psi = None
        self.confint = None

    def exposure_model(self, model, custom_model=None, bound=False, print_results=True):
        """Estimation of Pr(A=1|L), which is termed as g(A=1|L) in the literature

        Parameters
        ----------
        model : str
            Independent variables to predict the exposure. Example) 'var1 + var2 + var3'
        custom_model : optional
            Input for a custom model that is used in place of the logit model (default). The model must have the
            "fit()" and  "predict()" attributes. Both sklearn and supylearner are supported as custom models. In the
            background, TMLE will fit the custom model and generate the predicted probablities
        bound : float, list, optional
            Value between 0,1 to truncate predicted probabilities. Helps to avoid near positivity violations.
            Specifying this argument can improve finite sample performance for random positivity violations. However,
            inference becomes limited to the restricted population. Default is False, meaning no truncation of
            predicted probabilities occurs. Providing a single float assumes symmetric trunctation. A collection of
            floats can be provided for asymmetric trunctation
        print_results : bool, optional
            Whether to print the fitted model results. Default is True (prints results)
        """
        self._exp_model = self._exposure + ' ~ ' + model

        # Step 3) Estimation of g-model (exposure model)
        if custom_model is None:
            fitmodel = propensity_score(self.df, self._exp_model, print_results=print_results)
            self.g1W = fitmodel.predict(self.df)

        # User-specified prediction model
        else:
            data = patsy.dmatrix(model + ' - 1', self.df)
            try:
                fm = custom_model.fit(X=data, y=self.df[self._outcome])
            except TypeError:
                raise TypeError("Currently custom_model must have the 'fit' function with arguments 'X', 'y'. This "
                                "covers both sklearn and supylearner. If there is a predictive model you would "
                                "like to use, please open an issue at https://github.com/pzivich/zepid and I "
                                "can work on adding support")
            if print_results and hasattr(fm, 'summarize'):
                fm.summarize()
            if hasattr(fm, 'predict_proba'):
                self.g1W = fm.predict_proba(data)[:, 1]
            elif hasattr(fm, 'predict'):
                self.g1W = fm.predict(data)
            else:
                raise ValueError("Currently custom_model must have 'predict' or 'predict_proba' attribute")

        self.g0W = 1 - self.g1W
        if bound:  # Bounding predicted probabilities if requested
            self.g1W = self._bounding(self.g1W, bounds=bound)
        if bound:  # Bounding predicted probabilities if requested
            self.g0W = self._bounding(self.g0W, bounds=bound)

        self._fit_exposure_model = True

    def outcome_model(self, model, custom_model=None, print_results=True):
        """Estimation of E(Y|A,L), which is also written sometimes as Q(A,W) or Pr(Y=1|A,W)

        Parameters
        ----------
        model : str
            Independent variables to predict the exposure. Example) 'var1 + var2 + var3'
        custom_model : optional
            Input for a custom model that is used in place of the logit model (default). The model must have the
            "fit()" and  "predict()" attributes. Both sklearn and supylearner are supported as custom models. In the
            background, TMLE will fit the custom model and generate the predicted probablities
        print_results : bool, optional
            Whether to print the fitted model results. Default is True (prints results)
        """
        self._out_model = self._outcome + ' ~ ' + model

        # Step 1) Prediction for Q (estimation of Q-model)
        if custom_model is None:  # Logistic Regression model for predictions
            f = sm.families.family.Binomial()
            log = smf.glm(self._out_model, self.df, family=f).fit()

            if print_results:
                print('\n----------------------------------------------------------------')
                print('MODEL: ' + self._out_model)
                print('-----------------------------------------------------------------')
                print(log.summary())

            # Step 2) Estimation under the scenarios
            self.QAW = log.predict(self.df)
            dfx = self.df.copy()
            dfx[self._exposure] = 1
            self.QA1W = log.predict(dfx)
            dfx = self.df.copy()
            dfx[self._exposure] = 0
            self.QA0W = log.predict(dfx)

        # User-specified model
        else:
            data = patsy.dmatrix(model + ' - 1', self.df)
            try:
                fm = custom_model.fit(X=data, y=self.df[self._outcome])
            except TypeError:
                raise TypeError("Currently custom_model must have the 'fit' function with arguments 'X', 'y'. This "
                                "covers both sklearn and supylearner. If there is a predictive model you would "
                                "like to use, please open an issue at https://github.com/pzivich/zepid and I "
                                "can work on adding support")
            if print_results and hasattr(fm, 'summarize'):
                fm.summarize()
            if hasattr(fm, 'predict_proba'):
                self.QAW = fm.predict_proba(data)[:, 1]

                dfx = self.df.copy()
                dfx[self._exposure] = 1
                data = patsy.dmatrix(model + ' - 1', dfx)
                self.QA1W = fm.predict_proba(data)[:, 1]

                dfx = self.df.copy()
                dfx[self._exposure] = 0
                data = patsy.dmatrix(model + ' - 1', dfx)
                self.QA0W = fm.predict_proba(data)[:, 1]

            elif hasattr(fm, 'predict'):
                self.QAW = fm.predict(data)

                dfx = self.df.copy()
                dfx[self._exposure] = 1
                data = patsy.dmatrix(model + ' - 1', dfx)
                self.QA1W = fm.predict(data)

                dfx = self.df.copy()
                dfx[self._exposure] = 0
                data = patsy.dmatrix(model + ' - 1', dfx)
                self.QA0W = fm.predict(data)

            else:
                raise ValueError("Currently custom_model must have 'predict' or 'predict_proba' attribute")

        self._fit_outcome_model = True

    def fit(self):
        """Estimates psi based on the gAW and QAW. Confidence intervals come from the influence curve

        Returns
        -------
        TMLE gains Psi and confint attributes
        """
        if (self._fit_exposure_model is False) or (self._fit_outcome_model is False):
            raise ValueError('The exposure and outcome models must be specified before the psi estimate can '
                             'be generated')

        # Step 4) Calculating clever covariate (HAW)
        H1W = self.df[self._exposure] / self.g1W
        H0W = -(1 - self.df[self._exposure]) / (self.g0W)
        HAW = H1W + H0W

        # Step 5) Estimating TMLE
        f = sm.families.family.Binomial()
        log = sm.GLM(self.df[self._outcome], np.column_stack((H1W, H0W)), offset=np.log(probability_to_odds(self.QAW)),
                     family=f).fit()
        self._epsilon = log.params
        Qstar1 = logistic.cdf(np.log(probability_to_odds(self.QA1W)) + self._epsilon[0] * 1/self.g1W)
        Qstar0 = logistic.cdf(np.log(probability_to_odds(self.QA0W)) + self._epsilon[1] * -1/self.g0W)
        Qstar = log.predict(np.column_stack((H1W, H0W)), offset=np.log(probability_to_odds(self.QAW)))

        # Step 6) Calculating Psi
        if self.alpha == 0.05:  # Without this, won't match R exactly. R relies on 1.96, while I use SciPy
            zalpha = 1.96
        else:
            zalpha = norm.ppf(1 - self.alpha / 2, loc=0, scale=1)

        # p-values are not implemented (doing my part to enforce CL over p-values)
        if self._psi_correspond == 'risk_difference':
            self.psi = np.mean(Qstar1 - Qstar0)
            # Influence Curve for CL
            ic = HAW * (self.df[self._outcome] - Qstar) + (Qstar1 - Qstar0) - self.psi
            varIC = np.var(ic, ddof=1) / self.df.shape[0]
            self.confint = [self.psi - zalpha * math.sqrt(varIC),
                            self.psi + zalpha * math.sqrt(varIC)]

        elif self._psi_correspond == 'risk_ratio':
            self.psi = np.mean(Qstar1) / np.mean(Qstar0)
            # Influence Curve for CL
            ic = (1/np.mean(Qstar1)*(H1W * (self.df[self._outcome] - Qstar) + Qstar1 - np.mean(Qstar1)) -
                  (1/np.mean(Qstar0))*(-1*H0W * (self.df[self._outcome] - Qstar) + Qstar0 - np.mean(Qstar0)))
            varIC = np.var(ic, ddof=1) / self.df.shape[0]
            self.confint = [np.exp(np.log(self.psi) - zalpha * math.sqrt(varIC)),
                            np.exp(np.log(self.psi) + zalpha * math.sqrt(varIC))]

        elif self._psi_correspond == 'odds_ratio':
            self.psi = (np.mean(Qstar1) / (1 - np.mean(Qstar1))) / (np.mean(Qstar0) / (1 - np.mean(Qstar0)))
            # Influence Curve for CL
            ic = ((1/(np.mean(Qstar1)*(1-np.mean(Qstar1))) * (H1W*(self.df[self._outcome] - Qstar) + Qstar1)) -
                  (1/(np.mean(Qstar0)*(1 - np.mean(Qstar0))) * (-1*H0W*(self.df[self._outcome] - Qstar) + Qstar0)))
            seIC = math.sqrt(np.var(ic, ddof=1) / self.df.shape[0])
            self.confint = [np.exp(np.log(self.psi) - zalpha * seIC),
                            np.exp(np.log(self.psi) + zalpha * seIC)]

        else:
            raise ValueError('Specified parameter is not implemented. Please use one of the available options')

    def summary(self, decimal=3):
        """Prints summary of model results

        Parameters
        ----------
        decimal : int, optional
            Number of decimal places to display. Default is 3
        """
        if (self._fit_exposure_model is False) or (self._fit_exposure_model is False):
            raise ValueError('The exposure and outcome models must be specified before the psi estimate can '
                             'be generated')

        print('----------------------------------------------------------------------')
        print('Psi: ', round(float(self.psi), decimal))
        print(str(round(100 * (1 - self.alpha), 1)) + '% two-sided CI: (' + str(round(self.confint[0], decimal)), ',',
              str(round(self.confint[1], decimal)) + ')')
        print('----------------------------------------------------------------------')
        print('Psi corresponds to '+self._psi_correspond)
        print('----------------------------------------------------------------------')

    @staticmethod
    def _bounding(v, bounds):
        """Background function to perform bounding feature. Not intended for users to access

        v:
            -Values to be bounded
        bounds:
            -Percentile thresholds for bounds
        """
        if type(bounds) is float:  # Symmetric bounding
            if bounds < 0 or bounds > 1:
                raise ValueError('Bound value must be between (0, 1)')
            v = np.where(v < bounds, bounds, v)
            v = np.where(v > 1-bounds, 1-bounds, v)
        elif type(bounds) is str:  # Catching string inputs
            raise ValueError('Bounds must either be a float between (0, 1), or a collection of floats between (0, 1)')
        else:  # Asymmetric bounds
            if bounds[0] > bounds[1]:
                raise ValueError('Bound thresholds must be listed in ascending order')
            if len(bounds) > 2:
                warnings.warn('It looks like your specified bounds is more than two floats. Only the first two '
                              'specified bounds are used by the bound statement. So only ' +
                              str(bounds[0:2]) + ' will be used', UserWarning)
            if type(bounds[0]) is str or type(bounds[1]) is str:
                raise ValueError('Bounds must be floats between (0, 1)')
            if (bounds[0] < 0 or bounds[1] > 1) or (bounds[0] < 0 or bounds[1] > 1):
                raise ValueError('Both bound values must be between (0, 1)')
            v = np.where(v < bounds[0], bounds[0], v)
            v = np.where(v > bounds[1], bounds[1], v)
        print(np.max(v))
        return v

