#!/usr/bin/env python3
"""
Evaluate baseline ``trans_query.txt`` outputs against ``datasets/query_translation_eval.csv``.

Metrics (per row):
  - exact_match: normalized equality vs **reference-XML-adapted** gold (same identifier space
    as ``artifacts/baselines/claude/*.xml`` and verifyta on those models)
  - exact_match_dataset_csv: same vs raw ``ground_query`` from CSV
  - token_precision / recall / f1: vs adapted gold; ``*_dataset_csv`` vs CSV gold
  - semantic_correct (optional --verifyta): ``verdict(pred) == verdict(gold)`` on the same XML
  - csv_answer_agree (optional): whether ``verdict(pred)`` matches CSV ``answer`` (T/F)

Gold model layout:
  Pass ``--gold-models-dir D`` so that ``D/1.xml`` ... ``D/10.xml`` exist (``spec_id`` in CSV).

Usage:
  python scripts/eval_baseline_trans_query.py \\
    --csv datasets/query_translation_eval.csv \\
    --pred artifacts/baselines/gpt/trans_query.txt \\
    --gold-models-dir artifacts/_gold_modelbuild_runs_20260501_182558

  python scripts/eval_baseline_trans_query.py ... --verifyta --json-out artifacts/eval_gpt.json
"""

from __future__ import annotations

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

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

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.nl_uppaal_query import parse_verifyta_text_verdict

try:
    from frame.pipeline.model_checking_pipeline import ModelCheckingPipeline
except ImportError:
    ModelCheckingPipeline = None  # type: ignore[misc, assignment]


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


_TOK_RE = re.compile(
    r"""
    (A\s*\[\]|E\s*<>|A\s*<>|E\s*\[\])
    |(-->)
    |(&&|\|\||<=|>=|==|!=)
    |([(){}\[\],])
    |(\bforall\b|\bexists\b|\bimply\b|\bnot\b|\bdeadlock\b)
    |([A-Za-z_]\w*)
    |(\d+)
    """,
    re.VERBOSE | re.IGNORECASE,
)


def token_f1(pred: str, gold: str) -> tuple[float, float, float]:
    p = normalize_query(pred)
    g = normalize_query(gold)
    pt = [t for t in _TOK_RE.findall(p) for t in t if t]
    gt = [t for t in _TOK_RE.findall(g) for t in t if t]
    pc = Counter(pt)
    gc = Counter(gt)
    inter = sum((pc & gc).values())
    p_tot = sum(pc.values())
    g_tot = sum(gc.values())
    if p_tot == 0 and g_tot == 0:
        return 1.0, 1.0, 1.0
    if p_tot == 0 or g_tot == 0:
        return 0.0, 0.0, 0.0
    prec = inter / p_tot
    rec = inter / g_tot
    if prec + rec == 0:
        return 0.0, 0.0, 0.0
    return prec, rec, 2 * prec * rec / (prec + rec)


def _strip_inline_comment(line: str) -> str:
    if "//" in line:
        line = line.split("//", 1)[0]
    return line.strip()


def _strip_backticks(s: str) -> str:
    s = s.strip()
    if s.startswith("`") and s.endswith("`") and len(s) >= 2:
        return s[1:-1].strip()
    return s


def parse_trans_query_file(path: Path) -> list[str]:
    """
    Parse baseline text: numbered lines ``1. ...`` / ``2) ...``, optional ``Spec N`` headers,
    or one formula per non-empty line (Grok-style).
    """
    text = path.read_text(encoding="utf-8", errors="replace")
    raw_lines = text.splitlines()
    lines: list[str] = []
    for ln in raw_lines:
        ln = _strip_inline_comment(ln).strip()
        if not ln or re.match(r"^Spec\s+\d+\s*$", ln, re.IGNORECASE):
            continue
        lines.append(ln)

    numbered: list[tuple[int, str]] = []
    for ln in lines:
        m = re.match(r"^\s*(\d+)\s*[\.)]\s*(.+)$", ln)
        if m:
            numbered.append((int(m.group(1)), _strip_backticks(m.group(2))))

    # Prefer numbered format when it dominates (GPT / Claude exports).
    if len(numbered) >= max(3, len(lines) // 2):
        by_id: dict[int, str] = {}
        for idx, q in numbered:
            by_id[idx] = q
        order = sorted(by_id.keys())
        return [by_id[i] for i in order]

    out: list[str] = []
    for ln in lines:
        m = re.match(r"^\s*(\d+)\s*[\.)]\s*(.+)$", ln)
        if m:
            out.append(_strip_backticks(m.group(2)))
        else:
            out.append(_strip_backticks(ln))
    return out


def _csv_answer_to_verdict(ans: str) -> str | None:
    t = (ans or "").strip().upper()
    if t == "T":
        return "satisfied"
    if t == "F":
        return "not_satisfied"
    return None


def load_eval_csv(path: Path) -> list[dict[str, str]]:
    rows: list[dict[str, str]] = []
    with path.open(encoding="utf-8", newline="") as f:
        r = csv.DictReader(f)
        for row in r:
            if not row.get("spec_id"):
                continue
            rows.append(row)
    return rows


def resolve_model_path(spec_id: str, models_dir: Path | None) -> Path | None:
    if models_dir is None:
        return None
    p = models_dir / f"{int(spec_id)}.xml"
    return p if p.is_file() else None


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--csv", type=Path, default=ROOT / "datasets" / "query_translation_eval.csv")
    ap.add_argument("--pred", type=Path, required=True, help="Baseline trans_query.txt")
    ap.add_argument(
        "--gold-models-dir",
        type=Path,
        default=None,
        help="Directory with 1.xml .. 10.xml (must match spec_id in CSV for --verifyta)",
    )
    ap.add_argument("--verifyta", action="store_true", help="Run verifyta (needs pyuppaal + verifyta)")
    ap.add_argument("--json-out", type=Path, default=None)
    args = ap.parse_args()

    rows = load_eval_csv(args.csv.resolve())
    preds = parse_trans_query_file(args.pred.resolve())

    if len(preds) != len(rows):
        print(
            f"Count mismatch: CSV has {len(rows)} rows, parsed {len(preds)} predictions from {args.pred}",
            file=sys.stderr,
        )
        return 1

    if args.verifyta and args.gold_models_dir is None:
        print("--verifyta requires --gold-models-dir", file=sys.stderr)
        return 1
    if args.verifyta and ModelCheckingPipeline is None:
        print("Could not import ModelCheckingPipeline", file=sys.stderr)
        return 1

    checkers: dict[str, Any] = {}
    if args.verifyta and args.gold_models_dir:
        for r in rows:
            sid = str(int(r["spec_id"]))
            if sid in checkers:
                continue
            xm = resolve_model_path(sid, args.gold_models_dir.resolve())
            if xm is None:
                print(f"Missing model for spec_id={sid}: {args.gold_models_dir / (sid + '.xml')}", file=sys.stderr)
                return 1
            checkers[sid] = ModelCheckingPipeline(str(xm))

    per_row: list[dict[str, Any]] = []
    tot_exact = 0
    tot_exact_csv = 0
    tot_f1 = 0.0
    tot_f1_csv = 0.0
    sem_ok = 0
    sem_att = 0
    csv_agree = 0
    csv_att = 0
    gold_parse_fail = 0

    for i, row in enumerate(rows):
        gold_q = (row.get("ground_query") or "").strip()
        gold_ref = adapt_csv_gold_to_reference_xml(gold_q)
        nl = (row.get("nl_query") or "").strip()
        spec_id = str(int(row["spec_id"]))
        pred = preds[i]
        ex_ref = normalize_query(pred) == normalize_query(gold_ref)
        ex_csv = normalize_query(pred) == normalize_query(gold_q)
        if ex_ref:
            tot_exact += 1
        if ex_csv:
            tot_exact_csv += 1
        prec, rec, f1 = token_f1(pred, gold_ref)
        prec_c, rec_c, f1_c = token_f1(pred, gold_q)
        tot_f1 += f1
        tot_f1_csv += f1_c

        rec_out: dict[str, Any] = {
            "index": i + 1,
            "spec_id": spec_id,
            "nl_query": nl,
            "ground_query": gold_q,
            "ground_query_adapted_ref_xml": gold_ref,
            "pred_query": pred,
            "exact_match": ex_ref,
            "exact_match_dataset_csv": ex_csv,
            "token_precision": prec,
            "token_recall": rec,
            "token_f1": f1,
            "token_precision_dataset_csv": prec_c,
            "token_recall_dataset_csv": rec_c,
            "token_f1_dataset_csv": f1_c,
        }

        if args.verifyta:
            ch = checkers[spec_id]
            g_res, _gt, g_err = ch.verify(gold_ref)
            p_res, _pt, p_err = ch.verify(pred)
            g_ver = parse_verifyta_text_verdict(g_res, errors=g_err)
            p_ver = parse_verifyta_text_verdict(p_res, errors=p_err)
            sem = (
                g_ver == p_ver
                and g_ver in ("satisfied", "not_satisfied")
                and p_ver in ("satisfied", "not_satisfied")
            )
            if g_ver not in ("satisfied", "not_satisfied"):
                gold_parse_fail += 1
            if g_ver in ("satisfied", "not_satisfied") and p_ver in ("satisfied", "not_satisfied"):
                sem_att += 1
                if sem:
                    sem_ok += 1

            exp = _csv_answer_to_verdict(row.get("answer") or "")
            rec_out["gold_verdict"] = g_ver
            rec_out["pred_verdict"] = p_ver
            rec_out["semantic_correct"] = sem
            rec_out["gold_errors"] = (g_err or [])[:2]
            rec_out["pred_errors"] = (p_err or [])[:2]
            if exp is not None and p_ver in ("satisfied", "not_satisfied"):
                csv_att += 1
                if p_ver == exp:
                    csv_agree += 1
                rec_out["csv_answer_agree"] = p_ver == exp

        per_row.append(rec_out)

    n = len(rows)
    summary = {
        "csv": str(args.csv),
        "pred": str(args.pred),
        "gold_models_dir": str(args.gold_models_dir) if args.gold_models_dir else None,
        "n": n,
        "exact_match_rate": tot_exact / n,
        "exact_match_rate_dataset_csv": tot_exact_csv / n,
        "mean_token_f1": tot_f1 / n,
        "mean_token_f1_dataset_csv": tot_f1_csv / n,
        "verifyta": bool(args.verifyta),
    }
    if args.verifyta:
        summary["semantic_correct_rate"] = sem_ok / sem_att if sem_att else None
        summary["semantic_compared"] = sem_att
        summary["csv_answer_agree_rate"] = csv_agree / csv_att if csv_att else None
        summary["gold_verdict_fail_count"] = gold_parse_fail

    print(json.dumps(summary, indent=2))
    if args.json_out:
        args.json_out.parent.mkdir(parents=True, exist_ok=True)
        out_obj = {"summary": summary, "rows": per_row}
        args.json_out.write_text(json.dumps(out_obj, indent=2, ensure_ascii=False), encoding="utf-8")
        print(f"Wrote {args.json_out}", file=sys.stderr)

    return 0


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