#!/usr/bin/env python3
"""
MCR / MAC for **baseline-export UPPAAL XMLs** (e.g. ``artifacts/baselines/claude/1.xml``).

MCR (same spirit as ``aggregate_batch_metrics.py``):
  fraction of specs S01..S20 where the baseline ``{spec_id}.xml`` exists, is non-empty,
  and verifyta runs ``A[] not deadlock`` without engine/syntax errors.

MAC:
  For each smoke-OK spec, run **reference-adapted** gold queries from
  ``datasets/query_translation_eval.csv`` (see ``adapt_csv_gold_to_reference_xml``)
  on the **baseline** XML. Compare satisfy / not_satisfy to column ``answer`` (T/F).
  MAC = mean of per-spec accuracies (correct / 5), matching batch convention.

Sanity (``--check-gold-model``):
  On ``datasets/gold_models/S{{id}}/model.xml``, run the **raw** ``ground_query``
  from the CSV (identifiers match gold). If verdict ≠ ``answer``, the eval row or gold
  model is inconsistent — reported as warnings.

Usage
-----
  python scripts/eval_baseline_models_mcr_mac.py --baseline claude
  python scripts/eval_baseline_models_mcr_mac.py --baseline claude --check-gold-model
"""
from __future__ import annotations

import argparse
import csv
import json
import re
import sys
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))

from frame.evaluation.reference_xml_gold_adapt import adapt_csv_gold_to_reference_xml
from frame.pipeline.model_checking_pipeline import ModelCheckingPipeline
from frame.pipeline.nl_uppaal_query import parse_verifyta_text_verdict

EVAL_CSV = ROOT / "datasets" / "query_translation_eval.csv"
GOLD_MODELS = ROOT / "datasets" / "gold_models"
BASELINES = ROOT / "artifacts" / "baselines"
N_SPECS = 20


def _norm_ws(s: str) -> str:
    return re.sub(r"\s+", " ", (s or "").strip())


def _verdict_tf(v: str | None) -> str:
    if v == "satisfied":
        return "T"
    if v == "not_satisfied":
        return "F"
    return "?"


def _smoke_ok(checker: ModelCheckingPipeline) -> bool:
    _res, _tr, errs = checker.verify("A[] not deadlock")
    errors = [e for e in (errs or []) if e and str(e).strip()]
    return not errors


def load_rows(path: Path) -> list[dict[str, str]]:
    with path.open(encoding="utf-8", newline="") as f:
        return list(csv.DictReader(f))


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--baseline", required=True, help="Subfolder under artifacts/baselines (e.g. claude)")
    ap.add_argument("--csv", type=Path, default=EVAL_CSV)
    ap.add_argument(
        "--out-json",
        type=Path,
        default=None,
        help="Optional path to write detailed JSON (per-spec + per-row).",
    )
    ap.add_argument(
        "--check-gold-model",
        action="store_true",
        help="Verify raw ground_query on datasets/gold_models (sanity).",
    )
    args = ap.parse_args()

    base_dir = (BASELINES / args.baseline).resolve()
    if not base_dir.is_dir():
        print(f"Not a directory: {base_dir}", file=sys.stderr)
        return 1

    rows = load_rows(args.csv.resolve())
    by_spec: dict[int, list[dict[str, str]]] = {}
    for r in rows:
        by_spec.setdefault(int(r["spec_id"]), []).append(r)

    gold_checkers: dict[int, ModelCheckingPipeline] = {}
    gold_mismatch = 0
    if args.check_gold_model:
        for sid in range(1, N_SPECS + 1):
            gp = GOLD_MODELS / f"S{sid:02d}" / "model.xml"
            if not gp.is_file():
                print(f"[gold-model] missing {gp}", file=sys.stderr)
                continue
            gold_checkers[sid] = ModelCheckingPipeline(str(gp))

        for r in rows:
            sid = int(r["spec_id"])
            if sid not in gold_checkers:
                continue
            gq = _norm_ws(r.get("ground_query") or "")
            exp = (r.get("answer") or "").strip().upper()
            res, _t, errs = gold_checkers[sid].verify(gq)
            v = parse_verifyta_text_verdict(res, errors=errs)
            got = _verdict_tf(v)
            if got != exp:
                gold_mismatch += 1
                print(
                    f"[gold-model MISMATCH] S{sid:02d} expected {exp} got {got} "
                    f"(query …{gq[:60]}…)",
                    file=sys.stderr,
                )

    n_compile = 0
    per_spec_scores: list[tuple[int, str, float, int, int]] = []
    detail_specs: list[dict[str, Any]] = []

    for sid in range(1, N_SPECS + 1):
        xm = base_dir / f"{sid}.xml"
        spec_rows = by_spec.get(sid, [])
        spec_name = f"S{sid:02d}"

        if not xm.is_file() or xm.stat().st_size == 0:
            detail_specs.append(
                {
                    "spec_id": sid,
                    "xml": str(xm),
                    "smoke_ok": False,
                    "reason": "missing or empty XML",
                    "queries": [],
                }
            )
            continue

        bl_ch = ModelCheckingPipeline(str(xm))
        ok = _smoke_ok(bl_ch)
        if not ok:
            detail_specs.append(
                {
                    "spec_id": sid,
                    "xml": str(xm),
                    "smoke_ok": False,
                    "reason": "smoke verifyta error",
                    "queries": [],
                }
            )
            continue

        n_compile += 1
        q_correct = 0
        q_total = 0
        q_rows_out: list[dict[str, Any]] = []

        for r in spec_rows:
            raw_gold = _norm_ws(r.get("ground_query") or "")
            adapted = adapt_csv_gold_to_reference_xml(raw_gold)
            exp = (r.get("answer") or "").strip().upper()
            res, _t, errs = bl_ch.verify(adapted)
            v = parse_verifyta_text_verdict(res, errors=errs)
            got = _verdict_tf(v)
            errors = [e for e in (errs or []) if e and str(e).strip()]
            ok_row = got == exp and got in ("T", "F")
            if ok_row:
                q_correct += 1
            q_total += 1
            q_rows_out.append(
                {
                    "nl_query": (r.get("nl_query") or "")[:200],
                    "ground_query_adapted": adapted,
                    "expected": exp,
                    "verdict": got,
                    "correct": ok_row,
                    "verifyta_errors": errors[:3],
                }
            )

        acc = q_correct / q_total if q_total else 0.0
        per_spec_scores.append((sid, spec_name, acc, q_correct, q_total))
        detail_specs.append(
            {
                "spec_id": sid,
                "xml": str(xm),
                "smoke_ok": True,
                "q_correct": q_correct,
                "q_total": q_total,
                "accuracy": acc,
                "queries": q_rows_out,
            }
        )

    mcr = n_compile / N_SPECS
    mac = (
        sum(s[2] for s in per_spec_scores) / len(per_spec_scores)
        if per_spec_scores
        else 0.0
    )

    print(f"Baseline folder: {base_dir}")
    print(f"Eval CSV: {args.csv.resolve()}")
    print(f"MCR = {n_compile}/{N_SPECS} = {mcr:.4f} ({mcr*100:.1f}%)")
    print(
        f"MAC = mean(per-spec accuracy) over {len(per_spec_scores)} smoke-OK spec(s) "
        f"= {mac:.4f} ({mac * 100:.1f}%)"
        if per_spec_scores
        else "MAC = n/a (no smoke-OK specs)"
    )
    print()
    for sid, name, acc, qc, qt in per_spec_scores:
        print(f"  S{sid:02d}  {qc}/{qt}  ({acc:.2f})")
    if args.check_gold_model:
        print()
        print(f"Gold-model row checks: {gold_mismatch} mismatch(es) (expected 0)")

    if args.out_json:
        args.out_json.parent.mkdir(parents=True, exist_ok=True)
        payload = {
            "baseline": args.baseline,
            "mcr": mcr,
            "mac": mac,
            "n_compile": n_compile,
            "per_spec": [
                {"spec_id": a, "name": b, "accuracy": c, "q_correct": d, "q_total": e}
                for a, b, c, d, e in per_spec_scores
            ],
            "specs": detail_specs,
        }
        args.out_json.write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding="utf-8")
        print(f"\nWrote {args.out_json}")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())