#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Device profile of frnet-p150: one eager forward and one traced replay of the whole graph between Tracy
signposts, for one serve profile (``--variant ot128 | qt128``).

    ROOT=/home/ubuntu/experiments/tt-models
    $ROOT/bin/devrun -t 3600 -- python -m tracy -r -p -v --no-web-server --op-support-count 4000 \\
        -o $ROOT/generated/profiler/frnet_baseline code/scripts/profile_ops.py --variant ot128
    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 plan of ``--input``). The
device profiler buffer is flushed (``ttnn.ReadDeviceProfiler``) before and after each section. Inside the eager
section every module of the graph opens a signpost (``m:<module>``: ``encoder``, ``segmax``, ``frustum2pixel``,
``pixel2point``, ``stem``, ``point_stem``, ``fusion_stem``, ``layer<l>``, ``point_fusion<l>``, ``pixel_fusion<l>``,
``attention<l>``, ``decoder.up``, ``decoder.fuse1``, ``decoder.fuse2``, ``point_fuse``, ``head``) and every layer a
layer signpost (``c:<layer>``; ``--no-layer-signposts`` turns them off). Ops that the graph issues inline after a
module (the residual ``add`` + ``hardswish`` of a BasicBlock, ``addcmul`` after the attention, the softmax after the
logits) belong to the module signposted last. Through the identical op order the ops of the replay inherit the
eager attribution (``profile_summary.py``). The output folder must be absolute and outside the bundle (PLAN.md 5.1).
"""
from __future__ import annotations

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

from tt_frnet import FRNet
from tt_frnet.ttaw.profiling import read_device_profiler, signpost, signposted

CODE = Path(__file__).resolve().parents[1]
SAMPLES = CODE / "tt_frnet" / "samples"           # the shipped synthetic_<variant>.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]):
        self.inner, self._labels = inner, list(labels)

    def __call__(self, *args, **kwargs):
        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) -> Callable[[], None]:
    """Wrap the modules of a ``tt.model.TtFRNet`` (and the C19 gathers its forward calls) with signposting
    stand-ins; returns the function that restores them."""
    from tt_frnet.tt import model as tt_model

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

    def labels(module: str, layer: str) -> List[str]:
        return [f"m:{module}"] + ([f"c:{layer}"] if layers else [])

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

    def swap_item(seq, key, lab):
        inner = seq[key]
        seq[key] = _Signposted(inner, lab)
        undo.append(lambda: seq.__setitem__(key, inner))

    for fn_name in ("frustum2pixel", "pixel2point"):
        orig = getattr(tt_model, fn_name)

        def wrapped(*args, _orig=orig, _name=fn_name, **kwargs):
            signpost(f"m:{_name}")
            return _orig(*args, **kwargs)

        setattr(tt_model, fn_name, wrapped)
        undo.append(lambda _n=fn_name, _o=orig: setattr(tt_model, _n, _o))
    for i in range(len(graph.ffe)):
        swap_item(graph.ffe, i, labels("encoder", f"encoder.ffe{i}"))
    swap(graph, "compression", labels("encoder", "encoder.compression"))
    swap(graph, "segmax_encoder", labels("segmax", "segmax.encoder"))
    swap(graph, "segmax_points", labels("segmax", "segmax.points"))
    for i in range(len(graph.stem)):
        swap_item(graph.stem, i, labels("stem", f"backbone.stem{i}"))
    swap(graph, "point_stem", labels("point_stem", "backbone.point_stem"))
    swap(graph, "fusion_stem", labels("fusion_stem", "backbone.fusion_stem"))
    for st in graph.stages:
        l = st["index"]
        blocks = st["blocks"]
        for b, (c1, c2, ds) in enumerate(blocks):
            pre = f"backbone.layer{l}.block{b}"
            new = (_Signposted(c1, labels(f"layer{l}", f"{pre}.conv1")),
                   _Signposted(c2, labels(f"layer{l}", f"{pre}.conv2")),
                   None if ds is None else _Signposted(ds, labels(f"layer{l}", f"{pre}.downsample")))
            blocks[b] = new
            undo.append(lambda _bl=blocks, _b=b, _o=(c1, c2, ds): _bl.__setitem__(_b, _o))
        swap_item(st, "point_fusion", labels(f"point_fusion{l}", f"backbone.point_fusion{l}"))
        swap_item(st, "pixel_fusion", labels(f"pixel_fusion{l}", f"backbone.pixel_fusion{l}"))
        for j in range(len(st["attention"])):
            swap_item(st["attention"], j, labels(f"attention{l}", f"backbone.attention{l}.conv{j}"))
    for s in list(graph.ups):
        swap_item(graph.ups, s, labels("decoder.up", f"decoder.up{s}"))
    swap(graph, "fuse1", labels("decoder.fuse1", "decoder.fuse1"))
    swap(graph, "fuse2", labels("decoder.fuse2", "decoder.fuse2"))
    swap(graph, "point_fuse1", labels("point_fuse", "decoder.point_fuse1"))
    swap(graph, "point_fuse2", labels("point_fuse", "decoder.point_fuse2"))
    for i in range(len(graph.mlps)):
        swap_item(graph.mlps, i, labels("head", f"head.mlp{i}"))
    swap(graph, "logits", labels("head", "head.logits"))

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

    return restore


def calibration_for(path: Path, spec=None) -> dict:
    """The ``calibration`` of ``path`` (the rules of ``bench.py calibration_for``)."""
    if spec:
        if spec == "none":
            return {"ego_crop_box": None}
        if spec.startswith("preset:"):
            return {"preset": spec[len("preset:"):]}
        return json.loads(Path(spec).read_text())
    stem = path.name.split(".")[0]
    if (CODE / "tt_frnet" / "calib" / f"{stem}.json").is_file():
        return {"preset": stem}
    side = path.parent / f"{stem}.calib.json"
    return json.loads(side.read_text()) if side.is_file() else {"ego_crop_box": None}


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--input", default=None, help="default: the shipped synthetic sample of the variant")
    ap.add_argument("--calib", default=None, help="preset:<name> | <calibration .json> | none (bench.py rules)")
    ap.add_argument("--variant", default=None, help="ot128 | qt128 (default: FRNET_VARIANT or ot128)")
    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

    src = Path(a.input or SAMPLES / f"synthetic_{a.variant or os.environ.get('FRNET_VARIANT', 'ot128')}.npz")
    calib = calibration_for(src, a.calib)
    t0 = time.perf_counter()
    with FRNet.from_pretrained(variant=a.variant, dispatch=a.dispatch, num_command_queues=a.cqs) 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=str(src), calibration=calib)
        out = model(str(src), calibration=calib)       # one served frame: upload + replay + readback
        print(f"variant {model.variant}: N {prepared.plan.n} M {prepared.plan.m} (N_cap {prepared.plan.n_cap}, "
              f"M_cap {prepared.plan.m_cap}); {len(out.label_ids)} labelled points", flush=True)
        ttnn.synchronize_device(dev)
        read_device_profiler(dev)                      # warm-up / capture / first frame out of the buffer
        inputs = model.staging.write(prepared.plan)
        if a.eager:
            restore = module_signposts(model.graph, layers=a.layers)
            try:
                with signposted("eager"):
                    tr.run_eager("default", 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"):
            tr.replay("default", n=a.replays)
            ttnn.synchronize_device(dev)
        read_device_profiler(dev)
        print("trace:", json.dumps(tr.describe()), flush=True)


if __name__ == "__main__":
    main()