#!/usr/bin/env python3
"""Judge the sentence helper's answers with a large local model.

  python tools/ai/eval/judge.py ANSWERS.json OUT.json [--model mlx-community/gemma-4-31b-it-8bit] [lang...]

ANSWERS.json is what run.mjs wrote. For every "sentence" answer (the words
chosen and what the partner said), the judge says whether one of the three
suggestions says what the person meant, as a native speaker would say it.
For every "reply" answer, it counts the replies that are natural, grammatical
and fit what the partner said. The judge sees the person's intended sentence
(the item's "target") only for sentences.

This is one reader among others: languages are marked ready only when the
people reviewing agree (see ../README.md). Needs mlx-lm (Apple silicon).
"""

import argparse
import json
import re
import sys
from pathlib import Path

NAMES = None  # filled from the answers' languages when known

SENTENCE_RUBRIC = """You check suggestions made by an AAC app for a person who cannot speak. The person chose a few words (and their partner may have said something first); the app suggested three sentences in {lang}. You know what the person meant.

Partner said: {partner}
Words chosen: {words}
What the person meant: {target}
Suggestions:
1. {s1}
2. {s2}
3. {s3}

Is at least one suggestion a sentence a native {lang} speaker would naturally say, with correct grammar, that means what the person meant? Small differences are fine (a synonym, slightly more or fewer words, another tense or a gendered form, a polite variant). Not fine: a different or reversed meaning, speaking for the partner instead of the person, or broken grammar.

Answer with JSON only: {{"usable": true or false, "best": the number of the best suggestion or 0, "why": "a few words"}}"""

REPLY_RUBRIC = """You check reply suggestions made by an AAC app for a person who cannot speak. Their partner said something in {lang}; the app suggested replies the person might say back.

Partner said: {partner}
Replies:
{replies}

Count the replies that a native {lang} speaker would naturally say back to this, in the first person, with correct grammar, and that make sense as an answer.

Answer with JSON only: {{"good": a number from 0 to {n}, "why": "a few words"}}"""


def lines(text, n):
    out = []
    for line in text.splitlines():
        s = re.sub(r"^\s*(?:[-*•]|\d+[.)])\s*", "", line).strip().strip('"').strip()
        if s:
            out.append(s)
    return (out + ["(none)"] * n)[:n]


def parse(text):
    m = re.search(r"\{.*\}", text, re.S)
    if not m:
        return None
    try:
        return json.loads(m.group(0))
    except json.JSONDecodeError:
        return None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("answers")
    ap.add_argument("out")
    ap.add_argument("langs", nargs="*")
    ap.add_argument("--model", default="mlx-community/gemma-4-31b-it-8bit")
    ap.add_argument("--names", default=str(Path(__file__).resolve().parent.parent / "languages.json"))
    args = ap.parse_args()

    from mlx_lm import generate, load

    names = {c: l["name"] for c, l in json.loads(Path(args.names).read_text())["languages"].items()}
    data = json.loads(Path(args.answers).read_text())
    answers = [a for a in data["answers"] if not args.langs or a["lang"] in args.langs]
    model, tok = load(args.model)

    def ask(prompt):
        msgs = [{"role": "user", "content": prompt}]
        try:
            text = tok.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False)
        except TypeError:
            text = tok.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False)
        return generate(model, tok, prompt=text, max_tokens=80, verbose=False)

    out, done = [], 0
    for a in answers:
        lang = names.get(a["lang"], a["lang"])
        if a["kind"] == "sentence":
            s = lines(a["text"], 3)
            verdict = parse(ask(SENTENCE_RUBRIC.format(lang=lang, partner=a["partner"] or "(nothing)", words=a["words"],
                                                       target=a["target"], s1=s[0], s2=s[1], s3=s[2])))
            out.append({"lang": a["lang"], "kind": "sentence", "words": a["words"], "target": a["target"],
                        "usable": bool(verdict and verdict.get("usable")), "verdict": verdict})
        elif a["kind"] == "reply":
            r = lines(a["text"], 4)
            verdict = parse(ask(REPLY_RUBRIC.format(lang=lang, partner=a["partner"], n=4,
                                                    replies="\n".join(f"{i + 1}. {x}" for i, x in enumerate(r)))))
            good = verdict.get("good") if verdict else None
            out.append({"lang": a["lang"], "kind": "reply", "partner": a["partner"],
                        "good": good if isinstance(good, int) else None, "verdict": verdict})
        done += 1
        if done % 20 == 0:
            print(f"  {done}/{len(answers)}", file=sys.stderr, flush=True)

    summary = {}
    for code in sorted({o["lang"] for o in out}):
        sent = [o for o in out if o["lang"] == code and o["kind"] == "sentence"]
        reps = [o for o in out if o["lang"] == code and o["kind"] == "reply" and o["good"] is not None]
        summary[code] = {
            "sentences_usable": sum(o["usable"] for o in sent), "of": len(sent),
            "replies_good": sum(o["good"] for o in reps), "replies_of": 4 * len(reps),
        }
    Path(args.out).write_text(json.dumps({"judge": args.model, "summary": summary, "items": out},
                                         ensure_ascii=False, indent=1) + "\n")
    for code, s in summary.items():
        print(f"{code}: sentences {s['sentences_usable']}/{s['of']}, replies {s['replies_good']}/{s['replies_of']}")


if __name__ == "__main__":
    main()