"""The conformer that turns prosody features into a mel spectrogram.

Extracted verbatim from the training code -- `prosody_embeddings_to_mel` --
so the architecture here is the one the checkpoint was trained with rather
than a reimplementation of it. Only what inference touches: the model
classes and `generate()`, with no Trainer, dataset or metric code.

`generate()` is non-autoregressive. It sizes its output from the real input
length and takes a padding mask, so a phrase is predicted at its own
duration -- which is why this is the path the published numbers come from.
The ONNX export cannot do that: its input width is baked in at 2000 frames,
so every phrase is padded to ten seconds and the decoder attends over the
padding. Measured against the training pipeline's own predictions, 0.9985
correlation here against 0.9357 through the export.
"""

from __future__ import annotations

import math

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint as grad_checkpoint


class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer("pe", pe)

    def forward(self, x):
        # x is [B, T, D] (batch_first)
        return x + self.pe[: x.size(1)]


# ----------------------------
# Conformer Block
# ----------------------------


class ConformerBlock(nn.Module):
    """Single conformer block: FFN -> Self-Attention -> Conv -> FFN.

    Combines local convolutions with self-attention for both local
    detail and global context, at lower cost than a full transformer
    decoder (conv handles local, attention handles global).
    """
    def __init__(self, d_model, nhead=4, kernel_size=5, dim_feedforward=None, dropout=0.1):
        super().__init__()
        dim_ff = dim_feedforward or d_model * 4

        # First half-step FFN
        self.ffn1_norm = nn.LayerNorm(d_model)
        self.ffn1 = nn.Sequential(
            nn.Linear(d_model, dim_ff), nn.GELU(), nn.Dropout(dropout),
            nn.Linear(dim_ff, d_model), nn.Dropout(dropout),
        )

        # Self-attention
        self.attn_norm = nn.LayerNorm(d_model)
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
        self.attn_dropout = nn.Dropout(dropout)

        # Depthwise convolution module
        self.conv_norm = nn.LayerNorm(d_model)
        self.conv = nn.Sequential(
            nn.Conv1d(d_model, d_model, 1),  # pointwise
            nn.GELU(),
            nn.Conv1d(d_model, d_model, kernel_size, padding=kernel_size // 2, groups=d_model),  # depthwise
            nn.LayerNorm(d_model),  # NOTE: applied after transpose in forward
            nn.GELU(),
            nn.Conv1d(d_model, d_model, 1),  # pointwise
            nn.Dropout(dropout),
        )
        # Replace the LayerNorm inside Sequential with identity — we apply it manually
        self.conv_ln = self.conv[3]
        self.conv[3] = nn.Identity()

        # Second half-step FFN
        self.ffn2_norm = nn.LayerNorm(d_model)
        self.ffn2 = nn.Sequential(
            nn.Linear(d_model, dim_ff), nn.GELU(), nn.Dropout(dropout),
            nn.Linear(dim_ff, d_model), nn.Dropout(dropout),
        )

        self.final_norm = nn.LayerNorm(d_model)

    def forward(self, x, padding_mask=None):
        """
        Args:
            x: [B, T, D]
            padding_mask: [B, T] with True for padding
        Returns:
            [B, T, D]
        """
        # Half-step FFN
        x = x + 0.5 * self.ffn1(self.ffn1_norm(x))

        # Self-attention
        h = self.attn_norm(x)
        attn_out, _ = self.self_attn(h, h, h, key_padding_mask=padding_mask)
        x = x + self.attn_dropout(attn_out)

        # Conv module
        h = self.conv_norm(x)
        h = h.transpose(1, 2)  # [B, D, T]
        # Apply conv layers, but handle LayerNorm which needs [B, T, D]
        h = self.conv[0](h)  # pointwise
        h = self.conv[1](h)  # GELU
        h = self.conv[2](h)  # depthwise
        h = h.transpose(1, 2)  # [B, T, D]
        h = self.conv_ln(h)   # LayerNorm
        h = self.conv[4](h.transpose(1, 2))  # GELU (expects [B, D, T] from Conv1d but GELU is elementwise)
        h = self.conv[5](h)  # pointwise
        h = self.conv[6](h)  # dropout
        h = h.transpose(1, 2)  # [B, T, D]
        x = x + h

        # Half-step FFN
        x = x + 0.5 * self.ffn2(self.ffn2_norm(x))

        x = self.final_norm(x)

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

        return x


class ConformerDecoder(nn.Module):
    """Stack of conformer blocks as a decoder."""
    def __init__(self, d_model, num_layers=4, nhead=4, kernel_size=5, dim_feedforward=None, dropout=0.1):
        super().__init__()
        self.blocks = nn.ModuleList([
            ConformerBlock(d_model, nhead, kernel_size, dim_feedforward, dropout)
            for _ in range(num_layers)
        ])

    def forward(self, x, padding_mask=None):
        for block in self.blocks:
            x = block(x, padding_mask)
        return x


# ----------------------------
# Model
# ----------------------------


class WindowedConformerBlock(ConformerBlock):
    """ConformerBlock with windowed self-attention for O(T*W) memory instead of O(T²).

    Splits the sequence into non-overlapping windows of size W and applies
    self-attention independently within each window. The convolution module
    still provides cross-window local context.
    """
    def __init__(self, d_model, nhead=4, kernel_size=5, dim_feedforward=None, dropout=0.1, attn_window=256):
        super().__init__(d_model, nhead, kernel_size, dim_feedforward, dropout)
        self.attn_window = attn_window

    def forward(self, x, padding_mask=None):
        """Same as ConformerBlock but with windowed attention."""
        # Half-step FFN
        x = x + 0.5 * self.ffn1(self.ffn1_norm(x))

        # Windowed self-attention
        h = self.attn_norm(x)
        attn_out = self._windowed_attention(h, padding_mask)
        x = x + self.attn_dropout(attn_out)

        # Conv module (unchanged — provides cross-window context)
        h = self.conv_norm(x)
        h = h.transpose(1, 2)  # [B, D, T]
        h = self.conv[0](h)  # pointwise
        h = self.conv[1](h)  # GELU
        h = self.conv[2](h)  # depthwise
        h = h.transpose(1, 2)  # [B, T, D]
        h = self.conv_ln(h)   # LayerNorm
        h = self.conv[4](h.transpose(1, 2))  # GELU
        h = self.conv[5](h)  # pointwise
        h = self.conv[6](h)  # dropout
        h = h.transpose(1, 2)  # [B, T, D]
        x = x + h

        # Half-step FFN
        x = x + 0.5 * self.ffn2(self.ffn2_norm(x))

        x = self.final_norm(x)

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

        return x

    def _windowed_attention(self, x, padding_mask):
        """Self-attention within non-overlapping windows.

        Args:
            x: [B, T, D] — already normed
            padding_mask: [B, T] with True for padding, or None

        Returns:
            [B, T, D]
        """
        B, T, D = x.shape
        W = self.attn_window

        # Short sequences: full attention (no windowing needed)
        if T <= W:
            out, _ = self.self_attn(x, x, x, key_padding_mask=padding_mask)
            return out

        # Pad sequence to a multiple of W
        pad_len = (W - T % W) % W
        if pad_len > 0:
            x = torch.nn.functional.pad(x, (0, 0, 0, pad_len))  # pad time dim
            if padding_mask is not None:
                padding_mask = torch.nn.functional.pad(padding_mask, (0, pad_len), value=True)

        T_padded = x.size(1)
        num_windows = T_padded // W

        # Reshape: [B, num_windows, W, D] -> [B*num_windows, W, D]
        x_win = x.reshape(B * num_windows, W, D)

        # Reshape mask: [B, num_windows, W] -> [B*num_windows, W]
        mask_win = None
        if padding_mask is not None:
            mask_win = padding_mask.reshape(B * num_windows, W)
            # If an entire window is padding, MHA will produce NaN.
            # Detect all-padding windows and handle them.
            all_padding = mask_win.all(dim=1)  # [B*num_windows]
            if all_padding.any():
                # Temporarily unmask one position so MHA doesn't get empty input
                mask_win = mask_win.clone()
                mask_win[all_padding, 0] = False

        attn_out, _ = self.self_attn(x_win, x_win, x_win, key_padding_mask=mask_win)

        # Zero out windows that were entirely padding
        if padding_mask is not None and all_padding.any():
            attn_out[all_padding] = 0.0

        # Reshape back: [B*num_windows, W, D] -> [B, T_padded, D] -> [B, T, D]
        attn_out = attn_out.reshape(B, T_padded, D)[:, :T, :]
        return attn_out


class WindowedConformerDecoder(nn.Module):
    """Stack of windowed conformer blocks."""
    def __init__(self, d_model, num_layers=4, nhead=4, kernel_size=5, dim_feedforward=None, dropout=0.1, attn_window=256):
        super().__init__()
        self.blocks = nn.ModuleList([
            WindowedConformerBlock(d_model, nhead, kernel_size, dim_feedforward, dropout, attn_window)
            for _ in range(num_layers)
        ])

    def forward(self, x, padding_mask=None):
        for block in self.blocks:
            x = block(x, padding_mask)
        return x


# ----------------------------
# Convolutional Encoder
# ----------------------------


class ConvEncoder(nn.Module):
    """Multi-layer 1D convolutional encoder with stride-2 downsampling.

    Designed for smooth, low-dimensional prosodic signals (f0 + intensity).
    Uses dilated convolutions to capture multi-scale temporal context efficiently,
    with a final stride-2 layer to handle the 200 Hz → 100 Hz rate conversion.
    """
    def __init__(self, input_dim, d_model, num_layers=4, kernel_size=7, dropout=0.1):
        super().__init__()
        self.input_proj = nn.Conv1d(input_dim, d_model, kernel_size=1)

        layers = []
        for i in range(num_layers - 1):
            dilation = 2 ** i  # 1, 2, 4, 8... expanding receptive field
            padding = (kernel_size - 1) * dilation // 2
            layers.append(ConvBlock(d_model, kernel_size, dilation, padding, dropout))

        self.layers = nn.ModuleList(layers)

        # Final stride-2 conv for downsampling (200 Hz → 100 Hz)
        self.downsample = nn.Conv1d(
            d_model, d_model,
            kernel_size=4, stride=2, padding=1,
        )
        self.downsample_norm = nn.LayerNorm(d_model)

    def forward(self, x, padding_mask=None):
        """
        Args:
            x: [B, T, input_dim]
            padding_mask: [B, T] with True for padding positions

        Returns:
            out: [B, T//2, d_model]
            out_mask: [B, T//2] or None
        """
        # [B, T, C] → [B, C, T] for Conv1d
        h = x.transpose(1, 2)
        h = self.input_proj(h)

        # Zero out padded positions before conv
        if padding_mask is not None:
            h = h.masked_fill(padding_mask.unsqueeze(1), 0.0)

        for layer in self.layers:
            h = layer(h, padding_mask)

        # Stride-2 downsampling
        h = self.downsample(h)  # [B, d_model, T//2]
        h = h.transpose(1, 2)  # [B, T//2, d_model]
        h = self.downsample_norm(h)

        # Compute output padding mask
        out_mask = None
        if padding_mask is not None:
            T_out = h.size(1)
            # A downsampled position is padding if both source frames were padding
            T_in = padding_mask.size(1)
            if T_in % 2 != 0:
                padding_mask = torch.cat([
                    padding_mask,
                    torch.ones(padding_mask.size(0), 1, dtype=torch.bool, device=padding_mask.device)
                ], dim=1)
            mask_pairs = padding_mask.view(padding_mask.size(0), -1, 2)
            out_mask = mask_pairs.all(dim=-1)[:, :T_out]

        if out_mask is not None:
            h = h.masked_fill(out_mask.unsqueeze(-1), 0.0)

        return h, out_mask


class ConvBlock(nn.Module):
    """Residual dilated conv block with pre-norm."""
    def __init__(self, d_model, kernel_size, dilation, padding, dropout):
        super().__init__()
        self.norm = nn.LayerNorm(d_model)
        self.conv = nn.Conv1d(
            d_model, d_model,
            kernel_size=kernel_size,
            dilation=dilation,
            padding=padding,
        )
        self.activation = nn.GELU()
        self.dropout = nn.Dropout(dropout)
        self.proj = nn.Conv1d(d_model, d_model, kernel_size=1)

    def forward(self, x, padding_mask=None):
        """
        Args:
            x: [B, d_model, T]
            padding_mask: [B, T] with True for padding
        Returns:
            [B, d_model, T]
        """
        residual = x
        # Pre-norm (need to transpose for LayerNorm)
        h = x.transpose(1, 2)  # [B, T, D]
        h = self.norm(h)
        h = h.transpose(1, 2)  # [B, D, T]

        h = self.conv(h)
        h = self.activation(h)
        h = self.dropout(h)
        h = self.proj(h)
        h = self.dropout(h)

        out = residual + h

        if padding_mask is not None:
            out = out.masked_fill(padding_mask.unsqueeze(1), 0.0)

        return out


# ----------------------------
# Model
# ----------------------------


class FeaturesToMelConv(nn.Module):
    """Non-autoregressive conv-encoder + Conformer-decoder for f0+intensity to mel.

    Architecture rationale:
    - Conv encoder: F0 and intensity are smooth, slowly-varying 1D signals.
      Dilated convolutions naturally capture multi-scale temporal patterns
      (pitch contours, intensity envelopes) with the right inductive bias.
      The stride-2 final layer handles the 200→100 Hz rate conversion
      as part of the encoding, rather than as a separate step.
    - Conformer decoder: Refines the downsampled features using both local
      (depthwise conv) and global (self-attention) context to produce
      the mel spectrogram. Self-attention here is justified because the
      decoder operates at 100 Hz on d_model-dimensional features — a
      richer representation where global dependencies matter.

    Memory optimization: gradient checkpointing on decoder blocks (recomputes
    activations during backward instead of storing them).
    """
    def __init__(
        self,
        input_dim=2,
        mel_dim=14,
        d_model=256,
        nhead=4,
        num_encoder_layers=4,
        num_decoder_layers=4,
        encoder_kernel_size=7,
        conv_kernel_size=5,
        dim_feedforward=None,
        dropout=0.1,
        max_len=16000,
        loss_type="mse",
        length_ratio=0.5,
        gradient_checkpointing=True,
        decoder_only=False,
        attn_window=0,
    ):
        super().__init__()
        self.d_model = d_model
        self.mel_dim = mel_dim
        self.max_len = max_len
        self.length_ratio = length_ratio
        self.gradient_checkpointing = gradient_checkpointing
        self.decoder_only = decoder_only

        if not decoder_only:
            self.encoder = ConvEncoder(
                input_dim=input_dim,
                d_model=d_model,
                num_layers=num_encoder_layers,
                kernel_size=encoder_kernel_size,
                dropout=dropout,
            )
        else:
            # Decoder-only: simple linear projection + learned stride-2 downsample
            self.input_proj = nn.Linear(input_dim, d_model)
            self.downsample = nn.Linear(d_model * 2, d_model)

        self.pos_decoder = SinusoidalPositionalEncoding(d_model, max_len)

        dim_ff = dim_feedforward if dim_feedforward is not None else d_model * 4
        if attn_window > 0:
            self.decoder = WindowedConformerDecoder(
                d_model=d_model,
                num_layers=num_decoder_layers,
                nhead=nhead,
                kernel_size=conv_kernel_size,
                dim_feedforward=dim_ff,
                dropout=dropout,
                attn_window=attn_window,
            )
        else:
            self.decoder = ConformerDecoder(
                d_model=d_model,
                num_layers=num_decoder_layers,
                nhead=nhead,
                kernel_size=conv_kernel_size,
                dim_feedforward=dim_ff,
                dropout=dropout,
            )

        self.output_proj = nn.Linear(d_model, mel_dim)

        if loss_type == "mse":
            self.criterion = nn.MSELoss(reduction="none")
        elif loss_type == "l1":
            self.criterion = nn.L1Loss(reduction="none")
        elif loss_type == "smoothl1":
            self.criterion = nn.SmoothL1Loss(reduction="none")
        else:
            raise ValueError(f"Unknown loss_type={loss_type}")

    def _run_decoder(self, tgt, padding_mask):
        """Run decoder with optional gradient checkpointing per block."""
        if self.gradient_checkpointing and self.training:
            for block in self.decoder.blocks:
                tgt = grad_checkpoint(block, tgt, padding_mask, use_reentrant=False)
            return tgt
        else:
            return self.decoder(tgt, padding_mask=padding_mask)

    def _encode(self, inputs, input_padding_mask):
        """Encode inputs, returns (encoded, mask) at output frame rate."""
        if not self.decoder_only:
            return self.encoder(inputs, input_padding_mask)

        # Decoder-only path: project + pair-merge downsample
        B, T, _ = inputs.shape
        h = self.input_proj(inputs)  # [B, T, d_model]

        # Pad to even length
        if T % 2 != 0:
            h = torch.cat([h, torch.zeros(B, 1, self.d_model, device=h.device)], dim=1)
            if input_padding_mask is not None:
                input_padding_mask = torch.cat([
                    input_padding_mask,
                    torch.ones(B, 1, dtype=torch.bool, device=h.device)
                ], dim=1)
            T = T + 1

        # Merge pairs: [B, T, D] -> [B, T//2, 2*D] -> [B, T//2, D]
        h = h.view(B, T // 2, 2 * self.d_model)
        h = self.downsample(h)

        # Compute output mask
        out_mask = None
        if input_padding_mask is not None:
            mask_pairs = input_padding_mask.view(B, T // 2, 2)
            out_mask = mask_pairs.all(dim=-1)

        if out_mask is not None:
            h = h.masked_fill(out_mask.unsqueeze(-1), 0.0)

        return h, out_mask

    def forward(self, inputs, labels, input_padding_mask=None, label_padding_mask=None):
        """
        Args:
            inputs: [B, T_src, input_dim] - f0 + intensity at 200 Hz
            labels: [B, T_tgt, mel_dim] - ground truth mel at 100 Hz
            input_padding_mask: [B, T_src] with True for padding
            label_padding_mask: [B, T_tgt] with True for padding

        Returns:
            dict with 'loss' and 'logits'
        """
        B = labels.size(0)
        T_tgt = labels.size(1)

        # Encode + downsample
        enc_out, enc_mask = self._encode(inputs, input_padding_mask)

        # Align to target length
        if enc_out.size(1) < T_tgt:
            pad_len = T_tgt - enc_out.size(1)
            enc_out = torch.cat([enc_out, torch.zeros(B, pad_len, self.d_model, device=enc_out.device)], dim=1)
            if enc_mask is not None:
                enc_mask = torch.cat([enc_mask, torch.ones(B, pad_len, dtype=torch.bool, device=enc_out.device)], dim=1)
        elif enc_out.size(1) > T_tgt:
            enc_out = enc_out[:, :T_tgt, :]
            if enc_mask is not None:
                enc_mask = enc_mask[:, :T_tgt]

        # Add positional encoding for decoder
        tgt = enc_out * math.sqrt(self.d_model)
        tgt = self.pos_decoder(tgt)

        # Conformer decoder refines the features
        output = self._run_decoder(tgt, label_padding_mask)

        logits = self.output_proj(output)  # [B, T_tgt, mel_dim]

        # Compute loss
        loss = self.criterion(logits, labels)
        if label_padding_mask is not None:
            active_mask = ~label_padding_mask
            denom = (active_mask.sum() * logits.size(-1)).clamp_min(1).to(loss.dtype)
            loss = (loss * active_mask.unsqueeze(-1)).sum() / (denom + 1e-8)
        else:
            loss = loss.mean()

        return {"loss": loss, "logits": logits}

    def generate(self, inputs, input_padding_mask=None, max_steps=None, length_ratio=None, return_lengths=False):
        """Non-autoregressive parallel generation."""
        B = inputs.size(0)
        device = inputs.device

        if length_ratio is None:
            length_ratio = self.length_ratio

        if max_steps is None:
            if input_padding_mask is not None:
                input_lengths = (~input_padding_mask).sum(dim=1)
                target_lengths = (input_lengths.float() * length_ratio).round().long()
                T_gen = int(target_lengths.max().item())
            else:
                input_lengths = torch.full((B,), inputs.size(1), device=device, dtype=torch.long)
                target_lengths = (input_lengths.float() * length_ratio).round().long()
                T_gen = int(target_lengths[0].item())
        else:
            T_gen = max_steps
            target_lengths = torch.full((B,), T_gen, device=device, dtype=torch.long)

        T_gen = max(1, min(T_gen, self.max_len))

        output_padding_mask = torch.zeros(B, T_gen, dtype=torch.bool, device=device)
        for i in range(B):
            if target_lengths[i] < T_gen:
                output_padding_mask[i, target_lengths[i]:] = True

        enc_out, _ = self._encode(inputs, input_padding_mask)

        if enc_out.size(1) < T_gen:
            pad_len = T_gen - enc_out.size(1)
            enc_out = torch.cat([enc_out, torch.zeros(B, pad_len, self.d_model, device=device)], dim=1)
        elif enc_out.size(1) > T_gen:
            enc_out = enc_out[:, :T_gen, :]

        tgt = enc_out * math.sqrt(self.d_model)
        tgt = self.pos_decoder(tgt)

        output = self.decoder(tgt, padding_mask=output_padding_mask)
        generated = self.output_proj(output)

        if return_lengths:
            return generated, target_lengths, output_padding_mask
        return generated


# ----------------------------
# Dataset
# ----------------------------