from pathlib import Path

import os

import catenets.models as cate_models

from catenets.models.diffpo import DiffPOLearner

import numpy as np

import pandas as pd

import wandb 

import shap

from PIL import Image

import src.iterpretability.logger as log

from sklearn.metrics import roc_auc_score

from src.plotting import (

    plot_results_datasets_compare, 

    plot_expertise_metrics,

    plot_performance_metrics,

    plot_performance_metrics_f_cf,

)

from src.iterpretability.explain import Explainer

from src.iterpretability.datasets.data_loader import load

from src.iterpretability.simulators import (

    SimulatorBase,

    TYSimulator,

    TSimulator

)

from src.iterpretability.utils import (

    attribution_accuracy,

)

import time

# For contour plotting

from sklearn.metrics import mean_squared_error, roc_auc_score, f1_score

from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier

from sklearn.multioutput import MultiOutputRegressor, MultiOutputClassifier

from sklearn.linear_model import LinearRegression, Lasso, LogisticRegression, LogisticRegressionCV, LassoCV

from sklearn.model_selection import LeaveOneOut

from sklearn.metrics import auc

# Hydra for configuration

from omegaconf import DictConfig, OmegaConf

class ExperimentBase():

    """

    Base class for all experiments.

    """

    def __init__(self, cfg: DictConfig) -> None:

        # Store configuration

        self.cfg = cfg

        # General experiment settings

        self.seeds = cfg.seeds

        self.n_splits = cfg.n_splits

        self.experiment_name = cfg.experiment_name

        self.simulation_type = cfg.simulator.simulation_type

        self.evaluate_inference = cfg.evaluate_inference

        # Load data

        if self.simulation_type == "TY":

            self.X,self.feature_names = load(cfg.dataset, 

                                            train_ratio=1, 

                                            debug=cfg.debug, 

                                            sim_type=self.simulation_type,

                                            directory_path_=cfg.directory_path+'/',

                                            repo_path = cfg.repo_path,

                                            n_samples=cfg.n_samples)

            self.outcomes = None

            self.discrete_outcome = cfg.simulator.num_binary_outcome == cfg.simulator.dim_Y # Currently the only discrete outcome is binary

        elif self.simulation_type == "T":

            self.X, self.outcomes, self.feature_names = load(cfg.dataset, 

                                                            train_ratio=1, 

                                                            debug=cfg.debug, 

                                                            directory_path_=cfg.directory_path+'/', 

                                                            repo_path = cfg.repo_path,

                                                            sim_type=self.simulation_type)

            # Check whether all entries in outcomes are integers

            self.discrete_outcome = cfg.simulator.num_binary_outcome

            cfg.simulator.num_T = self.outcomes.shape[1]

        else:

            raise ValueError(f"Simulation type {self.cfg.simulator.simulation_type} not supported.")

        # Initialize learners

        self.discrete_treatment = True # Currently simulation only supports discrete treatment

        # Results path and directories

        self.file_name = f"{self.cfg.results_dictionary_prefix}_{self.cfg.dataset}_T{self.cfg.simulator.num_T}_Y{self.cfg.simulator.dim_Y}_{self.cfg.simulator.simulation_type}sim_numbin{self.cfg.simulator.num_binary_outcome}"

        self.results_path = Path(self.cfg.results_path) / Path(self.cfg.experiment_name) / Path(self.file_name)

        if not self.results_path.exists():

            self.results_path.mkdir(parents=True, exist_ok=True)

        # Variables set by some or all experiments

        self.learners = None

        self.baseline_learners = None

        self.pred_learners = None

        self.prog_learner = None

        self.select_learner = None

        self.explanations = None

        self.all_important_features = None

        self.pred_features = None

        self.prog_features = None

        self.select_features = None

        self.true_num_swaps = None

        self.swap_counter = None

        self.training_times = {}

    def get_learners(self, 

                     num_features: int, 

                     seed: int = 123) -> dict:

        """

        Get learners for the experiment, based on the configuration settings.

        """

        if self.cfg.simulator.dim_Y > 1 and "Torch_XLearner" in self.cfg.model_names:

            raise ValueError("Torch_XLearner only supports one outcome dimension.")

        elif self.cfg.simulator.dim_Y > 1 and "EconML_SparseLinearDRLearner" in self.cfg.model_names:

            raise ValueError("EconML_SparseLinearDRLearner only supports one outcome dimension.")

        # Let user know that torch and diffpo do not support multiple treatments

        if self.cfg.simulator.num_T > 2 and ("Torch" in self.cfg.model_names or "DiffPOLearner" in self.cfg.model_names):

            raise ValueError("Torch and DiffPO models do not support multiple treatments. Only the first treatment will be used.")

        binary_y = self.discrete_outcome and self.cfg.simulator.num_binary_outcome == 1

        learners = {

                        "EconML_CausalForestDML": cate_models.econml.EconMlEstimator2(self.cfg,

                                                                                model_name="EconML_CausalForestDML",

                                                                                discrete_treatment=self.discrete_treatment,

                                                                                discrete_outcome=self.discrete_outcome,

                                                                                seed=seed),

                        "EconML_DML": cate_models.econml.EconMlEstimator2(self.cfg,

                                                                        model_name="EconML_DML",

                                                                        discrete_treatment=self.discrete_treatment,

                                                                        discrete_outcome=self.discrete_outcome,

                                                                        seed=seed),

                        "EconML_DRLearner": cate_models.econml.EconMlEstimator2(self.cfg,

                                                                            model_name="EconML_DRLearner",

                                                                            discrete_treatment=self.discrete_treatment,

                                                                            discrete_outcome=self.discrete_outcome,

                                                                            seed=seed),

                        "EconML_DMLOrthoForest": cate_models.econml.EconMlEstimator2(self.cfg,

                                                                                model_name="EconML_DMLOrthoForest",

                                                                                discrete_treatment=self.discrete_treatment,

                                                                                discrete_outcome=self.discrete_outcome,

                                                                                seed=seed),

                        "EconML_DROrthoForest": cate_models.econml.EconMlEstimator2(self.cfg,

                                                                                model_name="EconML_DROrthoForest",

                                                                                discrete_treatment=self.discrete_treatment,

                                                                                discrete_outcome=self.discrete_outcome,

                                                                                seed=seed),

                        "EconML_ForestDRLearner": cate_models.econml.EconMlEstimator2(self.cfg,

                                                                                    model_name="EconML_ForestDRLearner",

                                                                                    discrete_treatment=self.discrete_treatment,

                                                                                    discrete_outcome=self.discrete_outcome,

                                                                                    seed=seed),

                        "EconML_LinearDML": cate_models.econml.EconMlEstimator2(self.cfg,

                                                                            model_name="EconML_LinearDML",

                                                                            discrete_treatment=self.discrete_treatment,

                                                                            discrete_outcome=self.discrete_outcome,

                                                                            seed=seed),

                        "EconML_LinearDRLearner": cate_models.econml.EconMlEstimator2(self.cfg,

                                                                                model_name="EconML_LinearDRLearner",

                                                                                discrete_treatment=self.discrete_treatment,

                                                                                discrete_outcome=self.discrete_outcome,

                                                                                seed=seed),

                        "EconML_SparseLinearDML": cate_models.econml.EconMlEstimator2(self.cfg,

                                                                                    model_name="EconML_SparseLinearDML",

                                                                                    discrete_treatment=self.discrete_treatment,

                                                                                    discrete_outcome=self.discrete_outcome,

                                                                                    seed=seed),

                        "EconML_SparseLinearDRLearner": cate_models.econml.EconMlEstimator2(self.cfg, 

                                                                                            model_name="EconML_SparseLinearDRLearner",

                                                                                            discrete_treatment=self.discrete_treatment,

                                                                                            discrete_outcome=self.discrete_outcome,

                                                                                            seed=seed), 

                        "EconML_SLearner_Lasso": cate_models.econml.EconMlEstimator2(self.cfg,

                                                                                    model_name="EconML_SLearner_Lasso",

                                                                                    discrete_treatment=self.discrete_treatment,

                                                                                    discrete_outcome=self.discrete_outcome,

                                                                                    seed=seed),

                        "EconML_TLearner_Lasso": cate_models.econml.EconMlEstimator2(self.cfg,

                                                                                    model_name="EconML_TLearner_Lasso",

                                                                                    discrete_treatment=self.discrete_treatment,

                                                                                    discrete_outcome=self.discrete_outcome,

                                                                                    seed=seed),

                        "EconML_XLearner_Lasso": cate_models.econml.EconMlEstimator2(self.cfg,

                                                                                    model_name="EconML_XLearner_Lasso",

                                                                                    discrete_treatment=self.discrete_treatment,

                                                                                    discrete_outcome=self.discrete_outcome,

                                                                                    seed=seed),

                        "Torch_SLearner": cate_models.torch.SLearner(num_features,

                                                                        binary_y=binary_y,

                                                                        **self.cfg.Torch_SLearner),

                        "Torch_TLearner": cate_models.torch.TLearner(num_features,

                                                                        binary_y=binary_y,

                                                                        **self.cfg.Torch_TLearner),

                        "Torch_XLearner": cate_models.torch.XLearner(num_features, 

                                                                     binary_y=binary_y, 

                                                                     **self.cfg.Torch_XLearner),

                        "Torch_DRLearner": cate_models.torch.DRLearner(num_features,

                                                                        binary_y=binary_y,

                                                                        **self.cfg.Torch_DRLearner),

                        "Torch_DragonNet": cate_models.torch.DragonNet(num_features,

                                                                        binary_y=binary_y,

                                                                        **self.cfg.Torch_DragonNet),

                        "Torch_DragonNet_2": cate_models.torch.DragonNet(num_features,

                                                                        binary_y=binary_y,

                                                                        **self.cfg.Torch_DragonNet_2),

                        "Torch_DragonNet_4": cate_models.torch.DragonNet(num_features,

                                                                        binary_y=binary_y,

                                                                        **self.cfg.Torch_DragonNet_4),                                                                                            

                        "Torch_ActionNet": cate_models.torch.ActionNet(num_features,

                                                                        binary_y=binary_y,

                                                                        **self.cfg.Torch_ActionNet),

                        # "Torch_FlexTENet": cate_models.torch.FlexTENet(num_features,

                        #                                                 binary_y=binary_y,

                        #                                                 **self.cfg.Torch_FlexTENet),

                        "Torch_PWLearner": cate_models.torch.PWLearner(num_features,

                                                                        binary_y=binary_y,

                                                                        **self.cfg.Torch_PWLearner),

                        "Torch_RALearner": cate_models.torch.RALearner(num_features,

                                                                        binary_y=binary_y,

                                                                        **self.cfg.Torch_RALearner),

                        "Torch_RLearner": cate_models.torch.RLearner(num_features,

                                                                        binary_y=binary_y,

                                                                        **self.cfg.Torch_RLearner),

                        "Torch_TARNet": cate_models.torch.TARNet(num_features,

                                                                        binary_y=binary_y,

                                                                        **self.cfg.Torch_TARNet),

                        "Torch_ULearner": cate_models.torch.ULearner(num_features,

                                                                        binary_y=binary_y,

                                                                        **self.cfg.Torch_ULearner),

                        "Torch_CFRNet_0.01": cate_models.torch.TARNet(num_features,

                                                                binary_y=binary_y,

                                                                **self.cfg.Torch_CRFNet_0_01),

                        "Torch_CFRNet_0.001": cate_models.torch.TARNet(num_features,

                                                                binary_y=binary_y,

                                                                **self.cfg.Torch_CRFNet_0_001),

                        "Torch_CFRNet_0.0001": cate_models.torch.TARNet(num_features,

                                                                binary_y=binary_y,

                                                                **self.cfg.Torch_CRFNet_0_0001),

                        "DiffPOLearner": DiffPOLearner(self.cfg,

                                                        num_features,

                                                        binary_y=binary_y,

                                                        ),

                    }

        # Deal with the case where the model is not available

        for name in self.cfg.model_names:

            if name not in learners:

                raise Exception(f"Unknown model name {name}.")

        # Only return learners from cfg.model_names

        return {k: v for k, v in learners.items() if k in self.cfg.model_names}

    def get_baseline_learners(self, 

                             seed: int = 123) -> dict:

        """

        Get baseline learners for the experiment, based on the configuration settings.

        These models will be used as a baseline for retrieving absolute outcomes from cate predictions. 

        """

        # Instantiate baseline learner for all possible treatment options

        baseline_learners = []

        if self.cfg.debug or self.cfg.dataset == "cytof_normalized_with_fastdrug" or self.cfg.dataset.startswith("melanoma"):

            cv = LeaveOneOut()

        else:

            cv = 5

        for t in range(self.cfg.simulator.num_T):

            if self.discrete_outcome == 1:

                base_model = LogisticRegressionCV(penalty='l1', solver='liblinear', cv=cv, random_state=seed)

                #base_model = RandomForestClassifier(n_estimators=100, random_state=seed)

                model = MultiOutputClassifier(base_model)

                baseline_learners.append(model)

            else:

                base_model = LassoCV(cv=cv, n_alphas=5, max_iter=50, random_state=seed)

                #base_model = RandomForestRegressor(n_estimators=100, random_state=seed)

                model = MultiOutputRegressor(base_model)

                baseline_learners.append(model)

        return baseline_learners

    def get_select_learner(self,

                            seed: int = 123) -> dict:

        """

        Get learners for feature selection.

        """

        # Instantiate baseline learner for all possible treatment options

        if self.cfg.debug:

            cv = LeaveOneOut()

        else:

            cv = 5

        select_learner = LogisticRegressionCV(penalty='l1', solver='liblinear', cv=cv, random_state=seed)

        #select_learner = RandomForestRegressor(n_estimators=100, max_depth=6)

        return select_learner

    def get_prog_learner(self,

                            seed: int = 123) -> dict:

        """

        Get learners for the prognostic part of the outcomes.

        """

        if self.cfg.debug:

            cv = LeaveOneOut()

        else:

            cv = 5

        base_model = LassoCV(cv=cv, n_alphas=5, max_iter=50, random_state=seed)

        #base_model = RandomForestRegressor(n_estimators=100, random_state=seed)

        prog_learner = MultiOutputRegressor(base_model)

        return prog_learner

    def get_pred_learners(self,

                            seed: int = 123) -> dict:

        """

        Get learners for the treatment-specific CATEs.

        """

        # Instantiate baseline learner for all possible treatment options

        pred_learners = []

        if self.cfg.debug or self.cfg.dataset.startswith("melanoma"):

            cv = LeaveOneOut()

        else:

            cv = 5

        #base_model = RandomForestRegressor(n_estimators=100, random_state=seed)

        for t in range(self.cfg.simulator.num_T):

            base_model = LassoCV(cv=cv, n_alphas=5, max_iter=50, random_state=seed)

            model = MultiOutputRegressor(base_model)

            pred_learners.append(model)

        return pred_learners

    def train_learners(self,

                        X_train: np.ndarray,

                        Y_train: np.ndarray,

                        T_train: np.ndarray,

                        outcomes_train: np.ndarray) -> None:

        """

        Train all learners.

        """

        self.training_times = {}

        for name in self.learners:

            # measure training time

            start_time = time.time()

            log.info(f"Fitting {name}.")

            if self.evaluate_inference:

                if not (name.startswith("EconML") or name.startswith("DiffPOLearner")):

                    raise ValueError("Only EconML models support inference.")

            if name == "DiffPOLearner":

                self.learners[name].train(X=X_train, y=Y_train, w=T_train, outcomes=outcomes_train) # fit before!!

            else:

                self.learners[name].train(X=X_train, y=Y_train, w=T_train) # fit before!!

            # measure training time

            end_time = time.time()

            self.training_times[name] = end_time - start_time

    def train_baseline_learners(self,

                                X_train: np.ndarray,

                                outcomes_train: np.ndarray,

                                T_train: np.ndarray) -> None:

        """

        Train all baseline learners.

        """

        for t in range(self.cfg.simulator.num_T):

            # Get all data points where treatment is t

            mask = T_train == t

            Y_train_t = outcomes_train[mask,t,:]

            X_train_t = X_train[mask,:]

            log.debug(

            f'Check baseline data for treatment {t}:'

            f'============================================'

            f'X_train: {X_train.shape}'

            f'\n{X_train}'

            f'\noutcomes_train: {outcomes_train.shape}'

            f'\n{outcomes_train}'

            f'\nT_train: {T_train.shape}'

            f'\n{T_train}'

            f'\nX_train_t: {X_train_t.shape}'

            f'\n{X_train_t}'

            f'\nY_train_t: {Y_train_t.shape}'

            f'\n{Y_train_t}'

            f'\n============================================\n\n'

            )

            # Check that there are data points for this treatment

            assert Y_train_t.shape[0] > 0, f"No data points for treatment {t}."

            log.info(f"Fitting baseline learner for treatment {t}.")

            if self.discrete_outcome:

                Y_train_t = Y_train_t.astype(bool)

            self.baseline_learners[t].fit(X_train_t, Y_train_t)

    def train_select_learner(self, 

                                X_train: np.ndarray,

                                T_train: np.ndarray) -> None:

        """

        Train the feature selection learner.

        """

        log.info(f"Fitting feature selection learner.")

        self.select_learner.fit(X_train, T_train)

    def train_prog_learner(self,

                            X_train: np.ndarray,

                            pred_outcomes_train: np.ndarray) -> None:

        """

        Train the model learning the prognostic part of the outcomes.

        Here the prognostic part is the average of all possible outcomes.

        """

        # pred_outcomes shape: n, num_T, dim_Y

        # X_train shape: n, dim_X

        # Compute the average of all possible outcomes

        prog_Y_train = np.mean(pred_outcomes_train, axis=1)

        # Train the model

        self.prog_learner.fit(X_train, prog_Y_train)

        log.debug(

            f'Check prog learner data:'

            f'============================================'

            f'X_train: {X_train.shape}'

            f'\n{X_train}'

            f'\npred_outcomes_train: {pred_outcomes_train.shape}'

            f'\n{pred_outcomes_train}'

            f'\nprog_Y_train: {prog_Y_train.shape}'

            f'\n{prog_Y_train}'

            f'\n============================================\n\n'

        )

    def train_pred_learner(self,

                            X_train: np.ndarray,

                            pred_cates_train: np.ndarray,

                            T_train: np.ndarray) -> None:

        """

        Train the model learning the CATEs.

        """

        for t in range(self.cfg.simulator.num_T):

            # Get treatment mask

            mask = T_train == t

            # For every patient with treatment t, we can simply compute the mean CATE

            # For others we need to add the cate for t, to get the desired cates (see biomarker attribution good notes)

            pred_Y_train = np.zeros((X_train.shape[0], self.cfg.simulator.dim_Y))

            pred_Y_train = pred_cates_train.mean(axis=1)

            pred_Y_train[~mask] -= pred_cates_train[~mask, t, :]

            # Train the model

            self.pred_learners[t].fit(X_train, pred_Y_train)

            log.debug(

                f'Check pred learner data for treatment {t}:'

                f'============================================'

                f'X_train: {X_train.shape}'

                f'\n{X_train[:10]}'

                f'\npred_cates_train: {pred_cates_train.shape}'

                f'\n{pred_cates_train[:10]}'

                f'\nT_train: {T_train.shape}'

                f'\n{T_train[:10]}'

                f'\npred_Y_train: {pred_Y_train.shape}'

                f'\n{pred_Y_train[:10]}'

                f'\n============================================\n\n'

            )

    def get_learner_explanations(self, 

                                 X: np.ndarray,

                                 return_explainer_names: bool = True,

                                 ignore_explainer_limit: bool = False,

                                 type: str = "pred") -> dict:

        """

        Get explanations for all learners.

        """

        learner_explainers = {}

        learner_explanations = {}

        explainer_names = {}

        for name in self.learners:

            log.info(f"Explaining {name}.")

            if ignore_explainer_limit:

                explainer_limit = X.shape[0]

            else:

                explainer_limit = self.cfg.explainer_limit

            if "EconML" in name:

                # EconML 

                if self.cfg.explainer_econml != "shap":

                    raise ValueError("Only shap is supported for EconML models.")

                if type == "pred":

                    cate_est = lambda X: self.learners[name].predict(X)

                    shap_values_avg = shap.Explainer(cate_est, X[:explainer_limit]).shap_values(X[:explainer_limit])

                    explainer_names[name] = "shap"

                    # shap_values = self.learners[name].explain(X[:explainer_limit], background_samples=None)

                    # treatment_names = self.learners[name].est.cate_treatment_names()

                    # output_names = self.learners[name].est.cate_output_names()

                    # # average absolute shaps over all treatment names

                    # shap_values_avg = np.zeros_like(shap_values[output_names[0]][treatment_names[0]].values)

                    # for output_name in output_names:

                    #     for treatment_name in treatment_names:

                    #         shap_values_avg += np.abs(shap_values[output_name][treatment_name].values)

                # Does not work yet!

                elif type == "prog":

                    if name == "EconML_SLearner_Lasso":

                        y0_est = lambda X: self.learners[name].est.overall_model.predict(np.hstack([X, np.ones((X.shape[0], 1)),np.zeros((X.shape[0], 1))]))

                        # pred outcomes shape should be: n, num_T, dim_Y

                    elif name == "EconML_TLearner_Lasso":

                        y0_est = lambda X: self.learners[name].est.models[0].predict(X)

                    elif name == "EconML_DML":

                        y0_est = lambda X: self.learners[name].predict_outcomes(X, outcomes=None)[:,0]

                    else:

                        raise ValueError(f"Model {name} not supported for prog learner explanations.")

                    shap_values = shap.Explainer(y0_est, X[:explainer_limit]).shap_values(X[:explainer_limit])

                    shap_values_avg = shap_values

                    # average absolute shaps over all treatment names

                    # shap_values_avg = np.zeros_like(shap_values[0])

                    # for i in range(len(shap_values)):

                    #     shap_values_avg += np.abs(shap_values[i])

                learner_explanations[name] = shap_values_avg

                explainer_names[name] = "shap"

            else:

                # Also doesn't work yet!

                if type == "prog" and name in ["Torch_SLearner", "Torch_TLearner", "Torch_TARNet", "Torch_CFRNet_0.001", "Torch_CFRNet_0.01","Torch_CFRNet_0.0001","Torch_DragonNet","Torch_DragonNet_2","Torch_DragonNet_4"]:

                    # model_to_explain = self.learners[name]._po_estimators[0]

                    # X_to_explain = self._repr_estimator(X).squeeze()

                    y0_est = lambda X: self.learners[name].predict(X, return_po=True)[1]

                    learner_explanations[name] = shap.Explainer(y0_est, X[:explainer_limit]).shap_values(X[:explainer_limit])

                    explainer_names[name] = "shap"  

                elif type == "prog" and not name in ["Torch_SLearner", "Torch_TLearner", "Torch_TARNet", "Torch_CFRNet_0.001", "Torch_CFRNet_0.01","Torch_CFRNet_0.0001","Torch_DragonNet","Torch_DragonNet_2","Torch_DragonNet_4"]:

                    learner_explanations[name] = None

                    explainer_names[name] = None

                    #raise ValueError(f"Model {name} not supported for prog learner explanations.")

                else:

                    cate_est = lambda X: self.learners[name].predict(X)

                    learner_explanations[name] = shap.Explainer(cate_est, X[:explainer_limit]).shap_values(X[:explainer_limit])

                    explainer_names[name] = "shap"

                    # model_to_explain = self.learners[name]

                    # X_to_explain = X

                    # learner_explainers[name] = Explainer(

                    #     model_to_explain,

                    #     feature_names=list(range(X.shape[1])),

                    #     explainer_list=[self.cfg.explainer_torch],

                    # )

                    # learner_explanations[name] = learner_explainers[name].explain(

                    #     X[: explainer_limit]

                    # )

                    # learner_explanations[name] = learner_explanations[name][self.cfg.explainer_torch]

                    # explainer_names[name] = self.cfg.explainer_torch

        # Check dimensions of explanations by looking at the shape of all explanation np arrays

        # for name in learner_explanations:

        #     log.debug(

        #         f"\nExplanations for {name} have shape {learner_explanations[name].shape}."

        #     )

        if return_explainer_names:

            return learner_explanations, explainer_names

        else:

            return learner_explanations

    def get_select_learner_explanations(self,

                                        X_reference: np.ndarray,

                                        X_to_explain: np.ndarray) -> np.ndarray:

        """

        Get explanations for the feature selection learner.

        Returns shape: n, dim_X, num_T.

        """

        explainer = shap.Explainer(self.select_learner, X_reference)

        shap_values = explainer(X_to_explain).values # Shape: n, dim_X, num_T

        return shap_values

    def get_prog_learner_explanations(self,

                                        X_reference: np.ndarray,

                                        X_to_explain: np.ndarray) -> np.ndarray:

        """

        Get explanations for the prognostic learner.

        Returns shape: n, dim_X, dim_Y.

        """

        # Get explanations for every outcome

        shap_values = np.zeros((X_to_explain.shape[0], X_to_explain.shape[1], self.cfg.simulator.dim_Y))

        for i in range(self.cfg.simulator.dim_Y):

            explainer = shap.Explainer(self.prog_learner.estimators_[i], X_reference, check_additivity=False)

            shap_values[:,:,i] = explainer(X_to_explain).values #, check_additivity=False for forest

        return shap_values

    def get_pred_learner_explanations(self,

                                        X_reference: np.ndarray,

                                        X_to_explain: np.ndarray) -> np.ndarray:

        """

        Get explanations for the treatment-specific CATEs (one vs. all).

        Returns shape: num_T, n, dim_X, dim_Y.

        """

        # Get explanations for every outcome and every reference treatment

        shap_values = np.zeros((self.cfg.simulator.num_T, X_to_explain.shape[0], X_to_explain.shape[1], self.cfg.simulator.dim_Y))

        shap_base_values = np.zeros((self.cfg.simulator.num_T, self.cfg.simulator.dim_Y))

        for t in range(self.cfg.simulator.num_T):

            for i in range(self.cfg.simulator.dim_Y):

                explainer = shap.Explainer(self.pred_learners[t].estimators_[i], X_reference)

                shap_explanation = explainer(X_to_explain)

                pred = self.pred_learners[t].predict(X_to_explain)

                shap_values[t,:,:,i] = explainer(X_to_explain).values #, check_additivity=False for forest

                shap_base_values[t,i] = shap_explanation.base_values[0]

        log.debug(

            f"\nExplanations for pred_learner have shape {shap_values.shape}."

            f"\n{shap_values[:,:10,:10,:]}"

        )

        return shap_values, shap_base_values

    def save_results(self, metrics_df: pd.DataFrame, compare_axis_values = None, plot_only: bool = False, save_df_only: bool = False, compare_axis=None, fig_type = "all") -> None:

        """

        Save results of the experiment.

        """

        file_name = self.file_name

        results_path = self.results_path

        if save_df_only:

            log.info(f"Saving intermediate results in {results_path}...")

            if not results_path.exists():

                results_path.mkdir(parents=True, exist_ok=True)

        else:

            log.info(f"Saving final results in {results_path}...")

            if not results_path.exists():

                results_path.mkdir(parents=True, exist_ok=True)

        if not plot_only:

            # Save metrics csv and metadata csv with the configuration

            if save_df_only:

                metrics_df.to_csv(

                    results_path / Path(file_name+"_checkpoint.csv")

                )

            else:

                metrics_df.to_csv(

                    results_path / Path(file_name+".csv")

                )

            OmegaConf.save(self.cfg, results_path / Path(file_name+".yaml"))

        # Choose figure size and legend position

        n_rows = len(self.cfg.metrics_to_plot)

        figsize = None #(10, 100)

        legend_position = 0.95

        log_x_axis = False

        # Save plots

        if self.cfg.plot_results and not save_df_only:

            if self.cfg.experiment_name == 'sample_size_sensitivity':

                compare_axis = "Propensity Scale"

                x_axis = "Sample Portion"

                x_label_name = r'$\omega_{\mathrm{ss}}$'

                x_values_to_plot = self.cfg.sample_sizes

            elif self.cfg.experiment_name == 'important_feature_num_sensitivity':

                compare_axis = "Feature Overlap"

                x_axis = "Num Important Features"

                x_label_name = r'$\omega_{\mathrm{IFs}}$'

                x_values_to_plot = self.cfg.important_feature_nums

            elif self.cfg.experiment_name == 'propensity_scale_sensitivity':

                # compare_axis = "Unbalancedness Exp"

                compare_axis = "Propensity Type"

                x_axis = "Propensity Scale"

                x_label_name = r'${\mathrm{\beta}}$'

                x_values_to_plot = self.cfg.propensity_scales

                log_x_axis = True

            elif self.cfg.experiment_name == 'predictive_scale_sensitivity':

                compare_axis = "Nonlinearity Scale"

                x_axis = "Predictive Scale"

                x_label_name = r'$\omega_{\mathrm{pds}}$'

                x_values_to_plot = self.cfg.predictive_scales

            elif self.cfg.experiment_name == 'treatment_space_sensitivity':

                compare_axis = "Binary Outcome"

                x_axis = "Treatment Options"

                x_label_name = r'$\omega_{\mathrm{Ts}}$'

                x_values_to_plot = self.cfg.num_Ts

            elif self.cfg.experiment_name == 'data_dimensionality_sensitivity':

                compare_axis = compare_axis

                x_axis = "Data Dimension"

                x_label_name = r'$Num Features$'

                x_values_to_plot = self.cfg.data_dims

            elif self.cfg.experiment_name == 'feature_overlap_sensitivity':

                compare_axis = "Overlap Type"

                x_axis = "Propensity Scale"

                x_label_name = r'$\omega_{\mathrm{pps}}$'

                x_values_to_plot = self.cfg.propensity_scales

            elif self.cfg.experiment_name == 'expertise_sensitivity':

                compare_axis = "Propensity Type"

                x_axis = "Alpha"

                x_label_name = r'$\alpha$'

                x_values_to_plot = self.cfg.alphas

            else:

                raise ValueError(f"Experiment {self.cfg.experiment_name} not supported for plotting.")

            if fig_type == "all":

                fig = plot_results_datasets_compare(

                            results_df=metrics_df,

                            model_names=self.cfg.model_names,

                            dataset=self.cfg.dataset,

                            compare_axis=compare_axis,

                            compare_axis_values=compare_axis_values,

                            x_axis=x_axis,

                            x_label_name=x_label_name,

                            x_values_to_plot=x_values_to_plot,

                            metrics_list=self.cfg.metrics_to_plot,

                            learners_list=self.cfg.model_names,

                            figsize=figsize,

                            legend_position=legend_position,

                            seeds_list=self.cfg.seeds,

                            n_splits=self.cfg.n_splits,

                            sharey="row",

                            legend_rows=1,

                            dim_X=self.X.shape[0],

                            log_x_axis=log_x_axis

                        )

            elif fig_type == "expertise":

                fig = plot_expertise_metrics(

                            results_df=metrics_df,

                            model_names=self.cfg.model_names,

                            dataset=self.cfg.dataset,

                            compare_axis=compare_axis,

                            compare_axis_values=compare_axis_values,

                            x_axis=x_axis,

                            x_label_name=x_label_name,

                            x_values_to_plot=x_values_to_plot,

                            metrics_list=self.cfg.metrics_to_plot,

                            learners_list=self.cfg.model_names,

                            figsize=figsize,

                            legend_position=legend_position,

                            seeds_list=self.cfg.seeds,

                            n_splits=self.cfg.n_splits,

                            sharey="row",

                            legend_rows=1,

                            dim_X=self.X.shape[0],

                            log_x_axis=log_x_axis

                        )

            elif fig_type == "performance":

                fig = plot_performance_metrics(

                            results_df=metrics_df,

                            model_names=self.cfg.model_names,

                            dataset=self.cfg.dataset,

                            compare_axis=compare_axis,

                            compare_axis_values=compare_axis_values,

                            x_axis=x_axis,

                            x_label_name=x_label_name,

                            x_values_to_plot=x_values_to_plot,

                            metrics_list=self.cfg.metrics_to_plot,

                            learners_list=self.cfg.model_names,

                            figsize=figsize,

                            legend_position=legend_position,

                            seeds_list=self.cfg.seeds,

                            n_splits=self.cfg.n_splits,

                            sharey="row",

                            legend_rows=1,

                            dim_X=self.X.shape[0],

                            log_x_axis=log_x_axis

                        )

            elif fig_type == "performance_f_cf":

                fig = plot_performance_metrics_f_cf(

                            results_df=metrics_df,

                            model_names=self.cfg.model_names,

                            dataset=self.cfg.dataset,

                            compare_axis=compare_axis,

                            compare_axis_values=compare_axis_values,

                            x_axis=x_axis,

                            x_label_name=x_label_name,

                            x_values_to_plot=x_values_to_plot,

                            metrics_list=self.cfg.metrics_to_plot,

                            learners_list=self.cfg.model_names,

                            figsize=figsize,

                            legend_position=legend_position,

                            seeds_list=self.cfg.seeds,

                            n_splits=self.cfg.n_splits,

                            sharey="row",

                            legend_rows=1,

                            dim_X=self.X.shape[0],

                            log_x_axis=log_x_axis

                        )

            # Save figure

            try:

                fig.savefig(results_path / f"{self.cfg.plot_name_prefix}.png", bbox_inches='tight')

            except:

                fig.savefig(results_path / f"{file_name}.png", bbox_inches='tight')

    def load_and_plot_results(self, compare_axis_values, fig_type: str = "all") -> None:

        """

        Load and plot results of the experiment.

        """

        file_name = self.file_name

        results_path = self.results_path

        log.info(f"Loading results from {results_path}...")

        if not results_path.exists():

            raise FileNotFoundError(f"Results path {results_path} does not exist.")

        # Load metrics

        metrics_df = pd.read_csv(results_path / Path(file_name+"_checkpoint.csv"))

        if self.cfg.experiment_name == "data_dimensionality_sensitivity":

            if self.cfg.compare_axis == "propensity":

                compare_axis = "Propensity Scale"

            elif self.cfg.compare_axis == "num_features":

                compare_axis = "Num Important Features"

            self.save_results(metrics_df, plot_only=True, compare_axis=compare_axis, compare_axis_values=compare_axis_values, fig_type=fig_type)

        else:

            self.save_results(metrics_df, plot_only=True, compare_axis_values=compare_axis_values, fig_type=fig_type)

    def compute_outcome_mse(self,

                            pred_outcomes: np.ndarray,

                            true_outcomes: np.ndarray,

                            T: np.ndarray = None) -> float:

        """

        Compute MSE for all outcomes.

        """

        # Only keep factual cates

        pred_outcomes_factual = np.zeros((pred_outcomes.shape[0], pred_outcomes.shape[2])) # n, dim_Y

        true_outcomes_factual = np.zeros((true_outcomes.shape[0], true_outcomes.shape[2]))

        pred_outcomes_factual_Y0 = np.zeros(((T==0).sum(), pred_outcomes.shape[2])) # n, dim_Y

        true_outcomes_factual_Y0 = np.zeros(((T==0).sum(), pred_outcomes.shape[2]))

        pred_outcomes_factual_Y1 = np.zeros(((T==1).sum(), pred_outcomes.shape[2])) # n, dim_Y

        true_outcomes_factual_Y1 = np.zeros(((T==1).sum(), pred_outcomes.shape[2]))

        pred_outcomes_cf = np.zeros((pred_outcomes.shape[0], pred_outcomes.shape[1]-1, pred_outcomes.shape[2])) # n, num_T-1, dim_Y

        true_outcomes_cf = np.zeros((true_outcomes.shape[0], true_outcomes.shape[1]-1, true_outcomes.shape[2]))

        pred_outcomes_cf_Y0 = np.zeros(((T==1).sum(), pred_outcomes.shape[1]-1, pred_outcomes.shape[2])) # n, dim_Y

        true_outcomes_cf_Y0 = np.zeros(((T==1).sum(), pred_outcomes.shape[1]-1, pred_outcomes.shape[2]))

        pred_outcomes_cf_Y1 = np.zeros(((T==0).sum(), pred_outcomes.shape[1]-1, pred_outcomes.shape[2])) # n, dim_Y

        true_outcomes_cf_Y1 = np.zeros(((T==0).sum(), pred_outcomes.shape[1]-1, pred_outcomes.shape[2]))

        counter_Y0 = 0

        counter_Y1 = 0

        for i in range(pred_outcomes.shape[0]):

            mask_factual = np.zeros(pred_outcomes.shape[1], dtype=bool)

            mask_cf = np.ones(pred_outcomes.shape[1], dtype=bool)

            mask_factual[T[i]] = True

            mask_cf[T[i]] = False

            pred_outcomes_factual[i,:] = pred_outcomes[i, mask_factual,:]

            true_outcomes_factual[i,:] = true_outcomes[i, mask_factual,:]

            pred_outcomes_cf[i,:,:] = pred_outcomes[i, mask_cf,:]

            true_outcomes_cf[i,:,:] = true_outcomes[i, mask_cf,:]

            if T[i] == 0:

                pred_outcomes_factual_Y0[counter_Y0,:] = pred_outcomes[i, mask_factual,:]

                true_outcomes_factual_Y0[counter_Y0,:] = true_outcomes[i, mask_factual,:]

                pred_outcomes_cf_Y1[counter_Y0,:,:] = pred_outcomes[i,mask_cf,:]

                true_outcomes_cf_Y1[counter_Y0,:,:] = true_outcomes[i,mask_cf,:]

                counter_Y0 += 1

            else:

                pred_outcomes_factual_Y1[counter_Y1,:] = pred_outcomes[i, mask_factual,:]

                true_outcomes_factual_Y1[counter_Y1,:] = true_outcomes[i, mask_factual,:]

                pred_outcomes_cf_Y0[counter_Y1,:,:] = pred_outcomes[i,mask_cf,:]

                true_outcomes_cf_Y0[counter_Y1,:,:] = true_outcomes[i,mask_cf,:]

                counter_Y1 += 1

        rmse_factual = mean_squared_error(true_outcomes_factual.reshape(-1), pred_outcomes_factual.reshape(-1), squared=False)

        rmse_factual_Y0 = mean_squared_error(true_outcomes_factual_Y0.reshape(-1), pred_outcomes_factual_Y0.reshape(-1), squared=False)

        rmse_factual_Y1 = mean_squared_error(true_outcomes_factual_Y1.reshape(-1), pred_outcomes_factual_Y1.reshape(-1), squared=False)

        rmse_cf = mean_squared_error(true_outcomes_cf.reshape(-1), pred_outcomes_cf.reshape(-1), squared=False)

        rmse_cf_Y0 = mean_squared_error(true_outcomes_cf_Y0.reshape(-1), pred_outcomes_cf_Y0.reshape(-1), squared=False)

        rmse_cf_Y1 = mean_squared_error(true_outcomes_cf_Y1.reshape(-1), pred_outcomes_cf_Y1.reshape(-1), squared=False)

        rmse_Y0 = mean_squared_error(np.concatenate([true_outcomes_factual_Y0.reshape(-1), true_outcomes_cf_Y0.reshape(-1)]), np.concatenate([pred_outcomes_factual_Y0.reshape(-1), pred_outcomes_cf_Y0.reshape(-1)]), squared=False)

        rmse_Y1 = mean_squared_error(np.concatenate([true_outcomes_factual_Y1.reshape(-1), true_outcomes_cf_Y1.reshape(-1)]), np.concatenate([pred_outcomes_factual_Y1.reshape(-1), pred_outcomes_cf_Y1.reshape(-1)]), squared=False)

        # Get variance of true outcomes, factual and cf

        # factual_std = np.std(true_outcomes_factual)

        # cf_std = np.std(true_outcomes_cf)

        factual_std = np.var(true_outcomes_factual)

        cf_std = np.var(true_outcomes_cf)

        log.debug(

            f"\nPred outcomes: {pred_outcomes.shape}: \n{pred_outcomes}"

            f"\n\nT: \n{T}"

            f"\n\nPred outcomes factual: {pred_outcomes_factual.shape}: \n{pred_outcomes_factual}"

        )

        return rmse_Y0, rmse_Y1, rmse_factual_Y0, rmse_factual_Y1, rmse_cf_Y0, rmse_cf_Y1, rmse_factual, rmse_cf, rmse_factual/factual_std, rmse_cf/cf_std, np.mean(true_outcomes_factual), np.mean(true_outcomes_cf), np.std(true_outcomes_factual), np.std(true_outcomes_cf)

    def compute_outcome_auroc(self,

                              pred_outcomes: np.ndarray,

                              true_outcomes: np.ndarray,

                              T: np.ndarray = None) -> float:

        """

        Compute AUROC for all outcomes.

        """

        # Only keep factual cates

        pred_outcomes_factual = np.zeros((pred_outcomes.shape[0], pred_outcomes.shape[2])) # n, dim_Y

        true_outcomes_factual = np.zeros((true_outcomes.shape[0], true_outcomes.shape[2]))

        pred_outcomes_cf = np.zeros((pred_outcomes.shape[0], pred_outcomes.shape[1]-1, pred_outcomes.shape[2])) # n, num_T-1, dim_Y

        true_outcomes_cf = np.zeros((true_outcomes.shape[0], pred_outcomes.shape[1]-1, true_outcomes.shape[2]))

        for i in range(pred_outcomes.shape[0]):

            mask_factual = np.zeros(pred_outcomes.shape[1], dtype=bool)

            mask_cf = np.ones(pred_outcomes.shape[1], dtype=bool)

            mask_factual[T[i]] = True

            mask_cf[T[i]] = False

            pred_outcomes_factual[i,:] = pred_outcomes[i, mask_factual,:]

            true_outcomes_factual[i,:] = true_outcomes[i, mask_factual,:]

            pred_outcomes_cf[i,:,:] = pred_outcomes[i, mask_cf,:]

            true_outcomes_cf[i,:,:] = true_outcomes[i, mask_cf,:]

        auroc_factual = roc_auc_score(true_outcomes_factual.reshape(-1), pred_outcomes_factual.reshape(-1))

        auroc_cf = roc_auc_score(true_outcomes_cf.reshape(-1), pred_outcomes_cf.reshape(-1))

        log.debug(

            f"\nPred outcomes: {pred_outcomes.shape}: \n{pred_outcomes}"

            f"\n\nT: \n{T}"

            f"\n\nPred outcomes factual: {pred_outcomes_factual.shape}: \n{pred_outcomes_factual}"

        )

        return auroc_factual, auroc_cf

        # if factual:

        #     # Only keep factual cates

        #     pred_outcomes_factual = np.zeros((pred_outcomes.shape[0], pred_outcomes.shape[2])) # dim_X, num_T, dim_Y

        #     true_outcomes_factual = np.zeros((true_outcomes.shape[0], true_outcomes.shape[2]))

        #     for i in range(pred_outcomes.shape[0]):

        #         mask = np.zeros(pred_outcomes.shape[1], dtype=bool)

        #         mask[T[i]] = True

        #         pred_outcomes_factual[i,:] = pred_outcomes[i, mask,:]

        #         true_outcomes_factual[i,:] = true_outcomes[i, mask,:]

        #     # f1 = f1_score(true_outcomes_factual.reshape(-1), pred_outcomes_factual.reshape(-1))

        #     f1 = roc_auc_score(true_outcomes_factual.reshape(-1), pred_outcomes_factual.clip(0,1).reshape(-1))

        # else:

        #     auroc = roc_auc_score(true_outcomes.reshape(-1), pred_outcomes.clip(0,1).reshape(-1))

        #     f1 = auroc

        #     #f1 = f1_score(true_outcomes.reshape(-1), pred_outcomes.reshape(-1))    

        # auroc = f1 #TODO: Change name to f1

        # return auroc

    def compute_overall_pehe(self,

                                pred_cates: np.ndarray,

                                true_cates: np.ndarray,

                                T: np.ndarray) -> float:

        """

        Compute average PEHE across all treatment options for all outcomes.

        """

        # Remove cates where basline and assigned treatment are the same - they evaluate to zero and are not informative

        if self.discrete_outcome:

            pred_cates = pred_cates.clip(-1,1)

        pehe_sum_total = 0

        pehe_normalized_sum_total = 0

        mean_sum_total = 0

        std_sum_total = 0

        for outcome_idx in range(pred_cates.shape[2]):

            pred_cates_curr = pred_cates[:,:,outcome_idx]

            true_cates_curr = true_cates[:,:,outcome_idx]

            counter = 0

            pehe_sum = 0

            pehe_normalized_sum = 0

            mean_sum = 0

            std_sum = 0

            for i in range(pred_cates.shape[1]):

                for j in range(i):

                    mask_j = T == j 

                    mask_i = T == i

                    n = np.sum(mask_j) + np.sum(mask_i)

                    counter += 1

                    pred_cates_curr_cate_j = pred_cates_curr[mask_j,i]

                    true_cates_curr_cate_j = true_cates_curr[mask_j,i]

                    # pred_cates_curr_cate_i = pred_cates_curr[mask_i,j]

                    # true_cates_curr_cate_i = true_cates_curr[mask_i,j]

                    pred_cates_curr_cate_i = -pred_cates_curr[mask_i,j] # Make sure to always record cate in same direction

                    true_cates_curr_cate_i = -true_cates_curr[mask_i,j] # TODO: Make sure this makes sense for multiple treatments & outcomes

                    pred_cates_curr_cate = np.concatenate((pred_cates_curr_cate_j, pred_cates_curr_cate_i)).reshape(-1)

                    true_cates_curr_cate = np.concatenate((true_cates_curr_cate_j, true_cates_curr_cate_i)).reshape(-1)

                    # Compute mean CATE

                    true_cates_mean = np.mean(true_cates_curr_cate)

                    mean_sum += true_cates_mean

                    # Compute std of CATE

                    true_cates_std = np.std(true_cates_curr_cate)

                    std_sum += true_cates_std

                    pehe_curr = mean_squared_error(true_cates_curr_cate, pred_cates_curr_cate)

                    pehe_sum += pehe_curr

                    pehe_normalized_sum += pehe_curr / np.var(true_cates_curr_cate)

                    log.debug(

                        f'Check pehe computation for outcome: {outcome_idx}, cate: {i}-{j}'

                        f'============================================'

                        f'\npred_cates: {pred_cates.shape}'

                        f'\n{pred_cates}'

                        f'\n\ntrue_cates: {true_cates.shape}'

                        f'\n{true_cates}'

                        f'\n\nT: {T.shape}'

                        f'\n{T}'

                        f'\n\npred_cates_curr: {pred_cates_curr.shape}'

                        f'\n{pred_cates_curr}'

                        f'\n\ntrue_cates_curr: {true_cates_mean.shape}'

                        f'\n{true_cates_curr}'

                        f'\n\nmask_i: {mask_i.shape}'

                        f'\n{mask_i}'

                        f'\n\nmask_j: {mask_j.shape}'

                        f'\n{mask_j}'

                        f'\n\pred_cates_curr_cate: {pred_cates_curr_cate.shape}'

                        f'\n{pred_cates_curr_cate}'

                        f'\n\ntrue_cates_curr_cate: {true_cates_curr_cate.shape}'

                        f'\n{true_cates_curr_cate}'

                        f'\n\ncounter is currently: {counter}'

                        f'\n\npehe_curr: {pehe_curr}'

                        f'\n\nmean_sum: {mean_sum}'

                        f'\n\nstd_sum: {std_sum}'

                        f'\n============================================\n\n'

                    )

            pehe = pehe_sum / counter

            pehe_normalized = pehe_normalized_sum / counter

            pehe_sum_total += pehe

            pehe_normalized_sum_total += pehe_normalized

            mean_sum_total += mean_sum / counter

            std_sum_total += std_sum / counter

        pehe_total = pehe_sum_total / pred_cates.shape[2]

        pehe_normalized_total = pehe_normalized_sum_total / pred_cates.shape[2]

        true_cates_mean_total = mean_sum_total / pred_cates.shape[2]

        true_cates_std_total = std_sum_total / pred_cates.shape[2]

        # pred_cates_filt = np.zeros((pred_cates.shape[0], pred_cates.shape[1]-1, pred_cates.shape[2])) # dim_X, num_T, dim_Y

        # true_cates_filt = np.zeros((true_cates.shape[0], true_cates.shape[1]-1, true_cates.shape[2]))

        # for i in range(pred_cates.shape[0]):

        #     mask = np.ones(pred_cates.shape[1], dtype=bool)

        #     mask[T[i]] = False

        #     pred_cates_filt[i,:,:] = pred_cates[i, mask,:]

        #     true_cates_filt[i,:,:] = true_cates[i, mask,:]

        # # Reshape to get all cates for each outcome in one dimension

        # pred_cates_filt = pred_cates_filt.reshape(-1)

        # true_cates_filt = true_cates_filt.reshape(-1)

        # # Compute PEHE 

        # pehe = mean_squared_error(true_cates_filt, pred_cates_filt)

        # Compute mean CATE

        # pred_cates_mean = np.mean(pred_cates_filt)

        # true_cates_mean = np.mean(true_cates_filt)

        # # Compute std of CATE

        # pred_cates_std = np.std(pred_cates_filt)