"""
Knowledge Distillation for CyMeta Models.

This module provides utilities for knowledge distillation from full-precision
teacher models to CyMeta compressed student models.
"""

import torch
import torch.nn as nn
from typing import Optional, Dict, Callable


class KnowledgeDistillation:
    """
    Knowledge distillation from teacher to student model.
    
    Implements standard knowledge distillation with temperature scaling and
    optional feature matching losses.
    
    Args:
        temperature: Temperature for softmax scaling (default: 4.0)
        alpha: Weight for distillation loss vs. task loss (default: 0.5)
        use_feature_matching: Whether to use feature matching loss (default: False)
    """
    
    def __init__(
        self,
        temperature: float = 4.0,
        alpha: float = 0.5,
        use_feature_matching: bool = False,
    ):
        self.temperature = temperature
        self.alpha = alpha
        self.use_feature_matching = use_feature_matching
    
    def compute_distillation_loss(
        self,
        student_logits: torch.Tensor,
        teacher_logits: torch.Tensor,
        hard_labels: Optional[torch.Tensor] = None,
        task_loss_fn: Optional[Callable] = None,
    ) -> Dict[str, torch.Tensor]:
        """
        Compute knowledge distillation loss.
        
        Args:
            student_logits: Student model logits
            teacher_logits: Teacher model logits
            hard_labels: Optional hard labels for task loss
            task_loss_fn: Optional task loss function (e.g., CrossEntropyLoss)
        
        Returns:
            Dictionary with loss components
        """
        # Soft targets with temperature
        teacher_probs = torch.softmax(teacher_logits / self.temperature, dim=-1)
        student_log_probs = torch.log_softmax(student_logits / self.temperature, dim=-1)
        
        # KL divergence loss
        distillation_loss = torch.nn.functional.kl_div(
            student_log_probs, teacher_probs, reduction="batchmean"
        ) * (self.temperature ** 2)
        
        # Task loss (if provided)
        task_loss = None
        if hard_labels is not None and task_loss_fn is not None:
            task_loss = task_loss_fn(student_logits, hard_labels)
        
        # Combined loss
        if task_loss is not None:
            total_loss = (1 - self.alpha) * task_loss + self.alpha * distillation_loss
        else:
            total_loss = distillation_loss
        
        return {
            "total_loss": total_loss,
            "distillation_loss": distillation_loss,
            "task_loss": task_loss,
        }
    
    def compute_feature_matching_loss(
        self,
        student_features: torch.Tensor,
        teacher_features: torch.Tensor,
    ) -> torch.Tensor:
        """
        Compute feature matching loss (MSE between intermediate features).
        
        Args:
            student_features: Student model intermediate features
            teacher_features: Teacher model intermediate features
        
        Returns:
            Feature matching loss
        """
        # Ensure same shape (may need projection)
        if student_features.shape != teacher_features.shape:
            # Use linear projection to match dimensions
            if student_features.dim() == 2:
                proj = nn.Linear(
                    student_features.shape[-1],
                    teacher_features.shape[-1],
                ).to(student_features.device)
                student_features = proj(student_features)
            else:
                # Handle other dimensions
                raise ValueError(
                    f"Cannot match feature shapes: {student_features.shape} vs {teacher_features.shape}"
                )
        
        return torch.nn.functional.mse_loss(student_features, teacher_features)

