import pickle

import pandas as pd

import numpy as np

from catenets.datasets import load as catenets_load

from src.iterpretability.datasets.news.process_news import process_news

from src.iterpretability.datasets.tcga.process_tcga import process_tcga

from src.utils import gen_unique_ids

import os

def remove_unbalanced(df, threshold=0.8):

    ub_col = []

    for col in df.select_dtypes(exclude='object').columns:

        if df[col].nunique() < 6:

            distribution = df[col].value_counts(normalize=True)

            if distribution.max() > threshold:

                ub_col.append(col)

    df = df.drop(columns=ub_col)

    return df

def load(dataset_name: str, train_ratio: float = 1.0, 

         directory_path_: str = None, 

         repo_path: str = None,

         debug=False,

         sim_type=None,

         n_samples=None):

    """

    Load the dataset

    """

    feature_names = None

    if "tcga" in dataset_name:

        try:

            tcga_dataset = pickle.load(

                open(

                    repo_path + "/data/tcga/" + str(dataset_name) + ".p",

                    "rb",

                )

            )

        except:

            process_tcga(

                max_num_genes=100, file_location=repo_path + "/data/tcga/"

            )

            tcga_dataset = pickle.load(

                open(

                    repo_path + "/data/tcga/" + str(dataset_name) + ".p",

                    "rb",

                )

            )

        # get gene names

        # ids = tcga_dataset["ids"]

        X_raw = tcga_dataset["rnaseq"][:5000]

    elif "news" in dataset_name:

        try:

            news_dataset = pickle.load(

                open(

                    "src/iterpretability/datasets/news/" + str(dataset_name) + ".p",

                    "rb",

                )

            )

        except:

            process_news(

                max_num_features=100, file_location="src/iterpretability/datasets/news/"

            )

            news_dataset = pickle.load(

                open(

                    "src/iterpretability/datasets/news/" + str(dataset_name) + ".p",

                    "rb",

                )

            )

        X_raw = news_dataset

    elif "twins" in dataset_name:

        # Total features  = 39

        X_raw, _, _, _, _, _ = catenets_load(dataset_name, train_ratio=1.0)

        # Remove columns which almost only contain one value from X_raw

        # Assuming X_raw is your numpy array

        # threshold = 2  # Adjust this value based on your definition of "almost only one value"

        # # Find the columns to keep

        # cols_to_keep = [i for i in range(X_raw.shape[1]) if np.unique(X_raw[:, i]).size > threshold and i != 3]

        # # Create a new array with only the columns to keep

        # X_raw = X_raw[:, cols_to_keep]

    elif "acic" in dataset_name:

        # Total features  = 55

        X_raw, _, _, _, _, _, _, _ = catenets_load("acic2016")

    elif dataset_name.startswith("depmap_drug_screen"):

        data = pd.read_csv(repo_path + "/data/DepMap_24Q2/real/"+dataset_name+".csv", index_col=0)

        outcomes = data[["LFC_az_628", "LFC_imatinib"]].to_numpy()

        outcomes = outcomes.reshape(outcomes.shape[0], outcomes.shape[1], 1)

        data = data.drop(["LFC_az_628", "LFC_imatinib"], axis=1)

        X_raw = data.to_numpy()

    elif dataset_name.startswith("depmap_crispr_screen"):

        data = pd.read_csv(repo_path + "/data/DepMap_24Q2/real/"+dataset_name+".csv", index_col=0)

        outcomes = data[["LFC_BRAF", "LFC_EGFR"]].to_numpy()

        outcomes = outcomes.reshape(outcomes.shape[0], outcomes.shape[1], 1)

        data = data.drop(["LFC_BRAF", "LFC_EGFR"], axis=1)

        X_raw = data.to_numpy()

    elif dataset_name.startswith("ovarian_semi_synthetic"):

        data = pd.read_csv(repo_path + "/data/DepMap_24Q2/real/"+dataset_name+".csv", index_col=0)

        outcomes = data[["pred_a0_y0", "pred_a1_y0"]].to_numpy()

        outcomes = outcomes.reshape(outcomes.shape[0], outcomes.shape[1], 1)

        data = data.drop(["pred_a0_y0", "pred_a1_y0"], axis=1)

        X_raw = data.to_numpy()

    elif dataset_name.startswith("melanoma_semi_synthetic"):

        data = pd.read_csv(repo_path + "/data/DepMap_24Q2/real/"+dataset_name+".csv", index_col=0)

        outcomes = data[["pred_a0_y2 (immuno)", "pred_a1_y2 (immuno)"]].to_numpy()

        outcomes = outcomes.reshape(outcomes.shape[0], outcomes.shape[1], 1)

        outcomes = outcomes.astype(int)

        data = data.drop(["pred_a0_y2 (immuno)", "pred_a1_y2 (immuno)"], axis=1)

        X_raw = data.to_numpy()

    elif dataset_name.startswith("toy_data"):

        if sim_type == "T":

            raise ValueError("Toy example data does not have treatment outcomes")

        data = pd.read_csv(repo_path + "/data/toy/"+dataset_name+".csv", index_col=0)

        X_raw = data.to_numpy()

    elif dataset_name.startswith("cytof"):

        data = pd.read_csv(directory_path_ + dataset_name + ".csv", index_col=0)

        feature_names = data.columns 

        # Also return true outcomes for TSimulation

        if sim_type == "T":

            all_treatment_outcomes_cols = [col for col in data.columns if col.startswith("09")]

            outcomes = data[all_treatment_outcomes_cols[:2]] 

            data = data.drop(all_treatment_outcomes_cols, axis=1)

            # Raise value error if outcomes contains string or nan values

            if outcomes.isnull().values.any():

                raise ValueError("Outcomes contain nan values")

            if outcomes.dtypes.apply(lambda x: x == 'object').any():

                raise ValueError("Outcomes contain string values")

            if type(outcomes) == pd.DataFrame:

                outcomes = outcomes.replace({False:0, True:1})

                outcomes = outcomes.to_numpy()

                outcomes = outcomes.reshape(outcomes.shape[0], outcomes.shape[1], 1)

        X_raw = data

    elif directory_path_ is not None and os.path.isfile(directory_path_ + dataset_name + ".csv"):

        data = pd.read_csv(directory_path_ + dataset_name + ".csv", index_col=0)

        all_treatment_outcomes_cols = [col for col in data.columns if col.startswith("09")][:5]

        # Also return true outcomes for TSimulation

        if sim_type == "T":

            outcomes = data[all_treatment_outcomes_cols]

            # Raise value error if outcomes contains string or nan values

            if outcomes.isnull().values.any():

                raise ValueError("Outcomes contain nan values")

            if outcomes.dtypes.apply(lambda x: x == 'object').any():

                raise ValueError("Outcomes contain string values")

            if type(outcomes) == pd.DataFrame:

                outcomes = outcomes.replace({False:0, True:1})

                outcomes = outcomes.to_numpy()

                outcomes = outcomes.reshape(outcomes.shape[0], outcomes.shape[1], 1)

        # Drop all columns starting with tre or out 

        X_raw = data.loc[:, ~data.columns.str.startswith('tre')]

        X_raw = X_raw.loc[:, ~X_raw.columns.str.startswith('out')]

        # One hot encode all categorical features

        X_raw = pd.get_dummies(X_raw)

        # Make sure there are no boolean variables

        X_raw = X_raw.replace({False:0, True:1})

        # Remove highly unbalanced features

        X_raw = remove_unbalanced(X_raw, 0.8)

        # Normalize all columns of the dataframe

        X_raw = X_raw.apply(lambda x: (x - x.min()) / (x.max() - x.min()))

        # X_raw.index = gen_unique_ids(X_raw.shape[0])

    else:

        raise Exception("Unknown dataset " + str(dataset_name) + "File:", directory_path_ + dataset_name + ".csv")

    debug_size = 200

    if type(X_raw) == pd.DataFrame:

        X_raw = X_raw.to_numpy()

    if train_ratio == 1.0:

        if debug:

            X_raw = X_raw[:int(debug_size*train_ratio),:]

            if sim_type == "T":

                outcomes = outcomes[:int(debug_size*train_ratio),:]

        if sim_type == "T":

            return X_raw, outcomes, feature_names

        else:

            return X_raw, feature_names

    else:

        X_raw_train = X_raw[: int(train_ratio * X_raw.shape[0])]

        X_raw_test = X_raw[int(train_ratio * X_raw.shape[0]) :]

        if debug:

            X_raw_train = X_raw_train[:int(debug_size*train_ratio),:]

            X_raw_test = X_raw_test[:int(debug_size*(1-train_ratio)),:]

            if sim_type == "T":

                outcomes_train = outcomes[:int(debug_size*train_ratio),:]

                outcomes_test = outcomes[int(debug_size*train_ratio):int(debug_size*(1-train_ratio)),:]

        if sim_type == "T":

            return X_raw_train, X_raw_test, outcomes_train, outcomes_test, feature_names

        else:

            return X_raw_train, X_raw_test, feature_names