from abc import ABC, abstractmethod

import torch
import torch.nn as nn

# ============================================================================
# ENHANCED CLASSIFICATION HEADS
# =============================================================================


class LinearHead(nn.Module):
    """Simple linear classification head"""

    def __init__(self, input_dim: int, num_classes: int, dropout: float = 0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Linear(input_dim, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.dropout(x)
        return self.classifier(x)


class MLPHead(nn.Module):
    """Multi-layer perceptron classification head"""

    def __init__(
        self,
        input_dim: int,
        num_classes: int,
        hidden_dim: int = 512,
        num_layers: int = 2,
        dropout: float = 0.1,
    ):
        super().__init__()

        layers = []
        current_dim = input_dim

        for _ in range(num_layers - 1):
            layers.extend(
                [
                    nn.Linear(current_dim, hidden_dim),
                    nn.GELU(),
                    nn.LayerNorm(hidden_dim),
                    nn.Dropout(dropout),
                ]
            )
            current_dim = hidden_dim

        layers.append(nn.Linear(current_dim, num_classes))
        self.classifier = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.classifier(x)


class AttentionHead(nn.Module):
    """Enhanced attention-based classification head"""

    def __init__(
        self,
        input_dim: int,
        num_classes: int,
        hidden_dim: int = 512,
        num_heads: int = 8,
        num_layers: int = 2,
        dropout: float = 0.1,
    ):
        super().__init__()

        self.input_projection = nn.Linear(input_dim, hidden_dim)

        # Multi-head attention layers
        self.attention_layers = nn.ModuleList(
            [
                nn.TransformerEncoderLayer(
                    d_model=hidden_dim,
                    nhead=num_heads,
                    dim_feedforward=hidden_dim * 4,
                    dropout=dropout,
                    activation="gelu",
                    batch_first=True,
                    norm_first=True,
                )
                for _ in range(num_layers)
            ]
        )

        # Global attention pooling
        self.global_query = nn.Parameter(torch.randn(1, 1, hidden_dim))
        self.global_attention = nn.MultiheadAttention(
            embed_dim=hidden_dim, num_heads=num_heads, dropout=dropout, batch_first=True
        )

        self.norm = nn.LayerNorm(hidden_dim)
        self.classifier = nn.Linear(hidden_dim, num_classes)

        # Initialize weights
        nn.init.normal_(self.global_query, std=0.02)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x shape: (batch_size, seq_len, input_dim) or (batch_size, input_dim)

        if len(x.shape) == 2:
            # Single vector input - expand to sequence
            x = x.unsqueeze(1)  # (batch_size, 1, input_dim)

        batch_size = x.size(0)

        # Project to hidden dimension
        x = self.input_projection(x)

        # Apply attention layers
        for layer in self.attention_layers:
            x = layer(x)

        # Global attention pooling
        query = self.global_query.expand(batch_size, -1, -1)
        attended, _ = self.global_attention(query=query, key=x, value=x)
        attended = self.norm(attended)

        # Classification
        attended = attended.squeeze(1)
        return self.classifier(attended)


# =============================================================================
# TEMPORAL PROCESSING MODULES
# =============================================================================


class TemporalProcessor(nn.Module, ABC):
    """Base class for temporal processing"""

    def __init__(self, input_dim: int):
        super().__init__()
        self.input_dim = input_dim

    @abstractmethod
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        raise NotImplementedError


class MeanPooling(TemporalProcessor):
    """Simple mean pooling across time dimension"""

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x.mean(dim=1)


class MaxPooling(TemporalProcessor):
    """Max pooling across time dimension"""

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x.max(dim=1)[0]


class LSTMProcessor(TemporalProcessor):
    """LSTM-based temporal processing"""

    def __init__(self, input_dim: int, hidden_dim: int = 512, num_layers: int = 2):
        super().__init__(input_dim)
        self.lstm = nn.LSTM(
            input_dim,
            hidden_dim,
            num_layers,
            batch_first=True,
            dropout=0.1 if num_layers > 1 else 0,
        )
        self.output_dim = hidden_dim

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        lstm_out, (h_n, c_n) = self.lstm(x)
        return h_n[-1]  # Take last hidden state


class AttentionProcessor(TemporalProcessor):
    """Attention-based temporal processing"""

    def __init__(self, input_dim: int, num_heads: int = 8):
        super().__init__(input_dim)
        self.attention = nn.MultiheadAttention(
            embed_dim=input_dim, num_heads=num_heads, batch_first=True
        )
        self.norm = nn.LayerNorm(input_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        attn_out, _ = self.attention(x, x, x)
        attn_out = self.norm(attn_out)
        return attn_out.mean(dim=1)


class AttentiveProbe(TemporalProcessor):
    """Attentive probe with M learned query vectors for temporal aggregation."""

    def __init__(
        self,
        input_dim: int,
        hidden_dim: int,
        num_queries: int = 8,
    ):
        super().__init__(input_dim)
        self.num_queries = num_queries
        self.hidden_dim = hidden_dim

        self.query_vectors = nn.Parameter(torch.randn(num_queries, input_dim))
        self.value_projection = nn.Linear(input_dim, hidden_dim)

        self.scale = input_dim**-0.5

        nn.init.normal_(self.query_vectors, std=0.02)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x: (B, T, D) where B=batch_size, T=sequence_length, D=input_dim
        Returns: (B, M*hidden_dim) where M=num_queries
        """
        batch_size, _, _ = x.shape

        values = self.value_projection(x)

        queries = self.query_vectors.unsqueeze(0).expand(batch_size, -1, -1)

        attention_scores = torch.matmul(queries, x.transpose(-2, -1)) * self.scale
        attention_weights = torch.softmax(attention_scores, dim=-1)

        attended_features = torch.matmul(attention_weights, values)

        return attended_features.flatten(1)
