# List of requirements
# torch~=1.13
# torchvision
# opencv-python
# scipy
# numpy
# tqdm
# timm
# einops
# scikit-video
# pillow
# logger
# diffusers
# transformers
# accelerate
# requests
# pycocoevalcap

import os
import torch
import cv2
import numpy as np
from PIL import Image
from transformers import CLIPProcessor, CLIPModel, AutoTokenizer
import time
import logging
from tqdm import tqdm
import argparse
import torchvision.transforms as transforms
from torchvision.transforms import Resize
from torchvision.utils import save_image
from diffusers import StableDiffusionXLPipeline
import requests
from transformers import AutoProcessor, Blip2ForConditionalGeneration
import ipdb
from pycocoevalcap.cider.cider import Cider
from pycocoevalcap.bleu.bleu import Bleu

def calculate_clip_score(video_path, text, model, tokenizer):
    # Load the video
    cap = cv2.VideoCapture(video_path)

    # Extract frames from the video 
    frames = []

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        resized_frame = cv2.resize(frame,(224,224))  # Resize the frame to match the expected input size
        frames.append(resized_frame)

    # Convert numpy arrays to tensors, change dtype to float, and resize frames
    tensor_frames = [torch.from_numpy(frame).permute(2, 0, 1).float() for frame in frames]

    # Initialize an empty tensor to store the concatenated features
    concatenated_features = torch.tensor([], device=device)

    # Generate embeddings for each frame and concatenate the features
    with torch.no_grad():
        for frame in tensor_frames:
            frame_input = frame.unsqueeze(0).to(device)  # Add batch dimension and move the frame to the device
            frame_features = model.get_image_features(frame_input)
            concatenated_features = torch.cat((concatenated_features, frame_features), dim=0)

    # Tokenize the text
    text_tokens = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=77)

    # Convert the tokenized text to a tensor and move it to the device
    text_input = text_tokens["input_ids"].to(device)

    # Generate text embeddings
    with torch.no_grad():
        text_features = model.get_text_features(text_input)

    # Calculate the cosine similarity scores
    concatenated_features = concatenated_features / concatenated_features.norm(p=2, dim=-1, keepdim=True)
    text_features = text_features / text_features.norm(p=2, dim=-1, keepdim=True)
    clip_score_frames = concatenated_features @ text_features.T
    # Calculate the average CLIP score across all frames, reflects temporal consistency 
    clip_score_frames_avg = clip_score_frames.mean().item()

    return clip_score_frames_avg

def calculate_clip_temp_score(video_path, model):
    # Load the video
    cap = cv2.VideoCapture(video_path)
    to_tensor = transforms.ToTensor()
    # Extract frames from the video 
    frames = []
    SD_images = []
    resize = transforms.Resize([224,224])
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        # resized_frame = cv2.resize(frame,(224,224))  # Resize the frame to match the expected input size
        frames.append(frame)
    
    tensor_frames = torch.stack([resize(torch.from_numpy(frame).permute(2, 0, 1).float()) for frame in frames])

    # tensor_frames = [extracted_frames[i] for i in range(extracted_frames.size()[0])]
    concatenated_frame_features = []

    # Generate embeddings for each frame and concatenate the features
    with torch.no_grad():  
        for frame in tensor_frames: # Too many frames in a video, must split before CLIP embedding, limited by the memory
            frame_input = frame.unsqueeze(0).to(device)  # Add batch dimension and move the frame to the device
            frame_feature = model.get_image_features(frame_input)
            concatenated_frame_features.append(frame_feature)

    concatenated_frame_features = torch.cat(concatenated_frame_features, dim=0)

    # Calculate the similarity scores
    clip_temp_score = []
    concatenated_frame_features = concatenated_frame_features / concatenated_frame_features.norm(p=2, dim=-1, keepdim=True)
    # ipdb.set_trace()

    for i in range(concatenated_frame_features.size()[0]-1):
        clip_temp_score.append(concatenated_frame_features[i].unsqueeze(0) @ concatenated_frame_features[i+1].unsqueeze(0).T)
    clip_temp_score=torch.cat(clip_temp_score, dim=0)
    # Calculate the average CLIP score across all frames, reflects temporal consistency 
    clip_temp_score_avg = clip_temp_score.mean().item()

    return clip_temp_score_avg

def compute_max(scorer, gt_prompts, pred_prompts):
    scores = []
    for pred_prompt in pred_prompts:
        for gt_prompt in gt_prompts:
            cand = {0: [pred_prompt]}
            ref = {0: [gt_prompt]}
            score, _ = scorer.compute_score(ref, cand)
            scores.append(score)
    return np.max(scores)

def calculate_blip_bleu(video_path, original_text, blip2_model, blip2_processor):
    # Load the video
    cap = cv2.VideoCapture(video_path)

    scorer_cider = Cider()
    bleu1 = Bleu(n=1)
    bleu2 = Bleu(n=2)
    bleu3 = Bleu(n=3)
    bleu4 = Bleu(n=4)

    # Extract frames from the video
    frames = []
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        resized_frame = cv2.resize(frame,(224,224))  # Resize the frame to match the expected input size
        frames.append(resized_frame)

    # Convert numpy arrays to tensors, change dtype to float, and resize frames
    tensor_frames = torch.stack([torch.from_numpy(frame).permute(2, 0, 1).float() for frame in frames])
    # Get five captions for one video
    Num = 5
    captions = []
    # for i in range(Num):
    N = len(tensor_frames)
    indices = torch.linspace(0, N - 1, Num).long()
    extracted_frames = torch.index_select(tensor_frames, 0, indices)
    for i in range(Num):
        frame = extracted_frames[i]
        inputs = blip2_processor(images=frame, return_tensors="pt").to(device, torch.float16)
        generated_ids = blip2_model.generate(**inputs)
        generated_text = blip2_processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
        captions.append(generated_text)


    original_text = [original_text]
    cider_score = (compute_max(scorer_cider, original_text, captions))
    bleu1_score = (compute_max(bleu1, original_text, captions))
    bleu2_score = (compute_max(bleu2, original_text, captions))
    bleu3_score = (compute_max(bleu3, original_text, captions))
    bleu4_score = (compute_max(bleu4, original_text, captions))

    blip_bleu_caps_avg = (bleu1_score + bleu2_score + bleu3_score + bleu4_score)/4
     
    return blip_bleu_caps_avg

def calculate_sd_score(video_path, text, pipe, model):
    # Load the video
    output_dir = "../../SDXL_Imgs"
    if not os.path.exists(output_dir):
        os.mkdir(output_dir)
    cap = cv2.VideoCapture(video_path)
    to_tensor = transforms.ToTensor()
    # Extract frames from the video 
    frames = []
    SD_images = []
    Num = 5
    resize = transforms.Resize([224,224])
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        # resized_frame = cv2.resize(frame,(224,224))  # Resize the frame to match the expected input size
        frames.append(frame)
    
    # Load SD imgs from local paths
    for i in range(Num): ## Num images for every prompt
        output_dir = "../../SDXL_Imgs"
        # ipdb.set_trace()
        SD_image_path = os.path.join(output_dir, f"{os.path.basename(video_path).split('.')[0]}_{i}.png")
        if os.path.exists(SD_image_path):
            image = Image.open(SD_image_path)
            # Convert the image to a tensor
            image = resize(to_tensor(image))
            SD_images.append(image.unsqueeze(0)) 
        else:
            image = pipe(text, height = 512, width= 512, num_inference_steps = 20).images[0]  #!!!!! same amount of SD images, but also can be mutiple times, TODO
            # Convert the image to a tensor
            image = resize(to_tensor(image))
            SD_images.append(image.unsqueeze(0)) 
            save_image(image,SD_image_path)

    tensor_frames = [resize(torch.from_numpy(frame).permute(2, 0, 1).float()) for frame in frames]
    SD_images = torch.cat(SD_images, 0)

    concatenated_frame_features = []
    concatenated_SDImg_features = []
    # Generate embeddings for each frame and concatenate the features
    with torch.no_grad():  
        for frame in tensor_frames: # Too many frames in a video, must split before CLIP embedding, limited by the memory
            frame_input = frame.unsqueeze(0).to(device)  # Add batch dimension and move the frame to the device
            frame_feature = model.get_image_features(frame_input)
            concatenated_frame_features.append(frame_feature)

        for i in range(SD_images.size()[0]):
            img = SD_images[i].unsqueeze(0).to(device)  # Add batch dimension and move the frame to the device
            SDImg_feature  = model.get_image_features(img)
            concatenated_SDImg_features.append(SDImg_feature)
    # ipdb.set_trace()
    concatenated_frame_features = torch.cat(concatenated_frame_features, dim=0)
    concatenated_SDImg_features = torch.cat(concatenated_SDImg_features, dim=0)

    # Calculate the similarity scores
    concatenated_frame_features = concatenated_frame_features / concatenated_frame_features.norm(p=2, dim=-1, keepdim=True)
    concatenated_SDImg_features = concatenated_SDImg_features / concatenated_SDImg_features.norm(p=2, dim=-1, keepdim=True)
    sd_score_frames = concatenated_frame_features @ concatenated_SDImg_features.T
    # Calculate the average CLIP score across all frames, reflects temporal consistency 
    sd_score_frames_avg = sd_score_frames.mean().item()

    return sd_score_frames_avg

def calculate_face_consistency_score(video_path, model):
    # Load the video
    cap = cv2.VideoCapture(video_path)
    to_tensor = transforms.ToTensor()
    # Extract frames from the video 
    frames = []
    SD_images = []
    resize = transforms.Resize([224,224])
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        # resized_frame = cv2.resize(frame,(224,224))  # Resize the frame to match the expected input size
        frames.append(frame)
    
    tensor_frames = [resize(torch.from_numpy(frame).permute(2, 0, 1).float()) for frame in frames]
    concatenated_frame_features = []

    # Generate embeddings for each frame and concatenate the features
    with torch.no_grad():  
        for frame in tensor_frames: # Too many frames in a video, must split before CLIP embedding, limited by the memory
            frame_input = frame.unsqueeze(0).to(device)  # Add batch dimension and move the frame to the device
            frame_feature = model.get_image_features(frame_input)
            concatenated_frame_features.append(frame_feature)

    concatenated_frame_features = torch.cat(concatenated_frame_features, dim=0)

    # Calculate the similarity scores
    concatenated_frame_features = concatenated_frame_features / concatenated_frame_features.norm(p=2, dim=-1, keepdim=True)
    face_consistency_score = concatenated_frame_features[1:] @ concatenated_frame_features[0].unsqueeze(0).T
    # Calculate the average CLIP score across all frames, reflects temporal consistency 
    face_consistency_score_avg = face_consistency_score.mean().item()

    return face_consistency_score_avg

def read_text_file(file_path):
    with open(file_path, 'r') as f:
        return f.read().strip()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--dir_videos", type=str, default='', help="Specify the path of generated videos")
    parser.add_argument("--metric", type=str, default='celebrity_id_score', help="Specify the metric to be used")
    args = parser.parse_args()

    dir_videos = args.dir_videos
    metric = args.metric

    dir_prompts =  '../../prompts/'
   
    video_paths = [os.path.join(dir_videos, x) for x in os.listdir(dir_videos)]
    prompt_paths = [os.path.join(dir_prompts, os.path.splitext(os.path.basename(x))[0]+'.txt') for x in video_paths]

     # Create the directory if it doesn't exist
    timestamp = time.strftime("%Y%m%d-%H%M%S")
    os.makedirs(f"../../results", exist_ok=True)
    # Set up logging
    log_file_path = f"../../results/{metric}_record.txt"
    # Delete the log file if it exists
    if os.path.exists(log_file_path):
        os.remove(log_file_path)
    # Set up logging
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)
    # File handler for writing logs to a file
    file_handler = logging.FileHandler(filename=f"../../results/{metric}_record.txt")
    file_handler.setFormatter(logging.Formatter("%(asctime)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S"))
    logger.addHandler(file_handler)
    # Stream handler for displaying logs in the terminal
    stream_handler = logging.StreamHandler()
    stream_handler.setFormatter(logging.Formatter("%(asctime)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S"))
    logger.addHandler(stream_handler)


    # Load pretrained models
    device = "cuda" if torch.cuda.is_available() else "cpu"

    
    if metric == 'blip_bleu': 
        blip2_processor = AutoProcessor.from_pretrained("../../checkpoints/blip2-opt-2.7b")
        blip2_model = Blip2ForConditionalGeneration.from_pretrained("../../checkpoints/blip2-opt-2.7b", torch_dtype=torch.float16).to(device)
    elif metric == 'sd_score':
        clip_model = CLIPModel.from_pretrained("../../checkpoints/clip-vit-base-patch32").to(device)
        clip_tokenizer = AutoTokenizer.from_pretrained("../../checkpoints/clip-vit-base-patch32")
        output_dir = "/apdcephfs/share_1290939/raphaelliu/Vid_Eval/Video_Gen/prompt700-release/SDXL_Imgs"
        SD_image_path = os.path.join(output_dir, f"{os.path.basename(os.path.basename(video_paths[0]).split('.')[0])}_0.png")
        # if os.path.exists(SD_image_path):
        #     pipe = None
        # else:
        pipe = StableDiffusionXLPipeline.from_pretrained(
            "../../checkpoints/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)
        pipe = pipe.to(device)
    else:
        clip_model = CLIPModel.from_pretrained("../../checkpoints/clip-vit-base-patch32").to(device)
        clip_tokenizer = AutoTokenizer.from_pretrained("../../checkpoints/clip-vit-base-patch32")
    
    # Calculate SD scores for all video-text pairs
    scores = []
    
    test_num = 10
    test_num = len(video_paths)
    count = 0
    for i in tqdm(range(len(video_paths))):
        video_path = video_paths[i]
        prompt_path = prompt_paths[i]
        if count == test_num:
            break
        else:
            text = read_text_file(prompt_path)
            # ipdb.set_trace()
            if metric == 'clip_score':
                score = calculate_clip_score(video_path, text, clip_model, clip_tokenizer)
            elif metric == 'blip_bleu': 
                score = calculate_blip_bleu(video_path, text, blip2_model, blip2_processor)
            elif metric == 'sd_score':
                score = calculate_sd_score(video_path, text, pipe,clip_model)
            elif metric == 'clip_temp_score':
                score = calculate_clip_temp_score(video_path,clip_model)
            elif metric == 'face_consistency_score':
                score = calculate_face_consistency_score(video_path,clip_model)
            count+=1
            scores.append(score)
            average_score = sum(scores) / len(scores)
            # count+=1
            logging.info(f"Vid: {os.path.basename(video_path)},  Current {metric}: {score}, Current avg. {metric}: {average_score},  ")
            
    # Calculate the average SD score across all video-text pairs
    logging.info(f"Final average {metric}: {average_score}, Total videos: {len(scores)}")