#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Dump the served model's debug outputs (class map, fp32 margins and logits) for structural A/B bit checks.

    bin/devrun -t 900 -- env SCENESEG_EPI=1 python -u code/scripts/diag/dump_outputs.py --out a.npz [--input img ...]
    python code/scripts/diag/dump_outputs.py --compare a.npz b.npz

Every knob comes from the environment (``SCENESEG_*``), so one process per configuration. Inputs: the shipped sample
and, when present, the PandaSet 019 frame of the workspace. The served (traced) class map and the debug trace's
margins / logits are stored per input; ``--compare`` reports bit-identity per array.
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

import numpy as np

CODE = Path(__file__).resolve().parents[2]
SAMPLE = CODE / "tt_sceneseg" / "samples" / "highway_normal_1.png"
PANDASET = (CODE.parents[2] / "research" / "sceneseg" / "public_data" / "shippable" / "pandaset_019_f40_front_camera"
            / "pandaset_019_f40_front_camera.jpg")


def compare(a: str, b: str) -> int:
    x, y = np.load(a), np.load(b)
    ok = True
    for k in sorted(set(x) | set(y)):
        if k not in x or k not in y:
            print(f"MISSING {k}")
            ok = False
            continue
        u, v = x[k], y[k]
        same = u.shape == v.shape and np.array_equal(u.view(np.uint8), v.view(np.uint8))
        extra = ""
        if not same and u.shape == v.shape:
            if u.dtype.kind == "f":
                d = np.abs(u.astype(np.float64) - v.astype(np.float64))
                extra = f" (max abs diff {d.max():.3g}, {int((u != v).sum())} of {u.size} differ)"
            else:
                extra = f" ({int((u != v).sum())} of {u.size} differ, agreement {(u == v).mean():.6f})"
        ok &= same
        print(f"{'SAME' if same else 'DIFF'} {k} {u.dtype} {u.shape}{extra}")
    print("ALL IDENTICAL" if ok else "NOT IDENTICAL")
    return 0 if ok else 1


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out")
    ap.add_argument("--input", action="append")
    ap.add_argument("--compare", nargs=2)
    a = ap.parse_args()
    if a.compare:
        return compare(*a.compare)
    from tt_sceneseg import SceneSeg

    inputs = [Path(p) for p in (a.input or [str(SAMPLE)] + ([str(PANDASET)] if PANDASET.is_file() else []))]
    res = {}
    with SceneSeg.from_pretrained() as m:
        m.runner.add_debug_variant()
        m.runner.capture()
        print("info:", m.extra_info().get("graph"), flush=True)
        for p in inputs:
            canvas = m._prepare(str(p))["prepared"].device_input()
            served = m.runner(canvas)
            dbg = m.runner.debug(canvas)
            res[f"{p.stem}.class_map"] = served["class_map"]
            for k in ("class_map", "margins", "logits"):
                res[f"{p.stem}.debug.{k}"] = dbg[k]
            print(f"{p.name}: class counts {np.bincount(served['class_map'].ravel(), minlength=3).tolist()}", flush=True)
    np.savez_compressed(a.out, **res)
    print("saved", a.out)
    return 0


if __name__ == "__main__":
    sys.exit(main())