# ------------------------------------------------------------------------
# RF-DETR
# Copyright (c) 2025 Roboflow. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------
"""Shared helpers for dataset augmentation configuration."""

from __future__ import annotations

from collections.abc import Callable
from typing import Any

# Transforms that include a horizontal-flip component. Applying these to keypoint
# data without swapping left/right joint pairs produces incorrect annotations.
HFLIP_TRANSFORM_NAMES: frozenset[str] = frozenset({"HorizontalFlip", "Flip", "D4"})

CONTAINER_TRANSFORM_NAMES: frozenset[str] = frozenset({"OneOf", "SomeOf", "Sequential"})

# Target-dict keys that describe the whole image rather than individual object instances.
# Both augmentation backends (torchvision-native in ``_torchvision.py`` and Albumentations in
# ``transforms.py``) treat these as global: they are never sliced by the per-instance keep mask
# when boxes are dropped. Shared here so the two pipelines stay in sync on this image-level
# subset while each keeps its own, deliberately different, per-instance-field policy (notably
# how each handles ``labels`` — see the per-instance filter helpers in both modules).
IMAGE_LEVEL_TARGET_FIELDS: frozenset[str] = frozenset({"orig_size", "size", "image_id"})


def resolve_keypoint_flip_pairs(args: Any, *, include_keypoints: bool) -> list[int] | None:
    """Resolve the ``keypoint_flip_pairs`` sentinel for augmentation-pipeline gating.

    ``None`` (not ``[]``) signals a detection pipeline to
    ``AlbumentationsWrapper.from_config``, which preserves horizontal-flip
    augmentations. A list (possibly empty) signals a keypoint pipeline; an empty
    list means "keypoint pipeline with no flip pairs defined", which strips hflip
    augmentations to prevent incorrect keypoint annotations. See #1243.

    Args:
        args: Argument namespace optionally carrying a ``keypoint_flip_pairs`` attribute.
        include_keypoints: Whether the augmentation config will be applied to keypoint data.

    Returns:
        ``None`` for detection-only pipelines, otherwise the resolved flip-pair
        list (possibly empty).
    """
    return (getattr(args, "keypoint_flip_pairs", []) or []) if include_keypoints else None


def _warn_keypoint_hflip_disabled(aug_name: str, warn: Callable[..., None]) -> None:
    """Emit the standard warning for a disabled keypoint horizontal flip."""
    warn(
        "Keypoint pipeline: '%s' performs a horizontal flip but no keypoint flip pairs "
        "were configured. The transform has been disabled to prevent incorrect keypoint "
        "annotations. Remove '%s' from your augmentation config or provide keypoint_flip_pairs.",
        aug_name,
        aug_name,
    )


def filter_keypoint_hflip_augmentations(
    config: Any,
    *,
    include_keypoints: bool = True,
    warn: Callable[..., None],
) -> Any:
    """Drop horizontal-flip transforms from keypoint augmentation configs.

    The helper preserves the input config shape: dictionary configs return
    dictionaries, and ordered list configs return lists. Container transforms are
    filtered recursively so nested hflip entries cannot survive inside
    ``OneOf``, ``SomeOf``, or ``Sequential``.

    Args:
        config: Augmentation config accepted by RF-DETR augmentation builders.
        include_keypoints: Whether the config will be applied to keypoint data.
        warn: Warning sink, typically ``logger.warning``.

    Returns:
        A filtered augmentation config with keypoint-unsafe hflip entries removed.
    """
    if not include_keypoints:
        return config

    if isinstance(config, list):
        return _filter_keypoint_hflip_entries(config, warn=warn)

    if isinstance(config, dict):
        return _filter_keypoint_hflip_dict(config, warn=warn)

    return config


def _filter_keypoint_hflip_entries(entries: list[dict[str, Any]], *, warn: Callable[..., None]) -> list[dict[str, Any]]:
    """Filter ordered single-key augmentation entries."""
    filtered_entries: list[dict[str, Any]] = []
    for entry in entries:
        if not isinstance(entry, dict) or len(entry) != 1:
            filtered_entries.append(entry)
            continue

        aug_name, params = next(iter(entry.items()))
        filtered_params = _filter_keypoint_hflip_item(aug_name, params, warn=warn)
        if isinstance(filtered_params, _DropAugmentation):
            continue
        filtered_entries.append({aug_name: filtered_params})
    return filtered_entries


def _filter_keypoint_hflip_dict(config: dict[str, Any], *, warn: Callable[..., None]) -> dict[str, Any]:
    """Filter a mapping-style augmentation config."""
    filtered_config: dict[str, Any] = {}
    for aug_name, params in config.items():
        filtered_params = _filter_keypoint_hflip_item(aug_name, params, warn=warn)
        if isinstance(filtered_params, _DropAugmentation):
            continue
        filtered_config[aug_name] = filtered_params
    return filtered_config


class _DropAugmentation:
    """Sentinel used when an augmentation entry should be removed."""


def _filter_keypoint_hflip_item(
    aug_name: str,
    params: Any,
    *,
    warn: Callable[..., None],
) -> Any | _DropAugmentation:
    """Filter one augmentation entry, recursing into supported containers."""
    if aug_name in HFLIP_TRANSFORM_NAMES:
        _warn_keypoint_hflip_disabled(aug_name, warn)
        return _DropAugmentation()

    if aug_name not in CONTAINER_TRANSFORM_NAMES:
        return params

    if isinstance(params, list):
        filtered_transforms = _filter_keypoint_hflip_entries(params, warn=warn)
        return filtered_transforms if filtered_transforms else _DropAugmentation()

    if not isinstance(params, dict) or "transforms" not in params:
        return params

    transforms = params["transforms"]
    if not isinstance(transforms, list):
        return params

    filtered_transforms = _filter_keypoint_hflip_entries(transforms, warn=warn)
    if not filtered_transforms:
        return _DropAugmentation()

    return {**params, "transforms": filtered_transforms}