#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Analisa o dataset de corredor (masks2) e mede % de pixels navegáveis (branco) vs não navegáveis (preto),
por split (train/val/test) e por grupo.

Como usar (de dentro da pasta split):
  python analyze_masks2_balance.py

Ou passando o caminho da pasta split:
  python analyze_masks2_balance.py --split_root "oak-d/dataset/split"
"""

import os
import argparse
from collections import defaultdict
import cv2
import numpy as np

IMG_EXTS = (".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp")

def is_image_file(p: str) -> bool:
    return p.lower().endswith(IMG_EXTS)

def iter_masks2_paths(split_root: str):
    """
    Estrutura esperada:
      split_root/
        train/group/<group_name>/masks2/*.png
        val/group/<group_name>/masks2/*.png
        test/group/<group_name>/masks2/*.png
    """
    for split in ("train", "val", "test"):
        split_dir = os.path.join(split_root, split, "group")
        if not os.path.isdir(split_dir):
            # tenta fallback caso não exista "group" (alguns datasets mudam)
            split_dir = os.path.join(split_root, split)
        if not os.path.isdir(split_dir):
            continue

        for group_name in sorted(os.listdir(split_dir)):
            group_dir = os.path.join(split_dir, group_name)
            if not os.path.isdir(group_dir):
                continue

            masks2_dir = os.path.join(group_dir, "masks2")
            if not os.path.isdir(masks2_dir):
                # se não existir, ainda assim reportamos vazio
                yield split, group_name, None
                continue

            files = [os.path.join(masks2_dir, f) for f in os.listdir(masks2_dir) if is_image_file(f)]
            if not files:
                yield split, group_name, None
                continue

            for fp in sorted(files):
                yield split, group_name, fp

def read_mask_gray(path: str):
    m = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
    return m

def compute_counts(mask: np.ndarray):
    """
    Considera navegável = pixels > 0 (funciona tanto para 0/1 quanto 0/255).
    Não navegável = pixels == 0
    """
    total = int(mask.size)
    white = int(np.count_nonzero(mask))          # >0
    black = total - white                        # ==0
    return black, white, total

def pct(a, b):
    return 0.0 if b == 0 else (100.0 * a / b)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--split_root", default=".", help="Caminho da pasta split (default: .)")
    ap.add_argument("--warn_thresh_all_white", type=float, default=98.0,
                    help="Aviso se % navegável >= esse valor (default 98)")
    args = ap.parse_args()

    split_root = os.path.abspath(args.split_root)

    # stats[split][group] = dict(counts...)
    stats = defaultdict(lambda: defaultdict(lambda: {
        "files": 0, "black": 0, "white": 0, "total": 0, "bad_read": 0, "missing_dir": 0
    }))

    seen_group_in_split = set()

    for split, group, fp in iter_masks2_paths(split_root):
        key = (split, group)
        seen_group_in_split.add(key)

        if fp is None:
            stats[split][group]["missing_dir"] += 1
            continue

        m = read_mask_gray(fp)
        if m is None:
            stats[split][group]["bad_read"] += 1
            continue

        b, w, t = compute_counts(m)
        st = stats[split][group]
        st["files"] += 1
        st["black"] += b
        st["white"] += w
        st["total"] += t

    # print relatório
    print(f"\n[OK] Análise masks2 em: {split_root}\n")
    header = f"{'split':<6} | {'grupo':<30} | {'arquivos':>7} | {'% preto':>8} | {'% branco':>8} | {'bad':>3} | {'missing':>7}"
    print(header)
    print("-" * len(header))

    # totais por split
    split_totals = {s: {"files": 0, "black": 0, "white": 0, "total": 0, "bad": 0, "missing": 0} for s in ("train","val","test")}

    for split in ("train", "val", "test"):
        if split not in stats:
            continue
        for group in sorted(stats[split].keys()):
            st = stats[split][group]
            total = st["total"]
            p_black = pct(st["black"], total)
            p_white = pct(st["white"], total)

            warn = ""
            if total > 0 and p_white >= args.warn_thresh_all_white:
                warn = "  <== ⚠ muito branco"

            print(f"{split:<6} | {group:<30} | {st['files']:>7} | {p_black:>7.2f}% | {p_white:>7.2f}% | {st['bad_read']:>3} | {st['missing_dir']:>7}{warn}")

            split_totals[split]["files"] += st["files"]
            split_totals[split]["black"] += st["black"]
            split_totals[split]["white"] += st["white"]
            split_totals[split]["total"] += st["total"]
            split_totals[split]["bad"] += st["bad_read"]
            split_totals[split]["missing"] += st["missing_dir"]

    print("\nTotais por split:")
    print(f"{'split':<6} | {'arquivos':>7} | {'% preto':>8} | {'% branco':>8} | {'bad':>3} | {'missing':>7}")
    print("-" * 60)
    for split in ("train", "val", "test"):
        tt = split_totals[split]
        total = tt["total"]
        p_black = pct(tt["black"], total)
        p_white = pct(tt["white"], total)
        if tt["files"] == 0 and tt["missing"] == 0 and tt["bad"] == 0:
            continue
        print(f"{split:<6} | {tt['files']:>7} | {p_black:>7.2f}% | {p_white:>7.2f}% | {tt['bad']:>3} | {tt['missing']:>7}")

    print("\nDicas rápidas:")
    print("- Se % branco ~ 95-100% em quase tudo, o modelo aprender 'navegável em tudo' é o esperado.")
    print("- Se bad_read > 0, tem máscara corrompida/caminho estranho.")
    print("- Se missing > 0, tem grupo sem pasta masks2 (ou estrutura diferente).")
    print("")

if __name__ == "__main__":
    main()