# 
BUFFER_SIZE = len(input_tensor_train)  #就是3万
# 批次设置为64
BATCH_SIZE = 64
steps_per_epoch = len(input_tensor_train)//BATCH_SIZE
# 
# 设置embedding维度
embedding_dim = 256
units = 1024
# 输入的
# 为什么要加1？ 因为word_index是从1开始的，而embedding是从0开始的，所以需要加1.
# 词典大小，输入的词典大小为：
vocab_inp_size = len(inp_lang.word_index)+1
print(vocab_inp_size)
#输出
# 输出的词典大小是多少
vocab_tar_size = len(targ_lang.word_index)+1
print(vocab_tar_size)
# 训练集
# 训练集的输入值和目标值
# 创建一个dataset是给定张量的维度。
# 所有输入张量的第一维必须有相同的大小。
# shuffle（）: 将所有数据重新排序。
dataset = tf.data.Dataset.from_tensor_slices((input_tensor_train, target_tensor_train)).shuffle(BUFFER_SIZE)
dataset = dataset.batch(BATCH_SIZE, drop_remainder=True)
# 一批次是64个样本数据。
# iter() 函数用来生成迭代器。
# next() 返回迭代器的下一个项目
example_input_batch, example_target_batch = next(iter(dataset))
example_input_batch,example_target_batch

next(iter(tf.data.Dataset.from_tensor_slices((input_tensor, target_tensor)).shuffle(BUFFER_SIZE)))

# 手法和之前的类似
# 继承自tf.keras.Model
class Encoder(tf.keras.Model):
    # vocab_size：词典大小是多少
    # embedding_dim: embedding的维度
    # encoding_units：编码的单元个数
    # batch_size：一批样本的大小
    def __init__(self, vocab_size, embedding_dim, encoding_units, batch_size):
        super(Encoder, self).__init__()
        self.batch_size = batch_size
        # 编码单元
        self.encoding_units = encoding_units
        # 创建Embedding层，
        # 词典大小，embedding的维度。
        self.embedding = keras.layers.Embedding(vocab_size, embedding_dim)
        #定义GRU层，gru是lstm变种，gru把遗忘门和输入门变为一个，因为遗忘门+输入门=1
        self.gru = keras.layers.GRU(self.encoding_units,
                                    # 默认为False，即只返回最后一个单元的output
                                    # 如果设置为True，表示返回所有单元的output
                                    return_sequences=True,
                                    # 是否返回最后的状态以添加到输出里边
                                    # 默认为False，即不返回最后一个单元的hidden_state
                                    # 如果设置为True，表示返回最后一个单元的hidden_state
                                    return_state=True,
                                    recurrent_initializer='glorot_uniform')
    # 初始化
    def call(self, x, hidden):
        # embedding的输出x
        x = self.embedding(x)
        # 初始状态 hidden
        output, state = self.gru(x, initial_state = hidden)
        return output, state
    # 隐藏状态
    def initialize_hidden_state(self):
        return tf.zeros((self.batch_size, self.encoding_units))
# embedding：特征维度 
# units：神经元个数
# embedding_dim：embedding维度，256.
# vocab_inp_size：输入单词的一个词典大小
encoder = Encoder(vocab_inp_size, embedding_dim, units, BATCH_SIZE)
# 获得初始化的hidden
# 初始化隐状态
# 二维列表： 默认的话是（64*1024） 即批次*编码单元个数
sample_hidden = encoder.initialize_hidden_state()
# 获得输出和隐含状态
# 对象运行调用了call方法，因为继承了tf.keras.model
# example_input_batch：转换为向量的训练集的一个批次48
sample_output, sample_hidden = encoder(example_input_batch, sample_hidden)
# 维度变化：
# 首先，句子经过tokenizer处理，句子的维度变为（,16）
# 
# 输出的16是长度，1024是状态的size
print('Encoder output shape: (batch size, sequence length, units) {}'.format(sample_output.shape))
print('Encoder Hidden state shape: (batch size, units) {}'.format(sample_hidden.shape))

#实现Attention机制
# 
class BahdanauAttention(tf.keras.Model):
    def __init__(self, units):
        super(BahdanauAttention, self).__init__()
        # 做三个全连接层
        # 做全连接
        # 做几个全连接层。
        # 激活函数 默认是没有的。
        self.W1 = tf.keras.layers.Dense(units)
        self.W2 = tf.keras.layers.Dense(units)
        self.V = tf.keras.layers.Dense(1)
    # query传的是decoder_ hidden
    # 实现公式
    # call传入的是sample_hidden
    def call(self, query, values):
        # deocoder_ hidden. shape: (batch_ size,units )
        # encoder outputs.shape: ( batch size, length,units )
        # 做维度扩展，扩展前后对比是下面两行
        # hidden shape == (batch_size, hidden size)
        # hidden_with_time_axis shape == (batch_size, 1, hidden size)
        # we are doing this to perform addition to calculate the score
        # 维度扩展
        # 运算需要，必须维度一致。
        hidden_with_time_axis = tf.expand_dims(query, 1)
        #接下来要实现Attention,Bahdanau方式的
        # beforeV:(batch_ size,length,units )
        # after V( batch_ size,length,1)
        # score shape == (batch_size, max_length, 1)
        # we get 1 at the last axis because we are applying score to self.V
        # the shape of the tensor before applying self.V is (batch_size, max_length, units)
        score = self.V(tf.nn.tanh(self.W1(values) + self.W2(hidden_with_time_axis)))
        # attention_weights shape == (batch_size, max_length, 1)
        attention_weights = tf.nn.softmax(score, axis=1)
        #先算加权，values就是encoder_outputs
        # context_vector.shape: ( batch size, length units )
        context_vector = attention_weights * values
        #再算平均，在length的维度去求和
        # context_vector shape after sum == (batch_size, hidden_size)
        context_vector = tf.reduce_sum(context_vector, axis=1)
        return context_vector, attention_weights
#     
attention_layer = BahdanauAttention(10)
attention_result, attention_weights = attention_layer(sample_hidden, sample_output)
print("Attention result shape: (batch size, units) {}".format(attention_result.shape))
print("Attention weights shape: (batch_size, sequence_length, 1) {}".format(attention_weights.shape))

Decoder output shape: (batch_size, vocab size) (64, 4935)
decoder_hidden.shape:  (64, 1024)
decoder_attention_weights.shape: (64, 16, 1)

# 我们用的优化器是adam
optimizer = keras.optimizers.Adam()
# 分类问题我们往往用SparseCategoricalCrossentropy，因为我们的fc是纯的输出，没有加softmax，
# 因此这里的from_logits为True，否则改为false，reduction是损失函数如何做聚合
# 交叉损失函数
# from_logits: 为True时，会将y_pred转化为概率（用softmax）,否则不进行转换，通常情况下用True结果更稳定；
# reduction：类型为tf.keras.losses.Reduction，对loss进行处理，默认是求平均；
loss_object = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction='none')
def loss_function(real, pred):
    # 是零的时候返回结果是True，因此要取反操作
    # tf.math.equal(real, 0)是padding的部分都是1，不是padding的部分都是零，因此我们要取反
    # 取反
    # tf.math.logical_not: 为Tensorflow中的逻辑NOT功能提供支持
    # 返回的mask是一个张量，数据类型为bool类型。
    mask = tf.math.logical_not(tf.math.equal(real, 0))
    loss_ = loss_object(real, pred)
    #将张量转换为新类型
    # 将张量转为float类型
    mask = tf.cast(mask, dtype=loss_.dtype)
    # padding部分的mask是零
    # 
    loss_ *= mask
    # 计算累计的平均损失
    # 用于计算张量tensor沿着指定的数轴（tensor的某一维度）上的的平均值，主要用作降维或者计算tensor（图像）的平均值。
    return tf.reduce_mean(loss_)
# 保存模型
checkpoint_dir = './8-1_checkpoints'
if not os.path.exists(checkpoint_dir):
    os.mkdir(checkpoint_dir)
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
checkpoint = tf.train.Checkpoint(optimizer=optimizer,
                                 encoder=encoder,
                                 decoder=decoder)

# 为了提高调用速度，把它变成图。
@tf.function
# 
def train_step(inp, targ, encoding_hidden):
    loss = 0
    # 
    with tf.GradientTape() as tape:
        #把输入给encoder，得到encoding_output, encoding_hidden
        encoding_output, encoding_hidden = encoder(inp, encoding_hidden)
        decoding_hidden = encoding_hidden
        decoding_input = tf.expand_dims([targ_lang.word_index['<start>']] * BATCH_SIZE, 1)
        #eg: <start> I am here <end>
        #1.<start>->I
        #2.I->am
        #3.am->here
        # 4. here ->< end>
        #对于here，我们相当于要把I am  的信息都要给过去
        # Teacher forcing - feeding the target as the next input
        for t in range(1, targ.shape[1]):
            # passing enc_output to the decoder
            #根据我们前面的原理解析，我们这里需要给3项信息
            predictions, decoding_hidden, _ = decoder(decoding_input, decoding_hidden, encoding_output)
            loss += loss_function(targ[:, t], predictions)
            # using teacher forcing
            decoding_input = tf.expand_dims(targ[:, t], 1)
    #这里是每个batch上平均的损失函数
    batch_loss = (loss / int(targ.shape[1]))
    variables = encoder.trainable_variables + decoder.trainable_variables
    #求梯度
    gradients = tape.gradient(loss, variables)
    #有了梯度以后，可以用optimizer去做apply
    optimizer.apply_gradients(zip(gradients, variables))
    return batch_loss

EPOCHS = 10
#这里运行时间比较久
for epoch in range(EPOCHS):
    start = time.time()
    # 第一次，全0的隐含状态
    encoding_hidden = encoder.initialize_hidden_state()
    total_loss = 0
    # 取多少次数据
    # 每次去取dataset.take(steps_per_epoch)这么多数据
    # 每训练100次打印一下损失
    for (batch, (inp, targ)) in enumerate(dataset.take(steps_per_epoch)):
        batch_loss = train_step(inp, targ, encoding_hidden)
        total_loss += batch_loss
        #这里增加打印
        if batch % 100 == 0:
            print('Epoch {} Batch {} Loss {:.4f}'.format(epoch + 1, batch, batch_loss.numpy()))
    # saving (checkpoint) the model every 2 epochs，保存模型
    # 满足条件才保存模型
    if (epoch + 1) % 2 == 0:
        checkpoint.save(file_prefix = checkpoint_prefix)
    print('Epoch {} Loss {:.4f}'.format(epoch + 1, total_loss / steps_per_epoch))
    print('Time taken for 1 epoch {} sec\n'.format(time.time() - start))