#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Device profile of ptv3-p150: per trace one eager forward (module / layer signposts) and one traced replay.

    ROOT=/home/ubuntu/experiments/tt-models
    $ROOT/bin/devrun -t 3600 -- python -m tracy -r -p -v --no-web-server --op-support-count 6000 \\
        -o $ROOT/generated/profiler/ptv3_baseline code/scripts/profile_ops.py --input <cloud>
    tt-perf-report <ops_perf_results_*.csv> --start-signpost trace_a --end-signpost trace_a_end --tracing-mode
    python code/scripts/profile_summary.py <the same csv> profile_summary.json --md profile_summary.md

Sections (signposts): ``eager_a`` .. ``eager_a_end`` (``TraceRunner.run_eager`` of ``encseg_<b>``: the same graph
and programs as trace A, nothing compiled), ``trace_a`` .. ``trace_a_end`` (``--replays`` replays of the captured
trace A), then the same for trace B (``eager_b`` / ``trace_b``, variant ``seg_det``). The device profiler buffer is
flushed (``ttnn.ReadDeviceProfiler``) before and after each section, after every replay and, inside the eager
sections, before every module signpost (``--no-flush-modules`` turns that off).

Module signposts (``m:``): ``embed``, ``enc<s>.pool``, ``enc<s>.block<b>``, ``dec<s>`` (up / skip linears, unpool
gather, sum), ``dec<s>.block0``, ``seg.head`` (classifier + softmax), ``seg.summary``; ``det.fusion``,
``det.point_proj_bev`` (point_proj + the 9 BEV candidate gathers + max), ``det.bev_refine``, ``det.stem``,
``det.block<i>``, ``det.shared``, ``det.heatmap`` (hidden + out conv + sigmoid), ``det.local_max``,
``query.<part>``. Layer signposts (``c:``) inside each PTv3 block: ``cpe`` (CPE table + gather-GEMM + LN + add),
``attn.in`` (LN, two-term sequence table, gather into patch order), ``attn.qkv`` (fp32 q / k / v projections),
``attn.core`` (head split, ``q k^T``, softmax, ``p v``, head merge), ``attn.out`` (proj + inverse gather + add),
``mlp`` (LN, fc1 + GELU, fc2, add). The replay has the same op sequence as the eager forward (checked by
``profile_summary.py``), so every replayed op inherits its module / layer. The output folder must be absolute and
outside the bundle (PLAN.md 5.1). (The baseline profile of OPT_BASELINE.md ran before ``det.block<i>`` got its
``c:conv2+residual`` signpost, so its ``conv1`` layer there is the whole block.)
"""
from __future__ import annotations

import argparse
import json
import time
from pathlib import Path
from typing import Callable, List, Sequence

import numpy as np

from tt_ptv3 import PTv3
from tt_ptv3.ttaw.profiling import read_device_profiler, signpost, signposted

CODE = Path(__file__).resolve().parents[1]
SAMPLES = CODE / "tt_ptv3" / "samples"


class _Signposted:
    """A callable stand-in that emits ``pre`` signposts (after an optional flush), calls the module, then emits
    ``post`` signposts; attributes forward to the module."""

    def __init__(self, inner, pre: Sequence[str] = (), flush: Callable[[], None] = None, post: Sequence[str] = ()):
        self.inner, self._pre, self._post, self._flush = inner, list(pre), list(post), flush

    def __call__(self, *args, **kwargs):
        if self._flush is not None and self._pre:
            self._flush()
        for label in self._pre:
            signpost(label)
        out = self.inner(*args, **kwargs)
        for label in self._post:
            signpost(label)
        return out

    def __getattr__(self, name):
        return getattr(self.inner, name)


def module_signposts(graph, flush: Callable[[], None] = None) -> Callable[[], None]:
    """Wrap the modules of a ``tt.model.TtPTv3`` with signposting stand-ins; returns the restore function."""
    undo: List[Callable[[], None]] = []

    def swap(obj, attr, pre=(), fl=flush, post=()):
        inner = getattr(obj, attr)
        had = attr in getattr(obj, "__dict__", {})
        setattr(obj, attr, _Signposted(inner, pre, fl, post))
        undo.append((lambda: setattr(obj, attr, inner)) if had else (lambda: delattr(obj, attr)))

    def swap_item(seq, i, pre=(), fl=flush):
        inner = seq[i]
        seq[i] = _Signposted(inner, pre, fl)
        undo.append(lambda: seq.__setitem__(i, inner))

    def block(blk, label):
        swap(blk, "attention", ["c:attn.in"], None)
        swap(blk, "q", ["c:attn.qkv"], None)
        swap(blk, "v", [], None, post=["c:attn.core"])
        swap(blk, "proj", ["c:attn.out"], None)
        swap(blk, "norm2", ["c:mlp"], None)

    swap(graph, "embed", ["m:embed"])
    for s in sorted(graph.pools):
        swap_item(graph.pools, s, [f"m:enc{s}.pool"])
    for s, bl in graph.blocks.items():
        for b in range(len(bl)):
            block(bl[b], f"enc{s}.block{b}")
            swap_item(bl, b, [f"m:enc{s}.block{b}", "c:cpe"])
    for s in sorted(graph.dec_up):
        swap_item(graph.dec_up, s, [f"m:dec{s}", "c:up"])
        swap_item(graph.dec_skip, s, ["c:skip+unpool+sum"], None)
    for s in sorted(graph.dec_block):
        block(graph.dec_block[s], f"dec{s}.block0")
        swap_item(graph.dec_block, s, [f"m:dec{s}.block0", "c:cpe"])
    swap(graph, "classifier", ["m:seg.head"])
    swap(graph, "seg_summary", ["m:seg.summary"])
    if graph.det is not None:
        swap(graph, "fusion_skip", ["m:det.fusion"])
        swap(graph, "point_proj", ["m:det.point_proj_bev"])
        swap(graph, "bev_refine", ["m:det.bev_refine"])
        swap(graph, "stem", ["m:det.stem"])
        for i, (c1, c2) in enumerate(list(graph.det_blocks)):
            w1 = _Signposted(c1, [f"m:det.block{i}", "c:conv1"], flush)
            if hasattr(c2, "conv_sum"):         # two-term form: tt.dense.two_term_residual calls conv2.conv_sum
                swap(c2, "conv_sum", ["c:conv2+residual"], None)
                w2 = c2
            else:
                w2 = _Signposted(c2, ["c:conv2+residual"], None)
            graph.det_blocks[i] = (w1, w2)
            undo.append(lambda i=i, c1=c1, c2=c2: graph.det_blocks.__setitem__(i, (c1, c2)))
        swap(graph, "shared", ["m:det.shared"])
        swap(graph, "heat_hidden", ["m:det.heatmap"])
        swap(graph, "local_max", ["m:det.local_max"])
        qh = graph.query_head
        swap(qh.select, "indices", ["m:query.topk"])
        swap(qh.select, "decode", ["m:query.decode"])
        for attr, label in (("init_queries", "m:query.init"), ("keys", "m:query.keys"),
                            ("decoder", "m:query.decoder"), ("predict", "m:query.predict"),
                            ("query_heat", "m:query.heat"), ("_indices_tile", "m:query.pack")):
            swap(qh, attr, [label])

    def restore() -> None:
        for fn in reversed(undo):
            fn()

    return restore


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--input", default=None, help="cloud (default: the shipped sample)")
    ap.add_argument("--dispatch", default=None, choices=["eth", "worker"])
    ap.add_argument("--cqs", type=int, default=None, choices=[1, 2])
    ap.add_argument("--variant", default=None, choices=["seg_det", "seg"])
    ap.add_argument("--replays", type=int, default=1, help="replays inside each trace section")
    ap.add_argument("--weights-dir", default=None)
    ap.add_argument("--no-eager", dest="eager", action="store_false")
    ap.add_argument("--no-flush-modules", dest="flush_modules", action="store_false")
    a = ap.parse_args()
    import ttnn

    from tt_ptv3 import io as tio
    from tt_ptv3.host import preprocess as pre
    from tt_ptv3.tt import inputs as tin

    src = a.input
    if src is None:
        src = next(str(SAMPLES / n) for n in ("test_pcd.npz", "pandaset_019_f40.npz") if (SAMPLES / n).is_file())
    t0 = time.perf_counter()
    with PTv3.from_pretrained(dispatch=a.dispatch, num_command_queues=a.cqs, variant=a.variant,
                              weights_dir=a.weights_dir) as model:
        tr, dev = model.runner, model.device
        print("loaded in %.1f s:" % (time.perf_counter() - t0), json.dumps(model.device_info, default=str), flush=True)
        out = model(src)                               # one served frame: upload + replays + readbacks
        print(json.dumps({"input": src, "detections": out.to_dict().get("num_detections"),
                          "frame": out.meta.get("frame")}, default=str), flush=True)
        frame = pre.prepare_frame(np.asarray(tio.load_points(src).points, np.float32))
        b, caps = model.bucket_for(frame.counts)       # the bucket the served frame ran in
        print(json.dumps({"bucket": b, "caps": list(caps)}), flush=True)
        tables = {f"{b}.{k}": v for k, v in tin.frame_inputs(frame, caps, det=model.with_det).items()}
        ttnn.synchronize_device(dev)
        read_device_profiler(dev)                      # warm-up / capture / first frame out of the buffer

        def flush():
            ttnn.synchronize_device(dev)
            read_device_profiler(dev)

        variants = [("a", f"encseg_{b}")] + ([("b", f"det_{b}")] if model.with_det else [])
        if a.eager:
            restore = module_signposts(model.graph, flush=flush if a.flush_modules else None)
            try:
                for tag, v in variants:
                    with signposted(f"eager_{tag}"):
                        tr.run_eager(v, inputs=tables if tag == "a" else None)
                        ttnn.synchronize_device(dev)
                    read_device_profiler(dev)
            finally:
                restore()
        tr.upload(tables)
        ttnn.synchronize_device(dev)
        read_device_profiler(dev)
        for tag, v in variants:
            with signposted(f"trace_{tag}"):
                for _ in range(a.replays):
                    tr.replay(v)
                    ttnn.synchronize_device(dev)
                    read_device_profiler(dev)
            read_device_profiler(dev)
        print("trace:", json.dumps(tr.describe(), default=str), flush=True)


if __name__ == "__main__":
    main()