# SPDX-FileCopyrightText: (c) 2026 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0

"""
SAM3 inference benchmark harness.

Emits results in a grep-parseable format so the experiment loop can pick them up:
    inference_speed <fps>
    accuracy <pcc_percent>
    peak_dram <bytes>

Runs either the ViT backbone only (``--mode backbone``) or the full E2E
pipeline (``--mode e2e``). In both modes we compute PCC of the TT-NN output
versus the torch CPU reference to report ``accuracy``.

``TT_FUSED=1`` (``tt_sam3/tt/fused_config.py``) is honoured exactly as the server honours it: the
knob is read once in :func:`main`, the device is opened with the same ``device_open_kwargs`` (a
trace region whenever the fused encoder trace is on -- ``end_trace_capture`` OOMs on a bare open),
``--mode backbone`` measures ``tt_vit_backbone_fused`` and ``--mode e2e`` builds the pipeline with
that config and applies the server's READY contract (trace captured during warm-up, zero legacy
fallbacks), so ``DEVICE_VALIDATION.md`` can use both modes as the per-step PCC gate.  ``TT_FUSED=0``
measures the legacy graph (the fused paths are the default since the device pass).  Extra output lines under the knob:
    fused <select_paths summary>
    encoder_trace_captured <0|1>      (e2e)
    encoder_fallbacks <n>             (e2e)
"""

from __future__ import annotations

import argparse
import os
import sys
import time

import torch

from tt_sam3.tt.device_open import open_device
from tt_sam3.tt.fused_config import FusedConfig, device_open_kwargs, select_paths
from tt_sam3.utils import comp_pcc


def peak_dram_bytes(device=None) -> int:
    try:
        import ttnn

        if device is not None and hasattr(ttnn, "get_memory_info"):
            return int(ttnn.get_memory_info(device).dram.peak_allocated_bytes)
    except Exception:
        pass
    return 0


def _select_backbone(fused: FusedConfig):
    """(preprocess, move_to_device, backbone_fn) for the legacy graph or, under ``TT_FUSED=1``, the
    fused graph the pipeline / server run (``tt_vit_backbone_fused`` with this config)."""
    if fused.enabled:
        from tt_sam3.tt.ttnn_sam3_vitdet import (
            move_backbone_params_to_device_fused,
            preprocess_vit_backbone_weights_fused,
            tt_vit_backbone_fused,
        )

        def preprocess(vit):
            return preprocess_vit_backbone_weights_fused(vit, fused)

        def move(params, device):
            return move_backbone_params_to_device_fused(params, device, fused)

        def backbone(pixel_values, params, device):
            return tt_vit_backbone_fused(pixel_values, params, device, fused)

        return preprocess, move, backbone

    from tt_sam3.tt.ttnn_sam3_vitdet import (
        move_backbone_params_to_device,
        preprocess_vit_backbone_weights,
        tt_vit_backbone,
    )

    return preprocess_vit_backbone_weights, move_backbone_params_to_device, tt_vit_backbone


def run_backbone_benchmark(
    iters: int, warmup: int, device_id: int, use_pretrained: bool, fused: FusedConfig, l1_small_size: int = 0
) -> dict:
    import ttnn

    from conftest import _build_sam3_cpu
    from sam3.model.vitdet import ViT

    preprocess, move, backbone_fn = _select_backbone(fused)

    model = _build_sam3_cpu(use_pretrained=use_pretrained)

    vit_backbone = None
    for _, module in model.named_modules():
        if isinstance(module, ViT):
            vit_backbone = module
            break
    assert vit_backbone is not None

    device = open_device(device_open_kwargs(fused, device_id=device_id, l1_small_size=l1_small_size),
                         os.environ.get("SAM3_DISPATCH", "auto"))[0]  # default ETH dispatch, 1 CQ, 12x10
    try:
        params = preprocess(vit_backbone)
        params = move(params, device)

        torch.manual_seed(42)
        pixel_values = torch.randn(1, 3, 1008, 1008)

        with torch.no_grad():
            ref_output = vit_backbone(pixel_values)
        ref_feat = ref_output[-1]

        for _ in range(warmup):
            with torch.no_grad():
                _ = backbone_fn(pixel_values, params, device)

        times = []
        tt_feat = None
        for _ in range(iters):
            t0 = time.perf_counter()
            with torch.no_grad():
                tt_output = backbone_fn(pixel_values, params, device)
            times.append(time.perf_counter() - t0)
            tt_feat = tt_output[-1]  # fused + SAM3_BB_BF16_READBACK=1: a bf16 channels_last view; .float() below

        _, pcc_val = comp_pcc(ref_feat.float(), tt_feat.float())

        avg_time = sum(times) / len(times)
        fps = 1.0 / avg_time
        dram = peak_dram_bytes(device)
    finally:
        ttnn.close_device(device)

    return {
        "inference_speed": fps,
        "accuracy": max(0.0, pcc_val) * 100.0,
        "peak_dram": dram,
        "avg_ms": avg_time * 1000,
        "per_run_ms": [t * 1000 for t in times],
    }


def run_e2e_benchmark(
    iters: int, warmup: int, device_id: int, use_pretrained: bool, fused: FusedConfig, l1_small_size: int = 0
) -> dict:
    import ttnn

    from conftest import _build_sam3_cpu
    from tt_sam3.tt.ttnn_sam3_model import TtSam3ImagePipeline, make_batched_datapoint, extract_predictions

    model = _build_sam3_cpu(use_pretrained=use_pretrained)

    torch.manual_seed(42)
    pixel_values = torch.randn(1, 3, 1008, 1008)
    text_prompts = ["object"]

    # CPU reference on the pristine (unpatched, fp32) model, BEFORE the device pipeline is built -- the
    # previous build-restore-rebuild sequence uploaded every device weight twice.
    with torch.inference_mode():
        ref_output = model(make_batched_datapoint(pixel_values, text_prompts))
    ref_preds = extract_predictions(ref_output)

    device = open_device(device_open_kwargs(fused, device_id=device_id, l1_small_size=l1_small_size),
                         os.environ.get("SAM3_DISPATCH", "auto"))[0]  # default ETH dispatch, 1 CQ, 12x10
    pipeline = None
    try:
        # same construction as the server: the FusedConfig read once in main() (legacy when TT_FUSED is unset)
        pipeline = TtSam3ImagePipeline(model, device, fused_config=fused)

        for _ in range(warmup):
            with torch.inference_mode():
                _ = pipeline.forward(make_batched_datapoint(pixel_values, text_prompts))

        if fused.enabled and fused.enc_trace and warmup > 0:
            # the server's READY contract: the encoder trace was captured in warm-up and nothing fell back
            if not pipeline.encoder_trace_captured or pipeline.encoder_fallbacks != 0:
                raise RuntimeError(
                    "TT_FUSED=1: warm-up did not run through the fused encoder trace "
                    f"(captured={pipeline.encoder_trace_captured}, legacy fallbacks={pipeline.encoder_fallbacks})"
                )

        times = []
        for _ in range(iters):
            t0 = time.perf_counter()
            with torch.inference_mode():
                tt_output = pipeline.forward(make_batched_datapoint(pixel_values, text_prompts))
            times.append(time.perf_counter() - t0)

        tt_preds = extract_predictions(tt_output)

        # extract_predictions returns the upstream keys (pred_masks / pred_logits); the old
        # "masks" / "scores" keys raised KeyError, so --mode e2e never produced a number.
        _, mask_pcc = comp_pcc(ref_preds["pred_masks"].float(), tt_preds["pred_masks"].float())
        _, score_pcc = comp_pcc(ref_preds["pred_logits"].float(), tt_preds["pred_logits"].float())

        avg_time = sum(times) / len(times)
        fps = 1.0 / avg_time
        dram = peak_dram_bytes(device)
        trace_captured = bool(pipeline.encoder_trace_captured)
        fallbacks = pipeline.encoder_fallbacks
    finally:
        if pipeline is not None:
            pipeline.restore()  # releases the encoder trace + persistent buffers before the device closes
            pipeline.backbone_params = pipeline._enc_params = pipeline._fused_encoder = None
            del pipeline
        ttnn.close_device(device)

    return {
        "inference_speed": fps,
        "accuracy": max(0.0, mask_pcc) * 100.0,
        "peak_dram": dram,
        "avg_ms": avg_time * 1000,
        "mask_pcc": mask_pcc,
        "score_pcc": score_pcc,
        "encoder_trace_captured": trace_captured,
        "encoder_fallbacks": fallbacks,
    }


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description="SAM3 TT-NN Benchmark")
    parser.add_argument("--mode", choices=["backbone", "e2e"], default="backbone")
    parser.add_argument("--iters", type=int, default=10)
    parser.add_argument("--warmup", type=int, default=2)
    parser.add_argument("--pretrained", action="store_true")
    parser.add_argument("--device-id", type=int, default=int(os.environ.get("SAM3_DEVICE", "0")))
    parser.add_argument(
        "--l1-small-size", type=int, default=int(os.environ.get("SAM3_L1_SMALL_SIZE", "0")),
        help="ttnn l1_small_size (0 = ttnn default, the served configuration; the old PCC tests used 32768)",
    )
    args = parser.parse_args(argv)

    torch.set_grad_enabled(False)
    from tt_sam3.tt.ttnn_sam3_model import _patch_pin_memory

    _patch_pin_memory()  # sam3's geometry encoder calls Tensor.pin_memory(), which raises on CPU-only torch (as the server)

    fused = FusedConfig.from_env()  # TT_FUSED read once, as the server does
    print(f"fused {select_paths(fused)}")

    if args.mode == "backbone":
        res = run_backbone_benchmark(args.iters, args.warmup, args.device_id, args.pretrained, fused, args.l1_small_size)
    else:
        res = run_e2e_benchmark(args.iters, args.warmup, args.device_id, args.pretrained, fused, args.l1_small_size)

    print(f"mode {args.mode}")
    print(f"inference_speed {res['inference_speed']:.4f}")
    print(f"accuracy {res['accuracy']:.4f}")
    print(f"peak_dram {res['peak_dram']}")
    print(f"avg_ms {res['avg_ms']:.1f}")
    if "mask_pcc" in res:
        print(f"mask_pcc {res['mask_pcc']:.6f}")
        print(f"score_pcc {res['score_pcc']:.6f}")
    if "encoder_trace_captured" in res:
        print(f"encoder_trace_captured {int(res['encoder_trace_captured'])}")
        print(f"encoder_fallbacks {res['encoder_fallbacks']}")
    return 0


if __name__ == "__main__":
    sys.exit(main())