"""Load Adult dataset and manage cross validation."""
import folktables 
from folktables import ACSDataSource
import pandas as pd
from sklearn.model_selection import StratifiedKFold
import numpy as np

from . import Dataset
from . import split

def load(sensitive=[],k=0):
    """Download if necessary folktables adult. Split and return train and test.

    :param sensitive: (Optional default=[]) List of sensitive attributes to include in the features. The sensitive attribute are "sex" and "race".
    :type sensitive: list of str
    :param k: (Optinal default=0) Corss validation step in {0,1,2,3,4}.
    :type k: int
    :return: Train and test split dataframes in a dictionary.
    :rtype: Dictionary
    """
    #states = ['AL', 'AK', 'AZ', 'AR', 'CA', 'CO', 'CT', 'DE', 'FL', 'GA', 'HI','ID', 'IL', 'IN', 'IA', 'KS', 'KY', 'LA', 'ME', 'MD', 'MA', 'MI', 'MN', 'MS', 'MO', 'MT', 'NE', 'NV', 'NH', 'NJ', 'NM', 'NY', 'NC', 'ND', 'OH', 'OK', 'OR', 'PA', 'RI', 'SC', 'SD', 'TN', 'TX', 'UT','VT', 'VA', 'WA', 'WV', 'WI', 'WY', 'PR']
    data_source = ACSDataSource(survey_year='2018', horizon='1-Year', survey='person')
    ca_data = data_source.get_data(states=["CA"],download=True)

    features=[
        'AGEP',
        'COW',
        'SCHL',
        'MAR',
        'OCCP',
        'POBP',
        'RELP',
        'WKHP',
        'PINCP'
    ]
    human_to_census = {"race":"RAC1P","sex":"SEX"}
    for s in sensitive:
        features += [human_to_census[s]]
    ACSIncome = folktables.BasicProblem(features=features,
        target='PINCP',
        target_transform=lambda x: x > 50000,    
        group=['SEX','RAC1P'],
        preprocess=folktables.adult_filter,
        postprocess=lambda x: np.nan_to_num(x, -1),
    )
    ca_features, ca_labels, ca_attrib = ACSIncome.df_to_pandas(ca_data)
    ca_features["PINCP"] = ca_features["PINCP"]>50000
    ca_features.rename({"RAC1P":"race","SEX":"sex"},axis=1,inplace=True)
    data_split = split.split_pandas(ca_features)
    out = {"train":Dataset(),"test":Dataset()}
    out["train"].update(data_split["train"])
    out["test"].update(data_split["test"])
    
    return out
