#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Augmenta imagens/máscaras *por grupo* usando RAW4 (float32, 0..1, R G IR B).

Novo fluxo (GAL5000):

Entrada:
  dataset/original/group/<grupo>/raws   (.raw, 4 canais float32)
  dataset/original/group/<grupo>/masks  (.png/.jpg)
  (opcional) dataset/original/group/<grupo>/masks2

Saída:
  dataset/augmented/group/<grupo>/previews  (RGB gerado do RAW blurado)
  dataset/augmented/group/<grupo>/raws      (RAW4 augmentado: geom + blur)
  dataset/augmented/group/<grupo>/masks
  dataset/augmented/group/<grupo>/masks2

Regras:
- GEOMETRIA (flip, shift, rotate) aplicada em RAW + masks (+ masks2).
- BLUR / MOTION BLUR aplicado no RAW, canal por canal, com o mesmo kernel.
- Preview é SEMPRE derivado do RAW já transformado (geom + blur).
- Máscaras não sofrem blur, apenas geom.

Uso:
  python _5_augmentation_raw.py --copies 5
  python _5_augmentation_raw.py --copies 5 --groups chao,chao_erva,cana
"""
import os
import json
import cv2
from PIL import Image
import albumentations as A
from albumentations import ReplayCompose
import argparse
import numpy as np

# ⚙️ Configurações
with open("config.json", "r", encoding="utf-8") as f:
    config = json.load(f)
MODELO = config.get("camera", ".")
RESOLUCAO = config.get("raw_size")

USE_MASKS2 = config.get("dual_head", False)

# Pastas base
DATASET_BASE = os.path.join("dataset")
ORIG_GROUP_ROOT = os.path.join(DATASET_BASE, "original", "group")
AUG_GROUP_ROOT  = os.path.join(DATASET_BASE, "augmented", "group")

# Fallback (modo antigo, sem grupos) – mantido por compat
ORIG_OLD_IMG  = os.path.join(DATASET_BASE, "original", "images")
ORIG_OLD_MSK  = os.path.join(DATASET_BASE, "original", "masks")
ORIG_OLD_MSK2 = os.path.join(DATASET_BASE, "original", "masks2")
AUG_OLD_IMG   = os.path.join(DATASET_BASE, "augmented", "images")
AUG_OLD_MSK   = os.path.join(DATASET_BASE, "augmented", "masks")
AUG_OLD_MSK2  = os.path.join(DATASET_BASE, "augmented", "masks2")

# Extensões aceitas
IMG_EXTS  = (".jpg", ".jpeg", ".png")
MSK_EXTS  = (".png", ".jpg", ".jpeg")
MSK2_EXTS = (".png", ".jpg", ".jpeg")
RAW_EXTS  = (".raw",)


def garantir_dir(p):
    os.makedirs(p, exist_ok=True)


# ===============================
# Pipelines de augmentations
# ===============================

# 1) Geométricas: aplicam em preview “temporário”, mask, mask2 e RAW.
#    RAW é passado como 'raw' em additional_targets para receber só geom.
train_geo = ReplayCompose([
    A.HorizontalFlip(p=0.5),
    A.ShiftScaleRotate(
        shift_limit=0.01,
        scale_limit=0.10,
        rotate_limit=5,
        border_mode=cv2.BORDER_REFLECT_101,
        interpolation=cv2.INTER_LINEAR,
        p=0.30
    ),
], additional_targets={
    'mask2': 'mask',
    'raw': 'mask',  # trata raw como “mask” para receber apenas geom
})


# ===============================
# Funções de blur físico no RAW
# ===============================

def motion_blur_kernel(ksize=5, angle=0.0):
    """
    Gera um kernel de motion blur 2D com tamanho ksize e ângulo em graus.
    Usado tanto para RAW quanto para preview.
    """
    ksize = int(ksize)
    if ksize < 3:
        ksize = 3
    if ksize % 2 == 0:
        ksize += 1

    kernel = np.zeros((ksize, ksize), dtype=np.float32)
    # linha horizontal no centro
    kernel[ksize // 2, :] = 1.0

    # rotaciona a linha para o ângulo desejado
    center = (ksize / 2.0 - 0.5, ksize / 2.0 - 0.5)
    M = cv2.getRotationMatrix2D(center, angle, 1.0)
    kernel = cv2.warpAffine(kernel, M, (ksize, ksize))

    s = kernel.sum()
    if s > 0:
        kernel /= s
    return kernel


def apply_coherent_blur_raw(raw, p_motion=0.20, p_gauss=0.20):
    """
    Aplica blur/motion blur de forma COERENTE no RAW, canal por canal.

    - Mesmo kernel em TODOS os canais.
    - Máscaras não passam aqui.
    - Se nada for sorteado, retorna o RAW original.
    """
    if raw is None:
        return None

    r = np.random.rand()
    if r >= (p_motion + p_gauss):
        # sem blur nesta amostra
        return raw

    # escolhe tamanho de kernel
    k_choices = [3, 5, 7]
    ksize = int(np.random.choice(k_choices))

    if r < p_motion:
        # Motion blur com ângulo pequeno (vibração / movimento leve)
        angle = float(np.random.uniform(-20.0, 20.0))
        kernel = motion_blur_kernel(ksize=ksize, angle=angle)
    else:
        # Gaussian blur (leve desfoque de foco)
        if ksize % 2 == 0:
            ksize += 1
        g1d = cv2.getGaussianKernel(ksize, 0)
        kernel = g1d @ g1d.T  # vira kernel 2D

    # Aplica kernel canal a canal
    if raw.ndim == 2:
        raw_blur = cv2.filter2D(raw, ddepth=-1, kernel=kernel)
    else:
        raw_blur = np.empty_like(raw)
        h, w, c = raw.shape
        for ch in range(c):
            raw_blur[..., ch] = cv2.filter2D(
                raw[..., ch],
                ddepth=-1,
                kernel=kernel,
                borderType=cv2.BORDER_REFLECT_101,
            )

    # Garante que continua em [0,1] para float
    if np.issubdtype(raw_blur.dtype, np.floating):
        raw_blur = np.clip(raw_blur, 0.0, 1.0)

    return raw_blur


# ===============================
# IO de imagens / RAW
# ===============================

def raw_to_preview_rgb(raw_np: np.ndarray) -> np.ndarray:
    """
    raw_np: (H,W,C) float32 em 0..1, C=4 ou 5.

    Layout esperado (após _load_raw):
      - C=4 -> [R,G,B,IR]
      - C=5 -> [R,G,B,IR,NDVI]  (se um dia tiver)

    Usamos sempre R,G,B para o preview e devolvemos (H,W,3) uint8.
    """
    assert raw_np.ndim == 3, "raw_np deve ser (H,W,C)"

    H, W, C = raw_np.shape
    if C not in (4, 5):
        raise RuntimeError(f"Esperado C=4 ou 5, veio {C} em raw_to_preview_rgb")

    # RAW já vem como [R,G,B,IR] do _load_raw
    r = raw_np[..., 0].astype(np.float32)
    g = raw_np[..., 1].astype(np.float32)
    b = raw_np[..., 2].astype(np.float32)

    # Monta RGB em (H,W,3) para trabalhar
    rgb = np.stack([r, g, b], axis=-1)  # (H,W,3)

    # Auto white balance simples (referenciado no G)
    means = rgb.reshape(-1, 3).mean(axis=0)  # (3,)
    ref   = means[1]  # canal G
    gains = ref / (means + 1e-6)
    gains = np.clip(gains, 0.5, 2.0)
    rgb_wb = rgb * gains  # broadcasting em (3,)

    # Força do WB
    alpha = 1.0  # teste 0.4..0.8 se quiser misturar com original
    if alpha >= 0.999:
        rgb = rgb_wb
    else:
        rgb_orig = rgb.copy()
        rgb = (1 - alpha) * rgb_orig + alpha * rgb_wb

    rgb = np.clip(rgb, 0.0, 1.0)

    # Gamma mais suave
    rgb = np.power(rgb, 1/1.9, out=rgb)

    # Converte pra uint8 OpenCV-friendly
    rgb8 = (rgb * 255.0).clip(0, 255).astype(np.uint8)  # (H,W,3)

    return rgb8


def load_rgb(path):
    im = cv2.imread(path, cv2.IMREAD_COLOR)
    if im is None:
        raise FileNotFoundError(path)
    return cv2.cvtColor(im, cv2.COLOR_BGR2RGB)


def save_rgb(path, arr_rgb):
    Image.fromarray(arr_rgb).save(path)


def _load_raw(path: str, raw_hw) -> np.ndarray:
    ext = os.path.splitext(path)[1].lower()
    if ext == ".npz":
        z = np.load(path)
        key = list(z.keys())[0]
        arr = z[key]
    elif ext == ".npy":
        arr = np.load(path)
    elif ext == ".raw":
        if raw_hw is None:
            raise RuntimeError(
                f"Para ler .raw preciso de raw_hw=(H,W). "
                f"Passe raw_hw ao criar o dataset."
            )
        H, W = raw_hw
        size_bytes = os.path.getsize(path)
        mosa_bytes = H * W                # mosaico uint8
        raw4_bytes = 4 * H * W * 4        # 4 canais * H * W * 4 bytes (float32)
        if size_bytes == mosa_bytes:
            # MODO ANTIGO: mosaico uint8
            arr = np.fromfile(path, dtype=np.uint8)
            raw2d = arr.reshape(H, W)
            if (H % 2) != 0 or (W % 2) != 0:
                raise RuntimeError(f"raw_hw deve ser par em H e W p/ mosaico 2x2 (R,G,IR,B), veio H={H}, W={W}")
            H2, W2 = H // 2, W // 2
            r_sub  = raw2d[0::2, 0::2]
            g_sub  = raw2d[0::2, 1::2]
            ir_sub = raw2d[1::2, 0::2]
            b_sub  = raw2d[1::2, 1::2]
            from PIL import Image
            def upsample(ch_2d: np.ndarray) -> np.ndarray:
                im = Image.fromarray(ch_2d)  # uint8
                im = im.resize((W, H), resample=Image.BILINEAR)
                return np.array(im)
            r_full  = upsample(r_sub)
            g_full  = upsample(g_sub)
            ir_full = upsample(ir_sub)
            b_full  = upsample(b_sub)
            # (H, W, 4) uint8 [R,G,IR,B]
            arr = np.stack([r_full, g_full, ir_full, b_full], axis=-1).astype(np.uint8)
            # REORGANIZA PARA CONTRATO INTERNO [R,G,B,IR]
            arr = arr[..., [0, 1, 3, 2]]
        elif size_bytes == raw4_bytes:
            # NOVO MODO: RAW4 float32 salvo pelo normalize
            arr_f32 = np.fromfile(path, dtype=np.float32)
            raw4 = arr_f32.reshape(4, H, W)          # (C,H,W) [R,G,IR,B]
            arr = np.transpose(raw4, (1, 2, 0))      # (H,W,4) [R,G,IR,B]
            # REORGANIZA PARA CONTRATO INTERNO: [R,G,B,IR]
            arr = arr[..., [0, 1, 3, 2]]
        else:
            raise RuntimeError(
                f"Tamanho inesperado em {path}: {size_bytes} bytes "
                f"(esperado {mosa_bytes} ou {raw4_bytes})"
            )
    else:
        raise RuntimeError(f"Extensão não suportada para RAW: {ext} ({path})")
    # Até aqui queremos (H,W,4) base
    if arr.ndim != 3:
        raise RuntimeError(f"RAW precisa ser 3D, veio {arr.shape} em {path}")
    if arr.shape[0] == 4 and arr.shape[-1] != 4:
        arr = np.transpose(arr, (1, 2, 0))  # (H,W,4)
    if arr.shape[-1] != 4:
        raise RuntimeError(f"Esperava 4 canais base, veio shape={arr.shape} em {path}")
    return arr  # (H,W,4)

def _scale_to_float01(raw: np.ndarray, max_value = None) -> np.ndarray:
    if raw.dtype == np.uint16:
        mv = float(max_value) if max_value is not None else 65535.0
    elif raw.dtype == np.uint8:
        mv = float(max_value) if max_value is not None else 255.0
    else:
        # float já vem normalizado ou não; se passar max_value, aplica
        mv = float(max_value) if max_value is not None else None
    raw_f = raw.astype(np.float32)
    if mv is not None and mv > 0:
        raw_f = raw_f / mv
    # clip pra não explodir
    raw_f = np.clip(raw_f, 0.0, 1.0)
    return raw_f


def save_raw_any(path, arr):
    arr = np.asarray(arr)

    if arr.ndim == 3 and arr.shape[-1] == 4:
        # arr está no contrato interno [R,G,B,IR]
        arr_cam = internal_rgbir_to_camera_raw4(arr)   # [R,G,IR,B]
        arr_cam = np.transpose(arr_cam, (2, 0, 1))     # (4,H,W)
        arr_cam.astype(np.float32).tofile(path)
    else:
        arr.tofile(path)


# ===============================
# Helpers de diretório / mapeamento
# ===============================

def list_groups(root):
    """
    Lista grupos válidos para o NOVO fluxo:
      - group/<g>/raws
      - group/<g>/masks
    (previews/images não são mais obrigatórios)
    """
    if not os.path.isdir(root):
        return []
    grupos = []
    for name in sorted(os.listdir(root)):
        gdir = os.path.join(root, name)
        if not os.path.isdir(gdir):
            continue
        has_masks = os.path.isdir(os.path.join(gdir, "masks"))
        has_raws  = os.path.isdir(os.path.join(gdir, "raws"))
        if has_masks and has_raws:
            grupos.append(name)
    return grupos


def map_by_base_priorizando_png(msk_dir, exts):
    """Mapeia arquivos por base (prioriza .png)."""
    by_base = {}
    if not os.path.isdir(msk_dir):
        return by_base
    for fname in os.listdir(msk_dir):
        f_lower = fname.lower()
        if not f_lower.endswith(exts):
            continue
        base, ext = os.path.splitext(fname)
        cand = os.path.join(msk_dir, fname)
        if base not in by_base:
            by_base[base] = cand
        else:
            cur_ext = os.path.splitext(by_base[base])[1].lower()
            if cur_ext != ".png" and ext.lower() == ".png":
                by_base[base] = cand
    return by_base


def map_raws_by_base(raw_dir):
    by_base = {}
    if not os.path.isdir(raw_dir):
        return by_base
    for fname in os.listdir(raw_dir):
        f_lower = fname.lower()
        if not f_lower.endswith(RAW_EXTS):
            continue
        base, _ = os.path.splitext(fname)
        by_base[base] = os.path.join(raw_dir, fname)
    return by_base


def ensure_aug_dirs(group_name=None, use_masks2=False, use_raw=False, is_preview_mode=True):
    """
    Cria diretórios de saída para o grupo ou modo antigo.

    is_preview_mode:
      True  -> saídas em 'previews' + 'raws'
      False -> saídas em 'images'  + (sem raws)
    """
    if group_name:
        base = os.path.join(AUG_GROUP_ROOT, group_name)
        if is_preview_mode:
            img_out = os.path.join(base, "previews")
        else:
            img_out = os.path.join(base, "images")
        msk_out = os.path.join(base, "masks")
        msk2_out = os.path.join(base, "masks2") if use_masks2 else None
        raw_out = os.path.join(base, "raws") if use_raw else None
    else:
        # modo legacy
        img_out = AUG_OLD_IMG
        msk_out = AUG_OLD_MSK
        msk2_out = AUG_OLD_MSK2 if use_masks2 else None
        raw_out = None

    garantir_dir(img_out)
    garantir_dir(msk_out)
    if use_masks2 and msk2_out:
        garantir_dir(msk2_out)
    if use_raw and raw_out:
        garantir_dir(raw_out)
    return img_out, msk_out, msk2_out, raw_out


# ===============================
# Núcleo: augmentação de um sample
# ===============================

def augment_sample(
    raw_path,
    msk_path,
    img_out_dir,
    msk_out_dir,
    copies,
    msk2_path=None,
    msk2_out_dir=None,
    raw_out_dir=None,
    aug_suffix="aug",
):
    """
    Faz a augmentação a partir de UM RAW4 + máscara (e opcionalmente máscara2).

    Passos:
      1) Carrega máscara (para saber HxW).
      2) Carrega RAW4 float32 usando HxW da máscara.
      3) Gera preview RGB a partir do RAW4 (para servir de "image" no Albumentations).
      4) Aplica GEOMETRIA (mesma em preview, mask, mask2 e raw).
      5) Aplica BLUR/MOTION BLUR no RAW geométrico, canal a canal, com o mesmo kernel.
      6) Gera preview final a partir do RAW blurado.
      7) Salva:
           - preview_n em img_out_dir
           - mask_n em msk_out_dir
           - raw_n em raw_out_dir
           - mask2_n em msk2_out_dir (se houver)
    """
    base_raw, raw_ext = os.path.splitext(os.path.basename(raw_path))
    base = base_raw

    base_msk, msk_ext = os.path.splitext(os.path.basename(msk_path))
    if base_msk != base:
        print(f"[WARN] RAW '{base}' e máscara '{base_msk}' com bases diferentes, mas seguindo mesmo assim.")

    msk2_ext = os.path.splitext(os.path.basename(msk2_path))[1] if msk2_path else None

    # 1) Máscaras
    msk = load_rgb(msk_path)
    msk2 = load_rgb(msk2_path) if msk2_path else None

    # 2) RAW4
    raw = _load_raw(raw_path, raw_hw=(RESOLUCAO[1], RESOLUCAO[0]))           # (H,W,4) uint8
    raw = _scale_to_float01(raw)        # float32 0..1

    # 3) Preview inicial derivado do RAW (antes de geom/blur, só referência)
    img = raw_to_preview_rgb(raw)

    gen = 0
    for i in range(copies):
        # 4) GEOMETRIA: mesma transformação em img, mask, mask2 e raw
        if msk2 is not None and msk2_out_dir:
            aug_geo = train_geo(image=img, mask=msk, mask2=msk2, raw=raw)
        else:
            aug_geo = train_geo(image=img, mask=msk, raw=raw)

        msk_g = aug_geo["mask"]
        raw_g = aug_geo.get("raw", None)
        msk2_g = aug_geo.get("mask2", None) if (msk2 is not None and msk2_out_dir) else None

        # 5) BLUR FÍSICO: aplica (ou não) o MESMO blur no RAW geométrico
        raw_b = apply_coherent_blur_raw(raw_g)

        # 6) Preview final derivado do RAW blurado
        img_b = raw_to_preview_rgb(raw_b)

        # 7) Salva cópias
        out_img = os.path.join(img_out_dir, f"{base}_{aug_suffix}_{i:02d}.png")
        out_msk = os.path.join(msk_out_dir, f"{base}_{aug_suffix}_{i:02d}{msk_ext}")
        save_rgb(out_img, img_b)
        save_rgb(out_msk, msk_g)

        if msk2_g is not None and msk2_out_dir:
            out_msk2 = os.path.join(msk2_out_dir, f"{base}_{aug_suffix}_{i:02d}{msk2_ext}")
            save_rgb(out_msk2, msk2_g)

        if raw_b is not None and raw_out_dir:
            out_raw = os.path.join(raw_out_dir, f"{base}_{aug_suffix}_{i:02d}{raw_ext}")
            save_raw_any(out_raw, raw_b)

        gen += 1

    return gen


# ===============================
# Processamento por grupo
# ===============================

def process_group(group_name, copies, limit=None, seed=42, aug_suffix="aug"):
    """
    Processa um grupo único usando APENAS:
      - group/<g>/raws
      - group/<g>/masks (+ masks2 opcional)

    Previews de entrada são ignoradas.
    As de saída são geradas do RAW4 (geom + blur).
    """
    gdir = os.path.join(ORIG_GROUP_ROOT, group_name)

    raw_dir          = os.path.join(gdir, "raws")
    msk_dir          = os.path.join(gdir, "masks")
    msk2_dir         = os.path.join(gdir, "masks2")

    if not (os.path.isdir(raw_dir) and os.path.isdir(msk_dir)):
        print(f"[WARN] Grupo '{group_name}' inválido (precisa de 'raws' e 'masks'). Pulando.")
        return 0

    use_raw = True
    use_masks2 = USE_MASKS2 and os.path.isdir(msk2_dir)

    raw_map  = map_raws_by_base(raw_dir)
    items = sorted(raw_map.items())

    if limit is not None and limit > 0:
        rng = np.random.default_rng(seed)
        if limit < len(items):
            idx = rng.choice(len(items), size=limit, replace=False)
            idx = sorted(idx.tolist())
            items = [items[i] for i in idx]

    msk_map  = map_by_base_priorizando_png(msk_dir, MSK_EXTS)
    msk2_map = map_by_base_priorizando_png(msk2_dir, MSK2_EXTS) if use_masks2 else {}

    img_out_dir, msk_out_dir, msk2_out_dir, raw_out_dir = ensure_aug_dirs(
        group_name,
        use_masks2=use_masks2,
        use_raw=use_raw,
        is_preview_mode=True
    )

    count = 0
    for base, raw_file in items:
        msk_file = msk_map.get(base)
        if not msk_file:
            print(f"[WARN] [{group_name}] Máscara não encontrada para RAW {base}, pulando.")
            continue

        msk2_file = msk2_map.get(base) if use_masks2 else None
        if use_masks2 and not msk2_file:
            print(f"[WARN] [{group_name}] mask2 não encontrada para {base}, gerando só RAW+preview+mask.")

        try:
            count += augment_sample(
                raw_path=raw_file,
                msk_path=msk_file,
                img_out_dir=img_out_dir,
                msk_out_dir=msk_out_dir,
                copies=copies,
                msk2_path=msk2_file,
                msk2_out_dir=msk2_out_dir,
                raw_out_dir=raw_out_dir,
                aug_suffix=aug_suffix,
            )
        except Exception as e:
            print(f"[ERRO] [{group_name}] {base}: {e}")
    print(f"[OK] Grupo '{group_name}' → {count} amostras geradas.")
    return count


# ===============================
# Fallback legacy (sem RAW)
# ===============================

def process_legacy(copies):
    """
    Fallback: modo sem grupos (original/images e original/masks).
    Mantido por compat, mas não usa RAW.
    """
    if not (os.path.isdir(ORIG_OLD_IMG) and os.path.isdir(ORIG_OLD_MSK)):
        print("[WARN] Modo legacy não encontrado. Nada a fazer.")
        return 0

    imgs = [f for f in os.listdir(ORIG_OLD_IMG) if os.path.splitext(f.lower())[1] in IMG_EXTS]
    msk_map = map_by_base_priorizando_png(ORIG_OLD_MSK, MSK_EXTS)
    use_masks2 = USE_MASKS2 and os.path.isdir(ORIG_OLD_MSK2)
    msk2_map = map_by_base_priorizando_png(ORIG_OLD_MSK2, MSK2_EXTS) if use_masks2 else {}

    img_out_dir, msk_out_dir, msk2_out_dir, _ = ensure_aug_dirs(
        group_name=None,
        use_masks2=use_masks2,
        use_raw=False,
        is_preview_mode=False
    )

    count = 0
    for img_file in sorted(imgs):
        base, img_ext = os.path.splitext(img_file)
        msk_file = msk_map.get(base)
        if not msk_file:
            print(f"[WARN] (legacy) Máscara não encontrada para {img_file}, pulando.")
            continue
        msk2_file = msk2_map.get(base) if use_masks2 else None
        if use_masks2 and not msk2_file:
            print(f"[WARN] (legacy) mask2 não encontrada para {img_file}, gerando só img+mask.")

        try:
            img_path = os.path.join(ORIG_OLD_IMG, img_file)
            img = load_rgb(img_path)
            base_msk, msk_ext = os.path.splitext(os.path.basename(msk_file))
            msk = load_rgb(msk_file)
            msk2 = load_rgb(msk2_file) if msk2_file else None

            for i in range(copies):
                if msk2 is not None and msk2_out_dir:
                    aug_geo = train_geo(image=img, mask=msk, mask2=msk2)
                else:
                    aug_geo = train_geo(image=img, mask=msk)

                img_g = aug_geo["image"]
                msk_g = aug_geo["mask"]
                msk2_g = aug_geo.get("mask2", None) if (msk2 is not None and msk2_out_dir) else None

                out_img = os.path.join(img_out_dir, f"{base}_aug_{i:02d}{img_ext}")
                out_msk = os.path.join(msk_out_dir, f"{base_msk}_aug_{i:02d}{msk_ext}")
                save_rgb(out_img, img_g)
                save_rgb(out_msk, msk_g)

                if msk2_g is not None and msk2_out_dir:
                    msk2_ext = os.path.splitext(msk2_file)[1]
                    out_msk2 = os.path.join(msk2_out_dir, f"{base_msk}_aug_{i:02d}{msk2_ext}")
                    save_rgb(out_msk2, msk2_g)

                count += 1

        except Exception as e:
            print(f"[ERRO] (legacy) {img_file}: {e}")
    print(f"[OK] Legacy → {count} amostras geradas.")
    return count

def internal_rgbir_to_camera_raw4(arr: np.ndarray) -> np.ndarray:
    """
    Converte do contrato interno (H,W,4) [R,G,B,IR]
    para o contrato bruto da câmera (H,W,4) [R,G,IR,B].
    """
    if arr is None:
        return None
    if arr.ndim != 3 or arr.shape[-1] != 4:
        raise RuntimeError(f"Esperado array (H,W,4), veio {arr.shape}")
    return arr[..., [0, 1, 3, 2]]


# ===============================
# main
# ===============================

def main(copies=5, groups_csv=None, limit=None, seed=42, suffix="aug"):
    total = 0
    if os.path.isdir(ORIG_GROUP_ROOT):
        grupos = list_groups(ORIG_GROUP_ROOT)
        if groups_csv:
            want = {g.strip() for g in groups_csv.split(",") if g.strip()}
            grupos = [g for g in grupos if g in want]
            if not grupos:
                print("[WARN] Nenhum grupo válido encontrado após filtro.")
        if not grupos:
            print("[WARN] Nenhum grupo encontrado em original/group. Tentando modo legacy...")
            total += process_legacy(copies)
        else:
            print(f"Grupos encontrados: {', '.join(grupos)}")
            for g in grupos:
                total += process_group(g, copies, limit=limit, seed=seed, aug_suffix=suffix)
    else:
        total += process_legacy(copies)

    print(f"\nAugmentation completed! Total: {total} amostras geradas.")


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description="Augmentação por grupos usando RAW4 (preview gerado do RAW).")
    ap.add_argument("--copies", type=int, default=5, help="Número de cópias augmentadas por imagem (default=5).")
    ap.add_argument("--groups", type=str, default=None, help="Lista de grupos separados por vírgula (ex: chao,erva_cana).")
    ap.add_argument("--limit", type=int, default=None, help="Quantidade máxima de imagens originais do grupo a augmentar.")
    ap.add_argument("--seed", type=int, default=42, help="Seed para seleção reproduzível quando usar --limit.")
    ap.add_argument("--suffix", type=str, default="aug", help="Sufixo usado no nome dos arquivos gerados.")
    args = ap.parse_args()
    main(
        copies=args.copies,
        groups_csv=args.groups,
        limit=args.limit,
        seed=args.seed,
        suffix=args.suffix,
    )