"""Specialized optimizers for LLM routing with subtle performance differences.

This module provides optimizers designed for the routing problem where:
- Performance differences between frontier LLMs are often small
- Much of the data is "easy" (most LLMs perform similarly)
- Some data is "hard" (only certain LLMs succeed)
- Risk of collapse to uniform predictions is high

Key optimizers and their properties:

1. **Muon** (Momentum with Orthogonalization)
   - Applies Newton-Schulz orthogonalization to momentum buffer
   - Prevents any single gradient direction from dominating
   - Implicit regularization toward orthogonal weight matrices
   - Best for: Preventing collapse, better conditioning

2. **SAM** (Sharpness-Aware Minimization)
   - Seeks parameters in flat loss regions (low sharpness)
   - Better generalization, especially on imbalanced data
   - Two forward passes per step (2x compute)
   - Best for: Hard queries where over-fitting to easy patterns hurts

3. **Lion** (Evolved Sign Momentum)
   - Uses sign of momentum for updates (more robust to outliers)
   - Lower memory than Adam (no second moment)
   - Often works with higher LR than Adam
   - Best for: Noisy gradients from small performance differences

4. **Prodigy** (D-Adaptation variant)
   - Auto-tunes learning rate during training
   - No LR search needed
   - Best for: When optimal LR is unknown

References:
- Muon: Jordan et al. (2024) "Muon: Momentum Orthogonalized Update"
- SAM: Foret et al. (2021) "Sharpness-Aware Minimization"
- Lion: Chen et al. (2023) "Symbolic Discovery of Optimization Algorithms"
- Prodigy: Mishchenko et al. (2023) "Prodigy: An Expeditiously Adaptive
  Parameter-Free Learner"
"""

from __future__ import annotations

import math
from typing import Callable, Iterable

import torch
from torch import Tensor
from torch.optim import Optimizer


def _newton_schulz_orthogonalize(
    matrix: Tensor,
    num_iterations: int = 5,
    eps: float = 1e-7,
) -> Tensor:
    """Apply Newton-Schulz iteration to orthogonalize a matrix.

    The Newton-Schulz iteration converges to the nearest orthogonal matrix:
        X_{k+1} = X_k (3I - X_k^T X_k) / 2

    This is equivalent to computing X @ (X^T X)^{-1/2} but without explicit
    matrix inversion, making it numerically stable.

    Args:
        matrix: Input matrix to orthogonalize (m, n) where m >= n.
        num_iterations: Number of Newton-Schulz iterations (5 is usually enough).
        eps: Small constant for numerical stability.

    Returns:
        Orthogonalized matrix with same shape as input.
    """
    # Normalize for numerical stability
    norm = matrix.norm()
    if norm < eps:
        return matrix
    X = matrix / (norm + eps)

    # Newton-Schulz iteration: X = X @ (3I - X^T X) / 2
    # For non-square matrices, we use X @ (3I - X^T @ X) / 2
    for _ in range(num_iterations):
        XtX = X.T @ X
        X = X @ (3.0 * torch.eye(XtX.size(0), device=X.device, dtype=X.dtype) - XtX) / 2.0

    return X * norm


class Muon(Optimizer):
    """Muon optimizer: Momentum with Orthogonalized Updates.

    Muon applies Newton-Schulz orthogonalization to the momentum buffer,
    projecting updates onto the Stiefel manifold (space of orthogonal matrices).
    This provides:

    1. Better conditioning: All singular values of the update are similar
    2. Implicit regularization: Weights tend toward orthogonal matrices
    3. Collapse prevention: No single direction dominates updates

    The orthogonalization is applied only to weight matrices (2D tensors) that
    are large enough to benefit. Biases and small tensors use standard momentum.

    For the routing problem, this helps because:
    - Small performance differences create small, potentially unbalanced gradients
    - Orthogonalization spreads learning across all output dimensions equally
    - Prevents collapse where model predicts uniform scores for all LLMs

    Args:
        params: Iterable of parameters to optimize.
        lr: Learning rate (default: 0.02, typically higher than Adam).
        momentum: Momentum factor (default: 0.95).
        weight_decay: Weight decay / L2 penalty (default: 0.01).
        ns_iterations: Newton-Schulz iterations for orthogonalization (default: 5).
        min_dim_for_orth: Minimum dimension to apply orthogonalization (default: 128).
            Smaller matrices don't benefit as much from orthogonalization.

    Example:
        >>> model = GATEncoderDecoderNet(...)
        >>> optimizer = Muon(model.parameters(), lr=0.02, momentum=0.95)
        >>> for batch in dataloader:
        ...     loss = model(batch)
        ...     loss.backward()
        ...     optimizer.step()
        ...     optimizer.zero_grad()
    """

    def __init__(
        self,
        params: Iterable[Tensor],
        lr: float = 0.02,
        momentum: float = 0.95,
        weight_decay: float = 0.01,
        ns_iterations: int = 5,
        min_dim_for_orth: int = 128,
    ) -> None:
        if lr < 0.0:
            raise ValueError(f"Invalid learning rate: {lr}")
        if momentum < 0.0 or momentum >= 1.0:
            raise ValueError(f"Invalid momentum: {momentum}")
        if weight_decay < 0.0:
            raise ValueError(f"Invalid weight_decay: {weight_decay}")

        defaults = {
            "lr": lr,
            "momentum": momentum,
            "weight_decay": weight_decay,
            "ns_iterations": ns_iterations,
            "min_dim_for_orth": min_dim_for_orth,
        }
        super().__init__(params, defaults)

    @torch.no_grad()
    def step(self, closure: Callable[[], float] | None = None) -> float | None:
        """Perform a single optimization step.

        Args:
            closure: A closure that reevaluates the model and returns the loss.

        Returns:
            Loss value if closure provided, else None.
        """
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            lr = group["lr"]
            momentum = group["momentum"]
            weight_decay = group["weight_decay"]
            ns_iterations = group["ns_iterations"]
            min_dim = group["min_dim_for_orth"]

            for p in group["params"]:
                if p.grad is None:
                    continue

                grad = p.grad
                if grad.is_sparse:
                    raise RuntimeError("Muon does not support sparse gradients")

                # Apply weight decay (decoupled, like AdamW)
                if weight_decay != 0:
                    p.mul_(1 - lr * weight_decay)

                # Get or initialize momentum buffer
                state = self.state[p]
                if "momentum_buffer" not in state:
                    state["momentum_buffer"] = torch.zeros_like(p)

                buf = state["momentum_buffer"]
                buf.mul_(momentum).add_(grad)

                # Decide whether to orthogonalize
                # Only for 2D weight matrices large enough to benefit
                should_orthogonalize = (
                    p.dim() == 2
                    and p.size(0) >= min_dim
                    and p.size(1) >= min_dim
                )

                if should_orthogonalize:
                    # Reshape if needed (for linear layers, weight is [out, in])
                    # Newton-Schulz expects tall matrix (m >= n)
                    if buf.size(0) < buf.size(1):
                        buf_t = buf.T
                        update = _newton_schulz_orthogonalize(buf_t, ns_iterations).T
                    else:
                        update = _newton_schulz_orthogonalize(buf, ns_iterations)
                else:
                    update = buf

                # Apply update
                p.add_(update, alpha=-lr)

        return loss


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

    SAM seeks parameters that lie in neighborhoods with uniformly low loss,
    rather than just low loss at a single point. This finds flatter minima
    which generalize better.

    The algorithm:
    1. Compute gradient at current parameters
    2. Move to adversarial point: w + rho * grad / ||grad||
    3. Compute gradient at adversarial point
    4. Update original parameters using adversarial gradient

    For the routing problem, SAM helps because:
    - Easy queries can create sharp minima that don't transfer to hard queries
    - Flat minima generalize better across query difficulty levels
    - Reduces over-confidence on easy patterns

    Trade-off: 2x forward/backward passes per step.

    Args:
        params: Iterable of parameters to optimize.
        base_optimizer: Base optimizer class (default: AdamW).
        rho: Neighborhood size for adversarial perturbation (default: 0.05).
        adaptive: Use adaptive SAM (scale rho by parameter norm).
        **kwargs: Arguments passed to base optimizer.

    Example:
        >>> model = GATEncoderDecoderNet(...)
        >>> optimizer = SAM(model.parameters(), base_optimizer=torch.optim.AdamW,
        ...                 lr=1e-3, weight_decay=0.01)
        >>> for batch in dataloader:
        ...     # First forward-backward (compute adversarial direction)
        ...     loss = model(batch)
        ...     loss.backward()
        ...     optimizer.first_step(zero_grad=True)
        ...     # Second forward-backward (at adversarial point)
        ...     loss = model(batch)
        ...     loss.backward()
        ...     optimizer.second_step(zero_grad=True)
    """

    def __init__(
        self,
        params: Iterable[Tensor],
        base_optimizer: type[Optimizer] = None,
        rho: float = 0.05,
        adaptive: bool = False,
        **kwargs,
    ) -> None:
        if base_optimizer is None:
            from torch.optim import AdamW
            base_optimizer = AdamW

        if rho < 0.0:
            raise ValueError(f"Invalid rho: {rho}")

        defaults = {"rho": rho, "adaptive": adaptive}
        super().__init__(params, defaults)

        # Create base optimizer with same param groups
        self.base_optimizer = base_optimizer(self.param_groups, **kwargs)
        self.param_groups = self.base_optimizer.param_groups

    @torch.no_grad()
    def first_step(self, zero_grad: bool = False) -> None:
        """First step: move to adversarial point."""
        grad_norm = self._grad_norm()
        for group in self.param_groups:
            rho = group["rho"]
            adaptive = group["adaptive"]
            scale = rho / (grad_norm + 1e-12)

            for p in group["params"]:
                if p.grad is None:
                    continue

                # Store original params
                self.state[p]["old_p"] = p.data.clone()

                # Compute epsilon (adversarial perturbation)
                if adaptive:
                    eps = (torch.abs(p) + 1e-12) * p.grad * scale
                else:
                    eps = p.grad * scale

                # Move to adversarial point
                p.add_(eps)

        if zero_grad:
            self.zero_grad()

    @torch.no_grad()
    def second_step(self, zero_grad: bool = False) -> None:
        """Second step: update from adversarial gradient."""
        # Restore original parameters
        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None:
                    continue
                p.data = self.state[p]["old_p"]

        # Apply base optimizer step with adversarial gradient
        self.base_optimizer.step()

        if zero_grad:
            self.zero_grad()

    @torch.no_grad()
    def step(self, closure: Callable[[], float] | None = None) -> float | None:
        """Standard step interface (requires closure for SAM)."""
        if closure is None:
            raise RuntimeError(
                "SAM requires closure. Use first_step() and second_step() "
                "for manual control, or provide a closure."
            )

        # First forward-backward
        with torch.enable_grad():
            loss = closure()
        self.first_step(zero_grad=True)

        # Second forward-backward at adversarial point
        with torch.enable_grad():
            closure()
        self.second_step(zero_grad=True)

        return loss

    def _grad_norm(self) -> Tensor:
        """Compute the gradient norm across all parameters."""
        shared_device = self.param_groups[0]["params"][0].device
        norm = torch.norm(
            torch.stack([
                (torch.abs(p) * p.grad if group["adaptive"] else 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: dict) -> None:
        """Load optimizer state."""
        super().load_state_dict(state_dict)
        self.base_optimizer.param_groups = self.param_groups


class Lion(Optimizer):
    """Lion optimizer: Evolved Sign Momentum.

    Lion was discovered through program search (AutoML). It uses only the
    sign of the momentum for updates, making it robust to gradient outliers.

    Key properties:
    - Memory efficient: No second moment (like SGD, unlike Adam)
    - Robust to outliers: Sign function ignores gradient magnitude
    - Often works with higher LR than Adam (10x is common)
    - Simpler than Adam with comparable or better results

    For the routing problem, Lion helps because:
    - Small performance differences create noisy, small gradients
    - Sign-based updates are robust to this noise
    - Lower memory allows larger batch sizes

    Args:
        params: Iterable of parameters to optimize.
        lr: Learning rate (default: 1e-4, try 3x-10x AdamW's LR).
        betas: Coefficients for momentum interpolation (default: (0.9, 0.99)).
        weight_decay: Weight decay / L2 penalty (default: 0.0).

    Example:
        >>> model = GATEncoderDecoderNet(...)
        >>> optimizer = Lion(model.parameters(), lr=3e-4, weight_decay=0.01)
    """

    def __init__(
        self,
        params: Iterable[Tensor],
        lr: float = 1e-4,
        betas: tuple[float, float] = (0.9, 0.99),
        weight_decay: float = 0.0,
    ) -> None:
        if lr < 0.0:
            raise ValueError(f"Invalid learning rate: {lr}")
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError(f"Invalid beta[0]: {betas[0]}")
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError(f"Invalid beta[1]: {betas[1]}")

        defaults = {"lr": lr, "betas": betas, "weight_decay": weight_decay}
        super().__init__(params, defaults)

    @torch.no_grad()
    def step(self, closure: Callable[[], float] | None = None) -> float | None:
        """Perform a single optimization step."""
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            lr = group["lr"]
            beta1, beta2 = group["betas"]
            weight_decay = group["weight_decay"]

            for p in group["params"]:
                if p.grad is None:
                    continue

                grad = p.grad
                if grad.is_sparse:
                    raise RuntimeError("Lion does not support sparse gradients")

                state = self.state[p]
                if "exp_avg" not in state:
                    state["exp_avg"] = torch.zeros_like(p)

                exp_avg = state["exp_avg"]

                # Weight decay (decoupled)
                if weight_decay != 0:
                    p.mul_(1 - lr * weight_decay)

                # Lion update: sign of interpolation between momentum and gradient
                update = exp_avg * beta1 + grad * (1 - beta1)
                p.add_(torch.sign(update), alpha=-lr)

                # Update momentum (different beta)
                exp_avg.mul_(beta2).add_(grad, alpha=1 - beta2)

        return loss


class Prodigy(Optimizer):
    """Prodigy optimizer: Parameter-free adaptive learning rate.

    Prodigy automatically estimates the optimal learning rate during training
    using D-Adaptation. It tracks the distance traveled in parameter space
    and adjusts the LR to maximize progress.

    Key properties:
    - No LR tuning needed (set lr=1.0 and let it adapt)
    - Estimates optimal LR from gradient statistics
    - Works well out of the box for many problems

    For the routing problem, Prodigy helps because:
    - Unknown optimal LR for subtle routing signals
    - Adapts to the loss landscape structure
    - Robust to misspecified initial LR

    Args:
        params: Iterable of parameters to optimize.
        lr: Learning rate multiplier (default: 1.0, the adaptation does the work).
        betas: Coefficients for Adam-style moments (default: (0.9, 0.999)).
        beta3: Coefficient for D estimate smoothing (default: None, uses beta2).
        weight_decay: Weight decay / L2 penalty (default: 0.0).
        d_coef: Coefficient for D estimate (default: 1.0).
        growth_rate: Maximum rate for D growth (default: inf).

    Example:
        >>> model = GATEncoderDecoderNet(...)
        >>> optimizer = Prodigy(model.parameters(), lr=1.0)  # lr=1.0 is fine!
    """

    def __init__(
        self,
        params: Iterable[Tensor],
        lr: float = 1.0,
        betas: tuple[float, float] = (0.9, 0.999),
        beta3: float | None = None,
        weight_decay: float = 0.0,
        d_coef: float = 1.0,
        growth_rate: float = float("inf"),
        eps: float = 1e-8,
    ) -> None:
        if lr < 0.0:
            raise ValueError(f"Invalid learning rate: {lr}")

        if beta3 is None:
            beta3 = betas[1] ** 0.5

        defaults = {
            "lr": lr,
            "betas": betas,
            "beta3": beta3,
            "weight_decay": weight_decay,
            "d_coef": d_coef,
            "growth_rate": growth_rate,
            "eps": eps,
            "d": 1e-6,  # Initial D estimate (will adapt)
            "k": 0,  # Step counter
        }
        super().__init__(params, defaults)

    @torch.no_grad()
    def step(self, closure: Callable[[], float] | None = None) -> float | None:
        """Perform a single optimization step with D-adaptation."""
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            lr = group["lr"]
            beta1, beta2 = group["betas"]
            beta3 = group["beta3"]
            weight_decay = group["weight_decay"]
            d_coef = group["d_coef"]
            growth_rate = group["growth_rate"]
            eps = group["eps"]
            d = group["d"]
            k = group["k"]

            # Accumulators for D estimation
            d_numerator = 0.0
            d_denominator = 0.0

            for p in group["params"]:
                if p.grad is None:
                    continue

                grad = p.grad
                if grad.is_sparse:
                    raise RuntimeError("Prodigy does not support sparse gradients")

                state = self.state[p]
                if "exp_avg" not in state:
                    state["exp_avg"] = torch.zeros_like(p)
                    state["exp_avg_sq"] = torch.zeros_like(p)
                    state["s"] = torch.zeros_like(p)  # Running sum of gradients
                    state["p0"] = p.clone()  # Initial parameters

                exp_avg = state["exp_avg"]
                exp_avg_sq = state["exp_avg_sq"]
                s = state["s"]
                p0 = state["p0"]

                # Update biased moments
                exp_avg.mul_(beta1).add_(grad, alpha=d * (1 - beta1))
                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=d * d * (1 - beta2))

                # Update gradient sum for D estimation
                s.add_(grad, alpha=d * lr)

                # D estimation numerators and denominators
                d_numerator += (d * lr * grad * (p - p0)).sum().item()
                d_denominator += (s * grad).sum().item()

                # Bias correction
                bias_correction1 = 1 - beta1 ** (k + 1)
                bias_correction2_sq = math.sqrt(1 - beta2 ** (k + 1))

                # Compute update
                denom = (exp_avg_sq.sqrt() / bias_correction2_sq).add_(eps)
                update = (exp_avg / bias_correction1) / denom

                # Weight decay
                if weight_decay != 0:
                    p.mul_(1 - lr * d * weight_decay)

                # Apply update
                p.add_(update, alpha=-lr)

            # Update D estimate
            if d_denominator > 0:
                d_new = d_coef * d_numerator / d_denominator
                # Clamp growth rate
                d_new = min(d_new, d * growth_rate)
                # Smooth update
                group["d"] = d * beta3 + d_new * (1 - beta3)

            group["k"] = k + 1

        return loss


def get_optimizer(
    name: str,
    params: Iterable[Tensor],
    lr: float,
    weight_decay: float = 0.01,
    **kwargs,
) -> Optimizer:
    """Factory function to create an optimizer by name.

    Args:
        name: Optimizer name (adamw, muon, sam, lion, prodigy).
        params: Model parameters to optimize.
        lr: Learning rate.
        weight_decay: Weight decay / L2 penalty.
        **kwargs: Additional optimizer-specific arguments.

    Returns:
        Configured optimizer instance.

    Raises:
        ValueError: If optimizer name is not recognized.

    Example:
        >>> optimizer = get_optimizer("muon", model.parameters(), lr=0.02)
    """
    name = name.lower()

    if name == "adamw":
        from torch.optim import AdamW
        betas = kwargs.pop("betas", (0.9, 0.999))
        return AdamW(params, lr=lr, weight_decay=weight_decay, betas=betas, **kwargs)

    elif name == "muon":
        momentum = kwargs.pop("momentum", 0.95)
        ns_iterations = kwargs.pop("ns_iterations", 5)
        min_dim = kwargs.pop("min_dim_for_orth", 128)
        return Muon(
            params,
            lr=lr,
            momentum=momentum,
            weight_decay=weight_decay,
            ns_iterations=ns_iterations,
            min_dim_for_orth=min_dim,
            **kwargs,
        )

    elif name == "sam":
        base_name = kwargs.pop("base_optimizer", "adamw")
        rho = kwargs.pop("rho", 0.05)
        adaptive = kwargs.pop("adaptive", False)
        if base_name == "adamw":
            from torch.optim import AdamW
            base_optimizer = AdamW
        elif base_name == "sgd":
            from torch.optim import SGD
            base_optimizer = SGD
        else:
            raise ValueError(f"Unknown base optimizer for SAM: {base_name}")
        return SAM(
            params,
            base_optimizer=base_optimizer,
            rho=rho,
            adaptive=adaptive,
            lr=lr,
            weight_decay=weight_decay,
            **kwargs,
        )

    elif name == "lion":
        betas = kwargs.pop("betas", (0.9, 0.99))
        return Lion(params, lr=lr, betas=betas, weight_decay=weight_decay, **kwargs)

    elif name == "prodigy":
        betas = kwargs.pop("betas", (0.9, 0.999))
        d_coef = kwargs.pop("d_coef", 1.0)
        return Prodigy(
            params,
            lr=lr,
            betas=betas,
            weight_decay=weight_decay,
            d_coef=d_coef,
            **kwargs,
        )

    else:
        raise ValueError(
            f"Unknown optimizer: {name}. "
            f"Supported: adamw, muon, sam, lion, prodigy"
        )


# Optimizer recommendations for the routing problem
OPTIMIZER_RECOMMENDATIONS = """
## Optimizer Selection Guide for LLM Routing

### Problem Characteristics
- Small performance differences between frontier LLMs
- Imbalanced difficulty: many "easy" queries, few "hard" queries
- Risk of collapse to uniform predictions

### Recommendations by Scenario

**Baseline (start here)**:
- AdamW with lr=1e-3, weight_decay=0.01
- Well-understood, battle-tested

**If model collapses to uniform predictions**:
- Try Muon: orthogonalization prevents direction dominance
- lr=0.02 (higher than Adam), momentum=0.95
- Or add uniformity loss to prevent embedding collapse

**If overfitting to easy queries**:
- Try SAM: finds flat minima that generalize to hard queries
- rho=0.05 (adversarial neighborhood size)
- 2x compute but better generalization

**If gradients are noisy/small**:
- Try Lion: sign-based updates ignore gradient magnitude
- lr=3e-4 to 1e-3 (3-10x AdamW's LR)
- Lower memory allows larger batches

**If you don't want to tune LR**:
- Try Prodigy: auto-adapts learning rate
- lr=1.0 (it will adapt)
- Good for exploration/prototyping

### Combining Techniques

You can combine optimizer choice with architectural fixes:
1. Muon + contrastive loss: best for collapse prevention
2. SAM + focal loss: best for easy/hard imbalance
3. Lion + larger batches: best for noisy signals
"""
