"""MUGEN: a unified motion-language model.

The model has three parts:

1. An adaptive-length motion autoencoder (ALAE). Cross-attention compresses a
   clip of any length into K continuous latent slots, and a second
   cross-attention stack expands those slots back to any requested number of
   frames. There is no codebook and no quantisation anywhere in the pipeline.
2. A GPT-2 language model. For text-to-motion it rolls out K hidden states
   after a `<MOT>` seed token; for captioning it reads the K slots back as
   continuous input embeddings.
3. A layer router and a calibrated latent head. The router gives each slot its
   own soft mixture over all twelve transformer layers, so a slot reads from
   the depth it needs instead of the final layer only. The head predicts a
   joint Gaussian over the whole flattened latent set with a rank-64 plus
   diagonal covariance, so one draw carries text-conditional variance that is
   correlated across slots.

Generating a motion therefore costs K language-model steps, one draw, and one
decoder pass. No iterative refinement, no residual stages, no denoising chain.

Usage:
    ```python
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    model = AutoModelForCausalLM.from_pretrained("zy22b/MUGEN", trust_remote_code=True).eval()
    tokenizer = AutoTokenizer.from_pretrained("zy22b/MUGEN")

    feats = model.generate_motion(["a person walks forward and waves."],
                                  lengths=[120], tokenizer=tokenizer)
    joints = model.features_to_joints(feats)          # (1, 120, 22, 3)
    caption = model.generate_caption(feats, tokenizer=tokenizer)
    ```
"""

from __future__ import annotations

import math
from typing import List, Optional, Sequence, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import GPT2Config, GPT2LMHeadModel, PreTrainedModel
from transformers.modeling_outputs import ModelOutput

try:
    from .configuration_mugen import MugenConfig
except ImportError:  # standalone / trust_remote_code loading
    from configuration_mugen import MugenConfig


# =====================================================================
# Motion autoencoder building blocks
# =====================================================================

class Swish(nn.Module):
    """x * sigmoid(x)."""

    def forward(self, x):
        return x * torch.sigmoid(x)


def _activation(name: str) -> nn.Module:
    if name == "relu":
        return nn.ReLU()
    if name == "gelu":
        return nn.GELU()
    if name == "silu":
        return Swish()
    raise ValueError(f"Unsupported activation: {name}")


class ResConv1DBlock(nn.Module):
    """Dilated residual convolution over the time axis."""

    def __init__(self, n_in, n_state, dilation=1, activation="gelu", norm=None):
        super().__init__()
        padding = dilation
        self.norm = norm
        if norm == "LN":
            self.norm1 = nn.LayerNorm(n_in)
            self.norm2 = nn.LayerNorm(n_in)
        elif norm == "GN":
            self.norm1 = nn.GroupNorm(32, n_in, eps=1e-6, affine=True)
            self.norm2 = nn.GroupNorm(32, n_in, eps=1e-6, affine=True)
        elif norm == "BN":
            self.norm1 = nn.BatchNorm1d(n_in, eps=1e-6, affine=True)
            self.norm2 = nn.BatchNorm1d(n_in, eps=1e-6, affine=True)
        else:
            self.norm1 = nn.Identity()
            self.norm2 = nn.Identity()

        self.activation1 = _activation(activation)
        self.activation2 = _activation(activation)
        self.conv1 = nn.Conv1d(n_in, n_state, 3, 1, padding, dilation)
        self.conv2 = nn.Conv1d(n_state, n_in, 1, 1, 0)

    def forward(self, x):
        x_orig = x
        if self.norm == "LN":
            x = self.norm1(x.transpose(-2, -1))
            x = self.activation1(x.transpose(-2, -1))
        else:
            x = self.norm1(x)
            x = self.activation1(x)
        x = self.conv1(x)
        if self.norm == "LN":
            x = self.norm2(x.transpose(-2, -1))
            x = self.activation2(x.transpose(-2, -1))
        else:
            x = self.norm2(x)
            x = self.activation2(x)
        x = self.conv2(x)
        return x + x_orig


class Resnet1D(nn.Module):
    """Stack of dilated residual convolutions."""

    def __init__(self, n_in, n_depth, dilation_growth_rate=1,
                 reverse_dilation=True, activation="gelu", norm=None):
        super().__init__()
        blocks = [
            ResConv1DBlock(n_in, n_in, dilation=dilation_growth_rate ** depth,
                           activation=activation, norm=norm)
            for depth in range(n_depth)
        ]
        if reverse_dilation:
            blocks = blocks[::-1]
        self.model = nn.Sequential(*blocks)

    def forward(self, x):
        return self.model(x)


class CrossAttentionBlock(nn.Module):
    """Pre-norm self-attention, cross-attention and feed-forward block.

    `memory_key_padding_mask` (True marks padding) is forwarded to the
    cross-attention. Passing `None` reproduces the unmasked block exactly.
    """

    def __init__(self, d_model, nhead, dim_feedforward, dropout=0.1, activation="gelu"):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
        self.cross_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        self.dropout = nn.Dropout(dropout)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.drop1 = nn.Dropout(dropout)
        self.drop2 = nn.Dropout(dropout)
        self.drop3 = nn.Dropout(dropout)
        self.activation = _activation(activation)

    def forward(self, target, memory, memory_key_padding_mask=None):
        x = self.norm1(target)
        x_sa, _ = self.self_attn(x, x, x, need_weights=False)
        target = target + self.drop1(x_sa)

        x = self.norm2(target)
        x_ca, _ = self.cross_attn(
            x, memory, memory, need_weights=False,
            key_padding_mask=memory_key_padding_mask,
        )
        target = target + self.drop2(x_ca)

        x = self.norm3(target)
        x_ff = self.linear2(self.dropout(self.activation(self.linear1(x))))
        target = target + self.drop3(x_ff)
        return target


class MotionEncoderBackbone(nn.Module):
    """Per-frame convolutional trunk feeding the cross-attention encoder."""

    def __init__(self, input_dim, hidden_dim, latent_dim, depth,
                 dilation_growth_rate, activation, norm, num_res_blocks):
        super().__init__()
        blocks = [
            nn.Conv1d(input_dim, hidden_dim, kernel_size=3, stride=1, padding=1),
            _activation(activation),
        ]
        for _ in range(num_res_blocks):
            blocks.append(Resnet1D(hidden_dim, depth, dilation_growth_rate,
                                   reverse_dilation=False, activation=activation, norm=norm))
        self.backbone = nn.Sequential(*blocks)
        self.output_proj = nn.Conv1d(hidden_dim, latent_dim, kernel_size=1)

    def forward(self, x):
        return self.output_proj(self.backbone(x))


class MotionDecoderRefiner(nn.Module):
    """Convolutional refiner turning decoded per-frame features into motion."""

    def __init__(self, latent_dim, hidden_dim, output_dim, depth,
                 dilation_growth_rate, activation, norm, num_res_blocks):
        super().__init__()
        blocks = [
            nn.Conv1d(latent_dim, hidden_dim, kernel_size=3, stride=1, padding=1),
            _activation(activation),
        ]
        for _ in range(num_res_blocks):
            blocks.append(Resnet1D(hidden_dim, depth, dilation_growth_rate,
                                   reverse_dilation=True, activation=activation, norm=norm))
        blocks.extend([
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, stride=1, padding=1),
            _activation(activation),
            nn.Conv1d(hidden_dim, output_dim, kernel_size=3, stride=1, padding=1),
        ])
        self.model = nn.Sequential(*blocks)

    def forward(self, x):
        return self.model(x)


def build_sine_position_encoding(length, dim, device, dtype=torch.float32):
    pe = torch.zeros(length, dim, device=device, dtype=dtype)
    position = torch.arange(length, device=device, dtype=dtype).unsqueeze(1)
    div_term = torch.exp(
        torch.arange(0, dim, 2, device=device, dtype=dtype) * (-math.log(10000.0) / dim)
    )
    pe[:, 0::2] = torch.sin(position * div_term)
    pe[:, 1::2] = torch.cos(position * div_term)
    return pe


class AdaptiveLengthAutoEncoder(nn.Module):
    """Compress a clip of any length into K latent slots, and back again.

    `encode` returns `(batch, K, latent_dim)` regardless of the input length.
    `decode` takes those slots plus a target frame count. Decoder query
    positions are spread evenly across the whole positional table, so a query
    encodes a frame's relative phase within the clip rather than an absolute
    frame index. That is what lets one decoder serve every clip length.
    """

    def __init__(self, config: MugenConfig):
        super().__init__()
        self.input_dim = config.motion_input_dim
        self.k = config.k_latent_slots
        self.latent_dim = config.latent_dim
        self.max_decode_len = config.alae_max_decode_len

        common = dict(
            depth=config.alae_depth,
            dilation_growth_rate=config.alae_dilation_growth_rate,
            activation=config.alae_activation,
            norm=config.alae_norm,
            num_res_blocks=config.alae_num_res_blocks,
        )
        self.encoder_backbone = MotionEncoderBackbone(
            input_dim=config.motion_input_dim,
            hidden_dim=config.alae_hidden_dim,
            latent_dim=config.latent_dim,
            **common,
        )
        self.latent_queries = nn.Parameter(torch.randn(self.k, config.latent_dim) * 0.02)

        block_kwargs = dict(
            d_model=config.latent_dim,
            nhead=config.alae_nhead,
            dim_feedforward=config.alae_dim_feedforward,
            dropout=config.alae_dropout,
            activation=config.alae_activation,
        )
        self.encoder_blocks = nn.ModuleList(
            [CrossAttentionBlock(**block_kwargs) for _ in range(config.alae_num_encoder_layers)]
        )
        self.encoder_norm = nn.LayerNorm(config.latent_dim)

        self.decoder_blocks = nn.ModuleList(
            [CrossAttentionBlock(**block_kwargs) for _ in range(config.alae_num_decoder_layers)]
        )
        self.decoder_norm = nn.LayerNorm(config.latent_dim)
        self.decoder_refiner = MotionDecoderRefiner(
            latent_dim=config.latent_dim,
            hidden_dim=config.alae_hidden_dim,
            output_dim=config.motion_input_dim,
            **common,
        )

        self.register_buffer(
            "decoder_query_pe",
            build_sine_position_encoding(self.max_decode_len, config.latent_dim,
                                         device=torch.device("cpu")),
            persistent=False,
        )

    def ensure_decode_capacity(self, target_len: int):
        target_len = int(target_len)
        if target_len <= self.max_decode_len:
            return
        self.max_decode_len = target_len
        self.register_buffer(
            "decoder_query_pe",
            build_sine_position_encoding(target_len, self.latent_dim,
                                         device=torch.device("cpu")),
            persistent=False,
        )

    def encode(self, motion: torch.Tensor) -> torch.Tensor:
        """`(batch, time, input_dim)` -> `(batch, K, latent_dim)`."""
        x_in = motion.permute(0, 2, 1).float()
        memory = self.encoder_backbone(x_in).permute(0, 2, 1)
        queries = self.latent_queries.unsqueeze(0).expand(memory.shape[0], -1, -1)
        latents = queries
        for block in self.encoder_blocks:
            latents = block(latents, memory)
        return self.encoder_norm(latents)

    def decode(self, latents: torch.Tensor, target_len: int) -> torch.Tensor:
        """`(batch, K, latent_dim)` -> `(batch, target_len, input_dim)`."""
        target_len = int(target_len)
        self.ensure_decode_capacity(target_len)
        if target_len > 1:
            idx = torch.linspace(0, self.max_decode_len - 1, target_len,
                                 device=latents.device).round().long()
        else:
            idx = torch.zeros(target_len, device=latents.device, dtype=torch.long)
        query_pe = self.decoder_query_pe.to(device=latents.device, dtype=latents.dtype)[idx]
        decoded = query_pe.unsqueeze(0).expand(latents.shape[0], -1, -1)
        for block in self.decoder_blocks:
            decoded = block(decoded, latents)
        decoded = self.decoder_norm(decoded).permute(0, 2, 1)
        return self.decoder_refiner(decoded).permute(0, 2, 1)

    def forward(self, motion: torch.Tensor, target_len: Optional[int] = None):
        if target_len is None:
            target_len = motion.shape[1]
        latents = self.encode(motion)
        return self.decode(latents, target_len), latents


# =====================================================================
# Quaternion helpers for joint recovery
# =====================================================================

def _qinv(q: torch.Tensor) -> torch.Tensor:
    mask = torch.ones_like(q)
    mask[..., 1:] = -mask[..., 1:]
    return q * mask


def _qrot(q: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
    original_shape = list(v.shape)
    q = q.contiguous().view(-1, 4)
    v = v.contiguous().view(-1, 3)
    qvec = q[:, 1:]
    uv = torch.cross(qvec, v, dim=1)
    uuv = torch.cross(qvec, uv, dim=1)
    return (v + 2 * (q[:, :1] * uv + uuv)).view(original_shape)


def recover_from_ric(features: torch.Tensor, num_joints: int) -> torch.Tensor:
    """HumanML3D features -> joint positions.

    `features` are the DENORMALISED 263-dimensional vectors; the return value
    is `(..., time, num_joints, 3)` in metres.
    """
    rot_vel = features[..., 0]
    r_rot_ang = torch.zeros_like(rot_vel)
    r_rot_ang[..., 1:] = rot_vel[..., :-1]
    r_rot_ang = torch.cumsum(r_rot_ang, dim=-1)

    r_rot_quat = torch.zeros(features.shape[:-1] + (4,),
                             dtype=features.dtype, device=features.device)
    r_rot_quat[..., 0] = torch.cos(r_rot_ang)
    r_rot_quat[..., 2] = torch.sin(r_rot_ang)

    r_pos = torch.zeros(features.shape[:-1] + (3,),
                        dtype=features.dtype, device=features.device)
    r_pos[..., 1:, [0, 2]] = features[..., :-1, 1:3]
    r_pos = _qrot(_qinv(r_rot_quat), r_pos)
    r_pos = torch.cumsum(r_pos, dim=-2)
    r_pos[..., 1] = features[..., 3]

    positions = features[..., 4:(num_joints - 1) * 3 + 4]
    positions = positions.view(positions.shape[:-1] + (-1, 3))
    positions = _qrot(
        _qinv(r_rot_quat[..., None, :]).expand(positions.shape[:-1] + (4,)), positions
    )
    positions[..., 0] += r_pos[..., 0:1]
    positions[..., 2] += r_pos[..., 2:3]
    return torch.cat([r_pos.unsqueeze(-2), positions], dim=-2)


# =====================================================================
# Model
# =====================================================================

class MugenOutput(ModelOutput):
    """Output of :meth:`MugenForConditionalGeneration.forward`.

    Attributes:
        latents: `(batch, K, latent_dim)` slots, in raw autoencoder space.
        reconstruction: `(batch, time, motion_input_dim)` when a motion was given.
        logits: language-model logits when `input_ids` was given.
    """

    latents: Optional[torch.FloatTensor] = None
    reconstruction: Optional[torch.FloatTensor] = None
    logits: Optional[torch.FloatTensor] = None


class MugenPreTrainedModel(PreTrainedModel):
    config_class = MugenConfig
    base_model_prefix = "mugen"
    supports_gradient_checkpointing = False

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, (nn.Linear, nn.Conv1d)):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)


class MugenForConditionalGeneration(MugenPreTrainedModel):
    """MUGEN: text-to-motion generation and motion-to-text understanding.

    The public API is three methods:

    * :meth:`generate_motion` turns text into motion features.
    * :meth:`generate_caption` turns motion into text.
    * :meth:`encode_motion` / :meth:`decode_motion` expose the latent interface
      both directions share.

    :meth:`features_to_joints` converts motion features into joint positions
    for rendering or measurement.
    """

    _tied_weights_keys = ["language_model.lm_head.weight"]
    _keys_to_ignore_on_load_missing = [r"motion_autoencoder\.decoder_query_pe"]

    def __init__(self, config: MugenConfig):
        super().__init__(config)

        self.k = int(config.k_latent_slots)
        self.latent_dim = int(config.latent_dim)

        # ---- part 1: the motion autoencoder ----
        self.motion_autoencoder = AdaptiveLengthAutoEncoder(config)

        # ---- part 2: the language model ----
        gpt2_config = GPT2Config(
            vocab_size=config.vocab_size,
            n_positions=config.n_positions,
            n_embd=config.n_embd,
            n_layer=config.n_layer,
            n_head=config.n_head,
            n_inner=config.n_inner,
            activation_function=config.activation_function,
            resid_pdrop=config.resid_pdrop,
            embd_pdrop=config.embd_pdrop,
            attn_pdrop=config.attn_pdrop,
            layer_norm_epsilon=config.layer_norm_epsilon,
            initializer_range=config.initializer_range,
            bos_token_id=config.bos_token_id,
            eos_token_id=config.eos_token_id,
            tie_word_embeddings=True,
        )
        self.language_model = GPT2LMHeadModel(gpt2_config)
        self.n_layers = int(config.n_layer)
        hidden_size = int(config.n_embd)

        # ---- part 3a: the layer router ----
        # A slot's routing logits are a bounded static table plus a bounded
        # text-conditional correction. Both halves are tanh-capped, so no logit
        # margin can grow large enough to saturate the routing softmax.
        self.static_router_logits = nn.Parameter(torch.zeros(self.k, self.n_layers))
        self.text_mem_proj = nn.Linear(hidden_size, self.latent_dim)
        self.stage2_latent_queries = nn.Parameter(torch.zeros(self.k, self.latent_dim))
        self.cross_attn_blocks = nn.ModuleList([
            CrossAttentionBlock(
                d_model=self.latent_dim,
                nhead=config.alae_nhead,
                dim_feedforward=config.alae_dim_feedforward,
                dropout=config.alae_dropout,
                activation=config.alae_activation,
            )
            for _ in range(int(config.num_cross_attn_layers))
        ])
        self.cross_attn_norm = nn.LayerNorm(self.latent_dim)
        self.layer_router = nn.Sequential(
            nn.Linear(self.latent_dim, int(config.router_hidden)),
            nn.GELU(),
            nn.Linear(int(config.router_hidden), self.n_layers),
        )

        # ---- part 3b: the calibrated latent head ----
        self.projector_norm = nn.LayerNorm(hidden_size)
        self.projector = nn.Linear(hidden_size, 2 * self.latent_dim)
        self.feedback_norm = nn.LayerNorm(hidden_size)
        self.latent_low_rank = int(config.latent_low_rank)
        if self.latent_low_rank > 0:
            self.latent_factors = nn.Parameter(
                torch.zeros(self.k * self.latent_dim, self.latent_low_rank)
            )
            self.factor_scale_head = nn.Linear(hidden_size, self.latent_low_rank)

        # ---- the understanding direction ----
        self.motion_in_projector = nn.Linear(self.latent_dim, hidden_size)

        # ---- normalisation statistics ----
        # latent_*: standardise the latent space the head predicts in.
        # feature_*: HumanML3D feature statistics, so the model can hand back
        # denormalised motion without the caller owning the dataset.
        self.register_buffer("latent_mean", torch.zeros(self.k, self.latent_dim))
        self.register_buffer("latent_std", torch.ones(self.k, self.latent_dim))
        self.register_buffer("feature_mean", torch.zeros(config.motion_input_dim))
        self.register_buffer("feature_std", torch.ones(config.motion_input_dim))

        self._tokenizer = None
        self.post_init()

    # ------------------------------------------------------------------
    # plumbing
    # ------------------------------------------------------------------
    def get_input_embeddings(self):
        return self.language_model.get_input_embeddings()

    def set_input_embeddings(self, value):
        self.language_model.set_input_embeddings(value)

    def get_output_embeddings(self):
        return self.language_model.get_output_embeddings()

    def set_output_embeddings(self, new_embeddings):
        self.language_model.set_output_embeddings(new_embeddings)

    def set_tokenizer(self, tokenizer):
        """Cache a tokenizer so the generate helpers can be called without one."""
        self._tokenizer = tokenizer
        return self

    def _get_tokenizer(self, tokenizer=None):
        if tokenizer is not None:
            return tokenizer
        if self._tokenizer is None:
            from transformers import AutoTokenizer
            name = self.config._name_or_path or "zy22b/MUGEN"
            self._tokenizer = AutoTokenizer.from_pretrained(name)
        return self._tokenizer

    # ------------------------------------------------------------------
    # the latent interface both directions share
    # ------------------------------------------------------------------
    @torch.no_grad()
    def encode_motion(self, motion: torch.Tensor, normalized: bool = False) -> torch.Tensor:
        """Motion -> K latent slots.

        Args:
            motion: `(batch, time, 263)` HumanML3D features.
            normalized: `True` if `motion` is already standardised by the
                dataset statistics. `False` (the default) normalises it here.

        Returns:
            `(batch, K, latent_dim)` slots in raw autoencoder space.
        """
        motion = motion.to(self.device, dtype=torch.float32)
        if not normalized:
            motion = (motion - self.feature_mean) / self.feature_std
        return self.motion_autoencoder.encode(motion)

    @torch.no_grad()
    def decode_motion(self, latents: torch.Tensor, length: int,
                      denormalize: bool = True) -> torch.Tensor:
        """K latent slots -> `(batch, length, 263)` motion features."""
        latents = latents.to(self.device, dtype=torch.float32)
        feats = self.motion_autoencoder.decode(latents, target_len=int(length))
        if denormalize:
            feats = feats * self.feature_std + self.feature_mean
        return feats

    def standardize_latents(self, latents: torch.Tensor) -> torch.Tensor:
        return (latents - self.latent_mean) / self.latent_std

    def destandardize_latents(self, latents: torch.Tensor) -> torch.Tensor:
        return latents * self.latent_std + self.latent_mean

    # ------------------------------------------------------------------
    # text -> latents
    # ------------------------------------------------------------------
    def _build_t2m_prompts(self, texts: Sequence[str]) -> List[str]:
        tpl = self.config.t2m_prompt_template
        return [tpl.format(text=t.strip(), mot=self.config.mot_token) for t in texts]

    def _routed_hidden(self, input_ids: torch.Tensor,
                       attention_mask: torch.Tensor) -> torch.Tensor:
        """Roll out K steps and return the routed per-slot hidden states.

        Returns `(batch, K, hidden_size)`. Each slot is a soft mixture over all
        `n_layers` transformer layers of its own rollout step, with mixture
        weights produced by the router from the prompt.
        """
        K = self.k
        device = input_ids.device
        prompt_len = input_ids.shape[1]

        def _positions(start: int, count: int) -> torch.Tensor:
            return torch.arange(start, start + count, dtype=torch.long,
                                device=device).unsqueeze(0)

        out = self.language_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=_positions(0, prompt_len),
            output_hidden_states=True,
            use_cache=True,
            return_dict=True,
        )

        # Router memory is the FULL last-layer token sequence plus its pad mask,
        # so the routing queries can attend to individual words.
        memory = self.text_mem_proj(out.hidden_states[-1])
        pad_mask = attention_mask == 0

        def _stack_all_layers(hidden_states) -> torch.Tensor:
            cols = [h[:, -1:, :] for h in hidden_states[1:]]
            return torch.stack(cols, dim=2)          # (B, 1, L, H)

        all_steps = [_stack_all_layers(out.hidden_states)]
        last_hiddens = [out.hidden_states[-1][:, -1:, :]]
        past = out.past_key_values
        cur_attn = attention_mask

        for step_idx in range(K - 1):
            next_input = self.feedback_norm(last_hiddens[-1])
            cur_attn = torch.cat(
                [cur_attn, torch.ones((cur_attn.shape[0], 1), dtype=cur_attn.dtype, device=device)],
                dim=1,
            )
            step = self.language_model(
                inputs_embeds=next_input,
                attention_mask=cur_attn,
                position_ids=_positions(prompt_len + step_idx, 1),
                past_key_values=past,
                output_hidden_states=True,
                use_cache=True,
                return_dict=True,
            )
            past = step.past_key_values
            all_steps.append(_stack_all_layers(step.hidden_states))
            last_hiddens.append(step.hidden_states[-1][:, -1:, :])

        H_all = torch.cat(all_steps, dim=1)           # (B, K, L, H)
        B = H_all.shape[0]

        latents = self.stage2_latent_queries.unsqueeze(0).expand(B, -1, -1)
        for block in self.cross_attn_blocks:
            latents = block(latents, memory, memory_key_padding_mask=pad_mask)
        latents = self.cross_attn_norm(latents)       # (B, K, D)

        s_s = float(self.config.router_static_scale)
        s_d = float(self.config.router_delta_scale)
        static = s_s * torch.tanh(self.static_router_logits / s_s)      # (K, L)
        delta = s_d * torch.tanh(self.layer_router(latents).float() / s_d)
        logits = static.unsqueeze(0).float() + delta                    # (B, K, L)

        tau = float(self.config.router_eval_tau)
        weights = F.softmax(logits / tau, dim=-1).to(H_all.dtype)
        return torch.einsum("bkl,bklh->bkh", weights, H_all)

    def _latent_distribution(self, hidden: torch.Tensor):
        """Routed hidden states -> `(mu, logvar, log_a)` in standardised space."""
        h = self.projector_norm(hidden)
        mu, logvar = self.projector(h).chunk(2, dim=-1)
        logvar = logvar.clamp(-10.0, 10.0)
        log_a = None
        if self.latent_low_rank > 0:
            log_a = self.factor_scale_head(h.mean(dim=1))               # (B, r)
        return mu, logvar, log_a

    def sample_latents(self, mu: torch.Tensor, logvar: torch.Tensor,
                       log_a: Optional[torch.Tensor], temperature: float,
                       generator: Optional[torch.Generator] = None) -> torch.Tensor:
        """One structured draw `z = mu + t * (U (a * eps1) + sigma * eps2)`.

        The rank-r factor term carries the variance that is correlated across
        slots; the diagonal term carries the rest. `temperature` scales the
        whole zero-mean perturbation, so 0.0 returns `mu` exactly.
        """
        B, K, D = mu.shape
        if temperature <= 0.0:
            return mu

        def _randn(like):
            if generator is None:
                return torch.randn_like(like)
            return torch.randn(like.shape, generator=generator,
                               device=like.device, dtype=like.dtype)

        # The factor noise is drawn BEFORE the diagonal noise, matching the
        # training implementation, so a given seed reproduces the same draw
        # here as it does there.
        factor_part = 0.0
        if log_a is not None:
            a = torch.exp(log_a.clamp(-6.0, 4.0))                       # (B, r)
            factor_part = torch.einsum(
                "dr,br->bd", self.latent_factors, a * _randn(a)
            )
        sigma = torch.exp(0.5 * logvar).reshape(B, K * D)
        diag_part = sigma * _randn(sigma)
        z = mu.reshape(B, K * D) + float(temperature) * (factor_part + diag_part)
        return z.reshape(B, K, D)

    @torch.no_grad()
    def text_to_latents(self, texts: Union[str, Sequence[str]], tokenizer=None,
                        temperature: Optional[float] = None,
                        generator: Optional[torch.Generator] = None) -> torch.Tensor:
        """Text -> K latent slots in raw autoencoder space."""
        if isinstance(texts, str):
            texts = [texts]
        tokenizer = self._get_tokenizer(tokenizer)
        prev_side = tokenizer.padding_side
        tokenizer.padding_side = "left"
        try:
            enc = tokenizer(
                self._build_t2m_prompts(texts),
                return_tensors="pt", padding=True, truncation=True, max_length=256,
            )
        finally:
            tokenizer.padding_side = prev_side
        enc = {k: v.to(self.device) for k, v in enc.items()}

        hidden = self._routed_hidden(enc["input_ids"], enc["attention_mask"])
        mu, logvar, log_a = self._latent_distribution(hidden)
        temp = float(self.config.eval_sample_temperature if temperature is None else temperature)
        z_norm = self.sample_latents(mu, logvar, log_a, temp, generator=generator)
        return self.destandardize_latents(z_norm)

    @torch.no_grad()
    def generate_motion(self, texts: Union[str, Sequence[str]],
                        lengths: Union[int, Sequence[int], None] = None,
                        tokenizer=None,
                        temperature: Optional[float] = None,
                        denormalize: bool = True,
                        generator: Optional[torch.Generator] = None) -> torch.Tensor:
        """Text -> motion features.

        Args:
            texts: one description, or a batch of them.
            lengths: target frame count per description (20 fps). A single int
                applies to the whole batch. Defaults to 120 frames (6 seconds).
                HumanML3D clips are multiples of 4 frames and at most 196, so
                stay in that range to stay in distribution.
            tokenizer: the tokenizer shipped with this repo. Loaded lazily when
                omitted.
            temperature: multiplier on the sampled perturbation. Defaults to
                `config.eval_sample_temperature` (1.0), the calibrated
                conditional distribution. 0.0 decodes the distribution mean,
                which is not the protocol the published numbers use.
            denormalize: return features in raw HumanML3D units. Keep this on
                unless you intend to feed the output back into the encoder.
            generator: optional `torch.Generator` for reproducible draws.

        Returns:
            `(batch, max_length, 263)`. When lengths differ within a batch, the
            tensor is padded with zeros and each row is valid up to its own
            length.
        """
        if isinstance(texts, str):
            texts = [texts]
        texts = list(texts)
        if lengths is None:
            lengths = [120] * len(texts)
        elif isinstance(lengths, int):
            lengths = [lengths] * len(texts)
        lengths = [int(v) for v in lengths]
        if len(lengths) != len(texts):
            raise ValueError(f"got {len(texts)} texts but {len(lengths)} lengths")

        latents = self.text_to_latents(texts, tokenizer=tokenizer,
                                       temperature=temperature, generator=generator)

        max_len = max(lengths)
        if len(set(lengths)) == 1:
            return self.decode_motion(latents, max_len, denormalize=denormalize)

        # Mixed lengths: the decoder's relative-phase queries mean a row must be
        # decoded at its own length, not cropped from the longest one.
        out = latents.new_zeros((len(texts), max_len, self.config.motion_input_dim))
        for i, length in enumerate(lengths):
            out[i, :length] = self.decode_motion(
                latents[i: i + 1], length, denormalize=denormalize
            )[0]
        return out

    # ------------------------------------------------------------------
    # motion -> text
    # ------------------------------------------------------------------
    def _m2t_prompt_embeds(self, latents: torch.Tensor, tokenizer):
        """`[instruction, <MOT>, K motion slots]` as input embeddings."""
        B = latents.shape[0]
        device = latents.device
        wte = self.language_model.get_input_embeddings()

        prefix_ids = torch.tensor(
            tokenizer(self.config.m2t_prompt_prefix, add_special_tokens=False)["input_ids"],
            dtype=torch.long, device=device,
        ).unsqueeze(0).expand(B, -1)
        mot_ids = torch.full((B, 1), int(self.config.mot_token_id),
                             dtype=torch.long, device=device)

        prompt = torch.cat(
            [wte(prefix_ids), wte(mot_ids), self.motion_in_projector(latents)], dim=1
        )
        attn = torch.ones(prompt.shape[:2], dtype=torch.long, device=device)
        return prompt, attn

    @torch.no_grad()
    def generate_caption(self, motion: Optional[torch.Tensor] = None,
                         latents: Optional[torch.Tensor] = None,
                         tokenizer=None, normalized: bool = False,
                         max_new_tokens: Optional[int] = None,
                         num_beams: Optional[int] = None) -> List[str]:
        """Motion -> one caption per clip.

        Pass either `motion` (`(batch, time, 263)` features) or `latents`
        (`(batch, K, latent_dim)`) if you already have the slots. Set
        `normalized=True` when `motion` is already standardised.
        """
        if (motion is None) == (latents is None):
            raise ValueError("pass exactly one of `motion` or `latents`")
        tokenizer = self._get_tokenizer(tokenizer)
        if latents is None:
            latents = self.encode_motion(motion, normalized=normalized)
        latents = latents.to(self.device, dtype=torch.float32)

        prompt, attn = self._m2t_prompt_embeds(latents, tokenizer)
        gen = self.language_model.generate(
            inputs_embeds=prompt,
            attention_mask=attn,
            max_new_tokens=int(self.config.m2t_max_new_tokens
                               if max_new_tokens is None else max_new_tokens),
            do_sample=False,
            num_beams=int(self.config.m2t_num_beams if num_beams is None else num_beams),
            pad_token_id=tokenizer.eos_token_id,
        )
        # Called with inputs_embeds and no input_ids, generate returns only the
        # newly produced tokens.
        return [t.strip() for t in tokenizer.batch_decode(gen, skip_special_tokens=True)]

    # ------------------------------------------------------------------
    # utilities
    # ------------------------------------------------------------------
    def features_to_joints(self, features: torch.Tensor,
                          normalized: bool = False) -> torch.Tensor:
        """Motion features -> `(batch, time, num_joints, 3)` joint positions in metres."""
        features = features.to(self.device, dtype=torch.float32)
        if normalized:
            features = features * self.feature_std + self.feature_mean
        return recover_from_ric(features, int(self.config.num_joints))

    def forward(self,
                motion: Optional[torch.Tensor] = None,
                input_ids: Optional[torch.LongTensor] = None,
                attention_mask: Optional[torch.Tensor] = None,
                target_length: Optional[int] = None,
                normalized: bool = False,
                return_dict: bool = True,
                **kwargs) -> MugenOutput:
        """Autoencode a motion, run the language model, or both.

        This is the plumbing entry point. For generation use
        :meth:`generate_motion` and :meth:`generate_caption`.
        """
        latents = reconstruction = logits = None
        if motion is not None:
            motion = motion.to(self.device, dtype=torch.float32)
            if not normalized:
                motion = (motion - self.feature_mean) / self.feature_std
            latents = self.motion_autoencoder.encode(motion)
            reconstruction = self.motion_autoencoder.decode(
                latents, target_len=int(target_length or motion.shape[1])
            )
        if input_ids is not None:
            logits = self.language_model(
                input_ids=input_ids, attention_mask=attention_mask, return_dict=True
            ).logits
        return MugenOutput(latents=latents, reconstruction=reconstruction, logits=logits)


MugenConfig.register_for_auto_class()
MugenForConditionalGeneration.register_for_auto_class("AutoModelForCausalLM")