#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""ttnn-visualizer memory + graph report of ONE eager forward of the served graph (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/transfusion-p150
    $ROOT/bin/devrun -t 2400 -- env PYTHONPATH=$PWD/code:$PYTHONPATH \\
        python code/scripts/visualizer_report.py $ROOT/generated/ttnn_visualizer/transfusion_baseline
    ttnn-visualizer --profiler-path $ROOT/generated/ttnn_visualizer/transfusion_baseline \\
        --performance-path $ROOT/generated/profiler/transfusion_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, weight upload), 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 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_transfusion" / "samples" / "test_pcd.npz"


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("out_dir")
    ap.add_argument("--variant", default=None)
    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_transfusion.api import TransFusion
    from tt_transfusion.device import close_device, describe_device, open_device
    from tt_transfusion.host.preprocess import TransFusionPreprocessor
    from tt_transfusion.reference.config import DEFAULT_VARIANT, TransFusionConfig
    from tt_transfusion.reference.weights import load_weights
    from tt_transfusion.tt.model import TtTransFusion, device_inputs
    from tt_transfusion.ttaw.tensors import to_host_tensor

    weights_path = TransFusion.resolve_weights(weights_dir=a.weights_dir)
    config = TransFusionConfig.from_weights(weights_path, a.variant or DEFAULT_VARIANT)
    prepared = TransFusionPreprocessor(config)(a.input)
    dev = open_device(allow_fallback=False)
    try:
        print("device:", json.dumps(describe_device(dev, 0)), flush=True)
        graph = TtTransFusion(dev, load_weights(weights_path, config), config)     # the default knobs / policy
        graph.query_head.prepare(dev)
        host = device_inputs(prepared, graph.num_pillars)
        feats = ttnn.allocate_tensor_on_device(ttnn.Shape(list(graph.features_shape)), ttnn.bfloat16,
                                               ttnn.ROW_MAJOR_LAYOUT, dev, ttnn.DRAM_MEMORY_CONFIG)
        index = ttnn.allocate_tensor_on_device(ttnn.Shape(list(graph.index_shape)), ttnn.uint32,
                                               ttnn.ROW_MAJOR_LAYOUT, dev, ttnn.DRAM_MEMORY_CONFIG)
        ttnn.copy_host_to_device_tensor(to_host_tensor(host["features"], "bfloat16", ttnn.ROW_MAJOR_LAYOUT), feats)
        ttnn.copy_host_to_device_tensor(to_host_tensor(host["canvas_index"], "uint32", ttnn.ROW_MAJOR_LAYOUT), index)

        def forward():
            out = graph.forward(feats, index).tensor            # ttaw.trace.Packed: the one output row
            ttnn.synchronize_device(dev)
            if out.is_allocated():
                ttnn.deallocate(out)

        t0 = time.perf_counter()
        forward()                                       # warm-up: JIT + weight upload (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 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)
        ttnn.deallocate(feats)
        ttnn.deallocate(index)
        graph.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": 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()