#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Normaliza/redimensiona imagens, máscaras, masks2 e labels mantendo a ESTRUTURA POR GRUPO.

Entradas via config.json -> camera, resolucao:
  - MODELO/dataset/original/group/<grupo>/{images,masks,(masks2),(labels)}
  - MODELO/dataset/augmented/group/<grupo>/{images,masks,(masks2),(labels)}

Saídas por resolução:
  - MODELO/dataset/<WxH>/group/<grupo>/{images,masks,(masks2),(labels)}

Fallback legado, se não houver group/:
  - original/{images,masks,(masks2),(labels)}
  - augmented/{images,masks,(masks2),(labels)}
  - saída: <WxH>/{images,masks,(masks2),(labels)}

Conversão de máscara:
  - Lê máscara RGB e converte para IDs via utils.converter_mask_rgb_para_ids.
  - Ignore conforme labelmap, usando índice 255 como fallback.

Labels:
  - Ativado por config['dual_head_label'].
  - Copia labels .json/.txt para a saída.
  - Se JSON, atualiza image/mask/mask2/label/base/source/normalized.
  - Se houver label_id, salva também um .npy com o inteiro para facilitar o Dataset no treino.
"""

import argparse
import os
import json
import cv2
from typing import Dict, List, Tuple, Optional

import numpy as np
from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids

# ===================== CONFIG =====================

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

MODELO = config["camera"]
MODEL_NAME = config["model_name"]

# Compatibilidade:
# - scripts antigos usam dual_head para masks2
# - scripts novos podem usar dual_head_mask e dual_head_label separados
USE_MASKS2 = bool(config.get("dual_head_mask", config.get("dual_head", False)))
USE_LABELS = bool(config.get("dual_head_label", False))

RESOLUCAO = tuple(config["resolucao"])  # [W, H]
pasta_base = os.path.join(MODELO, "dataset")
labelmap_path = os.path.join(pasta_base, "labelmap.txt")
mode = "single"
if USE_MASKS2:
    mode = "mask2"
elif USE_LABELS:
    mode = "label"
suffix = {
    "single": "_single",
    "mask2": "_dual_mask",
    "label": "_dual_label",
}[mode]
save_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME + suffix)

RESOLUCOES = {
    f"{RESOLUCAO[0]}x{RESOLUCAO[1]}": (RESOLUCAO[0], RESOLUCAO[1]),
}

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

GLOBAL_SUM = None
GLOBAL_SUMSQ = None
GLOBAL_PIXELS = 0


# ===================== HELPERS =====================

def infer_ignore_id(ignore_rgb, default_id=255):
    if isinstance(ignore_rgb, (list, tuple)):
        if len(ignore_rgb) == 1:
            try:
                return int(ignore_rgb[0])
            except Exception:
                return default_id
        if len(ignore_rgb) == 3:
            return default_id
    if isinstance(ignore_rgb, int):
        return ignore_rgb
    return default_id


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


def list_groups(root) -> List[str]:
    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 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: str, exts: Tuple[str, ...]) -> Dict[str, str]:
    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 map_masks_by_base(msk_dir: str) -> Dict[str, str]:
    return map_files_by_base(msk_dir, MSK_EXTS)


def map_masks2_by_base(msk2_dir: str) -> Dict[str, str]:
    return map_files_by_base(msk2_dir, MSK2_EXTS)


def map_labels_by_base(label_dir: str) -> Dict[str, str]:
    return map_files_by_base(label_dir, LABEL_EXTS)


def trocar_ext_para_png(nome: str) -> str:
    for ext in (".jpg", ".jpeg", ".png"):
        if nome.lower().endswith(ext):
            return nome[: -len(ext)] + ".png"
    return nome + ".png"


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


# ===================== NORMALIZAÇÃO =====================

def normalize_pair(
    caminho_rgb: str,
    caminho_mask: str,
    cor_para_id,
    ignore_id: int,
    out_img_dir: str,
    out_msk_dir: str,
    dim: Tuple[int, int],
    prefix: str = "",
):
    """Redimensiona e grava imagem + máscara ID."""
    img_bgr = cv2.imread(caminho_rgb)
    if img_bgr is None:
        print(f"[!] Erro ao ler imagem: {caminho_rgb}")
        return False, None, None
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

    nome = os.path.basename(caminho_rgb)
    nome_saida_img = f"{prefix}{nome}" if prefix else nome
    nome_saida_msk = trocar_ext_para_png(nome_saida_img)

    img_resized_rgb = cv2.resize(img_rgb, dim, interpolation=cv2.INTER_AREA)

    global GLOBAL_SUM, GLOBAL_SUMSQ, GLOBAL_PIXELS
    img_float = img_resized_rgb.astype(np.float32) / 255.0
    h, w, c = img_float.shape
    flat = img_float.reshape(-1, c).astype(np.float64)

    if GLOBAL_SUM is None:
        GLOBAL_SUM = np.zeros(c, dtype=np.float64)
        GLOBAL_SUMSQ = np.zeros(c, dtype=np.float64)

    GLOBAL_SUM += flat.sum(axis=0)
    GLOBAL_SUMSQ += (flat ** 2).sum(axis=0)
    GLOBAL_PIXELS += h * w

    garantir_dir(out_img_dir)
    out_img_path = os.path.join(out_img_dir, nome_saida_img)
    img_resized_bgr = cv2.cvtColor(img_resized_rgb, cv2.COLOR_RGB2BGR)
    cv2.imwrite(out_img_path, img_resized_bgr)

    out_msk_path = None
    if caminho_mask and os.path.isfile(caminho_mask):
        msk_bgr = cv2.imread(caminho_mask, cv2.IMREAD_COLOR)
        if msk_bgr is None:
            print(f"[!] Erro ao ler máscara: {caminho_mask}")
        else:
            msk_rgb = cv2.cvtColor(msk_bgr, cv2.COLOR_BGR2RGB)
            mask_ids = converter_mask_rgb_para_ids(msk_rgb, cor_para_id, ignore_id)
            mask_resized = cv2.resize(mask_ids, dim, interpolation=cv2.INTER_NEAREST)

            garantir_dir(out_msk_dir)
            out_msk_path = os.path.join(out_msk_dir, nome_saida_msk)
            cv2.imwrite(out_msk_path, mask_resized)

    return True, out_img_path, out_msk_path


def normalize_pair_mask2(
    caminho_rgb: str,
    caminho_mask2: str,
    out_msk2_dir: str,
    dim: Tuple[int, int],
    prefix: str = "",
):
    if not caminho_mask2 or not os.path.isfile(caminho_mask2):
        return False, None

    nome = os.path.basename(caminho_rgb)
    nome_saida = f"{prefix}{nome}" if prefix else nome
    nome_saida = trocar_ext_para_png(nome_saida)

    m2 = cv2.imread(caminho_mask2, cv2.IMREAD_UNCHANGED)
    if m2 is None:
        print(f"[!] Erro ao ler máscara2: {caminho_mask2}")
        return False, None

    if len(m2.shape) == 3:
        m2g = cv2.cvtColor(m2, cv2.COLOR_BGR2GRAY)
    else:
        m2g = m2

    _, m2bin = cv2.threshold(m2g, 127, 255, cv2.THRESH_BINARY)
    m2res = cv2.resize(m2bin, dim, interpolation=cv2.INTER_NEAREST)

    garantir_dir(out_msk2_dir)
    out_msk2_path = os.path.join(out_msk2_dir, nome_saida)
    cv2.imwrite(out_msk2_path, m2res)

    return True, out_msk2_path


def normalize_label(
    caminho_rgb: str,
    caminho_label: Optional[str],
    out_label_dir: Optional[str],
    out_img_path: Optional[str],
    out_msk_path: Optional[str],
    out_msk2_path: Optional[str],
    source_root: str,
    output_root: str,
    grupo: Optional[str],
    fonte_nome: str,
    prefix: str = "",
):
    """Copia/atualiza label global e salva label_id como .npy quando possível."""
    if not caminho_label or not out_label_dir or not os.path.isfile(caminho_label):
        return False

    nome = os.path.basename(caminho_rgb)
    nome_saida = f"{prefix}{nome}" if prefix else nome
    base_saida, _ = os.path.splitext(nome_saida)

    label_ext = os.path.splitext(caminho_label)[1].lower()
    garantir_dir(out_label_dir)

    out_label_path = os.path.join(out_label_dir, base_saida + label_ext)
    out_label_id_path = os.path.join(out_label_dir, base_saida + ".npy")

    label_id = None

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

        label_id = data.get("label_id")

        data["normalized"] = True
        data["source"] = fonte_nome
        data["group"] = grupo if grupo is not None else data.get("group")
        data["base"] = base_saida
        data["image"] = safe_rel(out_img_path, output_root) if out_img_path else None
        data["mask"] = safe_rel(out_msk_path, output_root) if out_msk_path else None
        data["label"] = safe_rel(out_label_path, output_root)
        data["source_label"] = safe_rel(caminho_label, source_root)

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

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

    else:
        with open(caminho_label, "r", encoding="utf-8") as f:
            txt = f.read().strip()
        with open(out_label_path, "w", encoding="utf-8") as f:
            f.write(txt)
        label_id = None

    if label_id is not None:
        try:
            np.save(out_label_id_path, np.array(int(label_id), dtype=np.int64))
        except Exception as e:
            print(f"[WARN] Não consegui salvar label_id npy para {caminho_label}: {e}")

    return True


# ===================== PROCESSAMENTO =====================

def process_group_root(
    fonte_root: str,
    fonte_nome: str,
    cor_para_id,
    ignore_id: int,
    groups_except: str = "",
    strict_label: bool = False,
):
    total = 0
    grupos = list_groups(fonte_root)
    if not grupos:
        return 0

    not_want = {g.strip() for g in groups_except.split(",") if g.strip()}
    grupos_desconsiderar = [g for g in grupos if g in not_want]

    for nome_res, dim in RESOLUCOES.items():
        out_root = os.path.join(pasta_base, nome_res, "group")

        for grupo in grupos:
            if grupo in grupos_desconsiderar:
                print(f"[WARN] Grupo desconsiderado não será processado: {grupo}")
                continue

            in_img_dir = os.path.join(fonte_root, grupo, "images")
            in_msk_dir = os.path.join(fonte_root, grupo, "masks")
            in_msk2_dir = os.path.join(fonte_root, grupo, "masks2")
            in_label_dir = os.path.join(fonte_root, grupo, "labels")

            if not (os.path.isdir(in_img_dir) and os.path.isdir(in_msk_dir)):
                print(f"[WARN] Grupo inválido sem images/masks: {grupo}")
                continue

            out_img_dir = os.path.join(out_root, grupo, "images")
            out_msk_dir = os.path.join(out_root, grupo, "masks")

            usar_masks2 = USE_MASKS2 and os.path.isdir(in_msk2_dir)
            usar_labels = USE_LABELS and os.path.isdir(in_label_dir)

            out_msk2_dir = os.path.join(out_root, grupo, "masks2") if usar_masks2 else None
            out_label_dir = os.path.join(out_root, grupo, "labels") if usar_labels else None

            msk_map = map_masks_by_base(in_msk_dir)
            msk2_map = map_masks2_by_base(in_msk2_dir) if usar_masks2 else {}
            label_map = map_labels_by_base(in_label_dir) if usar_labels else {}

            if USE_LABELS and not usar_labels:
                msg = f"[WARN] [{fonte_nome} | {grupo}] dual_head_label=true, mas labels/ não existe."
                if strict_label:
                    print(msg + " Pulando grupo.")
                    continue
                print(msg + " Seguindo sem labels.")

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

            for i, fname in enumerate(sorted(imgs), 1):
                base, _ = os.path.splitext(fname)
                base_norm = normalizar_base(base)

                caminho_rgb = os.path.join(in_img_dir, fname)
                caminho_mask = msk_map.get(base_norm)
                caminho_mask2 = msk2_map.get(base_norm) if usar_masks2 else None
                caminho_label = label_map.get(base_norm) if usar_labels else None

                if usar_labels and not caminho_label:
                    msg = f"[WARN] [{fonte_nome} | {grupo}] label não encontrado para {fname}."
                    if strict_label:
                        print(msg + " Pulando item.")
                        continue
                    print(msg + " Seguindo sem label.")

                ok, out_img_path, out_msk_path = normalize_pair(
                    caminho_rgb,
                    caminho_mask,
                    cor_para_id,
                    ignore_id,
                    out_img_dir,
                    out_msk_dir,
                    dim,
                    prefix=f"{fonte_nome}_",
                )

                out_msk2_path = None
                if usar_masks2 and out_msk2_dir:
                    if not caminho_mask2:
                        print(f"[WARN] [{fonte_nome} | {grupo}] masks2 existe, mas não achei mask2 p/ {fname}.")
                    else:
                        _, out_msk2_path = normalize_pair_mask2(
                            caminho_rgb,
                            caminho_mask2,
                            out_msk2_dir,
                            dim,
                            prefix=f"{fonte_nome}_",
                        )

                if usar_labels and caminho_label and out_label_dir:
                    normalize_label(
                        caminho_rgb=caminho_rgb,
                        caminho_label=caminho_label,
                        out_label_dir=out_label_dir,
                        out_img_path=out_img_path,
                        out_msk_path=out_msk_path,
                        out_msk2_path=out_msk2_path,
                        source_root=fonte_root,
                        output_root=os.path.join(pasta_base, nome_res),
                        grupo=grupo,
                        fonte_nome=fonte_nome,
                        prefix=f"{fonte_nome}_",
                    )

                if ok:
                    total += 1

                print(f"[{fonte_nome} | {grupo} | {nome_res}] {i}/{n} → {fname}")

    return total


def process_legacy_root(
    legacy_img: str,
    legacy_msk: str,
    fonte_nome: str,
    cor_para_id,
    ignore_id: int,
    strict_label: bool = False,
):
    if not (os.path.isdir(legacy_img) and os.path.isdir(legacy_msk)):
        return 0

    total = 0

    for nome_res, dim in RESOLUCOES.items():
        out_img_dir = os.path.join(pasta_base, nome_res, "images")
        out_msk_dir = os.path.join(pasta_base, nome_res, "masks")

        legacy_root = os.path.dirname(legacy_msk)
        legacy_msk2 = os.path.join(legacy_root, "masks2")
        legacy_labels = os.path.join(legacy_root, "labels")

        usar_masks2 = USE_MASKS2 and os.path.isdir(legacy_msk2)
        usar_labels = USE_LABELS and os.path.isdir(legacy_labels)

        out_msk2_dir = os.path.join(pasta_base, nome_res, "masks2") if usar_masks2 else None
        out_label_dir = os.path.join(pasta_base, nome_res, "labels") if usar_labels else None

        msk_map = map_masks_by_base(legacy_msk)
        msk2_map = map_masks2_by_base(legacy_msk2) if usar_masks2 else {}
        label_map = map_labels_by_base(legacy_labels) if usar_labels else {}

        if USE_LABELS and not usar_labels:
            msg = f"[WARN] [{fonte_nome} | legacy] dual_head_label=true, mas labels/ não existe."
            if strict_label:
                print(msg + " Pulando legacy.")
                return total
            print(msg + " Seguindo sem labels.")

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

        for i, fname in enumerate(sorted(imgs), 1):
            base, _ = os.path.splitext(fname)
            base_norm = normalizar_base(base)

            caminho_rgb = os.path.join(legacy_img, fname)
            caminho_mask = msk_map.get(base_norm)
            caminho_mask2 = msk2_map.get(base_norm) if usar_masks2 else None
            caminho_label = label_map.get(base_norm) if usar_labels else None

            if usar_labels and not caminho_label:
                msg = f"[WARN] [{fonte_nome} | legacy] label não encontrado para {fname}."
                if strict_label:
                    print(msg + " Pulando item.")
                    continue
                print(msg + " Seguindo sem label.")

            ok, out_img_path, out_msk_path = normalize_pair(
                caminho_rgb,
                caminho_mask,
                cor_para_id,
                ignore_id,
                out_img_dir,
                out_msk_dir,
                dim,
                prefix=f"{fonte_nome}_",
            )

            out_msk2_path = None
            if usar_masks2 and out_msk2_dir:
                if not caminho_mask2:
                    print(f"[WARN] [{fonte_nome} | legacy] masks2 existe, mas não achei mask2 p/ {fname}.")
                else:
                    _, out_msk2_path = normalize_pair_mask2(
                        caminho_rgb,
                        caminho_mask2,
                        out_msk2_dir,
                        dim,
                        prefix=f"{fonte_nome}_",
                    )

            if usar_labels and caminho_label and out_label_dir:
                normalize_label(
                    caminho_rgb=caminho_rgb,
                    caminho_label=caminho_label,
                    out_label_dir=out_label_dir,
                    out_img_path=out_img_path,
                    out_msk_path=out_msk_path,
                    out_msk2_path=out_msk2_path,
                    source_root=legacy_root,
                    output_root=os.path.join(pasta_base, nome_res),
                    grupo=None,
                    fonte_nome=fonte_nome,
                    prefix=f"{fonte_nome}_",
                )

            if ok:
                total += 1

            print(f"[{fonte_nome} | legacy | {nome_res}] {i}/{n} → {fname}")

    return total


# ===================== MAIN =====================

def main(args):
    cor_para_id, _colormap_rgb, _id_para_nome, ignore_rgb = carregar_labelmap_completo(labelmap_path)
    ignore_id = infer_ignore_id(ignore_rgb, default_id=255)

    print(f"[INFO] cor_para_id: {cor_para_id}")
    print(f"[INFO] classes: {_id_para_nome}")
    print(f"[INFO] dual_head_mask/masks2: {USE_MASKS2}")
    print(f"[INFO] dual_head_label/labels: {USE_LABELS}")

    total_geral = 0

    orig_group_root = os.path.join(pasta_base, "original", "group")
    if os.path.isdir(orig_group_root):
        total_geral += process_group_root(
            orig_group_root,
            "original",
            cor_para_id,
            ignore_id,
            groups_except=args.groups_except,
            strict_label=args.strict_label,
        )
    else:
        legacy_img = os.path.join(pasta_base, "original", "images")
        legacy_msk = os.path.join(pasta_base, "original", "masks")
        total_geral += process_legacy_root(
            legacy_img,
            legacy_msk,
            "original",
            cor_para_id,
            ignore_id,
            strict_label=args.strict_label,
        )

    aug_group_root = os.path.join(pasta_base, "augmented", "group")
    if os.path.isdir(aug_group_root):
        total_geral += process_group_root(
            aug_group_root,
            "augmented",
            cor_para_id,
            ignore_id,
            groups_except=args.groups_except,
            strict_label=args.strict_label,
        )
    else:
        legacy_img = os.path.join(pasta_base, "augmented", "images")
        legacy_msk = os.path.join(pasta_base, "augmented", "masks")
        total_geral += process_legacy_root(
            legacy_img,
            legacy_msk,
            "augmented",
            cor_para_id,
            ignore_id,
            strict_label=args.strict_label,
        )

    print(f"\n✅ Concluído! Total normalizados: {total_geral}")

    global GLOBAL_SUM, GLOBAL_SUMSQ, GLOBAL_PIXELS
    if GLOBAL_SUM is not None and GLOBAL_PIXELS > 0:
        mean = GLOBAL_SUM / GLOBAL_PIXELS
        var = (GLOBAL_SUMSQ / GLOBAL_PIXELS) - mean ** 2
        std = np.sqrt(np.maximum(var, 1e-6))

        mean_list = mean.tolist()
        std_list = std.tolist()
        channel_names = ["R", "G", "B"]

        stats = {
            "channels": channel_names[:len(mean_list)],
            "mean": mean_list,
            "std": std_list,
            "pixels_per_channel": int(GLOBAL_PIXELS),
        }

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

        print(f"📁 Stats salvos em: {stats_path}")
        print(f"  mean: {mean_list}")
        print(f"  std : {std_list}")
    else:
        print("⚠️ Nenhuma imagem processada, não há stats para salvar.")


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description="Normalização por grupos images/masks/masks2/labels.")
    ap.add_argument(
        "--groups-except",
        type=str,
        default="",
        help="Lista de grupos para não usar, separados por vírgula. Ex: chao,erva_cana",
    )
    ap.add_argument(
        "--strict-label",
        action="store_true",
        help="Se dual_head_label=true e faltar label, pula item/grupo.",
    )
    args = ap.parse_args()
    main(args)