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

import os
import re
import csv
import json
import shutil
import random
import argparse
from pathlib import Path


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

RESOLUCAO = tuple(config.get("resolucao"))

TENSOR_EXT = ".npy"
MASK_NPY_SUFFIX = ".npy"
AUX_MASK_DIRS = [
    "masks_vegetation",
    "masks_cana",
]

RE_ORIGINAL_PREFIX = re.compile(r"^original_(.+)$", re.IGNORECASE)
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)
MULTI_HEAD = bool(config.get("multi_head", False))


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


def limpar_dir(p):
    if os.path.isdir(p):
        shutil.rmtree(p)
    garantir(p)


def lista_grupos(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):
    if not os.path.isdir(tensor_dir):
        return []

    return sorted([
        f for f in os.listdir(tensor_dir)
        if f.lower().endswith(TENSOR_EXT)
    ])


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


def classify_source_and_family(filename_no_ext):
    """
    Mantém compatibilidade futura com augmentation.

    Exemplos:
      original_abc        -> source=original, family=abc
      augmented_abc_aug_00 -> source=augmented, family=abc
      abc_aug_00          -> source=augmented, family=abc
      abc                 -> source=unknown, family=abc

    No caso atual, sem aug, source=unknown é tratado como original.
    """
    m = RE_ORIGINAL_PREFIX.match(filename_no_ext)
    if m:
        return "original", m.group(1)

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

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

    return "unknown", filename_no_ext


def build_family_index(tensor_dir, mask_dir):
    """
    family -> {
      "original": tensor_name or None,
      "augmented": [tensor_name, ...],
      "all": [...]
    }

    Só indexa se houver mask .npy correspondente.
    """
    familias = {}
    tensors = listar_tensors(tensor_dir)

    for tensor_name in tensors:
        base_no_ext, _ = os.path.splitext(tensor_name)
        mask_npy_name = mask_npy_from_tensor_name(tensor_name)

        if not os.path.exists(os.path.join(mask_dir, mask_npy_name)):
            continue

        source, fam = classify_source_and_family(base_no_ext)

        d = familias.setdefault(
            fam,
            {
                "original": None,
                "augmented": [],
                "all": [],
            }
        )

        d["all"].append(tensor_name)

        if source == "original":
            d["original"] = tensor_name

        elif source == "augmented":
            d["augmented"].append(tensor_name)

        else:
            # Sem prefixo: no fluxo atual, é original.
            if d["original"] is None:
                d["original"] = tensor_name
            else:
                d["augmented"].append(tensor_name)

    return familias


def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
    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 copiar_optional(src_dir, dst_dir, base, ext):
    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_mask_dir_optional(src_group_dir, dst_group_dir, mask_dir_name, base):
    """
    Copia uma pasta auxiliar de máscara, como:
      masks_vegetation/
      masks_cana/

    Copia:
      <base>.npy obrigatório se existir
      <base>.png opcional se existir

    Retorna caminhos de destino ou None.
    """
    src_dir = os.path.join(src_group_dir, mask_dir_name)
    dst_dir = os.path.join(dst_group_dir, mask_dir_name)

    npy_src = os.path.join(src_dir, base + ".npy")
    png_src = os.path.join(src_dir, base + ".png")

    result = {
        "npy": None,
        "png": None,
    }

    if not os.path.exists(npy_src):
        return result

    garantir(dst_dir)

    npy_dst = os.path.join(dst_dir, base + ".npy")
    shutil.copy2(npy_src, npy_dst)
    result["npy"] = npy_dst

    if os.path.exists(png_src):
        png_dst = os.path.join(dst_dir, base + ".png")
        shutil.copy2(png_src, png_dst)
        result["png"] = png_dst

    return result


def copiar(
    nomes,
    src_group_dir,
    dst_group_dir,
    copy_meta_preview=True,
):
    """
    Copia:
      - tensor .npy obrigatório
      - mask .npy obrigatória
      - mask .png opcional
      - meta .json opcional
      - preview .png/.jpg/.jpeg opcional
    """
    src_tensor_dir = os.path.join(src_group_dir, "tensors")
    src_mask_dir = os.path.join(src_group_dir, "masks")
    src_meta_dir = os.path.join(src_group_dir, "metas")
    src_preview_dir = os.path.join(src_group_dir, "previews")

    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")

    garantir(dst_tensor_dir)
    garantir(dst_mask_dir)

    rows = []
    moved = 0

    for nome in nomes:
        base, _ = os.path.splitext(nome)

        tensor_src = os.path.join(src_tensor_dir, nome)
        mask_npy_name = mask_npy_from_tensor_name(nome)
        mask_npy_src = os.path.join(src_mask_dir, mask_npy_name)

        if not (os.path.exists(tensor_src) and os.path.exists(mask_npy_src)):
            continue

        tensor_dst = os.path.join(dst_tensor_dir, nome)
        mask_npy_dst = os.path.join(dst_mask_dir, mask_npy_name)

        shutil.copy2(tensor_src, tensor_dst)
        shutil.copy2(mask_npy_src, mask_npy_dst)

        mask_png_dst = None
        meta_dst = None
        preview_dst = None

        # Debug visual da máscara
        mask_png_src = os.path.join(src_mask_dir, base + ".png")
        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)

        if copy_meta_preview:
            meta_dst = copiar_optional(src_meta_dir, dst_meta_dir, base, ".json")

            for ext in (".png", ".jpg", ".jpeg"):
                cand = os.path.join(src_preview_dir, base + ext)
                if os.path.exists(cand):
                    garantir(dst_preview_dir)
                    preview_dst = os.path.join(dst_preview_dir, base + ext)
                    shutil.copy2(cand, preview_dst)
                    break

        aux_masks = {}
        for aux_dir in AUX_MASK_DIRS:
            aux_masks[aux_dir] = copiar_mask_dir_optional(
                src_group_dir=src_group_dir,
                dst_group_dir=dst_group_dir,
                mask_dir_name=aux_dir,
                base=base,
            )

        if MULTI_HEAD:
            for aux_dir in AUX_MASK_DIRS:
                aux_npy = aux_masks.get(aux_dir, {}).get("npy")
                if aux_npy is None:
                    raise RuntimeError(
                        f"multi_head=true, mas máscara auxiliar ausente: "
                        f"{aux_dir}/{base}.npy em {src_group_dir}"
                    )

        rows.append({
            "base": base,
            "tensor": tensor_dst,

            "mask_npy": mask_npy_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,
        })

        moved += 1

    return moved, rows


def split_group(
    group_name,
    src_root,
    dst_root,
    p_train,
    p_val,
    p_test,
    seed,
    mins,
    caps_map=None,
    copy_meta_preview=True,
):
    src_group_dir = os.path.join(src_root, group_name)
    src_tensor_dir = os.path.join(src_group_dir, "tensors")
    src_mask_dir = os.path.join(src_group_dir, "masks")

    familias = build_family_index(src_tensor_dir, src_mask_dir)
    familias_originais = [fam for fam, d in familias.items() if d["original"] is not None]
    total_familias = len(familias_originais)

    if total_familias == 0:
        print(f"[{group_name}] 0 famílias com original, pulando.")
        return {
            "train": 0,
            "val": 0,
            "test": 0,
            "familias": 0,
            "rows": [],
        }

    rng = random.Random(seed)
    rng.shuffle(familias_originais)

    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(familias_originais[:n_tr])
    fam_val = set(familias_originais[n_tr:n_tr + n_va])
    fam_test = set(familias_originais[n_tr + n_va:n_tr + n_va + n_te])

    if caps_map and group_name in caps_map:
        cap = 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"
            )

    nomes_train, nomes_val, nomes_test = [], [], []

    for fam, d in familias.items():
        if fam in fam_train:
            if d["original"]:
                nomes_train.append(d["original"])
            if d["augmented"]:
                nomes_train.extend(d["augmented"])

        elif fam in fam_val:
            if d["original"]:
                nomes_val.append(d["original"])

        elif fam in fam_test:
            if d["original"]:
                nomes_test.append(d["original"])

    rows_all = []

    split_defs = [
        ("train", nomes_train),
        ("val", nomes_val),
        ("test", nomes_test),
    ]

    counts = {}

    for split_name, nomes in split_defs:
        dst_group_dir = os.path.join(dst_root, split_name, "group", group_name)

        moved, rows = copiar(
            nomes=nomes,
            src_group_dir=src_group_dir,
            dst_group_dir=dst_group_dir,
            copy_meta_preview=copy_meta_preview,
        )

        counts[split_name] = moved

        for r in rows:
            r["split"] = split_name
            r["group"] = group_name
            rows_all.append(r)

    print(
        f"[{group_name}] famílias={total_familias} → "
        f"train={counts['train']}, val={counts['val']}, test={counts['test']}"
    )

    return {
        "train": counts["train"],
        "val": counts["val"],
        "test": counts["test"],
        "familias": total_familias,
        "rows": rows_all,
    }


def write_manifest(path, rows):
    garantir(os.path.dirname(path))

    fieldnames = [
        "split",
        "group",
        "base",
        "tensor",

        "mask_npy",
        "mask_png",

        "mask_vegetation_npy",
        "mask_vegetation_png",

        "mask_cana_npy",
        "mask_cana_png",

        "meta",
        "preview",
    ]

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


def write_summary(path, summary):
    garantir(os.path.dirname(path))

    with open(path, "w", encoding="utf-8") as f:
        json.dump(summary, f, ensure_ascii=False, indent=2)


def parse_cap_map(s):
    caps = {}

    if not s:
        return caps

    for item in s.split(","):
        if not item.strip():
            continue

        k, v = item.strip().split(":")
        caps[k.strip()] = int(v)

    return caps


def main():
    ap = argparse.ArgumentParser(
        description="Split estratificado por grupo sem vazamento para tensors/masks OAK-FCC-3."
    )

    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: 512x512.",
    )

    ap.add_argument(
        "--src-root",
        type=str,
        default=None,
        help="Raiz normalizada. Default: dataset/<resolucao>/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 separados por vírgula.",
    )

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

    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(
        "--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",
    )

    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

    src_root = args.src_root or os.path.join("dataset", f"{resolucao[0]}x{resolucao[1]}", "group")
    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),
    }

    caps_map = parse_cap_map(args.cap_train_families)

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

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

    grupos = lista_grupos(src_root)

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

    if not grupos:
        print(f"[WARN] Nenhum grupo encontrado em: {src_root}")
        return

    print("==========================================")
    print("Split OAK-FCC-3")
    print(f"SRC     : {src_root}")
    print(f"DST     : {dst_root}")
    print(f"Grupos  : {', '.join(grupos)}")
    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("==========================================")

    total_global = {
        "train": 0,
        "val": 0,
        "test": 0,
        "familias": 0,
    }

    all_rows = []
    summary_groups = {}

    for g in grupos:
        res = split_group(
            group_name=g,
            src_root=src_root,
            dst_root=dst_root,
            p_train=p_train,
            p_val=p_val,
            p_test=p_test,
            seed=args.seed,
            mins=mins,
            caps_map=caps_map,
            copy_meta_preview=not args.no_meta_preview,
        )

        for k in total_global.keys():
            total_global[k] += res.get(k, 0)

        all_rows.extend(res.get("rows", []))

        summary_groups[g] = {
            "train": res.get("train", 0),
            "val": res.get("val", 0),
            "test": res.get("test", 0),
            "familias": res.get("familias", 0),
        }

    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")

    write_manifest(manifest_path, all_rows)

    summary = {
        "src_root": src_root,
        "dst_root": dst_root,
        "resolution": list(resolucao),
        "proportions": {
            "train": p_train,
            "val": p_val,
            "test": p_test,
        },
        "mins": mins,
        "seed": args.seed,
        "groups": summary_groups,
        "total": total_global,
    }

    write_summary(summary_path, summary)

    print("\nResumo global:")
    print(f"  train: {total_global['train']}")
    print(f"  val:   {total_global['val']}")
    print(f"  test:  {total_global['test']}")
    print(f"  famílias: {total_global['familias']}")
    print(f"\nManifest: {manifest_path}")
    print(f"Summary : {summary_path}")
    print("\n✅ Split sem vazamento concluído!")


if __name__ == "__main__":
    main()