import cv2
import numpy as np
import torch

def extract_video_features(video_path, model, processor, device):
    """提取视频关键帧特征（智能抽帧版）

    Args:
        video_path: 视频文件路径
        model: CLIP模型
        processor: CLIP处理器
        device: 计算设备

    Returns:
        features: 形状为 (n_frames, feature_dim) 的特征数组
        timestamps: 对应帧的时间戳(秒)
    """
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise ValueError(f"无法打开视频文件: {video_path}")

    # 获取视频属性
    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    
    # 处理无效帧率
    if fps <= 0:
        fps = 30
        print(f"警告: 视频帧率无效，使用默认值 {fps}")

    # 智能抽帧参数设置
    base_interval = max(1, int(fps * 2))  # 基础间隔（2秒）
    min_interval = 1                      # 最小间隔（1帧）
    max_interval = max(1, int(fps * 5))   # 最大间隔（5秒）
    motion_threshold = 15                  # 运动阈值（像素平均变化）
    window_size = 3                       # 运动量平均窗口大小

    # 初始化变量
    prev_gray = None
    motion_history = []
    features = []
    timestamps = []
    current_index = 0

    while current_index < total_frames:
        # 设置当前帧位置
        cap.set(cv2.CAP_PROP_POS_FRAMES, current_index)
        ret, frame = cap.read()
        
        if not ret:
            break

        try:
            # 转换为灰度图计算运动量
            current_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            
            # 计算运动量（与前一帧的差异）
            if prev_gray is not None:
                frame_diff = cv2.absdiff(current_gray, prev_gray)  #  三个都是 (v_row, v_col)维
                motion_level = np.mean(frame_diff)
                motion_history.append(motion_level)
                
                # 保持窗口大小
                if len(motion_history) > window_size:
                    motion_history.pop(0)
                
                # 计算平均运动量
                avg_motion = np.mean(motion_history)
                
                # 动态调整抽帧间隔
                if avg_motion > motion_threshold:
                    next_interval = max(min_interval, base_interval - int(avg_motion/3))
                else:
                    next_interval = min(max_interval, base_interval + int(base_interval * 0.5))
            else:
                next_interval = base_interval  # 第一帧使用基础间隔
            
            # 转换为RGB并提取特征
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            inputs = processor(images=frame_rgb, return_tensors="pt").to(device)
            with torch.no_grad():
                frame_features = model.get_image_features(**inputs)
            
            # 保存结果
            features.append(frame_features.cpu().numpy().squeeze())
            timestamps.append(current_index / fps)
            
            # 更新状态
            prev_gray = current_gray
            current_index += next_interval
            
        except Exception as e:
            print(f"处理帧 {current_index} 时出错: {str(e)}")
            current_index += base_interval  # 出错时使用基础间隔
            continue

    cap.release()

    # 转换为numpy数组
    if len(features) > 0:
        features = np.vstack(features)
        timestamps = np.array(timestamps)
    else:
        features = np.empty((0, 512))  # 空数组
        timestamps = np.array([])

    return features, timestamps, total_frames, fps