train.py

# -*- coding: utf-8 -*-
"""
Training script for SAMS-VAE RNA binding prediction.
K-fold CV running ALL folds simultaneously on single GPU.
"""

import os
import argparse
import json
from datetime import datetime
from pathlib import Path
from multiprocessing import Process, set_start_method

# Limit CPU threads BEFORE importing torch
os.environ['OMP_NUM_THREADS'] = '4'
os.environ['MKL_NUM_THREADS'] = '4'
os.environ['OPENBLAS_NUM_THREADS'] = '4'
os.environ['NUMEXPR_NUM_THREADS'] = '4'

import numpy as np
import pandas as pd
import torch
torch.set_float32_matmul_precision('high')
torch.set_num_threads(4)  # Limit PyTorch CPU threads
# Anomaly detection disabled for speed (enable for debugging)
# torch.autograd.set_detect_anomaly(True)
import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
from pytorch_lightning.loggers import WandbLogger
from torch.utils.data import DataLoader
from sklearn.model_selection import KFold

from data import RNABindingDataset, collate_fn, load_embeddings, get_collate_fn
from model.model import LitModel


class MetricsTracker(pl.Callback):
    def __init__(self):
        self.history = []
        self.best_epoch = -1
        self.best_val_score = -1
        self.best_metrics = {}
        self.final_epoch = -1
        self.final_metrics = {}
    
    def on_validation_epoch_end(self, trainer, pl_module):
        metrics = {k: float(v) for k, v in trainer.callback_metrics.items()}
        epoch = trainer.current_epoch
        self.history.append({'epoch': epoch, 'metrics': metrics})
        val_score = metrics.get('val/f1', -1)  # Fixed to f1
        if val_score > self.best_val_score:
            self.best_val_score = val_score
            self.best_epoch = epoch
            self.best_metrics = metrics.copy()
        self.final_epoch = epoch
        self.final_metrics = metrics.copy()


def train_single_fold(fold_args):
    fold_idx = fold_args['fold_idx']
    gpu_id = fold_args['gpu_id']
    args = fold_args['args']
    train_df = fold_args['train_df']
    test_df = fold_args['test_df']
    rna_emb_dict = fold_args['rna_emb_dict']
    ion_to_idx = fold_args['ion_to_idx']
    rna_dim = fold_args['rna_dim']
    num_ions = fold_args['num_ions']
    num_ligands = fold_args.get('num_ligands', 1)
    ligand_to_idx = fold_args.get('ligand_to_idx', {})
    
    if args.random_split:
        train_indices = fold_args['train_indices']
        val_indices = fold_args['val_indices']
    else:
        split_values = fold_args['split_values']
        val_split = fold_args['val_split']
    
    if torch.cuda.is_available():
        torch.cuda.set_device(gpu_id)
    
    # Fix randomness - STRICT
    seed = args.seed + fold_idx
    import random
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    pl.seed_everything(seed, workers=True)
    torch.backends.cudnn.deterministic = False
    torch.backends.cudnn.benchmark = True  # Auto-tune for speed
    # torch.use_deterministic_algorithms(True, warn_only=True)  # Disabled for speed
    os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
    os.environ['PYTHONHASHSEED'] = str(seed)
    
    if args.random_split:
        fold_train_df = train_df.iloc[train_indices].reset_index(drop=True)
        fold_val_df = train_df.iloc[val_indices].reset_index(drop=True)
    else:
        train_split_values = [v for v in split_values if v != val_split]
        fold_train_df = train_df[train_df['split'].isin(train_split_values)].reset_index(drop=True)
        fold_val_df = train_df[train_df['split'] == val_split].reset_index(drop=True)
    
    train_dataset = RNABindingDataset(
        fold_train_df, args.path_pdbs, rna_emb_dict, ion_to_idx, ligand_to_idx,
        rna_dim, all_ligands=args.all_ligands,
    )
    
    # Always add RhoFold apo structures to training
    from data import RhoFoldApoDataset
    from torch.utils.data import ConcatDataset
    apo_dataset = RhoFoldApoDataset(
        rhofold_path=args.rhofold_path, rna_emb_dict=rna_emb_dict,
        rna_dim=rna_dim, num_ions=num_ions, num_ligands=num_ligands,
    )
    if len(apo_dataset) > 0:
        train_dataset = ConcatDataset([train_dataset, apo_dataset])
        print(f"[Fold {fold_idx}] Added {len(apo_dataset)} apo structures")
    
    val_dataset = RNABindingDataset(
        fold_val_df, args.path_pdbs, rna_emb_dict, ion_to_idx, ligand_to_idx,
        rna_dim, all_ligands=args.all_ligands,
    )
    test_dataset = RNABindingDataset(
        test_df, args.path_pdbs, rna_emb_dict, ion_to_idx, ligand_to_idx,
        rna_dim, all_ligands=args.all_ligands,
    )
    
    print(f"[Fold {fold_idx}] Train: {len(train_dataset)}, Val: {len(val_dataset)}, Test: {len(test_dataset)}")
    
    # avg_ion_vector is computed lazily in model during first val epoch
    
    ligand_to_smiles = fold_args.get('ligand_to_smiles', {})
    collate = get_collate_fn(ligand_to_idx, ligand_to_smiles)
    # Generator for reproducibility
    g = torch.Generator()
    g.manual_seed(seed)
    # num_workers for parallel data loading (persistent_workers keeps them alive)
    train_loader = DataLoader(train_dataset, args.batch_size, shuffle=True, collate_fn=collate, 
                              num_workers=4, generator=g, persistent_workers=True,
                              pin_memory=True, prefetch_factor=4)
    val_loader = DataLoader(val_dataset, args.batch_size, collate_fn=collate, 
                            num_workers=2, persistent_workers=True,
                            pin_memory=True, prefetch_factor=2)
    test_loader = DataLoader(test_dataset, args.batch_size, collate_fn=collate, 
                             num_workers=2, persistent_workers=True,
                             pin_memory=True, prefetch_factor=2)
    
    model = LitModel(
        num_ions=num_ions, num_ligands=num_ligands,
        rna_dim=rna_dim,
        lr=args.lr, weight_decay=args.weight_decay,
        focal_gamma=args.focal_gamma,
        mask_prior=args.mask_prior, mask_init_prob=args.mask_init_prob,
        embedding_prior_std=args.embedding_prior_std,
        max_epochs=args.epochs, dropout=args.dropout,
        alpha=args.alpha, beta=args.beta, gamma=args.gamma, delta=args.delta, epsilon=args.epsilon,
        kl_warmup=args.kl_warmup, kl_cycle=args.kl_cycle, max_kl_weight=args.max_kl_weight,
        coord_loss_type=args.loss_3d,
        latent_recon_weight=args.latent_recon_weight,
    )
    
    ckpt_path = f'checkpoints/{args.exp_name}/fold{fold_idx}'
    ckpt_dir = Path(ckpt_path)
    ckpt_dir.mkdir(parents=True, exist_ok=True)
    
    metrics_tracker = MetricsTracker()
    monitor_key = 'val/f1'
    callbacks = [
        ModelCheckpoint(dirpath=ckpt_dir, monitor=monitor_key, mode='max', save_top_k=1,
                        filename='best-{epoch}-{' + monitor_key + ':.4f}'),
        EarlyStopping(monitor=monitor_key, mode='max', patience=args.patience),
        metrics_tracker,
    ]
    
    # Wandb logger (always on unless --no_wandb)
    logger = None
    if not args.no_wandb:
        logger = WandbLogger(
            project=args.wandb_project,
            name=f"{args.exp_name}_fold{fold_idx}",
            save_dir=f'checkpoints/{args.exp_name}',
            config=vars(args),
            group=args.exp_name,
            tags=[f"fold{fold_idx}", args.rna_emb],
        )
    
    # Single GPU
    devices = [gpu_id] if torch.cuda.is_available() else 1
    
    trainer = pl.Trainer(
        max_epochs=args.epochs, accelerator='auto',
        devices=devices,
        logger=logger, callbacks=callbacks,
        gradient_clip_val=args.grad_clip, deterministic=True,
        enable_progress_bar=True, num_sanity_val_steps=0, check_val_every_n_epoch=1,
        log_every_n_steps=10,
        precision='16-mixed',  # Best setting hardcoded
    )
    
    trainer.fit(model, train_loader, val_loader)
    best_ckpt = trainer.checkpoint_callback.best_model_path
    
    trainer.test(model, test_loader, ckpt_path=best_ckpt)
    test_metrics = model.get_test_results()
    
    # Save test outputs for ensemble calculation
    test_outputs = model.get_test_outputs_for_ensemble()
    if test_outputs['logits'] is not None:
        torch.save(test_outputs, ckpt_dir / 'test_outputs.pt')
    
    # Print results
    all_metrics = ['auroc', 'auprc', 'mcc', 'f1', 'precision', 'recall', 'specificity']
    
    print(f"\n[Fold {fold_idx}] RESULTS (Best Epoch: {metrics_tracker.best_epoch})")
    print("=" * 80)
    
    # Train/Val @ Best Epoch
    print(f"\n[@ Best Epoch]")
    print(f"{'Metric':<15} {'Train':>15} {'Val':>15}")
    print("-" * 48)
    for metric in all_metrics:
        train_val = test_metrics.get(f'train_best/{metric}', 0)
        val_val = test_metrics.get(f'val_best/{metric}', 0)
        print(f"{metric:<15} {train_val:>15.4f} {val_val:>15.4f}")
    
    # Test
    print(f"\n[Test]")
    print(f"{'Metric':<15} {'Test (Main)':>15} {'Test (De Novo)':>15}")
    print("-" * 48)
    for metric in all_metrics:
        test_main = test_metrics.get(f'test/{metric}', 0)
        test_denovo = test_metrics.get(f'test_denovo/{metric}', 0)
        print(f"{metric:<15} {test_main:>15.4f} {test_denovo:>15.4f}")
    print("=" * 80)
    
    # Log final metrics to wandb
    if not args.no_wandb and logger:
        import wandb
        final_log = {"final/best_epoch": metrics_tracker.best_epoch}
        for k, v in test_metrics.items():
            final_log[f'final/{k}'] = v
        wandb.log(final_log)
        wandb.finish()
    
    result = {
        'fold_idx': fold_idx, 'best_epoch': int(metrics_tracker.best_epoch),
        'best_val_metrics': metrics_tracker.best_metrics,
        'test_metrics': test_metrics
    }
    
    with open(ckpt_dir / 'fold_result.json', 'w') as f:
        json.dump(result, f, indent=2)
    
    return result


def run_cv(args):
    # Fix randomness - STRICT
    import random
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    pl.seed_everything(args.seed, workers=True)
    torch.backends.cudnn.deterministic = False
    torch.backends.cudnn.benchmark = True  # Auto-tune for speed
    # torch.use_deterministic_algorithms(True, warn_only=True)  # Disabled for speed
    os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
    os.environ['PYTHONHASHSEED'] = str(args.seed)
    
    # Load data
    train_df = pd.read_csv(args.train_csv)
    test_df = pd.read_csv(args.test_csv)
    
    if args.random_split:
        n_splits = 4  # Always 4-fold CV for random split
        kf = KFold(n_splits=n_splits, shuffle=True, random_state=args.seed)
        fold_splits = list(kf.split(train_df))
    else:
        train_df['split'] = train_df['split'].astype(int)
        split_values = sorted(train_df['split'].unique())
        n_splits = len(split_values)
    
    rna_emb_dict, rna_dim = load_embeddings(args.rna_pkl)
    
    # Auto-detect ions from train/test CSV files
    from data import is_cation, get_ion_to_idx, get_ligand_to_idx, get_ligand_to_smiles
    ion_to_idx = get_ion_to_idx(args.train_csv, args.test_csv)
    num_ions = len(ion_to_idx)
    
    # Auto-build ligand_to_idx for non-ion ligands
    ligand_to_idx = get_ligand_to_idx(args.train_csv, args.test_csv)
    ligand_to_smiles = get_ligand_to_smiles(args.train_csv, args.test_csv)
    num_ligands = max(len(ligand_to_idx), 1)
    
    # Save ion_to_idx, ligand_to_idx, and ligand_to_smiles for inference
    ckpt_base = Path(f'checkpoints/{args.exp_name}')
    ckpt_base.mkdir(parents=True, exist_ok=True)
    with open(ckpt_base / 'ion_to_idx.json', 'w') as f:
        json.dump(ion_to_idx, f)
    with open(ckpt_base / 'ligand_to_idx.json', 'w') as f:
        json.dump(ligand_to_idx, f)
    with open(ckpt_base / 'ligand_to_smiles.json', 'w') as f:
        json.dump(ligand_to_smiles, f)
    # Save args for aggregate_results.py
    with open(ckpt_base / 'args.json', 'w') as f:
        json.dump(vars(args), f, indent=2)
    print(f"Saved ion_to_idx ({num_ions} ions) to {ckpt_base / 'ion_to_idx.json'}")
    print(f"Saved ligand_to_idx ({num_ligands} ligands) to {ckpt_base / 'ligand_to_idx.json'}")
    
    print(f"rna_dim={rna_dim}, num_ions={num_ions}, num_ligands={num_ligands}")
    
    folds_to_run = list(range(n_splits))
    
    # GPU assignment
    # Option 1: --gpu_idx explicitly specifies internal GPU for each fold
    # Option 2: Auto-distribute across CUDA_VISIBLE_DEVICES
    if args.gpu_idx:
        # Parse --gpu_idx "0,1,2,3,0,1,2,3,2,3"
        gpu_indices = [int(x) for x in args.gpu_idx.split(',')]
        if len(gpu_indices) < n_splits:
            # Extend by cycling
            gpu_indices = [gpu_indices[i % len(gpu_indices)] for i in range(n_splits)]
        fold_to_gpu = {f: gpu_indices[f] for f in range(n_splits)}
        print(f"Using manual GPU assignment: --gpu_idx {args.gpu_idx}")
    else:
        # Auto-distribute: fold i -> GPU (i % n_gpus)
        n_gpus = 1
        if torch.cuda.is_available():
            cuda_devices = os.environ.get('CUDA_VISIBLE_DEVICES', '')
            if cuda_devices:
                n_gpus = len(cuda_devices.split(','))
            else:
                n_gpus = torch.cuda.device_count()
            n_gpus = max(1, n_gpus)
        fold_to_gpu = {f: f % n_gpus for f in range(n_splits)}
    
    print(f"GPU allocation: {fold_to_gpu}")
    
    def make_fold_args(fold_idx):
        fold_args = {
            'fold_idx': fold_idx, 'gpu_id': fold_to_gpu.get(fold_idx, 0), 'args': args,
            'train_df': train_df, 'test_df': test_df,
            'rna_emb_dict': rna_emb_dict,
            'ion_to_idx': ion_to_idx, 'rna_dim': rna_dim, 'num_ions': num_ions,
            'ligand_to_idx': ligand_to_idx, 'ligand_to_smiles': ligand_to_smiles, 'num_ligands': num_ligands,
        }
        if args.random_split:
            train_indices, val_indices = fold_splits[fold_idx]
            fold_args['train_indices'] = train_indices
            fold_args['val_indices'] = val_indices
        else:
            fold_args['split_values'] = split_values
            fold_args['val_split'] = split_values[fold_idx]
        return fold_args
    
    # Run folds in rounds based on gpu_idx count
    if args.gpu_idx:
        folds_per_round = len(args.gpu_idx.split(','))
    else:
        folds_per_round = len(set(fold_to_gpu.values()))  # number of unique GPUs
    
    print(f"Running {folds_per_round} folds per round")
    
    for round_start in range(0, len(folds_to_run), folds_per_round):
        round_folds = folds_to_run[round_start:round_start + folds_per_round]
        print(f"\n=== Round: folds {round_folds} ===")
        
        processes = []
        for fold_idx in round_folds:
            gpu_id = fold_to_gpu[fold_idx]
            print(f"Starting fold {fold_idx} on GPU {gpu_id}...")
            p = Process(target=train_single_fold, args=(make_fold_args(fold_idx),))
            p.start()
            processes.append((fold_idx, gpu_id, p))
        
        # Wait for round to complete
        for fold_idx, gpu_id, p in processes:
            p.join()
            print(f"Fold {fold_idx} on GPU {gpu_id} completed.")
    
    # Aggregate results
    fold_results = []
    for fold_idx in range(n_splits):
        path = Path(f'checkpoints/{args.exp_name}/fold{fold_idx}/fold_result.json')
        if path.exists():
            with open(path) as f:
                fold_results.append(json.load(f))
    
    # Compute ensemble metrics (average logits across folds)
    ensemble_metrics = {}
    ensemble_metrics_denovo = {}
    ensemble_metrics_rna_only = {}
    ensemble_metrics_rna_avg_ion = {}
    try:
        all_logits = []
        all_logits_denovo = []
        all_logits_rna_only = []
        all_logits_rna_avg_ion = []
        all_labels = None
        
        for fold_idx in range(n_splits):
            output_path = Path(f'checkpoints/{args.exp_name}/fold{fold_idx}/test_outputs.pt')
            if output_path.exists():
                outputs = torch.load(output_path, map_location='cpu')
                all_logits.append(outputs['logits'])
                all_logits_denovo.append(outputs['logits_denovo'])
                if outputs.get('logits_rna_only') is not None:
                    all_logits_rna_only.append(outputs['logits_rna_only'])
                if outputs.get('logits_rna_avg_ion') is not None:
                    all_logits_rna_avg_ion.append(outputs['logits_rna_avg_ion'])
                if all_labels is None:
                    all_labels = outputs['labels']
        
        if len(all_logits) == n_splits and all_labels is not None:
            # Average logits across folds
            ensemble_logits = torch.stack(all_logits).mean(dim=0)
            ensemble_logits_denovo = torch.stack(all_logits_denovo).mean(dim=0)
            
            from model.loss import compute_metrics
            ensemble_metrics = compute_metrics(ensemble_logits, all_labels, threshold=0.5)
            ensemble_metrics_denovo = compute_metrics(ensemble_logits_denovo, all_labels, threshold=0.5)
            
            if len(all_logits_rna_only) == n_splits:
                ensemble_logits_rna_only = torch.stack(all_logits_rna_only).mean(dim=0)
                ensemble_metrics_rna_only = compute_metrics(ensemble_logits_rna_only, all_labels, threshold=0.5)
            
            if len(all_logits_rna_avg_ion) == n_splits:
                ensemble_logits_rna_avg_ion = torch.stack(all_logits_rna_avg_ion).mean(dim=0)
                ensemble_metrics_rna_avg_ion = compute_metrics(ensemble_logits_rna_avg_ion, all_labels, threshold=0.5)
            
            print(f"\n[Ensemble] Computed from {n_splits} folds")
    except Exception as e:
        print(f"[Warning] Could not compute ensemble metrics: {e}")
    
    base_metrics = ['auroc', 'auprc', 'mcc', 'f1', 'precision', 'recall', 'specificity']
    
    # Prepare CSV data
    csv_rows = []
    
    print(f"\n{'='*80}")
    print(f"FINAL CV RESULTS ({n_splits} folds)")
    print(f"{'='*110}")
    
    # Line 1: Train / Val / Val (RNA only) / Val (RNA+AvgIon)
    print(f"\n[Train / Val]")
    print(f"{'Metric':<15} {'Train':>22} {'Val':>22} {'Val (RNA only)':>22} {'Val (RNA+AvgIon)':>22}")
    print("-" * 106)
    
    for metric in base_metrics:
        train_vals = [r['test_metrics'].get(f'train_best/{metric}', 0) for r in fold_results if r]
        val_vals = [r['test_metrics'].get(f'val_best/{metric}', 0) for r in fold_results if r]
        val_rna_only_vals = [r['test_metrics'].get(f'val_rna_only/{metric}', 0) for r in fold_results if r]
        val_rna_avg_vals = [r['test_metrics'].get(f'val_rna_avg_ion/{metric}', 0) for r in fold_results if r]
        
        if train_vals and any(v > 0 for v in train_vals):
            train_str = f"{np.mean(train_vals):.4f} +/- {np.std(train_vals):.4f}"
        else:
            train_str = "N/A"
        
        if val_vals and any(v > 0 for v in val_vals):
            val_str = f"{np.mean(val_vals):.4f} +/- {np.std(val_vals):.4f}"
        else:
            val_str = "N/A"
        
        if val_rna_only_vals and any(v > 0 for v in val_rna_only_vals):
            val_rna_str = f"{np.mean(val_rna_only_vals):.4f} +/- {np.std(val_rna_only_vals):.4f}"
        else:
            val_rna_str = "N/A"
        
        if val_rna_avg_vals and any(v > 0 for v in val_rna_avg_vals):
            val_rna_avg_str = f"{np.mean(val_rna_avg_vals):.4f} +/- {np.std(val_rna_avg_vals):.4f}"
        else:
            val_rna_avg_str = "N/A"
        
        print(f"{metric:<15} {train_str:>22} {val_str:>22} {val_rna_str:>22} {val_rna_avg_str:>22}")
    
    # Line 2: Test (Main) / Test (De Novo) / Test (RNA only) / Test (RNA+AvgIon)
    print(f"\n[Test]")
    print(f"{'Metric':<15} {'Test (All info)':>22} {'Test (without 3D)':>22} {'Test (RNA only)':>22} {'Test (RNA+AvgIon)':>22}")
    print("-" * 106)
    
    for metric in base_metrics:
        main_vals = [r['test_metrics'].get(f'test/{metric}', 0) for r in fold_results if r]
        denovo_vals = [r['test_metrics'].get(f'test_denovo/{metric}', 0) for r in fold_results if r]
        test_rna_only = [r['test_metrics'].get(f'test_rna_only/{metric}', 0) for r in fold_results if r]
        test_rna_avg_ion = [r['test_metrics'].get(f'test_rna_avg_ion/{metric}', 0) for r in fold_results if r]
        
        if main_vals:
            main_str = f"{np.mean(main_vals):.4f} +/- {np.std(main_vals):.4f}"
        else:
            main_str = "N/A"
        
        if denovo_vals and any(v > 0 for v in denovo_vals):
            denovo_str = f"{np.mean(denovo_vals):.4f} +/- {np.std(denovo_vals):.4f}"
        else:
            denovo_str = "N/A"
        
        if test_rna_only and any(v > 0 for v in test_rna_only):
            test_rna_str = f"{np.mean(test_rna_only):.4f} +/- {np.std(test_rna_only):.4f}"
        else:
            test_rna_str = "N/A"
        
        if test_rna_avg_ion and any(v > 0 for v in test_rna_avg_ion):
            test_rna_avg_str = f"{np.mean(test_rna_avg_ion):.4f} +/- {np.std(test_rna_avg_ion):.4f}"
        else:
            test_rna_avg_str = "N/A"
        
        print(f"{metric:<15} {main_str:>22} {denovo_str:>22} {test_rna_str:>22} {test_rna_avg_str:>22}")
    
    # Line 3: Ensemble (Main) / Ensemble (De Novo) / Ensemble (RNA only) / Ensemble (RNA+AvgIon)
    if ensemble_metrics:
        print(f"\n[Ensemble ({n_splits} models)]")
        print(f"{'Metric':<15} {'Ens (All info)':>22} {'Ens (without 3D)':>22} {'Ens (RNA only)':>22} {'Ens (RNA+AvgIon)':>22}")
        print("-" * 106)
        
        for metric in base_metrics:
            ens_main = ensemble_metrics.get(metric, 0)
            ens_denovo = ensemble_metrics_denovo.get(metric, 0)
            ens_rna_only = ensemble_metrics_rna_only.get(metric, 0) if ensemble_metrics_rna_only else 0
            ens_rna_avg = ensemble_metrics_rna_avg_ion.get(metric, 0) if ensemble_metrics_rna_avg_ion else 0
            print(f"{metric:<15} {ens_main:>22.4f} {ens_denovo:>22.4f} {ens_rna_only:>22.4f} {ens_rna_avg:>22.4f}")
    
    # Coord RMSE
    print(f"\n[Coord RMSE (original space)]")
    print(f"{'Split':<15} {'RMSE (main)':>22} {'RMSE (zb_only)':>22}")
    print("-" * 60)
    
    for split in ['train_best', 'val_best', 'test']:
        # Main path RMSE (with skip connection)
        rmse_vals = [r['test_metrics'].get(f'{split}/coord_rmse', 0) for r in fold_results if r]
        if rmse_vals and any(v > 0 for v in rmse_vals):
            rmse_str = f"{np.mean(rmse_vals):.4f} +/- {np.std(rmse_vals):.4f}"
        else:
            rmse_str = "N/A"
        
        # z_b only path RMSE (for de novo capability)
        rmse_zb_vals = [r['test_metrics'].get(f'{split}/coord_rmse_zb_only', 0) for r in fold_results if r]
        if rmse_zb_vals and any(v > 0 for v in rmse_zb_vals):
            rmse_zb_str = f"{np.mean(rmse_zb_vals):.4f} +/- {np.std(rmse_zb_vals):.4f}"
        else:
            rmse_zb_str = "N/A"
        
        split_name = split.replace('_best', '').capitalize()
        print(f"{split_name:<15} {rmse_str:>22} {rmse_zb_str:>22}")
    
    # Additional Coord Metrics (Test only)
    coord_metric_names = ['chamfer_l1', 'chamfer_l2', 'tm_score', 'lddt', 'gdt_ts']
    
    print(f"\n[Additional Coord Metrics (Test)]")
    print(f"{'Metric':<15} {'Main':>22} {'zb_only':>22}")
    print("-" * 60)
    
    for metric_name in coord_metric_names:
        # Main
        main_vals = [r['test_metrics'].get(f'test/coord_{metric_name}', 0) for r in fold_results if r]
        if main_vals and any(v > 0 for v in main_vals):
            main_str = f"{np.mean(main_vals):.4f} +/- {np.std(main_vals):.4f}"
        else:
            main_str = "N/A"
        
        # zb_only
        zb_vals = [r['test_metrics'].get(f'test/coord_{metric_name}_zb', 0) for r in fold_results if r]
        if zb_vals and any(v > 0 for v in zb_vals):
            zb_str = f"{np.mean(zb_vals):.4f} +/- {np.std(zb_vals):.4f}"
        else:
            zb_str = "N/A"
        
        print(f"{metric_name:<15} {main_str:>22} {zb_str:>22}")
    
    print(f"{'='*80}")
    
    # Save results to CSV
    csv_data = []
    for metric in base_metrics:
        row = {'metric': metric, 'split': 'test'}
        for key in ['test', 'test_denovo', 'test_rna_only', 'test_rna_avg_ion']:
            vals = [r['test_metrics'].get(f'{key}/{metric}', 0) for r in fold_results if r]
            if vals and any(v > 0 for v in vals):
                row[f'{key}_mean'] = np.mean(vals)
                row[f'{key}_std'] = np.std(vals)
        csv_data.append(row)
    
    if ensemble_metrics:
        for metric in base_metrics:
            row = {'metric': metric, 'split': 'ensemble'}
            row['ensemble_mean'] = ensemble_metrics.get(metric, 0)
            row['ensemble_denovo_mean'] = ensemble_metrics_denovo.get(metric, 0)
            csv_data.append(row)
    
    csv_path = Path(f'checkpoints/{args.exp_name}/results.csv')
    pd.DataFrame(csv_data).to_csv(csv_path, index=False)
    print(f"Results saved to {csv_path}")
    
    return fold_results


def main():
    parser = argparse.ArgumentParser()
    # Data paths
    parser.add_argument('--train_csv', type=str, default='data/train_w_smiles.csv')
    parser.add_argument('--test_csv', type=str, default='data/test_w_smiles.csv')
    parser.add_argument('--path_pdbs', type=str, default='data/pdb_chain_files')
    parser.add_argument('--rhofold_path', type=str, default='rhofold_pdbs',
                        help='Path to RhoFold apo structures (always used)')
    parser.add_argument('--rna_emb', type=str, default='aido', choices=['aido', 'ernierna', 'rinalmo', 'rnamsm'])
    parser.add_argument('--exp_name', type=str, default='exp')
    parser.add_argument('--seed', type=int, default=42)
    
    # Model architecture (Official PTv3 config from GitHub)
    parser.add_argument('--dropout', type=float, default=0.3, help='PTv3 drop_path (official: 0.3)')
    
    # SAMS-VAE specific
    parser.add_argument('--mask_prior', type=float, default=0.001, help='Bernoulli prior for masks')
    parser.add_argument('--mask_init_prob', type=float, default=0.1, help='Initial mask probability')
    parser.add_argument('--embedding_prior_std', type=float, default=1.0, help='Prior std for embeddings')
    
    # Loss settings
    parser.add_argument('--loss_3d', type=str, default='chamfer_l2',
                        choices=['mse', 'chamfer_l1', 'chamfer_l2', 'chamfer_sq'],
                        help='3D coordinate loss: mse, chamfer_l1, chamfer_l2, chamfer_sq')
    
    # Training settings
    parser.add_argument('--only_bind_ligands', action='store_true')
    parser.add_argument('--epochs', type=int, default=300)
    parser.add_argument('--batch_size', type=int, default=32)
    parser.add_argument('--lr', type=float, default=5e-5)  # Lower lr for stability
    parser.add_argument('--weight_decay', type=float, default=0.02)
    parser.add_argument('--grad_clip', type=float, default=0.1)  # Very strict clipping
    parser.add_argument('--patience', type=int, default=25)
    parser.add_argument('--num_workers', type=int, default=4, help='DataLoader workers')
    
    # Memory optimization
    # Loss weights
    parser.add_argument('--alpha', type=float, default=1.0, help='3D coordinate loss weight')
    parser.add_argument('--beta', type=float, default=0.001, help='Basal KL weight')
    parser.add_argument('--gamma', type=float, default=0.001, help='Embedding KL weight')
    parser.add_argument('--delta', type=float, default=0.001, help='Mask KL weight')
    parser.add_argument('--epsilon', type=float, default=0.001, help='Molecular KL weight')
    parser.add_argument('--latent_recon_weight', type=float, default=0.1, help='Latent reconstruction weight')
    
    # KL annealing
    parser.add_argument('--kl_warmup', type=int, default=5000, help='KL warmup steps')
    parser.add_argument('--kl_cycle', type=int, default=10000, help='KL cycle length')
    parser.add_argument('--max_kl_weight', type=float, default=0.1, help='Maximum KL weight')
    
    # Other
    parser.add_argument('--focal_gamma', type=float, default=1.5, help='Focal loss gamma')
    parser.add_argument('--random_split', action='store_true', default=False,
                        help='Use random 4-fold CV instead of predefined splits')
    parser.add_argument('--gpu_idx', type=str, default=None,
                        help='Comma-separated GPU indices for each fold. E.g., "0,1,2,3" for 4-fold on 4 GPUs')
    
    # Wandb logging (always on by default)
    parser.add_argument('--no_wandb', action='store_true', help='Disable wandb logging')
    parser.add_argument('--wandb_project', type=str, default='rbs_final', help='Wandb project name')
    
    # CPU thread control
    parser.add_argument('--num_threads', type=int, default=4, help='Number of CPU threads per process')
    
    args = parser.parse_args()
    args.all_ligands = not args.only_bind_ligands
    
    # Apply thread limit
    os.environ['OMP_NUM_THREADS'] = str(args.num_threads)
    os.environ['MKL_NUM_THREADS'] = str(args.num_threads)
    os.environ['OPENBLAS_NUM_THREADS'] = str(args.num_threads)
    os.environ['NUMEXPR_NUM_THREADS'] = str(args.num_threads)
    torch.set_num_threads(args.num_threads)
    
    if args.rna_emb == 'aido':
        args.rna_pkl = 'data/aido_rna_emb.pkl'
    elif args.rna_emb == 'ernierna':
        args.rna_pkl = 'data/ernierna_rna_emb.pkl'
    elif args.rna_emb == 'rinalmo':
        args.rna_pkl = 'data/rinalmo_rna_emb.pkl'
    else:
        args.rna_pkl = 'data/rna_msm_emb.pkl'
    
    if args.exp_name == 'exp':
        args.exp_name = f"{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    
    print("=" * 80)
    print(f"RNA Binding Prediction - {args.exp_name}")
    print("=" * 80)
    print(f"Loss weights: alpha={args.alpha}, beta={args.beta}, gamma={args.gamma}, delta={args.delta}")
    print(f"3D Loss Type: {args.loss_3d}, Precision: 16-mixed (hardcoded)")
    print(f"Ligand Encoder: AttentiveFP (end-to-end from SMILES)")
    print(f"Wandb: {'disabled' if args.no_wandb else 'enabled (' + args.wandb_project + ')'}")
    print("=" * 80)
    
    run_cv(args)


if __name__ == '__main__':
    try:
        set_start_method('spawn')
    except RuntimeError:
        pass
    main()