#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
_12_benchmark_onnx.py

Benchmark PyTorch vs ONNX Runtime para SegFormer OAK-FCC-3 Multi-Head.

Mede:
  - PyTorch FP32
  - PyTorch AMP/FP16
  - ONNX Runtime CUDA ou CPU

Exemplos:

Benchmark ONNX cru 160x256:

python _12_benchmark_onnx.py --config config.json --max_samples 50 --warmup 10 --repeat 5 --device cuda --onnx_provider cuda

Benchmark ONNX resized 640x1024:

python _12_benchmark_onnx.py --config config.json --max_samples 50 --warmup 10 --repeat 5 --device cuda --onnx_provider cuda


TensorRT
python _12_benchmark_onnx.py --config config.json --max_samples 50 --warmup 10 --repeat 5 --device cuda --onnx_provider tensorrt --skip_torch_fp32 --skip_torch_amp

"""

from __future__ import annotations

import gc
import csv
import json
import time
import argparse
import importlib.util
from pathlib import Path
from typing import Optional, List, Dict, Tuple

import cv2
import numpy as np
import torch
import torch.nn as nn


# ============================================================
# Utils
# ============================================================

def load_json(path: str | Path) -> dict:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def save_json(path: str | Path, data: dict):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


def save_csv(path: str | Path, rows: List[dict]):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)

    if not rows:
        return

    keys = list(rows[0].keys())

    with path.open("w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=keys)
        w.writeheader()
        w.writerows(rows)


def resolve_path(path_like: Optional[str], base: Optional[Path] = None) -> Optional[Path]:
    if path_like is None:
        return None

    p = Path(path_like)
    if p.is_absolute():
        return p

    if base is None:
        base = Path.cwd()

    return (base / p).resolve()


def import_train_module(train_script_path: str | Path):
    train_script_path = Path(train_script_path)

    if not train_script_path.exists():
        raise FileNotFoundError(f"Script de treino não encontrado: {train_script_path}")

    spec = importlib.util.spec_from_file_location(
        "train_multihead_module",
        str(train_script_path.resolve())
    )

    if spec is None or spec.loader is None:
        raise RuntimeError(f"Não consegui importar o script: {train_script_path}")

    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def synchronize_if_cuda(device: torch.device):
    if device.type == "cuda":
        torch.cuda.synchronize()


def clear_cuda():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()


def percentile(values: List[float], p: float) -> float:
    if not values:
        return 0.0
    return float(np.percentile(np.asarray(values, dtype=np.float64), p))


def summarize_times(times_ms: List[float]) -> dict:
    arr = np.asarray(times_ms, dtype=np.float64)

    if arr.size == 0:
        return {
            "n": 0,
            "mean_ms": 0.0,
            "median_ms": 0.0,
            "min_ms": 0.0,
            "max_ms": 0.0,
            "p95_ms": 0.0,
            "p99_ms": 0.0,
            "fps_mean": 0.0,
            "fps_p95_latency": 0.0,
        }

    mean_ms = float(arr.mean())
    p95_ms = float(np.percentile(arr, 95))
    p99_ms = float(np.percentile(arr, 99))

    return {
        "n": int(arr.size),
        "mean_ms": mean_ms,
        "median_ms": float(np.median(arr)),
        "min_ms": float(arr.min()),
        "max_ms": float(arr.max()),
        "p95_ms": p95_ms,
        "p99_ms": p99_ms,
        "fps_mean": float(1000.0 / mean_ms) if mean_ms > 0 else 0.0,
        "fps_p95_latency": float(1000.0 / p95_ms) if p95_ms > 0 else 0.0,
    }


def resolve_model_artifact_paths(
    args,
    config: dict,
    config_dir: Path,
    channels: int,
) -> Tuple[Path, Path, str]:
    """
    Resolve checkpoint e ONNX.

    Se --checkpoint ou --onnx forem informados, usa os caminhos informados.
    Se ficarem vazios, monta a partir do config:

      backup/{modelo}/{model_name}/{fusion_mode}_raw{channels}/{ckpt_name}.pt
      backup/{modelo}/{model_name}/{fusion_mode}_raw{channels}/{ckpt_name}.onnx

    ckpt_name vem de:
      config["ckpt_test"] ou "best_score"
    """
    model = config.get("modelo", "segformer_b1")
    model_name = config.get("model_name", "target_teached")
    fusion_mode = config.get("fusion_mode", "stacked")
    ckpt_name = config.get("ckpt_test", "best_score")

    # Usa o número real de canais selecionados,
    # não necessariamente config["channels"].
    ch = int(channels)

    base_dir = config_dir / "backup" / model / model_name / f"{fusion_mode}_raw{ch}"

    if args.checkpoint:
        checkpoint_path = resolve_path(args.checkpoint, Path.cwd())
    else:
        checkpoint_path = base_dir / f"{ckpt_name}.pt"

    if args.onnx:
        onnx_path = resolve_path(args.onnx, Path.cwd())
    else:
        onnx_path = base_dir / f"{ckpt_name}.onnx"

    if checkpoint_path is None or not checkpoint_path.is_file():
        raise FileNotFoundError(
            f"Checkpoint não encontrado: {checkpoint_path}\n"
            f"Dica: informe --checkpoint ou ajuste config['ckpt_test']."
        )

    if onnx_path is None or not onnx_path.is_file():
        raise FileNotFoundError(
            f"ONNX não encontrado: {onnx_path}\n"
            f"Dica: informe --onnx ou ajuste config['ckpt_test']."
        )

    return checkpoint_path.resolve(), onnx_path.resolve(), str(ckpt_name)


# ============================================================
# Dataset / normalização
# ============================================================

def collect_tensor_samples(root: Path, max_samples: int = 50, start_idx: int = 0) -> List[Path]:
    tensor_paths: List[Path] = []

    direct = root / "tensors"
    if direct.is_dir():
        tensor_paths.extend(sorted(direct.glob("*.npy")))

    group_root = root / "group"
    if group_root.is_dir():
        for gdir in sorted(group_root.iterdir()):
            tdir = gdir / "tensors"
            if tdir.is_dir():
                tensor_paths.extend(sorted(tdir.glob("*.npy")))

    if not tensor_paths:
        tensor_paths.extend(sorted(root.glob("**/tensors/*.npy")))

    if not tensor_paths:
        raise RuntimeError(f"Nenhum tensor .npy encontrado em: {root}")

    start_idx = max(0, int(start_idx))
    selected = tensor_paths[start_idx:]

    if max_samples > 0:
        selected = selected[:int(max_samples)]

    return selected


def load_tensor(path: Path, channels: int, channel_indices: List[int]) -> np.ndarray:
    arr = np.load(str(path)).astype(np.float32)

    if arr.ndim != 3:
        raise RuntimeError(f"Tensor inválido {path}: shape={arr.shape}, esperado 3D")

    if arr.shape[0] in (3, 4, 5):
        chw = arr
    elif arr.shape[-1] in (3, 4, 5):
        chw = np.transpose(arr, (2, 0, 1))
    else:
        raise RuntimeError(f"Tensor com layout inesperado: {path} shape={arr.shape}")

    max_idx = max(channel_indices)
    if chw.shape[0] <= max_idx:
        raise RuntimeError(
            f"Tensor {path} tem {chw.shape[0]} canais, "
            f"mas precisa acessar índice {max_idx}."
        )

    chw = chw[channel_indices, :, :]

    finite = np.isfinite(chw)
    if finite.any():
        mx = float(np.nanmax(chw[finite]))
        if mx > 2.0 and mx <= 255.0:
            chw = chw / 255.0
        elif mx > 255.0:
            chw = chw / 65535.0

    chw = np.nan_to_num(chw, nan=0.0, posinf=1.0, neginf=0.0)
    return np.clip(chw, 0.0, 1.0).astype(np.float32)


def find_norm_stats(config: dict, config_dir: Path, save_dir: Path, explicit: Optional[str]) -> Optional[Path]:
    if explicit:
        return resolve_path(explicit, Path.cwd())

    W, H = config.get("resolucao", [1024, 640])
    dataset_path = config_dir / "dataset"

    candidates = [
        dataset_path / f"{int(W)}x{int(H)}" / "group" / "norm_stats.json",
        save_dir / "norm_stats.json",
        config_dir / "backup" / config.get("modelo", "segformer_b1") / config.get("model_name", "test") / config.get("stats_source_tag", "stacked_raw5") / "norm_stats.json",
    ]

    for p in candidates:
        if p.is_file():
            return p

    return candidates[0]


def load_norm_stats(
    path: Optional[Path],
    channel_indices: List[int],
    channel_names: List[str],
) -> Tuple[Optional[List[float]], Optional[List[float]], Optional[str]]:
    if path is None or not path.is_file():
        print("[NORM] Sem norm_stats. Usando tensor 0..1 sem padronização.")
        return None, None, None

    js = load_json(path)
    mean = js.get("mean")
    std = js.get("std")
    names = js.get("channels", [])

    if mean is None or std is None:
        raise RuntimeError(f"norm_stats inválido, faltando mean/std: {path}")

    max_idx = max(channel_indices)
    if len(mean) <= max_idx or len(std) <= max_idx:
        raise RuntimeError(
            f"norm_stats incompatível: precisa índices={channel_indices}, "
            f"mean={len(mean)} std={len(std)}"
        )

    mean_sel = [float(mean[i]) for i in channel_indices]
    std_sel = [float(std[i]) for i in channel_indices]

    if names:
        names_sel = [names[i] for i in channel_indices]
    else:
        names_sel = channel_names

    print(f"[NORM] usando {path}")
    print(f"[NORM] channels={names_sel}")
    print(f"[NORM] mean={mean_sel}")
    print(f"[NORM] std ={std_sel}")

    return mean_sel, std_sel, str(path)


def normalize_numpy_chw(chw: np.ndarray, mean: Optional[List[float]], std: Optional[List[float]]) -> np.ndarray:
    if mean is None or std is None:
        return chw.astype(np.float32)

    mean_np = np.asarray(mean, dtype=np.float32).reshape(-1, 1, 1)
    std_np = np.asarray(std, dtype=np.float32).reshape(-1, 1, 1)
    std_np = np.clip(std_np, 1e-6, None)

    return ((chw.astype(np.float32) - mean_np) / std_np).astype(np.float32)


def load_inputs_as_numpy(
    samples: List[Path],
    channels: int,
    channel_indices: List[int],
    mean: Optional[List[float]],
    std: Optional[List[float]],
    target_hw: Tuple[int, int],
    normalize_input: bool = True,
) -> List[np.ndarray]:
    H, W = target_hw
    xs = []

    for p in samples:
        chw01 = load_tensor(
            p,
            channels=channels,
            channel_indices=channel_indices,
        )

        if chw01.shape[-2:] != (H, W):
            hwc = np.transpose(chw01, (1, 2, 0))
            hwc = cv2.resize(hwc, (W, H), interpolation=cv2.INTER_LINEAR)
            chw01 = np.transpose(hwc, (2, 0, 1)).astype(np.float32)

        if normalize_input:
            chw = normalize_numpy_chw(chw01, mean=mean, std=std)
        else:
            chw = chw01.astype(np.float32, copy=False)

        x = np.expand_dims(chw, axis=0).astype(np.float32)
        xs.append(x)

    return xs


# ============================================================
# PyTorch
# ============================================================

class TorchTupleWrapper(nn.Module):
    def __init__(self, model: nn.Module, output_heads: List[str]):
        super().__init__()
        self.model = model
        self.output_heads = list(output_heads)

    def forward(self, pixel_values: torch.Tensor):
        outputs = self.model(pixel_values=pixel_values)
        return tuple(outputs[h] for h in self.output_heads)


@torch.inference_mode()
def benchmark_torch(
    model: nn.Module,
    inputs_np: List[np.ndarray],
    device: torch.device,
    warmup: int,
    repeat: int,
    amp: bool,
    label: str,
) -> Tuple[dict, List[dict]]:
    model.eval()

    times = []
    rows = []

    # Precarrega tensors na GPU para medir só inferência do modelo.
    inputs_t = [
        torch.from_numpy(x).to(device, non_blocking=True)
        for x in inputs_np
    ]

    if device.type == "cuda":
        torch.cuda.synchronize()

    print(f"\n[BENCH] {label} | warmup={warmup} repeat={repeat}")

    # Warmup
    for i in range(max(0, warmup)):
        x = inputs_t[i % len(inputs_t)]
        with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=amp and device.type == "cuda"):
            _ = model(x)

    synchronize_if_cuda(device)

    # Medição
    total_iter = len(inputs_t) * max(1, repeat)
    idx = 0

    for r in range(max(1, repeat)):
        for sample_idx, x in enumerate(inputs_t):
            synchronize_if_cuda(device)
            t0 = time.perf_counter()

            with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=amp and device.type == "cuda"):
                _ = model(x)

            synchronize_if_cuda(device)
            dt_ms = (time.perf_counter() - t0) * 1000.0

            times.append(dt_ms)
            rows.append({
                "engine": label,
                "repeat": r,
                "sample_idx": sample_idx,
                "iter_idx": idx,
                "latency_ms": dt_ms,
            })

            idx += 1

            if idx % 25 == 0 or idx == total_iter:
                print(f"  {idx:04d}/{total_iter:04d} | last={dt_ms:.2f}ms")

    summary = summarize_times(times)
    summary["engine"] = label

    return summary, rows


# ============================================================
# ONNX Runtime
# ============================================================

def create_onnx_session(onnx_path: Path, provider: str):
    import os

    trt_home = os.environ.get("TRT_HOME", r"C:\dev\TensorRT-10.10.0.31")

    for dll_dir in [
        os.path.join(trt_home, "lib"),
        os.path.join(trt_home, "bin"),
        r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\bin",
    ]:
        if os.path.isdir(dll_dir):
            try:
                os.add_dll_directory(dll_dir)
                print(f"[DLL] add_dll_directory: {dll_dir}")
            except Exception as e:
                print(f"[DLL][WARN] falha em {dll_dir}: {e}")
                
    try:
        import onnxruntime as ort
    except ImportError:
        raise ImportError(
            "onnxruntime não está instalado. Instale com:\n"
            "  pip install onnxruntime-gpu\n"
            "ou CPU:\n"
            "  pip install onnxruntime"
        )

    available = ort.get_available_providers()
    print(f"[ONNX] providers disponíveis: {available}")

    provider = provider.lower()

    sess_options = ort.SessionOptions()
    sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

    if provider == "cuda":
        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]

    elif provider == "cpu":
        providers = ["CPUExecutionProvider"]

    elif provider == "tensorrt":
        cache_dir = onnx_path.parent / "trt_cache"
        cache_dir.mkdir(parents=True, exist_ok=True)

        trt_options = {
            "device_id": 0,

            # FP16: o ponto principal do nosso teste.
            "trt_fp16_enable": True,

            # Cache: evita rebuild do engine a cada execução.
            "trt_engine_cache_enable": True,
            "trt_engine_cache_path": str(cache_dir),

            # Timing cache ajuda a acelerar builds futuros.
            "trt_timing_cache_enable": True,
            "trt_timing_cache_path": str(cache_dir),

            # Workspace. 4GB é razoável para RTX 3070, ajuste se faltar VRAM.
            "trt_max_workspace_size": 4 * 1024 * 1024 * 1024,
        }

        providers = [
            ("TensorrtExecutionProvider", trt_options),
            "CUDAExecutionProvider",
            "CPUExecutionProvider",
        ]

    else:
        providers = [provider]

    # Checagem de disponibilidade, lidando com provider tuple.
    requested_names = [
        p[0] if isinstance(p, tuple) else p
        for p in providers
    ]

    providers_ok = [
        p for p in providers
        if (p[0] if isinstance(p, tuple) else p) in available
    ]

    if not providers_ok:
        raise RuntimeError(
            f"Nenhum provider solicitado está disponível. "
            f"Solicitado={requested_names}, disponível={available}"
        )

    session = ort.InferenceSession(
        str(onnx_path),
        sess_options=sess_options,
        providers=providers_ok,
    )

    print(f"[ONNX] usando providers: {session.get_providers()}")

    active_providers = session.get_providers()

    if provider == "tensorrt" and "TensorrtExecutionProvider" not in active_providers:
        raise RuntimeError(
            "TensorRTExecutionProvider foi solicitado, mas não ficou ativo. "
            f"Providers ativos: {active_providers}. "
            "Provável causa: TensorRT não instalado, DLLs fora do PATH, "
            "ou versão incompatível com onnxruntime-gpu."
        )

    if provider == "cuda" and "CUDAExecutionProvider" not in active_providers:
        raise RuntimeError(
            "CUDAExecutionProvider foi solicitado, mas não ficou ativo. "
            f"Providers ativos: {active_providers}."
        )

    return session


def benchmark_onnx(
    session,
    inputs_np: List[np.ndarray],
    warmup: int,
    repeat: int,
    label: str,
) -> Tuple[dict, List[dict]]:
    input_name = session.get_inputs()[0].name

    times = []
    rows = []

    print(f"\n[BENCH] {label} | warmup={warmup} repeat={repeat}")

    # Warmup
    for i in range(max(0, warmup)):
        x = inputs_np[i % len(inputs_np)]
        _ = session.run(None, {input_name: x})

    # Medição
    total_iter = len(inputs_np) * max(1, repeat)
    idx = 0

    for r in range(max(1, repeat)):
        for sample_idx, x in enumerate(inputs_np):
            t0 = time.perf_counter()
            _ = session.run(None, {input_name: x})
            dt_ms = (time.perf_counter() - t0) * 1000.0

            times.append(dt_ms)
            rows.append({
                "engine": label,
                "repeat": r,
                "sample_idx": sample_idx,
                "iter_idx": idx,
                "latency_ms": dt_ms,
            })

            idx += 1

            if idx % 25 == 0 or idx == total_iter:
                print(f"  {idx:04d}/{total_iter:04d} | last={dt_ms:.2f}ms")

    summary = summarize_times(times)
    summary["engine"] = label

    return summary, rows


# ============================================================
# Main
# ============================================================

def main():
    parser = argparse.ArgumentParser()

    parser.add_argument("--config", default="config.json")
    parser.add_argument("--checkpoint", default="")
    parser.add_argument("--onnx", default="")
    parser.add_argument("--train-script", default="_8_train_multihead.py")
    parser.add_argument("--labelmap", default="dataset/labelmap.txt")

    parser.add_argument("--split_folder", default="val", choices=["train", "val", "test"])
    parser.add_argument("--root_override", default=None)
    parser.add_argument("--norm_stats", default=None)

    parser.add_argument("--max_samples", type=int, default=50)
    parser.add_argument("--start_idx", type=int, default=0)
    parser.add_argument("--warmup", type=int, default=10)
    parser.add_argument("--repeat", type=int, default=5)

    parser.add_argument("--device", default="cuda", choices=["cuda", "cpu"])
    parser.add_argument("--onnx_provider", default="cuda", choices=["cuda", "cpu", "tensorrt"])

    parser.add_argument("--skip_torch_fp32", action="store_true")
    parser.add_argument("--skip_torch_amp", action="store_true")
    parser.add_argument("--skip_onnx", action="store_true")

    parser.add_argument(
        "--onnx_has_norm",
        action="store_true",
        help="Use quando o ONNX já inclui normalização interna. Nesse caso o ONNX recebe tensor 0..1 cru.",
    )

    parser.add_argument("--out_dir", default=None)

    args = parser.parse_args()

    config_path = resolve_path(args.config, Path.cwd())
    train_script_path = resolve_path(args.train_script, Path.cwd())
    labelmap_path = resolve_path(args.labelmap, Path.cwd())

    if config_path is None or not config_path.is_file():
        raise FileNotFoundError(f"Config não encontrado: {config_path}")
    if train_script_path is None or not train_script_path.is_file():
        raise FileNotFoundError(f"Train script não encontrado: {train_script_path}")
    if labelmap_path is None or not labelmap_path.is_file():
        raise FileNotFoundError(f"Labelmap não encontrado: {labelmap_path}")

    config_dir = config_path.parent
    config = load_json(config_path)

    train_mod = import_train_module(train_script_path)

    W, H = config.get("resolucao", [1024, 640])
    W = int(W)
    H = int(H)

    backbone = config.get("backbone", "nvidia/mit-b1")
    input_channel_names = train_mod.get_input_channel_names(config)
    input_channel_indices = train_mod.get_input_channel_indices(config)
    channels = len(input_channel_names)

    checkpoint_path, onnx_path, ckpt_name = resolve_model_artifact_paths(
        args=args,
        config=config,
        config_dir=config_dir,
        channels=channels,
    )

    semantic_id2label, semantic_label2id, ignore_from_labelmap = train_mod.load_labelmap(
        str(labelmap_path)
    )

    heads_config = train_mod.build_heads_config(
        config,
        ignore_index=int(ignore_from_labelmap)
    )

    heads_config["semantic"]["num_classes"] = int(len(semantic_id2label))
    heads_config["semantic"]["ignore_index"] = int(ignore_from_labelmap)

    output_heads = list(heads_config.keys())

    save_dir = (
        config_dir
        / "backup"
        / config.get("modelo", "segformer_b1")
        / config.get("model_name", "test")
        / f"{config.get('fusion_mode', 'stacked')}_raw{channels}"
    )

    norm_stats_path = find_norm_stats(
        config=config,
        config_dir=config_dir,
        save_dir=save_dir,
        explicit=args.norm_stats,
    )

    mean, std, norm_stats_used = load_norm_stats(
        norm_stats_path,
        channel_indices=input_channel_indices,
        channel_names=input_channel_names,
    )

    if args.root_override:
        root = resolve_path(args.root_override, Path.cwd())
    else:
        root = (config_dir / "dataset" / "split" / args.split_folder).resolve()

    if root is None or not root.is_dir():
        raise FileNotFoundError(f"Root de dados não encontrado: {root}")

    samples = collect_tensor_samples(
        root=root,
        max_samples=args.max_samples,
        start_idx=args.start_idx,
    )

    if args.out_dir:
        out_dir = resolve_path(args.out_dir, Path.cwd())
    else:
        out_dir = onnx_path.parent / "benchmarks"

    assert out_dir is not None
    out_dir.mkdir(parents=True, exist_ok=True)

    use_cuda = args.device == "cuda" and torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")

    if args.device == "cuda" and not torch.cuda.is_available():
        print("[WARN] CUDA indisponível. Usando CPU no PyTorch.")

    print("==========================================")
    print("Benchmark PyTorch vs ONNX")
    print(f"Config       : {config_path}")
    print(f"Checkpoint   : {checkpoint_path}")
    print(f"ONNX         : {onnx_path}")
    print(f"Root         : {root}")
    print(f"Samples      : {len(samples)}")
    print(f"Warmup       : {args.warmup}")
    print(f"Repeat       : {args.repeat}")
    print(f"Backbone     : {backbone}")
    print(f"Input shape  : [1, {channels}, {H}, {W}]")
    print(f"Channels     : {input_channel_names} idx={input_channel_indices}")
    print(f"Heads        : {output_heads}")
    print(f"Device       : {device}")
    print(f"ONNX provider: {args.onnx_provider}")
    print(f"ONNX has norm: {args.onnx_has_norm}")
    print(f"Out dir      : {out_dir}")
    print("==========================================")

    if args.onnx_has_norm:
        print("\n[DATA] Carregando inputs 0..1 crus na RAM...")
    else:
        print("\n[DATA] Carregando inputs normalizados na RAM...")
    inputs_np = load_inputs_as_numpy(
        samples=samples,
        channels=channels,
        channel_indices=input_channel_indices,
        mean=mean,
        std=std,
        target_hw=(H, W),
        normalize_input=not args.onnx_has_norm,
    )
    print(f"[DATA] Inputs carregados: {len(inputs_np)}")

    summaries = []
    all_rows = []

    # ========================================================
    # PyTorch
    # ========================================================
    need_torch = not args.skip_torch_fp32 or not args.skip_torch_amp

    if need_torch:
        print("\n[MODEL] Montando PyTorch...")
        model = train_mod.build_model(
            backbone=backbone,
            channels=channels,
            heads_config=heads_config,
            semantic_id2label=semantic_id2label,
            semantic_label2id=semantic_label2id,
        )

        ckpt = torch.load(str(checkpoint_path), map_location="cpu", weights_only=False)
        if "model" not in ckpt:
            raise RuntimeError("Checkpoint não contém chave 'model'.")

        model.load_state_dict(ckpt["model"], strict=True)
        model.to(device)
        model.eval()

        torch_model = TorchTupleWrapper(
            model=model,
            output_heads=output_heads,
        ).to(device)
        torch_model.eval()

        clear_cuda()

        if not args.skip_torch_fp32:
            summary, rows = benchmark_torch(
                model=torch_model,
                inputs_np=inputs_np,
                device=device,
                warmup=args.warmup,
                repeat=args.repeat,
                amp=False,
                label="torch_fp32",
            )
            summaries.append(summary)
            all_rows.extend(rows)

        clear_cuda()

        if not args.skip_torch_amp:
            summary, rows = benchmark_torch(
                model=torch_model,
                inputs_np=inputs_np,
                device=device,
                warmup=args.warmup,
                repeat=args.repeat,
                amp=True,
                label="torch_amp_fp16",
            )
            summaries.append(summary)
            all_rows.extend(rows)

        del torch_model
        del model
        clear_cuda()

    # ========================================================
    # ONNX
    # ========================================================
    if not args.skip_onnx:
        print("\n[ONNX] Carregando sessão...")
        session = create_onnx_session(
            onnx_path=onnx_path,
            provider=args.onnx_provider,
        )

        summary, rows = benchmark_onnx(
            session=session,
            inputs_np=inputs_np,
            warmup=args.warmup,
            repeat=args.repeat,
            label=f"onnx_{args.onnx_provider}",
        )
        summaries.append(summary)
        all_rows.extend(rows)

    # ========================================================
    # Relatório
    # ========================================================
    print("\n========== RESUMO ==========")

    for s in summaries:
        print(
            f"{s['engine']:<16} "
            f"n={s['n']:<4} "
            f"mean={s['mean_ms']:.3f}ms "
            f"median={s['median_ms']:.3f}ms "
            f"p95={s['p95_ms']:.3f}ms "
            f"p99={s['p99_ms']:.3f}ms "
            f"fps_mean={s['fps_mean']:.2f} "
            f"fps_p95={s['fps_p95_latency']:.2f}"
        )

    base_name = f"{onnx_path.stem}_{args.onnx_provider}"
    report_json = out_dir / f"{base_name}_benchmark_report.json"
    report_csv = out_dir / f"{base_name}_benchmark_rows.csv"

    report = {
        "config": str(config_path),
        "checkpoint": str(checkpoint_path),
        "ckpt_name": ckpt_name,
        "onnx": str(onnx_path),
        "onnx_has_norm": bool(args.onnx_has_norm),
        "root": str(root),
        "samples": len(samples),
        "warmup": int(args.warmup),
        "repeat": int(args.repeat),
        "input_shape": [1, channels, H, W],
        "input_channel_names": input_channel_names,
        "input_channel_indices": input_channel_indices,
        "heads": output_heads,
        "norm_stats_used": norm_stats_used,
        "onnx_provider": args.onnx_provider,
        "device": str(device),
        "summaries": summaries,
    }

    save_json(report_json, report)
    save_csv(report_csv, all_rows)

    print(f"\n[OK] JSON salvo em: {report_json}")
    print(f"[OK] CSV salvo em : {report_csv}")
    print("\nBenchmark finalizado.")


if __name__ == "__main__":
    main()