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

"""
_5_augmentation_raw_oak.py

Augmentation multiespectral para OAK-FCC-3 usando os .bin RAW das câmeras,
máscaras e metadados do dataset.

Objetivo:
- Ler amostras em dataset/original/group/<grupo>/
- Carregar RGB/RE/NIR a partir dos .bin RAW10 packed
- Aplicar augmentation geométrico sincronizado em todos os canais + máscara
- Aplicar augmentation radiométrico/físico nos RAWs
- Salvar nova amostra em dataset/augmented/group/<grupo>/
- Gerar preview RGB novo para inspeção visual

Estrutura esperada de entrada:

dataset/original/group/<grupo>/
  bins/
  masks/
  metas/
  previews/          # opcional

Estrutura de saída:

dataset/augmented/group/<grupo>/
  bins/
  masks/
  metas/
  previews/

Exemplos:

python _5_augmentation_raw_oak.py --copies 3 --clear-dst

python _5_augmentation_raw_oak.py ^
  --copies 3 ^
  --groups chao,chao_cana,chao_erva,chao_cana_erva ^
  --clear-dst

python _5_augmentation_raw_oak.py ^
  --group-copies chao:1,chao_cana:3,chao_erva:3,chao_cana_erva:5 ^
  --clear-dst

python _5_augmentation_raw_oak.py --dry-run

Observações importantes:
- Este script NÃO augmenta preview PNG como fonte principal.
- O preview é somente consequência visual do RGB RAW augmentado.
- A geometria é sempre igual em RGB/RE/NIR/mask.
- Radiometria usa mesma base global + pequenas variações por banda/câmera.
- O script tenta ser tolerante a nomes de arquivo, mas espera que base da amostra
  esteja preservada nos nomes dos bins/masks/metas.
"""

import argparse
import copy
import json
import math
import os
import random
import shutil
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Any

import cv2
import numpy as np
from PIL import Image

try:
    from core.raw_processor_core import RawProcessorCore
    _HAS_RAW_PROCESSOR_CORE = True
except Exception:
    RawProcessorCore = None
    _HAS_RAW_PROCESSOR_CORE = False

_CORE_PREVIEW_CACHE = {}


# ============================================================
# Configuração base
# ============================================================

IMG_EXTS = (".png", ".jpg", ".jpeg")
META_EXTS = (".json",)
BIN_EXTS = (".bin",)

ROLE_ALIASES = {
    "rgb": ["rgb", "color", "cor", "cam_a", "cama", "CAM_A"],
    "re": ["re", "rededge", "red_edge", "red-edge", "cam_b", "camb", "CAM_B"],
    "nir": ["nir", "ir", "infra", "infrared", "cam_c", "camc", "CAM_C"],
}

CAM_KEYS = ["CAM_A", "CAM_B", "CAM_C", "cam_a", "cam_b", "cam_c", "cam0", "cam1", "cam2"]

DEFAULT_BIT_DEPTH = 10
DEFAULT_BAYER_PATTERN = "BGGR"
DEFAULT_RAW_FORMAT = "RAW10_PACKED"


# ============================================================
# Data classes
# ============================================================

@dataclass
class CameraBin:
    role: str
    cam_key: str
    path: Path
    width: int
    height: int
    bit_depth: int = DEFAULT_BIT_DEPTH
    raw_format: str = DEFAULT_RAW_FORMAT
    bayer_pattern: str = DEFAULT_BAYER_PATTERN
    data: Optional[np.ndarray] = None


@dataclass
class AugParams:
    seed: int

    # Geometria, igual para tudo
    flip_h: bool
    shift_x_frac: float
    shift_y_frac: float
    scale: float
    rotate_deg: float
    perspective: bool
    perspective_strength: float

    # Radiometria base
    exposure_mult: float
    gamma: float

    # Ganhos por papel espectral
    rgb_channel_gain: Tuple[float, float, float]
    re_gain: float
    nir_gain: float

    # Efeitos físicos leves
    shadow_enabled: bool
    shadow_strength: float
    shadow_angle_deg: float
    highlight_enabled: bool
    highlight_strength: float
    noise_enabled: bool
    noise_sigma_dn: float
    blur_enabled: bool
    blur_kernel: int


# ============================================================
# Utilitários gerais
# ============================================================

def read_config_model() -> str:
    """Tenta ler config.json para manter compatibilidade com seus scripts atuais."""
    cfg = Path("config.json")
    if not cfg.exists():
        return "."

    try:
        with cfg.open("r", encoding="utf-8") as f:
            config = json.load(f)
        return str(config.get("camera", "."))
    except Exception:
        return "."


def read_config_value(key: str, default=None):
    cfg = Path("config.json")
    if not cfg.exists():
        return default

    try:
        with cfg.open("r", encoding="utf-8") as f:
            config = json.load(f)
        return config.get(key, default)
    except Exception:
        return default


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


def clear_dir(path: Path) -> None:
    if path.exists():
        shutil.rmtree(path)
    ensure_dir(path)


def normalizar_base(stem: str) -> str:
    """Remove sufixos comuns para casar preview/mask/meta/bin."""
    suffixes = [
        "_rgb", "_RGB", "_Rgb", "_color", "_COLOR",
        "_re", "_RE", "_rededge", "_red_edge", "_nir", "_NIR",
        "_cam_a", "_cam_b", "_cam_c", "_CAM_A", "_CAM_B", "_CAM_C",
        "_cam0", "_cam1", "_cam2", "_CAM0", "_CAM1", "_CAM2",
        "_image", "_img", "_frame", "_preview", "_previews",
        "_mask", "_masks", "_seg", "_SEG", "_segment", "_segmentacao",
        "_meta", "_metadata",
    ]

    out = stem
    changed = True
    while changed:
        changed = False
        for sfx in suffixes:
            if out.endswith(sfx):
                out = out[: -len(sfx)]
                changed = True
                break
    return out


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

    groups = []
    for p in sorted(src_root.iterdir()):
        if not p.is_dir():
            continue
        if (p / "bins").is_dir() and (p / "masks").is_dir() and (p / "metas").is_dir():
            groups.append(p.name)
    return groups


def map_files_by_base(folder: Path, exts: Tuple[str, ...]) -> Dict[str, Path]:
    by_base: Dict[str, Path] = {}
    if not folder.is_dir():
        return by_base

    priority = {
        ".json": 0,
        ".png": 0,
        ".bin": 0,
        ".jpg": 1,
        ".jpeg": 2,
    }

    for p in sorted(folder.iterdir()):
        if not p.is_file():
            continue
        ext = p.suffix.lower()
        if ext not in exts:
            continue

        base = normalizar_base(p.stem)
        if base not in by_base:
            by_base[base] = p
        else:
            cur_ext = by_base[base].suffix.lower()
            if priority.get(ext, 99) < priority.get(cur_ext, 99):
                by_base[base] = p

    return by_base


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


def save_json(path: Path, data: Dict[str, Any]) -> None:
    ensure_dir(path.parent)
    with path.open("w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


def load_mask(path: Path) -> np.ndarray:
    im = cv2.imread(str(path), cv2.IMREAD_UNCHANGED)
    if im is None:
        raise FileNotFoundError(path)

    if im.ndim == 2:
        return im

    if im.shape[2] == 4:
        im = cv2.cvtColor(im, cv2.COLOR_BGRA2RGBA)
    else:
        im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
    return im


def save_mask(path: Path, mask: np.ndarray) -> None:
    ensure_dir(path.parent)
    if mask.ndim == 2:
        Image.fromarray(mask).save(path)
    else:
        Image.fromarray(mask).save(path)


# ============================================================
# RAW10 pack/unpack
# ============================================================

def unpack_raw10_packed(buf: bytes, width: int, height: int) -> np.ndarray:
    """
    Desempacota MIPI RAW10 packed no padrão mais comum usado por câmeras/DepthAI:

    Para cada grupo de 5 bytes:
      byte0 = bits [9:2] do pixel 0
      byte1 = bits [9:2] do pixel 1
      byte2 = bits [9:2] do pixel 2
      byte3 = bits [9:2] do pixel 3
      byte4 = bits [1:0] dos 4 pixels, empilhados em pares de bits

    Retorna uint16 com valores 0..1023.

    Importante:
    - A versão anterior tratava byte0..byte3 como bits baixos e byte4 como bits altos.
      Isso embaralha o RAW e gera exatamente aquele padrão de ruído/colorido sem imagem.
    """
    expected_pixels = width * height
    arr = np.frombuffer(buf, dtype=np.uint8)

    expected_bytes = (expected_pixels * 10 + 7) // 8
    if arr.size < expected_bytes:
        raise ValueError(
            f"RAW10 menor que esperado: bytes={arr.size}, esperado>={expected_bytes}, "
            f"shape={width}x{height}"
        )

    arr = arr[:expected_bytes]
    groups = arr.size // 5
    main = arr[: groups * 5].reshape(-1, 5).astype(np.uint16)

    p0 = (main[:, 0] << 2) | ((main[:, 4] >> 0) & 0x03)
    p1 = (main[:, 1] << 2) | ((main[:, 4] >> 2) & 0x03)
    p2 = (main[:, 2] << 2) | ((main[:, 4] >> 4) & 0x03)
    p3 = (main[:, 3] << 2) | ((main[:, 4] >> 6) & 0x03)

    out = np.empty(groups * 4, dtype=np.uint16)
    out[0::4] = p0
    out[1::4] = p1
    out[2::4] = p2
    out[3::4] = p3

    if out.size < expected_pixels:
        raise ValueError(f"RAW10 gerou pixels insuficientes: {out.size} < {expected_pixels}")

    return out[:expected_pixels].reshape(height, width)


def pack_raw10_packed(raw: np.ndarray) -> bytes:
    """
    Empacota uint16 0..1023 para MIPI RAW10 packed no mesmo padrão do unpack:

      byte0..byte3 = bits altos [9:2]
      byte4        = bits baixos [1:0] de p0,p1,p2,p3
    """
    flat = np.asarray(raw, dtype=np.uint16).reshape(-1)
    flat = np.clip(flat, 0, 1023).astype(np.uint16)

    pad = (-flat.size) % 4
    if pad:
        flat = np.pad(flat, (0, pad), mode="edge")

    p = flat.reshape(-1, 4)
    out = np.empty((p.shape[0], 5), dtype=np.uint8)

    out[:, 0] = ((p[:, 0] >> 2) & 0xFF).astype(np.uint8)
    out[:, 1] = ((p[:, 1] >> 2) & 0xFF).astype(np.uint8)
    out[:, 2] = ((p[:, 2] >> 2) & 0xFF).astype(np.uint8)
    out[:, 3] = ((p[:, 3] >> 2) & 0xFF).astype(np.uint8)
    out[:, 4] = (
        ((p[:, 0] & 0x03) << 0)
        | ((p[:, 1] & 0x03) << 2)
        | ((p[:, 2] & 0x03) << 4)
        | ((p[:, 3] & 0x03) << 6)
    ).astype(np.uint8)

    return out.reshape(-1).tobytes()


def pack_raw10_packed_array(raw: np.ndarray) -> np.ndarray:
    """
    RAW16 0..1023 -> ndarray uint8 RAW10 packed shape=(height, packed_width).
    Esse é o formato esperado pelo RawProcessorCore.decode_stream_cameras.
    """
    h, w = raw.shape[:2]
    packed = np.frombuffer(pack_raw10_packed(raw), dtype=np.uint8)
    packed_w = int(math.ceil(int(w) * 10 / 8))
    return packed.reshape(int(h), packed_w)


def read_raw_bin(path: Path, width: int, height: int, raw_format: str) -> np.ndarray:
    raw_format_u = (raw_format or "").upper()
    buf = path.read_bytes()

    if "RAW10" in raw_format_u or "PACKED" in raw_format_u:
        return unpack_raw10_packed(buf, width, height)

    # Fallback: RAW16 little-endian com valores possivelmente 0..1023.
    arr = np.frombuffer(buf, dtype=np.uint16)
    expected = width * height
    if arr.size < expected:
        raise ValueError(f"RAW16 menor que esperado em {path}: {arr.size} < {expected}")
    return arr[:expected].reshape(height, width)


def write_raw_bin(path: Path, raw: np.ndarray, raw_format: str) -> None:
    ensure_dir(path.parent)
    raw_format_u = (raw_format or "").upper()

    raw10 = np.clip(np.rint(raw), 0, 1023).astype(np.uint16)

    if "RAW10" in raw_format_u or "PACKED" in raw_format_u:
        path.write_bytes(pack_raw10_packed(raw10))
    else:
        path.write_bytes(raw10.astype(np.uint16).tobytes())


# ============================================================
# Metadados e descoberta de bins
# ============================================================

def _text_has_any(text: str, needles: List[str]) -> bool:
    t = text.lower()
    return any(n.lower() in t for n in needles)


def role_from_text(text: str) -> Optional[str]:
    for role, aliases in ROLE_ALIASES.items():
        if _text_has_any(text, aliases):
            return role
    return None


def normalize_role(role: Optional[str], cam_key: Optional[str] = None) -> Optional[str]:
    if role:
        r = str(role).lower().strip()
        if r in ("rgb", "color", "cor"):
            return "rgb"
        if r in ("re", "rededge", "red_edge", "red-edge"):
            return "re"
        if r in ("nir", "ir", "infrared", "infra"):
            return "nir"

    if cam_key:
        ck = str(cam_key).upper()
        # Convenção atual mais provável: CAM_A=RGB, CAM_B=RE, CAM_C=NIR.
        if ck == "CAM_A":
            return "rgb"
        if ck == "CAM_B":
            return "re"
        if ck == "CAM_C":
            return "nir"

    return None


def extract_camera_info(meta: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:
    """
    Tenta extrair camera_info de vários formatos possíveis.
    Retorna dict cam_key -> info.
    """
    candidates = []

    for key in ["camera_info", "cameras", "camera_infos", "payload_sources_info"]:
        if isinstance(meta.get(key), dict):
            candidates.append(meta[key])

    # Algumas versões guardam dentro de meta["capture"] ou meta["meta"]
    for parent_key in ["capture", "meta", "frame_meta"]:
        parent = meta.get(parent_key)
        if isinstance(parent, dict):
            for key in ["camera_info", "cameras", "camera_infos"]:
                if isinstance(parent.get(key), dict):
                    candidates.append(parent[key])

    if candidates:
        cam_info = candidates[0]
        out = {}
        for k, v in cam_info.items():
            if isinstance(v, dict):
                out[str(k)] = v
        return out

    # Fallback mínimo caso o meta já tenha campos diretos.
    out = {}
    for ck in CAM_KEYS:
        if isinstance(meta.get(ck), dict):
            out[ck] = meta[ck]
    return out


def get_int_any(d: Dict[str, Any], keys: List[str], default: Optional[int] = None) -> Optional[int]:
    for k in keys:
        if k in d and d[k] is not None:
            try:
                return int(d[k])
            except Exception:
                pass
    return default


def get_str_any(d: Dict[str, Any], keys: List[str], default: str = "") -> str:
    for k in keys:
        if k in d and d[k] is not None:
            return str(d[k])
    return default


def find_bin_for_camera(bins_dir: Path, base: str, role: str, cam_key: str) -> Optional[Path]:
    """
    Localiza o bin de uma câmera por base + role/cam_key.
    É tolerante a nomes tipo:
      base_CAM_A.bin
      base_rgb.bin
      base_cam0.bin
      base_arquivo_CAM_A.bin
    """
    if not bins_dir.is_dir():
        return None

    role_aliases = ROLE_ALIASES.get(role, [role])
    cam_key_variants = {cam_key, cam_key.upper(), cam_key.lower()}

    candidates = []
    for p in sorted(bins_dir.glob("*.bin")):
        name = p.stem
        name_l = name.lower()
        base_l = base.lower()
        if not name_l.startswith(base_l):
            continue

        score = 0
        if name_l == base_l:
            score += 1
        if any(v.lower() in name_l for v in cam_key_variants):
            score += 10
        if any(a.lower() in name_l for a in role_aliases):
            score += 12

        # Compatibilidade com cam0/cam1/cam2 caso necessário.
        if role == "rgb" and "cam0" in name_l:
            score += 3
        if role == "re" and "cam1" in name_l:
            score += 3
        if role == "nir" and "cam2" in name_l:
            score += 3

        if score > 0:
            candidates.append((score, p))

    if not candidates:
        return None

    candidates.sort(key=lambda x: x[0], reverse=True)
    return candidates[0][1]


def discover_sample_cameras(meta: Dict[str, Any], bins_dir: Path, base: str) -> Dict[str, CameraBin]:
    cam_info = extract_camera_info(meta)
    out: Dict[str, CameraBin] = {}

    # Primeiro caminho: usar camera_info do meta.
    for cam_key, info in cam_info.items():
        role = normalize_role(info.get("role") or info.get("papel") or info.get("type"), cam_key)
        if role not in ("rgb", "re", "nir"):
            role = role_from_text(cam_key) or role_from_text(json.dumps(info, ensure_ascii=False))

        if role not in ("rgb", "re", "nir"):
            continue

        width = get_int_any(info, ["width", "w", "sensor_width", "cols", "shape_w"])
        height = get_int_any(info, ["height", "h", "sensor_height", "rows", "shape_h"])

        if width is None or height is None:
            # Fallback para campos globais do meta.
            width = get_int_any(meta, ["width", "sensor_width", "raw_width", "w"])
            height = get_int_any(meta, ["height", "sensor_height", "raw_height", "h"])

        if width is None or height is None:
            continue

        bit_depth = get_int_any(info, ["bit_depth", "bits", "raw_bit_depth"], DEFAULT_BIT_DEPTH)
        raw_format = get_str_any(info, ["raw_format", "format", "encoding"], DEFAULT_RAW_FORMAT)
        bayer_pattern = get_str_any(info, ["bayer_pattern", "bayer", "pattern"], DEFAULT_BAYER_PATTERN)

        bin_path = find_bin_for_camera(bins_dir, base, role, cam_key)
        if bin_path is None:
            continue

        out[role] = CameraBin(
            role=role,
            cam_key=str(cam_key),
            path=bin_path,
            width=int(width),
            height=int(height),
            bit_depth=int(bit_depth or DEFAULT_BIT_DEPTH),
            raw_format=raw_format or DEFAULT_RAW_FORMAT,
            bayer_pattern=bayer_pattern or DEFAULT_BAYER_PATTERN,
        )

    # Segundo caminho: se o meta não ajudou, tenta descobrir por nome.
    if len(out) < 3:
        global_w = get_int_any(meta, ["width", "sensor_width", "raw_width", "w"])
        global_h = get_int_any(meta, ["height", "sensor_height", "raw_height", "h"])

        for role, cam_key in [("rgb", "CAM_A"), ("re", "CAM_B"), ("nir", "CAM_C")]:
            if role in out:
                continue
            if global_w is None or global_h is None:
                continue
            bin_path = find_bin_for_camera(bins_dir, base, role, cam_key)
            if bin_path is None:
                continue
            out[role] = CameraBin(
                role=role,
                cam_key=cam_key,
                path=bin_path,
                width=int(global_w),
                height=int(global_h),
                bit_depth=DEFAULT_BIT_DEPTH,
                raw_format=DEFAULT_RAW_FORMAT,
                bayer_pattern=DEFAULT_BAYER_PATTERN,
            )

    return out


# ============================================================
# Preview RGB simples a partir de RAW Bayer
# ============================================================

def _normalize_preview_channel(x: np.ndarray, lo_p: float = 1.0, hi_p: float = 99.5, gamma: float = 1.0 / 2.2) -> np.ndarray:
    """
    Normalização visual robusta para preview, não científica.
    Transforma um canal RAW/float em uint8 bonito para inspeção humana.
    """
    y = x.astype(np.float32)
    lo = float(np.percentile(y, lo_p))
    hi = float(np.percentile(y, hi_p))
    if hi <= lo:
        hi = lo + 1.0
    y = np.clip((y - lo) / (hi - lo), 0.0, 1.0)
    if gamma and abs(gamma - 1.0) > 1e-6:
        y = np.power(y, gamma)
    return np.clip(y * 255.0, 0, 255).astype(np.uint8)


def _get_bayer_cv2_code_like_core(bayer_pattern: str | None, algorithm: str = "ea") -> int:
    """
    Replica o mapeamento usado pelo RawProcessorCore._get_bayer_cv2_code.

    Observação importante:
    O mapeamento OpenCV parece invertido à primeira vista, mas é exatamente
    o contrato usado no core atual para gerar o RGB do treinamento/inferência.
    """
    p = str(bayer_pattern or DEFAULT_BAYER_PATTERN or "RGGB").upper()
    algo = str(algorithm or "ea").lower()

    if algo in ("bilinear", "linear", "fast", "normal"):
        code_map = {
            "BGGR": cv2.COLOR_BayerRG2RGB,
            "RGGB": cv2.COLOR_BayerBG2RGB,
            "GRBG": cv2.COLOR_BayerGR2RGB,
            "GBRG": cv2.COLOR_BayerGB2RGB,
        }
    elif algo in ("ea", "edge_aware", "edge-aware"):
        code_map = {
            "BGGR": cv2.COLOR_BayerRG2RGB_EA,
            "RGGB": cv2.COLOR_BayerBG2RGB_EA,
            "GRBG": cv2.COLOR_BayerGR2RGB_EA,
            "GBRG": cv2.COLOR_BayerGB2RGB_EA,
        }
    else:
        raise ValueError(f"demosaic_algorithm inválido: {algorithm}")

    if p not in code_map:
        raise ValueError(f"Padrão Bayer não suportado para demosaic: {p}")

    return code_map[p]


def demosaic_raw10_to_rgb01(raw: np.ndarray, bayer_pattern: str = DEFAULT_BAYER_PATTERN, algorithm: str = "ea") -> np.ndarray:
    """
    RAW Bayer 0..1023 -> RGB HWC float32 0..1 usando o mesmo mapeamento do core.
    """
    raw_u16 = np.clip(raw, 0, 1023).astype(np.uint16)
    cv_code = _get_bayer_cv2_code_like_core(bayer_pattern, algorithm=algorithm)
    rgb16 = cv2.cvtColor(raw_u16, cv_code)
    rgb = rgb16.astype(np.float32) / 1023.0
    return np.clip(rgb, 0.0, 1.0).astype(np.float32, copy=False)


def remosaic_rgb01_to_bayer_raw10(rgb: np.ndarray, bayer_pattern: str = DEFAULT_BAYER_PATTERN) -> np.ndarray:
    """
    RGB HWC float32 0..1 -> mosaico Bayer RAW10 uint16.

    Isso é necessário porque NÃO podemos aplicar warp diretamente no mosaico Bayer.
    O fluxo correto para augmentar RGB RAW é:
      raw Bayer -> demosaic RGB -> augmentation -> remosaic Bayer -> pack RAW10.

    Não é uma reconstrução óptica perfeita, mas preserva o contrato RAW10 Bayer para
    o normalize/RawProcessorCore e evita o xadrez colorido causado por interpolar
    diretamente o mosaico.
    """
    rgb = np.clip(rgb.astype(np.float32), 0.0, 1.0)
    h, w = rgb.shape[:2]
    raw = np.empty((h, w), dtype=np.float32)

    p = str(bayer_pattern or DEFAULT_BAYER_PATTERN or "RGGB").upper()

    if p == "RGGB":
        raw[0::2, 0::2] = rgb[0::2, 0::2, 0]  # R
        raw[0::2, 1::2] = rgb[0::2, 1::2, 1]  # G
        raw[1::2, 0::2] = rgb[1::2, 0::2, 1]  # G
        raw[1::2, 1::2] = rgb[1::2, 1::2, 2]  # B
    elif p == "BGGR":
        raw[0::2, 0::2] = rgb[0::2, 0::2, 2]  # B
        raw[0::2, 1::2] = rgb[0::2, 1::2, 1]  # G
        raw[1::2, 0::2] = rgb[1::2, 0::2, 1]  # G
        raw[1::2, 1::2] = rgb[1::2, 1::2, 0]  # R
    elif p == "GRBG":
        raw[0::2, 0::2] = rgb[0::2, 0::2, 1]  # G
        raw[0::2, 1::2] = rgb[0::2, 1::2, 0]  # R
        raw[1::2, 0::2] = rgb[1::2, 0::2, 2]  # B
        raw[1::2, 1::2] = rgb[1::2, 1::2, 1]  # G
    elif p == "GBRG":
        raw[0::2, 0::2] = rgb[0::2, 0::2, 1]  # G
        raw[0::2, 1::2] = rgb[0::2, 1::2, 2]  # B
        raw[1::2, 0::2] = rgb[1::2, 0::2, 0]  # R
        raw[1::2, 1::2] = rgb[1::2, 1::2, 1]  # G
    else:
        raise ValueError(f"Padrão Bayer não suportado para remosaic: {p}")

    return np.clip(np.rint(raw * 1023.0), 0, 1023).astype(np.uint16)


def rgb01_to_preview_rgb8(rgb: np.ndarray) -> np.ndarray:
    rgb = np.clip(rgb.astype(np.float32), 0.0, 1.0)
    rgb8 = np.zeros(rgb.shape, dtype=np.uint8)
    for c in range(3):
        rgb8[:, :, c] = _normalize_preview_channel(rgb[:, :, c], lo_p=1.0, hi_p=99.5, gamma=1.0 / 2.2)
    return rgb8


def debayer_raw10_to_rgb8(raw: np.ndarray, bayer_pattern: str = DEFAULT_BAYER_PATTERN) -> np.ndarray:
    """
    Preview visual a partir do RAW Bayer, usando o mesmo mapeamento Bayer do core.
    """
    rgb01 = demosaic_raw10_to_rgb01(raw, bayer_pattern=bayer_pattern, algorithm="ea")
    return rgb01_to_preview_rgb8(rgb01)


def _resolve_module_params_path(meta: Dict[str, Any]) -> Optional[str]:
    candidates = []

    for key in ("camera_params_json", "module_params_json"):
        v = meta.get(key)
        if v:
            candidates.append(str(v))

    stream_meta = meta.get("stream_meta") if isinstance(meta.get("stream_meta"), dict) else None
    if stream_meta:
        for key in ("camera_params_json", "module_params_json"):
            v = stream_meta.get(key)
            if v:
                candidates.append(str(v))

    cfg_module = read_config_value("module_params_json", None)
    if cfg_module:
        candidates.append(str(cfg_module))

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

    for c in candidates:
        p = Path(c)
        if p.is_file():
            return str(p)

    return None


def _build_processing_meta_for_core(meta: Dict[str, Any], cameras: Dict[str, CameraBin]) -> Dict[str, Any]:
    """
    Monta o meta que o RawProcessorCore espera no build_infer_tensor_from_stream.
    O validador faz isso a partir de stream_meta + camera_info. Aqui criamos/atualizamos
    esse pacote usando os bins augmentados.
    """
    stream_meta = copy.deepcopy(meta.get("stream_meta") if isinstance(meta.get("stream_meta"), dict) else {})

    camera_info = stream_meta.get("camera_info")
    if not isinstance(camera_info, dict):
        camera_info = copy.deepcopy(meta.get("camera_info") if isinstance(meta.get("camera_info"), dict) else {})

    camera_frames = stream_meta.get("camera_frames")
    if not isinstance(camera_frames, dict):
        camera_frames = copy.deepcopy(meta.get("camera_frames") if isinstance(meta.get("camera_frames"), dict) else {})

    for role, cam in cameras.items():
        info = dict(camera_info.get(cam.cam_key, {}) or {})
        info.update({
            "role": role,
            "width": int(cam.width),
            "height": int(cam.height),
            "bit_depth": int(cam.bit_depth),
            "raw_format": cam.raw_format or DEFAULT_RAW_FORMAT,
            "packed": True,
            "channels": 1,
            "bayer_pattern": cam.bayer_pattern or DEFAULT_BAYER_PATTERN,
        })
        camera_info[cam.cam_key] = info
        camera_frames[cam.cam_key] = dict(info)

    stream_meta["frame_type"] = "RAW_BRUTO"
    stream_meta["camera_info"] = camera_info
    stream_meta["camera_frames"] = camera_frames
    stream_meta["payload_sources"] = [cameras[r].cam_key for r in ("rgb", "re", "nir") if r in cameras]

    # A normalização radiométrica do core busca frame_controls no próprio meta.
    # Preservamos os controles originais se existirem.
    for key in ("frame_controls", "actual_camera_controls", "startup_camera_controls"):
        if key not in stream_meta and isinstance(meta.get(key), dict):
            stream_meta[key] = copy.deepcopy(meta[key])

    return stream_meta


def get_preview_core_cached(sensor_width: int, sensor_height: int, bayer: str, calib_path: str):
    """
    Reutiliza o RawProcessorCore entre amostras para evitar recarregar flat-field
    e recriar caches de gain/remap toda hora.
    """
    if not _HAS_RAW_PROCESSOR_CORE:
        return None

    key = (
        int(sensor_width),
        int(sensor_height),
        str(bayer).upper(),
        str(Path(calib_path).resolve()) if calib_path else "",
    )

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

    core = RawProcessorCore(
        sensor_width=int(sensor_width),
        sensor_height=int(sensor_height),
        bayer_pattern=str(bayer).upper(),
        calibration_json_path=calib_path,
    )

    _CORE_PREVIEW_CACHE[key] = core
    return core


def build_core_preview_from_augmented_raws(meta: Dict[str, Any], cameras: Dict[str, CameraBin]) -> Optional[np.ndarray]:
    """
    Gera o preview salvo passando os bins augmentados pelo RawProcessorCore,
    igual ao check_saved_files faz para montar o MULTISPEC final.

    Saída: RGB uint8 do tensor final [R,G,B], pronto para salvar com PIL.
    """
    if not _HAS_RAW_PROCESSOR_CORE:
        return None

    required = ("rgb", "re", "nir")
    if any(r not in cameras or cameras[r].data is None for r in required):
        return None

    sensor_width = int(meta.get("sensor_width") or cameras["rgb"].width)
    sensor_height = int(meta.get("sensor_height") or cameras["rgb"].height)
    bayer = str(meta.get("bayer_pattern") or cameras["rgb"].bayer_pattern or DEFAULT_BAYER_PATTERN)
    calib_path = _resolve_module_params_path(meta)

    if not calib_path:
        return None

    frame = {}
    for role in required:
        cam = cameras[role]
        frame[cam.cam_key] = pack_raw10_packed_array(cam.data)

    processing_meta = _build_processing_meta_for_core(meta, cameras)

    core = get_preview_core_cached(
        sensor_width=sensor_width,
        sensor_height=sensor_height,
        bayer=bayer,
        calib_path=calib_path,
    )

    if core is None:
        return None

    tensor = core.build_infer_tensor_from_stream(frame, processing_meta, 5)
    if tensor is None or not isinstance(tensor, np.ndarray) or tensor.ndim != 3 or tensor.shape[0] < 3:
        return None

    rgb_hwc = np.transpose(tensor[:3].astype(np.float32), (1, 2, 0))
    rgb8 = np.clip(rgb_hwc * 255.0, 0, 255).astype(np.uint8)
    return rgb8


def build_beauty_preview_from_augmented_raws(cameras: Dict[str, CameraBin]) -> np.ndarray:
    """
    Fallback visual a partir dos RAWs augmentados, sem passar pelo core.
    Preferimos build_core_preview_from_augmented_raws sempre que possível.
    """
    rgb_cam = cameras.get("rgb")
    if rgb_cam is None or rgb_cam.data is None:
        raise RuntimeError("Sem RAW RGB para gerar preview.")

    return debayer_raw10_to_rgb8(rgb_cam.data, rgb_cam.bayer_pattern)


def save_preview_from_augmented_raws(path: Path, cameras: Dict[str, CameraBin], meta: Optional[Dict[str, Any]] = None) -> str:
    ensure_dir(path.parent)

    method = "fallback_cam_a_debayer_preview"
    rgb8 = None

    if meta is not None:
        try:
            rgb8 = build_core_preview_from_augmented_raws(meta, cameras)
            if rgb8 is not None:
                method = "raw_processor_core_multispec_rgb_final"
        except Exception as e:
            print(f"[WARN] Preview via RawProcessorCore falhou em {path.name}: {e}. Usando fallback CAM_A.")
            rgb8 = None

    if rgb8 is None:
        rgb8 = build_beauty_preview_from_augmented_raws(cameras)

    Image.fromarray(rgb8).save(path)
    return method


# ============================================================
# Augmentation
# ============================================================

def sample_aug_params(rng: random.Random, seed_value: int) -> AugParams:
    # Geometria conservadora.
    perspective = rng.random() < 0.15
    blur_enabled = rng.random() < 0.15

    return AugParams(
        seed=seed_value,
        flip_h=rng.random() < 0.50,
        shift_x_frac=rng.uniform(-0.015, 0.015),
        shift_y_frac=rng.uniform(-0.015, 0.015),
        scale=rng.uniform(0.92, 1.08),
        rotate_deg=rng.uniform(-4.0, 4.0),
        perspective=perspective,
        perspective_strength=rng.uniform(0.002, 0.012) if perspective else 0.0,
        exposure_mult=rng.uniform(0.75, 1.30),
        gamma=rng.uniform(0.92, 1.08),
        rgb_channel_gain=(
            rng.uniform(0.92, 1.08),
            rng.uniform(0.92, 1.08),
            rng.uniform(0.92, 1.08),
        ),
        re_gain=rng.uniform(0.85, 1.20),
        nir_gain=rng.uniform(0.85, 1.20),
        shadow_enabled=rng.random() < 0.25,
        shadow_strength=rng.uniform(0.12, 0.35),
        shadow_angle_deg=rng.uniform(0.0, 180.0),
        highlight_enabled=rng.random() < 0.15,
        highlight_strength=rng.uniform(0.05, 0.18),
        noise_enabled=rng.random() < 0.30,
        noise_sigma_dn=rng.uniform(1.5, 6.0),
        blur_enabled=blur_enabled,
        blur_kernel=3 if blur_enabled else 1,
    )


def build_affine_matrix(width: int, height: int, p: AugParams) -> np.ndarray:
    cx = width * 0.5
    cy = height * 0.5
    M = cv2.getRotationMatrix2D((cx, cy), p.rotate_deg, p.scale)
    M[0, 2] += p.shift_x_frac * width
    M[1, 2] += p.shift_y_frac * height
    return M


def warp_image(img: np.ndarray, M: np.ndarray, out_size: Tuple[int, int], is_mask: bool) -> np.ndarray:
    interp = cv2.INTER_NEAREST if is_mask else cv2.INTER_LINEAR
    return cv2.warpAffine(
        img,
        M,
        out_size,
        flags=interp,
        borderMode=cv2.BORDER_REFLECT_101,
    )


def build_perspective_matrix(width: int, height: int, p: AugParams, rng: random.Random) -> Optional[np.ndarray]:
    if not p.perspective:
        return None

    s = p.perspective_strength
    dx = width * s
    dy = height * s

    src = np.float32([
        [0, 0],
        [width - 1, 0],
        [width - 1, height - 1],
        [0, height - 1],
    ])

    dst = src + np.float32([
        [rng.uniform(-dx, dx), rng.uniform(-dy, dy)],
        [rng.uniform(-dx, dx), rng.uniform(-dy, dy)],
        [rng.uniform(-dx, dx), rng.uniform(-dy, dy)],
        [rng.uniform(-dx, dx), rng.uniform(-dy, dy)],
    ])

    return cv2.getPerspectiveTransform(src, dst)


def warp_perspective(img: np.ndarray, H: np.ndarray, out_size: Tuple[int, int], is_mask: bool) -> np.ndarray:
    interp = cv2.INTER_NEAREST if is_mask else cv2.INTER_LINEAR
    return cv2.warpPerspective(
        img,
        H,
        out_size,
        flags=interp,
        borderMode=cv2.BORDER_REFLECT_101,
    )


def apply_geometry_raw(raw: np.ndarray, p: AugParams, rng: random.Random) -> np.ndarray:
    h, w = raw.shape[:2]

    out = raw
    if p.flip_h:
        out = cv2.flip(out, 1)

    M = build_affine_matrix(w, h, p)
    out = warp_image(out, M, (w, h), is_mask=False)

    H = build_perspective_matrix(w, h, p, rng)
    if H is not None:
        out = warp_perspective(out, H, (w, h), is_mask=False)

    return out


def apply_geometry_mask(mask: np.ndarray, p: AugParams, rng: random.Random) -> np.ndarray:
    h, w = mask.shape[:2]

    out = mask
    if p.flip_h:
        out = cv2.flip(out, 1)

    M = build_affine_matrix(w, h, p)
    out = warp_image(out, M, (w, h), is_mask=True)

    H = build_perspective_matrix(w, h, p, rng)
    if H is not None:
        out = warp_perspective(out, H, (w, h), is_mask=True)

    return out


def make_spatial_light_map(shape: Tuple[int, int], p: AugParams) -> np.ndarray:
    h, w = shape
    yy, xx = np.mgrid[0:h, 0:w].astype(np.float32)
    xx = (xx / max(w - 1, 1)) - 0.5
    yy = (yy / max(h - 1, 1)) - 0.5

    light = np.ones((h, w), dtype=np.float32)

    if p.shadow_enabled:
        theta = math.radians(p.shadow_angle_deg)
        direction = math.cos(theta) * xx + math.sin(theta) * yy
        direction = (direction - direction.min()) / max(direction.max() - direction.min(), 1e-6)
        shadow = 1.0 - p.shadow_strength * direction
        light *= shadow

    if p.highlight_enabled:
        # Mancha larga e suave, simulando região de sol/reflexo.
        cx = np.random.uniform(-0.25, 0.25)
        cy = np.random.uniform(-0.25, 0.25)
        sigma = np.random.uniform(0.20, 0.42)
        d2 = (xx - cx) ** 2 + (yy - cy) ** 2
        blob = np.exp(-d2 / (2 * sigma * sigma))
        light *= 1.0 + p.highlight_strength * blob

    return np.clip(light, 0.50, 1.45).astype(np.float32)


def apply_gamma_raw(raw: np.ndarray, gamma: float) -> np.ndarray:
    if abs(gamma - 1.0) < 1e-3:
        return raw
    x = np.clip(raw.astype(np.float32) / 1023.0, 0, 1)
    x = np.power(x, gamma)
    return x * 1023.0


def apply_bayer_channel_gains(raw: np.ndarray, gains_rgb: Tuple[float, float, float], pattern: str) -> np.ndarray:
    """
    Aplica ganhos aproximados por canal em mosaico Bayer.
    Para preview/treino, isso simula variação de balanço/cor no sensor RGB.
    """
    out = raw.astype(np.float32).copy()
    r_gain, g_gain, b_gain = gains_rgb
    pat = (pattern or DEFAULT_BAYER_PATTERN).upper()

    if pat == "RGGB":
        out[0::2, 0::2] *= r_gain
        out[0::2, 1::2] *= g_gain
        out[1::2, 0::2] *= g_gain
        out[1::2, 1::2] *= b_gain
    elif pat == "BGGR":
        out[0::2, 0::2] *= b_gain
        out[0::2, 1::2] *= g_gain
        out[1::2, 0::2] *= g_gain
        out[1::2, 1::2] *= r_gain
    elif pat == "GRBG":
        out[0::2, 0::2] *= g_gain
        out[0::2, 1::2] *= r_gain
        out[1::2, 0::2] *= b_gain
        out[1::2, 1::2] *= g_gain
    elif pat == "GBRG":
        out[0::2, 0::2] *= g_gain
        out[0::2, 1::2] *= b_gain
        out[1::2, 0::2] *= r_gain
        out[1::2, 1::2] *= g_gain
    else:
        out *= float(np.mean(gains_rgb))

    return out


def apply_radiometry_rgb01(rgb: np.ndarray, p: AugParams, rng_np: np.random.Generator, light_map_cache: Dict[Tuple[int, int], np.ndarray]) -> np.ndarray:
    """
    Radiometria para RGB já demosaicado.
    Evita mexer no mosaico Bayer diretamente.
    """
    h, w = rgb.shape[:2]
    x = np.clip(rgb.astype(np.float32), 0.0, 1.0)

    x *= np.float32(p.exposure_mult)

    gains = np.array(p.rgb_channel_gain, dtype=np.float32).reshape(1, 1, 3)
    x *= gains

    key = (h, w)
    if key not in light_map_cache:
        light_map_cache[key] = make_spatial_light_map((h, w), p)
    x *= light_map_cache[key][:, :, None]

    if abs(p.gamma - 1.0) > 1e-3:
        x = np.power(np.clip(x, 0.0, 1.0), p.gamma)

    if p.blur_enabled and p.blur_kernel >= 3:
        x = cv2.GaussianBlur(x, (p.blur_kernel, p.blur_kernel), 0)

    if p.noise_enabled:
        # Converte sigma DN RAW10 para escala 0..1.
        sigma01 = float(p.noise_sigma_dn) / 1023.0
        noise = rng_np.normal(0.0, sigma01, size=x.shape).astype(np.float32)
        x += noise

    return np.clip(x, 0.0, 1.0).astype(np.float32, copy=False)


def apply_radiometry_raw(raw: np.ndarray, role: str, p: AugParams, cam: CameraBin, rng_np: np.random.Generator, light_map_cache: Dict[Tuple[int, int], np.ndarray]) -> np.ndarray:
    h, w = raw.shape[:2]
    x = raw.astype(np.float32)

    # Base global de exposição: muda todo mundo junto.
    x *= p.exposure_mult

    # Variação por papel espectral/canal.
    # OBS: RGB não deve passar por aqui, porque RGB Bayer precisa ser demosaicado
    # antes de qualquer warp/radiometria por canal.
    if role == "re":
        x *= p.re_gain
    elif role == "nir":
        x *= p.nir_gain

    # Mesmo mapa espacial de luz por tamanho, para manter coerência física.
    key = (h, w)
    if key not in light_map_cache:
        light_map_cache[key] = make_spatial_light_map((h, w), p)
    x *= light_map_cache[key]

    # Gamma leve, opcional. Em RAW científico puro, gamma seria discutível.
    # Mantemos bem pequeno para simular resposta/exposição não ideal.
    x = apply_gamma_raw(x, p.gamma)

    if p.blur_enabled and p.blur_kernel >= 3:
        x = cv2.GaussianBlur(x, (p.blur_kernel, p.blur_kernel), 0)

    if p.noise_enabled:
        noise = rng_np.normal(0.0, p.noise_sigma_dn, size=x.shape).astype(np.float32)
        x += noise

    return np.clip(x, 0, 1023).astype(np.uint16)


# ============================================================
# Processamento de amostra/grupo
# ============================================================

def validate_mask_colors(mask_before: np.ndarray, mask_after: np.ndarray, name: str) -> None:
    if mask_before.ndim == 2 or mask_after.ndim == 2:
        before_colors = int(np.unique(mask_before.reshape(-1)).size)
        after_colors = int(np.unique(mask_after.reshape(-1)).size)
    else:
        before_colors = len(set(map(tuple, mask_before.reshape(-1, mask_before.shape[2]))))
        after_colors = len(set(map(tuple, mask_after.reshape(-1, mask_after.shape[2]))))

    if after_colors > max(before_colors * 3, 64):
        print(
            f"[WARN] {name}: máscara ganhou muitas cores. "
            f"antes={before_colors}, depois={after_colors}. Confira interpolação NEAREST."
        )


def save_aug_meta(
    meta: Dict[str, Any],
    out_path: Path,
    base: str,
    new_base: str,
    params: AugParams,
    cameras: Dict[str, CameraBin],
    preview_method: Optional[str] = None,
) -> None:
    out = copy.deepcopy(meta)
    out["synthetic"] = True
    out["augmentation"] = {
        "script": "_5_augmentation_raw_oak.py",
        "parent_base": base,
        "new_base": new_base,
        "params": asdict(params),
    }

    # Deixa o JSON compatível com o check_saved_files / validador.
    out["saved_payload_type"] = "raw_native_multi"
    out["saved_preview_method"] = preview_method or "unknown"
    out["saved_payload_paths"] = {}
    out["saved_payload_shapes"] = {}
    out["saved_payload_dtypes"] = {}

    stream_meta = copy.deepcopy(out.get("stream_meta") if isinstance(out.get("stream_meta"), dict) else {})
    camera_info = stream_meta.get("camera_info") if isinstance(stream_meta.get("camera_info"), dict) else {}
    camera_frames = stream_meta.get("camera_frames") if isinstance(stream_meta.get("camera_frames"), dict) else {}

    out["augmented_bins"] = {}

    for role, cam in cameras.items():
        packed_w = int(math.ceil(int(cam.width) * 10 / 8))
        fname = f"{new_base}_{cam.cam_key}.bin"

        out["saved_payload_paths"][cam.cam_key] = fname
        out["saved_payload_shapes"][cam.cam_key] = [int(cam.height), int(packed_w)]
        out["saved_payload_dtypes"][cam.cam_key] = "uint8"

        cam_info = dict(camera_info.get(cam.cam_key, {}) or {})
        cam_info.update({
            "role": role,
            "width": int(cam.width),
            "height": int(cam.height),
            "bit_depth": int(cam.bit_depth),
            "raw_format": cam.raw_format or DEFAULT_RAW_FORMAT,
            "packed": True,
            "channels": 1,
            "bayer_pattern": cam.bayer_pattern or DEFAULT_BAYER_PATTERN,
        })
        camera_info[cam.cam_key] = cam_info
        camera_frames[cam.cam_key] = dict(cam_info)

        out["augmented_bins"][role] = {
            "cam_key": cam.cam_key,
            "filename": fname,
            "role": role,
            "width": int(cam.width),
            "height": int(cam.height),
            "bit_depth": int(cam.bit_depth),
            "raw_format": cam.raw_format,
            "bayer_pattern": cam.bayer_pattern,
        }

    stream_meta["frame_type"] = "RAW_BRUTO"
    stream_meta["camera_info"] = camera_info
    stream_meta["camera_frames"] = camera_frames
    stream_meta["payload_sources"] = [cameras[r].cam_key for r in ("rgb", "re", "nir") if r in cameras]
    out["stream_meta"] = stream_meta

    save_json(out_path, out)


def process_one_sample(
    base: str,
    group_name: str,
    src_group: Path,
    dst_group: Path,
    mask_path: Path,
    meta_path: Path,
    copies: int,
    rng_global: random.Random,
    dry_run: bool = False,
) -> Tuple[int, int]:
    bins_dir = src_group / "bins"
    meta = load_json(meta_path)
    cameras = discover_sample_cameras(meta, bins_dir, base)

    required = ["rgb", "re", "nir"]
    missing = [r for r in required if r not in cameras]
    if missing:
        print(f"[WARN] [{group_name}] {base}: câmeras ausentes {missing}. Pulando.")
        return 0, 1

    mask = load_mask(mask_path)

    # Carrega todos os RAWs uma vez.
    for role, cam in cameras.items():
        cam.data = read_raw_bin(cam.path, cam.width, cam.height, cam.raw_format)

    generated = 0
    errors = 0

    out_bins = dst_group / "bins"
    out_masks = dst_group / "masks"
    out_metas = dst_group / "metas"
    out_previews = dst_group / "previews"

    for i in range(copies):
        seed_value = rng_global.randint(0, 2**31 - 1)
        rng = random.Random(seed_value)
        rng_np = np.random.default_rng(seed_value)
        params = sample_aug_params(rng, seed_value)
        new_base = f"{base}_aug_{i:02d}"

        try:
            if dry_run:
                generated += 1
                continue

            # Mesma geometria para mask e raws.
            mask_aug = apply_geometry_mask(mask, params, random.Random(seed_value + 1000))
            validate_mask_colors(mask, mask_aug, f"{new_base}_mask")

            light_map_cache: Dict[Tuple[int, int], np.ndarray] = {}
            cams_out: Dict[str, CameraBin] = {}

            for role in required:
                cam = cameras[role]
                assert cam.data is not None

                if role == "rgb":
                    # RGB RAW é mosaico Bayer. Não podemos aplicar warp direto nele,
                    # porque isso mistura pixels R/G/B antes do demosaic e cria ruído colorido.
                    # Fluxo correto offline:
                    #   RAW Bayer -> RGB linear -> geometria/radiometria -> mosaico Bayer -> RAW10
                    rgb01 = demosaic_raw10_to_rgb01(
                        cam.data,
                        bayer_pattern=cam.bayer_pattern,
                        algorithm="ea",
                    )

                    rgb_geom = rgb01
                    if params.flip_h:
                        rgb_geom = cv2.flip(rgb_geom, 1)

                    h_rgb, w_rgb = rgb_geom.shape[:2]
                    M_rgb = build_affine_matrix(w_rgb, h_rgb, params)
                    rgb_geom = warp_image(rgb_geom, M_rgb, (w_rgb, h_rgb), is_mask=False)

                    H_rgb = build_perspective_matrix(w_rgb, h_rgb, params, random.Random(seed_value + 1000))
                    if H_rgb is not None:
                        rgb_geom = warp_perspective(rgb_geom, H_rgb, (w_rgb, h_rgb), is_mask=False)

                    rgb_aug = apply_radiometry_rgb01(
                        rgb_geom,
                        params,
                        rng_np,
                        light_map_cache,
                    )

                    raw_aug = remosaic_rgb01_to_bayer_raw10(
                        rgb_aug,
                        bayer_pattern=cam.bayer_pattern,
                    )

                else:
                    raw_geom = apply_geometry_raw(cam.data, params, random.Random(seed_value + 1000))
                    raw_aug = apply_radiometry_raw(raw_geom, role, params, cam, rng_np, light_map_cache)

                out_cam = copy.deepcopy(cam)
                out_cam.path = out_bins / f"{new_base}_{cam.cam_key}.bin"
                out_cam.data = raw_aug
                cams_out[role] = out_cam

                write_raw_bin(out_cam.path, raw_aug, cam.raw_format)

            # Máscara com mesmo nome-base.
            mask_ext = mask_path.suffix.lower() if mask_path.suffix.lower() in IMG_EXTS else ".png"
            save_mask(out_masks / f"{new_base}{mask_ext}", mask_aug)

            # Preview visual usando o mesmo caminho do validador/check_saved_files:
            # bins augmentados -> RawProcessorCore -> MULTISPEC final -> RGB final.
            # Se o core falhar, cai no fallback CAM_A e avisa no console.
            preview_method = save_preview_from_augmented_raws(
                out_previews / f"{new_base}.png",
                cams_out,
                meta=meta,
            )

            # Meta novo, com mesmo stem do preview/mask para manter o layout dataset limpo.
            save_aug_meta(
                meta=meta,
                out_path=out_metas / f"{new_base}.json",
                base=base,
                new_base=new_base,
                params=params,
                cameras=cams_out,
                preview_method=preview_method,
            )

            generated += 1

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

    return generated, errors


def parse_group_copies(text: Optional[str]) -> Dict[str, int]:
    """
    Ex: chao:1,chao_cana:3,chao_erva:3,chao_cana_erva:5
    """
    out: Dict[str, int] = {}
    if not text:
        return out

    for part in text.split(","):
        part = part.strip()
        if not part:
            continue
        if ":" not in part:
            raise ValueError(f"group-copies inválido: {part}. Use grupo:N")
        g, n = part.split(":", 1)
        out[g.strip()] = int(n.strip())
    return out


def process_group(
    src_root: Path,
    dst_root: Path,
    group_name: str,
    copies: int,
    limit: Optional[int],
    seed: int,
    dry_run: bool,
) -> Tuple[int, int, int]:
    src_group = src_root / group_name
    dst_group = dst_root / group_name

    mask_map = map_files_by_base(src_group / "masks", IMG_EXTS)
    meta_map = map_files_by_base(src_group / "metas", META_EXTS)

    bases = sorted(set(mask_map.keys()) & set(meta_map.keys()))

    if limit is not None and limit > 0 and limit < len(bases):
        rng_select = random.Random(seed)
        bases = sorted(rng_select.sample(bases, limit))
        print(f"[INFO] [{group_name}] limit={limit}, amostras selecionadas={len(bases)}")

    if not bases:
        print(f"[WARN] [{group_name}] Nenhum par mask/meta encontrado.")
        return 0, 0, 0

    ensure_dir(dst_group / "bins")
    ensure_dir(dst_group / "masks")
    ensure_dir(dst_group / "metas")
    ensure_dir(dst_group / "previews")

    rng_global = random.Random(seed + abs(hash(group_name)) % 1000000)

    total_generated = 0
    total_errors = 0

    for idx, base in enumerate(bases, start=1):
        gen, err = process_one_sample(
            base=base,
            group_name=group_name,
            src_group=src_group,
            dst_group=dst_group,
            mask_path=mask_map[base],
            meta_path=meta_map[base],
            copies=copies,
            rng_global=rng_global,
            dry_run=dry_run,
        )
        total_generated += gen
        total_errors += err

        if idx % 25 == 0 or idx == len(bases):
            print(
                f"[INFO] [{group_name}] {idx}/{len(bases)} amostras | "
                f"gerados={total_generated} | erros={total_errors}"
            )

    print(
        f"[OK] Grupo '{group_name}' concluído: "
        f"originais={len(bases)} | copies={copies} | gerados={total_generated} | erros={total_errors}"
    )

    return len(bases), total_generated, total_errors


# ============================================================
# CLI
# ============================================================

def main() -> None:
    dataset_base = Path("dataset")

    default_src = dataset_base / "original" / "group"
    if not default_src.exists():
        # Compatibilidade com script antigo que usava "originals".
        alt = dataset_base / "originals" / "group"
        if alt.exists():
            default_src = alt

    default_dst = dataset_base / "augmented" / "group"

    ap = argparse.ArgumentParser(
        description="Augmentation multiespectral RAW10 packed para dataset OAK-FCC-3."
    )

    ap.add_argument("--copies", type=int, default=3, help="Cópias augmentadas por amostra, se --group-copies não sobrescrever.")
    ap.add_argument("--group-copies", type=str, default=None, help="Cópias por grupo. Ex: chao:1,chao_cana:3,chao_erva:3,chao_cana_erva:5")
    ap.add_argument("--groups", type=str, default=None, help="Lista de grupos separados por vírgula.")
    ap.add_argument("--src-root", type=str, default=str(default_src), help="Raiz dos grupos originais.")
    ap.add_argument("--dst-root", type=str, default=str(default_dst), help="Raiz dos grupos augmentados.")
    ap.add_argument("--limit", type=int, default=None, help="Quantidade máxima de amostras originais por grupo.")
    ap.add_argument("--seed", type=int, default=42, help="Seed geral para reproducibilidade.")
    ap.add_argument("--clear-dst", action="store_true", help="Apaga dst-root antes de gerar.")
    ap.add_argument("--dry-run", action="store_true", help="Simula sem salvar arquivos.")

    args = ap.parse_args()

    src_root = Path(args.src_root)
    dst_root = Path(args.dst_root)
    group_copies = parse_group_copies(args.group_copies)

    print("==============================================")
    print("Augmentation RAW OAK-FCC-3")
    print(f"MODEL/DATASET : {dataset_base}")
    print(f"SRC_ROOT      : {src_root}")
    print(f"DST_ROOT      : {dst_root}")
    print(f"copies        : {args.copies}")
    print(f"group_copies  : {group_copies if group_copies else '{}'}")
    print(f"limit         : {args.limit}")
    print(f"seed          : {args.seed}")
    print(f"dry_run       : {args.dry_run}")
    print("==============================================")

    if not src_root.exists():
        raise SystemExit(f"[ERRO] src-root não encontrado: {src_root}")

    if args.clear_dst and not args.dry_run:
        print(f"[INFO] Limpando destino: {dst_root}")
        clear_dir(dst_root)
    else:
        ensure_dir(dst_root)

    groups = list_groups(src_root)

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

    if not groups:
        raise SystemExit("[WARN] Nenhum grupo válido encontrado.")

    print(f"Grupos encontrados: {', '.join(groups)}")

    total_originals = 0
    total_generated = 0
    total_errors = 0

    for group_name in groups:
        copies = group_copies.get(group_name, args.copies)
        if copies <= 0:
            print(f"[INFO] [{group_name}] copies={copies}. Pulando.")
            continue

        n_orig, n_gen, n_err = process_group(
            src_root=src_root,
            dst_root=dst_root,
            group_name=group_name,
            copies=copies,
            limit=args.limit,
            seed=args.seed,
            dry_run=args.dry_run,
        )
        total_originals += n_orig
        total_generated += n_gen
        total_errors += n_err

    print("\n==============================================")
    print("Resumo final")
    print(f"Originais processados : {total_originals}")
    print(f"Amostras geradas      : {total_generated}")
    print(f"Erros/Pulos           : {total_errors}")
    print(f"Destino               : {dst_root}")
    print("==============================================")


if __name__ == "__main__":
    main()