from __future__ import annotations

import math
from dataclasses import dataclass

import torch
import torch.nn.functional as F
from torch import nn
from torch.utils.checkpoint import checkpoint

from livediffusion.config import ModelConfig


@dataclass(slots=True)
class PrefixCache:
    key_values: tuple[tuple[torch.Tensor, torch.Tensor], ...]
    pooled: torch.Tensor


MemoryCache = PrefixCache


def _video_positions(
    frames: int,
    height: int,
    width: int,
    time_offset: int,
    device: torch.device,
) -> torch.Tensor:
    time = torch.arange(time_offset, time_offset + frames, device=device)
    y = torch.arange(height, device=device)
    x = torch.arange(width, device=device)
    grid = torch.meshgrid(time, y, x, indexing="ij")
    return torch.stack(grid, dim=-1).reshape(-1, 3).float()


class Rotary3D(nn.Module):
    """3D RoPE with a shorter temporal base so time and space do not share a spectrum."""

    def __init__(
        self,
        head_size: int,
        time_base: float = 1_000.0,
        space_base: float = 10_000.0,
    ) -> None:
        super().__init__()
        self.rotary_size = (head_size // 6) * 6
        self.axis_size = self.rotary_size // 3
        axis_half = self.axis_size // 2
        steps = torch.arange(axis_half, dtype=torch.float32) / max(axis_half, 1)
        time_frequencies = torch.exp(-math.log(time_base) * steps)
        space_frequencies = torch.exp(-math.log(space_base) * steps)
        self.register_buffer(
            "frequencies",
            torch.stack((time_frequencies, space_frequencies, space_frequencies)),
            persistent=False,
        )

    @staticmethod
    def _rotate_half(inputs: torch.Tensor) -> torch.Tensor:
        pairs = inputs.reshape(*inputs.shape[:-1], -1, 2)
        first, second = pairs.unbind(dim=-1)
        return torch.stack((-second, first), dim=-1).flatten(-2)

    def forward(self, inputs: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
        angles = torch.cat(
            [
                positions[:, axis, None] * self.frequencies[axis][None]
                for axis in range(3)
            ],
            dim=-1,
        )
        cosine = angles.cos().repeat_interleave(2, dim=-1)[None, None]
        sine = angles.sin().repeat_interleave(2, dim=-1)[None, None]
        rotary, remainder = inputs.split(
            (self.rotary_size, inputs.shape[-1] - self.rotary_size),
            dim=-1,
        )
        rotated = rotary * cosine.to(rotary.dtype)
        rotated = rotated + self._rotate_half(rotary) * sine.to(rotary.dtype)
        return torch.cat((rotated, remainder), dim=-1)


class ScalarEmbedding(nn.Module):
    def __init__(self, hidden_size: int) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        self.mlp = nn.Sequential(
            nn.Linear(hidden_size, hidden_size),
            nn.SiLU(),
            nn.Linear(hidden_size, hidden_size),
        )

    def forward(self, value: torch.Tensor) -> torch.Tensor:
        value = value.reshape(-1).float()
        half = self.hidden_size // 2
        frequencies = torch.exp(
            -math.log(10_000.0)
            * torch.arange(half, device=value.device, dtype=torch.float32)
            / half
        )
        angles = value[:, None] * frequencies[None] * 1_000.0
        embedding = torch.cat((angles.cos(), angles.sin()), dim=-1)
        return self.mlp(embedding.to(dtype=self.mlp[0].weight.dtype))


class SwiGLU(nn.Module):
    def __init__(self, hidden_size: int, ratio: float) -> None:
        super().__init__()
        inner_size = int(hidden_size * ratio * 2 / 3)
        inner_size = ((inner_size + 63) // 64) * 64
        self.input = nn.Linear(hidden_size, inner_size * 2, bias=False)
        self.output = nn.Linear(inner_size, hidden_size, bias=False)

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        values, gates = self.input(inputs).chunk(2, dim=-1)
        return self.output(values * F.silu(gates))


class SelfAttention(nn.Module):
    def __init__(self, hidden_size: int, num_heads: int, qk_norm: bool) -> None:
        super().__init__()
        self.num_heads = num_heads
        self.head_size = hidden_size // num_heads
        self.qkv = nn.Linear(hidden_size, hidden_size * 3, bias=False)
        self.output = nn.Linear(hidden_size, hidden_size, bias=False)
        self.q_norm = (
            nn.RMSNorm(self.head_size, eps=1e-6)
            if qk_norm
            else nn.Identity()
        )
        self.k_norm = (
            nn.RMSNorm(self.head_size, eps=1e-6)
            if qk_norm
            else nn.Identity()
        )
        self.rotary = Rotary3D(self.head_size)

    def forward(self, inputs: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
        batch, tokens, channels = inputs.shape
        query, key, value = self.qkv(inputs).reshape(
            batch,
            tokens,
            3,
            self.num_heads,
            self.head_size,
        ).unbind(dim=2)
        query = self.rotary(self.q_norm(query).transpose(1, 2), positions)
        key = self.rotary(self.k_norm(key).transpose(1, 2), positions)
        value = value.transpose(1, 2)
        attended = F.scaled_dot_product_attention(query, key, value)
        return self.output(attended.transpose(1, 2).reshape(batch, tokens, channels))


class PrefixAttention(nn.Module):
    """Target attention over the joint clean-prefix and noisy-target sequence.

    Keys and values may use fewer heads than queries. Prefix and target keys are
    concatenated, so both sides share the same reduced head count.
    """

    def __init__(
        self,
        hidden_size: int,
        num_heads: int,
        qk_norm: bool,
        kv_heads: int = 0,
    ) -> None:
        super().__init__()
        self.num_heads = num_heads
        self.head_size = hidden_size // num_heads
        self.kv_heads = kv_heads or num_heads
        kv_size = self.kv_heads * self.head_size
        self.target_qkv = nn.Linear(hidden_size, hidden_size + kv_size * 2, bias=False)
        self.prefix_key_value = nn.Linear(hidden_size, kv_size * 2, bias=False)
        self.output = nn.Linear(hidden_size, hidden_size, bias=False)
        self.q_norm = (
            nn.RMSNorm(self.head_size, eps=1e-6)
            if qk_norm
            else nn.Identity()
        )
        self.k_norm = (
            nn.RMSNorm(self.head_size, eps=1e-6)
            if qk_norm
            else nn.Identity()
        )
        self.rotary = Rotary3D(self.head_size)

    def project_prefix(
        self,
        prefix: torch.Tensor,
        positions: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        batch, tokens, _ = prefix.shape
        key, value = self.prefix_key_value(prefix).reshape(
            batch,
            tokens,
            2,
            self.kv_heads,
            self.head_size,
        ).unbind(dim=2)
        key = self.rotary(self.k_norm(key).transpose(1, 2), positions)
        return key, value.transpose(1, 2)

    def forward_cached(
        self,
        target: torch.Tensor,
        target_positions: torch.Tensor,
        prefix_key_value: tuple[torch.Tensor, torch.Tensor],
    ) -> torch.Tensor:
        batch, tokens, channels = target.shape
        kv_size = self.kv_heads * self.head_size
        projected = self.target_qkv(target)
        query, key, value = projected.split((channels, kv_size, kv_size), dim=-1)
        query = query.reshape(batch, tokens, self.num_heads, self.head_size)
        key = key.reshape(batch, tokens, self.kv_heads, self.head_size)
        value = value.reshape(batch, tokens, self.kv_heads, self.head_size)
        query = self.rotary(self.q_norm(query).transpose(1, 2), target_positions)
        key = self.rotary(self.k_norm(key).transpose(1, 2), target_positions)
        value = value.transpose(1, 2)
        prefix_key, prefix_value = prefix_key_value
        key = torch.cat((prefix_key, key), dim=2)
        value = torch.cat((prefix_value, value), dim=2)
        attended = F.scaled_dot_product_attention(
            query,
            key,
            value,
            enable_gqa=self.kv_heads != self.num_heads,
        )
        return self.output(attended.transpose(1, 2).reshape(batch, tokens, channels))


@dataclass(slots=True)
class WorldMemoryState:
    plane: torch.Tensor
    registers: torch.Tensor


class WorldMemory(nn.Module):
    """Gated recurrent memory at latent resolution, plus global register state.

    The plane is a convolutional GRU over one cell per latent position, so an
    irreversible world edit stays represented after it scrolls out of the
    attention prefix. Registers carry frame-global state the plane cannot hold.
    """

    def __init__(
        self,
        latent_channels: int,
        hidden_size: int,
        memory_channels: int,
        read_patch: int,
        register_tokens: int,
    ) -> None:
        super().__init__()
        self.memory_channels = memory_channels
        self.read_patch = read_patch
        self.register_tokens = register_tokens
        self.cell = nn.Conv2d(
            latent_channels + memory_channels,
            memory_channels * 3,
            kernel_size=3,
            padding=1,
        )
        self.read = nn.Conv2d(
            memory_channels,
            hidden_size,
            kernel_size=read_patch,
            stride=read_patch,
        )
        self.read_norm = nn.RMSNorm(hidden_size, eps=1e-6)
        self.scale_embedding = nn.Parameter(torch.zeros(1, 1, hidden_size))
        self.register_update = (
            nn.Linear(memory_channels + hidden_size, hidden_size * 2)
            if register_tokens > 0
            else None
        )

    def initial_state(
        self,
        batch: int,
        height: int,
        width: int,
        registers: torch.Tensor,
        device: torch.device,
        dtype: torch.dtype,
    ) -> WorldMemoryState:
        plane = torch.zeros(
            batch,
            self.memory_channels,
            height,
            width,
            device=device,
            dtype=dtype,
        )
        return WorldMemoryState(plane, registers)

    def advance(
        self,
        state: WorldMemoryState,
        frames: torch.Tensor,
    ) -> WorldMemoryState:
        plane = state.plane
        for index in range(frames.shape[1]):
            latent = frames[:, index].to(plane.dtype)
            reset, update, candidate = self.cell(
                torch.cat((latent, plane), dim=1)
            ).chunk(3, dim=1)
            gate = torch.sigmoid(update)
            proposal = torch.tanh(candidate + torch.sigmoid(reset) * plane)
            plane = plane * (1.0 - gate) + proposal * gate
        registers = state.registers
        if self.register_update is not None:
            pooled = plane.mean(dim=(2, 3))
            summary = torch.cat(
                (pooled[:, None].expand(-1, registers.shape[1], -1), registers),
                dim=-1,
            )
            gate, proposal = self.register_update(summary).chunk(2, dim=-1)
            gate = torch.sigmoid(gate)
            registers = registers * (1.0 - gate) + torch.tanh(proposal) * gate
        return WorldMemoryState(plane, registers)

    def tokens(self, state: WorldMemoryState) -> torch.Tensor:
        read = self.read(state.plane).flatten(2).transpose(1, 2)
        return self.read_norm(read) + self.scale_embedding


class SharedModulation(nn.Module):
    """One adaLN projection shared by every block, with per-block offsets.

    Per-block modulation costs 6 * hidden^2 each; sharing recovers that budget
    for depth while keeping each block independently steerable.
    """

    def __init__(self, hidden_size: int, blocks: int) -> None:
        super().__init__()
        self.projection = nn.Linear(hidden_size, hidden_size * 6)
        self.offsets = nn.Parameter(torch.zeros(blocks, hidden_size * 6))

    def forward(self, condition: torch.Tensor, block_index: int) -> torch.Tensor:
        return self.projection(F.silu(condition)) + self.offsets[block_index]


class PrefixEncoderBlock(nn.Module):
    """Clean-context stack that runs before the joint prefix/target DiT."""

    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        hidden_size = config.hidden_size
        self.norm1 = nn.RMSNorm(hidden_size, eps=1e-6)
        self.attention = SelfAttention(
            hidden_size,
            config.num_heads,
            config.qk_norm,
        )
        self.norm2 = nn.RMSNorm(hidden_size, eps=1e-6)
        self.mlp = SwiGLU(hidden_size, config.mlp_ratio)
        self.attn_scale = nn.Parameter(torch.ones(hidden_size))
        self.mlp_scale = nn.Parameter(torch.ones(hidden_size))

    def forward(self, prefix: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
        prefix = prefix + self.attn_scale * self.attention(self.norm1(prefix), positions)
        return prefix + self.mlp_scale * self.mlp(self.norm2(prefix))


class PrefixDiTBlock(nn.Module):
    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        hidden_size = config.hidden_size
        self.prefix_norm1 = nn.RMSNorm(hidden_size, eps=1e-6)
        self.prefix_attention = SelfAttention(
            hidden_size,
            config.num_heads,
            config.qk_norm,
        )
        self.prefix_norm2 = nn.RMSNorm(hidden_size, eps=1e-6)
        self.prefix_mlp = SwiGLU(hidden_size, config.mlp_ratio)
        self.prefix_projection_norm = nn.RMSNorm(hidden_size, eps=1e-6)
        self.target_norm1 = nn.RMSNorm(
            hidden_size,
            eps=1e-6,
            elementwise_affine=False,
        )
        self.target_attention = PrefixAttention(
            hidden_size,
            config.num_heads,
            config.qk_norm,
            config.kv_heads,
        )
        self.target_norm2 = nn.RMSNorm(
            hidden_size,
            eps=1e-6,
            elementwise_affine=False,
        )
        self.target_mlp = SwiGLU(hidden_size, config.mlp_ratio)
        self.modulation = (
            None
            if config.shared_modulation
            else nn.Linear(hidden_size, hidden_size * 6)
        )

    @staticmethod
    def _modulate(
        inputs: torch.Tensor,
        shift: torch.Tensor,
        scale: torch.Tensor,
    ) -> torch.Tensor:
        return inputs * (1.0 + scale[:, None]) + shift[:, None]

    def advance_prefix(
        self,
        prefix: torch.Tensor,
        positions: torch.Tensor,
    ) -> torch.Tensor:
        prefix = prefix + self.prefix_attention(self.prefix_norm1(prefix), positions)
        return prefix + self.prefix_mlp(self.prefix_norm2(prefix))

    def prefix_key_value(
        self,
        prefix: torch.Tensor,
        positions: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        return self.target_attention.project_prefix(
            self.prefix_projection_norm(prefix),
            positions,
        )

    def advance_target(
        self,
        target: torch.Tensor,
        condition: torch.Tensor,
        target_positions: torch.Tensor,
        prefix_key_value: tuple[torch.Tensor, torch.Tensor],
        modulation: torch.Tensor | None = None,
    ) -> torch.Tensor:
        if modulation is None:
            if self.modulation is None:
                raise RuntimeError("Shared modulation must be supplied by the trunk.")
            modulation = self.modulation(F.silu(condition))
        shift1, scale1, gate1, shift2, scale2, gate2 = modulation.chunk(6, dim=-1)
        normalized = self._modulate(self.target_norm1(target), shift1, scale1)
        target = target + gate1[:, None] * self.target_attention.forward_cached(
            normalized,
            target_positions,
            prefix_key_value,
        )
        normalized = self._modulate(self.target_norm2(target), shift2, scale2)
        return target + gate2[:, None] * self.target_mlp(normalized)

    def forward(
        self,
        prefix: torch.Tensor,
        target: torch.Tensor,
        condition: torch.Tensor,
        prefix_positions: torch.Tensor,
        target_positions: torch.Tensor,
        modulation: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        prefix = self.advance_prefix(prefix, prefix_positions)
        key_value = self.prefix_key_value(prefix, prefix_positions)
        target = self.advance_target(
            target,
            condition,
            target_positions,
            key_value,
            modulation,
        )
        return prefix, target


class CausalLatentVideoDiT(nn.Module):
    """Multiscale clean-prefix/noisy-suffix latent video transformer."""

    def __init__(
        self,
        model_config: ModelConfig,
        latent_channels: int,
        history_frames: int = 12,
    ) -> None:
        super().__init__()
        self.model_config = model_config
        self.latent_channels = latent_channels
        self.history_frames = history_frames
        self.gradient_checkpointing = False
        hidden_size = model_config.hidden_size
        target_patch = model_config.target_patch_size
        coarse_patch = model_config.memory_patch_size
        fine_patch = model_config.fine_memory_patch_size
        self.target_patch_embed = nn.Conv3d(
            latent_channels,
            hidden_size,
            kernel_size=(1, target_patch, target_patch),
            stride=(1, target_patch, target_patch),
        )
        self.coarse_prefix_embed = nn.Conv3d(
            latent_channels,
            hidden_size,
            kernel_size=(1, coarse_patch, coarse_patch),
            stride=(1, coarse_patch, coarse_patch),
        )
        self.fine_prefix_embed = nn.Conv3d(
            latent_channels,
            hidden_size,
            kernel_size=(1, fine_patch, fine_patch),
            stride=(1, fine_patch, fine_patch),
        )
        self.motion_prefix_embed = nn.Conv3d(
            latent_channels,
            hidden_size,
            kernel_size=(1, fine_patch, fine_patch),
            stride=(1, fine_patch, fine_patch),
        )
        self.coarse_scale_embedding = nn.Parameter(torch.zeros(1, 1, hidden_size))
        self.fine_scale_embedding = nn.Parameter(torch.zeros(1, 1, hidden_size))
        self.register_tokens = (
            nn.Parameter(torch.randn(1, model_config.register_tokens, hidden_size) * 0.02)
            if model_config.register_tokens > 0
            else None
        )
        self.memory_blocks = nn.ModuleList(
            PrefixEncoderBlock(model_config)
            for _ in range(model_config.prefix_encoder_depth)
        )
        self.t_embedding = ScalarEmbedding(hidden_size)
        self.r_embedding = ScalarEmbedding(hidden_size)
        self.delta_embedding = ScalarEmbedding(hidden_size)
        self.time_projection = nn.Sequential(
            nn.Linear(hidden_size * 3, hidden_size),
            nn.SiLU(),
            nn.Linear(hidden_size, hidden_size),
        )
        self.prefix_condition = nn.Linear(hidden_size, hidden_size)
        self.action_projection = (
            nn.Sequential(
                nn.Linear(model_config.action_dim, hidden_size),
                nn.SiLU(),
                nn.Linear(hidden_size, hidden_size),
            )
            if model_config.action_dim > 0
            else None
        )
        self.context_projection = (
            nn.Linear(model_config.context_dim, hidden_size, bias=False)
            if model_config.context_dim > 0
            else None
        )
        self.blocks = nn.ModuleList(
            PrefixDiTBlock(model_config) for _ in range(model_config.depth)
        )
        self.shared_modulation = (
            SharedModulation(hidden_size, model_config.depth)
            if model_config.shared_modulation
            else None
        )
        self.world_memory = (
            WorldMemory(
                latent_channels,
                hidden_size,
                model_config.memory_channels,
                coarse_patch,
                model_config.register_tokens,
            )
            if model_config.memory_channels > 0
            else None
        )
        self.output_norm = nn.RMSNorm(
            hidden_size,
            eps=1e-6,
            elementwise_affine=False,
        )
        self.output_modulation = nn.Linear(hidden_size, hidden_size * 2)
        self.output_projection = nn.Linear(
            hidden_size,
            target_patch * target_patch * latent_channels,
        )
        self.outcome_norm = nn.RMSNorm(hidden_size, eps=1e-6)
        self.reward_head = nn.Linear(hidden_size, 1)
        self.terminal_head = nn.Linear(hidden_size, 1)
        self.reset_parameters()

    def enable_activation_checkpointing(self, enabled: bool = True) -> None:
        self.gradient_checkpointing = enabled

    def reset_parameters(self) -> None:
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
        for embedding in (
            self.target_patch_embed,
            self.coarse_prefix_embed,
            self.fine_prefix_embed,
            self.motion_prefix_embed,
        ):
            nn.init.normal_(embedding.weight, std=0.02)
        nn.init.normal_(self.coarse_scale_embedding, std=0.02)
        nn.init.normal_(self.fine_scale_embedding, std=0.02)
        for block in self.blocks:
            if block.modulation is not None:
                nn.init.zeros_(block.modulation.weight)
                nn.init.zeros_(block.modulation.bias)
        if self.shared_modulation is not None:
            nn.init.zeros_(self.shared_modulation.projection.weight)
            nn.init.zeros_(self.shared_modulation.projection.bias)
            nn.init.zeros_(self.shared_modulation.offsets)
        if self.world_memory is not None:
            nn.init.zeros_(self.world_memory.scale_embedding)
        nn.init.zeros_(self.output_modulation.weight)
        nn.init.zeros_(self.output_modulation.bias)
        nn.init.zeros_(self.output_projection.weight)
        nn.init.zeros_(self.output_projection.bias)
        nn.init.zeros_(self.reward_head.weight)
        nn.init.zeros_(self.reward_head.bias)
        nn.init.zeros_(self.terminal_head.weight)
        nn.init.zeros_(self.terminal_head.bias)

    @staticmethod
    def _patch(
        embedding: nn.Conv3d,
        video: torch.Tensor,
        time_offset: int,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        tokens = embedding(video.transpose(1, 2))
        _, _, frames, height, width = tokens.shape
        positions = _video_positions(
            frames,
            height,
            width,
            time_offset,
            video.device,
        )
        return tokens.flatten(2).transpose(1, 2), positions

    def _encode_prefix(
        self,
        history: torch.Tensor,
        context: torch.Tensor | None = None,
        history_actions: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if history.shape[1] < self.history_frames:
            raise ValueError(
                f"Expected at least {self.history_frames} history frames, "
                f"got {history.shape[1]}."
            )
        batch = history.shape[0]
        registers = (
            self.register_tokens.expand(batch, -1, -1)
            if self.register_tokens is not None
            else None
        )
        memory_tokens: torch.Tensor | None = None
        if self.world_memory is not None:
            # Every observed frame is folded into the recurrent state, including
            # frames older than the attention window.
            state = self.world_memory.initial_state(
                batch,
                history.shape[-2],
                history.shape[-1],
                registers
                if registers is not None
                else history.new_zeros(batch, 0, self.model_config.hidden_size),
                history.device,
                history.dtype,
            )
            state = self.world_memory.advance(state, history)
            memory_tokens = self.world_memory.tokens(state)
            if registers is not None:
                registers = state.registers
        window = history[:, -self.history_frames :]
        coarse, coarse_positions = self._patch(
            self.coarse_prefix_embed,
            window,
            0,
        )
        if self.action_projection is not None and history_actions is not None:
            window_actions = history_actions[:, -self.history_frames :]
            tokens_per_frame = coarse.shape[1] // self.history_frames
            coarse = coarse + self.action_projection(
                window_actions
            ).repeat_interleave(tokens_per_frame, dim=1)
        history = window
        fine_frames = self.model_config.fine_history_frames
        fine_history = history[:, -fine_frames:]
        previous = history[:, -fine_frames - 1 : -1]
        if previous.shape[1] != fine_frames:
            previous = torch.cat((fine_history[:, :1], fine_history[:, :-1]), dim=1)
        motion = fine_history - previous
        fine, fine_positions = self._patch(
            self.fine_prefix_embed,
            fine_history,
            self.history_frames - fine_frames,
        )
        motion_tokens, _ = self._patch(
            self.motion_prefix_embed,
            motion,
            self.history_frames - fine_frames,
        )
        coarse = coarse + self.coarse_scale_embedding
        fine = fine + motion_tokens + self.fine_scale_embedding
        prefix = torch.cat((coarse, fine), dim=1)
        positions = torch.cat((coarse_positions, fine_positions), dim=0)
        if memory_tokens is not None:
            coarse_patch = self.model_config.memory_patch_size
            memory_positions = _video_positions(
                1,
                history.shape[-2] // coarse_patch,
                history.shape[-1] // coarse_patch,
                -2,
                history.device,
            )
            prefix = torch.cat((memory_tokens, prefix), dim=1)
            positions = torch.cat((memory_positions, positions), dim=0)
        if registers is not None:
            register_count = registers.shape[1]
            register_positions = torch.stack(
                (
                    torch.full((register_count,), -1.0, device=history.device),
                    torch.zeros(register_count, device=history.device),
                    torch.arange(register_count, device=history.device, dtype=torch.float32),
                ),
                dim=-1,
            )
            prefix = torch.cat((registers, prefix), dim=1)
            positions = torch.cat((register_positions, positions), dim=0)
        if self.context_projection is not None and context is not None:
            projected = self.context_projection(context)
            prefix = torch.cat((prefix, projected), dim=1)
            positions = torch.cat(
                (
                    positions,
                    torch.zeros(projected.shape[1], 3, device=history.device),
                ),
                dim=0,
            )
        for block in self.memory_blocks:
            prefix = block(prefix, positions)
        return prefix, positions

    def _pool_prefix(self, prefix: torch.Tensor) -> torch.Tensor:
        register_count = self.model_config.register_tokens
        if register_count > 0:
            return prefix[:, :register_count].mean(dim=1)
        return prefix.mean(dim=1)

    def _condition(
        self,
        prefix: torch.Tensor,
        r: torch.Tensor,
        t: torch.Tensor,
    ) -> torch.Tensor:
        time = self.time_projection(
            torch.cat(
                (
                    self.t_embedding(t),
                    self.r_embedding(r),
                    self.delta_embedding(t - r),
                ),
                dim=-1,
            )
        )
        return time + self.prefix_condition(self._pool_prefix(prefix))

    def build_memory_cache(
        self,
        history: torch.Tensor,
        context: torch.Tensor | None = None,
        history_actions: torch.Tensor | None = None,
    ) -> PrefixCache:
        prefix, positions = self._encode_prefix(history, context, history_actions)
        pooled = self.prefix_condition(self._pool_prefix(prefix))
        key_values: list[tuple[torch.Tensor, torch.Tensor]] = []
        for block in self.blocks:
            prefix = block.advance_prefix(prefix, positions)
            key_values.append(block.prefix_key_value(prefix, positions))
        return PrefixCache(tuple(key_values), pooled)

    def predict_outcomes(
        self,
        history: torch.Tensor,
        actions: torch.Tensor | None = None,
        frame_count: int = 1,
        memory_cache: PrefixCache | None = None,
        history_actions: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if memory_cache is None:
            prefix, _ = self._encode_prefix(history, None, history_actions)
            pooled = self.prefix_condition(self._pool_prefix(prefix))
        else:
            pooled = memory_cache.pooled
        hidden = pooled[:, None].expand(-1, frame_count, -1)
        if self.action_projection is not None:
            if actions is None or actions.shape[1] != frame_count:
                raise ValueError("Outcome actions must match frame_count.")
            hidden = hidden + self.action_projection(actions)
        hidden = self.outcome_norm(hidden)
        return self.reward_head(hidden).squeeze(-1), self.terminal_head(hidden).squeeze(-1)

    def _unpatchify(
        self,
        patches: torch.Tensor,
        frames: int,
        height: int,
        width: int,
    ) -> torch.Tensor:
        batch = patches.shape[0]
        patch = self.model_config.target_patch_size
        patches = patches.reshape(
            batch,
            frames,
            height,
            width,
            patch,
            patch,
            self.latent_channels,
        )
        return patches.permute(0, 1, 6, 2, 4, 3, 5).reshape(
            batch,
            frames,
            self.latent_channels,
            height * patch,
            width * patch,
        )

    def forward(
        self,
        noisy_target: torch.Tensor,
        history: torch.Tensor,
        r: torch.Tensor,
        t: torch.Tensor,
        actions: torch.Tensor | None = None,
        context: torch.Tensor | None = None,
        memory_cache: PrefixCache | None = None,
        history_actions: torch.Tensor | None = None,
    ) -> torch.Tensor:
        target_volume = self.target_patch_embed(noisy_target.transpose(1, 2))
        batch, _, frames, height, width = target_volume.shape
        target = target_volume.flatten(2).transpose(1, 2)
        target_positions = _video_positions(
            frames,
            height,
            width,
            self.history_frames,
            noisy_target.device,
        )
        if memory_cache is None:
            prefix, prefix_positions = self._encode_prefix(
                history,
                context,
                history_actions,
            )
            condition = self._condition(prefix, r, t)
        else:
            prefix = None
            prefix_positions = None
            condition = self.time_projection(
                torch.cat(
                    (
                        self.t_embedding(t),
                        self.r_embedding(r),
                        self.delta_embedding(t - r),
                    ),
                    dim=-1,
                )
            ) + memory_cache.pooled
        if self.action_projection is not None:
            if actions is None or actions.shape[:2] != (batch, frames):
                raise ValueError("Actions must match [batch, target_frames, action_dim].")
            action_embedding = self.action_projection(actions)
            target = target + action_embedding.repeat_interleave(height * width, dim=1)
            condition = condition + action_embedding.mean(dim=1)
        for block_index, block in enumerate(self.blocks):
            modulation = (
                self.shared_modulation(condition, block_index)
                if self.shared_modulation is not None
                else None
            )
            if memory_cache is None:
                if prefix is None or prefix_positions is None:
                    raise RuntimeError("Prefix state is unavailable.")
                if self.training and self.gradient_checkpointing:
                    prefix, target = checkpoint(
                        block,
                        prefix,
                        target,
                        condition,
                        prefix_positions,
                        target_positions,
                        modulation,
                        use_reentrant=False,
                    )
                else:
                    prefix, target = block(
                        prefix,
                        target,
                        condition,
                        prefix_positions,
                        target_positions,
                        modulation,
                    )
            else:
                target = block.advance_target(
                    target,
                    condition,
                    target_positions,
                    memory_cache.key_values[block_index],
                    modulation,
                )
        shift, scale = self.output_modulation(F.silu(condition)).chunk(2, dim=-1)
        target = self.output_norm(target) * (1.0 + scale[:, None]) + shift[:, None]
        return self._unpatchify(
            self.output_projection(target),
            frames,
            height,
            width,
        )


CausalVideoDiT = CausalLatentVideoDiT