#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Device profile of BEVFormer's traced replays, one signposted range per trace variant (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 4000 \\
        -o $ROOT/generated/profiler/bevformer_baseline code/scripts/profile_ops.py --rounds 2 --eager
    tt-perf-report <reports/<ts>/ops_perf_results_<ts>.csv> --start-signpost image --end-signpost image_end
    python code/scripts/profile_summary.py $ROOT/generated/profiler/bevformer_baseline --json s.json --md s.md

ttnn-visualizer memory + graph report of one eager run of each variant, one report directory per variant (slow; never
used for timing; a plain graph capture: ``full_graph_capture``'s stack traces + detailed buffer tracing needed more than
58 GB of host RAM for the image + bev graph):

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

The model is loaded like the API (both variants captured); three frames of one stream (the PandaSet f30 sample when
present, else the synthetic frame) bring the history to the steady state (``use_prev_bev`` = 1); the profiler
buffer is flushed; then each variant is replayed once between its own signposts (``image`` .. ``image_end``,
``bev`` .. ``bev_end``, inside ``trace`` .. ``trace_end``) with ``ttnn.ReadDeviceProfiler`` after each (the device
buffer is sized by ``--op-support-count``; a frame is ~1,500 programs). ``--rounds 2`` repeats the window
(``trace2`` / ``image2`` ...). ``--eager`` adds one eager (untraced, program cache hot) run of each variant between
``eager_image`` / ``eager_bev`` signposts with module-stage signposts inside (labels only): ``img.stem``,
``img.layer1..4``, ``img.neck``, ``img.values``; ``enc.queries``, ``enc{i}.tsa``, ``enc{i}.sca``, ``enc{i}.ffn``,
``enc.out``; ``dec.value_proj``, ``dec{i}.self``, ``dec{i}.cross``, ``dec{i}.ffn``, ``dec{i}.reg``, ``dec.heads``,
``dec.pack``. ``profile_summary.py`` carries the labels over to the traced ops when the op sequences match.
"""
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))

VARIANTS = ("image", "bev")


class _Marked:
    """Transparent wrapper of a callable: emits a Tracy signpost, calls it, optionally emits another after."""

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

    def __call__(self, *a, **k):
        if self._name:
            self._signpost(self._name)
        out = self._obj(*a, **k)
        if self._after:
            self._signpost(self._after)
        return out

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


def mark_stages(tt, signpost):
    """Stage signposts at the module boundaries of the ``image`` and ``bev`` functions (one eager run). Returns an
    undo callable."""
    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))

    def set_item(seq, i, value):
        old = seq[i]
        seq[i] = value
        undo.append(lambda: seq.__setitem__(i, old))

    img = tt.image
    set_attr(img, "stem", _Marked(img.stem, "img.stem", signpost))
    for s, blocks in enumerate(img.stages.stages):
        set_item(blocks, 0, _Marked(blocks[0], f"img.layer{s + 1}", signpost))
    set_attr(img, "neck", _Marked(img.neck, "img.neck", signpost))
    set_attr(img, "values", _Marked(img.values, "img.values", signpost))
    enc = tt.encoder
    set_attr(enc, "bev_queries", _Marked(enc.bev_queries, "enc.queries", signpost))
    for i, L in enumerate(enc.layers):
        for key, label in (("tsa", "tsa"), ("sca", "sca"), ("ffn1", "ffn")):
            old = L[key]
            L[key] = _Marked(old, f"enc{i}.{label}", signpost)
            undo.append(lambda L=L, key=key, old=old: L.__setitem__(key, old))
    set_attr(tt, "encoder", _Marked(enc, "", signpost, after="enc.out"))
    dec = tt.decoder
    set_attr(dec.layers[0], "value_proj", _Marked(dec.layers[0].value_proj, "dec.value_proj", signpost))
    counter = {"self": 0, "cross": 0}

    def labelled(kind, fn):
        def wrapper(*a, **k):
            signpost(f"dec{counter[kind]}.{kind}")
            counter[kind] += 1
            return fn(*a, **k)
        return wrapper

    set_attr(dec, "self_attention", labelled("self", dec.self_attention))
    set_attr(dec, "cross_attention", labelled("cross", dec.cross_attention))
    for i, L in enumerate(dec.layers):
        set_attr(L, "ffn1", _Marked(L.ffn1, f"dec{i}.ffn", signpost))
        set_item(L.reg, 0, _Marked(L.reg[0], f"dec{i}.reg", signpost))
    set_attr(dec, "heads", _Marked(dec.heads, "dec.heads", signpost))
    set_attr(tt, "decoder", _Marked(dec, "", signpost, after="dec.pack"))

    def restore():
        for fn in reversed(undo):
            fn()
        counter.update(self=0, cross=0)
    return restore


def steady_request():
    """Kwargs of ``model(...)`` for the steady-state frames: the PandaSet f30 sample if present, else synthetic."""
    from bench import pandaset_dir

    from tt_bevformer import load_sample
    from tt_bevformer.tt.warmup import synthetic_request

    p = pandaset_dir() / "pandaset_019_f30.json"
    if p.is_file():
        req = load_sample(p)
        st = req.pop("stream", {}) or {}
        req["stream"] = {"id": "profile", "T_world_from_ego": st.get("T_world_from_ego")}
        return req, str(p)
    req = synthetic_request()
    req["stream"] = {"id": "profile"}
    return req, "synthetic"


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 "
                                                         "image + bev to this directory instead of profiling")
    ap.add_argument("--graph-variants", default="image,bev", help="variants captured by --graph-report, one report "
                                                                 "directory each (<dir>/<variant>)")
    ap.add_argument("--graph-full", action="store_true", help="ttnn.graph.full_graph_capture (Python stack traces + "
                    "detailed buffer tracing): above 58 GB of host RAM for this model, so the default is a plain "
                    "graph capture (ops, tensors, buffer allocations)")
    ap.add_argument("--max-rss-gb", type=float, default=30.0, help="abort the graph capture above this host RSS")
    ap.add_argument("--dispatch", default=None, choices=[None, "eth", "worker"])
    ap.add_argument("--num-cqs", type=int, 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_bevformer import BEVFormer
    from tt_bevformer.ttaw.profiling import read_device_profiler, signpost

    t0 = time.perf_counter()
    model = BEVFormer.from_pretrained(dispatch=a.dispatch, num_command_queues=a.num_cqs)
    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
    db_dir = None
    try:
        req, meta["input"] = steady_request()
        for _ in range(3):
            out = model(**req)
        meta["detections"] = len(out.scores)
        sync()
        if a.graph_report:
            import contextlib
            import os
            import threading

            def guard(stop):                         # the capture lives in host RAM: never let it take the host down
                page = os.sysconf("SC_PAGE_SIZE")
                while not stop.wait(1.0):
                    with open("/proc/self/statm") as fh:
                        rss = int(fh.read().split()[1]) * page
                    if rss > a.max_rss_gb * 2 ** 30:
                        print(f"graph capture: host RSS {rss / 2 ** 30:.1f} GB > {a.max_rss_gb} GB, abort", flush=True)
                        os._exit(3)

            @contextlib.contextmanager
            def plain_capture(path):
                prev = ttnn.CONFIG.enable_fast_runtime_mode
                ttnn.CONFIG.enable_fast_runtime_mode = False
                try:
                    ttnn.graph.begin_graph_capture()
                    yield
                    ttnn.graph.end_graph_capture_to_file(str(path))
                finally:
                    ttnn.CONFIG.enable_fast_runtime_mode = prev

            for v in VARIANTS:                                   # program cache hot, constants uploaded
                r.run_eager(v)
            sync()
            db_dir = []
            stop = threading.Event()
            threading.Thread(target=guard, args=(stop,), daemon=True).start()
            try:
                for v in [x.strip() for x in a.graph_variants.split(",") if x.strip()]:
                    out_dir = Path(a.graph_report).resolve() / v
                    out_dir.mkdir(parents=True, exist_ok=True)
                    json_path = out_dir / "graph_capture.json"
                    t1 = time.perf_counter()
                    cm = ttnn.graph.full_graph_capture(str(json_path)) if a.graph_full else plain_capture(json_path)
                    with cm:
                        r.run_eager(v)
                        sync()
                    meta[f"graph_capture_s.{v}"] = round(time.perf_counter() - t1, 1)
                    meta[f"graph_json_mb.{v}"] = round(json_path.stat().st_size / 1e6, 1)
                    print("graph captured:", v, meta[f"graph_capture_s.{v}"], "s,", meta[f"graph_json_mb.{v}"], "MB",
                          flush=True)
                    db_dir.append((out_dir, json_path))
            finally:
                stop.set()
        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}")
                for v in VARIANTS:
                    signpost(f"{v}{tag}")
                    r.replay(v)
                    sync()
                    signpost(f"{v}{tag}_end")
                    read_device_profiler(dev)
                signpost(f"trace{tag}_end")
            if a.eager:
                restore = mark_stages(model.tt, signpost)
                try:
                    signpost("eager")
                    for v in VARIANTS:
                        signpost(f"eager_{v}")
                        r.run_eager(v)
                        sync()
                        signpost(f"eager_{v}_end")
                        read_device_profiler(dev)
                    signpost("eager_end")
                finally:
                    restore()
        meta["trace"] = r.describe()
    finally:
        model.close()
    for out_dir, json_path in db_dir or []:
        t1 = time.perf_counter()
        db = ttnn.graph_report.import_report(json_path, out_dir)
        ttnn.save_config_to_json_file(out_dir / "config.json")
        cfg = json.loads((out_dir / "config.json").read_text())
        cfg.update({"enable_fast_runtime_mode": False, "enable_graph_report": True,
                    "enable_detailed_buffer_report": bool(a.graph_full), "root_report_path": str(out_dir.parent),
                    "report_name": out_dir.name})
        (out_dir / "config.json").write_text(json.dumps(cfg, indent=4) + "\n")
        meta.setdefault("graph_report", {})[out_dir.name] = {
            "dir": str(out_dir), "db": str(db), "import_s": round(time.perf_counter() - t1, 1),
            "db_mb": round(Path(db).stat().st_size / 1e6, 1)}
        json_path.unlink(missing_ok=True)
        print("graph report:", json.dumps(meta["graph_report"][out_dir.name]), flush=True)
    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())