core/evolution/train_ad_checker.py

"""
Applicability Domain Setup for Your Mixture DCN Model

Customized for your model that uses:
- InChI strings (not SMILES)
- 2-11 component mixtures
- Mole fractions
- GNN architecture from solvation_predictor
"""

import torch
import numpy as np
import pandas as pd
from typing import List, Tuple
from pathlib import Path
import sys
import os



from applicability_domain import ApplicabilityDomainChecker


class MixtureDCNEmbeddingExtractor:
    """
    Extract embeddings from your trained mixture DCN model ENSEMBLE.
    
    Uses all 10 models and averages their embeddings for more robust representation.
    """
    
    def __init__(self, model_dir: str, device='cpu'):
        """
        Args:
            model_dir: Directory containing all model .pt files
            device: 'cpu', 'cuda', or 'mps'
        """
        self.device = torch.device(device)
        
        # Import model loading function (aliases already set up above)
        from core.predictors.mixture.solvation_predictor.train.train import load_checkpoint, create_logger
        
        # Find all model files
        model_dir_path = Path(model_dir)
        self.model_files = sorted([f for f in model_dir_path.glob("*.pt")])
        
        if len(self.model_files) == 0:
            raise ValueError(f"No .pt model files found in {model_dir}")
        
        print(f"Found {len(self.model_files)} models in ensemble")
        
        # Initialize containers
        self.models = []
        self.hooks = []  # Initialize BEFORE loading models to prevent AttributeError
        self.embeddings = []
        
        # Load all models
        logging_obj = create_logger("embedding_extractor", ".")
        
        for i, model_path in enumerate(self.model_files):
            print(f"  Loading model {i+1}/{len(self.model_files)}: {model_path.name}")
            
            args = self._create_model_args(str(model_path))
            model = load_checkpoint(str(model_path), args, logger=logging_obj)
            model.to(self.device)
            model.eval()
            
            self.models.append(model)
        
        # Register hooks on all models
        self._register_hooks()
        
        print(f"✓ Loaded {len(self.models)} models")
        print(f"  Device: {self.device}")
    
    def _create_model_args(self, model_path, device=None):
        """Create minimal args needed for model loading / DataPoint / DataTensor."""
        _device = device  # capture for closure

        class ModelArgs:
            def __init__(self):
                self.model_path = model_path
                self.cuda = torch.cuda.is_available() or torch.backends.mps.is_available()

                # Model architecture
                self.depth = 4
                self.mpn_hidden = 200
                self.mpn_dropout = 0.0
                self.mpn_activation = "LeakyReLU"
                self.mpn_bias = False
                self.aggregation = "mean"

                self.ffn_hidden = 500
                self.ffn_num_layers = 4
                self.ffn_dropout = 0.0
                self.ffn_activation = "LeakyReLU"
                self.ffn_bias = True

                self.num_targets = 1
                self.attention = False
                self.postprocess = False
                self.num_features = 0
                self.f_mol_size = 2
                self.num_mols = 11
                self.max_num_mols = 11

                # Additional attributes from inp.py
                self.property = "solvation"
                self.solute = False
                self.add_hydrogens_to_solvent = False
                self.uncertainty = False
                self.ensemble_variance = False
                self.mix = False
                self.morgan_fingerprint = "None"
                self.morgan_bits = 16
                self.morgan_radius = 2

                # Device — propagated so DataPoint/DataTensor match the model's device
                self.device = _device

                # Attention parameters
                self.att_hidden = 200
                self.att_dropout = 0.0
                self.att_bias = False
                self.att_activation = "ReLU"
                self.att_normalize = "sigmoid"
                self.att_first_normalize = False

        return ModelArgs()

    @classmethod
    def from_models(cls, models: list, model_dir: str = ".", device=None):
        """
        Create an extractor from already-loaded models (avoids loading them twice).

        Args:
            models: List of loaded PyTorch model objects (e.g. from MixtureDCNPredictor)
            model_dir: Path used only for args metadata (no files are read)
            device: If None, inferred from the first model's parameters
        """
        instance = cls.__new__(cls)

        # Infer device from the model if not specified
        if device is None:
            try:
                device = next(models[0].parameters()).device
            except StopIteration:
                device = torch.device('cpu')

        instance.device = torch.device(device) if not isinstance(device, torch.device) else device
        instance.model_files = [Path(model_dir)]  # Needed only for _create_model_args path metadata
        instance.models = list(models)
        instance.hooks = []
        instance.embeddings = []

        instance._register_hooks()

        print(f"✓ Embedding extractor reusing {len(instance.models)} pre-loaded models")
        print(f"  Device: {instance.device}")
        return instance
    
                

    
    def _register_hooks(self):
        """Register hooks on all models to capture embeddings."""
        for model_idx, model in enumerate(self.models):
            # Create a separate embeddings list for each model
            model_embeddings = []
            
            # Pre-hook: captures the INPUT to a module (before any transformation).
            # Used on model.ffn to get the mole-fraction-weighted mixture vector
            # produced by mixture_forward() — one vector per mixture, not per molecule.
            def create_pre_hook_fn(emb_list):
                def hook_fn(module, input):
                    inp = input[0] if isinstance(input, tuple) else input
                    emb_list.append(inp.detach().cpu())
                return hook_fn

            # Register hook - prefer ffn pre-hook (true mixture-level representation)
            hook = None
            if hasattr(model, 'ffn'):
                # Captures the mole-fraction-weighted sum of MPN embeddings:
                # shape (N_mixtures, ffn_input_size) — fires once per forward pass
                hook = model.ffn.register_forward_pre_hook(
                    create_pre_hook_fn(model_embeddings)
                )
            elif hasattr(model, 'encoder'):
                hook = model.encoder.register_forward_hook(
                    create_pre_hook_fn(model_embeddings)
                )
            elif hasattr(model, 'predictor'):
                layers = list(model.children())
                hook = layers[-2].register_forward_hook(
                    create_pre_hook_fn(model_embeddings)
                )

            if hook is None:
                available = [name for name, _ in model.named_children()]
                raise RuntimeError(
                    f"Could not find a layer to hook on model {model_idx}. "
                    f"Available top-level children: {available}"
                )

            self.hooks.append(hook)
            self.embeddings.append(model_embeddings)
        
        print(f"  Registered hooks on {len(self.hooks)} models")
    
    def extract_embeddings_from_mixtures(self, 
                                        mixture_data: List[dict]) -> np.ndarray:
        """
        Extract embeddings from mixture dictionaries using ENSEMBLE AVERAGING.
        
        Args:
            mixture_data: List of dicts with keys:
                - 'inchis': List of InChI strings
                - 'fractions': List of mole fractions (N-1 for N components)
        
        Returns:
            embeddings: numpy array of shape (n_mixtures, embedding_dim)
                       (averaged across all models in ensemble)
        """
        from core.predictors.mixture.solvation_predictor.data.data import (
            DataPoint, DatapointList, MolencoderDatabase, DataTensor
        )

        # Clear previous embeddings
        for emb_list in self.embeddings:
            emb_list.clear()

        args = self._create_model_args(str(self.model_files[0]), device=self.device)

        with torch.no_grad():
            # Build DataPoints - one shared MolencoderDatabase caches mol graphs
            mol_encoder_db = MolencoderDatabase()
            datapoints = []
            for mix in mixture_data:
                dp = DataPoint(
                    smiles=mix['inchis'],  # Model accepts InChI in smiles field
                    targets=[0.0],         # Dummy target
                    features=[],
                    molefracs=mix['fractions'],
                    inp=args,
                    mol_encoders=mol_encoder_db,
                )
                datapoints.append(dp)

            data = DatapointList(datapoints)

            # Build tensors list: one DataTensor per molecule position (same as train.py)
            mol_encodings = [[] for _ in range(args.max_num_mols)]
            for d in datapoints:
                encoders = d.get_mol_encoder()
                # Pad short mixtures by repeating the first encoder
                while len(encoders) < args.max_num_mols:
                    encoders.append(encoders[0])
                for pos, enc in enumerate(encoders):
                    mol_encodings[pos].append(enc)

            tensors = [
                DataTensor(mol_enc_list, args, property=args.property)
                for mol_enc_list in mol_encodings
            ]

            # Forward pass through ALL models (hooks will capture MPN embeddings)
            for model in self.models:
                _ = model(data, tensors)
        
        # Collect embeddings from all models
        all_model_embeddings = []
        
        for model_idx, emb_list in enumerate(self.embeddings):
            if len(emb_list) > 0:
                # Concatenate batches for this model
                model_emb = torch.cat(emb_list, dim=0)
                all_model_embeddings.append(model_emb)
        
        if len(all_model_embeddings) == 0:
            raise RuntimeError("No embeddings captured! Check hook registration.")
        
        # Stack and average across models
        # Shape: (n_models, n_samples, embedding_dim)
        stacked_embeddings = torch.stack(all_model_embeddings, dim=0)
        
        # Average across models
        # Shape: (n_samples, embedding_dim)
        averaged_embeddings = stacked_embeddings.mean(dim=0)
        
        print(f"  Extracted embeddings from {len(all_model_embeddings)} models")
        print(f"  Shape per model: {all_model_embeddings[0].shape}")
        print(f"  Averaged shape: {averaged_embeddings.shape}")
        
        return averaged_embeddings.numpy()
    
    def __del__(self):
        """Remove all hooks when done."""
        for hook in self.hooks:
            hook.remove()


def load_training_data_for_ad(csv_path: str, max_samples: int = None) -> List[dict]:
    """
    Load training mixtures in format needed for embedding extraction.
    
    Args:
        csv_path: Path to formatted training CSV
        max_samples: Maximum number of samples to load (None = all)
    
    Returns:
        List of mixture dictionaries
    """
    # Try different encodings (same as your prediction script)
    encodings = ['utf-8', 'latin-1', 'iso-8859-1', 'cp1252']
    df = None
    
    for encoding in encodings:
        try:
            print(f"  Trying encoding: {encoding}")
            df = pd.read_csv(csv_path, encoding=encoding)
            print(f"  ✓ Successfully read with {encoding} encoding")
            break
        except UnicodeDecodeError:
            continue
    
    if df is None:
        raise ValueError(f"Could not read CSV with any of these encodings: {encodings}")
    
    print(f"  Loaded {len(df)} rows from CSV")
    
    mixtures = []
    
    for idx, row in df.iterrows():
        if max_samples and idx >= max_samples:
            break
        
        # Extract InChI strings (try multiple column name formats)
        inchis = []
        
        # Format 1: "fuel1 inchi", "fuel2 inchi", etc. (with space)
        for i in range(1, 12):
            col = f'fuel{i} inchi'
            if col in df.columns and pd.notna(row.get(col)):
                inchi = str(row[col]).strip()
                if inchi and inchi != 'nan':
                    inchis.append(inchi)
        
        # Format 2: "fuel1_inchi", "fuel2_inchi", etc. (with underscore)
        if len(inchis) == 0:
            for i in range(1, 12):
                col = f'fuel{i}_inchi'
                if col in df.columns and pd.notna(row.get(col)):
                    inchi = str(row[col]).strip()
                    if inchi and inchi != 'nan':
                        inchis.append(inchi)
        
        if len(inchis) == 0:
            continue
        
        # Extract mole fractions (N-1 for N components)
        fractions = []
        
        # Format 1: "molar fraction fuel 1", etc.
        for i in range(1, len(inchis)):
            col = f'molar fraction fuel {i}'
            if col in df.columns and pd.notna(row.get(col)):
                fractions.append(float(row[col]))
        
        # Format 2: "frac_fuel1 (molar)", etc. (from formatted data)
        if len(fractions) == 0:
            for i in range(1, len(inchis)):
                col = f'frac_fuel{i} (molar)'
                if col in df.columns and pd.notna(row.get(col)):
                    fractions.append(float(row[col]))
        
        # Only add if we have fractions
        if len(fractions) == len(inchis) - 1:
            mixtures.append({
                'inchis': inchis,
                'fractions': fractions
            })
    
    print(f"  ✓ Extracted {len(mixtures)} valid mixtures")
    
    return mixtures


def train_ad_checker_for_mixture_model(
    training_csv: str,
    model_dir: str,  # Changed from model_path to model_dir
    output_path: str = "models/mixture_ad_checker.pkl",
    nu: float = 0.02,
    device: str = 'cpu'
):
    """
    Train AD checker for your mixture DCN model ENSEMBLE.
    
    Args:
        training_csv: Path to formatted training data CSV
        model_dir: Directory containing all .pt model files
        output_path: Where to save AD checker
        nu: Outlier fraction for One-Class SVM
        device: 'cpu', 'cuda', or 'mps'
    """
    print("="*70)
    print("TRAINING APPLICABILITY DOMAIN CHECKER FOR MIXTURE DCN MODEL")
    print("="*70)
    
    # Step 1: Load training data
    print("\nStep 1: Loading training data...")
    mixtures = load_training_data_for_ad(training_csv)
    print(f"✓ Loaded {len(mixtures)} training mixtures")
    
    # Step 2: Create embedding extractor with ENSEMBLE
    print("\nStep 2: Loading model ensemble and creating embedding extractor...")
    extractor = MixtureDCNEmbeddingExtractor(model_dir, device=device)
    
    # Step 3: Extract embeddings (averaged across all models)
    print("\nStep 3: Extracting embeddings from training set...")
    print("  (This may take several minutes with 10 models...)")
    
    train_embeddings = extractor.extract_embeddings_from_mixtures(mixtures)
    print(f"✓ Extracted embeddings: shape {train_embeddings.shape}")
    print(f"  (Averaged across {len(extractor.models)} models)")
    
    # Step 4: Train One-Class SVM
    print("\nStep 4: Training One-Class SVM...")
    
    ad_checker = ApplicabilityDomainChecker(
        nu=nu,
        kernel='rbf',
        gamma='scale'
    )
    
    ad_checker.fit(train_embeddings)
    
    # Step 5: Validate on training set
    print("\nStep 5: Validating on training set...")
    
    train_in_domain = ad_checker.is_in_domain(train_embeddings)
    train_confidence = ad_checker.get_confidence_scores(train_embeddings)
    
    print(f"  In domain: {train_in_domain.sum()}/{len(train_in_domain)} ({train_in_domain.mean()*100:.1f}%)")
    print(f"  Mean confidence: {train_confidence.mean():.1f}%")
    print(f"  Confidence range: [{train_confidence.min():.1f}, {train_confidence.max():.1f}]")
    
    # Step 6: Save AD checker
    print("\nStep 6: Saving AD checker...")
    
    output_dir = Path(output_path).parent
    output_dir.mkdir(parents=True, exist_ok=True)
    
    ad_checker.save(output_path)
    
    print("\n" + "="*70)
    print("TRAINING COMPLETE!")
    print("="*70)
    print(f"\nAD checker saved to: {output_path}")
    print(f"Embedding extractor uses: {len(extractor.models)} models from {model_dir}")
    print("\nNext steps:")
    print("  1. Use this AD checker during evolution")
    print("  2. See mixture_evolution_with_ad.py for integration")
    print("="*70)
    
    return ad_checker, extractor


# =============================================================================
# MAIN SCRIPT
# =============================================================================

if __name__ == "__main__":
    # Configuration
    TRAINING_CSV = "data/database/mixture_training_dataset.csv"  # Your formatted training data
    MODEL_DIR = "models/mixture/dcn"  # Directory with all 10 models
    OUTPUT_PATH = "models/mixture_ad_checker.pkl"
    
    # Device selection
    if torch.cuda.is_available():
        device = 'cuda'
    elif torch.backends.mps.is_available():
        device = 'mps'
    else:
        device = 'cpu'
    
    print(f"Using device: {device}")
    
    # Check if model directory exists
    if not Path(MODEL_DIR).exists():
        print(f"\n✗ ERROR: Model directory not found: {MODEL_DIR}")
        print("Please update MODEL_DIR to point to your DCN models directory")
        sys.exit(1)
    
    model_files = list(Path(MODEL_DIR).glob("*.pt"))
    print(f"Found {len(model_files)} model files in {MODEL_DIR}")
    
    if len(model_files) == 0:
        print(f"\n✗ ERROR: No .pt model files found in {MODEL_DIR}")
        sys.exit(1)
    
    # Train AD checker
    try:
        ad_checker, extractor = train_ad_checker_for_mixture_model(
            training_csv=TRAINING_CSV,
            model_dir=MODEL_DIR,  # Now using directory
            output_path=OUTPUT_PATH,
            nu=0.02,  # 2% outliers
            device=device
        )
        
        print("\n✓ Setup complete! You can now use the AD checker in your evolution code.")
        
    except FileNotFoundError as e:
        print(f"\n✗ ERROR: File not found")
        print(f"  {e}")
        print("\nPlease update the paths at the bottom of this script:")
        print(f"  TRAINING_CSV = '{TRAINING_CSV}'")
        print(f"  MODEL_DIR = '{MODEL_DIR}'")
        
    except Exception as e:
        print(f"\n✗ ERROR: {type(e).__name__}")
        print(f"  {e}")
        import traceback
        traceback.print_exc()


# =============================================================================
# USAGE IN EVOLUTION
# =============================================================================

"""
Once trained, use in your mixture evolution like this:

from mixture_evolution_with_ad import MixtureAwareMolecularEvolutionWithAD

# Load AD checker
ad_checker = ApplicabilityDomainChecker.load('models/mixture_ad_checker.pkl')

# Load embedding extractor (loads ALL 10 models)
extractor = MixtureDCNEmbeddingExtractor(
    'models/mixture/dcn',  # Directory!
    device='cpu'
)

# During evolution, for each batch of new additives:
new_mixtures = [
    {
        'inchis': [additive_inchi] + base_fuel_inchis,
        'fractions': [additive_fraction] + base_fuel_fractions
    }
    for additive_inchi in new_additive_inchis
]

# Extract embeddings (averaged across 10 models!)
embeddings = extractor.extract_embeddings_from_mixtures(new_mixtures)

# Check AD
in_domain = ad_checker.is_in_domain(embeddings)
confidence = ad_checker.get_confidence_scores(embeddings)

# Filter
for i, additive in enumerate(new_additives):
    if not in_domain[i]:
        print(f"{additive}: OUT OF AD! (confidence: {confidence[i]:.1f}%)")
        continue  # Skip this molecule
    
    # Use this molecule (high confidence prediction)
    ...


WHY USE ALL 10 MODELS FOR EMBEDDINGS?
-------------------------------------

Ensemble averaging gives MORE ROBUST embeddings:
- Single model: Embedding might be noisy
- 10 models averaged: Smoother, more stable representation
- Better AD detection: Less false positives/negatives

The embeddings capture the "consensus" of what all 10 models learned!
"""