"""Hyper-Connection modules for improved residual connections.

Implements:
- HyperConnection: Basic hyper-connection with multiple streams
- ManifoldHyperConnection (mHC): Manifold-constrained version for stability

Based on: "mHC: Manifold-Constrained Hyper-Connections" (arXiv:2512.24880)
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, List, Callable


class HyperConnection(nn.Module):
    """Standard Hyper-Connection with multiple parallel streams.
    
    Extends residual connections by allowing multiple transformation paths
    with learnable weights.
    
    Standard residual: x = x + f(x)
    Hyper-connection: x = α₁·x + α₂·f(x) + α₃·g(x) + ...
    
    Args:
        dim: Hidden dimension.
        num_streams: Number of parallel streams. Defaults to 2.
        init_identity: Initialize to behave like identity. Defaults to True.
    """
    def __init__(
        self,
        dim: int,
        num_streams: int = 2,
        init_identity: bool = True,
    ):
        super().__init__()
        self.dim = dim
        self.num_streams = num_streams
        
        # Learnable stream weights
        self.stream_weights = nn.Parameter(torch.ones(num_streams))
        
        if init_identity:
            # Initialize so first weight (identity) dominates
            with torch.no_grad():
                self.stream_weights[0] = 1.0
                self.stream_weights[1:] = 0.1
    
    def forward(self, *streams: torch.Tensor) -> torch.Tensor:
        """Combine multiple streams with learnable weights.
        
        Args:
            *streams: Variable number of tensors, all same shape.
            
        Returns:
            Weighted combination of streams.
        """
        assert len(streams) == self.num_streams, \
            f"Expected {self.num_streams} streams, got {len(streams)}"
        
        # Normalize weights to sum to 1 for stability
        weights = F.softmax(self.stream_weights, dim=0)
        
        result = torch.zeros_like(streams[0])
        for w, stream in zip(weights, streams):
            result = result + w * stream
        
        return result


class ManifoldHyperConnection(nn.Module):
    """Manifold-Constrained Hyper-Connection (mHC).
    
    Projects the expanded residual stream onto a learned manifold to:
    1. Restore identity mapping property
    2. Improve training stability
    3. Enable better scalability
    
    Architecture:
        x_out = x + proj(concat([f(x), g(x), ...]))
        
    The projection learns to combine information from multiple streams
    while preserving the identity mapping through the residual.
    
    Args:
        dim: Hidden dimension.
        num_streams: Number of transformation streams. Defaults to 2.
        expansion_factor: Hidden expansion for projection. Defaults to 1.
        dropout: Dropout probability. Defaults to 0.0.
        use_gating: Use gating mechanism. Defaults to True.
    """
    def __init__(
        self,
        dim: int,
        num_streams: int = 2,
        expansion_factor: int = 1,
        dropout: float = 0.0,
        use_gating: bool = True,
    ):
        super().__init__()
        self.dim = dim
        self.num_streams = num_streams
        self.use_gating = use_gating
        
        expanded_dim = dim * num_streams
        hidden_dim = dim * expansion_factor
        
        # Projection from expanded streams back to model dim
        if expansion_factor > 1:
            self.projection = nn.Sequential(
                nn.Linear(expanded_dim, hidden_dim),
                nn.GELU(),
                nn.Dropout(dropout),
                nn.Linear(hidden_dim, dim),
            )
        else:
            self.projection = nn.Linear(expanded_dim, dim)
        
        # Optional gating for controlling contribution
        if use_gating:
            self.gate = nn.Sequential(
                nn.Linear(dim, dim),
                nn.Sigmoid(),
            )
            # Initialize gate to output small values initially
            nn.init.constant_(self.gate[0].bias, -2.0)
        
        self.dropout = nn.Dropout(dropout)
        self.layer_norm = nn.LayerNorm(dim)
        
        # Initialize projection to near-zero so identity dominates initially
        self._init_weights()
    
    def _init_weights(self):
        """Initialize weights for identity-preserving behavior."""
        if isinstance(self.projection, nn.Sequential):
            # Initialize last layer to small values
            nn.init.normal_(self.projection[-1].weight, std=0.02)
            nn.init.zeros_(self.projection[-1].bias)
        else:
            nn.init.normal_(self.projection.weight, std=0.02)
            nn.init.zeros_(self.projection.bias)
    
    def forward(
        self,
        identity: torch.Tensor,
        *streams: torch.Tensor,
    ) -> torch.Tensor:
        """Apply manifold-constrained hyper-connection.
        
        Args:
            identity: The original input (for residual).
            *streams: Transformed streams to combine.
            
        Returns:
            Output with residual connection.
            
        Example:
            >>> mhc = ManifoldHyperConnection(768, num_streams=2)
            >>> x = torch.randn(2, 128, 768)
            >>> attn_out = attention(x)
            >>> mlp_out = mlp(x)
            >>> out = mhc(x, attn_out, mlp_out)
        """
        assert len(streams) == self.num_streams, \
            f"Expected {self.num_streams} streams, got {len(streams)}"
        
        # Concatenate all streams
        combined = torch.cat(streams, dim=-1)
        
        # Project back to model dimension
        projected = self.projection(combined)
        projected = self.dropout(projected)
        
        # Optional gating
        if self.use_gating:
            gate = self.gate(identity)
            projected = gate * projected
        
        # Residual connection preserves identity mapping
        output = identity + projected
        
        return output


class AdaptiveHyperConnection(nn.Module):
    """Adaptive Hyper-Connection that learns stream weights per-layer.
    
    Uses input-dependent routing to determine how much each stream
    contributes to the output.
    
    Args:
        dim: Hidden dimension.
        num_streams: Number of streams.
    """
    def __init__(self, dim: int, num_streams: int = 2):
        super().__init__()
        self.dim = dim
        self.num_streams = num_streams
        
        # Router that produces per-stream weights
        self.router = nn.Sequential(
            nn.Linear(dim, dim // 4),
            nn.GELU(),
            nn.Linear(dim // 4, num_streams),
        )
        
        # Initialize router to produce roughly equal weights
        nn.init.zeros_(self.router[-1].weight)
        nn.init.zeros_(self.router[-1].bias)
    
    def forward(
        self,
        query: torch.Tensor,
        *streams: torch.Tensor,
    ) -> torch.Tensor:
        """Route between streams based on input.
        
        Args:
            query: Input to base routing decision on.
            *streams: Streams to combine.
            
        Returns:
            Input-dependent weighted combination.
        """
        assert len(streams) == self.num_streams
        
        # Compute routing weights from query
        # Use mean pooling if query has sequence dimension
        if query.dim() == 3:
            pooled = query.mean(dim=1)  # (batch, dim)
            weights = self.router(pooled)  # (batch, num_streams)
            weights = F.softmax(weights, dim=-1)
            weights = weights.unsqueeze(1).unsqueeze(-1)  # (batch, 1, num_streams, 1)
            
            # Stack and weight streams
            stacked = torch.stack(streams, dim=2)  # (batch, seq, num_streams, dim)
            result = (stacked * weights).sum(dim=2)
        else:
            weights = F.softmax(self.router(query), dim=-1)
            result = sum(w * s for w, s in zip(weights.unbind(-1), streams))
        
        return result
