import torch
import torch.nn as nn

class MuiltHeadAttention_v1(nn.Module):
    """ 
        输出多头下的上下文向量
    """
    def __init__(self,d_in,d_out,num_tokens,dropout_p,num_head,qkv_bias=False):
        '''
            d_in:定义输入向量的特征维度和
            d_out:经过嵌入后的嵌入向量维度
            num_tokens:序列长度
            dropout_p:随机丢弃率
            num_head:头数量
        '''
        super().__init__()
        # 1.添加断言,确保 输入向量的维度 可以整除 注意力头 的数量
        # 按照多个头拆分,这样每个头专注于一个类别或者某方面的细节
        assert d_out % num_head == 0
        self.num_head = num_head
        self.dim_head = d_out // self.num_head

        self.d_out = d_out
        self.W_query = nn.Linear(d_in, d_out,bias=qkv_bias)
        self.W_key   = nn.Linear(d_in, d_out,bias=qkv_bias)
        self.W_value = nn.Linear(d_in, d_out,bias=qkv_bias)
        self.register_buffer(
            'mask',
            torch.triu(
                torch.ones(num_tokens,num_tokens),
                diagonal=1
            )
        )
        self.dropout = nn.Dropout(dropout_p)
        # 2.线性层组合输出  为什么这个线性层可以组合输出呢???
        self.liner_lay = nn.Linear(d_out,d_out)

    def forward(self,x):
        batch_size,num_tokens,_ = x.shape
        # print(f"x:{x.shape}")
        # x:torch.Size([2, 6, 3])

        # 通过全连接层让多个头看到所有输入
        Q,K,V = self.W_query(x),self.W_key(x),self.W_value(x)
        # print(f"QKV: {Q.shape}、{K.shape}、{V.shape}")
        # QKV: torch.Size([2, 6, 2])、torch.Size([2, 6, 2])、torch.Size([2, 6, 2])

        # 3.使用view方法,将(batch_size,num_tokens,d_in) 展开成 (batch_size, num_tokens, num_heads, dim_head)
        # 只需执行一次矩阵乘法即可计算键（keys），对于查询（queries）和值（values）也是如此。
        Q,K,V = \
            Q.view(batch_size,num_tokens,self.num_head,self.dim_head),\
            K.view(batch_size,num_tokens,self.num_head,self.dim_head),\
            V.view(batch_size,num_tokens,self.num_head,self.dim_head),
        # print(f"QKV: {Q.shape}、{K.shape}、{V.shape}") 
        # QKV: torch.Size([2, 6, 2, 1])、torch.Size([2, 6, 2, 1])、torch.Size([2, 6, 2, 1])
        
        # 4.将(batch_size, num_tokens, num_heads, dim_head)
        # 转换成
        #   (batch_size, num_heads, num_tokens, dim_head)
        """
        为什么这种转置对于在不同注意力头之间正确对齐查询（queries）、键（keys）和值（values），并高效执行批量矩阵乘法至关重要。
        为了得到以head为单位的注意力分数,可以直接计算每个 head 的注意力得分

        原来的 Q ∈ (batch, seq_len, heads * dim),可以理解为“每个 token 拥有 heads 份不同的子表示”。
        """
        Q,K,V = Q.transpose(1,2),K.transpose(1,2),V.transpose(1,2)

        # 5.将第2维和第3维进行转换  并获取和应用掩码  因为多头下批量计算是需要转置最后两个维度
        # attn_scores = Q @ K.transpose(2, 3) 如下等同
        attn_scores = Q @ K.mT
        
        attn_scores.masked_fill_(
            self.mask.bool()[:num_tokens, :num_tokens],
            -torch.inf
        )
        attention_weight = torch.softmax(attn_scores / K.shape[-1] ** 0.5,dim=-1)
        # print(f"attention_weight: {attention_weight.shape}")
        # attention_weight: torch.Size([2, 2, 6, 6])

        attention_weight = self.dropout(attention_weight)

        # 6.需要对结果交换第1维 和 第2维
        context_vec = (attention_weight @ V).transpose(1, 2)
        # 7.重新展平回去 为 (b, num_tokens, d_out) 的形状
        context_vec = context_vec.contiguous().view(batch_size, num_tokens, self.d_out)
        # print(f"context_vec: {context_vec.shape}")

        # 8.输出投影层  组合多个头的输出后进行投影(可选层)
        # 融合这些不同视角的信息,这样可以得到更丰富、更综合的表示
        # 输出 = W * [头1输出, 头2输出, 头3输出, 头4输出] + b
        context_vec = self.liner_lay(context_vec)   
        return context_vec