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

"""
_6_normalize_oak_fcc3_originals.py

Converte dataset/originals/group para tensores finais de treino usando
RawProcessorCore + module_params.json.

Entrada:

dataset/originals/group/<grupo>/
  bins/
    <base>_CAM_A.bin
    <base>_CAM_B.bin
    <base>_CAM_C.bin
  metas/
    <base>.json
  masks/
    <base>.png
  previews/
    <base>.png

Saída:

dataset/<WxH>/group/<grupo>/
  tensors/
    <base>.npy      # float32 CHW conforme config["input_channels"]
  masks/
    <base>.npy      # uint8/uint16 HW com IDs de classe
    <base>.png      # debug visual dos IDs
  metas/
    <base>.json     # meta do tensor normalizado
  previews/
    <base>.png      # preview RGB do tensor final

Também salva:
  backup/<modelo>/<model_name>/<stats_source_tag>/norm_stats.json

Uso:

python _6_normalize_oak_fcc3_originals.py ^
  --src-root dataset/originals/group ^
  --out-root dataset ^
  --module-params calibration/module_params.json ^
  --clear-dst
"""

from __future__ import annotations

import argparse
import csv
import json
import os
import shutil
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import cv2
import numpy as np

from core.raw_processor_core import RawProcessorCore

from helpers import carregar_labelmap_completo, converter_mask_rgb_para_ids, _infer_ignore_id


# ============================================================
# Config
# ============================================================

with open("config.json", "r", encoding="utf-8") as f:
    config = json.load(f)

DEFAULT_RES = tuple(config.get("resolucao", [512, 512]))      # (W, H)
DEFAULT_RAW_SIZE = tuple(config.get("raw_size", [1280, 800])) # (W, H)
DEFAULT_MODULE_PARAMS = config.get("module_params_json")
DEFAULT_DATASET_BASE = "dataset"


# Ordem física produzida pelo RawProcessorCore. Estes cinco canais sempre
# existem antes da seleção/derivação dos canais que entrarão no modelo.
SOURCE_CHANNEL_NAMES = ["R", "G", "B", "RE", "NIR"]

# Canais derivados suportados pelo normalizador. Novos índices devem ser
# implementados em build_model_input_tensor() e adicionados aqui.
DERIVED_CHANNEL_NAMES = ["NDVI", "NDRE"]
SUPPORTED_INPUT_CHANNELS = SOURCE_CHANNEL_NAMES + DERIVED_CHANNEL_NAMES


def get_model_input_channels(cfg: dict) -> List[str]:
    """Resolve e valida, na ordem, os canais finais de entrada do modelo.

    Compatibilidade:
      - config novo: input_channels declara tudo explicitamente;
      - config antigo sem input_channels: usa os primeiros ``channels`` canais
        físicos e, quando use_ndvi=True, acrescenta NDVI.
    """
    configured = cfg.get("input_channels")

    if configured is None:
        raw_count = int(cfg.get("channels", len(SOURCE_CHANNEL_NAMES)))
        raw_count = max(1, min(len(SOURCE_CHANNEL_NAMES), raw_count))
        names = SOURCE_CHANNEL_NAMES[:raw_count]

        if bool(cfg.get("use_ndvi", False)) and "NDVI" not in names:
            names.append("NDVI")
    elif isinstance(configured, str):
        names = [x.strip().upper() for x in configured.split(",") if x.strip()]
    else:
        names = [str(x).strip().upper() for x in configured if str(x).strip()]

    if not names:
        raise RuntimeError("config.input_channels não pode ser vazio.")

    duplicates = sorted({name for name in names if names.count(name) > 1})
    if duplicates:
        raise RuntimeError(f"Canais duplicados em config.input_channels: {duplicates}")

    invalid = [name for name in names if name not in SUPPORTED_INPUT_CHANNELS]
    if invalid:
        raise RuntimeError(
            f"Canais inválidos em config.input_channels: {invalid}. "
            f"Suportados: {SUPPORTED_INPUT_CHANNELS}"
        )

    configured_count = cfg.get("channels")
    if configured_count is not None and int(configured_count) != len(names):
        raise RuntimeError(
            "config.channels incompatível com config.input_channels: "
            f"channels={configured_count}, input_channels={names} ({len(names)} canais)."
        )

    return names


MODEL_INPUT_CHANNELS = get_model_input_channels(config)


_CORE_CACHE: Dict[tuple, RawProcessorCore] = {}

def get_or_create_core(
    raw_size: Tuple[int, int],
    bayer_pattern: str,
    module_params_path: Optional[str],
) -> RawProcessorCore:
    raw_w, raw_h = raw_size
    bayer = str(bayer_pattern or "BGGR").upper()

    key = (
        int(raw_w),
        int(raw_h),
        bayer,
        str(Path(module_params_path).resolve()) if module_params_path else None,
    )

    core = _CORE_CACHE.get(key)
    if core is not None:
        return core

    print(
        "[CORE_CACHE] criando RawProcessorCore "
        f"raw={raw_w}x{raw_h} bayer={bayer} module={module_params_path}"
    )

    core = RawProcessorCore(
        sensor_width=int(raw_w),
        sensor_height=int(raw_h),
        bayer_pattern=bayer,
        calibration_json_path=module_params_path,
    )

    _CORE_CACHE[key] = core
    return core


@dataclass
class SampleBundle:
    group: str
    base: str
    meta_path: Path
    mask_path: Path
    preview_path: Optional[Path]
    bin_paths: Dict[str, Path]


# ============================================================
# Helpers gerais
# ============================================================

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


def maybe_clear_dir(path: Path):
    if path.exists():
        shutil.rmtree(path)
    ensure_dir(path)


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


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


def safe_rel(path: Path, root: Path) -> str:
    try:
        return str(path.resolve().relative_to(root.resolve())).replace("\\", "/")
    except Exception:
        return str(path).replace("\\", "/")


def list_groups(src_root: Path) -> List[str]:
    if not src_root.is_dir():
        return []

    groups = []
    for p in sorted(src_root.iterdir()):
        if not p.is_dir():
            continue

        if (p / "bins").is_dir() and (p / "metas").is_dir() and (p / "masks").is_dir():
            groups.append(p.name)

    return groups


def find_preview(previews_dir: Path, base: str) -> Optional[Path]:
    for ext in (".png", ".jpg", ".jpeg"):
        p = previews_dir / f"{base}{ext}"
        if p.exists():
            return p
    return None


def find_sample_bins(bins_dir: Path, base: str) -> Dict[str, Path]:
    out = {}

    for cam_id in ("CAM_A", "CAM_B", "CAM_C"):
        p = bins_dir / f"{base}_{cam_id}.bin"
        if p.exists():
            out[cam_id] = p

    return out


def collect_samples_from_group(group_dir: Path) -> List[SampleBundle]:
    group = group_dir.name

    metas_dir = group_dir / "metas"
    masks_dir = group_dir / "masks"
    previews_dir = group_dir / "previews"
    bins_dir = group_dir / "bins"

    samples = []

    if not metas_dir.is_dir():
        return samples

    for meta_path in sorted(metas_dir.glob("*.json")):
        base = meta_path.stem

        mask_path = masks_dir / f"{base}.png"
        if not mask_path.exists():
            print(f"[WARN] [{group}] sem mask para {base}. Pulando.")
            continue

        bin_paths = find_sample_bins(bins_dir, base)
        if not all(cam in bin_paths for cam in ("CAM_A", "CAM_B", "CAM_C")):
            print(f"[WARN] [{group}] bins incompletos para {base}: {list(bin_paths.keys())}. Pulando.")
            continue

        samples.append(
            SampleBundle(
                group=group,
                base=base,
                meta_path=meta_path,
                mask_path=mask_path,
                preview_path=find_preview(previews_dir, base),
                bin_paths=bin_paths,
            )
        )

    return samples


def copy_json_safe(obj):
    try:
        return json.loads(json.dumps(obj, default=str))
    except Exception:
        return None


# ============================================================
# Module params / tensor
# ============================================================

def resolve_module_params_path(meta: dict, meta_path: Path, cli_module_params: Optional[str]) -> Optional[str]:
    candidates = []

    if cli_module_params:
        candidates.append(cli_module_params)

    if meta.get("camera_params_json"):
        candidates.append(meta.get("camera_params_json"))

    if DEFAULT_MODULE_PARAMS:
        candidates.append(DEFAULT_MODULE_PARAMS)

    candidates.append("calibration/module_params.json")

    for c in candidates:
        if not c:
            continue

        p = Path(str(c))

        if p.is_file():
            return str(p)

        # relativo ao diretório atual
        p2 = Path.cwd() / p
        if p2.is_file():
            return str(p2)

        # relativo à pasta do meta
        p3 = meta_path.parent / p
        if p3.is_file():
            return str(p3)

    print("[WARN] module_params.json não encontrado. RawProcessorCore vai usar defaults.")
    return None


def load_frame_from_saved_bins(sample: SampleBundle, meta: dict) -> Dict[str, np.ndarray]:
    """
    Monta frame no contrato do RawProcessorCore.decode_stream_cameras:

      frame = {
        "CAM_A": ndarray RAW10 packed 2D,
        "CAM_B": ndarray RAW10 packed 2D,
        "CAM_C": ndarray RAW10 packed 2D,
      }

    Usa saved_payload_dtypes/saved_payload_shapes do JSON.
    """
    saved_dtypes = meta.get("saved_payload_dtypes", {}) or {}
    saved_shapes = meta.get("saved_payload_shapes", {}) or {}

    frame = {}

    for cam_id, bin_path in sample.bin_paths.items():
        dtype = saved_dtypes.get(cam_id)
        shape = saved_shapes.get(cam_id)

        if dtype is None or shape is None:
            raise RuntimeError(f"Faltam saved_payload_dtypes/shapes para {cam_id} em {sample.base}")

        arr = np.fromfile(str(bin_path), dtype=np.dtype(dtype)).reshape(tuple(shape))
        frame[cam_id] = arr

    return frame


def build_processing_meta(meta: dict) -> dict:
    """
    O RawProcessorCore precisa de:
      frame_type = RAW_BRUTO
      camera_info/camera_frames por CAM_A/B/C
      actual_camera_controls para radiometric_normalization
    """
    stream_meta = dict(meta.get("stream_meta", {}) or {})

    # Garante o frame_type esperado pelo core.
    stream_meta["frame_type"] = "RAW_BRUTO"

    # Em algumas capturas o camera_info está fora do stream_meta.
    if "camera_info" not in stream_meta and isinstance(meta.get("camera_info"), dict):
        stream_meta["camera_info"] = meta.get("camera_info")

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

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

    if meta.get("camera_controls") is not None:
        stream_meta["camera_controls"] = meta.get("camera_controls")

    return stream_meta


def build_tensor_from_sample(
    sample: SampleBundle,
    meta: dict,
    res: Tuple[int, int],
    raw_size: Tuple[int, int],
    module_params_path: Optional[str],
) -> Tuple[np.ndarray, np.ndarray, dict]:
    bayer = str(meta.get("bayer_pattern") or meta.get("bayer") or "BGGR").upper()

    core = get_or_create_core(
        raw_size=raw_size,
        bayer_pattern=bayer,
        module_params_path=module_params_path,
    )

    frame = load_frame_from_saved_bins(sample, meta)
    processing_meta = build_processing_meta(meta)

    source_tensor = core.build_infer_tensor_from_stream(
        frame=frame,
        meta=processing_meta,
        channels_expected=5,
        target_size=res,
    )

    source_tensor = np.ascontiguousarray(source_tensor.astype(np.float32, copy=False))

    if source_tensor.ndim != 3 or source_tensor.shape[0] != len(SOURCE_CHANNEL_NAMES):
        raise RuntimeError(
            "RawProcessorCore retornou tensor incompatível: "
            f"esperado CHW com {len(SOURCE_CHANNEL_NAMES)} canais "
            f"{SOURCE_CHANNEL_NAMES}, veio shape={source_tensor.shape}."
        )

    tensor = build_model_input_tensor(
        source_tensor=source_tensor,
        requested_channels=MODEL_INPUT_CHANNELS,
        derived_config=config.get("derived_channels", {}) or {},
    )

    info = {
        "module_params": module_params_path,
        "bayer_pattern": bayer,
        "fusion_result": copy_json_safe(getattr(core, "last_fusion_result", None)),
        "radiometric_normalization_result": copy_json_safe(getattr(core, "last_radiometric_normalization_result", None)),
        "patch_normalization_result": copy_json_safe(getattr(core, "last_patch_normalization_result", None)),
        "frame_quality": copy_json_safe(getattr(core, "last_frame_quality_result", None)),
    }

    # source_tensor é devolvido apenas para gerar o preview RGB. O .npy salvo
    # e as estatísticas pertencem sempre ao tensor final de entrada do modelo.
    return tensor, source_tensor, info


def build_model_input_tensor(
    source_tensor: np.ndarray,
    requested_channels: List[str],
    derived_config: Optional[dict] = None,
) -> np.ndarray:
    """Monta o CHW final combinando bandas físicas e índices espectrais.

    Os índices são calculados antes da normalização estatística do dataset:
      NDVI = (NIR - R)  / (NIR + R)
      NDRE = (NIR - RE) / (NIR + RE)

    Denominadores próximos de zero resultam em 0.0. O resultado é limitado
    por padrão a [-1, 1] e posteriormente receberá mean/std próprios.
    """
    derived_config = derived_config or {}
    eps = float(derived_config.get("epsilon", 1e-6))
    clip_min = float(derived_config.get("clip_min", -1.0))
    clip_max = float(derived_config.get("clip_max", 1.0))

    if not np.isfinite(eps) or eps <= 0.0:
        raise RuntimeError(f"derived_channels.epsilon inválido: {eps}")
    if not np.isfinite(clip_min) or not np.isfinite(clip_max) or clip_min >= clip_max:
        raise RuntimeError(
            "Faixa inválida para derived_channels: "
            f"clip_min={clip_min}, clip_max={clip_max}"
        )

    source = np.nan_to_num(
        np.asarray(source_tensor, dtype=np.float32),
        nan=0.0,
        posinf=1.0,
        neginf=0.0,
    )
    available = {
        name: source[index]
        for index, name in enumerate(SOURCE_CHANNEL_NAMES)
    }

    r = available["R"]
    re = available["RE"]
    nir = available["NIR"]

    def normalized_difference(a: np.ndarray, b: np.ndarray) -> np.ndarray:
        denominator = a + b
        result = np.zeros_like(a, dtype=np.float32)
        np.divide(
            a - b,
            denominator,
            out=result,
            where=np.abs(denominator) > eps,
        )
        return np.clip(result, clip_min, clip_max).astype(np.float32, copy=False)

    # Calcula somente o que foi solicitado para evitar trabalho desnecessário.
    if "NDVI" in requested_channels:
        available["NDVI"] = normalized_difference(nir, r)
    if "NDRE" in requested_channels:
        available["NDRE"] = normalized_difference(nir, re)

    tensor = np.stack(
        [available[name] for name in requested_channels],
        axis=0,
    )
    return np.ascontiguousarray(tensor.astype(np.float32, copy=False))



# ============================================================
# Máscara / preview
# ============================================================

def load_mask_ids_aligned_to_tensor(
    mask_path: Path,
    cor_para_id: dict,
    ignore_id: int,
    res: Tuple[int, int],
    fusion_result: dict | None,
) -> np.ndarray:
    bgr = cv2.imread(str(mask_path), cv2.IMREAD_COLOR)

    if bgr is None:
        raise RuntimeError(f"Falha ao abrir mask: {mask_path}")

    rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
    ids = converter_mask_rgb_para_ids(rgb, cor_para_id, ignore_id)

    # A mask foi feita no espaço RGB/CAM_A.
    # Então primeiro garantimos que ela está no mesmo tamanho da referência RGB.
    if isinstance(fusion_result, dict):
        ref_shape = fusion_result.get("ref_shape")

        if isinstance(ref_shape, list) and len(ref_shape) == 2:
            ref_h, ref_w = int(ref_shape[0]), int(ref_shape[1])

            if ids.shape[:2] != (ref_h, ref_w):
                ids = cv2.resize(
                    ids,
                    (ref_w, ref_h),
                    interpolation=cv2.INTER_NEAREST,
                )

        crop_box = fusion_result.get("crop_box")

        if crop_box is not None:
            x0, y0, x1, y1 = [int(v) for v in crop_box]

            h, w = ids.shape[:2]
            x0 = max(0, min(w - 1, x0))
            x1 = max(x0 + 1, min(w, x1))
            y0 = max(0, min(h - 1, y0))
            y1 = max(y0 + 1, min(h, y1))

            ids = ids[y0:y1, x0:x1]

    # Por fim, leva para a resolução final do tensor.
    ids_res = cv2.resize(
        ids,
        res,
        interpolation=cv2.INTER_NEAREST,
    )

    return ids_res


def tensor_to_preview_bgr(source_tensor: np.ndarray) -> np.ndarray:
    """Gera preview a partir do RGB físico, independente dos canais do modelo."""
    rgb = np.transpose(source_tensor[:3].astype(np.float32), (1, 2, 0))
    rgb_u8 = np.clip(rgb * 255.0, 0, 255).astype(np.uint8)
    return cv2.cvtColor(rgb_u8, cv2.COLOR_RGB2BGR)


def save_bgr(path: Path, bgr: np.ndarray):
    ensure_dir(path.parent)
    ok = cv2.imwrite(str(path), bgr)
    if not ok:
        raise RuntimeError(f"Falha ao salvar preview: {path}")


def get_class_ids_from_config() -> dict:
    heads = config.get("heads", {}) or {}
    semantic = heads.get("semantic", {}) or {}
    classes = semantic.get("classes", {}) or {}

    return {
        "chao": int(classes.get("chao", 0)),
        "cana": int(classes.get("cana", 1)),
        "erva": int(classes.get("erva", 2)),
    }


def build_head_masks_from_semantic(
    mask_ids: np.ndarray,
    class_ids: dict,
    ignore_id: int = 255,
) -> dict:
    """
    Gera máscaras auxiliares para multi-head a partir da máscara semântica alinhada.

    Entrada:
      mask_ids: HW com IDs semânticos.

    Saída:
      {
        "semantic": HW uint8,
        "vegetation": HW uint8 com 0/1/ignore,
        "cana": HW uint8 com 0/1/ignore
      }
    """
    chao_id = int(class_ids.get("chao", 0))
    cana_id = int(class_ids.get("cana", 1))
    erva_id = int(class_ids.get("erva", 2))

    semantic = mask_ids.astype(np.uint8, copy=True)

    ignore_mask = semantic == ignore_id

    vegetation = np.zeros_like(semantic, dtype=np.uint8)
    vegetation[(semantic == cana_id) | (semantic == erva_id)] = 1
    vegetation[ignore_mask] = ignore_id

    cana = np.zeros_like(semantic, dtype=np.uint8)
    cana[semantic == cana_id] = 1
    cana[ignore_mask] = ignore_id

    return {
        "semantic": semantic,
        "vegetation": vegetation,
        "cana": cana,
    }


def save_mask_npy_and_debug(mask: np.ndarray, npy_path: Path, png_path: Path):
    ensure_dir(npy_path.parent)
    ensure_dir(png_path.parent)

    mask = mask.astype(np.uint8, copy=False)
    np.save(str(npy_path), mask)

    # Debug visual:
    # 0 -> preto
    # 1 -> cinza claro
    # 2 -> mais claro, quando existir na semântica
    # 255 -> branco
    debug = mask.copy()
    debug_vis = np.zeros_like(debug, dtype=np.uint8)

    debug_vis[debug == 0] = 0
    debug_vis[debug == 1] = 120
    debug_vis[debug == 2] = 220
    debug_vis[debug == 255] = 255

    cv2.imwrite(str(png_path), debug_vis)


# ============================================================
# Stats
# ============================================================

class RunningStats:
    def __init__(self):
        self.sum = None
        self.sumsq = None
        self.pixels = 0

    def update(self, tensor: np.ndarray):
        c, h, w = tensor.shape

        if self.sum is None:
            self.sum = np.zeros(c, dtype=np.float64)
            self.sumsq = np.zeros(c, dtype=np.float64)
        elif len(self.sum) != c:
            raise RuntimeError(
                f"Quantidade de canais mudou durante o normalize: {len(self.sum)} -> {c}"
            )

        flat = tensor.reshape(c, -1).astype(np.float64)
        self.sum += flat.sum(axis=1)
        self.sumsq += (flat ** 2).sum(axis=1)
        self.pixels += h * w

    def result(self, channel_names: List[str]):
        if self.sum is None or self.pixels <= 0:
            return None

        mean = self.sum / self.pixels
        var = (self.sumsq / self.pixels) - mean ** 2
        std = np.sqrt(np.maximum(var, 1e-6))

        return {
            "channels": channel_names,
            "mean": mean.tolist(),
            "std": std.tolist(),
            "pixels": int(self.pixels),
        }


# ============================================================
# Processamento
# ============================================================

def process_group(
    group_name: str,
    src_root: Path,
    output_root: Path,
    res: Tuple[int, int],
    raw_size: Tuple[int, int],
    module_params_arg: Optional[str],
    cor_para_id: dict,
    ignore_id: int,
    stats: RunningStats,
    dataset_root: Path,
    skip_bad_quality: bool,
) -> List[dict]:
    group_dir = src_root / group_name
    samples = collect_samples_from_group(group_dir)

    out_group = output_root / group_name
    out_tensors = out_group / "tensors"

    out_masks = out_group / "masks"
    out_masks_vegetation = out_group / "masks_vegetation"
    out_masks_cana = out_group / "masks_cana"

    out_metas = out_group / "metas"
    out_previews = out_group / "previews"

    for d in (
        out_tensors,
        out_masks,
        out_masks_vegetation,
        out_masks_cana,
        out_metas,
        out_previews,
    ):
        ensure_dir(d)

    rows = []
    errors = 0
    skipped_quality = 0

    print(f"\n[GRUPO] {group_name} | samples={len(samples)}")

    for sample in samples:
        try:
            meta = load_json(sample.meta_path)
            module_params_path = resolve_module_params_path(meta, sample.meta_path, module_params_arg)

            tensor, source_tensor, processing_info = build_tensor_from_sample(
                sample=sample,
                meta=meta,
                res=res,
                raw_size=raw_size,
                module_params_path=module_params_path,
            )

            frame_quality = processing_info.get("frame_quality") or {}
            if skip_bad_quality and isinstance(frame_quality, dict):
                if frame_quality.get("usable_for_training") is False:
                    skipped_quality += 1
                    print(f"[SKIP-QUALITY] {sample.base}: {frame_quality.get('reasons')}")
                    continue

            mask_ids = load_mask_ids_aligned_to_tensor(
                mask_path=sample.mask_path,
                cor_para_id=cor_para_id,
                ignore_id=ignore_id,
                res=res,
                fusion_result=processing_info.get("fusion_result"),
            )

            if mask_ids.shape[:2] != tensor.shape[1:]:
                raise RuntimeError(
                    f"Shape mask/tensor incompatível: mask={mask_ids.shape}, tensor={tensor.shape}"
                )

            # Salva tensor e mask
            tensor_path = out_tensors / f"{sample.base}.npy"
            preview_path = out_previews / f"{sample.base}.png"
            meta_out_path = out_metas / f"{sample.base}.json"
            np.save(str(tensor_path), tensor)

            class_ids = get_class_ids_from_config()
            head_masks = build_head_masks_from_semantic(
                mask_ids=mask_ids,
                class_ids=class_ids,
                ignore_id=ignore_id,
            )

            mask_path = out_masks / f"{sample.base}.npy"
            mask_debug_path = out_masks / f"{sample.base}.png"

            mask_vegetation_path = out_masks_vegetation / f"{sample.base}.npy"
            mask_vegetation_debug_path = out_masks_vegetation / f"{sample.base}.png"

            mask_cana_path = out_masks_cana / f"{sample.base}.npy"
            mask_cana_debug_path = out_masks_cana / f"{sample.base}.png"

            save_mask_npy_and_debug(
                head_masks["semantic"],
                mask_path,
                mask_debug_path,
            )

            save_mask_npy_and_debug(
                head_masks["vegetation"],
                mask_vegetation_path,
                mask_vegetation_debug_path,
            )

            save_mask_npy_and_debug(
                head_masks["cana"],
                mask_cana_path,
                mask_cana_debug_path,
            )

            preview_bgr = tensor_to_preview_bgr(source_tensor)
            save_bgr(preview_path, preview_bgr)

            out_meta = {
                "schema": "oak_fcc3_normalized_tensor_v2",
                "source_group": sample.group,
                "source_base": sample.base,
                "source_meta": safe_rel(sample.meta_path, dataset_root),
                "source_bins": {k: safe_rel(v, dataset_root) for k, v in sample.bin_paths.items()},
                "source_mask": safe_rel(sample.mask_path, dataset_root),
                "source_preview": safe_rel(sample.preview_path, dataset_root) if sample.preview_path else None,

                "frame_type": "MULTISPEC",
                "saved_payload_type": "tensor_npy",
                "saved_tensor_path": safe_rel(tensor_path, dataset_root),
                "saved_mask_path": safe_rel(mask_path, dataset_root),
                "head_masks": {
                    "semantic": {
                        "path": safe_rel(mask_path, dataset_root),
                        "shape": list(head_masks["semantic"].shape),
                        "classes": {
                            "chao": int(class_ids["chao"]),
                            "cana": int(class_ids["cana"]),
                            "erva": int(class_ids["erva"]),
                        },
                        "ignore_index": int(ignore_id),
                    },
                    "vegetation": {
                        "path": safe_rel(mask_vegetation_path, dataset_root),
                        "shape": list(head_masks["vegetation"].shape),
                        "classes": {
                            "background": 0,
                            "vegetation": 1,
                        },
                        "positive_from_semantic": ["cana", "erva"],
                        "ignore_index": int(ignore_id),
                    },
                    "cana": {
                        "path": safe_rel(mask_cana_path, dataset_root),
                        "shape": list(head_masks["cana"].shape),
                        "classes": {
                            "not_cana": 0,
                            "cana": 1,
                        },
                        "positive_from_semantic": ["cana"],
                        "ignore_index": int(ignore_id),
                    },
                },
                "saved_preview_path": safe_rel(preview_path, dataset_root),
                "saved_payload_dtype": str(tensor.dtype),
                "saved_payload_shape": list(tensor.shape),
                "mask_shape": list(mask_ids.shape),
                "source_channels": SOURCE_CHANNEL_NAMES,
                "channels": MODEL_INPUT_CHANNELS,
                "derived_channels": {
                    "enabled": [
                        name for name in MODEL_INPUT_CHANNELS
                        if name in DERIVED_CHANNEL_NAMES
                    ],
                    "config": copy_json_safe(config.get("derived_channels", {}) or {}),
                },

                "resolution": {
                    "width": int(res[0]),
                    "height": int(res[1]),
                },
                "processing": processing_info,
                "camera_params_json": module_params_path,
                "source_capture_meta": {
                    "ts": meta.get("ts"),
                    "sensor_width": meta.get("sensor_width"),
                    "sensor_height": meta.get("sensor_height"),
                    "bayer_pattern": meta.get("bayer_pattern"),
                    "actual_camera_controls": meta.get("actual_camera_controls"),
                    "startup_camera_controls": meta.get("startup_camera_controls"),
                    "radiometric_last_result": meta.get("radiometric_last_result"),
                    "stream_meta": meta.get("stream_meta"),
                },
            }

            save_json(meta_out_path, out_meta)

            stats.update(tensor)

            rows.append({
                "group": sample.group,
                "base": sample.base,
                "tensor": str(tensor_path),
                "mask": str(mask_path),
                "mask_vegetation": str(mask_vegetation_path),
                "mask_cana": str(mask_cana_path),
                "meta": str(meta_out_path),
                "preview": str(preview_path),
                "quality_status": frame_quality.get("status") if isinstance(frame_quality, dict) else None,
            })

            print(f"[OK] {group_name}/{sample.base} tensor={list(tensor.shape)}")

        except Exception as e:
            errors += 1
            print(f"[ERRO] {group_name}/{sample.base}: {e}")

    print(f"[RESUMO] {group_name}: ok={len(rows)} | erros={errors} | skip_quality={skipped_quality}")
    return rows


def write_manifest(path: Path, rows: List[dict]):
    ensure_dir(path.parent)

    fieldnames = [
        "group",
        "base",
        "tensor",
        "mask",
        "mask_vegetation",
        "mask_cana",
        "meta",
        "preview",
        "quality_status",
    ]

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


def write_norm_stats(stats_path: Path, stats: dict):
    ensure_dir(stats_path.parent)

    with stats_path.open("w", encoding="utf-8") as f:
        json.dump(stats, f, ensure_ascii=False, indent=2)


def main():
    ap = argparse.ArgumentParser(
        description="Normaliza RAW_BRUTO OAK-FCC-3 para tensor MULTISPEC final de treino."
    )

    ap.add_argument(
        "--src-roots",
        default="dataset/original/group;dataset/augmented/group",
        help="Raízes de entrada separadas por ';'. Ex: dataset/original/group;dataset/augmented/group"
    )
    ap.add_argument("--out-root", default="dataset")
    ap.add_argument("--module-params", default=DEFAULT_MODULE_PARAMS)
    ap.add_argument("--res", default=f"{DEFAULT_RES[0]}x{DEFAULT_RES[1]}", help="Resolução final WxH.")
    ap.add_argument("--raw-size", default=f"{DEFAULT_RAW_SIZE[0]}x{DEFAULT_RAW_SIZE[1]}", help="Tamanho RAW real WxH.")
    ap.add_argument("--groups", default=None, help="Grupos separados por vírgula.")
    ap.add_argument("--clear-dst", action="store_true")
    ap.add_argument("--skip-bad-quality", action="store_true")
    ap.add_argument("--manifest", default="")
    ap.add_argument("--stats-out", default="")

    args = ap.parse_args()

    src_roots = [
        Path(x.strip())
        for x in str(args.src_roots).split(";")
        if x.strip()
    ]

    out_dataset_root = Path(args.out_root)

    valid_src_roots = []
    for src_root in src_roots:
        if src_root.is_dir():
            valid_src_roots.append(src_root)
        else:
            print(f"[WARN] src-root não encontrado, ignorando: {src_root}")

    if not valid_src_roots:
        raise SystemExit(f"[ERRO] Nenhum src-root válido encontrado: {src_roots}")

    try:
        res_w, res_h = [int(x) for x in args.res.lower().split("x")]
        raw_w, raw_h = [int(x) for x in args.raw_size.lower().split("x")]
    except Exception:
        raise SystemExit("[ERRO] Use --res WxH e --raw-size WxH. Ex: --res 512x512 --raw-size 1280x800")

    res = (res_w, res_h)
    raw_size = (raw_w, raw_h)

    output_root = out_dataset_root / f"{res_w}x{res_h}" / "group"

    if args.clear_dst:
        print(f"[INFO] Limpando destino: {output_root}")
        maybe_clear_dir(output_root)

    labelmap_path = out_dataset_root / "labelmap.txt"

    if not labelmap_path.exists():
        raise SystemExit(f"[ERRO] labelmap não encontrado: {labelmap_path}")

    cor_para_id, _, _, ignore_rgb = carregar_labelmap_completo(str(labelmap_path))
    ignore_id = _infer_ignore_id(ignore_rgb, 255)

    groups_by_root = []

    for src_root in valid_src_roots:
        groups = list_groups(src_root)

        if args.groups:
            want = {g.strip() for g in args.groups.split(",") if g.strip()}
            groups = [g for g in groups if g in want]

        if groups:
            groups_by_root.append((src_root, groups))

    if not groups_by_root:
        raise SystemExit("[ERRO] Nenhum grupo encontrado.")

    print("============================================")
    print("Normalize OAK-FCC-3")
    print(f"SRC          : {[str(x) for x in valid_src_roots]}")
    print(f"OUT          : {output_root}")
    print(f"MODULE PARAM : {args.module_params}")
    print(f"RES          : {res}")
    print(f"RAW SIZE     : {raw_size}")
    print(f"SOURCE CH    : {SOURCE_CHANNEL_NAMES}")
    print(f"MODEL CH     : {MODEL_INPUT_CHANNELS} ({len(MODEL_INPUT_CHANNELS)})")
    print("GROUPS       :")
    for root, groups in groups_by_root:
        print(f"  - {root}: {', '.join(groups)}")
    print(f"SKIP BAD     : {args.skip_bad_quality}")
    print("============================================")

    running_stats = RunningStats()
    all_rows = []

    for src_root, all_groups in groups_by_root:
        for group_name in all_groups:
            rows = process_group(
                group_name=group_name,
                src_root=src_root,
                output_root=output_root,
                res=res,
                raw_size=raw_size,
                module_params_arg=args.module_params,
                cor_para_id=cor_para_id,
                ignore_id=ignore_id,
                stats=running_stats,
                dataset_root=out_dataset_root,
                skip_bad_quality=args.skip_bad_quality,
            )
            all_rows.extend(rows)

    manifest_path = Path(args.manifest) if args.manifest else output_root / "normalize_manifest.csv"
    write_manifest(manifest_path, all_rows)

    stats = running_stats.result(MODEL_INPUT_CHANNELS)

    if stats is not None:
        stats_path = (
            Path(args.stats_out)
            if args.stats_out
            else Path("backup") / config.get("modelo", "modelo") / config.get("model_name", "model") / config.get("stats_source_tag", "oak_fcc3") / "norm_stats.json"
        )

        write_norm_stats(stats_path, stats)

        # Também salva uma cópia junto do dataset normalizado.
        write_norm_stats(output_root / "norm_stats.json", stats)

        print("\n📊 STATS:")
        print(json.dumps(stats, ensure_ascii=False, indent=2))
        print(f"[OK] norm_stats backup : {stats_path}")
        print(f"[OK] norm_stats dataset: {output_root / 'norm_stats.json'}")

    print("\n============================================")
    print("Normalize finalizado")
    print(f"Total samples: {len(all_rows)}")
    print(f"Manifest     : {manifest_path}")
    print("============================================")


if __name__ == "__main__":
    main()