from __future__ import annotations

"""Filter GTA-UAV-LR images by segmentation class coverage.

Reads palette-mode PNG segmentation masks, computes per-class pixel ratios,
and outputs a JSON meta file listing images that pass the filter
(i.e. background + water < threshold).

Classes (from manifest.json):
    0: background, 1: building, 2: road, 3: vegetation, 4: water, ...

Usage:
    python -m scripts.filter_segmentation [--threshold 0.9] [--output meta/seg_filter.json]
"""

import argparse
import json
import logging
from pathlib import Path

import coloredlogs
import numpy as np
from PIL import Image
from tqdm import tqdm

LOGGER = logging.getLogger("caption_test.filter_seg")

SEGM_ROOT = Path("/home/servml/Документы/datasets/GTA-UAV-LR-aug/segm")
EXCLUDE_CLASSES = {0, 4}  # background, water
DEFAULT_THRESHOLD = 0.90


def compute_class_ratios(mask_path: Path) -> dict[int, float]:
    """Load a palette-mode PNG mask and return per-class pixel ratios."""
    with Image.open(mask_path) as img:
        arr = np.array(img)
    total = arr.size
    unique, counts = np.unique(arr, return_counts=True)
    return {int(cls): float(cnt / total) for cls, cnt in zip(unique, counts)}


def scan_masks(
    segm_root: Path,
    exclude_classes: set[int],
    threshold: float,
) -> dict[str, dict]:
    """Scan all segmentation masks and classify as pass/fail.

    Returns:
        Dict keyed by relative image name (e.g. "drone/images/100_0001_0000000000.png")
        with values: {"ratios": {...}, "excluded_ratio": float, "pass": bool}
    """
    results: dict[str, dict] = {}
    subdirs = sorted(p for p in segm_root.iterdir() if p.is_dir())

    for subdir in subdirs:
        mask_files = sorted(subdir.rglob("*.png"))
        LOGGER.info("🔍 Scanning %s: %d masks", subdir.name, len(mask_files))

        for mask_path in tqdm(mask_files, desc=f"  📂 {subdir.name}", unit="img", leave=False):
            rel = mask_path.relative_to(segm_root)
            ratios = compute_class_ratios(mask_path)
            excluded_ratio = sum(ratios.get(c, 0.0) for c in exclude_classes)
            results[str(rel)] = {
                "ratios": ratios,
                "excluded_ratio": round(excluded_ratio, 6),
                "pass": excluded_ratio < threshold,
            }

    return results


def build_meta(
    results: dict[str, dict],
) -> dict[str, list[str]]:
    """Split results into passed and excluded lists."""
    passed = sorted(k for k, v in results.items() if v["pass"])
    excluded = sorted(k for k, v in results.items() if not v["pass"])
    return {"passed": passed, "excluded": excluded}


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Filter GTA-UAV-LR images by segmentation coverage.",
    )
    parser.add_argument(
        "--segm-root", type=str, default=str(SEGM_ROOT),
        help="Root of segmentation masks.",
    )
    parser.add_argument(
        "--threshold", type=float, default=DEFAULT_THRESHOLD,
        help="Exclude images where background+water >= threshold.",
    )
    parser.add_argument(
        "--output", type=str, default="meta/seg_filter.json",
        help="Output JSON meta file path.",
    )
    args = parser.parse_args()

    coloredlogs.install(
        level="INFO",
        logger=LOGGER,
        fmt="%(asctime)s %(name)s %(levelname)s %(message)s",
    )

    LOGGER.info("🚀 Starting segmentation filter (threshold=%.2f)", args.threshold)

    segm_root = Path(args.segm_root)
    results = scan_masks(segm_root, EXCLUDE_CLASSES, args.threshold)
    meta = build_meta(results)

    n_total = len(results)
    n_pass = len(meta["passed"])
    n_excl = len(meta["excluded"])
    LOGGER.info(
        "📊 total=%d  ✅ passed=%d (%.1f%%)  ❌ excluded=%d (%.1f%%)",
        n_total, n_pass, 100 * n_pass / max(n_total, 1),
        n_excl, 100 * n_excl / max(n_total, 1),
    )

    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    out = {
        "threshold": args.threshold,
        "exclude_classes": sorted(EXCLUDE_CLASSES),
        "total_images": n_total,
        "passed_count": n_pass,
        "excluded_count": n_excl,
        "passed": meta["passed"],
        "excluded": meta["excluded"],
    }
    with output_path.open("w", encoding="utf-8") as f:
        json.dump(out, f, indent=2)

    LOGGER.info("💾 Meta file saved to %s", output_path)


if __name__ == "__main__":
    main()