# 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]


@dataclass
class MotionBricksPoseTransformerOutput(BaseOutput):
    """
    Args:
        logits (`torch.Tensor` of shape `(batch_size, max_tokens, num_heads_codebook, codebook_size)`):
            The code logits of every token and head.
    """

    logits: torch.Tensor


class MotionBricksPoseTransformer(ModelMixin, ConfigMixin, AttentionMixin):
    """Masked token transformer that predicts the pose codes of every token in one pass, from the given poses, the
    root of every frame and the length of the motion. Each token covers `frames_per_token` frames and has one code
    per head."""

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

    @register_to_config
    def __init__(
        self,
        pose_dim: int = 268,
        root_dim: int = 4,
        num_heads_codebook: int = 5,
        codebook_size: int = 16,
        code_embed_dim: int = 32,
        token_dim: int = 640,
        root_embed_dim: int = 256,
        length_embed_dim: int = 128,
        hidden_dim: int = 1024,
        num_heads: int = 16,
        ff_dim: int = 2048,
        num_layers: int = 16,
        frames_per_token: int = 4,
        min_tokens: int = 6,
        max_tokens: int = 16,
    ):
        super().__init__()
        code_dim = num_heads_codebook * code_embed_dim
        # Each head has codebook_size codes and a mask code.
        self.code_embed = nn.Embedding(num_heads_codebook * (codebook_size + 1), code_embed_dim)
        self.code_proj = MotionBricksMLP(code_dim, token_dim, token_dim)
        self.pose_proj = nn.Linear(pose_dim, token_dim // frames_per_token)
        self.root_proj = nn.Linear(root_dim * frames_per_token, root_embed_dim)
        self.length_embed = nn.Embedding(max_tokens - min_tokens + 1, length_embed_dim)
        self.proj_in = nn.Linear(token_dim + root_embed_dim + length_embed_dim, hidden_dim)
        self.position_embed = nn.Parameter(torch.zeros(max_tokens, hidden_dim))
        self.layers = nn.ModuleList(
            [MotionBricksTransformerLayer(hidden_dim, num_heads, ff_dim) for _ in range(num_layers)]
        )
        self.proj_out = nn.Linear(hidden_dim, num_heads_codebook * codebook_size)

    def forward(
        self,
        pose: torch.Tensor,
        has_pose: torch.Tensor,
        root: torch.Tensor,
        num_tokens: torch.Tensor,
        return_dict: bool = True,
    ) -> MotionBricksPoseTransformerOutput | tuple:
        """
        Args:
            pose (`torch.Tensor` of shape `(batch_size, num_frames, pose_dim)`):
                Normalized pose features on the frame timeline.
            has_pose (`torch.BoolTensor` of shape `(batch_size, num_frames)`):
                The frames whose pose is given.
            root (`torch.Tensor` of shape `(batch_size, num_frames, root_dim)`):
                Normalized root of every frame: ground position and the cosine and sine of the heading.
            num_tokens (`torch.LongTensor` of shape `(batch_size,)`):
                The motion length in tokens.
            return_dict (`bool`, defaults to `True`):
                Whether to return a `MotionBricksPoseTransformerOutput` instead of a tuple.

        Returns:
            `MotionBricksPoseTransformerOutput` or `tuple`: The code logits; a one-element tuple when `return_dict` is
            `False`.
        """
        config = self.config
        batch_size = pose.shape[0]
        heads, size = config.num_heads_codebook, config.codebook_size
        mask_codes = size + torch.arange(heads, device=pose.device) * (size + 1)
        codes = self.code_embed(mask_codes).flatten()
        codes = self.code_proj(codes).expand(batch_size, config.max_tokens, -1)
        codes = codes.reshape(batch_size, -1, self.pose_proj.out_features)
        has_pose = has_pose[..., None].to(pose.dtype)
        frames = self.pose_proj(pose) * has_pose + codes * (1 - has_pose)
        frames = frames.reshape(batch_size, config.max_tokens, -1)
        root = self.root_proj(root.reshape(batch_size, config.max_tokens, -1))
        length = self.length_embed(num_tokens - config.min_tokens)[:, None].expand(-1, config.max_tokens, -1)
        hidden_states = F.relu(self.proj_in(torch.cat([frames, root, length], dim=-1))) + self.position_embed
        mask = key_mask(num_tokens, config.max_tokens)
        for layer in self.layers:
            hidden_states = layer(hidden_states, mask)
        logits = self.proj_out(hidden_states).unflatten(-1, (heads, size))

        if not return_dict:
            return (logits,)
        return MotionBricksPoseTransformerOutput(logits=logits)