"""
Human Prosody & Pitch (F0) Engine for ViuAI_TTS_200M.
Predicts frame durations, fundamental frequency (F0 pitch intonation),
and volume/energy dynamics to ensure 100% natural, expressive human prosody.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple, Optional


class ConvNeXtBlock(nn.Module):
    def __init__(self, dim: int, kernel_size: int = 7):
        super().__init__()
        self.dwconv = nn.Conv1d(dim, dim, kernel_size=kernel_size, padding=kernel_size // 2, groups=dim)
        self.norm = nn.LayerNorm(dim, eps=1e-6)
        self.pwconv1 = nn.Linear(dim, 4 * dim)
        self.act = nn.GELU()
        self.pwconv2 = nn.Linear(4 * dim, dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: [B, T, C]
        residual = x
        x = self.dwconv(x.transpose(1, 2)).transpose(1, 2)
        x = self.norm(x)
        x = self.pwconv1(x)
        x = self.act(x)
        x = self.pwconv2(x)
        return residual + x


class HumanProsodyEngine(nn.Module):
    """
    State-of-the-Art Prosody Modeling Engine.
    Predicts:
      1. Duration (speaking pace and micro-pauses)
      2. Pitch (F0) contour (emotional inflection, question vs statement tone)
      3. Energy contour (breathiness, volume variation)
    """
    def __init__(
        self,
        hidden_dim: int = 512,
        num_layers: int = 4,
        kernel_size: int = 7,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.shared_blocks = nn.ModuleList([
            ConvNeXtBlock(hidden_dim, kernel_size=kernel_size)
            for _ in range(num_layers)
        ])
        self.dropout = nn.Dropout(dropout)

        # Output projection heads
        self.duration_head = nn.Linear(hidden_dim, 1)
        self.pitch_head = nn.Linear(hidden_dim, 1)
        self.energy_head = nn.Linear(hidden_dim, 1)

        # Embeddings to project pitch and energy back into acoustic features
        self.pitch_proj = nn.Linear(1, hidden_dim)
        self.energy_proj = nn.Linear(1, hidden_dim)

    def forward(
        self,
        text_latents: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Args:
            text_latents: [B, T_text, hidden_dim]
        Returns:
            log_duration: [B, T_text]
            pitch_pred:   [B, T_text]
            energy_pred:  [B, T_text]
        """
        x = text_latents
        for block in self.shared_blocks:
            x = block(x)
            x = self.dropout(x)

        log_duration = self.duration_head(x).squeeze(-1)
        pitch_pred = self.pitch_head(x).squeeze(-1)
        energy_pred = self.energy_head(x).squeeze(-1)

        return log_duration, pitch_pred, energy_pred

    def get_prosody_features(self, pitch: torch.Tensor, energy: torch.Tensor) -> torch.Tensor:
        """Projects pitch and energy into hidden_dim vectors."""
        p = self.pitch_proj(pitch.unsqueeze(-1))
        e = self.energy_proj(energy.unsqueeze(-1))
        return p + e


# Maintain backward compatibility alias
DurationPredictor = HumanProsodyEngine