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

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

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

Fallback (modo legado, se não houver "group/"):
  - original/{images,masks} e augmented/{images,masks} -> <WxH>/{images,masks}

Conversão de máscara:
  - Lê máscara RGB e converte para IDs via utils.converter_mask_rgb_para_ids
  - Ignora classe "ignore" conforme labelmap (usa índice 255 como padrão quando necessário)
"""
import argparse
import os
import json
import cv2
from typing import Dict, List, Tuple

import numpy as np
from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids

# ⚙️ Configurações
with open("config.json", "r", encoding="utf-8") as f:
    config = json.load(f)
MODELO = config["camera"]
MODEL_NAME = config["model_name"]
USE_MASKS2 = config["dual_head"]
RESOLUCAO = tuple(config["resolucao"])  # [W, H] ou [width, height]
pasta_base = os.path.join(MODELO, "dataset")
labelmap_path = os.path.join(pasta_base, "labelmap.txt")
save_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME)

# Dimensões alvo (pode expandir para múltiplas se quiser)
RESOLUCOES = {
    f"{RESOLUCAO[0]}x{RESOLUCAO[1]}": (RESOLUCAO[0], RESOLUCAO[1]),
}

# Fontes a processar
FONTES = ["original", "augmented"]

# Extensões aceitas
IMG_EXTS = (".jpg", ".jpeg", ".png")
MSK_EXTS = (".png", ".jpg", ".jpeg")  # preferir .png
MSK2_EXTS = (".png", ".jpg", ".jpeg")  # idem

# === Acumuladores globais para mean/std dos canais RAW4 ===
GLOBAL_SUM   = None   # soma por canal
GLOBAL_SUMSQ = None   # soma dos quadrados por canal
GLOBAL_PIXELS = 0     # n de pixels por canal (H*W por imagem)

def infer_ignore_id(ignore_rgb, default_id=255):
    """
    Tenta inferir o ID de ignore a partir do valor retornado por carregar_labelmap_completo.
    - Se for [id] retorna id
    - Se for (R,G,B) retorna default_id (tipicamente 255)
    - Se for int, retorna direto
    """
    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]:
    """Lista grupos válidos com subpastas images e masks."""
    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 map_masks_by_base(msk_dir: str) -> Dict[str, str]:
    """Retorna {base: caminho_mask}, priorizando .png quando houver múltiplas por base."""
    by_base = {}
    if not os.path.isdir(msk_dir):
        return by_base
    for fname in os.listdir(msk_dir):
        f_lower = fname.lower()
        if not f_lower.endswith(MSK_EXTS):
            continue
        base, ext = os.path.splitext(fname)
        cand = os.path.join(msk_dir, fname)
        if base not in by_base:
            by_base[base] = cand
        else:
            cur_ext = os.path.splitext(by_base[base])[1].lower()
            if cur_ext != ".png" and ext.lower() == ".png":
                by_base[base] = cand
    return by_base

def map_masks2_by_base(msk2_dir: str) -> Dict[str, str]:
    """Retorna {base: caminho_mask2}, priorizando .png quando houver múltiplas por base."""
    by_base = {}
    if not os.path.isdir(msk2_dir):
        return by_base
    for fname in os.listdir(msk2_dir):
        f_lower = fname.lower()
        if not f_lower.endswith(MSK2_EXTS):
            continue
        base, ext = os.path.splitext(fname)
        cand = os.path.join(msk2_dir, fname)
        if base not in by_base:
            by_base[base] = cand
        else:
            cur_ext = os.path.splitext(by_base[base])[1].lower()
            if cur_ext != ".png" and ext.lower() == ".png":
                by_base[base] = cand
    return by_base

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 a imagem e a máscara (se houver)."""
    img_rgb = cv2.imread(caminho_rgb)
    if img_rgb is None:
        print(f"[!] Erro ao ler imagem: {caminho_rgb}")
        return False

    # Nome de saída com prefixo para distinguir fonte (ex: original_, augmented_)
    nome = os.path.basename(caminho_rgb)
    if prefix:
        nome_saida_img = f"{prefix}{nome}"
    else:
        nome_saida_img = nome
    nome_saida_msk = nome_saida_img
    for ext in (".jpg", ".jpeg", ".png"):
        if nome_saida_msk.lower().endswith(ext):
            nome_saida_msk = nome_saida_msk[: -len(ext)] + ".png"
            break

    # Redimensiona imagem
    img_resized = cv2.resize(img_rgb, dim, interpolation=cv2.INTER_AREA)

    global GLOBAL_SUM, GLOBAL_SUMSQ, GLOBAL_PIXELS
    img_float = img_resized.astype(np.float32) / 255.0 # 0-1
    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)
    cv2.imwrite(os.path.join(out_img_dir, nome_saida_img), img_resized)

    # Processa e redimensiona máscara (se existir)
    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)
            cv2.imwrite(os.path.join(out_msk_dir, nome_saida_msk), mask_resized)


    return True

def normalize_pair_mask2(caminho_rgb: str, caminho_mask2: str,
                         out_msk2_dir: str, dim: Tuple[int,int], prefix: str = ""):
    """
    Redimensiona e grava máscara2 (corredor binário), assumindo que ela já é uma máscara "pronta".
    - Se for RGB/BGR (3 canais): converte para cinza e faz threshold (0/255) antes de redimensionar.
    - Se for 1 canal: mantém, faz threshold (0/255).
    - Redimensiona com INTER_NEAREST.
    Saída sempre .png com o mesmo nome base/prefixo do arquivo de imagem.
    """
    if not caminho_mask2 or not os.path.isfile(caminho_mask2):
        return False

    nome = os.path.basename(caminho_rgb)
    nome_saida = f"{prefix}{nome}" if prefix else nome
    for ext in (".jpg", ".jpeg", ".png"):
        if nome_saida.lower().endswith(ext):
            nome_saida = nome_saida[: -len(ext)] + ".png"
            break

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

    if len(m2.shape) == 3:
        # BGR/RGB -> gray
        m2g = cv2.cvtColor(m2, cv2.COLOR_BGR2GRAY)
    else:
        m2g = m2

    # binariza para 0/255 (evita lixo de compressão)
    _, m2bin = cv2.threshold(m2g, 127, 255, cv2.THRESH_BINARY)
    m2res = cv2.resize(m2bin, dim, interpolation=cv2.INTER_NEAREST)
    garantir_dir(out_msk2_dir)
    cv2.imwrite(os.path.join(out_msk2_dir, nome_saida), m2res)
    return True

def process_group_root(fonte_root: str, fonte_nome: str, cor_para_id, ignore_id: int, groups_except: str = None):
    """Processa uma raiz do tipo .../<fonte>/group/ agrupando por cada subpasta de grupo."""
    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 nao sera 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")
            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)
            out_msk2_dir = os.path.join(out_root, grupo, "masks2") if usar_masks2 else None
            msk_map = map_masks_by_base(in_msk_dir)
            msk2_map = map_masks2_by_base(in_msk2_dir) if usar_masks2 else {}

            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)
                caminho_rgb = os.path.join(in_img_dir, fname)
                caminho_mask = msk_map.get(base)
                caminho_mask2 = msk2_map.get(base) if usar_masks2 else None
                ok = normalize_pair(
                    caminho_rgb, caminho_mask, cor_para_id, ignore_id,
                    out_img_dir, out_msk_dir, dim, prefix=f"{fonte_nome}_"
                )
                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} (vou seguir).")
                    else:
                        normalize_pair_mask2(
                            caminho_rgb, caminho_mask2,
                            out_msk2_dir, dim, 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):
    """Processa estrutura legado (sem grupos)."""
    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_msk2 = os.path.join(os.path.dirname(legacy_msk), "masks2")
        usar_masks2 = USE_MASKS2 and os.path.isdir(legacy_msk2)
        out_msk2_dir = os.path.join(pasta_base, nome_res, "masks2") if usar_masks2 else None

        msk_map = map_masks_by_base(legacy_msk)
        msk2_map = map_masks2_by_base(legacy_msk2) if usar_masks2 else {}
        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)
            caminho_rgb = os.path.join(legacy_img, fname)
            caminho_mask = msk_map.get(base)
            caminho_mask2 = msk2_map.get(base) if usar_masks2 else None
            ok = normalize_pair(
                caminho_rgb, caminho_mask, cor_para_id, ignore_id,
                out_img_dir, out_msk_dir, dim, prefix=f"{fonte_nome}_"
            )
            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} (vou seguir).")
                else:
                    normalize_pair_mask2(
                        caminho_rgb, caminho_mask2,
                        out_msk2_dir, dim, prefix=f"{fonte_nome}_"
                    )
            if ok:
                total += 1
            print(f"[{fonte_nome} | legacy | {nome_res}] {i}/{n} → {fname}")
    return total

def main(args):
    # === Labelmap ===
    # Espera tupla na ordem: (cor_para_id, colormap_rgb, id_para_nome, ignore_rgb)
    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(cor_para_id, _colormap_rgb, _id_para_nome)

    total_geral = 0
    # === ORIGINAL ===
    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)
    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)

    # === AUGMENTED ===
    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)
    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)

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

    # === calcula mean/std globais e salva em JSON ===
    global GLOBAL_SUM, GLOBAL_SUMSQ, GLOBAL_PIXELS
    if GLOBAL_SUM is not None and GLOBAL_PIXELS > 0:
        # média e variância por canal
        mean = (GLOBAL_SUM / GLOBAL_PIXELS)
        var  = (GLOBAL_SUMSQ / GLOBAL_PIXELS) - mean**2
        std  = np.sqrt(np.maximum(var, 1e-6))

        # Converte para list pra salvar em JSON
        mean_list = mean.tolist()
        std_list  = std.tolist()

        # Se quiser, você pode nomear os canais explicitamente
        # dependendo da convenção do raw4:
        channel_names = ["R", "G", "IR", "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("⚠️ Nenhum RAW processado, não há stats para salvar.")

if __name__ == "__main__":
    ap = argparse.ArgumentParser(description="Augmentação por grupos (images/masks)")
    ap.add_argument("--groups-except", type=str, default="", help="Lista de grupos para nao usar, separados por vírgula (ex: chao,erva_cana).")
    args = ap.parse_args()
    main(args)