#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""ttnn-visualizer memory + graph report of ONE eager forward of the served graph (default numerics, 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/sceneseg-p150
    $ROOT/bin/devrun -t 2400 -- env PYTHONPATH=$PWD/code:$PYTHONPATH HF_HUB_OFFLINE=1 \
        python code/scripts/visualizer_report.py $ROOT/generated/ttnn_visualizer/sceneseg_baseline
    ttnn-visualizer --profiler-path $ROOT/generated/ttnn_visualizer/sceneseg_baseline \
        --performance-path $ROOT/generated/profiler/sceneseg_baseline/reports/<ts>

What the pytest fixture of ttnn does (ttnn/graph_report.py ``_run_pytest_graph_report_fixture``), in a script:
warm-up forward (JIT, conv weight preparation), then the second forward 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 large and is deleted after the
import; ``--no-stack-traces`` leaves the Python stack traces out of the capture (a smaller JSON).
"""
from __future__ import annotations

import json
import sys
import time
from pathlib import Path


def capture(json_path: Path, stack_traces: bool):
    """``ttnn.graph.full_graph_capture`` (Python I/O, stack traces, detailed buffer tracing, slow dispatch), or the
    same without Python stack traces (``--no-stack-traces``: a much smaller graph JSON on a full disk)."""
    import contextlib

    import ttnn

    if stack_traces:
        return ttnn.graph.full_graph_capture(str(json_path))

    @contextlib.contextmanager
    def light():
        g = ttnn.graph
        prev_buf, prev_fast = g.is_detailed_buffer_tracing_enabled(), ttnn.CONFIG.enable_fast_runtime_mode
        g.enable_detailed_buffer_tracing()
        ttnn.CONFIG.enable_fast_runtime_mode = False
        try:
            g.begin_graph_capture(g.RunMode.NORMAL)
            yield
            g.end_graph_capture_to_file(str(json_path))
        finally:
            if g.is_graph_capture_active():
                try:
                    g._cpp_end_graph_capture()
                except Exception:
                    pass
            if not prev_buf:
                g.disable_detailed_buffer_tracing()
            ttnn.CONFIG.enable_fast_runtime_mode = prev_fast

    return light()


def main() -> None:
    stack_traces = "--no-stack-traces" not in sys.argv
    out_dir = Path([a for a in sys.argv[1:] if not a.startswith("--")][0]).resolve()
    out_dir.mkdir(parents=True, exist_ok=True)
    import torch
    import ttnn

    from tt_sceneseg.device import close_device, describe_device, open_device
    from tt_sceneseg.host.preprocess import prepare_image
    from tt_sceneseg.reference import config
    from tt_sceneseg.reference.weights import load_weights
    from tt_sceneseg.reference.assets import find_in, find_weights_dir
    from tt_sceneseg.tt.model import INPUT_SHAPE, TtSceneSeg

    sample = Path(__file__).resolve().parents[1] / "tt_sceneseg" / "samples" / "highway_normal_1.png"
    canvas = prepare_image(str(sample)).device_input()
    weights = load_weights(find_in(find_weights_dir(True), config.ONNX_FILE))
    dev = open_device(allow_fallback=False)
    try:
        print("device:", json.dumps(describe_device(dev, 0)), flush=True)
        model = TtSceneSeg(weights, config.DEFAULT_CHANNEL_ORDER)
        x = ttnn.from_torch(torch.from_numpy(canvas.reshape(INPUT_SHAPE)), dtype=ttnn.uint8,
                            layout=ttnn.ROW_MAJOR_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG)

        def forward():
            out = model.forward(x)
            ttnn.synchronize_device(dev)
            for t in [out["class_map"]]:
                if t.is_allocated():
                    ttnn.deallocate(t)

        t0 = time.perf_counter()
        forward()                                       # warm-up: JIT + weight preparation (host writes)
        print(f"warm-up forward {time.perf_counter() - t0:.1f} s", flush=True)
        json_path = out_dir / "graph_capture.json"
        t0 = time.perf_counter()
        with capture(json_path, stack_traces):
            forward()
        print(f"captured forward {time.perf_counter() - t0:.1f} s -> {json_path} "
              f"({json_path.stat().st_size / 1e6:.1f} MB)", flush=True)
        ttnn.deallocate(x)
        model.release()
    finally:
        close_device(dev)
    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": stack_traces,
                "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()