import torch
import torch.nn as nn
from torch.autograd import Function


class GradientReversalFunction(Function):
    @staticmethod
    def forward(ctx, input_tensor, lambda_value):
        ctx.lambda_value = lambda_value
        return input_tensor.view_as(input_tensor)

    @staticmethod
    def backward(ctx, grad_output):
        return -ctx.lambda_value * grad_output, None


class GradientReversalLayer(nn.Module):
    def __init__(self, lambda_value: float = 1.0):
        super().__init__()
        self.lambda_value = lambda_value

    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
        return GradientReversalFunction.apply(input_tensor, self.lambda_value)


class DomainDiscriminator(nn.Module):
    def __init__(self, input_size: int, hidden_size: int = 128):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(hidden_size, 2),
        )

    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
        return self.network(input_tensor)


class ConvolutionalFeatureExtractor(nn.Module):
    def __init__(self):
        super().__init__()
        self.block_1 = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.block_2 = nn.Sequential(
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.block_3 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )

    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
        x = self.block_1(input_tensor)
        x = self.block_2(x)
        x = self.block_3(x)
        return x


class AttentionPooling(nn.Module):
    def __init__(self, feature_size: int):
        super().__init__()
        self.attention_layer = nn.Linear(feature_size, 1)

    def forward(self, sequence_tensor: torch.Tensor) -> torch.Tensor:
        attention_scores = self.attention_layer(sequence_tensor).squeeze(-1)
        attention_weights = torch.softmax(attention_scores, dim=1).unsqueeze(-1)
        pooled_output = torch.sum(sequence_tensor * attention_weights, dim=1)
        return pooled_output


class TemporalAttentionPooling(nn.Module):
    def __init__(self, feature_size: int, attention_hidden_size: int = 128):
        super().__init__()
        self.score_network = nn.Sequential(
            nn.Linear(feature_size, attention_hidden_size),
            nn.Tanh(),
            nn.Linear(attention_hidden_size, 1),
        )

    def forward(self, sequence_tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        attention_scores = self.score_network(sequence_tensor).squeeze(-1)
        attention_weights = torch.softmax(attention_scores, dim=1).unsqueeze(-1)
        pooled_output = torch.sum(sequence_tensor * attention_weights, dim=1)
        return pooled_output, attention_weights


class CnnBiLstmAdversarialModel(nn.Module):
    def __init__(self, num_emotions: int = 7):
        super().__init__()
        self.cnn_feature_extractor = ConvolutionalFeatureExtractor()
        self.bi_lstm = nn.LSTM(
            input_size=128 * 16,
            hidden_size=128,
            num_layers=1,
            batch_first=True,
            bidirectional=True,
        )
        self.attention_pooling = AttentionPooling(feature_size=256)
        self.shared_projection = nn.Sequential(
            nn.Linear(256, 256),
            nn.LayerNorm(256),
            nn.ReLU(),
            nn.Dropout(0.3),
        )
        self.emotion_head = nn.Sequential(
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, num_emotions),
        )
        self.impairment_head = nn.Sequential(
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, 2),
        )
        self.gradient_reversal = GradientReversalLayer(lambda_value=1.0)
        self.domain_discriminator = DomainDiscriminator(256, hidden_size=128)

    def forward(
        self,
        mel_spectrogram: torch.Tensor,
        return_domain_output: bool = True,
    ) -> dict[str, torch.Tensor]:
        cnn_output = self.cnn_feature_extractor(mel_spectrogram)
        batch_size, channels, height, width = cnn_output.shape
        sequence_input = (
            cnn_output.permute(0, 3, 1, 2).contiguous().view(batch_size, width, channels * height)
        )

        lstm_output, _ = self.bi_lstm(sequence_input)
        shared_representation = self.attention_pooling(lstm_output)
        shared_representation = self.shared_projection(shared_representation)

        emotion_logits = self.emotion_head(shared_representation)
        impairment_logits = self.impairment_head(shared_representation)

        outputs = {
            "shared_representation": shared_representation,
            "emotion_logits": emotion_logits,
            "impairment_logits": impairment_logits,
        }

        if return_domain_output:
            reversed_features = self.gradient_reversal(shared_representation)
            outputs["domain_logits"] = self.domain_discriminator(reversed_features)

        return outputs


class CnnAttentionAdversarialModel(nn.Module):
    def __init__(self, num_emotions: int = 7):
        super().__init__()

        self.cnn_feature_extractor = ConvolutionalFeatureExtractor()
        self.attention_pooling = TemporalAttentionPooling(
            feature_size=128 * 16,
            attention_hidden_size=128,
        )

        self.shared_projection = nn.Sequential(
            nn.Linear(128 * 16, 256),
            nn.LayerNorm(256),
            nn.ReLU(),
            nn.Dropout(0.3),
        )

        self.emotion_head = nn.Sequential(
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, num_emotions),
        )

        self.impairment_head = nn.Sequential(
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, 2),
        )

        self.gradient_reversal = GradientReversalLayer(lambda_value=1.0)
        self.domain_discriminator = DomainDiscriminator(256, hidden_size=128)

    def forward(
        self,
        mel_spectrogram: torch.Tensor,
        return_domain_output: bool = True,
        return_attention: bool = False,
    ) -> dict[str, torch.Tensor]:
        cnn_output = self.cnn_feature_extractor(mel_spectrogram)
        batch_size, channels, height, width = cnn_output.shape
        sequence_input = (
            cnn_output.permute(0, 3, 1, 2).contiguous().view(batch_size, width, channels * height)
        )

        pooled_output, attention_weights = self.attention_pooling(sequence_input)
        shared_representation = self.shared_projection(pooled_output)

        emotion_logits = self.emotion_head(shared_representation)
        impairment_logits = self.impairment_head(shared_representation)

        outputs = {
            "shared_representation": shared_representation,
            "emotion_logits": emotion_logits,
            "impairment_logits": impairment_logits,
        }

        if return_domain_output:
            reversed_features = self.gradient_reversal(shared_representation)
            outputs["domain_logits"] = self.domain_discriminator(reversed_features)

        if return_attention:
            outputs["attention_weights"] = attention_weights

        return outputs


class MADA(nn.Module):
    """Notebook-aligned Multi-Adversarial Domain Adaptation model."""

    def __init__(
        self,
        in_dim: int = 418,
        hidden: int = 512,
        latent: int = 256,
        n_emotions: int = 7,
        dropout: float = 0.35,
    ):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(in_dim, hidden),
            nn.LayerNorm(hidden),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden, hidden),
            nn.LayerNorm(hidden),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden, latent),
            nn.LayerNorm(latent),
            nn.GELU(),
        )
        self.emotion_head = nn.Sequential(
            nn.Linear(latent, 256),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Dropout(dropout / 2),
            nn.Linear(128, n_emotions),
        )
        self.impairment_head = nn.Sequential(
            nn.Linear(latent, 64),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(64, 2),
        )
        self.grl = GradientReversalLayer(0.0)
        self.domain_disc = nn.Sequential(
            nn.Linear(latent, 128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, 2),
        )

    def set_lambda(self, lambda_value: float) -> None:
        self.grl.lambda_value = lambda_value

    def forward(
        self,
        input_tensor: torch.Tensor,
        return_domain_output: bool = True,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
        latent_representation = self.encoder(input_tensor)
        emotion_logits = self.emotion_head(latent_representation)
        impairment_logits = self.impairment_head(latent_representation)
        if not return_domain_output:
            return latent_representation, emotion_logits, impairment_logits, None
        domain_logits = self.domain_disc(self.grl(latent_representation))
        return latent_representation, emotion_logits, impairment_logits, domain_logits