"""Matcha-TTS inspired conditional flow matching for VITS latent features.

The flow-matching formulation and the masked 1-D U-Net layout are adapted from
Matcha-TTS (MIT License).  This implementation only depends on PyTorch and is
conditioned on Style-Bert-VITS2's aligned text prior and speaker embedding.
"""

import math
from typing import Optional

import torch
from torch import nn
from torch.nn import functional as F


def configure_matcha_only_training(model: nn.Module) -> int:
    """Freeze a synthesizer except for its Matcha branch.

    Returns the number of trainable parameters after configuration.
    """
    matcha = getattr(model, "matcha", None)
    if not getattr(model, "use_matcha", False) or matcha is None:
        raise ValueError("Matcha-only training requires an enabled Matcha branch")
    for parameter in model.parameters():
        parameter.requires_grad = False
    for parameter in matcha.parameters():
        parameter.requires_grad = True
    return sum(parameter.numel() for parameter in matcha.parameters())


class SinusoidalTimeEmbedding(nn.Module):
    def __init__(self, channels: int) -> None:
        super().__init__()
        if channels % 2:
            raise ValueError("Time embedding channels must be even")
        self.channels = channels

    def forward(self, time: torch.Tensor) -> torch.Tensor:
        if time.ndim == 0:
            time = time.unsqueeze(0)
        half = self.channels // 2
        scale = -math.log(10000) / max(half - 1, 1)
        frequencies = torch.exp(
            torch.arange(half, device=time.device, dtype=torch.float32) * scale
        ).to(time.dtype)
        angles = 1000 * time[:, None] * frequencies[None]
        return torch.cat((angles.sin(), angles.cos()), dim=-1)


class MaskedResBlock(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, time_channels: int) -> None:
        super().__init__()
        groups = math.gcd(out_channels, 8)
        self.conv1 = nn.Conv1d(in_channels, out_channels, 3, padding=1)
        self.norm1 = nn.GroupNorm(groups, out_channels)
        self.conv2 = nn.Conv1d(out_channels, out_channels, 3, padding=1)
        self.norm2 = nn.GroupNorm(groups, out_channels)
        self.time_proj = nn.Linear(time_channels, out_channels)
        self.residual = (
            nn.Identity()
            if in_channels == out_channels
            else nn.Conv1d(in_channels, out_channels, 1)
        )

    def forward(
        self, x: torch.Tensor, mask: torch.Tensor, time: torch.Tensor
    ) -> torch.Tensor:
        residual = self.residual(x * mask)
        x = F.mish(self.norm1(self.conv1(x * mask)))
        x = x + self.time_proj(F.mish(time)).unsqueeze(-1)
        x = self.norm2(self.conv2(x * mask))
        return F.mish(x + residual) * mask


class MaskedTransformerBlock(nn.Module):
    def __init__(
        self,
        channels: int,
        num_heads: int,
        dropout: float,
        use_diff_attention: bool,
    ) -> None:
        super().__init__()
        if channels % num_heads:
            raise ValueError("Matcha channels must be divisible by attention heads")
        self.norm1 = nn.LayerNorm(channels)
        self.use_diff_attention = use_diff_attention
        self.attention = (
            DifferentialAttentionV2(channels, num_heads, dropout)
            if use_diff_attention
            else nn.MultiheadAttention(
                channels, num_heads, dropout=dropout, batch_first=True
            )
        )
        self.norm2 = nn.LayerNorm(channels)
        self.feed_forward = nn.Sequential(
            nn.Linear(channels, channels * 4),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(channels * 4, channels),
        )

    def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        x = x.transpose(1, 2)
        valid = mask[:, 0].bool()
        hidden = self.norm1(x)
        if self.use_diff_attention:
            hidden = self.attention(hidden, valid)
        else:
            hidden, _ = self.attention(
                hidden, hidden, hidden, key_padding_mask=~valid, need_weights=False
            )
        x = x + hidden
        x = x + self.feed_forward(self.norm2(x))
        return x.transpose(1, 2) * mask


class DifferentialAttentionV2(nn.Module):
    """Differential Attention V2 without FlashAttention or custom kernels.

    Each logical head has two query heads that share one key/value head. Their
    contexts are combined as ``context_1 - sigmoid(lambda) * context_2``, where
    lambda is projected per token and logical head.
    """

    def __init__(self, channels: int, num_heads: int, dropout: float) -> None:
        super().__init__()
        if channels % num_heads:
            raise ValueError("Channels must be divisible by attention heads")
        self.num_heads = num_heads
        self.head_channels = channels // num_heads
        self.scale = self.head_channels**-0.5
        self.query_projection = nn.Linear(channels, channels * 2)
        self.key_projection = nn.Linear(channels, channels)
        self.value_projection = nn.Linear(channels, channels)
        self.lambda_projection = nn.Linear(channels, num_heads)
        self.output_projection = nn.Linear(channels, channels)
        self.attention_dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, valid: torch.Tensor) -> torch.Tensor:
        batch, length, channels = x.shape
        queries = self.query_projection(x).view(
            batch, length, self.num_heads, 2, self.head_channels
        )
        keys = self.key_projection(x).view(
            batch, length, self.num_heads, self.head_channels
        )
        values = self.value_projection(x).view(
            batch, length, self.num_heads, self.head_channels
        )

        queries = queries.permute(0, 2, 3, 1, 4)
        keys = keys.permute(0, 2, 1, 3)
        values = values.permute(0, 2, 1, 3)
        scores = torch.matmul(queries, keys.unsqueeze(2).transpose(-1, -2))
        scores = scores * self.scale
        scores = scores.masked_fill(~valid[:, None, None, None, :], float("-inf"))
        attention = self.attention_dropout(torch.softmax(scores, dim=-1))
        contexts = torch.matmul(attention, values.unsqueeze(2))

        lambda_value = torch.sigmoid(self.lambda_projection(x))
        lambda_value = lambda_value.permute(0, 2, 1).unsqueeze(-1)
        contexts = contexts[:, :, 0] - lambda_value * contexts[:, :, 1]
        contexts = contexts.permute(0, 2, 1, 3).reshape(batch, length, channels)
        contexts = self.output_projection(contexts)
        return contexts * valid.unsqueeze(-1).to(contexts.dtype)


class MatchaEstimator(nn.Module):
    """A compact masked 1-D U-Net velocity estimator."""

    def __init__(
        self,
        latent_channels: int,
        speaker_channels: int,
        channels: int,
        num_heads: int,
        dropout: float,
        use_diff_attention: bool,
    ) -> None:
        super().__init__()
        time_channels = channels * 4
        input_channels = latent_channels * 2 + speaker_channels
        self.time_embedding = SinusoidalTimeEmbedding(latent_channels)
        self.time_mlp = nn.Sequential(
            nn.Linear(latent_channels, time_channels),
            nn.SiLU(),
            nn.Linear(time_channels, time_channels),
        )
        self.down_block = MaskedResBlock(input_channels, channels, time_channels)
        self.down_attention = MaskedTransformerBlock(
            channels, num_heads, dropout, use_diff_attention
        )
        self.downsample = nn.Conv1d(channels, channels, 3, stride=2, padding=1)
        self.mid_block = MaskedResBlock(channels, channels, time_channels)
        self.mid_attention = MaskedTransformerBlock(
            channels, num_heads, dropout, use_diff_attention
        )
        self.up_block = MaskedResBlock(channels * 2, channels, time_channels)
        self.up_attention = MaskedTransformerBlock(
            channels, num_heads, dropout, use_diff_attention
        )
        self.output = nn.Conv1d(channels, latent_channels, 1)

    def forward(
        self,
        x: torch.Tensor,
        mask: torch.Tensor,
        mu: torch.Tensor,
        time: torch.Tensor,
        speaker: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        if time.ndim == 0:
            time = time.expand(x.shape[0])
        time_embedding = self.time_mlp(self.time_embedding(time))
        inputs = [x, mu]
        if speaker is not None:
            inputs.append(speaker.expand(-1, -1, x.shape[-1]))
        x = torch.cat(inputs, dim=1)

        skip = self.down_attention(self.down_block(x, mask, time_embedding), mask)
        down_mask = F.interpolate(mask, size=(skip.shape[-1] + 1) // 2, mode="nearest")
        x = self.downsample(skip * mask)
        x = self.mid_attention(self.mid_block(x, down_mask, time_embedding), down_mask)
        x = F.interpolate(x, size=skip.shape[-1], mode="nearest")
        x = torch.cat((x, skip), dim=1)
        x = self.up_attention(self.up_block(x, mask, time_embedding), mask)
        return self.output(x * mask) * mask


class MatchaFlow(nn.Module):
    """Conditional flow matching decoder operating in VITS prior space."""

    def __init__(
        self,
        latent_channels: int,
        speaker_channels: int,
        channels: int = 192,
        num_heads: int = 2,
        dropout: float = 0.05,
        sigma_min: float = 1e-4,
        use_diff_attention: bool = False,
    ) -> None:
        super().__init__()
        self.sigma_min = sigma_min
        self.estimator = MatchaEstimator(
            latent_channels,
            speaker_channels,
            channels,
            num_heads,
            dropout,
            use_diff_attention,
        )

    def compute_loss(
        self,
        target: torch.Tensor,
        mask: torch.Tensor,
        mu: torch.Tensor,
        speaker: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        batch = target.shape[0]
        time = torch.rand(batch, 1, 1, device=target.device, dtype=target.dtype)
        noise = torch.randn_like(target)
        sample = (1 - (1 - self.sigma_min) * time) * noise + time * target
        velocity = target - (1 - self.sigma_min) * noise
        prediction = self.estimator(
            sample, mask, mu, time.flatten(), speaker=speaker
        )
        error = (prediction - velocity).square() * mask
        return error.sum() / (mask.sum().clamp_min(1) * target.shape[1])

    def sample(
        self,
        mu: torch.Tensor,
        mask: torch.Tensor,
        n_timesteps: int,
        temperature: float = 1.0,
        speaker: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        if n_timesteps < 1:
            raise ValueError("Matcha inference requires at least one timestep")
        x = torch.randn_like(mu) * temperature
        dt = 1.0 / n_timesteps
        for step in range(n_timesteps):
            time = torch.full(
                (mu.shape[0],), step * dt, device=mu.device, dtype=mu.dtype
            )
            x = x + dt * self.estimator(x, mask, mu, time, speaker=speaker)
        return x * mask

    @torch.inference_mode()
    def forward(
        self,
        mu: torch.Tensor,
        mask: torch.Tensor,
        n_timesteps: int,
        temperature: float = 1.0,
        speaker: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Sample without building a graph (normal inference path)."""
        return self.sample(mu, mask, n_timesteps, temperature, speaker)