#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""ttnn-visualizer memory + graph report of ONE eager frame of the served graph (``sparse_<bucket>`` then ``dense``;
default precision policy, ETH dispatch, the device grid). Never use it for timing: graph capture and detailed buffer
tracing synchronise after every op.

    ROOT=/home/ubuntu/experiments/tt-models; cd $ROOT/bundles/bevfusion-p150
    $ROOT/bin/devrun -t 2400 -- env PYTHONPATH=$PWD/code:$PYTHONPATH \\
        python code/scripts/visualizer_report.py $ROOT/generated/ttnn_visualizer/bevfusion_baseline
    ttnn-visualizer --profiler-path $ROOT/generated/ttnn_visualizer/bevfusion_baseline \\
        --performance-path $ROOT/generated/profiler/bevfusion_baseline/reports/<ts>

What the pytest fixture of ttnn does (ttnn/graph_report.py ``_run_pytest_graph_report_fixture``), in a script: the
model is built without capturing traces (``warmup_variants="none"``), a warm-up eager frame (JIT, weight upload),
then the second eager frame (``TraceRunner.run_eager`` of both segments) inside ``ttnn.graph.full_graph_capture``
(Python I/O, stack traces, detailed buffer tracing), ``ttnn.graph_report.import_report`` -> ``db.sqlite``, and
``config.json`` (the settings the capture ran with). The intermediate graph JSON is deleted after the import.
"""
from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

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


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("out_dir")
    ap.add_argument("--input", default=str(SAMPLE))
    ap.add_argument("--weights-dir", default=None)
    a = ap.parse_args()
    out_dir = Path(a.out_dir).resolve()
    out_dir.mkdir(parents=True, exist_ok=True)
    import ttnn

    from tt_bevfusion import BEVFusion
    from tt_bevfusion.tt.model import sparse_variant

    json_path = out_dir / "graph_capture.json"
    with BEVFusion.from_pretrained(weights_dir=a.weights_dir, warmup_variants="none") as model:
        print("device:", json.dumps(model.device_info), flush=True)
        tr, dev, g = model.runner, model.device, model.graph
        prepared = model._prepare(points=a.input)
        plan = prepared.plan
        sv = sparse_variant(plan.bucket)
        inputs = g.frame_inputs(prepared.vfe, plan, plan.bucket)

        def frame():
            tr.run_eager(sv, inputs=inputs)
            tr.run_eager("dense")
            ttnn.synchronize_device(dev)

        t0 = time.perf_counter()
        frame()                                          # warm-up: JIT + weight upload (host writes)
        print(f"warm-up frame {time.perf_counter() - t0:.1f} s (bucket {plan.bucket})", flush=True)
        t0 = time.perf_counter()
        with ttnn.graph.full_graph_capture(str(json_path)):
            frame()
        print(f"captured frame {time.perf_counter() - t0:.1f} s -> {json_path} "
              f"({json_path.stat().st_size / 1e6:.1f} MB)", flush=True)
    t0 = 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())     # record what the capture ran with
    cfg.update({"enable_fast_runtime_mode": False, "enable_graph_report": True, "enable_graph_python_stack_traces": True,
                "enable_detailed_buffer_report": True, "root_report_path": str(out_dir.parent),
                "report_name": out_dir.name})
    (out_dir / "config.json").write_text(json.dumps(cfg, indent=4) + "\n")
    print(f"imported in {time.perf_counter() - t0:.1f} s -> {db} ({Path(db).stat().st_size / 1e6:.1f} MB)", flush=True)
    json_path.unlink(missing_ok=True)


if __name__ == "__main__":
    main()