"""Serving parity and repeat-run determinism (docs/classifier-plan.md §5.5 item 6, §6.4). Pure Python.

- ``serving_parity``: node12 endpoint logits vs pod reference logits on the same posts (the 1,000 test posts):
  per-field gated label flip rate and the largest calibrated gated probability difference; passes with <= 1% flips
  per field and a max difference <= 0.05 (twitter HANDOFF rule).
- ``determinism``: two runs of the same checkpoint on the same posts: max probability difference <= 1e-3.
Run as ``python -m training.parity serving|determinism A.jsonl B.jsonl --taxonomy DIR [--temperatures T.json]``.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Mapping

from .calibrate import gated, max_prob_diff
from .taxonomy import Taxonomy, argmax

FLIPS_MAX = 0.01
PROB_MAX = 0.05
DETERMINISM_MAX = 1e-3


def _paired(a_rows, b_rows) -> list[tuple[dict, dict]]:
    a = {r["id"]: r for r in a_rows}
    b = {r["id"]: r for r in b_rows}
    if len(a) != len(list(a_rows)) or len(b) != len(list(b_rows)):
        raise ValueError("duplicate ids")
    if set(a) != set(b):
        raise ValueError(f"id sets differ ({len(set(a) ^ set(b))} ids in only one side)")
    return [(a[i], b[i]) for i in sorted(a)]


def compare(a_rows, b_rows, tax: Taxonomy, temps: Mapping[str, float] | None = None) -> dict:
    pairs = _paired(list(a_rows), list(b_rows))
    flips = {f: 0 for f in tax.names}
    worst = 0.0
    for a, b in pairs:
        pa, pb = gated(a["logits"], tax, temps), gated(b["logits"], tax, temps)
        worst = max(worst, max_prob_diff(pa, pb))
        la, lb = argmax(pa), argmax(pb)
        for f in tax.names:
            flips[f] += la[f] != lb[f]
    n = len(pairs)
    return {"n": n, "flips": flips, "flip_rate": {f: (v / n if n else None) for f, v in flips.items()},
            "max_prob_diff": worst}


def serving_parity(reference_rows, served_rows, tax: Taxonomy, temps: Mapping[str, float] | None = None) -> dict:
    out = compare(reference_rows, served_rows, tax, temps)
    out["pass"] = bool(out["n"]) and all(v <= FLIPS_MAX for v in out["flip_rate"].values()) \
        and out["max_prob_diff"] <= PROB_MAX
    out["thresholds"] = {"flip_rate": FLIPS_MAX, "max_prob_diff": PROB_MAX}
    return out


def determinism(run_a_rows, run_b_rows, tax: Taxonomy, temps: Mapping[str, float] | None = None) -> dict:
    out = compare(run_a_rows, run_b_rows, tax, temps)
    out["argmax_identical"] = all(v == 0 for v in out["flips"].values())
    out["pass"] = bool(out["n"]) and out["max_prob_diff"] <= DETERMINISM_MAX
    out["threshold"] = DETERMINISM_MAX
    return out


def main(argv=None) -> int:
    from .data import load_jsonl

    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("kind", choices=["serving", "determinism"])
    p.add_argument("a", help="Reference (or first run) predictions JSONL")
    p.add_argument("b", help="Served (or second run) predictions JSONL")
    p.add_argument("--taxonomy", required=True, help="taxonomy.json or a directory holding it")
    p.add_argument("--temperatures")
    p.add_argument("--out")
    args = p.parse_args(argv)
    tax = Taxonomy.load(args.taxonomy)
    temps = json.loads(Path(args.temperatures).read_text()) if args.temperatures else None
    fn = serving_parity if args.kind == "serving" else determinism
    res = fn(load_jsonl(args.a), load_jsonl(args.b), tax, temps)
    text = json.dumps(res, indent=1)
    if args.out:
        Path(args.out).write_text(text + "\n")
    print(text)
    return 0 if res["pass"] else 1


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