# predictor_duration_speaker_context.py
# ==============================================================================
# Duration Predictor with Speaker ID + Context Window Conditioning
#
# Architecture:
#   - Text encoder (shared ConvNeXtV2TextEncoder from predictor_fsq_flow_convnext)
#     encodes the full context window (prev + curr + next concatenated) in one pass
#   - _extract_current_context_embeddings slices out only the current
#     sentence embeddings by token_start/end indices
#   - Stochastic Duration Predictor (Flow-based) with speaker conditioning
# ==============================================================================

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Optional
from Modules.transforms import piecewise_rational_quadratic_transform

# Import from meldataset
from meldataset import SPOKEN_IDS

# ==============================================================================
# 1. UTILS & BASIC LAYERS
# ==============================================================================

def length_to_mask(lengths: torch.Tensor, max_len: Optional[int] = None) -> torch.Tensor:
    max_len = int(max_len or lengths.max().item())
    rng = torch.arange(max_len, device=lengths.device).unsqueeze(0)
    mask = rng < lengths.unsqueeze(1)
    return mask


# ==============================================================================
# 2. SDP COMPONENTS (Flow Layers)
# ==============================================================================

class Flip(nn.Module):
    def forward(self, x, *args, reverse=False, **kwargs):
        x = torch.flip(x, [1])
        if not reverse:
            logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)
            return x, logdet
        else:
            return x

class ElementwiseAffine(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.channels = channels
        self.m = nn.Parameter(torch.zeros(channels, 1))
        self.logs = nn.Parameter(torch.zeros(channels, 1))

    def forward(self, x, x_mask, reverse=False, **kwargs):
        if not reverse:
            y = self.m + torch.exp(self.logs) * x
            y = y * x_mask
            logdet = torch.sum(self.logs * x_mask, [1, 2])
            return y, logdet
        else:
            x = (x - self.m) * torch.exp(-self.logs) * x_mask
            return x

class Log(nn.Module):
    def forward(self, x, x_mask, reverse=False, **kwargs):
        if not reverse:
            y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask
            logdet = torch.sum(-y, [1, 2])
            return y, logdet
        else:
            x = torch.exp(x) * x_mask
            return x

class LayerNorm(nn.Module):
    def __init__(self, channels, eps=1e-5):
        super().__init__()
        self.channels = channels
        self.eps = eps
        self.gamma = nn.Parameter(torch.ones(channels))
        self.beta = nn.Parameter(torch.zeros(channels))

    def forward(self, x):
        x = x.transpose(1, -1)
        x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
        return x.transpose(1, -1)

class DDSConv(nn.Module):
    def __init__(self, channels, kernel_size, n_layers, p_dropout=0.0):
        super().__init__()
        self.n_layers = n_layers
        self.drop = nn.Dropout(p_dropout)
        self.convs_sep = nn.ModuleList()
        self.linears = nn.ModuleList()
        self.norms_1 = nn.ModuleList()
        self.norms_2 = nn.ModuleList()
        for i in range(n_layers):
            dilation = kernel_size**i
            padding = (kernel_size * dilation - dilation) // 2
            self.convs_sep.append(nn.Conv1d(channels, channels, kernel_size, groups=channels, dilation=dilation, padding=padding))
            self.linears.append(nn.Linear(channels, channels))
            self.norms_1.append(LayerNorm(channels))
            self.norms_2.append(LayerNorm(channels))

    def forward(self, x, x_mask, g=None):
        if g is not None:
            x = x + g
        for i in range(self.n_layers):
            y = self.convs_sep[i](x * x_mask)
            y = self.norms_1[i](y)
            y = F.gelu(y)
            y = self.linears[i](y.mT).mT
            y = self.norms_2[i](y)
            y = F.gelu(y)
            y = self.drop(y)
            x = x + y
        return x * x_mask

class ConvFlow(nn.Module):
    def __init__(self, in_channels, filter_channels, kernel_size, n_layers, num_bins=16, tail_bound=5.0):
        super().__init__()
        self.filter_channels = filter_channels
        self.num_bins = num_bins
        self.tail_bound = tail_bound
        self.half_channels = in_channels // 2
        self.pre = nn.Linear(self.half_channels, filter_channels)
        self.convs = DDSConv(filter_channels, kernel_size, n_layers, p_dropout=0.0)
        self.proj = nn.Linear(filter_channels, self.half_channels * (num_bins * 3 - 1))

    def forward(self, x, x_mask, g=None, reverse=False):
        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
        h = self.pre(x0.mT).mT
        h = self.convs(h, x_mask, g=g)
        h = self.proj(h.mT).mT * x_mask
        
        b, c, t = x0.shape
        h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2)
        
        unnormalized_widths = h[..., : self.num_bins] / math.sqrt(self.filter_channels)
        unnormalized_heights = h[..., self.num_bins : 2 * self.num_bins] / math.sqrt(self.filter_channels)
        unnormalized_derivatives = h[..., 2 * self.num_bins :]

        x1, logabsdet = piecewise_rational_quadratic_transform(
            x1, unnormalized_widths, unnormalized_heights, unnormalized_derivatives, 
            inverse=reverse, tails="linear", tail_bound=self.tail_bound
        )
        x = torch.cat([x0, x1], 1) * x_mask
        logdet = torch.sum(logabsdet * x_mask, [1, 2])
        if not reverse:
            return x, logdet
        else:
            return x


# ==============================================================================
# 3. STOCHASTIC DURATION PREDICTOR WITH CONTEXT
# ==============================================================================

class StochasticDurationPredictorContext(nn.Module):
    """
    Stochastic Duration Predictor with speaker + context conditioning.
    """
    def __init__(self, in_channels, filter_channels, kernel_size, p_dropout, n_flows=4,
                 speaker_emb_dim=256):
        super().__init__()
        self.log_flow = Log()
        self.flows = nn.ModuleList()
        self.flows.append(ElementwiseAffine(2))
        for i in range(n_flows):
            self.flows.append(ConvFlow(2, filter_channels, kernel_size, n_layers=3))
            self.flows.append(Flip())

        self.pre = nn.Linear(in_channels, filter_channels)
        self.convs = DDSConv(filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout)
        self.proj = nn.Linear(filter_channels, filter_channels)

        self.post_pre = nn.Linear(1, filter_channels)
        self.post_convs = DDSConv(filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout)
        self.post_proj = nn.Linear(filter_channels, filter_channels)

        self.post_flows = nn.ModuleList()
        self.post_flows.append(ElementwiseAffine(2))
        for i in range(4):
            self.post_flows.append(ConvFlow(2, filter_channels, kernel_size, n_layers=3))
            self.post_flows.append(Flip())
        
        # Speaker conditioning
        if speaker_emb_dim > 0:
            self.cond_speaker = nn.Linear(speaker_emb_dim, filter_channels)

    def forward(self, x, x_mask, w=None, g=None, reverse=False, noise_scale=1.0):
        """
        Args:
            x: [B, C, T] text features
            x_mask: [B, 1, T] mask
            w: [B, 1, T] ground truth durations (training only)
            g: [B, speaker_emb_dim] speaker embedding
            reverse: inference mode
            noise_scale: sampling temperature
        """
        x = self.pre(x.mT).mT
        
        # Add speaker conditioning
        if g is not None and hasattr(self, 'cond_speaker'):
            g_expanded = g.unsqueeze(-1).expand(-1, -1, x.size(2))
            x = x + self.cond_speaker(g_expanded.mT).mT
        
        x = self.convs(x, x_mask)
        x = self.proj(x.mT).mT * x_mask

        if not reverse:
            # Training
            h_w = self.post_pre(w.mT).mT
            h_w = self.post_convs(h_w, x_mask)
            h_w = self.post_proj(h_w.mT).mT * x_mask
            
            e_q = torch.randn(w.size(0), 2, w.size(2)).to(device=x.device, dtype=x.dtype) * x_mask
            z_q = e_q
            logdet_tot_q = 0
            for flow in self.post_flows:
                z_q, logdet_q = flow(z_q, x_mask, g=(x + h_w))
                logdet_tot_q += logdet_q
            
            z_u, z1 = torch.split(z_q, [1, 1], 1)
            u = torch.sigmoid(z_u) * x_mask
            z0 = (w - u) * x_mask
            logdet_tot_q += torch.sum((F.logsigmoid(z_u) + F.logsigmoid(-z_u)) * x_mask, [1, 2])
            logq = torch.sum(-0.5 * (math.log(2 * math.pi) + (e_q**2)) * x_mask, [1, 2]) - logdet_tot_q

            logdet_tot = 0
            z0, logdet = self.log_flow(z0, x_mask)
            logdet_tot += logdet
            z = torch.cat([z0, z1], 1)
            for flow in self.flows:
                z, logdet = flow(z, x_mask, g=x, reverse=reverse)
                logdet_tot += logdet
            nll = torch.sum(0.5 * (math.log(2 * math.pi) + (z**2)) * x_mask, [1, 2]) - logdet_tot
            return nll + logq
        else:
            # Inference
            flows = list(reversed(self.flows))
            flows = flows[:-2] + [flows[-1]]
            z = torch.randn(x.size(0), 2, x.size(2)).to(device=x.device, dtype=x.dtype) * noise_scale
            for flow in flows:
                z = flow(z, x_mask, g=x, reverse=reverse)
            z0, z1 = torch.split(z, [1, 1], 1)
            return z0