#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Augmenta imagens, máscaras e opcionalmente labels globais por grupo.

Entrada via config_oak.json -> camera:

  MODELO/dataset/original/group/<grupo>/images/
  MODELO/dataset/original/group/<grupo>/masks/
  MODELO/dataset/original/group/<grupo>/masks2/    opcional, se dual_head_mask=true
  MODELO/dataset/original/group/<grupo>/labels/    opcional, se dual_head_label=true

Saída:

  MODELO/dataset/augmented/group/<grupo>/images/
  MODELO/dataset/augmented/group/<grupo>/masks/
  MODELO/dataset/augmented/group/<grupo>/masks2/    se dual_head_mask=true
  MODELO/dataset/augmented/group/<grupo>/labels/    se dual_head_label=true

Observação:
- A máscara e a mask2 recebem as mesmas transformações geométricas da imagem.
- O label global NÃO é transformado visualmente. Ele é copiado e atualizado para apontar para o novo base augmentado.

Uso:
  python _3_augmentation_grouped_with_labels.py --copies 5
  python _3_augmentation_grouped_with_labels.py --copies 5 --groups navegavel,naonavegavel_navegavel
  python _3_augmentation_grouped_with_labels.py --copies 5 --strict-label
"""

import os
import json
import argparse
from pathlib import Path

import cv2
from PIL import Image
import albumentations as A


# ====================== Configurações ======================

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

MODELO = config.get("camera", ".")
USE_MASKS2 = bool(config.get("dual_head_mask", False))
USE_LABELS = bool(config.get("dual_head_label", False))

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 legacy, 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")
ORIG_OLD_LABELS = os.path.join(DATASET_BASE, "original", "labels")

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")
AUG_OLD_LABELS = os.path.join(DATASET_BASE, "augmented", "labels")

IMG_EXTS = (".jpg", ".jpeg", ".png")
MSK_EXTS = (".png", ".jpg", ".jpeg")
MSK2_EXTS = (".png", ".jpg", ".jpeg")
LABEL_EXTS = (".json", ".txt")


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

train_tf = A.Compose(
    [
        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,
        ),

        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),
    ],
    additional_targets={"mask2": "mask"},
)


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


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


def normalizar_base(stem: str) -> str:
    sufixos = [
        "_rgb", "_RGB", "_Rgb",
        "_image", "_img", "_frame",
        "_mask", "_masks",
        "_seg", "_SEG", "_segment", "_segmentacao", "_Segmentacao",
        "_label", "_labels",
    ]

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


def map_files_by_base(folder, exts):
    by_base = {}
    if not os.path.isdir(folder):
        return by_base

    prioridade = {
        ".json": 0,
        ".png": 1,
        ".jpg": 2,
        ".jpeg": 3,
        ".txt": 4,
    }

    for fname in os.listdir(folder):
        lower = fname.lower()
        if not lower.endswith(exts):
            continue

        stem, ext = os.path.splitext(fname)
        base = normalizar_base(stem)
        cand = os.path.join(folder, fname)

        if base not in by_base:
            by_base[base] = cand
        else:
            cur_ext = os.path.splitext(by_base[base])[1].lower()
            if prioridade.get(ext.lower(), 99) < prioridade.get(cur_ext, 99):
                by_base[base] = cand

    return by_base


def list_groups(root):
    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
        if os.path.isdir(os.path.join(gdir, "images")) and os.path.isdir(os.path.join(gdir, "masks")):
            grupos.append(name)
    return grupos


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):
    garantir_dir(os.path.dirname(path))
    Image.fromarray(arr_rgb).save(path)


def ensure_aug_dirs(group_name=None, use_masks2=False, use_labels=False):
    if group_name:
        img_out = os.path.join(AUG_GROUP_ROOT, group_name, "images")
        msk_out = os.path.join(AUG_GROUP_ROOT, group_name, "masks")
        msk2_out = os.path.join(AUG_GROUP_ROOT, group_name, "masks2") if use_masks2 else None
        labels_out = os.path.join(AUG_GROUP_ROOT, group_name, "labels") if use_labels else None
    else:
        img_out = AUG_OLD_IMG
        msk_out = AUG_OLD_MSK
        msk2_out = AUG_OLD_MSK2 if use_masks2 else None
        labels_out = AUG_OLD_LABELS if use_labels else None

    garantir_dir(img_out)
    garantir_dir(msk_out)

    if use_masks2 and msk2_out:
        garantir_dir(msk2_out)

    if use_labels and labels_out:
        garantir_dir(labels_out)

    return img_out, msk_out, msk2_out, labels_out


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


def copiar_label_aug(label_path, out_label_path, new_base, out_img_path, out_msk_path, out_msk2_path=None, group_name=None):
    """
    Copia label global para a amostra augmentada.
    Se for JSON, atualiza campos úteis.
    Se for TXT, copia o conteúdo como está.
    """
    if not label_path or not out_label_path:
        return

    garantir_dir(os.path.dirname(out_label_path))
    ext = os.path.splitext(label_path)[1].lower()

    if ext == ".json":
        try:
            with open(label_path, "r", encoding="utf-8") as f:
                data = json.load(f)
        except Exception:
            data = {}

        data["base"] = new_base
        data["group"] = group_name if group_name is not None else data.get("group")
        data["image"] = safe_rel(out_img_path, DATASET_BASE)
        data["mask"] = safe_rel(out_msk_path, DATASET_BASE)
        data["label"] = safe_rel(out_label_path, DATASET_BASE)
        data["is_augmented"] = True
        data["source_label"] = safe_rel(label_path, DATASET_BASE)

        if out_msk2_path:
            data["mask2"] = safe_rel(out_msk2_path, DATASET_BASE)

        with open(out_label_path, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
    else:
        with open(label_path, "r", encoding="utf-8") as f:
            content = f.read()
        with open(out_label_path, "w", encoding="utf-8") as f:
            f.write(content)


def augment_pair(
    img_path,
    msk_path,
    img_out_dir,
    msk_out_dir,
    copies,
    msk2_path=None,
    msk2_out_dir=None,
    label_path=None,
    label_out_dir=None,
    group_name=None,
):
    base_img, img_ext = os.path.splitext(os.path.basename(img_path))
    _, 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
    label_ext = os.path.splitext(os.path.basename(label_path))[1] if label_path else None

    base = normalizar_base(base_img)

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

    gen = 0
    for i in range(copies):
        if msk2 is not None and msk2_out_dir:
            aug = train_tf(image=img, mask=msk, mask2=msk2)
        else:
            aug = train_tf(image=img, mask=msk)

        img_aug = aug["image"]
        msk_aug = aug["mask"]

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

        save_rgb(out_img, img_aug)
        save_rgb(out_msk, msk_aug)

        out_msk2 = None
        if msk2 is not None and msk2_out_dir:
            msk2_aug = aug["mask2"]
            out_msk2 = os.path.join(msk2_out_dir, f"{new_base}{msk2_ext}")
            save_rgb(out_msk2, msk2_aug)

        if label_path and label_out_dir:
            out_label = os.path.join(label_out_dir, f"{new_base}{label_ext}")
            copiar_label_aug(
                label_path=label_path,
                out_label_path=out_label,
                new_base=new_base,
                out_img_path=out_img,
                out_msk_path=out_msk,
                out_msk2_path=out_msk2,
                group_name=group_name,
            )

        gen += 1

    return gen


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


def process_group(group_name, copies, strict_label=False):
    img_dir = os.path.join(ORIG_GROUP_ROOT, group_name, "images")
    msk_dir = os.path.join(ORIG_GROUP_ROOT, group_name, "masks")
    msk2_dir = os.path.join(ORIG_GROUP_ROOT, group_name, "masks2")
    labels_dir = os.path.join(ORIG_GROUP_ROOT, group_name, "labels")

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

    imgs = [f for f in os.listdir(img_dir) if os.path.splitext(f.lower())[1] in IMG_EXTS]

    msk_map = map_files_by_base(msk_dir, MSK_EXTS)

    use_masks2 = USE_MASKS2 and os.path.isdir(msk2_dir)
    msk2_map = map_files_by_base(msk2_dir, MSK2_EXTS) if use_masks2 else {}

    use_labels = USE_LABELS and os.path.isdir(labels_dir)
    label_map = map_files_by_base(labels_dir, LABEL_EXTS) if use_labels else {}

    if USE_MASKS2 and not use_masks2:
        print(f"[WARN] [{group_name}] dual_head_mask=true, mas pasta masks2 não existe.")

    if USE_LABELS and not use_labels:
        msg = f"[WARN] [{group_name}] dual_head_label=true, mas pasta labels não existe."
        if strict_label:
            print(msg + " Pulando grupo.")
            return 0
        print(msg + " Gerando sem labels.")

    img_out_dir, msk_out_dir, msk2_out_dir, label_out_dir = ensure_aug_dirs(
        group_name,
        use_masks2=use_masks2,
        use_labels=use_labels,
    )

    count = 0
    sem_mask = 0
    sem_mask2 = 0
    sem_label = 0

    for img_file in sorted(imgs):
        base, _ = os.path.splitext(img_file)
        base_norm = normalizar_base(base)

        msk_file = msk_map.get(base_norm)
        if not msk_file:
            sem_mask += 1
            print(f"[WARN] [{group_name}] Máscara não encontrada para {img_file}, pulando.")
            continue

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

        label_file = label_map.get(base_norm) if use_labels else None
        if use_labels and not label_file:
            sem_label += 1
            msg = f"[WARN] [{group_name}] label não encontrado para {img_file}."
            if strict_label:
                print(msg + " Pulando item.")
                continue
            print(msg + " Gerando augmentation sem label.")

        try:
            count += augment_pair(
                img_path=os.path.join(img_dir, img_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,
                label_path=label_file,
                label_out_dir=label_out_dir,
                group_name=group_name,
            )
        except Exception as e:
            print(f"[ERRO] [{group_name}] {img_file}: {e}")

    print(
        f"[OK] Grupo '{group_name}' → {count} pares gerados. "
        f"sem_mask={sem_mask} | sem_mask2={sem_mask2} | sem_label={sem_label}"
    )
    return count


def process_legacy(copies, strict_label=False):
    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_files_by_base(ORIG_OLD_MSK, MSK_EXTS)

    use_masks2 = USE_MASKS2 and os.path.isdir(ORIG_OLD_MSK2)
    msk2_map = map_files_by_base(ORIG_OLD_MSK2, MSK2_EXTS) if use_masks2 else {}

    use_labels = USE_LABELS and os.path.isdir(ORIG_OLD_LABELS)
    label_map = map_files_by_base(ORIG_OLD_LABELS, LABEL_EXTS) if use_labels else {}

    if USE_LABELS and not use_labels and strict_label:
        print("[WARN] Legacy com dual_head_label=true, mas sem pasta labels. Pulando.")
        return 0

    img_out_dir, msk_out_dir, msk2_out_dir, label_out_dir = ensure_aug_dirs(
        group_name=None,
        use_masks2=use_masks2,
        use_labels=use_labels,
    )

    count = 0
    for img_file in sorted(imgs):
        base, _ = os.path.splitext(img_file)
        base_norm = normalizar_base(base)

        msk_file = msk_map.get(base_norm)
        if not msk_file:
            print(f"[WARN] (legacy) Máscara não encontrada para {img_file}, pulando.")
            continue

        msk2_file = msk2_map.get(base_norm) if use_masks2 else None
        label_file = label_map.get(base_norm) if use_labels else None

        if use_labels and not label_file and strict_label:
            print(f"[WARN] (legacy) Label não encontrado para {img_file}, pulando.")
            continue

        try:
            count += augment_pair(
                img_path=os.path.join(ORIG_OLD_IMG, img_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,
                label_path=label_file,
                label_out_dir=label_out_dir,
                group_name=None,
            )
        except Exception as e:
            print(f"[ERRO] (legacy) {img_file}: {e}")

    print(f"[OK] Legacy → {count} pares gerados.")
    return count


def main(copies=5, groups_csv=None, strict_label=False):
    total = 0

    print(f"[INFO] MODELO={MODELO}")
    print(f"[INFO] dual_head_mask={USE_MASKS2}")
    print(f"[INFO] dual_head_label={USE_LABELS}")

    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, strict_label=strict_label)
        else:
            print(f"Grupos encontrados: {', '.join(grupos)}")
            for g in grupos:
                total += process_group(g, copies, strict_label=strict_label)
    else:
        total += process_legacy(copies, strict_label=strict_label)

    print(f"\nAugmentation completed! Total: {total} pares gerados.")


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description="Augmentação por grupos com suporte a images/masks/masks2/labels.")
    ap.add_argument("--copies", type=int, default=5, help="Número de cópias augmentadas por imagem.")
    ap.add_argument("--groups", type=str, default=None, help="Lista de grupos separados por vírgula.")
    ap.add_argument("--strict-label", action="store_true", help="Se dual_head_label=true e faltar label, pula o item/grupo.")
    args = ap.parse_args()

    main(copies=args.copies, groups_csv=args.groups, strict_label=args.strict_label)