class Seq2SeqEncoder(d2l.Encoder):
    """用于序列到序列学习的 RNN 编码器"""
    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,dropout=0, **kwargs):
        super(Seq2SeqEncoder, self).__init__(**kwargs)
        self.embedding = nn.Embedding(vocab_size, embed_size)	# 嵌入层
        self.rnn = rnn.GRU(num_hiddens, num_layers, dropout=dropout)

    def forward(self, X, *args):
        # 输出'X'的形状：(`batch_size`, `num_steps`, `embed_size`)
        X = self.embedding(X)
        # 在RNN 模型中，第一个轴对应于时间步长
        X = X.swapaxes(0, 1)
        state = self.rnn.begin_state(batch_size=X.shape[1], ctx=X.ctx)
        output, state = self.rnn(X, state)
        # `output`的形状: (`num_steps`, `batch_size`, `num_hiddens`)
        # `state[0]`的形状: (`num_layers`, `batch_size`, `num_hiddens`)
        return output, state

# X_repeat 的形状: (n_test, n_train), 用于批量计算 $f(x)$，而不是一个个计算
# 每一行都包含着相同的测试输入，即同样的查询
X_repeat = x_test.repeat(n_train).reshape((-1, n_train))
# x_test 包含的是查询；x_train 包含的是键；y_train 包含的是值。
# attention_weights 的形状：(n_test, n_train), 
# 每一行包含的是在给定的每个键 (x_train) 对应的值 (y_train) 之间应该分配的注意力权重
# 当查询与键越接近时，注意力池化的注意力权重就越高
# 用于计算查询 (x_test) 对应的值 (y_hat) 的输出
# 注：由于输入的查询受所有的键和当前位置的值的影响，因此同样的查询输出不同的值
attention_weights = npx.softmax(-(X_repeat - x_train) ** 2 / 2)
# y_hat 的每个元素都是值的加权平均值，其中的权重是注意力权重
y_hat = np.dot(attention_weights, y_train)

class NWKernelRegression(nn.Block):
    def __init__(self, **kwargs):
        super(NWKernelRegression, self).__init__(**kwargs)
        self.attention_weights = None
        self.w = self.params.get('w', shape=(1,))

    def forward(self, queries, keys, values):
        # 'queries' 的形状: (查询的个数，“键－值”对的个数)
        queries = queries.repeat(keys.shape[1]).reshape((-1, keys.shape[1]))
        # 'attention_weights' 的形状: (查询的个数，“键－值”对的个数)
        self.attention_weights = npx.softmax(-((queries - keys) * self.w.data()) ** 2 / 2)
        # 'values' 的形状: (查询的个数，“键－值”对的个数)
        return npx.batch_dot(np.expand_dims(self.attention_weights, 1), 
                             np.expand_dims(values, -1)).reshape(-1)

net = NWKernelRegression()
net.initialize()
loss = gluon.loss.L2Loss()
trainer = gluon.Trainer(net.collect_params(), 'sgd', {'learning_rate': 0.5})
animator = d2l.Animator(xlabel='epoch', ylabel='loss', xlim=[1, 5])
for epoch in range(50):
    with autograd.record():
        l = loss(net(x_train, keys, values), y_train)
    l.backward()
    trainer.step(1)

def masked_softmax(X, valid_lens):
    """通过在最后一个轴上遮盖元素来执行 Softmax 操作"""
    # `X`: 3D tensor, `valid_lens`: 1D or 2D tensor
    if valid_lens is None:
        return npx.softmax(X)
    else:
        shape = X.shape
        if valid_lens.ndim == 1:
            valid_lens = valid_lens.repeat(shape[1])
        else:
            valid_lens = valid_lens.reshape(-1)
        # 在最后的轴上，被遮盖的元素使用一个非常大的负值替换，从而其 softmax (指数)输出为 0
        X = npx.sequence_mask(X.reshape(-1, shape[-1]), sequence_length=valid_lens, 
                              use_sequence_length=True, value=-1e6, axis=1)
        return npx.softmax(X).reshape(shape)

masked_softmax(np.random.uniform(size=(2, 2, 4)), d2l.tensor([2, 3]))    

class AdditiveAttention(nn.Block):
    """可加性注意力"""
    def __init__(self, num_hiddens, dropout, **kwargs):
        super(AdditiveAttention, self).__init__(**kwargs)
        # 使用' flatten=False '只转换最后一个轴，以便其他轴的形状保持不变
        self.W_k = nn.Dense(num_hiddens, use_bias=False, flatten=False)
        self.W_q = nn.Dense(num_hiddens, use_bias=False, flatten=False)
        self.w_v = nn.Dense(1, use_bias=False, flatten=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, queries, keys, values, valid_lens):
        queries, keys = self.W_q(queries), self.W_k(keys)
        # 在维度扩展后，
        # `queries` 的形状：(`batch_size`, 查询的个数, 1, `num_hidden`)
        # `key` 的形状：(`batch_size`, 1, “键－值”对的个数, `num_hiddens`)
        # 使用广播的方式进行求和
        features = np.expand_dims(queries, axis=2) + np.expand_dims(keys, axis=1)
        features = np.tanh(features)
        # `self.w_v` 仅有一个输出，因此从形状中移除最后那个维度。
        # `scores` 的形状：(`batch_size`, 查询的个数, “键-值”对的个数)
        scores = np.squeeze(self.w_v(features), axis=-1)
        self.attention_weights = masked_softmax(scores, valid_lens)
        # `values` 的形状：(`batch_size`, “键－值”对的个数, 值的维度)
        return npx.batch_dot(self.dropout(self.attention_weights), values)

attention = AdditiveAttention(num_hiddens=8, dropout=0.1)
attention.initialize()
attention(queries, keys, values, valid_lens)    

class DotProductAttention(nn.Block):
    """Scaled dot product attention."""
    def __init__(self, dropout, **kwargs):
        super(DotProductAttention, self).__init__(**kwargs)
        self.dropout = nn.Dropout(dropout)

    # `queries` 的形状：(`batch_size`, 查询的个数, `d`)
    # `keys` 的形状：(`batch_size`, “键－值”对的个数, `d`)
    # `values` 的形状：(`batch_size`, “键－值”对的个数, 值的维度)
    # `valid_lens` 的形状: (`batch_size`,) 或者 (`batch_size`, 查询的个数)
    def forward(self, queries, keys, values, valid_lens=None):
        d = queries.shape[-1]
        # 设置 `transpose_b=True` 为了交换 `keys` 的最后两个维度
        scores = npx.batch_dot(queries, keys, transpose_b=True) / math.sqrt(d)
        self.attention_weights = masked_softmax(scores, valid_lens)
        return npx.batch_dot(self.dropout(self.attention_weights), values)
    
attention = DotProductAttention(dropout=0.5)
attention.initialize()
attention(queries, keys, values, valid_lens)