"""Advanced training techniques from top ML researchers.

Implements:
1. SAM Optimizer (DeepMind/Google) - Sharpness-Aware Minimization
2. Curriculum Learning (Bengio) - Progressive difficulty training
3. Distributional Quality (Levine/Bellemare) - Predict quality distributions
4. Mixup (Zhang et al.) - Embedding interpolation augmentation
5. Self-Distillation (Hinton) - Online knowledge distillation
6. SWA (Izmailov) - Stochastic Weight Averaging
7. Consistency Regularization (Finn/Bachman) - Augmentation consistency

All techniques are configurable and default ON.
"""

from __future__ import annotations

import copy
import logging
import math
from typing import TYPE_CHECKING

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

if TYPE_CHECKING:
    from torch.optim import Optimizer

logger = logging.getLogger(__name__)


# =============================================================================
# 1. SAM Optimizer (Sharpness-Aware Minimization)
# =============================================================================

class SAM(torch.optim.Optimizer):
    """Sharpness-Aware Minimization optimizer.

    Finds parameters that lie in flat minima, improving generalization.
    Paper: https://arxiv.org/abs/2010.01412

    Usage:
        base_optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
        optimizer = SAM(model.parameters(), base_optimizer, rho=0.05)

        # Training step:
        loss = compute_loss()
        loss.backward()
        optimizer.first_step(zero_grad=True)

        loss_second = compute_loss()  # Re-compute with perturbed weights
        loss_second.backward()
        optimizer.second_step(zero_grad=True)
    """

    def __init__(
        self,
        params,
        base_optimizer: Optimizer,
        rho: float = 0.05,
        adaptive: bool = False,
    ):
        """Initialize SAM optimizer.

        Args:
            params: Model parameters to optimize.
            base_optimizer: Base optimizer (e.g., AdamW).
            rho: Neighborhood size for sharpness computation.
            adaptive: Use adaptive SAM (scales rho per parameter).
        """
        defaults = dict(rho=rho, adaptive=adaptive)
        super().__init__(params, defaults)

        self.base_optimizer = base_optimizer
        self.param_groups = self.base_optimizer.param_groups
        self.defaults.update(self.base_optimizer.defaults)

    @torch.no_grad()
    def first_step(self, zero_grad: bool = False):
        """Compute perturbation and apply to weights."""
        grad_norm = self._grad_norm()
        for group in self.param_groups:
            scale = group["rho"] / (grad_norm + 1e-12)

            for p in group["params"]:
                if p.grad is None:
                    continue
                self.state[p]["old_p"] = p.data.clone()
                e_w = (torch.pow(p, 2) if group["adaptive"] else 1.0) * p.grad * scale
                p.add_(e_w)  # Climb to sharpest direction

        if zero_grad:
            self.zero_grad()

    @torch.no_grad()
    def second_step(self, zero_grad: bool = False):
        """Restore original weights and apply base optimizer update."""
        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None:
                    continue
                p.data = self.state[p]["old_p"]  # Restore

        self.base_optimizer.step()  # Apply update

        if zero_grad:
            self.zero_grad()

    @torch.no_grad()
    def step(self, closure=None):
        """Standard step (for compatibility). Use first_step/second_step instead."""
        raise NotImplementedError("Use first_step() and second_step() for SAM")

    def _grad_norm(self):
        shared_device = self.param_groups[0]["params"][0].device
        norm = torch.norm(
            torch.stack([
                ((torch.abs(p) if group["adaptive"] else 1.0) * p.grad).norm(p=2).to(shared_device)
                for group in self.param_groups
                for p in group["params"]
                if p.grad is not None
            ]),
            p=2
        )
        return norm

    def load_state_dict(self, state_dict):
        super().load_state_dict(state_dict)
        self.base_optimizer.param_groups = self.param_groups


# =============================================================================
# 2. Curriculum Learning (Ground-Truth Based)
# =============================================================================

class CurriculumScheduler:
    """DEPRECATED: Standard curriculum is counterproductive for routing.

    In routing, "easy" queries (low perf_range) have NO learning signal
    because all LLMs perform similarly. "Hard" queries have strong signal.
    Training on easy first delays useful information.

    Use ImportanceWeighting instead.
    """

    def __init__(self, *args, **kwargs):
        logger.warning(
            "CurriculumScheduler is DEPRECATED for routing tasks. "
            "Use ImportanceWeighting instead - it weights hard queries higher."
        )


class ImportanceWeighting:
    """Weight loss by query difficulty (perf_range).

    Unlike curriculum learning, this approach:
    - Keeps ALL queries in training (no filtering)
    - Weights high perf_range queries MORE (strong signal)
    - Weights low perf_range queries LESS (weak signal)

    This is the correct approach for routing because:
    - High perf_range = LLM choice matters = strong learning signal
    - Low perf_range = LLM choice doesn't matter = weak/no signal
    """

    def __init__(
        self,
        weight_min: float = 0.5,
        weight_max: float = 2.0,
        normalize: bool = True,
    ):
        """Initialize importance weighting.

        Args:
            weight_min: Weight for lowest perf_range queries.
            weight_max: Weight for highest perf_range queries.
            normalize: Normalize weights to mean=1 within batch.
        """
        self.weight_min = weight_min
        self.weight_max = weight_max
        self.normalize = normalize

        logger.info(
            f"ImportanceWeighting: high-signal queries weighted {weight_max}x, "
            f"low-signal queries weighted {weight_min}x"
        )

    def compute_weights(self, perf_ranges: torch.Tensor) -> torch.Tensor:
        """Compute importance weights based on perf_range.

        Args:
            perf_ranges: (batch,) perf_range for each query.

        Returns:
            (batch,) importance weights.
        """
        # Normalize perf_ranges to [0, 1]
        if perf_ranges.max() > perf_ranges.min():
            normalized = (perf_ranges - perf_ranges.min()) / (perf_ranges.max() - perf_ranges.min())
        else:
            normalized = torch.ones_like(perf_ranges) * 0.5

        # Linear interpolation between weight_min and weight_max
        weights = self.weight_min + normalized * (self.weight_max - self.weight_min)

        # Optionally normalize to mean=1
        if self.normalize and weights.numel() > 0:
            weights = weights / weights.mean()

        return weights

    def apply_weights(
        self,
        loss_per_sample: torch.Tensor,
        perf_ranges: torch.Tensor,
    ) -> torch.Tensor:
        """Apply importance weights to per-sample losses.

        Args:
            loss_per_sample: (batch,) loss for each sample.
            perf_ranges: (batch,) perf_range for each query.

        Returns:
            Weighted mean loss.
        """
        weights = self.compute_weights(perf_ranges)
        weighted_loss = (loss_per_sample * weights).mean()
        return weighted_loss


# =============================================================================
# 3. Distributional Quality Prediction
# =============================================================================

class DistributionalHead(nn.Module):
    """Predict quality distribution, not just mean.

    Outputs quantiles [q10, q25, q50, q75, q90] for each edge,
    enabling risk-aware routing decisions.
    """

    def __init__(
        self,
        input_dim: int,
        quantiles: tuple[float, ...] = (0.1, 0.25, 0.5, 0.75, 0.9),
    ):
        super().__init__()
        self.quantiles = quantiles
        self.num_quantiles = len(quantiles)

        self.head = nn.Sequential(
            nn.Linear(input_dim, input_dim // 2),
            nn.ReLU(),
            nn.Linear(input_dim // 2, self.num_quantiles),
        )

        # Register quantiles as buffer
        self.register_buffer(
            "quantile_values",
            torch.tensor(quantiles, dtype=torch.float32)
        )

    def forward(self, edge_repr: torch.Tensor) -> torch.Tensor:
        """Predict quantiles for each edge.

        Args:
            edge_repr: (num_edges, input_dim) edge representations.

        Returns:
            (num_edges, num_quantiles) predicted quantiles.
        """
        return torch.sigmoid(self.head(edge_repr))  # Quantiles in [0, 1]

    def quantile_loss(
        self,
        predictions: torch.Tensor,
        targets: torch.Tensor,
    ) -> torch.Tensor:
        """Pinball loss for quantile regression.

        Args:
            predictions: (batch, num_quantiles) predicted quantiles.
            targets: (batch,) actual values.

        Returns:
            Scalar loss.
        """
        targets = targets.unsqueeze(1).expand_as(predictions)
        errors = targets - predictions

        # Pinball loss: asymmetric penalty based on quantile
        loss = torch.max(
            self.quantile_values * errors,
            (self.quantile_values - 1) * errors
        )
        return loss.mean()


# =============================================================================
# 4. Mixup for Embeddings
# =============================================================================

class EmbeddingMixup:
    """Mixup augmentation for query embeddings.

    Interpolates between pairs of embeddings and their labels,
    creating virtual training examples for better generalization.
    """

    def __init__(self, alpha: float = 0.2, prob: float = 0.5):
        """Initialize mixup.

        Args:
            alpha: Beta distribution parameter. Higher = more mixing.
            prob: Probability of applying mixup to a batch.
        """
        self.alpha = alpha
        self.prob = prob

    def __call__(
        self,
        embeddings: torch.Tensor,
        labels: torch.Tensor,
        edge_index: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, float]:
        """Apply mixup to embeddings and labels.

        Args:
            embeddings: (num_queries, embed_dim) query embeddings.
            labels: (num_queries, ...) labels (can be multi-dim for soft targets).
            edge_index: Optional edge index to permute consistently.

        Returns:
            Tuple of (mixed_embeddings, mixed_labels, lambda).
        """
        if np.random.random() > self.prob:
            return embeddings, labels, 1.0

        batch_size = embeddings.size(0)

        # Sample mixing coefficient
        if self.alpha > 0:
            lam = np.random.beta(self.alpha, self.alpha)
        else:
            lam = 1.0

        # Random permutation
        perm = torch.randperm(batch_size, device=embeddings.device)

        # Mix embeddings
        mixed_embeddings = lam * embeddings + (1 - lam) * embeddings[perm]

        # Mix labels
        if labels.dim() == 1:
            mixed_labels = lam * labels + (1 - lam) * labels[perm]
        else:
            mixed_labels = lam * labels + (1 - lam) * labels[perm]

        return mixed_embeddings, mixed_labels, lam


# =============================================================================
# 5. Self-Distillation (Online)
# =============================================================================

class OnlineDistillation:
    """Online self-distillation using EMA teacher.

    The EMA model acts as a slowly-updating teacher, and the student
    is trained to match its soft predictions.
    """

    def __init__(
        self,
        temperature: float = 3.0,
        distill_weight: float = 0.1,
    ):
        """Initialize online distillation.

        Args:
            temperature: Softmax temperature for soft targets.
            distill_weight: Weight for distillation loss.
        """
        self.temperature = temperature
        self.distill_weight = distill_weight

    def distillation_loss(
        self,
        student_logits: torch.Tensor,
        teacher_logits: torch.Tensor,
    ) -> torch.Tensor:
        """Compute KL divergence distillation loss.

        Args:
            student_logits: (batch, num_classes) student predictions.
            teacher_logits: (batch, num_classes) teacher predictions (detached).

        Returns:
            Scalar distillation loss.
        """
        student_probs = F.log_softmax(student_logits / self.temperature, dim=-1)
        teacher_probs = F.softmax(teacher_logits.detach() / self.temperature, dim=-1)

        loss = F.kl_div(student_probs, teacher_probs, reduction="batchmean")
        # Scale by T^2 as per Hinton et al.
        loss = loss * (self.temperature ** 2)

        return self.distill_weight * loss


# =============================================================================
# 6. Stochastic Weight Averaging (SWA)
# =============================================================================

class SWAWrapper:
    """Stochastic Weight Averaging for better generalization.

    Maintains a running average of model weights, which tends to
    find flatter minima with better generalization.
    """

    def __init__(
        self,
        model: nn.Module,
        swa_start: int = 50,
        swa_freq: int = 5,
        swa_lr: float | None = None,
    ):
        """Initialize SWA.

        Args:
            model: Model to average.
            swa_start: Epoch to start SWA.
            swa_freq: Frequency of weight averaging (epochs).
            swa_lr: Optional constant LR during SWA phase.
        """
        self.model = model
        self.swa_start = swa_start
        self.swa_freq = swa_freq
        self.swa_lr = swa_lr

        # Initialize SWA model
        self.swa_model = torch.optim.swa_utils.AveragedModel(model)
        self.swa_scheduler = None
        self.num_averaged = 0

    def update(self, epoch: int):
        """Update SWA model if appropriate.

        Args:
            epoch: Current epoch.
        """
        if epoch >= self.swa_start and (epoch - self.swa_start) % self.swa_freq == 0:
            self.swa_model.update_parameters(self.model)
            self.num_averaged += 1
            logger.debug(f"SWA update #{self.num_averaged} at epoch {epoch}")

    def get_swa_model(self) -> nn.Module:
        """Get the averaged model for evaluation."""
        return self.swa_model.module

    def finalize(self, train_loader=None):
        """Finalize SWA by updating batch norm statistics.

        Args:
            train_loader: DataLoader for BN update (optional).
        """
        if train_loader is not None and self.num_averaged > 0:
            torch.optim.swa_utils.update_bn(train_loader, self.swa_model)
        logger.info(f"SWA finalized with {self.num_averaged} averaged models")


# =============================================================================
# 7. Consistency Regularization
# =============================================================================

class ConsistencyRegularization:
    """Enforce consistent predictions under augmentation.

    Applies noise/dropout to inputs and penalizes prediction differences,
    encouraging robust representations.
    """

    def __init__(
        self,
        noise_std: float = 0.1,
        consistency_weight: float = 0.1,
        use_kl: bool = True,
    ):
        """Initialize consistency regularization.

        Args:
            noise_std: Standard deviation of Gaussian noise.
            consistency_weight: Weight for consistency loss.
            use_kl: Use KL divergence (True) or MSE (False).
        """
        self.noise_std = noise_std
        self.consistency_weight = consistency_weight
        self.use_kl = use_kl

    def augment(self, embeddings: torch.Tensor) -> torch.Tensor:
        """Apply augmentation to embeddings.

        Args:
            embeddings: (batch, embed_dim) embeddings.

        Returns:
            Augmented embeddings.
        """
        noise = torch.randn_like(embeddings) * self.noise_std
        return embeddings + noise

    def consistency_loss(
        self,
        pred_orig: torch.Tensor,
        pred_aug: torch.Tensor,
    ) -> torch.Tensor:
        """Compute consistency loss between original and augmented predictions.

        Args:
            pred_orig: Predictions on original input.
            pred_aug: Predictions on augmented input.

        Returns:
            Scalar consistency loss.
        """
        if self.use_kl:
            # KL divergence (asymmetric - anchor on original)
            p_orig = F.softmax(pred_orig.detach(), dim=-1)
            log_p_aug = F.log_softmax(pred_aug, dim=-1)
            loss = F.kl_div(log_p_aug, p_orig, reduction="batchmean")
        else:
            # MSE
            loss = F.mse_loss(pred_aug, pred_orig.detach())

        return self.consistency_weight * loss


# =============================================================================
# Config Helper
# =============================================================================

def get_advanced_training_config(config: dict) -> dict:
    """Extract advanced training config with defaults.

    All techniques default to ON.
    """
    return {
        # SAM
        "use_sam": config.get("use_sam", True),
        "sam_rho": config.get("sam_rho", 0.05),
        "sam_adaptive": config.get("sam_adaptive", False),

        # Curriculum
        "use_curriculum": config.get("use_curriculum", True),
        "curriculum_warmup_fraction": config.get("curriculum_warmup_fraction", 0.1),
        "curriculum_start_percentile": config.get("curriculum_start_percentile", 0.3),
        "curriculum_end_percentile": config.get("curriculum_end_percentile", 1.0),
        "curriculum_schedule": config.get("curriculum_schedule", "linear"),

        # Distributional
        "use_distributional": config.get("use_distributional", True),
        "distributional_quantiles": config.get("distributional_quantiles", (0.1, 0.25, 0.5, 0.75, 0.9)),
        "distributional_weight": config.get("distributional_weight", 0.1),

        # Mixup
        "use_mixup": config.get("use_mixup", True),
        "mixup_alpha": config.get("mixup_alpha", 0.2),
        "mixup_prob": config.get("mixup_prob", 0.5),

        # Self-distillation
        "use_self_distillation": config.get("use_self_distillation", True),
        "distillation_temperature": config.get("distillation_temperature", 3.0),
        "distillation_weight": config.get("distillation_weight", 0.1),

        # SWA
        "use_swa": config.get("use_swa", True),
        "swa_start_epoch": config.get("swa_start_epoch", 50),
        "swa_freq": config.get("swa_freq", 5),
        "swa_lr": config.get("swa_lr", None),

        # Consistency
        "use_consistency": config.get("use_consistency", True),
        "consistency_noise_std": config.get("consistency_noise_std", 0.1),
        "consistency_weight": config.get("consistency_weight", 0.1),
        "consistency_use_kl": config.get("consistency_use_kl", True),
    }


def log_advanced_config(adv_config: dict):
    """Log enabled advanced training techniques."""
    enabled = []
    if adv_config["use_sam"]:
        enabled.append(f"SAM(rho={adv_config['sam_rho']})")
    if adv_config["use_curriculum"]:
        enabled.append(f"Curriculum(start={adv_config['curriculum_start_percentile']:.0%})")
    if adv_config["use_distributional"]:
        enabled.append(f"Distributional(w={adv_config['distributional_weight']})")
    if adv_config["use_mixup"]:
        enabled.append(f"Mixup(α={adv_config['mixup_alpha']})")
    if adv_config["use_self_distillation"]:
        enabled.append(f"Distill(T={adv_config['distillation_temperature']})")
    if adv_config["use_swa"]:
        enabled.append(f"SWA(start={adv_config['swa_start_epoch']})")
    if adv_config["use_consistency"]:
        enabled.append(f"Consistency(σ={adv_config['consistency_noise_std']})")

    if enabled:
        logger.info(f"Advanced training: {', '.join(enabled)}")
    else:
        logger.info("Advanced training: None enabled")
