import os
import torch
import numpy as np
import math
import argparse
from decord import VideoReader
from diffusers import AutoencoderKLCogVideoX
from safetensors.torch import save_file
import tqdm
import random

def encode_video(video, vae):
    video = video[None].permute(0, 2, 1, 3, 4).contiguous()
    video = video.to(vae.device, dtype=vae.dtype)
    latent_dist = vae.encode(video).latent_dist
    latent = latent_dist.sample() * vae.config.scaling_factor
    return latent

def add_dashed_rays_to_video(video_tensor, num_perp_samples=50, density_decay=0.075):
    T, C, H, W = video_tensor.shape
    max_length = int((H**2 + W**2) ** 0.5) + 10
    center = torch.tensor([W / 2, H / 2])
    theta = torch.rand(1).item() * 2 * math.pi
    direction = torch.tensor([math.cos(theta), math.sin(theta)])
    direction = direction / direction.norm()
    d_perp = torch.tensor([-direction[1], direction[0]])
    half_len = max(H, W) // 2
    positions = torch.linspace(-half_len, half_len, num_perp_samples)
    perp_coords = center[None, :] + positions[:, None] * d_perp[None, :]
    x0, y0 = perp_coords[:, 0], perp_coords[:, 1]
    steps = []
    dist = 0
    while dist < max_length:
        steps.append(dist)
        dist += 1.0 + density_decay * dist
    steps = torch.tensor(steps)
    S = len(steps)
    dxdy = direction[None, :] * steps[:, None]
    all_xy = perp_coords[:, None, :] + dxdy[None, :, :]
    all_xy = all_xy.reshape(-1, 2)
    all_x = all_xy[:, 0].round().long()
    all_y = all_xy[:, 1].round().long()
    valid = (0 <= all_x) & (all_x < W) & (0 <= all_y) & (all_y < H)
    all_x = all_x[valid]
    all_y = all_y[valid]
    x0r = x0.round().long().clamp(0, W - 1)
    y0r = y0.round().long().clamp(0, H - 1)
    frame0 = video_tensor[0]
    base_colors = frame0[:, y0r, x0r]
    base_colors = base_colors.repeat_interleave(S, dim=1)[:, valid]
    video_out = video_tensor.clone()
    offsets = [(0, 0), (0, 1), (1, 0), (1, 1)]
    for dxo, dyo in offsets:
        ox = all_x + dxo
        oy = all_y + dyo
        inside = (0 <= ox) & (ox < W) & (0 <= oy) & (oy < H)
        ox = ox[inside]
        oy = oy[inside]
        colors = base_colors[:, inside]
        for c in range(C):
            video_out[1:, c, oy, ox] = colors[c][None, :].expand(T - 1, -1)
    return video_out

def main(args):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    vae = AutoencoderKLCogVideoX.from_pretrained(args.pretrained_model_path, subfolder="vae")
    vae.requires_grad_(False)
    vae = vae.to(device, dtype=torch.float16)

    masked_video_path = os.path.join(args.video_root, "masked_videos")
    source_video_path = os.path.join(args.video_root, "videos")
    joint_latent_path = os.path.join(args.video_root, "joint_latents")
    os.makedirs(joint_latent_path, exist_ok=True)

    all_video_names = sorted(os.listdir(source_video_path))
    video_names = all_video_names[args.start_idx : args.end_idx]

    for video_name in tqdm.tqdm(video_names, desc=f"GPU {args.gpu_id}"):
        masked_video_file = os.path.join(masked_video_path, video_name)
        source_video_file = os.path.join(source_video_path, video_name)
        output_file = os.path.join(joint_latent_path, video_name.replace('.mp4', '.safetensors'))

        if not os.path.exists(masked_video_file):
            print(f"Skipping {video_name}, masked video not found.")
            continue
        if os.path.exists(output_file):
            continue

        try:
            vr = VideoReader(source_video_file)
            video = torch.from_numpy(vr.get_batch(np.arange(49)).asnumpy()).permute(0, 3, 1, 2).contiguous()
            video = (video / 255.0) * 2 - 1
            source_latent = encode_video(video, vae)

            vr = VideoReader(masked_video_file)
            video = torch.from_numpy(vr.get_batch(np.arange(49)).asnumpy()).permute(0, 3, 1, 2).contiguous()
            video = (video / 255.0) * 2 - 1
            video = add_dashed_rays_to_video(video)
            masked_latent = encode_video(video, vae)

            source_latent = source_latent.to("cpu")
            masked_latent = masked_latent.to("cpu")
            cated_latent = torch.cat([source_latent, masked_latent], dim=2)
            save_file({'joint_latents': cated_latent}, output_file)

        except Exception as e:
            print(f"[GPU {args.gpu_id}] Error processing {video_name}: {e}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--video_root", type=str, required=True)
    parser.add_argument("--pretrained_model_path", type=str, required=True)
    parser.add_argument("--start_idx", type=int, required=True)
    parser.add_argument("--end_idx", type=int, required=True)
    parser.add_argument("--gpu_id", type=int, required=True)
    args = parser.parse_args()

    os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_id)
    main(args)