#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Augmenta imagens/máscaras *por grupo*, agora com suporte a RAW.

Modos suportados:

1) Modo novo (RAW):
   MODELO/dataset/original/group/<grupo>/previews
   MODELO/dataset/original/group/<grupo>/raws
   MODELO/dataset/original/group/<grupo>/masks
   (opcional) MODELO/dataset/original/group/<grupo>/masks2

   Saída:
   MODELO/dataset/augmented/group/<grupo>/previews
   MODELO/dataset/augmented/group/<grupo>/raws
   MODELO/dataset/augmented/group/<grupo>/masks
   MODELO/dataset/augmented/group/<grupo>/masks2 (se existir)

2) Modo antigo (sem RAW, só images/masks):
   MODELO/dataset/original/group/<grupo>/images
   MODELO/dataset/original/group/<grupo>/masks

   Saída:
   MODELO/dataset/augmented/group/<grupo>/images
   MODELO/dataset/augmented/group/<grupo>/masks

Fallback legacy:
   MODELO/dataset/original/{images,masks}
   MODELO/dataset/augmented/{images,masks}

Transformações:
- Geométricas (HorizontalFlip, ShiftScaleRotate) → aplicam em preview, mask, mask2 e RAW.
- Fotométricas (brightness/contrast, gamma, blur, noise, flare, etc.) → só em preview RGB.

Uso:
  python _5_augmentation.py --copies 5
  python _5_augmentation.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", ".")
USE_MASKS2 = config.get("dual_head", False)

# Pastas base
DATASET_BASE = os.path.join(MODELO, "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)
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")   # prioriza PNG
MSK2_EXTS = (".png", ".jpg", ".jpeg")
RAW_EXTS  = (".raw",)

# Config de RAW (mosaico 1 canal; R G / IR B é tratado depois no loader)
RAW_DTYPE   = np.uint16
RAW_CHANNELS = 1  # tratamos como grayscale (H, W)
# Se RAW tiver tamanho diferente, ajuste no loader depois.
# Aqui assumimos que altura/largura batem com preview/mask.


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


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

# Geométricas: aplicam em preview, mask, mask2 e RAW
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',   # tratamos RAW como "mask" pra NÃO sofrer fotométricas
})

# Fotométricas: só preview RGB
train_photo = A.Compose([
    A.OneOf([
        A.RandomBrightnessContrast(0.2, 0.2, p=1.0),
        A.HueSaturationValue(hue_shift_limit=5, sat_shift_limit=20, val_shift_limit=15, p=1.0),
        A.RandomGamma(gamma_limit=(90, 110), p=1.0),
    ], p=0.70),

    A.OneOf([
        A.MotionBlur(blur_limit=3, p=1.0),
        A.GaussianBlur(blur_limit=3, p=1.0),
    ], p=0.20),

    A.OneOf([
        A.GaussNoise(var_limit=(5.0, 15.0), p=1.0),
        A.ImageCompression(quality_lower=50, quality_upper=85, p=1.0),
    ], p=0.20),

    A.RandomShadow(p=0.10),
    A.RandomSunFlare(p=0.10),
    A.ChannelShuffle(p=0.05),
    A.CoarseDropout(max_holes=6, max_height=16, max_width=16, p=0.10),
])


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_grayscale(path, like_shape=None):
    """
    Carrega RAW como (H, W) em uint8 ou uint16 detectando pelo tamanho do arquivo.
    """
    if like_shape is None:
        raise RuntimeError("like_shape é obrigatório para RAW (preciso do HxW).")

    h, w = like_shape[:2]
    npx = h * w

    fsize = os.path.getsize(path)
    if fsize == npx:            # 1 byte por pixel
        dtype = np.uint8
    elif fsize == npx * 2:      # 2 bytes por pixel
        dtype = np.uint16
    else:
        raise RuntimeError(
            f"Tamanho inesperado para RAW {path}: {fsize} bytes "
            f"(esperado {npx} (u8) ou {npx*2} (u16) para {w}x{h})."
        )

    data = np.fromfile(path, dtype=dtype)
    if data.size != npx:
        raise RuntimeError(f"RAW {path}: size={data.size} != {npx} (HxW)")

    return data.reshape((h, w))


def save_raw_grayscale(path, arr):
    np.asarray(arr).tofile(path)


def list_groups(root):
    """Lista grupos válidos (tem subpasta 'masks' e pelo menos 'images' ou 'previews')."""
    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_imgs  = os.path.isdir(os.path.join(gdir, "images")) or os.path.isdir(os.path.join(gdir, "previews"))
        if has_masks and has_imgs:
            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


def augment_sample(
    img_path,
    msk_path,
    img_out_dir,
    msk_out_dir,
    copies,
    msk2_path=None,
    msk2_out_dir=None,
    raw_path=None,
    raw_out_dir=None,
):
    base_img, img_ext = os.path.splitext(os.path.basename(img_path))
    base = base_img

    base_msk, msk_ext = os.path.splitext(os.path.basename(msk_path))
    msk2_ext = os.path.splitext(os.path.basename(msk2_path))[1] if msk2_path else None
    raw_ext = os.path.splitext(os.path.basename(raw_path))[1] if raw_path else None

    img = load_rgb(img_path)
    msk = load_rgb(msk_path)
    msk2 = load_rgb(msk2_path) if msk2_path else None

    if raw_path is not None and raw_out_dir is not None:
        # usa shape do preview pra validar RAW
        raw = load_raw_grayscale(raw_path, like_shape=img.shape)
    else:
        raw = None

    gen = 0
    for i in range(copies):
        # 1) GEOM: mesma transformação em img, mask, mask2 e raw
        if raw is not None:
            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)
        else:
            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"]
        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

        # 2) FOTO: só preview RGB
        img_p = train_photo(image=img_g)["image"]

        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}_aug_{i:02d}{msk_ext}")
        save_rgb(out_img, img_p)
        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_{i:02d}{msk2_ext}")
            save_rgb(out_msk2, msk2_g)

        if raw_g is not None and raw_out_dir:
            out_raw = os.path.join(raw_out_dir, f"{base}_aug_{i:02d}{raw_ext}")
            save_raw_grayscale(out_raw, raw_g)

        gen += 1

    return gen


def process_group(group_name, copies):
    """
    Processa um grupo único.
    Suporta:
      - group/<g>/previews + raws + masks (+ masks2)
      - group/<g>/images + masks (+ masks2)  [modo antigo]
    """
    gdir = os.path.join(ORIG_GROUP_ROOT, group_name)

    # prioridade para novo modo (previews/raws)
    img_dir_previews = os.path.join(gdir, "previews")
    img_dir_images   = os.path.join(gdir, "images")
    raw_dir          = os.path.join(gdir, "raws")
    msk_dir          = os.path.join(gdir, "masks")
    msk2_dir         = os.path.join(gdir, "masks2")

    use_preview_mode = os.path.isdir(img_dir_previews)
    img_dir = img_dir_previews if use_preview_mode else img_dir_images

    if not (os.path.isdir(img_dir) and os.path.isdir(msk_dir)):
        print(f"[WARN] Grupo '{group_name}' inválido (sem images/previews ou masks). Pulando.")
        return 0

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

    imgs = [f for f in os.listdir(img_dir) if os.path.splitext(f.lower())[1] in IMG_EXTS]
    msk_map = map_by_base_priorizando_png(msk_dir, MSK_EXTS)
    raw_map = map_raws_by_base(raw_dir) if use_raw else {}
    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=use_preview_mode
    )

    count = 0
    for img_file in sorted(imgs):
        base, _ = os.path.splitext(img_file)
        msk_file = msk_map.get(base)
        if not msk_file:
            print(f"[WARN] [{group_name}] Máscara não encontrada para {img_file}, pulando.")
            continue

        raw_file = raw_map.get(base) if use_raw else None
        if use_raw and not raw_file:
            print(f"[WARN] [{group_name}] RAW não encontrado para {img_file}, gerando só preview/mask (+mask2).")

        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 {img_file}, gerando só img+mask(+raw).")

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


def process_legacy(copies):
    """Fallback: modo sem grupos (original/images e original/masks)."""
    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, _ = 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:
            count += augment_sample(
                os.path.join(ORIG_OLD_IMG, img_file),
                msk_file,
                img_out_dir,
                msk_out_dir,
                copies=copies,
                msk2_path=msk2_file,
                msk2_out_dir=msk2_out_dir,
                raw_path=None,
                raw_out_dir=None,
            )
        except Exception as e:
            print(f"[ERRO] (legacy) {img_file}: {e}")
    print(f"[OK] Legacy → {count} amostras geradas.")
    return count


def main(copies=5, groups_csv=None):
    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)
    else:
        total += process_legacy(copies)

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


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description="Augmentação por grupos (preview/raw/masks ou images/masks)")
    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).")
    args = ap.parse_args()
    main(copies=args.copies, groups_csv=args.groups)