from __future__ import annotations

import base64
import copy
import logging
import math
import os
import sys
import time
import warnings
from functools import lru_cache
from io import BytesIO
from typing import Optional

import requests
import torch
import torchvision
from packaging import version
from PIL import Image
from torchvision import io, transforms
from torchvision.transforms import InterpolationMode

import torch.nn.functional as F
# import time
import numpy as np
from PIL import Image
import os
from decord import VideoReader
from decord import cpu

logger = logging.getLogger(__name__)

IMAGE_FACTOR = 28
MIN_PIXELS = 4 * 28 * 28
MAX_PIXELS = 16384 * 28 * 28
MAX_RATIO = 200

VIDEO_MIN_PIXELS = 128 * 28 * 28
VIDEO_MAX_PIXELS = 768 * 28 * 28
FRAME_FACTOR = 2
FPS = 2.0
FPS_MIN_FRAMES = 4
FPS_MAX_FRAMES = 768

# Set the maximum number of video token inputs.
# Here, 128K represents the maximum number of input tokens for the VLLM model.
# Remember to adjust it according to your own configuration.
VIDEO_TOTAL_PIXELS = int(float(os.environ.get('VIDEO_MAX_PIXELS', 128000 * 28 * 28 * 0.9)))
logger.info(f"set VIDEO_TOTAL_PIXELS: {VIDEO_TOTAL_PIXELS}")


def round_by_factor(number: int, factor: int) -> int:
    """Returns the closest integer to 'number' that is divisible by 'factor'."""
    return round(number / factor) * factor


def ceil_by_factor(number: int, factor: int) -> int:
    """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'."""
    return math.ceil(number / factor) * factor


def floor_by_factor(number: int, factor: int) -> int:
    """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'."""
    return math.floor(number / factor) * factor


def smart_resize(
    height: int, width: int, factor: int = IMAGE_FACTOR, min_pixels: int = MIN_PIXELS, max_pixels: int = MAX_PIXELS
) -> tuple[int, int]:
    """
    Rescales the image so that the following conditions are met:

    1. Both dimensions (height and width) are divisible by 'factor'.

    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].

    3. The aspect ratio of the image is maintained as closely as possible.
    """
    if max(height, width) / min(height, width) > MAX_RATIO:
        raise ValueError(
            f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}"
        )
    h_bar = max(factor, round_by_factor(height, factor))
    w_bar = max(factor, round_by_factor(width, factor))
    if h_bar * w_bar > max_pixels:
        beta = math.sqrt((height * width) / max_pixels)
        h_bar = floor_by_factor(height / beta, factor)
        w_bar = floor_by_factor(width / beta, factor)
    elif h_bar * w_bar < min_pixels:
        beta = math.sqrt(min_pixels / (height * width))
        h_bar = ceil_by_factor(height * beta, factor)
        w_bar = ceil_by_factor(width * beta, factor)
    return h_bar, w_bar


def to_rgb(pil_image: Image.Image) -> Image.Image:
    if pil_image.mode == 'RGBA':
        white_background = Image.new("RGB", pil_image.size, (255, 255, 255))
        white_background.paste(pil_image, mask=pil_image.split()[3])  # Use alpha channel as mask
        return white_background
    else:
        return pil_image.convert("RGB")


def fetch_image(ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACTOR) -> Image.Image:
    if "image" in ele:
        image = ele["image"]
    else:
        image = ele["image_url"]
    image_obj = None
    if isinstance(image, Image.Image):
        image_obj = image
    elif image.startswith("http://") or image.startswith("https://"):
        # fix memory leak issue while using BytesIO
        with requests.get(image, stream=True) as response:
            response.raise_for_status()
            with BytesIO(response.content) as bio:
                image_obj = copy.deepcopy(Image.open(bio))
    elif image.startswith("file://"):
        image_obj = Image.open(image[7:])
    elif image.startswith("data:image"):
        if "base64," in image:
            _, base64_data = image.split("base64,", 1)
            data = base64.b64decode(base64_data)
            # fix memory leak issue while using BytesIO
            with BytesIO(data) as bio:
                image_obj = copy.deepcopy(Image.open(bio))
    else:
        image_obj = Image.open(image)
    if image_obj is None:
        raise ValueError(f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}")
    image = to_rgb(image_obj)
    ## resize
    if "resized_height" in ele and "resized_width" in ele:
        resized_height, resized_width = smart_resize(
            ele["resized_height"],
            ele["resized_width"],
            factor=size_factor,
        )
    else:
        width, height = image.size
        min_pixels = ele.get("min_pixels", MIN_PIXELS)
        max_pixels = ele.get("max_pixels", MAX_PIXELS)
        resized_height, resized_width = smart_resize(
            height,
            width,
            factor=size_factor,
            min_pixels=min_pixels,
            max_pixels=max_pixels,
        )
    image = image.resize((resized_width, resized_height))

    return image


def smart_nframes(
    ele: dict,
    total_frames: int,
    video_fps: int | float,
) -> int:
    """calculate the number of frames for video used for model inputs.

    Args:
        ele (dict): a dict contains the configuration of video.
            support either `fps` or `nframes`:
                - nframes: the number of frames to extract for model inputs.
                - fps: the fps to extract frames for model inputs.
                    - min_frames: the minimum number of frames of the video, only used when fps is provided.
                    - max_frames: the maximum number of frames of the video, only used when fps is provided.
        total_frames (int): the original total number of frames of the video.
        video_fps (int | float): the original fps of the video.

    Raises:
        ValueError: nframes should in interval [FRAME_FACTOR, total_frames].

    Returns:
        int: the number of frames for video used for model inputs.
    """
    assert not ("fps" in ele and "nframes" in ele), "Only accept either `fps` or `nframes`"
    if "nframes" in ele:
        nframes = round_by_factor(ele["nframes"], FRAME_FACTOR)
    else:
        fps = ele.get("fps", FPS)
        min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES), FRAME_FACTOR)
        max_frames = floor_by_factor(ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)), FRAME_FACTOR)
        nframes = total_frames / video_fps * fps
        if nframes > total_frames:
            logger.warning(f"smart_nframes: nframes[{nframes}] > total_frames[{total_frames}]")
        nframes = min(min(max(nframes, min_frames), max_frames), total_frames)
        nframes = floor_by_factor(nframes, FRAME_FACTOR)
    if not (FRAME_FACTOR <= nframes and nframes <= total_frames):
        raise ValueError(f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}.")
    return nframes


def _read_video_torchvision(
    ele: dict,
) -> (torch.Tensor, float):
    """read video using torchvision.io.read_video

    Args:
        ele (dict): a dict contains the configuration of video.
        support keys:
            - video: the path of video. support "file://", "http://", "https://" and local path.
            - video_start: the start time of video.
            - video_end: the end time of video.
    Returns:
        torch.Tensor: the video tensor with shape (T, C, H, W).
    """
    video_path = ele["video"]
    if version.parse(torchvision.__version__) < version.parse("0.19.0"):
        if "http://" in video_path or "https://" in video_path:
            warnings.warn("torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0.")
        if "file://" in video_path:
            video_path = video_path[7:]
    st = time.time()
    video, audio, info = io.read_video(
        video_path,
        start_pts=ele.get("video_start", 0.0),
        end_pts=ele.get("video_end", None),
        pts_unit="sec",
        output_format="TCHW",
    )
    total_frames, video_fps = video.size(0), info["video_fps"]
    logger.info(f"torchvision:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s")
    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
    idx = torch.linspace(0, total_frames - 1, nframes).round().long()
    sample_fps = nframes / max(total_frames, 1e-6) * video_fps
    video = video[idx]
    return video, sample_fps


def is_decord_available() -> bool:
    import importlib.util

    return importlib.util.find_spec("decord") is not None


def calculate_video_frame_range(
    ele: dict,
    total_frames: int,
    video_fps: float,
) -> tuple[int, int, int]:
    """
    Calculate the start and end frame indices based on the given time range.

    Args:
        ele (dict): A dictionary containing optional 'video_start' and 'video_end' keys (in seconds).
        total_frames (int): Total number of frames in the video.
        video_fps (float): Frames per second of the video.

    Returns:
        tuple: A tuple containing (start_frame, end_frame, frame_count).

    Raises:
        ValueError: If input parameters are invalid or the time range is inconsistent.
    """
    # Validate essential parameters
    if video_fps <= 0:
        raise ValueError("video_fps must be a positive number")
    if total_frames <= 0:
        raise ValueError("total_frames must be a positive integer")

    # Get start and end time in seconds
    video_start = ele.get("video_start", None)
    video_end = ele.get("video_end", None)
    if video_start is None and video_end is None:
        return 0, total_frames - 1, total_frames

    max_duration = total_frames / video_fps
    # Process start frame
    if video_start is not None:
        video_start_clamped = max(0.0, min(video_start, max_duration))
        start_frame = math.ceil(video_start_clamped * video_fps)
    else:
        start_frame = 0
    # Process end frame
    if video_end is not None:
        video_end_clamped = max(0.0, min(video_end, max_duration))
        end_frame = math.floor(video_end_clamped * video_fps)
        end_frame = min(end_frame, total_frames - 1)
    else:
        end_frame = total_frames - 1

    # Validate frame order
    if start_frame >= end_frame:
        raise ValueError(
            f"Invalid time range: Start frame {start_frame} (at {video_start_clamped if video_start is not None else 0}s) "
            f"exceeds end frame {end_frame} (at {video_end_clamped if video_end is not None else max_duration}s). "
            f"Video duration: {max_duration:.2f}s ({total_frames} frames @ {video_fps}fps)"
        )

    logger.info(f"calculate video frame range: {start_frame=}, {end_frame=}, {total_frames=} from {video_start=}, {video_end=}, {video_fps=:.3f}")
    return start_frame, end_frame, end_frame - start_frame + 1

# New
def _read_video_decord_dynamic(
    ele: dict,
) -> (torch.Tensor, float):
    """read video using decord.VideoReader

    Args:
        ele (dict): a dict contains the configuration of video.
        support keys:
            - video: the path of video. support "file://", "http://", "https://" and local path.
            - video_start: the start time of video.
            - video_end: the end time of video.
    Returns:
        torch.Tensor: the video tensor with shape (T, C, H, W).
    """
    import decord
    # import pdb; pdb.set_trace()
    video_path = ele["video"]
    st = time.time()
    # vr = VideoReader(video_path, ctx=cpu(0), num_threads=256)
    vr = VideoReader(video_path, ctx=cpu(0))
    total_frames, video_fps = len(vr), vr.get_avg_fps()
    start_frame, end_frame, total_frames = calculate_video_frame_range(
            ele,
            total_frames,
            video_fps,
        )
    # print(total_frames)
    if total_frames < 64:
        # video_fps = vr.get_avg_fps()
        
        nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
        idx = torch.linspace(start_frame, end_frame, nframes).round().long().tolist()
        video = vr.get_batch(idx).asnumpy()
        video = torch.tensor(video).permute(0, 3, 1, 2)  # Convert to TCHW format
        logger.info(f"decord:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s")
        sample_fps = nframes / max(total_frames, 1e-6) * video_fps
        real_idx = None
    else:
        step = int(round(video_fps / 5))
        nframes = total_frames // round(video_fps)
        idx = torch.linspace(0, len(vr)-1, nframes).round().long().tolist()
        frames = vr.get_batch(idx).asnumpy()
        small_frames = resize_frames_torch(frames, max_area=244*244)
        # import pdb; pdb.set_trace()
        video, sel_idx = select_keyframes(small_frames, frames)
        # import pdb; pdb.set_trace()
        real_idx = [idx[i] for i in sel_idx]
        video = torch.from_numpy(video).permute(0, 3, 1, 2).float()
        logger.info(f"decord:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s")
        sample_fps = 1
        # 如果大于768帧，那么就取前750个
        if video.shape[0] > 768:
           video = video[:768,:,:,:]
        # print("video.shape", video.shape)
        # print("total_frams", total_frames)
    log_file = "video_frame_log_mean_all.txt"
    with open(log_file, "a") as f:
        f.write(f"{video_path}, total_frames={total_frames}, used_frames={video.shape[0]}\n")
    return video, sample_fps, real_idx

def _read_video_decord(
    ele: dict,
) -> (torch.Tensor, float):
    """read video using decord.VideoReader

    Args:
        ele (dict): a dict contains the configuration of video.
        support keys:
            - video: the path of video. support "file://", "http://", "https://" and local path.
            - video_start: the start time of video.
            - video_end: the end time of video.
    Returns:
        torch.Tensor: the video tensor with shape (T, C, H, W).
    """
    import decord
    video_path = ele["video"]
    st = time.time()
    vr = decord.VideoReader(video_path)
    total_frames, video_fps = len(vr), vr.get_avg_fps()
    start_frame, end_frame, total_frames = calculate_video_frame_range(
        ele,
        total_frames,
        video_fps,
    )
    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
    idx = torch.linspace(start_frame, end_frame, nframes).round().long().tolist()
    video = vr.get_batch(idx).asnumpy()
    video = torch.tensor(video).permute(0, 3, 1, 2)  # Convert to TCHW format
    logger.info(f"decord:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s")
    sample_fps = nframes / max(total_frames, 1e-6) * video_fps
    return video, sample_fps


def is_torchcodec_available() -> bool:
    """Check if torchcodec is available and properly installed."""
    try:
        import importlib.util
        if importlib.util.find_spec("torchcodec") is None:
            return False
        from torchcodec.decoders import VideoDecoder
        return True
    except (ImportError, AttributeError, Exception):
        return False


def _read_video_torchcodec(
    ele: dict,
) -> (torch.Tensor, float):
    """read video using torchcodec.decoders.VideoDecoder

    Args:
        ele (dict): a dict contains the configuration of video.
        support keys:
            - video: the path of video. support "file://", "http://", "https://" and local path.
            - video_start: the start time of video.
            - video_end: the end time of video.
    Returns:
        torch.Tensor: the video tensor with shape (T, C, H, W).
    """
    from torchcodec.decoders import VideoDecoder
    TORCHCODEC_NUM_THREADS = int(os.environ.get('TORCHCODEC_NUM_THREADS', 8))
    logger.info(f"set TORCHCODEC_NUM_THREADS: {TORCHCODEC_NUM_THREADS}")
    video_path = ele["video"]
    st = time.time()
    decoder = VideoDecoder(video_path, num_ffmpeg_threads=TORCHCODEC_NUM_THREADS)
    video_fps = decoder.metadata.average_fps
    total_frames = decoder.metadata.num_frames
    start_frame, end_frame, total_frames = calculate_video_frame_range(
        ele,
        total_frames,
        video_fps,
    )
    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
    idx = torch.linspace(start_frame, end_frame, nframes).round().long().tolist()
    sample_fps = nframes / max(total_frames, 1e-6) * video_fps
    video = decoder.get_frames_at(indices=idx).data
    logger.info(f"torchcodec:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s")
    return video, sample_fps


VIDEO_READER_BACKENDS = {
    # "decord": _read_video_decord,
    "decord": _read_video_decord_dynamic,
    "torchvision": _read_video_torchvision,
    "torchcodec": _read_video_torchcodec,
}

FORCE_QWENVL_VIDEO_READER = os.getenv("FORCE_QWENVL_VIDEO_READER", None)


@lru_cache(maxsize=1)
def get_video_reader_backend() -> str:
    if FORCE_QWENVL_VIDEO_READER is not None:
        video_reader_backend = FORCE_QWENVL_VIDEO_READER
    elif is_torchcodec_available():
        video_reader_backend = "torchcodec"
    elif is_decord_available():
        video_reader_backend = "decord"
    else:
        video_reader_backend = "torchvision"
    print(f"qwen-vl-utils using {video_reader_backend} to read video.", file=sys.stderr)
    return video_reader_backend



def fetch_video(ele: dict, image_factor: int = IMAGE_FACTOR, return_video_sample_fps: bool = False) -> torch.Tensor | list[Image.Image]:
    if isinstance(ele["video"], str):
        # 使用decord
        video_reader_backend = get_video_reader_backend()
        try:
            video, sample_fps, real_idx = VIDEO_READER_BACKENDS[video_reader_backend](ele)
        except Exception as e:
            logger.warning(f"video_reader_backend {video_reader_backend} error, use torchvision as default, msg: {e}")
            video, sample_fps = VIDEO_READER_BACKENDS["torchvision"](ele)
        
        nframes, _, height, width = video.shape
        min_pixels = ele.get("min_pixels", VIDEO_MIN_PIXELS)
        total_pixels = ele.get("total_pixels", VIDEO_TOTAL_PIXELS)
        max_pixels = max(min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR), int(min_pixels * 1.05))
        max_pixels_supposed = ele.get("max_pixels", max_pixels)
        if max_pixels_supposed > max_pixels:
            logger.warning(f"The given max_pixels[{max_pixels_supposed}] exceeds limit[{max_pixels}].")
        max_pixels = min(max_pixels_supposed, max_pixels)
        if "resized_height" in ele and "resized_width" in ele:
            resized_height, resized_width = smart_resize(
                ele["resized_height"],
                ele["resized_width"],
                factor=image_factor,
            )
        else:
            resized_height, resized_width = smart_resize(
                height,
                width,
                factor=image_factor,
                min_pixels=min_pixels,
                max_pixels=max_pixels,
            )

        video = transforms.functional.resize(
            video,
            [resized_height, resized_width],
            interpolation=InterpolationMode.BICUBIC,
            antialias=True,
        ).float()
        if return_video_sample_fps:
            return video, sample_fps, real_idx
        return video, real_idx
    else:
        assert isinstance(ele["video"], (list, tuple))
        process_info = ele.copy()
        process_info.pop("type", None)
        process_info.pop("video", None)
        images = [
            fetch_image({"image": video_element, **process_info}, size_factor=image_factor)
            for video_element in ele["video"]
        ]
        nframes = ceil_by_factor(len(images), FRAME_FACTOR)
        if len(images) < nframes:
            images.extend([images[-1]] * (nframes - len(images)))
        if return_video_sample_fps:
            return images, process_info.pop("fps", 2.0)
        return images


def extract_vision_info(conversations: list[dict] | list[list[dict]]) -> list[dict]:
    vision_infos = []
    if isinstance(conversations[0], dict):
        conversations = [conversations]
    for conversation in conversations:
        for message in conversation:
            if isinstance(message["content"], list):
                for ele in message["content"]:
                    if (
                        "image" in ele
                        or "image_url" in ele
                        or "video" in ele
                        or ele.get("type","") in ("image", "image_url", "video")
                    ):
                        vision_infos.append(ele)
    return vision_infos


def process_vision_info(
    conversations: list[dict] | list[list[dict]],
    return_video_kwargs: bool = False,
) -> tuple[list[Image.Image] | None, list[torch.Tensor | list[Image.Image]] | None, Optional[dict]]:
    # import pdb; pdb.set_trace()
    vision_infos = extract_vision_info(conversations)
    ## Read images or videos
    image_inputs = []
    video_inputs = []
    video_sample_fps_list = []
    for vision_info in vision_infos:
        if "image" in vision_info or "image_url" in vision_info:
            image_inputs.append(fetch_image(vision_info))
        elif "video" in vision_info:
            video_input, video_sample_fps, real_idx = fetch_video(vision_info, return_video_sample_fps=True)
            video_sample_fps_list.append(video_sample_fps)
            video_inputs.append(video_input)
        else:
            raise ValueError("image, image_url or video should in content.")
    if len(image_inputs) == 0:
        image_inputs = None
    if len(video_inputs) == 0:
        video_inputs = None
    if return_video_kwargs:
        return image_inputs, video_inputs, real_idx, {'fps': video_sample_fps_list}
    return image_inputs, video_inputs




# def compute_entropy(frame):
#     """Compute entropy across RGB channels and average."""
#     entropy_total = 0
#     for c in range(3):
#         channel = frame[:, :, c]
#         hist, _ = np.histogram(channel, bins=256, range=(0, 256), density=True)
#         hist = hist[hist > 0]
#         entropy = -np.sum(hist * np.log2(hist))
#         entropy_total += entropy
#     return entropy_total / 3

# def compute_mutual_information(f1, f2):
#     """Estimate mutual information between two grayscale images."""
#     f1_gray = np.dot(f1[..., :3], [0.2989, 0.5870, 0.1140]).astype(np.uint8)
#     f2_gray = np.dot(f2[..., :3], [0.2989, 0.5870, 0.1140]).astype(np.uint8)
#     joint_hist = np.histogram2d(f1_gray.ravel(), f2_gray.ravel(), bins=256)[0]
#     joint_prob = joint_hist / joint_hist.sum()
#     p1 = joint_prob.sum(axis=1)
#     p2 = joint_prob.sum(axis=0)
#     H1 = -np.sum(p1 * np.log2(p1 + 1e-12))
#     H2 = -np.sum(p2 * np.log2(p2 + 1e-12))
#     H_joint = -np.sum(joint_prob * np.log2(joint_prob + 1e-12))
#     return H1 + H2 - H_joint

# def dynamic_frame_selection(frames):
#     """Select frames based on information density score."""
#     D_scores = []
#     for i in range(1, len(frames)):
#         H = compute_entropy(frames[i])
#         MI = compute_mutual_information(frames[i-1], frames[i])
#         D = H - MI
#         D_scores.append(D)

#     threshold = np.mean(D_scores) + np.std(D_scores)
#     selected = [frames[0]]  # Always keep first frame
#     for i, D in enumerate(D_scores):
#         if D > threshold:
#             selected.append(frames[i+1])
#     selected_tensor = np.stack(selected, axis=0)
#     return selected


# New

def compute_entropy(image: Image.Image) -> float:
    image_np = np.array(image)
    entropy_sum = 0.0
    for c in range(3):
        channel = image_np[..., c]
        histogram, _ = np.histogram(channel, bins=256, range=(0, 255), density=True)
        histogram += 1e-9
        entropy = -np.sum(histogram * np.log2(histogram))
        entropy_sum += entropy
    return entropy_sum / 3.0

def compute_mutual_information(f1, f2):
    """Estimate mutual information between two grayscale images."""
    f1_gray = np.dot(f1[..., :3], [0.2989, 0.5870, 0.1140]).astype(np.uint8)
    f2_gray = np.dot(f2[..., :3], [0.2989, 0.5870, 0.1140]).astype(np.uint8)
    joint_hist = np.histogram2d(f1_gray.ravel(), f2_gray.ravel(), bins=256)[0]
    joint_prob = joint_hist / joint_hist.sum()
    p1 = joint_prob.sum(axis=1)
    p2 = joint_prob.sum(axis=0)
    H1 = -np.sum(p1 * np.log2(p1 + 1e-12))
    H2 = -np.sum(p2 * np.log2(p2 + 1e-12))
    H_joint = -np.sum(joint_prob * np.log2(joint_prob + 1e-12))
    return H1 + H2 - H_joint

def normalize_array(arr):
    arr_min = arr.min()
    arr_max = arr.max()
    return (arr - arr_min) / (arr_max - arr_min + 1e-9)


def select_keyframes(frames, raw_frames, base_threshold=0.03, 
                      alpha=1.01, max_iter=10, min_frames=32):
    # device = clip_model.device
    frame_features = []
    sel_frames = frames
    selected_indices = []
    entropies = np.array([compute_entropy(f) for f in sel_frames])
    entropies = entropies[:-1]
    MI = np.array([compute_mutual_information(sel_frames[i], sel_frames[i + 1]) for i in range(len(sel_frames) - 1)])
    entropies = normalize_array(entropies)
    MI = normalize_array(MI)
    D_scores = entropies - MI
    threshold = np.mean(D_scores)
    #  + 3 * np.std(D_scores)
    #  + 0.25*np.std(D_scores)
    # selected = [raw_frames[0]]  # Always keep first frame
    # for i, D in enumerate(D_scores):
    #     if D > threshold:
    #         selected.append(raw_frames[i+1])
    # selected = np.stack(selected, axis=0)
    # return selected
    selected_frames = [raw_frames[0]]  # Always keep first frame
    selected_indices.append(0)
    for i, D in enumerate(D_scores):
        if D > threshold:
            selected_frames.append(raw_frames[i + 1])   # +1是因为D是[i]对应[i+1]帧
            selected_indices.append(i + 1)

    selected_frames = np.stack(selected_frames, axis=0)
    
    return selected_frames, selected_indices
    # for m in range(max_iter):
    #     print(f'---the current iteration is {m}---')
    #     sel_features = frame_features[selected_indices]
    #     sel_frames = [frames[i] for i in selected_indices]

    #     with torch.no_grad():   
    #         similarities = torch.mm(sel_features, text_feature.T).squeeze(1).numpy()

    #     entropies = np.array([compute_entropy(f) for f in sel_frames])
    #     similarities = normalize_array(similarities)
    #     entropies = normalize_array(entropies)

    #     ecrs_scores = similarities * entropies
    #     ecrs_scores = normalize_array(ecrs_scores)

    #     threshold = base_threshold * (alpha ** m)
    #     candidates = [selected_indices[i] for i, score in enumerate(ecrs_scores) if score > threshold]

    #     if not candidates or set(candidates) == set(selected_indices):
    #         break
    #     selected_indices = candidates

    # if len(selected_indices) < min_frames:
    #     all_entropies = np.array([compute_entropy(f) for f in frames])
    #     with torch.no_grad():
    #         all_similarities = torch.mm(frame_features, text_feature.T).squeeze(1).numpy()
    #     all_ecrs = normalize_array(all_similarities) * normalize_array(all_entropies)
    #     sorted_indices = np.argsort(all_ecrs)[::-1]
    #     selected_indices = sorted(set(selected_indices + sorted_indices[:min_frames].tolist()))
    
    # selected_frames = [frames[i] for i in selected_indices]
    
    # print(f'------the selected_frames is {selected_indices}------')
    # return selected_frames, selected_indices

# def dynamic_sample(frames_np):  # frames: List[np.ndarray] RGB
#     # Step 1: Convert to torch tensor and normalize
#     frames_tensor = torch.from_numpy(np.stack(frames_np)).permute(0, 3, 1, 2).float().cuda() / 255.0  # [N, 3, H, W]

#     # Step 2: Gray + Entropy + MI
#     gray = rgb_to_grayscale_batch(frames_tensor)  # [N, H, W]
#     entropy = compute_batch_entropy(gray)         # [N]
#     mutual_info = compute_batch_mutual_information(gray)  # [N-1]
#     D = entropy[1:] - mutual_info                 # [N-1]

#     # Step 3: Select based on D > mean + std
#     threshold = D.mean() + D.std()
#     selected_indices = [0] + [i + 1 for i, d in enumerate(D) if d > threshold]

#     # Step 4: Slice frames
#     selected = frames_tensor[selected_indices]  # [T', 3, H, W]
#     return selected
# def resize_frames_torch(frames_np, max_area=244*244):
#     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

#     # 转换为 PyTorch tensor: (N, C, H, W), float32 in [0, 1]
#     frames = torch.from_numpy(frames_np).permute(0, 3, 1, 2).float().to(device) / 255.0

#     resized_frames = []
#     for frame in frames:
#         c, h, w = frame.shape
#         area = h * w
#         if area > max_area:
#             scale = (max_area / area) ** 0.5
#             new_h = int(h * scale)
#             new_w = int(w * scale)
#             resized = F.interpolate(frame.unsqueeze(0), size=(new_h, new_w), mode='bilinear', align_corners=False)
#             resized = resized.squeeze(0)
#         else:
#             resized = frame

#         # 转回 numpy array, (H, W, C), uint8
#         resized_np = (resized.permute(1, 2, 0).clamp(0, 1).cpu().numpy() * 255).astype(np.uint8)
#         resized_frames.append(resized_np)

#     return resized_frames

def resize_frames_torch(frames_np, max_area=244*244, batch_size=1024):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    resized_frames = []

    num_frames = len(frames_np)
    for i in range(0, num_frames, batch_size):
        batch_np = frames_np[i:i+batch_size]  # (B, H, W, C)
        batch = torch.from_numpy(batch_np).permute(0, 3, 1, 2).float() / 255.0  # (B, C, H, W)
        batch = batch.to(device)

        b, c, h, w = batch.shape
        area = h * w
        if area > max_area:
            scale = (max_area / area) ** 0.5
            new_h = int(h * scale)
            new_w = int(w * scale)
            batch_resized = F.interpolate(batch, size=(new_h, new_w), mode='bilinear', align_corners=False)
        else:
            batch_resized = batch

        # 转回 numpy array, 每个 (H, W, C), uint8
        batch_resized_np = (batch_resized.permute(0, 2, 3, 1).clamp(0, 1).cpu().numpy() * 255).astype(np.uint8)
        resized_frames.extend(batch_resized_np)

    return resized_frames

