'''
python video_filter.py input_video.mp4 mask_area_ratios.json 0.1 filtered_video_01.mp4

python video_filter.py input_video.mp4 mask_area_ratios.json 0.0 filtered_video_00.mp4
'''

import argparse
import cv2
import json
import os
import numpy as np
from moviepy.editor import ImageSequenceClip

def process_video(video_path, json_path, threshold, output_path):
    # 读取视频
    if os.path.exists("input_video.mp4"):
        video_path = "input_video.mp4"

    cap = cv2.VideoCapture(video_path)
    input_fps = cap.get(cv2.CAP_PROP_FPS)  # 获取输入视频的帧率
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))  # 获取总帧数

    # 读取所有帧
    all_frames = []
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  # 转换为RGB
        all_frames.append(frame_rgb)
    cap.release()

    # 读取JSON文件
    with open(json_path, 'r') as f:
        mask_area_ratios = json.load(f)

    # 获取mask_area_ratios中的帧索引和对应的ratio
    frame_indices = [int(key.split("_")[1]) for key in mask_area_ratios.keys()]
    ratios = [mask_area_ratios[key] for key in mask_area_ratios.keys()]

    # 创建一个长度为 total_frames 的数组来存储对应的 ratio
    all_ratios = [None] * len(all_frames)

    # 将 mask_area_ratios 的帧索引和对应的 ratio 映射到 all_ratios
    for frame_index, ratio in zip(frame_indices, ratios):
        scaled_index = int(frame_index * (len(all_frames) - 1) / max(frame_indices))
        all_ratios[scaled_index] = ratio

    # 填充缺失的 ratio 值，使用最邻近值
    for i in range(len(all_ratios)):
        if all_ratios[i] is None:
            # 找到最近的非空值
            left_index = i - 1
            right_index = i + 1
            while left_index >= 0 and all_ratios[left_index] is None:
                left_index -= 1
            while right_index < len(all_ratios) and all_ratios[right_index] is None:
                right_index += 1

            if left_index >= 0 and right_index < len(all_ratios):
                # 使用最近的左边或右边的值
                all_ratios[i] = all_ratios[left_index] if all_ratios[left_index] is not None else all_ratios[right_index]
            elif left_index >= 0:
                all_ratios[i] = all_ratios[left_index]
            elif right_index < len(all_ratios):
                all_ratios[i] = all_ratios[right_index]

    # 筛选符合条件的帧
    frames = []
    for frame_idx, ratio in enumerate(all_ratios):
        if ratio >= threshold:
            frames.append(all_frames[frame_idx])

    # 保存符合条件的帧为视频
    if frames:
        clip = ImageSequenceClip(frames, fps=input_fps)  # 使用输入视频的帧率
        clip.write_videofile(output_path, codec="libx264")
        print(f"Output video saved to {output_path}")
    else:
        print("No frames meet the threshold condition.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Filter video frames based on mask area ratios.")
    parser.add_argument("video_path", type=str, nargs='?', default="input_video.mp4", help="Path to the input video file. Default: input_video.mp4")
    parser.add_argument("json_path", type=str, nargs='?', default="mask_area_ratios.json", help="Path to the JSON file containing mask area ratios. Default: mask_area_ratios.json")
    parser.add_argument("threshold", type=float, nargs='?', default=0.25, help="Threshold for mask area ratio. Default: 0.25")
    parser.add_argument("output_path", type=str, nargs='?', default="output_seg_video.mp4", help="Path to save the output video. Default: output_seg_video.mp4")

    args = parser.parse_args()

    process_video(args.video_path, args.json_path, args.threshold, args.output_path)