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

# ==================== 1. 本地路径配置 ====================
base_model_path = "./Qwen3-4B-Thinking-2507"
lora_path = "./QiMing-Polaris-Qwen3-4B-Thinking-2507_burden_trained_lora"

# 测试 Prompt
prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
What is the 'Burden-based Training' method?

### Input:


### Response:
"""

print("🚀 启动逐层“活跃候选词数量”测算工具...")

# 2. 加载分词器和准备 Input
tokenizer = AutoTokenizer.from_pretrained(base_model_path, trust_remote_code=True)
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")


# ==================== Logit Lens 候选词测算函数 ====================
def count_active_words_per_layer(model, p_thresh=0.001):
    """
    p_thresh = 0.001 代表统计概率 > 0.1% 的所有活跃候选词数量
    """
    residual_outputs = {}
    hooks = []

    def make_hook(layer_idx):
        def hook(module, input_tensor, output_tensor):
            out = (
                output_tensor[0] if isinstance(output_tensor, tuple) else output_tensor
            )
            residual_outputs[layer_idx] = out.detach()

        return hook

    # 定位模型结构
    if hasattr(model, "model") and hasattr(model.model, "layers"):
        layers = model.model.layers
        final_norm = model.model.norm
        lm_head = model.lm_head
    elif hasattr(model, "base_model"):
        layers = model.base_model.model.model.layers
        final_norm = model.base_model.model.model.norm
        lm_head = model.base_model.model.lm_head
    else:
        raise AttributeError("无法定位模型层级结构")

    # 注册 Hook
    for i, layer in enumerate(layers):
        hooks.append(layer.register_forward_hook(make_hook(i)))

    # 前向传播
    with torch.no_grad():
        _ = model(**inputs)

    # 及时清理 Hook
    for h in hooks:
        h.remove()

    # 逐层计算 Logit Lens 映射
    metrics = []
    for i in range(len(layers)):
        # 拿到当前层最后一个 Token 的向量
        h_l = residual_outputs[i][0, -1, :].to(dtype=final_norm.weight.dtype)

        # 强制把当前层的向量过 Final Norm + LM Head，映射到整个词表上
        norm_h = final_norm(h_l)
        logits = lm_head(norm_h)
        probs = F.softmax(logits, dim=-1)

        # 1. 统计概率 > 0.1% 的活跃候选词个数
        active_count = (probs > p_thresh).sum().item()

        # 2. 计算有效候选词数 (e^Entropy)
        log_probs = F.log_softmax(logits, dim=-1)
        entropy = -(probs * log_probs).sum().item()
        effective_count = math.exp(entropy)  # 指数熵，代表“等效候选词个数”

        metrics.append(
            {
                "layer": i,
                "active_words": active_count,
                "effective_words": effective_count,
                "entropy": entropy,
            }
        )

    return metrics


# ==================== 3. 测量 Base 模型 ====================
print(f"\n🔍 正在测算 1/2: 原始模型 [{base_model_path}]...")
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_path,
    torch_dtype=torch.bfloat16,
    device_map="cuda",
    trust_remote_code=True,
)
base_results = count_active_words_per_layer(base_model)


# ==================== 4. 挂载 FT LoRA 并测量 ====================
print(f"🔍 正在测算 2/2: FT 负重训练模型 [Base + {lora_path}]...")
ft_model = PeftModel.from_pretrained(base_model, lora_path)
ft_results = count_active_words_per_layer(ft_model)


# ==================== 5. 打印对比报告 ====================
print("\n" + "=" * 95)
print(
    f"{'层数':<6} | {'[Base] 活跃词数(>0.1%)':<20} | {'[FT负重] 活跃词数(>0.1%)':<20} | {'词数差值 (FT-Base)':<18} | {'找词/砍词趋势'}"
)
print("=" * 95)

for i in range(len(base_results)):
    b_cnt = base_results[i]["active_words"]
    ft_cnt = ft_results[i]["active_words"]
    diff = ft_cnt - b_cnt

    if diff > 0:
        trend = f"⬆️ FT 广搜找词更多 (+{diff})"
    elif diff < 0:
        trend = f"⬇️ FT 猛减砍词 ({diff})"
    else:
        trend = "➡️ 词数持平"

    print(f"L-{i:<3} | {b_cnt:<20} | {ft_cnt:<20} | {diff:<+18} | {trend}")

print("=" * 95)
print("💡 数据解读说明:")
print("1. [活跃词数(>0.1%)] 代表这一层模型脑子里存留的候选词个数。")
print(
    "2. 如果浅层显示 '⬆️ FT 广搜找词更多'，证明 FT 在浅层把关联词全捞出来了（前面找全）；"
)
print(
    "3. 如果中深层显示 '⬇️ FT 猛减砍词'，证明 FT 在中深层下狠手把无用词砍光了（向下猛减）！"
)
print("=" * 95)