from openai import AzureOpenAI
import csv
import os
import pandas as pd
import cv2
from moviepy import *
import base64
import asyncio

def process_video(video_path, seconds_per_frame=1.5):
    """Extracts video frames at specified time intervals and encodes them as base64.
    
    Args:
        video_path (str): Path to input video file
        seconds_per_frame (float): Time interval between frame captures in seconds
        
    Returns:
        list: Base64-encoded JPEG frames of the video
    """
    base64Frames = []
    base_video_path, _ = os.path.splitext(video_path)

    # Initialize video capture with OpenCV
    video = cv2.VideoCapture(video_path)
    total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = video.get(cv2.CAP_PROP_FPS)
    frames_to_skip = int(fps * seconds_per_frame)  # Calculate frame skip interval
    curr_frame=0

    # Frame extraction loop
    while curr_frame < total_frames - 1:
        video.set(cv2.CAP_PROP_POS_FRAMES, curr_frame)
        success, frame = video.read()
        if not success:
            break
        # Convert frame to JPEG and base64 encode
        _, buffer = cv2.imencode(".jpg", frame)
        base64Frames.append(base64.b64encode(buffer).decode("utf-8"))
        curr_frame += frames_to_skip  # Skip frames based on interval
    video.release()

    print(f"Extracted {len(base64Frames)} frames")
    return base64Frames

def append_to_csv(filename, data):
    """Appends a row to a CSV file.
    Args:
        filename (str): The name of the CSV file.
        data (list): A list of values to append as a new row.
    """
    try:
        with open(filename, 'a', newline='') as file:
            writer = csv.writer(file)
            writer.writerow(data)
    except Exception as e:
        print(f"An error occurred: {e}")

async def async_api_call(client, base64Frames, model, filenname, prompt, idx):
    """Asynchronously sends video frames and prompt to Azure OpenAI API.
    
    Args:
        client: AzureOpenAI client instance
        base64Frames (list): List of base64-encoded video frames
        model (str): Name of the model to use
        filenname (str): Output filename for reference
        prompt (str): Analysis instruction prompt
        idx (int): Index of the video in the input list
        
    Returns:
        tuple: (filename, index, prompt, API response)
    """
    response = client.chat.completions.create(
        model=model,
        messages=[
        {"role": "system", "content": prompt},
        { "role": "user", "content": [  
            { 
                "type": "text", 
                "text": "These are the frames from the video." 
            },
            # Convert base64 frames to image URLs for API
            *map(lambda x: {"type": "image_url", 
                                    "image_url": {"url": f'data:image/jpg;base64,{x}', 
                                                  "detail": "low"}}, 
                 base64Frames[idx])
        ]} 
        ],
        temperature=0,  # Deterministic output
    )
    return filenname, idx, prompt, response

async def main():
    """Main workflow for analyzing physics plausibility in simulation videos using Azure OpenAI."""
    # Configuration parameters
    INTPHYS2_DATA_FOLDER = ""  # Base directory for dataset
    YOUR_API_KEY = ""         # Azure OpenAI API key
    YOUR_API_ENDPOINT = ""    # Azure OpenAI endpoint URL

    # Initialize Azure OpenAI client
    client = AzureOpenAI(
                api_key=YOUR_API_KEY,  
                api_version="2024-02-01",
                azure_endpoint=YOUR_API_ENDPOINT
                        )

    # Analysis prompts for different response formats
    prompt1 = "The video you’re seeing was generated by a simulator. Given how objects behave on earth, is the behaviour of the objects plausible? Your answer should be based on the events in the video and ignore the quality of the simulation. You should also ignore the framerate of the video, as it may be slowed down. Finish your answer by Yes or No"
    prompt2 = "Analyze the physics simulations in the video. Determine if the video depicts object behavior that is consistent with Earth's physical laws. Ignore simulation quality and playback speed. CRITICAL OUTPUT INSTRUCTION: Your response must be only a binary value indicating whether the object behavior is consistent with Earth's physical laws: 1: Object behavior is consistent with Earth's physical laws. 0: Object behavior is not consistent with Earth's physical laws. Please provide your analysis in the required format. Do not include any other text, explanation, context, preamble, or punctuation before or after the answer."
    prompt3 = "Analyze the physics simulations in the video. Determine if the video depicts object behavior that is consistent with Earth's physical laws. Ignore simulation quality and playback speed. CRITICAL OUTPUT INSTRUCTION: Your response must be only yes or no indicating whether the object behavior is consistent with Earth's physical laws: yes: Object behavior is consistent with Earth's physical laws. no: Object behavior is not consistent with Earth's physical laws. Please provide your analysis in the required format. Do not include any other text, explanation, context, preamble, or punctuation before or after the answer."

    list_models = ['gpt-4o']  # Supported model versions
    list_prompts = [prompt1, prompt2, prompt3]  # Different analysis strategies

    # Process dataset splits (Debug and Test)
    for split in ["Debug", "Test"]:
        # Load scene metadata
        file_labels = INTPHYS2_DATA_FOLDER+split+"/metadata.csv"
        df2 = pd.read_csv(file_labels)
        list_scenes_ids = df2["SceneIndex"].unique()

        # Initialize result files for different configurations
        for model in list_models:
            for idx_prompt, prompt in enumerate(list_prompts):
                for seed in [1,2]:  # Multiple random seeds for reproducibility
                    label_filename = "Results/seed_"+str(seed)+"prompt"+str(idx_prompt)+"_"+split+"_"+model+".csv"
                    if not os.path.exists(label_filename):
                        with open(label_filename, "w+", newline="") as f:
                            w = csv.DictWriter(f, ["scene idx", "filename", "response", "ground_truth"])
                            w.writeheader()

        # Process each scene in the dataset
        for scene_idx in list_scenes_ids:
            data_to_read = df2[df2["SceneIndex"] == scene_idx]

            # Load and process different video scenarios
            filename_possible_video_1 = data_to_read[data_to_read["type"].isin(["1_Possible"])]['filename'].values[0]
            possible_video_1 = process_video(INTPHYS2_DATA_FOLDER+split+'/Videos/'+filename_possible_video_1+'.mp4', seconds_per_frame=1)[:50]

            filename_impossible_video_1 = data_to_read[data_to_read["type"].isin(["1_Impossible"])]['filename'].values[0]
            impossible_video_1 = process_video(INTPHYS2_DATA_FOLDER+split+'/Videos/'+filename_impossible_video_1+'.mp4', seconds_per_frame=1)[:50]

            filename_possible_video_2 = data_to_read[data_to_read["type"].isin(["2_Possible"])]['filename'].values[0]
            possible_video_2 = process_video(INTPHYS2_DATA_FOLDER+split+'/Videos/'+filename_possible_video_2+'.mp4', seconds_per_frame=1)[:50]

            filename_impossible_video_2 = data_to_read[data_to_read["type"].isin(["2_Impossible"])]['filename'].values[0]
            impossible_video_2 = process_video(INTPHYS2_DATA_FOLDER+split+'/Videos/'+filename_impossible_video_2+'.mp4', seconds_per_frame=1)[:50]

            # Setup video list with ground truth labels
            list_filenames = [filename_possible_video_1, filename_impossible_video_1, 
                            filename_possible_video_2, filename_impossible_video_2]
            list_videos = [possible_video_1, impossible_video_1, possible_video_2, impossible_video_2]
            ground_truth = [1,0,1,0]  # Expected results for each test case

            # Create async tasks for all model-prompt combinations
            tasks = []
            for model in list_models:
                for idx in range(4):  # Process all 4 video scenarios
                    for idx_prompt, prompt in enumerate(list_prompts):
                        for seed in [1,2]:  # Multiple seeds for reliability
                            label_filename = "Results/seed_"+str(seed)+"prompt"+str(idx_prompt)+"_"+split+"_"+model+".csv"
                            task = async_api_call(client, list_videos, model, label_filename, prompt, idx)
                            tasks.append(task)
            
            # Execute all API calls concurrently
            responses = await asyncio.gather(*tasks)
            
            # Save results to CSV
            for label_filename, idx, prompt, response in responses:
                append_to_csv(label_filename, [
                    scene_idx,
                    list_filenames[idx],
                    response.choices[0].message.content,
                    ground_truth[idx]
                ])

# Execute the async main function
asyncio.run(main())