"""Shared generator setup and few-step rollout utilities.

This module contains only generator behavior used by distribution matching.  Loss
construction lives in :mod:`model.distribution_matching`.
"""

from __future__ import annotations

from typing import Optional

import torch
import torch.distributed as dist

from model.base import SelfForcingModel


class FewStepGenerationModel(SelfForcingModel):
    """Initialize the causal generator and expose deterministic replay helpers."""

    def __init__(self, args, device):
        super().__init__(args, device)
        self.causal_generator = bool(getattr(args, "causal", True))
        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.num_training_frames = getattr(args, "num_training_frames", 21)

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

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

        self.inference_pipeline = None
        self.guidance_scale = float(getattr(args, "guidance_scale", 1.0))

        if not hasattr(self, "denoising_step_list"):
            raise ValueError(
                "few-step generation requires denoising_step_list")
        denoising_steps = self.denoising_step_list.to(device=device)
        if denoising_steps.numel() > 0 and denoising_steps[-1].item() == 0:
            denoising_steps = denoising_steps[:-1]
        self.generation_denoising_step_list = denoising_steps

    @staticmethod
    def _apply_gradient_mask(
        latents: torch.Tensor,
        gradient_mask: Optional[torch.Tensor],
    ) -> torch.Tensor:
        if gradient_mask is None:
            return latents
        frame_mask = gradient_mask.flatten(2).any(dim=2).view(
            latents.shape[0], latents.shape[1], 1, 1, 1)
        return torch.where(frame_mask, latents, latents.detach())

    def _predict_x0_with_cfg(
        self,
        noisy_latent: torch.Tensor,
        conditional_dict: dict,
        unconditional_dict: dict,
        timestep: torch.Tensor,
    ) -> torch.Tensor:
        _, x0_cond = self.generator(
            noisy_image_or_video=noisy_latent,
            conditional_dict=conditional_dict,
            timestep=timestep,
        )
        if self.guidance_scale == 1.0:
            return x0_cond

        _, x0_uncond = self.generator(
            noisy_image_or_video=noisy_latent,
            conditional_dict=unconditional_dict,
            timestep=timestep,
        )
        return x0_uncond + self.guidance_scale * (x0_cond - x0_uncond)

    def _run_few_step_generator(
        self,
        image_or_video_shape,
        conditional_dict: dict,
        unconditional_dict: dict,
        exit_index: Optional[int] = None,
    ) -> torch.Tensor:
        """Run detached rollout steps followed by one differentiated exit step."""
        latent = torch.randn(
            image_or_video_shape,
            device=self.device,
            dtype=self.dtype,
        )
        batch_size, num_frames = latent.shape[:2]
        denoising_steps = self.generation_denoising_step_list.to(
            device=self.device)
        num_steps = len(denoising_steps)

        if exit_index is None:
            if num_steps <= 1:
                exit_index = 0
            else:
                index_tensor = torch.randint(
                    0, num_steps, (1,), device=self.device)
                if dist.is_available() and dist.is_initialized():
                    dist.broadcast(index_tensor, src=0)
                exit_index = int(index_tensor.item())

        for index in range(num_steps):
            current_timestep = denoising_steps[index]
            timestep = torch.ones(
                [batch_size, num_frames],
                device=self.device,
                dtype=current_timestep.dtype,
            ) * current_timestep
            if index < exit_index:
                with torch.no_grad():
                    x0_pred = self._predict_x0_with_cfg(
                        noisy_latent=latent,
                        conditional_dict=conditional_dict,
                        unconditional_dict=unconditional_dict,
                        timestep=timestep,
                    )
                    next_timestep = denoising_steps[index + 1]
                    latent = self.scheduler.add_noise(
                        x0_pred.flatten(0, 1),
                        torch.randn_like(x0_pred.flatten(0, 1)),
                        next_timestep * torch.ones(
                            [batch_size * num_frames],
                            device=self.device,
                            dtype=next_timestep.dtype,
                        ),
                    ).unflatten(0, x0_pred.shape[:2]).to(self.dtype)
            else:
                return self._predict_x0_with_cfg(
                    noisy_latent=latent,
                    conditional_dict=conditional_dict,
                    unconditional_dict=unconditional_dict,
                    timestep=timestep,
                )
        return latent