"""From-scratch optimizers, one per category of Table 1 in
*Fantastic Pretraining Optimizers and Where to Find Them* (Wen et al., 2025).

    category                          implementation      state / param
    --------------------------------  ------------------  -------------
    1. AdamW                          AdamW               2x
    2. Variance-reduced AdamW variant  MARS               3x
    3. Memory-efficient                Lion               1x
    4. Matrix-based                    Muon               1x (2D) / 2x
    5. Hessian-based (bonus)           Sophia-G           2x

``torch.optim.Optimizer`` is used only as a base class; no ``torch.optim``
update rule is called anywhere.
"""

from .adamw import AdamW
from .lion import Lion
from .mars import MARS
from .muon import Muon
from .sophia import SophiaG

OPTIMIZERS = {
    "adamw": AdamW,
    "mars": MARS,
    "lion": Lion,
    "muon": Muon,
    "sophia": SophiaG,
}

CATEGORY = {
    "adamw": "1. AdamW",
    "mars": "2. Variance-reduced AdamW variant",
    "lion": "3. Memory-efficient",
    "muon": "4. Matrix-based",
    "sophia": "5. Hessian-based (bonus)",
}

# Peak learning rates are the values selected by src/part2/sweep.py; the
# remaining hyperparameters are each method's published defaults and were not
# swept.  Note how far apart the selected rates are -- Muon tolerates 500x
# Sophia's -- which is exactly why every optimizer needed its own grid.
DEFAULT_HPARAMS = {
    "adamw":  dict(lr=1e-3,  betas=(0.9, 0.95),  weight_decay=0.1),
    "mars":   dict(lr=3e-3,  betas=(0.95, 0.99), weight_decay=0.1, gamma=0.025),
    "lion":   dict(lr=3e-4,  betas=(0.9, 0.99),  weight_decay=0.6),
    "muon":   dict(lr=5e-2,  momentum=0.95,      weight_decay=0.1, adamw_lr=1e-3),
    "sophia": dict(lr=1e-4,  betas=(0.965, 0.99), rho=0.05, weight_decay=0.1),
}

__all__ = ["AdamW", "MARS", "Lion", "Muon", "SophiaG",
           "OPTIMIZERS", "CATEGORY", "DEFAULT_HPARAMS"]