#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Write the CPU fp32 goldens of the TT port with the bundle's own reference (``tt_bevformer.reference``, port form:
host geometry, rebatched SCA, rotation by gather) and cross-check them against the research goldens.

Large goldens (internal validation data, never shipped) go to ``research/bevformer/goldens/`` (``--out``):

==============================================  ==================================================================
``pandaset_019/f30.npz``, ``f35.npz``, ``f40``  the shipped sample as a stateful sequence (reset at f30): outputs,
                                                ``bev_embed``, inputs, host geometry; f35 (history + ego motion)
                                                also every per-module tap and the network input ``image``
``sample0/f1.npz``, ``f2``, ``f3``              the bevdet_vendor nuScenes key-frame (PLAN.md 2.11 detection
                                                gates on frames 1-2; frame 3 = +2 m / +3 deg synthetic motion);
                                                f1 also the image, backbone.layer4, neck and transformer taps
``nuscenes_scene-0103/kf01.npz``                mini_val, teacher-forced with the public golden kf00 bev_embed;
                                                image + transformer taps (in-domain, real 4.3 m / -1.0 deg motion)
==============================================  ==================================================================

Taps above 20 M elements (backbone.layer1-3) are stored as float16 (their PCC against float32 changes by < 1e-7);
everything else float32; ``__meta__`` records the dtypes, the generator, the bundle commit, the weights digest and
the PCCs against the research goldens. Small goldens (PandaSet-derived, CC BY 4.0) go to the bundle:
``staging_samples_pandaset/tests/goldens/pandaset_019_f30-40.json`` (detections, counts, PCCs; git-ignored until the PandaSet sample ships),
``pandaset_019_f30-40_last_layer.npz`` (last-layer logits / boxes) and the smoke test's stored reference
``staging_samples_pandaset/samples/pandaset_019_f30.reference.json`` (``to_dict()`` of a fresh-stream request of frame 30).

    cd bundles/bevformer-p150 && OMP_NUM_THREADS=4 PYTHONPATH=$PWD/code \\
        /home/ubuntu/experiments/tt-models/tools/research-venv/bin/python code/scripts/make_goldens.py --threads 4
"""
from __future__ import annotations

import argparse
import fnmatch
import hashlib
import json
import subprocess
import sys
import time
from pathlib import Path

import numpy as np

HERE = Path(__file__).resolve().parent
CODE = HERE.parent
sys.path.insert(0, str(CODE))

from tt_bevformer import __version__  # noqa: E402
from tt_bevformer.host import postprocess as hpost  # noqa: E402
from tt_bevformer.host.inputs import load_frame  # noqa: E402
from tt_bevformer.reference.pipeline import BEVFormerReference  # noqa: E402
from tt_bevformer.tests import paths  # noqa: E402
from tt_bevformer.ttaw.golden import save_goldens  # noqa: E402
from tt_bevformer.ttaw.metrics import pcc  # noqa: E402

FP16_ABOVE = 20_000_000
OUTS = ("bev_embed", "outputs_classes", "outputs_coords")
GEOMETRY = ("ref_cam", "sca_rows", "sca_counts", "sca_gather", "sca_scale", "shift", "rot_index")
ALL_TAPS = ["backbone.layer1", "backbone.layer2", "backbone.layer3", "backbone.layer4", "neck", "enc.*", "ref_cam",
            "encoder.*", "decoder.*", "head.*"]
TRANSFORMER_TAPS = ["backbone.layer4", "neck", "enc.*", "ref_cam", "encoder.layer0.*", "encoder.layer0",
                    "encoder.layer1", "encoder.layer2", "decoder.*", "head.*"]


def _git_head(path: Path) -> str:
    """``<short sha>`` of the bundle, with ``+dirty`` when the working tree differs from it."""
    try:
        head = subprocess.run(["git", "-C", str(path), "rev-parse", "--short", "HEAD"], capture_output=True,
                              text=True, check=True).stdout.strip()
        dirty = subprocess.run(["git", "-C", str(path), "status", "--porcelain", "--untracked-files=no"],
                               capture_output=True, text=True, check=True).stdout.strip()
        return head + ("+dirty" if dirty else "")
    except (OSError, subprocess.CalledProcessError):
        return "unknown"


def _select(taps, patterns):
    return {k: v for k, v in taps.items() if any(fnmatch.fnmatch(k, p) for p in patterns)}


def _pack(res, taps_patterns, with_image, meta):
    """Golden tensors of one frame + the dtype record."""
    t = {name: getattr(res, name) for name in OUTS}
    for k, v in res.inputs.items():
        if k == "image" and not with_image:
            continue
        if k == "prev_bev":
            meta["prev_bev_sha256"] = hashlib.sha256(np.ascontiguousarray(v).tobytes()).hexdigest()
            continue
        t[f"input.{k}"] = v
    for k in GEOMETRY:
        t[f"geometry.{k}"] = res.geometry[k]
    dtypes = {}
    for k, v in _select(res.taps, taps_patterns).items():
        arr = np.asarray(v)
        if arr.size > FP16_ABOVE:
            arr = arr.astype(np.float16)
        t[f"tap.{k}"] = arr
        dtypes[k] = str(arr.dtype)
    meta["tap_dtypes"] = dtypes
    return t


def _compare(res, golden_path: Path, report: dict, key: str) -> None:
    if golden_path is None or not Path(golden_path).is_file():
        return
    out = {}
    with np.load(golden_path) as g:
        for name in OUTS:
            for gk in (name, f"{key}_{name}"):
                if gk in g.files:
                    out[name] = pcc(getattr(res, name), g[gk])
        for k in g.files:
            if k.startswith("tap_") and k[4:] in res.taps:
                out["tap." + k[4:]] = pcc(res.taps[k[4:]].reshape(g[k].shape), g[k])
    report[key] = out
    print(f"  {key}: " + ", ".join(f"{k} {v:.10f}" for k, v in out.items()), flush=True)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", type=Path, default=paths.golden_dir())
    ap.add_argument("--threads", type=int, default=4)
    ap.add_argument("--only", nargs="*", default=["pandaset", "sample0", "nuscenes"])
    a = ap.parse_args()
    import torch

    torch.set_num_threads(a.threads)
    wpath = paths.weights_file()
    if wpath is None:
        raise SystemExit("weights not found (BEVFORMER_WEIGHTS_DIR)")
    ref = BEVFormerReference(wpath)
    base_meta = {"generator": "bundles/bevformer-p150/code/scripts/make_goldens.py", "bundle_version": __version__,
                 "bundle_commit": _git_head(CODE.parent), "weights": str(wpath),
                 "weights_sha256": ref.weights.sha256, "reference_form": "port (BEVFormerNet.forward_port)",
                 "fp16_taps_above_elements": FP16_ABOVE}
    report: dict = {}
    t_all = time.time()

    if "pandaset" in a.only:
        pub = paths.public_data()
        ref.reset()
        small = {"frames": [], "doc": "PandaSet 019 f30-40 (shipped sample, CC BY 4.0 + PandaSet Dataset Terms): "
                                      "CPU fp32 reference (port form), stateful sequence reset at f30; Autoware "
                                      "post-processing (score > 0.2)"}
        last = {}
        for f in (30, 35, 40):
            t0 = time.time()
            fr = load_frame(paths.PANDASET_SAMPLES / f"pandaset_019_f{f}.json")
            res = ref.step(fr, port=True, taps=(f == 35), workers=4)
            meta = dict(base_meta, frame=f, sequence="pandaset_019_f30-40", reset=bool(fr.stream["reset"]))
            g = pub / "pandaset" / "goldens" / "ps019_f30-40" / f"f{f}.npz" if pub else None
            _compare(res, g, report, f"pandaset_019/f{f}")
            meta["pcc_vs_research"] = report.get(f"pandaset_019/f{f}", {})
            save_goldens(a.out / "pandaset_019" / f"f{f}.npz", _pack(res, ALL_TAPS if f == 35 else [], f == 35, meta),
                         meta, compress=True)
            dets = res.detections
            small["frames"].append({"frame": f, "use_prev_bev": float(res.inputs["use_prev_bev"][0]),
                                    "n_detections": len(dets), "pcc_vs_research": meta["pcc_vs_research"],
                                    "detections": [{k: (round(v, 5) if isinstance(v, float) else v)
                                                    for k, v in d.items()} for d in dets.to_records()]})
            last[f"f{f}_outputs_classes"] = res.outputs_classes[-1, 0]
            last[f"f{f}_outputs_coords"] = res.outputs_coords[-1, 0]
            if f == 30:
                body = hpost.to_output(dets, model="bevformer-p150").to_dict()
                (paths.PANDASET_SAMPLES / "pandaset_019_f30.reference.json").write_text(json.dumps(body, indent=1) + "\n")
            print(f"pandaset f{f}: {len(dets)} detections, {time.time() - t0:.0f} s", flush=True)
        paths.PANDASET_GOLDENS.mkdir(parents=True, exist_ok=True)
        (paths.PANDASET_GOLDENS / "pandaset_019_f30-40.json").write_text(json.dumps(small, indent=1) + "\n")
        np.savez_compressed(paths.PANDASET_GOLDENS / "pandaset_019_f30-40_last_layer.npz", **last)

    if "sample0" in a.only:
        g0 = paths.research_file("golden", "sample0_golden.npz")
        if g0 is not None:
            with np.load(g0) as g:
                gs = {k: g[k] for k in g.files}
            prev = gs["f1_prev_bev"]
            for f in (1, 2, 3):
                t0 = time.time()
                res = ref.run_frame(gs["image"], prev, gs[f"f{f}_use_prev_bev"], gs[f"f{f}_can_bus"], gs["lidar2img"],
                                    port=True, taps=(f == 1))
                meta = dict(base_meta, frame=f, sequence="sample0 (bevdet_vendor nuScenes key-frame; internal only)")
                _compare(res, g0, report, f"f{f}")
                report[f"sample0/f{f}"] = report.pop(f"f{f}")
                meta["pcc_vs_research"] = report[f"sample0/f{f}"]
                meta["n_detections"] = len(res.detections)
                save_goldens(a.out / "sample0" / f"f{f}.npz",
                             _pack(res, TRANSFORMER_TAPS if f == 1 else [], f == 1, meta), meta, compress=True)
                prev = res.bev_embed
                print(f"sample0 f{f}: {len(res.detections)} detections, {time.time() - t0:.0f} s", flush=True)

    if "nuscenes" in a.only and paths.public_data() is not None:
        P = paths.import_public_common()
        pub = paths.public_data()
        t0 = time.time()
        x = P.load_frame_inputs(str(pub / "nuscenes" / "inputs" / "scene-0103" / "kf01.npz"))
        with np.load(pub / "nuscenes" / "goldens" / "scene-0103" / "kf00.npz") as g:
            prev = g["bev_embed"]
        res = ref.run_frame(x["image"], prev, x["use_prev_bev"], x["can_bus"], x["lidar2img"], port=True, taps=True)
        meta = dict(base_meta, frame="kf01", sequence="nuScenes mini_val scene-0103 (CC BY-NC-SA 4.0; internal only)",
                    teacher_forced="prev_bev = public-data golden kf00 bev_embed")
        _compare(res, pub / "nuscenes" / "goldens" / "scene-0103" / "kf01.npz", report, "nuscenes_scene-0103/kf01")
        meta["pcc_vs_research"] = report["nuscenes_scene-0103/kf01"]
        save_goldens(a.out / "nuscenes_scene-0103" / "kf01.npz", _pack(res, TRANSFORMER_TAPS, True, meta), meta,
                     compress=True)
        print(f"nuscenes kf01: {len(res.detections)} detections, {time.time() - t0:.0f} s", flush=True)

    a.out.mkdir(parents=True, exist_ok=True)
    worst = min((v for r in report.values() for v in r.values()), default=float("nan"))
    (a.out / "report.json").write_text(json.dumps({"meta": base_meta, "pcc_vs_research": report,
                                                   "worst_pcc": worst, "seconds": time.time() - t_all}, indent=1))
    print(f"done in {time.time() - t_all:.0f} s; worst PCC vs research {worst:.10f}")
    return 0


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