# ==============================================================================
# FSQ Latent Predictor with:
#   1. ConvNeXt V2 Text Encoder with Language Embeddings
#   2. Duration-based Monotonic Attention (s2s attention from durations)
#   3. Speaker Embedding Conditioning
#   4. Downsampling from 100Hz to 50Hz (codec rate)
#   5. DiT-based Flow Matching (Asymmetric: 18-layer encoder, 4-layer decoder)
#   6. CTC Loss on intermediate condition encoder representations
#   7. Context-aware encoding (Prev + Curr + Next via ConvNeXt convolutions)
#   8. Velocity Direction Loss (Cosine Similarity)
# ==============================================================================

import math
from typing import Optional, List, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
from torch import Tensor

from Modules.predictor_duration_speaker_context import StochasticDurationPredictorContext
from meldataset import SPOKEN_IDS


# ==============================================================================
# GLOBAL RESPONSE NORMALIZATION (GRN) - ConvNeXt V2 key component
# ==============================================================================

class GRN(nn.Module):
    """Global Response Normalization from ConvNeXt V2."""
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.gamma = nn.Parameter(torch.zeros(1, dim, 1))
        self.beta = nn.Parameter(torch.zeros(1, dim, 1))
        self.eps = eps

    def forward(self, x: Tensor) -> Tensor:
        gx = torch.norm(x, p=2, dim=-1, keepdim=True)
        nx = gx / (gx.mean(dim=1, keepdim=True) + self.eps)
        return self.gamma * (x * nx) + self.beta + x

class StyleAdaLN1D(nn.Module):
    """AdaLN for 1D sequences with style/language conditioning."""
    def __init__(self, channels: int, style_dim: int, eps: float = 1e-6):
        super().__init__()
        self.channels = channels
        self.eps = eps
        self.fc = nn.Linear(style_dim, channels * 2)
        
        nn.init.zeros_(self.fc.weight)
        nn.init.zeros_(self.fc.bias)

    def forward(self, x: Tensor, s: Tensor) -> Tensor:
        h = self.fc(s)
        h = h.view(h.size(0), h.size(1), 1)
        gamma, beta = torch.chunk(h, chunks=2, dim=1)
        
        x = x.transpose(1, 2)
        x = F.layer_norm(x, (self.channels,), eps=self.eps)
        x = x.transpose(1, 2)
        
        return (1 + gamma) * x + beta

class ConvNeXtV2Block(nn.Module):
    def __init__(
        self,
        dim: int,
        kernel_size: int = 7,
        expansion_ratio: float = 4.0,
        drop_path: float = 0.0,
        style_dim: int = 64,
    ):
        super().__init__()
        hidden_dim = int(dim * expansion_ratio)
        padding = kernel_size // 2
        
        self.dwconv = nn.Conv1d(dim, dim, kernel_size, padding=padding, groups=dim)
        self.adaln = StyleAdaLN1D(dim, style_dim)
        self.pwconv1 = nn.Linear(dim, hidden_dim)
        self.act = nn.GELU()
        self.grn = GRN(hidden_dim)
        self.pwconv2 = nn.Linear(hidden_dim, dim)
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        
    def _mask_bct(self, x: Tensor, mask_bt: Optional[Tensor]) -> Tensor:
        # x: [B, C, T], mask_bt: [B, T] (bool)
        if mask_bt is None:
            return x
        return x * mask_bt.unsqueeze(1).to(dtype=x.dtype)

    def _mask_btc(self, x: Tensor, mask_bt: Optional[Tensor]) -> Tensor:
        # x: [B, T, C], mask_bt: [B, T] (bool)
        if mask_bt is None:
            return x
        return x * mask_bt.unsqueeze(-1).to(dtype=x.dtype)

    def forward(
        self,
        x: Tensor,
        style: Tensor,
        mask: Optional[Tensor] = None,
    ) -> Tensor:
        residual = x
        # Important: mask after depthwise conv so padded positions
        # cannot influence GRN normalization across time.
        x = self.dwconv(x)
        x = self._mask_bct(x, mask)
        x = self.adaln(x, style)
        x = self._mask_bct(x, mask)

        x = x.transpose(1, 2)
        x = self.pwconv1(x)
        x = self.act(x)
        x = self._mask_btc(x, mask)

        x = x.transpose(1, 2)
        x = self.grn(x)
        x = self._mask_bct(x, mask)
        x = x.transpose(1, 2)

        x = self.pwconv2(x)
        x = x.transpose(1, 2)

        x = residual + self.drop_path(x)
        x = self._mask_bct(x, mask)
        return x

class DropPath(nn.Module):
    """Stochastic depth (per-sample)."""
    def __init__(self, drop_prob: float = 0.0):
        super().__init__()
        self.drop_prob = drop_prob

    def forward(self, x: Tensor) -> Tensor:
        if self.drop_prob == 0.0 or not self.training:
            return x
        keep = 1.0 - self.drop_prob
        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
        mask = x.new_empty(shape).bernoulli_(keep)
        return x * mask / keep


# ==============================================================================
# CONVNEXT V2 TEXT ENCODER WITH LANGUAGE EMBEDDINGS
# ==============================================================================

class ConvNeXtV2TextEncoder(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        d_model: int = 512,
        num_layers: int = 6,
        kernel_size: int = 7,
        expansion_ratio: float = 4.0,
        num_languages: int = 10,
        language_dim: int = 64,
        drop_path_rate: float = 0.1,
        dropout: float = 0.1,
        pad_token_id: int = 0,
    ):
        super().__init__()
        self.d_model = d_model
        self.pad_token_id = pad_token_id
        self.language_dim = language_dim
        
        # Use non-learnable padding index to avoid PAD leakage into convs
        self.token_embedding = nn.Embedding(vocab_size, d_model, padding_idx=pad_token_id)
        self.language_embedding = nn.Embedding(num_languages, language_dim)
        self.input_proj = nn.Conv1d(d_model, d_model, 1)
        
        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, num_layers)]
        self.blocks = nn.ModuleList([
            ConvNeXtV2Block(
                dim=d_model,
                kernel_size=kernel_size,
                expansion_ratio=expansion_ratio,
                drop_path=dpr[i],
                style_dim=language_dim,
            )
            for i in range(num_layers)
        ])
        
        self.norm = nn.LayerNorm(d_model, eps=1e-6)
        self.dropout = nn.Dropout(dropout)
        
    def forward(
        self,
        phonemes: Tensor,
        language_ids: Tensor,
        mask: Optional[Tensor] = None,
    ) -> Tensor:
        B, T = phonemes.shape
        x = self.token_embedding(phonemes)
        lang_emb = self.language_embedding(language_ids)
        x = self.dropout(x)
        
        x = x.transpose(1, 2)
        x = self.input_proj(x)

        # Apply boolean length mask as early as possible and pass it into blocks
        if mask is not None:
            x = x * mask.unsqueeze(1).to(dtype=x.dtype)

        for block in self.blocks:
            x = block(x, lang_emb, mask=mask)
        
        x = x.transpose(1, 2)
        x = self.norm(x)
        
        if mask is not None:
            x = x * mask.unsqueeze(-1).float()
            
        return x


# ==============================================================================
# DURATION-BASED MONOTONIC ATTENTION (S2S Attention from Durations)
# ==============================================================================

def create_duration_attention_matrix(
    durations: Tensor,
    target_length: int,
    text_mask: Optional[Tensor] = None,
    token_mask: Optional[Tensor] = None,
    sigma: float = 1.0,
) -> Tensor:
    B, T_text = durations.shape
    device = durations.device
    
    cum_dur = torch.cumsum(durations, dim=1)
    centers = cum_dur - durations / 2
    
    frames = torch.arange(target_length, device=device, dtype=durations.dtype)
    frames = frames.view(1, -1, 1)
    centers = centers.unsqueeze(1)
    
    weights = torch.exp(-((frames - centers) ** 2) / (2 * sigma ** 2))

    valid_mask = None
    if text_mask is not None and token_mask is not None:
        valid_mask = text_mask & token_mask
    elif text_mask is not None:
        valid_mask = text_mask
    elif token_mask is not None:
        valid_mask = token_mask
        
    if valid_mask is not None:
        weights = weights * valid_mask.unsqueeze(1).to(weights.dtype)
    
    weights = weights / (weights.sum(dim=-1, keepdim=True) + 1e-8)
    return weights


def jitter_durations_lognormal(
    durations: torch.Tensor,
    valid_mask: torch.Tensor,
    sigma: float = 0.20,
    clip: float = 0.60,
) -> torch.Tensor:
    d = durations

    eps = torch.randn_like(d) * sigma
    eps = eps.clamp(-clip, clip)

    d_pos = d.clamp_min(1.0)
    d_jit = torch.exp(torch.log(d_pos) + eps).clamp_min(1.0)

    return torch.where(valid_mask, d_jit, d)


def duration_based_upsample(
    text_emb: Tensor,
    durations: Tensor,
    target_length: Optional[int] = None,
    text_mask: Optional[Tensor] = None,
    token_mask: Optional[Tensor] = None,
    sigma: float = 1.0,
) -> Tuple[Tensor, Tensor]:
    B, T_text, C = text_emb.shape
    
    if target_length is None:
        target_length = int(durations.sum(dim=-1).max().item())
        
    attention = create_duration_attention_matrix(
        durations,
        target_length,
        text_mask=text_mask,
        token_mask=token_mask,
        sigma=sigma,
    )
    
    upsampled = torch.bmm(attention, text_emb)
    return upsampled, attention


# ==============================================================================
# GAUSSIAN UPSAMPLE (direct 50Hz, numerically stable)
# ==============================================================================

def gaussian_upsample(
    token_feats: Tensor,               # [B, D, N]
    durations: Tensor,                 # [B, N] (int or float) - can be fractional
    T: Tensor,                         # [B] target frames per sample (50Hz codec lengths)
    sigma_sq: float = 3.0,
    token_mask: Optional[Tensor] = None,
) -> Tuple[Tensor, Tensor]:
    """
    Numerically-stable Gaussian upsampling directly to target frame rate.
    Rescales durations to exactly fill T[b] frames per sample (no round-trip).
    Returns:
        frame_feats: [B, D, max_T]
        weights:     [B, N, max_T]
    """
    B, D, N = token_feats.shape
    device = token_feats.device
    T_list = [int(T[b].item()) for b in range(B)]
    max_T = max(T_list) if T_list else 0

    out_list: List[Tensor] = []
    weights_list: List[Tensor] = []

    for b in range(B):
        Tb = T_list[b]
        if token_mask is not None:
            # Use actual nonzero indices so that zero-duration tokens (BOS, EOS)
            # at any position are truly excluded — not just those at the tail.
            # Using sum() + prefix slice would include BOS (index 0, duration=0)
            # which would Gaussian-bleed its embedding into frame 0.
            valid_idx = token_mask[b].nonzero(as_tuple=True)[0]  # true sparse gather
            Nb = valid_idx.shape[0]
        else:
            valid_idx = None
            Nb = N

        if Nb <= 0 or Tb <= 0:
            out_list.append(torch.zeros(D, max_T, device=device, dtype=token_feats.dtype))
            weights_list.append(torch.zeros(N, max_T, device=device, dtype=token_feats.dtype))
            continue

        if valid_idx is not None:
            E = token_feats[b, :, valid_idx]          # [D, Nb] — only valid tokens
            d = durations[b, valid_idx].float()       # [Nb]
        else:
            E = token_feats[b, :, :Nb]               # [D, Nb]
            d = durations[b, :Nb].float()             # [Nb]

        # Rescale durations to exactly fill Tb frames
        d_sum = d.sum().clamp(min=1e-8)
        scale = Tb / d_sum
        l = d * scale                    # [Nb] fractional frames at target rate

        # Phoneme centers in target-frame coordinates
        c = torch.cumsum(l, dim=0) - 0.5 * l  # [Nb]
        t = torch.arange(Tb, device=device).float().unsqueeze(0)  # [1, Tb]

        # Stable softmax along phoneme axis
        dist2 = (t - c.unsqueeze(1)) ** 2     # [Nb, Tb]
        logits = (-dist2 / float(sigma_sq)).to(torch.float32)
        w = F.softmax(logits, dim=0).to(dist2.dtype)  # [Nb, Tb]

        F_bt = E @ w  # [D, Tb]

        if Tb < max_T:
            F_pad = torch.zeros(D, max_T, device=device, dtype=F_bt.dtype)
            F_pad[:, :Tb] = F_bt
            F_bt = F_pad

        out_list.append(F_bt)

        w_full = torch.zeros(N, max_T, device=device, dtype=w.dtype)
        if valid_idx is not None:
            w_full[valid_idx, :Tb] = w   # scatter to original positions
        else:
            w_full[:Nb, :Tb] = w
        weights_list.append(w_full)

    frame_feats = torch.stack(out_list, dim=0)    # [B, D, max_T]
    weights = torch.stack(weights_list, dim=0)    # [B, N, max_T]
    return frame_feats, weights


# ==============================================================================
# SINUSOIDAL TIME EMBEDDING
# ==============================================================================

class SinusoidalTimeEmbedding(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        assert dim % 2 == 0
        self.dim = dim

    def forward(self, t: Tensor) -> Tensor:
        half = self.dim // 2
        device = t.device
        exponents = torch.arange(half, device=device, dtype=t.dtype)
        freqs = torch.exp(-math.log(10000.0) * exponents / half)
        args = t[:, None] * freqs[None, :]
        return torch.cat([torch.sin(args), torch.cos(args)], dim=-1)


class TimeMLPEmbedding(nn.Module):
    def __init__(self, out_features: int, time_embed_dim: int = 128):
        super().__init__()
        self.net = nn.Sequential(
            SinusoidalTimeEmbedding(time_embed_dim),
            nn.Linear(time_embed_dim, out_features),
            nn.SiLU(),
            nn.Linear(out_features, out_features),
        )

    def forward(self, t: Tensor) -> Tensor:
        return self.net(t)


# ==============================================================================
# ADAPTIVE LAYER NORM (DiT Style)
# ==============================================================================

class AdaLNZero(nn.Module):
    def __init__(self, dim: int, cond_dim: int):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.linear = nn.Linear(cond_dim, dim * 2)
        
        nn.init.zeros_(self.linear.weight)
        nn.init.zeros_(self.linear.bias)

    def forward(self, x: Tensor, cond: Tensor) -> Tensor:
        x = self.norm(x)
        gamma_beta = self.linear(cond)
        gamma, beta = gamma_beta.chunk(2, dim=-1)
        return x * (1 + gamma.unsqueeze(1)) + beta.unsqueeze(1)


# ==============================================================================
# MULTI-HEAD ATTENTION WITH RoPE
# ==============================================================================

class RotaryEmbedding(nn.Module):
    def __init__(self, dim: int, base: float = 10000.0):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def get_cos_sin(self, seq_len: int, device: torch.device, dtype: torch.dtype):
        t = torch.arange(seq_len, device=device, dtype=dtype)
        freqs = torch.outer(t, self.inv_freq.to(dtype))
        emb = torch.cat([freqs, freqs], dim=-1)
        return emb.cos(), emb.sin()


def apply_rope(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
    x_even = x[..., ::2]
    x_odd = x[..., 1::2]
    x_rot = torch.stack([-x_odd, x_even], dim=-1).flatten(-2)
    cos = cos.unsqueeze(0).unsqueeze(0)
    sin = sin.unsqueeze(0).unsqueeze(0)
    return (x * cos) + (x_rot * sin)


class MultiheadSelfAttention(nn.Module):
    def __init__(
        self,
        dim: int,
        num_heads: int = 8,
        dropout: float = 0.0,
        use_rope: bool = True,
    ):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim ** -0.5
        
        self.qkv = nn.Linear(dim, dim * 3)
        self.out_proj = nn.Linear(dim, dim)
        self.dropout = nn.Dropout(dropout)
        
        self.rope = RotaryEmbedding(self.head_dim) if use_rope else None

    def forward(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor:
        B, T, D = x.shape
        
        qkv = self.qkv(x)
        q, k, v = qkv.chunk(3, dim=-1)
        
        q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        
        if self.rope is not None:
            cos, sin = self.rope.get_cos_sin(T, x.device, x.dtype)
            q = apply_rope(q, cos, sin)
            k = apply_rope(k, cos, sin)
        
        attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
        
        if mask is not None:
            mask = mask.unsqueeze(1).unsqueeze(2)
            attn = attn.masked_fill(~mask, float('-inf'))
        
        attn = F.softmax(attn, dim=-1)
        attn = self.dropout(attn)
        
        out = torch.matmul(attn, v)
        out = out.transpose(1, 2).contiguous().view(B, T, D)
        return self.out_proj(out)


# ==============================================================================
# DIT BLOCK FOR FLOW MATCHING
# ==============================================================================

class DiTBlock(nn.Module):
    def __init__(
        self,
        dim: int,
        num_heads: int = 8,
        mlp_ratio: float = 4.0,
        cond_dim: int = 512,
        dropout: float = 0.1,
        use_rope: bool = True,
    ):
        super().__init__()
        self.adaln1 = AdaLNZero(dim, cond_dim)
        self.attn = MultiheadSelfAttention(dim, num_heads, dropout, use_rope)
        
        self.adaln2 = AdaLNZero(dim, cond_dim)
        hidden_dim = int(dim * mlp_ratio)
        self.mlp = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(dropout),
        )
        
    def forward(self, x: Tensor, cond: Tensor, mask: Optional[Tensor] = None) -> Tensor:
        x = x + self.attn(self.adaln1(x, cond), mask)
        x = x + self.mlp(self.adaln2(x, cond))
        return x


# ==============================================================================
# DIT FLOW MATCHING MODEL
# ==============================================================================

class ConditionEncoder(nn.Module):
    def __init__(
        self,
        dim: int,
        num_layers: int,
        num_heads: int,
        mlp_ratio: float,
        dropout: float,
        cond_dim: int,
        ctc_layer_idx: int = -1, 
    ):
        super().__init__()
        self.num_layers = num_layers
        self.input_proj = nn.Linear(cond_dim, dim)
        self.blocks = nn.ModuleList([
            DiTBlock(
                dim=dim,
                num_heads=num_heads,
                mlp_ratio=mlp_ratio,
                cond_dim=dim,
                dropout=dropout,
            )
            for _ in range(num_layers)
        ])
        self.norm = nn.LayerNorm(dim)
        
        self.ctc_layer_idx = ctc_layer_idx if ctc_layer_idx >= 0 else max(0, num_layers // 2 - 1)

    def forward(
        self,
        cond: Tensor,
        adaln_cond: Tensor,
        mask: Optional[Tensor] = None,
        return_intermediate: bool = False,
    ) -> Tuple[Tensor, Optional[Tensor]]:
        x = self.input_proj(cond)
        ctc_hidden = None
        
        for i, block in enumerate(self.blocks):
            x = block(x, adaln_cond, mask)
            if return_intermediate and i == self.ctc_layer_idx:
                ctc_hidden = x
        
        final_hidden = self.norm(x)
        return final_hidden, ctc_hidden


class VelocityDecoder(nn.Module):
    def __init__(
        self,
        dim: int,
        fsq_dim: int,
        num_layers: int,
        num_heads: int,
        mlp_ratio: float,
        dropout: float,
    ):
        super().__init__()
        self.dim = dim
        self.x_proj = nn.Linear(fsq_dim, dim)
        
        self.cond_global_proj = nn.Sequential(
            nn.Linear(dim, dim),
            nn.SiLU(),
            nn.Linear(dim, dim),
        )
        
        self.blocks = nn.ModuleList([
            DiTBlock(
                dim=dim,
                num_heads=num_heads,
                mlp_ratio=mlp_ratio,
                cond_dim=dim,
                dropout=dropout,
            )
            for _ in range(num_layers)
        ])
        self.norm = nn.LayerNorm(dim)
        self.out_proj = nn.Linear(dim, fsq_dim)

        nn.init.zeros_(self.out_proj.weight)
        nn.init.zeros_(self.out_proj.bias)

    def forward(
        self,
        x_t: Tensor,
        cond_hidden: Tensor,
        adaln_cond: Tensor,
        mask: Optional[Tensor] = None,
    ) -> Tensor:
        if mask is not None:
            mask_expanded = mask.unsqueeze(-1).float()
            h_global = (cond_hidden * mask_expanded).sum(dim=1) / (mask_expanded.sum(dim=1) + 1e-8)
        else:
            h_global = cond_hidden.mean(dim=1)
        
        h_global = self.cond_global_proj(h_global)
        combined_adaln = adaln_cond + h_global
        
        x = self.x_proj(x_t)
        x = x + cond_hidden 

        for block in self.blocks:
            x = block(x, combined_adaln, mask)

        x = self.norm(x)
        return self.out_proj(x)


class DiTFlowMatchingModel(nn.Module):
    def __init__(
        self,
        fsq_dim: int = 6,
        cond_dim: int = 512,
        hidden_dim: int = 512,
        cond_layers: int = 6,
        vel_layers: int = 4,
        num_heads: int = 8,
        mlp_ratio: float = 4.0,
        dropout: float = 0.1,
        time_embed_dim: int = 128,
        speaker_dim: int = 128,
        language_dim: int = 64,
    ):
        super().__init__()
        self.fsq_dim = fsq_dim
        self.hidden_dim = hidden_dim
        
        self.time_embed = TimeMLPEmbedding(hidden_dim, time_embed_dim)
        self.spk_lang_proj = nn.Sequential(
            nn.Linear(speaker_dim + language_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim),
        )

        self.condition_encoder = ConditionEncoder(
            dim=hidden_dim,
            num_layers=cond_layers,
            num_heads=num_heads,
            mlp_ratio=mlp_ratio,
            dropout=dropout,
            cond_dim=cond_dim,
        )
        self.velocity_decoder = VelocityDecoder(
            dim=hidden_dim,
            fsq_dim=fsq_dim,
            num_layers=vel_layers,
            num_heads=num_heads,
            mlp_ratio=mlp_ratio,
            dropout=dropout,
        )

    def _build_spk_lang_cond(self, speaker_emb: Tensor, language_emb: Tensor) -> Tensor:
        spk_lang = torch.cat([speaker_emb, language_emb], dim=-1)
        return self.spk_lang_proj(spk_lang)

    def _build_adaln_cond(self, t: Tensor, spk_lang_emb: Tensor) -> Tensor:
        t_emb = self.time_embed(t * 1000)
        return t_emb + spk_lang_emb

    def forward(
        self,
        x_t: Tensor,
        t: Tensor,
        cond: Tensor,
        speaker_emb: Tensor,
        language_emb: Tensor,
        mask: Optional[Tensor] = None,
        return_hidden: bool = False,
    ) -> Tuple[Tensor, Optional[Tensor]]:
        spk_lang_emb = self._build_spk_lang_cond(speaker_emb, language_emb)
        adaln_cond = self._build_adaln_cond(t, spk_lang_emb)

        cond_hidden, ctc_hidden = self.condition_encoder(
            cond, adaln_cond, mask, return_intermediate=return_hidden
        )

        v_pred = self.velocity_decoder(x_t, cond_hidden, adaln_cond, mask)
        
        if return_hidden:
            hidden_for_ctc = ctc_hidden if ctc_hidden is not None else cond_hidden
            return v_pred, hidden_for_ctc
        return v_pred, None


# ==============================================================================
# LIGHTWEIGHT CTC PREDICTOR
# ==============================================================================

class CTCPredictor(nn.Module):
    def __init__(
        self,
        input_dim: int,
        hidden_dim: int = 256,
        n_phonemes: int = 178,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.n_phonemes = n_phonemes
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, n_phonemes + 1),
        )
        
    def forward(self, x: Tensor) -> Tensor:
        return F.log_softmax(self.net(x), dim=-1)


# ==============================================================================
# MAIN FSQ FLOW PREDICTOR
# ==============================================================================

class FSQFlowPredictorConvNeXt(nn.Module):
    def __init__(
        self,
        vocab_size: int = 178,
        text_dim: int = 512,
        text_num_layers: int = 6,
        text_kernel_size: int = 7,
        num_languages: int = 10,
        language_dim: int = 64,
        
        n_speakers: int = 11,
        speaker_dim: int = 128,
        
        fsq_levels: List[int] = [4, 4, 4, 4, 4, 4],
        
        hidden_dim: int = 512,
        cond_layers: int = 6, 
        vel_layers: int = 4,   
        flow_num_heads: int = 8,
        
        n_phonemes: int = 178,
        
        duration_n_flows: int = 6,
        duration_kernel_size: int = 3,
        duration_p_dropout: float = 0.1,
        
        dropout: float = 0.1,
        codec_stride: int = 2,
        cfg_dropout_prob: float = 0.15, 
        
        pad_token_id: int = 0,
    ):
        super().__init__()
        self.fsq_levels = fsq_levels
        self.fsq_dim = len(fsq_levels)
        self.codec_stride = codec_stride
        self.text_dim = text_dim
        self.speaker_dim = speaker_dim
        self.language_dim = language_dim
        self.pad_token_id = pad_token_id
        self.n_phonemes = n_phonemes
        self.cfg_dropout_prob = cfg_dropout_prob
        
        self.text_encoder = ConvNeXtV2TextEncoder(
            vocab_size=vocab_size,
            d_model=text_dim,
            num_layers=text_num_layers,
            kernel_size=text_kernel_size,
            num_languages=num_languages,
            language_dim=language_dim,
            dropout=dropout,
            pad_token_id=pad_token_id,
        )
        
        self.speaker_embedding = nn.Embedding(n_speakers, speaker_dim)
        nn.init.normal_(self.speaker_embedding.weight, mean=0, std=0.5)

        self.language_embedding_dit = nn.Embedding(num_languages, language_dim)
        nn.init.normal_(self.language_embedding_dit.weight, mean=0, std=0.5)
        
        cond_dim = text_dim
        
        self.flow_model = DiTFlowMatchingModel(
            fsq_dim=self.fsq_dim,
            cond_dim=cond_dim,
            hidden_dim=hidden_dim,
            cond_layers=cond_layers,
            vel_layers=vel_layers,
            num_heads=flow_num_heads,
            dropout=dropout,
            speaker_dim=speaker_dim,
            language_dim=language_dim,
        )
        
        self.ctc_predictor = CTCPredictor(
            input_dim=hidden_dim,
            hidden_dim=hidden_dim // 2,
            n_phonemes=n_phonemes,
            dropout=dropout,
        )

        # ------------------------------------------------------------------
        # DURATION PREDICTOR HEAD (pinned after ConvNeXt encoder)
        # Uses the same ConvNeXt text features + speaker embeddings.
        # ------------------------------------------------------------------
        self.duration_predictor = StochasticDurationPredictorContext(
            in_channels=text_dim,
            filter_channels=text_dim // 2,
            kernel_size=duration_kernel_size,
            p_dropout=duration_p_dropout,
            n_flows=duration_n_flows,
            speaker_emb_dim=speaker_dim + language_dim,  # cat(spk, lang)
        )

        # Spoken-phoneme mask: only compute duration loss on "voiced" tokens
        self.register_buffer('spoken_mask_tensor', torch.zeros(vocab_size, dtype=torch.bool))
        for _sid in SPOKEN_IDS:
            if _sid < vocab_size:
                self.spoken_mask_tensor[_sid] = True

        self.register_buffer('fsq_levels_tensor', torch.tensor(fsq_levels, dtype=torch.float32))
        basis = torch.cumprod(torch.tensor([1] + fsq_levels[:-1]), dim=0)
        self.register_buffer('fsq_basis', basis)
        
    def indices_to_continuous(self, indices: Tensor) -> Tensor:
        if indices.dim() == 3:
            indices = indices.squeeze(1)
            
        z_q = []
        remainder = indices
        for i in range(self.fsq_dim):
            val = remainder % int(self.fsq_levels[i])
            remainder = remainder // int(self.fsq_levels[i])
            z_q.append(val)
        
        z_q = torch.stack(z_q, dim=-1).float()
        half_l = (self.fsq_levels_tensor - 1) / 2
        return z_q - half_l
    
    def continuous_to_indices(self, z_continuous: Tensor) -> Tensor:
        half_l = (self.fsq_levels_tensor - 1) / 2
        z_shifted = z_continuous + half_l
        z_ind = z_shifted.round()
        z_ind = torch.clamp(z_ind, torch.zeros_like(self.fsq_levels_tensor), self.fsq_levels_tensor - 1)
        
        z_ind_long = z_ind.long()
        return (z_ind_long * self.fsq_basis).sum(dim=-1)

    def _extract_current_context_embeddings(
        self, ctx_encoded: Tensor, token_start_indices: Tensor, token_end_indices: Tensor, target_text_len: int,
    ) -> Tensor:
        B, T_ctx, _ = ctx_encoded.shape
        device = ctx_encoded.device
        text_emb = torch.zeros(B, target_text_len, self.text_dim, device=device, dtype=ctx_encoded.dtype)
        copy_len = max(target_text_len, 0)
        if copy_len == 0 or T_ctx == 0:
            return text_emb

        starts = token_start_indices.to(device=device, dtype=torch.long).unsqueeze(1)
        ends = token_end_indices.to(device=device, dtype=torch.long).unsqueeze(1)
        rel_pos = torch.arange(copy_len, device=device, dtype=torch.long).unsqueeze(0)
        src_pos = starts + rel_pos

        valid = (src_pos >= 0) & (src_pos < ends) & (src_pos < T_ctx)
        safe_src_pos = src_pos.clamp(min=0, max=T_ctx - 1)

        batch_idx = torch.arange(B, device=device, dtype=torch.long).unsqueeze(1)
        gathered = ctx_encoded[batch_idx, safe_src_pos]
        gathered = gathered * valid.unsqueeze(-1).to(gathered.dtype)

        text_emb[:, :copy_len, :] = gathered
        return text_emb

    def forward(
        self,
        phonemes: Tensor,
        phoneme_lengths: Tensor,
        durations: Tensor,
        speaker_ids: Tensor,
        language_ids: Tensor,
        gt_fsq_embeddings: Tensor,
        codec_lengths: Tensor,
        ctx_tokens: Optional[Tensor] = None,
        ctx_lens: Optional[Tensor] = None,
        token_start_indices: Optional[Tensor] = None,
        token_end_indices: Optional[Tensor] = None,
        t_override: Optional[Tensor] = None,
        noise_override: Optional[Tensor] = None,
        return_debug: bool = False,
    ) -> dict:
        B = phonemes.shape[0]
        device = phonemes.device
        T_text = phonemes.shape[1]
        
        text_mask = self._length_to_mask(phoneme_lengths, T_text)
        durations = durations.float() * text_mask.to(durations.dtype)
        upsample_mask = durations > 0
        
        if ctx_tokens is not None and token_start_indices is not None and token_end_indices is not None and ctx_lens is not None:
            T_ctx = ctx_tokens.shape[1]
            ctx_mask = self._length_to_mask(ctx_lens, T_ctx)
            ctx_encoded = self.text_encoder(ctx_tokens, language_ids, ctx_mask)

            text_emb = self._extract_current_context_embeddings(
                ctx_encoded, token_start_indices, token_end_indices, T_text,
            )
            text_emb = text_emb * text_mask.unsqueeze(-1).float()
        else:
            text_emb = self.text_encoder(phonemes, language_ids, text_mask)

     
        durations_corr = durations

        spk_emb = self.speaker_embedding(speaker_ids)
        lang_emb_dit = self.language_embedding_dit(language_ids)

        # ------------------------------------------------------------------
        # DURATION PREDICTOR HEAD (after ConvNeXt, before CFG dropout)
        # Uses the original (non-dropped) speaker embeddings.
        # ------------------------------------------------------------------
        is_spoken = self.spoken_mask_tensor[phonemes]               # [B, T]
        dur_valid_mask = text_mask & upsample_mask & is_spoken      # [B, T]
        dur_mask_1 = dur_valid_mask.unsqueeze(1)                     # [B, 1, T]

        text_emb_t = text_emb.transpose(1, 2)                       # [B, C, T]
        w_dur = durations.float().unsqueeze(1) * dur_mask_1.float() # [B, 1, T]
        if self.training:
            w_dur = w_dur + torch.rand_like(w_dur) * dur_mask_1.float()

        dur_loss_raw = self.duration_predictor(
            text_emb_t, dur_mask_1, w=w_dur,
            g=torch.cat([spk_emb, lang_emb_dit], dim=-1)
        )  # [B]  (NLL + logq per sample)
        denom_dur = dur_mask_1.sum().clamp_min(1).float()
        loss_dur = (dur_loss_raw / denom_dur).mean()
        # ------------------------------------------------------------------

        # Direct 50Hz upsampling: gaussian_upsample to codec_lengths (no round-trip)
        token_feats = text_emb.transpose(1, 2)  # [B, D, T_text]
        frame_feats, attn_weights = gaussian_upsample(
            token_feats, durations_corr, codec_lengths,
            sigma_sq=3.0, token_mask=upsample_mask,
        )
        T_codec = gt_fsq_embeddings.shape[1]
        cond = frame_feats.transpose(1, 2)  # [B, max_codec_len, D]
        # Trim/pad in case of off-by-one between codec_lengths.max() and gt length
        if cond.shape[1] > T_codec:
            cond = cond[:, :T_codec, :]
        elif cond.shape[1] < T_codec:
            cond = F.pad(cond, (0, 0, 0, T_codec - cond.shape[1]))

        codec_mask = self._length_to_mask(codec_lengths, T_codec)
        cond = cond * codec_mask.unsqueeze(-1).to(cond.dtype)

        # CFG dropout mask for CTC loss
        cfg_drop_mask = None

        if self.training and self.cfg_dropout_prob > 0:
            cfg_drop_mask = torch.rand(B, device=device) < self.cfg_dropout_prob
            if cfg_drop_mask.any():
                cond = cond.clone()
                spk_emb = spk_emb.clone()
                lang_emb_dit = lang_emb_dit.clone()
                cond[cfg_drop_mask] = 0.0
                spk_emb[cfg_drop_mask] = 0.0
                lang_emb_dit[cfg_drop_mask] = 0.0
        
        if t_override is not None:
            t = t_override.to(device=device, dtype=gt_fsq_embeddings.dtype)
        else:
            logit_normal_m = 0.0
            logit_normal_s = 1.0
            t_normal = torch.randn(B, device=device) * logit_normal_s + logit_normal_m
            t = torch.sigmoid(t_normal)

        if noise_override is not None:
            noise = noise_override.to(device=device, dtype=gt_fsq_embeddings.dtype)
        else:
            noise = torch.randn_like(gt_fsq_embeddings)

        if t.dim() != 1 or t.shape[0] != B:
            raise ValueError(f"Expected t to have shape [B], got {tuple(t.shape)} with B={B}")
        if noise.shape != gt_fsq_embeddings.shape:
            raise ValueError(
                f"Expected noise to have shape {tuple(gt_fsq_embeddings.shape)}, got {tuple(noise.shape)}"
            )
        
        t_view = t.view(-1, 1, 1)
        x_t = (1 - t_view) * noise + t_view * gt_fsq_embeddings
        v_target = gt_fsq_embeddings - noise
        
        # Enable return_hidden=True for CTC Loss extraction
        v_pred, hidden_for_ctc = self.flow_model(
            x_t, t, cond, spk_emb, lang_emb_dit, codec_mask, return_hidden=True,
        )
        
        # Fix: compute losses only over valid codec frames to avoid pad leakage
        valid = codec_mask  # [B, T] bool
        denom = valid.sum().clamp_min(1).to(dtype=v_pred.dtype)

        per_frame_mse = (v_pred - v_target).pow(2).sum(dim=-1)  # [B, T]
        loss_flow = per_frame_mse.masked_select(valid).sum() / (denom * float(self.fsq_dim))

        # Velocity Direction Loss (Cosine Similarity) computed only on valid frames
        v_pred_flat = v_pred.reshape(-1, self.fsq_dim)
        v_target_flat = v_target.reshape(-1, self.fsq_dim)
        mask_flat = valid.reshape(-1)

        if mask_flat.any():
            v_pred_sel = v_pred_flat[mask_flat]
            v_target_sel = v_target_flat[mask_flat]
            cos_sim = F.cosine_similarity(v_pred_sel, v_target_sel, dim=-1)
            loss_dir = (1.0 - cos_sim).mean()
        else:
            loss_dir = loss_flow.new_zeros(())
        
        # CTC Loss using DiT intermediate representations
        # Only compute CTC loss on samples that were NOT dropped by CFG
        phoneme_logits = self.ctc_predictor(hidden_for_ctc)
        phoneme_logits = phoneme_logits.transpose(0, 1)
        
        dur_mask = upsample_mask
        filtered_lengths = dur_mask.sum(dim=-1).long()
        max_filtered_len = filtered_lengths.max().item()
        
        filtered_phonemes = torch.zeros(B, max_filtered_len, dtype=torch.long, device=device)
        for b in range(B):
            valid_indices = dur_mask[b].nonzero(as_tuple=True)[0]
            filtered_phonemes[b, :len(valid_indices)] = phonemes[b, valid_indices]
        
        # If CFG dropout is active, exclude fully-dropped samples (cond=0) from CTC loss
        # spk_lang_drop samples keep their text cond intact so CTC is still valid for them
        if cfg_drop_mask is not None and cfg_drop_mask.any():
            # Only compute CTC on non-dropped samples
            valid_mask = ~cfg_drop_mask
            if valid_mask.any():
                valid_indices = valid_mask.nonzero(as_tuple=True)[0]
                phoneme_logits_valid = phoneme_logits[:, valid_indices, :]
                filtered_phonemes_valid = filtered_phonemes[valid_indices]
                input_lengths_valid = codec_lengths[valid_indices]
                target_lengths_valid = filtered_lengths[valid_indices]
                
                loss_ctc = F.ctc_loss(
                    phoneme_logits_valid, filtered_phonemes_valid, 
                    input_lengths_valid, target_lengths_valid,
                    blank=self.n_phonemes, reduction='mean', zero_infinity=True,
                )
            else:
                # All samples dropped - set CTC loss to zero with proper dtype/device
                loss_ctc = loss_flow.new_zeros(1).squeeze()
        else:
            # No CFG dropout or no samples dropped - compute normally
            input_lengths = codec_lengths
            target_lengths = filtered_lengths
            
            loss_ctc = F.ctc_loss(
                phoneme_logits, filtered_phonemes, input_lengths, target_lengths,
                blank=self.n_phonemes, reduction='mean', zero_infinity=True,
            )
        
        result = {
            'loss_flow': loss_flow,
            'loss_dir': loss_dir,
            'loss_ctc': loss_ctc,
            'loss_dur': loss_dur,
            'loss_total': loss_flow + loss_dir + 0.1 * loss_ctc,
            'v_pred': v_pred,
            'v_target': v_target,
            'attention': attn_weights,
        }
        if return_debug:
            result['t'] = t
            result['x_t'] = x_t

        return result
    
    @torch.no_grad()
    def inference(
        self,
        phonemes: Tensor,
        phoneme_lengths: Tensor,
        durations: Tensor,
        speaker_ids: Tensor,
        language_ids: Tensor,
        n_steps: int = 20,
        temperature: float = 0.8,
        cfg_strength: float = 1.0,
        ctx_tokens: Optional[Tensor] = None,
        ctx_lens: Optional[Tensor] = None,
        token_start_indices: Optional[Tensor] = None,
        token_end_indices: Optional[Tensor] = None,
    ) -> dict:
        B = phonemes.shape[0]
        device = phonemes.device
        T_text = phonemes.shape[1]
        
        text_mask = self._length_to_mask(phoneme_lengths, T_text)
        durations = durations.float() * text_mask.to(durations.dtype)
        upsample_mask = durations > 0
        
        if ctx_tokens is not None and token_start_indices is not None and token_end_indices is not None and ctx_lens is not None:
            T_ctx = ctx_tokens.shape[1]
            ctx_mask = self._length_to_mask(ctx_lens, T_ctx)
            ctx_encoded = self.text_encoder(ctx_tokens, language_ids, ctx_mask)
            text_emb = self._extract_current_context_embeddings(
                ctx_encoded, token_start_indices, token_end_indices, T_text,
            )
            text_emb = text_emb * text_mask.unsqueeze(-1).float()
        else:
            text_emb = self.text_encoder(phonemes, language_ids, text_mask)

        durations_corr = durations
        frame_lengths = durations_corr.sum(dim=-1).long()
        codec_lengths = (frame_lengths + self.codec_stride - 1) // self.codec_stride
        max_codec_len = int(codec_lengths.max().item())

        spk_emb = self.speaker_embedding(speaker_ids)
        lang_emb_dit = self.language_embedding_dit(language_ids)
        spk_lang_emb = self.flow_model._build_spk_lang_cond(spk_emb, lang_emb_dit)

        # Direct 50Hz upsampling: gaussian_upsample to codec_lengths (no round-trip)
        token_feats = text_emb.transpose(1, 2)  # [B, D, T_text]
        frame_feats, attn = gaussian_upsample(
            token_feats, durations_corr, codec_lengths,
            sigma_sq=3.0, token_mask=upsample_mask,
        )
        cond = frame_feats.transpose(1, 2)  # [B, max_codec_len, D]

        codec_mask = self._length_to_mask(codec_lengths, max_codec_len)
        cond = cond * codec_mask.unsqueeze(-1).to(cond.dtype)

        x = torch.randn(B, max_codec_len, self.fsq_dim, device=device) * temperature
        
        # Pre-create zero conditions once if using CFG (optimization)
        if cfg_strength != 1.0:
            zero_cond = torch.zeros_like(cond)
            zero_spk_emb = torch.zeros_like(spk_emb)
            zero_lang_emb = torch.zeros_like(lang_emb_dit)
            zero_spk_lang_emb = self.flow_model._build_spk_lang_cond(zero_spk_emb, zero_lang_emb)
        
        # --- Sigmoid (Logit-Normal matched) Timestep Schedule ---
        logits = torch.linspace(-3.0, 3.0, steps=n_steps + 1, device=device)
        t_steps = torch.sigmoid(logits)
        
        # Enforce strict boundary conditions so the solver starts at 0 and ends perfectly at 1
        t_steps[0] = 0.0
        t_steps[-1] = 1.0
        
        for step in range(n_steps):
            # Calculate dynamic step size (dt)
            t_curr = t_steps[step]
            t_next = t_steps[step + 1]
            dt = t_next - t_curr  
            
            # --- Euler Step: Calculate v at (x, t_curr) ---
            t_tensor = torch.full((B,), t_curr, device=device)
            adaln_cond = self.flow_model._build_adaln_cond(t_tensor, spk_lang_emb)
            cond_hidden, _ = self.flow_model.condition_encoder(cond, adaln_cond, codec_mask, return_intermediate=False)
            v_cond = self.flow_model.velocity_decoder(x, cond_hidden, adaln_cond, codec_mask)
            
            if cfg_strength != 1.0:
                zero_adaln_cond = self.flow_model._build_adaln_cond(t_tensor, zero_spk_lang_emb)
                cond_hidden_uncond, _ = self.flow_model.condition_encoder(zero_cond, zero_adaln_cond, codec_mask, return_intermediate=False)
                v_uncond = self.flow_model.velocity_decoder(x, cond_hidden_uncond, zero_adaln_cond, codec_mask)
                
                # CFG formula
                v = v_uncond + cfg_strength * (v_cond - v_uncond)
            else:
                v = v_cond
            
            # Advance x directly using the Euler method
            x = x + v * dt
        
        pred_indices = self.continuous_to_indices(x)
        
        return {
            'pred_embeddings': x,
            'pred_indices': pred_indices,
            'attention': attn,
        }
    
    def _length_to_mask(self, lengths: Tensor, max_len: int) -> Tensor:
        rng = torch.arange(max_len, device=lengths.device).unsqueeze(0)
        return rng < lengths.unsqueeze(1)

    @torch.no_grad()
    def infer_durations(
        self,
        phonemes: Tensor,
        phoneme_lengths: Tensor,
        speaker_ids: Tensor,
        language_ids: Tensor,
        ctx_tokens: Optional[Tensor] = None,
        ctx_lens: Optional[Tensor] = None,
        token_start_indices: Optional[Tensor] = None,
        token_end_indices: Optional[Tensor] = None,
        noise_scale: float = 0.667,
        length_scale: float = 1.0,
    ) -> dict:
        """
        Run the stochastic duration predictor in reverse to obtain predicted integer
        durations from phoneme tokens + speaker embedding.

        Returns:
            durations:      [B, T]  predicted integer frame counts per phoneme
            frame_lengths:  [B]     total 100Hz frame count per sample
            codec_lengths:  [B]     total 50Hz codec frame count per sample
        """
        B = phonemes.shape[0]
        device = phonemes.device
        T_text = phonemes.shape[1]

        text_mask = self._length_to_mask(phoneme_lengths, T_text)

        # Encode text (same path as forward)
        if (ctx_tokens is not None and token_start_indices is not None
                and token_end_indices is not None and ctx_lens is not None):
            T_ctx = ctx_tokens.shape[1]
            ctx_mask = self._length_to_mask(ctx_lens, T_ctx)
            ctx_encoded = self.text_encoder(ctx_tokens, language_ids, ctx_mask)
            text_emb = self._extract_current_context_embeddings(
                ctx_encoded, token_start_indices, token_end_indices, T_text,
            )
            text_emb = text_emb * text_mask.unsqueeze(-1).float()
        else:
            text_emb = self.text_encoder(phonemes, language_ids, text_mask)

        spk_emb = self.speaker_embedding(speaker_ids)          # [B, speaker_dim]
        lang_emb_dit = self.language_embedding_dit(language_ids) # [B, language_dim]

        # Build spoken-token mask for SDP
        is_spoken = self.spoken_mask_tensor[phonemes]       # [B, T]
        dur_valid_mask = text_mask & is_spoken              # [B, T]
        dur_mask_1 = dur_valid_mask.unsqueeze(1)            # [B, 1, T]

        text_emb_t = text_emb.transpose(1, 2)              # [B, C, T]

        # SDP reverse pass → log durations
        log_w = self.duration_predictor(
            text_emb_t, dur_mask_1,
            g=torch.cat([spk_emb, lang_emb_dit], dim=-1),
            reverse=True, noise_scale=noise_scale,
        )  # [B, 1, T]

        w = torch.exp(log_w) * dur_mask_1.float() * length_scale
        w = torch.where(dur_mask_1.bool(), torch.clamp(w, min=1.0), w)
        w = torch.ceil(w).squeeze(1)    # [B, T]  integer durations

        frame_lengths = w.sum(dim=-1).long()                                        # [B]
        codec_lengths = (frame_lengths + self.codec_stride - 1) // self.codec_stride  # [B]

        return {
            'durations':      w,               # [B, T]
            'frame_lengths':  frame_lengths,   # [B]
            'codec_lengths':  codec_lengths,   # [B]
        }