"""Loss functions for SalTTS training."""

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


class ProsodyLoss(nn.Module):
    """Multi-task loss for prosody prediction."""

    def __init__(
        self,
        pitch_weight: float = 1.0,
        energy_weight: float = 1.0,
        duration_weight: float = 1.0,
    ):
        """Initialize prosody loss.

        Args:
            pitch_weight: Weight for pitch loss
            energy_weight: Weight for energy loss
            duration_weight: Weight for duration loss
        """
        super().__init__()
        self.pitch_weight = pitch_weight
        self.energy_weight = energy_weight
        self.duration_weight = duration_weight

    def forward(
        self,
        pred_pitch: torch.Tensor,
        target_pitch: torch.Tensor,
        pred_energy: torch.Tensor,
        target_energy: torch.Tensor,
        pred_duration: torch.Tensor,
        target_duration: torch.Tensor,
    ) -> torch.Tensor:
        """Compute prosody loss."""
        pitch_loss = F.mse_loss(pred_pitch, target_pitch)
        energy_loss = F.mse_loss(pred_energy, target_energy)
        duration_loss = F.mse_loss(pred_duration, target_duration)

        total_loss = (
            self.pitch_weight * pitch_loss
            + self.energy_weight * energy_loss
            + self.duration_weight * duration_loss
        )

        return total_loss


class PerceptualLoss(nn.Module):
    """Perceptual loss for audio quality."""

    def __init__(self):
        super().__init__()

    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """Compute perceptual loss."""
        # Simple L1 loss as placeholder
        return F.l1_loss(pred, target)


class FlowNLLLoss(nn.Module):
    """Negative log-likelihood loss for flow models."""

    def __init__(self):
        super().__init__()

    def forward(self, z: torch.Tensor, log_det: torch.Tensor) -> torch.Tensor:
        """Compute flow NLL loss."""
        nll = 0.5 * torch.sum(z**2, dim=[1, 2]) - log_det
        return nll.mean()


class AdversarialLoss(nn.Module):
    """Adversarial loss for GAN training."""

    def __init__(self, loss_type: str = "hinge"):
        """Initialize adversarial loss.

        Args:
            loss_type: Loss type (hinge, lsgan, vanilla)
        """
        super().__init__()
        self.loss_type = loss_type

    def discriminator_loss(
        self,
        real_logits: torch.Tensor,
        fake_logits: torch.Tensor,
    ) -> torch.Tensor:
        """Compute discriminator loss."""
        if self.loss_type == "hinge":
            real_loss = F.relu(1.0 - real_logits).mean()
            fake_loss = F.relu(1.0 + fake_logits).mean()
        else:
            real_loss = F.binary_cross_entropy_with_logits(
                real_logits, torch.ones_like(real_logits)
            )
            fake_loss = F.binary_cross_entropy_with_logits(
                fake_logits, torch.zeros_like(fake_logits)
            )

        return real_loss + fake_loss

    def generator_loss(self, fake_logits: torch.Tensor) -> torch.Tensor:
        """Compute generator loss."""
        if self.loss_type == "hinge":
            return -fake_logits.mean()
        else:
            return F.binary_cross_entropy_with_logits(
                fake_logits, torch.ones_like(fake_logits)
            )
