import numpy as np
import os
import random
import pickle


class MimicDataLoader(object):

    def __init__(self, data_folder, batch_size=1, max_sequence=10, max_iter=None, split = 0.75, train_keep=1):
        super(MimicDataLoader, self).__init__()
        self.data_folder = data_folder
        self.batch_size = batch_size
        self.num_step = max_sequence
        self.max_iter = max_iter
        self.num_iter = 0
        self.input_map=pickle.load(open(data_folder+'/dig_map.pkl','rb'))
        self.ouput_map = pickle.load(open(data_folder + '/proc_map.pkl', 'rb'))
        self.all_input = pickle.load(open(data_folder+'/dig_input.pkl','rb'))
        self.all_output = pickle.load(open(data_folder + '/proc_output.pkl', 'rb'))
        self.output_size = self.all_output.shape[1]
        if len(np.shape(self.all_output))>1:
            self.all_output = np.argmax(self.all_output, axis=1)
        print(self.all_output[:10])
        print(self.all_output.shape)
        self.num_samples=self.all_input.shape[0]
        print('num samples {}'.format(self.num_samples))
        lindex=list(range(self.num_samples))
        # random.shuffle(lindex)
        self.train_data_indexes = lindex[:int(self.num_samples*split*train_keep)]
        self.test_data_indexes = lindex[int(self.num_samples*split):]
        self.is_training=True
        self.data_offset=0
        self.input_size=self.all_input.shape[1]


        print('num train samples: {}'.format(len(self.train_data_indexes)))
        print('train index: {} ...'.format(self.train_data_indexes[:10]))
        print('num test samples: {}'.format(len(self.test_data_indexes)))
        print('test index: {} ...'.format(self.test_data_indexes[:10]))
        print('num classes: {}'.format(self.output_size))
        print('num steps per episode: {}'.format(self.num_step))
        print('batch size: {}'.format(self.batch_size))

    def fetch_all(self):
        train_x=[]
        train_y=[]
        test_x=[]
        test_y=[]
        for ind in self.train_data_indexes:
            train_x.append(self.all_input[ind])
            train_y.append(self.all_output[ind])
        for ind in self.test_data_indexes:
            test_x.append(self.all_input[ind])
            test_y.append(self.all_output[ind])
        return np.asarray(train_x), np.asarray(train_y), np.asarray(test_x), np.asarray(test_y)

    def fetch_batch(self, is_training=True):

        if is_training:
            list_index=self.train_data_indexes
        else:
            list_index=self.test_data_indexes

        indexes = np.zeros((self.batch_size, self.num_step), dtype=np.int32)
        for i in range(self.batch_size):
            indexes[i, :] = np.random.choice(len(list_index), self.num_step, replace=False)
        # print('-------------')
        # print(indexes[:10])
        all_inputs=[]
        all_outputs=[]

        for s in range(self.num_step):
            example_inputs = np.zeros((self.batch_size, self.input_size))
            example_outputs = np.zeros(self.batch_size)
            for b in range(self.batch_size):
                example_inputs[b,:]=self.all_input[list_index[indexes[b,s]]]
                all_inputs.append(example_inputs.astype('float32'))
                example_outputs[b] = self.all_output[list_index[indexes[b, s]]]
                all_outputs.append(example_outputs.astype('int32'))

        return all_inputs, all_outputs

    def fetch_batch_full(self, is_training, is_rand=True):

        if is_training:
            list_index=self.train_data_indexes
        else:
            list_index=self.test_data_indexes

        num_t = len(list_index)

        indexes = np.zeros((self.batch_size, num_t),dtype=np.int32)
        for i in range(self.batch_size):
            if is_rand:
                indexes[i, :] = np.random.choice(len(list_index), num_t, replace=False)
            else:
                indexes[i, :] = np.asarray(list(range(len(list_index))))

        # indexes = np.zeros((self.batch_size, num_t), dtype=np.int32)
        # for i in range(self.batch_size):
        #     indexes[i,:]=np.arange(num_t)

        all_inputs=[]
        all_outputs=[]

        for s in range(num_t):
            example_inputs = np.zeros((self.batch_size, self.input_size))
            example_outputs = np.zeros(self.batch_size)
            for b in range(self.batch_size):
                example_inputs[b,:]=self.all_input[list_index[indexes[b,s]]]
                all_inputs.append(example_inputs.astype('float32'))
                example_outputs[b] = self.all_output[list_index[indexes[b, s]]]
                all_outputs.append(example_outputs.astype('int32'))

        return all_inputs, all_outputs

    # indexes just have shape (batch,)
    def predict_index2data(self, list_index, indexes):
        all_inputs = []
        all_outputs = []
        for s in range(self.num_step):
            example_inputs = np.zeros((self.batch_size, self.input_size))
            example_outputs = np.zeros(self.batch_size)
            for b in range(self.batch_size):
                example_inputs[b,  :] = self.all_input[list_index[indexes[b]]]
                all_inputs.append(example_inputs)
                example_outputs[b] = self.all_output[list_index[indexes[b]]]
                all_outputs.append(example_outputs)


        return all_inputs, all_outputs


    def predict_online_index2data(self, list_index, indexes):
        all_inputs = []
        all_outputs = []
        for s in range(self.num_step):
            example_inputs = np.zeros((self.batch_size, self.input_size))
            example_outputs = np.zeros(self.batch_size)
            for b in range(self.batch_size):
                example_inputs[b,  :] = self.all_input[list_index[indexes[b]]]
                all_inputs.append(example_inputs)
                example_outputs[b] = self.all_output[list_index[indexes[b]]]
                all_outputs.append(example_outputs)


        return all_inputs, all_outputs




