"""Parity fixtures: inputs + expected outputs of a released model, so any port (Go, Rust, C) is proven bit-exact.

A fixture is a JSON file next to the model:
  flows:  {"model": "<npz>", "names": [...], "rows": [[float...], ...], "raw": [[float...], ...]}
  http:   {"model": "<npz>", "texts": [...], "deepest": [...], "buckets": [[idx...]...], "logits": [[...]...]}
Tolerance for a port: raw scores within 1e-4 (float32 summation order), bucket ids exactly equal.
"""
from __future__ import annotations

import json

import numpy as np


def export_trees(model, x: np.ndarray, path: str, model_file: str) -> None:
    x = np.asarray(x, np.float32)
    raw = model.raw(x)
    with open(path, "w", encoding="utf-8") as f:
        json.dump({"kind": "oblivious", "model": model_file, "names": model.names, "classes": model.classes,
                   "rows": x.tolist(), "raw": raw.tolist(), "tol": 1e-4}, f)


def export_http(model, texts: list[str], deepest, path: str, model_file: str) -> None:
    from dibaguard.hashing import ngram_buckets

    idx, doc = ngram_buckets(texts, model.buckets, model.orders, model.max_bytes)
    buckets = [sorted(idx[doc == i].tolist()) for i in range(len(texts))]
    z = model.logits(texts, deepest)
    with open(path, "w", encoding="utf-8") as f:
        json.dump({"kind": "http_tier1", "model": model_file, "texts": texts, "deepest": [int(d) for d in deepest],
                   "buckets": buckets, "logits": z.tolist(), "tol": 1e-3}, f, ensure_ascii=False)


def npz_to_flat(npz_path: str, out_path: str) -> dict:
    """Write the model as raw little-endian arrays + a JSON header (for ports without an npz reader)."""
    z = np.load(npz_path)
    header, offset = {}, 0
    with open(out_path, "wb") as f:
        for k in z.files:
            a = np.ascontiguousarray(z[k])
            if a.dtype == np.uint8 and k == "info":
                header["info"] = json.loads(bytes(a).decode())
                continue
            if a.dtype.byteorder == ">":
                a = a.byteswap().view(a.dtype.newbyteorder("<"))
            f.write(a.tobytes())
            header[k] = {"dtype": a.dtype.str, "shape": list(a.shape), "offset": offset, "nbytes": a.nbytes}
            offset += a.nbytes
    with open(out_path + ".json", "w", encoding="utf-8") as f:
        json.dump(header, f)
    return header