src/training/trainer.py

#!/usr/bin/env python3
"""
DCA-Net Trainer with AMP, DataParallel, and early stopping on AUC-ROC.
"""

import torch
import torch.nn as nn
import torch.optim as optim
from torch.amp import GradScaler, autocast
from pathlib import Path
import time
import logging
import numpy as np
from tqdm import tqdm
from sklearn.metrics import roc_auc_score, average_precision_score, precision_recall_fscore_support, confusion_matrix

from src.training.losses import DCANetLoss


class Trainer:
    """Training loop for DCA-Net.
    
    Features:
      - Mixed precision (AMP) for training only
      - DataParallel (multi-GPU)
      - Gradient clipping
      - CosineAnnealingWarmRestarts scheduler
      - Checkpoint save/load
      - Early stopping on AUC-ROC
    """

    def __init__(self, model, config, train_loader, val_loader, logger=None):
        self.config = config
        self.logger = logger or logging.getLogger('dca-net')
        self.train_cfg = config.get('training', {})
        self.log_cfg = config.get('logging', {})
        
        # Device setup
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # Multi-GPU
        if (self.train_cfg.get('use_data_parallel', False) and 
                torch.cuda.device_count() > 1):
            device_ids = self.train_cfg.get('device_ids', [0, 1])
            available = list(range(torch.cuda.device_count()))
            device_ids = [d for d in device_ids if d in available]
            self.logger.info(f"Using DataParallel with GPUs: {device_ids}")
            model = nn.DataParallel(model, device_ids=device_ids)
        
        self.model = model.to(self.device)
        self.train_loader = train_loader
        self.val_loader = val_loader

        # Loss
        loss_weights = self.train_cfg.get('loss_weights', {})
        pos_weight = float(self.config.get('data', {}).get('positive_negative_ratio', 15.0))
        
        self.criterion = DCANetLoss(
            bce_weight=loss_weights.get('bce', 0.4),
            focal_weight=loss_weights.get('focal', 0.4),
            uncertainty_weight=loss_weights.get('uncertainty', 0.2),
            focal_gamma=self.train_cfg.get('focal_gamma', 2.0),
            focal_alpha=self.train_cfg.get('focal_alpha', 0.9375),  # 15/16 for 1:15 ratio
            label_smoothing=self.train_cfg.get('label_smoothing', 0.1),
            pos_weight=pos_weight
        )

        # Optimizer
        lr = self.train_cfg.get('learning_rate', 5e-5)
        wd = self.train_cfg.get('weight_decay', 1e-5)
        self.optimizer = optim.AdamW(
            self.model.parameters(), lr=lr, weight_decay=wd
        )

        # Warmup + Scheduler
        self.warmup_epochs = self.train_cfg.get('warmup_epochs', 5)
        T0 = self.train_cfg.get('scheduler_T0', 15)
        Tmult = self.train_cfg.get('scheduler_Tmult', 2)
        self.scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(
            self.optimizer, T_0=T0, T_mult=Tmult
        )
        self.base_lr = lr

        # Mixed precision — ONLY used during training, NOT validation
        self.use_amp = self.train_cfg.get('use_amp', True) and torch.cuda.is_available()
        self.scaler = GradScaler('cuda', enabled=self.use_amp)
        self.grad_clip = self.train_cfg.get('gradient_clip', 0.5)
        self.accum_steps = self.train_cfg.get('gradient_accumulation_steps', 1)

        # Early stopping — based on AUC-ROC (higher = better)
        self.patience = self.train_cfg.get('early_stopping_patience', 15)
        self.best_val_auc = 0.0
        self.best_val_loss = float('inf')
        self.epochs_no_improve = 0

        # Checkpoint
        self.ckpt_dir = Path(self.log_cfg.get('checkpoint_dir', 'results/checkpoints'))
        self.ckpt_dir.mkdir(parents=True, exist_ok=True)
        self.save_interval = self.log_cfg.get('save_interval', 5)
        self.log_interval = self.log_cfg.get('log_interval', 10)

        # TensorBoard
        self.writer = None
        if self.log_cfg.get('use_tensorboard', False):
            try:
                from torch.utils.tensorboard import SummaryWriter
                log_dir = Path(self.log_cfg.get('log_dir', 'logs'))
                log_dir.mkdir(parents=True, exist_ok=True)
                self.writer = SummaryWriter(log_dir=str(log_dir))
            except ImportError:
                self.logger.warning("TensorBoard not available, skipping")

        self.global_step = 0
        self.start_epoch = 0

        # Curriculum learning state
        self._curriculum = self.train_cfg.get('curriculum', None)
        self._current_curriculum_stage = 1 if self._curriculum else None

    def train_epoch(self, epoch):
        """Train for one epoch with gradient accumulation and warmup."""
        self.model.train()
        total_loss = 0.0
        num_batches = 0
        nan_batches = 0

        # Linear warmup: scale LR from 10% to 100% over warmup_epochs
        if epoch < self.warmup_epochs:
            warmup_factor = 0.1 + 0.9 * (epoch / self.warmup_epochs)
            for pg in self.optimizer.param_groups:
                pg['lr'] = self.base_lr * warmup_factor

        # Custom modern progress bar
        bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]'
        pbar = tqdm(
            self.train_loader, 
            desc=f"Epoch {epoch+1}", 
            leave=False,
            bar_format=bar_format,
            ascii=" █",             # Solid blocks instead of hashes
            colour='white'          # Clean white bar
        )
        self.optimizer.zero_grad()

        for batch_idx, (nodule, context, labels) in enumerate(pbar):
            nodule = nodule.to(self.device)
            context = context.to(self.device)
            labels = labels.to(self.device)

            with autocast('cuda', enabled=self.use_amp):
                logits = self.model(nodule, context)
                loss, loss_dict = self.criterion(logits, labels)
                loss = loss / self.accum_steps  # Scale loss for accumulation

            if torch.isnan(loss) or torch.isinf(loss):
                nan_batches += 1
                if nan_batches <= 3:  # Only log first few
                    self.logger.warning(f"NaN/Inf loss at train batch {batch_idx}, skipping")
                self.optimizer.zero_grad()  # Clear any stale gradients
                continue

            self.scaler.scale(loss).backward()

            # Step optimizer every accum_steps batches (or at end of epoch)
            if (batch_idx + 1) % self.accum_steps == 0 or (batch_idx + 1) == len(self.train_loader):
                if self.grad_clip > 0:
                    self.scaler.unscale_(self.optimizer)
                    nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip)

                self.scaler.step(self.optimizer)
                self.scaler.update()
                self.optimizer.zero_grad()

            total_loss += loss.item() * self.accum_steps  # Unscale for logging
            num_batches += 1
            self.global_step += 1

            pbar.set_postfix({
                'loss': f"{loss.item() * self.accum_steps:.4f}",
                'lr': f"{self.optimizer.param_groups[0]['lr']:.2e}"
            })

            # Logging
            if batch_idx % self.log_interval == 0:
                self.logger.info(
                    f"Epoch {epoch+1} | Batch {batch_idx}/{len(self.train_loader)} | "
                    f"Loss: {loss.item() * self.accum_steps:.4f} | BCE: {loss_dict['bce']:.4f} | "
                    f"Focal: {loss_dict['focal']:.4f} | Unc: {loss_dict['uncertainty']:.4f}"
                )
                if self.writer:
                    for k, v in loss_dict.items():
                        self.writer.add_scalar(f'train/{k}', v, self.global_step)

        if epoch >= self.warmup_epochs:
            self.scheduler.step()
        avg_loss = total_loss / max(num_batches, 1)

        if nan_batches > 0:
            self.logger.warning(
                f"Epoch {epoch+1}: {nan_batches}/{len(self.train_loader)} batches had NaN loss"
            )

        return avg_loss, nan_batches

    def _compute_metrics(self, all_preds, all_labels):
        """Compute metrics for imbalanced binary classification."""
        preds_arr = np.array(all_preds)
        labels_arr = np.array(all_labels)
        pred_binary = (preds_arr > 0.5).astype(int)

        metrics = {}

        # AUC-ROC
        try:
            if len(np.unique(labels_arr)) > 1:
                metrics['auc_roc'] = roc_auc_score(labels_arr, preds_arr)
            else:
                metrics['auc_roc'] = 0.0
        except Exception:
            metrics['auc_roc'] = 0.0

        # Average Precision (PR-AUC)
        try:
            if len(np.unique(labels_arr)) > 1:
                metrics['avg_precision'] = average_precision_score(labels_arr, preds_arr)
            else:
                metrics['avg_precision'] = 0.0
        except Exception:
            metrics['avg_precision'] = 0.0

        # Precision, Recall, F1
        precision, recall, f1, _ = precision_recall_fscore_support(
            labels_arr, pred_binary, average='binary', zero_division=0
        )
        metrics['precision'] = precision
        metrics['recall'] = recall
        metrics['f1'] = f1

        # Confusion matrix
        try:
            tn, fp, fn, tp = confusion_matrix(labels_arr, pred_binary, labels=[0, 1]).ravel()
            metrics['sensitivity'] = tp / (tp + fn) if (tp + fn) > 0 else 0.0
            metrics['specificity'] = tn / (tn + fp) if (tn + fp) > 0 else 0.0
            metrics['accuracy'] = (tp + tn) / (tp + tn + fp + fn)
            metrics['tp'] = int(tp)
            metrics['fp'] = int(fp)
            metrics['tn'] = int(tn)
            metrics['fn'] = int(fn)
        except Exception:
            metrics['sensitivity'] = 0.0
            metrics['specificity'] = 0.0
            metrics['accuracy'] = (pred_binary == labels_arr).mean()

        return metrics

    @torch.no_grad()
    def validate(self, epoch):
        """Run validation on single GPU to avoid DataParallel eval-mode errors.
        
        DataParallel + MultiheadAttention in eval mode causes CUDA misaligned
        address errors. Solution: unwrap to model.module for validation.
        """
        # Unwrap DataParallel for validation (avoids CUDA misaligned address)
        if isinstance(self.model, nn.DataParallel):
            val_model = self.model.module
        else:
            val_model = self.model
        
        val_model.eval()
        total_loss = 0.0
        num_batches = 0
        all_preds = []
        all_labels = []

        for nodule, context, labels in self.val_loader:
            nodule = nodule.to(self.device)
            context = context.to(self.device)
            labels = labels.to(self.device)

            logits = val_model(nodule, context)
            
            # Replace NaN logits with 0 (uncertain) to prevent metric crashes
            if torch.isnan(logits).any():
                logits = torch.nan_to_num(logits, nan=0.0)
            
            loss, _ = self.criterion(logits, labels)

            if not torch.isnan(loss):
                total_loss += loss.item()
                num_batches += 1

            probs = torch.sigmoid(logits.squeeze(-1))
            probs = torch.clamp(probs, 1e-7, 1.0 - 1e-7)
            all_preds.extend(probs.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())

        avg_loss = total_loss / max(num_batches, 1)

        # Compute metrics
        metrics = self._compute_metrics(all_preds, all_labels)

        # Debug: prediction distribution (safe for NaN arrays)
        preds_arr = np.array(all_preds)
        labels_arr = np.array(all_labels)
        nan_count = np.isnan(preds_arr).sum()
        
        # Track prediction std for collapse detection
        valid_preds = preds_arr[~np.isnan(preds_arr)]
        metrics['_preds_std'] = float(valid_preds.std()) if len(valid_preds) > 0 else 0.0
        
        if nan_count > 0:
            self.logger.warning(
                f"Epoch {epoch+1} | Val preds contain {nan_count} NaN values!"
            )
        else:
            self.logger.info(
                f"Epoch {epoch+1} | Val preds: min={preds_arr.min():.4f} "
                f"max={preds_arr.max():.4f} mean={preds_arr.mean():.4f} | "
                f"Labels: {int(labels_arr.sum())}/{len(labels_arr)} positive"
            )

        self.logger.info(
            f"Epoch {epoch+1} | Val Loss: {avg_loss:.4f} | "
            f"AUC: {metrics['auc_roc']:.4f} | "
            f"Sens: {metrics['sensitivity']:.4f} | "
            f"Spec: {metrics['specificity']:.4f} | "
            f"F1: {metrics['f1']:.4f} | "
            f"Acc: {metrics['accuracy']:.4f}"
        )

        if self.writer:
            self.writer.add_scalar('val/loss', avg_loss, epoch)
            for k, v in metrics.items():
                if isinstance(v, (int, float)):
                    self.writer.add_scalar(f'val/{k}', v, epoch)

        return avg_loss, metrics

    def save_checkpoint(self, epoch, val_loss, val_metrics=None, is_best=False):
        """Save model checkpoint."""
        model_state = (self.model.module.state_dict() 
                      if isinstance(self.model, nn.DataParallel) 
                      else self.model.state_dict())
        
        checkpoint = {
            'epoch': epoch,
            'model_state_dict': model_state,
            'optimizer_state_dict': self.optimizer.state_dict(),
            'scheduler_state_dict': self.scheduler.state_dict(),
            'scaler_state_dict': self.scaler.state_dict(),
            'val_loss': val_loss,
            'val_metrics': val_metrics,
            'best_val_auc': self.best_val_auc,
            'global_step': self.global_step,
            'config': self.config,
        }

        # Save latest
        path = self.ckpt_dir / 'last.pth'
        torch.save(checkpoint, path)

        if is_best:
            best_path = self.ckpt_dir / 'best.pth'
            torch.save(checkpoint, best_path)
            auc_str = f"{val_metrics['auc_roc']:.4f}" if val_metrics else "N/A"
            self.logger.info(f"New best model saved (AUC: {auc_str}, val_loss: {val_loss:.4f})")

        if (epoch + 1) % self.save_interval == 0:
            periodic_path = self.ckpt_dir / f'epoch_{epoch+1}.pth'
            torch.save(checkpoint, periodic_path)

    def load_checkpoint(self, checkpoint_path):
        """Load model checkpoint."""
        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
        
        if isinstance(self.model, nn.DataParallel):
            self.model.module.load_state_dict(ckpt['model_state_dict'])
        else:
            self.model.load_state_dict(ckpt['model_state_dict'])
        
        self.optimizer.load_state_dict(ckpt['optimizer_state_dict'])
        self.scheduler.load_state_dict(ckpt['scheduler_state_dict'])
        self.scaler.load_state_dict(ckpt['scaler_state_dict'])
        self.start_epoch = ckpt['epoch'] + 1
        self.global_step = ckpt['global_step']
        self.best_val_loss = ckpt.get('val_loss', float('inf'))
        self.best_val_auc = ckpt.get('best_val_auc', 0.0)
        
        self.logger.info(
            f"Resumed from epoch {self.start_epoch} "
            f"(val_loss: {self.best_val_loss:.4f}, best_auc: {self.best_val_auc:.4f})"
        )

    def _rebuild_train_loader(self, stage):
        """Rebuild training DataLoader for curriculum stage transition."""
        from src.data.dataset import create_data_loaders
        self.logger.info(f"Rebuilding train loader for curriculum stage {stage}...")
        train_loader, _, _ = create_data_loaders(self.config, curriculum_stage=stage)
        self.train_loader = train_loader
        self._current_curriculum_stage = stage
        self.logger.info(f"  New train batches: {len(self.train_loader)}")

    def train(self, num_epochs=None, dry_run=False):
        """Full training loop."""
        if num_epochs is None:
            num_epochs = self.train_cfg.get('num_epochs', 60)

        self.logger.info(f"\n{'='*60}")
        self.logger.info(f"Starting training: {num_epochs} epochs")
        self.logger.info(f"Device: {self.device}")
        self.logger.info(f"AMP (training only): {self.use_amp}")
        self.logger.info(f"Early stopping: patience={self.patience}, metric=AUC-ROC")
        if self._curriculum:
            self.logger.info(
                f"Curriculum learning: stage1={self._curriculum.get('stage1_epochs',0)} eps, "
                f"stage2={self._curriculum.get('stage2_epochs',0)} eps, "
                f"stage3={self._curriculum.get('stage3_epochs',0)} eps"
            )
        self.logger.info(f"{'='*60}\n")

        # Initialize curriculum: rebuild with stage 1 if curriculum is enabled
        if self._curriculum and self._current_curriculum_stage == 1 and self.start_epoch == 0:
            self._rebuild_train_loader(stage=1)

        if dry_run:
            num_epochs = 1
            self.logger.info("DRY RUN MODE: running 2 batches only")

        for epoch in range(self.start_epoch, self.start_epoch + num_epochs):
            # Check curriculum stage transitions
            if self._curriculum:
                s1_end = self._curriculum.get('stage1_epochs', 0)
                s2_end = s1_end + self._curriculum.get('stage2_epochs', 0)
                if epoch == s1_end and self._current_curriculum_stage < 2:
                    self.logger.info(f"\n>>> Curriculum Stage 2: Adding medium-difficulty samples (epoch {epoch+1})")
                    self._rebuild_train_loader(stage=2)
                elif epoch == s2_end and self._current_curriculum_stage < 3:
                    self.logger.info(f"\n>>> Curriculum Stage 3: Using ALL samples (epoch {epoch+1})")
                    self._rebuild_train_loader(stage=3)
            start = time.time()

            if dry_run:
                self.model.train()
                batch_count = 0
                for nodule, context, labels in self.train_loader:
                    nodule = nodule.to(self.device)
                    context = context.to(self.device)
                    labels = labels.to(self.device)

                    self.optimizer.zero_grad()
                    with autocast('cuda', enabled=self.use_amp):
                        logits = self.model(nodule, context)
                        loss, loss_dict = self.criterion(logits, labels)
                    self.scaler.scale(loss).backward()
                    self.scaler.step(self.optimizer)
                    self.scaler.update()

                    batch_count += 1
                    self.logger.info(
                        f"Dry run batch {batch_count}/2 | Loss: {loss.item():.4f}"
                    )
                    if batch_count >= 2:
                        break

                self.logger.info("\nDry run complete! Model can train successfully.")
                return

            # Normal training
            train_result = self.train_epoch(epoch)
            train_loss = train_result[0]
            nan_count = train_result[1] if len(train_result) > 1 else 0

            # NaN recovery: even a few NaN batches can corrupt weights
            if nan_count >= 3:
                self.logger.warning(
                    f"\nWARNING: {nan_count} NaN batches detected — weights likely corrupted!"
                )
                recovery_ckpt = self.ckpt_dir / 'best.pth'
                if not recovery_ckpt.exists():
                    recovery_ckpt = self.ckpt_dir / 'last.pth'
                if recovery_ckpt.exists():
                    self.logger.info(f"Recovering from: {recovery_ckpt}")
                    self.load_checkpoint(recovery_ckpt)
                    # Lower learning rate to prevent recurrence
                    for pg in self.optimizer.param_groups:
                        pg['lr'] *= 0.5
                    self.logger.info(
                        f"Reduced LR to {self.optimizer.param_groups[0]['lr']:.2e}"
                    )
                    continue
                else:
                    self.logger.error("No checkpoint to recover from!")
                    break

            val_loss, val_metrics = self.validate(epoch)
            elapsed = time.time() - start

            val_auc = val_metrics.get('auc_roc', 0.0)

            self.logger.info(
                f"Epoch {epoch+1}/{self.start_epoch + num_epochs} | "
                f"Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f} | "
                f"Val AUC: {val_auc:.4f} | Time: {elapsed:.1f}s"
            )

            # Detect model collapse: all predictions are constant (e.g., all 0.5)
            preds_std = val_metrics.get('_preds_std', 1.0)
            if preds_std < 0.001 and epoch > self.warmup_epochs:
                self._collapse_count = getattr(self, '_collapse_count', 0) + 1
                self.logger.warning(
                    f"Model collapse detected! Val preds are constant "
                    f"(std={preds_std:.6f}). Collapse count: {self._collapse_count}"
                )
                # Auto-recover if we have a good checkpoint and haven't looped too many times
                if self._collapse_count <= 3:
                    recovery_ckpt = self.ckpt_dir / 'best.pth'
                    if recovery_ckpt.exists():
                        self.logger.info(f"Auto-recovering from: {recovery_ckpt}")
                        self.load_checkpoint(recovery_ckpt)
                        for pg in self.optimizer.param_groups:
                            pg['lr'] *= 0.5
                        self.logger.info(
                            f"Reduced LR to {self.optimizer.param_groups[0]['lr']:.2e}"
                        )
                        continue
                else:
                    self.logger.warning(
                        "Multiple collapse recoveries failed — continuing with current weights"
                    )
            else:
                self._collapse_count = 0  # Reset on healthy epoch

            # Early stopping on AUC-ROC (min_delta=0.001 per spec §2.4)
            is_best = val_auc > (self.best_val_auc + 0.001)
            if is_best:
                self.best_val_auc = val_auc
                self.best_val_loss = val_loss
                self.epochs_no_improve = 0
            else:
                self.epochs_no_improve += 1

            self.save_checkpoint(epoch, val_loss, val_metrics=val_metrics, is_best=is_best)

            if self.epochs_no_improve >= self.patience:
                self.logger.info(
                    f"\nEarly stopping at epoch {epoch+1} "
                    f"(no AUC improvement for {self.patience} epochs)"
                )
                break

        if self.writer:
            self.writer.close()
        
        self.logger.info(
            f"\nTraining complete! Best AUC: {self.best_val_auc:.4f} | "
            f"Best val loss: {self.best_val_loss:.4f}"
        )