model.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Tuple, Optional

# Import all modules
from encoders import MultiModalEncoder
from memory import DualTimescaleAdapter, AdaptiveMixing, EpisodicMemory, WriteGate
from backbone import TransformerBackbone, ProjectionLayer
from mamba_backbone import MambaBackbone, BiMambaBackbone, HybridMambaTransformer, get_mamba_backbone
from adaptation import FiLM, SimpleNormalizingFlow
from heads import PredictionHeads


class MultiModalTransformerModel(nn.Module):
    """
    Complete multi-modal transformer model with all components
    """
    
    def __init__(self, config: dict):
        super().__init__()
        
        # Store configuration
        self.config = config
        
        # ============================================
        # Initialize all modules
        # ============================================
        
        # Multi-modal encoder
        self.encoder = MultiModalEncoder(config)
        
        # Dual-timescale adapters
        self.adapter_fast = DualTimescaleAdapter(
            config['input_dim_x'],
            config['hidden_dim'],
            config['memory_dim'],
            'fast',
            config.get('lstm_layers', 2)
        )
        
        self.adapter_slow = DualTimescaleAdapter(
            config['input_dim_x'],
            config['hidden_dim'],
            config['memory_dim'],
            'slow',
            config.get('lstm_layers', 2)
        )
        
        # Adaptive mixing module
        self.adaptive_mixing = AdaptiveMixing(
            config['z_dim'],
            config['hidden_dim']
        )
        
        # Projection layer for backbone input
        self.W_proj = ProjectionLayer(config['z_dim'], config['d_model'])
        
        # Backbone selection (Transformer or Mamba)
        backbone_type = config.get('backbone_type', 'transformer')
        
        if backbone_type == 'transformer':
            self.backbone = TransformerBackbone(
                config['d_model'],
                config['nhead'],
                config['num_layers'],
                config['dim_feedforward'],
                config.get('dropout', 0.1)
            )
        elif backbone_type in ['mamba', 'bimamba', 'hybrid']:
            # Mamba-specific parameters
            mamba_config = {
                'd_model': config['d_model'],
                'num_layers': config['num_layers'],
                'd_state': config.get('d_state', 16),
                'd_conv': config.get('d_conv', 4),
                'expand': config.get('expand', 2),
                'dropout': config.get('dropout', 0.1)
            }
            
            # Add hybrid-specific parameters if needed
            if backbone_type == 'hybrid':
                mamba_config['num_heads'] = config['nhead']
                mamba_config['dim_feedforward'] = config['dim_feedforward']
                mamba_config['mamba_layers'] = config.get('mamba_layers', None)
            
            self.backbone = get_mamba_backbone(backbone_type, **mamba_config)
        else:
            raise ValueError(f"Unknown backbone type: {backbone_type}")
        
        # Episodic memory
        self.episodic_memory = EpisodicMemory(
            config['memory_size'],
            config['d_model'],
            config.get('memory_heads', 8)
        )
        
        # Normalizing flow
        self.normalizing_flow = SimpleNormalizingFlow(
            config['d_model'],
            config.get('num_flows', 4)
        )
        
        # Platform adaptation (FiLM)
        self.film = FiLM(
            config['num_platforms'],
            config['d_model'],
            use_embedding=config.get('use_film_embedding', True)
        )
        
        # Prediction heads
        self.heads = PredictionHeads(
            config['d_model'],
            config['memory_dim'],
            config['z_dim'],
            config['output_dim_y'],
            config['output_dim_n'],
            config.get('head_hidden_dims', [256, 128]),
            config.get('dropout', 0.1)
        )
        
        # Write gate
        self.write_gate = WriteGate(
            config['z_dim'],
            config['memory_dim'],
            config['hidden_dim']
        )
        
        # Loss weight parameters (learnable)
        self.register_buffer('lambda_dis', torch.tensor(config.get('lambda_dis', 0.1)))
        self.register_buffer('lambda_rep', torch.tensor(config.get('lambda_rep', 0.05)))
        
        # Initialize parameters
        self._initialize_parameters()
    
    def _initialize_parameters(self):
        """Initialize model parameters"""
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.Embedding):
                nn.init.normal_(module.weight, std=0.02)
            elif isinstance(module, nn.LSTM):
                for name, param in module.named_parameters():
                    if 'weight' in name:
                        nn.init.xavier_uniform_(param)
                    elif 'bias' in name:
                        nn.init.zeros_(param)
    
    def forward(self, o: torch.Tensor, p: torch.Tensor, t: torch.Tensor,
                x_history: torch.Tensor, platform_id: torch.Tensor,
                epoch_progress: float = 0.0) -> Dict[str, torch.Tensor]:
        """
        Forward pass through the complete model
        
        Args:
            o: Observation tensor (batch, input_dim_o)
            p: Platform tensor (batch, input_dim_p)
            t: Time tensor (batch, input_dim_t)
            x_history: History tensor (batch, seq_len, input_dim_x)
            platform_id: Platform IDs (batch,)
            epoch_progress: Training progress (0-1)
            
        Returns:
            Dictionary containing predictions and intermediate outputs
        """
        batch_size = o.shape[0]
        
        # ============================================
        # Step 1: Encode and disentangle
        # ============================================
        z, z_c, z_p, z_t = self.encoder(o, p, t)
        
        # Compute disentanglement loss
        dis_loss = self.encoder.compute_disentanglement_loss(
            z, z_c, z_p, z_t,
            beta=self.config.get('beta_kl', 1.0),
            lambda_dis=self.lambda_dis.item()
        )
        
        # ============================================
        # Step 2: Dual-timescale memory
        # ============================================
        m_fast = self.adapter_fast(x_history)
        m_slow = self.adapter_slow(x_history)
        
        # Adaptive mixing
        M, alpha = self.adaptive_mixing(z_c, z_t, m_fast, m_slow)
        
        # ============================================
        # Step 3: Transformer backbone processing
        # ============================================
        z_proj = self.W_proj(z)
        s_t = self.backbone(z_proj)
        
        # ============================================
        # Step 4: Episodic memory retrieval
        # ============================================
        r = self.episodic_memory.retrieve(s_t)
        
        # Combine backbone output with retrieved memory
        b = F.layer_norm(s_t + r, [self.config['d_model']])
        
        # ============================================
        # Step 5: Normalizing flow
        # ============================================
        b_flow = self.normalizing_flow(b)
        
        # ============================================
        # Step 6: Platform adaptation (FiLM)
        # ============================================
        h = self.film(b_flow, platform_id)
        
        # ============================================
        # Step 7: Prediction heads
        # ============================================
        y_pred, n_pred, head_outputs = self.heads(
            h, m_fast, m_slow, M, z_c,
            use_learned_weight=self.config.get('use_learned_weight', True)
        )
        
        # ============================================
        # Step 8: Compute write gate
        # ============================================
        g_write = self.write_gate(z, M, epoch_progress)
        
        # ============================================
        # Step 9: Update episodic memory
        # ============================================
        self.episodic_memory.write(s_t, g_write)
        
        # ============================================
        # Prepare outputs
        # ============================================
        outputs = {
            # Primary outputs
            'y_pred': y_pred,
            'n_pred': n_pred,
            
            # Losses
            'dis_loss': dis_loss,
            
            # Intermediate features (for analysis/debugging)
            'z': z,
            'z_c': z_c,
            'z_p': z_p,
            'z_t': z_t,
            'M': M,
            'm_fast': m_fast,
            'm_slow': m_slow,
            'alpha': alpha,
            's_t': s_t,
            'h': h,
            'g_write': g_write,
            
            # Head outputs
            **head_outputs
        }
        
        return outputs
    
    def compute_loss(self, outputs: Dict[str, torch.Tensor],
                     y_true: torch.Tensor, n_true: torch.Tensor,
                     replay_loss: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]:
        """
        Compute all losses
        
        Args:
            outputs: Model outputs dictionary
            y_true: True target values for regression
            n_true: True target values for classification
            replay_loss: Optional replay loss from buffer
            
        Returns:
            Dictionary containing all losses
        """
        # MAPE loss for regression
        mape_loss = F.l1_loss(outputs['y_pred'], y_true)
        
        # Cross-entropy loss for classification
        ce_loss = F.cross_entropy(outputs['n_pred'], n_true)
        
        # Combined prediction loss
        pred_loss = mape_loss + ce_loss
        
        # Get disentanglement loss
        dis_loss = outputs['dis_loss']
        
        # Total loss
        total_loss = pred_loss + self.lambda_dis * dis_loss
        
        # Add replay loss if provided
        if replay_loss is not None:
            total_loss += self.lambda_rep * replay_loss
        
        losses = {
            'total': total_loss,
            'pred': pred_loss,
            'mape': mape_loss,
            'ce': ce_loss,
            'dis': dis_loss,
            'replay': replay_loss if replay_loss is not None else torch.tensor(0.0)
        }
        
        return losses
    
    def get_memory_stats(self) -> dict:
        """Get statistics about memory modules"""
        return {
            'episodic': self.episodic_memory.get_memory_stats(),
            'alpha_mean': self.adaptive_mixing.mixing_network[0].weight.data.mean().item(),
            'write_gate_mean': self.write_gate.gate_network[0].weight.data.mean().item()
        }
    
    @torch.no_grad()
    def extract_features(self, o: torch.Tensor, p: torch.Tensor, 
                         t: torch.Tensor) -> Dict[str, torch.Tensor]:
        """
        Extract features without forward pass (for analysis)
        
        Args:
            o, p, t: Input tensors
            
        Returns:
            Dictionary of extracted features
        """
        z, z_c, z_p, z_t = self.encoder(o, p, t)
        
        return {
            'z': z,
            'z_c': z_c,
            'z_p': z_p,
            'z_t': z_t
        }