#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Device profile of sceneseg-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 --op-support-count 6000 \\
        -o $ROOT/generated/profiler/sceneseg_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 <ops_perf_results_*.csv> summary.json --md summary.md

Sections (signposts): ``eager`` .. ``eager_end``, one eager forward after the warm-up (``SceneSegRunner.run_eager``:
the same graph and programs as the trace, nothing compiled), and ``trace`` .. ``trace_end``, one replay of the
captured trace (the served device path; the uploaded input is the squashed sample). The device profiler buffer is
flushed (``ttnn.ReadDeviceProfiler``) before and after each section. Inside the eager section every block of the
graph opens a signpost (``m:<block>``: ``stem``, ``s<stage>.b<block>``, ``context``, ``SceneNeck.<up>``,
``SceneSegHead.<up>``, ``head_tail``, ``class_decision``) and every conv / linear / transposed conv layer opens one
before (``c:<key>``) and after (``c:<key>|post``: the glue that follows it, e.g. the SE gate multiply after ``fc2``
or the residual add after a project conv) its call; ``--no-layer-signposts`` turns them off. Through the identical
op order, the ops of the replay are attributed to the same layers. One frame issues more programs than the
profiler's default 1000-program buffer, so run with ``--op-support-count`` above the program count of the largest
section. The output folder must be absolute and outside the bundle (PLAN.md 5.1).
"""
from __future__ import annotations

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

from tt_sceneseg import SceneSeg
from tt_sceneseg.reference import assets
from tt_sceneseg.ttaw.profiling import read_device_profiler, signpost, signposted

SAMPLE = Path(__file__).resolve().parents[1] / "tt_sceneseg" / "samples" / "highway_normal_1.png"


class _Signposted:
    """A callable stand-in for a layer that emits ``c:<name>`` before and ``c:<name>|post`` after each call (and
    around ``terms``, the un-summed call of a transposed conv)."""

    def __init__(self, inner, name: str):
        self.inner, self._name = inner, name

    def __call__(self, *args, **kwargs):
        signpost(f"c:{self._name}")
        out = self.inner(*args, **kwargs)
        signpost(f"c:{self._name}|post")
        return out

    def terms(self, *args, **kwargs):
        signpost(f"c:{self._name}")
        out = self.inner.terms(*args, **kwargs)
        signpost(f"c:{self._name}|post")
        return out

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


def layer_signposts(m) -> Callable[[], None]:
    """Shadow the block methods and the layers of a ``tt.model.TtSceneSeg`` instance with versions that emit Tracy
    signposts; returns the function that removes them again."""
    names: List[str] = []

    def wrap(attr: str, label: Callable) -> None:
        fn = getattr(m, attr)

        @functools.wraps(fn)
        def inner(*args, **kwargs):
            signpost(label(*args))
            return fn(*args, **kwargs)

        setattr(m, attr, inner)
        names.append(attr)

    wrap("stem", lambda *a: "m:stem")
    wrap("mbconv", lambda s, b, *a: f"m:s{s}.b{b}")
    wrap("se", lambda key, i, *a: f"c:{key}.{i}.se_sum")
    wrap("head_conv", lambda *a: "m:head_conv")
    wrap("context", lambda *a: "m:context")
    wrap("up_block", lambda mod, block, *a: f"m:{mod}.{block[0]}")
    wrap("head_tail", lambda *a: "m:head_tail")
    wrap("channel_major", lambda *a: "m:class_decision")
    saved = {}
    for attr in ("convs", "linears", "deconvs"):
        d = getattr(m, attr)
        saved[attr] = dict(d)
        for k in list(d):
            d[k] = _Signposted(d[k], k)
    for attr in ("stem_conv", "margin_conv", "logit_conv"):
        layer = getattr(m, attr)
        saved[attr] = layer
        setattr(m, attr, _Signposted(layer, layer.name))
    # S1 sub-pixel blocks: the Z conv(s) and the composed skip conv; the epilogue follows as "<last conv>|post"
    s1_saved = []
    for blk in getattr(m, "s1", {}).values():
        s1_saved.append((blk, list(blk.zconvs), blk.sconv))
        blk.zconvs = [_Signposted(zc, zc.name) for zc in blk.zconvs]
        if blk.sconv is not None:
            blk.sconv = _Signposted(blk.sconv, blk.sconv.name)

    def undo() -> None:
        for blk, zcs, sc in s1_saved:
            blk.zconvs, blk.sconv = zcs, sc
        for attr in names:
            delattr(m, attr)
        for attr, v in saved.items():
            if isinstance(v, dict):
                getattr(m, attr).clear()
                getattr(m, attr).update(v)
            else:
                setattr(m, attr, v)

    return undo


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--input", default=str(SAMPLE))
    ap.add_argument("--variant", default=None)
    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("--no-eager", dest="eager", action="store_false")
    ap.add_argument("--no-layer-signposts", dest="layers", action="store_false")
    a = ap.parse_args()
    import ttnn

    weights = assets.find_weights_dir()
    t0 = time.perf_counter()
    with SceneSeg.from_pretrained(variant=a.variant, dispatch=a.dispatch, num_command_queues=a.cqs,
                                  weights_dir=str(weights) if weights else None) as model:
        sr, dev = model.runner, model.device
        tr, var = sr.runner, sr.variant
        print("loaded in %.1f s:" % (time.perf_counter() - t0), json.dumps(model.device_info, default=str),
              "knobs", json.dumps(model.knobs.as_dict()), flush=True)
        canvas = model._prepare(a.input)["prepared"].device_input()
        sr(canvas)                                   # one served frame: upload + replay + readback
        ttnn.synchronize_device(dev)
        read_device_profiler(dev)                    # warm-up / capture / first frame out of the buffer
        if a.eager:
            undo = layer_signposts(sr.model) if a.layers else (lambda: None)
            try:
                with signposted("eager"):
                    sr.run_eager(canvas)
                    ttnn.synchronize_device(dev)
            finally:
                undo()
            read_device_profiler(dev)
        tr.upload({"image": sr._upload_value(canvas)})
        ttnn.synchronize_device(dev)
        read_device_profiler(dev)
        with signposted("trace"):
            tr.replay(var, n=a.replays)
            ttnn.synchronize_device(dev)
        read_device_profiler(dev)
        print("trace:", json.dumps(tr.describe(), default=str), flush=True)


if __name__ == "__main__":
    main()