core/evolution/embedding_extractor

"""
Embedding Extractor for Mixture DCN GNN Model

Extracts the latent embeddings (before the MLP head) from your trained
mixture DCN predictor. These embeddings represent the molecule in the
learned latent space.
"""

import torch
import numpy as np
from typing import List, Optional
from torch_geometric.data import Data, Batch


class EmbeddingExtractor:
    """
    Extract embeddings from mixture DCN GNN model.
    
    The architecture is typically:
        Graph Conv layers → Pooling → [EMBEDDING] → MLP → DCN prediction
                                       ^^^^^^^^^^
                                       Extract this!
    """
    
    def __init__(self, model, device='cpu'):
        """
        Args:
            model: Your trained mixture DCN model
            device: 'cpu' or 'cuda'
        """
        self.model = model
        self.device = device
        self.model.to(device)
        self.model.eval()
        
        # Hook to capture embeddings
        self.embeddings = []
        self._register_hook()
    
    def _register_hook(self):
        """
        Register a forward hook to capture embeddings.
        
        You'll need to identify which layer is the embedding layer.
        Common patterns:
        - model.graph_conv_layers → model.pool → [model.embedding] → model.mlp
        - model.encoder → [embedding] → model.decoder
        
        Adjust based on your actual architecture!
        """
        # OPTION 1: If your model has an explicit embedding layer
        if hasattr(self.model, 'embedding_layer'):
            self.hook = self.model.embedding_layer.register_forward_hook(
                self._hook_fn
            )
        
        # OPTION 2: If MLP head is a separate module
        elif hasattr(self.model, 'mlp_head') or hasattr(self.model, 'fc'):
            # Hook the layer BEFORE the MLP head
            # This is usually the pooling layer or the last graph conv
            if hasattr(self.model, 'pool'):
                self.hook = self.model.pool.register_forward_hook(
                    self._hook_fn
                )
            else:
                # Find the last layer before MLP
                layers = list(self.model.children())
                self.hook = layers[-2].register_forward_hook(
                    self._hook_fn
                )
        
        # OPTION 3: Generic - hook the layer before final prediction
        else:
            # You may need to manually identify this
            # Example: if your model is Sequential-like
            layers = list(self.model.children())
            # Hook second-to-last layer
            self.hook = layers[-2].register_forward_hook(
                self._hook_fn
            )
    
    def _hook_fn(self, module, input, output):
        """Capture the output of the hooked layer."""
        # Detach and move to CPU to save memory
        if isinstance(output, tuple):
            # Some layers return (output, additional_info)
            output = output[0]
        
        self.embeddings.append(output.detach().cpu())
    
    def extract_embeddings_from_graphs(self, graph_list: List[Data]) -> np.ndarray:
        """
        Extract embeddings for a list of PyTorch Geometric graphs.
        
        Args:
            graph_list: List of PyG Data objects (molecule graphs)
        
        Returns:
            embeddings: numpy array of shape (n_molecules, embedding_dim)
        """
        self.embeddings = []
        
        with torch.no_grad():
            # Batch the graphs
            batch = Batch.from_data_list(graph_list).to(self.device)
            
            # Forward pass (hook will capture embeddings)
            _ = self.model(batch)
        
        # Concatenate all captured embeddings
        embeddings = torch.cat(self.embeddings, dim=0)
        
        return embeddings.numpy()
    
    def extract_embeddings_from_smiles(self, 
                                      smiles_list: List[str],
                                      featurizer) -> np.ndarray:
        """
        Extract embeddings from SMILES strings.
        
        Args:
            smiles_list: List of SMILES
            featurizer: Function to convert SMILES → PyG Data object
        
        Returns:
            embeddings: numpy array of shape (n_molecules, embedding_dim)
        """
        # Convert SMILES to graphs
        graphs = [featurizer(smiles) for smiles in smiles_list]
        
        # Extract embeddings
        return self.extract_embeddings_from_graphs(graphs)
    
    def __del__(self):
        """Remove hook when done."""
        if hasattr(self, 'hook'):
            self.hook.remove()


# =============================================================================
# USAGE EXAMPLE
# =============================================================================

"""
STEP 1: Identify your model architecture
-----------------------------------------

You need to know where the embedding layer is. Common patterns:

Pattern A: Explicit embedding
    self.graph_conv = GCN(...)
    self.pool = GlobalMeanPool()
    self.embedding = Linear(hidden_dim, embedding_dim)  ← Hook here!
    self.mlp_head = MLP(embedding_dim, 1)

Pattern B: Pooling as embedding
    self.graph_conv = GCN(...)
    self.pool = GlobalMeanPool()  ← Hook here! (output IS the embedding)
    self.mlp_head = MLP(hidden_dim, 1)

Pattern C: Sequential
    self.layers = Sequential(
        GCN(...),
        GlobalMeanPool(),
        Linear(hidden_dim, embedding_dim),  ← Hook here!
        ReLU(),
        Linear(embedding_dim, 1)
    )


STEP 2: Load your model and extract embeddings
----------------------------------------------

from mixture_dcn_model import load_trained_model

# Load your trained model
model = load_trained_model('path/to/model.pth')

# Create extractor
extractor = EmbeddingExtractor(model, device='cuda')

# Extract embeddings for training set
train_embeddings = extractor.extract_embeddings_from_smiles(
    train_smiles_list,
    featurizer=your_featurizer_function
)

print(f"Embeddings shape: {train_embeddings.shape}")
# Output: (n_train_samples, embedding_dim)


STEP 3: Use these embeddings for One-Class SVM
----------------------------------------------

See applicability_domain.py for next steps!
"""


# =============================================================================
# DEBUGGING: Find the right layer to hook
# =============================================================================

def print_model_structure(model):
    """
    Print model structure to help identify embedding layer.
    
    Usage:
        model = load_trained_model(...)
        print_model_structure(model)
    """
    print("Model Structure:")
    print("="*70)
    
    for i, (name, module) in enumerate(model.named_modules()):
        if name:  # Skip the root module
            print(f"{i}: {name}")
            print(f"   Type: {type(module).__name__}")
            
            # If it's a Linear layer, show dimensions
            if hasattr(module, 'in_features') and hasattr(module, 'out_features'):
                print(f"   Shape: ({module.in_features}, {module.out_features})")
            
            print()


"""
EXAMPLE OUTPUT:

Model Structure:
======================================================================
1: graph_conv
   Type: GCN
   
2: pool
   Type: GlobalMeanPool
   
3: embedding_layer       ← THIS IS WHAT WE WANT!
   Type: Linear
   Shape: (256, 128)     ← 128-dim embeddings
   
4: mlp_head
   Type: Sequential
   
5: mlp_head.0
   Type: Linear
   Shape: (128, 64)
   
6: mlp_head.1
   Type: ReLU
   
7: mlp_head.2
   Type: Linear
   Shape: (64, 1)        ← Final prediction
"""