# AUTOGENERATED! DO NOT EDIT! File to edit: ../02_estimators.ipynb.

# %% auto 0
__all__ = ['CausalForestEstimator', 'OPENBTITEEstimator', 'GPEstimator', 'SKLiftEstimator']

# %% ../02_estimators.ipynb 2
#| output: false
#| output: false
from .base import *
from econml.dml import CausalForestDML
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestRegressor
from sklearn.linear_model import LinearRegression, LogisticRegression
import econml
import sklift
# import causalml
# import pymc as pm 
# import pymc_bart as pmb

# %% ../02_estimators.ipynb 3
class CausalForestEstimator(BaseITEEstimator):
    def fit(self, **kwargs):
        if self.model is None:
            self.model = CausalForestDML()
        self.model.fit(Y=kwargs["y_training"],
                       T=kwargs["t_training"],
                       X=kwargs["X_training"])

    def predict(self, **kwargs):
        if 'return_counterfactuals' in kwargs:
            raise ValueError("Causal Forest does not support counterfactual predictions out of the box")
        preds = self.model.effect_inference(kwargs["X"])
        if "return_mean" in kwargs:
            out = preds.pred
        else:
            out = (preds.pred, preds.var)
        return out

# %% ../02_estimators.ipynb 4
class OPENBTITEEstimator(BaseITEEstimator):
    """Modified ITE estimator for OPENBT

    The predictions are transposed so the uncertainty sampler can calculate uncertianty easily"""
    def predict(self, **kwargs):
        X = kwargs["X"]
        if self.ps_model is not None:
            ps_scores = self.ps_model.predict_proba(X)
            X = np.hstack((X, ps_scores[:,1].reshape((-1, 1))))
        X0 = np.concatenate((X,
                             np.zeros(X.shape[0]).reshape((-1,1))),axis=1)
        X1 = np.concatenate((X,
                             np.ones(X.shape[0]).reshape((-1,1))),axis=1)
        preds0 = self.model.predict(X0)
        preds1 = self.model.predict(X1)
        if "return_mean" in kwargs:
            if kwargs["return_mean"]:
                out = preds1["mmean"] - preds0["mmean"]
        else:
            out = preds1["mdraws"].T - preds0["mdraws"].T
        if "return_per_cf" in kwargs:
            if kwargs["return_per_cf"]:
                return {"pred1": preds1["mdraws"].T , "pred0":preds0["mdraws"].T}
        return out

# %% ../02_estimators.ipynb 6
class GPEstimator(BaseITEEstimator):
    def predict(self, **kwargs):
        if 'return_mean' in kwargs:
            pred0 = self.model.predict(kwargs["X"])
            pred1 = self.m1.predict(kwargs["X"])
            ite = pred1 - pred0
        else:
            draws0 = self.model.sample_y(kwargs["X"], n_samples=100)
            draws1 = self.m1.sample_y(kwargs["X"], n_samples=100)
            ite = draws1 - draws0
        return ite

# %% ../02_estimators.ipynb 8
class SKLiftEstimator(BaseITEEstimator):
    def __init__(self, sk_model):
        self.model = sk_model
    
    def fit(self, **kwargs):
        self.model.fit(X=kwargs["X_training"],
                       y=kwargs["y_training"],
                       treatment=kwargs["t_training"])
    
    def predict(self, **kwargs):
        if 'return_counterfactuals' in kwargs:
            raise ValueError("Causal Forest does not support counterfactual predictions out of the box")
        out = self.model.predict(kwargs["X_test"])
        return out
