# Diffusers-derived portions: Attention and AttnProcessor, adapted via F5-TTS from:
# https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py
# https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention.py
# These are modified adaptations for the audio/text model interfaces in this
# file. The following copyright and license notice applies to those portions
# only; other portions retain their respective licenses and attributions.
#
# Copyright 2024 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.
#
# A copy is included at docs/licenses/Apache-2.0.txt in the WaveGen package source.
# See docs/THIRD_PARTY_NOTICES.md for the scope and provenance of these portions.

"""
ein notation:
b - batch
n - sequence
nt - text sequence
nw - raw wave length
d - dimension
"""
# ruff: noqa: F722 F821

from __future__ import annotations

import math
from typing import Optional

import torch
import torch.nn.functional as F
import torchaudio
from torch import nn
from x_transformers.x_transformers import apply_rotary_pos_emb

from wavegen_v_spec.model.utils import is_package_available, lens_to_mask


# sinusoidal position embedding


class SinusPositionEmbedding(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, x, scale=1000):
        device = x.device
        half_dim = self.dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)
        emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)
        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
        return emb


# convolutional position embedding


class ConvPositionEmbedding(nn.Module):
    def __init__(self, dim, kernel_size=31, groups=16):
        super().__init__()
        assert kernel_size % 2 != 0
        self.conv1d = nn.Sequential(
            nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2),
            nn.Mish(),
            nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2),
            nn.Mish(),
        )
        self.layer_need_mask_idx = [i for i, layer in enumerate(self.conv1d) if isinstance(layer, nn.Conv1d)]

    def forward(self, x: float["b n d"], mask: bool["b n"] | None = None):
        if mask is not None:
            mask = mask.unsqueeze(1)  # [B 1 N]
        x = x.permute(0, 2, 1)  # [B D N]

        if mask is not None:
            x = x.masked_fill(~mask, 0.0)
        for i, block in enumerate(self.conv1d):
            x = block(x)
            if mask is not None and i in self.layer_need_mask_idx:
                x = x.masked_fill(~mask, 0.0)

        x = x.permute(0, 2, 1)  # [B N D]

        return x


# rotary positional embedding related


def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, theta_rescale_factor=1.0):
    # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
    # has some connection to NTK literature
    # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/
    # https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py
    theta *= theta_rescale_factor ** (dim / (dim - 2))
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(end, device=freqs.device)  # type: ignore
    freqs = torch.outer(t, freqs).float()  # type: ignore
    freqs_cos = torch.cos(freqs)  # real part
    freqs_sin = torch.sin(freqs)  # imaginary part
    return torch.cat([freqs_cos, freqs_sin], dim=-1)


def get_pos_embed_indices(start, length, max_pos, scale=1.0):
    # length = length if isinstance(length, int) else length.max()
    scale = scale * torch.ones_like(start, dtype=torch.float32)  # in case scale is a scalar
    pos = (
        start.unsqueeze(1)
        + (torch.arange(length, device=start.device, dtype=torch.float32).unsqueeze(0) * scale.unsqueeze(1)).long()
    )
    # avoid extra long error.
    pos = torch.where(pos < max_pos, pos, max_pos - 1)
    return pos


# Global Response Normalization layer (Instance Normalization ?)


class GRN(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.gamma = nn.Parameter(torch.zeros(1, 1, dim))
        self.beta = nn.Parameter(torch.zeros(1, 1, dim))

    def forward(self, x):
        Gx = torch.norm(x, p=2, dim=1, keepdim=True)
        Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
        return self.gamma * (x * Nx) + self.beta + x


# ConvNeXt-V2 Block https://github.com/facebookresearch/ConvNeXt-V2/blob/main/models/convnextv2.py
# ref: https://github.com/bfs18/e2_tts/blob/main/rfwave/modules.py#L108


class ConvNeXtV2Block(nn.Module):
    def __init__(
        self,
        dim: int,
        intermediate_dim: int,
        dilation: int = 1,
    ):
        super().__init__()
        padding = (dilation * (7 - 1)) // 2
        self.dwconv = nn.Conv1d(
            dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation
        )  # depthwise conv
        self.norm = nn.LayerNorm(dim, eps=1e-6)
        self.pwconv1 = nn.Linear(dim, intermediate_dim)  # pointwise/1x1 convs, implemented with linear layers
        self.act = nn.GELU()
        self.grn = GRN(intermediate_dim)
        self.pwconv2 = nn.Linear(intermediate_dim, dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = x.transpose(1, 2)  # b n d -> b d n
        x = self.dwconv(x)
        x = x.transpose(1, 2)  # b d n -> b n d
        x = self.norm(x)
        x = self.pwconv1(x)
        x = self.act(x)
        x = self.grn(x)
        x = self.pwconv2(x)
        return residual + x


# RMSNorm


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
        # Probe the operator directly. Parsing ``torch.__version__[:3]``
        # misclassifies versions such as 2.11 as 2.1 and incorrectly selects
        # the manual fallback even though native RMSNorm is available.
        self.native_rms_norm = hasattr(F, "rms_norm")

    def forward(self, x):
        if self.native_rms_norm:
            # AMP keeps the master parameter in fp32 while q/k activations
            # are bf16. Match their dtypes for native/fused dispatch while
            # retaining gradient accumulation into the fp32 master weight.
            weight = self.weight
            if weight.dtype != x.dtype:
                weight = weight.to(dtype=x.dtype)
            return F.rms_norm(
                x,
                normalized_shape=(x.shape[-1],),
                weight=weight,
                eps=self.eps,
            )

        input_dtype = x.dtype
        variance = x.to(torch.float32).pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(variance + self.eps)
        return (x * self.weight).to(dtype=input_dtype)


# AdaLayerNorm
# return with modulated x for attn input, and params for later mlp modulation


class AdaLayerNorm(nn.Module):
    def __init__(self, dim):
        super().__init__()

        self.silu = nn.SiLU()
        self.linear = nn.Linear(dim, dim * 6)

        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)

    def forward(self, x, emb=None):
        emb = self.linear(self.silu(emb))
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1)

        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp


# AdaLayerNorm for final layer
# return only with modulated x for attn input, cuz no more mlp modulation


class AdaLayerNorm_Final(nn.Module):
    def __init__(self, dim):
        super().__init__()

        self.silu = nn.SiLU()
        self.linear = nn.Linear(dim, dim * 2)

        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)

    def forward(self, x, emb):
        emb = self.linear(self.silu(emb))
        scale, shift = torch.chunk(emb, 2, dim=1)

        x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
        return x


# FeedForward


class FeedForward(nn.Module):
    def __init__(self, dim, dim_out=None, mult=4, dropout=0.0, approximate: str = "none"):
        super().__init__()
        inner_dim = int(dim * mult)
        dim_out = dim_out if dim_out is not None else dim

        activation = nn.GELU(approximate=approximate)
        project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation)
        self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out))

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


# Attention with possible joint part
# modified from diffusers/src/diffusers/models/attention_processor.py


class Attention(nn.Module):
    def __init__(
        self,
        processor: AttnProcessor,
        dim: int,
        heads: int = 8,
        dim_head: int = 64,
        dropout: float = 0.0,
        context_dim: Optional[int] = None,  # if not None -> joint attention
        context_pre_only: bool = False,
        qk_norm: Optional[str] = None,
    ):
        super().__init__()

        if not hasattr(F, "scaled_dot_product_attention"):
            raise ImportError("Attention equires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")

        self.processor = processor

        self.dim = dim
        self.heads = heads
        self.inner_dim = dim_head * heads
        self.dropout = dropout

        self.context_dim = context_dim
        self.context_pre_only = context_pre_only

        self.to_q = nn.Linear(dim, self.inner_dim)
        self.to_k = nn.Linear(dim, self.inner_dim)
        self.to_v = nn.Linear(dim, self.inner_dim)

        if qk_norm is None:
            self.q_norm = None
            self.k_norm = None
        elif qk_norm == "rms_norm":
            self.q_norm = RMSNorm(dim_head, eps=1e-6)
            self.k_norm = RMSNorm(dim_head, eps=1e-6)
        else:
            raise ValueError(f"Unimplemented qk_norm: {qk_norm}")

        if self.context_dim is not None:
            self.to_q_c = nn.Linear(context_dim, self.inner_dim)
            self.to_k_c = nn.Linear(context_dim, self.inner_dim)
            self.to_v_c = nn.Linear(context_dim, self.inner_dim)
            if qk_norm is None:
                self.c_q_norm = None
                self.c_k_norm = None
            elif qk_norm == "rms_norm":
                self.c_q_norm = RMSNorm(dim_head, eps=1e-6)
                self.c_k_norm = RMSNorm(dim_head, eps=1e-6)

        self.to_out = nn.ModuleList([])
        self.to_out.append(nn.Linear(self.inner_dim, dim))
        self.to_out.append(nn.Dropout(dropout))

        if self.context_dim is not None and not self.context_pre_only:
            self.to_out_c = nn.Linear(self.inner_dim, context_dim)

    def forward(
        self,
        x: float["b n d"],  # noised input x
        c: float["b n d"] = None,  # context c
        mask: bool["b n"] | None = None,
        rope=None,  # rotary position embedding for x
        c_rope=None,  # rotary position embedding for c
    ) -> torch.Tensor:
        if c is not None:
            return self.processor(self, x, c=c, mask=mask, rope=rope, c_rope=c_rope)
        else:
            return self.processor(self, x, mask=mask, rope=rope)


# Attention processor

if is_package_available("flash_attn"):
    from flash_attn import flash_attn_func, flash_attn_varlen_func
    from flash_attn.bert_padding import pad_input, unpad_input


class AttnProcessor:
    def __init__(
        self,
        pe_attn_head: int | None = None,  # number of attention head to apply rope, None for all
        attn_backend: str = "torch",  # "torch" or "flash_attn"
        attn_mask_enabled: bool = True,
    ):
        if attn_backend == "flash_attn":
            assert is_package_available("flash_attn"), "Please install flash-attn first."

        self.pe_attn_head = pe_attn_head
        self.attn_backend = attn_backend
        self.attn_mask_enabled = attn_mask_enabled

    def __call__(
        self,
        attn: Attention,
        x: float["b n d"],  # noised input x
        mask: bool["b n"] | None = None,
        rope=None,  # rotary position embedding
    ) -> torch.FloatTensor:
        batch_size = x.shape[0]

        # `sample` projections
        query = attn.to_q(x)
        key = attn.to_k(x)
        value = attn.to_v(x)

        # attention
        inner_dim = key.shape[-1]
        head_dim = inner_dim // attn.heads
        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)

        # qk norm
        if attn.q_norm is not None:
            query = attn.q_norm(query)
        if attn.k_norm is not None:
            key = attn.k_norm(key)

        # apply rotary position embedding
        if rope is not None:
            freqs, xpos_scale = rope
            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)

            if self.pe_attn_head is not None:
                pn = self.pe_attn_head
                query[:, :pn, :, :] = apply_rotary_pos_emb(query[:, :pn, :, :], freqs, q_xpos_scale)
                key[:, :pn, :, :] = apply_rotary_pos_emb(key[:, :pn, :, :], freqs, k_xpos_scale)
            else:
                query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
                key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)

        if self.attn_backend == "torch":
            # mask. e.g. inference got a batch with different target durations, mask out the padding
            if self.attn_mask_enabled and mask is not None:
                attn_mask = mask
                attn_mask = attn_mask.unsqueeze(1).unsqueeze(1)  # 'b n -> b 1 1 n'
                attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2])
            else:
                attn_mask = None
            x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
            x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)

        elif self.attn_backend == "flash_attn":
            query = query.transpose(1, 2)  # [b, h, n, d] -> [b, n, h, d]
            key = key.transpose(1, 2)
            value = value.transpose(1, 2)
            if self.attn_mask_enabled and mask is not None:
                query, indices, q_cu_seqlens, q_max_seqlen_in_batch, _ = unpad_input(query, mask)
                key, _, k_cu_seqlens, k_max_seqlen_in_batch, _ = unpad_input(key, mask)
                value, _, _, _, _ = unpad_input(value, mask)
                x = flash_attn_varlen_func(
                    query,
                    key,
                    value,
                    q_cu_seqlens,
                    k_cu_seqlens,
                    q_max_seqlen_in_batch,
                    k_max_seqlen_in_batch,
                )
                x = pad_input(x, indices, batch_size, q_max_seqlen_in_batch)
                x = x.reshape(batch_size, -1, attn.heads * head_dim)
            else:
                x = flash_attn_func(query, key, value, dropout_p=0.0, causal=False)
                x = x.reshape(batch_size, -1, attn.heads * head_dim)

        x = x.to(query.dtype)

        # linear proj
        x = attn.to_out[0](x)
        # dropout
        x = attn.to_out[1](x)

        if mask is not None:
            mask = mask.unsqueeze(-1)
            x = x.masked_fill(~mask, 0.0)

        return x


# Joint Attention processor for MM-DiT
# modified from diffusers/src/diffusers/models/attention_processor.py




# DiT Block


class DiTBlock(nn.Module):
    def __init__(
        self,
        dim,
        heads,
        dim_head,
        ff_mult=4,
        dropout=0.1,
        qk_norm=None,
        pe_attn_head=None,
        attn_backend="torch",  # "torch" or "flash_attn"
        attn_mask_enabled=True,
    ):
        super().__init__()

        self.attn_norm = AdaLayerNorm(dim)
        self.attn = Attention(
            processor=AttnProcessor(
                pe_attn_head=pe_attn_head,
                attn_backend=attn_backend,
                attn_mask_enabled=attn_mask_enabled,
            ),
            dim=dim,
            heads=heads,
            dim_head=dim_head,
            dropout=dropout,
            qk_norm=qk_norm,
        )

        self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")

    def forward(self, x, t, mask=None, rope=None):  # x: noised input, t: time embedding
        # pre-norm & modulation for attention input
        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t)

        # attention
        attn_output = self.attn(x=norm, mask=mask, rope=rope)

        # process attention output for input x
        x = x + gate_msa.unsqueeze(1) * attn_output

        norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
        ff_output = self.ff(norm)
        x = x + gate_mlp.unsqueeze(1) * ff_output

        return x






# time step conditioning embedding


class TimestepEmbedding(nn.Module):
    def __init__(self, dim, freq_embed_dim=256):
        super().__init__()
        self.time_embed = SinusPositionEmbedding(freq_embed_dim)
        self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim))

    def forward(self, timestep: float["b"]):
        time_hidden = self.time_embed(timestep)
        time_hidden = time_hidden.to(timestep.dtype)
        time = self.time_mlp(time_hidden)  # b d
        return time

# auxiliary loss for mel spectrogram
class MelSpectrogramLoss(nn.Module):
    """
    General PyTorch implementation of the Multi-Scale Mel Spectrogram Loss from DAC.
    Reference: https://github.com/descriptinc/lyrebird-audiotools
    """
    def __init__(
        self,
        sample_rate: int = 24000,
        n_mels: List[int] = [5, 10, 20, 40, 80, 160, 320],
        window_lengths: List[int] = [32, 64, 128, 256, 512, 1024, 2048],
        mel_fmin: List[float] = [0, 0, 0, 0, 0, 0, 0],
        mel_fmax: List[Optional[float]] = [None, None, None, None, None, None, None],
        pow: float = 1.0,
        clamp_eps: float = 1e-5,
        weight: float = 1.0,
    ):
        super().__init__()
        self.sample_rate = sample_rate
        self.weight = weight
        self.pow = pow
        self.clamp_eps = clamp_eps

        
        assert len(n_mels) == len(window_lengths), "n_mels and window_lengths must have the same length"
        
        self.mel_transforms = nn.ModuleList()
        
        for i, (nm, wl) in enumerate(zip(n_mels, window_lengths)):
            fmin = mel_fmin[i] if i < len(mel_fmin) else 0.0
            fmax = mel_fmax[i] if i < len(mel_fmax) else None
            
            # DAC default logic: hop_length = window_length // 4
            hop_length = wl // 4
            
            transform = torchaudio.transforms.MelSpectrogram(
                sample_rate=sample_rate,
                n_fft=wl,        # using window_length as n_fft usually
                win_length=wl,
                hop_length=hop_length,
                n_mels=nm,
                f_min=fmin,
                f_max=fmax,
                power=1.0,       # computing Magnitude Mel (power=1.0) initially
                normalized=False,
                center=True,
                pad_mode="reflect"
            )
            self.mel_transforms.append(transform)

    def forward(
        self,
        x_pred: torch.Tensor,
        x_true: torch.Tensor,
        frame_mask: torch.Tensor | None = None,
        frame_lengths: torch.Tensor | None = None,
    ):
        """
        Args:
            x_pred: [B, T] or [B, 1, T] Estimated Waveform
            x_true: [B, T] or [B, 1, T] Ground Truth Waveform
        Returns:
            Weighted scalar loss
        """
        # Ensure correct shape [B, T] for torchaudio, or [B, 1, T] is also fine but usually squeeze
        if x_pred.ndim == 3 and x_pred.shape[1] == 1:
            x_pred = x_pred.squeeze(1)
        if x_true.ndim == 3 and x_true.shape[1] == 1:
            x_true = x_true.squeeze(1)
            
        total_loss = x_pred.new_tensor(0.0)

        use_mask = frame_mask is not None
        if use_mask:
            span_starts, span_ends = self._get_span_bounds_from_mask(frame_mask, frame_lengths)
            span_starts = span_starts.to(device=x_pred.device)
            span_ends = span_ends.to(device=x_pred.device)
        
        for mel_transform in self.mel_transforms:
            x_mels = mel_transform(x_pred)
            y_mels = mel_transform(x_true)

            mel_mask = None
            if use_mask:
                hop = int(mel_transform.hop_length)
                mel_t = x_mels.shape[-1]
                mel_mask = torch.zeros(
                    (x_mels.shape[0], mel_t),
                    device=x_mels.device,
                    dtype=torch.bool,
                )
                for i in range(x_mels.shape[0]):
                    s = int(span_starts[i].item())
                    e = int(span_ends[i].item())
                    if e <= s:
                        continue
                    ms = s // hop
                    me = (e + hop - 1) // hop
                    ms = max(0, min(ms, mel_t))
                    me = max(0, min(me, mel_t))
                    if me > ms:
                        mel_mask[i, ms:me] = True

                if not mel_mask.any():
                    continue

                mel_mask = mel_mask.unsqueeze(1)
            
            # Log Magnitude Loss
            # Formula: L1( log10(x^pow + eps), log10(y^pow + eps) )
            x_log = x_mels.clamp(min=self.clamp_eps).pow(self.pow).log10()
            y_log = y_mels.clamp(min=self.clamp_eps).pow(self.pow).log10()

            diff_full = (x_log - y_log).abs()

            if mel_mask is None:
                diff = diff_full.mean(dim=(1, 2))
            else:
                diff = []
                for i in range(diff_full.shape[0]):
                    m = mel_mask[i]
                    if m.any():
                        diff.append(diff_full[i].masked_select(m).mean())
                    else:
                        diff.append(diff_full.new_tensor(0.0))
                diff = torch.stack(diff, dim=0)

            total_loss += diff.mean()

        return total_loss * self.weight

    @staticmethod
    def _get_span_bounds_from_mask(mask: torch.Tensor, lengths: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]:
        """Convert boolean frame mask [B, T] into per-sample [start, end) bounds.

        This helper assumes one contiguous masked span per sample (matching current
        random-span masking strategy). If a sample has no masked frame, start=end=0.
        """
        if mask.dtype != torch.bool:
            mask = mask.bool()

        bsz, t = mask.shape
        device = mask.device

        if lengths is None:
            lengths = torch.full((bsz,), t, device=device, dtype=torch.long)
        else:
            lengths = lengths.to(device=device, dtype=torch.long)

        starts = torch.zeros((bsz,), device=device, dtype=torch.long)
        ends = torch.zeros((bsz,), device=device, dtype=torch.long)

        for i in range(bsz):
            li = int(lengths[i].item())
            row = mask[i, :li]
            idx = row.nonzero(as_tuple=False).squeeze(-1)
            if idx.numel() == 0:
                continue
            starts[i] = idx[0]
            ends[i] = idx[-1] + 1

        return starts, ends

    @staticmethod
    def _aligned_random_span_mask(
        lengths: torch.Tensor,
        frac_lengths: torch.Tensor,
        align_to: int,
        max_length: int | None = None,
    ) -> torch.Tensor:
        """Sample one contiguous mask span per sample with boundary alignment.

        Args:
            lengths: [B] valid sequence lengths in model frames.
            frac_lengths: [B] desired masked length fraction.
            align_to: span start/end alignment in model frames.
            max_length: optional mask width. If None, use lengths.max().
        Returns:
            bool mask with shape [B, max_length].
        """
        assert align_to >= 1
        device = lengths.device
        lengths = lengths.to(dtype=torch.long)
        frac_lengths = frac_lengths.to(device=device, dtype=torch.float32)

        max_len = int(lengths.max().item()) if max_length is None else int(max_length)
        mask = torch.zeros((lengths.shape[0], max_len), device=device, dtype=torch.bool)

        for i in range(lengths.shape[0]):
            li = int(lengths[i].item())
            if li <= 0:
                continue

            target = int((frac_lengths[i].item() * li))
            target = max(1, min(target, li))

            span_len = ((target + align_to - 1) // align_to) * align_to
            span_len = min(span_len, li)

            max_start = li - span_len
            if max_start <= 0:
                start = 0
            else:
                candidate = torch.randint(0, max_start + 1, (1,), device=device).item()
                start = (candidate // align_to) * align_to
                start = min(start, max_start)

            end = start + span_len
            mask[i, start:end] = True

        return mask