#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Device profile of transfusion-p150: one eager forward and one traced replay of the whole graph between Tracy
signposts.

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

Sections (signposts): ``eager`` .. ``eager_end``, one eager forward after the warm-up (``TraceRunner.run_eager``:
the same graph and programs as the trace, nothing compiled), and ``trace`` .. ``trace_end``, ``--replays`` replays
of the captured ``default`` trace (the served device path; the uploaded inputs are the voxelized ``--input``). The
device profiler buffer is flushed (``ttnn.ReadDeviceProfiler``) before and after each section, and inside the eager
section after every module (``--flush-modules``, default on: a frame is well above the profiler buffer's 1,000
programs). Inside the eager section every module of the graph opens a signpost (``m:pfn``, ``m:scatter``,
``m:backbone.block<i>``, ``m:neck``, ``m:head.shared``, ``m:head.heatmap.hidden``, ``m:head.heatmap``,
``m:head.sigmoid``, ``m:head.local_max``, ``m:query.topk``, ``m:query.decode``, ``m:query.init``, ``m:query.keys``,
``m:query.decoder``, ``m:query.predict``, ``m:query.heat``, ``m:query.pack``) and every backbone conv a layer
signpost (``c:backbone.block<i>.conv<j>``; ``--no-layer-signposts`` turns them off), so the ops of the eager
forward, and through the identical op order those of the replay, can be attributed to modules
(``profile_summary.py``). The trace section is flushed every replay; with ``TT_METAL_PROFILER_MID_RUN_DUMP=1`` the
profiler also dumps mid-trace (needed when one replay exceeds the buffer). The output folder must be absolute and
outside the bundle (PLAN.md 5.1).
"""
from __future__ import annotations

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

from tt_transfusion import TransFusion
from tt_transfusion.ttaw.profiling import read_device_profiler, signpost, signposted

SAMPLE = Path(__file__).resolve().parents[1] / "tt_transfusion" / "samples" / "test_pcd.npz"


class _Signposted:
    """A callable stand-in for a module that emits its signposts, then calls the module (attributes forward)."""

    def __init__(self, inner, labels: Sequence[str], flush: Callable[[], None] = None):
        self.inner, self._labels, self._flush = inner, list(labels), flush

    def __call__(self, *args, **kwargs):
        if self._flush is not None:
            self._flush()
        for label in self._labels:
            signpost(label)
        return self.inner(*args, **kwargs)

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


def module_signposts(graph, layers: bool = True, flush: Callable[[], None] = None) -> Callable[[], None]:
    """Wrap the modules of a ``tt.model.TtTransFusion`` (and the module-level ops its forward calls) with
    signposting stand-ins; returns the function that restores them. ``flush`` runs before each module signpost
    (device profiler read-out: one eager frame holds more programs than the buffer)."""
    from tt_transfusion.tt import model as tt_model

    undo: List[Callable[[], None]] = []

    def swap(obj, attr, labels, fl=flush):
        inner = getattr(obj, attr)
        setattr(obj, attr, _Signposted(inner, labels, fl))
        undo.append(lambda: setattr(obj, attr, inner))

    def swap_method(obj, attr, labels, fl=flush):
        """An instance attribute that shadows a bound method (removed again on restore)."""
        inner = getattr(obj, attr)
        setattr(obj, attr, _Signposted(inner, labels, fl))
        undo.append(lambda: delattr(obj, attr))

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

    def swap_fn(module, name, label):
        inner = getattr(module, name)

        def wrapped(*args, **kwargs):
            if flush is not None:
                flush()
            signpost(label)
            return inner(*args, **kwargs)

        setattr(module, name, wrapped)
        undo.append(lambda: setattr(module, name, inner))

    swap(graph, "pfn", ["m:pfn"])
    swap_fn(tt_model, "scatter_rows", "m:scatter")
    for i, block in enumerate(graph.backbone.blocks):
        for j in range(len(block)):
            labels = ([f"m:backbone.block{i}"] if j == 0 else []) + ([f"c:backbone.block{i}.conv{j}"] if layers else [])
            swap_item(block, j, labels, flush if j == 0 else None)
    swap(graph, "neck", ["m:neck"])
    swap(graph, "shared", ["m:head.shared"])
    swap(graph, "hm_hidden", ["m:head.heatmap.hidden"])
    swap(graph, "hm_out", ["m:head.heatmap"])
    swap_fn(tt_model, "sigmoid_heat", "m:head.sigmoid")
    swap(graph, "local_max", ["m:head.local_max"])
    qh = graph.query_head
    swap_method(qh.select, "indices", ["m:query.topk"])
    swap_method(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_method(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=str(SAMPLE))
    ap.add_argument("--dispatch", default=None, choices=["eth", "worker"])
    ap.add_argument("--cqs", type=int, default=None, choices=[1, 2])
    ap.add_argument("--replays", type=int, default=1, help="replays inside the trace section")
    ap.add_argument("--weights-dir", default=None)
    ap.add_argument("--no-eager", dest="eager", action="store_false")
    ap.add_argument("--no-layer-signposts", dest="layers", action="store_false")
    ap.add_argument("--no-flush-modules", dest="flush_modules", action="store_false")
    a = ap.parse_args()
    import ttnn

    from tt_transfusion.tt.model import stage_inputs

    t0 = time.perf_counter()
    with TransFusion.from_pretrained(dispatch=a.dispatch, num_command_queues=a.cqs,
                                     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), flush=True)
        prepared = model._prepare(points=a.input)
        out = model(a.input)                           # one served frame: upload + replay + readback
        print(f"{out.to_dict()['num_detections']} detections, {prepared.num_pillars} pillars", flush=True)
        ttnn.synchronize_device(dev)
        read_device_profiler(dev)                      # warm-up / capture / first frame out of the buffer
        inputs = stage_inputs(model.staging, prepared)
        variant = getattr(model.staging, "variant", "default")     # PFN_BUCKETS: the frame's capacity variant
        print(f"variant {variant}", flush=True)
        if a.eager:
            def flush():
                ttnn.synchronize_device(dev)
                read_device_profiler(dev)

            restore = module_signposts(model.graph, layers=a.layers, flush=flush if a.flush_modules else None)
            try:
                with signposted("eager"):
                    tr.run_eager(variant, inputs=inputs)
                    ttnn.synchronize_device(dev)
            finally:
                restore()
            read_device_profiler(dev)
        tr.upload(inputs)
        ttnn.synchronize_device(dev)
        read_device_profiler(dev)
        with signposted("trace"):
            for _ in range(a.replays):
                tr.replay(variant)
                ttnn.synchronize_device(dev)
                read_device_profiler(dev)
        read_device_profiler(dev)
        print("trace:", json.dumps(tr.describe()), flush=True)


if __name__ == "__main__":
    main()