import collections

import numpy as np

import tensorflow as tf

TRAIN_DATA = "./data/ptb.train.txt"  # 训练数据路径

TEST_DATA = "./data/ptb.test.txt"  # 测试数据路径

EVAL_DATA = "./data/ptb.valid.txt"  # 验证数据路径

HIDDEN_SIZE = 300  # 隐藏层中cell的个数

NUM_LAYERS = 2  # 深度循环神经网络中LSTM结构的层数

VOCAB_SIZE = 10000  # 词典规模

TRAIN_BATCH_SIZE = 20  # 训练数据batch的大小

TRAIN_NUM_STEP = 35  # 训练数据的截断长度，也可以看作是序列的长度

EVAL_BATCH_SIZE = 1

EVAL_NUM_STEP = 35

NUM_EPOCH = 30  # 使用训练数据的轮数

LSTM_KEEP_PROB = 0.9  # LSTM节点不被dropout的概率

EMBEDDING_KEEP_PROB = 0.9  # 词向量不被dropout的概率

MAX_GRAD_NORM = 5  # 用于控制梯度膨胀的梯度大小上限

SHARE_EMB_AND_SOFTMAX = True  # 在softmax层和词向量层之间共享参数

def read_data(file_path):

    """

    读取文件数据，将文本中的词转换成词空间中对应的索引

    :param file_path: 文件路径

    :return: 由数值取代后的文本词列表

    """

    # 采用TensorFlow中的读取文件的方法去读取文件

    with tf.gfile.GFile(file_path, "r") as f:

        # 将文本读取出来，并且进行分词，将换行符替换成<eos>,eos的意思就是end of sentence

        word_list = f.read().replace("\n", "<eos>").split()

        # 对分词后的列表进行统计，统计每个单词出现的数量, 返回的数据类型Counter({'jiang': 2, 'zhang': 1})

        counter = collections.Counter(word_list)

        # 对每个词按照词频排序，对于词频相同的按词本身进行排序，返回的数据类型[('jiang', 2), ('zhang', 1)]

        count_pairs = sorted(counter.items(), key=lambda x: (-x[1], x[0]))

        # 取出单词元组，返回的数据类型words=('jiang', 'zhang')

        words, _ = list(zip(*count_pairs))

        # 将上面排序后的词（无重复的词空间）标记索引数值，返回的数据类型{'jiang': 0, 'zhang': 1}

        word_to_id = dict(zip(words, range(len(words))))

        # 将文本中所有的单词用索引数值取代，组成新的列表

        id_list = [word_to_id[word] for word in word_list if word in word_to_id]

    return id_list

def make_batches(id_list, batch_size, num_steps):

    """

    将原始的数据集转换成mini-batch进行训练

    :param id_list: 原始的数值文本列表

    :param batch_size:

    :param num_steps: 一个样本的序列长度

    :return: 整个样本转换后的batchs数据

    """

    # 计算总的batch数量。每个batch包含的单词数量是batch_size * num_steps，batch_size为一个batch中样本的数量，

    # num_steps为一个样本的序列长度，因此num_batchs表示整个训练文本能分成的batch的数量

    num_batches = (len(id_list) - 1) // (batch_size * num_steps)

    # 根据上面分配好的num_batchs, batch_size, num_steps参数来构建数据集，先取出能整除的序列长度

    data = np.array(id_list[: num_batches * batch_size * num_steps])

    # reshape取出来的序列（一维数组）成二维数组，因为是序列数据，所以行为batch_size，列为num_batchs * num_steps，之后训练在横向上分割

    data = np.reshape(data, [batch_size, num_batches * num_steps])

    # 将数据在axis=1的轴上分割，分割的数量就是num_batchs

    data_batches = np.split(data, num_batches, axis=1)

    # 因为是根据前面的词预测后面的词，因此输出的值要往后移一位，其余操作和上面的一致

    label = np.array(id_list[1: num_batches * batch_size * num_steps + 1])

    label = np.reshape(label, [batch_size, num_batches * num_steps])

    label_batches = np.split(label, num_batches, axis=1)

    return list(zip(data_batches, label_batches))

class PTBModel(object):

    def __init__(self, is_training, batch_size, num_steps):

        # 记录使用的batch大小和截断长度（也就是一个样本的序列长度）

        self.batch_size = batch_size

        self.num_steps = num_steps

        # 定义每一步的输入和预期输出。两者的维度都是[batch_size, num_steps]，batch_size时间上就是指一个batch中样本的数量

        self.input_data = tf.placeholder(tf.int32, [batch_size, num_steps])

        self.targets = tf.placeholder(tf.int32, [batch_size, num_steps])

        # 定义dropout的值，训练时取0.9，否则取1.0，表示不做dropout

        dropout_keep_prob = LSTM_KEEP_PROB if is_training else 1.0

        # 定义lstm cell的结构，dropout相当于装饰器直接包裹在lstm_cell上,此时的cell是垂直方向的，所以for循环中的值是NUM_LAYERS

        lstm_cells = [tf.nn.rnn_cell.DropoutWrapper(tf.nn.rnn_cell.BasicLSTMCell(HIDDEN_SIZE),

                                                    output_keep_prob=dropout_keep_prob)

                      for i in range(NUM_LAYERS)]

        # 组合成多层循环

        cell = tf.nn.rnn_cell.MultiRNNCell(lstm_cells)

        # 初始化最初的状态值，即全为0的向量。这个量只在每个epoch初始化第一个batch时使用

        # #返回[batch_size, 2*len(cells)],或者[batch_size, s]，至于为什么是2 * ，这是因为初始值有两个h0和C0

        self.initial_state = cell.zero_state(batch_size, tf.float32)

        # 定义单词的词向量矩阵，这里用get_variable，之后在测试时就可以实现参数共享了

        # VOCAB_SIZE是指词向量空间中词的个数（在这里起始是len(words)的长度，也等于10000），HIDDEN_SIZE是值词嵌入之后的词向量长度

        embedding = tf.get_variable("embedding", [VOCAB_SIZE, HIDDEN_SIZE])

        # 将输入的单词转化为词向量，相当于将每个序列中的单词按照索引（之前转化为了数值在这里就很方便的），直接将序列中的每个词在已经训练好的

        # 词空间中寻找对应的向量，此空间应该是二维的，输出的结果应该是三维的，也就是batch_size * num_steps * HIDDEN_SIZE

        inputs = tf.nn.embedding_lookup(embedding, self.input_data)

        # 只在训练时使用dropout来训练词向量

        if is_training:

            inputs = tf.nn.dropout(inputs, EMBEDDING_KEEP_PROB)

        # 定义输出列表。在这里先将不同时刻LSTM结构的输出收集起来，再一起提供给softmax层，在这里是实现时间序上的cell输出

        outputs = []

        state = self.initial_state

        with tf.variable_scope("RNN"):

            for time_step in range(num_steps):

                if time_step > 0:

                    # 实现在同一个variable_scope下共享参数

                    tf.get_variable_scope().reuse_variables()

                # 从这里就可以看出之前embedding输出的是三维，我们根据时间序取出词来进行训练,

                # cell_output shape=(batch_size, HIDDEN_SIZE)

                cell_output, state = cell(inputs[:, time_step, :], state)

                outputs.append(cell_output)

        # 把输出队列展开成[batch_size, num_steps * hidden_size]，然后再reshape成[batch_size * num_steps, hidden_size]

        output = tf.reshape(tf.concat(outputs, 1), [-1, HIDDEN_SIZE])

        # softmax层：将RNN在每个位置的输出转化为各个单词的logits

        # weight的shape是[HIDDEN_SIZE, VOCAB_SIZE]

        if SHARE_EMB_AND_SOFTMAX:

            weight = tf.transpose(embedding)

        else:

            weight = tf.get_variable("weight", [HIDDEN_SIZE, VOCAB_SIZE])

        bias = tf.get_variable("bias", [VOCAB_SIZE])

        # 算出最终输出的logits，用于之后的交叉熵和softmax计算

        logits = tf.matmul(output, weight) + bias

        # 定义交叉熵损失函数和平均损失函数，返回的loss是和labels、logits相同的shape

        loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=tf.reshape(self.targets, [-1]),

                                                              logits=logits)

        # 求平均损失

        self.cost = tf.reduce_sum(loss) / batch_size

        # 该状态用来存储训练完一个batch之后的状态，在训练下一个batch时，会将该状态作为初始状态，这个在run_epoch函数中控制的

        self.final_state = state

        # 只在训练模型时定义反向传播操作

        if not is_training:

            return

        # 拿到所有的训练参数，用于之后的梯度下降更新参数

        trainable_variables = tf.trainable_variables()

        # 梯度截断控制梯度大小，是为了避免梯度弥散和梯度爆炸，将梯度控制在某一范围内

        grads, _ = tf.clip_by_global_norm(

            tf.gradients(self.cost, trainable_variables), MAX_GRAD_NORM

        )

        optimizer = tf.train.GradientDescentOptimizer(learning_rate=1.0)

        self.train_op = optimizer.apply_gradients(zip(grads, trainable_variables))

def run_epoch(session, model, batches, train_op, output_log, step):

    """

    训练模型和模型预测。使用给定的模型model在数据data上运行train_op并返回在全部数据上的perplexity值。

    :param session:

    :param model:

    :param batches:

    :param train_op:

    :param output_log: 判断是训练过程还是其他过程

    :param step:

    :return:

    """

    # 计算平均perplexity的辅助变量，perplexity表示在模型生成一句话时下一个词有perplexity个合理的选择，认为perplexity小于100都是比较好的

    # 结果，该值越小，说明模型越好，也可以认为模型预测的精确度越高

    total_costs = 0.0  # 存储损失值

    iters = 0

    state = session.run(model.initial_state)

    # 训练一个epoch

    for x, y in batches:

        # 可以在run函数中通过字典的形式给模型输入数据，也可以直接读取出模型的值，session.run（）方法通过驱动三个operation来驱动整个图

        # 其中train_op是只是用来驱动训练过程的

        cost, state, _ = session.run([model.cost, model.final_state, model.train_op],

                                     feed_dict={model.input_data: x, model.targets: y, model.initial_state: state}

                                     )

        total_costs += cost

        iters += model.num_steps

        if output_log and step % 100 == 0:

            print("After {} steps, perplexity id {}".format(step, np.exp(total_costs / iters)))

        step += 1

    # pplx是语言模型perplexity指标

    pplx = np.exp(total_costs / iters)

    return step, pplx

def main():

    with tf.Graph().as_default():

        # 定义初始化函数, 用于决定之后的variable_scope中的变量的初始值取值范围

        initializer = tf.random_uniform_initializer(-0.05, 0.05)

        # 定义训练用的循环神经网络模型

        with tf.name_scope("train"):

            train_batches = make_batches(read_data(TRAIN_DATA), TRAIN_BATCH_SIZE, TRAIN_NUM_STEP)
　　　　　　　# 利用variable_scope()和get_variable()来实现变量共享

            with tf.variable_scope("language_model", reuse=None, initializer=initializer):

                train_model = PTBModel(True, TRAIN_BATCH_SIZE, TRAIN_NUM_STEP)

        # 定义测试用的循环神经网络模型，它与train_model共享参数，但没有dropout，可以通过is_training来控制

        with tf.name_scope('test'):

            eval_batches = make_batches(read_data(EVAL_DATA), EVAL_BATCH_SIZE, EVAL_NUM_STEP)

            test_batches = make_batches(read_data(TEST_DATA), EVAL_BATCH_SIZE, EVAL_NUM_STEP)

            # 设置同样的名称language_model来实现共享变量，变量共享和name_scope无关

            with tf.variable_scope("language_model", reuse=True, initializer=initializer):

                eval_model = PTBModel(False, EVAL_BATCH_SIZE, EVAL_NUM_STEP)

        # 训练模型

        with tf.Session() as session:

            tf.global_variables_initializer().run()

            step = 0

            for i in range(NUM_EPOCH):

                print("In iteration: {}".format(i + 1))

                step, train_pplx = run_epoch(session, train_model, train_batches, train_model.train_op, True, step)

                print("Epoch: {} Train Perplexity: {}".format(i+1, train_pplx))

                _, eval_pplx = run_epoch(session, eval_model, eval_batches, tf.no_op(), False, 0)

                print("Epoch: {} Eval Perplexity: {}".format(i+1, eval_pplx))

                _, test_pplx = run_epoch(session, eval_model, test_batches, tf.no_op(), False, 0)

                print("Epoch: {} Test Perplexity: {}".format(i+1, test_pplx))