# Copyright (c) PyWhy contributors. All rights reserved.
# Licensed under the MIT License.

"""Doubly Robust IV for Heterogeneous Treatment Effects.

An Doubly Robust machine learning approach to estimation of heterogeneous
treatment effect with an endogenous treatment and an instrument.

Implements the DRIV algorithm for estimating CATE with IVs from the paper:

Machine Learning Estimation of Heterogeneous Treatment Effects with Instruments
Vasilis Syrgkanis, Victor Lei, Miruna Oprescu, Maggie Hei, Keith Battocchi, Greg Lewis
https://arxiv.org/abs/1905.10176
"""

import numpy as np
from sklearn.model_selection import KFold, train_test_split, StratifiedKFold
from econml.utilities import hstack
from sklearn.base import clone
import copy
from itertools import tee


class _BaseDRIV:

    """
    The _BaseDRIV algorithm for estimating CATE with IVs. It is the parent of the
    two public classes {DRIV, ProjectedDRIV}
    """

    def __init__(self, nuisance_models,
                 model_effect,
                 cov_clip=.1,
                 n_splits=3,
                 binary_instrument=False, binary_treatment=False,
                 opt_reweighted=False):
        """
        Parameters
        ----------
        nuisance_models : dictionary of nuisance models, with {'name_of_model' : EstimatorObject, ...}
        model_effect : model to estimate second stage effect model from doubly robust target
        cov_clip : clipping of the covariate for regions with low "overlap",
            so as to reduce variance
        n_splits : number of splits to use in cross-fitting
        binary_instrument : whether to stratify cross-fitting splits by instrument
        binary_treatment : whether to stratify cross-fitting splits by treatment
        opt_reweighted : whether to reweight the samples to minimize variance. If True then
            model_effect.fit must accept sample_weight as a kw argument (WeightWrapper from
            utilities can be used for any linear model to enable sample_weights). If True then
            assumes the model_effect is flexible enough to fit the true CATE model. Otherwise,
            it method will return a biased projection to the model_effect space, biased
            to give more weight on parts of the feature space where the instrument is strong.
        """
        for n_name, n_model in nuisance_models.items():
            setattr(self, n_name, [clone(n_model, safe=False)
                                   for _ in range(n_splits)])
        self.nuisance_model_names = list(nuisance_models.keys())
        self.model_effect = clone(model_effect, safe=False)
        self.cov_clip = cov_clip
        self.n_splits = n_splits
        self.binary_instrument = binary_instrument
        self.binary_treatment = binary_treatment
        self.opt_reweighted = opt_reweighted
        self.stored_final_data = False

    def _check_inputs(self, y, T, X, Z):
        """ Checks dimension of inputs and reshapes inputs. Implemented at child
        """
        raise("Child class needs to implement this method")

    def _nuisance_estimates(self, y, T, X, Z):
        """ Estimates the nuisance quantities for each sample with cross-fitting.
        Implemented at child
        """
        raise("Child class needs to implement this method")

    def _get_split_enum(self, y, T, X, Z):
        """
        Returns an enumerator over the splits for cross-fitting
        """
        # We do a three way split, as typically a preliminary theta estimator would require
        # many samples. So having 2/3 of the sample to train model_theta seems appropriate.
        if self.n_splits == 1:
            splits = [(np.arange(X.shape[0]), np.arange(X.shape[0]))]
        # TODO. Deal with multi-class instrument
        elif self.binary_instrument or self.binary_treatment:
            group = 2 * T * self.binary_treatment + Z.flatten() * self.binary_instrument
            splits = StratifiedKFold(
                n_splits=self.n_splits, shuffle=True).split(X, group)
        else:
            splits = KFold(n_splits=self.n_splits, shuffle=True).split(X)

        return splits

    def fit(self, y, T, X, Z, store_final=False):
        """
        Parameters
        ----------
        y : outcome
        T : treatment (single dimensional)
        X : features/controls
        Z : instrument (single dimensional)
        store_final : whether to store nuisance data that are used in the final for
            refitting the final stage later on
        """
        y, T, X, Z = self._check_inputs(y, T, X, Z)

        prel_theta, res_t, res_y, res_z, cov = self._nuisance_estimates(
            y, T, X, Z)

        # Estimate final model of theta(X) by minimizing the square loss:
        # (prel_theta(X) + (Y_res - prel_theta(X) * T_res) * Z_res / cov[T,Z | X] - theta(X))^2
        # We clip the covariance so that it is bounded away from zero, so as to reduce variance
        # at the expense of some small bias. For points with very small covariance we revert
        # to the model-based preliminary estimate and do not add the correction term.
        cov_sign = np.sign(cov)
        cov_sign[cov_sign == 0] = 1
        clipped_cov = cov_sign * np.clip(np.abs(cov),
                                         self.cov_clip, np.inf)
        theta_dr = prel_theta + \
            (res_y - prel_theta * res_t) * res_z / clipped_cov
        if self.opt_reweighted:
            self.model_effect.fit(X, theta_dr, sample_weight=clipped_cov**2)
        else:
            self.model_effect.fit(X, theta_dr)

        if store_final:
            self.X = X
            self.theta_dr = theta_dr
            self.stored_final_data = True
            self.clipped_cov = clipped_cov

        return self

    def refit_final(self, model_effect, opt_reweighted=None):
        """
        Change the final effect model and refit the final stage.
        Parameters
        ----------
        model_effect : an instance of the new effect model to be fitted in the final stage
        opt_reweighted : whether to weight samples for variance reduction in the final model fitting
        """
        if not self.stored_final_data:
            raise AttributeError(
                "Estimator is not yet fit with store_data=True")
        if opt_reweighted is not None:
            self.opt_reweighted = opt_reweighted
        self.model_effect = model_effect
        if self.opt_reweighted:
            self.model_effect.fit(self.X, self.theta_dr,
                                  sample_weight=self.clipped_cov**2)
        else:
            self.model_effect.fit(self.X, self.theta_dr)
        return self

    def effect(self, X):
        """
        Parameters
        ----------
        X : features
        """
        return self.model_effect.predict(X)

    @property
    def effect_model(self):
        return self.model_effect

    @property
    def fitted_nuisances(self):
        nuisance_dict = {}
        for n_name in self.nuisance_model_names:
            nuisance_dict[n_name] = getattr(self, n_name)
        return nuisance_dict

    @property
    def coef_(self):
        if not hasattr(self.effect_model, 'coef_'):
            raise AttributeError("Effect model is not linear!")
        return self.effect_model.coef_

    @property
    def intercept_(self):
        if not hasattr(self.effect_model, 'intercept_'):
            raise AttributeError("Effect model is not linear!")
        return self.effect_model.intercept_


class DRIV(_BaseDRIV):
    """
    Implements the DRIV algorithm
    """

    def __init__(self, model_Y_X, model_T_X, model_Z_X,
                 prel_model_effect, model_TZ_X,
                 model_effect,
                 cov_clip=.1,
                 n_splits=3,
                 binary_instrument=False, binary_treatment=False,
                 opt_reweighted=False):
        """
        Parameters
        ----------
        model_Y_X : model to predict E[Y | X]
        model_T_X : model to predict E[T | X]. In alt_fit, this model is also
            used to predict E[T | X, Z]
        model_Z_X : model to predict E[Z | X]
        model_theta : model that estimates a preliminary version of the CATE
            (e.g. via DMLIV or other method)
        model_TZ_X : model to estimate E[T * Z | X]
        model_effect : model to estimate second stage effect model from doubly robust target
        cov_clip : clipping of the covariate for regions with low "overlap",
            so as to reduce variance
        n_splits : number of splits to use in cross-fitting
        binary_instrument : whether to stratify cross-fitting splits by instrument
        binary_treatment : whether to stratify cross-fitting splits by treatment
        opt_reweighted : whether to reweight the samples to minimize variance. If True then
            model_effect.fit must accept sample_weight as a kw argument (WeightWrapper from
            utilities can be used for any linear model to enable sample_weights). If True then
            assumes the model_effect is flexible enough to fit the true CATE model. Otherwise,
            it method will return a biased projection to the model_effect space, biased
            to give more weight on parts of the feature space where the instrument is strong.
        """
        nuisance_models = {'prel_model_effect': prel_model_effect,
                           'model_TZ_X': model_TZ_X,
                           'model_T_X': model_T_X,
                           'model_Z_X': model_Z_X,
                           'model_Y_X': model_Y_X}
        super(DRIV, self).__init__(nuisance_models, model_effect,
                                   cov_clip=cov_clip,
                                   n_splits=n_splits,
                                   binary_instrument=binary_instrument, binary_treatment=binary_treatment,
                                   opt_reweighted=opt_reweighted)
        return

    def _check_inputs(self, y, T, X, Z):
        if len(Z.shape) > 1 and Z.shape[1] > 1:
            raise AssertionError(
                "Can only accept single dimensional instrument")
        if len(T.shape) > 1 and T.shape[1] > 1:
            raise AssertionError(
                "Can only accept single dimensional treatment")
        if len(y.shape) > 1 and y.shape[1] > 1:
            raise AssertionError("Can only accept single dimensional outcome")
        Z = Z.flatten()
        T = T.flatten()
        y = y.flatten()
        return y, T, X, Z

    def _nuisance_estimates(self, y, T, X, Z):
        n_samples = y.shape[0]
        prel_theta = np.zeros(n_samples)
        res_t = np.zeros(n_samples)
        res_y = np.zeros(n_samples)
        res_z = np.zeros(n_samples)
        cov = np.zeros(n_samples)

        splits = self._get_split_enum(y, T, X, Z)
        for idx, (train, test) in enumerate(splits):
            # Estimate preliminary theta in cross fitting manner
            prel_theta[test] = self.prel_model_effect[idx].fit(
                y[train], T[train], X[train], Z[train]).effect(X[test]).flatten()
            # Estimate p(X) = E[T | X] in cross fitting manner
            self.model_T_X[idx].fit(X[train], T[train])
            pr_t_test = self.model_T_X[idx].predict(X[test])
            # Estimate r(Z) = E[Z | X] in cross fitting manner
            self.model_Z_X[idx].fit(X[train], Z[train])
            pr_z_test = self.model_Z_X[idx].predict(X[test])
            # Calculate residual T_res = T - p(X) and Z_res = Z - r(X)
            res_t[test] = T[test] - pr_t_test
            res_z[test] = Z[test] - pr_z_test
            # Estimate residual Y_res = Y - q(X) = Y - E[Y | X] in cross fitting manner
            res_y[test] = y[test] - \
                self.model_Y_X[idx].fit(X[train], y[train]).predict(X[test])
            # Estimate cov[T, Z | X] = E[(T-p(X))*(Z-r(X)) | X] = E[T*Z | X] - E[T |X]*E[Z | X]
            cov[test] = self.model_TZ_X[idx].fit(
                X[train], T[train] * Z[train]).predict(X[test]) - pr_t_test * pr_z_test

        return prel_theta, res_t, res_y, res_z, cov


class ProjectedDRIV(_BaseDRIV):
    """
    This is a slight variant of DRIV where we use E[T|Z, X] as
    the instrument as opposed to Z. The rest is the same as the normal
    fit.
    """

    def __init__(self, model_Y_X, model_T_X, model_T_XZ,
                 prel_model_effect, model_TZ_X,
                 model_effect,
                 cov_clip=.1,
                 n_splits=3,
                 binary_instrument=False, binary_treatment=False,
                 opt_reweighted=False):
        """
        Parameters
        ----------
        model_Y_X : model to predict E[Y | X]
        model_T_X : model to predict E[T | X]. In alt_fit, this model is also
            used to predict E[T | X, Z]
        model_T_XZ : model to predict E[T | X, Z]
        model_theta : model that estimates a preliminary version of the CATE
            (e.g. via DMLIV or other method)
        model_TZ_X : model to estimate cov[T, E[T|X,Z] | X] = E[(T-E[T|X]) * (E[T|X,Z] - E[T|X]) | X].
        model_effect : model to estimate second stage effect model from doubly robust target
        cov_clip : clipping of the covariate for regions with low "overlap",
            so as to reduce variance
        n_splits : number of splits to use in cross-fitting
        binary_instrument : whether to stratify cross-fitting splits by instrument
        binary_treatment : whether to stratify cross-fitting splits by treatment
        opt_reweighted : whether to reweight the samples to minimize variance. If True then
            model_effect.fit must accept sample_weight as a kw argument (WeightWrapper from
            utilities can be used for any linear model to enable sample_weights). If True then
            assumes the model_effect is flexible enough to fit the true CATE model. Otherwise,
            it method will return a biased projection to the model_effect space, biased
            to give more weight on parts of the feature space where the instrument is strong.
        """
        nuisance_models = {'prel_model_effect': prel_model_effect,
                           'model_Y_X': model_Y_X,
                           'model_T_X': model_T_X,
                           'model_T_XZ': model_T_XZ,
                           'model_TZ_X': model_TZ_X}
        super(ProjectedDRIV, self).__init__(nuisance_models, model_effect,
                                            cov_clip=cov_clip,
                                            n_splits=n_splits,
                                            binary_instrument=binary_instrument, binary_treatment=binary_treatment,
                                            opt_reweighted=opt_reweighted)
        return

    def _check_inputs(self, y, T, X, Z):

        if len(T.shape) > 1 and T.shape[1] > 1:
            raise AssertionError(
                "Can only accept single dimensional treatment")
        if len(y.shape) > 1 and y.shape[1] > 1:
            raise AssertionError("Can only accept single dimensional outcome")
        if len(Z.shape) == 1:
            Z = Z.reshape(-1, 1)
        if (Z.shape[1] > 1) and self.binary_instrument:
            raise AssertionError(
                "Binary instrument flag is True, but instrument is multi-dimensional")
        T = T.flatten()
        y = y.flatten()

        return y, T, X, Z

    def _nuisance_estimates(self, y, T, X, Z):

        n_samples = y.shape[0]
        prel_theta = np.zeros(n_samples)
        res_t = np.zeros(n_samples)
        res_y = np.zeros(n_samples)
        res_z = np.zeros(n_samples)
        cov = np.zeros(n_samples)
        proj_t = np.zeros(n_samples)

        splits = self._get_split_enum(y, T, X, Z)

        # TODO. The solution below is not really a valid cross-fitting
        # as the test data are used to create the proj_t on the train
        # which in the second train-test loop is used to create the nuisance
        # cov on the test data. Hence the T variable of some sample
        # is implicitly correlated with its cov nuisance, through this flow
        # of information. However, this seems a rather weak correlation.
        # The more kosher would be to do an internal nested cv loop for the T_XZ
        # model.
        splits, splits_one = tee(splits)
        # Estimate h(X, Z) = E[T | X, Z] in cross fitting manner
        for idx, (train, test) in enumerate(splits_one):
            self.model_T_XZ[idx].fit(hstack([X[train], Z[train]]), T[train])
            proj_t[test] = self.model_T_XZ[idx].predict(
                hstack([X[test], Z[test]]))

        for idx, (train, test) in enumerate(splits):
            # Estimate preliminary theta in cross fitting manner
            prel_theta[test] = self.prel_model_effect[idx].fit(
                y[train], T[train], X[train], Z[train]).effect(X[test]).flatten()
            # Estimate p(X) = E[T | X] in cross fitting manner
            self.model_T_X[idx].fit(X[train], T[train])
            pr_t_test = self.model_T_X[idx].predict(X[test])
            # Calculate residual T_res = T - p(X) and Z_res = h(Z, X) - p(X)
            res_t[test] = T[test] - pr_t_test
            res_z[test] = proj_t[test] - pr_t_test
            # Estimate residual Y_res = Y - q(X) = Y - E[Y | X] in cross fitting manner
            res_y[test] = y[test] - \
                self.model_Y_X[idx].fit(X[train], y[train]).predict(X[test])
            # Estimate cov[T, E[T|X,Z] | X] = E[T * E[T|X,Z]] - E[T|X]^2
            cov[test] = self.model_TZ_X[idx].fit(
                X[train], T[train] * proj_t[train]).predict(X[test]) - pr_t_test**2

        return prel_theta, res_t, res_y, res_z, cov

##############################################################################
# Classes for the DRIV implementation for the special case of intent-to-treat
# A/B test
##############################################################################


class _IntentToTreatDRIV(_BaseDRIV):
    """
    Helper class for the DRIV algorithm for the intent-to-treat A/B test setting
    """

    def __init__(self, model_Y_X, model_T_XZ,
                 prel_model_effect,
                 model_effect,
                 cov_clip=.1,
                 n_splits=3,
                 opt_reweighted=False):
        """
        """
        nuisance_models = {'model_Y_X': model_Y_X,
                           'model_T_XZ': model_T_XZ,
                           'prel_model_effect': prel_model_effect}
        super(_IntentToTreatDRIV, self).__init__(nuisance_models, model_effect,
                                                 cov_clip=cov_clip,
                                                 n_splits=n_splits,
                                                 binary_instrument=True, binary_treatment=True,
                                                 opt_reweighted=opt_reweighted)
        return

    def _check_inputs(self, y, T, X, Z):
        if len(Z.shape) > 1 and Z.shape[1] > 1:
            raise AssertionError(
                "Can only accept single dimensional instrument")
        if len(T.shape) > 1 and T.shape[1] > 1:
            raise AssertionError(
                "Can only accept single dimensional treatment")
        if len(y.shape) > 1 and y.shape[1] > 1:
            raise AssertionError("Can only accept single dimensional outcome")
        Z = Z.flatten()
        T = T.flatten()
        y = y.flatten()
        return y, T, X, Z

    def _nuisance_estimates(self, y, T, X, Z):
        n_samples = y.shape[0]
        prel_theta = np.zeros(n_samples)
        res_t = np.zeros(n_samples)
        res_y = np.zeros(n_samples)
        delta = np.zeros(n_samples)

        splits = self._get_split_enum(y, T, X, Z)
        for idx, (train, test) in enumerate(splits):
            # Estimate preliminary theta in cross fitting manner
            prel_theta[test] = self.prel_model_effect[idx].fit(
                y[train], T[train], X[train], Z[train]).effect(X[test]).flatten()
            # Estimate p(X) = E[T | X] in cross fitting manner
            self.model_T_XZ[idx].fit(hstack([X[train], Z[train].reshape(-1, 1)]), T[train])
            Z_one = np.ones((Z[test].shape[0], 1))
            Z_zero = np.zeros((Z[test].shape[0], 1))
            pr_t_test_one = self.model_T_XZ[idx].predict(hstack([X[test], Z_one]))
            pr_t_test_zero = self.model_T_XZ[idx].predict(hstack([X[test], Z_zero]))
            delta[test] = (pr_t_test_one - pr_t_test_zero) / 2
            pr_t_test = (pr_t_test_one + pr_t_test_zero) / 2
            res_t[test] = T[test] - pr_t_test
            # Estimate residual Y_res = Y - q(X) = Y - E[Y | X] in cross fitting manner
            res_y[test] = y[test] - \
                self.model_Y_X[idx].fit(X[train], y[train]).predict(X[test])

        return prel_theta, res_t, res_y, 2 * Z - 1, delta


class _DummyCATE:
    """
    A dummy cate effect model that always returns zero effect
    """

    def __init__(self):
        return

    def fit(self, y, T, X, Z):
        return self

    def effect(self, X):
        return np.zeros(X.shape[0])


class IntentToTreatDRIV(_IntentToTreatDRIV):
    """
    Implements the DRIV algorithm for the intent-to-treat A/B test setting
    """

    def __init__(self, model_Y_X, model_T_XZ,
                 flexible_model_effect,
                 final_model_effect=None,
                 cov_clip=.1,
                 n_splits=3,
                 opt_reweighted=False):
        """
        Parameters
        ----------
        model_Y_X : model to predict E[Y | X]
        model_T_XZ : model to predict E[T | X, Z]
        flexible_model_effect : a flexible model for a preliminary version of the CATE, must accept
            sample_weight at fit time.
        final_model_effect : a final model for the CATE and projections. If None, then
            flexible_model_effect is also used as a final model
        cov_clip : clipping of the covariate for regions with low "overlap",
            so as to reduce variance
        n_splits : number of splits to use in cross-fitting
        opt_reweighted : whether to reweight the samples to minimize variance. If True then
            final_model_effect.fit must accept sample_weight as a kw argument (WeightWrapper from
            utilities can be used for any linear model to enable sample_weights). If True then
            assumes the final_model_effect is flexible enough to fit the true CATE model. Otherwise,
            it method will return a biased projection to the model_effect space, biased
            to give more weight on parts of the feature space where the instrument is strong.
        """
        prel_model_effect = _IntentToTreatDRIV(clone(model_Y_X, safe=False),
                                               clone(model_T_XZ, safe=False),
                                               _DummyCATE(),
                                               clone(flexible_model_effect, safe=False),
                                               cov_clip=1e-7, n_splits=1, opt_reweighted=True)
        if final_model_effect is None:
            final_model_effect = clone(flexible_model_effect, safe=False)
        super(IntentToTreatDRIV, self).__init__(model_Y_X, model_T_XZ, prel_model_effect,
                                                final_model_effect,
                                                cov_clip=cov_clip,
                                                n_splits=n_splits,
                                                opt_reweighted=opt_reweighted)
        return
