# 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.modeling_utils import ModelMixin
from diffusers.utils import BaseOutput


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 MotionBricksPoseDecoderOutput(BaseOutput):
    """
    Args:
        features (`torch.Tensor` of shape `(batch_size, num_frames, feature_dim)`):
            The normalized features of every frame.
    """

    features: torch.Tensor


class MotionBricksPoseDecoder(ModelMixin, ConfigMixin):
    """Decoder of the pose VQ-VAE: looks up the codes of every token and decodes them to normalized frame features,
    conditioned on the given poses and the root velocity.

    The features of a frame are the local root (heading rate, ground velocity, height) followed by the body: joint
    positions relative to the root, 6D global rotations, joint velocities and foot contacts. Their normalization
    statistics and the neutral joint positions of the skeleton are stored as buffers.
    """

    _no_split_modules = ["MotionBricksConditionedDecoder"]
    _repeated_blocks = ["MotionBricksResnetBlock"]

    @register_to_config
    def __init__(
        self,
        num_heads_codebook: int = 5,
        codebook_size: int = 16,
        code_embed_dim: int = 32,
        pose_dim: int = 268,
        velocity_dim: int = 2,
        feature_dim: int = 365,
        num_joints: int = 30,
        width: int = 512,
        depth: int = 4,
        dilation_growth_rate: int = 3,
        frames_per_token: int = 4,
        max_tokens: int = 16,
    ):
        super().__init__()
        self.codebook = nn.Parameter(torch.zeros(num_heads_codebook, codebook_size, code_embed_dim))
        self.decoder = MotionBricksConditionedDecoder(
            num_heads_codebook * code_embed_dim,
            feature_dim,
            target_channels=pose_dim,
            external_channels=velocity_dim,
            width=width,
            depth=depth,
            dilation_growth_rate=dilation_growth_rate,
            num_tokens=max_tokens,
            num_frames=max_tokens * frames_per_token,
        )
        self.register_buffer("feature_mean", torch.zeros(feature_dim))
        self.register_buffer("feature_std", torch.ones(feature_dim))
        self.register_buffer("neutral_joints", torch.zeros(num_joints, 3))

    def denormalize(self, features: torch.Tensor) -> torch.Tensor:
        # The buffers can sit on another device while the model is offloaded.
        std = self.feature_std.to(features.device, features.dtype)
        return features * std + self.feature_mean.to(features.device, features.dtype)

    def forward(
        self,
        codes: torch.LongTensor,
        num_tokens: torch.Tensor,
        pose: torch.Tensor,
        has_pose: torch.Tensor,
        velocity: torch.Tensor,
        return_dict: bool = True,
    ) -> MotionBricksPoseDecoderOutput | tuple:
        """
        Args:
            codes (`torch.LongTensor` of shape `(batch_size, max_tokens, num_heads_codebook)`):
                The pose code of every token and head.
            num_tokens (`torch.LongTensor` of shape `(batch_size,)`):
                The motion length in tokens.
            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.
            velocity (`torch.Tensor` of shape `(batch_size, num_frames, velocity_dim)`):
                Normalized ground velocity of the root.
            return_dict (`bool`, defaults to `True`):
                Whether to return a `MotionBricksPoseDecoderOutput` instead of a tuple.

        Returns:
            `MotionBricksPoseDecoderOutput` or `tuple`: The normalized frame features; a one-element tuple when
            `return_dict` is `False`.
        """
        heads = torch.arange(self.config.num_heads_codebook, device=codes.device)
        embeds = self.codebook[heads, codes].flatten(2)
        token_mask = (torch.arange(codes.shape[1], device=codes.device)[None] < num_tokens[:, None]).to(embeds.dtype)
        features = self.decoder(embeds.transpose(1, 2), token_mask, pose, has_pose, velocity).transpose(1, 2)

        if not return_dict:
            return (features,)
        return MotionBricksPoseDecoderOutput(features=features)