""" Code for data loader """

import numpy as np

import os, sys, copy

import random

import tensorflow as tf

from sklearn.model_selection import StratifiedKFold

from tensorflow.python.platform import flags

import tqdm

import pickle as pkl

FLAGS = flags.FLAGS

PADDING_ID = 1016 # make the padding id as the number of group code

                  # maximum of group code index is 1015, start from 0

N_WORDS = 1017

TIMESTEPS = 21 # choice by statistics

TASKS = ["AD", "PD", "DM", "AM", "MCI"]

class DataLoader(object):

    '''

    Data Loader capable of generating batches of ohsu data.

    '''

    def __init__(self, source, target, true_target, n_tasks, n_samples_per_task, meta_batch_size):

        """

        Args:

            source:             source tasks

            target:             simulated target task(s)

            true_target:        true target task (to test)

            n_tasks:            number of tasks including both source and simulated target tasks

            n_samples_per_task: number samples to generate per task in one batch

            meta_batch_size:    size of meta batch size (e.g. number of functions)

        """

        ### load data: training

        self.intmd_path = 'intermediate/'

        self.source = source

        self.target = target

        self.timesteps = TIMESTEPS

        self.code_size = 0

        # self.code_size = N_WORDS-1 # set the code_size as the number of all the possible codes

        #                            # in order to use in pretrain

        self.task_code_size = dict() # maintain a dictionary for icd codes, disease : code list

        print ("The selected timesteps is: ", self.timesteps)

        self.data_to_show = dict()

        self.label_to_show = dict()

        self.ratio_t = 0.8

        self.pat_reduce = False

        self.code_set = set()

        self.data_s, self.data_t, self.label_s, self.label_t = self.load_data()

        ## load data: validate & test

        self.true_target = true_target

        if FLAGS.method == "mlp":

            data_tt, label_tt = self.load_data_vector(self.true_target[0]) # only 1 true target, index is 0

        elif FLAGS.method == "rnn" or FLAGS.method == "cnn":

            data_tt, label_tt = self.load_data_matrix(self.true_target[0])

            # compute code_size

            self.code_size = max([cz for cz in self.task_code_size.values()])

            print ("The code_size is: ", self.code_size)

            # make data the same size matrices

            data_tt, label_tt = self.get_data_prepared(data_tt, label_tt)

            for i in range(len(self.source)):

                self.data_s[i], self.label_s[i] = self.get_data_prepared(self.data_s[i], self.label_s[i])

            for i in range(len(self.target)):

                self.data_t[i], self.label_t[i] = self.get_data_prepared(self.data_t[i], self.label_t[i])

        # cross validation for true target

        self.n_fold = 5

        self.get_cross_val(data_tt, label_tt, n_fold=self.n_fold)

        ### set model params

        self.meta_batch_size = meta_batch_size

        self.n_samples_per_task = n_samples_per_task # in one meta batch

        self.n_tasks = n_tasks

        self.n_words = N_WORDS

        ## generate finetune data

        self.tt_sample, self.tt_label = dict(), dict()

        self.tt_sample_val, self.tt_label_val = dict(), dict()

        for ifold in range(self.n_fold): # generate n-fold cv data for finetuning

            self.tt_sample[ifold], self.tt_label[ifold] = self.generate_finetune_data(is_training=True, ifold=ifold)

            self.tt_sample_val[ifold], self.tt_label_val[ifold] = self.generate_finetune_data(is_training=False, ifold=ifold)

        self.episode = self.generate_meta_idx_batches(is_training=True)

        self.episode_val = dict()

        for ifold in range(self.n_fold): # true target validation

            self.episode_val[ifold] = self.generate_meta_idx_batches(is_training=False, ifold=ifold)

    def get_cross_val(self, X, y, n_fold=5):

        '''split the true target into train (might be useful in finetunning) and test (for evaluation)'''

        self.data_tt_tr, self.data_tt_val = dict(), dict()

        self.label_tt_tr, self.label_tt_val = dict(), dict()

        skf = StratifiedKFold(n_splits = n_fold, random_state = 99991)

        ifold = 0

        print ("split the true target ...")

        for train_index, test_index in skf.split(X, y):

            self.data_tt_tr[ifold], self.data_tt_val[ifold] = X[train_index], X[test_index]

            self.label_tt_tr[ifold], self.label_tt_val[ifold] = y[train_index], y[test_index]

            ifold+=1

    def load_data_matrix(self, task):

        '''load data sequential vectors for cnn or rnn. One matrix per sample'''

        X_pos, y_pos = [], []

        X_neg, y_neg = [], []

        with open(self.intmd_path + task + '.pos.pkl', 'rb') as f:

            X_pos_mat, y_pos_mat = pkl.load(f)

            f.close()

        with open(self.intmd_path + task + '.neg.pkl', 'rb') as f:

            X_neg_mat, y_neg_mat = pkl.load(f)

            f.close()

        print ("The number of positive samles in task %s is: " %task, len(y_pos_mat))

        print ("The number of negative samles in task %s is: " %task, len(y_neg_mat))

        for s, array in X_pos_mat.items():

             X_pos.append(array) # X_pos_mat[s] size: seq_len x n_words

             y_pos.append(y_pos_mat[s])

        for s, array in X_neg_mat.items():

             X_neg.append(array)

             y_neg.append(y_neg_mat[s])

        return (X_pos, X_neg), (y_pos, y_neg)

    def get_fixed_timesteps(self, X_pos, X_neg):

        '''delete the first several timesteps according to the selected number'''

        # postives:

        for i in range(len(X_pos)):

            timesteps = X_pos[i].shape[0]

            if timesteps > self.timesteps:

                X_pos[i] = X_pos[i][timesteps-self.timesteps:, :]

        # negatives:

        for i in range(len(X_neg)):

            timesteps = X_neg[i].shape[0]

            if timesteps > self.timesteps:

                X_neg[i] = X_neg[i][timesteps-self.timesteps:, :]

        return (X_pos, X_neg)

    def get_fixed_codesize(self, X_pos, X_neg):

        '''delete the -1 values according to the code size'''

        # postives:

        for i in range(len(X_pos)):

            code_size = X_pos[i].shape[1]

            if code_size > self.code_size:

                X_pos[i] = X_pos[i][:, :self.code_size]

        # negatives:

        for i in range(len(X_neg)):

            code_size = X_neg[i].shape[1]

            if code_size > self.code_size:

                X_neg[i] = X_neg[i][:, :self.code_size]

        return (X_pos, X_neg)

    def get_feed_records(self, X):

        '''generate ehrs as a 3d tensor that can be used to feed networks'''

        n_samples = len(X)

        X_new = np.zeros([n_samples, self.timesteps, self.code_size], dtype="int32") + PADDING_ID

        for i in range(n_samples):

            timesteps = X[i].shape[0]

            X_new[i, self.timesteps-timesteps:, :] = X[i]

        return X_new

    def get_data_prepared(self, data, label):

        X_pos, X_neg = data

        y_pos, y_neg = label

        X_pos, X_neg = self.get_fixed_timesteps(X_pos, X_neg)

        X_pos, X_neg = self.get_fixed_codesize(X_pos, X_neg)

        X_pos = self.get_feed_records(X_pos)

        X_neg = self.get_feed_records(X_neg)

        # concatenate pos and neg

        data, label = np.concatenate((X_pos, X_neg), axis=0), np.concatenate((y_pos, y_neg), axis=0)

        return data, label

    def load_data(self):

        '''load data vectors or matrices for samples with labels'''

        data_s, label_s = dict(), dict()

        data_t, label_t = dict(), dict()

        self.dim_input = [TIMESTEPS, N_WORDS]

        for i in range(len(self.source)):

            data_s[i], label_s[i] = self.load_data_matrix(self.source[i])

        for i in range(len(self.target)):

            data_t[i], label_t[i] = self.load_data_matrix(self.target[i])

        return data_s, data_t, label_s, label_t

    def generate_finetune_data(self, is_training=True, ifold=0):

        ''' get finetuning samples and labels'''

        try:

            if is_training:

                sample = self.data_tt_tr[ifold]

                label = self.label_tt_tr[ifold]

            else:

                sample = self.data_tt_val[ifold]

                label = self.label_tt_val[ifold]

        except:

            print ("Error: split training and validate first!")

        return sample, label

    def generate_meta_batches(self, is_training=True, ifold=0):

        ''' get samples and the corresponding labels with episode for batching'''

        if is_training: # training

            prefix = "metatrain"

            data_s = self.data_s

            data_t = self.data_t

            label_s = self.label_s

            label_t = self.label_t

            self.n_total_batches = FLAGS.n_total_batches

        else: # test & eval, say, true target task is used here

            try:

                prefix = "metaval" + str(ifold)

                data_s = self.data_s

                label_s = self.label_s

                data_t = self.data_tt_val[ifold]

                label_t = self.label_tt_val[ifold]

                self.n_total_batches = int(len(label_t)/self.n_samples_per_task)

            except:

                print ("Error: split training and validate first!")

        # check if the meta batch file dumped

        if os.path.isfile(self.intmd_path + "meta.batch." + prefix + ".pkl"):

            print ('meta batch file exits')

            with open(self.intmd_path + "meta.batch." + prefix + ".pkl", 'rb') as f:

                sample, label = pkl.load(f)

                f.close()

        else:

            # generate episode

            sample, label = [], []

            s_dict, t_dict = dict(), dict()

            for i in range(len(self.source)):

                s_dict[i] = range(len(self.label_s[i]))

            for i in range(len(self.target)):

                t_dict[i] = range(len(self.label_t[i]))

            batch_count = 0

            for _ in tqdm.tqdm(range(self.n_total_batches), 'generating meta batches'): # progress bar

                # i.e., sample 16 patients from selected tasks

                # len of spl and lbl: 4 * 16

                spl, lbl = [], [] # samples and labels in one episode

                for i in range(len(self.source)): # fetch from source tasks olderly

                    ### do not keep pos/neg ratio

                    s_idx = random.sample(s_dict[i], self.n_samples_per_task)

                    spl.extend(data_s[i][s_idx])

                    lbl.extend(label_s[i][s_idx])

                ### do not keep pos/neg ratio

                if is_training:

                    t_idx = random.sample(t_dict[0], self.n_samples_per_task)

                    spl.extend(data_t[0][t_idx])

                    lbl.extend(label_t[0][t_idx])

                else:

                    spl.extend(data_t[batch_count*self.n_samples_per_task:(batch_count+1)*self.n_samples_per_task])

                    lbl.extend(label_t[batch_count*self.n_samples_per_task:(batch_count+1)*self.n_samples_per_task])

                batch_count += 1

                # add meta_batch

                sample.append(spl)

                label.append(lbl)

        print ("batch counts: ", batch_count)

        sample = np.array(sample, dtype="float32")

        label = np.array(label, dtype="float32")

        return sample, label

    def generate_meta_idx_batches(self, is_training=True, ifold=0):

        ''' get samples and the corresponding labels with episode for batching'''

        if is_training: # training

            prefix = "metatrain"

            data_s = self.data_s

            data_t = self.data_t

            label_s = self.label_s

            label_t = self.label_t

            self.n_total_batches = FLAGS.n_total_batches

        else: # test & eval, say, true target task is used here

            try:

                prefix = "metaval" + str(ifold)

                data_s = self.data_s

                label_s = self.label_s

                data_t = self.data_tt_val[ifold]

                label_t = self.label_tt_val[ifold]

                self.n_total_batches = int(len(label_t)/self.n_samples_per_task)

                print (data_t.shape)

                print (label_t.shape)

                print (len(label_t))

            except:

                print ("Error: split training and validate first!")

        # generate episode

        episode = []

        s_dict, t_dict = dict(), dict()

        for i in range(len(self.source)):

            s_dict[i] = range(len(self.label_s[i]))

        for i in range(len(self.target)):

            t_dict[i] = range(len(self.label_t[i]))

        batch_count = 0

        for _ in tqdm.tqdm(range(self.n_total_batches), 'generating meta batches'): # progress bar

            # i.e., sample 16 patients from selected tasks

            # len of spl and lbl: 4 * 16

            idx = [] # index in one episode

            for i in range(len(self.source)): # fetch from source tasks olderly

                ### do not keep pos/neg ratio

                s_idx = random.sample(s_dict[i], self.n_samples_per_task)

                idx.extend(s_idx)

            ### do not keep pos/neg ratio

            if is_training:

                t_idx = random.sample(t_dict[0], self.n_samples_per_task)

                idx.extend(t_idx)

            else:

                t_idx = range(batch_count*self.n_samples_per_task, (batch_count+1)*self.n_samples_per_task)

                idx.extend(t_idx)

            batch_count += 1

            # add meta_batch

            episode.append(idx)

        print ("batch counts: ", batch_count)

        return episode