#!/usr/bin/env python3
"""候选学生（GGUF 量化 / uncensored 微调 / 任意 HF 目录）vs bf16 原版 的特征对齐测试。

这是"换学生"的照妖镜——适配器训练时的输入是 bf16 原版的 hidden_states[-1]，
候选版隐藏态的任何漂移都会直接传导到生成的 conditioning。

用法:
    python check_alignment.py <bf16 原版目录> <候选: *.gguf 文件 或 HF 目录>
    [--prompts 3] [--device cuda]

判定:
    cos > 0.99   合格（量化级噪声，直接用）
    0.95-0.99    有可见漂移，评估后可用（建议先跑一次端到端生成对比）
    < 0.95       漂移显著，需用候选学生重提特征 + 重训/微调适配器

说明:
    GGUF 走 transformers 内置加载器（>=4.45，AutoModel.from_pretrained(".gguf")），
    隐藏态语义与 bf16 原版完全一致（都是 transformers 路径），测的是量化+微调漂移本身。
"""
from __future__ import annotations

import argparse
import os
import sys

import torch

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)

from student import StudentTextEncoder, load_student, find_language_model  # noqa: E402
from adapter_clip import H3_HIDDEN  # noqa: E402


def load_candidate(path: str, device: torch.device, baseline_mem: float = 0.0):
    from transformers import AutoConfig, AutoModel

    if path.endswith(".gguf"):
        from gguf_qwen35 import load_gguf_model_quantized
        model, info = load_gguf_model_quantized(path)
        print(f"    [candidate] 自定义 GGUF 加载器(量化驻留): {info}", flush=True)
        model = model.to(torch.bfloat16).to(device)  # GGMLTensor 的 to 忽略 dtype, 只移设备
        if device.type == "cuda":
            used = torch.cuda.memory_allocated() / 2**30 - baseline_mem
            print(f"    [candidate] 净显存: {used:.2f} GiB", flush=True)
        return model
    # HF 目录
    cfg = AutoConfig.from_pretrained(path)
    model = AutoModel.from_config(cfg)
    try:
        sd = {k[len("model."):] if k.startswith("model.") else k: v
              for k, v in torch.load(os.path.join(path, "pytorch_model.bin"), map_location="cpu").items()}
        missing, _ = model.load_state_dict(sd, strict=False)
        assert not missing, f"missing: {missing[:5]}"
    except FileNotFoundError:
        model = AutoModel.from_pretrained(path, torch_dtype=torch.bfloat16)
    return model.to(device)


@torch.no_grad()
def hidden_fn(lm, tok, prompt: str, device: torch.device) -> torch.Tensor:
    ids = tok(prompt, add_special_tokens=False)["input_ids"]
    out = lm(input_ids=torch.tensor([ids], device=device), output_hidden_states=True)
    return out.hidden_states[-1][0].float()  # [S, D] 与训练同路径


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("reference", help="bf16 原版 Qwen3.5-4B HF 目录")
    ap.add_argument("candidate", help="候选: .gguf 或 HF 目录")
    ap.add_argument("--prompts", type=int, default=3)
    ap.add_argument("--device", default="")
    args = ap.parse_args()

    from transformers import AutoTokenizer
    dev = torch.device(args.device) if args.device else torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"device: {dev}")

    print("[1/2] loading reference (bf16)...", flush=True)
    ref_model = load_student(args.reference).to(torch.bfloat16).to(dev)
    ref_lm = find_language_model(ref_model)
    ref_tok = AutoTokenizer.from_pretrained(args.reference)

    print("[2/2] loading candidate...", flush=True)
    baseline = torch.cuda.memory_allocated() / 2**30 if dev.type == "cuda" else 0.0
    cand_model = load_candidate(args.candidate, dev, baseline)
    cand_lm = find_language_model(cand_model)
    # .gguf 候选无自带 tokenizer: 复用参考模型的（同一 Qwen3.5-4B 词表）
    cand_tok = ref_tok if args.candidate.endswith(".gguf") else AutoTokenizer.from_pretrained(args.candidate)

    demo = [
        "A cinematic shot of a tiny robot repairing a broken music box in a giant's attic",
        "A surfer riding a massive neon wave at sunset, spray frozen in the air, dramatic backlight",
        "Paper collage birds flocking over a misty city skyline, soft morning light",
    ]
    print(f"\ncosine 对比 {args.prompts} 条 prompt（hidden_states[-1], [S, 2560]）:")
    worst = 1.0
    for i, p in enumerate(demo[: args.prompts]):
        h_ref = hidden_fn(ref_lm, ref_tok, p, dev)
        h_can = hidden_fn(cand_lm, cand_tok, p, dev)
        n = min(h_ref.shape[0], h_can.shape[0])
        if h_ref.shape[0] != h_can.shape[0]:
            print(f"  [注意] prompt {i}: 长度不同 ref={h_ref.shape[0]} cand={h_can.shape[0]}（截断到 {n} 比较）")
        a, b = h_ref[:n].flatten(), h_can[:n].flatten()
        cos = (a * b).sum() / (a.norm() * b.norm())
        worst = min(worst, cos.item())
        print(f"  prompt {i}: cos={cos.item():.4f}")

    print(f"\n最差 cos = {worst:.4f} -> ", end="")
    if worst > 0.99:
        print("合格（量化级噪声，直接用）✅")
    elif worst >= 0.95:
        print("有可见漂移，评估后可用（建议先端到端生成对比）⚠️")
    else:
        print("漂移显著，需用候选学生重提特征 + 重训/微调适配器 ❌")


if __name__ == "__main__":
    main()