'''
This file contains a class which will load the different
attributes of the patients of a given dataset.
Firstly, NSCLC-Radiogenomics
'''

from copy import deepcopy
import pandas as pd
import numpy as np
import os
from itertools import chain
from sklearn.preprocessing import LabelBinarizer
import nibabel as nib


class Dataset(object):
    '''
    Class which contains a list of all patients,
    images, genomics, recurrence, survival, and mutation information.

    If the corresponding data is unavailable, 'NA' will
    be used.
    '''

    def __init__(self, config, dataset='NSCLC-Radiogenomics'):
        self.dataset_info = dataset
        self.config = config
        self.get_patient_list()

        # these are list of files which should be read appropriately
        self.image_list = {}
        self.images = {}
        self.seg_list = {}
        self.feature_list = {}

        # these are already loaded attributes
        self.genomics_list = {}
        self.egfr_mutation = {}
        self.recurrence_bool = {}
        self.recurrence_value = {}
        self.survival_bool = {}
        self.survival_value = {}
        self.durations = {}

        #TODO: include this:       self.last_known_alive = []
        self.clinical_list = {}

        self.load_all()

    def get_patient_list(self):

        self.data_location = self.config.location
        self.patient_list = list(pd.read_csv(self.config.clinical)['Case ID'])
        return

    def set_patient_list(self, patient_list):
        self.patient_list = patient_list

    def load_all(self):
        self.load_images()
        self.load_segmentations()
        self.load_pyradiomics()
        self.load_genomics()
        self.load_recurrence()
        self.load_survival()
        self.load_egfr_mutation()
#         self.load_clinical()
        self.load_densenet_features()

    def load_images(self):
        for patientID in self.patient_list:
            image_path = self.config.images + patientID + '.nii'
            if os.path.exists(image_path):
                self.image_list[patientID] = image_path
            else:
                self.image_list[patientID] = 'N/A'
        return

    def get_images_cropped(self):
        self.images = {}
        for patientID in self.patient_list:
            image_path = self.config.cropped + patientID + '_cropped_nodule.nii'
            if os.path.exists(image_path):
                self.images[patientID] = nib.load(image_path).get_fdata()
            else:
                self.images[patientID] = 'N/A'
        return

    def get_pyradiomics(self, patient_list):
        features=[]
        for patientID in patient_list:
            feature_path = self.config.pyradiomics + patientID + '_dilated.npz'
            features.append(np.load(feature_path)['arr_0'])
        return features

    def get_densenet_features(self, patient_list):
        features=[]
        for patientID in patient_list:
            feature_path = self.config.densenet + patientID + '_densenet.npy'
            loaded = np.load(feature_path)
            features.append(loaded)

        features = np.array(features)
        features = np.squeeze(features)
        return features

    def get_genomics(self, patient_list):
        genomics = pd.read_csv(self.config.genomics, index_col=False)
        genomics.set_index('Unnamed: 0.1', inplace=True)
        genomics = genomics.drop('Unnamed: 0', axis=1)
        genomics = genomics.transpose()

        # TODO: we can add code here to normalize the genomics data
        genomics_list = []
        for id in patient_list:
            genomics_list.append(list(genomics.loc[id]))
        return genomics_list, genomics.columns


    def get_clinical(self, patient_list):
        clinical = pd.read_csv(self.config.clinical)

        list_of_variables = clinical.columns.values

        predictors_labels = list(chain([list_of_variables[0]], [list_of_variables[2]], list_of_variables[6:8], list_of_variables[9:22], list_of_variables[23:24])) #:30]))
        predictors = clinical[predictors_labels]

        predictors.set_index('Case ID', inplace=True)

        predictors['Smoking status'].replace(self.config.smoking_dict, inplace=True)
        predictors['Pack Years'].replace({'N/A': 0, 'Not Collected': 40}, inplace=True)
        predictors['%GG'].replace(self.config.gg_dict, inplace=True)
        for idx in range(10, 17):
            predictors[list_of_variables[idx]].replace(self.config.location_dict, inplace=True)

        encoder = LabelBinarizer()
        for idx in range(17, 24):
            if idx == 22:
                continue
            predictors[list_of_variables[idx]] = encoder.fit_transform(predictors[list_of_variables[idx]])
        predictors.fillna(0, inplace=True)

        clinical_data = []

        for id in patient_list:
            clinical_data.append([int(x) for x in list(predictors.loc[id])])

        return clinical_data

    def load_segmentations(self):
        self.seg_list = {}
        for patientID in self.patient_list:
            seg_path = self.config.segs + patientID + '.nii.gz'
            if os.path.exists(seg_path):
                self.seg_list[patientID] = seg_path
            else:
                self.seg_list[patientID] = 'N/A'
        return

    def load_pyradiomics(self):

        self.feature_list = {}
        for patientID in self.patient_list:
            feature_path = self.config.pyradiomics + patientID + '.npz'
            if os.path.exists(feature_path):
                self.feature_list[patientID] = feature_path
            else:
                self.feature_list[patientID] = 'N/A'
        return

    def load_densenet_features(self):
        self.densenet_features = {}
        for patientID in self.patient_list:
            feature_path = self.config.densenet + patientID + '_densenet.npy'
            if os.path.exists(feature_path):
                temp = np.load(feature_path)
                if np.size(temp) == 1:
                    self.densenet_features[patientID] = 'N/A'
                else:
                    self.densenet_features[patientID] = feature_path
            else:
                self.densenet_features[patientID] = 'N/A'
        return

    def load_genomics(self):
        self.genomics_list = {}
        genomics = pd.read_csv(self.config.genomics, index_col=False)
        genomics.set_index('Unnamed: 0.1', inplace=True)
        genomics = genomics.drop('Unnamed: 0', axis=1)
        genomics = genomics.transpose()

        for id in self.patient_list:
            if id in genomics.index.values:
               self.genomics_list[id] = list(genomics.loc[id])
            else:
                self.genomics_list[id] = 'N/A'
        return

    def load_recurrence(self):
        #TODO: Include the location information as well

        self.recurrence_value = {}
        self.recurrence_bool = {}
        self.durations = {}
        recurrence = pd.read_csv(self.config.recurrence, index_col=False)
        recurrence.set_index('Case ID', inplace=True)

        for id in self.patient_list:
            if id in recurrence.index.values:
                curr_patient = recurrence.loc[id]
                value = curr_patient['Recurrence']
                self.recurrence_bool[id] = value
                self.recurrence_value[id] = curr_patient['Days']

        return

    def load_survival(self):

        self.survival_value = {}
        self.survival_bool = {}
        recurrence = pd.read_csv(self.config.clinical, index_col=False)
        recurrence.set_index('Case ID', inplace=True)

        for id in self.patient_list:
            if id in recurrence.index.values:
                curr_patient = recurrence.loc[id]
                value = curr_patient['Survival Status']
                mapped_value = self.config.survival_mapping[value]
                self.survival_bool[id] = mapped_value

                if mapped_value == 1:
                    self.survival_value[id] = curr_patient['Time to Death (days)']
                else:
                    self.survival_value[id] = 'N/A'
            else:
                self.survival_bool[id] = 'N/A'
                self.survival_value[id] = 'N/A'
        return

    def load_egfr_mutation(self):
        self.egfr_mutation = {}
        egfr = pd.read_csv(self.config.clinical, index_col=False)

        egfr.set_index('Case ID', inplace=True)

        for id in self.patient_list:
            if id in egfr.index.values:
                value = egfr.loc[id]['EGFR mutation status']
                mapped_value = self.config.mutation_mapping[value]
                self.egfr_mutation[id] = mapped_value

            else:
                self.egfr_mutation[id] = 'N/A'
        return

    def load_clinical(self):
        self.clinical_list = {}
        clinical = pd.read_csv(self.config.clinical)

        list_of_variables = clinical.columns.values

        # clinical = clinical.loc[49:]

        # for id in range(len(list_of_variables)):
        #     print(id, list_of_variables[id])

        predictors_labels = list(chain([list_of_variables[0]], [list_of_variables[2]], list_of_variables[6:8], list_of_variables[9:22], list_of_variables[23:24])) #:30]))
        predictors = clinical[predictors_labels]

        predictors.set_index('Case ID', inplace=True)

        predictors['Smoking status'].replace(self.config.smoking_dict, inplace=True)
        predictors['Pack Years'].replace({'N/A': 0, 'Not Collected': 40}, inplace=True)
        predictors['%GG'].replace(self.config.gg_dict, inplace=True)
        for idx in range(10, 17):
            predictors[list_of_variables[idx]].replace(self.config.location_dict, inplace=True)

        encoder = LabelBinarizer()
        for idx in range(17, 24):
            if idx == 22:
                continue
            predictors[list_of_variables[idx]] = encoder.fit_transform(predictors[list_of_variables[idx]])
        predictors.fillna(0, inplace=True)

        print(predictors.columns.values)
        for id in self.patient_list:
            self.clinical_list[id] = [int(x) for x in list(predictors.loc[id])]
        return

    def select_subset_patients(self, to_select, replace_list=False):
        '''

        :param dataset: dataset of type Dataset
        :param to_select: list of features to subselect
        :return: updated dataset
        '''

        patient_list = deepcopy(self.patient_list)
        for id in self.patient_list:
            remove_bool = False
            for attr in to_select:
                if attr == 'pyradiomics':
                    if self.feature_list[id] == 'N/A':
                        remove_bool = True
                if attr == 'gene_expressions':
                    if self.genomics_list[id] == 'N/A':
                        remove_bool = True
                if attr == 'clinical':
                    if self.clinical_list[id] == 'N/A':
                        remove_bool = True
                if attr == 'recurrence':
                    if self.recurrence_bool[id] == 'N/A':
                        remove_bool = True
                if attr == 'densenet':
                    if self.densenet_features[id] == 'N/A':
                        remove_bool = True
                if attr == 'egfr':
                    if self.egfr_mutation[id] == 'N/A':
                        remove_bool = True
            if remove_bool is True:
                patient_list.remove(id)

        if replace_list == True:
            self.set_patient_list(patient_list)
            self.load_all()

        return patient_list

if __name__ == '__main__':
    nrg = Dataset()