from pathlib import Path

import os

import catenets.models as cate_models

import numpy as np

import pandas as pd

import src.iterpretability.logger as log

from src.iterpretability.explain import Explainer

from src.iterpretability.datasets.data_loader import load

from src.iterpretability.synthetic_simulate import (

    SyntheticSimulatorLinear,

    SyntheticSimulatorModulatedNonLinear,

)

from src.iterpretability.utils import (

    attribution_accuracy,

    compute_pehe,

)

# For contour plotting

import umap 

from sklearn.decomposition import PCA

from sklearn.manifold import TSNE

from sklearn.linear_model import LogisticRegression

import matplotlib.pyplot as plt

import matplotlib.tri as tri

import imageio

import torch

import shap

def get_learners(model_list, X_train, Y_train, n_iter, batch_size, batch_norm, discrete_treatment=True, discrete_outcome=False):

    learners = {

                    "TLearner": cate_models.torch.TLearner(

                        X_train.shape[1],

                        binary_y=(len(np.unique(Y_train)) == 2),

                        n_layers_out=2,

                        n_units_out=100,

                        batch_size=batch_size,

                        n_iter=n_iter,

                        batch_norm=batch_norm,

                        nonlin="relu",

                    ),

                    "SLearner": cate_models.torch.SLearner(

                        X_train.shape[1],

                        binary_y=(len(np.unique(Y_train)) == 2),

                        n_layers_out=2,

                        n_units_out=100,

                        n_iter=n_iter,

                        batch_size=batch_size,

                        batch_norm=batch_norm,

                        nonlin="relu",

                    ),

                    "TARNet": cate_models.torch.TARNet(

                        X_train.shape[1],

                        binary_y=(len(np.unique(Y_train)) == 2),

                        n_layers_r=1,

                        n_layers_out=1,

                        n_units_out=100,

                        n_units_r=100,

                        batch_size=batch_size,

                        n_iter=n_iter,

                        batch_norm=batch_norm,

                        nonlin="relu",

                    ),

                    "DRLearner": cate_models.torch.DRLearner(

                        X_train.shape[1],

                        binary_y=(len(np.unique(Y_train)) == 2),

                        n_layers_out=2,

                        n_units_out=100,

                        n_iter=n_iter,

                        batch_size=batch_size,

                        batch_norm=batch_norm,

                        nonlin="relu",

                    ),

                    "XLearner": cate_models.torch.XLearner(

                        X_train.shape[1],

                        binary_y=(len(np.unique(Y_train)) == 2),

                        n_layers_out=2,

                        n_units_out=100,

                        n_iter=n_iter,

                        batch_size=batch_size,

                        batch_norm=batch_norm,

                        nonlin="relu",

                    ),

                    "CFRNet_0.01": cate_models.torch.TARNet(

                        X_train.shape[1],

                        binary_y=(len(np.unique(Y_train)) == 2),

                        n_layers_r=1,

                        n_layers_out=1,

                        n_units_out=100,

                        n_units_r=100,

                        batch_size=batch_size,

                        n_iter=n_iter,

                        batch_norm=batch_norm,

                        nonlin="relu",

                        penalty_disc=0.01,

                    ),

                    "CFRNet_0.001": cate_models.torch.TARNet(

                        X_train.shape[1],

                        binary_y=(len(np.unique(Y_train)) == 2),

                        n_layers_r=1,

                        n_layers_out=1,

                        n_units_out=100,

                        n_units_r=100,

                        batch_size=batch_size,

                        n_iter=n_iter,

                        batch_norm=batch_norm,

                        nonlin="relu",

                        penalty_disc=0.001,

                    ),

                    "CFRNet_0.0001": cate_models.torch.TARNet(

                        X_train.shape[1],

                        binary_y=(len(np.unique(Y_train)) == 2),

                        n_layers_r=1,

                        n_layers_out=1,

                        n_units_out=100,

                        n_units_r=100,

                        batch_size=batch_size,

                        n_iter=n_iter,

                        batch_norm=batch_norm,

                        nonlin="relu",

                        penalty_disc=0.0001,

                    ),

                    "EconML_CausalForestDML": cate_models.econml.EconMlEstimator(

                        model_name="EconML_CausalForestDML",

                        # cv_model_selection=3,

                        discrete_treatment=discrete_treatment,

                        discrete_outcome=discrete_outcome,

                    ),

                    "EconML_DMLOrthoForest": cate_models.econml.EconMlEstimator(

                        model_name="EconML_DMLOrthoForest",

                        # cv_model_selection=3,

                        discrete_treatment=discrete_treatment,

                        discrete_outcome=discrete_outcome,

                    ),

                    "EconML_SparseLinearDML": cate_models.econml.EconMlEstimator(

                        model_name="EconML_SparseLinearDML",

                        # cv_model_selection=3,

                        discrete_treatment=discrete_treatment,

                        discrete_outcome=discrete_outcome,

                    ),

                    "EconML_SparseLinearDRLearner": cate_models.econml.EconMlEstimator(

                        model_name="EconML_SparseLinearDRLearner",

                        # cv_model_selection=3,

                        discrete_treatment=discrete_treatment,

                        discrete_outcome=discrete_outcome,

                    ),

                    "EconML_LinearDRLearner": cate_models.econml.EconMlEstimator(

                        model_name="EconML_LinearDRLearner",

                        # cv_model_selection=3,

                        discrete_treatment=discrete_treatment,

                        discrete_outcome=discrete_outcome,

                    ),

                    "EconML_DRLearner": cate_models.econml.EconMlEstimator(

                        model_name="EconML_DRLearner",

                        # cv_model_selection=3,

                        discrete_treatment=discrete_treatment,

                        discrete_outcome=discrete_outcome,

                    ),

                    "EconML_XLearner": cate_models.econml.EconMlEstimator(

                        model_name="EconML_XLearner",

                        # cv_model_selection=3,

                        discrete_treatment=discrete_treatment,

                        discrete_outcome=discrete_outcome,

                    ),

                    "EconML_SLearner": cate_models.econml.EconMlEstimator(

                        model_name="EconML_SLearner",

                        # cv_model_selection=3,

                        discrete_treatment=discrete_treatment,

                        discrete_outcome=discrete_outcome,

                    ),

                    "EconML_TLearner": cate_models.econml.EconMlEstimator(

                        model_name="EconML_TLearner",

                        # cv_model_selection=3,

                        discrete_treatment=discrete_treatment,

                        discrete_outcome=discrete_outcome,

                    ),

                    "EconML_SparseLinearDRIV": cate_models.econml.EconMlEstimator(

                        model_name="EconML_SparseLinearDRIV",

                        # cv_model_selection=3,

                        discrete_treatment=discrete_treatment,

                        discrete_outcome=discrete_outcome,

                    ),

                }

    for name in model_list:

        if name not in learners:

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

    # Only return the learners that are in the model_list

    learners = {name: learners[name] for name in model_list}

    return learners

def get_learner_explanations(learners, X_test, X_train, Y_train, W_train, explainer_limit, explainer_list, return_learners=False, already_trained=False):

    learner_explainers = {}

    learner_explanations = {}

    for name in learners:

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

        if not already_trained:

            learners[name].fit(X=X_train, y=Y_train, w=W_train)

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

        if "EconML" in name:

            shap_values = learners[name].est.shap_values(X_test[:explainer_limit], background_samples=None)

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

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

            output_name = output_names[0]

            treatment_name = treatment_names[0]

            learner_explanations[name] = {"kernel_shap" : shap_values[output_name][treatment_name].values} 

        else:

            learner_explainers[name] = Explainer(

                learners[name],

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

                explainer_list=explainer_list,

            )

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

                X_test[: explainer_limit]

            )

    if return_learners:

        return learner_explanations, learners

    else:

        return learner_explanations

class PredictiveSensitivity:

    """

    Sensitivity analysis for predictive scale.

    """

    def __init__(

        self,

        n_units_hidden: int = 50,

        n_layers: int = 1,

        penalty_orthogonal: float = 0.01,

        batch_size: int = 1024,

        batch_norm: bool = False,

        n_iter: int = 1000,

        seed: int = 42,

        explainer_limit: int = 1000,

        save_path: Path = Path.cwd(),

        propensity_type: str = "pred",

        predictive_scales: list = [1e-3, 1e-2, 1e-1, 0.5, 1, 2],

        num_interactions: int = 1,

        synthetic_simulator_type: str = "linear",

        selection_type: str = "random",

        non_linearity_scale: float = 0,

        model_list: list = ["TLearner"]

    ) -> None:

        self.n_units_hidden = n_units_hidden

        self.n_layers = n_layers

        self.penalty_orthogonal = penalty_orthogonal

        self.batch_size = batch_size

        self.batch_norm = batch_norm

        self.n_iter = n_iter

        self.seed = seed

        self.explainer_limit = explainer_limit

        self.save_path = save_path

        self.predictive_scales = predictive_scales

        self.propensity_type = propensity_type

        self.num_interactions = num_interactions

        self.synthetic_simulator_type = synthetic_simulator_type

        self.selection_type = selection_type

        self.non_linearity_scale = non_linearity_scale

        self.model_list = model_list

    def run(

        self,

        dataset: str = "tcga_10",

        train_ratio: float = 0.8,

        num_important_features: int = 2,

        binary_outcome: bool = False,

        random_feature_selection: bool = True,

        explainer_list: list = [

            "feature_ablation",

            "feature_permutation",

            "integrated_gradients",

            "shapley_value_sampling",

        ],

        debug: bool = False,

        directory_path_: str = None,

    ) -> None:

        log.info(

            f"Using dataset {dataset} with num_important features = {num_important_features}."

        )

        X_raw_train, X_raw_test = load(dataset, train_ratio=train_ratio, debug=debug, directory_path_=directory_path_)

        if self.synthetic_simulator_type == "linear":

            sim = SyntheticSimulatorLinear(

                X_raw_train,

                num_important_features=num_important_features,

                random_feature_selection=random_feature_selection,

                seed=self.seed,

            )

        elif self.synthetic_simulator_type == "nonlinear":

            sim = SyntheticSimulatorModulatedNonLinear(

                X_raw_train,

                num_important_features=num_important_features,

                non_linearity_scale=self.non_linearity_scale,

                seed=self.seed,

                selection_type=self.selection_type,

            )

        else:

            raise Exception("Unknown simulator type.")

        explainability_data = []

        for predictive_scale in self.predictive_scales:

            log.info(f"Now working with predictive_scale = {predictive_scale}...")

            (

                X_train,

                W_train,

                Y_train,

                po0_train,

                po1_train,

                propensity_train,

            ) = sim.simulate_dataset(

                X_raw_train,

                predictive_scale=predictive_scale,

                binary_outcome=binary_outcome,

                treatment_assign=self.propensity_type,

            )

            X_test, W_test, Y_test, po0_test, po1_test, _ = sim.simulate_dataset(

                X_raw_test,

                predictive_scale=predictive_scale,

                binary_outcome=binary_outcome,

                treatment_assign=self.propensity_type,

            )

            log.info("Fitting and explaining learners...")

            learners = get_learners(

                model_list=self.model_list,

                X_train=X_train,

                Y_train=Y_train,

                n_iter=self.n_iter,

                batch_size=self.batch_size,

                batch_norm=self.batch_norm,

                discrete_outcome=binary_outcome

            )

            learner_explanations = get_learner_explanations(learners, 

                                                            X_test, X_train, Y_train, W_train, 

                                                            self.explainer_limit, explainer_list)

            all_important_features = sim.get_all_important_features(with_selective=True)

            pred_features = sim.get_predictive_features()

            prog_features = sim.get_prognostic_features()

            select_features = sim.get_selective_features()

            cate_test = sim.te(X_test)

            for explainer_name in explainer_list:

                for learner_name in learners:

                    attribution_est = np.abs(

                        learner_explanations[learner_name][explainer_name]

                    )

                    acc_scores_all_features = attribution_accuracy(

                        all_important_features, attribution_est

                    )

                    acc_scores_predictive_features = attribution_accuracy(

                        pred_features, attribution_est

                    )

                    acc_scores_prog_features = attribution_accuracy(

                        prog_features, attribution_est

                    )

                    acc_scores_selective_features = attribution_accuracy(

                        select_features, attribution_est

                    )

                    cate_pred = learners[learner_name].predict(X=X_test)

                    pehe_test = compute_pehe(cate_true=cate_test, cate_pred=cate_pred)

                    explainability_data.append(

                        [

                            predictive_scale,

                            learner_name,

                            explainer_name,

                            acc_scores_all_features,

                            acc_scores_predictive_features,

                            acc_scores_prog_features,

                            acc_scores_selective_features,

                            pehe_test,

                            np.mean(cate_test),

                            np.var(cate_test),

                            pehe_test / np.sqrt(np.var(cate_test)),

                        ]

                    )

        metrics_df = pd.DataFrame(

            explainability_data,

            columns=[

                "Predictive Scale",

                "Learner",

                "Explainer",

                "All features ACC",

                "Pred features ACC",

                "Prog features ACC",

                "Select features ACC",

                "PEHE",

                "CATE true mean",

                "CATE true var",

                "Normalized PEHE",

            ],

        )

        results_path = self.save_path / "results/predictive_sensitivity"

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

        if not results_path.exists():

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

        metrics_df.to_csv(

            results_path / f"predictive_scale_{dataset}_{num_important_features}_"

            f"{self.synthetic_simulator_type}_random_{random_feature_selection}_"

            f"binary_{binary_outcome}-seed{self.seed}.csv"

        )

class NonLinearitySensitivity:

    """

    Sensitivity analysis for nonlinearity in prognostic and predictive functions.

    """

    def __init__(

        self,

        n_units_hidden: int = 50,

        n_layers: int = 1,

        penalty_orthogonal: float = 0.01,

        batch_size: int = 1024,

        batch_norm: bool = False,

        n_iter: int = 1000,

        seed: int = 42,

        explainer_limit: int = 1000,

        save_path: Path = Path.cwd(),

        propensity_type: str = "pred",

        nonlinearity_scales: list = [0.0, 0.2, 0.5, 0.7, 1.0],

        selection_type: str = "random",

        predictive_scale: float = 1,

        synthetic_simulator_type: str = "random",

        model_list: list = ["TLearner"]

    ) -> None:

        self.n_units_hidden = n_units_hidden

        self.n_layers = n_layers

        self.penalty_orthogonal = penalty_orthogonal

        self.batch_size = batch_size

        self.batch_norm = batch_norm

        self.n_iter = n_iter

        self.seed = seed

        self.explainer_limit = explainer_limit

        self.save_path = save_path

        self.propensity_type = propensity_type

        self.nonlinearity_scales = nonlinearity_scales

        self.selection_type = selection_type

        self.predictive_scale = predictive_scale

        self.synthetic_simulator_type = synthetic_simulator_type

        self.model_list = model_list

    def run(

        self,

        dataset: str = "tcga_100",

        num_important_features: int = 15,

        explainer_list: list = [

            "feature_ablation",

            "feature_permutation",

            "integrated_gradients",

            "shapley_value_sampling",

        ],

        train_ratio: float = 0.8,

        binary_outcome: bool = False,

        debug=False,

        directory_path_: str = None,

    ) -> None:

        log.info(

            f"Using dataset {dataset} with num_important features = {num_important_features}."

        )

        X_raw_train, X_raw_test = load(dataset, train_ratio=train_ratio, debug=debug, directory_path_=directory_path_)

        explainability_data = []

        for nonlinearity_scale in self.nonlinearity_scales:

            log.info(f"Now working with a nonlinearity scale {nonlinearity_scale}...")

            if self.synthetic_simulator_type == "linear":

                raise Exception("Linear simulator not supported for nonlinearity sensitivity.")

            elif self.synthetic_simulator_type == "nonlinear":

                sim = SyntheticSimulatorModulatedNonLinear(

                    X_raw_train,

                    num_important_features=num_important_features,

                    non_linearity_scale=nonlinearity_scale,

                    seed=self.seed,

                    selection_type=self.selection_type

                )

            else:

                raise Exception("Unknown simulator type.")

            (

                X_train,

                W_train,

                Y_train,

                po0_train,

                po1_train,

                propensity_train,

            ) = sim.simulate_dataset(

                X_raw_train,

                predictive_scale=self.predictive_scale,

                binary_outcome=binary_outcome,

                treatment_assign=self.propensity_type,

            )

            X_test, W_test, Y_test, po0_test, po1_test, _ = sim.simulate_dataset(

                X_raw_test,

                predictive_scale=self.predictive_scale,

                binary_outcome=binary_outcome,

                treatment_assign=self.propensity_type,

            )

            log.info("Fitting and explaining learners...")

            learners = get_learners(

                model_list=self.model_list,

                X_train=X_train,

                Y_train=Y_train,

                n_iter=self.n_iter,

                batch_size=self.batch_size,

                batch_norm=False,

                discrete_outcome=binary_outcome

            )

            learner_explanations = get_learner_explanations(learners, 

                                                            X_test, X_train, Y_train, W_train, 

                                                            self.explainer_limit, explainer_list)

            all_important_features = sim.get_all_important_features(with_selective=True)

            pred_features = sim.get_predictive_features()

            prog_features = sim.get_prognostic_features()

            select_features = sim.get_selective_features()

            cate_test = sim.te(X_test)

            for explainer_name in explainer_list:

                for learner_name in learners:

                    attribution_est = np.abs(

                        learner_explanations[learner_name][explainer_name]

                    )

                    acc_scores_all_features = attribution_accuracy(

                        all_important_features, attribution_est

                    )

                    acc_scores_predictive_features = attribution_accuracy(

                        pred_features, attribution_est

                    )

                    acc_scores_prog_features = attribution_accuracy(

                        prog_features, attribution_est

                    )

                    acc_scores_selective_features = attribution_accuracy(

                        select_features, attribution_est

                    )

                    cate_pred = learners[learner_name].predict(X=X_test)

                    pehe_test = compute_pehe(cate_true=cate_test, cate_pred=cate_pred)

                    explainability_data.append(

                        [

                            nonlinearity_scale,

                            learner_name,

                            explainer_name,

                            acc_scores_all_features,

                            acc_scores_predictive_features,

                            acc_scores_prog_features,

                            acc_scores_selective_features,

                            pehe_test,

                            np.mean(cate_test),

                            np.var(cate_test),

                            pehe_test / np.sqrt(np.var(cate_test)),

                        ]

                    )

        metrics_df = pd.DataFrame(

            explainability_data,

            columns=[

                "Nonlinearity Scale",

                "Learner",

                "Explainer",

                "All features ACC",

                "Pred features ACC",

                "Prog features ACC",

                "Select features ACC",

                "PEHE",

                "CATE true mean",

                "CATE true var",

                "Normalized PEHE",

            ],

        )

        results_path = (

            self.save_path

            / f"results/nonlinearity_sensitivity/{self.synthetic_simulator_type}"

        )

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

        if not results_path.exists():

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

        metrics_df.to_csv(

            results_path

            / f"{dataset}_{num_important_features}_binary_{binary_outcome}-seed{self.seed}.csv"

        )

class PropensitySensitivity:

    """

    Sensitivity analysis for confounding.

    """

    def __init__(

        self,

        n_units_hidden: int = 50,

        n_layers: int = 1,

        penalty_orthogonal: float = 0.01,

        batch_size: int = 1024,

        batch_norm: bool = False,

        n_iter: int = 1000,

        seed: int = 42,

        explainer_limit: int = 1000,

        save_path: Path = Path.cwd(),

        num_interactions: int = 1,

        synthetic_simulator_type: str = "linear",

        nonlinearity_scale: float = 0,

        selection_type: str = "random",

        propensity_type: str = "pred",

        propensity_scales: list = [0, 0.5, 1, 2, 5, 10],

        model_list: list = ["TLearner"]

    ) -> None:

        self.n_units_hidden = n_units_hidden

        self.n_layers = n_layers

        self.penalty_orthogonal = penalty_orthogonal

        self.batch_size = batch_size

        self.batch_norm = batch_norm

        self.n_iter = n_iter

        self.seed = seed

        self.explainer_limit = explainer_limit

        self.save_path = save_path

        self.num_interactions = num_interactions

        self.synthetic_simulator_type = synthetic_simulator_type

        self.nonlinearity_scale = nonlinearity_scale

        self.selection_type = selection_type

        self.propensity_type = propensity_type

        self.propensity_scales = propensity_scales

        self.model_list = model_list

    def run(

        self,

        dataset: str = "tcga_10",

        train_ratio: float = 0.8,

        num_important_features: int = 2,

        binary_outcome: bool = False,

        random_feature_selection: bool = True,

        predictive_scale: float = 1,

        nonlinearity_scale: float = 0.5,

        explainer_list: list = [

            "feature_ablation",

            "feature_permutation",

            "integrated_gradients",

            "shapley_value_sampling",

        ],

        debug: bool = False,

        directory_path_: str = None,

    ) -> None:

        log.info(

            f"Using dataset {dataset} with num_important features = {num_important_features} and predictive scale {predictive_scale}."

        )

        X_raw_train, X_raw_test = load(dataset, train_ratio=train_ratio, debug=debug, directory_path_=directory_path_)

        if self.synthetic_simulator_type == "linear":

            sim = SyntheticSimulatorLinear(

                X_raw_train,

                num_important_features=num_important_features,

                random_feature_selection=random_feature_selection,

                seed=self.seed,

            )

        elif self.synthetic_simulator_type == "nonlinear":

            sim = SyntheticSimulatorModulatedNonLinear(

                X_raw_train,

                num_important_features=num_important_features,

                non_linearity_scale=self.nonlinearity_scale,

                seed=self.seed,

                selection_type=self.selection_type,

            )

        else:

            raise Exception("Unknown simulator type.")

        explainability_data = []

        for propensity_scale in self.propensity_scales:

            log.info(f"Now working with propensity_scale = {propensity_scale}...")

            (

                X_train,

                W_train,

                Y_train,

                po0_train,

                po1_train,

                propensity_train,

            ) = sim.simulate_dataset(

                X_raw_train,

                predictive_scale=predictive_scale,

                binary_outcome=binary_outcome,

                treatment_assign=self.propensity_type,

                prop_scale=propensity_scale,

            )

            X_test, W_test, Y_test, po0_test, po1_test, _ = sim.simulate_dataset(

                X_raw_test,

                predictive_scale=predictive_scale,

                binary_outcome=binary_outcome,

                treatment_assign=self.propensity_type,

                prop_scale=propensity_scale,

            )

            log.info("Fitting and explaining learners...")

            learners = get_learners(

                model_list=self.model_list,

                X_train=X_train,

                Y_train=Y_train,

                n_iter=self.n_iter,

                batch_size=self.batch_size,

                batch_norm=self.batch_norm,

                discrete_outcome=binary_outcome

            )

            learner_explanations = get_learner_explanations(learners, 

                                                            X_test, X_train, Y_train, W_train, 

                                                            self.explainer_limit, explainer_list)

            all_important_features = sim.get_all_important_features(with_selective=False)

            pred_features = sim.get_predictive_features()

            prog_features = sim.get_prognostic_features()

            cate_test = sim.te(X_test)

            for explainer_name in explainer_list:

                for learner_name in learners:

                    attribution_est = np.abs(

                        learner_explanations[learner_name][explainer_name]

                    )

                    acc_scores_all_features = attribution_accuracy(

                        all_important_features, attribution_est

                    )

                    acc_scores_predictive_features = attribution_accuracy(

                        pred_features, attribution_est

                    )

                    acc_scores_prog_features = attribution_accuracy(

                        prog_features, attribution_est

                    )

                    cate_pred = learners[learner_name].predict(X=X_test)

                    pehe_test = compute_pehe(cate_true=cate_test, cate_pred=cate_pred)

                    explainability_data.append(

                        [

                            propensity_scale,

                            learner_name,

                            explainer_name,

                            acc_scores_all_features,

                            acc_scores_predictive_features,

                            acc_scores_prog_features,

                            pehe_test,

                            np.mean(cate_test),

                            np.var(cate_test),

                            pehe_test / np.sqrt(np.var(cate_test)),

                        ]

                    )

        metrics_df = pd.DataFrame(

            explainability_data,

            columns=[

                "Propensity Scale",

                "Learner",

                "Explainer",

                "All features ACC",

                "Pred features ACC",

                "Prog features ACC",

                "PEHE",

                "CATE true mean",

                "CATE true var",

                "Normalized PEHE",

            ],

        )

        results_path = (

            self.save_path

            / f"results/propensity_sensitivity/{self.synthetic_simulator_type}"

        )

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

        if not results_path.exists():

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

        metrics_df.to_csv(

            results_path / f"propensity_scale_{dataset}_{num_important_features}_"

            f"proptype_{self.propensity_type}_"

            f"predscl_{predictive_scale}_"

            f"nonlinscl_{nonlinearity_scale}_"

            f"trainratio_{train_ratio}_"

            f"binary_{binary_outcome}-seed{self.seed}.csv"

        )

class CohortSizeSensitivity:

    """

    Sensitivity analysis for varying numbers of samples. This experiment will generate a .csv with the recorded metrics.

    It will also generate a gif, showing the progression on dimensionality-reduced spaces.

    """

    def __init__(

        self,

        n_units_hidden: int = 50,

        n_layers: int = 1,

        penalty_orthogonal: float = 0.01,

        batch_size: int = 1024,

        batch_norm: bool = False,

        n_iter: int = 1000,

        seed: int = 42,

        explainer_limit: int = 1000,

        save_path: Path = Path.cwd(),

        propensity_type: str = "selective",

        cohort_sizes: list = [0.5, 0.7, 1.0],

        nonlinearity_scale: float = 0.5,

        predictive_scale: float = 1,

        synthetic_simulator_type: str = "random",

        selection_type: str = "random",

        model_list: list = ["TLearner"],

        num_cube_samples: int = 1000,

        dim_reduction_method: str = "umap",

        dim_reduction_on_important_features: bool = True,

        visualize_progression: bool = True,

    ) -> None:

        self.n_units_hidden = n_units_hidden

        self.n_layers = n_layers

        self.penalty_orthogonal = penalty_orthogonal

        self.batch_size = batch_size

        self.batch_norm = batch_norm

        self.n_iter = n_iter

        self.seed = seed

        self.explainer_limit = explainer_limit

        self.save_path = save_path

        self.propensity_type = propensity_type

        self.cohort_sizes = cohort_sizes

        self.nonlinearity_scale = nonlinearity_scale

        self.predictive_scale = predictive_scale

        self.synthetic_simulator_type = synthetic_simulator_type

        self.selection_type = selection_type

        self.model_list = model_list

        self.num_cube_samples = num_cube_samples

        self.dim_reduction_method = dim_reduction_method

        self.dim_reduction_on_important_features = dim_reduction_on_important_features

        self.visualize_progression = visualize_progression

    def run(

        self,

        dataset: str = "tcga_100",

        num_important_features: int = 15,

        explainer_list: list = [

            "feature_ablation",

            "feature_permutation",

            "integrated_gradients",

            "shapley_value_sampling",

        ],

        train_ratio: float = 0.8,

        binary_outcome: bool = False,

        debug=False,

        directory_path_: str = None,

    ) -> None:

        # Log setting

        log.info(

            f"Using dataset {dataset} with num_important features = {num_important_features}."

        )

        # Load data

        X_raw_train_full, X_raw_test_full = load(dataset, train_ratio=train_ratio, debug=debug, directory_path_=directory_path_)

        explainability_data = []

        # Simulate treatment and outcome for train and test

        sim = SyntheticSimulatorModulatedNonLinear(

                X_raw_train_full,

                num_important_features=num_important_features,

                non_linearity_scale=self.nonlinearity_scale,

                seed=self.seed,

                selection_type=self.selection_type

            )

        (

            X_train_full,

            W_train_full,

            Y_train_full,

            po0_train_full,

            po1_train_full,

            propensity_train_full

        ) = sim.simulate_dataset(

                X_raw_train_full,

                predictive_scale=self.predictive_scale,

                binary_outcome=binary_outcome,

                treatment_assign=self.propensity_type,

            )

        X_test_full, W_test_full, Y_test_full, po0_test_full, po1_test_full, _ = sim.simulate_dataset(

            X_raw_test_full,

            predictive_scale=self.predictive_scale,

            binary_outcome=binary_outcome,

            treatment_assign=self.propensity_type,

        )

        # Retrieve important features

        all_important_features = sim.get_all_important_features(with_selective=True)

        pred_features = sim.get_predictive_features()

        prog_features = sim.get_prognostic_features()

        select_features = sim.get_selective_features()

        # Code for sampling from hypercube -> does not work well because data lives in very small part of that hypercube

        # # Sample from a hypercube grid for the X_raw_train_full dataset - making a complete hypercube grid would be too large

        # # Sample for important features only as we know these are the ones that matter and will make sampling from a hypercube more meaningful

        # X_raw_train_full_important = X_raw_train_full[:, all_important_features]

        # min_vals = X_raw_train_full_important.min(axis=0)

        # max_vals = X_raw_train_full_important.max(axis=0)

        # # Add some relative padding to the min and max values

        # min_vals = min_vals - 0.1 * np.abs(min_vals)

        # max_vals = max_vals + 0.1 * np.abs(max_vals)

        # # Sample points

        # grid_samples = np.zeros((self.num_cube_samples, X_raw_train_full.shape[1]))

        # grid_samples[:, all_important_features] = np.random.uniform(min_vals, max_vals, (self.num_cube_samples, X_raw_train_full_important.shape[1]))

        if self.visualize_progression:

            # Use full training set with added noisy samples as focused samples of space

            std_dev = np.std(X_raw_train_full, axis=0)

            grid_samples = X_raw_train_full

            grid_samples = np.vstack([grid_samples,

                                        X_raw_train_full + std_dev*np.random.normal(0, 1, X_raw_train_full.shape),

                                        X_raw_train_full + 0.1*std_dev*np.random.normal(0, 1, X_raw_train_full.shape),

                                        X_raw_train_full + 0.1*std_dev*np.random.normal(0, 1, X_raw_train_full.shape),

                                        # X_raw_train_full + 0.3*std_dev*np.random.normal(0, 1, X_raw_train_full.shape),

                                        X_raw_train_full + 0.1*std_dev*np.random.normal(0, 1, X_raw_train_full.shape),])

            # Reduce samples to two dimensions for plotting using umap

            if self.dim_reduction_method == "umap":

                reducer = umap.UMAP(min_dist=1, n_neighbors=30, spread=1)

                reducer_shap = umap.UMAP(min_dist=3, n_neighbors=40, spread=4)

                reducer_shap_prop = umap.UMAP(min_dist=2, n_neighbors=30, spread=3)

            elif self.dim_reduction_method == "pca":

                reducer = PCA(n_components=2)

                reducer_shap = PCA(n_components=2)

                reducer_shap_prop = PCA(n_components=2)

            elif self.dim_reduction_method == "tsne":

                raise Exception("t-SNE not supported for this analysis. Does not offer .transform() method.")

            else:

                raise Exception("Unknown dimensionality reduction method.")

            # Fit on grid samples and training data

            if self.dim_reduction_on_important_features:

                grid_samples_2d = reducer.fit_transform(grid_samples[:, all_important_features])

                train_samples_2d = reducer.transform(X_raw_train_full[:, all_important_features])

            else:

                grid_samples_2d = reducer.fit_transform(grid_samples)

                train_samples_2d = reducer.transform(X_raw_train_full)

            # Get model learners (here only one) and explanations for grid samples

            learners = get_learners(

                    model_list=self.model_list,

                    X_train=X_train_full,

                    Y_train=Y_train_full,

                    n_iter=self.n_iter,

                    batch_size=self.batch_size,

                    batch_norm=False,

                    discrete_outcome=binary_outcome

                )

            learner_explanations, learners = get_learner_explanations(learners, 

                                                                grid_samples, X_train_full, Y_train_full, W_train_full, 

                                                                grid_samples.shape[0], explainer_list,

                                                                return_learners=True)

            learner_explanations_train = get_learner_explanations(learners,

                                                                X_train_full, X_train_full, Y_train_full, W_train_full,

                                                                X_train_full.shape[0], explainer_list,

                                                                already_trained=True)

            # Get shap for grid samples

            shap_values_grid = learner_explanations[self.model_list[0]][explainer_list[0]]

            shap_values_train = learner_explanations_train[self.model_list[0]][explainer_list[0]]

            # Perform logistic regression for propensity

            est_prop = LogisticRegression().fit(X_train_full, W_train_full)

            explainer_prop = shap.LinearExplainer(est_prop, X_train_full)

            shap_values_prop_grid = explainer_prop.shap_values(grid_samples)

            shap_values_prop_train = explainer_prop.shap_values(X_train_full)

            # Fit umap reducer for shap grid samples 

            if self.dim_reduction_on_important_features:

                shap_values_grid_2d = reducer_shap.fit_transform(shap_values_grid[:, all_important_features])

                shap_values_train_2d = reducer_shap.transform(shap_values_train[:, all_important_features])

                shap_values_prop_grid_2d = reducer_shap_prop.fit_transform(shap_values_prop_grid[:, all_important_features])

                shap_values_prop_train_2d = reducer_shap_prop.transform(shap_values_prop_train[:, all_important_features])

            else:

                shap_values_grid_2d = reducer_shap.fit_transform(shap_values_grid)

                shap_values_train_2d = reducer_shap.transform(shap_values_train)

                shap_values_prop_grid_2d = reducer_shap_prop.fit_transform(shap_values_prop_grid)

                shap_values_prop_train_2d = reducer_shap_prop.transform(shap_values_prop_train)

            # Initialize variable for storing frames

            frames = []  # To store each frame for the GIF

        cohort_size_full = X_train_full.shape[0]

        for cohort_size_perc in self.cohort_sizes:

            cohort_size = int(cohort_size_perc * cohort_size_full)

            # Get a subset of the training data

            X_train = X_train_full[:cohort_size]

            W_train = W_train_full[:cohort_size]

            Y_train = Y_train_full[:cohort_size]

            po0_train = po0_train_full[:cohort_size]

            po1_train = po1_train_full[:cohort_size]

            propensity_train = propensity_train_full[:cohort_size]

            cohort_size_train = X_train.shape[0]

            # Get subsets for test data

            X_test = X_test_full[:cohort_size]

            W_test = W_test_full[:cohort_size]

            Y_test = Y_test_full[:cohort_size]

            po0_test = po0_test_full[:cohort_size]

            po1_test = po1_test_full[:cohort_size]

            log.info(f"Now working with a cohort size of {cohort_size}/{X_train_full.shape[0]}...")

            log.info("Fitting and explaining learners...")

            learners = get_learners(

                model_list=self.model_list,

                X_train=X_train,

                Y_train=Y_train,

                n_iter=self.n_iter,

                batch_size=self.batch_size,

                batch_norm=False,

                discrete_outcome=binary_outcome

            )

            # Get learners and explanations for training data

            learner_explanations, learners = get_learner_explanations(learners, 

                                                            X_train, X_train, Y_train, W_train, 

                                                            X_train.shape[0], explainer_list,

                                                            return_learners=True)

            if self.visualize_progression:

                # Make logistic regression for propensity

                est_prop = LogisticRegression().fit(X_train, W_train)

                explainer_prop = shap.LinearExplainer(est_prop, X_train)

                # Set up plot

                fig, axs = plt.subplots(2, 4, figsize=(15, 10))

                eff_grid = sim.te(grid_samples)

                prop_grid = sim.prop(grid_samples)

                # Get cate estimator

                est_eff = learners[self.model_list[0]]

                cate_pred_train = est_eff.predict(X=X_train)

                # Get predictions for the first model

                p_eff_grid = est_eff.predict(X=grid_samples)

                p_prop_grid = est_prop.predict_proba(grid_samples)[:, 1]

                outcomes = [p_prop_grid, prop_grid, p_eff_grid, eff_grid]

                titles = ['prop(x)', 'prop_true(x)', 'cate(x)', 'cate_true(x)']

                # # Get X_Train in 2d

                # shap_values_train = learner_explanations[self.model_list[0]][explainer_list[0]]

                # shap_values_prop_train = explainer_prop.shap_values(X_train)

                # if self.dim_reduction_on_important_features:

                #     X_train_2d = reducer.transform(X_train[:, all_important_features])

                # else:

                #     X_train_2d = reducer.transform(X_train)

                # # Get shap values of X_Train in 2d

                # if self.dim_reduction_on_important_features:

                #     shap_values_train_2d = reducer_shap.transform(shap_values_train[:, all_important_features])

                #     shap_values_prop_train_2d = reducer_shap_prop.transform(shap_values_prop_train[:, all_important_features])

                # else:

                #     shap_values_train_2d = reducer_shap.transform(shap_values_train)

                #     shap_values_prop_train_2d = reducer_shap_prop.transform(shap_values_prop_train)

                # If there are any non-finite elements in the outcomes, remove them from the grid_samples and the other outcomes and raise a warning

                for i, outcome in enumerate(outcomes):

                    if type(outcome) == torch.Tensor:

                        outcome = outcome.cpu().detach().numpy()

                    outcome = np.array(outcome)

                    if not np.all(np.isfinite(outcome)):

                        print("-----------")

                        print(i, np.sum(~np.isfinite(outcome)))

                        log.warning(f'Found non-finite elements in outcomes. Removing {np.sum(~np.isfinite(outcome))} elements.')

                        mask = np.isfinite(outcome)

                        grid_samples_2d = grid_samples_2d[mask]

                        for j in range(len(outcomes)):

                            outcomes[j] = outcomes[j][mask]

                        break

                for j, outcome in enumerate(outcomes):

                    if type(outcome) == torch.Tensor:

                        outcome = outcome.cpu().detach().numpy()

                    # Plot settings

                    cmap = "viridis"