CAFA 6 Protein Function Prediction - Complete Advanced Solution
Features: Hyperparameter Tuning, Multiple ESM Models, Structure Integration, Adaptive Thresholds
 

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split, KFold
from sklearn.preprocessing import MultiLabelBinarizer
from sklearn.metrics import f1_score, precision_score, recall_score
import networkx as nx
from collections import defaultdict, deque
import warnings
import json
from typing import Dict, List, Tuple, Optional
warnings.filterwarnings('ignore')

# ==================== HYPERPARAMETER TUNING WITH OPTUNA ====================

try:
    import optuna
    from optuna.trial import Trial
    OPTUNA_AVAILABLE = True
except ImportError:
    OPTUNA_AVAILABLE = False
    print("Optuna not available. Install with: pip install optuna")

try:
    import wandb
    WANDB_AVAILABLE = True
except ImportError:
    WANDB_AVAILABLE = False
    print("W&B not available. Install with: pip install wandb")


class HyperparameterTuner:
    
    
    def __init__(self, model_class, X_train, y_train, X_val, y_val, 
                 device='cuda', use_wandb=False, project_name='cafa6-tuning'):
        self.model_class = model_class
        self.X_train = X_train
        self.y_train = y_train
        self.X_val = X_val
        self.y_val = y_val
        self.device = device
        self.use_wandb = use_wandb and WANDB_AVAILABLE
        self.project_name = project_name
        self.best_params = None
        self.best_score = 0
        
    def objective(self, trial: 'Trial') -> float:
        
        # Suggest hyperparameters
        params = {
            'learning_rate': trial.suggest_float('learning_rate', 1e-5, 1e-2, log=True),
            'weight_decay': trial.suggest_float('weight_decay', 1e-6, 1e-2, log=True),
            'dropout': trial.suggest_float('dropout', 0.1, 0.5),
            'hidden_dim_1': trial.suggest_categorical('hidden_dim_1', [512, 768, 1024, 1536]),
            'hidden_dim_2': trial.suggest_categorical('hidden_dim_2', [256, 384, 512, 768]),
            'hidden_dim_3': trial.suggest_categorical('hidden_dim_3', [128, 256, 384]),
            'batch_size': trial.suggest_categorical('batch_size', [16, 32, 64, 128]),
            'num_heads': trial.suggest_categorical('num_heads', [4, 8, 16]),
            'use_batch_norm': trial.suggest_categorical('use_batch_norm', [True, False]),
            'activation': trial.suggest_categorical('activation', ['relu', 'gelu', 'silu']),
        }
        
        # Initialize W&B run if enabled
        if self.use_wandb:
            wandb.init(
                project=self.project_name,
                config=params,
                reinit=True,
                name=f"trial_{trial.number}"
            )
        
        # Create model with suggested parameters
        model = self.model_class(
            esm_dim=self.X_train.shape[1],
            hidden_dims=[params['hidden_dim_1'], params['hidden_dim_2'], params['hidden_dim_3']],
            num_labels=self.y_train.shape[1],
            num_heads=params['num_heads'],
            dropout=params['dropout'],
            use_batch_norm=params['use_batch_norm'],
            activation=params['activation']
        ).to(self.device)
        
        # Setup training
        optimizer = torch.optim.AdamW(
            model.parameters(),
            lr=params['learning_rate'],
            weight_decay=params['weight_decay']
        )
        
        criterion = nn.BCEWithLogitsLoss()
        
        # Create data loaders
        train_dataset = torch.utils.data.TensorDataset(
            torch.FloatTensor(self.X_train),
            torch.FloatTensor(self.y_train)
        )
        val_dataset = torch.utils.data.TensorDataset(
            torch.FloatTensor(self.X_val),
            torch.FloatTensor(self.y_val)
        )
        
        train_loader = DataLoader(train_dataset, batch_size=params['batch_size'], shuffle=True)
        val_loader = DataLoader(val_dataset, batch_size=params['batch_size'], shuffle=False)
        
        # Training loop
        best_val_f1 = 0
        patience = 0
        max_patience = 5
        num_epochs = 30
        
        for epoch in range(num_epochs):
            # Train
            model.train()
            train_loss = 0
            for batch_X, batch_y in train_loader:
                batch_X, batch_y = batch_X.to(self.device), batch_y.to(self.device)
                
                optimizer.zero_grad()
                outputs = model(batch_X)
                loss = criterion(outputs, batch_y)
                loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
                optimizer.step()
                
                train_loss += loss.item()
            
            # Validate
            model.eval()
            val_preds = []
            val_labels = []
            
            with torch.no_grad():
                for batch_X, batch_y in val_loader:
                    batch_X = batch_X.to(self.device)
                    outputs = model(batch_X)
                    probs = torch.sigmoid(outputs)
                    val_preds.append(probs.cpu().numpy())
                    val_labels.append(batch_y.numpy())
            
            val_preds = np.vstack(val_preds)
            val_labels = np.vstack(val_labels)
            
            # Calculate F1
            preds_binary = (val_preds > 0.5).astype(int)
            val_f1 = f1_score(val_labels, preds_binary, average='micro', zero_division=0)
            
            # Log metrics
            if self.use_wandb:
                wandb.log({
                    'epoch': epoch,
                    'train_loss': train_loss / len(train_loader),
                    'val_f1': val_f1
                })
            
            # Early stopping
            if val_f1 > best_val_f1:
                best_val_f1 = val_f1
                patience = 0
            else:
                patience += 1
                if patience >= max_patience:
                    break
            
            # Report intermediate value for pruning
            trial.report(best_val_f1, epoch)
            
            if trial.should_prune():
                if self.use_wandb:
                    wandb.finish()
                raise optuna.TrialPruned()
        
        if self.use_wandb:
            wandb.log({'best_val_f1': best_val_f1})
            wandb.finish()
        
        return best_val_f1
    
    def tune(self, n_trials=50, timeout=None):
        
        if not OPTUNA_AVAILABLE:
            print("Optuna not available. Skipping hyperparameter tuning.")
            return None
        
        print(f"\n{'='*70}")
        print(f"STARTING HYPERPARAMETER TUNING ({n_trials} trials)")
        print(f"{'='*70}\n")
        
        study = optuna.create_study(
            direction='maximize',
            pruner=optuna.pruners.MedianPruner(n_startup_trials=5, n_warmup_steps=5)
        )
        
        study.optimize(self.objective, n_trials=n_trials, timeout=timeout)
        
        # Save best parameters
        self.best_params = study.best_params
        self.best_score = study.best_value
        
        print(f"\n{'='*70}")
        print("HYPERPARAMETER TUNING COMPLETE")
        print(f"{'='*70}")
        print(f"Best F1 Score: {self.best_score:.4f}")
        print("\nBest Parameters:")
        for param, value in self.best_params.items():
            print(f"  {param}: {value}")
        print(f"{'='*70}\n")
        
        # Save results
        with open('best_hyperparameters.json', 'w') as f:
            json.dump({
                'best_params': self.best_params,
                'best_score': self.best_score,
                'n_trials': len(study.trials)
            }, f, indent=2)
        
        return self.best_params


# ==================== MULTI-MODEL ESM EXTRACTOR ====================

class MultiESMExtractor:
  
    
    # Available ESM-2 models with their specs
    ESM_MODELS = {
        'esm2_t6_8M_UR50D': {'params': '8M', 'layers': 6, 'embed_dim': 320},
        'esm2_t12_35M_UR50D': {'params': '35M', 'layers': 12, 'embed_dim': 480},
        'esm2_t30_150M_UR50D': {'params': '150M', 'layers': 30, 'embed_dim': 640},
        'esm2_t33_650M_UR50D': {'params': '650M', 'layers': 33, 'embed_dim': 1280},
        'esm2_t36_3B_UR50D': {'params': '3B', 'layers': 36, 'embed_dim': 2560},
        'esm2_t48_15B_UR50D': {'params': '15B', 'layers': 48, 'embed_dim': 5120},
    }
    
    def __init__(self, model_names=['esm2_t33_650M_UR50D'], device='cuda', 
                 pooling_strategy='mean', layer_selection='last'):
        """
        Args:
            model_names: list of ESM model names to use
            device: 'cuda' or 'cpu'
            pooling_strategy: 'mean', 'max', 'cls', or 'attention'
            layer_selection: 'last', 'second_to_last', or list of layer indices
        """
        self.model_names = model_names
        self.device = device
        self.pooling_strategy = pooling_strategy
        self.layer_selection = layer_selection
        self.models = {}
        self.alphabets = {}
        self.embedding_dims = []
        
        self._load_models()
    
    def _load_models(self):
        """Load all specified ESM models"""
        try:
            import esm
            
            for model_name in self.model_names:
                if model_name not in self.ESM_MODELS:
                    print(f"Warning: {model_name} not recognized. Skipping.")
                    continue
                
                print(f"Loading {model_name} ({self.ESM_MODELS[model_name]['params']} parameters)...")
                
                model, alphabet = esm.pretrained.load_model_and_alphabet(model_name)
                model = model.to(self.device)
                model.eval()
                
                self.models[model_name] = model
                self.alphabets[model_name] = alphabet
                self.embedding_dims.append(self.ESM_MODELS[model_name]['embed_dim'])
                
                print(f"✓ Loaded successfully. Embedding dim: {self.ESM_MODELS[model_name]['embed_dim']}")
        
        except ImportError:
            print("WARNING: ESM not installed. Using random embeddings.")
            print("Install with: pip install fair-esm")
            self.models = None
    
    def _apply_pooling(self, embeddings, mask=None):
        
        if self.pooling_strategy == 'mean':
            if mask is not None:
                # Masked mean pooling
                masked_embeddings = embeddings * mask.unsqueeze(-1)
                return masked_embeddings.sum(dim=1) / mask.sum(dim=1, keepdim=True)
            else:
                return embeddings.mean(dim=1)
        
        elif self.pooling_strategy == 'max':
            return embeddings.max(dim=1)[0]
        
        elif self.pooling_strategy == 'cls':
            return embeddings[:, 0, :]  # CLS token
        
        elif self.pooling_strategy == 'attention':
            # Learnable attention pooling
            attention_weights = torch.softmax(embeddings.mean(dim=-1), dim=1)
            return (embeddings * attention_weights.unsqueeze(-1)).sum(dim=1)
    
    def extract_embeddings(self, sequences, batch_size=8, 
                          combine_method='concat'):
     
        if self.models is None:
            # Return random embeddings
            total_dim = sum(self.embedding_dims) if combine_method == 'concat' else self.embedding_dims[0]
            return np.random.randn(len(sequences), total_dim).astype(np.float32)
        
        all_model_embeddings = []
        
        for model_name, model in self.models.items():
            print(f"\nExtracting embeddings with {model_name}...")
            alphabet = self.alphabets[model_name]
            batch_converter = alphabet.get_batch_converter()
            
            model_embeddings = []
            
            for i in range(0, len(sequences), batch_size):
                batch_seqs = sequences[i:i+batch_size]
                data = [(f"protein_{j}", seq) for j, seq in enumerate(batch_seqs)]
                
                batch_labels, batch_strs, batch_tokens = batch_converter(data)
                batch_tokens = batch_tokens.to(self.device)
                
                with torch.no_grad():
                    # Get representations from specified layers
                    if self.layer_selection == 'last':
                        repr_layers = [self.ESM_MODELS[model_name]['layers']]
                    elif self.layer_selection == 'second_to_last':
                        repr_layers = [self.ESM_MODELS[model_name]['layers'] - 1]
                    else:
                        repr_layers = self.layer_selection
                    
                    results = model(batch_tokens, repr_layers=repr_layers)
                    
                    # Get embeddings and pool
                    layer_idx = repr_layers[0]
                    embeddings = results['representations'][layer_idx]
                    
                    # Create mask (exclude padding and special tokens)
                    mask = (batch_tokens != alphabet.padding_idx)
                    mask[:, 0] = False  # Exclude CLS
                    mask[:, -1] = False  # Exclude EOS
                    
                    # Apply pooling
                    pooled = self._apply_pooling(embeddings[:, 1:-1, :])  # Exclude special tokens
                    model_embeddings.append(pooled.cpu().numpy())
                
                if (i // batch_size + 1) % 10 == 0:
                    print(f"  Processed {min(i+batch_size, len(sequences))}/{len(sequences)}")
            
            model_embeddings = np.vstack(model_embeddings)
            all_model_embeddings.append(model_embeddings)
        
        # Combine embeddings from multiple models
        if combine_method == 'concat':
            combined = np.hstack(all_model_embeddings)
        elif combine_method == 'avg':
            combined = np.mean(all_model_embeddings, axis=0)
        elif combine_method == 'weighted_avg':
            # Weight by model size
            weights = [self.ESM_MODELS[name]['layers'] for name in self.model_names]
            weights = np.array(weights) / sum(weights)
            combined = sum(w * emb for w, emb in zip(weights, all_model_embeddings))
        
        return combined
    
    def compare_models(self, sequences, labels, sample_size=100):
       
        print(f"\n{'='*70}")
        print("COMPARING ESM MODELS")
        print(f"{'='*70}\n")
        
        # Sample data
        if len(sequences) > sample_size:
            indices = np.random.choice(len(sequences), sample_size, replace=False)
            sample_seqs = [sequences[i] for i in indices]
            sample_labels = labels[indices]
        else:
            sample_seqs = sequences
            sample_labels = labels
        
        results = {}
        
        # Test each model individually
        for model_name in self.model_names:
            single_extractor = MultiESMExtractor([model_name], self.device)
            embeddings = single_extractor.extract_embeddings(sample_seqs, batch_size=8)
            
            # Quick evaluation
            X_train, X_val, y_train, y_val = train_test_split(
                embeddings, sample_labels, test_size=0.2, random_state=42
            )
            
            # Train simple model
            from sklearn.multioutput import MultiOutputClassifier
            from sklearn.linear_model import LogisticRegression
            
            clf = MultiOutputClassifier(LogisticRegression(max_iter=100))
            clf.fit(X_train, y_train)
            
            # Evaluate
            val_preds = clf.predict(X_val)
            f1 = f1_score(y_val, val_preds, average='micro', zero_division=0)
            
            results[model_name] = {
                'embedding_dim': embeddings.shape[1],
                'f1_score': f1
            }
            
            print(f"{model_name}:")
            print(f"  Embedding Dim: {embeddings.shape[1]}")
            print(f"  F1 Score: {f1:.4f}")
            print()
        
        # Recommend best model
        best_model = max(results.items(), key=lambda x: x[1]['f1_score'])
        print(f"✓ Best Model: {best_model[0]} (F1: {best_model[1]['f1_score']:.4f})")
        print(f"{'='*70}\n")
        
        return results


# ==================== PROTEIN STRUCTURE FEATURES ====================

class ProteinStructureFeatureExtractor:
     
    
    def __init__(self, structure_source='alphafold'):
        """
        Args:
            structure_source: 'pdb', 'alphafold', or 'predicted'
        """
        self.structure_source = structure_source
        
        try:
            from Bio.PDB import PDBParser, DSSP, PPBuilder
            from Bio.PDB.Polypeptide import PPBuilder
            self.parser = PDBParser(QUIET=True)
            self.dssp_available = True
        except ImportError:
            print("Biopython not available. Install with: pip install biopython")
            self.dssp_available = False
    
    def extract_from_pdb(self, pdb_file):
        """Extract structural features from PDB file"""
        if not self.dssp_available:
            return None
        
        from Bio.PDB import DSSP
        
        structure = self.parser.get_structure('protein', pdb_file)
        model = structure[0]
        
        features = {
            'secondary_structure': [],
            'solvent_accessibility': [],
            'phi_psi_angles': [],
            'contact_map': None,
            'distance_matrix': None
        }
        
        # Extract secondary structure using DSSP
        try:
            dssp = DSSP(model, pdb_file)
            for residue in dssp:
                features['secondary_structure'].append(self._encode_ss(residue[2]))
                features['solvent_accessibility'].append(residue[3])
                features['phi_psi_angles'].append((residue[4], residue[5]))
        except:
            pass
        
        # Calculate contact map
        features['contact_map'] = self._calculate_contact_map(model)
        
        # Calculate distance matrix
        features['distance_matrix'] = self._calculate_distance_matrix(model)
        
        return features
    
    def _encode_ss(self, ss_char):
     
        ss_types = {'H': 0, 'B': 1, 'E': 2, 'G': 3, 'I': 4, 'T': 5, 'S': 6, '-': 7}
        encoding = np.zeros(8)
        encoding[ss_types.get(ss_char, 7)] = 1
        return encoding
    
    def _calculate_contact_map(self, model, threshold=8.0):
         
        residues = list(model.get_residues())
        n = len(residues)
        contact_map = np.zeros((n, n))
        
        for i, res1 in enumerate(residues):
            for j, res2 in enumerate(residues[i+1:], i+1):
                try:
                    ca1 = res1['CA']
                    ca2 = res2['CA']
                    dist = ca1 - ca2
                    if dist < threshold:
                        contact_map[i, j] = 1
                        contact_map[j, i] = 1
                except:
                    pass
        
        return contact_map
    
    def _calculate_distance_matrix(self, model):
       
        residues = list(model.get_residues())
        n = len(residues)
        dist_matrix = np.zeros((n, n))
        
        for i, res1 in enumerate(residues):
            for j, res2 in enumerate(residues):
                try:
                    ca1 = res1['CA']
                    ca2 = res2['CA']
                    dist_matrix[i, j] = ca1 - ca2
                except:
                    dist_matrix[i, j] = 999.0
        
        return dist_matrix
    
    def extract_structure_embedding(self, features, method='graph'):
      
        if method == 'graph':
            # Graph-based embedding using contact map
            if features['contact_map'] is not None:
                # Simple graph features
                contact_map = features['contact_map']
                embedding = [
                    np.sum(contact_map) / contact_map.size,  # Contact density
                    np.mean(np.sum(contact_map, axis=0)),  # Avg contacts per residue
                    np.std(np.sum(contact_map, axis=0)),  # Contact variability
                ]
                
                # Add secondary structure content
                if len(features['secondary_structure']) > 0:
                    ss_array = np.array(features['secondary_structure'])
                    ss_content = np.mean(ss_array, axis=0)
                    embedding.extend(ss_content.tolist())
                
                return np.array(embedding)
        
        return None
    
    def augment_with_structure(self, sequence_embeddings, structure_files):
    
        structure_embeddings = []
        
        for pdb_file in structure_files:
            if pdb_file is not None:
                features = self.extract_from_pdb(pdb_file)
                struct_emb = self.extract_structure_embedding(features)
                structure_embeddings.append(struct_emb if struct_emb is not None else np.zeros(11))
            else:
                structure_embeddings.append(np.zeros(11))
        
        structure_embeddings = np.array(structure_embeddings)
        
        # Concatenate
        augmented = np.hstack([sequence_embeddings, structure_embeddings])
        
        return augmented


# ==================== ADAPTIVE THRESHOLD OPTIMIZATION ====================

class AdaptiveThresholdOptimizer:
 
    
    def __init__(self, go_hierarchy=None):
        self.go_hierarchy = go_hierarchy
        self.global_threshold = 0.5
        self.subontology_thresholds = {}
        self.term_thresholds = {}
        self.threshold_method = 'subontology'  # 'global', 'subontology', or 'per_term'
    
    def optimize_global(self, predictions, labels, term_weights=None):
        """Find single optimal threshold"""
        thresholds = np.arange(0.05, 0.95, 0.05)
        best_f1 = 0
        best_threshold = 0.5
        
        for threshold in thresholds:
            binary_preds = (predictions >= threshold).astype(int)
            f1 = self._calculate_weighted_f1(binary_preds, labels, term_weights)
            
            if f1 > best_f1:
                best_f1 = f1
                best_threshold = threshold
        
        self.global_threshold = best_threshold
        self.threshold_method = 'global'
        
        print(f"Global optimal threshold: {best_threshold:.3f} (F1: {best_f1:.4f})")
        
        return best_threshold, best_f1
    
    def optimize_per_subontology(self, predictions, labels, term_list, 
                                 term_weights=None):
      
        if self.go_hierarchy is None:
            print("GO hierarchy not available. Using global threshold.")
            return self.optimize_global(predictions, labels, term_weights)
        
        print("\nOptimizing per-subontology thresholds...")
        
        ontologies = {}
        for ontology in ['MF', 'BP', 'CC']:
            indices = [i for i, term in enumerate(term_list)
                      if self.go_hierarchy.subontologies.get(term) == ontology]
            
            if len(indices) > 0:
                sub_preds = predictions[:, indices]
                sub_labels = labels[:, indices]
                sub_weights = term_weights[indices] if term_weights is not None else None
                
                thresholds = np.arange(0.05, 0.95, 0.05)
                best_f1 = 0
                best_threshold = 0.5
                
                for threshold in thresholds:
                    binary_preds = (sub_preds >= threshold).astype(int)
                    f1 = self._calculate_weighted_f1(binary_preds, sub_labels, sub_weights)
                    
                    if f1 > best_f1:
                        best_f1 = f1
                        best_threshold = threshold
                
                self.subontology_thresholds[ontology] = best_threshold
                ontologies[ontology] = (best_threshold, best_f1, len(indices))
                
                print(f"  {ontology}: threshold={best_threshold:.3f}, F1={best_f1:.4f}, terms={len(indices)}")
        
        self.threshold_method = 'subontology'
        
        return self.subontology_thresholds
    
    def optimize_per_term(self, predictions, labels, term_weights=None,
                         min_samples=10):
        """Find optimal threshold for each GO term (computationally expensive)"""
        print("\nOptimizing per-term thresholds (this may take a while)...")
        
        n_terms = predictions.shape[1]
        self.term_thresholds = np.full(n_terms, 0.5)
        
        thresholds = np.arange(0.1, 0.9, 0.1)
        
        for term_idx in range(n_terms):
            # Check if enough positive samples
            if np.sum(labels[:, term_idx]) < min_samples:
                continue
            
            best_f1 = 0
            best_threshold = 0.5
            
            term_preds = predictions[:, term_idx]
            term_labels = labels[:, term_idx]
            
            for threshold in thresholds:
                binary_preds = (term_preds >= threshold).astype(int)
                
                # Calculate F1 for this term
                tp = np.sum(binary_preds * term_labels)
                fp = np.sum(binary_preds * (1 - term_labels))
                fn = np.sum((1 - binary_preds) * term_labels)
                
                precision = tp / (tp + fp + 1e-10)
                recall = tp / (tp + fn + 1e-10)
                f1 = 2 * precision * recall / (precision + recall + 1e-10)
                
                if f1 > best_f1:
                    best_f1 = f1
                    best_threshold = threshold
            
            self.term_thresholds[term_idx] = best_threshold
            
            if (term_idx + 1) % 100 == 0:
                print(f"  Optimized {term_idx + 1}/{n_terms} terms")
        
        self.threshold_method = 'per_term'
        
        print(f"Per-term optimization complete.")
        print(f"  Mean threshold: {np.mean(self.term_thresholds):.3f}")
        print(f"  Std threshold: {np.std(self.term_thresholds):.3f}")
        
        return self.term_thresholds
    
    def apply_thresholds(self, predictions, term_list=None):
       
        if self.threshold_method == 'global':
            return (predictions >= self.global_threshold).astype(float) * predictions
        
        elif self.threshold_method == 'subontology' and term_list is not None:
            thresholded = predictions.copy()
            
            for ontology, threshold in self.subontology_thresholds.items():
                indices = [i for i, term in enumerate(term_list)
                          if self.go_hierarchy.subontologies.get(term) == ontology]
                
                mask = predictions[:, indices] >= threshold
                thresholded[:, indices] = predictions[:, indices] * mask
            
            return thresholded
        
        elif self.threshold_method == 'per_term':
            mask = predictions >= self.term_thresholds[np.newaxis, :]
            return predictions * mask
        
        else:
            # Default to 0.5
            return (predictions >= 0.5).astype(float) * predictions
    
    def _calculate_weighted_f1(self, predictions, labels, weights=None):
         
        if weights is None:
            weights = np.ones(predictions.shape[1])
        
        weighted_tp = np.sum((predictions * labels) * weights[np.newaxis, :])
        weighted_fp = np.sum((predictions * (1 - labels)) * weights[np.newaxis, :])
        weighted_fn = np.sum(((1 - predictions) * labels) * weights[np.newaxis, :])
        
        precision = weighted_tp / (weighted_tp + weighted_fp + 1e-10)
        recall = weighted_tp / (weighted_tp + weighted_fn + 1e-10)
        f1 = 2 * precision * recall / (precision + recall + 1e-10)
        
        return f1
    
    def save_thresholds(self, filepath='optimized_thresholds.json'):
        """Save optimized thresholds"""
        data = {
            'method': self.threshold_method,
            'global_threshold': float(self.global_threshold),
            'subontology_thresholds': {k: float(v) for k, v in self.subontology_thresholds.items()},
        }
        
        if self.threshold_method == 'per_term':
            data['term_thresholds'] = self.term_thresholds.tolist()
        
        with open(filepath, 'w') as f:
            json.dump(data, f, indent=2)
        
        print(f"Thresholds saved to {filepath}")
    
    def load_thresholds(self, filepath='optimized_thresholds.json'):
        """Load saved thresholds"""
        with open(filepath, 'r') as f:
            data = json.load(f)
        
        self.threshold_method = data['method']
        self.global_threshold = data['global_threshold']
        self.subontology_thresholds = data['subontology_thresholds']
        
        if 'term_thresholds' in data:
            self.term_thresholds = np.array(data['term_thresholds'])
        
        print(f"Thresholds loaded from {filepath}")


# ==================== ENHANCED MODEL WITH FLEXIBLE ARCHITECTURE ====================

class FlexibleProteinPredictor(nn.Module):
   
    
    def __init__(self, esm_dim=1280, hidden_dims=[1024, 512, 256],
                 num_labels=1000, num_heads=8, dropout=0.3,
                 use_batch_norm=True, activation='relu',
                 use_attention=True, use_residual=True):
        super(FlexibleProteinPredictor, self).__init__()
        
        self.use_attention = use_attention
        self.use_residual = use_residual
        
        # Activation function
        self.activation_fn = self._get_activation(activation)
        
        # Input projection
        self.input_proj = nn.Linear(esm_dim, hidden_dims[0])
        
        if use_batch_norm:
            self.input_norm = nn.BatchNorm1d(hidden_dims[0])
        else:
            self.input_norm = nn.LayerNorm(hidden_dims[0])
        
        # Multi-head attention (optional)
        if use_attention:
            self.attention = nn.MultiheadAttention(
                hidden_dims[0], num_heads, dropout=dropout, batch_first=True
            )
            self.attn_norm = nn.LayerNorm(hidden_dims[0])
        
        # Feed-forward layers
        layers = []
        for i in range(len(hidden_dims) - 1):
            layers.append(nn.Linear(hidden_dims[i], hidden_dims[i+1]))
            
            if use_batch_norm:
                layers.append(nn.BatchNorm1d(hidden_dims[i+1]))
            else:
                layers.append(nn.LayerNorm(hidden_dims[i+1]))
            
            layers.append(self._get_activation(activation))
            layers.append(nn.Dropout(dropout))
        
        self.ffn = nn.Sequential(*layers)
        
        # Output layer
        self.output = nn.Linear(hidden_dims[-1], num_labels)
        
        self.dropout = nn.Dropout(dropout)
    
    def _get_activation(self, name):
   
        activations = {
            'relu': nn.ReLU(),
            'gelu': nn.GELU(),
            'silu': nn.SiLU(),
            'tanh': nn.Tanh(),
            'leaky_relu': nn.LeakyReLU()
        }
        return activations.get(name, nn.ReLU())
    
    def forward(self, x):
        # Input projection
        x = self.input_proj(x)
        x = self.input_norm(x)
        x = self.activation_fn(x)
        
        # Attention (optional)
        if self.use_attention:
            x_unsqueezed = x.unsqueeze(1)
            attn_out, _ = self.attention(x_unsqueezed, x_unsqueezed, x_unsqueezed)
            attn_out = attn_out.squeeze(1)
            
            if self.use_residual:
                x = x + self.dropout(attn_out)
            else:
                x = attn_out
            
            x = self.attn_norm(x)
        
        # Feed-forward
        x = self.ffn(x)
        
        # Output
        return self.output(x)


# ==================== GO HIERARCHY (from previous version) ====================

class GOHierarchy:
  
    
    def __init__(self):
        self.graph = nx.DiGraph()
        self.subontologies = {}
        self.information_content = {}
        
    def compute_information_content(self, term_frequencies):
       
        total = sum(term_frequencies.values())
        for term, freq in term_frequencies.items():
            if freq > 0:
                self.information_content[term] = -np.log(freq / total)
            else:
                self.information_content[term] = 10.0
    
    def get_term_weights(self, terms):
     
        weights = []
        for term in terms:
            weights.append(self.information_content.get(term, 1.0))
        return np.array(weights)


# ==================== COMPLETE PIPELINE WITH ALL FEATURES ====================

def run_complete_pipeline_with_tuning():
    """
    Complete pipeline demonstrating all advanced features:
    1. Hyperparameter tuning with Optuna
    2. Multiple ESM model comparison
    3. Structure feature integration
    4. Adaptive threshold optimization
    """
    
    print("="*80)
    print("CAFA 6 COMPLETE ADVANCED PIPELINE")
    print("="*80)
    
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"Device: {device}\n")
    
    # ========== Step 1: Generate Sample Data ==========
    print("[1/8] Creating sample data...")
    
    train_sequences = [
        'MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEKAVQVKVKALPDAQFEVVHSLAKWKRQTLGQHDFSAGEGLYTHMKALRPDEDRLSALLRSGDGSP',
        'MAFSAEDVLKEYDRRRRMEALLLSLYYPNDRKLLDYKEWSPPRVQVECPKAPVEWNNPPSEKGLIVGHFSGIKYKGEKAQASEVDVNKMCCWVSKFKDAMRRYQGIQTCKIPGKVLSDLDAKIKAYNLTVEGVEGFVRYSRVTKQHVAAFLKELRHSKQYENVNLIHYILTDKRVDIQHLEKDLVKDFKALVESAHRMRQGHMINVKYILYQLLKKHGHGPDGPDILTVKTGSKGVLYDDSFRKIYTDLGWKFTPL',
    ] * 400
    
    train_labels = [
        ['GO:0003677', 'GO:0005515', 'GO:0005634'],
        ['GO:0004672', 'GO:0005524', 'GO:0005737', 'GO:0006468'],
    ] * 400
    
    mlb = MultiLabelBinarizer()
    mlb.fit([term for label_list in train_labels for term in label_list])
    y_encoded = mlb.transform(train_labels)
    
    print(f"Training samples: {len(train_sequences)}")
    print(f"Unique GO terms: {len(mlb.classes_)}")
    
    # ========== Step 2: ESM Model Comparison ==========
    print("\n[2/8] Comparing ESM models...")
    
    # For demo, use simulated embeddings
    # In production: use real ESM models
    print("Using simulated ESM embeddings for demo...")
    train_embeddings = np.random.randn(len(train_sequences), 1280).astype(np.float32)
    
    # Uncomment for real ESM comparison:
    # esm_extractor = MultiESMExtractor(
    #     model_names=['esm2_t12_35M_UR50D', 'esm2_t33_650M_UR50D'],
    #     device=device
    # )
    # results = esm_extractor.compare_models(train_sequences, y_encoded, sample_size=100)
    
    # ========== Step 3: Add Structure Features (Optional) ==========
    print("\n[3/8] Structure feature extraction (simulated)...")
    
    # Simulate structure features
    structure_features = np.random.randn(len(train_sequences), 11).astype(np.float32)
    combined_features = np.hstack([train_embeddings, structure_features])
    
    print(f"Combined feature dim: {combined_features.shape[1]}")
    
    # Uncomment for real structure features:
    # struct_extractor = ProteinStructureFeatureExtractor()
    # pdb_files = ['path/to/protein1.pdb', ...]
    # combined_features = struct_extractor.augment_with_structure(train_embeddings, pdb_files)
    
    # ========== Step 4: Split Data ==========
    print("\n[4/8] Splitting data...")
    
    X_train, X_val, y_train, y_val = train_test_split(
        combined_features, y_encoded,
        test_size=0.15,
        random_state=42
    )
    
    print(f"Train: {X_train.shape}, Val: {X_val.shape}")
    
    # ========== Step 5: Hyperparameter Tuning ==========
    print("\n[5/8] Hyperparameter tuning...")
    
    if OPTUNA_AVAILABLE:
        tuner = HyperparameterTuner(
            FlexibleProteinPredictor,
            X_train, y_train,
            X_val, y_val,
            device=device,
            use_wandb=False  # Set to True if W&B configured
        )
        
        best_params = tuner.tune(n_trials=20)  # Increase for better results
    else:
        print("Optuna not available. Using default hyperparameters.")
        best_params = {
            'learning_rate': 0.001,
            'weight_decay': 0.01,
            'dropout': 0.3,
            'hidden_dim_1': 1024,
            'hidden_dim_2': 512,
            'hidden_dim_3': 256,
            'batch_size': 32,
            'num_heads': 8,
            'use_batch_norm': True,
            'activation': 'relu'
        }
    
    # ========== Step 6: Train Best Model ==========
    print("\n[6/8] Training final model with best hyperparameters...")
    
    model = FlexibleProteinPredictor(
        esm_dim=X_train.shape[1],
        hidden_dims=[best_params['hidden_dim_1'], best_params['hidden_dim_2'], best_params['hidden_dim_3']],
        num_labels=y_train.shape[1],
        num_heads=best_params['num_heads'],
        dropout=best_params['dropout'],
        use_batch_norm=best_params['use_batch_norm'],
        activation=best_params['activation']
    ).to(device)
    
    # Training loop (simplified)
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=best_params['learning_rate'],
        weight_decay=best_params['weight_decay']
    )
    criterion = nn.BCEWithLogitsLoss()
    
    train_dataset = torch.utils.data.TensorDataset(
        torch.FloatTensor(X_train),
        torch.FloatTensor(y_train)
    )
    train_loader = DataLoader(train_dataset, batch_size=best_params['batch_size'], shuffle=True)
    
    for epoch in range(20):
        model.train()
        total_loss = 0
        
        for batch_X, batch_y in train_loader:
            batch_X, batch_y = batch_X.to(device), batch_y.to(device)
            
            optimizer.zero_grad()
            outputs = model(batch_X)
            loss = criterion(outputs, batch_y)
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
        
        if (epoch + 1) % 5 == 0:
            print(f"Epoch {epoch+1}: Loss={total_loss/len(train_loader):.4f}")
    
    # ========== Step 7: Adaptive Threshold Optimization ==========
    print("\n[7/8] Optimizing prediction thresholds...")
    
    # Get validation predictions
    model.eval()
    with torch.no_grad():
        val_tensor = torch.FloatTensor(X_val).to(device)
        val_preds = torch.sigmoid(model(val_tensor)).cpu().numpy()
    
    # Setup GO hierarchy
    go_hierarchy = GOHierarchy()
    term_frequencies = defaultdict(int)
    for terms in train_labels:
        for term in terms:
            term_frequencies[term] += 1
    go_hierarchy.compute_information_content(term_frequencies)
    
    # Map terms to subontologies (simplified)
    for i, term in enumerate(mlb.classes_):
        if i % 3 == 0:
            go_hierarchy.subontologies[term] = 'MF'
        elif i % 3 == 1:
            go_hierarchy.subontologies[term] = 'BP'
        else:
            go_hierarchy.subontologies[term] = 'CC'
    
    term_weights = go_hierarchy.get_term_weights(mlb.classes_)
    
    # Optimize thresholds
    threshold_optimizer = AdaptiveThresholdOptimizer(go_hierarchy)
    
    print("\n--- Method 1: Global Threshold ---")
    global_t, global_f1 = threshold_optimizer.optimize_global(val_preds, y_val, term_weights)
    
    print("\n--- Method 2: Per-Subontology Thresholds ---")
    subont_thresholds = threshold_optimizer.optimize_per_subontology(
        val_preds, y_val, list(mlb.classes_), term_weights
    )
    
    # Save thresholds
    threshold_optimizer.save_thresholds('optimized_thresholds.json')
    
    # ========== Step 8: Generate Predictions ==========
    print("\n[8/8] Generating final predictions...")
    
    # Test data (using training data for demo)
    test_features = combined_features[:50]
    test_tensor = torch.FloatTensor(test_features).to(device)
    
    with torch.no_grad():
        test_preds = torch.sigmoid(model(test_tensor)).cpu().numpy()
    
    # Apply adaptive thresholds
    thresholded_preds = threshold_optimizer.apply_thresholds(
        test_preds,
        list(mlb.classes_)
    )
    
    # Generate submission
    test_df = pd.DataFrame({
        'protein_id': [f'TEST_{i}' for i in range(50)],
        'sequence': train_sequences[:50]
    })
    
    submission_data = []
    for i, protein_id in enumerate(test_df['protein_id']):
        pred_probs = thresholded_preds[i]
        significant_idx = np.where(pred_probs > 0)[0]
        sorted_idx = significant_idx[np.argsort(-pred_probs[significant_idx])]
        
        for idx in sorted_idx[:500]:
            submission_data.append([
                protein_id,
                mlb.classes_[idx],
                round(float(pred_probs[idx]), 6)
            ])
    
    submission_df = pd.DataFrame(submission_data)
    submission_df.to_csv('advanced_optimized_submission.tsv', 
                        sep='\t', index=False, header=False)
    
    print(f"\n{'='*80}")
    print("PIPELINE COMPLETE!")
    print(f"{'='*80}")
    print(f"✓ Hyperparameters optimized and saved")
    print(f"✓ Model trained with best configuration")
    print(f"✓ Adaptive thresholds calculated")
    print(f"✓ Submission generated: advanced_optimized_submission.tsv")
    print(f"{'='*80}\n")
    
    return model, threshold_optimizer, submission_df


if __name__ == "__main__":
    run_complete_pipeline_with_tuning()


# ==================== USAGE EXAMPLES ====================

"""
HYPERPARAMETER TUNING EXAMPLE:
-------------------------------

tuner = HyperparameterTuner(
    FlexibleProteinPredictor,
    X_train, y_train, X_val, y_val,
    device='cuda',
    use_wandb=True,
    project_name='cafa6-optimization'
)

best_params = tuner.tune(n_trials=100, timeout=3600)


ESM MODEL COMPARISON EXAMPLE:
------------------------------

esm_extractor = MultiESMExtractor(
    model_names=[
        'esm2_t12_35M_UR50D',
        'esm2_t33_650M_UR50D',
        'esm2_t36_3B_UR50D'
    ],
    device='cuda',
    pooling_strategy='mean'
)

# Compare models
results = esm_extractor.compare_models(sequences, labels, sample_size=200)

# Use best model
embeddings = esm_extractor.extract_embeddings(
    sequences,
    batch_size=8,
    combine_method='concat'
)


STRUCTURE FEATURES EXAMPLE:
----------------------------

struct_extractor = ProteinStructureFeatureExtractor(source='alphafold')

# Extract from PDB
features = struct_extractor.extract_from_pdb('protein.pdb')

# Augment embeddings
augmented = struct_extractor.augment_with_structure(
    sequence_embeddings,
    pdb_files
)


ADAPTIVE THRESHOLDS EXAMPLE:
-----------------------------

optimizer = AdaptiveThresholdOptimizer(go_hierarchy)

# Method 1: Global
global_t, f1 = optimizer.optimize_global(val_preds, val_labels, weights)

# Method 2: Per-subontology
subont_t = optimizer.optimize_per_subontology(
    val_preds, val_labels, term_list, weights
)

# Method 3: Per-term (slow but most accurate)
term_t = optimizer.optimize_per_term(val_preds, val_labels, weights)

# Apply
thresholded = optimizer.apply_thresholds(test_preds, term_list)

# Save/load
optimizer.save_thresholds('thresholds.json')
optimizer.load_thresholds('thresholds.json')
"""