"""5.11 用留出的例子评估 LoRA。

对 data/pairs_test.jsonl 里训练时没见过的例子，比较三样东西离“你改过的版本”有多近：
  input    不改，直接用默认输出
  fewshot  不训练，把 3 组“默认输出 → 改过的版本”写进上下文，让同一个底座模型改
  lora     用训好的 LoRA 改，提示里不放例子
并生成一份盲评表（A/B 顺序随机打乱，答案另存），供你自己判断哪个更像你写的。
加 --judge 时，再请一个 API 模型做一遍盲评，并检查有没有改动原意（事实漂移）。

用法：
    python evaluate.py --backend mlx                      # Mac，配合 train_mlx.py
    python evaluate.py --backend torch                    # NVIDIA，配合 train_peft.py
    python evaluate.py --backend mlx --judge gpt-5        # 另外用 API 模型盲评（需要 AI_BUILDER_TOKEN）
"""
import argparse
import difflib
import json
import os
import random
import re
import statistics
from datetime import datetime
from pathlib import Path

from common import HERE, fewshot_messages, read_pairs, user_message


def similarity(a, b):
    """逐字比较的相似度，0 到 1，1 表示完全相同。"""
    return difflib.SequenceMatcher(None, a, b, autojunk=False).ratio()


def mlx_generator(model_id, adapter):
    from mlx_lm import generate, load
    model, tok = load(model_id, adapter_path=adapter)

    def gen(messages, max_tokens):
        prompt = tok.apply_chat_template(messages, add_generation_prompt=True, enable_thinking=False)
        return generate(model, tok, prompt=prompt, max_tokens=max_tokens)
    return gen


def torch_generator(model_id, adapter):
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer
    device = 'cuda' if torch.cuda.is_available() else ('mps' if torch.backends.mps.is_available() else 'cpu')
    tok = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16).to(device)
    if adapter:
        from peft import PeftModel
        model = PeftModel.from_pretrained(model, adapter)
    model.eval()

    def gen(messages, max_tokens):
        text = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
        ids = tok(text, return_tensors='pt').to(device)
        out = model.generate(**ids, max_new_tokens=max_tokens, do_sample=False, pad_token_id=tok.eos_token_id)
        return tok.decode(out[0][ids['input_ids'].shape[1]:], skip_special_tokens=True)
    return gen


JUDGE = '''下面是同一段话的两个改写版本 A 和 B，以及一段“参考原文”。请判断：
1. A 和 B 哪个在用词、句式和语气上更接近参考原文的写法？只能选一个。
2. A 和 B 各自有没有改动原意：增加了参考原文里没有的信息，或者丢掉、改错了信息？

参考原文：{ref}

A：{a}

B：{b}

只输出一行 JSON：{{"closer": "A" 或 "B", "a_changed_meaning": true/false, "b_changed_meaning": true/false}}'''


def judge(model, items):
    from openai import OpenAI
    client = OpenAI(base_url='https://space.ai-builders.com/backend/v1', api_key=os.environ['AI_BUILDER_TOKEN'])
    verdicts = []
    for it in items:
        r = client.chat.completions.create(model=model, max_tokens=4000, messages=[{'role': 'user', 'content': JUDGE.format(
            ref=it['original'], a=it['A'], b=it['B'])}])
        m = re.search(r'\{.*\}', r.choices[0].message.content or '', re.S)
        verdicts.append(json.loads(m.group(0)) if m else None)
    return verdicts


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--backend', choices=['mlx', 'torch'], required=True)
    ap.add_argument('--model', default='Qwen/Qwen3-1.7B')
    ap.add_argument('--adapter', help='LoRA 目录，默认 adapters/mlx 或 adapters/peft')
    ap.add_argument('--data', help='放 pairs_*.jsonl 的目录，默认 data/')
    ap.add_argument('--max-tokens', type=int, default=400)
    ap.add_argument('--judge', help='用哪个 AI Builder Space 模型做盲评，例如 gpt-5；不给就跳过')
    ap.add_argument('--markers', default='伊,的说,罢,么', help='你的写法里有代表性的词，逗号分隔，统计各版本里出现几次')
    ap.add_argument('--out', default='results')
    args = ap.parse_args()
    adapter = args.adapter or str(HERE / 'adapters' / ('mlx' if args.backend == 'mlx' else 'peft'))
    make = mlx_generator if args.backend == 'mlx' else torch_generator

    test, shots = read_pairs('test', args.data), read_pairs('train', args.data)[:3]
    print(f'{len(test)} 条留出例子，底座 {args.model}，{args.backend}，{datetime.now():%Y-%m-%d %H:%M}')
    base = make(args.model, None)
    fewshot = [base(fewshot_messages(p['ai'], shots), args.max_tokens).strip() for p in test]
    print('  基线（3 组例子写进上下文）生成完毕')
    del base
    tuned = make(args.model, adapter)
    lora = [tuned([{'role': 'user', 'content': user_message(p['ai'])}], args.max_tokens).strip() for p in test]
    print('  LoRA 生成完毕')

    sims = {name: [similarity(o, p['original']) for o, p in zip(outs, test)]
            for name, outs in [('input', [p['ai'] for p in test]), ('fewshot', fewshot), ('lora', lora)]}
    wins = sum(l > f for l, f in zip(sims['lora'], sims['fewshot']))
    edit = {name: statistics.mean(similarity(o, p['ai']) for o, p in zip(outs, test))
            for name, outs in [('fewshot', fewshot), ('lora', lora)]}
    closer = {name: sum(x > y for x, y in zip(sims[name], sims['input'])) for name in ['fewshot', 'lora']}
    print('\n## 离“改过的版本”有多近（逐字相似度，0 到 1，越高越近）\n')
    print('| | 与改过的版本的相似度 | 与默认输出的相似度（越低改得越多） | 比不改更接近改过的版本的条数 |')
    print('|---|---|---|---|')
    print(f'| 默认输出本身（不改） | {statistics.mean(sims["input"]):.3f} | 1.000 | — |')
    for name, label in [('fewshot', '基线：3 组例子写进上下文'), ('lora', 'LoRA')]:
        print(f'| {label} | {statistics.mean(sims[name]):.3f} | {edit[name]:.3f} | {closer[name]}/{len(test)} |')
    print(f'\n逐条比较，LoRA 比基线更接近改过的版本的有 {wins}/{len(test)} 条。')
    if args.markers:
        print('\n## 标志性用词出现几次（默认输出 / 基线 / LoRA / 改过的版本）\n')
        for m in args.markers.split(','):
            c = lambda texts: sum(t.count(m) for t in texts)
            print(f'- {m}：{c([p["ai"] for p in test])} / {c(fewshot)} / {c(lora)} / {c([p["original"] for p in test])}')

    rnd = random.Random(0)
    items, key = [], []
    for i, (p, f, l) in enumerate(zip(test, fewshot, lora), 1):
        flip = rnd.random() < 0.5
        a, b = (l, f) if flip else (f, l)
        items.append(dict(n=i, ai=p['ai'], original=p['original'], A=a, B=b))
        key.append(dict(n=i, lora='A' if flip else 'B'))
    out = Path(args.out)
    out.mkdir(exist_ok=True)
    stem = f'{args.backend}_{datetime.now():%Y%m%d_%H%M}'
    sheet = ['# 盲评表', '', '每条看 A 和 B，哪个更像你会写的？把你的选择记下来，再对照 '
             f'`{stem}_key.json`（里面写着哪个是 LoRA）。也留意有没有哪一版改动了原意。', '']
    for it in items:
        sheet += [f'## {it["n"]}', '', f'**默认输出**：{it["ai"]}', '', f'**A**：{it["A"]}', '', f'**B**：{it["B"]}', '']
    (out / f'{stem}_blind.md').write_text('\n'.join(sheet), encoding='utf-8')
    (out / f'{stem}_key.json').write_text(json.dumps(key, indent=1), encoding='utf-8')
    result = dict(model=args.model, backend=args.backend, similarity={k: statistics.mean(v) for k, v in sims.items()},
                  similarity_to_input=edit, closer_than_input=closer, lora_closer_than_fewshot=wins, n=len(test),
                  outputs=[dict(ai=p['ai'], original=p['original'], fewshot=f, lora=l) for p, f, l in zip(test, fewshot, lora)])

    if args.judge:
        verdicts = judge(args.judge, items)
        ok = [(v, k) for v, k in zip(verdicts, key) if v]
        lora_pref = sum(v['closer'] == k['lora'] for v, k in ok)
        drift = {'lora': sum(v['a_changed_meaning'] if k['lora'] == 'A' else v['b_changed_meaning'] for v, k in ok),
                 'fewshot': sum(v['b_changed_meaning'] if k['lora'] == 'A' else v['a_changed_meaning'] for v, k in ok)}
        print(f'\n## {args.judge} 盲评（参考原文在场，A/B 顺序随机）\n')
        print(f'- 判 LoRA 更接近原文写法：{lora_pref}/{len(ok)}')
        print(f'- 判为改动了原意：LoRA {drift["lora"]}/{len(ok)}，基线 {drift["fewshot"]}/{len(ok)}')
        result['judge'] = dict(model=args.judge, lora_preferred=lora_pref, n=len(ok), changed_meaning=drift, verdicts=verdicts)
    (out / f'{stem}_result.json').write_text(json.dumps(result, ensure_ascii=False, indent=1), encoding='utf-8')
    print(f'\n盲评表：{out}/{stem}_blind.md；逐条输出：{out}/{stem}_result.json')


if __name__ == '__main__':
    main()