from pathlib import Path
from collections import defaultdict
import argparse
import math

IMG_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}

def is_float(s: str) -> bool:
    try:
        float(s)
        return True
    except:
        return False

def parse_label_line(line: str):
    """
    Suporta:
    - YOLO det:  cls xc yc w h
    - YOLO seg:  cls x1 y1 x2 y2 ...
    Retorna dict com:
      cls (int), kind ("det"|"seg"|None), n_points (int|None)
    """
    parts = line.strip().split()
    if len(parts) < 2:
        return None

    if not is_float(parts[0]):
        return None
    cls = int(float(parts[0]))

    nums = []
    for p in parts[1:]:
        if not is_float(p):
            return None
        nums.append(float(p))

    # det clássico: 4 números
    if len(nums) == 4:
        return {"cls": cls, "kind": "det", "n_points": None}

    # seg: pares (x,y)
    if len(nums) >= 6 and (len(nums) % 2 == 0):
        return {"cls": cls, "kind": "seg", "n_points": len(nums) // 2}

    # caso estranho
    return {"cls": cls, "kind": "unknown", "n_points": None}

def analyze_split(labels_dir: Path):
    stats = {
        "images_total": 0,
        "images_bg": 0,
        "images_with_obj": 0,

        "images_per_class": defaultdict(int),     # quantas imagens têm a classe
        "instances_per_class": defaultdict(int),  # quantas instâncias (linhas) por classe

        "kind_counts": defaultdict(int),          # det/seg/unknown
        "poly_points": [],                        # lista de n_points (para seg)
        "weird_lines": 0,
        "empty_label_files": 0,
    }

    if not labels_dir.exists():
        return stats

    label_files = sorted(labels_dir.glob("*.txt"))
    stats["images_total"] = len(label_files)

    for lf in label_files:
        text = lf.read_text(encoding="utf-8", errors="ignore").strip()
        if not text:
            stats["images_bg"] += 1
            stats["empty_label_files"] += 1
            continue

        stats["images_with_obj"] += 1

        classes_in_image = set()

        for line in text.splitlines():
            parsed = parse_label_line(line)
            if parsed is None:
                stats["weird_lines"] += 1
                continue

            cls = parsed["cls"]
            kind = parsed["kind"]
            stats["kind_counts"][kind] += 1

            classes_in_image.add(cls)
            stats["instances_per_class"][cls] += 1

            if kind == "seg" and parsed["n_points"] is not None:
                stats["poly_points"].append(parsed["n_points"])

        for cls in classes_in_image:
            stats["images_per_class"][cls] += 1

    return stats

def summarize(stats):
    out = []
    out.append(f"images_total     : {stats['images_total']}")
    out.append(f"images_bg        : {stats['images_bg']} ({pct(stats['images_bg'], stats['images_total'])})")
    out.append(f"images_with_obj  : {stats['images_with_obj']} ({pct(stats['images_with_obj'], stats['images_total'])})")
    out.append(f"empty_label_files: {stats['empty_label_files']}")
    out.append(f"weird_lines      : {stats['weird_lines']}")

    if stats["kind_counts"]:
        out.append("label_kinds      : " + ", ".join(f"{k}={v}" for k, v in sorted(stats["kind_counts"].items())))

    # imagens por classe
    if stats["images_per_class"]:
        out.append("images_per_class : " + ", ".join(f"c{c}={n}" for c, n in sorted(stats["images_per_class"].items())))
    else:
        out.append("images_per_class : (nenhuma)")

    # instâncias por classe
    if stats["instances_per_class"]:
        out.append("inst_per_class   : " + ", ".join(f"c{c}={n}" for c, n in sorted(stats["instances_per_class"].items())))
    else:
        out.append("inst_per_class   : (nenhuma)")

    # polígonos
    if stats["poly_points"]:
        pts = stats["poly_points"]
        out.append(f"seg_poly_points  : count={len(pts)} min={min(pts)} mean={sum(pts)/len(pts):.2f} max={max(pts)}")
    else:
        out.append("seg_poly_points  : (n/a)")

    return "\n".join(out)

def pct(a, b):
    if b == 0:
        return "n/a"
    return f"{(100.0*a/b):.1f}%"

def merge_stats(a, b):
    """merge b into a"""
    a["images_total"] += b["images_total"]
    a["images_bg"] += b["images_bg"]
    a["images_with_obj"] += b["images_with_obj"]
    a["empty_label_files"] += b["empty_label_files"]
    a["weird_lines"] += b["weird_lines"]

    for k, v in b["images_per_class"].items():
        a["images_per_class"][k] += v
    for k, v in b["instances_per_class"].items():
        a["instances_per_class"][k] += v
    for k, v in b["kind_counts"].items():
        a["kind_counts"][k] += v
    a["poly_points"].extend(b["poly_points"])
    return a

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", type=str, required=True,
                    help="Pasta base que contém labels/train labels/val labels/test (ou labels direto)")
    args = ap.parse_args()

    base = Path(args.base)
    labels = base / "labels"

    splits = []
    # padrão YOLO: labels/train, labels/val, labels/test
    if (labels / "train").exists() or (labels / "val").exists() or (labels / "test").exists():
        splits = ["train", "val", "test"]
        split_dirs = {s: labels / s for s in splits}
    else:
        # fallback: base/labels direto
        split_dirs = {"all": labels}

    global_stats = {
        "images_total": 0,
        "images_bg": 0,
        "images_with_obj": 0,
        "images_per_class": defaultdict(int),
        "instances_per_class": defaultdict(int),
        "kind_counts": defaultdict(int),
        "poly_points": [],
        "weird_lines": 0,
        "empty_label_files": 0,
    }

    for name, d in split_dirs.items():
        st = analyze_split(d)
        print("===================================")
        print(f"SPLIT: {name}  ({d})")
        print(summarize(st))
        print("===================================")
        merge_stats(global_stats, st)

    print("\n========== GLOBAL ==========")
    print(summarize(global_stats))

if __name__ == "__main__":
    main()