import tensorflow as tf
tf.enable_eager_execution()

import numpy as np
import os
import time

path_to_file = tf.keras.utils.get_file('shakespeare.txt', 'https://storage.googleapis.com/download.tensorflow.org/data/shakespeare.txt')
print(path_to_file)

text = open(path_to_file, 'rb').read().decode(encoding='utf-8')
# length of text is the number of characters in it
print ('Length of text: {} characters'.format(len(text)))

# Take a look at the first 1000 characters in text
print(text[:250])

# The unique characters in the file
vocab = sorted(set(text)) # sorted保证了集合的顺序
print ('{} unique characters'.format(len(vocab)))

# Creating a mapping from unique characters to indices
char2idx = {u:i for i, u in enumerate(vocab)}
# vocab是有序集合，转化为数组后其下标自然就是序号，但是不如char2idx结构直观
# 如果模仿char2idx也很简单：idx2char = {i:u for i,u in enumerate(vocab)}
idx2char = np.array(vocab)

text_as_int = np.array([char2idx[c] for c in text])
text_as_int[:10]

# 取出char2idx前20个元素的奇怪写法。zip方法返回成对的元组，range(20)提供了序号。
for char,_ in zip(char2idx, range(20)):
    print('{:6s} ---> {:4d}'.format(repr(char), char2idx[char]))

# Show how the first 13 characters from the text are mapped to integers
print ('{} ---- characters mapped to int ---- > {}'.format(text[:13], text_as_int[:13]))

# The maximum length sentence we want for a single input in characters
# 每次喂入RNN的字符数。注意和后面的BATCH_SIZE的区别以及匹配
# 为了更好的观察数据，初始的时候seq_length可以设置为10，但是执行时要恢复为100或者
# 更大的数。当然，也可以测试不同的seq_length下的结果
seq_length = 100
examples_per_epoch = len(text)//seq_length

char_dataset = tf.data.Dataset.from_tensor_slices(text_as_int)
print(char_dataset)
for i in char_dataset.take(5):
    print(idx2char[i.numpy()])

# sequences也是一个Dataset对象，但是经过了batch操作进行数据分组，每一个batch的数据
# 长度是seq_length+1（101）.sequences用来创建输入文本和目标文本（长度为seq_length）
# 注意：这里的batch操作和训练模型时的BATCH_SIZE没有关系，这里的batch操作纯粹
# 为了按照指定的尺寸切分数据
sequences = char_dataset.batch(seq_length+1, drop_remainder=True)
# repl函数的意义相当于Java的toString方法
# 注意，这里的item已经是tensor了，通过numpy()方法转化为numpy矩阵（向量）
# numpy数组（List）的强大之处：允许接受一个list作为索引参数，因此idx2char[item.numpy()]即为根据item
# 的数字为索引获得字符构造出一个字符串
for item in sequences.take(5):
    print(repr(''.join(idx2char[item.numpy()])))

def split_input_target(chunk):
    input_text = chunk[:-1] # 不包括-1即最后一个字符，总共100个字符。这就是为什么chunk的长度是101的原因
    target_text = chunk[1:]
    return input_text, target_text
 
# 注意到，sequences已经是被batch过的了，因此这里的map是针对每个batch的数据来进行的
# 此时dataset的结果已经比较复杂了，所谓的nested structure of tensors
# print（dateset）的结果显示其shape为：shapes: ((10,), (10,))
# 即，dataset是一个tuple，tuple的每个数据又包含两个tuple，每个tuple是seq_length
# 长度的向量。其中第一个tuple是input_example，第二个tuple是target_example
dataset = sequences.map(split_input_target)
print(dataset)

# 将take的参数设为2能看的更清楚
for input_example, target_example in  dataset.take(1):
  print ('Input data: ', repr(''.join(idx2char[input_example.numpy()])))
  print ('Target data:', repr(''.join(idx2char[target_example.numpy()])))

for i, (input_idx, target_idx) in enumerate(zip(input_example[:5], target_example[:5])):
    print("Step {:4d}".format(i))
    print("  input: {} ({:s})".format(input_idx, repr(idx2char[input_idx])))
    print("  expected output: {} ({:s})".format(target_idx, repr(idx2char[target_idx])))

# Batch size
BATCH_SIZE = 32

# steps_per_epoch说明每次迭代整个训练数据集需要的**步数**。
# 使用model.fit时，如果传入的数据集是Dataset对象，必须显式声明steps_per_epoch参数
# 道理很简单，否则tensorflow不知道以多大的步长（或者说，总共需要多少步）循环迭代给定的Dataset。
# 因为传入fit函数的Dataset只是经过了seq_length分组的input_text和target_text，并没有指定训练时
# 使用多大的步长来迭代整个Dataset。
# steps_per_epoch = len(text)//seq_length//BATCH_SIZE
steps_per_epoch = examples_per_epoch//BATCH_SIZE

# Buffer size to shuffle the dataset
# (TF data is designed to work with possibly infinite sequences,
# so it doesn't attempt to shuffle the entire sequence in memory. Instead,
# it maintains a buffer in which it shuffles elements).
# 只在BUFFER_SIZE内进行shuffle操作，即每次最多取BUFFER_SIZE个样本进行shuffle操作。
BUFFER_SIZE = 10000

# 此时的dataset是已经batch化了的，即shape为(batch_size, seqence_length)
dataset = dataset.shuffle(BUFFER_SIZE).batch(BATCH_SIZE, drop_remainder=True)
dataset

# Length of the vocabulary in chars
# 这是输入层和输出层的维度。
# 每一个字符都需要进行one-hot编码，因此每一个输入都是vocab_size维度的向量
# 同样的，每一个预测的输出也是vocab_size维度的向量
vocab_size = len(vocab)

# The embedding dimension 
#embedding_dim = 256
embedding_dim = 256

# Number of RNN units
#rnn_units = 1024
rnn_units = 512

if tf.test.is_gpu_available():
  #rnn = tf.keras.layers.CuDNNGRU
  # CuDNNLSTM的执行速度大概是普通的LSTM的5倍以上
  lstm = tf.keras.layers.CuDNNLSTM
else:
  import functools
  #rnn = functools.partial(tf.keras.layers.GRU, recurrent_activation='sigmoid')
  lstm = functools.partial(tf.keras.layers.LSTM, recurrent_activation='sigmoid')

def build_model(vocab_size, embedding_dim, rnn_units, batch_size):
  model = tf.keras.Sequential([
    tf.keras.layers.Embedding(vocab_size, embedding_dim, 
                              batch_input_shape=[batch_size, None]),
    # 替换rnn为LSTM
    lstm(rnn_units,
        return_sequences=True, 
        recurrent_initializer='glorot_uniform',
        stateful=True),
    tf.keras.layers.Dense(vocab_size) # 这里不需要激活函数？softmax？
  ])
  return model

model = build_model(
  vocab_size = len(vocab), 
  embedding_dim=embedding_dim, 
  rnn_units=rnn_units, 
  batch_size=BATCH_SIZE)
model.summary()

# input_example_batch是dataset的一个batch，这里是32个seq_length的input_text
# 由于喂入的数据的shape是(32,seq_length)，输出example_batch_prediction的
# shape自然就是(32,seq_length,65)
for input_example_batch, target_example_batch in dataset.take(1): 
  example_batch_predictions = model(input_example_batch)
  print(example_batch_predictions.shape, "# (batch_size, sequence_length, vocab_size)")

# 检查第0批数据。
# example_batch_predictions[0]的shape为(100,65)
# sampled_indices的shape为(100,1)
sampled_indices = tf.random.categorical(example_batch_predictions[0], num_samples=1)
sampled_indices = tf.squeeze(sampled_indices,axis=-1).numpy() # 去掉维度为1的列
sampled_indices

print("Input: \n", repr("".join(idx2char[input_example_batch[0].numpy()])))
print("Next Char Predictions: \n", repr("".join(idx2char[sampled_indices])))

def loss(labels, logits):
  return tf.keras.losses.sparse_categorical_crossentropy(labels, logits, from_logits=True)

example_batch_loss  = loss(target_example_batch, example_batch_predictions)
print("Prediction shape: ", example_batch_predictions.shape, " # (batch_size, sequence_length, vocab_size)") 
print("scalar_loss:      ", example_batch_loss.numpy().mean())

# Directory where the checkpoints will be saved
checkpoint_dir = './training_checkpoints'
# Name of the checkpoint files
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt_{epoch}")

checkpoint_callback=tf.keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_prefix,
    save_weights_only=True)

ckpt = tf.train.latest_checkpoint(checkpoint_dir)
if ckpt != None:
  print("load model from checkpoint")
  model = build_model(vocab_size, embedding_dim, rnn_units, batch_size=BATCH_SIZE)
  model.load_weights(ckpt)
  model.build(tf.TensorShape([1, None])) # 参数为batch_input_shape
  model.summary()

model.compile(
    optimizer = tf.train.AdamOptimizer(),
    loss = loss)

EPOCHS = 3

history = model.fit(dataset.repeat(), epochs=EPOCHS, 
                    steps_per_epoch=steps_per_epoch, 
                    callbacks=[checkpoint_callback])

import matplotlib.pyplot as plt
history_dict = history.history
history_dict.keys()
loss=history_dict['loss']
epochs = range(1, len(loss) + 1)
# "bo" is for "blue dot"
plt.plot(epochs, loss, 'bo', label='Training loss')
plt.title('Training loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()

plt.show()

ckpt = tf.train.latest_checkpoint(checkpoint_dir)
model = build_model(vocab_size, embedding_dim, rnn_units, batch_size=1)
model.load_weights(ckpt)
# 在这里调用build多此一举了，因为在build_model中已经传入了batch_size=1这个参数用于构建Input_shape
#model.build(tf.TensorShape([1, None])) 
model.summary()

def generate_text(model, start_string):
  # Evaluation step (generating text using the learned model)

  # Number of characters to generate
  num_generate = 1000

  # Converting our start string to numbers (vectorizing) 
  input_eval = [char2idx[s] for s in start_string]
  # 构造维度合适的输入数据：[batch_size,seq_length]
  # 这里batch_size=1，因此只需要将start_string扩展一维即可
  input_eval = tf.expand_dims(input_eval, 0)

  # Empty string to store our results
  text_generated = []

  # Low temperatures results in more predictable text.
  # Higher temperatures results in more surprising text.
  # Experiment to find the best setting.
  temperature = 1.

  # Here batch size == 1
  model.reset_states()
  for i in range(num_generate):
      # why not call model.predict()?
      # predictions的shape：(1,None,65)
      predictions = model(input_eval)
      #print("predictions.shape:",predictions.shape)
      
      # remove the batch dimension
      # 此时的predictions的shape：(None,65)
      predictions = tf.squeeze(predictions, 0)
      #print(predictions.shape)

      # using a multinomial distribution to predict the word returned by the model
      predictions = predictions / temperature
      #print(predictions)
      predicted_id = tf.random.categorical(predictions, num_samples=1)[-1,0].numpy()
      #print("predicted_id:",predicted_id)
      
      # We pass the predicted word as the next input to the model
      # along with the previous hidden state
      input_eval = tf.expand_dims([predicted_id], 0)
      
      text_generated.append(idx2char[predicted_id])

  return (start_string + ''.join(text_generated))

print(generate_text(model, start_string=u"First:"))