from transformers import Qwen2_5_VLForConditionalGeneration, AutoTokenizer, AutoProcessor
from qwen_vl_utils import process_vision_info
import time
import csv
import os
import pandas as pd
import time
import torch

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}")

def async_api_call(processor, base64Frames, model, filenname, prompt, idx):
    """Processes video frames and generates text response using the Qwen-VL model.
    
    Args:
        processor: Model processor for input preparation
        base64Frames: Video frames in base64 encoding
        model: Loaded Qwen-VL model instance
        filenname (str): Filename for reference
        prompt (str): Instruction prompt for analysis
        idx (int): Index identifier for the video
        
    Returns:
        tuple: (filename, index, prompt, model response)
    """
    # Construct conversation messages with video data
    messages = [
        {"role": "system", "content": prompt},
        {
            "role": "user",
            "content": [
                { 
                    "type": "text", 
                    "text": "These are the frames from the video." 
                },
                {
                    "type": "video",
                    "video": base64Frames,
                    "fps": 1.0,
                },
                {"type": "text", "text": "Describe this video."},
            ],
        }
    ]
    
    # Prepare inputs for model
    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    # Process vision information from messages
    image_inputs, video_inputs, video_kwargs = process_vision_info(messages, return_video_kwargs=True)
    
    # Create model inputs
    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
        **video_kwargs,
    )
    inputs = inputs.to("cuda")  # Move inputs to GPU
    # Generate model response
    generated_ids = model.generate(**inputs, max_new_tokens=512, temperature=0.0, do_sample=False)
    
    # Post-process output
    generated_ids_trimmed = [
        out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    output_text = processor.batch_decode(
        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )
    return filenname, idx, prompt, output_text

def main():

    # Configuration paths
    INTPHYS2_DATA_FOLDER = "" # Base data directory to replace

    # Define different analysis prompts
    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 of prompts for different analysis
    list_prompts = [prompt1, prompt2, prompt3]

    # Initialize model and processor
    model_name = 'Qwen2_5_VL'
    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
         "Qwen/Qwen2.5-VL-72B-Instruct",
         torch_dtype=torch.bfloat16,
         attn_implementation="flash_attention_2",
         device_map="auto",
    )

    # Initialize processor for multimodal inputs
    processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-72B-Instruct")
    # Process different dataset splits
    for split in ["Debug", "Test"]:
        # Load scene labels metadata
        file_labels = INTPHYS2_DATA_FOLDER+split+"/metadata.csv"
        df2 = pd.read_csv(file_labels)
        list_scenes_ids = df2["SceneIndex"].unique()
        # Prepare result files for different prompts and seeds
        for idx_prompt, prompt in enumerate(list_prompts):
            for seed in [1,2]:
                label_filename = "Results/seed_"+str(seed)+"prompt"+str(idx_prompt)+"_"+split+"_"+model_name+".csv"
                if not os.path.exists(label_filename):
                    # Create new result file with header
                    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]
            # Construct video file paths for different test cases
            filename_possible_video_1 =  INTPHYS2_DATA_FOLDER+split+'/Videos/' + data_to_read[data_to_read["type"].isin(["1_Possible"])]['filename'].values[0]+'.mp4'
            filename_impossible_video_1 =  INTPHYS2_DATA_FOLDER+split+'/Videos/' + data_to_read[data_to_read["type"].isin(["1_Impossible"])]['filename'].values[0]+'.mp4'
            filename_possible_video_2 =  INTPHYS2_DATA_FOLDER+split+'/Videos/' + data_to_read[data_to_read["type"].isin(["2_Possible"])]['filename'].values[0]+'.mp4'
            filename_impossible_video_2 =  INTPHYS2_DATA_FOLDER+split+'/Videos/' + data_to_read[data_to_read["type"].isin(["2_Impossible"])]['filename'].values[0]+'.mp4'
            # List of test cases with corresponding ground truths
            list_filenames = [filename_possible_video_1, filename_impossible_video_1, 
                            filename_possible_video_2, filename_impossible_video_2]
            ground_truth = [1,0,1,0]  # Expected results for each test case
            responses = []
            # Process each video with different prompts and seeds
            for idx in range(4):  # Iterate through 4 test cases
                for idx_prompt, prompt in enumerate(list_prompts):
                    for seed in [1,2]:  # Two different random seeds
                        label_filename = "Results/seed_"+str(seed)+"prompt"+str(idx_prompt)+"_"+split+"_"+model_name+".csv"
                        # Make api call to get model's answer
                        task = async_api_call(processor, list_filenames[idx], model, label_filename, prompt, idx) 
                        responses.append(task)
            
            # Save collected responses to CSV
            for label_filename, idx, prompt, response in responses:
                append_to_csv(label_filename, [scene_idx, list_filenames[idx], response[0], ground_truth[idx]])

main()
