import sys
import os
import time

import cPickle as pickle
import tensorflow as tf
import numpy as np

import utils.reader as reader
import utils.evaluation as eva


def train(conf, _model):
    
    if conf['rand_seed'] is not None:
        np.random.seed(conf['rand_seed'])

    if not os.path.exists(conf['save_path']):
        os.makedirs(conf['save_path'])

    # load data
    print('starting loading data')
    print(time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time())))
    train_data, val_data, test_data = pickle.load(open(conf["data_path"], 'rb'))
    print('lyang test: val_data: ', len(val_data), type(val_data),
          len(val_data['y']))
    print('lyang test: val_data[y]: ', val_data['y'][0:2])
    print('lyang test: val_data[c]: ', val_data['c'][0:2])
    print('lyang test: val_data[r]: ', val_data['r'][0:2])
    print('lyang test: val_data[qids]: ', val_data['qids'][0:2])
    print('lyang test: val_data[dids]: ', val_data['dids'][0:2])
    print('map id to words ...')
    id2word = reader.read_dict('../data/' + conf["data_name"]+ '/word2id')
    response_ids = val_data['r'][0:1][0]
    context_ids = val_data['c'][0:1][0]
    print('lyang test: val_data[c]: ',
          [id2word[str(id)] for id in context_ids], val_data.keys())
    print('lyang test: val_data[r]: ',
          [id2word[str(id)] for id in response_ids], val_data.keys())

    print('finish loading data')
    print(time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time())))

    print('init intent_dict...')
    conf['intent_dict'] = reader.read_intent(conf['intent_vec_path']) if conf[
        'model_name'] != 'dam' else None
    print('lyang test len(conf[intent_dict])', len(conf['intent_dict']))
    val_batches = reader.build_batches(val_data, conf)

    # check the example 0 and 1 in batch 0
    print('intent of val_batches context: ', val_batches['turns_intent'][0][0:2])
    print('intent of val_batches response: ',val_batches['response_intent'][0][0:2])

    print("finish building valid batches")
    print(time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time())))

    # refine conf
    batch_num = len(train_data['y']) / conf["batch_size"]
    print('batch_size: ', conf["batch_size"])
    print('total number of batches in one epoch: ', batch_num)
    val_batch_num = len(val_batches["response"])

    conf["train_steps"] = conf["num_scan_data"] * batch_num # total number of training steps epoch_num * batch_num
    conf["save_step"] = max(1, batch_num / 10) # at most save 10 times
    conf["print_step"] = max(1, batch_num / 100) # at most print 100 times

    print('configurations:')
    conf_copy = {}
    for k in conf:
        if k != 'intent_dict':
            conf_copy[k] = conf[k]
    print(conf_copy)

    print('model sucess')
    print(time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time())))

    _graph = _model.build_graph()
    print('build graph sucess')
    print(time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time())))

    with tf.Session(graph=_graph) as sess:
        # train_writer = tf.summary.FileWriter(
        #     conf["save_path"] + "tensorboard_log/", sess.graph)
        # merge = tf.summary.merge_all()  # for tensorboard
        _model.init.run();
        if conf["init_model"]:
            _model.saver.restore(sess, conf["init_model"])
            print("sucess init %s" %conf["init_model"])

        average_loss = 0.0
        batch_index = 0
        step = 0
        best_result = [0, 0, 0, 0]

        for step_i in xrange(conf["num_scan_data"]):
            #for batch_index in rng.permutation(range(batch_num)):
            print('starting shuffle train data')
            shuffle_train = reader.unison_shuffle(train_data)
            train_batches = reader.build_batches(shuffle_train, conf)
            print('finish building train data')
            for batch_index in range(batch_num):

                feed = {
                    _model.turns: train_batches["turns"][batch_index], 
                    _model.tt_turns_len: train_batches["tt_turns_len"][batch_index],
                    _model.every_turn_len: train_batches["every_turn_len"][batch_index],
                    _model.response: train_batches["response"][batch_index], 
                    _model.response_len: train_batches["response_len"][batch_index],
                    _model.label: train_batches["label"][batch_index],
                }
                if conf['model_name'] != 'dam':
                    feed[_model.turns_intent] = train_batches["turns_intent"][batch_index]
                    feed[_model.response_intent] = train_batches["response_intent"][batch_index]

                batch_index = (batch_index + 1) % batch_num;

                _, curr_loss = sess.run([_model.g_updates, _model.loss], feed_dict = feed)
                # print loss and metrics into tensorboard log
                # train_writer.add_summary(summ, global_step=step)

                average_loss += curr_loss

                step += 1

                if step % conf["print_step"] == 0 and step > 0:
                    g_step, lr = sess.run([_model.global_step, _model.learning_rate])
                    print('step: %s lr: %s, epoch: %s ' %(g_step, lr, step_i))
                    print("step: " + str(g_step)+ " processed current epoch: [" \
                          + str(step * 1.0 / batch_num) + "] loss: " + \
                          str(average_loss / conf["print_step"]))
                    average_loss = 0

                if step % conf["save_step"] == 0 and step > 0:
                    index = step / conf['save_step']
                    score_file_path = conf['save_path'] + 'score.' + str(index)
                    score_file = open(score_file_path, 'w')
                    print('save step: %s' %index)
                    print(time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time())))

                    for batch_index in xrange(val_batch_num):
                
                        feed = { 
                            _model.turns: val_batches["turns"][batch_index],
                            _model.tt_turns_len: val_batches["tt_turns_len"][batch_index],
                            _model.every_turn_len: val_batches["every_turn_len"][batch_index],
                            _model.response: val_batches["response"][batch_index],
                            _model.response_len: val_batches["response_len"][batch_index],
                            _model.label: val_batches["label"][batch_index]
                        }
                        if conf['model_name'] != 'dam':
                            feed[_model.turns_intent] = \
                            val_batches["turns_intent"][batch_index]
                            feed[_model.response_intent] = \
                            val_batches["response_intent"][batch_index]
                
                        scores = sess.run(_model.logits, feed_dict = feed)
                    
                        for i in xrange(conf["batch_size"]):
                            score_file.write(
                                str(scores[i]) + '\t' + 
                                str(val_batches["label"][batch_index][i]) + '\n')
                    score_file.close()

                    #write evaluation result
                    result = eva.evaluate(score_file_path)
                    result_file_path = conf["save_path"] + "result." + str(index)
                    with open(result_file_path, 'w') as out_file:
                        for m in result:
                            out_file.write(str(m) + '\n')
                    print('finish evaluation')
                    # lyang: also print metrics in log file
                    print('save step:\t{:d}\t[current metrics (r2@1 r10@1 r10@2 r10@5 map)]\t{:f}\t{:f}\t{:f}\t{:f}\t{:f}'.format(
                        index, result[0], result[1], result[2], result[3], result[4]))
                    print(time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time())))
                    # lyang: also print metrics in tensorboard log file
                    # metrics = tf.Summary(value=[
                    #     tf.Summary.Value(tag="R10at1", simple_value=result[1]),
                    # ])
                    # metrics.value.add(tag="MAP", simple_value=result[4])
                    # # metrics.value.add(tag="R10at2", simple_value=result[2])
                    # # metrics.value.add(tag="R10at5", simple_value=result[3])
                    # train_writer.add_summary(metrics, global_step=step)

                    if result[1] + result[2] > best_result[1] + best_result[2]: # save model only when find a model better than previously best model
                        best_result = result
                        _save_path = _model.saver.save(sess, conf["save_path"] + "model.ckpt." + str(step / conf["save_step"]))
                        print("succ saving model in " + _save_path)
                        print(time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time())))



                

