#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Conta a porcentagem de pixels por CLASSE (segundo o labelmap) dentro da ROI,
para cada grupo em dataset/split/train/group.

- Usa utils.carregar_labelmap_completo(labelmap_path) para obter:
    colormap_rgb, classes (id->nome) e ignore_rgb
- Ignora pixels com o valor de "ignore" do labelmap
- Normaliza a porcentagem SOMENTE sobre classes válidas (sem ignore)

Uso:
  python _12_check_percent_class_labelmap.py
"""

import argparse
import os, json
import numpy as np
from PIL import Image
from utils import carregar_labelmap_completo

# -------------- Config --------------
with open("config.json", "r", encoding="utf-8") as f:
    config = json.load(f)
MODELO = config["camera"]
W, H = config["resolucao"][0], config["resolucao"][1]
ROI_INICIO = config["roi_inicio"]
ROI_TAMANHO = config["roi_tamanho"]

pasta_base = os.path.join(MODELO, "dataset")
labelmap_path = os.path.join(pasta_base, "labelmap.txt")

# Limite de amostras por grupo (para rodar rápido). Ajuste se quiser.
MAX_SAMPLES_PER_GROUP = 1000

# -------------- Utils --------------
def infer_ignore_id(ignore_rgb, default_id=255):
    """
    Converte o 'ignore' do labelmap (que pode vir como [id] ou (R,G,B) ou int)
    para um ID inteiro que devemos ignorar nas máscaras de IDs.
    """
    # pode vir como lista/tupla com 1 elemento (id) ou 3 (cor)
    if isinstance(ignore_rgb, (list, tuple)):
        if len(ignore_rgb) == 1 and isinstance(ignore_rgb[0], (int, np.integer)):
            return int(ignore_rgb[0])
        if len(ignore_rgb) == 3:
            return default_id
    # pode vir como inteiro
    if isinstance(ignore_rgb, (int, np.integer)):
        return int(ignore_rgb)
    return default_id

def roi_slice(h):
    y_fim = int((1.0 - ROI_TAMANHO) * h)
    y_ini = int(ROI_INICIO * h)
    if y_ini <= y_fim:
        y_fim, y_ini = max(0, h - int(ROI_TAMANHO * h)), h
    return slice(y_fim, y_ini)

def main(args):
    # -------------- Labelmap --------------
    _, colormap_rgb, classes, ignore_rgb = carregar_labelmap_completo(labelmap_path)
    ignore_id = infer_ignore_id(ignore_rgb, default_id=255)

    # 'classes' esperado como dict: id -> nome
    # Ordena por id para imprimir de forma estável
    class_ids_sorted = sorted(classes.keys())
    class_names_sorted = [classes[cid] for cid in class_ids_sorted]

    root = os.path.join(pasta_base, "split", args.split, "group")
    #root = os.path.join(pasta_base, "576x320", "group")

    # -------------- Coleta --------------
    if not os.path.isdir(root):
        raise SystemExit(f"Nenhum diretório encontrado em {root}")

    grupos = [g for g in os.listdir(root) if os.path.isdir(os.path.join(root, g))]

    for g in sorted(grupos):
        mdir = os.path.join(root, g, "masks")
        if not os.path.isdir(mdir):
            continue

        totals = {cid: 0 for cid in class_ids_sorted}
        n = 0

        for fname in os.listdir(mdir):
            if not fname.lower().endswith(".png"):
                continue
            m = np.array(Image.open(os.path.join(mdir, fname)).convert("L"))
            rs = roi_slice(m.shape[0])
            roi = m[rs, :]

            # Acumula só das classes válidas do labelmap (ignorando 'ignore' e outros valores)
            for cid in class_ids_sorted:
                totals[cid] += int((roi == cid).sum())

            n += 1
            if n >= MAX_SAMPLES_PER_GROUP:
                break

        s = sum(totals.values())
        s = s if s > 0 else 1  # evita div/0

        # Monta string dinâmica "nome=xx.xx%"
        parts = []
        for cid in class_ids_sorted:
            name = classes[cid]
            perc = totals[cid] / s
            parts.append(f"{name}={perc:6.2%}")

        print(f"{g:16s}  " + "  ".join(parts) + f"  (amostras={n})")

if __name__ == "__main__":
    ap = argparse.ArgumentParser(description="Validacao de classes no dataset")
    ap.add_argument("--split", type=str, default="train", help="Pasta para verificar (ex: train,val,test).")
    args = ap.parse_args()
    main(args)