"""Knowledge Distillation trainer for transferring knowledge from teacher to student.

Supports multiple teacher sources via TeacherWrapper abstraction:
- HuggingFace models
- Local composennent models
- Custom models via wrapper

Based on: "Distilling the Knowledge in a Neural Network" (Hinton et al., 2015)
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from typing import Optional, Union, List, Callable, Protocol, Any
from abc import ABC, abstractmethod
from tqdm.auto import tqdm
import time


class TeacherModel(ABC):
    """Abstract interface for teacher models (Dependency Injection pattern).
    
    Implement this interface to use any model as a teacher:
    - HuggingFace transformers
    - Local composennent models
    - API-based models (OpenAI, etc.)
    
    Example:
        >>> class MyTeacher(TeacherModel):
        ...     def get_logits(self, input_ids, attention_mask):
        ...         return self.model(input_ids, attention_mask).logits
    """
    
    @abstractmethod
    def get_logits(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Get teacher's output logits.
        
        Args:
            input_ids: Token IDs (batch, seq_len).
            attention_mask: Attention mask (batch, seq_len).
            
        Returns:
            Logits tensor (batch, seq_len, vocab_size).
        """
        pass
    
    @abstractmethod
    def to(self, device: str) -> "TeacherModel":
        """Move teacher to device."""
        pass
    
    @abstractmethod
    def eval(self) -> "TeacherModel":
        """Set teacher to evaluation mode."""
        pass


class HuggingFaceTeacher(TeacherModel):
    """Wrapper for HuggingFace models as teacher.
    
    Example:
        >>> from transformers import AutoModelForCausalLM
        >>> hf_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-1.5B")
        >>> teacher = HuggingFaceTeacher(hf_model)
    """
    
    def __init__(self, model):
        """Initialize with a HuggingFace model.
        
        Args:
            model: HuggingFace model with forward() returning logits.
        """
        self.model = model
        self._device = "cpu"
    
    @classmethod
    def from_pretrained(cls, model_name: str, **kwargs) -> "HuggingFaceTeacher":
        """Load a HuggingFace model by name.
        
        Args:
            model_name: HuggingFace model identifier (e.g., "Qwen/Qwen2.5-1.5B").
            **kwargs: Additional args for from_pretrained().
            
        Returns:
            HuggingFaceTeacher instance.
        """
        try:
            from transformers import AutoModelForCausalLM
        except ImportError:
            raise ImportError(
                "HuggingFace transformers required. Install with: pip install transformers"
            )
        
        model = AutoModelForCausalLM.from_pretrained(model_name, **kwargs)
        return cls(model)
    
    def get_logits(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        with torch.no_grad():
            outputs = self.model(
                input_ids=input_ids,
                attention_mask=attention_mask,
            )
            # Handle different output formats
            if hasattr(outputs, "logits"):
                return outputs.logits
            elif isinstance(outputs, tuple):
                return outputs[0]
            else:
                return outputs
    
    def to(self, device: str) -> "HuggingFaceTeacher":
        self.model = self.model.to(device)
        self._device = device
        return self
    
    def eval(self) -> "HuggingFaceTeacher":
        self.model.eval()
        return self


class ComposennentTeacher(TeacherModel):
    """Wrapper for composennent models as teacher.
    
    Example:
        >>> from composennent.nlp.transformers import GPT
        >>> model = GPT.load("large_model.pt")
        >>> teacher = ComposennentTeacher(model)
    """
    
    def __init__(self, model: nn.Module):
        """Initialize with a composennent model.
        
        Args:
            model: Any composennent model (GPT, BERT, etc.).
        """
        self.model = model
        self._device = "cpu"
    
    @classmethod
    def from_checkpoint(cls, model_class, path: str, device: str = "cpu", **kwargs) -> "ComposennentTeacher":
        """Load from a checkpoint.
        
        Args:
            model_class: The model class (e.g., GPT).
            path: Path to checkpoint.
            device: Device to load on.
            **kwargs: Additional model kwargs.
            
        Returns:
            ComposennentTeacher instance.
        """
        model = model_class.load(path, device=device, **kwargs)
        teacher = cls(model)
        teacher._device = device
        return teacher
    
    def get_logits(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        with torch.no_grad():
            outputs = self.model(
                input_ids=input_ids,
                attention_mask=attention_mask,
            )
            # Handle different output formats
            if isinstance(outputs, dict):
                return outputs.get("logits", outputs.get("output"))
            elif isinstance(outputs, tuple):
                return outputs[0]
            else:
                return outputs
    
    def to(self, device: str) -> "ComposennentTeacher":
        self.model = self.model.to(device)
        self._device = device
        return self
    
    def eval(self) -> "ComposennentTeacher":
        self.model.eval()
        return self


class DistillationTrainer:
    """Trainer for knowledge distillation.
    
    Trains a smaller student model to mimic a larger teacher model's outputs.
    Uses a combination of:
    - Soft targets: KL divergence between temperature-scaled distributions
    - Hard targets: Cross-entropy with ground truth labels
    
    Args:
        student: Student model to train.
        teacher: Teacher model (frozen).
        tokenizer: Tokenizer for the student model.
        temperature: Softmax temperature for distillation. Higher = softer. Defaults to 2.0.
        alpha: Balance between soft (0) and hard (1) loss. Defaults to 0.5.
        lr: Learning rate. Defaults to 1e-4.
        device: Training device. Defaults to "cuda".
        
    Example:
        >>> teacher = HuggingFaceTeacher.from_pretrained("Qwen/Qwen2.5-1.5B")
        >>> student = GPT(vocab_size=32000, latent_dim=512, ...)
        >>> trainer = DistillationTrainer(student, teacher, tokenizer)
        >>> trainer.train(dataloader, epochs=3)
    """
    
    def __init__(
        self,
        student: nn.Module,
        teacher: TeacherModel,
        tokenizer,
        temperature: float = 2.0,
        alpha: float = 0.5,
        lr: float = 1e-4,
        device: str = "cuda",
        use_amp: bool = True,
    ):
        self.student = student.to(device)
        self.teacher = teacher.to(device).eval()
        self.tokenizer = tokenizer
        self.temperature = temperature
        self.alpha = alpha
        self.device = device
        self.use_amp = use_amp
        
        # Freeze teacher
        if hasattr(self.teacher, 'model'):
            for param in self.teacher.model.parameters():
                param.requires_grad = False
        
        self.optimizer = torch.optim.AdamW(self.student.parameters(), lr=lr)
        self.scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
    
    def distillation_loss(
        self,
        student_logits: torch.Tensor,
        teacher_logits: torch.Tensor,
        labels: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Compute combined distillation loss.
        
        Args:
            student_logits: Student output (batch, seq, vocab).
            teacher_logits: Teacher output (batch, seq, vocab).
            labels: Ground truth labels (batch, seq). Optional.
            
        Returns:
            Combined loss scalar.
        """
        T = self.temperature
        
        # Soft targets - KL divergence with temperature
        soft_student = F.log_softmax(student_logits / T, dim=-1)
        soft_teacher = F.softmax(teacher_logits / T, dim=-1)
        
        soft_loss = F.kl_div(
            soft_student,
            soft_teacher,
            reduction="batchmean",
        ) * (T * T)  # Scale by T^2 as per Hinton
        
        if labels is not None and self.alpha > 0:
            # Hard targets - standard cross-entropy
            # Shift for next-token prediction
            shift_logits = student_logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            
            hard_loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=-100,
            )
            
            # Combined loss
            loss = (1 - self.alpha) * soft_loss + self.alpha * hard_loss
        else:
            loss = soft_loss
        
        return loss
    
    def train_step(self, batch: dict) -> float:
        """Single training step.
        
        Args:
            batch: Dict with input_ids, attention_mask, and optionally labels.
            
        Returns:
            Loss value.
        """
        input_ids = batch["input_ids"].to(self.device)
        attention_mask = batch.get("attention_mask")
        if attention_mask is not None:
            attention_mask = attention_mask.to(self.device)
        labels = batch.get("labels")
        if labels is not None:
            labels = labels.to(self.device)
        
        self.optimizer.zero_grad()
        
        with torch.cuda.amp.autocast(enabled=self.use_amp):
            # Get teacher logits (no grad)
            teacher_logits = self.teacher.get_logits(input_ids, attention_mask)
            
            # Get student logits
            student_output = self.student(input_ids, attention_mask=attention_mask)
            if isinstance(student_output, dict):
                student_logits = student_output.get("logits", student_output.get("output"))
            elif isinstance(student_output, tuple):
                student_logits = student_output[0]
            else:
                student_logits = student_output
            
            # Compute loss
            loss = self.distillation_loss(student_logits, teacher_logits, labels)
        
        self.scaler.scale(loss).backward()
        self.scaler.step(self.optimizer)
        self.scaler.update()
        
        return loss.item()
    
    def train(
        self,
        dataloader: DataLoader,
        epochs: int = 3,
        log_interval: int = 10,
        save_path: Optional[str] = None,
        callbacks: Optional[List] = None,
    ):
        """Main training loop.
        
        Args:
            dataloader: Training data.
            epochs: Number of epochs.
            log_interval: Log every N steps.
            save_path: Path to save checkpoints.
            callbacks: Optional list of callbacks.
        """
        callbacks = callbacks or []
        
        self.student.train()
        print(f"Starting distillation for {epochs} epochs...")
        print(f"Temperature: {self.temperature}, Alpha: {self.alpha}")
        start_time = time.time()
        
        for callback in callbacks:
            callback.on_train_begin()
        
        for epoch in range(epochs):
            pbar = tqdm(dataloader, desc=f"Epoch {epoch+1}/{epochs}")
            epoch_loss = 0.0
            
            for step, batch in enumerate(pbar):
                loss = self.train_step(batch)
                epoch_loss += loss
                
                if step % log_interval == 0:
                    pbar.set_postfix({"loss": f"{loss:.4f}"})
            
            avg_loss = epoch_loss / len(dataloader)
            print(f"Epoch {epoch+1} - Avg Loss: {avg_loss:.4f}")
            
            for callback in callbacks:
                callback.on_epoch_end(epoch, {"avg_loss": avg_loss})
            
            if save_path and hasattr(self.student, "save"):
                self.student.save(f"{save_path}_ep{epoch+1}.pt")
        
        for callback in callbacks:
            callback.on_train_end()
        
        total_time = time.time() - start_time
        print(f"Distillation finished in {total_time/60:.2f} minutes.")


def create_teacher(
    source: str,
    model_name_or_path: str,
    device: str = "cuda",
    **kwargs,
) -> TeacherModel:
    """Factory function to create a teacher from various sources.
    
    Args:
        source: "huggingface", "composennent", or "local".
        model_name_or_path: Model identifier or path.
        device: Device to load on.
        **kwargs: Additional loading args.
        
    Returns:
        TeacherModel instance.
        
    Example:
        >>> teacher = create_teacher("huggingface", "Qwen/Qwen2.5-1.5B")
        >>> teacher = create_teacher("composennent", "my_model.pt", model_class=GPT)
    """
    if source == "huggingface":
        return HuggingFaceTeacher.from_pretrained(model_name_or_path, **kwargs).to(device)
    elif source == "composennent":
        model_class = kwargs.pop("model_class")
        return ComposennentTeacher.from_checkpoint(
            model_class, model_name_or_path, device=device, **kwargs
        )
    elif source == "local":
        # Assume it's a torch checkpoint with model state dict
        model_class = kwargs.pop("model_class")
        model = model_class(**kwargs.pop("model_config", {}))
        model.load_state_dict(torch.load(model_name_or_path, map_location=device))
        return ComposennentTeacher(model).to(device)
    else:
        raise ValueError(f"Unknown source: {source}. Use 'huggingface', 'composennent', or 'local'.")
