# -*- coding: utf-8 -*-

import paddle
import paddle.fluid as fluid
import json
import codecs
import numpy
import link_model
import time
import os
import sys

reload(sys)
sys.setdefaultencoding("utf-8")

id2char, char2id = json.load(open('data_path/all_chars_me.json'))
print("Finish loaded char_dict!")

def get_train_data(filename):
    def reader():
        with codecs.open(filename, encoding="utf-8") as f:
            for l in f:
                _ = json.loads(l)
                text = _['text']
                kb_dict = _['kb_dict']
                label = _['label']
                new_slot = text + kb_dict
                text_slot = [char2id.get(c, 1) for c in new_slot]
                # kb_dict_slot = [char2id.get(c, 1) for c in kb_dict]
                target_slot = numpy.zeros(2)
                target_slot[int(label)] = 1
                sample_result = [text_slot, target_slot]
                if sample_result is None:
                    continue
                yield tuple(sample_result)

    print("Finish loaded train_data!")
    return reader

def del_file(path):
    ls = os.listdir(path)
    for i in ls:
        c_path = os.path.join(path, i)
        if os.path.isdir(c_path):
            del_file(c_path)
        else:
            os.remove(c_path)

char_len = len(char2id)
embedding_size = 64
batch_size = 128

def train():
    # train_data = get_train_data("data_path/all_train_deal.json")
    train_data = get_train_data("data_path/link_train_data.json")
    dev_data = get_train_data("data_path/link_dev_data.json")
    text = fluid.layers.data(name="text", shape=[1], dtype='int64', lod_level=1)
    desc = fluid.layers.data(name="desc", shape=[1], dtype='int64', lod_level=1)
    target = fluid.layers.data(
        name='target', shape=[2], dtype='float32', lod_level=0)

    feature_out = link_model.bilstm(text=text,desc=desc, embdding_size=embedding_size, char_len=char_len)

    entropy_cost = fluid.layers.cross_entropy(input=feature_out, label=target, soft_label=True)
    avg_cost = fluid.layers.mean(entropy_cost)

    # valid program
    valid_program = fluid.default_main_program().clone(for_test=True)

    optimizer = fluid.optimizer.AdamOptimizer(learning_rate=0.001)
    optimizer.minimize(avg_cost)

    train_batch_reader = paddle.batch(
        reader=paddle.reader.shuffle(train_data, buf_size=6144),
        batch_size=batch_size
    )

    valid_batch_reader = paddle.batch(
        reader=paddle.reader.shuffle(dev_data, buf_size=2048), batch_size=batch_size
    )

    # place = fluid.CPUPlace()
    place = fluid.CUDAPlace(0)

    feeder = fluid.DataFeeder(feed_list=[text,desc, target], place=place)
    exe = fluid.Executor(place)

    save_dirname = './output/link_model'

    exe.run(fluid.default_startup_program())
    start_time = time.time()
    batch_id = 0
    best_precision = 0.0
    for pass_id in range(200):
        pass_start_time = time.time()
        cost_sum, cost_counter = 0, 0
        for data in train_batch_reader():
            # print(data[1][1])
            cost = exe.run(program=fluid.default_main_program(), feed=feeder.feed(data),
                           fetch_list=[avg_cost])
            # ss= numpy.array(feature_outs)
            # sss= numpy.array(kb_dict_slot)
            # print("feature_outs::::::")
            # print(ss.shape)
            # print(ss[1])
            # print("targets::::::")
            # print(sss.shape)
            # print(sss[1])
            # print("======================")
            # cost = numpy.array(cost)
            cost = cost[0]
            cost_sum += cost
            cost_counter += 1
            if batch_id % 40 == 0 and batch_id != 0:
                print >> sys.stderr, "batch %d finished, second per batch: %02f, avg_cost: %f" % (
                    batch_id, (time.time() - start_time) / batch_id, cost)

            if batch_id % 200 == 0 and batch_id != 0:
                all_target_count_0 = 0.0
                all_target_count_1 = 0.0
                all_correct_count_0 = 0.0
                all_correct_count_1 = 0.0
                for valid_data in valid_batch_reader():
                    targets, feature_outs = exe.run(program=valid_program, feed=feeder.feed(valid_data),
                                                    fetch_list=[target, feature_out],
                                                    return_numpy=False)

                    label_tag_scores = numpy.array(feature_outs)
                    target_label_tag_scores = numpy.array(targets)
                    for r in range(len(target_label_tag_scores)):
                        tag_pre = label_tag_scores[r].argmax()
                        tag_target = target_label_tag_scores[r].argmax()
                        if tag_target == 1:
                            all_target_count_1 += 1
                            if tag_target == tag_pre:
                                all_correct_count_1 += 1
                        else:
                            all_target_count_0 += 1
                            if tag_target == tag_pre:
                                all_correct_count_0 += 1

                precision_1 = all_correct_count_1 / all_target_count_1
                precision_0 = all_correct_count_0 / all_target_count_0
                print(all_correct_count_1)
                print(all_correct_count_0)
                print(all_target_count_1)
                print(all_target_count_0)
                print('======================')

                if precision_1 > best_precision:
                    save_path = os.path.join(save_dirname, 'best')
                    if os.path.exists(save_path) == False:
                        os.makedirs(save_path)
                    del_file(save_path)
                    with open(save_path + "/best_precision.txt", "w") as f:
                        f.write("precision_1:" + str(precision_1) + "\tprecision_0:" + str(precision_0))
                    fluid.io.save_inference_model(save_path, ['text','desc'],
                                                  [feature_out], exe, main_program=valid_program,
                                                  params_filename='params')
                    best_precision = precision_1
                print >> sys.stderr, "Valid end, precision_1: %f,\tprecision_0: %f" % (precision_1, precision_0)

            batch_id = batch_id + 1
        pass_avg_cost = cost_sum / cost_counter if cost_counter > 0 else 0.0
        print >> sys.stderr, "%d pass end, cost time: %02f, avg_cost: %f" % (
            pass_id, time.time() - pass_start_time, pass_avg_cost)

    else:
        # pass times complete and the training is over
        save_path = os.path.join(save_dirname, 'final')
        fluid.io.save_inference_model(save_path, ['text','desc'],
                                      [feature_out], exe, params_filename='params')
    return

train()