import torch
import torch.nn.functional as F
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（使用你的 Alpaca 标准格式）
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")


# ==================== 审计核心 Hook 函数 ====================
def audit_model_interference(model):
    """通过注册 Hook 测量模型每一层的残差流变化量 Δh 以及夹角余弦 cos(h, Δh)"""
    residual_inputs = {}
    residual_outputs = {}
    hooks = []

    def make_hook(layer_idx):
        def hook(module, input_tensor, output_tensor):
            residual_inputs[layer_idx] = input_tensor[0].detach()
            # 如果输出是 tuple（如 Qwen 架构），取第 0 项向量
            out = output_tensor[0] if isinstance(output_tensor, tuple) else output_tensor
            residual_outputs[layer_idx] = out.detach()
        return hook

    # 兼容 PEFT 挂载前后的模型层获取
    if hasattr(model, "model") and hasattr(model.model, "layers"):
        layers = model.model.layers
    elif hasattr(model, "base_model"):
        layers = model.base_model.model.model.layers
    else:
        raise AttributeError("无法自动定位模型的 layers 结构")

    # 注册 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()

    # 计算各层干涉数据
    metrics = []
    for i in range(len(layers)):
        h_l = residual_inputs[i][0, -1, :].float()      # 当前层最后一个 Token 的向量
        h_next = residual_outputs[i][0, -1, :].float()  # 下一层向量
        delta_h = h_next - h_l                    # 本层新增的向量变化量 Δh_l

        cos_sim = F.cosine_similarity(h_l, delta_h, dim=0).item()
        delta_norm = delta_h.norm().item()

        metrics.append({
            "layer": i,
            "norm": delta_norm,
            "cos_sim": cos_sim
        })

    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_metrics = audit_model_interference(base_model)


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


# ==================== 5. 打印对比报告 ====================
print("\n" + "="*100)
print(f"{'层数':<6} | {'[原始 Base] cos(h,Δh)':<22} | {'[FT 负重] cos(h,Δh)':<22} | {'余弦变化(Diff)':<14} | {'干涉趋势变化'}")
print("="*100)

for i in range(len(base_metrics)):
    b_cos = base_metrics[i]["cos_sim"]
    ft_cos = ft_metrics[i]["cos_sim"]
    diff = ft_cos - b_cos

    b_type = "🔴减法" if b_cos < 0 else "🟢加法"
    ft_type = "🔴减法" if ft_cos < 0 else "🟢加法"

    # 判断趋势变化
    if ft_cos < b_cos:
        trend = "⬇️ 负向干涉增强 (做减法/抵消变强)"
    elif ft_cos > b_cos:
        trend = "⬆️ 正向叠加增强 (做加法)"
    else:
        trend = "➡️ 无变化"

    print(f"L-{i:<3} | {b_cos:<8.4f} ({b_type})   | {ft_cos:<8.4f} ({ft_type})   | {diff:<+10.4f}     | {trend}")

print("="*100)
print("💡 结果解读指南:")
print("1. [FT 负重] 的 cos(h,Δh) 数值越小或越负，说明 FT 训练在该层施加的‘相消干涉（减法/抵消）’越强。")
print("2. 如果变化趋势显示 '⬇️ 负向干涉增强'，证明挂载 FT LoRA 后，该层正在主动计算反向向量去抵消杂音噪声！")
print("="*100)