import torch
import torch.nn as nn
from typing import Dict, Any, Tuple
from ares_compression.layers import GatedMultiLobeRouter

class SpecialistLobe(nn.Module):
    """
    A specialist knowledge lobe within the MAS-RON architecture.
    """
    def __init__(self, name: str, dimension: int, hidden_dim: int, device: str = "cuda"):
        super().__init__()
        self.name = name
        self.dimension = dimension
        
        # Specialist weights (only active when routed here)
        self.projector = nn.Linear(hidden_dim, dimension)
        self.processor = nn.Sequential(
            nn.Linear(dimension, dimension),
            nn.ReLU(),
            nn.Linear(dimension, hidden_dim)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Project hidden state into the high-dimensional sparse space
        sparse_repr = self.projector(x)
        # Process and project back
        return self.processor(sparse_repr)

class MultiLobeSystem(nn.Module):
    def __init__(self, d_model: int, device: str = "cuda"):
        super().__init__()
        # router is in layers.py
        self.router = GatedMultiLobeRouter(d_model, num_lobes=2, device=device)
        
        # Lobe 0: General Knowledge (Reduced for local memory stability)
        self.general_lobe = SpecialistLobe("General", 10000, d_model, device)
        
        # Lobe 1: Code Specialist (High Capacity)
        self.code_lobe = SpecialistLobe("Code", 20000, d_model, device)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # [batch, seq_len, num_lobes]
        routing_weights = self.router(x)
        
        # Compute lobe outputs
        out_gen = self.general_lobe(x)
        out_code = self.code_lobe(x)
        
        # Weighted sum based on routing
        # [batch, seq_len, d_model]
        combined = (routing_weights[:, :, 0].unsqueeze(-1) * out_gen + 
                    routing_weights[:, :, 1].unsqueeze(-1) * out_code)
        
        return x + combined
