# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F

from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.models.attention import AttentionMixin, AttentionModuleMixin
from diffusers.models.attention_dispatch import dispatch_attention_fn
from diffusers.models.modeling_utils import ModelMixin
from diffusers.utils import BaseOutput


class MotionBricksMLP(nn.Module):
    """Two linear layers, each followed by LeakyReLU(0.01), then a linear projection."""

    def __init__(self, in_features: int, hidden_features: int, out_features: int):
        super().__init__()
        self.layers = nn.ModuleList(
            [nn.Linear(in_features, hidden_features), nn.Linear(hidden_features, hidden_features)]
        )
        self.proj_out = nn.Linear(hidden_features, out_features)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        for layer in self.layers:
            hidden_states = F.leaky_relu(layer(hidden_states), 0.01)
        return self.proj_out(hidden_states)


class MotionBricksAttnProcessor:
    _attention_backend = None
    _parallel_config = None

    def __call__(
        self,
        attn: "MotionBricksAttention",
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        batch_size, seq_len, _ = hidden_states.shape
        shape = (batch_size, seq_len, attn.heads, -1)
        query = attn.to_q(hidden_states).view(shape)
        key = attn.to_k(hidden_states).view(shape)
        value = attn.to_v(hidden_states).view(shape)
        hidden_states = dispatch_attention_fn(
            query,
            key,
            value,
            attn_mask=attention_mask,
            backend=self._attention_backend,
            parallel_config=self._parallel_config,
        )
        hidden_states = hidden_states.flatten(2, 3).to(query.dtype)
        return attn.to_out[0](hidden_states)


class MotionBricksAttention(nn.Module, AttentionModuleMixin):
    """Multi-head self-attention. `attention_mask` is a boolean `(batch_size, 1, 1, seq_len)` mask of the keys to
    attend to."""

    _default_processor_cls = MotionBricksAttnProcessor
    _available_processors = [MotionBricksAttnProcessor]

    def __init__(self, dim: int, heads: int):
        super().__init__()
        self.heads = heads
        self.to_q = nn.Linear(dim, dim)
        self.to_k = nn.Linear(dim, dim)
        self.to_v = nn.Linear(dim, dim)
        self.to_out = nn.ModuleList([nn.Linear(dim, dim), nn.Dropout(0.0)])
        self.set_processor(MotionBricksAttnProcessor())

    def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
        return self.processor(self, hidden_states, attention_mask)


class MotionBricksTransformerLayer(nn.Module):
    """Post-norm transformer encoder layer with a ReLU feed-forward."""

    def __init__(self, dim: int, heads: int, ff_dim: int):
        super().__init__()
        self.attn = MotionBricksAttention(dim, heads)
        self.norm1 = nn.LayerNorm(dim)
        self.ff_in = nn.Linear(dim, ff_dim)
        self.ff_out = nn.Linear(ff_dim, dim)
        self.norm2 = nn.LayerNorm(dim)

    def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
        hidden_states = self.norm1(hidden_states + self.attn(hidden_states, attention_mask))
        return self.norm2(hidden_states + self.ff_out(F.relu(self.ff_in(hidden_states))))


def key_mask(num_tokens: torch.Tensor, max_tokens: int, prefix: int = 0) -> torch.Tensor:
    """Boolean `(batch_size, 1, 1, prefix + max_tokens)` attention mask: the `prefix` positions, then the first
    `num_tokens` of `max_tokens` token positions."""
    positions = torch.arange(max_tokens, device=num_tokens.device)
    valid = positions[None] < num_tokens[:, None]
    valid = F.pad(valid, (prefix, 0), value=True)
    return valid[:, None, None]


class MotionBricksConv1d(nn.Conv1d):
    """Stride-1 `nn.Conv1d` computed as one matrix product over the shifted inputs, faster than cuDNN at these
    sizes."""

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        kernel, dilation, padding = self.kernel_size[0], self.dilation[0], self.padding[0]
        length = hidden_states.shape[-1]
        if kernel > 1:
            padded = F.pad(hidden_states, (padding, padding))
            shifted = [padded[..., i * dilation : i * dilation + length] for i in range(kernel)]
            hidden_states = torch.stack(shifted, dim=2).flatten(1, 2)
        weight = self.weight.reshape(self.out_channels, -1)
        return torch.matmul(weight, hidden_states) + self.bias[:, None]


class MotionBricksResnetBlock(nn.Module):
    def __init__(self, channels: int, dilation: int):
        super().__init__()
        self.conv1 = MotionBricksConv1d(channels, channels, 3, padding=dilation, dilation=dilation)
        self.conv2 = MotionBricksConv1d(channels, channels, 1)

    def forward(self, hidden_states: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        residual = hidden_states * mask
        hidden_states = self.conv1(F.relu(residual) * mask)
        hidden_states = self.conv2(F.relu(hidden_states) * mask)
        return (hidden_states + residual) * mask


class MotionBricksUpBlock(nn.Module):
    def __init__(self, channels: int, depth: int, dilation_growth_rate: int):
        super().__init__()
        dilations = [dilation_growth_rate**i for i in reversed(range(depth))]
        self.resnets = nn.ModuleList([MotionBricksResnetBlock(channels, dilation) for dilation in dilations])
        self.conv = MotionBricksConv1d(channels, channels, 3, padding=1)


class MotionBricksConditionedDecoder(nn.Module):
    """Dilated 1D convolution decoder that upsamples tokens to frames, conditioned on per-frame targets and
    external features."""

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        target_channels: int,
        external_channels: int,
        width: int = 512,
        depth: int = 4,
        dilation_growth_rate: int = 3,
        num_tokens: int = 16,
        num_frames: int = 64,
    ):
        super().__init__()
        self.width = width
        self.num_frames = num_frames
        lengths = [num_tokens, num_tokens * 2]
        self.conv_in = MotionBricksConv1d(in_channels, width, 3, padding=1)
        self.target_projections = nn.ModuleList(
            [nn.Linear(target_channels, width * length // num_frames) for length in lengths]
        )
        self.external_projections = nn.ModuleList(
            [nn.Linear(width + external_channels * num_frames // length, width) for length in lengths]
        )
        self.up_blocks = nn.ModuleList([MotionBricksUpBlock(width, depth, dilation_growth_rate) for _ in lengths])
        self.conv_mid = MotionBricksConv1d(width, width, 3, padding=1)
        self.conv_out = MotionBricksConv1d(width, out_channels, 3, padding=1)

    def forward(
        self,
        hidden_states: torch.Tensor,
        token_mask: torch.Tensor,
        target: torch.Tensor,
        target_mask: torch.Tensor,
        external: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            hidden_states: `(batch_size, in_channels, num_tokens)` token features.
            token_mask: `(batch_size, num_tokens)` float mask of the valid tokens.
            target: `(batch_size, num_frames, target_channels)` per-frame targets.
            target_mask: `(batch_size, num_frames)` boolean mask of the frames with a target.
            external: `(batch_size, num_frames, external_channels)` per-frame conditioning.

        Returns:
            `(batch_size, out_channels, num_frames)` frame features.
        """
        batch_size = hidden_states.shape[0]
        mask = token_mask[:, None]
        hidden_states = F.relu(self.conv_in(hidden_states * mask))
        for block, target_projection, external_projection in zip(
            self.up_blocks, self.target_projections, self.external_projections
        ):
            length = hidden_states.shape[-1]
            projected = F.relu(target_projection(target))
            grouped = hidden_states.transpose(1, 2).reshape(batch_size, self.num_frames, -1)
            grouped = torch.where(target_mask[..., None], projected, grouped)
            grouped = grouped.reshape(batch_size, length, self.width)
            conditioned = torch.cat([grouped, external.reshape(batch_size, length, -1)], dim=-1)
            hidden_states = F.relu(external_projection(conditioned).transpose(1, 2))
            for resnet in block.resnets:
                hidden_states = resnet(hidden_states, mask)
            hidden_states = F.interpolate(hidden_states, scale_factor=2, mode="nearest")
            mask = mask.repeat_interleave(2, dim=-1)
            hidden_states = block.conv(hidden_states * mask)
        hidden_states = F.relu(self.conv_mid(hidden_states * mask)) * mask
        return self.conv_out(hidden_states) * mask


@dataclass
class MotionBricksRootOutput(BaseOutput):
    """
    Args:
        num_tokens (`torch.LongTensor` of shape `(batch_size,)`):
            The chosen motion length in tokens; the motion spans `num_tokens * frames_per_token` frames.
        root (`torch.Tensor` of shape `(batch_size, num_frames, 5)`):
            The normalized global root of every frame: position, then cosine and sine of the heading.
    """

    num_tokens: torch.LongTensor
    root: torch.Tensor


class MotionBricksRootModel(ModelMixin, ConfigMixin, AttentionMixin):
    """Chooses the length of a transition and predicts its root trajectory from `context_frames` context poses and
    `target_frames` target poses. Stores the normalization statistics of the guide features as buffers."""

    _no_split_modules = ["MotionBricksTransformerLayer", "MotionBricksConditionedDecoder"]
    _repeated_blocks = ["MotionBricksTransformerLayer"]
    _skip_layerwise_casting_patterns = ["norm", "embed"]

    @register_to_config
    def __init__(
        self,
        pose_dim: int = 268,
        global_root_dim: int = 5,
        local_root_dim: int = 4,
        pose_embed_dim: int = 256,
        root_embed_dim: int = 64,
        hidden_dim: int = 512,
        num_heads: int = 16,
        ff_dim: int = 2048,
        num_shared_layers: int = 6,
        num_token_layers: int = 2,
        context_frames: int = 4,
        target_frames: int = 4,
        frames_per_token: int = 4,
        min_tokens: int = 6,
        max_tokens: int = 16,
        decoder_depth: int = 4,
        decoder_dilation_growth_rate: int = 3,
    ):
        super().__init__()
        guide_frames = context_frames + target_frames
        num_lengths = max_tokens - min_tokens + 1
        guide_dim = pose_embed_dim + 2 * root_embed_dim

        self.pose_proj = nn.Linear(pose_dim, pose_embed_dim)
        self.local_root_proj = nn.Linear(local_root_dim, root_embed_dim)
        self.global_root_proj = nn.Linear(global_root_dim, root_embed_dim)
        self.missing_pose_embed = nn.Parameter(torch.zeros(pose_embed_dim))
        self.missing_local_root_embed = nn.Parameter(torch.zeros(root_embed_dim))
        self.missing_global_root_embed = nn.Parameter(torch.zeros(root_embed_dim))
        self.context_proj = MotionBricksMLP(guide_dim, hidden_dim, hidden_dim)
        self.target_proj = MotionBricksMLP(guide_dim, hidden_dim, hidden_dim)
        self.guide_position_embed = nn.Parameter(torch.zeros(guide_frames, hidden_dim))

        # Index num_lengths + 1 asks for a prediction; the length classes are min_tokens..max_tokens and one unused.
        self.length_query_embed = nn.Embedding(num_lengths + 2, hidden_dim)
        self.shared_layers = nn.ModuleList(
            [MotionBricksTransformerLayer(hidden_dim, num_heads, ff_dim) for _ in range(num_shared_layers)]
        )
        self.length_head = nn.Linear(hidden_dim, num_lengths + 1)

        self.length_embed = nn.Embedding(num_lengths + 1, hidden_dim)
        self.token_position_embed = nn.Parameter(torch.zeros(max_tokens, hidden_dim))
        self.token_layers = nn.ModuleList(
            [MotionBricksTransformerLayer(hidden_dim, num_heads, ff_dim) for _ in range(num_token_layers)]
        )

        self.missing_frame_embed = nn.Parameter(torch.zeros(hidden_dim))
        self.decoder = MotionBricksConditionedDecoder(
            hidden_dim,
            global_root_dim,
            target_channels=global_root_dim,
            external_channels=hidden_dim,
            width=hidden_dim,
            depth=decoder_depth,
            dilation_growth_rate=decoder_dilation_growth_rate,
            num_tokens=max_tokens,
            num_frames=max_tokens * frames_per_token,
        )

        self.register_buffer("global_root_mean", torch.zeros(global_root_dim))
        self.register_buffer("global_root_std", torch.ones(global_root_dim))
        self.register_buffer("local_root_mean", torch.zeros(local_root_dim))
        self.register_buffer("local_root_std", torch.ones(local_root_dim))
        self.register_buffer("pose_mean", torch.zeros(pose_dim))
        self.register_buffer("pose_std", torch.ones(pose_dim))

    def _statistics(self, value: torch.Tensor, name: str) -> tuple[torch.Tensor, torch.Tensor]:
        # The buffers can sit on another device while the model is offloaded.
        mean, std = getattr(self, f"{name}_mean"), getattr(self, f"{name}_std")
        return mean.to(value.device, value.dtype), std.to(value.device, value.dtype)

    def normalize(self, value: torch.Tensor, name: str) -> torch.Tensor:
        """Normalize `global_root`, `local_root` or `pose` features with this model's statistics."""
        mean, std = self._statistics(value, name)
        return (value - mean) / std

    def denormalize(self, value: torch.Tensor, name: str) -> torch.Tensor:
        mean, std = self._statistics(value, name)
        return value * std + mean

    def place_guides(self, guides: torch.Tensor, num_tokens: torch.Tensor, fill) -> torch.Tensor:
        """Place `(batch_size, guide_frames, ...)` guide values on the frame timeline: the context poses at the start
        and the target poses at the last frames of the chosen length. Other frames hold `fill`."""
        config = self.config
        shape = (guides.shape[0], config.max_tokens * config.frames_per_token, *guides.shape[2:])
        if isinstance(fill, torch.Tensor):
            timeline = fill.to(guides.dtype).expand(shape).clone()
        else:
            timeline = torch.full(shape, fill, dtype=guides.dtype, device=guides.device)
        timeline[:, : config.context_frames] = guides[:, : config.context_frames]
        end = num_tokens * config.frames_per_token
        frames = end[:, None] - config.target_frames + torch.arange(config.target_frames, device=end.device)
        batch = torch.arange(guides.shape[0], device=guides.device)[:, None]
        timeline[batch, frames] = guides[:, config.context_frames :]
        return timeline

    def forward(
        self,
        global_root: torch.Tensor,
        local_root: torch.Tensor,
        pose: torch.Tensor,
        has_global_root: torch.Tensor,
        has_local_root: torch.Tensor,
        has_pose: torch.Tensor,
        allowed_num_tokens: torch.Tensor,
        return_dict: bool = True,
    ) -> MotionBricksRootOutput | tuple:
        """
        Args:
            global_root (`torch.Tensor` of shape `(batch_size, guide_frames, global_root_dim)`):
                Normalized global root of the guide poses.
            local_root (`torch.Tensor` of shape `(batch_size, guide_frames, local_root_dim)`):
                Normalized local root of the guide poses.
            pose (`torch.Tensor` of shape `(batch_size, guide_frames, pose_dim)`):
                Normalized pose features of the guide poses.
            has_global_root (`torch.BoolTensor` of shape `(batch_size, guide_frames)`):
                Which guide poses have a global root.
            has_local_root (`torch.BoolTensor` of shape `(batch_size, guide_frames)`):
                Which guide poses have a local root.
            has_pose (`torch.BoolTensor` of shape `(batch_size, guide_frames)`):
                Which guide poses have pose features.
            allowed_num_tokens (`torch.BoolTensor` of shape `(batch_size, max_tokens - min_tokens + 1)`):
                The lengths the model may choose, from `min_tokens`.
            return_dict (`bool`, defaults to `True`):
                Whether to return a `MotionBricksRootOutput` instead of a tuple.

        Returns:
            `MotionBricksRootOutput` or `tuple`: The chosen length and the normalized root; a tuple of the two when
            `return_dict` is `False`.
        """
        config = self.config
        has_pose_f = has_pose[..., None].to(pose.dtype)
        has_local_f = has_local_root[..., None].to(pose.dtype)
        has_global_f = has_global_root[..., None].to(pose.dtype)
        pose_embeds = self.pose_proj(pose) * has_pose_f + self.missing_pose_embed * (1 - has_pose_f)
        local_embeds = self.local_root_proj(local_root) * has_local_f + self.missing_local_root_embed * (
            1 - has_local_f
        )
        global_embeds = self.global_root_proj(global_root) * has_global_f + self.missing_global_root_embed * (
            1 - has_global_f
        )
        guides = torch.cat([pose_embeds, local_embeds, global_embeds], dim=-1)
        guides = torch.cat(
            [
                self.context_proj(guides[:, : config.context_frames]),
                self.target_proj(guides[:, config.context_frames :]),
            ],
            dim=1,
        )
        positioned = guides + self.guide_position_embed

        batch_size = guides.shape[0]
        num_lengths = config.max_tokens - config.min_tokens + 1
        query = self.length_query_embed(torch.full((batch_size, 1), num_lengths + 1, device=guides.device))
        hidden_states = torch.cat([query, positioned], dim=1)
        for layer in self.shared_layers:
            hidden_states = layer(hidden_states)
        logits = self.length_head(hidden_states[:, 0])
        allowed = F.pad(allowed_num_tokens, (0, 1), value=False)
        length = logits.masked_fill(~allowed, float("-inf")).argmax(dim=-1)
        num_tokens = length + config.min_tokens

        tokens = self.token_position_embed.expand(batch_size, -1, -1)
        hidden_states = torch.cat([self.length_embed(length)[:, None], positioned, tokens], dim=1)
        mask = key_mask(num_tokens, config.max_tokens, prefix=1 + guides.shape[1])
        for layer in self.token_layers:
            hidden_states = layer(hidden_states, mask)
        tokens = hidden_states[:, -config.max_tokens :]

        has_guide = has_pose | has_local_root | has_global_root
        frame_guides = torch.where(has_guide[..., None], guides, self.missing_frame_embed)
        external = self.place_guides(frame_guides, num_tokens, fill=self.missing_frame_embed)
        target = self.place_guides(global_root, num_tokens, fill=0)
        target_mask = self.place_guides(has_global_root, num_tokens, fill=False)
        token_mask = (torch.arange(config.max_tokens, device=num_tokens.device)[None] < num_tokens[:, None]).to(
            tokens.dtype
        )
        root = self.decoder(tokens.transpose(1, 2), token_mask, target, target_mask, external).transpose(1, 2)

        if not return_dict:
            return (num_tokens, root)
        return MotionBricksRootOutput(num_tokens=num_tokens, root=root)