# Copyright(C) 2025 Advanced Micro Devices, Inc. All rights reserved.  
import argparse
from copy import deepcopy
import os
import pickle
from omegaconf import OmegaConf
import torch
import torchvision
from tqdm import tqdm
from ode_solver.ddim_solver import DDIMSolver
from scheduler.t2v_turbo_scheduler import T2VTurboScheduler
from utils.utils import instantiate_from_config
from utils.common_utils import (
    extract_into_tensor,
    get_transform,
    get_predicted_noise,
    get_predicted_original_sample,
    load_model_checkpoint,
    read_video_to_tensor,
    compute_temp_loss,
)
from reward_fn import get_reward_fn


SAMPLE_SIZE = (320, 512)


def reverse_ddim_loop(latents, unet, context, solver, num_inference_steps, device):
    intermediate_latents = []
    for i in range(num_inference_steps):
        index = torch.full((1,), i, device=device, dtype=torch.long)
        ts = solver.ddim_timesteps[index].long()
        pred_noise = unet(latents, ts, **context)

        latents = solver.ddim_reverse_step(latents, pred_noise, ts)
        intermediate_latents.append(latents)

    return intermediate_latents


def get_temp_attn_prob(unet, latent, ts, context):
    attention_prob = {}
    model_output = unet(latent, ts, **context)
    for name, module in unet.named_modules():
        if (
            name.startswith("output_blocks.3.2")
            or name.startswith("output_blocks.4.2")
            or name.startswith("output_blocks.5.2")
            or name.startswith("output_blocks.6.2")
            or name.startswith("output_blocks.7.2")
            or name.startswith("output_blocks.8.2")
            or name.startswith("output_blocks.9.2")
            or name.startswith("output_blocks.10.2")
            or name.startswith("output_blocks.11.2")
        ) and name.endswith("blocks.0.attn1"):
            attention_prob[name] = module.attention_probs
    return model_output, attention_prob


def get_motion_prior_score(
    unet,
    latents,
    ts,
    example_latent,
    original_context,
    inference_context,
    temp_loss_scale,
):
    attention_prob = {}
    attention_prob_example = {}
    with torch.no_grad():
        _, attention_prob_example = get_temp_attn_prob(
            unet, example_latent, ts, original_context
        )
    with torch.set_grad_enabled(True):
        latents.requires_grad_(True)
        cond_teacher_output, attention_prob = get_temp_attn_prob(
            unet, latents, ts, inference_context
        )
        loss = temp_loss_scale * compute_temp_loss(
            attention_prob, attention_prob_example
        )
        score = torch.autograd.grad(loss, latents)[0].detach()

    return score, cond_teacher_output


@torch.no_grad()
def main(
    unet,
    vae,
    text_encoder,
    scheduler,
    pretrained_t2v,
    solver,
    reward_fn,
    args,
    device,
    dtype,
):
    pixel_transforms = get_transform(sample_size=SAMPLE_SIZE)

    video = read_video_to_tensor(
        args.ref_video_path, sample_fps=args.fps, sample_frames=args.n_frames
    )
    video = pixel_transforms(video)

    latents = vae.encode(video.to(device, dtype)).sample()
    latents = latents.unsqueeze(0)
    latents = latents.permute(0, 2, 1, 3, 4)
    vae_scale_factor = 0.18215
    latents = latents * vae_scale_factor

    inference_prompt_embeds = text_encoder([args.inference_prompt]).to(device, dtype)
    original_prompt_embeds = text_encoder([args.ref_prompt]).to(device, dtype)
    uncond_prompt_embeds = text_encoder([""]).to(device, dtype)

    inference_context = {
        "context": torch.cat([inference_prompt_embeds], 1),
        "fps": args.fps,
    }
    original_context = {
        "context": torch.cat([original_prompt_embeds], 1),
        "fps": args.fps,
    }
    uncond_context = {"context": torch.cat([uncond_prompt_embeds], 1), "fps": args.fps}

    alpha_schedule = torch.sqrt(scheduler.alphas_cumprod).to(device)
    sigma_schedule = torch.sqrt(1 - scheduler.alphas_cumprod).to(device)

    # Reversed DDIM
    unet.to(dtype)
    num_inference_steps = len(solver.ddim_timesteps)
    original_latents = latents.to(dtype)
    if not os.path.exists(args.intermediate_latents_path):
        print("Starting reversed DDIM...")
        latents = deepcopy(original_latents)
        intermediate_latents = reverse_ddim_loop(
            latents, unet, original_context, solver, num_inference_steps, device
        )
        print("Finished reversed DDIM.")
        with open("resources/intermediate_latents.pkl", "wb") as f:
            pickle.dump(intermediate_latents, f)
    else:
        with open(args.intermediate_latents_path, "rb") as f:
            intermediate_latents = pickle.load(f)
    # DDIM
    # latents = torch.randn_like(intermediate_latents[-1])
    noise = torch.randn_like(intermediate_latents[-1])
    latents = scheduler.add_noise(original_latents, noise, solver.ddim_timesteps[-1])
    if not os.path.exists(args.motion_intermediate_path):
        motion_intermediate = []
    else:
        with open(args.motion_intermediate_path, "rb") as f:
            motion_intermediate = pickle.load(f)

    sampled_latents = []
    for i in tqdm(range(num_inference_steps - 1, -1, -1)):
        index = torch.full((1,), i, device=device, dtype=torch.long)
        ts = solver.ddim_timesteps[index]
        if i > num_inference_steps - args.percentage * num_inference_steps:
            idx = num_inference_steps - 1 - i
            if len(motion_intermediate) > idx:
                results = motion_intermediate[idx]
                cond_pred_x_0 = results["cond_pred_x_0"].to(device, dtype)
                cond_pred_noise = results["cond_pred_noise"].to(device, dtype)
                uncond_pred_x_0 = results["uncond_pred_x_0"].to(device, dtype)
                uncond_pred_noise = results["uncond_pred_noise"].to(device, dtype)
                score = results["score"].to(device, dtype)
            else:
                score, cond_pred_noise = get_motion_prior_score(
                    unet,
                    latents,
                    ts,
                    intermediate_latents[i],
                    original_context,
                    inference_context,
                    args.temp_loss_scale,
                )
                with torch.set_grad_enabled(False):
                    cond_pred_x_0 = get_predicted_original_sample(
                        cond_pred_noise,
                        ts,
                        latents,
                        "epsilon",
                        alpha_schedule,
                        sigma_schedule,
                    )
                    uncond_pred_noise = unet(latents, ts, **uncond_context)
                    uncond_pred_x_0 = get_predicted_original_sample(
                        uncond_pred_noise,
                        ts,
                        latents,
                        "epsilon",
                        alpha_schedule,
                        sigma_schedule,
                    )
                motion_intermediate.append(
                    {
                        "cond_pred_x_0": cond_pred_x_0.detach().cpu(),
                        "uncond_pred_x_0": uncond_pred_x_0.detach().cpu(),
                        "cond_pred_noise": cond_pred_noise.detach().cpu(),
                        "uncond_pred_noise": uncond_pred_noise.detach().cpu(),
                        "score": score.detach().cpu(),
                    }
                )

            if args.use_rm_guide and args.reward_scale > 0.0:
                with torch.set_grad_enabled(False):
                    cond_pred_noise = unet(latents, ts, **inference_context)
                    uncond_pred_noise = unet(latents, ts, **uncond_context)
                    pred_noise = cond_pred_noise + args.guidance_scale * (
                        cond_pred_noise - uncond_pred_noise
                    )
                    cond_pred_x_0 = get_predicted_original_sample(
                        cond_pred_noise,
                        ts,
                        latents,
                        "epsilon",
                        alpha_schedule,
                        sigma_schedule,
                    )
                    uncond_pred_x_0 = get_predicted_original_sample(
                        uncond_pred_noise,
                        ts,
                        latents,
                        "epsilon",
                        alpha_schedule,
                        sigma_schedule,
                    )

                with torch.set_grad_enabled(True):
                    latents.requires_grad_(True)
                    pred_x_0 = get_predicted_original_sample(
                        pred_noise,
                        ts,
                        latents,
                        "epsilon",
                        alpha_schedule,
                        sigma_schedule,
                    )

                    idx = torch.randint(0, args.n_frames, (args.reward_batch_size,))

                    selected_latents = (
                        pred_x_0[:, :, idx].to(vae.dtype) / vae_scale_factor
                    )
                    selected_latents = selected_latents.permute(0, 2, 1, 3, 4)
                    selected_latents = selected_latents.reshape(
                        args.reward_batch_size, *selected_latents.shape[2:]
                    )
                    decoded_imgs = vae.decode(selected_latents)
                    decoded_imgs = (decoded_imgs / 2 + 0.5).clamp(0, 1)
                    expert_rewards = reward_fn(decoded_imgs, args.inference_prompt)
                    reward_loss = -expert_rewards.mean() * args.reward_scale
                    score += torch.autograd.grad(reward_loss, latents)[0].detach()
            else:
                score += torch.zeros_like(latents)
        else:
            with torch.set_grad_enabled(False):
                cond_pred_noise = unet(latents, ts, **inference_context)
                uncond_pred_noise = unet(latents, ts, **uncond_context)
                pred_noise = cond_pred_noise + args.guidance_scale * (
                    cond_pred_noise - uncond_pred_noise
                )
                cond_pred_x_0 = get_predicted_original_sample(
                    cond_pred_noise,
                    ts,
                    latents,
                    "epsilon",
                    alpha_schedule,
                    sigma_schedule,
                )
                uncond_pred_x_0 = get_predicted_original_sample(
                    uncond_pred_noise,
                    ts,
                    latents,
                    "epsilon",
                    alpha_schedule,
                    sigma_schedule,
                )
            score = torch.zeros_like(latents)

        with torch.set_grad_enabled(False):
            pred_x_0 = cond_pred_x_0 + args.guidance_scale * (
                cond_pred_x_0 - uncond_pred_x_0
            )
            pred_noise = cond_pred_noise + args.guidance_scale * (
                cond_pred_noise - uncond_pred_noise
            )
            alphas = extract_into_tensor(alpha_schedule, ts, score.shape)
            pred_noise -= (1 - alphas) ** (0.5) * score
            latents = solver.ddim_step(pred_x_0, pred_noise, index).to(device, dtype)

        sampled_latents.append(pred_x_0)

    with open(args.motion_intermediate_path, "wb") as f:
        pickle.dump(motion_intermediate, f)

    videos = pretrained_t2v.decode_first_stage_2DAE(latents)
    videos = (videos.clamp(-1.0, 1.0) + 1.0) / 2.0
    videos = (videos * 255).to(torch.uint8).permute(0, 2, 1, 3, 4).cpu().numpy()

    torchvision.io.write_video(
        args.output_video_path,
        torch.from_numpy(videos)[0].permute(0, 2, 3, 1),
        fps=args.fps,
        video_codec="h264",
        options={"crf": "10"},
    )

    if args.save_intermediate_videos:
        os.makedirs("motion_intermediates", exist_ok=True)
        for i, latents in enumerate(sampled_latents):
            videos = pretrained_t2v.decode_first_stage_2DAE(latents)
            videos = (videos.clamp(-1.0, 1.0) + 1.0) / 2.0
            videos = (videos * 255).to(torch.uint8).permute(0, 2, 1, 3, 4).cpu().numpy()

            torchvision.io.write_video(
                args.output_video_path.replace(".mp4", f"_{i}.mp4").replace(
                    "samples/", "motion_intermediates/"
                ),
                torch.from_numpy(videos)[0].permute(0, 2, 3, 1),
                fps=args.fps,
                video_codec="h264",
                options={"crf": "10"},
            )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Motion Clone Sampling")
    # Add model name as parameter
    parser.add_argument(
        "--sampling_steps", type=int, default=200, help="Number of sampling steps"
    )
    parser.add_argument(
        "--guidance_scale", type=float, default=7.5, help="Guidance scale"
    )
    parser.add_argument("--n_frames", type=int, default=16, help="Number of frames")
    parser.add_argument("--fps", type=int, default=10, help="Frames per second")
    parser.add_argument(
        "--percentage",
        type=float,
        default=0.6,
        help="Percentage of steps to apply guidance",
    )
    parser.add_argument(
        "--temp_loss_scale",
        type=float,
        default=100,
        help="Temperature loss scale",
    )
    parser.add_argument(
        "--ref_video_path",
        type=str,
        default="resources/sample_fox.mp4",
        help="Path to reference video",
    )
    parser.add_argument(
        "--inference_prompt",
        type=str,
        default="A cat turning its heads in the bedroom.",
        help="Inference prompt for the video",
    )
    parser.add_argument(
        "--ref_prompt",
        type=str,
        default="A fox turning its heads in the woods.",
        help="Inference prompt for the video",
    )
    parser.add_argument(
        "--output_video_path",
        type=str,
        default="samples/motion_prior_sample_video.mp4",
        help="Path to output video",
    )
    parser.add_argument(
        "--sample_config_path",
        type=str,
        default="configs/inference_t2v_512_v2.0_fp16.yaml",
        help="Path to output video",
    )
    parser.add_argument(
        "--vc2_model_path",
        type=str,
        default="model_cache/VideoCrafter2_model.ckpt",
        help="Path to output video",
    )
    parser.add_argument(
        "--use_rm_guide",
        action="store_true",
        default=False,
        help="Use reward model to guide the generation",
    )
    parser.add_argument(
        "--reward_batch_size",
        type=int,
        default=12,
        help="Batch size (per device) for optimizing the text-image RM.",
    )
    parser.add_argument(
        "--reward_scale",
        type=float,
        default=1.0,
        help="The scale of the reward loss",
    )
    parser.add_argument(
        "--intermediate_latents_path",
        type=str,
        default="resources/intermediate_latents.pkl",
        help="Path to intermediate latents",
    )
    parser.add_argument(
        "--motion_intermediate_path",
        type=str,
        default="resources/motion_intermediate.pkl",
        help="Path to intermediate motion guidance results",
    )
    parser.add_argument(
        "--save_intermediate_videos",
        action="store_true",
        default=False,
        help="Save intermediate videos",
    )

    args = parser.parse_args()

    device = torch.device("cuda")
    dtype = torch.float16
    config = OmegaConf.load(args.sample_config_path)
    model_config = config.pop("model", OmegaConf.create())
    pretrained_t2v = instantiate_from_config(model_config)
    pretrained_t2v = load_model_checkpoint(pretrained_t2v, args.vc2_model_path)
    unet = pretrained_t2v.model.diffusion_model.to(device, dtype)
    vae = pretrained_t2v.first_stage_model.to(device, dtype)
    text_encoder = pretrained_t2v.cond_stage_model.to(device, dtype)

    unet.requires_grad_(False).eval()
    vae.requires_grad_(False).eval()
    text_encoder.requires_grad_(False).eval()

    scheduler = T2VTurboScheduler(
        linear_start=model_config["params"]["linear_start"],
        linear_end=model_config["params"]["linear_end"],
    )

    solver = DDIMSolver(
        scheduler.alphas_cumprod.numpy(),
        ddim_timesteps=args.sampling_steps,
        use_scale=False,
    ).to(device, dtype)

    if args.use_rm_guide:
        reward_fn = get_reward_fn("hpsv2", precision="fp16")
    else:
        reward_fn = None

    main(
        unet,
        vae,
        text_encoder,
        scheduler,
        pretrained_t2v,
        solver,
        reward_fn,
        args,
        device,
        dtype,
    )
