"""Optimizer + LR schedule construction.

Two details that materially affect quality and that most tutorials get wrong:
  1. No weight decay on 1-D params (norms, biases) or the embedding table.
     Decaying norm/embedding weights quietly hurts.
  2. The LR schedule includes linear warmup; starting at peak LR on random
     weights diverges.
"""

from __future__ import annotations

import math

import torch


def _adamw(groups, lr, betas, eps):
    """AdamW with fused kernel when available, falling back if unsupported."""
    fused = torch.cuda.is_available()
    try:
        return torch.optim.AdamW(groups, lr=lr, betas=betas, eps=eps, fused=fused)
    except (RuntimeError, TypeError):
        return torch.optim.AdamW(groups, lr=lr, betas=betas, eps=eps)


def build_adamw(model, lr=3e-4, weight_decay=0.1,
                betas=(0.9, 0.95), eps=1e-8) -> torch.optim.AdamW:
    decay, no_decay = [], []
    for name, p in model.named_parameters():
        if not p.requires_grad:
            continue
        # 2-D matmul weights get decay; norms/biases (1-D) and embeddings don't.
        if p.ndim < 2 or "embed" in name:
            no_decay.append(p)
        else:
            decay.append(p)
    groups = [
        {"params": decay, "weight_decay": weight_decay},
        {"params": no_decay, "weight_decay": 0.0},
    ]
    return _adamw(groups, lr, betas, eps)


@torch.no_grad()
def zeropower_via_newtonschulz5(G, steps=5, eps=1e-7):
    """Orthogonalize a 2-D gradient via the quintic Newton-Schulz iteration
    (Keller Jordan). Pushes all singular values toward 1 in ~5 matmuls, so the
    update weights every direction equally instead of being dominated by large
    singular directions. Runs in bf16, as in the reference implementation.
    """
    assert G.ndim == 2
    a, b, c = 3.4445, -4.7750, 2.0315
    X = G.bfloat16()
    X = X / (X.norm() + eps)
    transposed = G.size(0) > G.size(1)
    if transposed:
        X = X.T
    for _ in range(steps):
        A = X @ X.T
        B = b * A + c * (A @ A)
        X = a * X + B @ X
    if transposed:
        X = X.T
    return X.to(G.dtype)


class Muon(torch.optim.Optimizer):
    """Muon: momentum + orthogonalized update for 2-D parameters only.

    Use ONLY for interior matrices (attn/MLP weights). Embeddings, norms, biases
    must stay on AdamW (see build_optimizer). Muon's natural LR is ~0.02, much
    higher than Adam's, because the orthogonalized update has ~unit scale.
    """

    def __init__(self, params, lr=0.02, momentum=0.95, nesterov=True, ns_steps=5):
        super().__init__(params, dict(lr=lr, momentum=momentum,
                                      nesterov=nesterov, ns_steps=ns_steps))

    @torch.no_grad()
    def step(self):
        for group in self.param_groups:
            lr, mom, nesterov = group["lr"], group["momentum"], group["nesterov"]
            for p in group["params"]:
                if p.grad is None:
                    continue
                g = p.grad
                state = self.state[p]
                if "m" not in state:
                    state["m"] = torch.zeros_like(g)
                buf = state["m"]
                buf.mul_(mom).add_(g)
                g = g.add(buf, alpha=mom) if nesterov else buf
                g = zeropower_via_newtonschulz5(g, steps=group["ns_steps"])
                # scale so the update RMS matches across differently-shaped matrices
                scale = max(1.0, p.size(0) / p.size(1)) ** 0.5
                p.add_(g, alpha=-lr * scale)


class HybridOptimizer:
    """Presents several optimizers as one (unified step/zero_grad/state_dict and
    a concatenated param_groups so a single LR scheduler drives all of them)."""

    def __init__(self, optimizers):
        self.optimizers = optimizers

    @property
    def param_groups(self):
        return [g for o in self.optimizers for g in o.param_groups]

    def zero_grad(self, set_to_none=True):
        for o in self.optimizers:
            o.zero_grad(set_to_none=set_to_none)

    def step(self):
        for o in self.optimizers:
            o.step()

    def state_dict(self):
        return {"opts": [o.state_dict() for o in self.optimizers]}

    def load_state_dict(self, sd):
        for o, s in zip(self.optimizers, sd["opts"]):
            o.load_state_dict(s)


def build_optimizer(model, name="adamw", lr=3e-4, weight_decay=0.1,
                    betas=(0.9, 0.95), eps=1e-8, muon_lr=0.02):
    """Dispatch: 'adamw' (default) or 'muon' (Muon on 2-D interior weights +
    AdamW on embeddings/norms/biases)."""
    if name == "adamw":
        return build_adamw(model, lr, weight_decay, betas, eps)
    if name != "muon":
        raise ValueError(f"unknown optimizer: {name}")

    muon_p, adamw_decay, adamw_nodecay = [], [], []
    for n, p in model.named_parameters():
        if not p.requires_grad:
            continue
        if p.ndim == 2 and "embed" not in n:     # interior matrices -> Muon
            muon_p.append(p)
        elif p.ndim < 2 or "embed" in n:          # norms/biases/embeddings -> AdamW
            adamw_nodecay.append(p)
        else:
            adamw_decay.append(p)
    adamw = _adamw(
        [{"params": adamw_decay, "weight_decay": weight_decay},
         {"params": adamw_nodecay, "weight_decay": 0.0}],
        lr, betas, eps)
    return HybridOptimizer([Muon(muon_p, lr=muon_lr), adamw])


class WarmupCosine:
    """Linear warmup then cosine decay to min_lr_ratio * peak.

    Works on any object exposing `param_groups` (torch optimizers AND our
    HybridOptimizer, which torch's LambdaLR rejects). Applies one multiplier to
    every group, scaling each group's own base lr (so Muon's 0.02 and AdamW's
    3e-4 warm up/decay together). state_dict captures the step for exact resume.
    """

    def __init__(self, optimizer, warmup_steps, total_steps, min_lr_ratio=0.1):
        self.opt = optimizer
        self.warmup_steps = warmup_steps
        self.total_steps = total_steps
        self.min_lr_ratio = min_lr_ratio
        self.base_lrs = [g["lr"] for g in optimizer.param_groups]
        self.last_step = -1
        self.step()                       # apply step 0

    def _scale(self, step):
        if step < self.warmup_steps:
            return (step + 1) / max(1, self.warmup_steps)
        if step >= self.total_steps:
            return self.min_lr_ratio
        progress = (step - self.warmup_steps) / max(1, self.total_steps - self.warmup_steps)
        cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
        return self.min_lr_ratio + (1 - self.min_lr_ratio) * cosine

    def step(self):
        self.last_step += 1
        s = self._scale(self.last_step)
        for group, base in zip(self.opt.param_groups, self.base_lrs):
            group["lr"] = base * s

    def get_last_lr(self):
        return [g["lr"] for g in self.opt.param_groups]

    def state_dict(self):
        return {"last_step": self.last_step, "base_lrs": self.base_lrs}

    def load_state_dict(self, sd):
        self.last_step = sd["last_step"]
        self.base_lrs = sd["base_lrs"]
        # re-apply so lrs match the restored step
        s = self._scale(self.last_step)
        for group, base in zip(self.opt.param_groups, self.base_lrs):
            group["lr"] = base * s


def cosine_warmup_scheduler(optimizer, warmup_steps, total_steps, min_lr_ratio=0.1):
    return WarmupCosine(optimizer, warmup_steps, total_steps, min_lr_ratio)


class WarmupStableDecay:
    """Linear warmup -> hold at peak -> linear decay to min_lr_ratio * peak.

    MiniCPM-style WSD: empirically beats cosine at sub-1B scale, and intermediate
    checkpoints stay useful because they're sampled at peak LR (no slow-roll decay).
    `stable_share` is the fraction of the post-warmup steps spent at peak LR
    (rest is the decay phase). 0.8 is MiniCPM's reported sweet spot.

    Same interface contract as WarmupCosine so train.py just dispatches by name
    and the rest of the stack (resume, multi-group LRs, HybridOptimizer) is
    unchanged.
    """

    def __init__(self, optimizer, warmup_steps, total_steps,
                 stable_share=0.8, min_lr_ratio=0.1):
        assert 0.0 < stable_share < 1.0, "stable_share must be in (0, 1)"
        self.opt = optimizer
        self.warmup_steps = warmup_steps
        self.total_steps = total_steps
        self.stable_share = stable_share
        self.min_lr_ratio = min_lr_ratio
        post_warm = max(1, total_steps - warmup_steps)
        self.decay_start = warmup_steps + int(stable_share * post_warm)
        self.base_lrs = [g["lr"] for g in optimizer.param_groups]
        self.last_step = -1
        self.step()                       # apply step 0

    def _scale(self, step):
        if step < self.warmup_steps:
            return (step + 1) / max(1, self.warmup_steps)
        if step < self.decay_start:
            return 1.0
        if step >= self.total_steps:
            return self.min_lr_ratio
        progress = (step - self.decay_start) / max(
            1, self.total_steps - self.decay_start)
        return 1.0 - progress * (1.0 - self.min_lr_ratio)

    def step(self):
        self.last_step += 1
        s = self._scale(self.last_step)
        for group, base in zip(self.opt.param_groups, self.base_lrs):
            group["lr"] = base * s

    def get_last_lr(self):
        return [g["lr"] for g in self.opt.param_groups]

    def state_dict(self):
        return {"last_step": self.last_step, "base_lrs": self.base_lrs}

    def load_state_dict(self, sd):
        self.last_step = sd["last_step"]
        self.base_lrs = sd["base_lrs"]
        s = self._scale(self.last_step)
        for group, base in zip(self.opt.param_groups, self.base_lrs):
            group["lr"] = base * s


def wsd_scheduler(optimizer, warmup_steps, total_steps,
                  stable_share=0.8, min_lr_ratio=0.1):
    return WarmupStableDecay(optimizer, warmup_steps, total_steps,
                             stable_share, min_lr_ratio)


def build_scheduler(name, optimizer, warmup_steps, total_steps,
                    stable_share=0.8, min_lr_ratio=0.1):
    """Dispatch by name: 'cosine' (v1 default) or 'wsd' (MiniCPM)."""
    if name == "cosine":
        return cosine_warmup_scheduler(optimizer, warmup_steps, total_steps,
                                       min_lr_ratio)
    if name == "wsd":
        return wsd_scheduler(optimizer, warmup_steps, total_steps,
                             stable_share, min_lr_ratio)
    raise ValueError(f"unknown lr_schedule: {name}")