"""
This script demonstrates how to generate a video using the CogVideoX model with the Hugging Face `diffusers` pipeline.
The script supports different types of video generation, including text-to-video (t2v), image-to-video (i2v),
and video-to-video (v2v), depending on the input data and different weight.

- text-to-video: THUDM/CogVideoX-5b, THUDM/CogVideoX-2b or THUDM/CogVideoX1.5-5b
- video-to-video: THUDM/CogVideoX-5b, THUDM/CogVideoX-2b or THUDM/CogVideoX1.5-5b
- image-to-video: THUDM/CogVideoX-5b-I2V or THUDM/CogVideoX1.5-5b-I2V

Running the Script:
To run the script, use the following command with appropriate arguments:

```bash
$ python cli_demo.py --prompt "A girl riding a bike." --model_path THUDM/CogVideoX1.5-5b --generate_type "t2v"
```

Additional options are available to specify the model path, guidance scale, number of inference steps, video generation type, and output paths.
"""

import argparse
from typing import Literal

import torch
from diffusers import (
    # CogVideoXPipeline,
    # CogVideoXDPMScheduler,
    CogVideoXImageToVideoPipeline,
    CogVideoXVideoToVideoPipeline,
)
from scheduler import CogVideoXDPMScheduler
from pipeline import CogVideoXPipeline

from diffusers.utils import export_to_video, load_image, load_video
import json 
import re 
import random 
import numpy as np 
import cv2 
from einops import rearrange, repeat
from torchvision.utils import flow_to_image
import os 
from tqdm import tqdm 
import torchvision.transforms.v2 as tr
import imageio
from torchvision.transforms import InterpolationMode
from diffusers.image_processor import VaeImageProcessor
from PIL import Image 


def vis_tensor(x, h, w, output_dir, min_=None, max_=None):
    """
    Args:
        x (torch.Tensor): of shape [T, C, H, W]
    """
    x = x.mean(dim=1)
    if min_ is None:
        min_ = x.min()
    if max_ is None:
        max_ = x.max()
    x = (x - min_) / (max_ - min_)
    x = x.clamp(0, 1)
    x = tr.Resize([h, w], interpolation=InterpolationMode.NEAREST)(x)
    x = (x * 255).to(torch.uint8).cpu().numpy()
    x = [xx for xx in x]
    imageio.mimwrite(output_dir, x, fps=8, loop=0)
    
    
def set_random_seed(seed):
    """Set random seed for reproducability."""
    if seed is not None:
        global _GLOBAL_RANDOM_SEED
        _GLOBAL_RANDOM_SEED = seed
        assert seed > 0
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.enabled = True # False
        torch.backends.cuda.matmul.allow_tf32 = False # if set it to True will be much faster but not accurate
        try:
            import deepspeed
            if deepspeed.checkpointing.is_configured():
                mpu.model_parallel_cuda_manual_seed(seed)
        except ImportError:
            pass

def sigma_matrix2(sig_x, sig_y, theta):
    """Calculate the rotated sigma matrix (two dimensional matrix).
    Args:
        sig_x (float):
        sig_y (float):
        theta (float): Radian measurement.
    Returns:
        ndarray: Rotated sigma matrix.
    """
    d_matrix = np.array([[sig_x**2, 0], [0, sig_y**2]])
    u_matrix = np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]])
    return np.dot(u_matrix, np.dot(d_matrix, u_matrix.T))


def mesh_grid(kernel_size):
    """Generate the mesh grid, centering at zero.
    Args:
        kernel_size (int):
    Returns:
        xy (ndarray): with the shape (kernel_size, kernel_size, 2)
        xx (ndarray): with the shape (kernel_size, kernel_size)
        yy (ndarray): with the shape (kernel_size, kernel_size)
    """
    ax = np.arange(-kernel_size // 2 + 1.0, kernel_size // 2 + 1.0)
    xx, yy = np.meshgrid(ax, ax)
    xy = np.hstack(
        (
            xx.reshape((kernel_size * kernel_size, 1)),
            yy.reshape(kernel_size * kernel_size, 1),
        )
    ).reshape(kernel_size, kernel_size, 2)
    return xy, xx, yy


def pdf2(sigma_matrix, grid):
    """Calculate PDF of the bivariate Gaussian distribution.
    Args:
        sigma_matrix (ndarray): with the shape (2, 2)
        grid (ndarray): generated by :func:`mesh_grid`,
            with the shape (K, K, 2), K is the kernel size.
    Returns:
        kernel (ndarrray): un-normalized kernel.
    """
    inverse_sigma = np.linalg.inv(sigma_matrix)
    kernel = np.exp(-0.5 * np.sum(np.dot(grid, inverse_sigma) * grid, 2))
    return kernel


def bivariate_Gaussian(kernel_size, sig_x, sig_y, theta, grid=None, isotropic=True):
    """Generate a bivariate isotropic or anisotropic Gaussian kernel.
    In the isotropic mode, only `sig_x` is used. `sig_y` and `theta` is ignored.
    Args:
        kernel_size (int):
        sig_x (float):
        sig_y (float):
        theta (float): Radian measurement.
        grid (ndarray, optional): generated by :func:`mesh_grid`,
            with the shape (K, K, 2), K is the kernel size. Default: None
        isotropic (bool):
    Returns:
        kernel (ndarray): normalized kernel.
    """
    if grid is None:
        grid, _, _ = mesh_grid(kernel_size)
    if isotropic:
        sigma_matrix = np.array([[sig_x**2, 0], [0, sig_x**2]])
    else:
        sigma_matrix = sigma_matrix2(sig_x, sig_y, theta)
    kernel = pdf2(sigma_matrix, grid)
    kernel = kernel / np.sum(kernel)
    return kernel


def read_points(files, video_len=16, reverse=False):
    output_points = []
    files = files.split(",")
    for file in files:
        with open(file, "r") as f:
            lines = f.readlines()
        points = []
        for line in lines:
            x, y = line.strip().split(",")
            points.append((int(x), int(y)))
        if reverse:
            points = points[::-1]

        if len(points) > video_len:
            skip = len(points) // video_len
            points = points[::skip]
        points = points[:video_len]
        output_points.append(points)

    return output_points


size = 99
sigma = 10
blur_kernel = bivariate_Gaussian(size, sigma, sigma, 0, grid=None, isotropic=True)
blur_kernel = blur_kernel / blur_kernel[size // 2, size // 2]


def parse_frames(input_text):
    frames = []
    # Improved pattern for parsing frames and captions
    pattern = re.compile(
        r'Frame_(\d+):\s*(\[\[.*?\]\]|\[.*?\])\s*,\s*caption:\s*(.*?)(?=(?:Frame_\d+:|\Z))',
        re.DOTALL
    )
    # pattern2 = re.compile(
    #     r'Frame_(\d+):\s*\[(.*?)\],\s*caption:\s*(.*?)(?=(?:Frame_\d+:|\Z))',
    #     re.DOTALL
    # )
    matches = pattern.finditer(input_text)
    # matches2 = pattern2.finditer(input_text)
    # matches = list(matches) + list(matches2)
    for match in matches:
        frame_number = int(match.group(1))
        objects_list_str = match.group(2).strip()
        caption = match.group(3).strip()
        
        if "Reasoning:" in caption:
            caption = caption.split("Reasoning:")[0].strip("*").strip("\n")

        # Parse the objects list string into a Python list
        try:
            objects_list = json.loads(f"{objects_list_str}")
        except json.JSONDecodeError as e:
            print(f"Error parsing JSON in frame {frame_number}: {e}")
            continue

        frames.append({
            'frame_number': frame_number,
            'objects': objects_list,
            'caption': caption
        })
    return frames

def process_traj_data(points_files, num_frames, video_size, device="cpu"):
    optical_flow = np.zeros((num_frames, video_size[0], video_size[1], 2), dtype=np.float32)
    processed_points = []
    xy_range = 256
    h, w = video_size
    for points in points_files:
        # multi_points = read_points(points_file, video_len=num_frames)
        # for points in multi_points:
        points = process_points(points, num_frames)
        points = [[int(w * x / xy_range), int(h * y / xy_range)] for x, y in points]
        optical_flow = get_flow(points, optical_flow, video_len=num_frames)
        processed_points.append(points)

    for i in range(1, num_frames):
        optical_flow[i] = cv2.filter2D(optical_flow[i], -1, blur_kernel)

    optical_flow = torch.tensor(optical_flow).to(device)

    return optical_flow, processed_points


def process_points(points, frames):
    defualt_points = [[512, 512]] * frames

    if len(points) < 2:
        return defualt_points
    elif len(points) >= frames:
        skip = len(points) // frames
        return points[::skip][: frames - 1] + points[-1:]
    else:
        insert_num = frames - len(points)
        insert_num_dict = {}
        interval = len(points) - 1
        n = insert_num // interval
        m = insert_num % interval
        for i in range(interval):
            insert_num_dict[i] = n
        for i in range(m):
            insert_num_dict[i] += 1

        res = []
        for i in range(interval):
            insert_points = []
            x0, y0 = points[i]
            x1, y1 = points[i + 1]

            delta_x = x1 - x0
            delta_y = y1 - y0
            for j in range(insert_num_dict[i]):
                x = x0 + (j + 1) / (insert_num_dict[i] + 1) * delta_x
                y = y0 + (j + 1) / (insert_num_dict[i] + 1) * delta_y
                insert_points.append([int(x), int(y)])

            res += points[i : i + 1] + insert_points
        res += points[-1:]
        return res


def get_flow(points, optical_flow, video_len):
    for i in range(video_len - 1):
        p = points[i]
        p1 = points[i + 1]
        if p[0] < 0:
            continue
        elif p[0] > 0 and p1[0] < 0:
            continue
        # print(p1, p)
        optical_flow[i + 1, p[1]-1, p[0]-1, 0] = p1[0] - p[0]
        optical_flow[i + 1, p[1]-1, p[0]-1, 1] = p1[1] - p[1]

    return optical_flow


def generate_video(
    prompt: str,
    model_path: str,
    lora_path: str = None,
    lora_rank: int = 128,
    num_frames: int = 81,
    width: int = 1360,
    height: int = 768,
    output_path: str = "./output.mp4",
    image_or_video_path: str = "",
    num_inference_steps: int = 50,
    guidance_scale: float = 6.0,
    num_videos_per_prompt: int = 1,
    dtype: torch.dtype = torch.bfloat16,
    generate_type: str = Literal["t2v", "i2v", "v2v"],  # i2v: image to video, v2v: video to video
    seed: int = 42,
    fps: int = 8,
    latent_path: str = ""
):
    """
    Generates a video based on the given prompt and saves it to the specified path.

    Parameters:
    - prompt (str): The description of the video to be generated.
    - model_path (str): The path of the pre-trained model to be used.
    - lora_path (str): The path of the LoRA weights to be used.
    - lora_rank (int): The rank of the LoRA weights.
    - output_path (str): The path where the generated video will be saved.
    - num_inference_steps (int): Number of steps for the inference process. More steps can result in better quality.
    - num_frames (int): Number of frames to generate. CogVideoX1.0 generates 49 frames for 6 seconds at 8 fps, while CogVideoX1.5 produces either 81 or 161 frames, corresponding to 5 seconds or 10 seconds at 16 fps.
    - width (int): The width of the generated video, applicable only for CogVideoX1.5-5B-I2V
    - height (int): The height of the generated video, applicable only for CogVideoX1.5-5B-I2V
    - guidance_scale (float): The scale for classifier-free guidance. Higher values can lead to better alignment with the prompt.
    - num_videos_per_prompt (int): Number of videos to generate per prompt.
    - dtype (torch.dtype): The data type for computation (default is torch.bfloat16).
    - generate_type (str): The type of video generation (e.g., 't2v', 'i2v', 'v2v').·
    - seed (int): The seed for reproducibility.
    - fps (int): The frames per second for the generated video.
    """

    # 1.  Load the pre-trained CogVideoX pipeline with the specified precision (bfloat16).
    # add device_map="balanced" in the from_pretrained function and remove the enable_model_cpu_offload()
    # function to use Multi GPUs.

    set_random_seed(seed)
    image = None
    video = None

    CACHE_DIR = ""
    if generate_type == "i2v":
        pipe = CogVideoXImageToVideoPipeline.from_pretrained(model_path, torch_dtype=dtype, cache_dir=CACHE_DIR)
        image = load_image(image=image_or_video_path)
    elif generate_type == "t2v":
        pipe = CogVideoXPipeline.from_pretrained(model_path, torch_dtype=dtype, cache_dir=CACHE_DIR, timeout=30)
    else:
        pipe = CogVideoXVideoToVideoPipeline.from_pretrained(model_path, torch_dtype=dtype, cache_dir=CACHE_DIR)
        video = load_video(image_or_video_path)

    # If you're using with lora, add this code
    if lora_path:
        pipe.load_lora_weights(lora_path, weight_name="pytorch_lora_weights.safetensors", adapter_name="test_1", cache_dir=CACHE_DIR)
        pipe.fuse_lora(lora_scale=1 / lora_rank)

    # 2. Set Scheduler.
    # Can be changed to `CogVideoXDPMScheduler` or `CogVideoXDDIMScheduler`.
    # We recommend using `CogVideoXDDIMScheduler` for CogVideoX-2B.
    # using `CogVideoXDPMScheduler` for CogVideoX-5B / CogVideoX-5B-I2V.

    pipe.scheduler = CogVideoXDPMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")

    # 3. Enable CPU offload for the model.
    # turn off if you have multiple GPUs or enough GPU memory(such as H100) and it will cost less time in inference
    # and enable to("cuda")

    pipe.to("cuda")
    pipe.enable_sequential_cpu_offload()
    # pipe.vae.enable_slicing()
    pipe.vae.enable_tiling()
    pipe.vae.requires_grad_(False)
    
        
    with open("T2V-CompBench/prompts/xxx", "r") as file:
        lines = file.readlines()
    
    # Remove trailing newline characters from each line
    eval_prompts = [line.strip() for line in lines]
    
    if not os.path.exists(output_path):
        os.mkdir(output_path)
    
    skill_name = "xxx"

    if not os.path.exists(f"{output_path}/{skill_name}"):
        os.mkdir(f"{output_path}/{skill_name}")
    
    for i, prompt in enumerate(tqdm(eval_prompts[0:])):
        index = i + 1
        
        flag = True  
        image_process = VaeImageProcessor()
        images = []
        
        for i in range(49):
            image = np.array(Image.open(f"{latent_path}/{skill_name}/{index:04d}/{i}.jpg").resize((720, 480)))
            images.append(image)
            
        frame_images = torch.tensor(np.array(images)).unsqueeze(0)
        frame_images = (rearrange(frame_images / 255.0 * 2 - 1, "B T H W C -> B C T H W").contiguous().to(torch.bfloat16))
        
        vae = pipe.vae 
        noise_images = vae.encode(frame_images).latent_dist.sample() * vae.config.scaling_factor
        noise_images = noise_images.permute(0, 2, 1, 3, 4).contiguous()
        
        noise = torch.randn_like(noise_images)
        pipe.scheduler.set_timesteps(num_inference_steps)
        all_steps = list(pipe.scheduler.timesteps)
        
        for noise_step in [699]:
            noisy_video_latents = pipe.scheduler.add_noise(noise_images, noise, torch.tensor(noise_step))       
            
            if flag:
                start_step = all_steps.index(noise_step)
                timesteps = all_steps[start_step:]
                
                prefix = noisy_video_latents
            else:
                start_step = 0
                timesteps = all_steps
                prefix = None 
            # 4. Generate the video frames based on the prompt.
            # `num_frames` is the Number of frames to generate.
            if generate_type == "i2v":
                video_generate = pipe(
                    height=height,
                    width=width,
                    prompt=prompt,
                    image=image,
                    # The path of the image, the resolution of video will be the same as the image for CogVideoX1.5-5B-I2V, otherwise it will be 720 * 480
                    num_videos_per_prompt=num_videos_per_prompt,  # Number of videos to generate per prompt
                    num_inference_steps=num_inference_steps - start_step,  # Number of inference steps
                    num_frames=num_frames,  # Number of frames to generate
                    use_dynamic_cfg=True,  # This id used for DPM scheduler, for DDIM scheduler, it should be False
                    guidance_scale=guidance_scale,
                    generator=torch.Generator().manual_seed(seed),  # Set the seed for reproducibility
                ).frames[0]
            elif generate_type == "t2v":
                video_generate = pipe(
                    prompt=prompt,
                    num_videos_per_prompt=num_videos_per_prompt,
                    num_inference_steps=len(timesteps),
                    timesteps=timesteps,
                    num_frames=num_frames,
                    use_dynamic_cfg=True,
                    guidance_scale=guidance_scale,
                    generator=torch.Generator().manual_seed(seed),
                    latents=prefix
                ).frames[0]
            else:
                video_generate = pipe(
                    prompt=prompt,
                    video=video,  # The path of the video to be used as the background of the video
                    num_videos_per_prompt=num_videos_per_prompt,
                    num_inference_steps=num_inference_steps,
                    num_frames=num_frames,
                    use_dynamic_cfg=True,
                    guidance_scale=guidance_scale,
                    generator=torch.Generator().manual_seed(seed),  # Set the seed for reproducibility
                ).frames[0]
            export_to_video(video_generate, os.path.join(output_path, skill_name, f'{index:04d}_{noise_step}_{start_step}.mp4'), fps=fps)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate a video from a text prompt using CogVideoX")
    parser.add_argument("--prompt", type=str, default=None, help="The description of the video to be generated")
    parser.add_argument(
        "--image_or_video_path",
        type=str,
        default=None,
        help="The path of the image to be used as the background of the video",
    )
    parser.add_argument(
        "--model_path", type=str, default="THUDM/CogVideoX-5B", help="Path of the pre-trained model use"
    )
    parser.add_argument("--lora_path", type=str, default=None, help="The path of the LoRA weights to be used")
    parser.add_argument("--lora_rank", type=int, default=128, help="The rank of the LoRA weights")
    parser.add_argument("--output_path", type=str, default="./output.mp4", help="The path save generated video")
    parser.add_argument("--guidance_scale", type=float, default=6.0, help="The scale for classifier-free guidance")
    parser.add_argument("--num_inference_steps", type=int, default=50, help="Inference steps")
    parser.add_argument("--num_frames", type=int, default=49, help="Number of steps for the inference process")
    parser.add_argument("--width", type=int, default=1360, help="Number of steps for the inference process")
    parser.add_argument("--height", type=int, default=768, help="Number of steps for the inference process")
    parser.add_argument("--fps", type=int, default=8, help="Number of steps for the inference process")
    parser.add_argument("--num_videos_per_prompt", type=int, default=1, help="Number of videos to generate per prompt")
    parser.add_argument("--generate_type", type=str, default="t2v", help="The type of video generation")
    parser.add_argument("--dtype", type=str, default="bfloat16", help="The data type for computation")
    parser.add_argument("--seed", type=int, default=42, help="The seed for reproducibility")
    parser.add_argument("--latent_path", type=str, default="")

    args = parser.parse_args()
    dtype = torch.float16 if args.dtype == "float16" else torch.bfloat16
    generate_video(
        prompt=args.prompt,
        model_path=args.model_path,
        lora_path=args.lora_path,
        lora_rank=args.lora_rank,
        output_path=args.output_path,
        num_frames=args.num_frames,
        width=args.width,
        height=args.height,
        image_or_video_path=args.image_or_video_path,
        num_inference_steps=args.num_inference_steps,
        guidance_scale=args.guidance_scale,
        num_videos_per_prompt=args.num_videos_per_prompt,
        dtype=dtype,
        generate_type=args.generate_type,
        seed=args.seed,
        fps=args.fps,
        latent_path=args.latent_path
    )
