#加了循环输入
#加了jieba分词
#更改保存参数路径
#更改文档名全加后缀new
import tensorflow as tf
from tensorflow.python.layers.core import Dense

import numpy as np
import time

import jieba
import json


# 超参数
# Number of Epochs
epochs = 10000
# Batch Size
batch_size = 200
# RNN Size
rnn_size = 50
# Number of Layers
num_layers = 2
# Embedding Size
encoding_embedding_size = 15
decoding_embedding_size = 15
# Learning Rate
learning_rate = 0.001



set_words=[]

source = open("data/source_new_final.txt",'w',encoding='utf-8')
target = open("data/target_new_final.txt",'w',encoding='utf-8')

test=open("data/text_new_perfectworld.txt",'w',encoding='utf-8')

jiebacut_text=open('data/完美世界new_final.txt',encoding='utf-8').read()
seg_list=jieba.cut(jiebacut_text,cut_all=False)
liststr="".join(seg_list)


test.write(liststr)
test.close()

with open("data/text_new_perfectworld.txt",'r',encoding='utf-8') as f:
    lines = f.readlines()
    for line in lines:
        line=line.strip().split("，")
        for l in range(len(line)-1):
            if line[l]=="" or line[l+1]=="" or line[l]=="”" or line[l+1]=="”" or line[l]=="/" or line[l+1]=="/"or line[l]=="/”/" or line[l+1]=="/”/"or line[l]=="/）/" or line[l+1]=="/）/":
                continue
            source.write(line[l]+'\n')
            target.write(line[l+4]+'\n')
source.close()
target.close()

with open('data/source_new_final.txt','r',encoding='utf-8') as f:
    source_data = f.read()

with open('data/target_new_final.txt','r',encoding='utf-8') as f:
    target_data = f.read()

print(source_data.split('\n')[:10])
print(target_data.split('\n')[:10])




def extract_character_vocab(data):

    global set_words
    special_words = ['<PAD>','<UNK>','<GO>','<EOS>']
    for line in data.split('\n'):

        set_words.extend(jieba.cut(line,cut_all=False))
    set_words=list(set(set_words))
    int_to_vocab = {idx:word for idx,word in enumerate(special_words + set_words)}
    vocab_to_int = {word:idx for idx,word in int_to_vocab.items()}

    return int_to_vocab,vocab_to_int


source_int_to_letter,source_letter_to_int = extract_character_vocab(source_data+target_data)
target_int_to_letter,target_letter_to_int = extract_character_vocab(source_data+target_data)


with open('data/source_int_to_letter_new_final.json','w',encoding='utf-8') as f:
    json.dump(source_int_to_letter,f)
with open('data/source_letter_to_int_new_final.json','w',encoding='utf-8') as f:
    json.dump(source_letter_to_int,f)
with open('data/target_int_to_letter_new_final.json','w',encoding='utf-8') as f:
    json.dump(target_int_to_letter,f)
with open('data/target_letter_to_int_new_final.json','w',encoding='utf-8') as f:
    json.dump(target_letter_to_int,f)

with open('data/source_new_final.txt','r',encoding='utf-8') as f:
    source_data = f.read()

with open('data/target_new_final.txt','r',encoding='utf-8') as f:
    target_data = f.read()


with open('data/source_int_to_letter_new_final.json','r',encoding='utf-8') as load_f:
    source_int_to_letter=json.load(load_f)
with open('data/source_letter_to_int_new_final.json','r',encoding='utf-8') as load_f:
    source_letter_to_int=json.load(load_f)
with open('data/target_int_to_letter_new_final.json','r',encoding='utf-8') as load_f:
    target_int_to_letter=json.load(load_f)
with open('data/target_letter_to_int_new_final.json','r',encoding='utf-8') as load_f:
    target_letter_to_int=json.load(load_f)


source_int=[[source_letter_to_int.get(letter,source_letter_to_int['<UNK>'])
            for letter in jieba.cut(line,cut_all=False)] for line in source_data.split('\n')]
#print(source_int)

target_int = [[target_letter_to_int.get(letter, target_letter_to_int['<UNK>'])
               for letter in jieba.cut(line,cut_all=False)] + [target_letter_to_int['<EOS>']] for line in target_data.split('\n')]

print(target_int)


# 输入层
def get_inputs():

    inputs = tf.placeholder(tf.int32,[None,None],name='inputs')
    targets = tf.placeholder(tf.int32,[None,None],name='targets')
    learning_rate = tf.placeholder(tf.float32,name='learning_rate')

    # 定义target序列最大长度（之后target_sequence_length和source_sequence_length会作为feed_dict的参数）
    target_sequence_length = tf.placeholder(tf.int32,(None,),name='target_sequence_length')
    max_target_sequence_length = tf.reduce_max(target_sequence_length,name='max_target_len')
    source_sequence_length = tf.placeholder(tf.int32,(None,),name='source_sequence_length')

    return inputs,targets,learning_rate,target_sequence_length,max_target_sequence_length,source_sequence_length



def get_encoder_layer(input_data,rnn_size,num_layers,source_sequence_length,source_vocab_size,encoding_embedding_size):

    encoder_embed_input = tf.contrib.layers.embed_sequence(input_data,source_vocab_size,encoding_embedding_size)

    def get_lstm_cell(rnn_size):
        lstm_cell = tf.contrib.rnn.LSTMCell(rnn_size,initializer=tf.random_uniform_initializer(-0.1,0.1,seed=2))
        return lstm_cell

    cell =  tf.contrib.rnn.MultiRNNCell([get_lstm_cell(rnn_size) for _ in range(num_layers)])

    encoder_output , encoder_state = tf.nn.dynamic_rnn(cell,encoder_embed_input,sequence_length=source_sequence_length,dtype=tf.float32)

    return encoder_output,encoder_state



def process_decoder_input(data,vocab_to_int,batch_size):

    ending = tf.strided_slice(data,[0,0],[batch_size,-1],[1,1])
    decoder_input = tf.concat([tf.fill([batch_size,1],vocab_to_int['<GO>']),ending],1)

    return decoder_input


def decoding_layer(target_letter_to_int,decoding_embedding_size,num_layers,rnn_size,
                   target_sequence_length,max_target_sequence_length,encoder_state,decoder_input):

    # 1. Embedding
    target_vocab_size = len(target_letter_to_int)
    decoder_embeddings = tf.Variable(tf.random_uniform([target_vocab_size,decoding_embedding_size]))
    decoder_embed_input = tf.nn.embedding_lookup(decoder_embeddings,decoder_input)###实际上tf.nn.embedding_lookup的作用就是找到要寻找的embedding data中的对应的行下的vector。

    # 构造Decoder中的RNN单元
    def get_decoder_cell(rnn_size):
        decoder_cell = tf.contrib.rnn.LSTMCell(rnn_size,initializer=tf.random_uniform_initializer(-0.1,0.1,seed=2))
        return decoder_cell

    cell = tf.contrib.rnn.MultiRNNCell([get_decoder_cell(rnn_size) for _ in range(num_layers)])

    # Output全连接层
    # target_vocab_size定义了输出层的大小
    output_layer = Dense(target_vocab_size,kernel_initializer=tf.truncated_normal_initializer(mean=0.1,stddev=0.1))##tf.truncated_normal_initializer 从截断的正态分布中输出随机值。生成的值服从具有指定平均值和标准偏差的正态分布，如果生成的值大于平均值2个标准偏差的值则丢弃重新选择。


    # 4. Training decoder
    with tf.variable_scope("decode"):
        training_helper = tf.contrib.seq2seq.TrainingHelper(inputs = decoder_embed_input,
                                                            sequence_length = target_sequence_length,
                                                            time_major = False)


        training_decoder = tf.contrib.seq2seq.BasicDecoder(cell,training_helper,encoder_state,output_layer)
        training_decoder_output,_,_ = tf.contrib.seq2seq.dynamic_decode(training_decoder,impute_finished=True,
                                                                        maximum_iterations = max_target_sequence_length)

    
    # 5. Predicting decoder
    # 与training共享参数
    with tf.variable_scope("decode",reuse=True):
        # 创建一个常量tensor并复制为batch_size的大小
        start_tokens = tf.tile(tf.constant([target_letter_to_int['<GO>']],dtype=tf.int32),[batch_size],name='start_token')
        predicting_helper = tf.contrib.seq2seq.GreedyEmbeddingHelper(decoder_embeddings,start_tokens,target_letter_to_int['<EOS>'])

        predicting_decoder = tf.contrib.seq2seq.BasicDecoder(cell,
                                                             predicting_helper,
                                                             encoder_state,
                                                             output_layer)
        predicting_decoder_output,_,_ = tf.contrib.seq2seq.dynamic_decode(predicting_decoder,impute_finished = True,
                                                                          maximum_iterations = max_target_sequence_length)


    return training_decoder_output,predicting_decoder_output

# 上面已经构建完成Encoder和Decoder，下面将这两部分连接起来，构建seq2seq模型
def seq2seq_model(input_data,targets,lr,target_sequence_length,max_target_sequence_length,
                  source_sequence_length,source_vocab_size,target_vocab_size,encoder_embedding_size,
                  decoder_embedding_size,rnn_size,num_layers):

    _,encoder_state = get_encoder_layer(input_data,
                                        rnn_size,
                                        num_layers,
                                        source_sequence_length,
                                        source_vocab_size,
                                        encoding_embedding_size)

    decoder_input = process_decoder_input(targets,target_letter_to_int,batch_size)

    training_decoder_output,predicting_decoder_output = decoding_layer(target_letter_to_int,
                                                                       decoding_embedding_size,
                                                                       num_layers,
                                                                       rnn_size,
                                                                       target_sequence_length,
                                                                       max_target_sequence_length,
                                                                       encoder_state,
                                                                       decoder_input)

    return training_decoder_output,predicting_decoder_output






# 构造graph
train_graph = tf.Graph()

with train_graph.as_default():
    input_data, targets, lr, target_sequence_length, max_target_sequence_length, source_sequence_length = get_inputs()

    training_decoder_output, predicting_decoder_output = seq2seq_model(input_data,
                                                                       targets,
                                                                       lr,
                                                                       target_sequence_length,
                                                                       max_target_sequence_length,
                                                                       source_sequence_length,
                                                                       len(source_letter_to_int),
                                                                       len(target_letter_to_int),
                                                                       encoding_embedding_size,
                                                                       decoding_embedding_size,
                                                                       rnn_size,
                                                                       num_layers)

    training_logits = tf.identity(training_decoder_output.rnn_output,'logits')
    #training_logits = training_logits + 1e-8
    #training_logits = tf.clip_by_value(training_logits,1e-8,1)貌似限制在两个数之间
    predicting_logits = tf.identity(predicting_decoder_output.sample_id,name='predictions')

    #mask是权重的意思
    #tf.sequence_mask([1, 3, 2], 5)  # [[True, False, False, False, False],
                                #  [True, True, True, False, False],
                                #  [True, True, False, False, False]]
    masks = tf.sequence_mask(target_sequence_length,max_target_sequence_length,dtype=tf.float32,name="masks")

    # logits: A Tensor of shape [batch_size, sequence_length, num_decoder_symbols] and dtype float.
    # The logits correspond to the prediction across all classes at each timestep.
    #targets: A Tensor of shape [batch_size, sequence_length] and dtype int.
    # The target represents the true class at each timestep.
    #weights: A Tensor of shape [batch_size, sequence_length] and dtype float.
    # weights constitutes the weighting of each prediction in the sequence. When using weights as masking,
    # set all valid timesteps to 1 and all padded timesteps to 0, e.g. a mask returned by tf.sequence_mask.
    with tf.name_scope("optimization"):
        cost = tf.contrib.seq2seq.sequence_loss(
            training_logits,
            targets,
            masks
        )

        optimizer = tf.train.AdamOptimizer(lr)

        # minimize函数用于添加操作节点，用于最小化loss，并更新var_list.
        # 该函数是简单的合并了compute_gradients()与apply_gradients()函数返回为一个优化更新后的var_list，
        # 如果global_step非None，该操作还会为global_step做自增操作

        #这里将minimize拆解为了以下两个部分：

        # 对var_list中的变量计算loss的梯度 该函数为函数minimize()的第一部分，返回一个以元组(gradient, variable)组成的列表
        gradients = optimizer.compute_gradients(cost)
        capped_gradients = [(tf.clip_by_value(grad, -5., 5.), var) for grad, var in gradients if grad is not None]
        # 将计算出的梯度应用到变量上，是函数minimize()的第二部分，返回一个应用指定的梯度的操作Operation，对global_step做自增操作
        train_op = optimizer.apply_gradients(capped_gradients)


def pad_sentence_batch(sentence_batch,pad_int):

    max_sentence = max([len(sentence) for sentence in sentence_batch])
    return [sentence + [pad_int] * (max_sentence - len(sentence)) for sentence in sentence_batch]


def get_batches(targets,sources,batch_size,source_pad_int,target_pad_int):

    for batch_i in range(0,len(sources)//batch_size):
        start_i = batch_i * batch_size
        sources_batch = sources[start_i : start_i + batch_size]
        targets_batch = targets[start_i : start_i + batch_size]

        pad_sources_batch = np.array(pad_sentence_batch(sources_batch,source_pad_int))
        pad_targets_batch = np.array(pad_sentence_batch(targets_batch,target_pad_int))

        targets_lengths = []
        for target in targets_batch:
            targets_lengths.append(len(target))

        source_lengths = []
        for source in sources_batch:
            source_lengths.append(len(source))

        yield pad_targets_batch,pad_sources_batch,targets_lengths,source_lengths



# Train
train_source = source_int[batch_size:]
train_target = target_int[batch_size:]

# 留出一个batch进行验证
valid_source = source_int[:batch_size]
valid_target = target_int[:batch_size]

(valid_targets_batch, valid_sources_batch, valid_targets_lengths, valid_sources_lengths) = next(get_batches(valid_target, valid_source, batch_size,
                           source_letter_to_int['<PAD>'],
                           target_letter_to_int['<PAD>']))

display_step = 50

checkpoint = "data/checkpoint_my3/trained_model.ckpt"

with tf.Session(graph=train_graph) as sess:
    sess.run(tf.global_variables_initializer())
    print()
    for epoch_i in range(1,epochs+1):
        if epoch_i%100==0:
                saver = tf.train.Saver()
                saver.save(sess, checkpoint)
                print('Model Trained and Saved')

        for batch_i,(targets_batch, sources_batch, targets_lengths, sources_lengths) in enumerate(get_batches(
            train_target,train_source,batch_size,source_letter_to_int['<PAD>'],
                           target_letter_to_int['<PAD>']
        )):
            _,loss = sess.run([train_op,cost],feed_dict={
                input_data:sources_batch,
                targets:targets_batch,
                lr:learning_rate,
                target_sequence_length:targets_lengths,
                source_sequence_length:sources_lengths
            })

            if batch_i % display_step == 0:
                # 计算validation loss
                validation_loss = sess.run(
                    [cost],
                    {input_data: valid_sources_batch,
                     targets: valid_targets_batch,
                     lr: learning_rate,
                     target_sequence_length: valid_targets_lengths,
                     source_sequence_length: valid_sources_lengths})

                print('Epoch {:>3}/{} Batch {:>4}/{} - Training Loss: {:>6.3f}  - Validation loss: {:>6.3f}'
                      .format(epoch_i,
                              epochs,
                              batch_i,
                              len(train_source) // batch_size,
                              loss,
                              validation_loss[0]))

    saver = tf.train.Saver()
    saver.save(sess, checkpoint)
    print('Model Trained and Saved')


# 预测
def source_to_seq(text):
    text_word=[]
    sequence_length = 15
    #seg_list=jieba.cut(text,cut_all=False)
    #liststr="/".join(seg_list)
    #text_word.extend(liststr.split('/'))
    text_word.extend(jieba.cut(text,cut_all=False))
    return [source_letter_to_int.get(word,source_letter_to_int['<UNK>']) for word in text_word] + [source_letter_to_int['<PAD>']] * (sequence_length - len(text_word))


for sen in range(100):
    input_word = input()
    if input_word=="end":
        break
    text = source_to_seq(input_word)



    checkpoint = "data/checkpoint_my3/trained_model.ckpt"
    loaded_graph = tf.Graph()

    with tf.Session(graph=loaded_graph) as sess:
        loader = tf.train.import_meta_graph(checkpoint+'.meta')
        loader.restore(sess,checkpoint)

        input_data = loaded_graph.get_tensor_by_name('inputs:0')
        logits = loaded_graph.get_tensor_by_name('predictions:0')
        source_sequence_length = loaded_graph.get_tensor_by_name('source_sequence_length:0')
        target_sequence_length = loaded_graph.get_tensor_by_name('target_sequence_length:0')

        answer_logits = sess.run(logits, {input_data: [text] * batch_size,
                                          target_sequence_length: [20] * batch_size,
                                          source_sequence_length: [len(input_word)] * batch_size})[0]

        pad = source_letter_to_int["<PAD>"]

        print('原始输入:', input_word)

        print('\nSource')
        print('  Word 编号:    {}'.format([i for i in text]))
        print('  Input Words: {}'.format(" ".join([source_int_to_letter[str(i)] for i in text])))

        print('\nTarget')
        print('  Word 编号:       {}'.format([i for i in answer_logits if i != pad]))
        print('  Response Words: {}'.format(" ".join([target_int_to_letter[str(i)] for i in answer_logits if i != pad])))