#!/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 and
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/frnet-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/frnet_baseline [--variant ot128]
    ttnn-visualizer --profiler-path $ROOT/generated/ttnn_visualizer/frnet_baseline \\
        --performance-path $ROOT/generated/profiler/frnet_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. The inputs are the
device plan of ``--input`` (default: the shipped synthetic sample of the variant with its calibration preset).
"""
from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

CODE = Path(__file__).resolve().parents[1]
SAMPLES = CODE / "tt_frnet" / "samples"           # the shipped synthetic_<variant>.npz


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("out_dir")
    ap.add_argument("--variant", default="ot128")
    ap.add_argument("--input", default=None, help="default: the shipped synthetic sample of the variant")
    ap.add_argument("--calib", default=None, help="preset:<name> | <calibration .json> | none (default: the preset "
                                                    "named like the input, else crop disabled)")
    ap.add_argument("--json-dir", default=None, help="where the intermediate graph JSON goes (default: out_dir; "
                                                      "it can be several GB, e.g. /dev/shm/<name>); deleted after the import")
    a = ap.parse_args()
    out_dir = Path(a.out_dir).resolve()
    out_dir.mkdir(parents=True, exist_ok=True)
    import ttnn

    from tt_frnet.api import FRNet
    from tt_frnet.device import close_device, describe_device, open_device
    from tt_frnet.host.frame import FRNetPreprocessor
    from tt_frnet.reference.config import FRNetConfig
    from tt_frnet.reference.weights import load_weights
    from tt_frnet.tt.layers import free
    from tt_frnet.tt.model import TtFRNet, default_policy, frame_arrays
    from tt_frnet.ttaw.tensors import to_host_tensor

    src = Path(a.input or SAMPLES / f"synthetic_{a.variant}.npz")
    if a.calib == "none":
        calib = {"ego_crop_box": None}
    elif a.calib and a.calib.startswith("preset:"):
        calib = {"preset": a.calib[len("preset:"):]}
    elif a.calib:
        calib = json.loads(Path(a.calib).read_text())
    else:
        stem = src.name.split(".")[0]
        calib = {"preset": stem} if (CODE / "tt_frnet" / "calib" / f"{stem}.json").is_file() else {"ego_crop_box": None}
    weights_path = FRNet.resolve_weights()
    config = FRNetConfig.from_model_dir(weights_path, a.variant)
    weights = load_weights(weights_path, a.variant)
    pre = FRNetPreprocessor(config, pre_norm=(weights.pre_norm_scale, weights.pre_norm_shift))
    prepared = pre(str(src), calib)
    if prepared.plan is None:
        raise SystemExit(f"the frame is skipped: {prepared.skipped}")
    dev = open_device(allow_fallback=False)
    try:
        print("device:", json.dumps(describe_device(dev, 0)), flush=True)
        graph = TtFRNet(dev, weights, config, policy=default_policy())
        arrays = frame_arrays(prepared.plan, graph.m_cap)
        inp = {}
        for key, (shape, dtype) in graph.input_specs().items():
            tdt = ttnn.float32 if dtype == "float32" else ttnn.uint32
            t = ttnn.allocate_tensor_on_device(ttnn.Shape(list(shape)), tdt, ttnn.ROW_MAJOR_LAYOUT, dev,
                                               ttnn.DRAM_MEMORY_CONFIG)
            ttnn.copy_host_to_device_tensor(to_host_tensor(arrays[key].reshape(shape), dtype, ttnn.ROW_MAJOR_LAYOUT), t)
            inp[key] = t

        def forward():
            out = graph.forward(inp)
            ttnn.synchronize_device(dev)
            free(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_dir = Path(a.json_dir).resolve() if a.json_dir else out_dir
        json_dir.mkdir(parents=True, exist_ok=True)
        json_path = json_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)
        free(*inp.values())
        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()