"""The NNC subset of extensions/prior-carriage.md §12, measured against an open implementation
of ISO/IEC 15938-17 (§12.4). Not part of the corpus and not run in CI: it runs that
implementation's own code, which this repository does not carry.

For each prior of live-conformance/prior-forms/, it requires that the other implementation's
decoder reads our NNC file to the integers our safetensors file holds, that its encoder,
given those integers, writes data units byte-identical to ours, and that our reader
(prior_forms_verify.py) reads its bitstream to the same integers.

    python3 nnc_interop_check.py <its-source-tree> <wai/live-conformance> <wai/live-conformance/prior-forms>

with the other implementation's compiled arithmetic coder importable (PYTHONPATH). The exact
setup used is in the corpus README.
"""
import json, os, struct, sys
import numpy as np
SRC, OURS, CORPUS = sys.argv[1], sys.argv[2], sys.argv[3]
sys.path.insert(0, SRC)
from nnc_core import coder  # noqa: E402  (imports their compiled arithmetic coder)
from nnc_core.coder import baseline  # noqa: E402

# Harness glue: their decoder keeps dq_flag only when the MPS signals quantisation, which an
# integer-only bitstream does not; give it the dict it writes 0 into. The bits are untouched.
_orig = baseline.decode


def _decode(decoder, approx_data, *a):
    approx_data.setdefault("dq_flag", {})
    return _orig(decoder, approx_data, *a)


baseline.decode = _decode
sys.path.insert(0, OURS)
import prior_forms_verify as V  # noqa: E402

NAMES = ["generic", "title", "seeded", "hyper_g_s", "hyper_h_s", "world_t", "world_g_s"]


def units(b):
    out, at = [], 0
    while at < len(b):
        size = (int.from_bytes(b[at:at + 2], "big") & 0x7FFF) if b[at] & 0x80 == 0 else (int.from_bytes(b[at:at + 4], "big") & 0x7FFFFFFF)
        u = b[at:at + size]
        out.append(u)
        at += size
    return out


def ndu_name(u):
    off = 4 if u[0] & 0x80 == 0 else 6  # size field, then 1 header byte, then 1 byte of NDU flags
    z = u.index(0, off)
    return u[off:z].decode()


report = {}
all_ok = True
for n in NAMES:
    ours_nnc = open(os.path.join(CORPUS, "files", n + ".nnc"), "rb").read()
    st = V.st_read(open(os.path.join(CORPUS, "files", n + ".safetensors"), "rb").read())
    want = {k: np.array(V.ints_of(dt, data), dtype=np.int64).reshape(shape) for k, (dt, shape, data) in st.t.items()}
    # B: their decoder on our bitstream
    mi = {'parameter_type': {}, 'parameter_dimensions': {}, 'parameter_index': {}, 'block_identifier': {},
          'topology_storage_format': None, 'topology_compression_format': None,
          'performance_maps': {"mps": {}, "lps": {}},
          'performance_map_flags': {k: {} for k in ["mps_sparsification_flag", "lps_sparsification_flag", "mps_pruning_flag", "lps_pruning_flag",
                                                    "mps_unification_flag", "lps_unification_flag", "mps_decomposition_performance_map_flag", "lps_decomposition_performance_map_flag"]}}
    dec = coder.decode(bytearray(ours_nnc), mi, {})
    got = dec["parameters"]
    b_ok = set(got) == set(want) and all(np.array_equal(np.asarray(got[k], dtype=np.int64), want[k]) for k in want)
    # C: their encoder on the same integers, unit by unit against ours
    names = sorted(want, key=lambda x: x.encode())
    params = {k: want[k].astype(np.int32) for k in names}
    model_info = {"parameter_index": {k: i for i, k in enumerate(names)}, "parameter_type": {k: "unspecified" for k in names},
                  "parameter_dimensions": {k: list(params[k].shape) for k in names}, "block_identifier": {},
                  "topology_storage_format": None}
    approx = {"parameters": params, "approx_method": {k: "skip" for k in names}, "dq_flag": {k: 0 for k in names},
              "scan_order": {k: 0 for k in names if params[k].ndim > 1}, "compressed_parameter_types": {}}
    theirs = bytes(coder.encode({"cabac_unary_length_minus1": 10, "param_opt_flag": 0}, model_info, approx))
    ours_ndus = {ndu_name(u): u for u in units(ours_nnc) if (u[2 if u[0] & 0x80 == 0 else 4] >> 2) == 5}
    their_ndus = {ndu_name(u): u for u in units(theirs) if (u[2 if u[0] & 0x80 == 0 else 4] >> 2) == 5}
    same = sorted(k for k in ours_ndus if their_ndus.get(k) == ours_ndus[k])
    c_ok = set(ours_ndus) == set(their_ndus) and len(same) == len(ours_ndus)
    # D: our reader on their bitstream
    try:
        s = V.nnc_read(theirs)
        d_ok = set(s.t) == set(want) and all(V.ints_of("I32", s.t[k][2]) == list(want[k].reshape(-1)) for k in want)
        d_err = None
    except V.Refuse as r:
        d_ok, d_err = False, r.code
    report[n] = {"their_decoder_reads_ours": b_ok, "data_units_identical": f"{len(same)} of {len(ours_ndus)}", "our_reader_reads_theirs": d_ok,
                 "our_reader_error": d_err, "elements": int(sum(v.size for v in want.values())), "their_bytes": len(theirs), "our_bytes": len(ours_nnc)}
    all_ok = all_ok and b_ok and c_ok and d_ok
print(json.dumps(report, indent=1))
print("ALL_OK" if all_ok else "MISMATCH")