"""
Training Utilities for CyMeta Models.

This module provides training utilities for jointly training meta-dictionaries,
index maps, and gating networks, including support for knowledge distillation.
"""

import torch
import torch.nn as nn
from typing import Optional, Dict, List, Callable
from torch.utils.data import DataLoader

try:
    from tqdm import tqdm
except ImportError:
    # Fallback if tqdm is not available
    def tqdm(iterable, desc=None):
        return iterable


class CyMetaTrainer:
    """
    Trainer for CyMeta compressed models.
    
    Supports:
    - Joint training of dictionaries, index maps, and gating networks
    - Knowledge distillation from full-precision teacher models
    - Fine-tuning of compressed models
    
    Args:
        model: CyMeta model to train
        device: Device for training (default: 'cpu')
        optimizer: Optimizer (default: None, will create Adam)
        learning_rate: Learning rate (default: 0.001)
    """
    
    def __init__(
        self,
        model: nn.Module,
        device: str = "cpu",
        optimizer: Optional[torch.optim.Optimizer] = None,
        learning_rate: float = 0.001,
    ):
        self.model = model.to(device)
        self.device = device
        
        if optimizer is None:
            self.optimizer = torch.optim.Adam(
                self.model.parameters(),
                lr=learning_rate,
            )
        else:
            self.optimizer = optimizer
        
        self.history: List[Dict] = []
    
    def train_epoch(
        self,
        dataloader: DataLoader,
        criterion: Callable,
        teacher_model: Optional[nn.Module] = None,
        distillation_weight: float = 0.5,
        temperature: float = 4.0,
        verbose: bool = True,
    ) -> Dict:
        """
        Train for one epoch.
        
        Args:
            dataloader: DataLoader for training data
            criterion: Loss function
            teacher_model: Optional teacher model for knowledge distillation
            distillation_weight: Weight for distillation loss (default: 0.5)
            temperature: Temperature for distillation (default: 4.0)
            verbose: Whether to show progress bar (default: True)
        
        Returns:
            Dictionary with training metrics
        """
        self.model.train()
        total_loss = 0.0
        total_distillation_loss = 0.0
        num_batches = 0
        
        iterator = tqdm(dataloader, desc="Training") if verbose else dataloader
        
        for batch in iterator:
            # Move batch to device
            if isinstance(batch, (list, tuple)):
                inputs = batch[0].to(self.device)
                targets = batch[1].to(self.device)
            else:
                inputs = batch.to(self.device)
                targets = None
            
            self.optimizer.zero_grad()
            
            # Forward pass
            outputs = self.model(inputs)
            
            # Compute loss
            if targets is not None:
                loss = criterion(outputs, targets)
            else:
                # Self-supervised or other loss
                loss = criterion(outputs)
            
            # Knowledge distillation
            distillation_loss = 0.0
            if teacher_model is not None:
                with torch.no_grad():
                    teacher_outputs = teacher_model(inputs)
                
                # Soft targets with temperature
                if hasattr(teacher_outputs, "logits"):
                    teacher_logits = teacher_outputs.logits
                else:
                    teacher_logits = teacher_outputs
                
                student_logits = outputs if isinstance(outputs, torch.Tensor) else outputs.logits
                
                # Distillation loss (KL divergence)
                teacher_probs = torch.softmax(teacher_logits / temperature, dim=-1)
                student_logs = torch.log_softmax(student_logits / temperature, dim=-1)
                distillation_loss = torch.nn.functional.kl_div(
                    student_logs, teacher_probs, reduction="batchmean"
                ) * (temperature ** 2)
                
                # Combined loss
                loss = (1 - distillation_weight) * loss + distillation_weight * distillation_loss
            
            # Backward pass
            loss.backward()
            self.optimizer.step()
            
            total_loss += loss.item()
            if distillation_loss > 0:
                total_distillation_loss += distillation_loss.item()
            num_batches += 1
        
        avg_loss = total_loss / num_batches if num_batches > 0 else 0.0
        avg_distillation_loss = total_distillation_loss / num_batches if num_batches > 0 else 0.0
        
        metrics = {
            "loss": avg_loss,
            "distillation_loss": avg_distillation_loss if distillation_loss > 0 else None,
        }
        
        self.history.append(metrics)
        return metrics
    
    def evaluate(
        self,
        dataloader: DataLoader,
        criterion: Callable,
        verbose: bool = True,
    ) -> Dict:
        """
        Evaluate the model.
        
        Args:
            dataloader: DataLoader for evaluation data
            criterion: Loss function
            verbose: Whether to show progress bar (default: True)
        
        Returns:
            Dictionary with evaluation metrics
        """
        self.model.eval()
        total_loss = 0.0
        num_batches = 0
        
        iterator = tqdm(dataloader, desc="Evaluating") if verbose else dataloader
        
        with torch.no_grad():
            for batch in iterator:
                # Move batch to device
                if isinstance(batch, (list, tuple)):
                    inputs = batch[0].to(self.device)
                    targets = batch[1].to(self.device)
                else:
                    inputs = batch.to(self.device)
                    targets = None
                
                # Forward pass
                outputs = self.model(inputs)
                
                # Compute loss
                if targets is not None:
                    loss = criterion(outputs, targets)
                else:
                    loss = criterion(outputs)
                
                total_loss += loss.item()
                num_batches += 1
        
        avg_loss = total_loss / num_batches if num_batches > 0 else 0.0
        
        return {"loss": avg_loss}
    
    def train(
        self,
        train_loader: DataLoader,
        num_epochs: int,
        criterion: Callable,
        val_loader: Optional[DataLoader] = None,
        teacher_model: Optional[nn.Module] = None,
        distillation_weight: float = 0.5,
        temperature: float = 4.0,
        verbose: bool = True,
    ) -> List[Dict]:
        """
        Train the model for multiple epochs.
        
        Args:
            train_loader: DataLoader for training data
            num_epochs: Number of training epochs
            criterion: Loss function
            val_loader: Optional DataLoader for validation data
            teacher_model: Optional teacher model for knowledge distillation
            distillation_weight: Weight for distillation loss
            temperature: Temperature for distillation
            verbose: Whether to show progress
        
        Returns:
            List of training history dictionaries
        """
        self.history = []
        
        for epoch in range(num_epochs):
            if verbose:
                print(f"\nEpoch {epoch + 1}/{num_epochs}")
            
            # Train
            train_metrics = self.train_epoch(
                train_loader,
                criterion,
                teacher_model=teacher_model,
                distillation_weight=distillation_weight,
                temperature=temperature,
                verbose=verbose,
            )
            
            # Validate
            val_metrics = None
            if val_loader is not None:
                val_metrics = self.evaluate(val_loader, criterion, verbose=verbose)
            
            # Log
            if verbose:
                print(f"Train Loss: {train_metrics['loss']:.6f}")
                if val_metrics:
                    print(f"Val Loss: {val_metrics['loss']:.6f}")
            
            epoch_history = {
                "epoch": epoch + 1,
                "train": train_metrics,
                "val": val_metrics,
            }
            self.history.append(epoch_history)
        
        return self.history
    
    def save_checkpoint(self, filepath: str) -> None:
        """Save training checkpoint."""
        checkpoint = {
            "model_state_dict": self.model.state_dict(),
            "optimizer_state_dict": self.optimizer.state_dict(),
            "history": self.history,
        }
        torch.save(checkpoint, filepath)
    
    def load_checkpoint(self, filepath: str) -> None:
        """Load training checkpoint."""
        checkpoint = torch.load(filepath, map_location=self.device)
        self.model.load_state_dict(checkpoint["model_state_dict"])
        self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        self.history = checkpoint.get("history", [])

