train_pretrained_finetune.py

"""
Fine-tune Pretrained Model on BBBP Dataset

This script takes the pretrained encoder (from ZINC 250k) and fine-tunes
it on the BBBP dataset for BBB permeability prediction.

Comparison:
1. Baseline: Standard model trained from scratch
2. Pretrained: This model - pretrained on ZINC, then fine-tuned on BBBP
3. Quantum: Model with quantum descriptors (separate comparison)
"""

import torch
import torch.nn as nn
import torch.optim as optim
from torch_geometric.loader import DataLoader
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score, accuracy_score, precision_score, recall_score, f1_score
import os
import time
from tqdm import tqdm

from advanced_bbb_model import AdvancedHybridBBBNet
from mol_to_graph import mol_to_graph, batch_smiles_to_graphs
from pretrain_zinc import transfer_weights_to_bbb_model


def load_bbbp_data():
    """Load BBBP dataset and convert to graphs"""

    print("Loading BBBP dataset...")

    data_path = "data/BBBP.csv"
    df = pd.read_csv(data_path)

    # Get columns
    smiles_col = 'smiles' if 'smiles' in df.columns else 'SMILES'
    target_col = 'p_np' if 'p_np' in df.columns else 'BBB_permeability'

    smiles_list = df[smiles_col].tolist()
    y_list = df[target_col].tolist()

    print(f"Total samples: {len(smiles_list)}")

    # Convert to graphs
    graphs = batch_smiles_to_graphs(smiles_list, y_list)

    print(f"Valid graphs: {len(graphs)}")
    print(f"Features per node: {graphs[0].x.shape[1]}")

    return graphs


def train_model(model, train_loader, val_loader, epochs=100, lr=0.0001,
                patience=30, device='cpu', class_weight=3.24, model_name='model'):
    """
    Train the model with early stopping
    """
    model = model.to(device)

    # Loss with class weights
    pos_weight = torch.tensor([class_weight]).to(device)
    criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)

    optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=1e-4)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.5,
                                                      patience=10, min_lr=1e-6)

    best_auc = 0
    best_epoch = 0
    no_improve = 0

    history = {
        'train_loss': [], 'val_loss': [],
        'train_auc': [], 'val_auc': [],
        'val_accuracy': []
    }

    print(f"\nTraining {model_name}...")
    print("=" * 60)

    for epoch in range(1, epochs + 1):
        # Training
        model.train()
        train_loss = 0
        train_preds = []
        train_labels = []

        for batch in train_loader:
            batch = batch.to(device)
            optimizer.zero_grad()

            out = model(batch.x, batch.edge_index, batch.batch)
            out_flat = out.view(-1)
            y_flat = batch.y.view(-1)

            loss = criterion(out_flat, y_flat)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()

            train_loss += loss.item()
            train_preds.extend(torch.sigmoid(out_flat).detach().cpu().numpy())
            train_labels.extend(y_flat.cpu().numpy())

        train_loss /= len(train_loader)
        train_auc = roc_auc_score(train_labels, train_preds)

        # Validation
        model.eval()
        val_loss = 0
        val_preds = []
        val_labels = []

        with torch.no_grad():
            for batch in val_loader:
                batch = batch.to(device)
                out = model(batch.x, batch.edge_index, batch.batch)
                out_flat = out.view(-1)
                y_flat = batch.y.view(-1)

                loss = criterion(out_flat, y_flat)
                val_loss += loss.item()
                val_preds.extend(torch.sigmoid(out_flat).cpu().numpy())
                val_labels.extend(y_flat.cpu().numpy())

        val_loss /= len(val_loader)
        val_auc = roc_auc_score(val_labels, val_preds)
        val_acc = accuracy_score(val_labels, [1 if p > 0.5 else 0 for p in val_preds])

        # Update scheduler
        scheduler.step(val_auc)
        current_lr = optimizer.param_groups[0]['lr']

        # Save history
        history['train_loss'].append(train_loss)
        history['val_loss'].append(val_loss)
        history['train_auc'].append(train_auc)
        history['val_auc'].append(val_auc)
        history['val_accuracy'].append(val_acc)

        # Check for improvement
        if val_auc > best_auc:
            best_auc = val_auc
            best_epoch = epoch
            no_improve = 0

            # Save best model
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'val_auc': val_auc,
                'val_accuracy': val_acc,
            }, f'models/best_{model_name}.pth')

            print(f"Epoch {epoch}/{epochs} | Train AUC: {train_auc:.4f} | Val AUC: {val_auc:.4f} | "
                  f"Val Acc: {val_acc*100:.1f}% | LR: {current_lr:.6f} | *BEST*")
        else:
            no_improve += 1
            if epoch % 10 == 0:
                print(f"Epoch {epoch}/{epochs} | Train AUC: {train_auc:.4f} | Val AUC: {val_auc:.4f} | "
                      f"Val Acc: {val_acc*100:.1f}% | LR: {current_lr:.6f}")

        # Early stopping
        if no_improve >= patience:
            print(f"\nEarly stopping at epoch {epoch}. Best AUC: {best_auc:.4f} at epoch {best_epoch}")
            break

    return model, history, best_auc


def evaluate_model(model, test_loader, device='cpu'):
    """Evaluate model on test set"""
    model.eval()
    all_preds = []
    all_labels = []

    with torch.no_grad():
        for batch in test_loader:
            batch = batch.to(device)
            out = model(batch.x, batch.edge_index, batch.batch)
            preds = torch.sigmoid(out).view(-1).cpu().numpy()
            labels = batch.y.view(-1).cpu().numpy()
            all_preds.extend(preds)
            all_labels.extend(labels)

    # Calculate metrics
    binary_preds = [1 if p > 0.5 else 0 for p in all_preds]

    metrics = {
        'auc': roc_auc_score(all_labels, all_preds),
        'accuracy': accuracy_score(all_labels, binary_preds),
        'precision': precision_score(all_labels, binary_preds),
        'recall': recall_score(all_labels, binary_preds),
        'f1': f1_score(all_labels, binary_preds),
    }

    return metrics


def main():
    """Main fine-tuning function"""

    os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
    os.makedirs('models', exist_ok=True)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")

    # Training parameters
    epochs = 150
    lr = 0.00005  # Lower LR for fine-tuning
    patience = 40
    batch_size = 32
    class_weight = 3.24

    results = {}

    # Load BBBP data
    print("\n" + "=" * 60)
    print("LOADING BBBP DATA")
    print("=" * 60)

    graphs = load_bbbp_data()

    # Split data (same split for fair comparison)
    train_graphs, temp_graphs = train_test_split(graphs, test_size=0.2, random_state=42)
    val_graphs, test_graphs = train_test_split(temp_graphs, test_size=0.5, random_state=42)

    print(f"Train: {len(train_graphs)}, Val: {len(val_graphs)}, Test: {len(test_graphs)}")

    # Create loaders
    train_loader = DataLoader(train_graphs, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_graphs, batch_size=batch_size)
    test_loader = DataLoader(test_graphs, batch_size=batch_size)

    # ========================================
    # Model 1: Baseline (no pretraining)
    # ========================================
    print("\n" + "=" * 60)
    print("TRAINING BASELINE MODEL (no pretraining)")
    print("=" * 60)

    model_baseline = AdvancedHybridBBBNet(
        num_node_features=15,
        hidden_channels=128,
        num_heads=8,
        dropout=0.3
    )

    model_baseline, history_baseline, best_auc_baseline = train_model(
        model_baseline, train_loader, val_loader,
        epochs=epochs, lr=0.0001, patience=patience,
        device=device, class_weight=class_weight,
        model_name='baseline_model_comparison'
    )

    # Evaluate baseline
    metrics_baseline = evaluate_model(model_baseline, test_loader, device)
    results['baseline'] = {
        'best_val_auc': best_auc_baseline,
        'test_metrics': metrics_baseline
    }

    print(f"\nBaseline Model Test Results:")
    print(f"  AUC: {metrics_baseline['auc']:.4f}")
    print(f"  Accuracy: {metrics_baseline['accuracy']*100:.2f}%")
    print(f"  F1 Score: {metrics_baseline['f1']:.4f}")

    # ========================================
    # Model 2: Pretrained + Fine-tuned
    # ========================================
    print("\n" + "=" * 60)
    print("TRAINING PRETRAINED MODEL (ZINC 250k → BBBP)")
    print("=" * 60)

    # Check if pretrained weights exist
    pretrained_path = 'models/pretrained_encoder.pth'
    if not os.path.exists(pretrained_path):
        print(f"ERROR: Pretrained weights not found at {pretrained_path}")
        print("Please run pretrain_zinc.py first!")
        return results

    # Create model and load pretrained weights
    model_pretrained = AdvancedHybridBBBNet(
        num_node_features=15,
        hidden_channels=128,
        num_heads=8,
        dropout=0.3
    )

    # Transfer pretrained weights
    print("\nTransferring pretrained weights...")
    checkpoint = torch.load(pretrained_path, map_location='cpu', weights_only=False)
    pretrained_dict = checkpoint['model_state_dict']
    model_dict = model_pretrained.state_dict()

    # Only transfer matching layers
    transferred = []
    for name, param in pretrained_dict.items():
        if name in model_dict and model_dict[name].shape == param.shape:
            model_dict[name] = param
            transferred.append(name)

    model_pretrained.load_state_dict(model_dict)
    print(f"Transferred {len(transferred)} layers from pretrained model")

    # Fine-tune with lower learning rate
    model_pretrained, history_pretrained, best_auc_pretrained = train_model(
        model_pretrained, train_loader, val_loader,
        epochs=epochs, lr=lr, patience=patience,  # Lower LR for fine-tuning
        device=device, class_weight=class_weight,
        model_name='pretrained_finetuned_model'
    )

    # Evaluate pretrained
    metrics_pretrained = evaluate_model(model_pretrained, test_loader, device)
    results['pretrained'] = {
        'best_val_auc': best_auc_pretrained,
        'test_metrics': metrics_pretrained
    }

    print(f"\nPretrained Model Test Results:")
    print(f"  AUC: {metrics_pretrained['auc']:.4f}")
    print(f"  Accuracy: {metrics_pretrained['accuracy']*100:.2f}%")
    print(f"  F1 Score: {metrics_pretrained['f1']:.4f}")

    # ========================================
    # Summary
    # ========================================
    print("\n" + "=" * 60)
    print("COMPARISON SUMMARY")
    print("=" * 60)

    print("\nBaseline Model (no pretraining):")
    print(f"  Best Val AUC: {results['baseline']['best_val_auc']:.4f}")
    print(f"  Test AUC: {results['baseline']['test_metrics']['auc']:.4f}")
    print(f"  Test Accuracy: {results['baseline']['test_metrics']['accuracy']*100:.2f}%")
    print(f"  Test F1: {results['baseline']['test_metrics']['f1']:.4f}")

    print("\nPretrained Model (ZINC 250k pretraining):")
    print(f"  Best Val AUC: {results['pretrained']['best_val_auc']:.4f}")
    print(f"  Test AUC: {results['pretrained']['test_metrics']['auc']:.4f}")
    print(f"  Test Accuracy: {results['pretrained']['test_metrics']['accuracy']*100:.2f}%")
    print(f"  Test F1: {results['pretrained']['test_metrics']['f1']:.4f}")

    # Calculate improvement
    baseline_auc = results['baseline']['test_metrics']['auc']
    pretrained_auc = results['pretrained']['test_metrics']['auc']
    improvement = (pretrained_auc - baseline_auc) / baseline_auc * 100

    print(f"\n{'='*60}")
    print(f"IMPROVEMENT FROM PRETRAINING: {improvement:+.2f}%")
    print(f"{'='*60}")

    # Save results
    np.save('models/pretrained_comparison_results.npy', results)
    print("\nResults saved to models/pretrained_comparison_results.npy")

    return results


if __name__ == "__main__":
    main()