"""Check release/vela3_coreml.py (fresh import, release folder only) against the fp32 torch release code:
decisions on typed-decisions TEST (choice + noul, data/ref.jsonl), PII spans on the 25 hand-written demo sections, and
unsupported-claim spans on a few context / answer pairs. Also times each Core ML function on the GPU."""
import json, sys, time
import numpy as np
import pyarrow.parquet as pq

sys.path.insert(0, "release")
from vela3_coreml import Vela3CoreML  # noqa: E402

m = Vela3CoreML("release")
# 1) decisions
ref = {x["id"]: x for x in map(json.loads, open("data/ref.jsonl"))}
nq = flips = 0; dmax = 0.0; margins = []
for r in pq.read_table("data/typed/all/test-00000-of-00001.parquet").to_pylist():
    R = ref[r["id"]]; state, qs = json.loads(r["state"]), json.loads(r["questions"])
    qs = {k: q for k, q in qs.items() if q["type"] in ("choice", "noul")}
    for k, an in m.decide(state, qs).items():
        rb = R["answers"][k]; nq += 1
        if an["type"] == "noul":
            dmax = max(dmax, abs(an["noul"] - rb["noul"])); flips += (an["noul"] >= 0.5) != (rb["noul"] >= 0.5)
        else:
            pr = rb["probabilities"]; dmax = max(dmax, max(abs(an["probabilities"][o] - pr[o]) for o in pr)); flips += an["choice"] != rb["choice"]
print(f"decisions: {nq} questions, {flips} flips vs fp32 torch, max |dp| {dmax:.4f}")

# 2) + 3) spans, torch reference computed fresh
sys.path.insert(0, "demo")
from doc import RAW, TAG  # noqa: E402
sections = [TAG.sub(r"\2", l) for l in RAW.strip().splitlines() if l.strip() and not l.startswith("#")]
sys.path.insert(0, "."); from span_engine import PIIEngine  # noqa: E402
t = PIIEngine("torch").v
key = lambda ss: {(s["label"], s["start"], s["end"]) for s in ss}
pii_diff = [i for i, s in enumerate(sections) if key(m.pii(s)) != key(t.span_answer(s, {"type": "span", "criteria": json.load(open("src/vela3_config.json"))["span"]["router_label_sets"]["pii"]})[0]["spans"])]
n_spans = sum(len(m.pii(s)) for s in sections)
print(f"pii: {len(sections)} sections, {n_spans} spans, sections whose span set differs from torch: {pii_diff}")
pairs = [("Adults may take up to 4 g of paracetamol in 24 hours.", "Adults may take up to 6 g of paracetamol in 24 hours."),
         ("The Eiffel Tower is 330 m tall and was finished in 1889.", "The Eiffel Tower, finished in 1899, is 330 m tall."),
         ("Refunds are issued within 5 business days to the original card.", "Refunds are issued within 5 business days to the original card."),
         ("The meeting moved from Tuesday to Thursday at 3 pm in room B.", "The meeting is now on Thursday at 4 pm in room C.")]
for c, a in pairs:
    got = [s["text"] for s in m.unsupported(c, a)]
    exp = [s["text"] for s in t.span_answer({"context": c, "answer": a}, {"type": "span", "criteria": ["unsupported"]})[0]["spans"]]
    print(f"unsupported: coreml {got} | torch {exp} | {'same' if got == exp else 'DIFF'}")

# 4) latency per Core ML function (GPU), median of 20 after warm-up
for name, enc in (("decision", m.dec_enc), ("span_sets", m.span_enc)):
    row = {}
    for L in enc.lengths:
        ids = list(np.random.randint(5, 49000, size=L - 4))
        for _ in range(3): enc(ids)
        enc.last_ms.clear(); [enc(ids) for _ in range(20)]
        row[f"L{L}"] = round(float(np.median([ms for _, ms in enc.last_ms])), 1)
    print(f"latency {name} (ms, CPU_AND_GPU):", row)