"""Offline scorer for eval_baseline.py dumps.

Reads a dump JSON (from eval_baseline.py) and computes metrics under
multiple slot-matching strategies without re-running inference.

Strategies:
  exact:      pred dict == gold dict (case-sensitive)
  ci:         case-insensitive keys/values
  subset:     every gold key+value present in pred (tolerates extra pred slots)
  ci-subset:  case-insensitive subset (the V32-baseline-equivalent semantic)
  intent_snap: also snap pred intent to nearest 50-enum (post-hoc constraint)

Usage:
  uv run python poc/llm-finetune/training/score_baseline.py \\
      poc/llm-finetune/training/eval_dump_q3b_adapter_v7_r16.json \\
      --strategy ci-subset --by intent

  uv run python poc/llm-finetune/training/score_baseline.py <dump> --strategy all
"""
from __future__ import annotations
import argparse
import json
import re
from collections import defaultdict
from difflib import get_close_matches
from pathlib import Path

INTENTS_50 = sorted(json.loads(Path("/Users/jean-patricksmith/digital/kingly/apps/production/naac/poc/deberta_intent/checkpoints-base/label_mapping.json").read_text())["intent2id"].keys())
INTENT_SET = set(INTENTS_50)


def normalize(d: dict, ci: bool) -> dict:
    if not isinstance(d, dict):
        return {}
    if ci:
        return {str(k).upper(): str(v).strip().lower() for k, v in d.items() if v is not None and str(v) != ""}
    return {str(k): str(v) for k, v in d.items() if v is not None}


def slot_match(pred_slots: dict, gold_slots: dict, strategy: str) -> bool:
    if strategy == "exact":
        return pred_slots == gold_slots
    if strategy == "ci":
        return normalize(pred_slots, ci=True) == normalize(gold_slots, ci=True)
    if strategy == "subset":
        return all(pred_slots.get(k) == v for k, v in gold_slots.items())
    if strategy == "ci-subset":
        p = normalize(pred_slots, ci=True)
        g = normalize(gold_slots, ci=True)
        return all(p.get(k) == v for k, v in g.items())
    raise ValueError(f"unknown strategy: {strategy}")


def snap_intent(intent_str: str | None) -> str | None:
    if not intent_str or intent_str in INTENT_SET:
        return intent_str
    matches = get_close_matches(intent_str, INTENTS_50, n=1, cutoff=0.4)
    return matches[0] if matches else intent_str


def score(rows: list[dict], strategy: str, intent_snap: bool = False) -> dict:
    schema_ok = intent_em = slots_em = 0
    invented_intents = 0
    by_intent = defaultdict(lambda: [0, 0, 0])
    by_scenario = defaultdict(lambda: [0, 0, 0])
    failures = []

    for r in rows:
        gold_intent = r["gold_intent"]
        gold_params = r["gold_parameters"] or {}
        parsed = r["parsed"]

        if parsed:
            schema_ok += 1

        pred_intent = (parsed or {}).get("intent")
        if pred_intent and pred_intent not in INTENT_SET:
            invented_intents += 1
            if intent_snap:
                pred_intent = snap_intent(pred_intent)

        is_intent = bool(parsed and pred_intent == gold_intent)
        is_slots = bool(parsed and slot_match((parsed or {}).get("slots") or {}, gold_params, strategy))

        intent_em += int(is_intent); slots_em += int(is_slots)
        by_intent[gold_intent][0] += 1
        by_intent[gold_intent][1] += int(is_intent)
        by_intent[gold_intent][2] += int(is_slots)
        by_scenario[r["scenario_id"]][0] += 1
        by_scenario[r["scenario_id"]][1] += int(is_intent)
        by_scenario[r["scenario_id"]][2] += int(is_slots)

        if not (is_intent and is_slots) and len(failures) < 20:
            failures.append({
                "scenario": r["scenario_id"],
                "turn": r["turn_index"],
                "text": r["text"][:120],
                "gold": {"intent": gold_intent, "slots": gold_params},
                "pred": parsed,
                "intent_ok": is_intent, "slots_ok": is_slots,
            })

    n = len(rows)
    return {
        "n": n, "strategy": strategy, "intent_snap": intent_snap,
        "schema_valid": schema_ok, "intent_em": intent_em, "slots_em": slots_em,
        "schema_pct": schema_ok / max(n, 1), "intent_pct": intent_em / max(n, 1), "slots_pct": slots_em / max(n, 1),
        "invented_intents": invented_intents,
        "by_intent": dict(by_intent), "by_scenario": dict(by_scenario),
        "failures": failures,
    }


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("dump", help="JSON dump from eval_baseline.py")
    p.add_argument("--strategy", choices=["exact", "ci", "subset", "ci-subset", "all"], default="ci-subset")
    p.add_argument("--snap-intent", action="store_true", help="post-hoc snap pred intent to nearest 50-enum")
    p.add_argument("--by", choices=["intent", "scenario", "both", "none"], default="intent")
    p.add_argument("--top", type=int, default=20)
    p.add_argument("--show-fails", action="store_true")
    p.add_argument("--out", default=None)
    args = p.parse_args()

    dump = json.loads(Path(args.dump).read_text())
    rows = dump["rows"]
    print(f"loaded {len(rows)} rows from {args.dump}")
    print(f"adapter: {dump.get('adapter')}\n")

    strategies = ["exact", "ci", "subset", "ci-subset"] if args.strategy == "all" else [args.strategy]
    results = []
    for strat in strategies:
        r = score(rows, strat, intent_snap=args.snap_intent)
        results.append(r)
        print(f"=== strategy={strat}{'+snap' if args.snap_intent else ''} ===")
        print(f"  schema_valid: {r['schema_valid']}/{r['n']} = {r['schema_pct']:.1%}")
        print(f"  intent_em:    {r['intent_em']}/{r['n']} = {r['intent_pct']:.1%}    (V32 baseline: 87.4%)")
        print(f"  slots_em:     {r['slots_em']}/{r['n']} = {r['slots_pct']:.1%}    (V32 baseline: 80.6%)")
        print(f"  invented_intents: {r['invented_intents']}")

        if args.by in ("intent", "both"):
            print(f"  --- per-intent (top {args.top}) ---")
            for intent, (n_i, i_em, s_em) in sorted(r["by_intent"].items(), key=lambda kv: -kv[1][0])[:args.top]:
                print(f"    {intent:30s} n={n_i:4d}  intent={i_em}/{n_i} ({i_em/n_i:.0%})  slots={s_em}/{n_i} ({s_em/n_i:.0%})")
        if args.by in ("scenario", "both"):
            print(f"  --- per-scenario (top {args.top}) ---")
            for scn, (n_s, i_em, s_em) in sorted(r["by_scenario"].items(), key=lambda kv: -kv[1][0])[:args.top]:
                print(f"    {scn:30s} n={n_s:4d}  intent={i_em}/{n_s} ({i_em/n_s:.0%})  slots={s_em}/{n_s} ({s_em/n_s:.0%})")
        if args.show_fails:
            print(f"  --- first {len(r['failures'])} failures ---")
            for f in r["failures"]:
                print(f"    [{f['scenario']}/{f['turn']}] intent_ok={f['intent_ok']} slots_ok={f['slots_ok']}")
                print(f"      text: {f['text']}")
                print(f"      gold: {f['gold']}")
                print(f"      pred: {f['pred']}")
        print()

    if args.out:
        Path(args.out).write_text(json.dumps({"adapter": dump.get("adapter"), "results": results}, indent=2))
        print(f"saved scoring report to {args.out}")


if __name__ == "__main__":
    main()