import copy
from pipeline import SelfForcingTrainingPipeline
import torch.nn.functional as F
from typing import Tuple, Optional
import torch
from model.base import SelfForcingModel
import torch.distributed as dist

class CGAN_DMD(SelfForcingModel):
    def __init__(self, args, device):
        """
        Initialize the GAN module.
        This class is self-contained and compute generator and fake score losses
        in the forward pass.
        """
        super().__init__(args, device)
        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
        self.same_step_across_blocks = getattr(args, "same_step_across_blocks", True)
        self.concat_time_embeddings = getattr(args, "concat_time_embeddings", False)
        self.num_class = args.num_class
        self.gan_loss_weight = args.gan_loss_weight
        self.relativistic_discriminator = getattr(args, "relativistic_discriminator", False)
        self.num_training_frames = getattr(args, "num_training_frames", 21)
        # self.gamma = getattr(args, "gamma", 10)

        if self.num_frame_per_block > 1:
            self.generator.model.num_frame_per_block = self.num_frame_per_block

        # self.fake_score.adding_cls_branch(
        #     atten_dim=1536, num_class=args.num_class, time_embed_dim=1536 if self.concat_time_embeddings else 0)
        # self.fake_score.model.requires_grad_(True)

        # add the classifier on student
        self.fake_score.adding_cls_branch(
            atten_dim=1536, num_class=args.num_class, time_embed_dim=1536 if self.concat_time_embeddings else 0)
            # atten_dim=5120, num_class=args.num_class, time_embed_dim=5120 if self.concat_time_embeddings else 0, ffn_dim=13824, num_heads=40)
        self.fake_score.model.requires_grad_(True)

        self.independent_first_frame = getattr(args, "independent_first_frame", False)
        if self.independent_first_frame:
            self.generator.model.independent_first_frame = True
        if args.gradient_checkpointing:
            self.generator.enable_gradient_checkpointing()
            self.fake_score.enable_gradient_checkpointing()

        # this will be init later with fsdp-wrapped modules
        self.inference_pipeline: SelfForcingTrainingPipeline = None

        # Step 2: Initialize all dmd hyperparameters
        self.num_train_timestep = args.num_train_timestep
        self.min_step = int(0.02 * self.num_train_timestep)
        self.max_step = int(0.98 * self.num_train_timestep)
        if hasattr(args, "real_guidance_scale"):
            self.real_guidance_scale = args.real_guidance_scale
            self.fake_guidance_scale = args.fake_guidance_scale
        else:
            self.real_guidance_scale = args.guidance_scale
            self.fake_guidance_scale = 0.0
        self.timestep_shift = getattr(args, "timestep_shift", 1.0)
        self.critic_timestep_shift = getattr(args, "critic_timestep_shift", self.timestep_shift)
        self.ts_schedule = getattr(args, "ts_schedule", True)
        self.ts_schedule_max = getattr(args, "ts_schedule_max", False)
        self.min_score_timestep = getattr(args, "min_score_timestep", 0)

        self.gan_g_weight = getattr(args, "gan_g_weight", 1e-2)
        self.gan_d_weight = getattr(args, "gan_d_weight", 1e-2)
        self.r1_weight = getattr(args, "r1_weight", 0.00)
        self.r2_weight = getattr(args, "r2_weight", 0.00)
        self.r1_sigma = getattr(args, "r1_sigma", 0.01)
        self.r2_sigma = getattr(args, "r2_sigma", 0.01)

        if getattr(self.scheduler, "alphas_cumprod", None) is not None:
            self.scheduler.alphas_cumprod = self.scheduler.alphas_cumprod.to(device)
        else:
            self.scheduler.alphas_cumprod = None
    
    @staticmethod
    def ZeroCenteredGradientPenalty(Samples, Critics):
        Gradient, = torch.autograd.grad(outputs=Critics.sum(), inputs=Samples, create_graph=True)
        return Gradient.square().sum()

    def generate_and_sync_list(self, num_blocks, num_denoising_steps, device):
        rank = dist.get_rank() if dist.is_initialized() else 0

        if rank == 0:
            # Generate random indices
            indices = torch.randint(
                low=0,
                high=num_denoising_steps,
                size=(num_blocks,),
                device=device
            )
            # if self.last_step_only:
            #     indices = torch.ones_like(indices) * (num_denoising_steps - 1)
        else:
            indices = torch.empty(num_blocks, dtype=torch.long, device=device)

        dist.broadcast(indices, src=0)  # Broadcast the random indices to all ranks
        return indices.tolist()


    def _run_cls_pred_branch(self,
                             noisy_image_or_video: torch.Tensor,
                             conditional_dict: dict,
                             timestep: torch.Tensor) -> torch.Tensor:
        """
            Run the classifier prediction branch on the generated image or video.
            Input:
                - image_or_video: a tensor with shape [B, F, C, H, W].
            Output:
                - cls_pred: a tensor with shape [B, 1, 1, 1, 1] representing the feature map for classification.
        """
        _, _, noisy_logit = self.fake_score(
            noisy_image_or_video=noisy_image_or_video,
            conditional_dict=conditional_dict,
            timestep=timestep,
            classify_mode=True,
            concat_time_embeddings=self.concat_time_embeddings
        )

        return noisy_logit

    def _compute_kl_grad(
        self, noisy_image_or_video: torch.Tensor,
        estimated_clean_image_or_video: torch.Tensor,
        timestep: torch.Tensor,
        conditional_dict: dict, unconditional_dict: dict,
        normalization: bool = True
    ) -> Tuple[torch.Tensor, dict]:
        """
        Compute the KL grad (eq 7 in https://arxiv.org/abs/2311.18828).
        Input:
            - noisy_image_or_video: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.
            - estimated_clean_image_or_video: a tensor with shape [B, F, C, H, W] representing the estimated clean image or video.
            - timestep: a tensor with shape [B, F] containing the randomly generated timestep.
            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
            - normalization: a boolean indicating whether to normalize the gradient.
        Output:
            - kl_grad: a tensor representing the KL grad.
            - kl_log_dict: a dictionary containing the intermediate tensors for logging.
        """
        # Step 1: Compute the fake score
        _, pred_fake_image_cond = self.fake_score(
            noisy_image_or_video=noisy_image_or_video,
            conditional_dict=conditional_dict,
            timestep=timestep
        )

        if self.fake_guidance_scale != 0.0:
            _, pred_fake_image_uncond = self.fake_score(
                noisy_image_or_video=noisy_image_or_video,
                conditional_dict=unconditional_dict,
                timestep=timestep
            )
            pred_fake_image = pred_fake_image_cond + (
                pred_fake_image_cond - pred_fake_image_uncond
            ) * self.fake_guidance_scale
        else:
            pred_fake_image = pred_fake_image_cond

        # Step 2: Compute the real score
        # We compute the conditional and unconditional prediction
        # and add them together to achieve cfg (https://arxiv.org/abs/2207.12598)
        _, pred_real_image_cond = self.real_score(
            noisy_image_or_video=noisy_image_or_video,
            conditional_dict=conditional_dict,
            timestep=timestep
        )

        _, pred_real_image_uncond = self.real_score(
            noisy_image_or_video=noisy_image_or_video,
            conditional_dict=unconditional_dict,
            timestep=timestep
        )

        pred_real_image = pred_real_image_cond + (
            pred_real_image_cond - pred_real_image_uncond
        ) * self.real_guidance_scale

        # Step 3: Compute the DMD gradient (DMD paper eq. 7).
        grad = (pred_fake_image - pred_real_image)

        if normalization:
            # Step 4: Gradient normalization (DMD paper eq. 8).
            p_real = (estimated_clean_image_or_video - pred_real_image)
            normalizer = torch.abs(p_real).mean(dim=[1, 2, 3, 4], keepdim=True)
            grad = grad / normalizer
        grad = torch.nan_to_num(grad)

        return grad, {
            "dmdtrain_gradient_norm": torch.mean(torch.abs(grad)).detach(),
            "timestep": timestep.detach()
        }

    def compute_distribution_matching_loss(
        self,
        image_or_video: torch.Tensor,
        conditional_dict: dict,
        unconditional_dict: dict,
        gradient_mask: Optional[torch.Tensor] = None,
        denoised_timestep_from: int = 0,
        denoised_timestep_to: int = 0
    ) -> Tuple[torch.Tensor, dict]:
        """
        Compute the DMD loss (eq 7 in https://arxiv.org/abs/2311.18828).
        Input:
            - image_or_video: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.
            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
            - gradient_mask: a boolean tensor with the same shape as image_or_video indicating which pixels to compute loss .
        Output:
            - dmd_loss: a scalar tensor representing the DMD loss.
            - dmd_log_dict: a dictionary containing the intermediate tensors for logging.
        """
        original_latent = image_or_video

        batch_size, num_frame = image_or_video.shape[:2]

        with torch.no_grad():
            # Step 1: Randomly sample timestep based on the given schedule and corresponding noise
            min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep
            max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep
            timestep = self._get_timestep(
                min_timestep,
                max_timestep,
                batch_size,
                num_frame,
                self.num_frame_per_block,
                uniform_timestep=True
            )

            if self.timestep_shift > 1:
                timestep = self.timestep_shift * \
                    (timestep / 1000) / \
                    (1 + (self.timestep_shift - 1) * (timestep / 1000)) * 1000
            timestep = timestep.clamp(self.min_step, self.max_step)

            noise = torch.randn_like(image_or_video)
            noisy_latent = self.scheduler.add_noise(
                image_or_video.flatten(0, 1),
                noise.flatten(0, 1),
                timestep.flatten(0, 1)
            ).detach().unflatten(0, (batch_size, num_frame))

            # Step 2: Compute the KL grad
            grad, dmd_log_dict = self._compute_kl_grad(
                noisy_image_or_video=noisy_latent,
                estimated_clean_image_or_video=original_latent,
                timestep=timestep,
                conditional_dict=conditional_dict,
                unconditional_dict=unconditional_dict
            )

        if gradient_mask is not None:
            dmd_loss = 0.5 * F.mse_loss(original_latent.double(
            )[gradient_mask], (original_latent.double() - grad.double()).detach()[gradient_mask], reduction="mean")
        else:
            dmd_loss = 0.5 * F.mse_loss(original_latent.double(
            ), (original_latent.double() - grad.double()).detach(), reduction="mean")
        return dmd_loss, dmd_log_dict


    def generator_loss(
        self,
        image_or_video_shape,
        conditional_dict: dict,
        unconditional_dict: dict,
        clean_latent: torch.Tensor,
        initial_latent: torch.Tensor = None
    ) -> Tuple[torch.Tensor, dict]:
        """
        Generate image/videos from noise and compute the DMD loss.
        The noisy input to the generator is backward simulated.
        This removes the need of any datasets during distillation.
        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
        Input:
            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
        Output:
            - loss: a scalar tensor representing the generator loss.
            - generator_log_dict: a dictionary containing the intermediate tensors for logging.
        """
        # Step 1: Unroll generator to obtain fake videos, and obtain the prediction
        # TODO: run_generator should output preditions too, 1 prediction per chunk
        denoise_steps = self.generate_and_sync_list(1, self.num_class, device=conditional_dict["prompt_embeds"].device)[0]
        pred_image, gradient_mask, denoised_timestep_from, denoised_timestep_to = self._run_generator(
            image_or_video_shape=image_or_video_shape,
            conditional_dict=conditional_dict,
            initial_latent=initial_latent,
            denoise_steps=denoise_steps
        )

        if denoise_steps != len(self.inference_pipeline.denoising_step_list)-1:
            # not last step
            gan_loss, log_dict = self.gan_loss(pred_image, image_or_video_shape, conditional_dict, unconditional_dict, None, denoise_steps)
        else:
            gan_loss = 0
            log_dict = {}

        # Step 2: Compute the DMD loss
        dmd_loss, dmd_log_dict = self.compute_distribution_matching_loss(
            image_or_video=pred_image,
            conditional_dict=conditional_dict,
            unconditional_dict=unconditional_dict,
            gradient_mask=gradient_mask,
            denoised_timestep_from=denoised_timestep_from,
            denoised_timestep_to=denoised_timestep_to
        )
        
        loss = self.gan_loss_weight * gan_loss+dmd_loss


        dmd_log_dict["dmd_loss"] = dmd_loss
        if len(log_dict) != 0:
            dmd_log_dict["gan_loss"] = gan_loss
            dmd_log_dict.update(log_dict)

        return loss, dmd_log_dict


    # def gan_loss(
    #     self,
    #     image_or_video,
    #     image_or_video_shape,
    #     conditional_dict: dict,
    #     unconditional_dict: dict,
    #     # clean_latent: torch.Tensor,
    #     # initial_latent: torch.Tensor = None,
    #     denoise_steps: int=None,
    # ) -> Tuple[torch.Tensor, dict]:
    #     """
    #     Generate image/videos from noise and compute the DMD loss.
    #     The noisy input to the generator is backward simulated.
    #     This removes the need of any datasets during distillation.
    #     See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
    #     Input:
    #         - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
    #         - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
    #         - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
    #         - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
    #     Output:
    #         - loss: a scalar tensor representing the generator loss.
    #         - generator_log_dict: a dictionary containing the intermediate tensors for logging.
    #     """
    #     # Step 1: Unroll generator to obtain fake videos, and obtain the prediction
    #     # TODO: run_generator should output preditions too, 1 prediction per chunk
    #     pred_image = image_or_video

    #     # Step 2: Get timestep and add noise to generated/real latents
    #     min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep
    #     max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep
    
        
    #     critic_timestep = self._get_timestep(
    #         min_timestep,
    #         max_timestep,
    #         image_or_video_shape[0],
    #         image_or_video_shape[1],
    #         self.num_frame_per_block,
    #         uniform_timestep=True
    #     )

    #     if self.critic_timestep_shift > 1:
    #         critic_timestep = self.critic_timestep_shift * \
    #             (critic_timestep / 1000) / (1 + (self.critic_timestep_shift - 1) * (critic_timestep / 1000)) * 1000

    #     critic_timestep = critic_timestep.clamp(self.min_step, self.max_step)

    #     # use raw image
    #     # critic_timestep = critic_timestep * 0

    #     # add noise
    #     critic_noise = torch.randn_like(pred_image)
    #     noisy_fake_latent = self.scheduler.add_noise(
    #         pred_image.flatten(0, 1),
    #         critic_noise.flatten(0, 1),
    #         critic_timestep.flatten(0, 1)
    #     ).unflatten(0, image_or_video_shape[:2])

        
    #     # noisy_fake_latent_next.requires_grad_(True)

    #     logit = self._run_cls_pred_branch(noisy_fake_latent,conditional_dict, critic_timestep)[0][denoise_steps+1]

    #     gan_loss = F.softplus(-logit.float()).mean() * self.gan_g_weight

    #     log_dict = {"gan_loss": gan_loss}
    #     gan_loss = gan_loss
        

    #     return gan_loss, log_dict

    


    def gan_loss(
        self,
        real_image_or_video,
        image_or_video_shape,
        conditional_dict: dict,
        unconditional_dict: dict,
        # clean_latent: torch.Tensor,
        initial_latent: torch.Tensor = None,
        denoise_steps: int=None,
    ) -> Tuple[torch.Tensor, dict]:
        """
        Generate image/videos from noise and train the critic with generated samples.
        The noisy input to the generator is backward simulated.
        This removes the need of any datasets during distillation.
        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
        Input:
            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
        Output:
            - loss: a scalar tensor representing the generator loss.
            - critic_log_dict: a dictionary containing the intermediate tensors for logging.
        """

        # Step 1: Run generator on backward simulated noisy input
        with torch.no_grad():
            generated_image, _, denoised_timestep_from, denoised_timestep_to = self._run_generator(
                image_or_video_shape=image_or_video_shape,
                conditional_dict=conditional_dict,
                initial_latent=initial_latent,
                denoise_steps=denoise_steps+1,
            )

            # real_image_or_video, _, denoised_timestep_from, denoised_timestep_to = self._run_generator(
            #     image_or_video_shape=image_or_video_shape,
            #     conditional_dict=conditional_dict,
            #     initial_latent=initial_latent,
            #     denoise_steps=denoise_steps+1,
            # )

        # Step 2: Get timestep and add noise to generated/real latents
        min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep
        max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep
        critic_timestep = self._get_timestep(
            min_timestep,
            max_timestep,
            image_or_video_shape[0],
            image_or_video_shape[1],
            self.num_frame_per_block,
            uniform_timestep=True
        )

        # # no gan noise
        # critic_timestep = critic_timestep * 0

        if self.critic_timestep_shift > 1:
            critic_timestep = self.critic_timestep_shift * \
                (critic_timestep / 1000) / (1 + (self.critic_timestep_shift - 1) * (critic_timestep / 1000)) * 1000

        critic_timestep = critic_timestep.clamp(self.min_step, self.max_step)

        critic_noise = torch.randn_like(generated_image)
        noisy_fake_latent = self.scheduler.add_noise(
            generated_image.flatten(0, 1),
            critic_noise.flatten(0, 1),
            critic_timestep.flatten(0, 1)
        ).unflatten(0, image_or_video_shape[:2])

        # Step 4: Compute the real GAN discriminator loss
        noisy_real_latent = self.scheduler.add_noise(
            real_image_or_video.flatten(0, 1),
            critic_noise.flatten(0, 1),
            critic_timestep.flatten(0, 1)
        ).unflatten(0, image_or_video_shape[:2])

        conditional_dict_cloned = copy.deepcopy(conditional_dict)
        conditional_dict_cloned["prompt_embeds"] = torch.concatenate(
            (conditional_dict_cloned["prompt_embeds"], conditional_dict_cloned["prompt_embeds"]), dim=0)
        # _, _, noisy_logit = self.fake_score(
        #     noisy_image_or_video=torch.concatenate((noisy_fake_latent, noisy_real_latent), dim=0),
        #     conditional_dict=conditional_dict_cloned,
        #     timestep=torch.concatenate((critic_timestep, critic_timestep), dim=0),
        #     classify_mode=True,
        #     concat_time_embeddings=self.concat_time_embeddings
        # )
        noisy_logit = self._run_cls_pred_branch(torch.concatenate((noisy_fake_latent, noisy_real_latent), dim=0), conditional_dict_cloned, torch.concatenate((critic_timestep, critic_timestep), dim=0))
        noisy_fake_logit, noisy_real_logit = noisy_logit.chunk(2, dim=0)

        # process logits
        noisy_fake_logit = noisy_fake_logit[0][denoise_steps+1]
        noisy_real_logit = noisy_real_logit[0][denoise_steps+1]

        # # process logits using all discriminator
        # noisy_fake_logit = noisy_fake_logit[0][1:]
        # noisy_real_logit = noisy_real_logit[0][1:]


        if not self.relativistic_discriminator:
            gan_D_loss = F.softplus(-noisy_real_logit.float()).mean() + F.softplus(noisy_fake_logit.float()).mean()
        else:
            relative_real_logit = noisy_real_logit - noisy_fake_logit
            gan_D_loss = F.softplus(-relative_real_logit.float()).mean()
        gan_D_loss = gan_D_loss * self.gan_d_weight

        critic_log_dict = {
            "critic_timestep": critic_timestep.detach(),
            'noisy_real_logit': noisy_fake_logit.detach(),
            'noisy_fake_logit': noisy_real_logit.detach(),
            'logit_diff': (-noisy_real_logit+noisy_fake_logit).detach(),
        }

        return gan_D_loss, critic_log_dict



    def discriminator_loss(
        self,
        image_or_video_shape,
        conditional_dict: dict,
        unconditional_dict: dict,
        clean_latent: torch.Tensor,
        real_image_or_video: torch.Tensor,
        initial_latent: torch.Tensor = None
    ) -> Tuple[torch.Tensor, dict]:
        """
        Generate image/videos from noise and train the critic with generated samples.
        The noisy input to the generator is backward simulated.
        This removes the need of any datasets during distillation.
        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
        Input:
            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
        Output:
            - loss: a scalar tensor representing the generator loss.
            - critic_log_dict: a dictionary containing the intermediate tensors for logging.
        """

        # Step 1: Run generator on backward simulated noisy input
        denoise_steps = self.inference_pipeline.generate_and_sync_list(1, len(self.inference_pipeline.denoising_step_list)-1, device=conditional_dict["prompt_embeds"].device)[0]
        with torch.no_grad():
            generated_image, _, denoised_timestep_from, denoised_timestep_to = self._run_generator(
                image_or_video_shape=image_or_video_shape,
                conditional_dict=conditional_dict,
                initial_latent=initial_latent,
                denoise_steps=denoise_steps,
            )

            real_image_or_video, _, denoised_timestep_from, denoised_timestep_to = self._run_generator(
                image_or_video_shape=image_or_video_shape,
                conditional_dict=conditional_dict,
                initial_latent=initial_latent,
                denoise_steps=denoise_steps+1,
            )

        # Step 2: Get timestep and add noise to generated/real latents
        min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep
        max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep
        critic_timestep = self._get_timestep(
            min_timestep,
            max_timestep,
            image_or_video_shape[0],
            image_or_video_shape[1],
            self.num_frame_per_block,
            uniform_timestep=True
        )

        # # no gan noise
        # critic_timestep = critic_timestep * 0

        if self.critic_timestep_shift > 1:
            critic_timestep = self.critic_timestep_shift * \
                (critic_timestep / 1000) / (1 + (self.critic_timestep_shift - 1) * (critic_timestep / 1000)) * 1000

        critic_timestep = critic_timestep.clamp(self.min_step, self.max_step)

        critic_noise = torch.randn_like(generated_image)
        noisy_fake_latent = self.scheduler.add_noise(
            generated_image.flatten(0, 1),
            critic_noise.flatten(0, 1),
            critic_timestep.flatten(0, 1)
        ).unflatten(0, image_or_video_shape[:2])

        # Step 4: Compute the real GAN discriminator loss
        noisy_real_latent = self.scheduler.add_noise(
            real_image_or_video.flatten(0, 1),
            critic_noise.flatten(0, 1),
            critic_timestep.flatten(0, 1)
        ).unflatten(0, image_or_video_shape[:2])

        conditional_dict_cloned = copy.deepcopy(conditional_dict)
        conditional_dict_cloned["prompt_embeds"] = torch.concatenate(
            (conditional_dict_cloned["prompt_embeds"], conditional_dict_cloned["prompt_embeds"]), dim=0)
        # _, _, noisy_logit = self.fake_score(
        #     noisy_image_or_video=torch.concatenate((noisy_fake_latent, noisy_real_latent), dim=0),
        #     conditional_dict=conditional_dict_cloned,
        #     timestep=torch.concatenate((critic_timestep, critic_timestep), dim=0),
        #     classify_mode=True,
        #     concat_time_embeddings=self.concat_time_embeddings
        # )
        noisy_logit = self._run_cls_pred_branch(torch.concatenate((noisy_fake_latent, noisy_real_latent), dim=0), conditional_dict_cloned, torch.concatenate((critic_timestep, critic_timestep), dim=0))
        noisy_fake_logit, noisy_real_logit = noisy_logit.chunk(2, dim=0)

        # process logits
        # noisy_fake_logit = noisy_fake_logit[0][denoise_steps+1]
        # noisy_real_logit = noisy_real_logit[0][denoise_steps+1]


        if not self.relativistic_discriminator:
            gan_D_loss = F.softplus(-noisy_real_logit[0][denoise_steps+1].float()).mean() + F.softplus(noisy_fake_logit[0][denoise_steps+1].float()).mean()
        else:
            relative_real_logit = noisy_real_logit[0][denoise_steps+1] - noisy_fake_logit[0][denoise_steps+1]
            gan_D_loss = F.softplus(-relative_real_logit.float()).mean()
        gan_D_loss = gan_D_loss * self.gan_d_weight

        # R1 regularization
        if self.r1_weight > 0.:
            noisy_real_latent_perturbed = noisy_real_latent.clone()
            epison_real = self.r1_sigma * torch.randn_like(noisy_real_latent_perturbed)
            noisy_real_latent_perturbed = noisy_real_latent_perturbed + epison_real
            noisy_real_logit_perturbed = self._run_cls_pred_branch(
                noisy_image_or_video=noisy_real_latent_perturbed,
                conditional_dict=conditional_dict,
                timestep=critic_timestep
            )

            r1_grad = (noisy_real_logit_perturbed - noisy_real_logit) / self.r1_sigma
            # add gan d weight
            r1_loss = self.r1_weight * torch.mean((r1_grad)**2) * self.gan_d_weight
        else:
            r1_loss = torch.zeros_like(gan_D_loss)

        # R2 regularization
        if self.r2_weight > 0.:
            noisy_fake_latent_perturbed = noisy_fake_latent.clone()
            epison_generated = self.r2_sigma * torch.randn_like(noisy_fake_latent_perturbed)
            noisy_fake_latent_perturbed = noisy_fake_latent_perturbed + epison_generated
            noisy_fake_logit_perturbed = self._run_cls_pred_branch(
                noisy_image_or_video=noisy_fake_latent_perturbed,
                conditional_dict=conditional_dict,
                timestep=critic_timestep
            )

            r2_grad = (noisy_fake_logit_perturbed - noisy_fake_logit) / self.r2_sigma
            # add gan d weight
            r2_loss = self.r2_weight * torch.mean((r2_grad)**2) * self.gan_d_weight
        else:
            r2_loss = torch.zeros_like(gan_D_loss)

        critic_log_dict = {
            "critic_timestep": critic_timestep.detach(),
            'noisy_real_logit': noisy_real_logit.detach(),
            'noisy_fake_logit': noisy_fake_logit.detach(),
            'logit_diff': (noisy_real_logit[0][denoise_steps+1]-noisy_fake_logit[0][denoise_steps+1]).detach(),
            "r1_loss": r1_loss,
            "r2_loss": r2_loss,
        }

        gan_D_loss = gan_D_loss + 0.5 * (r1_loss + r2_loss)

        return gan_D_loss, critic_log_dict




    # def discriminator_loss(
    #     self,
    #     image_or_video_shape,
    #     conditional_dict: dict,
    #     unconditional_dict: dict,
    #     clean_latent: torch.Tensor,
    #     real_image_or_video: torch.Tensor,
    #     initial_latent: torch.Tensor = None
    # ) -> Tuple[torch.Tensor, dict]:
    #     """
    #     Generate image/videos from noise and train the critic with generated samples.
    #     The noisy input to the generator is backward simulated.
    #     This removes the need of any datasets during distillation.
    #     See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
    #     Input:
    #         - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
    #         - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
    #         - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
    #         - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
    #     Output:
    #         - loss: a scalar tensor representing the generator loss.
    #         - critic_log_dict: a dictionary containing the intermediate tensors for logging.
    #     """

    #     # Step 1: Run generator on backward simulated noisy input
    #     # with torch.no_grad():
    #     denoise_steps = self.inference_pipeline.generate_and_sync_list(1, len(self.inference_pipeline.denoising_step_list)-1, device=conditional_dict["prompt_embeds"].device)[0]
    #     generated_image, _, denoised_timestep_from, denoised_timestep_to = self._run_generator(
    #         image_or_video_shape=image_or_video_shape,
    #         conditional_dict=conditional_dict,
    #         initial_latent=initial_latent,
    #         denoise_steps=denoise_steps,
    #     )

    #     denoise_steps_next = denoise_steps+1
    #     generated_image_next, _, denoised_timestep_from, denoised_timestep_to = self._run_generator(
    #         image_or_video_shape=image_or_video_shape,
    #         conditional_dict=conditional_dict,
    #         initial_latent=initial_latent,
    #         denoise_steps=denoise_steps_next,
    #     )

    #     # Step 2: Get timestep and add noise to generated/real latents
    #     min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep
    #     max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep
    #     critic_timestep = self._get_timestep(
    #         min_timestep,
    #         max_timestep,
    #         image_or_video_shape[0],
    #         image_or_video_shape[1],
    #         self.num_frame_per_block,
    #         uniform_timestep=True
    #     )

    #     if self.critic_timestep_shift > 1:
    #         critic_timestep = self.critic_timestep_shift * \
    #             (critic_timestep / 1000) / (1 + (self.critic_timestep_shift - 1) * (critic_timestep / 1000)) * 1000

    #     critic_timestep = critic_timestep.clamp(self.min_step, self.max_step)

    #     # critic_timestep = critic_timestep * 0

    #     critic_noise = torch.randn_like(generated_image)
    #     noisy_latent = self.scheduler.add_noise(
    #         generated_image.flatten(0, 1),
    #         critic_noise.flatten(0, 1),
    #         critic_timestep.flatten(0, 1)
    #     ).unflatten(0, image_or_video_shape[:2])

    #     critic_noise = torch.randn_like(generated_image)
    #     noisy_latent_next = self.scheduler.add_noise(
    #         generated_image_next.flatten(0, 1),
    #         critic_noise.flatten(0, 1),
    #         critic_timestep.flatten(0, 1)
    #     ).unflatten(0, image_or_video_shape[:2])

    #     # # Step 4: Compute the real GAN discriminator loss
    #     # noisy_real_latent = self.scheduler.add_noise(
    #     #     real_image_or_video.flatten(0, 1),
    #     #     critic_noise.flatten(0, 1),
    #     #     critic_timestep.flatten(0, 1)
    #     # ).unflatten(0, image_or_video_shape[:2])

    #     # conditional_dict_cloned = copy.deepcopy(conditional_dict)
    #     # conditional_dict_cloned["prompt_embeds"] = torch.concatenate(
    #     #     (conditional_dict_cloned["prompt_embeds"], conditional_dict_cloned["prompt_embeds"]), dim=0)
    #     # _, _, noisy_logit = self.fake_score(
    #     #     noisy_image_or_video=torch.concatenate((noisy_fake_latent, noisy_real_latent), dim=0),
    #     #     conditional_dict=conditional_dict_cloned,
    #     #     timestep=torch.concatenate((critic_timestep, critic_timestep), dim=0),
    #     #     classify_mode=True,
    #     #     concat_time_embeddings=self.concat_time_embeddings
    #     # )
    #     # noisy_fake_logit, noisy_real_logit = noisy_logit.chunk(2, dim=0)

    #     logit = self._run_cls_pred_branch(noisy_latent, conditional_dict, critic_timestep)[0]
    #     logit_next = self._run_cls_pred_branch(noisy_latent_next, conditional_dict, critic_timestep)[0]
    #     logit[0] = 0
    #     logit_next[0] = 0
    #     logit = F.softmax(logit)[denoise_steps_next]
    #     logit_next = F.softmax(logit_next)[denoise_steps_next]
    #     if not self.relativistic_discriminator:
    #         gan_D_loss = F.softplus(-noisy_latent.float()).mean()
    #     else:
    #         relative_real_logit = logit_next - logit
    #         gan_D_loss = F.softplus(-relative_real_logit.float()).mean()

    #     gan_D_loss = gan_D_loss * self.gan_d_weight

    #     discriminator_log_dict = {
    #         # "critic_timestep": critic_timestep.detach(),
    #         'logit': logit.detach(),
    #         "logit_diff": relative_real_logit.detach(),
    #         'denoise_steps': denoise_steps,
    #     }

    #     return gan_D_loss, discriminator_log_dict


    def critic_loss(
        self,
        image_or_video_shape,
        conditional_dict: dict,
        unconditional_dict: dict,
        clean_latent: torch.Tensor,
        initial_latent: torch.Tensor = None
    ) -> Tuple[torch.Tensor, dict]:
        """
        Generate image/videos from noise and train the critic with generated samples.
        The noisy input to the generator is backward simulated.
        This removes the need of any datasets during distillation.
        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
        Input:
            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
        Output:
            - loss: a scalar tensor representing the generator loss.
            - critic_log_dict: a dictionary containing the intermediate tensors for logging.
        """

        # Step 1: Run generator on backward simulated noisy input
        with torch.no_grad():
            generated_image, _, denoised_timestep_from, denoised_timestep_to = self._run_generator(
                image_or_video_shape=image_or_video_shape,
                conditional_dict=conditional_dict,
                initial_latent=initial_latent
            )

        # Step 2: Compute the fake prediction
        min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep
        max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep
        critic_timestep = self._get_timestep(
            min_timestep,
            max_timestep,
            image_or_video_shape[0],
            image_or_video_shape[1],
            self.num_frame_per_block,
            uniform_timestep=True
        )

        if self.timestep_shift > 1:
            critic_timestep = self.timestep_shift * \
                (critic_timestep / 1000) / (1 + (self.timestep_shift - 1) * (critic_timestep / 1000)) * 1000

        critic_timestep = critic_timestep.clamp(self.min_step, self.max_step)

        critic_noise = torch.randn_like(generated_image)
        noisy_generated_image = self.scheduler.add_noise(
            generated_image.flatten(0, 1),
            critic_noise.flatten(0, 1),
            critic_timestep.flatten(0, 1)
        ).unflatten(0, image_or_video_shape[:2])

        _, pred_fake_image = self.fake_score(
            noisy_image_or_video=noisy_generated_image,
            conditional_dict=conditional_dict,
            timestep=critic_timestep
        )

        # Step 3: Compute the denoising loss for the fake critic
        if self.args.denoising_loss_type == "flow":
            from utils.wan_wrapper import WanDiffusionWrapper
            flow_pred = WanDiffusionWrapper._convert_x0_to_flow_pred(
                scheduler=self.scheduler,
                x0_pred=pred_fake_image.flatten(0, 1),
                xt=noisy_generated_image.flatten(0, 1),
                timestep=critic_timestep.flatten(0, 1)
            )
            pred_fake_noise = None
        else:
            flow_pred = None
            pred_fake_noise = self.scheduler.convert_x0_to_noise(
                x0=pred_fake_image.flatten(0, 1),
                xt=noisy_generated_image.flatten(0, 1),
                timestep=critic_timestep.flatten(0, 1)
            ).unflatten(0, image_or_video_shape[:2])

        denoising_loss = self.denoising_loss_func(
            x=generated_image.flatten(0, 1),
            x_pred=pred_fake_image.flatten(0, 1),
            noise=critic_noise.flatten(0, 1),
            noise_pred=pred_fake_noise,
            alphas_cumprod=self.scheduler.alphas_cumprod,
            timestep=critic_timestep.flatten(0, 1),
            flow_pred=flow_pred
        )

        # Step 5: Debugging Log
        critic_log_dict = {
            "critic_timestep": critic_timestep.detach()
        }

        return denoising_loss, critic_log_dict

