
# coding: utf-8

# In[1]:


import tensorflow as tf
import numpy as np
from sklearn.metrics import mean_squared_error
from sklearn.metrics import r2_score
from sklearn import metrics
from math import sqrt
import time
import random
from RNNUtil import RNNUtil
from RNNCell import RNNCell
from DataPreprocessor import DataPreprocessor

import csv


# In[2]:


# generic class for running student data
# pilot dataset is assistments
# generic rnn parameters ...
# input vector size => input does not have to be one-hot
# output vector size
# state vector size => for each hidden layer => len(state_size) is number of hidden layers
# cell type for each hidden layer
# activation => for each hidden layer
# dropout probability => for each hidden layer
# batch size
# num_steps => number of maximum time steps
# training or testing
# inject prediction function => TODO
# inject loss function => TODO
# inject trainer => TODO

class Generic_RNN(object):
    # use this method to create rnn
    @staticmethod
    def make_rnn_param_config(input_size, 
                              output_size, 
                              state_size, 
                              cell_type,
                              activation,
                              keep_in, 
                              keep_out, 
                              keep_states, 
                              num_steps,
                              is_training):
        config = {
            'input_size': input_size, # int
            'output_size': output_size, # int
            'state_size': state_size, # list of ints
            'cell_type': cell_type, # list of types
            'activation': activation, # list of strings
            'keep_in': keep_in, # list of floats
            'keep_out': keep_out, # list of floats
            'keep_states': keep_states, # list of floats
            'num_steps': num_steps, # int
            'is_training': is_training # boolean
        }
        
        return config
    
    # method to pre-check the validity of config
    # only checks types and length
    @staticmethod
    def check_rnn_param_config_validity(config):
        input_size = config.get('input_size', None)
        output_size = config.get('output_size', None)
        state_size = config.get('state_size', None)
        cell_type = config.get('cell_type', None)
        activation = config.get('activation', None)
        keep_in = config.get('keep_in', None)
        keep_out = config.get('keep_out', None)
        keep_states = config.get('keep_states', None)
        num_steps = config.get('num_steps', None)
        is_training = config.get('is_training', None)
        
        is_valid = RNNUtil.is_type_length_valid(var=state_size, var_type_list=[list])
        
        hidden_layer_num = len(state_size)
        
        is_valid = (
            is_valid
            and RNNUtil.is_type_length_valid(var=input_size, var_type_list=[int])
            and RNNUtil.is_type_length_valid(var=output_size, var_type_list=[int]) 
            and RNNUtil.is_type_length_valid(var=cell_type, var_type_list=[list], length=hidden_layer_num) 
            and RNNUtil.is_type_length_valid(var=activation, var_type_list=[list], length=hidden_layer_num)
            and RNNUtil.is_type_length_valid(var=keep_in, var_type_list=[list], length=hidden_layer_num) 
            and RNNUtil.is_type_length_valid(var=keep_out, var_type_list=[list], length=hidden_layer_num)
            and RNNUtil.is_type_length_valid(var=keep_states, var_type_list=[list], length=hidden_layer_num)
            and RNNUtil.is_type_length_valid(var=num_steps, var_type_list=[int])
            and RNNUtil.is_type_length_valid(var=is_training, var_type_list=[bool])
        )
        
        return is_valid
    
    def __init__(self, rnn_param_config):
        is_valid = __class__.check_rnn_param_config_validity(rnn_param_config)
        if (not is_valid):
            print('rnn_param_config is not valid')
            print(rnn_param_config)
            exit(1)
        
        self.set_rnn_params(rnn_param_config)
        self.set_graph()
        
    def set_rnn_params(self, param_config):
        self.param_config = param_config
        self.num_features = param_config.get('input_size', None)
        self.num_classes = param_config.get('output_size', None)
        self.cell_type = param_config.get('cell_type', None)
        self.state_size = param_config.get('state_size', None)
        self.num_steps = param_config.get('num_steps', None)
        self.is_training = param_config.get('is_training', None)
        self.activation_str = param_config.get('activation', None)
        self.keep_in = param_config.get('keep_in', None)
        self.keep_states = param_config.get('keep_states', None)
        self.keep_out = param_config.get('keep_out', None)
    
    def set_graph(self):
        logits = self.build_feedforward_graph()
    
        if (self.is_training):
            # prediction and loss nodes
            self.pred, self.loss = pred, loss = self.predict_and_loss(logits)
                        
            # cost
            self.cost = cost = self.get_cost_from_loss(loss, 'reduce_sum')
            
            # TODO: more oop trainer
            # get optimizer
            starting_learning_rate = 0.1
            init = tf.constant(dtype=tf.int32, value=0)
            global_step = tf.get_variable(name='global_step', dtype=tf.int32, trainable=False, 
                                          initializer=init)
            learning_rate = tf.train.exponential_decay(starting_learning_rate, global_step, 3000, 0.96, staircase=True)
            epsilon = 0.1

            self.optimizer = RNNUtil.get_trainer('adamOptimizer', learning_rate, epsilon)
    
            # cap gradient
            max_grad_norm = 20.0
            grads_and_vars = self.optimizer.compute_gradients(cost)
            grads_and_vars = [(tf.clip_by_norm(g, max_grad_norm), v) for g, v in grads_and_vars if g is not None]

            # optimization
            self.eval_op = self.optimizer.apply_gradients(grads_and_vars, name = 'train_op', global_step = global_step)
        
        else:
            self.pred = pred = self.predict(logits)
            self.eval_op = tf.no_op()
        
    # build rnn inputs, sequence length, cells, cell, outputs, states, weight, bias, logits    
    def build_feedforward_graph(self):
        # inputs: [batch_size, num_steps, num_features]
        self.inputs = inputs = tf.placeholder(dtype = tf.float32, 
                                              shape = [None, self.num_steps, self.num_features])
        
        # seq_len: [batch_size]
        self.seq_len = seq_len = tf.placeholder(dtype = tf.int32, shape = [None])
        
        hidden_layer_num = len(self.state_size)
        
        # cells refer to individual cells
        self.cells = cells =  [RNNCell.makeRNNCell(self.cell_type[layer_index], 
                                                   self.state_size[layer_index],
                                                   self.is_training,
                                                   activation_str = self.activation_str[layer_index],
                                                   keep_in = self.keep_in[layer_index],
                                                   keep_states = self.keep_states[layer_index],
                                                   keep_out = self.keep_out[layer_index]) 
                               for layer_index in range(hidden_layer_num)]
        
        # cell is multi rnn cell
        self.multi_cell = multi_cell = tf.contrib.rnn.MultiRNNCell(cells = cells,
                                                                   state_is_tuple = True)
        
        # outputs and states
        # outputs: [batch_size, num_steps, state_size[-1]]
        self.outputs, self.states = outputs, states = tf.nn.dynamic_rnn(cell = multi_cell, 
                                                                        inputs = self.inputs,
                                                                        sequence_length = self.seq_len,
                                                                        dtype = tf.float32)
        
        # outputs: num_steps x [batch_size, state_size[-1]]
        outputs = [tf.squeeze(output, axis = 1) 
                   for output in tf.split(value = outputs, 
                                          num_or_size_splits = self.num_steps, 
                                          axis = 1)]
        
        # weight and bias
        # weight: [state_size[-1], num_classes]
        # bias: [num_classes]
        self.weight = weight = tf.get_variable('weight', [self.state_size[-1], self.num_classes])
        self.bias = bias = tf.get_variable('bias', [self.num_classes])
    
        # produce logit outputs
        # logits: num_steps x [batch_size, num_classes]
        logits = [tf.matmul(outputs[i], weight) + bias for i in range(self.num_steps)]
        
        # stack logits: [num_steps, batch_size, num_classes]
        logits = tf.stack(logits)
        
        # transpose logits: [batch_size, num_steps, num_classes]
        logits = tf.transpose(logits, [1, 0, 2])
        
        return logits
            
    # need to input target_id
    def predict(self, logits):
        # reshape logits: [batch_size x num_steps x num_classes]
        logits = tf.reshape(logits, [-1])

        # target_id[batch_i x num_steps + step_i] = batch_i x num_steps x num_classes + step_i x num_classes + class_i
        self.target_id = target_id =tf.placeholder(dtype = tf.int32, shape = [None])
        selected_logits = tf.gather(params = logits, indices = target_id)
        
        return tf.nn.sigmoid(selected_logits)
    
    # need to input target_id, target_correctness
    def predict_and_loss(self, logits):
        # reshape logits: [batch_size x num_steps x num_classes]
        logits = tf.reshape(logits, [-1])

        # target_id[batch_i x num_steps + step_i] = batch_i x num_steps x num_classes + step_i x num_classes + class_i
        self.target_id = target_id = tf.placeholder(dtype = tf.int32, shape = [None])
        selected_logits = tf.gather(params = logits, indices = target_id)
        
        # define skill_seq_len
        # denominator for each selected_logits element
        # its size is the same as target_id
        # logits[target_id[i]] = logits[target_id[i]] / skill_seq_len[i]
        # => skill_seq_len[batch_i * num_steps + step_i] = N: number of skills for each interaction
        self.skill_seq_len = skill_seq_len = tf.placeholder(dtype = tf.float32, shape = [None])
        
        self.target_correctness = target_correctness = tf.placeholder(dtype = tf.float32, shape = [None])
        unweighted_loss = tf.nn.sigmoid_cross_entropy_with_logits(logits = selected_logits,
                                                                  labels = target_correctness)
        weighted_loss = tf.divide(unweighted_loss, skill_seq_len)

        return tf.nn.sigmoid(selected_logits), weighted_loss
    
    def get_cost_from_loss(self, loss, policy):
        if ('reduce_mean' == policy):
            cost = tf.reduce_mean(loss)
        elif ('reduce_sum' == policy):
            cost = tf.reduce_sum(loss)
        else:
            print('{} policy not yet realized'.format(policy))
            exit(1)
        
        return cost          


# In[3]:


# RNN using Datapreprocessor
class RNN_Datapreprocessor(object):
    def __init__(self):
        self.model_name = 'DKT'
    
    # run once for specific dataset config
    def run(self, run_config, rnn_param_config, datapreprocessor_config, data_config, ext):
        # run_config includes ...
        # num_epochs: default to 150
        # init_scale: default to 0.05
        # batch_size: default to 100
        
        # rnn_param_config includes every config except input_size, output_size, num_steps, is_training

        # set datapreprocessor
        self.datapreprocessor = self.load_datapreprocessor(datapreprocessor_config)
        self.data_config = data_config
        
        # load rnn data
        self.train = self.load_rnn_data(data_config, is_training=True, ext=ext)
        self.test = self.load_rnn_data(data_config, is_training=False, ext=ext)
        
        #self.train_students, train_num_steps, train_num_skills = self.read_data_from_csv_file('../data/csv_rnn_data/0910_b_train.csv')
        #self.test_students, test_num_steps, test_num_skills = self.read_data_from_csv_file('../data/csv_rnn_data/0910_b_test.csv')

        if ('pkl' == ext):
            self.train_students = self.train['students']
            train_num_steps = int(self.train['num_steps'])
            train_num_skills = int(self.train['num_skills'])
            
            self.test_students = self.test['students']
            test_num_steps = int(self.test['num_steps'])
            test_num_skills = int(self.test['num_skills'])
            
            train_rnn_param_config = Generic_RNN.make_rnn_param_config(input_size=train_num_skills*2, 
                                                                       output_size=train_num_skills,
                                                                       num_steps=train_num_steps-1, 
                                                                       is_training=True, 
                                                                       state_size=rnn_param_config.get('state_size', [200]),
                                                                       cell_type=rnn_param_config.get('cell_type', ['LSTM']),
                                                                       activation=rnn_param_config.get('activation', ['tanh']),
                                                                       keep_in=rnn_param_config.get('keep_in', [1.]), 
                                                                       keep_out=rnn_param_config.get('keep_out', [0.6]), 
                                                                       keep_states=rnn_param_config.get('keep_states', [1.]))
            
            test_rnn_param_config = Generic_RNN.make_rnn_param_config(input_size=test_num_skills*2, 
                                                                      output_size=test_num_skills,
                                                                      num_steps=test_num_steps-1, 
                                                                      is_training=False, 
                                                                      state_size=rnn_param_config.get('state_size', [200]),
                                                                      cell_type=rnn_param_config.get('cell_type', ['LSTM']),
                                                                      activation=rnn_param_config.get('activation', ['tanh']),
                                                                      keep_in=rnn_param_config.get('keep_in', [1.]), 
                                                                      keep_out=rnn_param_config.get('keep_out', [1.]), 
                                                                      keep_states=rnn_param_config.get('keep_states', [1.]))
                                        
        else:
            print('loading from {} not yet realized'.format(ext))
            exit(1)
        
        # start graph
        with tf.Graph().as_default():
            session_conf = tf.ConfigProto(allow_soft_placement = True,
                                          log_device_placement = False)
            session_conf.gpu_options.allow_growth = True
            
            # start session
            with tf.Session(config = session_conf) as sess:
                init_scale = run_config.get('init_scale', 0.05)
                num_epochs = run_config.get('num_epochs', 150)
                batch_size = run_config.get('batch_size', 100)
                shuffle_every_epoch = run_config.get('shuffle_every_epoch', True)
                result_path = self.get_result_path(run_config)
                print('result_path: ', result_path)
                
                initializer = tf.random_uniform_initializer(-init_scale, init_scale)
                
                # set up train and test models
                with tf.variable_scope('model', reuse = None, initializer = initializer):
                    self.train_model = Generic_RNN(train_rnn_param_config)
                
                with tf.variable_scope('model', reuse = True, initializer = initializer):
                    self.test_model = Generic_RNN(test_rnn_param_config)
                
                # run epochs
                self.run_epochs(sess, batch_size, num_epochs, shuffle_every_epoch, result_path, eval_interval=5)
                
    def run_epochs(self, sess, batch_size, num_epochs, shuffle_every_epoch, result_file_path, eval_interval=5):
        # initialize
        sess.run(tf.global_variables_initializer())

        # saver
        saver = tf.train.Saver(tf.global_variables())
        train_students = self.train_students
        
        if (not shuffle_every_epoch):
            random.shuffle(train_students)
        
        for epoch in range(num_epochs):
            if (shuffle_every_epoch):
                random.shuffle(train_students)
            
            rmse, auc, r2 = self.run_epoch(sess, batch_size, train_students, is_training=True)
            print("Epoch: %d Train Metrics:\n rmse: %.3f \t auc: %.3f \t r2: %.3f \n" % (epoch + 1, rmse, auc, r2))
            
            with open(result_file_path, "a+") as f:
                f.write("Epoch: %d Train Metrics:\n rmse: %.3f \t auc: %.3f \t r2: %.3f \n" % (epoch + 1, rmse, auc, r2))
                f.write("\n")

                print("*"*10)
                f.write("\n")
            
            # save testing results
            if (0 == (epoch + 1) % eval_interval):
                save_path = saver.save(sess, self.model_name)
                print('*' * 10)
                print('Start to test model')
                
                test_students = self.test_students
                # shuffle test students
                #random.shuffle(test_students)
                
                rmse, auc, r2 = self.run_epoch(sess, batch_size, test_students, is_training=False)
                print("Epoch: %d Test Metrics:\n rmse: %.3f \t auc: %.3f \t r2: %.3f" % (epoch + 1, rmse, auc, r2))
                
                with open(result_file_path, "a+") as f:
                    f.write("Epoch: %d Test Metrics:\n rmse: %.3f \t auc: %.3f \t r2: %.3f" % (epoch + 1, rmse, auc, r2))
                    f.write("\n")

                    print("*"*10)
                    f.write("\n")
    
    # run rnn for one_hot and not_one_hot
    def run_epoch(self, sess, batch_size, students, is_training):
        if (is_training):
            model = self.train_model
        else:
            model = self.test_model
        
        batch_start_i = 0
        pred_labels = []
        actual_labels = []

        not_one_hot = (not self.data_config.get('one_hot', None)) and self.data_config.get('allow_multi_skills', None)
        is_one_hot = not not_one_hot
        
        while (batch_start_i + batch_size < len(students)):
            if (batch_start_i + batch_size < len(students)):
                mini_batch_size = batch_size
            else:
                mini_batch_size = len(students) - batch_start_i
            mini_batch_size = 100
            
            pred, actual = self.run_mini_batch(sess, 
                                               mini_batch_size,
                                               batch_start_i, 
                                               model, 
                                               students, 
                                               actual_labels, 
                                               pred_labels, 
                                               is_training,
                                               is_one_hot)
            pred_labels.extend(pred)
            actual_labels.extend(actual)
            
            batch_start_i += batch_size
        
        # print pred labels
        #print('len(actual_labels): ', len(actual_labels))
        #print('len(pred_lables): ', len(pred_labels))
        rmse = sqrt(mean_squared_error(actual_labels, pred_labels))
        fpr, tpr, thresholds = metrics.roc_curve(actual_labels, pred_labels, pos_label = 1)
        auc = metrics.auc(fpr, tpr)

        # calculate r2
        r2 = r2_score(actual_labels, pred_labels)

        return rmse, auc, r2
    
    # return pred labels and actual labels to append
    def run_mini_batch(self, 
                       sess, 
                       batch_size, 
                       batch_start_i, 
                       model, 
                       students, 
                       actual_labels, 
                       pred_labels, 
                       is_training, 
                       is_one_hot=True):

        actual_labels = []
        pred_labels = []
        
        input_x = np.zeros((batch_size, model.num_steps, model.num_features))
        target_id = []
        target_correctness = []
        
        # seq_len[student_i]: number of steps for each student
        # skill_seq_len: number of skills for each interaction. constructed as 1d
        seq_len = np.empty(dtype = np.int32, shape = [batch_size])
        skill_seq_len = []
        
        for student_i in range(batch_size):
            student = students[batch_start_i + student_i]
            seq_len[student_i] = steps = int(student[0][0]) - 1
            skill_ids = student[1]
            correctness = student[2]

            # one_hot
            if (is_one_hot):
                # skill ids is a list of integers for one_hot
                for step_i in range(steps):
                    skill_id = int(skill_ids[step_i])
                    is_correct = int(correctness[step_i])
                    # student_i x num_steps x num_classes + step_i x num_classes + class_i
                    target_id.append(student_i * model.num_steps * model.num_classes 
                                     + step_i * model.num_classes + int(skill_ids[step_i + 1]))
                    skill_seq_len.append(1)
                    
                    target_correctness.append(int(correctness[step_i + 1]))
                    actual_labels.append(int(correctness[step_i + 1]))

                    if (is_correct):
                        input_x[student_i, step_i, skill_id] = 1
                    else:
                        input_x[student_i, step_i, skill_id + model.num_classes] = 1

            # not_one_hot
            else:
                # skill ids is a list of lists of integers for not_one_hot
                # target correctness should be duplicated with the length same as multiple skills
                for step_i in range(steps):
                    skill_ids_input = skill_ids[step_i]
                    is_correct = int(correctness[step_i])
                    # student_i x num_steps x num_classes + step_i x num_classes + k-th skill to predict
                    skill_ids_to_predict = skill_ids[step_i + 1]
                    
                    for skill_id_to_predict in skill_ids_to_predict:
                        target_id.append(student_i * model.num_steps * model.num_classes 
                                         + step_i * model.num_classes + int(skill_id_to_predict))
                        
                        skill_seq_len.append(len(skill_ids_to_predict))
                        
                        target_correctness.append(int(correctness[step_i + 1]))
                        actual_labels.append(int(correctness[step_i + 1]))
                        
                    for skill_id in skill_ids_input:
                        try:
                            if (is_correct):
                                input_x[student_i, step_i, skill_id] = 1
                            else:
                                input_x[student_i, step_i, skill_id + model.num_classes] = 1
                        except:
                            print('student_i: ', student_i)
                            print('step_i: ', step_i)
                            print('skill_id: ', skill_id)
                            print('num_classes: ', model.num_classes)
                            print('input_x.shape: ', input_x.shape)
        
        feed_dict = {
            model.inputs: input_x, model.target_id: target_id,
            model.seq_len: seq_len
        }

        if (is_training):
            feed_dict[model.target_correctness] = target_correctness
            feed_dict[model.skill_seq_len] = skill_seq_len

        pred, _ = sess.run([model.pred, model.eval_op], feed_dict = feed_dict)

        # since prediction is always float, does not have to consider one_hot or not
        for p in pred:
            pred_labels.append(p)
        
        return pred_labels, actual_labels        
                
    def load_datapreprocessor(self, datapreprocessor_config):
        dataset = datapreprocessor_config.get('dataset', 'Assistments')
        version = datapreprocessor_config.get('version', '2009')
        return DataPreprocessor(dataset, version)
        
    
    def load_rnn_data(self, data_config, is_training, ext='pkl'):
        if ('datapreprocessor' not in self.__dict__):
            print('please set datapreprocessor first')
            exit(1)
        else:
            self.datapreprocessor.set_config(data_config)
            return self.datapreprocessor.load_rnn_data(is_training=is_training, ext=ext)
    
    def read_data_from_csv_file(self, fileName):
        inputs = []
        targets = []
        rows = []
        max_skill_num = 0
        max_num_problems = 0
        with open(fileName, "r") as csvfile:
            reader = csv.reader(csvfile, delimiter=',')
            for row in reader:
                rows.append(row)
        index = 0
        i = 0
        print ("the number of rows is " + str(len(rows)))
        tuple_rows = []
        #turn list to tuple
        while(index < len(rows)-1):
            problems_num = int(rows[index][0])
            tmp_max_skill = max(map(int, rows[index+1]))
            if(tmp_max_skill > max_skill_num):
                max_skill_num = tmp_max_skill
            if(problems_num <= 2):
                index += 3
            else:
                if problems_num > max_num_problems:
                    max_num_problems = problems_num
                tup = (rows[index], rows[index+1], rows[index+2])
                tuple_rows.append(tup)
                index += 3
        #shuffle the tuple

        random.shuffle(tuple_rows)