#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
_9_test_infer_multihead.py

Teste/visualização do SegFormer OAK-FCC-3 Multi-Head.

Contrato esperado após normalize + split:

  dataset/split/val/group/<grupo>/
    tensors/<base>.npy              # CHW float32 [R,G,B,RE,NIR]
    masks/<base>.npy                # semantic: 0=chao, 1=cana, 2=erva, 255=ignore
    masks_vegetation/<base>.npy     # vegetation: 0=background, 1=vegetation, 255=ignore
    masks_cana/<base>.npy           # cana: 0=not_cana, 1=cana, 255=ignore
    metas/<base>.json
    previews/<base>.png

Mostra:
  - RGB preview do tensor
  - GT/pred/overlay semantic
  - GT/pred/overlay vegetation
  - GT/pred/overlay cana
  - alvo operacional = vegetation == 1 AND cana == 0
  - mapas de confiança das heads binárias

Exemplo:

python .\_9_test_infer_multihead.py ^
  --config config.json ^
  --split_folder val ^
  --ckpt backup\segformer_b1\test_multi\stacked_raw5_multihead\best_score.pt

Controles:
  D / seta direita : próxima amostra
  A / seta esquerda: amostra anterior
  S               : salvar painel atual em --out_dir
  SPACE           : alterna modo compacto/detalhado
  Q / ESC         : sair
"""

from __future__ import annotations

import argparse
import copy
import json
import time
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional, Sequence, Tuple

import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import SegformerForSemanticSegmentation

from core.raw_processor_core import RawProcessorCore


# ============================================================
# Config default
# ============================================================

DEFAULT_HEADS = {
    "semantic": {
        "enabled": True,
        "type": "multiclass",
        "num_classes": 3,
        "mask_dir": "masks",
        "classes": {"chao": 0, "cana": 1, "erva": 2},
        "ignore_index": 255,
    },
    "vegetation": {
        "enabled": True,
        "type": "binary",
        "num_classes": 2,
        "mask_dir": "masks_vegetation",
        "classes": {"background": 0, "vegetation": 1},
        "ignore_index": 255,
    },
    "cana": {
        "enabled": True,
        "type": "binary",
        "num_classes": 2,
        "mask_dir": "masks_cana",
        "classes": {"not_cana": 0, "cana": 1},
        "ignore_index": 255,
    },
    "target": {
        "enabled": True,
        "type": "binary",
        "num_classes": 2,
        "mask_dir": "__derived_target__",
        "classes": {
            "background": 0,
            "target": 1,
        },
        "ignore_index": 255,
        "derived_from": ["vegetation", "cana"],
    },
}

SEMANTIC_COLORS_RGB = {
    0: (85, 85, 85),      # chao
    1: (0, 190, 0),       # cana
    2: (230, 55, 55),     # erva
}

BINARY_COLORS_RGB = {
    0: (30, 30, 30),
    1: (0, 220, 80),
}

CANA_COLORS_RGB = {
    0: (30, 30, 30),
    1: (40, 210, 255),
}

TARGET_COLORS_RGB = {
    0: (30, 30, 30),
    1: (255, 70, 30),
}


@dataclass
class SampleItem:
    group: str
    base: str
    tensor_path: Optional[Path]
    masks: Dict[str, Optional[Path]]
    meta_path: Optional[Path] = None
    preview_path: Optional[Path] = None

    # novo
    source_kind: str = "tensor"  # "tensor" ou "raw_native_multi"
    raw_group: Optional[dict] = None


# ============================================================
# Util
# ============================================================

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


def ensure_dir(path: Path):
    path.mkdir(parents=True, exist_ok=True)


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 merge_dict(dst: dict, src: dict) -> dict:
    out = copy.deepcopy(dst)

    def rec(a, b):
        for k, v in b.items():
            if isinstance(v, dict) and isinstance(a.get(k), dict):
                rec(a[k], v)
            else:
                a[k] = v

    if isinstance(src, dict):
        rec(out, src)
    return out


def build_heads_config(config: dict, ignore_index: int) -> Dict[str, dict]:
    cfg = merge_dict(DEFAULT_HEADS, config.get("heads", {}) or {})
    active = {}
    for name, hcfg in cfg.items():
        if not bool(hcfg.get("enabled", True)):
            continue
        hcfg.setdefault("ignore_index", ignore_index)
        hcfg["ignore_index"] = int(hcfg.get("ignore_index", ignore_index))
        hcfg["num_classes"] = int(hcfg.get("num_classes", 2))
        hcfg["mask_dir"] = str(hcfg.get("mask_dir", "masks"))
        active[name] = hcfg

    for required in ("semantic", "vegetation", "cana"):
        if required not in active:
            raise RuntimeError(f"Head obrigatória ausente no config: {required}")

    return active


def load_labelmap(labelmap_path: Path) -> Tuple[Dict[int, str], Dict[str, int], int, Dict[int, Tuple[int, int, int]]]:
    try:
        from helpers import carregar_labelmap_completo, _infer_ignore_id

        _cor_para_id, _colormap_rgb, id_para_nome, ignore_rgb = carregar_labelmap_completo(str(labelmap_path))
        ignore_id = int(_infer_ignore_id(ignore_rgb, 255))

        id2label = {
            int(k): str(v)
            for k, v in id_para_nome.items()
            if str(v).lower() not in ("ignore", "void", "background_ignore")
        }
        label2id = {v.lower(): k for k, v in id2label.items()}

        colormap_rgb = {}
        if isinstance(_colormap_rgb, dict):
            for k, v in _colormap_rgb.items():
                ik = int(k)
                if ik in id2label:
                    colormap_rgb[ik] = tuple(map(int, v[:3]))
        elif isinstance(_colormap_rgb, (list, tuple)):
            for ik, v in enumerate(_colormap_rgb):
                if ik in id2label:
                    colormap_rgb[ik] = tuple(map(int, v[:3]))

        return id2label, label2id, ignore_id, colormap_rgb

    except Exception as e:
        print(f"[WARN] Não consegui usar helpers.carregar_labelmap_completo: {e}")
        print("[WARN] Usando parser simples do labelmap.")

    id2label: Dict[int, str] = {}
    colormap_rgb: Dict[int, Tuple[int, int, int]] = {}
    ignore_id = 255
    next_id = 0

    with labelmap_path.open("r", encoding="utf-8") as f:
        for raw_line in f:
            s = raw_line.strip()
            if not s or s.startswith("#"):
                continue

            name = None
            color = None
            cid = None

            if ":" in s and not s.split(":", 1)[0].strip().isdigit():
                name_part, rest = s.split(":", 1)
                name = name_part.strip()
                color_txt = rest.split("::", 1)[0].strip().strip(":")
                rgb_parts = [p.strip() for p in color_txt.split(",") if p.strip()]
                if len(rgb_parts) >= 3:
                    color = tuple(int(float(p)) for p in rgb_parts[:3])
            else:
                parts = s.replace(",", " ").replace(":", " ").split()
                if len(parts) >= 2 and parts[0].isdigit():
                    cid = int(parts[0])
                    name = parts[1]
                    if len(parts) >= 5:
                        color = tuple(int(float(p)) for p in parts[2:5])
                elif len(parts) >= 1:
                    name = parts[0]

            if not name:
                continue

            if name.lower() in ("ignore", "void", "background_ignore"):
                ignore_id = 255
                continue

            if cid is None:
                cid = next_id
            next_id = max(next_id, cid + 1)

            id2label[int(cid)] = str(name)
            if color is not None:
                colormap_rgb[int(cid)] = color

    if not id2label:
        id2label = {0: "chao", 1: "cana", 2: "erva"}

    label2id = {v.lower(): k for k, v in id2label.items()}
    for cid in id2label:
        colormap_rgb.setdefault(cid, SEMANTIC_COLORS_RGB.get(cid, (255, 255, 255)))

    return id2label, label2id, int(ignore_id), colormap_rgb


def find_raw_dataset_layout_root(path: Path) -> Optional[Path]:
    """
    Detecta layout bruto:
      root/metas
      root/previews
      root/bins
      root/masks opcional

    Aceita root, root/metas, root/bins, root/previews ou arquivo dentro deles.
    """
    p = path.resolve()
    candidates = []

    if p.is_file():
        candidates.append(p.parent)
        candidates.append(p.parent.parent)
    else:
        candidates.append(p)
        candidates.append(p.parent)

    for c in candidates:
        if not c:
            continue

        if c.name.lower() in ("metas", "metadata", "jsons", "previews", "bins", "masks"):
            root = c.parent
        else:
            root = c

        if (root / "metas").is_dir() and (root / "previews").is_dir() and (root / "bins").is_dir():
            return root

    return None


def resolve_raw_sibling_file(root: Path, subdir: str, stem: str, exts: Tuple[str, ...]) -> Optional[Path]:
    folder = root / subdir
    if not folder.is_dir():
        return None

    for ext in exts:
        p = folder / f"{stem}{ext}"
        if p.exists():
            return p

    return None


def resolve_raw_capture_group_from_json(json_path: Path, dataset_root: Path) -> dict:
    """
    Resolve uma captura bruta:
      metas/<base>.json
      previews/<base>.png
      bins/<arquivos das cameras>
      masks/<base> opcional
    """
    json_path = json_path.resolve()
    meta = load_json(json_path)
    base_name = json_path.stem
    bins_dir = dataset_root / "bins"

    preview_path = resolve_raw_sibling_file(
        dataset_root,
        "previews",
        base_name,
        (".png", ".jpg", ".jpeg"),
    )

    mask_path = resolve_raw_sibling_file(
        dataset_root,
        "masks",
        base_name,
        (".npy", ".png", ".tif", ".tiff"),
    )

    group = {
        "json": json_path,
        "png": preview_path,
        "mask": mask_path,
        "final_raw": None,
        "cameras": {},
        "dataset_root": dataset_root,
    }

    if "saved_payload_paths" in meta:
        for cam_id, fname in meta["saved_payload_paths"].items():
            fname_path = Path(fname)

            candidates = [
                bins_dir / fname_path.name,
                json_path.parent / fname,
                dataset_root / fname,
                bins_dir / f"{base_name}_{cam_id}.bin",
                bins_dir / f"{base_name}_{cam_id}.raw",
                bins_dir / f"{base_name}_{cam_id.lower()}.bin",
                bins_dir / f"{base_name}_{cam_id.lower()}.raw",
            ]

            found = next((c for c in candidates if c.exists()), None)
            if found is not None:
                group["cameras"][cam_id] = found
            else:
                print(f"[WARN] bin não encontrado para {cam_id}: {fname}")

    return group


def build_multispec_from_raw_native_multi_for_infer(group: dict, meta: dict) -> Tuple[np.ndarray, dict]:
    """
    Gera tensor MULTISPEC [R,G,B,RE,NIR] em CHW float32 0..1
    a partir dos bins RAW_BRUTO salvos.
    """
    if meta.get("saved_payload_type") != "raw_native_multi":
        raise RuntimeError(
            f"Captura bruta não suportada aqui: saved_payload_type={meta.get('saved_payload_type')}"
        )

    stream_meta = meta.get("stream_meta", {}) or {}
    saved_dtypes = meta.get("saved_payload_dtypes", {}) or {}
    saved_shapes = meta.get("saved_payload_shapes", {}) or {}

    frame = {}

    for cam_id, path in group["cameras"].items():
        saved_dtype = saved_dtypes.get(cam_id)
        saved_shape = saved_shapes.get(cam_id)

        if saved_dtype is None or saved_shape is None:
            raise RuntimeError(f"Faltam dtype/shape para {cam_id} no JSON: {group['json']}")

        arr = np.fromfile(str(path), dtype=np.dtype(saved_dtype)).reshape(tuple(saved_shape))
        frame[cam_id] = arr

    if not frame:
        raise RuntimeError(f"Nenhum bin de câmera encontrado para: {group['json']}")

    sensor_width = int(meta.get("sensor_width", 1280))
    sensor_height = int(meta.get("sensor_height", 800))
    bayer = meta.get("bayer_pattern", "RGGB")

    calib_path = meta.get("camera_params_json") or "calibration/module_params.json"

    if calib_path and not os.path.isfile(calib_path):
        json_dir = Path(group["json"]).parent
        alt = json_dir / calib_path

        if alt.exists():
            calib_path = str(alt)
        else:
            raise RuntimeError(
                f"module_params.json não encontrado: {calib_path}. "
                f"Sem ele eu não gero MULTISPEC, porque faltaria fusion/radiometria/flatfield."
            )

    core = RawProcessorCore(
        sensor_width=sensor_width,
        sensor_height=sensor_height,
        bayer_pattern=bayer,
        calibration_json_path=calib_path,
    )

    processing_meta = dict(stream_meta)

    if meta.get("actual_camera_controls") is not None:
        processing_meta["actual_camera_controls"] = meta.get("actual_camera_controls")

    if meta.get("startup_camera_controls") is not None:
        processing_meta["startup_camera_controls"] = meta.get("startup_camera_controls")

    tensor = core.build_infer_tensor_from_stream(frame, processing_meta, 5)

    if tensor is None:
        raise RuntimeError(f"RawProcessorCore retornou tensor None para: {group['json']}")

    tensor = np.asarray(tensor, dtype=np.float32)
    tensor = np.nan_to_num(tensor, nan=0.0, posinf=1.0, neginf=0.0)
    tensor = np.clip(tensor, 0.0, 1.0)

    processing_info = {
        "fusion_result": getattr(core, "last_fusion_result", None),
        "fusion_config_used": getattr(core, "fusion_config", None),
        "radiometric_normalization_result": getattr(core, "last_radiometric_normalization_result", None),
        "patch_normalization_result": getattr(core, "last_patch_normalization_result", None),
        "frame_quality": getattr(core, "last_frame_quality_result", None),
    }

    return tensor, processing_info


DEFAULT_CHANNEL_ORDER = ["R", "G", "B", "RE", "NIR"]

def get_input_channel_names(config: dict, channels_override: Optional[int] = None) -> List[str]:
    if "input_channels" in config:
        names = [str(c).upper() for c in config["input_channels"]]
    else:
        n = int(channels_override or config.get("channels", 5))
        names = DEFAULT_CHANNEL_ORDER[:n]

    invalid = [c for c in names if c not in DEFAULT_CHANNEL_ORDER]
    if invalid:
        raise RuntimeError(f"Canais inválidos em input_channels: {invalid}")

    return names


def get_input_channel_indices(config: dict, channels_override: Optional[int] = None) -> List[int]:
    names = get_input_channel_names(config, channels_override)
    return [DEFAULT_CHANNEL_ORDER.index(c) for c in names]


# ============================================================
# Dataset
# ============================================================

def collect_samples(root: Path, heads_config: Dict[str, dict], require_masks: bool = False) -> Tuple[List[SampleItem], bool]:
    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")))

    # ============================================================
    # MODO 1: dataset já normalizado com tensors/*.npy
    # ============================================================
    if tensor_paths:
        samples: List[SampleItem] = []
        has_any_gt = False

        for tp in tensor_paths:
            base = tp.stem
            group_name = tp.parent.parent.name if tp.parent.name == "tensors" else "default"
            group_dir = tp.parent.parent if tp.parent.name == "tensors" else tp.parent

            masks: Dict[str, Optional[Path]] = {}
            missing = []

            for head_name, hcfg in heads_config.items():
                mask_dir = group_dir / str(hcfg.get("mask_dir", "masks"))
                mask_path = mask_dir / f"{base}.npy"

                if mask_path.exists():
                    masks[head_name] = mask_path
                    has_any_gt = True
                else:
                    masks[head_name] = None
                    missing.append(f"{head_name}:{mask_path}")

            if require_masks and missing:
                raise RuntimeError(f"Masks ausentes para {tp}: {missing}")

            meta_path = group_dir / "metas" / f"{base}.json"
            preview_path = group_dir / "previews" / f"{base}.png"

            samples.append(SampleItem(
                group=group_name,
                base=base,
                tensor_path=tp,
                masks=masks,
                meta_path=meta_path if meta_path.exists() else None,
                preview_path=preview_path if preview_path.exists() else None,
                source_kind="tensor",
                raw_group=None,
            ))

        return samples, has_any_gt

    # ============================================================
    # MODO 2: dataset bruto com bins/metas/previews/masks
    # ============================================================
    raw_root = find_raw_dataset_layout_root(root)

    if raw_root is None:
        raise RuntimeError(
            f"Nenhum tensor .npy encontrado e também não detectei layout bruto "
            f"com bins/metas/previews em: {root}"
        )

    meta_paths = sorted((raw_root / "metas").glob("*.json"))
    if not meta_paths:
        raise RuntimeError(f"Nenhum meta .json encontrado em: {raw_root / 'metas'}")

    samples: List[SampleItem] = []
    has_any_gt = False

    for mp in meta_paths:
        base = mp.stem
        meta = load_json(mp)

        if meta.get("saved_payload_type") != "raw_native_multi":
            print(f"[WARN] pulando {mp.name}: saved_payload_type={meta.get('saved_payload_type')}")
            continue

        raw_group = resolve_raw_capture_group_from_json(mp, raw_root)

        masks: Dict[str, Optional[Path]] = {}

        # Tenta masks específicas multi-head primeiro.
        for head_name, hcfg in heads_config.items():
            mask_dir = raw_root / str(hcfg.get("mask_dir", "masks"))
            mask_path = None

            mask_dir_name = str(hcfg.get("mask_dir", "masks"))

            if mask_dir_name == "__derived_target__" or bool(hcfg.get("derived", False)):
                masks[head_name] = None
                continue

            for ext in (".npy", ".png", ".tif", ".tiff"):
                p = mask_dir / f"{base}{ext}"
                if p.exists():
                    mask_path = p
                    break

            masks[head_name] = mask_path
            if mask_path is not None:
                has_any_gt = True

        # Se só existir masks/<base>.png ou .npy semantic, conecta na head semantic.
        if masks.get("semantic") is None:
            semantic_mask = resolve_raw_sibling_file(
                raw_root,
                "masks",
                base,
                (".npy", ".png", ".tif", ".tiff"),
            )
            if semantic_mask is not None:
                masks["semantic"] = semantic_mask
                has_any_gt = True

        if require_masks and not any(p is not None for p in masks.values()):
            raise RuntimeError(f"Mask ausente para captura bruta: {mp}")

        preview_path = resolve_raw_sibling_file(
            raw_root,
            "previews",
            base,
            (".png", ".jpg", ".jpeg"),
        )

        samples.append(SampleItem(
            group=raw_root.name,
            base=base,
            tensor_path=None,
            masks=masks,
            meta_path=mp,
            preview_path=preview_path,
            source_kind="raw_native_multi",
            raw_group=raw_group,
        ))

    if not samples:
        raise RuntimeError(f"Nenhuma captura raw_native_multi válida encontrada em: {raw_root}")

    return samples, has_any_gt


def load_tensor(
    path: Path,
    channels: int,
    channel_indices: Optional[List[int]] = None,
) -> 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")

    # Normaliza para CHW.
    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}")

    if channel_indices is not None:
        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}. Shape={chw.shape}"
            )
        chw = chw[channel_indices, :, :]
    else:
        if chw.shape[0] < channels:
            raise RuntimeError(
                f"Tensor {path} tem {chw.shape[0]} canais, "
                f"mas config pediu {channels}."
            )
        chw = chw[:channels, :, :]

    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 load_sample_tensor(
    sample: SampleItem,
    channels: int,
    channel_indices: Optional[List[int]] = None,
) -> np.ndarray:
    """
    Carrega tensor de uma amostra.
    - Se for tensor pronto: lê .npy.
    - Se for RAW_BRUTO: gera MULTISPEC em tempo real a partir dos bins.
    """
    if sample.source_kind == "tensor":
        if sample.tensor_path is None:
            raise RuntimeError(f"Sample tensor sem tensor_path: {sample.base}")
        return load_tensor(
            sample.tensor_path,
            channels=channels,
            channel_indices=channel_indices,
        )

    if sample.source_kind == "raw_native_multi":
        if sample.raw_group is None or sample.meta_path is None:
            raise RuntimeError(f"Sample raw sem raw_group/meta_path: {sample.base}")

        meta = load_json(sample.meta_path)
        tensor, _processing_info = build_multispec_from_raw_native_multi_for_infer(sample.raw_group, meta)

        if tensor.ndim != 3:
            raise RuntimeError(f"Tensor RAW gerado inválido: {sample.base} shape={tensor.shape}")

        if channel_indices is not None:
            max_idx = max(channel_indices)
            if tensor.shape[0] <= max_idx:
                raise RuntimeError(
                    f"Tensor RAW tem {tensor.shape[0]} canais, "
                    f"mas precisa acessar índice {max_idx}."
                )
            tensor = tensor[channel_indices, :, :]
        else:
            if tensor.shape[0] < channels:
                raise RuntimeError(
                    f"Tensor RAW com canais incompatíveis: {sample.base} "
                    f"shape={tensor.shape}, esperado pelo menos {channels}"
                )
            tensor = tensor[:channels, :, :]

        return tensor.astype(np.float32)

    raise RuntimeError(f"source_kind desconhecido: {sample.source_kind}")


def load_mask(path: Optional[Path]) -> Optional[np.ndarray]:
    if path is None:
        return None
    if path.suffix.lower() == ".npy":
        mask = np.load(str(path))
    else:
        mask = cv2.imread(str(path), cv2.IMREAD_UNCHANGED)
        if mask is None:
            raise RuntimeError(f"Falha ao ler mask: {path}")
        if mask.ndim == 3:
            mask = mask[:, :, 0]
    return mask.astype(np.int64)


# ============================================================
# Normalização
# ============================================================

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

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

    if mean is None or std is None:
        raise RuntimeError(f"norm_stats inválido, faltando mean/std: {path}")
    if channel_indices is not None:
        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 = [mean[i] for i in channel_indices]
        std = [std[i] for i in channel_indices]

        if names:
            names = [names[i] for i in channel_indices]
        elif channel_names:
            names = channel_names

    else:
        if len(mean) != channels or len(std) != channels:
            raise RuntimeError(
                f"norm_stats incompatível com channels={channels}: "
                f"mean={len(mean)} std={len(std)}"
            )

    print(f"[NORM] usando {path}")
    print(f"[NORM] channels={names}")
    print(f"[NORM] mean={mean}")
    print(f"[NORM] std ={std}")
    return list(map(float, mean)), list(map(float, std))


# ============================================================
# Modelo multi-head, igual ao treino
# ============================================================

def patch_segformer_encoder_input_channels(segformer_encoder: nn.Module, in_ch: int):
    if in_ch == 3:
        return segformer_encoder

    proj = segformer_encoder.encoder.patch_embeddings[0].proj

    if proj.in_channels == in_ch:
        return segformer_encoder

    old_weight = proj.weight.data.clone()
    old_bias = proj.bias.data.clone() if proj.bias is not None else None

    new_proj = nn.Conv2d(
        in_channels=in_ch,
        out_channels=proj.out_channels,
        kernel_size=proj.kernel_size,
        stride=proj.stride,
        padding=proj.padding,
        dilation=proj.dilation,
        groups=proj.groups,
        bias=proj.bias is not None,
        padding_mode=proj.padding_mode,
    )

    with torch.no_grad():
        if in_ch <= old_weight.shape[1]:
            new_proj.weight.copy_(old_weight[:, :in_ch, :, :])
        else:
            new_proj.weight[:, :old_weight.shape[1], :, :].copy_(old_weight)
            extra = in_ch - old_weight.shape[1]
            mean_w = old_weight.mean(dim=1, keepdim=True)
            new_proj.weight[:, old_weight.shape[1]:, :, :].copy_(mean_w.repeat(1, extra, 1, 1))

        if old_bias is not None:
            new_proj.bias.copy_(old_bias)

    segformer_encoder.encoder.patch_embeddings[0].proj = new_proj
    print(f"[MODEL] patch input channels: 3 -> {in_ch}")
    return segformer_encoder


def replace_segformer_decode_classifier(decode_head: nn.Module, num_classes: int):
    old = decode_head.classifier
    if not isinstance(old, nn.Conv2d):
        raise RuntimeError(f"decode_head.classifier não é Conv2d: {type(old)}")

    new = nn.Conv2d(
        in_channels=old.in_channels,
        out_channels=int(num_classes),
        kernel_size=old.kernel_size,
        stride=old.stride,
        padding=old.padding,
        dilation=old.dilation,
        groups=old.groups,
        bias=old.bias is not None,
        padding_mode=old.padding_mode,
    )
    decode_head.classifier = new
    return decode_head


class MultiHeadSegFormer(nn.Module):
    def __init__(
        self,
        backbone: str,
        channels: int,
        heads_config: Dict[str, dict],
        semantic_id2label: Dict[int, str],
        semantic_label2id: Dict[str, int],
    ):
        super().__init__()

        semantic_classes = int(heads_config["semantic"].get("num_classes", len(semantic_id2label)))

        base = SegformerForSemanticSegmentation.from_pretrained(
            backbone,
            num_labels=semantic_classes,
            id2label={int(k): str(v) for k, v in semantic_id2label.items()},
            label2id={str(k): int(v) for k, v in semantic_label2id.items()},
            ignore_mismatched_sizes=True,
        )

        patch_segformer_encoder_input_channels(base.segformer, channels)
        base.config.num_channels = int(channels)

        self.segformer = base.segformer
        self.decode_heads = nn.ModuleDict()
        self.heads_config = heads_config

        for head_name, hcfg in heads_config.items():
            h = copy.deepcopy(base.decode_head)
            h = replace_segformer_decode_classifier(h, int(hcfg["num_classes"]))
            self.decode_heads[head_name] = h

        self.config = base.config

    def forward(self, pixel_values: torch.Tensor, head_names=None) -> Dict[str, torch.Tensor]:
        outputs = self.segformer(
            pixel_values=pixel_values,
            output_hidden_states=True,
            return_dict=True,
        )

        hidden_states = outputs.hidden_states

        if head_names is None:
            selected = list(self.decode_heads.keys())
        else:
            selected = [h for h in head_names if h in self.decode_heads]

        return {
            head_name: self.decode_heads[head_name](hidden_states)
            for head_name in selected
        }


class MultiHeadTester:
    def __init__(
        self,
        config: dict,
        ckpt_path: Path,
        device: torch.device,
        channels: int,
        heads_config: Dict[str, dict],
        semantic_id2label: Dict[int, str],
        semantic_label2id: Dict[str, int],
        mean: Optional[Sequence[float]],
        std: Optional[Sequence[float]],
        use_amp: bool = True,
    ):
        self.config = config
        self.ckpt_path = ckpt_path
        self.device = device
        self.channels = channels
        self.heads_config = heads_config
        self.use_amp = use_amp and device.type == "cuda"
        self.runtime_mode = str(config.get("runtime_mode", "all")).lower()

        self.mean = None if mean is None else torch.tensor(mean, dtype=torch.float32).view(1, channels, 1, 1).to(device)
        self.std = None if std is None else torch.tensor(std, dtype=torch.float32).view(1, channels, 1, 1).to(device)

        backbone = config.get("backbone", config.get("pretrained_model", "nvidia/mit-b1"))
        print(f"[MODEL] backbone={backbone}")
        print(f"[MODEL] ckpt={ckpt_path}")

        self.model = MultiHeadSegFormer(
            backbone=backbone,
            channels=channels,
            heads_config=heads_config,
            semantic_id2label=semantic_id2label,
            semantic_label2id=semantic_label2id,
        )

        self._load_checkpoint(ckpt_path)
        self.model.to(device)
        self.model.eval()

    def _load_checkpoint(self, ckpt_path: Path):
        ckpt = torch.load(str(ckpt_path), map_location="cpu", weights_only=False)

        if isinstance(ckpt, dict):
            for key in ("model", "model_state", "model_state_dict", "state_dict"):
                if key in ckpt and isinstance(ckpt[key], dict):
                    state = ckpt[key]
                    break
            else:
                state = ckpt
        else:
            raise RuntimeError(f"Checkpoint em formato inesperado: {type(ckpt)}")

        clean = {}
        for k, v in state.items():
            nk = k
            for prefix in ("module.", "model."):
                if nk.startswith(prefix):
                    nk = nk[len(prefix):]
            clean[nk] = v

        missing, unexpected = self.model.load_state_dict(clean, strict=False)
        print(f"[MODEL] load_state_dict strict=False | missing={len(missing)} unexpected={len(unexpected)}")
        if missing:
            print("[MODEL] primeiros missing:", missing[:8])
        if unexpected:
            print("[MODEL] primeiros unexpected:", unexpected[:8])

    def _normalize(self, x: torch.Tensor) -> torch.Tensor:
        if self.mean is not None and self.std is not None:
            return (x - self.mean) / torch.clamp(self.std, min=1e-6)
        return x

    @torch.inference_mode()
    def infer(self, chw_01: np.ndarray) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray], float]:
        x = torch.from_numpy(chw_01).unsqueeze(0).to(self.device, non_blocking=True)
        x = self._normalize(x)

        h, w = int(chw_01.shape[1]), int(chw_01.shape[2])

        if self.device.type == "cuda":
            torch.cuda.synchronize()
        t0 = time.perf_counter()

        with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=self.use_amp):
            head_names = None
            if self.runtime_mode in ("target_direct", "target_head"):
                head_names = ["target"]
            elif self.runtime_mode in ("operational", "target_op"):
                head_names = ["vegetation", "cana"]

            logits_by_head = self.model(pixel_values=x, head_names=head_names)

            preds = {}
            probs = {}

            for head_name, logits in logits_by_head.items():
                logits = F.interpolate(logits, size=(h, w), mode="bilinear", align_corners=False)
                prob = torch.softmax(logits, dim=1)[0]
                pred = torch.argmax(prob, dim=0)
                preds[head_name] = pred.detach().cpu().numpy().astype(np.uint8)
                probs[head_name] = prob.detach().cpu().numpy().astype(np.float32)

        if self.device.type == "cuda":
            torch.cuda.synchronize()
        t_ms = (time.perf_counter() - t0) * 1000.0

        return preds, probs, t_ms


class OnnxMultiHeadTester:
    def __init__(
        self,
        config: dict,
        onnx_path: Path,
        provider: str,
        channels: int,
        heads_config: Dict[str, dict],
        mean: Optional[Sequence[float]],
        std: Optional[Sequence[float]],
        trt_home: Optional[str] = None,
        trt_fp16: bool = True,
    ):
        self.config = config
        self.onnx_path = onnx_path
        self.provider = provider.lower()
        self.channels = int(channels)
        self.heads_config = heads_config
        self.runtime_mode = str(config.get("runtime_mode", "all")).lower()

        self.mean = None if mean is None else np.asarray(mean, dtype=np.float32).reshape(1, channels, 1, 1)
        self.std = None if std is None else np.asarray(std, dtype=np.float32).reshape(1, channels, 1, 1)

        print(f"[ONNX_MODEL] onnx={onnx_path}")
        print(f"[ONNX_MODEL] provider={provider}")

        self.session = self._create_session(
            onnx_path=onnx_path,
            provider=provider,
            trt_home=trt_home,
            trt_fp16=trt_fp16,
        )

        self.input_name = self.session.get_inputs()[0].name
        self.output_names = [o.name for o in self.session.get_outputs()]
        print(f"[ONNX_MODEL] input={self.input_name}")
        print(f"[ONNX_MODEL] outputs={self.output_names}")

    def _create_session(
        self,
        onnx_path: Path,
        provider: str,
        trt_home: Optional[str],
        trt_fp16: bool,
    ):
        try:
            import onnxruntime as ort
        except ImportError:
            raise ImportError(
                "onnxruntime não está instalado. Use:\n"
                "  pip install onnxruntime-gpu"
            )

        provider = provider.lower()

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

            dll_dirs = [
                os.path.join(trt_home, "lib"),
                os.path.join(trt_home, "bin"),
            ]

            cuda_home = os.environ.get("CUDA_PATH")
            if cuda_home:
                dll_dirs.append(os.path.join(cuda_home, "bin"))

            # fallback comum que vocês estão usando
            dll_dirs.append(r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.3\bin")
            dll_dirs.append(r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\bin")

            for dll_dir in dll_dirs:
                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}")

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

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

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

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

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

            trt_options = {
                "device_id": 0,
                "trt_fp16_enable": bool(trt_fp16),
                "trt_engine_cache_enable": True,
                "trt_engine_cache_path": str(cache_dir),
                "trt_timing_cache_enable": True,
                "trt_timing_cache_path": str(cache_dir),
                "trt_max_workspace_size": 4 * 1024 * 1024 * 1024,
            }

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

        else:
            raise RuntimeError(f"Provider ONNX inválido: {provider}")

        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 ONNX disponível. Pedido={providers}, disponíveis={available}")

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

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

        if provider == "tensorrt" and "TensorrtExecutionProvider" not in active:
            raise RuntimeError(f"TensorRT solicitado, mas não ficou ativo. Providers ativos: {active}")

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

        return session

    def _normalize(self, x: np.ndarray) -> np.ndarray:
        if self.mean is not None and self.std is not None:
            return ((x - self.mean[0]) / np.clip(self.std[0], 1e-6, None)).astype(np.float32)
        return x.astype(np.float32)

    def _selected_heads(self) -> Optional[List[str]]:
        if self.runtime_mode in ("target_direct", "target_head"):
            return ["target"]
        if self.runtime_mode in ("operational", "target_op"):
            return ["vegetation", "cana"]
        return None

    @staticmethod
    def _softmax_np(logits: np.ndarray, axis: int = 1) -> np.ndarray:
        x = logits.astype(np.float32)
        x = x - np.max(x, axis=axis, keepdims=True)
        e = np.exp(x)
        return e / np.clip(np.sum(e, axis=axis, keepdims=True), 1e-12, None)

    @staticmethod
    def _resize_logits_nchw(logits: np.ndarray, target_hw: Tuple[int, int]) -> np.ndarray:
        n, c, h, w = logits.shape
        th, tw = target_hw

        if (h, w) == (th, tw):
            return logits

        out = np.empty((n, c, th, tw), dtype=np.float32)
        for bi in range(n):
            for ci in range(c):
                out[bi, ci] = cv2.resize(
                    logits[bi, ci].astype(np.float32),
                    (tw, th),
                    interpolation=cv2.INTER_LINEAR,
                )
        return out

    def _map_outputs(self, outputs: List[np.ndarray]) -> Dict[str, np.ndarray]:
        raw = {
            name: arr.astype(np.float32)
            for name, arr in zip(self.output_names, outputs)
        }

        mapped = {}
        for head in self.heads_config.keys():
            candidates = [
                head,
                f"{head}_logits",
                f"output_{head}",
            ]

            found = None
            for c in candidates:
                if c in raw:
                    found = c
                    break

            if found is not None:
                mapped[head] = raw[found]

        return mapped

    def infer(self, chw_01: np.ndarray) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray], float]:
        h, w = int(chw_01.shape[1]), int(chw_01.shape[2])

        x = self._normalize(chw_01)
        x = np.expand_dims(x, axis=0).astype(np.float32)

        t0 = time.perf_counter()
        outputs = self.session.run(None, {self.input_name: x})
        t_ms = (time.perf_counter() - t0) * 1000.0

        logits_by_head = self._map_outputs(outputs)

        selected = self._selected_heads()
        if selected is not None:
            logits_by_head = {
                hname: logits
                for hname, logits in logits_by_head.items()
                if hname in selected
            }

        preds = {}
        probs = {}

        for head_name, logits in logits_by_head.items():
            logits = self._resize_logits_nchw(logits, (h, w))
            prob = self._softmax_np(logits, axis=1)[0]
            pred = np.argmax(prob, axis=0).astype(np.uint8)

            preds[head_name] = pred
            probs[head_name] = prob.astype(np.float32)

        return preds, probs, t_ms


# ============================================================
# Checkpoints / paths
# ============================================================

def infer_experiment_tag(config: dict, channels: int) -> str:
    fusion_mode = config.get("fusion_mode", "stacked")
    return f"{fusion_mode}_raw{channels}"


def infer_save_dir(config: dict, config_dir: Path, channels: int) -> Path:
    model_name = config.get("model_name", "test_multi")
    modelo_folder = config.get("modelo", "segformer_b1")
    exp_tag = infer_experiment_tag(config, channels)
    return (config_dir / "backup" / modelo_folder / model_name / exp_tag).resolve()


def find_checkpoint(save_dir: Path, preferred: Optional[str] = None) -> Path:
    if preferred is not None:
        ckpt = Path(preferred)
        if not ckpt.is_absolute():
            ckpt_cwd = (Path.cwd() / ckpt).resolve()
            ckpt_save = (save_dir / ckpt).resolve()
            ckpt = ckpt_cwd if ckpt_cwd.is_file() else ckpt_save
        if not ckpt.is_file():
            raise FileNotFoundError(f"Checkpoint não encontrado: {ckpt}")
        return ckpt

    candidates = [
        save_dir / "best_score.pt",
        save_dir / "best_target.pt",
        save_dir / "best_cana_head.pt",
        save_dir / "best_semantic_miou.pt",
        save_dir / "last.pt",
    ]
    for c in candidates:
        if c.is_file():
            return c

    raise FileNotFoundError("Nenhum checkpoint encontrado. Procurei:\n" + "\n".join(str(c) for c in candidates))


def find_onnx_model(save_dir: Path, ckpt_path: Path, preferred: Optional[str] = None) -> Path:
    """
    Resolve o .onnx.

    Se preferred for informado, usa ele.
    Caso contrário, usa o mesmo stem do checkpoint:
      best_score.pt -> best_score.onnx
    """
    if preferred:
        p = Path(preferred)
        if not p.is_absolute():
            p_cwd = (Path.cwd() / p).resolve()
            p_save = (save_dir / p).resolve()
            p = p_cwd if p_cwd.is_file() else p_save

        if not p.is_file():
            raise FileNotFoundError(f"ONNX não encontrado: {p}")

        return p.resolve()

    p = ckpt_path.with_suffix(".onnx")

    if not p.is_file():
        alt = save_dir / f"{ckpt_path.stem}.onnx"
        p = alt

    if not p.is_file():
        raise FileNotFoundError(
            f"ONNX não encontrado para checkpoint {ckpt_path.name}. Procurei:\n"
            f"  {ckpt_path.with_suffix('.onnx')}\n"
            f"  {save_dir / (ckpt_path.stem + '.onnx')}\n"
            f"Informe manualmente com --onnx."
        )

    return p.resolve()


# ============================================================
# Visualização
# ============================================================

def tensor_to_preview_rgb(chw: np.ndarray, gamma: float = 0.85) -> np.ndarray:
    c, h, w = chw.shape
    if c >= 3:
        rgb = np.transpose(chw[:3], (1, 2, 0)).copy()
    else:
        one = chw[0]
        rgb = np.stack([one, one, one], axis=-1)

    rgb = np.nan_to_num(rgb, nan=0.0, posinf=1.0, neginf=0.0)
    lo = np.percentile(rgb, 1.0)
    hi = np.percentile(rgb, 99.0)
    if hi > lo:
        rgb = (rgb - lo) / (hi - lo)
    rgb = np.clip(rgb, 0.0, 1.0)
    if gamma and gamma > 0:
        rgb = np.power(rgb, gamma)
    return (rgb * 255.0).astype(np.uint8)


def ids_to_rgb(mask: np.ndarray, colormap_rgb: Dict[int, Tuple[int, int, int]], ignore_id: int = 255) -> np.ndarray:
    h, w = mask.shape[:2]
    out = np.zeros((h, w, 3), dtype=np.uint8)
    for cid, color in colormap_rgb.items():
        out[mask == cid] = color
    out[mask == ignore_id] = (0, 0, 0)
    return out


def prob_to_heat_rgb(prob01: np.ndarray) -> np.ndarray:
    p = np.clip(prob01, 0.0, 1.0)
    u8 = (p * 255.0).astype(np.uint8)
    bgr = cv2.applyColorMap(u8, cv2.COLORMAP_TURBO)
    return cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)


def overlay_rgb(base_rgb: np.ndarray, mask_rgb: np.ndarray, alpha: float) -> np.ndarray:
    return cv2.addWeighted(base_rgb, 1.0 - alpha, mask_rgb, alpha, 0.0)


def put_label(img_rgb: np.ndarray, title: str, subtitle: str = "") -> np.ndarray:
    out = img_rgb.copy()
    cv2.putText(out, title, (10, 26), cv2.FONT_HERSHEY_SIMPLEX, 0.68, (0, 0, 0), 4, cv2.LINE_AA)
    cv2.putText(out, title, (10, 26), cv2.FONT_HERSHEY_SIMPLEX, 0.68, (255, 255, 255), 2, cv2.LINE_AA)
    if subtitle:
        cv2.putText(out, subtitle, (10, 52), cv2.FONT_HERSHEY_SIMPLEX, 0.48, (0, 0, 0), 3, cv2.LINE_AA)
        cv2.putText(out, subtitle, (10, 52), cv2.FONT_HERSHEY_SIMPLEX, 0.48, (0, 255, 90), 1, cv2.LINE_AA)
    return out


def resize_panel(img: np.ndarray, size: Tuple[int, int]) -> np.ndarray:
    w, h = size
    if img.shape[1] == w and img.shape[0] == h:
        return img
    return cv2.resize(img, (w, h), interpolation=cv2.INTER_NEAREST)


def compose_grid(panels: List[Tuple[str, np.ndarray, str]], cols: int = 3, max_width: int = 1800) -> np.ndarray:
    if not panels:
        return np.zeros((480, 640, 3), dtype=np.uint8)

    base_h, base_w = panels[0][1].shape[:2]
    labeled = []
    for title, img, subtitle in panels:
        img = resize_panel(img, (base_w, base_h))
        labeled.append(put_label(img, title, subtitle))

    rows = []
    blank = np.zeros_like(labeled[0])
    for i in range(0, len(labeled), cols):
        row_imgs = labeled[i:i + cols]
        while len(row_imgs) < cols:
            row_imgs.append(blank.copy())
        rows.append(np.hstack(row_imgs))

    canvas = np.vstack(rows)

    if canvas.shape[1] > max_width:
        scale = max_width / canvas.shape[1]
        canvas = cv2.resize(canvas, (int(canvas.shape[1] * scale), int(canvas.shape[0] * scale)), interpolation=cv2.INTER_AREA)

    return canvas


def class_percent(mask: np.ndarray, class_id: int, ignore_id: int = 255) -> float:
    valid = mask != ignore_id
    den = int(valid.sum())
    if den <= 0:
        return 0.0
    return float(((mask == class_id) & valid).sum() * 100.0 / den)


def compare_pred_equal_percent(a: Optional[np.ndarray], b: Optional[np.ndarray]) -> Optional[float]:
    if a is None or b is None:
        return None

    if a.shape != b.shape:
        b = cv2.resize(
            b.astype(np.uint8),
            (a.shape[1], a.shape[0]),
            interpolation=cv2.INTER_NEAREST,
        )

    return float(np.mean(a == b) * 100.0)


def diff_mask_rgb(a: Optional[np.ndarray], b: Optional[np.ndarray]) -> Optional[np.ndarray]:
    if a is None or b is None:
        return None

    if a.shape != b.shape:
        b = cv2.resize(
            b.astype(np.uint8),
            (a.shape[1], a.shape[0]),
            interpolation=cv2.INTER_NEAREST,
        )

    diff = (a != b).astype(np.uint8) * 255
    rgb = np.zeros((diff.shape[0], diff.shape[1], 3), dtype=np.uint8)
    rgb[:, :, 0] = diff  # vermelho em RGB
    return rgb


# ============================================================
# Métricas numpy
# ============================================================

def confusion_matrix_np(pred: np.ndarray, gt: np.ndarray, num_classes: int, ignore_id: int) -> np.ndarray:
    if pred.shape != gt.shape:
        pred = cv2.resize(pred.astype(np.uint8), (gt.shape[1], gt.shape[0]), interpolation=cv2.INTER_NEAREST)

    valid = gt != ignore_id
    valid &= gt >= 0
    valid &= gt < num_classes
    gt_v = gt[valid].astype(np.int64)
    pred_v = pred[valid].astype(np.int64)
    pred_v = np.clip(pred_v, 0, num_classes - 1)

    cm = np.bincount(num_classes * gt_v + pred_v, minlength=num_classes * num_classes)
    return cm.reshape(num_classes, num_classes).astype(np.int64)


def metrics_from_cm(cm: np.ndarray) -> Tuple[np.ndarray, float, float]:
    tp = np.diag(cm).astype(np.float64)
    fp = cm.sum(axis=0).astype(np.float64) - tp
    fn = cm.sum(axis=1).astype(np.float64) - tp
    denom = tp + fp + fn
    iou = np.divide(tp, denom, out=np.zeros_like(tp), where=denom > 0)
    miou = float(np.mean(iou)) if len(iou) else 0.0
    acc = float(tp.sum() / max(cm.sum(), 1))
    return iou, miou, acc


def operational_target_mask(veg_mask: np.ndarray, cana_mask: np.ndarray, ignore_id: int = 255) -> np.ndarray:
    out = np.zeros_like(veg_mask, dtype=np.uint8)
    ignore = (veg_mask == ignore_id) | (cana_mask == ignore_id)
    out[(veg_mask == 1) & (cana_mask == 0)] = 1
    out[ignore] = ignore_id
    return out


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

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="config.json")
    parser.add_argument("--split_folder", default="val", choices=["train", "val", "test"])
    parser.add_argument("--root_override", default=None)
    parser.add_argument("--test_folder", default=None)
    parser.add_argument("--ckpt", default=None)
    parser.add_argument("--norm_stats", default=None)
    parser.add_argument("--channels", type=int, default=None)
    parser.add_argument("--resize_w", type=int, default=None)
    parser.add_argument("--resize_h", type=int, default=None)
    parser.add_argument("--alpha", type=float, default=0.45)
    parser.add_argument("--ignore_id", type=int, default=None)
    parser.add_argument("--no_amp", action="store_true")
    parser.add_argument("--require_masks", action="store_true")
    parser.add_argument("--out_dir", default="outputs_test_multihead")
    parser.add_argument("--start_idx", type=int, default=0)
    parser.add_argument("--max_width", type=int, default=1800)
    parser.add_argument("--runtime_mode", default="all", choices=["all", "target_direct", "operational"])
    parser.add_argument("--onnx", default="")
    parser.add_argument("--onnx_provider", default="", choices=["", "cpu", "cuda", "tensorrt"])
    parser.add_argument("--trt_home", default=None)
    parser.add_argument("--trt_no_fp16", action="store_true")
    args = parser.parse_args()

    config_path = resolve_path(args.config, Path.cwd())
    if config_path is None or not config_path.is_file():
        raise FileNotFoundError(f"Config não encontrado: {config_path}")

    config_dir = config_path.parent
    config = load_json(config_path)
    config["runtime_mode"] = args.runtime_mode

    input_channel_names = get_input_channel_names(config, args.channels)
    input_channel_indices = get_input_channel_indices(config, args.channels)
    channels = len(input_channel_names)

    print(f"Input channels: {input_channel_names} idx={input_channel_indices}")
    res = config.get("resolucao", [1024, 640])
    default_w, default_h = int(res[0]), int(res[1])
    target_w = int(args.resize_w or default_w)
    target_h = int(args.resize_h or default_h)

    dataset_path = config_dir / "dataset"
    labelmap_path = dataset_path / "labelmap.txt"
    semantic_id2label, semantic_label2id, labelmap_ignore_id, loaded_colormap_rgb = load_labelmap(labelmap_path)

    ignore_id = int(args.ignore_id if args.ignore_id is not None else labelmap_ignore_id)
    heads_config = build_heads_config(config, ignore_index=ignore_id)
    heads_config["semantic"]["num_classes"] = int(len(semantic_id2label))
    heads_config["semantic"]["ignore_index"] = int(ignore_id)

    semantic_cmap = dict(SEMANTIC_COLORS_RGB)
    semantic_cmap.update({int(k): tuple(map(int, v)) for k, v in loaded_colormap_rgb.items()})

    save_dir = infer_save_dir(config, config_dir, channels)
    ckpt_path = find_checkpoint(save_dir, args.ckpt)

    if args.norm_stats is not None:
        norm_stats_path = resolve_path(args.norm_stats, Path.cwd())
    else:
        # O treino multihead usa stats do dataset normalizado, mas também tentamos alguns fallbacks.
        candidates = [
            dataset_path / f"{default_w}x{default_h}" / "group" / "norm_stats.json",
            save_dir / "norm_stats.json",
            config_dir / "backup" / config.get("modelo", "segformer_b1") / config.get("model_name", "test_multi") / config.get("stats_source_tag", "stacked_raw5") / "norm_stats.json",
        ]
        norm_stats_path = next((p for p in candidates if p.is_file()), candidates[0])

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

    if args.test_folder is not None:
        root = resolve_path(args.test_folder, Path.cwd())
    elif args.root_override is not None:
        root = resolve_path(args.root_override, Path.cwd())
    else:
        root = (dataset_path / "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, has_gt = collect_samples(root, heads_config=heads_config, require_masks=args.require_masks)
    n = len(samples)

    print("==========================================")
    print("Teste SegFormer OAK-FCC-3 Multi-Head")
    print(f"Root        : {root}")
    print(f"Samples     : {n}")
    print(f"GT          : {'sim' if has_gt else 'não'}")
    print(f"Resolution  : {target_w}x{target_h}")
    print(f"Channels    : {channels}")
    print(f"Semantic    : {semantic_id2label}")
    print(f"Ignore index: {ignore_id}")
    print("Heads:")
    for name, hcfg in heads_config.items():
        print(f"  - {name}: classes={hcfg['num_classes']} mask_dir={hcfg['mask_dir']}")
    print("==========================================")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    tester = MultiHeadTester(
        config=config,
        ckpt_path=ckpt_path,
        device=device,
        channels=channels,
        heads_config=heads_config,
        semantic_id2label=semantic_id2label,
        semantic_label2id=semantic_label2id,
        mean=mean,
        std=std,
        use_amp=not args.no_amp,
    )

    onnx_tester = None
    onnx_path = None

    if args.onnx_provider:
        onnx_path = find_onnx_model(
            save_dir=save_dir,
            ckpt_path=ckpt_path,
            preferred=args.onnx,
        )

        onnx_tester = OnnxMultiHeadTester(
            config=config,
            onnx_path=onnx_path,
            provider=args.onnx_provider,
            channels=channels,
            heads_config=heads_config,
            mean=mean,
            std=std,
            trt_home=args.trt_home,
            trt_fp16=not args.trt_no_fp16,
        )

    if onnx_tester is not None:
        print(f"ONNX        : {onnx_path}")
        print(f"ONNX provider: {args.onnx_provider}")

    out_dir = Path(args.out_dir)
    ensure_dir(out_dir)

    idx = max(0, min(args.start_idx, n - 1))
    detailed = False

    cms_total = {
        name: np.zeros((int(hcfg["num_classes"]), int(hcfg["num_classes"])), dtype=np.int64)
        for name, hcfg in heads_config.items()
    }
    cm_target_total = np.zeros((2, 2), dtype=np.int64)
    visited = set()

    win_name = "OAK-FCC-3 MultiHead Test | D/A navega | S salva | SPACE detalhado | Q sai"
    cv2.namedWindow(win_name, cv2.WINDOW_NORMAL)
    win_name_onnx = None
    if onnx_tester is not None:
        win_name_onnx = "ONNX/TensorRT MultiHead Test | comparação visual"
        cv2.namedWindow(win_name_onnx, cv2.WINDOW_NORMAL)

    while True:
        sample = samples[idx]
        chw = load_sample_tensor(
            sample,
            channels=channels,
            channel_indices=input_channel_indices,
        )

        if (chw.shape[2], chw.shape[1]) != (target_w, target_h):
            hwc = np.transpose(chw, (1, 2, 0))
            hwc = cv2.resize(hwc, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
            chw = np.transpose(hwc, (2, 0, 1)).astype(np.float32)

        preds, probs, t_inf = tester.infer(chw)
        preview_rgb = tensor_to_preview_rgb(chw)

        onnx_preds = None
        onnx_probs = None
        t_onnx = None

        if onnx_tester is not None:
            onnx_preds, onnx_probs, t_onnx = onnx_tester.infer(chw)

        first_pred = next(iter(preds.values()))
        pred_h, pred_w = first_pred.shape[:2]

        pred_sem = preds.get("semantic")
        pred_veg = preds.get("vegetation")
        pred_cana = preds.get("cana")

        # Target operacional antigo: vegetation AND not cana
        pred_target_op = None
        if pred_veg is not None and pred_cana is not None:
            pred_target_op = operational_target_mask(pred_veg, pred_cana, ignore_id=ignore_id)

        # Target direta nova: saída da 4ª cabeça
        pred_target_head = preds.get("target")

        # Fallback visual: se não tiver head target, mostra operacional
        pred_target = pred_target_head if pred_target_head is not None else pred_target_op

        prob_veg = probs["vegetation"][1] if "vegetation" in probs and probs["vegetation"].shape[0] > 1 else None
        prob_cana = probs["cana"][1] if "cana" in probs and probs["cana"].shape[0] > 1 else None

        prob_target_op = None
        if prob_veg is not None and prob_cana is not None:
            prob_target_op = np.clip(prob_veg * (1.0 - prob_cana), 0.0, 1.0)

        prob_target_head = None
        if "target" in probs:
            prob_target_head = probs["target"][1] if probs["target"].shape[0] > 1 else probs["target"][0]

        prob_target = prob_target_head if prob_target_head is not None else prob_target_op

        gt_masks = {name: load_mask(path) for name, path in sample.masks.items()}
        for name, gt in list(gt_masks.items()):
            if gt is not None and gt.shape != (pred_h, pred_w):
                gt_masks[name] = cv2.resize(
                    gt.astype(np.uint8),
                    (pred_w, pred_h),
                    interpolation=cv2.INTER_NEAREST,
                )

        gt_target = None
        if gt_masks.get("target") is not None:
            gt_target = gt_masks["target"]
        elif gt_masks.get("vegetation") is not None and gt_masks.get("cana") is not None:
            gt_target = operational_target_mask(
                gt_masks["vegetation"],
                gt_masks["cana"],
                ignore_id=ignore_id,
            )
            gt_masks["target"] = gt_target

        # Métricas da amostra.
        metric_lines = []
        sample_metrics = {}
        for head_name, pred in preds.items():
            gt = gt_masks.get(head_name)
            if gt is not None:
                cm = confusion_matrix_np(pred, gt, int(heads_config[head_name]["num_classes"]), ignore_id)
                iou, miou, acc = metrics_from_cm(cm)
                sample_metrics[head_name] = {"iou": iou, "miou": miou, "acc": acc}
                metric_lines.append(f"{head_name}: mIoU={miou:.3f} acc={acc:.3f}")

        if gt_target is not None:
            cm_t = confusion_matrix_np(pred_target, gt_target, 2, ignore_id)
            iou_t, miou_t, acc_t = metrics_from_cm(cm_t)
            sample_metrics["target_op"] = {"iou": iou_t, "miou": miou_t, "acc": acc_t}
            metric_lines.append(f"target_op: IoU_alvo={iou_t[1]:.3f} acc={acc_t:.3f}")
            if "target" in sample_metrics:
                iou_head = sample_metrics["target"]["iou"]
                metric_lines.append(f"target_head: IoU_alvo={iou_head[1]:.3f}")

        if idx not in visited:
            for head_name, pred in preds.items():
                gt = gt_masks.get(head_name)
                if gt is not None:
                    cms_total[head_name] += confusion_matrix_np(pred, gt, int(heads_config[head_name]["num_classes"]), ignore_id)
            if gt_target is not None and pred_target_op is not None:
                cm_target_total += confusion_matrix_np(pred_target_op, gt_target, 2, ignore_id)
            visited.add(idx)

        # Visuals.
        panels: List[Tuple[str, np.ndarray, str]] = []
        if pred_sem is not None:
            pred_sem_rgb = ids_to_rgb(pred_sem, semantic_cmap, ignore_id)
            if gt_masks.get("semantic") is not None:
                panels.append(("GT semantic", ids_to_rgb(gt_masks["semantic"], semantic_cmap, ignore_id), "chao/cana/erva"))
            panels.append(("Pred semantic", pred_sem_rgb, f"erva={class_percent(pred_sem, 2, ignore_id):.1f}% cana={class_percent(pred_sem, 1, ignore_id):.1f}%"))
            panels.append(("Overlay semantic", overlay_rgb(preview_rgb, pred_sem_rgb, args.alpha), ""))

        if pred_veg is not None:
            pred_veg_rgb = ids_to_rgb(pred_veg, BINARY_COLORS_RGB, ignore_id)
            if gt_masks.get("vegetation") is not None:
                panels.append(("GT vegetation", ids_to_rgb(gt_masks["vegetation"], BINARY_COLORS_RGB, ignore_id), "0=fundo 1=veg"))
            panels.append(("Pred vegetation", pred_veg_rgb, f"veg={class_percent(pred_veg, 1, ignore_id):.1f}%"))
            if prob_veg is not None:
                panels.append(("P vegetation", prob_to_heat_rgb(prob_veg), f"mean={float(prob_veg.mean()):.3f}"))

        if pred_cana is not None:
            pred_cana_rgb = ids_to_rgb(pred_cana, CANA_COLORS_RGB, ignore_id)
            if gt_masks.get("cana") is not None:
                panels.append(("GT cana", ids_to_rgb(gt_masks["cana"], CANA_COLORS_RGB, ignore_id), "0=not_cana 1=cana"))
            panels.append(("Pred cana", pred_cana_rgb, f"cana={class_percent(pred_cana, 1, ignore_id):.1f}%"))
            if prob_cana is not None:
                panels.append(("P cana", prob_to_heat_rgb(prob_cana), f"mean={float(prob_cana.mean()):.3f}"))

        if detailed:
            if gt_target is not None:
                panels.append(("GT target", ids_to_rgb(gt_target, TARGET_COLORS_RGB, ignore_id), "derivado"))

            if pred_target_head is not None:
                panels.append(("Pred target HEAD", ids_to_rgb(pred_target_head, TARGET_COLORS_RGB, ignore_id), f"alvo={class_percent(pred_target_head, 1, ignore_id):.1f}%"))
                if prob_target_head is not None:
                    panels.append(("P target HEAD", prob_to_heat_rgb(prob_target_head), f"mean={float(prob_target_head.mean()):.3f}"))

            if pred_target_op is not None:
                panels.append(("Pred target OP", ids_to_rgb(pred_target_op, TARGET_COLORS_RGB, ignore_id), f"alvo={class_percent(pred_target_op, 1, ignore_id):.1f}%"))
                if prob_target_op is not None:
                    panels.append(("P target OP", prob_to_heat_rgb(prob_target_op), f"mean={float(prob_target_op.mean()):.3f}"))
        else:
            if pred_target_head is not None:
                rgb = ids_to_rgb(pred_target_head, TARGET_COLORS_RGB, ignore_id)
                panels.append(("Pred target HEAD", rgb, f"alvo={class_percent(pred_target_head, 1, ignore_id):.1f}%"))
                panels.append(("Overlay target HEAD", overlay_rgb(preview_rgb, rgb, args.alpha), "head direta"))
            elif pred_target_op is not None:
                rgb = ids_to_rgb(pred_target_op, TARGET_COLORS_RGB, ignore_id)
                panels.append(("Pred target OP", rgb, f"alvo={class_percent(pred_target_op, 1, ignore_id):.1f}%"))
                panels.append(("Overlay target OP", overlay_rgb(preview_rgb, rgb, args.alpha), "veg & !cana"))

        canvas = compose_grid(panels, cols=3, max_width=args.max_width)

        header_h = 78
        header = np.zeros((header_h, canvas.shape[1], 3), dtype=np.uint8)
        header[:] = (25, 25, 25)
        source_name = sample.tensor_path.name if sample.tensor_path is not None else f"{sample.base}.json [{sample.source_kind}]"
        h1 = f"idx {idx + 1}/{n} | {source_name} | inf={t_inf:.1f}ms | {'detalhado' if detailed else 'compacto'}"
        h2 = " | ".join(metric_lines[:3]) if metric_lines else "sem GT"
        cv2.putText(header, h1, (12, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (235, 235, 235), 1, cv2.LINE_AA)
        cv2.putText(header, h2, (12, 58), cv2.FONT_HERSHEY_SIMPLEX, 0.52, (0, 255, 120), 1, cv2.LINE_AA)
        canvas = np.vstack([header, canvas])

        onnx_canvas = None
        if onnx_tester is not None and onnx_preds is not None and onnx_probs is not None:
            onnx_pred_sem = onnx_preds.get("semantic")
            onnx_pred_veg = onnx_preds.get("vegetation")
            onnx_pred_cana = onnx_preds.get("cana")

            onnx_pred_target_op = None
            if onnx_pred_veg is not None and onnx_pred_cana is not None:
                onnx_pred_target_op = operational_target_mask(
                    onnx_pred_veg,
                    onnx_pred_cana,
                    ignore_id=ignore_id,
                )

            onnx_pred_target_head = onnx_preds.get("target")
            onnx_pred_target = onnx_pred_target_head if onnx_pred_target_head is not None else onnx_pred_target_op

            onnx_prob_veg = onnx_probs["vegetation"][1] if "vegetation" in onnx_probs and onnx_probs["vegetation"].shape[0] > 1 else None
            onnx_prob_cana = onnx_probs["cana"][1] if "cana" in onnx_probs and onnx_probs["cana"].shape[0] > 1 else None

            onnx_prob_target_op = None
            if onnx_prob_veg is not None and onnx_prob_cana is not None:
                onnx_prob_target_op = np.clip(onnx_prob_veg * (1.0 - onnx_prob_cana), 0.0, 1.0)

            onnx_prob_target_head = None
            if "target" in onnx_probs:
                onnx_prob_target_head = onnx_probs["target"][1] if onnx_probs["target"].shape[0] > 1 else onnx_probs["target"][0]

            onnx_prob_target = onnx_prob_target_head if onnx_prob_target_head is not None else onnx_prob_target_op

            eq_sem = compare_pred_equal_percent(pred_sem, onnx_pred_sem)
            eq_veg = compare_pred_equal_percent(pred_veg, onnx_pred_veg)
            eq_cana = compare_pred_equal_percent(pred_cana, onnx_pred_cana)
            eq_target = compare_pred_equal_percent(pred_target, onnx_pred_target)

            onnx_panels: List[Tuple[str, np.ndarray, str]] = []

            if onnx_pred_sem is not None:
                onnx_sem_rgb = ids_to_rgb(onnx_pred_sem, semantic_cmap, ignore_id)
                onnx_panels.append((
                    "ONNX semantic",
                    onnx_sem_rgb,
                    "" if eq_sem is None else f"igual PT={eq_sem:.3f}%"
                ))
                onnx_panels.append((
                    "ONNX overlay semantic",
                    overlay_rgb(preview_rgb, onnx_sem_rgb, args.alpha),
                    ""
                ))

                d = diff_mask_rgb(pred_sem, onnx_pred_sem)
                if d is not None:
                    onnx_panels.append(("Diff semantic", d, "vermelho=diferente"))

            if onnx_pred_veg is not None:
                onnx_veg_rgb = ids_to_rgb(onnx_pred_veg, BINARY_COLORS_RGB, ignore_id)
                onnx_panels.append((
                    "ONNX vegetation",
                    onnx_veg_rgb,
                    "" if eq_veg is None else f"igual PT={eq_veg:.3f}%"
                ))

                if onnx_prob_veg is not None:
                    onnx_panels.append((
                        "ONNX P vegetation",
                        prob_to_heat_rgb(onnx_prob_veg),
                        f"mean={float(onnx_prob_veg.mean()):.3f}"
                    ))

            if onnx_pred_cana is not None:
                onnx_cana_rgb = ids_to_rgb(onnx_pred_cana, CANA_COLORS_RGB, ignore_id)
                onnx_panels.append((
                    "ONNX cana",
                    onnx_cana_rgb,
                    "" if eq_cana is None else f"igual PT={eq_cana:.3f}%"
                ))

                if onnx_prob_cana is not None:
                    onnx_panels.append((
                        "ONNX P cana",
                        prob_to_heat_rgb(onnx_prob_cana),
                        f"mean={float(onnx_prob_cana.mean()):.3f}"
                    ))

            if onnx_pred_target is not None:
                onnx_target_rgb = ids_to_rgb(onnx_pred_target, TARGET_COLORS_RGB, ignore_id)
                title = "ONNX target HEAD" if onnx_pred_target_head is not None else "ONNX target OP"

                onnx_panels.append((
                    title,
                    onnx_target_rgb,
                    "" if eq_target is None else f"igual PT={eq_target:.3f}%"
                ))
                onnx_panels.append((
                    "ONNX overlay target",
                    overlay_rgb(preview_rgb, onnx_target_rgb, args.alpha),
                    f"inf={t_onnx:.1f}ms"
                ))

                if onnx_prob_target is not None:
                    onnx_panels.append((
                        "ONNX P target",
                        prob_to_heat_rgb(onnx_prob_target),
                        f"mean={float(onnx_prob_target.mean()):.3f}"
                    ))

                d = diff_mask_rgb(pred_target, onnx_pred_target)
                if d is not None:
                    onnx_panels.append(("Diff target", d, "vermelho=diferente"))

            onnx_canvas = compose_grid(onnx_panels, cols=3, max_width=args.max_width)

            header_h_onnx = 78
            header_onnx = np.zeros((header_h_onnx, onnx_canvas.shape[1], 3), dtype=np.uint8)
            header_onnx[:] = (18, 18, 35)

            h1_onnx = f"ONNX {args.onnx_provider} | {source_name} | inf={t_onnx:.1f}ms"
            h2_parts = []
            if eq_sem is not None:
                h2_parts.append(f"sem={eq_sem:.3f}%")
            if eq_veg is not None:
                h2_parts.append(f"veg={eq_veg:.3f}%")
            if eq_cana is not None:
                h2_parts.append(f"cana={eq_cana:.3f}%")
            if eq_target is not None:
                h2_parts.append(f"target={eq_target:.3f}%")
            h2_onnx = "igual PyTorch: " + " | ".join(h2_parts) if h2_parts else "comparação indisponível"

            cv2.putText(header_onnx, h1_onnx, (12, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (235, 235, 235), 1, cv2.LINE_AA)
            cv2.putText(header_onnx, h2_onnx, (12, 58), cv2.FONT_HERSHEY_SIMPLEX, 0.52, (0, 255, 120), 1, cv2.LINE_AA)

            onnx_canvas = np.vstack([header_onnx, onnx_canvas])

        cv2.imshow(win_name, cv2.cvtColor(canvas, cv2.COLOR_RGB2BGR))
        if onnx_canvas is not None and win_name_onnx is not None:
            cv2.imshow(win_name_onnx, cv2.cvtColor(onnx_canvas, cv2.COLOR_RGB2BGR))
        
        k = cv2.waitKey(0) & 0xFF

        if k in (ord("q"), ord("Q"), 27):
            break
        elif k in (ord("d"), ord("D"), 83):
            idx = (idx + 1) % n
        elif k in (ord("a"), ord("A"), 81):
            idx = (idx - 1 + n) % n
        elif k == ord(" "):
            detailed = not detailed
        elif k in (ord("s"), ord("S")):
            out_path = out_dir / f"multihead_pytorch_{idx:05d}_{sample.base}.png"
            cv2.imwrite(str(out_path), cv2.cvtColor(canvas, cv2.COLOR_RGB2BGR))
            print(f"[SAVE] {out_path}")

            if onnx_canvas is not None:
                out_path_onnx = out_dir / f"multihead_onnx_{args.onnx_provider}_{idx:05d}_{sample.base}.png"
                cv2.imwrite(str(out_path_onnx), cv2.cvtColor(onnx_canvas, cv2.COLOR_RGB2BGR))
                print(f"[SAVE] {out_path_onnx}")

    cv2.destroyAllWindows()

    if visited:
        print("\n========== RESUMO DOS SAMPLES VISITADOS ==========")
        print(f"visitados={len(visited)}/{n}")
        for head_name, cm in cms_total.items():
            if cm.sum() <= 0:
                continue
            iou, miou, acc = metrics_from_cm(cm)
            print(f"\n[{head_name}] acc={acc:.4f} mIoU={miou:.4f}")
            for i, v in enumerate(iou):
                print(f"  IoU {i} = {v:.4f}")

        if cm_target_total.sum() > 0:
            iou, miou, acc = metrics_from_cm(cm_target_total)
            print(f"\n[target operacional] acc={acc:.4f} mIoU={miou:.4f}")
            print(f"  IoU background = {iou[0]:.4f}")
            print(f"  IoU alvo       = {iou[1]:.4f}")


if __name__ == "__main__":
    main()