import os
import argparse
from omegaconf import OmegaConf

import torch
from src import ViewPipeline, ViewTransformer3DModel
from utils.infer_utils import *
from utils import save, save_cube_video
from decord import VideoReader, cpu
from diffusers.utils import export_to_video, load_image, load_video
import torch.nn.functional as F

def generate_video(model_path, transformer_id, prompt, video_path, direction, rotate, output_dir, H, W, L, fps):
    transformer = ViewTransformer3DModel.from_pretrained(
        transformer_id, torch_dtype=torch.bfloat16, 
    )
    pipe = ViewPipeline.from_pretrained(model_path, transformer=transformer, torch_dtype=torch.bfloat16).to("cuda")
 
    if video_path is not None:
        vid = VideoReader(video_path, ctx=cpu(0), width=W, height=H)
        vid = vid.get_batch(range(len(vid))).asnumpy()
        vid = torch.from_numpy(vid).permute(0,3,1,2).to(device="cuda", dtype=torch.bfloat16)
        save(vid[:L,...].permute(1,0,2,3)/255.0, os.path.join(output_dir,"input.mp4"), fps=fps,rescale=False)
        print("input_shape:",vid[:L,...].shape)
    else:
        vid = None
        direction = "null"
    
    video = pipe(
        prompt=prompt,
        video=vid,
        selected_cube=direction,
        height = H,
        width = W,
        num_frames = L,
        rotate = rotate,
    ).frames[0]

    export_to_video(video,os.path.join(output_dir,"viewpoint.mp4"),fps=fps)
 
    cubemap = convert_to_cubemap2(video)
 
    save_cube_video(cubemap.permute(1,0,2,3), os.path.join(output_dir,"cubemap") ,fps=fps,rescale=False)
   
    equirec = cube2equi(cubemap)
    save(equirec.permute(1,0,2,3), os.path.join(output_dir,"equi.mp4"), fps=fps,rescale=False)
 
   

def main():
    parser = argparse.ArgumentParser(description="Generate video using Wan2.1 and LoRA weights")
    parser.add_argument("-H", type=int, default=256)
    parser.add_argument("-W", type=int, default=256)
    parser.add_argument("-L", type=int, default=49)
    parser.add_argument("--fps", type=int, default=16, help="Frames per second for the output video")
    
    parser.add_argument("--config", type=str, default="./configs/VanGogh.yaml")

    args = parser.parse_args()
    config = OmegaConf.load(args.config)
  
    os.makedirs(config.output_dir,exist_ok=True)
    torch.manual_seed(config.seed)
    print("generating:",config.prompt)
    generate_video(
        config.model_path, 
        config.transformer_id, 
        config.prompt, 
        config.get("video", None),
        config.get("direction", "null"),
        config.get("rotate",False),
        config.output_dir, 
        args.H,
        args.W,
        args.L,
        args.fps
    )


if __name__ == "__main__":
    main()
