import os

import pandas as pd

from imblearn.under_sampling import RandomUnderSampler

from optparse import OptionParser

from sklearn.model_selection import train_test_split

#This will make a train/validation/test split 80/20/20

def resample_data(t):

    t = t[['HADM_ID', 'text', 'readm_30d']]

    label = t.pop('readm_30d')

    rus = RandomUnderSampler(random_state=42)

    X, y = rus.fit_resample(t, label.astype('category'))

    ids = pd.Series(X[:, 0])

    texts = pd.Series(X[:, 1])

    df = pd.DataFrame()

    df['readm_30d'] = pd.Series(y)

    df['HADM_ID'] = ids

    df['text'] = texts

    return df

def split_data(admissions, ratio):

    # Do some limited preprocessing

    X = admissions[['HADM_ID', 'text']]

    y = admissions['readm_30d']

    # Create a stratified train test split to preserver distribution.

    X_train, X_test, y_train, y_test = train_test_split(X, y,stratify=y, test_size=ratio, random_state=42)

    train = pd.merge(X_train, y_train, left_index=True, right_index=True)

    test = pd.merge(X_test, y_test, left_index=True, right_index=True)

    return train, test

def main(input_data, output_dir, ratio):

    # read the dataset from file.

    print("Reading raw data")

    data = pd.read_csv(input_data)

    # split into training and testing

    print("Splitting into training and testing")

    train, test = split_data(data, ratio)

    # split into train and validation

    print("spliting train into train and validation")

    train, validation = split_data(train, ratio)

    # undersample the train

    print("Undersampling the train")

    train = resample_data(train)

    # now save the files

    if not os.path.exists(output_dir):

        os.makedirs(output_dir)

    train.to_csv(os.path.join(output_dir, "train.csv"), index=None)

    test.to_csv(os.path.join(output_dir, "test.csv"), index=None)

    validation.to_csv(os.path.join(output_dir, "validation.csv"), index=None)

if __name__ == "__main__":

    parser = OptionParser()

    parser.add_option("--input", help="specify the input data")

    parser.add_option("--output_dir", help="specify the output location")

    parser.add_option("--ratio", help="specify the proportion to keep for testing", type="float")

    (options, args) = parser.parse_args()

    # load the data

    main(options.input, options.output_dir, options.ratio)