"""Swift parity fixtures for FluidUse's Vela3Manager, from vela3_coreml.py (the same Core ML encoders on the GPU, the
release's own rendering, thresholds and span decoder):
  tokens     ids + offsets of varied texts (Unicode, emoji, CJK, newlines, URLs, special tokens)
  decisions  typed-decisions TEST choice / yes-no questions: rendered state, encoder sequence, probabilities
  spans      PII and unsupported-claim calls: encoder sequence, spans, threshold
python fixtures.py <release dir> <typed-decisions parquet> <out.json> [n_rows]"""
import json, sys
sys.path.insert(0, sys.argv[1])
import pyarrow.parquet as pq
from vela3_coreml import Vela3CoreML, _SpanRun
from ub_d3_common import to_row

m = Vela3CoreML(sys.argv[1])
tok = m.decision.tok.t
TEXTS = [
    "Hi, my name is Marisol Okafor and I'm writing about a double charge.", " hello  world", "Héllo 😀 日本語 x",
    "a\n\n  b\tc", "cost: $4,111.00!!", "<|endoftext|>hi there", "Call +1 (415) 555-0193 or mail marisol.okafor@example.com.",
    "See https://example.com/a?b=c&d=e, then www.example.org.", "Ünïcödé façade — naïve café; Zürich, São Paulo.",
    "東京都千代田区1-1、山田太郎 (090-1234-5678)", "Привет, меня зовут Иван Петров.", "مرحبا اسمي أحمد", "tab\tseparated\tvalues\r\nCRLF",
    "IBAN GB82 WEST 1234 5698 7654 32, SSN 912-44-0387", "   leading and trailing   ", "emoji 👩‍👩‍👧 family 🇨🇦 flag", "",
    "hello <|mask|> world", "x \t<|mask|>", "\n\n<|mask|> y", "a<|mask|>  b", "a" * 3000, "Ab" * 1500 + " end",
]
tokens = []
for s in TEXTS:
    e = tok.encode(s, add_special_tokens=False)
    tokens.append(dict(text=s, ids=e.ids, offsets=[list(o) for o in e.offsets]))

decisions = []
n_rows = int(sys.argv[4]) if len(sys.argv) > 4 else 40
for r in pq.read_table(sys.argv[2]).to_pylist()[:n_rows]:
    state, qs = json.loads(r["state"]), json.loads(r["questions"])
    for k, q in qs.items():
        if q["type"] not in ("choice", "noul"):
            continue
        row = to_row("q", state, q)
        qq, st = row["questions"][0], row["parts"][0]["text"]
        rec = m.decision.encode(state, q)
        ans = m.decide(state, {k: q})[k]
        if q["type"] == "noul":
            names, probs = ["no", "yes"], [1 - ans["noul"], ans["noul"]]
        else:
            names, probs = list(ans["probabilities"]), list(ans["probabilities"].values())
        decisions.append(dict(state=st, id=k, type=q["type"], instructions=qq["text"],
                              options=[[n, d] for n, d in qq["options"].items()], ids=rec["ids"], names=names, probabilities=probs))

captured = []
orig = _SpanRun.__call__
def cap(self, x, mm, *a):
    for b in range(x.shape[0]):
        captured.append([int(v) for v in x[b, :int(mm[b].sum())]])
    return orig(self, x, mm, *a)
_SpanRun.__call__ = cap
span_cases = []
PII = [
    "Hi, I'm Dana Whitfield (dana@example.com, 312-555-0110), 4410 Elm Ridge Drive, Dayton, OH 45429.",
    "Patient: Henrik Lindqvist, 58 years old, date of birth 11/09/1967. Referred by Dr. Amara Nwosu at Lakeside Family Medicine. "
    "Social security number on the insurance form: 912-44-0387.",
    "Wire 18,400 euros to GB82 WEST 1234 5698 7654 32 on February 19; the card on file is 4111 1111 1111 1111.",
    "Login from 203.0.113.47 for admin@intranet.example; site priyaraman.dev; Ms. Raman is Portuguese and 34.",
    "東京都千代田区の山田太郎さん（090-1234-5678）に連絡してください。", "Nothing personal here, just the weather report for tomorrow.",
    "Ünïcödé: Zoë Ågren lives at Kungsgatan 12, 111 43 Stockholm, phone +46 8 555 012 34.",
    "Contact: jose.garcia@example.org — José García, Calle Mayor 5, 28013 Madrid 🇪🇸",
]
# long inputs: windowed PII (2 and 4 windows of the 1,024-token pass)
PII += [" ".join(f"Record {i}: " + PII[i % 4] for i in range(n)) for n in (40, 90)]
for s in PII:
    captured.clear(); sp = m.v.span_answer(s, {"type": "span", "criteria": json.load(open(sys.argv[1] + "/vela3_config.json"))["span"]["router_label_sets"]["pii"]})[0]
    span_cases.append(dict(kind="pii", text=s, ids=list(captured), threshold=sp["threshold"],
                           spans=[dict(label=x["label"], start=x["start"], end=x["end"], probability=x["probability"]) for x in sp["spans"]]))
HALU = [("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.")]
# long inputs: a long context is cut (answer intact); a long answer is also read in windows
LONG_CTX = " ".join(f"Note {i}: the warehouse in Lyon shipped {i * 7} boxes on day {i}." for i in range(120))
HALU += [(LONG_CTX, "The warehouse in Lyon shipped 70 boxes on day 10 and 999 boxes on day 11."),
         (LONG_CTX, " ".join(f"On day {i} it shipped {i * 7 + (1 if i % 9 == 0 else 0)} boxes." for i in range(150)))]
for c, a in HALU:
    captured.clear(); sp = m.v.span_answer({"context": c, "answer": a}, {"type": "span", "criteria": ["unsupported"]})[0]
    span_cases.append(dict(kind="unsupported", context=c, answer=a, ids=list(captured), threshold=sp["threshold"],
                           spans=[dict(label=x["label"], start=x["start"], end=x["end"], probability=x["probability"]) for x in sp["spans"]]))
json.dump(dict(tokens=tokens, decisions=decisions, spans=span_cases), open(sys.argv[3], "w"), ensure_ascii=False)
print(len(tokens), "token texts,", len(decisions), "decisions,", len(span_cases), "span calls ->", sys.argv[3])