src/models/drug_encoder.py

"""
drug_encoder.py — GINE-based molecular graph encoder + Morgan fingerprint fusion.

Uses Graph Isomorphism Network with Edge features (GINE) for learned
molecular representations, combined with Morgan fingerprints (ECFP4)
for knowledge-driven chemical similarity encoding.

v2: adds DropNode/DropEdge graph augmentation for cold-start robustness,
    and a MorganFPEncoder for dual-channel drug representation.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GINEConv, BatchNorm
from torch_geometric.data import Batch


class GINELayer(nn.Module):
    """Single GINE layer with batch norm and residual connection."""
    
    def __init__(self, hidden_dim: int, edge_dim: int, dropout: float = 0.2):
        super().__init__()
        # GINE uses an MLP as the update function
        mlp = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim * 2),
            nn.ReLU(),
            nn.Linear(hidden_dim * 2, hidden_dim),
        )
        self.conv = GINEConv(nn=mlp, edge_dim=edge_dim)
        self.bn = BatchNorm(hidden_dim)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x, edge_index, edge_attr):
        h = self.conv(x, edge_index, edge_attr)
        h = self.bn(h)
        h = F.relu(h)
        h = self.dropout(h)
        # residual connection
        return h + x


class DrugEncoder(nn.Module):
    """
    GINE-based drug molecular graph encoder.
    
    Produces per-atom representations suitable for cross-attention
    with protein residue representations.
    
    Architecture:
        Input projection → N × GINE layers → atom-level output
    
    We intentionally do NOT apply global readout (mean/sum pooling)
    here — the cross-attention module needs individual atom vectors
    to compute residue-level interaction maps.
    """
    
    def __init__(self,
                 num_atom_features: int = 52,
                 edge_dim: int = 16,
                 hidden_dim: int = 256,
                 num_layers: int = 3,
                 dropout: float = 0.2,
                 jk: str = 'last'):
        """
        Args:
            num_atom_features: input atom feature dimension
            edge_dim: bond feature dimension
            hidden_dim: GINE hidden dimension
            num_layers: number of GINE layers
            dropout: dropout rate
            jk: jumping knowledge mode ('last' or 'cat')
        """
        super().__init__()
        self.num_layers = num_layers
        self.jk = jk
        
        # input projection
        self.input_proj = nn.Sequential(
            nn.Linear(num_atom_features, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
        )
        
        # edge feature projection (match edge_attr dim to hidden_dim for GINE)
        self.edge_proj = nn.Linear(edge_dim, hidden_dim)
        
        # GINE layers
        self.layers = nn.ModuleList([
            GINELayer(hidden_dim, hidden_dim, dropout)
            for _ in range(num_layers)
        ])
        
        # jumping knowledge output
        if jk == 'cat':
            self.jk_proj = nn.Linear(hidden_dim * num_layers, hidden_dim)
        
        self.output_dim = hidden_dim
    
    def forward(self, drug_batch: Batch) -> tuple:
        """
        Args:
            drug_batch: PyG Batch of molecular graphs
        
        Returns:
            atom_repr: (total_atoms, hidden_dim) — per-atom representations
            batch_index: (total_atoms,) — which graph each atom belongs to
        """
        x = drug_batch.x
        edge_index = drug_batch.edge_index
        edge_attr = drug_batch.edge_attr
        batch_index = drug_batch.batch
        
        # project inputs
        h = self.input_proj(x)
        edge_features = self.edge_proj(edge_attr) if edge_attr.size(0) > 0 \
            else edge_attr
        
        # message passing
        layer_outputs = []
        for layer in self.layers:
            h = layer(h, edge_index, edge_features)
            layer_outputs.append(h)
        
        # jumping knowledge
        if self.jk == 'cat':
            h = torch.cat(layer_outputs, dim=-1)
            h = self.jk_proj(h)
        else:
            h = layer_outputs[-1]
        
        return h, batch_index


class GraphAugmentation(nn.Module):
    """
    Stochastic graph augmentation for regularisation during training.

    DropNode: randomly masks node features (zero-out) with probability p_node.
    DropEdge: randomly removes edges with probability p_edge.

    These augmentations prevent the GNN from memorising specific molecular
    graphs, which is critical for cold-start generalisation where test
    molecules are unseen during training.

    Reference: Rong et al., "DropEdge: Towards Deep Graph ConvNets", ICLR 2020
    """

    def __init__(self, p_node: float = 0.1, p_edge: float = 0.15):
        super().__init__()
        self.p_node = p_node
        self.p_edge = p_edge

    def forward(self, drug_batch: Batch) -> Batch:
        if not self.training:
            return drug_batch

        # DropNode: zero-out whole node feature vectors
        if self.p_node > 0:
            mask = torch.rand(drug_batch.x.size(0), 1,
                              device=drug_batch.x.device) > self.p_node
            drug_batch.x = drug_batch.x * mask.float()

        # DropEdge: remove random edges
        if self.p_edge > 0 and drug_batch.edge_index.size(1) > 0:
            n_edges = drug_batch.edge_index.size(1)
            keep = torch.rand(n_edges, device=drug_batch.edge_index.device) > self.p_edge
            drug_batch.edge_index = drug_batch.edge_index[:, keep]
            if drug_batch.edge_attr is not None and drug_batch.edge_attr.size(0) == n_edges:
                drug_batch.edge_attr = drug_batch.edge_attr[keep]

        return drug_batch


class MorganFPEncoder(nn.Module):
    """
    Encoder for Morgan (ECFP4) molecular fingerprints.

    Morgan fingerprints capture circular substructure patterns and provide
    a chemistry-prior-based global drug representation that inherently
    generalises to unseen molecules — unlike GNN features that depend on
    message passing over specific graph topologies.

    This encoder projects the binary fingerprint to a dense vector and
    is fused with the GNN representation for dual-channel drug encoding.
    """

    def __init__(self, input_dim: int = 1024, hidden_dim: int = 256,
                 dropout: float = 0.2):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim * 2),
            nn.LayerNorm(hidden_dim * 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
        )
        self.output_dim = hidden_dim

    def forward(self, fp: torch.Tensor) -> torch.Tensor:
        """
        Args:
            fp: (B, input_dim) binary Morgan fingerprint
        Returns:
            (B, hidden_dim) dense drug representation
        """
        return self.encoder(fp)