# AUTOGENERATED! DO NOT EDIT! File to edit: ../99_misc.ipynb.

# %% auto 0
__all__ = ['get_ihdp_dict']

# %% ../99_misc.ipynb 2
from .base import *
from .models import *
from .estimators import *
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from copy import deepcopy

# %% ../99_misc.ipynb 3
def get_ihdp_dict(i = 1, test_size = 0.9, seperate_pool_test = False, pool_size=0.8, seed = None):
    """
    Method to query the IHDP data from the AMlab github repo

    Arguments
    ---------
    i : int = 1
        Which IHDP dataset to get (between 1 and 10)
    test_size : 0.9
        Fraction of the data used for test set
    seperate_pool_test:
        Whether the pool and the test set are different
    pool_size :
        Fraction of test data when pool vs test split is done
    seed :
        Seed when splitting data
    """
    seed = 1005 if seed is None else seed
    df = pd.read_csv(
    f"https://raw.githubusercontent.com/AMLab-Amsterdam/CEVAE/master/datasets/IHDP/csv/ihdp_npci_{i}.csv",
    names = ["treatment", "y_factual", "y_cfactual", "mu0", "mu1"] + [f'x{x}' for x in range(25)])

    X = df.loc[:,"x0":].to_numpy()
    t = df["treatment"].to_numpy()
    #t = np.zeros_like(t)
    y = df["y_factual"].to_numpy()
    y1 = np.where(df["treatment"] == 1,
                   df['y_factual'],
                   df['y_cfactual'])
    y0 = np.where(df["treatment"] == 0,
                   df['y_factual'],
                   df['y_cfactual'])
    ite = np.where(df["treatment"] == 1,
                   df['y_factual'] - df["y_cfactual"],
                   df['y_cfactual'] - df["y_factual"])
    if seperate_pool_test:
        X_train, X_pool, t_train, t_pool, y_train, y_pool, ite_train, ite_pool, y1_train, y1_pool, y0_train, y0_pool = train_test_split(
    X, t, y, ite, y1, y0,  test_size=test_size, random_state=seed)
        X_pool, X_test, t_pool, t_test, y_pool,y_test, ite_pool, ite_test, y1_pool,y1_test, y0_pool, y0_test = train_test_split(
    X_pool, t_pool, y_pool, ite_pool, y1_pool, y0_pool, test_size=pool_size, random_state=seed)
        ds = {"X_training": X_train,
             "y_training": y_train,
             "t_training": t_train,
             "X_pool": X_pool,
             "y_pool": y_pool,
             "t_pool": t_pool,
             "y1_pool": y1_pool,
             "y0_pool":y0_pool,
             "X_test": X_test,
             "y_test": y_test,
              "t_test": t_test,
              "ite_training" : ite_train,
              "ite_pool" : ite_pool,
              "ite_test": ite_test
             }
    else:
        X_train, X_test, t_train, t_test, y_train, y_test, ite_train, ite_test, y1_train, y1_test, y0_train, y0_test = train_test_split(
    X, t, y, ite, y1, y0,  test_size=test_size, random_state=1005)
        ds = {"X_training": X_train,
     "y_training": y_train,
     "t_training": t_train,
     "X_pool": deepcopy(X_test),
     "y_pool": deepcopy(y_test),
     "t_pool": deepcopy(t_test),
     "y1_pool": y1_test,
     "y0_pool":y0_test,
     "X_test": X_test,
     "y_test": y_test,
      "t_test": t_test,
      "ite_test": ite_test
     }
    return ds
