import os

# ─────────────────────────────────────────────────────────────────────────────
# ✏️  EDIT THIS PATH
# ─────────────────────────────────────────────────────────────────────────────
DATASET_DIR = r"D:\merged_dataset"
NC          = 20   # total number of classees
# ─────────────────────────────────────────────────────────────────────────────

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

def verify_split(split):
    img_dir = os.path.join(DATASET_DIR, split, "images")
    lbl_dir = os.path.join(DATASET_DIR, split, "labels")

    errors   = []
    warnings = []

    if not os.path.isdir(img_dir):
        print(f"  [SKIP] {split}/images folder not found — skipping")
        return

    img_files = {os.path.splitext(f)[0]: f
                 for f in os.listdir(img_dir)
                 if os.path.splitext(f)[1].lower() in IMG_EXTS}

    lbl_files = {os.path.splitext(f)[0]: f
                 for f in os.listdir(lbl_dir)
                 if f.endswith(".txt")} if os.path.isdir(lbl_dir) else {}

    total_images = len(img_files)
    total_labels = len(lbl_files)
    total_boxes  = 0

    # ── Check 1: every image has a label ─────────────────────────────────────
    for stem in img_files:
        if stem not in lbl_files:
            errors.append(f"  [NO LABEL]  {img_files[stem]}")

    # ── Check 2: every label has an image ────────────────────────────────────
    for stem in lbl_files:
        if stem not in img_files:
            warnings.append(f"  [NO IMAGE]  {lbl_files[stem]}")

    # ── Check 3–7: validate label content ────────────────────────────────────
    for stem, lbl_fname in lbl_files.items():
        lbl_path = os.path.join(lbl_dir, lbl_fname)
        try:
            with open(lbl_path, "r") as f:
                lines = [l.strip() for l in f.readlines() if l.strip()]
        except Exception as e:
            errors.append(f"  [READ ERROR] {lbl_fname}: {e}")
            continue

        if len(lines) == 0:
            warnings.append(f"  [EMPTY]     {lbl_fname} — no annotations")
            continue

        for i, line in enumerate(lines, 1):
            parts = line.split()

            # Check 4: must have exactly 5 values
            if len(parts) != 5:
                errors.append(
                    f"  [BAD FORMAT] {lbl_fname} line {i}: "
                    f"expected 5 values, got {len(parts)} → '{line}'"
                )
                continue

            try:
                cls_id = int(parts[0])
                x, y, w, h = float(parts[1]), float(parts[2]), \
                              float(parts[3]), float(parts[4])
            except ValueError:
                errors.append(
                    f"  [NOT NUMERIC] {lbl_fname} line {i}: '{line}'"
                )
                continue

            # Check 3: class ID in range
            if cls_id < 0 or cls_id >= NC:
                errors.append(
                    f"  [BAD CLASS ID] {lbl_fname} line {i}: "
                    f"class={cls_id} (valid range 0–{NC-1})"
                )

            # Check 5: bbox values in [0, 1]
            for val_name, val in [("x", x), ("y", y), ("w", w), ("h", h)]:
                if not (0.0 <= val <= 1.0):
                    errors.append(
                        f"  [OUT OF RANGE] {lbl_fname} line {i}: "
                        f"{val_name}={val} (must be 0.0–1.0)"
                    )

            total_boxes += 1

    # ── Print split report ────────────────────────────────────────────────────
    status = "PASS" if not errors else "FAIL"
    print(f"\n  [{status}] {split}/")
    print(f"         images : {total_images}")
    print(f"         labels : {total_labels}")
    print(f"         boxes  : {total_boxes}")

    if warnings:
        print(f"         warnings ({len(warnings)}):")
        for w in warnings[:20]:
            print(w)
        if len(warnings) > 20:
            print(f"         ... and {len(warnings)-20} more warnings")

    if errors:
        print(f"         errors ({len(errors)}):")
        for e in errors[:30]:
            print(e)
        if len(errors) > 30:
            print(f"         ... and {len(errors)-30} more errors")
    else:
        print("         No errors found!")

    return len(errors)


def main():
    print("=" * 70)
    print("  Label verification report")
    print(f"  Dataset : {DATASET_DIR}")
    print(f"  nc      : {NC} classes (valid IDs: 0 – {NC-1})")
    print("=" * 70)

    total_errors = 0
    for split in ("train", "valid", "test"):
        result = verify_split(split)
        if result:
            total_errors += result

    print("\n" + "=" * 70)
    if total_errors == 0:
        print("  ALL CHECKS PASSED — dataset is ready for training!")
    else:
        print(f"  TOTAL ERRORS: {total_errors} — fix these before training.")
    print("=" * 70)

    # ── Class ID distribution ─────────────────────────────────────────────────
    print("\n  Class ID distribution across entire dataset:")
    class_counts = {i: 0 for i in range(NC)}
    for split in ("train", "valid", "test"):
        lbl_dir = os.path.join(DATASET_DIR, split, "labels")
        if not os.path.isdir(lbl_dir):
            continue
        for fname in os.listdir(lbl_dir):
            if not fname.endswith(".txt"):
                continue
            with open(os.path.join(lbl_dir, fname), "r") as f:
                for line in f:
                    parts = line.strip().split()
                    if parts:
                        try:
                            class_counts[int(parts[0])] += 1
                        except (ValueError, KeyError):
                            pass

    class_names = [
        'Mask','can','cellphone','electronics','gbottle','glove','metal',
        'misc','net','pbag','pbottle','plastic','rod','sunglasses','tire',
        'Microplastic','fiber','film','fragment','pallet'
    ]

    print(f"\n  {'ID':>3}  {'Class':<15} {'Boxes':>8}")
    print(f"  {'─'*3}  {'─'*15} {'─'*8}")
    for i in range(NC):
        name  = class_names[i] if i < len(class_names) else f"class_{i}"
        count = class_counts[i]
        flag  = "  ← ZERO annotations!" if count == 0 else ""
        print(f"  {i:>3}  {name:<15} {count:>8}{flag}")


if __name__ == "__main__":
    main()