#!/usr/bin/env python3
"""
End-to-end query-translation table metrics on ``datasets/query_translation_eval.csv``.

For each row uses ``artifacts/baselines/claude/{spec_id}.xml`` as the *reference model*.

Metrics
---------
- **parse_success**: ``verifyta(pred)`` yields satisfied / not_satisfied.
- **EM**: normalized exact match vs CSV ``ground_query``.
- **Faithfulness** (verdict agreement): among rows where a *reference formula* verifies,
  fraction where ``verdict(pred) == verdict(ref)``.

Reference formula: ``adapt_csv_gold_to_model(gold)``; if that fails to verify, try in order
the GPT, Grok, and Claude baseline lines for that row.

Optional MChatbot column: ``--mchatbot`` runs ``translate_natural_language_to_uppaal_query``
per row (needs LLM env configured).

Usage
-----
  python scripts/run_query_translation_table_experiment.py
  python scripts/run_query_translation_table_experiment.py --mchatbot --mchatbot-model gpt-4o-mini
  python scripts/run_query_translation_table_experiment.py --mchatbot --llm-sleep-seconds 3
"""

from __future__ import annotations

import argparse
import csv
import json
import os
import re
import sys
from dataclasses import dataclass
from pathlib import Path

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


def default_llm_inter_request_sleep_s() -> float:
    raw = os.getenv("LLM_INTER_REQUEST_SLEEP_SECONDS", "2.0").strip()
    try:
        return max(0.0, float(raw))
    except ValueError:
        return 2.0


from frame.evaluation.natural2ctl_query_normalize import normalize_natural2ctl_gold_query
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


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


def adapt_csv_gold_to_model(gold: str) -> str:
    """Alias for :func:`adapt_csv_gold_to_reference_xml` (benchmark sheet naming)."""
    return adapt_csv_gold_to_reference_xml(gold)


def parse_pred_lines(path: Path) -> list[str]:
    text = path.read_text(encoding="utf-8", errors="replace")
    out: list[str] = []
    for ln in text.splitlines():
        ln = ln.strip()
        if not ln or re.match(r"^Spec\s+\d+$", ln, re.I):
            continue
        m = re.match(r"^\d+\s*[\.)]\s*(.+)$", ln)
        if m:
            s = m.group(1).strip().strip("`")
        else:
            s = ln.strip().strip("`")
        if s:
            out.append(s)
    return out


def verdict_of(ch: ModelCheckingPipeline, q: str) -> str:
    res, _t, err = ch.verify(q.strip())
    return parse_verifyta_text_verdict(res, errors=err)


def resolve_reference(
    ch: ModelCheckingPipeline,
    gold_adapt: str,
    *fallback_lines: str,
) -> tuple[str | None, str | None]:
    for cand in (gold_adapt, *fallback_lines):
        if not cand.strip():
            continue
        v = verdict_of(ch, cand)
        if v in ("satisfied", "not_satisfied"):
            return cand.strip(), v
    return None, None


@dataclass
class ModelStats:
    name: str
    parse_ok: int
    em_ok: int
    faith_ok: int
    faith_den: int
    n: int

    @property
    def parse_rate(self) -> float:
        return self.parse_ok / self.n if self.n else 0.0

    @property
    def em_rate(self) -> float:
        return self.em_ok / self.n if self.n else 0.0

    @property
    def faith_rate(self) -> float | None:
        return self.faith_ok / self.faith_den if self.faith_den else None


def evaluate_preds(
    name: str,
    preds: list[str],
    rows: list[dict[str, str]],
    checkers: dict[int, ModelCheckingPipeline],
    gpt_preds: list[str],
    grok_preds: list[str],
    claude_preds: list[str],
) -> ModelStats:
    parse_ok = em_ok = faith_ok = faith_den = 0
    n = len(rows)
    for i, row in enumerate(rows):
        sid = int(row["spec_id"])
        gold = row["ground_query"].strip()
        pred = preds[i].strip()
        ch = checkers[sid]

        if verdict_of(ch, pred) in ("satisfied", "not_satisfied"):
            parse_ok += 1
        if normalize_query(pred) == normalize_query(gold):
            em_ok += 1

        g_ad = adapt_csv_gold_to_model(gold)
        ref_q, ref_v = resolve_reference(
            ch,
            g_ad,
            gpt_preds[i].strip(),
            grok_preds[i].strip(),
            claude_preds[i].strip(),
        )
        if ref_q is None or ref_v is None:
            continue
        faith_den += 1
        pv = verdict_of(ch, pred)
        if pv == ref_v and pv in ("satisfied", "not_satisfied"):
            faith_ok += 1

    return ModelStats(name=name, parse_ok=parse_ok, em_ok=em_ok, faith_ok=faith_ok, faith_den=faith_den, n=n)


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--csv", type=Path, default=ROOT / "datasets" / "query_translation_eval.csv")
    ap.add_argument("--claude-models", type=Path, default=ROOT / "artifacts" / "baselines" / "claude")
    ap.add_argument("--mchatbot", action="store_true")
    ap.add_argument("--mchatbot-model", type=str, default="gpt-4o-mini")
    ap.add_argument(
        "--llm-sleep-seconds",
        type=float,
        default=None,
        help="Sleep after each MChatbot LLM call (default: env LLM_INTER_REQUEST_SLEEP_SECONDS or 2.0).",
    )
    ap.add_argument(
        "--oracle-adapt-row",
        action="store_true",
        help="Add row using adapt(csv gold) as predictions (ceiling on this reference XML).",
    )
    ap.add_argument("--json-out", type=Path, default=None)
    args = ap.parse_args()

    llm_sleep = (
        float(args.llm_sleep_seconds)
        if args.llm_sleep_seconds is not None
        else default_llm_inter_request_sleep_s()
    )

    rows = list(csv.DictReader(args.csv.open(encoding="utf-8", newline="")))
    n = len(rows)
    if n != 40:
        print(f"Expected 40 CSV rows, got {n}", file=sys.stderr)

    checkers: dict[int, ModelCheckingPipeline] = {}
    for row in rows:
        sid = int(row["spec_id"])
        if sid not in checkers:
            xm = (args.claude_models / f"{sid}.xml").resolve()
            if not xm.is_file():
                print(f"Missing model {xm}", file=sys.stderr)
                return 1
            checkers[sid] = ModelCheckingPipeline(str(xm))

    gpt_preds = parse_pred_lines(ROOT / "artifacts" / "baselines" / "gpt" / "trans_query.txt")
    grok_preds = parse_pred_lines(ROOT / "artifacts" / "baselines" / "grok" / "trans_query.txt")
    claude_preds = parse_pred_lines(ROOT / "artifacts" / "baselines" / "claude" / "trans_query.txt")

    for label, arr in ("GPT", gpt_preds), ("Grok", grok_preds), ("Claude", claude_preds):
        if len(arr) != n:
            print(f"{label} pred count {len(arr)} != {n}", file=sys.stderr)
            return 1

    stats_list: list[ModelStats] = [
        evaluate_preds("GPT-5.3", gpt_preds, rows, checkers, gpt_preds, grok_preds, claude_preds),
        evaluate_preds("Grok-4", grok_preds, rows, checkers, gpt_preds, grok_preds, claude_preds),
        evaluate_preds("Claude", claude_preds, rows, checkers, gpt_preds, grok_preds, claude_preds),
    ]

    if args.oracle_adapt_row:
        adapted_preds = [adapt_csv_gold_to_model(r["ground_query"]) for r in rows]
        stats_list.append(
            evaluate_preds(
                "Oracle (adapt CSV gold)",
                adapted_preds,
                rows,
                checkers,
                gpt_preds,
                grok_preds,
                claude_preds,
            )
        )

    if args.mchatbot:
        try:
            import importlib.util
            import time

            from frame.prompts import COMPLEX_QUERY_LLM_PROMPT
            from frame.rag_component.llm import LLM
            from frame.pipeline.nl_uppaal_query import translate_natural_language_to_uppaal_query

            bpath = ROOT / "scripts" / "benchmark_gold_bundles_query_translation.py"
            spec = importlib.util.spec_from_file_location("_qt_bench", bpath)
            if spec is None or spec.loader is None:
                raise RuntimeError("cannot load benchmark module")
            bmod = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(bmod)
            schema_from_uppaal_xml = bmod.schema_from_uppaal_xml

            llm = LLM(system_prompt=COMPLEX_QUERY_LLM_PROMPT, model_name=args.mchatbot_model, max_tokens=1200)
            mc_preds: list[str] = []
            for row in rows:
                sid = int(row["spec_id"])
                xm = (args.claude_models / f"{sid}.xml").resolve()
                nl = row["nl_query"].strip()
                schema = schema_from_uppaal_xml(xm)
                pred, _parsed, _raw = translate_natural_language_to_uppaal_query(
                    nl, schema, llm=llm, xml_path=str(xm)
                )
                mc_preds.append((pred or "").strip())
                if llm_sleep > 0:
                    time.sleep(llm_sleep)
            stats_list.append(
                evaluate_preds(
                    f"MChatbot ({args.mchatbot_model})",
                    mc_preds,
                    rows,
                    checkers,
                    gpt_preds,
                    grok_preds,
                    claude_preds,
                )
            )
        except Exception as e:
            print(f"MChatbot run failed: {e}", file=sys.stderr)
            return 1

    out = {
        "n_rows": n,
        "reference_models": str(args.claude_models),
        "metrics_note": "Faithfulness = verifyta verdict match vs reference (adapted CSV gold, else GPT, Grok, Claude).",
        "models": [],
    }
    print("\n=== Query translation experiment (40 NL queries) ===\n")
    print(f"{'Model':<32} {'parse':>8} {'EM':>8} {'Faith':>8}  (Faith denominator = rows with ref verify)\n")
    for s in stats_list:
        fr = s.faith_rate
        out["models"].append(
            {
                "name": s.name,
                "parse_success": round(s.parse_rate, 4),
                "exact_match": round(s.em_rate, 4),
                "faithfulness": None if fr is None else round(fr, 4),
                "faithfulness_denominator": s.faith_den,
            }
        )
        print(
            f"{s.name:<32} {s.parse_rate:8.2f} {s.em_rate:8.2f} "
            f"{(fr if fr is not None else float('nan')):8.2f}   (n_ref={s.faith_den})"
        )

    if args.json_out:
        args.json_out.write_text(json.dumps(out, indent=2), encoding="utf-8")
        print(f"\nWrote {args.json_out}")

    print(
        "\nLaTeX (copy rows; caption: 40 natural-language queries per model; scale if you merge splits):\n\n"
        r"\begin{tabular}{|l|c|c|c|}"
        "\n\\hline\n"
        r"\textbf{Model} & \textbf{parse\_success} & \textbf{EM} & \textbf{Faithfulness} \\"
        "\n\\hline"
    )
    for m in out["models"]:
        f = m["faithfulness"]
        fs = f"{f:.2f}" if f is not None else "---"
        row_name = m["name"].replace("&", r"\&")
        print(f"{row_name} & {m['parse_success']:.2f} & {m['exact_match']:.2f} & {fs} \\\\")
    print("\\hline\n\\end{tabular}")
    return 0


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