#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""ttnn-visualizer memory + graph report of ONE eager forward of the served graph (trace A ``encseg`` and, variant
``seg_det``, trace B ``det``; default knobs, 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/ptv3-p150
    $ROOT/bin/devrun -t 3600 -- env PYTHONPATH=$PWD/code:$PYTHONPATH \\
        python code/scripts/visualizer_report.py $ROOT/generated/ttnn_visualizer/ptv3_baseline --input <cloud>
    ttnn-visualizer --profiler-path $ROOT/generated/ttnn_visualizer/ptv3_baseline \\
        --performance-path $ROOT/generated/profiler/ptv3_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 loaded as served (traces captured, every program compiled), one eager forward of each trace variant
(``TraceRunner.run_eager``) runs 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

import numpy as np

CODE = Path(__file__).resolve().parents[1]
SAMPLES = CODE / "tt_ptv3" / "samples"


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("out_dir")
    ap.add_argument("--variant", default=None, choices=["seg_det", "seg"])
    ap.add_argument("--input", default=None)
    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_ptv3 import PTv3
    from tt_ptv3 import io as tio
    from tt_ptv3.host import preprocess as pre
    from tt_ptv3.tt import inputs as tin

    src = a.input or next(str(SAMPLES / n) for n in ("test_pcd.npz",)
                          if (SAMPLES / n).is_file())
    json_path = out_dir / "graph_capture.json"
    with PTv3.from_pretrained(variant=a.variant, weights_dir=a.weights_dir) as model:
        print("device:", json.dumps(model.device_info, default=str), flush=True)
        tr = model.runner
        frame = pre.prepare_frame(np.asarray(tio.load_points(src).points, np.float32))
        b, caps = model.bucket_for(frame.counts)
        tables = {f"{b}.{k}": v for k, v in tin.frame_inputs(frame, caps, det=model.with_det).items()}
        variants = [f"encseg_{b}"] + ([f"det_{b}"] if model.with_det else [])

        def forward():
            for v in variants:
                tr.run_eager(v, inputs=tables if v.startswith("encseg") else None)
            ttnn.synchronize_device(model.device)

        t0 = time.perf_counter()
        forward()                                       # warm (programs are compiled by the load already)
        print(f"eager forward {time.perf_counter() - t0:.1f} s", flush=True)
        t0 = time.perf_counter()
        with ttnn.graph.full_graph_capture(str(json_path)):
            forward()
        print(f"captured forward {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()