src/training/trainer.py

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.amp import GradScaler, autocast
from tqdm import tqdm
import numpy as np
from pathlib import Path
import json
import time
import logging
from typing import Dict, Tuple, Optional, List, Union
from sklearn.metrics import confusion_matrix, classification_report

from .losses import get_loss_function
from .metrics import MetricTracker

logger = logging.getLogger(__name__)


class Trainer:
    """Main trainer class for model training with multi-task support"""

    def __init__(
            self,
            model: nn.Module,
            config: Dict,
            device: torch.device,
            experiment_dir: Path,
            logger: Optional[logging.Logger] = None,
            multi_task: bool = False
    ):
        self.model = model
        self.config = config
        self.device = device
        self.exp_dir = experiment_dir
        self.logger = logger or logging.getLogger(__name__)
        self.multi_task = multi_task

        # Initialize tracking
        self.current_epoch = 0
        self.best_val_metric = 0
        self.best_epoch = 0
        self.metrics_history = {'train': [], 'val': []}

        # Setup components
        self._setup_loss()
        self._setup_amp()

    def _setup_loss(self):
        """Setup loss function"""
        if self.multi_task:
            # Multi-task loss setup - will be set externally
            self.logger.info("Multi-task loss will be set externally")
            self.criterion = None
        else:
            # Single-task loss setup (original)
            loss_config = self.config.get('loss', {})
            loss_name = loss_config.get('name', 'cross_entropy')
            loss_params = loss_config.get('params', {})
            self.criterion = get_loss_function(loss_name, **loss_params)

    def _setup_amp(self):
        """Setup automatic mixed precision"""
        self.use_amp = self.config.get('use_amp', True) and torch.cuda.is_available()
        self.scaler = GradScaler('cuda') if self.use_amp else None

    def _calculate_accuracy(self, outputs: Union[torch.Tensor, Dict], labels: Union[torch.Tensor, Dict]) -> Dict[str, float]:
        """Calculate accuracy for single-task or multi-task"""
        if self.multi_task:
            # Multi-task accuracy calculation
            accuracies = {}
            num_tasks = 0
            total_acc = 0
            
            # Decade accuracy
            _, decade_pred = outputs['decade'].max(1)
            decade_acc = decade_pred.eq(labels['decade']).float().mean().item()
            accuracies['decade_accuracy'] = decade_acc
            total_acc += decade_acc
            num_tasks += 1
            
            # Cluster accuracy
            _, cluster_pred = outputs['cluster'].max(1)
            cluster_acc = cluster_pred.eq(labels['cluster']).float().mean().item()
            accuracies['cluster_accuracy'] = cluster_acc
            total_acc += cluster_acc
            num_tasks += 1
            
            # Device accuracy (if present)
            if 'device' in outputs and 'device' in labels:
                _, device_pred = outputs['device'].max(1)
                device_acc = device_pred.eq(labels['device']).float().mean().item()
                accuracies['device_accuracy'] = device_acc
                total_acc += device_acc
                num_tasks += 1
            
            # Overall accuracy (average of all tasks)
            accuracies['accuracy'] = total_acc / num_tasks
            
            return accuracies
        else:
            # Single-task accuracy calculation
            _, predicted = outputs.max(1)
            accuracy = predicted.eq(labels).float().mean().item()
            return {'accuracy': accuracy}

    def train_epoch(
            self,
            dataloader: DataLoader,
            optimizer: torch.optim.Optimizer,
            scheduler: Optional[torch.optim.lr_scheduler._LRScheduler] = None,
            epoch: int = 0
    ) -> Dict[str, float]:
        """Train for one epoch"""
        self.model.train()
        metric_tracker = MetricTracker()

        pbar = tqdm(dataloader, desc=f'Epoch {epoch} - Training')

        for batch_idx, (images, labels, _) in enumerate(pbar):
            images = images.to(self.device)
            
            # Handle labels for multi-task vs single-task
            if self.multi_task:
                # labels is a dict with 'decade' and 'cluster' keys
                if isinstance(labels, dict):
                    processed_labels = {}
                    for k, v in labels.items():
                        if isinstance(v, torch.Tensor):
                            try:
                                processed_labels[k] = v.to(self.device)
                            except Exception as e:
                                self.logger.error(f"Error moving label '{k}' to device: {e}")
                                self.logger.error(f"Label key: {k}, Value type: {type(v)}, Device: {self.device}")
                                raise
                        elif v == 'device':
                            # Skip if v is the string 'device' - this seems to be metadata
                            self.logger.warning(f"Skipping non-tensor label '{k}' with value '{v}'")
                            continue
                        else:
                            self.logger.error(f"Label '{k}' is not a tensor: {type(v)}, value: {v}")
                            raise TypeError(f"Expected tensor for label '{k}', got {type(v)}")
                    labels = processed_labels
                else:
                    self.logger.error(f"Labels is not a dict in multi-task mode: {type(labels)}")
                    raise ValueError("Multi-task mode requires labels to be a dictionary")
            else:
                # labels is a tensor
                labels = labels.to(self.device)

            # Forward pass
            optimizer.zero_grad()

            with autocast('cuda', enabled=self.use_amp):
                outputs = self.model(images)
                
                if self.multi_task:
                    # Multi-task loss calculation
                    losses = self.criterion(outputs, labels)
                    loss = losses['total_loss']
                else:
                    # Single-task loss calculation
                    loss = self.criterion(outputs, labels)

            # Backward pass
            if self.use_amp:
                self.scaler.scale(loss).backward()

                # Gradient clipping
                if self.config.get('gradient_clip_val', 0) > 0:
                    self.scaler.unscale_(optimizer)
                    torch.nn.utils.clip_grad_norm_(
                        self.model.parameters(),
                        self.config['gradient_clip_val']
                    )

                self.scaler.step(optimizer)
                self.scaler.update()
            else:
                loss.backward()

                # Gradient clipping
                if self.config.get('gradient_clip_val', 0) > 0:
                    torch.nn.utils.clip_grad_norm_(
                        self.model.parameters(),
                        self.config['gradient_clip_val']
                    )

                optimizer.step()

            # Update metrics
            if self.multi_task:
                # Track individual task losses
                metric_tracker.update('total_loss', losses['total_loss'].item())
                metric_tracker.update('decade_loss', losses['decade_loss'].item())
                metric_tracker.update('cluster_loss', losses['cluster_loss'].item())
                if 'device_loss' in losses:
                    metric_tracker.update('device_loss', losses['device_loss'].item())
                
                # Track accuracies
                accuracies = self._calculate_accuracy(outputs, labels)
                for acc_name, acc_value in accuracies.items():
                    metric_tracker.update(acc_name, acc_value)
                
                # Update progress bar
                postfix_dict = {
                    'loss': f'{metric_tracker.avg("total_loss"):.4f}',
                    'dec_acc': f'{metric_tracker.avg("decade_accuracy") * 100:.1f}%',
                    'cls_acc': f'{metric_tracker.avg("cluster_accuracy") * 100:.1f}%'
                }
                if 'device_accuracy' in accuracies:
                    postfix_dict['dev_acc'] = f'{metric_tracker.avg("device_accuracy") * 100:.1f}%'
                pbar.set_postfix(postfix_dict)
            else:
                # Single-task metrics
                metric_tracker.update('loss', loss.item())
                accuracies = self._calculate_accuracy(outputs, labels)
                metric_tracker.update('accuracy', accuracies['accuracy'])
                
                # Update progress bar
                pbar.set_postfix({
                    'loss': f'{metric_tracker.avg("loss"):.4f}',
                    'acc': f'{metric_tracker.avg("accuracy") * 100:.2f}%'
                })

            # Step scheduler if it's batch-wise
            if scheduler and self.config.get('scheduler_step', 'epoch') == 'batch':
                scheduler.step()

        return metric_tracker.get_averages()

    def validate(
            self,
            dataloader: DataLoader,
            epoch: int = 0
    ) -> Tuple[Dict[str, float], Dict[str, np.ndarray], Dict[str, np.ndarray]]:
        """Validate the model"""
        self.model.eval()
        metric_tracker = MetricTracker()

        if self.multi_task:
            all_predictions = {'decade': [], 'cluster': []}
            all_labels = {'decade': [], 'cluster': []}
            # Always initialize device lists for multi-task since we have device labels
            all_predictions['device'] = []
            all_labels['device'] = []
        else:
            all_predictions = []
            all_labels = []

        with torch.no_grad():
            for batch_idx, (images, labels, _) in enumerate(tqdm(dataloader, desc=f'Epoch {epoch} - Validation')):
                try:
                    images = images.to(self.device)
                except Exception as e:
                    self.logger.error(f"Error moving images to device at batch {batch_idx}: {e}")
                    self.logger.error(f"Device: {self.device}, Images type: {type(images)}")
                    raise
                
                # Handle labels
                if self.multi_task:
                    # Ensure labels is a dictionary and all values are tensors
                    if isinstance(labels, dict):
                        processed_labels = {}
                        for k, v in labels.items():
                            if isinstance(v, torch.Tensor):
                                try:
                                    processed_labels[k] = v.to(self.device)
                                except Exception as e:
                                    self.logger.error(f"Error moving label '{k}' to device: {e}")
                                    self.logger.error(f"Label key: {k}, Value type: {type(v)}, Device: {self.device}")
                                    self.logger.error(f"Tensor shape: {v.shape if hasattr(v, 'shape') else 'N/A'}")
                                    raise
                            elif v == 'device':
                                # Skip if v is the string 'device' - this seems to be metadata
                                self.logger.warning(f"Skipping non-tensor label '{k}' with value '{v}'")
                                continue
                            else:
                                self.logger.error(f"Label '{k}' is not a tensor: {type(v)}, value: {v}")
                                raise TypeError(f"Expected tensor for label '{k}', got {type(v)}")
                        labels = processed_labels
                    else:
                        self.logger.error(f"Labels is not a dict in multi-task mode: {type(labels)}")
                        self.logger.error(f"Labels content: {labels}")
                        raise ValueError("Multi-task mode requires labels to be a dictionary")
                else:
                    labels = labels.to(self.device)

                outputs = self.model(images)
                
                # Calculate loss
                if self.multi_task:
                    losses = self.criterion(outputs, labels)
                    loss = losses['total_loss']
                    
                    # Update loss metrics
                    metric_tracker.update('total_loss', losses['total_loss'].item())
                    metric_tracker.update('decade_loss', losses['decade_loss'].item())
                    metric_tracker.update('cluster_loss', losses['cluster_loss'].item())
                    if 'device_loss' in losses:
                        metric_tracker.update('device_loss', losses['device_loss'].item())
                    
                    # Update accuracy metrics
                    accuracies = self._calculate_accuracy(outputs, labels)
                    for acc_name, acc_value in accuracies.items():
                        metric_tracker.update(acc_name, acc_value)
                    
                    # Collect predictions and labels
                    _, decade_pred = outputs['decade'].max(1)
                    _, cluster_pred = outputs['cluster'].max(1)
                    
                    all_predictions['decade'].extend(decade_pred.cpu().numpy())
                    all_predictions['cluster'].extend(cluster_pred.cpu().numpy())
                    all_labels['decade'].extend(labels['decade'].cpu().numpy())
                    all_labels['cluster'].extend(labels['cluster'].cpu().numpy())
                    
                    # Collect device predictions if present
                    if 'device' in outputs and 'device' in labels:
                        _, device_pred = outputs['device'].max(1)
                        all_predictions['device'].extend(device_pred.cpu().numpy())
                        all_labels['device'].extend(labels['device'].cpu().numpy())
                    
                else:
                    loss = self.criterion(outputs, labels)
                    metric_tracker.update('loss', loss.item())
                    
                    accuracies = self._calculate_accuracy(outputs, labels)
                    metric_tracker.update('accuracy', accuracies['accuracy'])
                    
                    # Collect predictions and labels
                    _, predicted = outputs.max(1)
                    all_predictions.extend(predicted.cpu().numpy())
                    all_labels.extend(labels.cpu().numpy())

        # Convert to numpy arrays
        if self.multi_task:
            predictions = {k: np.array(v) for k, v in all_predictions.items()}
            labels_np = {k: np.array(v) for k, v in all_labels.items()}
        else:
            predictions = np.array(all_predictions)
            labels_np = np.array(all_labels)

        metrics = metric_tracker.get_averages()
        
        return metrics, predictions, labels_np

    def _log_classification_reports(
        self, 
        predictions: Union[np.ndarray, Dict[str, np.ndarray]], 
        labels: Union[np.ndarray, Dict[str, np.ndarray]],
        class_names: Optional[Dict[str, List[str]]] = None
    ):
        """Log detailed classification reports"""
        if self.multi_task:
            # Multi-task classification reports
            tasks = ['decade', 'cluster']
            if 'device' in predictions:
                tasks.append('device')
                
            for task in tasks:
                self.logger.info(f"\n{task.capitalize()} Classification Report:")
                task_class_names = class_names.get(task) if class_names else None
                
                # Check if we have multiple classes for this task
                unique_labels = np.unique(labels[task])
                if len(unique_labels) <= 1:
                    self.logger.info(f"Only {len(unique_labels)} class(es) found for {task}. Skipping classification report.")
                    continue
                
                try:
                    report = classification_report(
                        labels[task], 
                        predictions[task],
                        target_names=task_class_names,
                        output_dict=False
                    )
                    self.logger.info(f"\n{report}")
                except Exception as e:
                    self.logger.warning(f"Failed to generate classification report for {task}: {e}")
                    self.logger.info(f"Labels shape: {labels[task].shape}, Predictions shape: {predictions[task].shape}")
                    if task_class_names is not None:
                        self.logger.info(f"Class names type: {type(task_class_names)}, Class names: {task_class_names}")
        else:
            # Single-task classification report
            self.logger.info("\nClassification Report:")
            task_class_names = class_names if isinstance(class_names, list) else None
            
            # Check if we have multiple classes
            unique_labels = np.unique(labels)
            if len(unique_labels) <= 1:
                self.logger.info(f"Only {len(unique_labels)} class(es) found. Skipping classification report.")
                return
            
            try:
                report = classification_report(
                    labels, 
                    predictions,
                    target_names=task_class_names,
                    output_dict=False
                )
                self.logger.info(f"\n{report}")
            except Exception as e:
                self.logger.warning(f"Failed to generate classification report: {e}")
                self.logger.info(f"Labels shape: {labels.shape}, Predictions shape: {predictions.shape}")
                if task_class_names is not None:
                    self.logger.info(f"Class names type: {type(task_class_names)}, Class names: {task_class_names}")

    def train(
            self,
            train_loader: DataLoader,
            val_loader: DataLoader,
            optimizer: torch.optim.Optimizer,
            scheduler: Optional[torch.optim.lr_scheduler._LRScheduler] = None,
            start_epoch: int = 0,
            class_names: Optional[Dict[str, List[str]]] = None
    ) -> Dict:
        """Main training loop"""
        self.logger.info("Starting training...")
        start_time = time.time()

        num_epochs = self.config.get('epochs', 30)

        for epoch in range(start_epoch, num_epochs):
            self.current_epoch = epoch

            # Train
            train_metrics = self.train_epoch(
                train_loader, optimizer, scheduler, epoch
            )

            # Validate
            val_metrics, predictions, labels = self.validate(
                val_loader, epoch
            )

            # Step scheduler if it's epoch-wise
            if scheduler and self.config.get('scheduler_step', 'epoch') == 'epoch':
                if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
                    # Use appropriate metric for scheduler
                    if self.multi_task:
                        # For multi-task, use total_loss if available, otherwise use accuracy
                        monitor_metric = 'total_loss' if 'total_loss' in val_metrics else 'accuracy'
                    else:
                        monitor_metric = 'loss'
                    
                    if monitor_metric in val_metrics:
                        scheduler.step(val_metrics[monitor_metric])
                    else:
                        self.logger.warning(f"Monitor metric '{monitor_metric}' not found in validation metrics")
                else:
                    scheduler.step()

            # Log metrics
            current_lr = optimizer.param_groups[0]['lr']
            
            if self.multi_task:
                log_msg = (
                    f"Epoch {epoch}/{num_epochs} - "
                    f"Train Loss: {train_metrics['total_loss']:.4f} "
                    f"(Dec: {train_metrics['decade_loss']:.4f}, Cls: {train_metrics['cluster_loss']:.4f}"
                )
                if 'device_loss' in train_metrics:
                    log_msg += f", Dev: {train_metrics['device_loss']:.4f}"
                log_msg += f"), Train Acc: {train_metrics['accuracy'] * 100:.2f}% "
                log_msg += f"(Dec: {train_metrics['decade_accuracy'] * 100:.2f}%, Cls: {train_metrics['cluster_accuracy'] * 100:.2f}%"
                if 'device_accuracy' in train_metrics:
                    log_msg += f", Dev: {train_metrics['device_accuracy'] * 100:.2f}%"
                log_msg += f"), Val Loss: {val_metrics['total_loss']:.4f}, "
                log_msg += f"Val Acc: {val_metrics['accuracy'] * 100:.2f}% "
                log_msg += f"(Dec: {val_metrics['decade_accuracy'] * 100:.2f}%, Cls: {val_metrics['cluster_accuracy'] * 100:.2f}%"
                if 'device_accuracy' in val_metrics:
                    log_msg += f", Dev: {val_metrics['device_accuracy'] * 100:.2f}%"
                log_msg += f"), LR: {current_lr:.6f}"
                self.logger.info(log_msg)
            else:
                self.logger.info(
                    f"Epoch {epoch}/{num_epochs} - "
                    f"Train Loss: {train_metrics['loss']:.4f}, "
                    f"Train Acc: {train_metrics['accuracy'] * 100:.2f}%, "
                    f"Val Loss: {val_metrics['loss']:.4f}, "
                    f"Val Acc: {val_metrics['accuracy'] * 100:.2f}%, "
                    f"LR: {current_lr:.6f}"
                )

            # Log detailed classification report every few epochs
            if epoch % self.config.get('log_report_every', 5) == 0:
                self._log_classification_reports(predictions, labels, class_names)

            # Save metrics history
            self.metrics_history['train'].append({
                'epoch': epoch,
                **train_metrics,
                'lr': current_lr
            })
            self.metrics_history['val'].append({
                'epoch': epoch,
                **val_metrics
            })

            # Check if best model
            monitor_metric = self.config.get('monitor_metric', 'accuracy')
            
            # Ensure the monitor metric exists in val_metrics
            if monitor_metric not in val_metrics:
                available_metrics = list(val_metrics.keys())
                self.logger.warning(f"Monitor metric '{monitor_metric}' not found. Available metrics: {available_metrics}")
                # Fallback to a sensible default
                if 'accuracy' in val_metrics:
                    monitor_metric = 'accuracy'
                elif 'total_loss' in val_metrics:
                    monitor_metric = 'total_loss'
                else:
                    monitor_metric = available_metrics[0] if available_metrics else 'accuracy'
                self.logger.info(f"Using fallback monitor metric: {monitor_metric}")
            
            val_metric = val_metrics[monitor_metric]
            
            # For loss metrics, lower is better
            if 'loss' in monitor_metric:
                is_best = val_metric < self.best_val_metric if self.best_val_metric != 0 else True
            else:
                is_best = val_metric > self.best_val_metric

            if is_best:
                self.best_val_metric = val_metric
                self.best_epoch = epoch
                self.logger.info(f"New best model! {monitor_metric}: {val_metric:.4f}")

            # Save checkpoint
            self.save_checkpoint(
                optimizer, scheduler, epoch, val_metrics, is_best, class_names
            )

            # Early stopping
            if self.config.get('early_stopping', 0) > 0:
                epochs_without_improvement = epoch - self.best_epoch
                if epochs_without_improvement >= self.config['early_stopping']:
                    self.logger.info(f"Early stopping triggered after {epochs_without_improvement} epochs")
                    break

        # Training complete
        total_time = time.time() - start_time
        self.logger.info(f"Training complete in {total_time / 60:.2f} minutes")
        self.logger.info(
            f"Best {self.config.get('monitor_metric', 'accuracy')}: {self.best_val_metric:.4f} at epoch {self.best_epoch}")

        return {
            'best_metric': self.best_val_metric,
            'best_epoch': self.best_epoch,
            'total_time': total_time,
            'metrics_history': self.metrics_history
        }

    def save_checkpoint(
            self,
            optimizer: torch.optim.Optimizer,
            scheduler: Optional[torch.optim.lr_scheduler._LRScheduler],
            epoch: int,
            val_metrics: Dict[str, float],
            is_best: bool = False,
            class_names: Optional[Union[List[str], Dict[str, List[str]]]] = None
    ):
        """Save model checkpoint"""
        checkpoint = {
            'epoch': epoch,
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'scheduler_state_dict': scheduler.state_dict() if scheduler else None,
            'val_metrics': val_metrics,
            'best_metric': self.best_val_metric,
            'config': self.config,
            'metrics_history': self.metrics_history,
            'multi_task': self.multi_task,  # Save multi-task flag
            'class_names': class_names  # Save class names for inference
        }

        # Save last checkpoint
        checkpoint_path = self.exp_dir / 'checkpoints' / 'last_checkpoint.pth'
        checkpoint_path.parent.mkdir(exist_ok=True)
        torch.save(checkpoint, checkpoint_path)

        # Save best checkpoint
        if is_best:
            best_path = self.exp_dir / 'checkpoints' / 'best_checkpoint.pth'
            torch.save(checkpoint, best_path)

        # Save periodic checkpoint
        if self.config.get('save_every', 0) > 0 and epoch % self.config['save_every'] == 0:
            periodic_path = self.exp_dir / 'checkpoints' / f'checkpoint_epoch_{epoch}.pth'
            torch.save(checkpoint, periodic_path)

    def load_checkpoint(self, checkpoint_path: Path) -> Dict:
        """Load checkpoint and restore training state"""
        checkpoint = torch.load(checkpoint_path, map_location=self.device)

        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.metrics_history = checkpoint.get('metrics_history', {'train': [], 'val': []})
        self.best_val_metric = checkpoint.get('best_metric', 0)
        self.current_epoch = checkpoint['epoch']
        self.multi_task = checkpoint.get('multi_task', False)  # Load multi-task flag

        self.logger.info(f"Loaded checkpoint from epoch {self.current_epoch}")
        self.logger.info(f"Multi-task mode: {self.multi_task}")

        return checkpoint


class DistributedTrainer(Trainer):
    """Trainer for distributed training across multiple GPUs"""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.world_size = torch.distributed.get_world_size() if torch.distributed.is_initialized() else 1
        self.rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0

    def train_epoch(self, dataloader, optimizer, scheduler=None, epoch=0):
        """Override to handle distributed sampling"""
        if hasattr(dataloader.sampler, 'set_epoch'):
            dataloader.sampler.set_epoch(epoch)

        return super().train_epoch(dataloader, optimizer, scheduler, epoch)

    def save_checkpoint(self, *args, **kwargs):
        """Only save checkpoint on main process"""
        if self.rank == 0:
            super().save_checkpoint(*args, **kwargs)


def collate_multitask_fn(batch):
    """Custom collate function for multi-task learning"""
    images, labels, metadata = zip(*batch)
    
    # Stack images
    images = torch.stack(images)
    
    # Handle labels - check if multi-task or single task
    if isinstance(labels[0], dict):
        # Multi-task: separate decade, cluster, and device labels
        decade_labels = torch.tensor([label['decade'] for label in labels])
        cluster_labels = torch.tensor([label['cluster'] for label in labels])
        labels_dict = {
            'decade': decade_labels,
            'cluster': cluster_labels
        }
        # Add device labels if present
        if 'device' in labels[0]:
            device_labels = torch.tensor([label['device'] for label in labels])
            labels_dict['device'] = device_labels
        labels = labels_dict
    else:
        # Single task: just decade labels
        labels = torch.tensor(labels)
    
    return images, labels, metadata