# SPDX-License-Identifier: Apache-2.0
"""Where the e2e tt(image) time goes beyond the trace (round 11 opt tool): timestamps inside the call
(start -> _exec_trace entry = host prep + band uploads, -> _read_fused entry, _read_fused (band readback on
CQ1 + C++ tail), -> return (host finish)), medians over N calls. Device opened as served (worker, MOGE_2CQ)."""
import os
import statistics
import sys
import time

sys.path.insert(0, "/home/ttuser/experiments/tt-models/tools")
import torch


def main():
    from open_device_12x10 import open_device, close_device
    import ttnn
    import tt_moge  # noqa: F401
    from tt_moge.benchmark import load_image
    from tt_moge.reference.load_pretrained import load_moge2
    from tt_moge.tt.ttnn_moge import TtMoGe
    torch.set_grad_enabled(False)
    n = int(os.environ.get("N", 40))
    ref = load_moge2()
    image = load_image(os.path.join(os.path.dirname(__file__), "..", "..", "..", "media", "source.png")).float()
    dev = open_device(grid="12x10", dispatch=os.environ.get("MOGE_DISPATCH", "worker"), l1_small_size=32768,
                      trace_region_size=1500000000,
                      num_command_queues=2 if os.environ.get("MOGE_2CQ", "0") == "1" else 1)
    try:
        tt = TtMoGe(ref, dev)
        tt(image, num_tokens=1800)
        print(f"grouped={tt._grouped} segments={1 + len(tt._dec_segs)} scale_trace={tt._dec_tid2 is not None}")
        marks = {}
        ex, rd = tt._exec_trace, tt._read_fused

        def ex_w(*a, **k):
            marks["exec"] = time.perf_counter()
            r = ex(*a, **k)
            marks["exec_ret"] = time.perf_counter()
            return r

        def rd_w(*a, **k):
            marks["read"] = time.perf_counter()
            r = rd(*a, **k)
            marks["read_ret"] = time.perf_counter()
            return r
        tt._exec_trace, tt._read_fused = ex_w, rd_w
        rows = []
        for _ in range(n):
            ttnn.synchronize_device(dev)
            if tt._grouped:
                tt._grp_times = []
            t0 = time.perf_counter()
            tt(image, num_tokens=1800)
            t1 = time.perf_counter()
            if tt._grouped:
                for i, t in enumerate(tt._grp_times):
                    marks[f"g{i}"] = t - marks["exec"]
            rows.append(dict({k: v for k, v in marks.items() if k[0] == "g"},pre=marks["exec"] - t0, exec=marks["exec_ret"] - marks["exec"],
                             gap=marks["read"] - marks["exec_ret"], read=marks["read_ret"] - marks["read"],
                             finish=t1 - marks["read_ret"], e2e=t1 - t0))
        for k in rows[0]:
            v = [r[k] * 1e3 for r in rows]
            print(f"{k:8s} median {statistics.median(v):7.2f} ms  min {min(v):7.2f}")
    finally:
        close_device(dev)


if __name__ == "__main__":
    main()