#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Device profile of StreamPETR's traced ``frame`` replay between signposts (OPT_BASELINE.md).

Tracy + device profiler (absolute ``-o`` under ``$ROOT/generated/profiler``, never inside the bundle):

    bin/devrun -t 3600 -- python -m tracy -r -p -v --no-web-server --op-support-count 6000 \\
        -o $ROOT/generated/profiler/streampetr_baseline code/scripts/profile_ops.py --rounds 2 --eager
    python code/scripts/profile_summary.py $ROOT/generated/profiler/streampetr_baseline --json ... --md ...
    tt-perf-report <reports/<ts>/ops_perf_results_<ts>.csv> --start-signpost frame --end-signpost frame_end --arch p150b

ttnn-visualizer memory + graph report of one eager frame (slow; never used for timing):

    bin/devrun -t 3600 -- python code/scripts/profile_ops.py --graph-report $ROOT/generated/ttnn_reports/streampetr_baseline
    ttnn-visualizer --profiler-path $ROOT/generated/ttnn_reports/streampetr_baseline \\
        --performance-path $ROOT/generated/profiler/streampetr_baseline/reports/<ts>

The model is loaded like the API (the ``frame`` trace captured), four frames of the PandaSet sample (or synthetic
frames when it is absent) bring the stream's temporal memory to the steady state, the profiler buffer is flushed,
then the ``frame`` trace is replayed once between the signposts ``frame`` .. ``frame_end`` (inside ``trace`` ..
``trace_end``) and ``ttnn.ReadDeviceProfiler`` reads the buffer. ``--rounds 2`` repeats the window (``frame2``) to
check that the profile is stable; ``--eager`` adds one eager (untraced, program cache hot) frame between
``eager_frame`` / ``eager_frame_end`` with module-stage signposts (``bb.stem`` ... ``head.update``; the op sequence
is matched op for op against the traced one by ``profile_summary.py``).
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

HERE = Path(__file__).resolve()
sys.path.insert(0, str(HERE.parents[1]))
sys.path.insert(0, str(HERE.parent))


class _Marked:
    """Transparent wrapper of a callable: emits a Tracy signpost, then calls it (host-side only)."""

    def __init__(self, obj, name: str, signpost):
        self._obj, self._name, self._signpost = obj, name, signpost

    def __call__(self, *a, **k):
        self._signpost(self._name)
        return self._obj(*a, **k)

    def __getattr__(self, item):
        return getattr(self._obj, item)


def mark_stages(tt, signpost):
    """Signposts at the module boundaries of the ``frame`` function (one eager run). Returns an undo callable.
    Labels: bb.stem, bb.stage2 .. bb.stage5 (each from its max-pool on), bb.neck; head.tokens, head.cross_kv,
    head.temporal, head.queries, head.dec<L>.self / .cross / .ffn (L = 0..5; each label starts at that block's first
    projection, so the preceding residual add belongs to the previous label), head.branches, head.select,
    head.update."""
    undo = []

    def set_attr(obj, name, value):
        had = name in obj.__dict__
        old = obj.__dict__.get(name)
        setattr(obj, name, value)
        undo.append(lambda: setattr(obj, name, old) if had else obj.__dict__.pop(name, None))

    bb, hd = tt.backbone, tt.head
    stem0 = bb.stem[0]
    bb.stem[0] = _Marked(stem0, "bb.stem", signpost)
    undo.append(lambda: bb.stem.__setitem__(0, stem0))
    b20 = bb.stages[2][0]
    bb.stages[2][0] = _Marked(b20, "bb.stage2", signpost)
    undo.append(lambda: bb.stages[2].__setitem__(0, b20))
    pool, count = bb.pool, [2]

    def marked_pool(x):
        count[0] += 1
        signpost(f"bb.stage{count[0]}")
        return pool(x)
    set_attr(bb, "pool", marked_pool)
    set_attr(bb, "lateral0", _Marked(bb.lateral0, "bb.neck", signpost))
    for name in ("tokens", "cross_kv", "temporal", "queries", "branches", "select", "update"):
        set_attr(hd, name, _Marked(getattr(hd, name), f"head.{name}", signpost))
    for li, L in enumerate(hd.layers):
        for key, lab in (("self.q", "self"), ("cross.q", "cross"), ("ffn.0", "ffn")):
            old = L[key]
            L[key] = _Marked(old, f"head.dec{li}.{lab}", signpost)
            undo.append(lambda L=L, key=key, old=old: L.__setitem__(key, old))

    def restore():
        for fn in reversed(undo):
            fn()
    return restore


def steady_requests(model):
    """Five frames of one stream: the PandaSet sample f40-43 (+ f40 of the next cycle) when present, else synthetic
    frames; the first four run through the API, the fifth gives the eager / graph-report inputs."""
    import bench

    d = bench.pandaset_dir()
    paths = [d / f"pandaset_019_f{n}.json" for n in (40, 41, 42, 43)]
    frames = bench.manifest_frames(paths) if all(p.is_file() for p in paths) else bench.synthetic_frames("pandaset_019")
    seq = bench.Sequencer(frames, "profile")
    return [seq(i) for i in range(5)]


def frame_inputs(model, req):
    """Host values of one frame for ``run_eager`` (the API's ``_forward`` without the replay)."""
    from tt_streampetr.tt.params import age_input

    prepared = model._prepare(**req)
    plan = prepared["plan"]
    params = model.tt.geometry_params(plan.M)
    params.update(prepared["calib"]["params"])
    inputs = {"img": model._img_input(prepared), "ages": age_input(plan.pre["pre_memory_timestamp"], device_shape=True)}
    return inputs, params


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--rounds", type=int, default=1)
    ap.add_argument("--eager", action="store_true")
    ap.add_argument("--graph-report", default=None, help="write a ttnn-visualizer memory/graph report of one eager "
                                                         "frame to this directory instead of profiling")
    ap.add_argument("--dispatch", default=None, choices=[None, "eth", "worker"])
    ap.add_argument("--num-cqs", type=int, default=None)
    ap.add_argument("--weights-dir", default=None)
    ap.add_argument("--json", default=None, help="write the run's metadata (config, trace description)")
    a = ap.parse_args()
    import ttnn

    from tt_streampetr import StreamPETR
    from tt_streampetr.ttaw.profiling import read_device_profiler, signpost

    t0 = time.perf_counter()
    model = StreamPETR.from_pretrained(dispatch=a.dispatch, num_command_queues=a.num_cqs, weights_dir=a.weights_dir)
    meta = {"config": model.device_info, "knobs": model.knobs.as_dict(), "load_s": round(time.perf_counter() - t0, 1)}
    print("config:", json.dumps(meta["config"]), flush=True)
    dev, r = model.device, model.runner
    sync = lambda: ttnn.synchronize_device(dev)  # noqa: E731
    try:
        reqs = steady_requests(model)
        for req in reqs[:4]:                          # steady state: a cold start, then three carried frames
            model(**req)
        sync()
        inputs, params = frame_inputs(model, reqs[4])
        if a.graph_report:
            out_dir = Path(a.graph_report)
            out_dir.mkdir(parents=True, exist_ok=True)
            r.run_eager("frame", inputs, params)      # program cache hot: nothing compiles inside the capture
            sync()
            t1 = time.perf_counter()
            with ttnn.graph.full_graph_capture(str(out_dir / "graph_capture.json")):
                r.run_eager("frame", inputs, params)
                sync()
            meta["graph_capture_s"] = round(time.perf_counter() - t1, 1)
            from ttnn import graph_report

            db = graph_report.import_report(out_dir, out_dir)
            ttnn.save_config_to_json_file(str(out_dir / "config.json"))
            meta["graph_report"] = {"dir": str(out_dir), "db": str(db)}
            raw = out_dir / "graph_capture.json"
            if raw.is_file():
                meta["graph_report"]["raw_json_bytes"] = raw.stat().st_size
                raw.unlink()
            print("graph report:", json.dumps(meta["graph_report"]), flush=True)
        else:
            read_device_profiler(dev)                 # drop the load / warm-up data
            for k in range(1, a.rounds + 1):
                tag = "" if k == 1 else str(k)
                signpost(f"trace{tag}")
                signpost(f"frame{tag}")
                r.replay("frame")
                sync()
                signpost(f"frame{tag}_end")
                read_device_profiler(dev)
                signpost(f"trace{tag}_end")
            if a.eager:
                restore = mark_stages(model.tt, signpost)
                try:
                    signpost("eager_frame")
                    r.run_eager("frame", inputs, params)
                    sync()
                    signpost("eager_frame_end")
                    read_device_profiler(dev)
                finally:
                    restore()
        meta["trace"] = r.describe()
    finally:
        model.close()
    if a.json:
        Path(a.json).parent.mkdir(parents=True, exist_ok=True)
        Path(a.json).write_text(json.dumps(meta, indent=1, default=str) + "\n")
    return 0


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