import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer

# 探针1: 使用小模型（GPT-2）
model_name = "gpt2"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 探针2: 简单但可观测的文本
tracer_text = "The cat sat on the mat"

print(f"\n[实验设置]")
print(f"输入文本：\"{tracer_text}\"")

# ============================================================
# 步骤1：Tokenize & Embedding
# ============================================================
input_ids = tokenizer.encode(tracer_text, return_tensors="pt")
tokens = tokenizer.convert_ids_to_tokens(input_ids[0])

embedding_layer = model.transformer.wte
embeddings = embedding_layer(input_ids)

print(f"Token数量：{input_ids.shape[1]}")
print(f"Embedding shape: {embeddings.shape}")  # [1, 6, 768]

# ============================================================
# 步骤2：手动计算第一层的Attention
# ============================================================
layer0 = model.transformer.h[0]
attn = layer0.attn

# 获取Q、K、V的投影权重
c_attn_weight = attn.c_attn.weight  # [768, 2304] = [hidden, 3*hidden]
c_attn_bias = attn.c_attn.bias      # [2304]

# 投影到Q、K、V空间
qkv = embeddings @ c_attn_weight.T + c_attn_bias
qkv = qkv.reshape(1, 6, 3, 12, 64)  # [batch, seq, 3, heads, head_dim]
q, k, v = qkv.unbind(2)

print(f"\n[Q、K、V shape]")
print(f"Q: {q.shape}  [batch, seq_len, num_heads, head_dim]")

# ============================================================
# 步骤3：计算Attention Score (QK^T)
# ============================================================
# 取第0个头
q_head0 = q[0, :, 0, :]  # [6, 64]
k_head0 = k[0, :, 0, :]  # [6, 64]
v_head0 = v[0, :, 0, :]  # [6, 64]

# 计算attention scores
attn_scores = q_head0 @ k_head0.T / (64 ** 0.5)  # [6, 6]

print(f"\n[Attention Scores shape: {attn_scores.shape}]")
print(f"说明：[seq_len, seq_len] - 每个token对其他token的关注度\n")

# 打印attention scores矩阵
print(f"Attention Scores矩阵（Head 0）：")
print(f"{'Token':<10} ", end="")
for token in tokens:
    print(f"{token:<10} ", end="")
print()
for i, token_row in enumerate(tokens):
    print(f"{token_row:<10} ", end="")
    for j, score in enumerate(attn_scores[i]):
        print(f"{score:>6.2f}  ", end="")
    print()

# ============================================================
# 步骤4：应用Causal Mask
# ============================================================
seq_len = input_ids.shape[1]
mask = torch.tril(torch.ones(seq_len, seq_len))

attn_scores_masked = attn_scores.masked_fill(mask == 0, float('-inf'))
attn_weights = F.softmax(attn_scores_masked, dim=-1)

print(f"\n[Causal Mask后的Attention Weights]")

# ============================================================
# 步骤5：加权求和V
# ============================================================
output_head0 = attn_weights @ v_head0  # [6, 64]

print(f"Output shape: {output_head0.shape}")
print(f"说明：[seq_len, head_dim] - 每个token的上下文表示")

# ============================================================
# 复杂度分析
# ============================================================
print(f"\n[计算复杂度分析]")
T = seq_len
print(f"序列长度T = {T}")
print(f"QK^T:     O(T² × d_k) = O({T}² × 64) = {T**2 * 64}")
print(f"@V:       O(T² × d_v) = O({T}² × 64) = {T**2 * 64}")
print(f"总复杂度: O(T²)")
print(f"\n如果T=100:  {T**2:,} 次计算")
print(f"如果T=1000: {(T*10)**2:,} 次计算（慢100倍）")

# 步骤1: 运行探针
python attention_tracer.py

# 步骤2: 测量不同序列长度的计算时间
python -c "
import time
import torch
from transformers import GPT2Model, GPT2Tokenizer

model = GPT2Model.from_pretrained('gpt2')
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

for length in [10, 50, 100, 200]:
    text = 'A ' * length
    inputs = tokenizer(text, return_tensors='pt')

    start = time.time()
    with torch.no_grad():
        _ = model(**inputs)
    elapsed = time.time() - start

    print(f'长度={length:3d}, 耗时={elapsed:.3f}s')
"