"""

ACIC2016 dataset

"""

# stdlib

import random

from pathlib import Path

from typing import Any, Tuple

import glob

# third party

import numpy as np

import pandas as pd

from sklearn.preprocessing import OneHotEncoder, StandardScaler

from sklearn.model_selection import train_test_split

import catenets.logger as log

from .network import download_if_needed

np.random.seed(0)

random.seed(0)

FILE_ID = "0B7pG5PPgj6A3N09ibmFwNWE1djA"

PREPROCESSED_FILE_ID = "1iOfEAk402o3jYBs2Prfiz6oaailwWcR5"

NUMERIC_COLS = [

    0,

    3,

    4,

    16,

    17,

    18,

    20,

    21,

    22,

    24,

    24,

    25,

    30,

    31,

    32,

    33,

    39,

    40,

    41,

    53,

    54,

]

N_NUM_COLS = len(NUMERIC_COLS)

def get_acic_covariates(

        fn_csv: Path, keep_categorical: bool = False, preprocessed: bool = True

) -> np.ndarray:

    X = pd.read_csv(fn_csv)

    if not keep_categorical:

        X = X.drop(columns=["x_2", "x_21", "x_24"])

    else:

        # encode categorical features

        feature_list = []

        for cols_ in X.columns:

            if type(X.loc[X.index[0], cols_]) not in [np.int64, np.float64]:

                enc = OneHotEncoder(drop="first")

                enc.fit(np.array(X[[cols_]]).reshape((-1, 1)))

                for k in range(len(list(enc.get_feature_names()))):

                    X[cols_ + list(enc.get_feature_names())[k]] = enc.transform(

                        np.array(X[[cols_]]).reshape((-1, 1))

                    ).toarray()[:, k]

                feature_list.append(cols_)

        X.drop(feature_list, axis=1, inplace=True)

    if preprocessed:

        X_t = X.values

    else:

        scaler = StandardScaler()

        X_t = scaler.fit_transform(X)

    return X_t

def preprocess_simu(

        fn_csv: Path,

        n_0: int = 2000,

        n_1: int = 200,

        n_test: int = 500,

        error_sd: float = 1,

        sp_lin: float = 0.6,

        sp_nonlin: float = 0.3,

        prop_gamma: float = 0,

        prop_omega: float = 0,

        ate_goal: float = 0,

        inter: bool = True,

        i_exp: int = 0,

        keep_categorical: bool = False,

        preprocessed: bool = True,

) -> Tuple:

    X = get_acic_covariates(

        fn_csv, keep_categorical=keep_categorical, preprocessed=preprocessed

    )

    np.random.seed(i_exp)

    # shuffle indices

    n_total, n_cov = X.shape

    ind = np.arange(n_total)

    np.random.shuffle(ind)

    ind_test = ind[-n_test:]

    ind_1 = ind[n_0: (n_0 + n_1)]

    # create treatment indicator (treatment assignment does not matter in test set)

    w = np.zeros(n_total).reshape((-1, 1))

    w[ind_1] = 1

    # create dgp

    coeffs_ = [0, 1]

    # sample baseline coefficients

    beta_0 = np.random.choice(coeffs_, size=n_cov, replace=True, p=[1 - sp_lin, sp_lin])

    intercept = np.random.choice([x for x in np.arange(-1, 1.25, 0.25)])

    # sample treatment effect coefficients

    gamma = np.random.choice(

        coeffs_, size=n_cov, replace=True, p=[1 - prop_gamma, prop_gamma]

    )

    omega = np.random.choice(

        [0, 1], replace=True, size=n_cov, p=[prop_omega, 1 - prop_omega]

    )

    # simulate mu_0 and mu_1

    mu_0 = (intercept + np.dot(X, beta_0)).reshape((-1, 1))

    mu_1 = (intercept + np.dot(X, gamma + beta_0 * omega)).reshape((-1, 1))

    if sp_nonlin > 0:

        coefs_sq = [0, 0.1]

        beta_sq = np.random.choice(

            coefs_sq, size=N_NUM_COLS, replace=True, p=[1 - sp_nonlin, sp_nonlin]

        )

        omega = np.random.choice(

            [0, 1], replace=True, size=N_NUM_COLS, p=[prop_omega, 1 - prop_omega]

        )

        X_sq = X[:, NUMERIC_COLS] ** 2

        mu_0 = mu_0 + np.dot(X_sq, beta_sq).reshape((-1, 1))

        mu_1 = mu_1 + np.dot(X_sq, beta_sq * omega).reshape((-1, 1))

        if inter:

            # randomly add some interactions

            ind_c = np.arange(n_cov)

            np.random.shuffle(ind_c)

            inter_list = list()

            for i in range(0, n_cov - 2, 2):

                inter_list.append(X[:, ind_c[i]] * X[:, ind_c[i + 1]])

            X_inter = np.array(inter_list).T

            n_inter = X_inter.shape[1]

            beta_inter = np.random.choice(

                coefs_sq, size=n_inter, replace=True, p=[1 - sp_nonlin, sp_nonlin]

            )

            omega = np.random.choice(

                [0, 1], replace=True, size=n_inter, p=[prop_omega, 1 - prop_omega]

            )

            mu_0 = mu_0 + np.dot(X_inter, beta_inter).reshape((-1, 1))

            mu_1 = mu_1 + np.dot(X_inter, beta_inter * omega).reshape((-1, 1))

    ate = np.mean(mu_1 - mu_0)

    mu_1 = mu_1 - ate + ate_goal

    y = (

            w * mu_1

            + (1 - w) * mu_0

            + np.random.normal(0, error_sd, n_total).reshape((-1, 1))

    )

    X_train, y_train, w_train, mu_0_train, mu_1_train = (

        X[ind[: (n_0 + n_1)], :],

        y[ind[: (n_0 + n_1)]],

        w[ind[: (n_0 + n_1)]],

        mu_0[ind[: (n_0 + n_1)]],

        mu_1[ind[: (n_0 + n_1)]],

    )

    X_test, y_test, w_test, mu_0_t, mu_1_t = (

        X[ind_test, :],

        y[ind_test],

        w[ind_test],

        mu_0[ind_test],

        mu_1[ind_test],

    )

    return (

        X_train,

        w_train,

        y_train,

        np.asarray([mu_0_train, mu_1_train]).squeeze().T,

        X_test,

        w_test,

        y_test,

        np.asarray([mu_0_t, mu_1_t]).squeeze().T,

    )

def get_acic_orig_filenames(data_path: Path, simu_num: int) -> list:

    return sorted(glob.glob((data_path / ("data_cf_all/" + str(simu_num) +

                                          '/zymu_*.csv')).__str__()))

def get_acic_orig_outcomes(data_path: Path,

                           simu_num: int,

                           i_exp: int) -> Tuple:

    file_list = get_acic_orig_filenames(data_path=data_path,

                                        simu_num=simu_num)

    out = pd.read_csv(file_list[i_exp])

    w = out['z']

    y = w * out['y1'] + (1 - w) * out['y0']

    mu_0, mu_1 = out['mu0'], out['mu1']

    return y.values, w.values, mu_0.values, mu_1.values

def preprocess_acic_orig(fn_csv: Path,

                         data_path: Path,

                         preprocessed: bool = False,

                         keep_categorical: bool = True,

                         simu_num: int = 1,

                         i_exp: int = 0,

                         train_size: int = 4000,

                         random_split: bool = False

                         )-> Tuple:

    X = get_acic_covariates(

        fn_csv, keep_categorical=keep_categorical, preprocessed=preprocessed

    )

    y, w, mu_0, mu_1 = get_acic_orig_outcomes(data_path=data_path, simu_num=simu_num, i_exp=i_exp)

    if not random_split:

        X_train, y_train, w_train, mu_0_train, mu_1_train = X[:train_size, :], y[:train_size], \

                                                            w[:train_size], mu_0[:train_size], \

                                                            mu_1[:train_size]

        X_test, y_test, w_test, mu_0_test, mu_1_test = X[train_size:, :], y[train_size:], \

                                                       w[train_size:], mu_0[train_size:], \

                                                       mu_1[train_size:]

    else:

        X_train, X_test, y_train, y_test, w_train, w_test, \

        mu_0_train, mu_0_test, mu_1_train, mu_1_test = train_test_split(X, y, w, mu_0, mu_1,

                                                                        test_size=1 - train_size,

                                                                        random_state=i_exp)

    return (

        X_train,

        w_train,

        y_train,

        np.asarray([mu_0_train, mu_1_train]).squeeze().T,

        X_test,

        w_test,

        y_test,

        np.asarray([mu_0_test, mu_1_test]).squeeze().T,

    )

def preprocess(fn_csv: Path,

               data_path: Path,

               preprocessed: bool = True,

               original_acic_outcomes: bool = False,

               **kwargs: Any,

               ) -> Tuple:

    if not original_acic_outcomes:

        return preprocess_simu(fn_csv=fn_csv, preprocessed=preprocessed, **kwargs)

    else:

        return preprocess_acic_orig(fn_csv=fn_csv, preprocessed=preprocessed,

                                    data_path=data_path, **kwargs)

def load(

        data_path: Path,

        preprocessed: bool = True,

        original_acic_outcomes: bool = False,

        **kwargs: Any,

) -> Tuple:

    """

    ACIC2016 dataset dataloader.

        - Download the dataset if needed.

        - Load the dataset.

        - Preprocess the data.

        - Return train/test split.

    Parameters

    ----------

    data_path: Path

        Path to the CSV. If it is missing, it will be downloaded.

    preprocessed: bool

        Switch between the raw and preprocessed versions of the dataset.

    original_acic_outcomes: bool

        Switch between new simulations (Inductive bias paper) and original acic outcomes

    Returns

    -------

    train_x: array or pd.DataFrame

        Features in training data.

    train_t: array or pd.DataFrame

        Treatments in training data.

    train_y: array or pd.DataFrame

        Observed outcomes in training data.

    train_potential_y: array or pd.DataFrame

        Potential outcomes in training data.

    test_x: array or pd.DataFrame

        Features in testing data.

    test_potential_y: array or pd.DataFrame

        Potential outcomes in testing data.

    """

    if preprocessed:

        csv = data_path / "x_trans.csv"

        download_if_needed(csv, file_id=PREPROCESSED_FILE_ID)

    else:

        arch = data_path / "data_cf_all.tar.gz"

        download_if_needed(

            arch, file_id=FILE_ID, unarchive=True, unarchive_folder=data_path

        )

        csv = data_path / "data_cf_all/x.csv"

    log.debug(f"load dataset {csv}")

    return preprocess(csv, data_path=data_path, preprocessed=preprocessed,

                      original_acic_outcomes=original_acic_outcomes,

                      **kwargs)