"""Eval H1 adapters on the canonical 929-row held-out eval_set.jsonl.

This is the FROZEN production target — never touched during training.
"""
import json
import sys
from pathlib import Path
from mlx_lm import load, generate

INTENTS = sorted(json.loads(open("/Users/jean-patricksmith/digital/kingly/apps/production/naac/poc/deberta_intent/checkpoints-base/label_mapping.json").read())["intent2id"].keys())
SLOTS = ["ALTIMETER_SETTING", "ALTITUDE", "APPROACH_TYPE", "CALL_SIGN", "CLOCK_POSITION",
         "DIRECTION", "DISTANCE", "FACILITY", "FIX", "FREQUENCY", "HEADING", "PATTERN_LEG",
         "ROUTE", "RUNWAY", "SPEED", "TAXIWAY", "TIME", "TRANSPONDER_CODE", "TURN_DIRECTION"]

SYSTEM_PROMPT = (
    "You are an ATC parser. Parse the air traffic control transmission into a JSON object.\n"
    "OUTPUT FORMAT: {\"intent\": <one-of-enum>, \"slots\": {<SLOT_TYPE>: <value>}}\n"
    f"\nINTENT MUST BE ONE OF: {', '.join(INTENTS)}\n"
    f"\nSLOT TYPES (UPPERCASE): {', '.join(SLOTS)}\n"
    "\nOutput ONLY the JSON object. No prose, no code fences."
)

import re

ADAPTER = sys.argv[1] if len(sys.argv) > 1 else "poc/llm-finetune/training/h1_adapter_v3"
SRC = Path(sys.argv[2] if len(sys.argv) > 2 else "poc/deberta_intent/data/eval_set.jsonl")
BASE = sys.argv[3] if len(sys.argv) > 3 else (
    "mlx-community/Qwen3-4B-Instruct-2507-4bit" if "q3b" in ADAPTER.lower()
    else "mlx-community/Qwen2.5-32B-Instruct-4bit"
)


def parse(text):
    text = text.strip()
    text = re.sub(r'<think>.*?</think>\s*', '', text, flags=re.DOTALL).strip()
    if text.startswith("```"):
        text = text.split("```")[1]
        if text.startswith("json"):
            text = text[4:]
    try:
        return json.loads(text.strip())
    except Exception:
        return None


def main():
    print(f"loading {BASE} + {ADAPTER}")
    model, tok = load(BASE, adapter_path=ADAPTER)
    rows = [json.loads(l) for l in SRC.read_text().splitlines() if l.strip()]
    print(f"held-out rows: {len(rows)}")

    schema_ok = intent_em = slots_em = 0
    by_intent = {}
    for i, row in enumerate(rows):
        text = row.get("text", "")
        gold_intent = row.get("intent", "")
        gold_slots = {k: v.get("value", "") for k, v in (row.get("slots") or {}).items() if isinstance(v, dict)}
        msgs = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": text},
        ]
        prompt = tok.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False)
        out = generate(model, tok, prompt=prompt, max_tokens=200, verbose=False)
        parsed = parse(out)
        is_schema = bool(parsed and "intent" in parsed and "slots" in parsed)
        is_intent = bool(parsed and parsed.get("intent") == gold_intent)
        is_slots = bool(parsed and parsed.get("slots") == gold_slots)
        schema_ok += int(is_schema); intent_em += int(is_intent); slots_em += int(is_slots)
        by_intent.setdefault(gold_intent, [0, 0, 0])  # n, intent, slots
        by_intent[gold_intent][0] += 1
        by_intent[gold_intent][1] += int(is_intent)
        by_intent[gold_intent][2] += int(is_slots)
        if (i + 1) % 100 == 0:
            print(f"  ...{i+1}/{len(rows)}: intent={intent_em/(i+1):.1%}, slots={slots_em/(i+1):.1%}")

    n = len(rows)
    print(f"\n=== HOLD-OUT (eval_set.jsonl, 929 rows, NEVER trained on) ===")
    print(f"adapter:      {ADAPTER}")
    print(f"schema_valid: {schema_ok}/{n} = {schema_ok/n:.1%}")
    print(f"intent_em:    {intent_em}/{n} = {intent_em/n:.1%}")
    print(f"slots_em:     {slots_em}/{n} = {slots_em/n:.1%}")
    print(f"\n=== per-intent (top 15 by count) ===")
    for intent, (n_i, i_em, s_em) in sorted(by_intent.items(), key=lambda kv: -kv[1][0])[:15]:
        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 __name__ == "__main__":
    main()