#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
_7_split_multi_source.py

Split estratificado por grupo/família para OAK-FCC-3, com suporte a múltiplas
raízes de dataset e proteção contra vazamento de dados sintéticos/augmentados.

Por que este script existe?
---------------------------
No fluxo atual temos, por exemplo:

  dataset/1024x640/group              -> dados reais normalizados
  dataset/copypaste/group             -> dados sintéticos copy/paste pós-normalização

A validação precisa continuar 100% real. Então este script:
  - usa dados reais para decidir o split por família
  - manda aug/copy-paste/sintéticos somente para TRAIN
  - opcionalmente só inclui sintéticos cuja família real caiu no TRAIN
  - copia tensores, máscaras, masks auxiliares, metas, previews e visuals

Estrutura esperada:
  <root>/<grupo>/tensors/*.npy
  <root>/<grupo>/masks/*.npy
  <root>/<grupo>/masks_vegetation/*.npy
  <root>/<grupo>/masks_cana/*.npy
  <root>/<grupo>/metas/*.json
  <root>/<grupo>/previews/*.png
  <root>/<grupo>/visuals/*.png

Uso recomendado para real + copy/paste:
---------------------------------------
python _7_split_multi_source.py ^
  --src-roots dataset/1024x640/group,dataset/copypaste/group ^
  --train-only-roots dataset/copypaste/group ^
  --train 0.7 --val 0.3 --test 0.0 ^
  --synthetic-train-only ^
  --synthetic-respect-family-split ^
  --clear-dst

Com isso:
  - dados reais vão para train/val conforme split
  - copy/paste vai apenas para train
  - copy/paste derivado de família que caiu em val/test é ignorado por padrão
"""

from __future__ import annotations

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


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

try:
    with open("config.json", "r", encoding="utf-8") as f:
        CONFIG = json.load(f)
except Exception:
    CONFIG = {}

RESOLUCAO = tuple(CONFIG.get("resolucao", [1024, 640]))
MULTI_HEAD = bool(CONFIG.get("multi_head", False))

TENSOR_EXT = ".npy"
MASK_NPY_SUFFIX = ".npy"

AUX_MASK_DIRS = [
    "masks_vegetation",
    "masks_cana",
]

COPY_OPTIONAL_DIRS = [
    "metas",
    "previews",
    "visuals",
]

IGNORE_FILENAMES = {
    "normalize_manifest.csv",
    "copy_paste_manifest.csv",
    "copy_paste_rejected.csv",
    "copy_paste_summary.json",
    "audit_summary.json",
    "audit_health.json",
}


# ============================================================
# Regex de família/origem
# ============================================================

RE_ORIGINAL_PREFIX = re.compile(r"^original_(.+)$", re.IGNORECASE)

# augmented_abc_augx_00 -> family abc
RE_AUGMENTED_FAMILY = re.compile(
    r"^augmented_(.+?)(?:_aug[a-zA-Z0-9]*_\d+)?$",
    re.IGNORECASE,
)
RE_AUG_SUFFIX = re.compile(r"_aug[a-zA-Z0-9]*_\d+$", re.IGNORECASE)

# copy-paste: abc_cp_00 -> family abc
RE_COPYPASTE_SUFFIX = re.compile(r"(.+)_cp_\d+$", re.IGNORECASE)

# grupos copy/paste: chao_cana_copypaste -> split_group chao_cana
RE_GROUP_COPYPASTE_SUFFIX = re.compile(r"(.+)_copypaste$", re.IGNORECASE)


# ============================================================
# Utilidades
# ============================================================

def garantir(p: str | Path) -> None:
    os.makedirs(str(p), exist_ok=True)


def limpar_dir(p: str | Path) -> None:
    if os.path.isdir(str(p)):
        shutil.rmtree(str(p))
    garantir(p)


def norm_path(p: str | Path) -> str:
    return os.path.normcase(os.path.abspath(str(p)))


def parse_csv_list(s: Optional[str]) -> List[str]:
    if not s:
        return []
    return [x.strip() for x in str(s).split(",") if x.strip()]


def parse_map(s: str, value_type=int) -> Dict[str, Any]:
    out: Dict[str, Any] = {}
    if not s:
        return out
    for item in s.split(","):
        item = item.strip()
        if not item:
            continue
        if ":" not in item:
            raise ValueError(f"Item de mapa inválido: {item}. Use grupo:valor")
        k, v = item.split(":", 1)
        out[k.strip()] = value_type(v.strip())
    return out


def load_json_safe(path: str | Path) -> Dict[str, Any]:
    try:
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return {}


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


def lista_grupos(root: str | Path) -> List[str]:
    root = str(root)
    if not os.path.isdir(root):
        return []

    out = []
    for g in sorted(os.listdir(root)):
        gdir = os.path.join(root, g)
        if not os.path.isdir(gdir):
            continue
        if (
            os.path.isdir(os.path.join(gdir, "tensors"))
            and os.path.isdir(os.path.join(gdir, "masks"))
        ):
            out.append(g)
    return out


def listar_tensors(tensor_dir: str | Path) -> List[str]:
    if not os.path.isdir(str(tensor_dir)):
        return []
    return sorted([
        f for f in os.listdir(str(tensor_dir))
        if f.lower().endswith(TENSOR_EXT)
    ])


def mask_npy_from_tensor_name(tensor_name: str) -> str:
    base, _ = os.path.splitext(tensor_name)
    return base + MASK_NPY_SUFFIX


def base_no_ext(filename: str) -> str:
    return os.path.splitext(filename)[0]


def normalize_group_for_split(group_name: str) -> str:
    """
    chao_cana_copypaste -> chao_cana
    outros ficam iguais.
    """
    m = RE_GROUP_COPYPASTE_SUFFIX.match(group_name)
    if m:
        return m.group(1)
    return group_name


def output_group_name(source_group: str, split_group: str, is_synthetic: bool, mode: str) -> str:
    """
    mode:
      source: mantém nome original do grupo da fonte
      base: usa grupo normalizado do split
      suffix: sintético vai para <base>_synthetic, real fica base
    """
    if mode == "source":
        return source_group
    if mode == "base":
        return split_group
    if mode == "suffix":
        if is_synthetic:
            return f"{split_group}_synthetic"
        return split_group
    return source_group


def classify_source_and_family(filename_no_ext: str, meta: Optional[Dict[str, Any]] = None) -> Tuple[str, str, bool]:
    """
    Retorna:
      source: original | augmented | copypaste | synthetic | unknown
      family: id da família para impedir vazamento
      is_synthetic: True para aug/copy/synthetic

    Regras:
      original_abc           -> original, abc
      augmented_abc_aug_00   -> augmented, abc
      abc_aug_00             -> augmented, abc
      abc_cp_00              -> copypaste, abc
      meta.synthetic == true -> synthetic/copypaste, tenta family pelo receiver/base
      abc                    -> unknown, abc, tratado como original real
    """
    meta = meta or {}

    m = RE_ORIGINAL_PREFIX.match(filename_no_ext)
    if m:
        return "original", m.group(1), False

    m = RE_AUGMENTED_FAMILY.match(filename_no_ext)
    if m:
        return "augmented", m.group(1), True

    if RE_AUG_SUFFIX.search(filename_no_ext):
        fam = RE_AUG_SUFFIX.sub("", filename_no_ext)
        return "augmented", fam, True

    m = RE_COPYPASTE_SUFFIX.match(filename_no_ext)
    if m:
        return "copypaste", m.group(1), True

    # Meta do copy/paste v1/v2.
    cp = meta.get("copy_paste_augmentation")
    if isinstance(cp, dict):
        recv = cp.get("receiver")
        if isinstance(recv, dict):
            fam = str(recv.get("base") or filename_no_ext)
            return "copypaste", fam, True
        return "copypaste", filename_no_ext, True

    if bool(meta.get("synthetic", False)):
        return "synthetic", filename_no_ext, True

    return "unknown", filename_no_ext, False


@dataclass
class Sample:
    src_root: str
    source_root_label: str
    source_root_train_only: bool

    source_group: str
    split_group: str
    output_group_default: str

    tensor_path: str
    tensor_name: str
    base: str

    mask_path: str
    source: str
    family: str
    is_synthetic: bool

    meta_path: Optional[str]
    preview_path: Optional[str]
    visual_path: Optional[str]


# ============================================================
# Coleta de amostras
# ============================================================

def collect_samples_from_root(
    src_root: str | Path,
    train_only_roots_norm: set[str],
    label: str = "",
    groups_filter: Optional[set[str]] = None,
) -> List[Sample]:
    src_root = str(src_root)
    src_root_norm = norm_path(src_root)
    root_train_only = src_root_norm in train_only_roots_norm
    root_label = label or os.path.basename(os.path.normpath(src_root)) or src_root

    samples: List[Sample] = []

    for group_name in lista_grupos(src_root):
        if groups_filter and group_name not in groups_filter and normalize_group_for_split(group_name) not in groups_filter:
            continue

        group_dir = os.path.join(src_root, group_name)
        tensor_dir = os.path.join(group_dir, "tensors")
        mask_dir = os.path.join(group_dir, "masks")
        meta_dir = os.path.join(group_dir, "metas")
        preview_dir = os.path.join(group_dir, "previews")
        visual_dir = os.path.join(group_dir, "visuals")

        split_group = normalize_group_for_split(group_name)

        for tensor_name in listar_tensors(tensor_dir):
            base = base_no_ext(tensor_name)
            mask_name = mask_npy_from_tensor_name(tensor_name)
            tensor_path = os.path.join(tensor_dir, tensor_name)
            mask_path = os.path.join(mask_dir, mask_name)

            if not os.path.exists(mask_path):
                continue

            meta_path = os.path.join(meta_dir, base + ".json")
            if not os.path.exists(meta_path):
                meta_path = None

            meta = load_json_safe(meta_path) if meta_path else {}
            source, family, is_synthetic = classify_source_and_family(base, meta)

            preview_path = None
            for ext in (".png", ".jpg", ".jpeg"):
                cand = os.path.join(preview_dir, base + ext)
                if os.path.exists(cand):
                    preview_path = cand
                    break

            visual_path = None
            for ext in (".png", ".jpg", ".jpeg"):
                cand = os.path.join(visual_dir, base + "_debug" + ext)
                if os.path.exists(cand):
                    visual_path = cand
                    break
                cand2 = os.path.join(visual_dir, base + ext)
                if os.path.exists(cand2):
                    visual_path = cand2
                    break

            samples.append(Sample(
                src_root=src_root,
                source_root_label=root_label,
                source_root_train_only=root_train_only,

                source_group=group_name,
                split_group=split_group,
                output_group_default=group_name,

                tensor_path=tensor_path,
                tensor_name=tensor_name,
                base=base,

                mask_path=mask_path,
                source=source,
                family=family,
                is_synthetic=is_synthetic,

                meta_path=meta_path,
                preview_path=preview_path,
                visual_path=visual_path,
            ))

    return samples


def collect_all_samples(
    src_roots: List[str],
    train_only_roots: List[str],
    groups: Optional[List[str]] = None,
) -> List[Sample]:
    train_only_norm = {norm_path(p) for p in train_only_roots}
    groups_filter = set(groups) if groups else None

    all_samples: List[Sample] = []
    for i, root in enumerate(src_roots):
        label = f"root{i}"
        samples = collect_samples_from_root(
            src_root=root,
            train_only_roots_norm=train_only_norm,
            label=label,
            groups_filter=groups_filter,
        )
        all_samples.extend(samples)

    return all_samples


# ============================================================
# Split por família real
# ============================================================

def allocate_counts(n: int, p_train: float, p_val: float, p_test: float, min_train: int, min_val: int, min_test: int) -> Tuple[int, int, int]:
    n_train = int(round(n * p_train))
    n_val = int(round(n * p_val))
    n_test = n - n_train - n_val

    if n_test < 0:
        excesso = -n_test
        take_train = min(excesso, max(0, n_train))
        n_train -= take_train
        excesso -= take_train
        if excesso > 0:
            take_val = min(excesso, max(0, n_val))
            n_val -= take_val
            excesso -= take_val
        n_test = 0

    min_sum = min_train + min_val + min_test

    if n >= min_sum:
        n_train = max(n_train, min_train)
        n_val = max(n_val, min_val)
        n_test = max(n_test, min_test)

        total = n_train + n_val + n_test
        while total > n:
            if n_test > min_test:
                n_test -= 1
            elif n_val > min_val:
                n_val -= 1
            elif n_train > min_train:
                n_train -= 1
            else:
                break
            total = n_train + n_val + n_test

        while total < n:
            if n_train - min_train <= n_val - min_val:
                n_train += 1
            else:
                n_val += 1
            total = n_train + n_val + n_test
    else:
        n_train = min(n, max(1, min_train))
        resto = n - n_train
        n_val = max(0, min(resto, min_val))
        n_test = max(0, resto - n_val)

    diff = n - (n_train + n_val + n_test)
    if diff != 0:
        if diff > 0:
            take = min(diff, n - n_train)
            n_train += take
            diff -= take
            if diff > 0:
                n_val += diff
        else:
            diff = -diff
            take = min(diff, n_test)
            n_test -= take
            diff -= take
            if diff > 0:
                n_val -= diff

    return n_train, n_val, n_test


def build_family_split(
    samples: List[Sample],
    p_train: float,
    p_val: float,
    p_test: float,
    seed: int,
    mins: Dict[str, int],
    caps_map: Optional[Dict[str, int]] = None,
) -> Tuple[Dict[str, Dict[str, str]], Dict[str, Dict[str, int]]]:
    """
    Usa apenas amostras reais/originais para decidir a partição das famílias.

    Retorna:
      family_split[split_group][family] = train|val|test
      group_summary[split_group] = contagens de famílias
    """
    caps_map = caps_map or {}

    families_by_group: Dict[str, set[str]] = {}

    for s in samples:
        if s.source_root_train_only:
            continue
        if s.is_synthetic:
            continue
        if s.source in ("augmented", "copypaste", "synthetic"):
            continue

        families_by_group.setdefault(s.split_group, set()).add(s.family)

    family_split: Dict[str, Dict[str, str]] = {}
    group_summary: Dict[str, Dict[str, int]] = {}

    for group_name in sorted(families_by_group.keys()):
        fams = sorted(families_by_group[group_name])
        rng = random.Random(seed)
        rng.shuffle(fams)

        total_familias = len(fams)
        if total_familias == 0:
            group_summary[group_name] = {"familias": 0, "train_families": 0, "val_families": 0, "test_families": 0}
            continue

        n_tr, n_va, n_te = allocate_counts(
            total_familias,
            p_train,
            p_val,
            p_test,
            mins["train"],
            mins["val"],
            mins["test"],
        )

        fam_train = set(fams[:n_tr])
        fam_val = set(fams[n_tr:n_tr + n_va])
        fam_test = set(fams[n_tr + n_va:n_tr + n_va + n_te])

        if group_name in caps_map:
            cap = int(caps_map[group_name])
            if len(fam_train) > cap:
                fam_list = list(fam_train)
                rng.shuffle(fam_list)
                kept = set(fam_list[:cap])
                dropped = set(fam_list[cap:])
                fam_train = kept
                print(
                    f"[{group_name}] cap-train-families={cap} → "
                    f"mantidas {len(kept)} famílias, descartadas {len(dropped)} do TRAIN"
                )

        group_map: Dict[str, str] = {}
        for f in fam_train:
            group_map[f] = "train"
        for f in fam_val:
            group_map[f] = "val"
        for f in fam_test:
            group_map[f] = "test"

        family_split[group_name] = group_map
        group_summary[group_name] = {
            "familias": total_familias,
            "train_families": len(fam_train),
            "val_families": len(fam_val),
            "test_families": len(fam_test),
        }

    return family_split, group_summary


# ============================================================
# Cópia
# ============================================================

def copy_optional_file(src_path: Optional[str], dst_dir: str, dst_base: str, suffix: str = "") -> Optional[str]:
    if not src_path or not os.path.exists(src_path):
        return None

    garantir(dst_dir)
    ext = os.path.splitext(src_path)[1]
    dst = os.path.join(dst_dir, dst_base + suffix + ext)
    shutil.copy2(src_path, dst)
    return dst


def copy_named_optional(src_dir: str, dst_dir: str, base: str, ext: str) -> Optional[str]:
    src = os.path.join(src_dir, base + ext)
    if not os.path.exists(src):
        return None
    garantir(dst_dir)
    dst = os.path.join(dst_dir, base + ext)
    shutil.copy2(src, dst)
    return dst


def copiar_sample(
    s: Sample,
    dst_root: str,
    split_name: str,
    out_group_name: str,
    copy_meta_preview: bool = True,
    copy_visuals: bool = True,
) -> Optional[Dict[str, Any]]:
    dst_group_dir = os.path.join(dst_root, split_name, "group", out_group_name)

    dst_tensor_dir = os.path.join(dst_group_dir, "tensors")
    dst_mask_dir = os.path.join(dst_group_dir, "masks")
    dst_meta_dir = os.path.join(dst_group_dir, "metas")
    dst_preview_dir = os.path.join(dst_group_dir, "previews")
    dst_visual_dir = os.path.join(dst_group_dir, "visuals")

    garantir(dst_tensor_dir)
    garantir(dst_mask_dir)

    tensor_dst = os.path.join(dst_tensor_dir, s.tensor_name)
    mask_dst = os.path.join(dst_mask_dir, os.path.basename(s.mask_path))

    if not (os.path.exists(s.tensor_path) and os.path.exists(s.mask_path)):
        return None

    shutil.copy2(s.tensor_path, tensor_dst)
    shutil.copy2(s.mask_path, mask_dst)

    base = s.base

    # mask debug PNG opcional
    mask_png_src = os.path.join(os.path.dirname(s.mask_path), base + ".png")
    mask_png_dst = None
    if os.path.exists(mask_png_src):
        mask_png_dst = os.path.join(dst_mask_dir, base + ".png")
        shutil.copy2(mask_png_src, mask_png_dst)

    # aux masks
    aux_masks: Dict[str, Dict[str, Optional[str]]] = {}
    src_group_dir = os.path.join(s.src_root, s.source_group)

    for aux_dir in AUX_MASK_DIRS:
        aux_src_dir = os.path.join(src_group_dir, aux_dir)
        aux_dst_dir = os.path.join(dst_group_dir, aux_dir)

        aux_npy_src = os.path.join(aux_src_dir, base + ".npy")
        aux_png_src = os.path.join(aux_src_dir, base + ".png")

        aux_npy_dst = None
        aux_png_dst = None

        if os.path.exists(aux_npy_src):
            garantir(aux_dst_dir)
            aux_npy_dst = os.path.join(aux_dst_dir, base + ".npy")
            shutil.copy2(aux_npy_src, aux_npy_dst)

            if os.path.exists(aux_png_src):
                aux_png_dst = os.path.join(aux_dst_dir, base + ".png")
                shutil.copy2(aux_png_src, aux_png_dst)

        aux_masks[aux_dir] = {
            "npy": aux_npy_dst,
            "png": aux_png_dst,
        }

        if MULTI_HEAD and aux_npy_dst is None:
            raise RuntimeError(
                f"multi_head=true, mas máscara auxiliar ausente: "
                f"{aux_dir}/{base}.npy em {src_group_dir}"
            )

    meta_dst = None
    preview_dst = None
    visual_dst = None

    if copy_meta_preview:
        if s.meta_path and os.path.exists(s.meta_path):
            garantir(dst_meta_dir)
            meta_dst = os.path.join(dst_meta_dir, base + ".json")
            shutil.copy2(s.meta_path, meta_dst)

        if s.preview_path and os.path.exists(s.preview_path):
            garantir(dst_preview_dir)
            preview_dst = os.path.join(dst_preview_dir, base + os.path.splitext(s.preview_path)[1])
            shutil.copy2(s.preview_path, preview_dst)

    if copy_visuals and s.visual_path and os.path.exists(s.visual_path):
        garantir(dst_visual_dir)
        visual_dst = os.path.join(dst_visual_dir, os.path.basename(s.visual_path))
        shutil.copy2(s.visual_path, visual_dst)

    return {
        "split": split_name,
        "group": out_group_name,
        "source_group": s.source_group,
        "split_group": s.split_group,
        "base": base,
        "family": s.family,
        "source": s.source,
        "is_synthetic": int(bool(s.is_synthetic)),
        "source_root": s.src_root,
        "source_root_label": s.source_root_label,
        "source_root_train_only": int(bool(s.source_root_train_only)),

        "tensor": tensor_dst,
        "mask_npy": mask_dst,
        "mask_png": mask_png_dst,

        "mask_vegetation_npy": aux_masks.get("masks_vegetation", {}).get("npy"),
        "mask_vegetation_png": aux_masks.get("masks_vegetation", {}).get("png"),

        "mask_cana_npy": aux_masks.get("masks_cana", {}).get("npy"),
        "mask_cana_png": aux_masks.get("masks_cana", {}).get("png"),

        "meta": meta_dst,
        "preview": preview_dst,
        "visual_debug": visual_dst,
    }


def decide_sample_split(
    s: Sample,
    family_split: Dict[str, Dict[str, str]],
    args,
) -> Tuple[Optional[str], str]:
    """
    Retorna (split, reason)
      split None = ignorado
    """
    group_map = family_split.get(s.split_group, {})
    assigned = group_map.get(s.family)

    is_train_only_candidate = (
        s.source_root_train_only
        or (bool(args.synthetic_train_only) and s.is_synthetic)
        or s.source in ("augmented", "copypaste", "synthetic")
    )

    if is_train_only_candidate:
        if bool(args.synthetic_respect_family_split):
            if assigned is None:
                if bool(args.allow_orphan_synthetic_train):
                    return "train", "train_only_orphan_allowed"
                return None, "skip_train_only_orphan_no_real_family"

            if assigned != "train":
                return None, f"skip_train_only_family_assigned_{assigned}"

            return "train", "train_only_family_train"

        return "train", "train_only_forced"

    # real/original
    if assigned is None:
        return None, "skip_real_no_family_assignment"

    return assigned, f"real_assigned_{assigned}"


# ============================================================
# Manifestos
# ============================================================

def write_manifest(path: str, rows: List[Dict[str, Any]]) -> None:
    garantir(os.path.dirname(path))

    fieldnames = [
        "split",
        "group",
        "source_group",
        "split_group",
        "base",
        "family",
        "source",
        "is_synthetic",
        "source_root",
        "source_root_label",
        "source_root_train_only",

        "tensor",
        "mask_npy",
        "mask_png",

        "mask_vegetation_npy",
        "mask_vegetation_png",

        "mask_cana_npy",
        "mask_cana_png",

        "meta",
        "preview",
        "visual_debug",
    ]

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


def write_skipped(path: str, rows: List[Dict[str, Any]]) -> None:
    garantir(os.path.dirname(path))
    fieldnames = [
        "source_root",
        "source_group",
        "split_group",
        "base",
        "family",
        "source",
        "is_synthetic",
        "reason",
    ]
    with open(path, "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore")
        w.writeheader()
        w.writerows(rows)


# ============================================================
# Main
# ============================================================

def main() -> None:
    ap = argparse.ArgumentParser(
        description="Split multi-source OAK-FCC-3 sem vazamento, com sintéticos/aug apenas no train."
    )

    ap.add_argument("--train", type=float, default=0.70)
    ap.add_argument("--val", type=float, default=0.29)
    ap.add_argument("--test", type=float, default=0.01)
    ap.add_argument("--seed", type=int, default=42)

    ap.add_argument("--min-train", type=int, default=1)
    ap.add_argument("--min-val", type=int, default=1)
    ap.add_argument("--min-test", type=int, default=0)

    ap.add_argument(
        "--resolucao",
        type=str,
        default=None,
        help="Sobrescreve resolução no formato WxH. Ex: 1024x640.",
    )

    ap.add_argument(
        "--src-root",
        type=str,
        default=None,
        help="Compatibilidade: uma raiz normalizada. Default: dataset/<resolucao>/group",
    )

    ap.add_argument(
        "--src-roots",
        type=str,
        default="",
        help="Lista de raízes separadas por vírgula. Ex: dataset/1024x640/group,dataset/copypaste/group",
    )

    ap.add_argument(
        "--train-only-roots",
        type=str,
        default="",
        help="Raízes que devem entrar somente no TRAIN. Ex: dataset/copypaste/group",
    )

    ap.add_argument(
        "--dst-root",
        type=str,
        default="dataset/split",
        help="Raiz do split.",
    )

    ap.add_argument(
        "--groups",
        type=str,
        default=None,
        help="Lista de grupos/split_groups separados por vírgula.",
    )

    ap.add_argument(
        "--cap-train-families",
        type=str,
        default="",
        help="Mapa 'grupo:cap,...' para limitar famílias reais no TRAIN. Ex: 'chao:50'",
    )

    ap.add_argument(
        "--clear-dst",
        action="store_true",
        help="Apaga dst-root antes de copiar.",
    )

    ap.add_argument(
        "--no-meta-preview",
        action="store_true",
        help="Não copia metas/previews para o split.",
    )

    ap.add_argument(
        "--no-visuals",
        action="store_true",
        help="Não copia pasta visuals/debug.",
    )

    ap.add_argument(
        "--synthetic-train-only",
        action="store_true",
        default=True,
        help="Força samples synthetic/augmented/copypaste para TRAIN apenas.",
    )

    ap.add_argument(
        "--allow-synthetic-val",
        dest="synthetic_train_only",
        action="store_false",
        help="Permite sintéticos no val/test. Não recomendado.",
    )

    ap.add_argument(
        "--synthetic-respect-family-split",
        action="store_true",
        default=True,
        help="Só inclui sintético no TRAIN se a família real correspondente caiu no TRAIN.",
    )

    ap.add_argument(
        "--no-synthetic-respect-family-split",
        dest="synthetic_respect_family_split",
        action="store_false",
        help="Inclui sintéticos no TRAIN mesmo sem checar a família real. Mais arriscado.",
    )

    ap.add_argument(
        "--allow-orphan-synthetic-train",
        action="store_true",
        default=False,
        help="Permite sintético no TRAIN mesmo sem família real encontrada.",
    )

    ap.add_argument(
        "--output-group-mode",
        choices=["source", "base", "suffix"],
        default="source",
        help="Nome do grupo de saída. source mantém chao_cana_copypaste; base junta em chao_cana; suffix cria <base>_synthetic.",
    )

    ap.add_argument(
        "--manifest",
        type=str,
        default="",
        help="CSV de manifesto. Default: <dst-root>/split_manifest.csv",
    )

    ap.add_argument(
        "--summary",
        type=str,
        default="",
        help="JSON de resumo. Default: <dst-root>/split_summary.json",
    )

    ap.add_argument(
        "--skipped",
        type=str,
        default="",
        help="CSV de samples ignorados. Default: <dst-root>/split_skipped.csv",
    )

    args = ap.parse_args()

    if args.resolucao:
        try:
            w, h = args.resolucao.lower().split("x")
            resolucao = (int(w), int(h))
        except Exception:
            resolucao = RESOLUCAO
    else:
        resolucao = RESOLUCAO

    default_src = os.path.join("dataset", f"{resolucao[0]}x{resolucao[1]}", "group")

    src_roots = parse_csv_list(args.src_roots)
    if args.src_root:
        src_roots.insert(0, args.src_root)
    if not src_roots:
        src_roots = [default_src]

    # Remove duplicatas preservando ordem.
    seen = set()
    src_roots_unique = []
    for r in src_roots:
        nr = norm_path(r)
        if nr not in seen:
            seen.add(nr)
            src_roots_unique.append(r)
    src_roots = src_roots_unique

    train_only_roots = parse_csv_list(args.train_only_roots)

    for r in src_roots:
        if not os.path.isdir(r):
            raise SystemExit(f"[ERRO] src-root não encontrado: {r}")

    dst_root = args.dst_root

    soma = args.train + args.val + args.test
    if soma <= 0:
        raise ValueError("Soma de proporções deve ser > 0.")

    p_train = args.train / soma
    p_val = args.val / soma
    p_test = args.test / soma

    mins = {
        "train": max(0, args.min_train),
        "val": max(0, args.min_val),
        "test": max(0, args.min_test),
    }

    groups = parse_csv_list(args.groups) if args.groups else None
    caps_map = parse_map(args.cap_train_families, int)

    if args.clear_dst:
        print(f"[INFO] Limpando destino: {dst_root}")
        limpar_dir(dst_root)
    else:
        garantir(dst_root)

    samples = collect_all_samples(
        src_roots=src_roots,
        train_only_roots=train_only_roots,
        groups=groups,
    )

    if not samples:
        print("[WARN] Nenhuma amostra encontrada.")
        return

    family_split, family_summary = build_family_split(
        samples=samples,
        p_train=p_train,
        p_val=p_val,
        p_test=p_test,
        seed=args.seed,
        mins=mins,
        caps_map=caps_map,
    )

    print("==========================================")
    print("Split OAK-FCC-3 Multi-source")
    print("SRC_ROOTS:")
    for r in src_roots:
        marker = " [TRAIN_ONLY]" if norm_path(r) in {norm_path(x) for x in train_only_roots} else ""
        print(f"  - {r}{marker}")
    print(f"DST       : {dst_root}")
    print(f"Split     : train={p_train:.3f}, val={p_val:.3f}, test={p_test:.3f}")
    print(f"Mínimos   : train={mins['train']} val={mins['val']} test={mins['test']}")
    print(f"Seed      : {args.seed}")
    print(f"Synthetic : train_only={args.synthetic_train_only} respect_family={args.synthetic_respect_family_split}")
    print("==========================================")

    print("\nFamílias reais por split_group:")
    for g, info in family_summary.items():
        print(
            f"[{g}] famílias={info['familias']} → "
            f"train={info['train_families']}, val={info['val_families']}, test={info['test_families']}"
        )

    all_rows: List[Dict[str, Any]] = []
    skipped_rows: List[Dict[str, Any]] = []

    totals = {
        "train": 0,
        "val": 0,
        "test": 0,
        "skipped": 0,
        "synthetic_train": 0,
        "real_train": 0,
        "real_val": 0,
        "real_test": 0,
    }

    by_group: Dict[str, Dict[str, int]] = {}

    for s in samples:
        split_name, reason = decide_sample_split(s, family_split, args)

        if split_name is None:
            totals["skipped"] += 1
            skipped_rows.append({
                "source_root": s.src_root,
                "source_group": s.source_group,
                "split_group": s.split_group,
                "base": s.base,
                "family": s.family,
                "source": s.source,
                "is_synthetic": int(s.is_synthetic),
                "reason": reason,
            })
            continue

        out_group = output_group_name(
            source_group=s.source_group,
            split_group=s.split_group,
            is_synthetic=s.is_synthetic,
            mode=args.output_group_mode,
        )

        row = copiar_sample(
            s=s,
            dst_root=dst_root,
            split_name=split_name,
            out_group_name=out_group,
            copy_meta_preview=not args.no_meta_preview,
            copy_visuals=not args.no_visuals,
        )

        if row is None:
            totals["skipped"] += 1
            skipped_rows.append({
                "source_root": s.src_root,
                "source_group": s.source_group,
                "split_group": s.split_group,
                "base": s.base,
                "family": s.family,
                "source": s.source,
                "is_synthetic": int(s.is_synthetic),
                "reason": "copy_failed_missing_tensor_or_mask",
            })
            continue

        all_rows.append(row)

        totals[split_name] += 1
        if split_name == "train" and s.is_synthetic:
            totals["synthetic_train"] += 1
        elif split_name == "train":
            totals["real_train"] += 1
        elif split_name == "val":
            totals["real_val"] += 1
        elif split_name == "test":
            totals["real_test"] += 1

        gsum = by_group.setdefault(out_group, {"train": 0, "val": 0, "test": 0, "synthetic_train": 0, "skipped": 0})
        gsum[split_name] = gsum.get(split_name, 0) + 1
        if split_name == "train" and s.is_synthetic:
            gsum["synthetic_train"] += 1

    manifest_path = args.manifest or os.path.join(dst_root, "split_manifest.csv")
    summary_path = args.summary or os.path.join(dst_root, "split_summary.json")
    skipped_path = args.skipped or os.path.join(dst_root, "split_skipped.csv")

    write_manifest(manifest_path, all_rows)
    write_skipped(skipped_path, skipped_rows)

    summary = {
        "schema": "oak_fcc3_multi_source_split_v1",
        "src_roots": src_roots,
        "train_only_roots": train_only_roots,
        "dst_root": dst_root,
        "resolution": list(resolucao),
        "proportions": {
            "train": p_train,
            "val": p_val,
            "test": p_test,
        },
        "mins": mins,
        "seed": args.seed,
        "multi_head": MULTI_HEAD,
        "options": {
            "synthetic_train_only": args.synthetic_train_only,
            "synthetic_respect_family_split": args.synthetic_respect_family_split,
            "allow_orphan_synthetic_train": args.allow_orphan_synthetic_train,
            "output_group_mode": args.output_group_mode,
        },
        "family_summary": family_summary,
        "groups": by_group,
        "total": totals,
        "samples_seen": len(samples),
        "manifest": manifest_path,
        "skipped": skipped_path,
    }

    save_json(summary_path, summary)

    print("\nResumo global:")
    print(f"  train total:       {totals['train']}")
    print(f"    real train:      {totals['real_train']}")
    print(f"    synthetic train: {totals['synthetic_train']}")
    print(f"  val real:          {totals['real_val']}")
    print(f"  test real:         {totals['real_test']}")
    print(f"  skipped:           {totals['skipped']}")
    print(f"\nManifest: {manifest_path}")
    print(f"Skipped : {skipped_path}")
    print(f"Summary : {summary_path}")

    print("\n✅ Split multi-source sem vazamento concluído!")


if __name__ == "__main__":
    main()