# stdlib

import random

from typing import Tuple

import src.iterpretability.logger as log

# third party

import numpy as np

import torch

from scipy.special import expit

from scipy.stats import zscore

from omegaconf import DictConfig, OmegaConf

from src.iterpretability.utils import enable_reproducible_results

from abc import ABC, abstractmethod

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# For computing the propensities from scores

from scipy.special import softmax

from scipy.stats import zscore

from sklearn.model_selection import train_test_split

EPS = 0

class SimulatorBase(ABC):

    """

    Base class for simulators.

    """

    @abstractmethod

    def simulate(self, X: np.ndarray, outcomes: np.ndarray = None) -> Tuple:

        raise NotImplementedError

    @abstractmethod

    def get_simulated_data(self, train_ratio: float) -> Tuple:

        raise NotImplementedError

    @property

    @abstractmethod

    def selective_features(self) -> np.ndarray:

        raise NotImplementedError

    @property

    @abstractmethod

    def prognostic_features(self) -> np.ndarray:

        raise NotImplementedError

    @property

    @abstractmethod

    def predictive_features(self) -> np.ndarray:

        raise NotImplementedError

class TYSimulator(SimulatorBase):

    """

    Data generation process class for simulating treatment selection and outcomes (and effects)

    """

    nonlinear_fcts = [

                #lambda x: np.abs(x),

                lambda x: np.exp(-(x**2) / 2),

                #  lambda x: 1 / (1 + x**2),

                # lambda x: np.sqrt(x)*(1+x),

                #lambda x: np.cos(5*x),

                #lambda x: x**2,

                # lambda x: np.arctan(x),

                # lambda x: np.tanh(x),

                # lambda x: np.sin(x),

                # lambda x: np.log(1 + x**2),

                #lambda x: np.sqrt(0.02 + x**2),

                #lambda x: np.cosh(x),

            ]

    def __init__(

        self,

        # Data dimensionality

        dim_X: int,

        # Seed

        seed: int = 42,

        # Simulation type

        simulation_type: str = "ty",

        # Dimensionality of treatments and outcome

        num_binary_outcome: int = 0,

        outcome_unbalancedness_ratio: float = 0,

        standardize_outcome: bool = False,

        num_T: int = 3,

        dim_Y: int = 3,

        # Scale parameters

        predictive_scale: float = 1,

        prognostic_scale: float = 1,

        propensity_scale: float = 1,

        unbalancedness_exp: float = 0,

        nonlinearity_scale: float = 1,

        propensity_type: str = "prog_pred",

        alpha: float = 0.5,

        enforce_balancedness: bool = False,

        # Control

        include_control: bool = False,

        # Important features

        num_pred_features: int = 5,

        num_prog_features: int = 5,

        num_select_features: int = 5,

        feature_type_overlap: str = "sel_none",

        treatment_feature_overlap: bool = False,

        # Feature selection

        random_feature_selection: bool = False,

        nonlinearity_selection_type: bool = True,

        # Noise

        noise: bool = True,

        noise_std: float = 0.1,

    ) -> None:

        # Number of features

        self.dim_X = dim_X

        # Make sure results are reproducible by setting seed for np, torch, random

        self.seed = seed

        enable_reproducible_results(seed=self.seed)

        # Simulation type

        self.simulation_type = simulation_type

        # Store dimensions

        self.num_binary_outcome = num_binary_outcome

        self.outcome_unbalancedness_ratio = outcome_unbalancedness_ratio

        self.standardize_outcome = standardize_outcome

        self.num_T = num_T

        self.dim_Y = dim_Y

        # Scale parameters

        self.predictive_scale = predictive_scale

        self.prognostic_scale = prognostic_scale

        self.propensity_scale = propensity_scale

        self.unbalancedness_exp = unbalancedness_exp

        self.nonlinearity_scale = nonlinearity_scale

        self.propensity_type = propensity_type

        self.alpha = alpha

        self.enforce_balancedness = enforce_balancedness

        # Control

        self.include_control = include_control

        # Important features

        self.num_pred_features = num_pred_features

        self.num_prog_features = num_prog_features

        self.num_select_features = num_select_features

        self.num_important_features = self.num_T*(num_pred_features + num_select_features) + num_prog_features

        self.feature_type_overlap = feature_type_overlap

        self.treatment_feature_overlap = treatment_feature_overlap

        # Feature selection

        self.random_feature_selection = random_feature_selection

        self.nonlinearity_selection_type = nonlinearity_selection_type

        # Noise

        self.noise = noise

        self.noise_std = noise_std

        # Setup variables

        self.nonlinearities = None

        self.prog_mask, self.pred_masks, self.select_masks = None, None, None

        self.prog_weights, self.pred_weights, self.select_weights = None, None, None

        # Setup

        self.setup()

        # Simulation variables

        self.X = None

        self.prog_scores, self.pred_scores, self.select_scores = None, None, None

        self.select_scores_pred_overlap = None

        self.select_scores_prog_overlap = None

        self.propensities, self.outcomes, self.T, self.Y = None, None, None, None

    def get_simulated_data(self):

        """

        Extract results and split into training and test set. Include counterfactual outcomes and propensities.

        """

        return self.X, self.T, self.Y, self.outcomes, self.propensities

        ## OLD CODE

        # Split data

        # train_size = int(train_ratio * self.X.shape[0])

        # if self.num_binary_outcome > 0:

        #     (

        #         X_train, X_test, 

        #         Y_train, Y_test, 

        #         T_train, T_test,

        #         outcomes_train, outcomes_test,

        #         propensities_train, propensities_test,

        #     ) = train_test_split(self.X, self.Y, self.T, self.outcomes, self.propensities, train_size=train_size, stratify=self.Y)

        # else:

        #     X_train, X_test = self.X[:train_size], self.X[train_size:]

        #     T_train, T_test = self.T[:train_size], self.T[train_size:]

        #     Y_train, Y_test = self.Y[:train_size], self.Y[train_size:]

        #     outcomes_train, outcomes_test = self.outcomes[:train_size,:,:], self.outcomes[train_size:,:,:]

        #     propensities_train, propensities_test = self.propensities[:train_size], self.propensities[train_size:]

        # if train_ratio == 1:

        #     return self.X, self.T, self.Y, self.outcomes, self.propensities

        # return X_train, X_test, T_train, T_test, Y_train, Y_test, outcomes_train, outcomes_test, propensities_train, propensities_test

    def simulate(self, X, outcomes=None) -> Tuple:

        """

        Simulate treatment and outcome for a dataset based on the configuration.

        """

        log.debug(

            f'Simulating treatment and outcome for a dataset with:'

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

            f'\nDim X: {self.dim_X}'

            f'\nDim T: {self.num_T}'

            f'\nDim Y: {self.dim_Y}'

            f'\nPredictive Scale: {self.predictive_scale}'

            f'\nPrognostic Scale: {self.prognostic_scale}'

            f'\nPropensity Scale: {self.propensity_scale}'

            f'\nUnbalancedness Exponent: {self.unbalancedness_exp}'

            f'\nNonlinearity Scale: {self.nonlinearity_scale}'

            f'\nNum Pred Features: {self.num_pred_features}'

            f'\nNum Prog Features: {self.num_prog_features}'

            f'\nNum Select Features: {self.num_select_features}'

            f'\nFeature Overlap: {self.treatment_feature_overlap}'

            f'\nRandom Feature Selection: {self.random_feature_selection}'

            f'\nNonlinearity Selection Type: {self.nonlinearity_selection_type}'

            f'\nNoise: {self.noise}'

            f'\nNoise Std: {self.noise_std}'

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

        )

        # 1. Store data with min max scaling to range [0, 1]

        self.X = X

        # self.X = (X - X.min(axis=0)) / (X.max(axis=0) - X.min(axis=0) + EPS)

        # 2. Compute scores for prognostic, predictive, and selective features

        self.compute_scores()

        # 3. Compute factual and counterfactual outcomes based on the data and the predictive and prognostic scores

        self.compute_all_outcomes()

        # 4. Compute propensities based on the data and the selective scores

        self.compute_propensities()

        # 5. Sample treatment assignment based on the propensities

        self.sample_T()

        # 6. Extract the outcome based on the treatment assignment

        self.extract_Y()

        return None

    def setup(self) -> None:

        """

        Setup the simulator by defining variables which remain the same across simulations with different samples but the same configuration.

        """

        # 1. Sample nonlinearities used 

        num_nonlinearities = 2 + self.dim_Y # Different non-linearities for each outcome (predictive), same for all treatments

        self.nonlinearities = self.sample_nonlinearities(num_nonlinearities)

        # 2. Set important feature masks - determine which features should be used for treatment selection, outcome prediction

        self.sample_important_feature_masks()

        # 3. Sample weights for features

        self.sample_uniform_weights()

    def get_true_cates(self, 

                       X: np.ndarray, 

                       T: np.ndarray, 

                       outcomes: np.ndarray) -> np.ndarray:

        """

        Compute true CATEs for each treatment based on the data and the outcomes.

        Always use the selected treatment as the base treatment.

        """

        # Compute CATEs for each treatment

        cates = np.zeros((X.shape[0], self.num_T, self.dim_Y))

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

            for j in range(self.num_T):

                cates[i,j,:] = outcomes[i,j,:] - outcomes[i,int(T[i]),:]

        log.debug(

            f'\nCheck if true CATEs are computed correctly:'

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

            f'\nOutcomes: {outcomes.shape}'

            f'\n{outcomes}'

            f'\n\nTreatment Assignment: {T.shape}'

            f'\n{T}'

            f'\n\nTrue CATEs: {cates.shape}'

            f'\n{cates}'

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

        )

        return cates

    def extract_Y(self) -> None:

        """

        Extract the outcome based on the treatment assignment.

        """

        self.Y = self.outcomes[np.arange(self.X.shape[0]), self.T] 

        log.debug(

            f'\nCheck if outcomes are extracted correctly:'

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

            f'\nOutcomes'

            f'\n{self.outcomes}'

            f'\n{self.outcomes.shape}'

            f'\n\nTreatment Assignment'

            f'\n{self.T}'

            f'\n{self.T.shape}'

            f'\n\nExtracted Outcomes'

            f'\n{self.Y}'

            f'\n{self.Y.shape}'

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

        )

        return None

    def compute_all_outcomes_toy(self) -> None:

        # Compute outcomes for each treatment and outcome

        outcomes = np.zeros((self.X.shape[0], self.num_T, self.dim_Y))

        X0 = self.X[:,0]

        X1 = self.X[:,1]

        k=20

        nonlinearity = lambda x: 1 / (1 + np.exp(-k * (x - 0.5))) #logistic

        if self.propensity_type.startswith("toy1") or self.propensity_type.startswith("toy3") or self.propensity_type.startswith("toy4"):

            fun_y0 = lambda X0, X1: X0

            fun_y1 = lambda X0, X1: 1-X0

        elif self.propensity_type.startswith("toy2"):

            fun_y0 = lambda X0, X1: X0

            fun_y1 = lambda X0, X1: 1-X1

        elif self.propensity_type.startswith("toy6"):

            fun_y0 = lambda X0, X1: X0

            fun_y1 = lambda X0, X1: X1

        elif self.propensity_type.startswith("toy5"):

            fun_y0 = lambda X0, X1: np.sin(X0*10*np.pi)

            fun_y1 = lambda X0, X1: np.sin((1-X0)*10*np.pi)

        elif self.propensity_type.startswith("toy7"):

            fun_y0 = lambda X0, X1: nonlinearity(X0)-nonlinearity(X1)

            fun_y1 = lambda X0, X1: nonlinearity(X0)+nonlinearity(X1)

        elif self.propensity_type.startswith("toy8"):

            fun_y0 = lambda X0, X1: X0

            fun_y1 = lambda X0, X1: 1-X0

        Y = np.array([fun_y0(X0, X1),fun_y1(X0, X1)]).T

        if self.propensity_type.endswith("nonlinear"):

            Y = nonlinearity(Y)

        Y = zscore(Y, axis=None)

        outcomes[:,:,0] = Y

        return outcomes

    def compute_all_outcomes(self) -> None:

        """

        Compute factual and counterfactual outcomes based on the data and the predictive and prognostic scores.

        """

        if self.propensity_type.startswith("toy"):

            outcomes = self.compute_all_outcomes_toy()

        else:

            # Compute outcomes for each treatment and outcome

            outcomes = np.zeros((self.X.shape[0], self.num_T, self.dim_Y))

            for i in range(self.num_T):

                for j in range(self.dim_Y):

                    if self.include_control and i == 0:

                        outcomes[:,i,j] = self.prognostic_scale*self.prog_scores[:,j]

                    else:

                        outcomes[:,i,j] = self.prognostic_scale*self.prog_scores[:,j] + self.predictive_scale*self.pred_scores[:,i,j]

        # Add gaussian noise to outcomes

        if self.noise:

            outcomes = outcomes + np.random.normal(0, self.noise_std, size=outcomes.shape)

        # Create binary outcomes and introduce unbalancedness

        if int(self.num_binary_outcome) > 0:

            for j in range(self.num_binary_outcome):

                scores = zscore(outcomes[:,:,j], axis=0)

                prob = expit(scores)

                outcomes[:,:,j] = prob > self.outcome_unbalancedness_ratio

        self.outcomes = outcomes

        # Standardize outcomes

        if self.standardize_outcome:

            # normalize outcomes per outcome

            self.outcomes = zscore(self.outcomes, axis=0)

        log.debug(

            f'\nCheck if outcomes are computed correctly:'

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

            f'\nProg Scores'

            f'\n{self.prog_scores}'

            f'\n{self.prog_scores.shape}'

            f'\n\nPred Scores'

            f'\n{self.pred_scores}'

            f'\n{self.pred_scores.shape}'

            f'\n\nOutcomes'

            f'\n{self.outcomes}'

            f'\n{self.outcomes.shape}'

            f'\n\nMean Outcomes'

            f'\n{self.outcomes.mean(axis=0)}'

            f'\n\nVariance Outcomes'

            f'\n{self.outcomes.var(axis=0)}'

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

        )

        return None

    def sample_T(self) -> None:

        """

        Sample treatment assignment based on the propensities.

        """

        # Sample from the resulting categorical distribution per row

        self.T = np.array([np.random.choice([tre for tre in range(self.propensities.shape[1])], p=row) for row in self.propensities])

        log.debug(

            f'\nCheck if treatment assignment is sampled correctly:'

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

            f'\nPropensities'

            f'\n{self.propensities}'

            f'\n{self.propensities.shape}'

            f'\n\nTreatment Assignment'

            f'\n{self.T}'

            f'\n{self.T.shape}'

            f'\n\nUnique Treatment Counts'

            f'\n{np.unique(self.T, return_counts=True)}'

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

        )

        return None

    def get_unbalancedness_weights(self, size: int) -> np.ndarray:

        """

        Create weights for introducing unbalancedness for class probabilities.

        """

        # Sample initial distribution of treatment assignment 

        unb_weights = np.random.uniform(0, 1, size=size) 

        unb_weights = unb_weights / unb_weights.sum()

        # Standardize the weights and make sure that a treatment doesn't completely disappear for small unbalancedness exponents

        min_val = unb_weights.min()

        range_val = unb_weights.max() - min_val

        unb_weights = (unb_weights - min_val) / range_val

        unb_weights = 0.01 + unb_weights * 0.98

        return unb_weights

    def compute_propensity_scores_toy(self) -> np.ndarray:

        X0 = self.X[:,0]

        X1 = self.X[:,1]

        if self.propensity_type.startswith("toy1"):

            fun_t0 = lambda X0, X1: X0

            fun_t1 = lambda X0, X1: 1-X0

        elif self.propensity_type.startswith("toy2"):

            fun_t0 = lambda X0, X1: X0

            fun_t1 = lambda X0, X1: 1-X1

        elif self.propensity_type.startswith("toy3"):

            fun_t0 = lambda X0, X1: X1

            fun_t1 = lambda X0, X1: 1-X1

        elif self.propensity_type.startswith("toy4"):

            fun_t0 = lambda X0, X1: np.sin(X0*10*np.pi)

            fun_t1 = lambda X0, X1: np.sin((1-X0)*10*np.pi)

        elif self.propensity_type.startswith("toy5"):

            fun_t0 = lambda X0, X1: 1-X0

            fun_t1 = lambda X0, X1: X0

        elif self.propensity_type.startswith("toy6"):

            fun_t0 = lambda X0, X1: 1-X0

            fun_t1 = lambda X0, X1: X0

        elif self.propensity_type.startswith("toy7"):

            fun_t0 = lambda X0, X1: 1-X0

            fun_t1 = lambda X0, X1: X0

        elif self.propensity_type.startswith("toy8"):

            fun_t0 = lambda X0, X1: 1-X0

            fun_t1 = lambda X0, X1: X0

        scores = np.array([fun_t0(X0, X1),fun_t1(X0, X1)]).T

        return scores

    def compute_propensities(self) -> None:

        """

        Compute propensities based on the data and the selective scores.

        """

        select_scores_pred_overlap = zscore(self.select_scores_pred_overlap, axis=0) # Comment for Predictive Epertise

        select_scores_prog_overlap = zscore(self.select_scores_prog_overlap, axis=0) # Comment for Predictive Epertise

        select_scores_none = zscore(self.select_scores, axis=0) # Comment for Predictive Epertise

        select_scores_pred = np.zeros((self.X.shape[0], self.num_T))

        select_scores_pred_flipped = np.zeros((self.X.shape[0], self.num_T))

        select_scores_prog = np.zeros((self.X.shape[0], self.num_T))

        select_scores_tre = np.zeros((self.X.shape[0], self.num_T))

        select_scores_pred[:,0] = self.outcomes[:,0,0] - self.outcomes[:,1,0]

        select_scores_pred[:,1] = self.outcomes[:,1,0] - self.outcomes[:,0,0]

        select_scores_pred_flipped[:,0] = self.outcomes[:,1,0] - self.outcomes[:,0,0]

        select_scores_pred_flipped[:,1] = self.outcomes[:,0,0] - self.outcomes[:,1,0]

        select_scores_prog[:,0] = self.outcomes[:,0,0]

        select_scores_prog[:,1] = -self.outcomes[:,0,0]

        select_scores_tre[:,0] = -self.outcomes[:,1,0]

        select_scores_tre[:,1] = self.outcomes[:,1,0]

        if self.propensity_type == "prog_tre":

            scores = self.alpha * select_scores_tre + (1 - self.alpha) * select_scores_prog

        # Standardize all scores

        select_scores_pred = zscore(select_scores_pred, axis=0)

        select_scores_pred_flipped = zscore(select_scores_pred_flipped, axis=0)

        select_scores_prog = zscore(select_scores_prog, axis=0)

        select_scores_tre = zscore(select_scores_tre, axis=0)

        if self.propensity_type == "prog_pred":

            scores = self.alpha * select_scores_pred + (1 - self.alpha) * select_scores_prog

        elif self.propensity_type == "prog_tre":

            pass

        elif self.propensity_type == "none_prog":

            scores = self.alpha * select_scores_prog + (1 - self.alpha) * select_scores_none

        elif self.propensity_type == "none_pred":

            scores = self.alpha * select_scores_pred + (1 - self.alpha) * select_scores_none

        elif self.propensity_type == "none_tre":

            scores = self.alpha * select_scores_tre + (1 - self.alpha) * select_scores_none

        elif self.propensity_type == "none_pred_flipped":

            scores = self.alpha * select_scores_pred_flipped + (1 - self.alpha) * select_scores_none

        elif self.propensity_type == "pred_pred_flipped":

            scores = self.alpha * select_scores_pred_flipped + (1 - self.alpha) * select_scores_pred

        elif self.propensity_type == "none_pred_overlap":

            scores = self.alpha * select_scores_pred_overlap + (1 - self.alpha) * select_scores_none

        elif self.propensity_type == "none_prog_overlap":

            scores = self.alpha * select_scores_prog_overlap + (1 - self.alpha) * select_scores_none

        elif self.propensity_type == "pred_overalp_prog_overlap":

            scores = self.alpha * select_scores_prog_overlap + (1 - self.alpha) * select_scores_pred_overlap

        elif self.propensity_type == "rct_none":

            scores = select_scores_none

        elif self.propensity_type.startswith("toy"):

            scores = self.compute_propensity_scores_toy()

        else:

            raise ValueError(f"Unknown propensity type {self.propensity_type}.")

        if self.enforce_balancedness:

            scores = zscore(scores, axis=0)

        if self.propensity_type == "rct_none":

            scores = self.alpha * select_scores_none

        # Introduce unbalancedness and manipulate unbalancedness weights for comparable experiments with different seeds

        unb_weights = self.get_unbalancedness_weights(size=scores.shape[1])

        # Apply the softmax function to each row to get probabilities

        p = softmax(self.propensity_scale*scores, axis=1)

        # Scale probabilities to introduce unbalancedness

        p = p * (1 - unb_weights) ** self.unbalancedness_exp

        # Make sure rows add up to one again

        row_sums = p.sum(axis=1, keepdims=True)

        p = p / row_sums

        self.propensities = p

        log.debug(

            f'\nCheck if propensities are computed correctly:'

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

            f'\nSelect Scores'

            f'\n{self.select_scores}'

            f'\n{self.select_scores.shape}'

            f'\n\nPropensities'

            f'\n{self.propensities}'

            f'\n{self.propensities.shape}'

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

        )

        return None

    def compute_scores(self) -> None:

        """

        Compute scores for prognostic, predictive, and selective features based on the data and the feature weights.

        """

        # Each column of the score matrix corresponds to the score for a specific outcome. Rows correspond to samples.

        prog_lin = self.X @ self.prog_weights.T

        select_lin = self.X @ self.select_weights.T

        select_lin_pred = self.X @ self.select_weights_pred.T

        select_lin_prog = self.X @ self.select_weights_prog.T

        log.debug(

            f'\nCheck if linear scores are computed correctly for selective features:'

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

            f'\nself.X'

            f'\n{self.X}'

            f'\n{self.X.shape}'

            f'\n\nSelect Weights'

            f'\n{self.select_weights}'

            f'\n{self.select_weights.shape}'

            f'\n\nSelect Lin'

            f'\n{select_lin}'

            f'\n{select_lin.shape}'

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

        )

        # Compute scores for predictive and selective features for each treatment and outcome

        pred_lin = np.zeros((self.X.shape[0], self.num_T, self.dim_Y))

        # This creates a score for each treatment and outcome for each sample

        for i in range(self.num_T):

            pred_lin[:,i,:] = self.X @ self.pred_weights[i].T

        # Introduce non-linearity and get final scores

        prog_scores = (1 - self.nonlinearity_scale) * prog_lin + self.nonlinearity_scale * self.nonlinearities[0](prog_lin)

        select_scores = (1 - self.nonlinearity_scale) * select_lin + self.nonlinearity_scale * self.nonlinearities[1](select_lin)

        select_scores_pred_overlap = (1 - self.nonlinearity_scale) * select_lin_pred + self.nonlinearity_scale * self.nonlinearities[1](select_lin_pred)

        select_scores_prog_overlap = (1 - self.nonlinearity_scale) * select_lin_prog + self.nonlinearity_scale * self.nonlinearities[1](select_lin_prog)

        pred_scores = np.zeros((self.X.shape[0], self.num_T, self.dim_Y))

        for i in range(self.dim_Y):

            pred_scores[:,:,i] = (1 - self.nonlinearity_scale) * pred_lin[:,:,i] + self.nonlinearity_scale * self.nonlinearities[i+2](pred_lin[:,:,i])

        log.debug(

            f'\nCheck if all scores are computed correctly for predictive features:'

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

            f'\nself.X'

            f'\n{self.X}'

            f'\n{self.X.shape}'

            f'\n\nPred Weights'

            f'\n{self.pred_weights}'

            f'\n{self.pred_weights.shape}'

            f'\n\nPred Lin'

            f'\n{pred_lin}'

            f'\n{pred_lin.shape}'

            f'\n\nPred Scores'

            f'\n{pred_scores}'

            f'\n{pred_scores.shape}'

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

        )

        self.prog_scores = prog_scores

        self.select_scores = select_scores

        self.select_scores_pred_overlap = select_scores_pred_overlap

        self.select_scores_prog_overlap = select_scores_prog_overlap

        self.pred_scores = pred_scores

        return None

    @property

    def weights(self) -> Tuple:

        """

        Return weights for prognostic, predictive, and selective features.

        """

        return self.prog_weights, self.pred_weights, self.select_weights

    def sample_uniform_weights(self) -> None:

        """

        sample uniform weights for the features.

        """

        if self.propensity_type.startswith("toy"):

            self.prog_weights = np.zeros((self.dim_Y, self.dim_X))

            self.pred_weights = np.zeros((self.num_T, self.dim_Y, self.dim_X))

            self.select_weights = np.zeros((self.num_T, self.dim_X))

            self.select_weights_pred = np.zeros((self.num_T, self.dim_X))

            self.select_weights_prog = np.zeros((self.num_T, self.dim_X))

            return None

        # Sample weights for prognostic features, a weight for every outcome

        prog_weights = np.random.uniform(-1, 1, size=(self.dim_Y, self.dim_X)) * self.prog_mask

        # Sample weights for predictive and selective features, a weight for every dimension for every treatment and outcome

        pred_weights = np.random.uniform(-1, 1, size=(self.num_T, self.dim_Y, self.dim_X))

        select_weights = np.random.uniform(-1, 1, size=(self.num_T, self.dim_X))

        select_weights_pred = select_weights.copy()

        select_weights_prog = select_weights.copy()

        # # Sample weights for prognostic features, a weight for every outcome

        # prog_weights = np.random.uniform(0, 1, size=(self.dim_Y, self.dim_X)) * self.prog_mask

        # # Sample weights for predictive and selective features, a weight for every dimension for every treatment and outcome

        # pred_weights = np.random.uniform(0, 1, size=(self.num_T, self.dim_Y, self.dim_X))

        # select_weights = np.random.uniform(0, 1, size=(self.num_T, self.dim_X))

        # # Make sure treatments are different

        # pred_weights[0] = -pred_weights[0]

        # select_weights[0] = -select_weights[0]

        # # Ones as weights

        # prog_weights = np.ones((self.dim_Y, self.dim_X)) * self.prog_mask#/ self.prog_mask.sum()

        # pred_weights = np.ones((self.num_T, self.dim_Y, self.dim_X))  #/ self.pred_masks.sum(axis=1, keepdims=True)

        # select_weights = np.ones((self.num_T, self.dim_X))  #/ self.select_masks.sum(axis=1, keepdims=True)

        # Mask weights for features that are not important

        for i in range(self.num_T):

            pred_weights[i] = pred_weights[i] * self.pred_masks[:,i]

            select_weights[i] = select_weights[i] * self.select_masks[:,i]

            select_weights_pred[i] = select_weights_pred[i] * self.select_masks_pred[:,i]

            select_weights_prog[i] = select_weights_prog[i] * self.select_masks_prog[:,i]

        # for i in range(self.num_T):

        #     row_sums = pred_weights[i].sum(axis=1, keepdims=True)

        #     pred_weights[i] = pred_weights[i] / row_sums

        #     row_sums = select_weights[i].sum()

        #     select_weights[i] = select_weights[i] / row_sums

        # # Make sure that prog weights sum to one per outcome

        # row_sums = prog_weights.sum(axis=1, keepdims=True)

        # prog_weights = prog_weights / row_sums

        log.debug(

            f'\nCheck if masks are applied correctly:'

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

            f'\nSelect Weights'

            f'\n{select_weights}'

            f'\n{select_weights.shape}'

            f'\n\nSelect Masks'

            f'\n{self.select_masks}'

            f'\n{self.select_masks.shape}'

            f'\n\nPred Weights'

            f'\n{pred_weights}'

            f'\n{pred_weights.shape}'

            f'\n\nPred Masks'

            f'\n{self.pred_masks}'

            f'\n{self.pred_masks.shape}'

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

        )

        self.prog_weights = prog_weights

        self.pred_weights = pred_weights

        self.select_weights = select_weights

        self.select_weights_pred = select_weights_pred

        self.select_weights_prog = select_weights_prog

        return None

    @property

    def all_important_features(self) -> np.ndarray:

        """

        Return all important feature indices.

        """

        all_important_features = np.union1d(self.predictive_features, self.prognostic_features)

        all_important_features = np.union1d(all_important_features, self.selective_features)

        log.debug(

            f'\nCheck if all important features are computed correctly:'

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

            f'\nProg Features'

            f'\n{self.prognostic_features}'

            f'\n\nPred Features'

            f'\n{self.predictive_features}'

            f'\n\nSelect Features'

            f'\n{self.selective_features}'

            f'\n\nAll Important Features'

            f'\n{all_important_features}'

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

        )

        return all_important_features

    @property

    def prognostic_features(self) -> np.ndarray:

        """

        Return prognostic feature indices.

        """

        prog_features = np.where((self.prog_mask).astype(np.int32) != 0)

        return prog_features

    @property

    def predictive_features(self) -> np.ndarray:

        """

        Return predictive feature indices.

        """

        pred_features = np.where((self.pred_masks.sum(axis=1)).astype(np.int32) != 0)

        return pred_features

    @property

    def selective_features(self) -> np.ndarray:

        """

        Return selective feature indices.

        """

        select_features = np.where((self.select_masks.sum(axis=1)).astype(np.int32) != 0)

        return select_features

    def sample_important_feature_masks(self) -> None:

        """

        Pick features that are important for treatment selection, outcome prediction, and prognostic prediction based on the configuration.

        """

        if self.propensity_type.startswith("toy"):

            self.prog_mask = np.zeros(shape=(self.dim_X))

            self.pred_masks = np.zeros(shape=(self.dim_X, self.num_T))

            self.select_masks = np.zeros(shape=(self.dim_X, self.num_T))

            self.prog_mask[0] = 1

            self.pred_masks[0,0] = 1

            self.pred_masks[1,1] = 1

            self.select_masks[0,0] = 1

            self.select_masks[1,1] = 1

            return None

        # Get indices for features and shuffle if random_feature_selection is True

        all_indices = np.arange(self.dim_X)

        n = self.num_pred_features

        if self.random_feature_selection:

            np.random.shuffle(all_indices)

        # Initialize masks

        prog_mask = np.zeros(shape=(self.dim_X))

        pred_masks = np.zeros(shape=(self.dim_X, self.num_T))

        select_masks = np.zeros(shape=(self.dim_X, self.num_T))

        # Handle case with feature overlap

        if self.feature_type_overlap == "sel_pred":

            prog_indices = all_indices[:n]

            prog_mask[prog_indices] = 1

            if self.treatment_feature_overlap:

                assert 2*n <= int(self.dim_X)

                pred_indices = np.array(self.num_T * [all_indices[n:2*n]])

                select_indices = np.array(self.num_T * [all_indices[n:2*n]])

                prog_mask[prog_indices] = 1

                pred_masks[pred_indices] = 1

                select_masks[select_indices] = 1

            else:

                assert n*(1+self.num_T) <= int(self.dim_X)

                for i in range(self.num_T):

                    pred_indices = all_indices[(i+1)*n: (i+2)*n]

                    select_indices = all_indices[(i+1)*n: (i+2)*n]

                    pred_masks[pred_indices,i] = 1

                    select_masks[select_indices,i] = 1

        elif self.feature_type_overlap == "sel_prog":

            if self.treatment_feature_overlap:

                assert 2*n <= int(self.dim_X)

                prog_indices = all_indices[:n]

                prog_mask[prog_indices] = 1

                pred_indices = np.array(self.num_T * [all_indices[n:2*n]])

                select_indices = np.array(self.num_T * [all_indices[:n]])

                prog_mask[prog_indices] = 1

                pred_masks[pred_indices] = 1

                select_masks[select_indices] = 1

            else:

                assert 2*n*self.num_T <= int(self.dim_X)

                prog_indices = all_indices[:n*self.num_T:self.num_T]

                prog_mask[prog_indices] = 1

                for i in range(self.num_T):

                    select_indices = all_indices[i*n: (i+1)*n]

                    pred_indices = all_indices[(i+self.num_T+1)*n: (i+self.num_T+2)*n]

                    pred_masks[pred_indices,i] = 1

                    select_masks[select_indices,i] = 1

        elif self.feature_type_overlap == "sel_none":

            prog_indices = all_indices[:n]

            prog_mask[prog_indices] = 1

            if self.treatment_feature_overlap:

                assert 3*n <= int(self.dim_X)

                pred_indices = np.array(self.num_T * [all_indices[n:2*n]])

                select_indices = np.array(self.num_T * [all_indices[2*n:3*n]])

                prog_mask[prog_indices] = 1

                pred_masks[pred_indices] = 1

                select_masks[select_indices] = 1

            else:

                #assert n+2*n*self.num_T <= int(self.dim_X)

                for i in range(1,self.num_T+1):

                    select_indices = all_indices[i*n: (i+1)*n]

                    pred_indices = all_indices[(i+self.num_T)*n: (i+self.num_T+1)*n]

                    pred_masks[pred_indices,i-1] = 1

                    select_masks[select_indices,i-1] = 1

        # # Handle case with feature overlap

        # if self.feature_overlap:

        #     assert max(self.num_pred_features, self.num_prog_features, self.num_select_features) <= int(self.dim_X)

        #     prog_indices = all_indices[:self.num_prog_features]

        #     pred_indices = np.array(self.num_T * [all_indices[:self.num_pred_features]])

        #     select_indices = np.array(self.num_T * [all_indices[:self.num_select_features]])

        #     prog_mask[prog_indices] = 1

        #     pred_masks[pred_indices] = 1

        #     select_masks[select_indices] = 1

        # # Handle case without feature overlap

        # else:

        #     assert (self.num_prog_features + self.num_T * (self.num_pred_features + self.num_select_features)) <= int(self.dim_X)

        #     prog_indices = all_indices[:self.num_prog_features]

        #     prog_mask[prog_indices] = 1

        #     pred_indices = all_indices[self.num_prog_features : (self.num_prog_features + self.num_T*self.num_pred_features)]

        #     select_indices = all_indices[(self.num_prog_features + self.num_T*self.num_pred_features):(self.num_prog_features + self.num_T*(self.num_pred_features+self.num_select_features))]

        #     # Mask features for every treatment

        #     for i in range(self.num_T):

        #         pred_masks[pred_indices[i*self.num_pred_features:(i+1)*self.num_pred_features],i] = 1

        #         select_masks[select_indices[i*self.num_select_features:(i+1)*self.num_select_features],i] = 1

        self.prog_mask = prog_mask

        self.pred_masks = pred_masks

        self.select_masks = select_masks

        self.select_masks_pred = pred_masks.copy()

        self.select_masks_prog = select_masks.copy()

        log.debug(

            f'\nCheck if important features are sampled correctly:'

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

            f'\nProg Indices'

            f'\n{prog_indices}'

            f'\n\nPred Indices'

            f'\n{pred_indices}'

            f'\n\nSelect Indices'

            f'\n{select_indices}'

            f'\n\nProg Mask'

            f'\n{prog_mask}'

            f'\n\nPred Masks'

            f'\n{pred_masks}'

            f'\n\nSelect Masks'

            f'\n{select_masks}'

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

        )

        return None

    def sample_nonlinearities(self, num_nonlinearities: int):

        """

        Sample non-linearities for each outcome.

        """

        if self.nonlinearity_selection_type == "random":

            # pick num_nonlinearities 

            return random.choices(population=self.nonlinear_fcts, k=num_nonlinearities)

        else:

            raise ValueError(f"Unknown nonlinearity selection type {self.selection_type}.")

class TSimulator(SimulatorBase):

    """

    Data generation process class for simulating treatment selection only, when counterfactual outcomes are available (as for in-vitro/pharmacoscopy data).

    """

    nonlinear_fcts = [

                lambda x: np.abs(x),

                lambda x: np.exp(-(x**2) / 2),

                lambda x: 1 / (1 + x**2),

                lambda x: np.cos(x),

                lambda x: np.arctan(x),

                lambda x: np.tanh(x),

                lambda x: np.sin(x),

                lambda x: np.log(1 + x**2),

                lambda x: np.sqrt(1 + x**2),

                lambda x: np.cosh(x),

            ]

    def __init__(

        self,

        # Data dimensionality

        dim_X: int,

        # Seed

        seed: int = 42,

        # Simulation type

        simulation_type: str = "T",

        # Dimensionality of treatments and outcome

        num_binary_outcome: int = 0,

        standardize_outcome: bool = False,

        standardize_per_outcome: bool = False,

        num_T: int = 3,

        dim_Y: int = 3,

        # Scale parameters

        propensity_scale: float = 1,

        unbalancedness_exp: float = 0,

        nonlinearity_scale: float = 1,

        propensity_type: str = "prog_pred",

        alpha: float = 0.5,

        enforce_balancedness: bool = False,

        # Important features

        num_select_features: int = 5,

        treatment_feature_overlap: bool = False,

        # Feature selection

        random_feature_selection: bool = True,

        nonlinearity_selection_type: bool = True,

    ) -> None:

        # Number of features

        self.dim_X = dim_X

        # Make sure results are reproducible by setting seed for np, torch, random

        self.seed = seed

        enable_reproducible_results(seed=self.seed)

        # Simulation type

        self.simulation_type = simulation_type

        # Store dimensions

        self.num_binary_outcome = num_binary_outcome

        self.standardize_outcome = standardize_outcome

        self.standardize_per_outcome = standardize_per_outcome

        self.num_T = num_T

        self.dim_Y = dim_Y

        # Scale parameters

        self.propensity_scale = propensity_scale

        self.unbalancedness_exp = unbalancedness_exp

        self.nonlinearity_scale = nonlinearity_scale

        self.propensity_type = propensity_type

        self.alpha = alpha

        self.enforce_balancedness = enforce_balancedness

        # Important features

        self.num_select_features = num_select_features

        self.treatment_feature_overlap = treatment_feature_overlap

        self.num_important_features = num_select_features

        # Feature selection

        self.random_feature_selection = random_feature_selection

        self.nonlinearity_selection_type = nonlinearity_selection_type

        # Setup variables

        self.nonlinearities = None

        self.select_masks = None

        self.select_weights = None

        # Setup

        self.setup()

        # Simulation variables

        self.X = None

        self.select_scores = None

        self.propensities, self.outcomes, self.T, self.Y = None, None, None, None

    def get_simulated_data(self, train_ratio: float = 0.8):

        """

        Extract results and split into training and test set. Include counterfactual outcomes and propensities.

        """

        return self.X, self.T, self.Y, self.outcomes, self.propensities

        # Split data

        # train_size = int(train_ratio * self.X.shape[0])

        # X_train, X_test = self.X[:train_size], self.X[train_size:]

        # T_train, T_test = self.T[:train_size], self.T[train_size:]

        # Y_train, Y_test = self.Y[:train_size], self.Y[train_size:]

        # outcomes_train, outcomes_test = self.outcomes[:train_size,:,:], self.outcomes[train_size:,:,:]

        # propensities_train, propensities_test = self.propensities[:train_size], self.propensities[train_size:]

        # if train_ratio == 1:

        #     return self.X, self.T, self.Y, self.outcomes, self.propensities

        # return X_train, X_test, T_train, T_test, Y_train, Y_test, outcomes_train, outcomes_test, propensities_train, propensities_test

    def simulate(self, X, outcomes=None) -> Tuple:

        """

        Simulate treatment and outcome for a dataset based on the configuration.

        """

        log.debug(

            f'Simulating treatment and outcome for a dataset with:'

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

            f'\nDim X: {self.dim_X}'

            f'\nDim T: {self.num_T}'

            f'\nDim Y: {self.dim_Y}'

            f'\nPropensity Scale: {self.propensity_scale}'

            f'\nUnbalancedness Exponent: {self.unbalancedness_exp}'

            f'\nNonlinearity Scale: {self.nonlinearity_scale}'

            f'\nNum Select Features: {self.num_select_features}'

            f'\nFeature Overlap: {self.treatment_feature_overlap}'

            f'\nRandom Feature Selection: {self.random_feature_selection}'

            f'\nNonlinearity Selection Type: {self.nonlinearity_selection_type}'

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

        )

        # 1. Store data

        self.X = X

        # 2. Compute scores for prognostic, predictive, and selective features

        self.compute_scores()

        # 3. Retrieve factual and counterfactual outcomes based on the data and the predictive and prognostic scores

        self.outcomes = outcomes

        assert self.outcomes.shape == (self.X.shape[0], self.num_T, self.dim_Y)

        if self.standardize_outcome:

            if self.standardize_per_outcome:

                self.outcomes = zscore(self.outcomes, axis=0) #, axis=None) # add axis=None to make problem easier again

            else:

                self.outcomes = zscore(self.outcomes, axis=None) #, axis=None) # add axis=None to make problem easier again

        log.debug(

            f'\nCheck if outcomes are processed correctly:'

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

            f'\n\nOutcomes'

            f'\n{self.outcomes}'

            f'\n{self.outcomes.shape}'

            f'\n\nMean Outcomes'

            f'\n{self.outcomes.mean(axis=0)}'

            f'\n\nVariance Outcomes'