"""
Novel Hierarchical Glycan Architecture

This module provides components for hierarchical representation learning
that leverages our unique multimodal data (sequence + MS + 3D structure).

Key advantages over GIFFLAR (which only uses graph structure):
1. We have 3D atomic structure from PDB files → structural supervision
2. We have MS fragmentation patterns → mass-based supervision  
3. We have WURCS with explicit residue boundaries → hierarchical attention

Our novel approach: Cross-Modal Hierarchical Learning
- Token level: Learn atomic WURCS patterns
- Residue level: Pool tokens → monosaccharide representations
- Glycan level: Pool residues → glycan representation
- Cross-modal: Align sequence residues with 3D structure residues using contrastive learning
"""

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


class MonoTypePredictor(nn.Module):
    """
    Prediction head for monosaccharide type classification.
    
    Given a masked residue's pooled representation, predict the monosaccharide type.
    This is used with MonosaccharideMaskingStrategy during pre-training.
    """
    
    # Same type vocabulary as in masking.py
    MONO_TYPES = [
        '<UNK>', 'Glc', 'Gal', 'Man', 'Fuc', 'Xyl', 'Rha', 'Ara',
        'GlcNAc', 'GalNAc', 'ManNAc', 'GlcA', 'GalA', 'IdoA',
        'Neu5Ac', 'Neu5Gc', 'Kdn', 'GlcN', 'GalN', 'Hex', 'HexNAc',
        'dHex', 'Pent', 'Sia', 'GlcS', 'GalS', 'Ido', 'All', 'Alt', 'Gul', 'Tal'
    ]
    
    def __init__(self, hidden_size: int, dropout: float = 0.1):
        super().__init__()
        self.num_types = len(self.MONO_TYPES)
        
        self.predictor = nn.Sequential(
            nn.Linear(hidden_size, hidden_size),
            nn.GELU(),
            nn.LayerNorm(hidden_size),
            nn.Dropout(dropout),
            nn.Linear(hidden_size, self.num_types)
        )
        
        # Type embeddings (can be used to enrich residue representations)
        self.type_embeddings = nn.Embedding(self.num_types, hidden_size)
    
    def forward(self, residue_hidden: torch.Tensor) -> torch.Tensor:
        """
        Predict monosaccharide type from residue representation.
        
        Args:
            residue_hidden: (batch, hidden_size) or (batch, num_residues, hidden_size)
        
        Returns:
            logits: (batch, num_types) or (batch, num_residues, num_types)
        """
        return self.predictor(residue_hidden)
    
    def get_type_embedding(self, type_ids: torch.Tensor) -> torch.Tensor:
        """Get embeddings for monosaccharide type IDs."""
        return self.type_embeddings(type_ids)


class HierarchicalGlycanEncoder(nn.Module):
    """
    Novel hierarchical encoder for glycan sequences.
    
    Architecture:
    1. Token Encoder: Standard transformer on WURCS tokens
    2. Residue Encoder: Pool tokens within residues → residue transformer
    3. Glycan Encoder: Pool residues → final glycan representation
    
    This explicitly models the glycan hierarchy (atoms → monosaccharides → glycan)
    which is similar to how GNNs work but within a transformer framework.
    """
    
    def __init__(
        self,
        hidden_size: int = 768,
        num_residue_layers: int = 2,
        num_heads: int = 8,
        dropout: float = 0.1,
        max_residues: int = 32,
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.max_residues = max_residues
        
        # Residue position embedding
        self.residue_pos_embedding = nn.Embedding(max_residues, hidden_size)
        
        # Residue-level transformer
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=hidden_size,
            nhead=num_heads,
            dim_feedforward=hidden_size * 4,
            dropout=dropout,
            activation='gelu',
            batch_first=True
        )
        self.residue_transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_residue_layers)
        
        # Glycan-level attention pooling
        self.glycan_attention = nn.Linear(hidden_size, 1)
        
        # Optional: branch-aware attention (if glycan has branches)
        self.branch_embedding = nn.Embedding(2, hidden_size)  # 0=linear, 1=branch point
    
    def forward(
        self,
        token_hidden: torch.Tensor,      # (batch, seq_len, hidden)
        residue_ids: torch.Tensor,        # (batch, seq_len)
        attention_mask: torch.Tensor,     # (batch, seq_len)
        is_branch_point: torch.Tensor = None  # (batch, max_residues) optional
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Hierarchical encoding: tokens → residues → glycan.
        
        Returns:
            residue_hidden: (batch, max_residues, hidden) - Residue-level representations
            glycan_hidden: (batch, hidden) - Glycan-level representation
        """
        batch_size = token_hidden.size(0)
        device = token_hidden.device
        
        # Step 1: Pool tokens within each residue
        residue_hidden = torch.zeros(batch_size, self.max_residues, self.hidden_size, device=device)
        residue_mask = torch.zeros(batch_size, self.max_residues, dtype=torch.bool, device=device)
        
        for b in range(batch_size):
            unique_residues = torch.unique(residue_ids[b])
            # Only consider real residues (>= 0)
            real_residues = unique_residues[unique_residues >= 0]
            
            for i, rid in enumerate(real_residues[:self.max_residues]):
                token_mask = (residue_ids[b] == rid) & (attention_mask[b] == 1)
                if token_mask.sum() > 0:
                    # Mean pool tokens in this residue
                    residue_hidden[b, i] = token_hidden[b, token_mask].mean(dim=0)
                    residue_mask[b, i] = True
        
        # Add residue position embeddings
        positions = torch.arange(self.max_residues, device=device).unsqueeze(0).expand(batch_size, -1)
        residue_hidden = residue_hidden + self.residue_pos_embedding(positions)
        
        # Add branch embeddings if provided
        if is_branch_point is not None:
            residue_hidden = residue_hidden + self.branch_embedding(is_branch_point.long())
        
        # Step 2: Apply residue-level transformer
        # Create attention mask for transformer (True = ignore)
        residue_attn_mask = ~residue_mask
        residue_hidden = self.residue_transformer(
            residue_hidden, 
            src_key_padding_mask=residue_attn_mask
        )
        
        # Step 3: Attention pooling to get glycan representation
        scores = self.glycan_attention(residue_hidden).squeeze(-1)  # (batch, max_residues)
        scores = scores.masked_fill(~residue_mask, -1e9)
        weights = F.softmax(scores, dim=1).unsqueeze(-1)  # (batch, max_residues, 1)
        glycan_hidden = (residue_hidden * weights).sum(dim=1)  # (batch, hidden)
        
        return residue_hidden, glycan_hidden


class CrossModalResidueLoss(nn.Module):
    """
    Cross-modal contrastive loss at residue level.
    
    This is our NOVEL contribution: align sequence residue representations
    with 3D structure residue representations. GIFFLAR doesn't have access
    to 3D structure, so this is our unique advantage.
    
    For each residue in a glycan with 3D structure:
    - Positive: The same residue from sequence and structure should be similar
    - Negative: Different residues should be dissimilar
    """
    
    def __init__(self, hidden_size: int, temperature: float = 0.07):
        super().__init__()
        self.temperature = temperature
        
        # Project both modalities to shared space
        self.seq_proj = nn.Linear(hidden_size, hidden_size)
        self.struct_proj = nn.Linear(hidden_size, hidden_size)
    
    def forward(
        self,
        seq_residue_hidden: torch.Tensor,    # (batch, num_residues, hidden)
        struct_residue_hidden: torch.Tensor,  # (batch, num_residues, hidden)
        residue_mask: torch.Tensor            # (batch, num_residues) - which residues are valid
    ) -> torch.Tensor:
        """
        Compute cross-modal contrastive loss.
        
        For each valid residue, its sequence representation should be
        similar to its structure representation (positive pair) and
        dissimilar to other residues (negative pairs).
        """
        batch_size, num_residues, hidden = seq_residue_hidden.shape
        
        # Project to shared space
        seq_proj = F.normalize(self.seq_proj(seq_residue_hidden), dim=-1)
        struct_proj = F.normalize(self.struct_proj(struct_residue_hidden), dim=-1)
        
        total_loss = 0.0
        valid_count = 0
        
        for b in range(batch_size):
            valid_indices = residue_mask[b].nonzero(as_tuple=True)[0]
            if len(valid_indices) < 2:
                continue
            
            # Get valid residue representations
            seq_valid = seq_proj[b, valid_indices]      # (n_valid, hidden)
            struct_valid = struct_proj[b, valid_indices]  # (n_valid, hidden)
            
            # Compute similarity matrix
            sim_matrix = torch.matmul(seq_valid, struct_valid.T) / self.temperature
            
            # Labels: diagonal should be positive pairs
            labels = torch.arange(len(valid_indices), device=sim_matrix.device)
            
            # Cross-entropy loss (both directions)
            loss_seq_to_struct = F.cross_entropy(sim_matrix, labels)
            loss_struct_to_seq = F.cross_entropy(sim_matrix.T, labels)
            
            total_loss += (loss_seq_to_struct + loss_struct_to_seq) / 2
            valid_count += 1
        
        if valid_count > 0:
            return total_loss / valid_count
        return torch.tensor(0.0, device=seq_residue_hidden.device)


# Test the modules
if __name__ == '__main__':
    batch_size = 2
    seq_len = 50
    hidden_size = 768
    max_residues = 8
    
    # Create dummy data
    token_hidden = torch.randn(batch_size, seq_len, hidden_size)
    residue_ids = torch.zeros(batch_size, seq_len, dtype=torch.long)
    # Simulate 4 residues per glycan
    for b in range(batch_size):
        for r in range(4):
            residue_ids[b, 1 + r*10 : 1 + (r+1)*10] = r
        residue_ids[b, 41:] = -2  # Linkages
    residue_ids[:, 0] = -1   # START
    residue_ids[:, -1] = -1  # END
    attention_mask = torch.ones(batch_size, seq_len)
    
    # Test HierarchicalGlycanEncoder
    encoder = HierarchicalGlycanEncoder(hidden_size=hidden_size, max_residues=max_residues)
    residue_hidden, glycan_hidden = encoder(token_hidden, residue_ids, attention_mask)
    print(f"✓ HierarchicalGlycanEncoder: residue={residue_hidden.shape}, glycan={glycan_hidden.shape}")
    
    # Test MonoTypePredictor
    predictor = MonoTypePredictor(hidden_size=hidden_size)
    logits = predictor(residue_hidden[:, :4, :])  # First 4 residues
    print(f"✓ MonoTypePredictor: logits={logits.shape}")
    
    # Test CrossModalResidueLoss
    struct_hidden = torch.randn(batch_size, max_residues, hidden_size)
    residue_mask = torch.zeros(batch_size, max_residues, dtype=torch.bool)
    residue_mask[:, :4] = True  # 4 valid residues
    
    loss_fn = CrossModalResidueLoss(hidden_size=hidden_size)
    loss = loss_fn(residue_hidden, struct_hidden, residue_mask)
    print(f"✓ CrossModalResidueLoss: {loss.item():.4f}")
    
    print("\n✓ All novel architecture components working!")