#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Device diagnostic: per-module error of the device graph on one catalogued frame, next to the CPU emulation.

    cd bundles/frnet-p150 && bin/devrun -t 1200 -- env PYTHONPATH=code python code/scripts/diag_frame_taps.py \\
        --frame nuscenes_scene-0916_40.ot128 --out /home/ubuntu/experiments/tt-models/logs/frnet/m1b/diag_0916_40.json

The exact fp32 TT dataflow (``reference/rewrites.py FRNetPlanTorch``) is the reference of every tap; the CPU
emulation (``--emulation``, default the ``plan_exact_points`` row of ``code/scripts/precision_study.py``: bf16 conv
operands, gathers and segment tables, exact point layers) and the device ``taps`` variant (one replay; knobs from the
``FRNET_*`` environment) are compared with it tap by tap: where the device's ``1 - PCC`` grows faster than the
emulation's is where the device loses accuracy that the emulation does not model. Also reported: ``pred_probs`` PCC
and raw-point argmax agreement vs ONNX Runtime for the device and the emulation. The CPU work runs first (in the
device lock: ~1-2 min at OT128).
"""
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]))
sys.path.insert(0, str(HERE.parent))

TAPS = ("encoder.pf3", "encoder.voxel_feats", "backbone.stem", "backbone.point_stem", "backbone.fusion_stem",
        *(f"backbone.{k}{i}" for i in range(1, 5) for k in ("layer", "point_fusion", "pixel_fusion", "attention",
                                                            "stage")),
        "decoder.fuse1", "decoder.fuse2", "decoder.point_fuse1", "decoder.point_fuse2",
        *(f"head.mlp{i}" for i in range(4)), "head.logits", "pred_probs")


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--frame", required=True)
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--emulation", default="plan_exact_points")
    ap.add_argument("--threads", type=int, default=4)
    ap.add_argument("--cpu-only", action="store_true", help="stop after the CPU part (a dry run without a device)")
    args = ap.parse_args()

    import torch

    torch.set_num_threads(args.threads)
    from precision_study import CONFIGS
    from tt_frnet.host.preprocess import build_plan
    from tt_frnet.reference.config import FRNetConfig
    from tt_frnet.reference.rewrites import Emulation, FRNetPlanTorch
    from tt_frnet.reference.weights import load_weights
    from tt_frnet.tests import _data
    from tt_frnet.ttaw.golden import TapRegistry
    from tt_frnet.ttaw.metrics import argmax_agreement, pcc

    variant = args.frame.rpartition(".")[2]
    wdir = _data.weights_dir()
    cfg = FRNetConfig.from_model_dir(wdir, variant)
    w = load_weights(wdir, variant)
    inputs, ort_probs, _ = _data.load_frame(args.frame)
    plan = build_plan(inputs, cfg, pre_norm=(w.pre_norm_scale, w.pre_norm_shift))
    n, n_raw = plan.n, inputs.num_points_raw
    H, W = plan.frustum_hw
    t0 = time.time()
    exact = TapRegistry(include=list(TAPS))
    FRNetPlanTorch(w)(plan, taps=exact)
    emu = TapRegistry(include=list(TAPS))
    emu_probs = plan.unsort(FRNetPlanTorch(w, emulation=Emulation.parse(CONFIGS[args.emulation]))(plan, taps=emu))
    cpu_s = time.time() - t0

    def rows(name, a):
        """point taps: the first N sorted rows; voxel taps: the first M rows."""
        a = np.asarray(a)
        if a.shape[0] == plan.n_cap:
            return a[:n]
        return a[: plan.m] if a.shape[0] in (cfg.m_cap, cfg.m_cap + 1) else a

    if args.cpu_only:
        ort = np.asarray(ort_probs, np.float32)
        print(f"cpu {cpu_s:.0f} s; emulation pred_probs pcc {pcc(emu_probs, ort):.6f} argmax "
              f"{argmax_agreement(emu_probs[:n_raw], ort[:n_raw]):.5f}")
        return 0
    from tt_frnet.device import device_session
    from tt_frnet.tt.model import TtFRNet, build_runner, frame_arrays

    report = {"frame": args.frame, "N": n, "M": plan.m, "emulation": args.emulation, "cpu_s": round(cpu_s, 1),
              "taps": {}}
    with device_session() as device:
        graph = TtFRNet(device, w, cfg)
        arrays = frame_arrays(plan, cfg.m_cap)
        runner = build_runner(graph, variants=("taps",), name="frnet_diag", init=arrays)
        try:
            runner.capture()
            dev = runner("taps", inputs=arrays)
            report["knobs"] = graph.knobs.as_dict()
            for name in TAPS:
                if name not in dev or name not in exact:
                    continue
                ref = np.asarray(exact[name], np.float32)
                d = np.asarray(dev[name], np.float32)
                if ref.ndim == 4:
                    _, c, h, ww = ref.shape
                    d = d.reshape(-1, d.shape[-1])[: h * ww, :c].reshape(1, h, ww, c).transpose(0, 3, 1, 2)
                    e = np.asarray(emu[name], np.float32)
                else:
                    ref = rows(name, ref)[:, : ref.shape[-1]]
                    d = rows(name, d.reshape(-1, d.shape[-1]))[:, : ref.shape[-1]]
                    e = rows(name, np.asarray(emu[name], np.float32))[:, : ref.shape[-1]]
                p_dev, p_emu = float(pcc(d, ref)), float(pcc(e, ref))
                ratio = (1 - p_dev) / max(1 - p_emu, 1e-12)
                report["taps"][name] = {"pcc_dev": p_dev, "pcc_emu": p_emu, "err_ratio": ratio}
                print(f"{name:26s} dev {p_dev:.6f} emu {p_emu:.6f}  (1-pcc) ratio {ratio:6.2f}", flush=True)
            probs = plan.unsort(np.asarray(dev["pred_probs"], np.float32).reshape(-1, 32)[:, :27])
        finally:
            runner.release()
            graph.release()
    ort = np.asarray(ort_probs, np.float32)
    report["device"] = {"pred_probs_pcc": float(pcc(probs, ort)),
                        "argmax_agreement_raw": float(argmax_agreement(probs[:n_raw], ort[:n_raw]))}
    report["emulation_result"] = {"pred_probs_pcc": float(pcc(emu_probs, ort)),
                                  "argmax_agreement_raw": float(argmax_agreement(emu_probs[:n_raw], ort[:n_raw]))}
    print("device", report["device"], "emulation", report["emulation_result"])
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(report, indent=1))
    return 0


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