#!/usr/bin/env python3
"""exp12 — LLM-судьи (слой 2, ВТОРИЧНО). Кросс-семейные, семейство не судит свои выходы.

Судьи: gemini-3.1-pro-preview + grok-4.3 (кросс-семейные, D13.3г) + gpt-5-mini (нейтральный).
Self-family exclusion: gemini НЕ судит A5; grok НЕ судит A3/A4/A6; gpt-5-mini судит ВСЕ.
Одна комбинированная реплика на (судья, глава, повтор): слепые «Текст N» (перешаффл каждый
повтор = свап позиций), судья возвращает JSON {ranking:[N…best→worst], fidelity:{N:1-5}}.
≥3 повтора → медиана ранга (стиль) и медиана 1-5 (верность, БИЛИНГВАЛЬНО против zh).

Запуск: eval/.venv/bin/python eval/exp12_judges.py
Выход: diag/judges.json (+ summary).
"""
from __future__ import annotations
import json, re, sys, random, time
from pathlib import Path
from statistics import median

sys.path.insert(0, str(Path(__file__).resolve().parent))
from refusal_bench import call_provider, load_env_file
from exp12_pack import arm_chapter_text, POOL, CHAPTERS
load_env_file()

DIAG = Path("/home/ubuntu/books/gu-zhenren/diag")
CAP_USD = 3.0
REPEATS = 3

PRICES = {"gemini-3.1-pro-preview": (2.0, 12.0), "grok-4.3": (1.25, 2.50), "gpt-5-mini": (0.25, 2.0)}

JUDGES = {
    "gemini": dict(name="gemini", base_url="https://generativelanguage.googleapis.com/v1beta/openai",
                   model="gemini-3.1-pro-preview", api_key_env="GEMINI_API_KEY",
                   temperature=0.2, max_tokens=8000, exclude={"A5"}),
    "grok": dict(name="grok", base_url="https://api.x.ai/v1", model="grok-4.3",
                 api_key_env="XAI_API_KEY", temperature=0.2, max_tokens=4000,
                 extra_body={"reasoning_effort": "none"}, exclude={"A3", "A4", "A6"}),
    "gpt5mini": dict(name="gpt5mini", base_url="https://api.openai.com/v1", model="gpt-5-mini",
                     api_key_env="OPENAI_API_KEY", max_tokens=6000, reasoning_params=True,
                     extra_body={"reasoning_effort": "minimal"}, exclude=set()),
}

SYS = ("Ты — литературный критик и билингвальный редактор (китайский→русский). "
       "Тебе дают ИСХОДНИК на китайском и несколько вариантов его русского перевода одной сцены. "
       "Оцени два РАЗДЕЛЬНЫХ измерения:\n"
       "1) СТИЛЬ — качество русского как ХУДОЖЕСТВЕННОЙ прозы (живость, естественность, "
       "отсутствие канцелярита и калек с китайского), без оглядки на точность.\n"
       "2) ВЕРНОСТЬ — насколько точно передан смысл ИСХОДНИКА (без потерь, отсебятины, искажений).\n"
       "Верни СТРОГО JSON без пояснений, комментариев и текста вне JSON. Значения fidelity — "
       "ТОЛЬКО целое число 1..5, без скобок и пояснений. Формат ровно такой:\n"
       "{\"ranking\":[3,1,2],\"fidelity\":{\"1\":4,\"2\":5,\"3\":3}}\n"
       "ranking — номера всех текстов от лучшего к худшему по СТИЛЮ; fidelity — оценка верности "
       "каждого текста 1..5 (5=безупречно точно).")


def cost_of(model, usage):
    inp = usage.get("prompt_tokens", 0) or 0
    comp = usage.get("completion_tokens", 0) or 0
    total = usage.get("total_tokens", 0) or 0
    reason = usage.get("reasoning_tokens", 0) or 0
    pin, pout = PRICES.get(model, (0, 0))
    if model == "gemini-3.1-pro-preview":
        reason = max(0, total - inp - comp)
    return (inp * pin + (comp + reason) * pout) / 1_000_000


def parse_json(text):
    """Lenient: prefer strict JSON, else regex-extract ranking[] + integer fidelity even amid prose."""
    if not text:
        return None
    m = re.search(r"\{.*\}", text, re.S)
    if m:
        try:
            obj = json.loads(m.group(0))
            if "ranking" in obj:
                return obj
        except Exception:
            pass
    out = {}
    rm = re.search(r'"?ranking"?\s*:\s*\[([0-9,\s]+)\]', text, re.S)
    if rm:
        out["ranking"] = [int(x) for x in re.findall(r"\d+", rm.group(1))]
    fm = re.search(r'"?fidelity"?\s*:\s*\{(.*)\}', text, re.S)
    if fm:
        fid = {}
        for k, v in re.findall(r'"?(\d+)"?\s*:\s*(\d+)', fm.group(1)):
            fid[k] = int(v)
        if fid:
            out["fidelity"] = fid
    return out if "ranking" in out else None


def main():
    chunks_json = json.loads((DIAG / "chunks.json").read_text(encoding="utf-8"))
    # chapter texts per arm
    ch_texts = {ch: {arm: arm_chapter_text(arm, ch, chunks_json) for arm in POOL} for ch in CHAPTERS}
    sources = {ch: "\n\n".join(chunks_json[f"{ch}/{ci}"]["source"] for ci in (0, 1)) for ch in CHAPTERS}

    raw = []
    spent = 0.0
    stopped = False
    for jkey, jspec in JUDGES.items():
        if stopped:
            break
        for ch in CHAPTERS:
            arms = [a for a in POOL if a not in jspec["exclude"] and ch_texts[ch][a]]
            for rep in range(REPEATS):
                if spent > CAP_USD * 0.85:
                    print(f"!! judges near cap ${spent:.3f} — stop"); stopped = True; break
                order = arms[:]
                random.Random(f"{jkey}-{ch}-{rep}").shuffle(order)  # position swap
                labels = {i + 1: arm for i, arm in enumerate(order)}
                parts = [f"ИСХОДНИК (zh):\n{sources[ch]}", ""]
                for i, arm in enumerate(order):
                    parts.append(f"=== Текст {i+1} ===\n{ch_texts[ch][arm]}\n")
                user = "\n".join(parts)
                t0 = time.time()
                text, err, usage = call_provider(jspec, SYS, user, timeout=240)
                c = cost_of(jspec["model"], usage); spent += c
                obj = parse_json(text or "")
                ok = bool(obj and "ranking" in obj)
                print(f"  [{jkey} ch{ch} r{rep}] {time.time()-t0:.0f}s ${c:.4f} cum=${spent:.3f} "
                      f"{'OK' if ok else 'PARSE-FAIL:'+str(err or (text or '')[:60])}")
                raw.append({"judge": jkey, "chapter": ch, "rep": rep, "labels": labels,
                            "result": obj, "ok": ok, "cost": round(c, 5),
                            "raw": (text or "")[:1200], "err": err})

    # aggregate
    # style: per (judge,chapter) collect rank of each arm across repeats; normalize rank/N; median
    agg = {}
    for r in raw:
        if not r["ok"]:
            continue
        j, ch, labels, res = r["judge"], r["chapter"], r["labels"], r["result"]
        ranking = res.get("ranking", [])
        n = len(ranking)
        for pos, num in enumerate(ranking):
            arm = labels.get(int(num)) if str(num).isdigit() else labels.get(num)
            if not arm:
                continue
            agg.setdefault((j, arm), {"style_norm": [], "fid": []})
            agg[(j, arm)]["style_norm"].append(pos / max(1, n - 1))  # 0=best..1=worst
        fid = res.get("fidelity", {}) or {}
        for num, sc in fid.items():
            arm = labels.get(int(num)) if str(num).isdigit() else labels.get(num)
            if arm and isinstance(sc, (int, float)):
                agg.setdefault((j, arm), {"style_norm": [], "fid": []})
                agg[(j, arm)]["fid"].append(float(sc))

    summary = {}
    for arm in POOL:
        row = {"arm": arm}
        for j in JUDGES:
            a = agg.get((j, arm))
            if a and a["style_norm"]:
                row[f"{j}_style"] = round(median(a["style_norm"]), 3)
            if a and a["fid"]:
                row[f"{j}_fid"] = round(median(a["fid"]), 2)
        # combined style = mean of available judges' normalized median (lower=better)
        sv = [row[f"{j}_style"] for j in JUDGES if f"{j}_style" in row]
        fv = [row[f"{j}_fid"] for j in JUDGES if f"{j}_fid" in row]
        if sv:
            row["style_combined_norm"] = round(sum(sv) / len(sv), 3)
        if fv:
            row["fid_combined"] = round(sum(fv) / len(fv), 2)
        summary[arm] = row

    (DIAG / "judges.json").write_text(json.dumps(
        {"raw": raw, "summary": summary, "total_cost": round(spent, 4), "repeats": REPEATS},
        ensure_ascii=False, indent=1), encoding="utf-8")

    print(f"\n=== JUDGE SUMMARY (style_norm: 0=best..1=worst; fid 1..5) total ${spent:.4f} ===")
    order = sorted(summary.values(), key=lambda r: r.get("style_combined_norm", 9))
    for r in order:
        print(f"  {r['arm']}: style={r.get('style_combined_norm','–')} "
              f"fid={r.get('fid_combined','–')}  "
              f"(gemini s={r.get('gemini_style','–')}/f={r.get('gemini_fid','–')}, "
              f"grok s={r.get('grok_style','–')}/f={r.get('grok_fid','–')}, "
              f"gpt5 s={r.get('gpt5mini_style','–')}/f={r.get('gpt5mini_fid','–')})")


if __name__ == "__main__":
    main()
