"""

Helper methods for reading data

Leon Zheng

"""

import pandas as pd

import numpy as np

from sklearn import preprocessing

def read_input(file_radiomics, file_clinical):

    """

    Read radiomics and clinical feature and return dataframe.

    :param file_radiomics: filename

    :param file_clinical: filename

    :return:

    """

    radiomics = pd.read_csv(file_radiomics, index_col=0)

    clinical = pd.read_csv(file_clinical, index_col=0)

    clinical = cleaning_clinical(clinical)

    input = pd.concat([radiomics, clinical], axis=1)

    return input

def normalizing_input(x_train, x_test):

    all_x = pd.concat([x_train, x_test])

    x_val = all_x.values

    min_max_scaler = preprocessing.MinMaxScaler()

    x_scaled = min_max_scaler.fit_transform(x_val)

    df = pd.DataFrame(x_scaled, index=all_x.index, columns=all_x.columns)

    return df.loc[x_train.index], df.loc[x_test.index]

def cleaning_clinical(clinical):

    """

    Cleaning the clinical dataframe.

    :param clinical:

    :return:

    """

    # Encoding label for SourceDataset

    le = preprocessing.LabelEncoder()

    le.fit(clinical['SourceDataset'])

    le.transform(clinical['SourceDataset'])

    clinical['SourceDataset'] = le.transform(clinical['SourceDataset'])

    # Cleaning Histology

    clinical.replace("NSCLC NOS (not otherwise specified)", "nos", inplace=True)

    clinical.replace("Adenocarcinoma", "adenocarcinoma", inplace=True)

    clinical.replace("Squamous cell carcinoma", "squamous cell carcinoma", inplace=True)

    # Dummies for Histology

    clinical = pd.get_dummies(clinical)

    # Fill age nan

    clinical['age'].fillna((clinical['age'].mean()), inplace=True)

    return clinical

def read_output(file_output):

    """

    Return dataframe for event + survival time.

    :param file_output:

    :return:

    """

    output = pd.read_csv(file_output, index_col=0, header=0)

    return output

def clean_clinical_data(file, newfile):

    """

    Cleaning clinical data.

    :param file: clinical data file

    :param newfile: path for the cleaned data file.

    :return:

    """

    clinical = pd.read_csv(file, index_col=0)

    cleaned = cleaning_clinical(clinical)

    cleaned.to_csv(newfile)

def y_dataframe_to_rsf_input(y_df):

    """

    Input for random survival forest.

    :param y_df: event + survival time dataframe.

    :return:

    """

    y_array = []

    Y = y_df.to_numpy()

    for y in Y:

        tuple = (bool(y[1]), y[0])

        y_array.append(tuple)

    return np.array(y_array, dtype = [(f'{y_df.columns[1]}', np.bool), (f'{y_df.columns[0]}', np.float)])

def load_owkin_data(radiomics_path_train="data/train/features/radiomics.csv",

                    clinical_path_train="data/train/features/clinical_data.csv",

                    label_path_train='data/train/y_train.csv',

                    radiomics_path_test="data/test/features/radiomics.csv",

                    clinical_path_test="data/test/features/clinical_data.csv"):

    """

    Load Owkin data: return PyRadiomics + clinical features of training set in dataframe,

                    event + time of training set in dataframe, and

                    PyRadiomics + clinical features of testing set in dataframe,

    """

    input_train = read_input(radiomics_path_train, clinical_path_train)

    output_train = read_output(label_path_train)

    input_test = read_input(radiomics_path_test, clinical_path_test)

    return input_train, output_train, input_test

if __name__ == '__main__':

    input_train, output_train, input_test = load_owkin_data()

    print(input_train, output_train, input_test)