# SPDX-License-Identifier: Apache-2.0
"""Device bring-up of the PTv3 graph on a (cropped) frame: build the graph at a capacity bucket fitted to the frame,
capture the tap variants, replay them on the frame's tables and compare every module tap with the fp32 CPU
reference fed the identical host tables (PCC over the frame's real rows); optionally replay vs eager bit check.

    bin/devrun -t 1800 -- python code/scripts/bringup.py --crop 8 [--det] [--no-teacher] [--eager-check]

PLAN.md 4.4: the first run of a new config is a small crop under ``TT_METAL_WATCHER=2``. Writes a JSON report
(``--json``). Device-opening: run it through ``bin/devrun``.
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

import numpy as np

HERE = Path(__file__).resolve()
sys.path.insert(0, str(HERE.parents[1]))


def _rows_nchw(a):
    a = np.asarray(a, np.float32)
    return a[0].transpose(1, 2, 0).reshape(-1, a.shape[1])


def compare(dev_taps, ref_taps, counts, *, det_names=()):
    from tt_ptv3.ttaw.metrics import error_stats

    out = {}
    for name, ref in ref_taps.items():
        if name not in dev_taps:
            continue
        d = np.asarray(dev_taps[name], np.float32)
        r = np.asarray(ref, np.float32)
        if r.ndim == 4:                       # det NCHW maps -> rows
            r = _rows_nchw(r)
        d = d.reshape(-1, d.shape[-1])
        n = r.shape[0]
        d = d[:n, :r.shape[1]]
        out[name] = {k: float(v) for k, v in error_stats(d, r).items()}
    return out


def main(argv=None) -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--frame", default="sample")
    ap.add_argument("--crop", type=float, default=8.0, help="keep |x|, |y| < crop metres (0: the whole frame)")
    ap.add_argument("--det", action="store_true")
    ap.add_argument("--no-teacher", action="store_true")
    ap.add_argument("--qhead", action="store_true", help="also run the C26 head on the reference's dense maps")
    ap.add_argument("--eager-check", action="store_true")
    ap.add_argument("--mode", default="alias-first")
    ap.add_argument("--margin", type=int, default=0)
    ap.add_argument("--json", default=None)
    ap.add_argument("--threads", type=int, default=2)
    args = ap.parse_args(argv)

    from tt_ptv3.host import preprocess as pre
    from tt_ptv3.reference.pipeline import Reference, onnx_dir_default, set_threads
    from tt_ptv3.tests._data import load_frame
    from tt_ptv3.tt import inputs as tin

    set_threads(args.threads)
    pts = load_frame(args.frame)
    if args.crop > 0:
        keep = (np.abs(pts[:, 0]) < args.crop) & (np.abs(pts[:, 1]) < args.crop)
        pts = pts[keep]
    frame = pre.prepare_frame(pts, neighbour_mode=args.mode, capacities=None)
    caps = tin.fitted_capacities(frame.counts, args.margin)
    print(f"points {len(pts)} counts {frame.counts} caps {caps}", flush=True)
    t0 = time.time()
    ref = Reference(onnx_dir_default(), det=args.det)
    taps = {}
    res = ref.run(None, frame=frame, det=args.det, taps=taps)
    print(f"reference {time.time() - t0:.1f} s", flush=True)

    from tt_ptv3.device import close_device, describe_device, open_device
    from tt_ptv3.tt.model import TtPTv3, build_runner, teacher_indices, unpack_seg

    dev = open_device(allow_fallback=False)
    report = {"frame": args.frame, "crop": args.crop, "counts": frame.counts, "caps": list(caps),
              "device": describe_device(dev)}
    try:
        graph = TtPTv3(dev, ref.weights, det=args.det)
        variants = ["taps"] + (["dettaps"] + ([] if args.no_teacher else ["detteacher"]) if args.det else [])
        variants += ["qhead"] if (args.det and args.qhead) else []
        runner = build_runner(graph, {"b": caps}, variants=variants)
        t0 = time.time()
        runner.capture()
        report["capture_s"] = time.time() - t0
        print(f"capture {report['capture_s']:.1f} s", flush=True)
        inputs = {f"b.{k}": v for k, v in tin.frame_inputs(frame, caps, det=args.det).items()}
        if args.det and not args.no_teacher:
            inputs["b.topk"] = teacher_indices(taps["det.topk_indices"], graph.kp)
        t0 = time.time()
        dev_taps = runner("taps_b", inputs=inputs)
        report["replay_taps_s"] = time.time() - t0
        seg = unpack_seg(dev_taps["seg"], frame.counts[0], caps[0])
        report["seg_label_agreement"] = float((seg["labels"] == res.pred_labels).mean())
        p = np.asarray(res.pred_probs, np.float64)
        ent = -(np.where(p > 0, p * np.log(np.where(p > 0, p, 1)), 0)).sum(1) / np.log(15)
        report["seg_prob_max_abs"] = float(np.abs(seg["prob"] - p.max(1)).max())
        report["seg_entropy_max_abs"] = float(np.abs(seg["entropy"] - ent).max())
        report["taps"] = compare(dev_taps, taps, frame.counts)
        if args.eager_check:
            eager = runner.run_eager("taps_b")
            eager_np = {k: np.asarray(v) for k, v in eager.items()} if isinstance(eager, dict) else eager
            report["replay_vs_eager_seg_equal"] = bool(np.array_equal(np.asarray(dev_taps["seg"]),
                                                                     np.asarray(eager_np["seg"])))
        if args.det:
            d = runner("dettaps_b")
            report["det_taps"] = compare(d, taps, frame.counts)
            from tt_ptv3.ttaw.ops.topk import selection_metrics
            idx = np.rint(np.asarray(d["indices"], np.float64).reshape(-1)[:graph.num_proposals]).astype(np.int64)
            report["det_topk_shared"] = int(len(set(idx.tolist()) & set(taps["det.topk_indices"].tolist())))
            if not args.no_teacher:
                t = runner("detteacher_b")
                rep = {"det.norm1": t["det.decoder.self_attn"], "det.norm2": t["det.decoder.cross_attn"],
                       "det.norm3": t["det.decoder.out"]}
                det_out = graph.assemble({"heads": t["heads"], "query_heat": t["query_heat"],
                                          "indices": t["indices"]})
                for name in ("center", "height", "dim", "rot", "vel", "heatmap"):
                    rep[f"det.head.{name}"] = det_out[name][0].T
                report["det_teacher"] = compare(rep, taps, frame.counts)
                report["det_teacher"]["query_heatmap_score"] = compare(
                    {"q": det_out["query_heatmap_score"][0].T}, {"q": res.det["query_heatmap_score"][0].T},
                    frame.counts)["q"]
            if args.qhead:
                qin = {"b.lidar": _rows_nchw(taps["det.shared"]).reshape(1, 1, -1, 128),
                       "b.topk": teacher_indices(taps["det.topk_indices"], graph.kp)}
                h = np.zeros((1, 1, graph.cells, 32), np.float32)
                h[0, 0, :, :graph.num_det_classes] = _rows_nchw(taps["det.heat_nms"])
                qin["b.heatin"] = h
                q = runner("qhead_b", inputs=qin)
                rep = {"det.norm1": q["det.decoder.self_attn"], "det.norm2": q["det.decoder.cross_attn"],
                       "det.norm3": q["det.decoder.out"]}
                det_out = graph.assemble({"heads": q["heads"], "query_heat": q["query_heat"], "indices": q["indices"]})
                for name in ("center", "height", "dim", "rot", "vel", "heatmap"):
                    rep[f"det.head.{name}"] = det_out[name][0].T
                report["det_qhead_golden_inputs"] = compare(rep, taps, frame.counts)
        runner.release()
    finally:
        close_device(dev)
    worst = {k: v["pcc"] for k, v in report["taps"].items()}
    for k in sorted(worst, key=lambda k: worst[k])[:8]:
        print(f"  {k:28s} pcc {worst[k]:.6f}")
    print(json.dumps({k: v for k, v in report.items() if k not in ("taps", "det_taps", "det_teacher", "det_qhead_golden_inputs")}, indent=1))
    for sec in ("det_taps", "det_teacher", "det_qhead_golden_inputs"):
        for k, v in report.get(sec, {}).items():
            print(f"  {sec} {k:26s} pcc {v['pcc']:.6f} max_abs {v['max_abs']:.4g}")
    if args.json:
        Path(args.json).write_text(json.dumps(report, indent=1))
    return 0


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