"""5.2 / 5.4 base 模型和后训练模型：同一家族、同一尺寸，同一组提示。

base 模型（Qwen3.5-xB-Base）只做过预训练，提示原样当成一段文字喂给它，它接着往下写。
后训练模型（Qwen3.5-xB）做过 SFT 和偏好对齐，按它被训练时的对话格式喂进去。
每个提示先用温度 0 各生成一次，再用温度 0.7 各采样几次，看现象是否稳定。

用法：
    python compare.py                      # 默认 2B，中文提示
    python compare.py --size 0.8B          # 内存小的机器
    python compare.py --size 9B --lang en  # 大一号、英文提示（9B 需要约 20 GB 显存或内存）
    python compare.py --prompts my_prompts.json
两个模型依次加载，同一时刻只占一个模型的内存。Mac 上装了 mlx-lm 就用 MLX，否则用 transformers（CUDA / MPS / CPU 自动选择）。
"""
import argparse
import gc
import json
from datetime import datetime
from pathlib import Path


HERE = Path(__file__).resolve().parent


def pick_backend(choice):
    """Mac 上装了 mlx-lm 就用 MLX（快得多）；否则用 transformers，自动选 CUDA / MPS / CPU。"""
    if choice in ('auto', 'mlx'):
        try:
            import mlx_lm  # noqa: F401
            return 'mlx'
        except ImportError:
            if choice == 'mlx':
                raise SystemExit('没装 mlx-lm：uv pip install -r requirements-mac.txt')
    import torch
    if torch.cuda.is_available():
        return 'cuda'
    return 'mps' if torch.backends.mps.is_available() else 'cpu'


def load(repo, backend):
    """返回一个函数 gen(text, temperature, n) -> n 条续写。"""
    if backend == 'mlx':
        import mlx.core as mx
        from mlx_lm import generate, load as mlx_load
        from mlx_lm.sample_utils import make_sampler
        model, tok = mlx_load(repo)

        def gen(text, temperature, n, max_new_tokens):
            sampler = make_sampler(temp=temperature, top_p=0.95 if temperature else 0.0)
            return [generate(model, tok, prompt=text, max_tokens=max_new_tokens, sampler=sampler) for _ in range(n)]
        return tok, gen, lambda: mx.clear_cache()

    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer
    tok = AutoTokenizer.from_pretrained(repo)
    model = AutoModelForCausalLM.from_pretrained(repo, dtype=torch.float32 if backend == 'cpu' else torch.bfloat16).to(backend)
    model.eval()

    def gen(text, temperature, n, max_new_tokens):
        ids = tok(text, return_tensors='pt').to(backend)
        kw = dict(do_sample=True, temperature=temperature, top_p=0.95, top_k=0) if temperature else dict(do_sample=False)
        out = model.generate(**ids, max_new_tokens=max_new_tokens, num_return_sequences=n, pad_token_id=tok.eos_token_id, **kw)
        return [tok.decode(o[ids['input_ids'].shape[1]:], skip_special_tokens=True) for o in out]

    def free():
        nonlocal model
        del model
        gc.collect()
        if backend == 'cuda':
            torch.cuda.empty_cache()
        elif backend == 'mps':
            torch.mps.empty_cache()
    return tok, gen, free


def run_model(repo, prompts, chat, backend, max_new_tokens, samples, seed):
    tok, gen, free = load(repo, backend)
    if backend == 'mlx':
        import mlx.core as mx
        mx.random.seed(seed)
    else:
        import torch
        torch.manual_seed(seed)
    out = {}
    for p in prompts:
        if chat:      # 后训练模型：按对话格式，关闭思考模式
            text = tok.apply_chat_template([{'role': 'user', 'content': p['text']}], tokenize=False,
                                           add_generation_prompt=True, enable_thinking=False)
        else:         # base 模型：原样当成一段文字
            text = p['text']
        greedy = gen(text, 0, 1, max_new_tokens)[0].strip()
        sampled = gen(text, 0.7, samples, max_new_tokens) if samples else []
        out[p['id']] = dict(greedy=greedy, samples=[x.strip() for x in sampled])
        print(f'  {repo} · {p["id"]} 完成')
    free()
    return out


def one_line(s, limit=400):
    s = s.replace('\n', ' ⏎ ')
    return s if len(s) <= limit else s[:limit] + '……（截断）'


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--size', default='2B', help='0.8B / 2B / 4B / 9B')
    ap.add_argument('--lang', choices=['zh', 'en'], default='zh')
    ap.add_argument('--prompts', help='提示文件（默认 prompts_<lang>.json）')
    ap.add_argument('--max-new-tokens', type=int, default=150)
    ap.add_argument('--samples', type=int, default=3, help='温度 0.7 下每个提示再采样几次')
    ap.add_argument('--seed', type=int, default=0)
    ap.add_argument('--backend', choices=['auto', 'mlx', 'torch'], default='auto')
    ap.add_argument('--out', default='results')
    args = ap.parse_args()

    device = pick_backend(args.backend)
    prompts = json.loads(Path(args.prompts or HERE / f'prompts_{args.lang}.json').read_text(encoding='utf-8'))
    base_repo, inst_repo = f'Qwen/Qwen3.5-{args.size}-Base', f'Qwen/Qwen3.5-{args.size}'
    print(f'{base_repo} 对比 {inst_repo}，设备 {device}，{datetime.now():%Y-%m-%d %H:%M}')
    res = {'base': run_model(base_repo, prompts, False, device, args.max_new_tokens, args.samples, args.seed),
           'instruct': run_model(inst_repo, prompts, True, device, args.max_new_tokens, args.samples, args.seed)}

    lines = [f'# {base_repo} vs {inst_repo}（{device}，{datetime.now():%Y-%m-%d}）', '',
             f'温度 0 各生成一次，最多 {args.max_new_tokens} 个 token；再在温度 0.7 下各采样 {args.samples} 次。⏎ 表示换行。']
    for p in prompts:
        lines += ['', f'## {p["kind"]}：{p["id"]}', '', f'说明：{p["note"]}', '', f'提示：{one_line(p["text"])}', '']
        for who, label in [('base', 'base'), ('instruct', '后训练')]:
            r = res[who][p['id']]
            lines.append(f'- **{label}（温度 0）**：{one_line(r["greedy"])}')
            for i, s in enumerate(r['samples'], 1):
                lines.append(f'    - 采样 {i}：{one_line(s, 200)}')
    report = '\n'.join(lines) + '\n'
    print('\n' + report)
    Path(args.out).mkdir(exist_ok=True)
    stem = f'{args.size}_{args.lang}_{device}_{datetime.now():%Y%m%d_%H%M}'
    (Path(args.out) / f'{stem}.md').write_text(report, encoding='utf-8')
    (Path(args.out) / f'{stem}.json').write_text(json.dumps(dict(prompts=prompts, results=res), ensure_ascii=False, indent=1), encoding='utf-8')
    print(f'报告已存到 {args.out}/{stem}.md')


if __name__ == '__main__':
    main()