"""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


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) -> None:
        super().__init__()
        if channels % num_heads:
            raise ValueError("Matcha channels must be divisible by attention heads")
        self.norm1 = nn.LayerNorm(channels)
        self.attention = 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)
        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 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,
    ) -> 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)
        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)
        self.up_block = MaskedResBlock(channels * 2, channels, time_channels)
        self.up_attention = MaskedTransformerBlock(channels, num_heads, dropout)
        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,
    ) -> None:
        super().__init__()
        self.sigma_min = sigma_min
        self.estimator = MatchaEstimator(
            latent_channels, speaker_channels, channels, num_heads, dropout
        )

    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])

    @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:
        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