import sys
import subprocess

# 检查cv2是否已安装
try:
    import cv2
except ImportError:
    print("正在安装OpenCV...")
    subprocess.check_call([sys.executable, "-m", "pip", "install", "opencv-python"])
    import cv2  # 再次尝试导入

import gradio as gr
import pandas as pd
import os
from datetime import datetime
import json
import random
import glob
import re
import numpy as np

# 修改1：使用绝对路径，避免相对路径引起的访问问题
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
BASE_DIR = os.path.join(CURRENT_DIR, "selected")  # 视频组的基础目录
RESULTS_DIR = os.path.join(CURRENT_DIR, "user_study_results")
SCALED_VIDEOS_DIR = os.path.join(CURRENT_DIR, "scaled_videos")  # 放缩后的视频存储目录
MAX_EDGE_SIZE = 512  # 最大边尺寸
NUM_GROUPS_TO_SELECT = 10  # 随机选择的视频组数量

os.makedirs(RESULTS_DIR, exist_ok=True)
os.makedirs(SCALED_VIDEOS_DIR, exist_ok=True)

# 修改2：添加文件权限检查和修复函数
def ensure_file_permissions(file_path):
    """确保文件具有正确的读取权限"""
    if os.path.exists(file_path):
        # 添加读取权限 (Linux/Mac)
        if os.name != 'nt':  # 非Windows系统
            os.chmod(file_path, 0o644)  # 设置为644权限
        return True
    return False

# 缩放并准备所有视频，这个只需一次性完成
def prepare_all_videos():
    if not os.path.exists(BASE_DIR):
        raise ValueError(f"基础目录不存在: {BASE_DIR}")
    
    all_groups = [d for d in os.listdir(BASE_DIR) if os.path.isdir(os.path.join(BASE_DIR, d)) and not d.startswith('.')]
    
    for group in all_groups:
        group_path = os.path.join(BASE_DIR, group)
        video_files = glob.glob(f"{group_path}/*.mp4")
        
        for video_path in video_files:
            video_name = os.path.basename(video_path)
            scaled_video_dir = os.path.join(SCALED_VIDEOS_DIR, group)
            os.makedirs(scaled_video_dir, exist_ok=True)
            scaled_video_path = os.path.join(scaled_video_dir, video_name)
            scale_video(video_path, scaled_video_path)
            # 确保文件权限正确
            ensure_file_permissions(scaled_video_path)
    
    print("所有视频已准备完成")

# 生成8位随机数作为用户ID
def generate_user_id():
    return str(random.randint(10000000, 99999999))

# 获取所有视频组并随机选择10个
def get_random_video_groups():
    if not os.path.exists(BASE_DIR):
        raise ValueError(f"基础目录不存在: {BASE_DIR}")
    all_groups = [d for d in os.listdir(BASE_DIR) if os.path.isdir(os.path.join(BASE_DIR, d)) and not d.startswith('.')]

    random.shuffle(all_groups)
    return all_groups
    # # 确保我们有足够的组可以选择
    # if len(all_groups) <= NUM_GROUPS_TO_SELECT:
    #     return sorted(all_groups)
    
    # # 随机选择指定数量的组
    # selected_groups = random.sample(all_groups, NUM_GROUPS_TO_SELECT)
    # selected_groups = sorted(selected_groups)
    # return selected_groups

# 缩放视频到指定的最大边尺寸
def scale_video(input_path, output_path):
    # 如果已经存在缩放后的视频，直接返回路径并确保权限正确
    if os.path.exists(output_path):
        ensure_file_permissions(output_path)
        return output_path
    
    try:
        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        
        # 打开输入视频
        cap = cv2.VideoCapture(input_path)
        if not cap.isOpened():
            print(f"无法打开视频文件: {input_path}")
            return input_path
        
        # 获取原始视频的属性
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = cap.get(cv2.CAP_PROP_FPS)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        
        # 计算新的尺寸
        if width >= height:
            new_width = MAX_EDGE_SIZE
            new_height = int(height * (MAX_EDGE_SIZE / width))
        else:
            new_height = MAX_EDGE_SIZE
            new_width = int(width * (MAX_EDGE_SIZE / height))
        
        # 获取视频编解码器信息
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 使用MP4编码
        
        # 创建输出视频写入器
        out = cv2.VideoWriter(output_path, fourcc, fps, (new_width, new_height))
        
        # 逐帧读取、缩放并写入
        frame_count = 0
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            # 缩放帧
            resized_frame = cv2.resize(frame, (new_width, new_height), interpolation=cv2.INTER_AREA)
            
            # 写入输出视频
            out.write(resized_frame)
            
            frame_count += 1
            if frame_count % 20 == 0:  # 每20帧打印一次进度
                print(f"处理视频 {os.path.basename(input_path)}: {frame_count}/{total_frames} 帧 ({frame_count/total_frames*100:.1f}%)")
        
        # 释放资源
        cap.release()
        out.release()
        
        # 确保输出文件有正确的权限
        ensure_file_permissions(output_path)
        
        return output_path
    except Exception as e:
        print(f"处理视频时出错: {input_path}, 错误: {e}")
        # 如果缩放失败，返回原始视频路径
        return input_path

# 为每个组生成随机排序及映射，并返回缩放后的视频路径
def generate_video_mappings(groups):
    mappings = {}
    for group in groups:
        group_path = os.path.join(BASE_DIR, group)
        video_files = glob.glob(f"{group_path}/*.mp4")
        
        if not video_files:
            print(f"警告: 组 {group} 中没有找到视频文件")
            continue
            
        gt_video = next((v for v in video_files if os.path.basename(v) == "gt.mp4"), None)
        method_videos = [v for v in video_files if os.path.basename(v) != "gt.mp4"]
        
        if not gt_video:
            print(f"警告: 组 {group} 中没有找到gt.mp4文件")
            # 如果没有gt.mp4，使用第一个可用的视频作为输入视频
            if method_videos:
                gt_video = method_videos[0]
                method_videos = method_videos[1:]
            else:
                continue  # 跳过这个组
                
        random.shuffle(method_videos)
        
        # 原始视频路径
        original_video_paths = [gt_video] + method_videos
        
        # 获取缩放后的视频路径
        scaled_video_paths = []
        for video_path in original_video_paths:
            if video_path:
                video_name = os.path.basename(video_path)
                scaled_video_dir = os.path.join(SCALED_VIDEOS_DIR, group)
                os.makedirs(scaled_video_dir, exist_ok=True)
                scaled_video_path = os.path.join(scaled_video_dir, video_name)
                # 处理视频并确保权限正确
                scaled_path = scale_video(video_path, scaled_video_path)
                scaled_video_paths.append(scaled_path)
                print(f"视频路径: {scaled_path}, 文件存在: {os.path.exists(scaled_path)}")
        
        display_order = ["GT"] + [os.path.splitext(os.path.basename(v))[0] for v in method_videos]
        mappings[group] = {"video_paths": scaled_video_paths, "display_order": display_order}
    return mappings

# 提交单个组的评分结果
def submit_group_rating(user_id, group, rating, method_mapping, group_num):
    if not rating:
        return "<span style='color:red'>请选择一个方法后再提交</span>"
    
    # 创建用户结果文件路径
    user_results_file = os.path.join(RESULTS_DIR, f"{user_id}_results.json")
    
    # 读取现有结果（如果有）
    if os.path.exists(user_results_file):
        with open(user_results_file, "r") as f:
            try:
                results = json.load(f)
            except:
                results = {
                    "user_id": user_id,
                    "timestamp_start": datetime.now().isoformat(),
                    "ratings": {},
                    "anonymous_ratings": {}
                }
    else:
        results = {
            "user_id": user_id,
            "timestamp_start": datetime.now().isoformat(),
            "ratings": {},
            "anonymous_ratings": {}
        }
    
    # 更新结果 - 同时保存真实方法名和匿名方法名
    if rating == "其他":
        real_method = "OTHER"
    else:
        real_method = method_mapping.get(rating)
    
    results["ratings"][group] = real_method
    results["anonymous_ratings"][group] = rating
    results["timestamp_last_update"] = datetime.now().isoformat()
    
    # 保存结果
    with open(user_results_file, "w") as f:
        json.dump(results, f, indent=2, ensure_ascii=False)
    
    return f"<span style='color:green'>组{group_num} 的评分已提交成功！您选择了 {rating}</span>"

# 查看所有结果的函数
def view_all_results(user_id, groups):
    user_results_file = os.path.join(RESULTS_DIR, f"{user_id}_results.json")
    if not os.path.exists(user_results_file):
        return "<span style='color:orange'>您还没有提交任何评分</span>"
    
    with open(user_results_file, "r") as f:
        results = json.load(f)
    
    anonymous_ratings = results.get("anonymous_ratings", {})
    if not anonymous_ratings:
        return "<span style='color:orange'>您还没有提交任何评分</span>"
    
    output = "<h4>您的评分结果:</h4><ul>"
    for group_idx, group in enumerate(groups):
        method = anonymous_ratings.get(group, "未评分")
        output += f"<li>组{group_idx+1}: {method}</li>"
    output += "</ul>"
    
    output += "<p>感谢您参与此次评价研究！</p>"
    
    return output

# 创建Gradio接口，移动随机选择和ID生成的逻辑到接口内部
def create_interface():
    # 每次加载应用时重新生成用户ID
    user_id = generate_user_id()
    
    # 每次加载应用时重新随机选择视频组
    groups = get_random_video_groups()
    video_mappings = generate_video_mappings(groups)
    
    with gr.Blocks(title="4D视频生成效果评价") as demo:
        gr.Markdown("# 4D视频生成效果评价")
        gr.Markdown(f"请观看每组中的视频。输入图像/视频，输出不同方法生成的大角度视频。请选择您认为生成**物理一致性**效果的方法。")

        # 使用Gradio状态保存用户ID，但每次刷新时会重新生成
        user_id_state = gr.State(value=user_id)
        # id_display = gr.Markdown(value=f"**您的用户ID: {user_id}** (请记下此ID，以便研究人员跟踪您的参与)")
        
        with gr.Tabs():
            for group_idx, group in enumerate(groups):
                if group not in video_mappings:
                    continue  # 跳过没有视频的组
                    
                mapping = video_mappings[group]
                video_paths = mapping["video_paths"]
                
                if not video_paths:
                    continue  # 跳过没有视频的组
                    
                real_method_names = mapping["display_order"]
                anon_method_names = [f"方法{i}" for i in range(1, len(real_method_names))]
                # 添加"其他"选项
                anon_method_names.append("其他")
                method_mapping = dict(zip(anon_method_names[:-1], real_method_names[1:]))  # 不包含"其他"选项

                with gr.TabItem(f"组 {group_idx+1}"):
                    gr.Markdown(f"### 组 {group_idx+1} 视频比较")

                    with gr.Row():
                        with gr.Column(scale=1):
                            gr.Video(value=video_paths[0], label="输入视频", autoplay=True, height=300)
                    
                    # 2x2 网格布局显示方法视频
                    method_videos_count = len(anon_method_names) - 1  # 减1是因为"其他"选项没有对应视频
                    rows_needed = (method_videos_count + 1) // 2  # 向上取整，确保所有视频都能显示
                    
                    for row in range(rows_needed):
                        with gr.Row():
                            for col in range(2):  # 每行2个视频
                                video_index = row * 2 + col + 1  # 计算当前视频索引
                                if video_index < len(video_paths):  # 确保索引有效
                                    with gr.Column(scale=1):
                                        gr.Video(
                                            value=video_paths[video_index], 
                                            label=anon_method_names[video_index-1], 
                                            autoplay=False, 
                                            height=300
                                        )

                    rating = gr.Radio(choices=anon_method_names, label="请选择生成效果最好的方法")
                    result_output = gr.HTML(label="提交结果")
                    
                    # 为每个组创建单独的提交按钮
                    submit_button = gr.Button(f"提交组 {group_idx+1} 评分")
                    
                    # 创建闭包函数以正确传递method_mapping字典和组编号
                    def create_submit_fn(group_name, method_map, group_idx):
                        def submit_fn(user_id, rating):
                            return submit_group_rating(user_id, group_name, rating, method_map, group_idx+1)
                        return submit_fn
                    
                    submit_button.click(
                        fn=create_submit_fn(group, method_mapping, group_idx),
                        inputs=[user_id_state, rating],
                        outputs=result_output
                    )

        # 添加查看完整结果的功能
        with gr.Row():
            view_results_button = gr.Button("查看我的所有评分")
            all_results_output = gr.HTML(label="评分汇总")
        
        # 使用闭包来确保正确传递groups参数
        def create_view_results_fn(current_groups):
            def view_fn(user_id):
                return view_all_results(user_id, current_groups)
            return view_fn
            
        view_results_button.click(
            fn=create_view_results_fn(groups), 
            inputs=[user_id_state], 
            outputs=all_results_output
        )
        
    return demo

# 启动Gradio应用
if __name__ == "__main__":
    # 确保所有必要的目录都有正确的访问权限
    prepare_all_videos()  # 预处理所有视频
    
    # 打印调试信息
    print(f"BASE_DIR: {BASE_DIR}")
    print(f"SCALED_VIDEOS_DIR: {SCALED_VIDEOS_DIR}")
    
    demo = create_interface()
    # 修改启动参数，解决文件访问问题
    demo.launch(
        server_name="0.0.0.0",  # 允许外部访问
        server_port=7860,  # 固定端口
        share=False,  # 设置为True可以生成公共链接
    )