# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1

import math
import re
from typing import List, Tuple

IMAGE_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = {
    # Our desired 256 resolution is the one below (commented).
    # Desired: "256": {"1,1": (336, 336), "4,3": (384, 288), "3,4": (288, 384), "16,9": (448, 256), "9,16": (256, 448)},
    "256": {
        "1,1": (256, 256),
        "4,3": (320, 256),
        "3,4": (256, 320),
        "16,9": (320, 192),
        "9,16": (192, 320),
    },
    "480": {"1,1": (640, 640), "4,3": (736, 544), "3,4": (544, 736), "16,9": (832, 480), "9,16": (480, 832)},
    # 704 resolutions are nicely divisible by 32
    "704": {"1,1": (960, 960), "4,3": (1088, 832), "3,4": (832, 1088), "16,9": (1280, 704), "9,16": (704, 1280)},
    "720": {"1,1": (960, 960), "4,3": (1104, 832), "3,4": (832, 1104), "16,9": (1280, 720), "9,16": (720, 1280)},
    # 768 for arena.ai
    "768": {"1,1": (1024, 1024), "4,3": (1184, 880), "3,4": (880, 1184), "16,9": (1360, 768), "9,16": (768, 1360)},
    "1080": {"1,1": (1440, 1440), "4,3": (1664, 1248), "3,4": (1248, 1664), "16,9": (1920, 1080), "9,16": (1080, 1920)},
    "1280": {"1,1": (1712, 1712), "4,3": (1968, 1472), "3,4": (1472, 1968), "16,9": (2272, 1280), "9,16": (1280, 2272)},
    "2048": {
        "1,1": (2728, 2728),
        "4,3": (3160, 2368),
        "3,4": (2368, 3160),
        "16,9": (3640, 2048),
        "9,16": (2048, 3640),
    },
    "gt_2048": {
        "1,1": (5464, 5464),
        "4,3": (6304, 4728),
        "3,4": (4728, 6304),
        "16,9": (7280, 4096),
        "9,16": (4096, 7280),
    },
}

VIDEO_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = {
    # Our desired 256 resolution is the one below (commented).
    # Desired: "256": {"1,1": (336, 336), "4,3": (384, 288), "3,4": (288, 384), "16,9": (448, 256), "9,16": (256, 448)},
    "256": {
        "1,1": (256, 256),
        "4,3": (320, 256),
        "3,4": (256, 320),
        "16,9": (320, 192),
        "9,16": (192, 320),
    },
    # "2,3" (h/w=1.5, W:H=2:3): added for the BEHAVIOR-1K R1Pro concat_view composite
    # (native 1080x720), so it snaps to an exact-aspect 512x768 canvas with no
    # reflection padding. Scoped to the 480 tier that b1k trains at.
    "480": {
        "1,1": (640, 640),
        "4,3": (736, 544),
        "3,4": (544, 736),
        "16,9": (832, 480),
        "9,16": (480, 832),
        "2,3": (512, 768),
    },
    # 704 resolutions are nicely divisible by 32
    "704": {"1,1": (960, 960), "4,3": (1088, 832), "3,4": (832, 1088), "16,9": (1280, 704), "9,16": (704, 1280)},
    "720": {"1,1": (960, 960), "4,3": (1104, 832), "3,4": (832, 1104), "16,9": (1280, 720), "9,16": (720, 1280)},
    # 768 for arena.ai
    "768": {"1,1": (1024, 1024), "4,3": (1184, 880), "3,4": (880, 1184), "16,9": (1360, 768), "9,16": (768, 1360)},
    "1080": {"1,1": (1440, 1440), "4,3": (1664, 1248), "3,4": (1248, 1664), "16,9": (1920, 1080), "9,16": (1080, 1920)},
    "1280": {"1,1": (1712, 1712), "4,3": (1968, 1472), "3,4": (1472, 1968), "16,9": (2272, 1280), "9,16": (1280, 2272)},
    "2048": {
        "1,1": (2728, 2728),
        "4,3": (3160, 2368),
        "3,4": (2368, 3160),
        "16,9": (3640, 2048),
        "9,16": (2048, 3640),
    },
    "gt_2048": {
        "1,1": (5464, 5464),
        "4,3": (6304, 4728),
        "3,4": (4728, 6304),
        "16,9": (7280, 4096),
        "9,16": (4096, 7280),
    },
}


def get_aspect_ratios_from_wdinfos(wdinfos: list[str]) -> list[str]:
    aspect_ratios = []
    for wdinfo in wdinfos:
        aspect_ratio_match = re.search(r"aspect_ratio_(\d+_\d+)", wdinfo)
        aspect_ratios.append(aspect_ratio_match.group(1))

    return aspect_ratios


def get_wdinfos_w_aspect_ratio(wdinfos: list[str]) -> List[Tuple[str, str]]:
    aspect_ratios = get_aspect_ratios_from_wdinfos(wdinfos)

    # return a list of (wdinfo_path, aspect_ratio) pairs
    return [(wdinfo, aspect_ratio.replace("_", ",")) for wdinfo, aspect_ratio in zip(wdinfos, aspect_ratios)]


def parse_frame_range_from_wdinfo(wdinfo: str) -> tuple[int, int | float] | None:
    """
    Parse frame range from wdinfo path.

    Args:
        wdinfo: wdinfo path string containing a frames_X_Y pattern, where Y may be ``inf``

    Returns:
        Tuple of (min_frames, max_frames) if found, where max_frames is
        ``math.inf`` for unbounded ``frames_X_inf`` and ``frames_gt_X`` buckets.
        ``frames_gt_X`` is treated as starting at X. Returns None if no frame
        range is present.

    Example:
        >>> parse_frame_range_from_wdinfo("wdinfo/v4/tv_drama/resolution_720/aspect_ratio_16_9/frames_300_400/wdinfo.json")
        (300, 400)
        >>> parse_frame_range_from_wdinfo("wdinfo/v4/tv_drama/resolution_720/aspect_ratio_16_9/frames_3700_inf/wdinfo.json")
        (3700, inf)
    """
    match = re.search(r"frames_(\d+)_(\d+|inf)(?!\w)", wdinfo)
    if match:
        max_frames = math.inf if match.group(2) == "inf" else int(match.group(2))
        return (int(match.group(1)), max_frames)
    match = re.search(r"frames_gt_(\d+)", wdinfo)
    if match:
        # Wdinfo generation uses frames_gt_X for the bucket that starts at X, so treat it as >= X.
        return (int(match.group(1)), math.inf)
    return None


def normalize_skip_frame_ranges(
    skip_frame_range: str | list[str] | None,
) -> set[tuple[int, int | float]]:
    """Normalize ``skip_frame_range`` into a set of (min_frames, max_frames) buckets.

    Args:
        skip_frame_range: A single bucket string like ``"300_400"`` or ``"3700_inf"``,
            a list of such strings, or None. Each string identifies the frame-range
            bucket (e.g. ``frames_300_400``) that should be skipped.

    Returns:
        Set of (min_frames, max_frames) tuples to skip. An ``inf`` upper bound is
        represented by ``math.inf``. Empty if ``skip_frame_range`` is None.
    """
    if skip_frame_range is None:
        return set()

    if isinstance(skip_frame_range, str):
        skip_frame_range = [skip_frame_range]

    skip_buckets: set[tuple[int, int | float]] = set()
    for bucket in skip_frame_range:
        match = re.fullmatch(r"(\d+)_(\d+|inf)", bucket.strip())
        if match is None:
            raise ValueError(
                f"Invalid skip_frame_range entry {bucket!r}. "
                "Expected the form '<min>_<max>', e.g. '300_400' or '3700_inf'."
            )
        max_frames = math.inf if match.group(2) == "inf" else int(match.group(2))
        skip_buckets.add((int(match.group(1)), max_frames))

    return skip_buckets


def filter_wdinfos_by_frame_range(
    wdinfos: list[str],
    min_frames: int | None = None,
    max_frames: int | None = None,
    skip_frame_range: str | list[str] | None = None,
) -> list[str]:
    """
    Filter wdinfo files based on frame range.

    The frame range in wdinfo path (e.g., frames_300_400) represents videos
    with frames between those values. This function filters wdinfo files
    based on the wdinfo's upper bound (wdinfo_max):
    - min_frames is EXCLUSIVE: wdinfo_max must be > min_frames
    - max_frames is INCLUSIVE: wdinfo_max must be <= max_frames
    - an ``inf`` upper bound is treated as infinity and excluded by any finite max_frames

    Additionally, any wdinfo whose frame-range bucket matches an entry in
    ``skip_frame_range`` is excluded.

    Args:
        wdinfos: List of wdinfo paths
        min_frames: Minimum number of frames (exclusive). If None, no lower bound.
        max_frames: Maximum number of frames (inclusive). If None, no upper bound.
        skip_frame_range: Frame-range bucket(s) to exclude, e.g. ``"300_400"`` to
            drop the ``frames_300_400`` bucket or ``"3700_inf"`` to drop the
            ``frames_3700_inf`` bucket. Accepts a single string or a list of strings.
            If None, no bucket is skipped.

    Returns:
        Filtered list of wdinfo paths

    Example:
        >>> wdinfos = [
        ...     "wdinfo/frames_400_500/wdinfo.json",
        ...     "wdinfo/frames_500_600/wdinfo.json",
        ...     "wdinfo/frames_600_700/wdinfo.json",
        ... ]
        >>> filter_wdinfos_by_frame_range(wdinfos, min_frames=500, max_frames=600)
        ['wdinfo/frames_500_600/wdinfo.json']
        # frames_400_500 excluded because wdinfo_max (500) <= min_frames (500)
        # frames_500_600 included because wdinfo_max (600) > min_frames (500) AND <= max_frames (600)
        # frames_600_700 excluded because wdinfo_max (700) > max_frames (600)

        >>> filter_wdinfos_by_frame_range(wdinfos, skip_frame_range="500_600")
        ['wdinfo/frames_400_500/wdinfo.json', 'wdinfo/frames_600_700/wdinfo.json']
        # frames_500_600 excluded because its bucket matches skip_frame_range
    """
    skip_buckets = normalize_skip_frame_ranges(skip_frame_range)

    if min_frames is None and max_frames is None and not skip_buckets:
        return wdinfos

    filtered = []
    for wdinfo in wdinfos:
        frame_range = parse_frame_range_from_wdinfo(wdinfo)
        if frame_range is None:
            # If no frame range in path, include by default
            filtered.append(wdinfo)
            continue

        wdinfo_min, wdinfo_max = frame_range

        # Skip explicitly excluded buckets (matched on the full (min, max) bucket).
        if (wdinfo_min, wdinfo_max) in skip_buckets:
            continue

        # Filter based on wdinfo's upper bound (wdinfo_max):
        # - min_frames is exclusive: wdinfo_max must be > min_frames
        # - max_frames is inclusive: wdinfo_max must be <= max_frames
        include = True
        if min_frames is not None and wdinfo_max <= min_frames:
            include = False
        if max_frames is not None and wdinfo_max > max_frames:
            include = False

        if include:
            filtered.append(wdinfo)

    return filtered