#!/usr/bin/env python3
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
"""Clean ground 2D dataset labels and corrupt image pairs.

This script matches the duplicate-label definition used by
``ultralytics.data.utils.verify_image_label_ground``:

    [class_id, x_center, y_center, width, height, difficulty1 + difficulty2]

Rows that map to the same parsed tuple are duplicates; the first raw row is kept.
By default the script is a dry run. Pass ``--apply`` to rewrite label files,
delete corrupt image/label pairs, and update split list files.
"""

from __future__ import annotations

import argparse
import os
from collections import OrderedDict
from datetime import datetime
from pathlib import Path

import numpy as np
from PIL import Image

try:
    from tqdm import tqdm
except ImportError:
    tqdm = None

IMG_FORMATS = {
    "avif",
    "bmp",
    "dng",
    "heic",
    "heif",
    "jp2",
    "jpeg",
    "jpeg2000",
    "jpg",
    "mpo",
    "png",
    "tif",
    "tiff",
    "webp",
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--data",
        default="ultralytics/cfg/datasets/mono2d_ground.yaml",
        help="Ground 2D dataset YAML path.",
    )
    parser.add_argument(
        "--splits",
        nargs="+",
        default=["train", "val"],
        help="Dataset YAML split keys to scan, for example: train val test.",
    )
    parser.add_argument(
        "--apply",
        action="store_true",
        help="Actually modify files. Without this flag, only prints what would change.",
    )
    parser.add_argument(
        "--no-update-lists",
        action="store_true",
        help="Do not remove corrupt entries from split list files.",
    )
    parser.add_argument(
        "--no-backup",
        action="store_true",
        help="Do not create .bak files before rewriting labels/list files.",
    )
    parser.add_argument(
        "--report",
        default="",
        help="Optional report file path. Defaults to no report file.",
    )
    parser.add_argument(
        "--progress-interval",
        type=int,
        default=100,
        help="Print progress every N entries when tqdm is not installed. Set 0 to disable fallback progress.",
    )
    return parser.parse_args()


def read_lines(path: Path) -> list[str]:
    return path.read_text(encoding="utf-8").splitlines()


def strip_yaml_comment(value: str) -> str:
    """Strip simple YAML comments while preserving path strings."""
    return value.split("#", 1)[0].strip()


def parse_scalar(value: str) -> object:
    value = strip_yaml_comment(value)
    if not value:
        return None
    if value.lower() in {"true", "false"}:
        return value.lower() == "true"
    if value.startswith("[") and value.endswith("]"):
        items = [x.strip() for x in value[1:-1].split(",") if x.strip()]
        parsed = []
        for item in items:
            try:
                parsed.append(float(item))
            except ValueError:
                parsed.append(item.strip("'\""))
        return parsed
    try:
        return int(value)
    except ValueError:
        try:
            return float(value)
        except ValueError:
            return value.strip("'\"")


def load_ground_yaml(path: Path) -> dict[str, object]:
    """Load the small YAML subset used by mono2d_ground.yaml without external dependencies."""
    data: dict[str, object] = {}
    current_map: str | None = None
    for raw_line in read_lines(path):
        if not raw_line.strip() or raw_line.lstrip().startswith("#"):
            continue
        indent = len(raw_line) - len(raw_line.lstrip(" "))
        line = raw_line.strip()
        if indent == 0:
            key, sep, value = line.partition(":")
            if not sep:
                continue
            parsed = parse_scalar(value)
            if parsed is None:
                data[key] = {}
                current_map = key
            else:
                data[key] = parsed
                current_map = None
        elif current_map and isinstance(data.get(current_map), dict):
            key, sep, value = line.partition(":")
            if sep:
                data[current_map][key.strip()] = parse_scalar(value)
    return data


def img2label_path(img_path: str) -> str:
    sa, sb = f"{os.sep}images{os.sep}", f"{os.sep}labels{os.sep}"
    return sb.join(img_path.rsplit(sa, 1)).rsplit(".", 1)[0] + ".txt"


def write_lines(path: Path, lines: list[str], backup: bool) -> None:
    if backup and path.exists():
        stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        backup_path = path.with_name(f"{path.name}.bak.{stamp}")
        backup_path.write_text(path.read_text(encoding="utf-8"), encoding="utf-8")
    path.write_text("\n".join(lines) + ("\n" if lines else ""), encoding="utf-8")


def split_entries(split_path: Path) -> list[str]:
    return [line.strip() for line in read_lines(split_path) if line.strip()]


def expand_split_entry(entry: str, split_path: Path) -> str:
    """Mirror YOLOGroundDataset.get_img_files list-file expansion."""
    if entry.startswith("./"):
        return entry.replace("./", str(split_path.parent) + os.sep, 1)
    return entry


def resolve_split_path(value: str | None, data_dir: Path) -> Path | None:
    if not value:
        return None
    path = Path(value)
    return path if path.is_absolute() else data_dir / path


def ground_image_path_from_label_path(path: str, image_root: Path, label_path: bool = False) -> Path:
    """Mirror YOLOGroundDataset._ground_image_path_from_label_path."""
    parts = Path(path).parts
    detection_root_index = next(
        (i for i, part in enumerate(parts) if part.lower().startswith(("detection2d", "2ddetection"))),
        None,
    )
    if detection_root_index is not None and detection_root_index + 1 < len(parts):
        rel_parts = list(parts[detection_root_index + 1 :])
    else:
        rel_path = Path(path) if not Path(path).is_absolute() else Path(*parts[1:])
        rel_parts = list(rel_path.parts)

    if label_path:
        image_path = image_root.joinpath(*rel_parts).with_suffix(".jpg")
        if image_path.exists():
            return image_path
        image_rel_parts = ["images" if part == "labels" else part for part in rel_parts]
        image_path = image_root.joinpath(*image_rel_parts).with_suffix(".jpg")
        if image_path.exists():
            return image_path
        for suffix in ("png", *sorted(IMG_FORMATS - {"jpg", "jpeg", "png"}), "jpeg"):
            candidate = image_path.with_suffix(f".{suffix}")
            if candidate.exists():
                return candidate
        return image_path

    return image_root.joinpath(*rel_parts)


def label_path_for_entry(entry: str) -> Path:
    suffix = Path(entry).suffix.lower().lstrip(".")
    if suffix == "txt":
        return Path(entry)
    return Path(img2label_path(entry))


def parse_known_label_key(parts: list[str], class_map: dict[str, int]) -> tuple[float, ...] | None:
    if len(parts) < 7 or parts[0] not in class_map:
        return None
    coords = [float(x) for x in parts[1:5]]
    difficulty = float(parts[5]) + float(parts[6])
    row = np.array([class_map[parts[0]], *coords, difficulty], dtype=np.float32)
    return tuple(float(x) for x in row)


def deduplicate_label_file(label_path: Path, class_map: dict[str, int]) -> tuple[int, list[str], list[str]]:
    """Return duplicate count, cleaned raw lines, and warnings."""
    if not label_path.exists():
        return 0, [], []

    raw_lines = read_lines(label_path)
    seen: OrderedDict[tuple[float, ...], int] = OrderedDict()
    keep = [True] * len(raw_lines)
    warnings = []

    for idx, line in enumerate(raw_lines):
        stripped = line.strip()
        if not stripped:
            continue
        parts = stripped.split()
        try:
            key = parse_known_label_key(parts, class_map)
        except ValueError as exc:
            warnings.append(f"{label_path}:{idx + 1}: cannot parse label ({exc})")
            continue
        if key is None:
            continue
        if key in seen:
            keep[idx] = False
        else:
            seen[key] = idx

    cleaned = [line for line, should_keep in zip(raw_lines, keep) if should_keep]
    return len(raw_lines) - len(cleaned), cleaned, warnings


def exif_size(img: Image.Image) -> tuple[int, int]:
    size = img.size
    if img.format == "JPEG":
        try:
            if exif := img.getexif():
                if exif.get(274, None) in {6, 8}:
                    size = size[1], size[0]
        except Exception:
            pass
    return size


def image_corruption_reason(image_path: Path) -> str | None:
    try:
        with Image.open(image_path) as im:
            im.verify()
            shape = exif_size(im)
            if shape[0] <= 9 or shape[1] <= 9:
                return f"image size {shape} <10 pixels"
            if (im.format or "").lower() not in IMG_FORMATS:
                return f"invalid image format {im.format}"
        if image_path.suffix.lower() in {".jpg", ".jpeg"}:
            with image_path.open("rb") as f:
                f.seek(-2, 2)
                if f.read() != b"\xff\xd9":
                    return "corrupt JPEG end marker"
    except Exception as exc:
        return str(exc)
    return None


def unlink_if_exists(path: Path) -> bool:
    if path.exists():
        path.unlink()
        return True
    return False


def progress_iter(items: list[str], desc: str, interval: int):
    total = len(items)
    if tqdm is not None:
        yield from tqdm(items, total=total, desc=desc, unit="entry")
        return

    for idx, item in enumerate(items, 1):
        if interval > 0 and (idx == 1 or idx % interval == 0 or idx == total):
            print(f"{desc}: {idx}/{total}", flush=True)
        yield item


def main() -> int:
    args = parse_args()
    data_file = Path(args.data)
    data = load_ground_yaml(data_file)
    data_dir = data_file.parent
    image_root = Path(data["path"])
    class_map = data.get("class_map", {})
    if not class_map:
        raise SystemExit(f"{data_file} does not contain class_map; this cleaner is for ground 2D labels.")

    update_lists = args.apply and not args.no_update_lists
    backup = not args.no_backup
    report_lines = []
    label_files_seen: set[Path] = set()
    stats = {
        "images": 0,
        "labels": 0,
        "duplicate_rows": 0,
        "labels_rewritten": 0,
        "corrupt_images": 0,
        "images_removed": 0,
        "labels_removed": 0,
        "list_entries_removed": 0,
    }

    for split in args.splits:
        split_path = resolve_split_path(data.get(split), data_dir)
        if split_path is None:
            continue
        if not split_path.exists():
            report_lines.append(f"[missing split] {split}: {split_path}")
            continue

        entries = split_entries(split_path)
        kept_entries = []
        removed_from_split = 0

        for entry in progress_iter(entries, f"Scanning {split}", args.progress_interval):
            expanded_entry = expand_split_entry(entry, split_path)
            suffix = Path(expanded_entry).suffix.lower().lstrip(".")
            if suffix in IMG_FORMATS:
                image_path = ground_image_path_from_label_path(expanded_entry, image_root, label_path=False)
                label_path = label_path_for_entry(expanded_entry)
            elif suffix == "txt":
                label_path = Path(expanded_entry)
                image_path = ground_image_path_from_label_path(expanded_entry, image_root, label_path=True)
            else:
                report_lines.append(f"[unsupported entry] {split}: {entry}")
                kept_entries.append(entry)
                continue

            stats["images"] += 1
            reason = image_corruption_reason(image_path)
            if reason:
                stats["corrupt_images"] += 1
                removed_from_split += 1
                report_lines.append(f"[corrupt] {image_path} | {reason}")
                report_lines.append(f"          label: {label_path}")
                if args.apply:
                    if unlink_if_exists(image_path):
                        stats["images_removed"] += 1
                    if unlink_if_exists(label_path):
                        stats["labels_removed"] += 1
                continue

            kept_entries.append(entry)
            if label_path in label_files_seen:
                continue
            label_files_seen.add(label_path)
            if label_path.exists():
                stats["labels"] += 1
            dup_count, cleaned_lines, warnings = deduplicate_label_file(label_path, class_map)
            report_lines.extend(f"[warning] {warning}" for warning in warnings)
            if dup_count:
                stats["duplicate_rows"] += dup_count
                report_lines.append(f"[duplicates] {label_path}: remove {dup_count}")
                if args.apply:
                    write_lines(label_path, cleaned_lines, backup=backup)
                    stats["labels_rewritten"] += 1

        if update_lists and removed_from_split:
            write_lines(split_path, kept_entries, backup=backup)
            stats["list_entries_removed"] += removed_from_split
            report_lines.append(f"[list updated] {split_path}: removed {removed_from_split} corrupt entries")

    mode = "APPLIED" if args.apply else "DRY RUN"
    summary = [
        f"Mode: {mode}",
        f"Images checked: {stats['images']}",
        f"Label files checked: {stats['labels']}",
        f"Duplicate label rows found: {stats['duplicate_rows']}",
        f"Label files rewritten: {stats['labels_rewritten']}",
        f"Corrupt images found: {stats['corrupt_images']}",
        f"Images removed: {stats['images_removed']}",
        f"Labels removed with corrupt images: {stats['labels_removed']}",
        f"Split list entries removed: {stats['list_entries_removed']}",
    ]

    output = "\n".join(summary + (["", *report_lines] if report_lines else []))
    print(output)
    if args.report:
        report_path = Path(args.report)
        report_path.parent.mkdir(parents=True, exist_ok=True)
        report_path.write_text(output + "\n", encoding="utf-8")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())