"""
Training & Knowledge Distillation Pipelines for Q-TensorFormer (`src/training.py`).

Provides:
  - `LanguageModelTrainer`: Deterministic causal LM trainer with AdamW, cosine LR schedule,
    gradient clipping, validation tracking, and full training history logging.
  - `DistillationTrainer`: Teacher-student knowledge distillation trainer for `scripts/distill.py`.
"""

import math
import time
from typing import Any, Dict, List, Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader

from .config import TrainingConfig
from .metrics import evaluate_model


class LanguageModelTrainer:
    """
    Canonical causal language modeling trainer for Q-TensorFormer and DenseBaseline.
    """

    def __init__(
        self,
        model: nn.Module,
        train_config: Optional[TrainingConfig] = None,
        device: str = "cpu",
        pad_token_id: int = 0,
    ):
        self.model = model.to(device)
        self.cfg = train_config or TrainingConfig()
        self.device = device
        self.pad_token_id = pad_token_id

        self.optimizer = torch.optim.AdamW(
            self.model.parameters(),
            lr=self.cfg.learning_rate,
            weight_decay=self.cfg.weight_decay,
        )

    def _get_lr(self, step: int, total_steps: int) -> float:
        base_lr = self.cfg.learning_rate
        min_lr = base_lr * self.cfg.lr_min_factor
        warmup = min(self.cfg.warmup_steps, max(1, total_steps // 5))
        if step < warmup:
            return base_lr * float(step + 1) / float(max(1, warmup))
        if self.cfg.lr_scheduler == "cosine":
            progress = float(step - warmup) / float(max(1, total_steps - warmup))
            progress = min(1.0, max(0.0, progress))
            return min_lr + 0.5 * (base_lr - min_lr) * (1.0 + math.cos(math.pi * progress))
        if self.cfg.lr_scheduler == "linear":
            progress = float(step - warmup) / float(max(1, total_steps - warmup))
            progress = min(1.0, max(0.0, progress))
            return base_lr - progress * (base_lr - min_lr)
        return base_lr

    def train(
        self,
        train_loader: DataLoader,
        val_loader: Optional[DataLoader] = None,
        verbose: bool = True,
    ) -> Dict[str, Any]:
        torch.manual_seed(self.cfg.seed)
        total_steps = max(1, self.cfg.max_epochs * len(train_loader))

        initial_val = (
            evaluate_model(self.model, val_loader, device=self.device, pad_token_id=self.pad_token_id)
            if val_loader is not None
            else evaluate_model(self.model, train_loader, device=self.device, pad_token_id=self.pad_token_id)
        )

        epoch_history: List[Dict[str, Any]] = []
        global_step = 0
        t_start = time.perf_counter()

        for epoch in range(1, self.cfg.max_epochs + 1):
            self.model.train()
            epoch_nll = 0.0
            epoch_tokens = 0
            t_ep = time.perf_counter()

            self.optimizer.zero_grad()
            for b_idx, batch in enumerate(train_loader):
                if len(batch) == 3:
                    x, y, attn_mask = batch
                    attn_mask = attn_mask.to(self.device)
                else:
                    x, y = batch
                    attn_mask = (x != self.pad_token_id).long().to(self.device)

                x = x.to(self.device)
                y = y.to(self.device)

                lr = self._get_lr(global_step, total_steps)
                for pg in self.optimizer.param_groups:
                    pg["lr"] = lr

                try:
                    out = self.model(x, attention_mask=attn_mask)
                except TypeError:
                    out = self.model(x)

                logits = out[0] if isinstance(out, tuple) else (out.logits if hasattr(out, "logits") else out)
                vocab_size = logits.size(-1)

                valid_mask = y != self.pad_token_id
                n_valid = int(valid_mask.sum().item())
                if n_valid == 0:
                    continue

                loss = F.cross_entropy(
                    logits.reshape(-1, vocab_size),
                    y.reshape(-1),
                    ignore_index=self.pad_token_id,
                    reduction="mean",
                )

                scaled_loss = loss / max(1, self.cfg.gradient_accumulation_steps)
                scaled_loss.backward()

                if (b_idx + 1) % max(1, self.cfg.gradient_accumulation_steps) == 0 or (b_idx + 1) == len(train_loader):
                    if self.cfg.max_grad_norm > 0:
                        torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.cfg.max_grad_norm)
                    self.optimizer.step()
                    self.optimizer.zero_grad()

                epoch_nll += float(loss.item()) * n_valid
                epoch_tokens += n_valid
                global_step += 1

            train_loss = epoch_nll / max(1, epoch_tokens)
            train_ppl = math.exp(min(train_loss, 20.0))
            val_metrics = (
                evaluate_model(self.model, val_loader, device=self.device, pad_token_id=self.pad_token_id)
                if val_loader is not None
                else {"loss": train_loss, "perplexity": train_ppl, "accuracy": 0.0}
            )
            ep_time = time.perf_counter() - t_ep

            rec = {
                "epoch": epoch,
                "train_loss": round(train_loss, 6),
                "train_perplexity": round(train_ppl, 4),
                "val_loss": val_metrics["loss"],
                "val_perplexity": val_metrics["perplexity"],
                "val_accuracy": val_metrics["accuracy"],
                "lr": round(self.optimizer.param_groups[0]["lr"], 8),
                "epoch_time_s": round(ep_time, 3),
            }
            epoch_history.append(rec)
            if verbose:
                print(
                    f"  Epoch {epoch:02d}/{self.cfg.max_epochs:02d} | "
                    f"train_loss={rec['train_loss']:.4f} (ppl={rec['train_perplexity']:.2f}) | "
                    f"val_loss={rec['val_loss']:.4f} (ppl={rec['val_perplexity']:.2f}, acc={rec['val_accuracy']:.3f}) | "
                    f"{rec['epoch_time_s']:.2f}s"
                )

        total_time = time.perf_counter() - t_start
        final_val = (
            evaluate_model(self.model, val_loader, device=self.device, pad_token_id=self.pad_token_id)
            if val_loader is not None
            else evaluate_model(self.model, train_loader, device=self.device, pad_token_id=self.pad_token_id)
        )

        return {
            "initial_val": initial_val,
            "final_val": final_val,
            "epochs": epoch_history,
            "total_steps": global_step,
            "total_training_time_s": round(total_time, 3),
            "loss_improved": final_val["loss"] < initial_val["loss"],
            "perplexity_reduction_ratio": round(
                initial_val["perplexity"] / max(1e-6, final_val["perplexity"]), 4
            ),
        }


class DistillationTrainer:
    """
    Knowledge distillation trainer for `scripts/distill.py`.
    Minimizes `alpha * KL(student || teacher) * T^2 + (1 - alpha) * CE(student, y)`.
    """

    def __init__(
        self,
        student: nn.Module,
        teacher: nn.Module,
        train_config: Optional[TrainingConfig] = None,
        config: Optional[Any] = None,
        train_loader: Optional[DataLoader] = None,
        val_loader: Optional[DataLoader] = None,
        test_loader: Optional[DataLoader] = None,
        output_dir: Optional[str] = None,
        temperature: float = 2.0,
        alpha: float = 0.7,
        device: str = "cpu",
        pad_token_id: int = 0,
        **kwargs,
    ):
        self.student = student.to(device)
        self.teacher = teacher.to(device)
        self.teacher.eval()
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.test_loader = test_loader
        self.output_dir = output_dir

        if train_config is not None:
            self.cfg = train_config
        elif config is not None and hasattr(config, "training"):
            self.cfg = config.training
        elif config is not None and isinstance(config, TrainingConfig):
            self.cfg = config
        else:
            self.cfg = TrainingConfig()

        self.temperature = float(temperature)
        self.alpha = float(alpha)
        self.device = device
        self.pad_token_id = pad_token_id

        self.optimizer = torch.optim.AdamW(
            self.student.parameters(),
            lr=self.cfg.learning_rate,
            weight_decay=self.cfg.weight_decay,
        )

    def train(
        self,
        train_loader: Optional[DataLoader] = None,
        val_loader: Optional[DataLoader] = None,
    ) -> Dict[str, Any]:
        tl = train_loader if train_loader is not None else self.train_loader
        vl = val_loader if val_loader is not None else self.val_loader
        if tl is None:
            raise ValueError("train_loader must be provided either in __init__ or train().")

        torch.manual_seed(self.cfg.seed)
        history: List[Dict[str, Any]] = []
        for epoch in range(1, self.cfg.max_epochs + 1):
            self.student.train()
            total_loss = 0.0
            n_batches = 0
            for batch in tl:
                x, y = batch[0].to(self.device), batch[1].to(self.device)
                self.optimizer.zero_grad()

                with torch.no_grad():
                    t_out = self.teacher(x)
                    t_logits = t_out[0] if isinstance(t_out, tuple) else t_out

                s_out = self.student(x)
                s_logits = s_out[0] if isinstance(s_out, tuple) else s_out

                vocab_size = s_logits.size(-1)
                ce_loss = F.cross_entropy(
                    s_logits.reshape(-1, vocab_size),
                    y.reshape(-1),
                    ignore_index=self.pad_token_id,
                )
                T = self.temperature
                kd_loss = F.kl_div(
                    F.log_softmax(s_logits / T, dim=-1),
                    F.softmax(t_logits / T, dim=-1),
                    reduction="batchmean",
                ) * (T * T)

                loss = self.alpha * kd_loss + (1.0 - self.alpha) * ce_loss
                loss.backward()
                if self.cfg.max_grad_norm > 0:
                    torch.nn.utils.clip_grad_norm_(self.student.parameters(), self.cfg.max_grad_norm)
                self.optimizer.step()

                total_loss += float(loss.item())
                n_batches += 1

            val_metrics = (
                evaluate_model(self.student, vl, device=self.device, pad_token_id=self.pad_token_id)
                if vl is not None
                else {"loss": total_loss / max(1, n_batches), "perplexity": 0.0}
            )
            history.append({
                "epoch": epoch,
                "train_loss": round(total_loss / max(1, n_batches), 4),
                "val_loss": val_metrics["loss"],
                "val_perplexity": val_metrics["perplexity"],
            })
        return {"epochs": history}