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

"""
_6_normalize_oak_fcc3_originals.py

Converte dataset/originals/group para tensores finais de treino usando
RawProcessorCore + module_params.json.

Entrada:

dataset/originals/group/<grupo>/
  bins/
    <base>_CAM_A.bin
    <base>_CAM_B.bin
    <base>_CAM_C.bin
  metas/
    <base>.json
  masks/
    <base>.png
  previews/
    <base>.png

Saída:

dataset/<WxH>/group/<grupo>/
  tensors/
    <base>.npy      # float32 CHW [R,G,B,RE,NIR]
  masks/
    <base>.npy      # uint8/uint16 HW com IDs de classe
    <base>.png      # debug visual dos IDs
  metas/
    <base>.json     # meta do tensor normalizado
  previews/
    <base>.png      # preview RGB do tensor final

Também salva:
  backup/<modelo>/<model_name>/<stats_source_tag>/norm_stats.json

Uso:

python _6_normalize_oak_fcc3_originals.py ^
  --src-root dataset/originals/group ^
  --out-root dataset ^
  --module-params calibration/module_params.json ^
  --clear-dst
"""

from __future__ import annotations

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

import cv2
import numpy as np

from core.raw_processor_core import RawProcessorCore

from helpers import carregar_labelmap_completo, converter_mask_rgb_para_ids, _infer_ignore_id


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

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

DEFAULT_RES = tuple(config.get("resolucao", [512, 512]))      # (W, H)
DEFAULT_RAW_SIZE = tuple(config.get("raw_size", [1280, 800])) # (W, H)
DEFAULT_MODULE_PARAMS = config.get("module_params_json")
DEFAULT_DATASET_BASE = "dataset"


CHANNEL_NAMES = ["R", "G", "B", "RE", "NIR"]


_CORE_CACHE: Dict[tuple, RawProcessorCore] = {}

def get_or_create_core(
    raw_size: Tuple[int, int],
    bayer_pattern: str,
    module_params_path: Optional[str],
) -> RawProcessorCore:
    raw_w, raw_h = raw_size
    bayer = str(bayer_pattern or "BGGR").upper()

    key = (
        int(raw_w),
        int(raw_h),
        bayer,
        str(Path(module_params_path).resolve()) if module_params_path else None,
    )

    core = _CORE_CACHE.get(key)
    if core is not None:
        return core

    print(
        "[CORE_CACHE] criando RawProcessorCore "
        f"raw={raw_w}x{raw_h} bayer={bayer} module={module_params_path}"
    )

    core = RawProcessorCore(
        sensor_width=int(raw_w),
        sensor_height=int(raw_h),
        bayer_pattern=bayer,
        calibration_json_path=module_params_path,
    )

    _CORE_CACHE[key] = core
    return core


@dataclass
class SampleBundle:
    group: str
    base: str
    meta_path: Path
    mask_path: Path
    preview_path: Optional[Path]
    bin_paths: Dict[str, Path]


# ============================================================
# Helpers gerais
# ============================================================

def ensure_dir(path: Path):
    path.mkdir(parents=True, exist_ok=True)


def maybe_clear_dir(path: Path):
    if path.exists():
        shutil.rmtree(path)
    ensure_dir(path)


def load_json(path: Path) -> dict:
    with path.open("r", encoding="utf-8") as f:
        return json.load(f)


def save_json(path: Path, data: dict):
    ensure_dir(path.parent)
    with path.open("w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


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


def list_groups(src_root: Path) -> List[str]:
    if not src_root.is_dir():
        return []

    groups = []
    for p in sorted(src_root.iterdir()):
        if not p.is_dir():
            continue

        if (p / "bins").is_dir() and (p / "metas").is_dir() and (p / "masks").is_dir():
            groups.append(p.name)

    return groups


def find_preview(previews_dir: Path, base: str) -> Optional[Path]:
    for ext in (".png", ".jpg", ".jpeg"):
        p = previews_dir / f"{base}{ext}"
        if p.exists():
            return p
    return None


def find_sample_bins(bins_dir: Path, base: str) -> Dict[str, Path]:
    out = {}

    for cam_id in ("CAM_A", "CAM_B", "CAM_C"):
        p = bins_dir / f"{base}_{cam_id}.bin"
        if p.exists():
            out[cam_id] = p

    return out


def collect_samples_from_group(group_dir: Path) -> List[SampleBundle]:
    group = group_dir.name

    metas_dir = group_dir / "metas"
    masks_dir = group_dir / "masks"
    previews_dir = group_dir / "previews"
    bins_dir = group_dir / "bins"

    samples = []

    if not metas_dir.is_dir():
        return samples

    for meta_path in sorted(metas_dir.glob("*.json")):
        base = meta_path.stem

        mask_path = masks_dir / f"{base}.png"
        if not mask_path.exists():
            print(f"[WARN] [{group}] sem mask para {base}. Pulando.")
            continue

        bin_paths = find_sample_bins(bins_dir, base)
        if not all(cam in bin_paths for cam in ("CAM_A", "CAM_B", "CAM_C")):
            print(f"[WARN] [{group}] bins incompletos para {base}: {list(bin_paths.keys())}. Pulando.")
            continue

        samples.append(
            SampleBundle(
                group=group,
                base=base,
                meta_path=meta_path,
                mask_path=mask_path,
                preview_path=find_preview(previews_dir, base),
                bin_paths=bin_paths,
            )
        )

    return samples


def copy_json_safe(obj):
    try:
        return json.loads(json.dumps(obj, default=str))
    except Exception:
        return None


# ============================================================
# Module params / tensor
# ============================================================

def resolve_module_params_path(meta: dict, meta_path: Path, cli_module_params: Optional[str]) -> Optional[str]:
    candidates = []

    if cli_module_params:
        candidates.append(cli_module_params)

    if meta.get("camera_params_json"):
        candidates.append(meta.get("camera_params_json"))

    if DEFAULT_MODULE_PARAMS:
        candidates.append(DEFAULT_MODULE_PARAMS)

    candidates.append("calibration/module_params.json")

    for c in candidates:
        if not c:
            continue

        p = Path(str(c))

        if p.is_file():
            return str(p)

        # relativo ao diretório atual
        p2 = Path.cwd() / p
        if p2.is_file():
            return str(p2)

        # relativo à pasta do meta
        p3 = meta_path.parent / p
        if p3.is_file():
            return str(p3)

    print("[WARN] module_params.json não encontrado. RawProcessorCore vai usar defaults.")
    return None


def load_frame_from_saved_bins(sample: SampleBundle, meta: dict) -> Dict[str, np.ndarray]:
    """
    Monta frame no contrato do RawProcessorCore.decode_stream_cameras:

      frame = {
        "CAM_A": ndarray RAW10 packed 2D,
        "CAM_B": ndarray RAW10 packed 2D,
        "CAM_C": ndarray RAW10 packed 2D,
      }

    Usa saved_payload_dtypes/saved_payload_shapes do JSON.
    """
    saved_dtypes = meta.get("saved_payload_dtypes", {}) or {}
    saved_shapes = meta.get("saved_payload_shapes", {}) or {}

    frame = {}

    for cam_id, bin_path in sample.bin_paths.items():
        dtype = saved_dtypes.get(cam_id)
        shape = saved_shapes.get(cam_id)

        if dtype is None or shape is None:
            raise RuntimeError(f"Faltam saved_payload_dtypes/shapes para {cam_id} em {sample.base}")

        arr = np.fromfile(str(bin_path), dtype=np.dtype(dtype)).reshape(tuple(shape))
        frame[cam_id] = arr

    return frame


def build_processing_meta(meta: dict) -> dict:
    """
    O RawProcessorCore precisa de:
      frame_type = RAW_BRUTO
      camera_info/camera_frames por CAM_A/B/C
      actual_camera_controls para radiometric_normalization
    """
    stream_meta = dict(meta.get("stream_meta", {}) or {})

    # Garante o frame_type esperado pelo core.
    stream_meta["frame_type"] = "RAW_BRUTO"

    # Em algumas capturas o camera_info está fora do stream_meta.
    if "camera_info" not in stream_meta and isinstance(meta.get("camera_info"), dict):
        stream_meta["camera_info"] = meta.get("camera_info")

    if meta.get("actual_camera_controls") is not None:
        stream_meta["actual_camera_controls"] = meta.get("actual_camera_controls")

    if meta.get("startup_camera_controls") is not None:
        stream_meta["startup_camera_controls"] = meta.get("startup_camera_controls")

    if meta.get("camera_controls") is not None:
        stream_meta["camera_controls"] = meta.get("camera_controls")

    return stream_meta


def build_tensor_from_sample(
    sample: SampleBundle,
    meta: dict,
    res: Tuple[int, int],
    raw_size: Tuple[int, int],
    module_params_path: Optional[str],
) -> Tuple[np.ndarray, dict]:
    bayer = str(meta.get("bayer_pattern") or meta.get("bayer") or "BGGR").upper()

    core = get_or_create_core(
        raw_size=raw_size,
        bayer_pattern=bayer,
        module_params_path=module_params_path,
    )

    frame = load_frame_from_saved_bins(sample, meta)
    processing_meta = build_processing_meta(meta)

    tensor = core.build_infer_tensor_from_stream(
        frame=frame,
        meta=processing_meta,
        channels_expected=5,
        target_size=res,
    )

    tensor = np.ascontiguousarray(tensor.astype(np.float32, copy=False))

    info = {
        "module_params": module_params_path,
        "bayer_pattern": bayer,
        "fusion_result": copy_json_safe(getattr(core, "last_fusion_result", None)),
        "patch_normalization_result": copy_json_safe(getattr(core, "last_patch_normalization_result", None)),
        "frame_quality": copy_json_safe(getattr(core, "last_frame_quality_result", None)),
    }

    return tensor, info



# ============================================================
# Máscara / preview
# ============================================================

def load_mask_ids_aligned_to_tensor(
    mask_path: Path,
    cor_para_id: dict,
    ignore_id: int,
    res: Tuple[int, int],
    fusion_result: dict | None,
) -> np.ndarray:
    bgr = cv2.imread(str(mask_path), cv2.IMREAD_COLOR)

    if bgr is None:
        raise RuntimeError(f"Falha ao abrir mask: {mask_path}")

    rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
    ids = converter_mask_rgb_para_ids(rgb, cor_para_id, ignore_id)

    # A mask foi feita no espaço RGB/CAM_A.
    # Então primeiro garantimos que ela está no mesmo tamanho da referência RGB.
    if isinstance(fusion_result, dict):
        ref_shape = fusion_result.get("ref_shape")

        if isinstance(ref_shape, list) and len(ref_shape) == 2:
            ref_h, ref_w = int(ref_shape[0]), int(ref_shape[1])

            if ids.shape[:2] != (ref_h, ref_w):
                ids = cv2.resize(
                    ids,
                    (ref_w, ref_h),
                    interpolation=cv2.INTER_NEAREST,
                )

        crop_box = fusion_result.get("crop_box")

        if crop_box is not None:
            x0, y0, x1, y1 = [int(v) for v in crop_box]

            h, w = ids.shape[:2]
            x0 = max(0, min(w - 1, x0))
            x1 = max(x0 + 1, min(w, x1))
            y0 = max(0, min(h - 1, y0))
            y1 = max(y0 + 1, min(h, y1))

            ids = ids[y0:y1, x0:x1]

    # Por fim, leva para a resolução final do tensor.
    ids_res = cv2.resize(
        ids,
        res,
        interpolation=cv2.INTER_NEAREST,
    )

    return ids_res


def tensor_to_preview_bgr(tensor: np.ndarray) -> np.ndarray:
    rgb = np.transpose(tensor[:3].astype(np.float32), (1, 2, 0))
    rgb_u8 = np.clip(rgb * 255.0, 0, 255).astype(np.uint8)
    return cv2.cvtColor(rgb_u8, cv2.COLOR_RGB2BGR)


def save_bgr(path: Path, bgr: np.ndarray):
    ensure_dir(path.parent)
    ok = cv2.imwrite(str(path), bgr)
    if not ok:
        raise RuntimeError(f"Falha ao salvar preview: {path}")


def get_class_ids_from_config() -> dict:
    heads = config.get("heads", {}) or {}
    semantic = heads.get("semantic", {}) or {}
    classes = semantic.get("classes", {}) or {}

    return {
        "chao": int(classes.get("chao", 0)),
        "cana": int(classes.get("cana", 1)),
        "erva": int(classes.get("erva", 2)),
    }


def build_head_masks_from_semantic(
    mask_ids: np.ndarray,
    class_ids: dict,
    ignore_id: int = 255,
) -> dict:
    """
    Gera máscaras auxiliares para multi-head a partir da máscara semântica alinhada.

    Entrada:
      mask_ids: HW com IDs semânticos.

    Saída:
      {
        "semantic": HW uint8,
        "vegetation": HW uint8 com 0/1/ignore,
        "cana": HW uint8 com 0/1/ignore
      }
    """
    chao_id = int(class_ids.get("chao", 0))
    cana_id = int(class_ids.get("cana", 1))
    erva_id = int(class_ids.get("erva", 2))

    semantic = mask_ids.astype(np.uint8, copy=True)

    ignore_mask = semantic == ignore_id

    vegetation = np.zeros_like(semantic, dtype=np.uint8)
    vegetation[(semantic == cana_id) | (semantic == erva_id)] = 1
    vegetation[ignore_mask] = ignore_id

    cana = np.zeros_like(semantic, dtype=np.uint8)
    cana[semantic == cana_id] = 1
    cana[ignore_mask] = ignore_id

    return {
        "semantic": semantic,
        "vegetation": vegetation,
        "cana": cana,
    }


def save_mask_npy_and_debug(mask: np.ndarray, npy_path: Path, png_path: Path):
    ensure_dir(npy_path.parent)
    ensure_dir(png_path.parent)

    mask = mask.astype(np.uint8, copy=False)
    np.save(str(npy_path), mask)

    # Debug visual:
    # 0 -> preto
    # 1 -> cinza claro
    # 2 -> mais claro, quando existir na semântica
    # 255 -> branco
    debug = mask.copy()
    debug_vis = np.zeros_like(debug, dtype=np.uint8)

    debug_vis[debug == 0] = 0
    debug_vis[debug == 1] = 120
    debug_vis[debug == 2] = 220
    debug_vis[debug == 255] = 255

    cv2.imwrite(str(png_path), debug_vis)


# ============================================================
# Stats
# ============================================================

class RunningStats:
    def __init__(self):
        self.sum = None
        self.sumsq = None
        self.pixels = 0

    def update(self, tensor: np.ndarray):
        c, h, w = tensor.shape

        if self.sum is None:
            self.sum = np.zeros(c, dtype=np.float64)
            self.sumsq = np.zeros(c, dtype=np.float64)

        flat = tensor.reshape(c, -1).astype(np.float64)
        self.sum += flat.sum(axis=1)
        self.sumsq += (flat ** 2).sum(axis=1)
        self.pixels += h * w

    def result(self, channel_names: List[str]):
        if self.sum is None or self.pixels <= 0:
            return None

        mean = self.sum / self.pixels
        var = (self.sumsq / self.pixels) - mean ** 2
        std = np.sqrt(np.maximum(var, 1e-6))

        return {
            "channels": channel_names,
            "mean": mean.tolist(),
            "std": std.tolist(),
            "pixels": int(self.pixels),
        }


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

def process_group(
    group_name: str,
    src_root: Path,
    output_root: Path,
    res: Tuple[int, int],
    raw_size: Tuple[int, int],
    module_params_arg: Optional[str],
    cor_para_id: dict,
    ignore_id: int,
    stats: RunningStats,
    dataset_root: Path,
    skip_bad_quality: bool,
) -> List[dict]:
    group_dir = src_root / group_name
    samples = collect_samples_from_group(group_dir)

    out_group = output_root / group_name
    out_tensors = out_group / "tensors"

    out_masks = out_group / "masks"
    out_masks_vegetation = out_group / "masks_vegetation"
    out_masks_cana = out_group / "masks_cana"

    out_metas = out_group / "metas"
    out_previews = out_group / "previews"

    for d in (
        out_tensors,
        out_masks,
        out_masks_vegetation,
        out_masks_cana,
        out_metas,
        out_previews,
    ):
        ensure_dir(d)

    rows = []
    errors = 0
    skipped_quality = 0

    print(f"\n[GRUPO] {group_name} | samples={len(samples)}")

    for sample in samples:
        try:
            meta = load_json(sample.meta_path)
            module_params_path = resolve_module_params_path(meta, sample.meta_path, module_params_arg)

            tensor, processing_info = build_tensor_from_sample(
                sample=sample,
                meta=meta,
                res=res,
                raw_size=raw_size,
                module_params_path=module_params_path,
            )

            frame_quality = processing_info.get("frame_quality") or {}
            if skip_bad_quality and isinstance(frame_quality, dict):
                if frame_quality.get("usable_for_training") is False:
                    skipped_quality += 1
                    print(f"[SKIP-QUALITY] {sample.base}: {frame_quality.get('reasons')}")
                    continue

            mask_ids = load_mask_ids_aligned_to_tensor(
                mask_path=sample.mask_path,
                cor_para_id=cor_para_id,
                ignore_id=ignore_id,
                res=res,
                fusion_result=processing_info.get("fusion_result"),
            )

            if mask_ids.shape[:2] != tensor.shape[1:]:
                raise RuntimeError(
                    f"Shape mask/tensor incompatível: mask={mask_ids.shape}, tensor={tensor.shape}"
                )

            # Salva tensor e mask
            tensor_path = out_tensors / f"{sample.base}.npy"
            preview_path = out_previews / f"{sample.base}.png"
            meta_out_path = out_metas / f"{sample.base}.json"
            np.save(str(tensor_path), tensor)

            class_ids = get_class_ids_from_config()
            head_masks = build_head_masks_from_semantic(
                mask_ids=mask_ids,
                class_ids=class_ids,
                ignore_id=ignore_id,
            )

            mask_path = out_masks / f"{sample.base}.npy"
            mask_debug_path = out_masks / f"{sample.base}.png"

            mask_vegetation_path = out_masks_vegetation / f"{sample.base}.npy"
            mask_vegetation_debug_path = out_masks_vegetation / f"{sample.base}.png"

            mask_cana_path = out_masks_cana / f"{sample.base}.npy"
            mask_cana_debug_path = out_masks_cana / f"{sample.base}.png"

            save_mask_npy_and_debug(
                head_masks["semantic"],
                mask_path,
                mask_debug_path,
            )

            save_mask_npy_and_debug(
                head_masks["vegetation"],
                mask_vegetation_path,
                mask_vegetation_debug_path,
            )

            save_mask_npy_and_debug(
                head_masks["cana"],
                mask_cana_path,
                mask_cana_debug_path,
            )

            preview_bgr = tensor_to_preview_bgr(tensor)
            save_bgr(preview_path, preview_bgr)

            out_meta = {
                "schema": "oak_fcc3_normalized_tensor_v1",
                "source_group": sample.group,
                "source_base": sample.base,
                "source_meta": safe_rel(sample.meta_path, dataset_root),
                "source_bins": {k: safe_rel(v, dataset_root) for k, v in sample.bin_paths.items()},
                "source_mask": safe_rel(sample.mask_path, dataset_root),
                "source_preview": safe_rel(sample.preview_path, dataset_root) if sample.preview_path else None,

                "frame_type": "MULTISPEC",
                "saved_payload_type": "tensor_npy",
                "saved_tensor_path": safe_rel(tensor_path, dataset_root),
                "saved_mask_path": safe_rel(mask_path, dataset_root),
                "head_masks": {
                    "semantic": {
                        "path": safe_rel(mask_path, dataset_root),
                        "shape": list(head_masks["semantic"].shape),
                        "classes": {
                            "chao": int(class_ids["chao"]),
                            "cana": int(class_ids["cana"]),
                            "erva": int(class_ids["erva"]),
                        },
                        "ignore_index": int(ignore_id),
                    },
                    "vegetation": {
                        "path": safe_rel(mask_vegetation_path, dataset_root),
                        "shape": list(head_masks["vegetation"].shape),
                        "classes": {
                            "background": 0,
                            "vegetation": 1,
                        },
                        "positive_from_semantic": ["cana", "erva"],
                        "ignore_index": int(ignore_id),
                    },
                    "cana": {
                        "path": safe_rel(mask_cana_path, dataset_root),
                        "shape": list(head_masks["cana"].shape),
                        "classes": {
                            "not_cana": 0,
                            "cana": 1,
                        },
                        "positive_from_semantic": ["cana"],
                        "ignore_index": int(ignore_id),
                    },
                },
                "saved_preview_path": safe_rel(preview_path, dataset_root),
                "saved_payload_dtype": str(tensor.dtype),
                "saved_payload_shape": list(tensor.shape),
                "mask_shape": list(mask_ids.shape),
                "channels": CHANNEL_NAMES,

                "resolution": {
                    "width": int(res[0]),
                    "height": int(res[1]),
                },
                "processing": processing_info,
                "camera_params_json": module_params_path,
                "source_capture_meta": {
                    "ts": meta.get("ts"),
                    "sensor_width": meta.get("sensor_width"),
                    "sensor_height": meta.get("sensor_height"),
                    "bayer_pattern": meta.get("bayer_pattern"),
                    "actual_camera_controls": meta.get("actual_camera_controls"),
                    "startup_camera_controls": meta.get("startup_camera_controls"),
                    "radiometric_last_result": meta.get("radiometric_last_result"),
                    "stream_meta": meta.get("stream_meta"),
                },
            }

            save_json(meta_out_path, out_meta)

            stats.update(tensor)

            rows.append({
                "group": sample.group,
                "base": sample.base,
                "tensor": str(tensor_path),
                "mask": str(mask_path),
                "mask_vegetation": str(mask_vegetation_path),
                "mask_cana": str(mask_cana_path),
                "meta": str(meta_out_path),
                "preview": str(preview_path),
                "quality_status": frame_quality.get("status") if isinstance(frame_quality, dict) else None,
            })

            print(f"[OK] {group_name}/{sample.base} tensor={list(tensor.shape)}")

        except Exception as e:
            errors += 1
            print(f"[ERRO] {group_name}/{sample.base}: {e}")

    print(f"[RESUMO] {group_name}: ok={len(rows)} | erros={errors} | skip_quality={skipped_quality}")
    return rows


def write_manifest(path: Path, rows: List[dict]):
    ensure_dir(path.parent)

    fieldnames = [
        "group",
        "base",
        "tensor",
        "mask",
        "mask_vegetation",
        "mask_cana",
        "meta",
        "preview",
        "quality_status",
    ]

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


def write_norm_stats(stats_path: Path, stats: dict):
    ensure_dir(stats_path.parent)

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


def main():
    ap = argparse.ArgumentParser(
        description="Normaliza RAW_BRUTO OAK-FCC-3 para tensor MULTISPEC final de treino."
    )

    ap.add_argument(
        "--src-roots",
        default="dataset/original/group;dataset/augmented/group",
        help="Raízes de entrada separadas por ';'. Ex: dataset/original/group;dataset/augmented/group"
    )
    ap.add_argument("--out-root", default="dataset")
    ap.add_argument("--module-params", default=DEFAULT_MODULE_PARAMS)
    ap.add_argument("--res", default=f"{DEFAULT_RES[0]}x{DEFAULT_RES[1]}", help="Resolução final WxH.")
    ap.add_argument("--raw-size", default=f"{DEFAULT_RAW_SIZE[0]}x{DEFAULT_RAW_SIZE[1]}", help="Tamanho RAW real WxH.")
    ap.add_argument("--groups", default=None, help="Grupos separados por vírgula.")
    ap.add_argument("--clear-dst", action="store_true")
    ap.add_argument("--skip-bad-quality", action="store_true")
    ap.add_argument("--manifest", default="")
    ap.add_argument("--stats-out", default="")

    args = ap.parse_args()

    src_roots = [
        Path(x.strip())
        for x in str(args.src_roots).split(";")
        if x.strip()
    ]

    out_dataset_root = Path(args.out_root)

    valid_src_roots = []
    for src_root in src_roots:
        if src_root.is_dir():
            valid_src_roots.append(src_root)
        else:
            print(f"[WARN] src-root não encontrado, ignorando: {src_root}")

    if not valid_src_roots:
        raise SystemExit(f"[ERRO] Nenhum src-root válido encontrado: {src_roots}")

    try:
        res_w, res_h = [int(x) for x in args.res.lower().split("x")]
        raw_w, raw_h = [int(x) for x in args.raw_size.lower().split("x")]
    except Exception:
        raise SystemExit("[ERRO] Use --res WxH e --raw-size WxH. Ex: --res 512x512 --raw-size 1280x800")

    res = (res_w, res_h)
    raw_size = (raw_w, raw_h)

    output_root = out_dataset_root / f"{res_w}x{res_h}" / "group"

    if args.clear_dst:
        print(f"[INFO] Limpando destino: {output_root}")
        maybe_clear_dir(output_root)

    labelmap_path = out_dataset_root / "labelmap.txt"

    if not labelmap_path.exists():
        raise SystemExit(f"[ERRO] labelmap não encontrado: {labelmap_path}")

    cor_para_id, _, _, ignore_rgb = carregar_labelmap_completo(str(labelmap_path))
    ignore_id = _infer_ignore_id(ignore_rgb, 255)

    groups_by_root = []

    for src_root in valid_src_roots:
        groups = list_groups(src_root)

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

        if groups:
            groups_by_root.append((src_root, groups))

    if not groups_by_root:
        raise SystemExit("[ERRO] Nenhum grupo encontrado.")

    print("============================================")
    print("Normalize OAK-FCC-3")
    print(f"SRC          : {[str(x) for x in valid_src_roots]}")
    print(f"OUT          : {output_root}")
    print(f"MODULE PARAM : {args.module_params}")
    print(f"RES          : {res}")
    print(f"RAW SIZE     : {raw_size}")
    print("GROUPS       :")
    for root, groups in groups_by_root:
        print(f"  - {root}: {', '.join(groups)}")
    print(f"SKIP BAD     : {args.skip_bad_quality}")
    print("============================================")

    running_stats = RunningStats()
    all_rows = []

    for src_root, all_groups in groups_by_root:
        for group_name in all_groups:
            rows = process_group(
                group_name=group_name,
                src_root=src_root,
                output_root=output_root,
                res=res,
                raw_size=raw_size,
                module_params_arg=args.module_params,
                cor_para_id=cor_para_id,
                ignore_id=ignore_id,
                stats=running_stats,
                dataset_root=out_dataset_root,
                skip_bad_quality=args.skip_bad_quality,
            )
            all_rows.extend(rows)

    manifest_path = Path(args.manifest) if args.manifest else output_root / "normalize_manifest.csv"
    write_manifest(manifest_path, all_rows)

    stats = running_stats.result(CHANNEL_NAMES)

    if stats is not None:
        stats_path = (
            Path(args.stats_out)
            if args.stats_out
            else Path("backup") / config.get("modelo", "modelo") / config.get("model_name", "model") / config.get("stats_source_tag", "oak_fcc3") / "norm_stats.json"
        )

        write_norm_stats(stats_path, stats)

        # Também salva uma cópia junto do dataset normalizado.
        write_norm_stats(output_root / "norm_stats.json", stats)

        print("\n📊 STATS:")
        print(json.dumps(stats, ensure_ascii=False, indent=2))
        print(f"[OK] norm_stats backup : {stats_path}")
        print(f"[OK] norm_stats dataset: {output_root / 'norm_stats.json'}")

    print("\n============================================")
    print("Normalize finalizado")
    print(f"Total samples: {len(all_rows)}")
    print(f"Manifest     : {manifest_path}")
    print("============================================")


if __name__ == "__main__":
    main()