from __future__ import annotations
import torch


def parameter_group_name(name: str, parameter: torch.nn.Parameter) -> str:
    """Return the frozen AdamW group for one named parameter."""
    if getattr(parameter, "_no_weight_decay", False):
        return "no_decay"
    if parameter.ndim < 2 or name.endswith("bias"):
        return "no_decay"
    return "decay"


def optimizer_group_manifest(model) -> dict:
    rows = {"decay": [], "no_decay": []}
    for name, parameter in model.named_parameters():
        if parameter.requires_grad:
            rows[parameter_group_name(name, parameter)].append(name)
    return {
        group: {
            "names": names,
            "tensors": len(names),
            "parameters": sum(dict(model.named_parameters())[name].numel() for name in names),
        }
        for group, names in rows.items()
    }


def build_adamw(model, lr=1e-3, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.1, fused=True):
    decay, no_decay = [], []
    for name, p in model.named_parameters():
        if not p.requires_grad:
            continue
        if parameter_group_name(name, p) == "no_decay":
            no_decay.append(p)
        else:
            decay.append(p)
    groups = [
        {"params": decay, "weight_decay": weight_decay},
        {"params": no_decay, "weight_decay": 0.0},
    ]
    kwargs = dict(lr=lr, betas=betas, eps=eps)
    if fused and torch.cuda.is_available():
        kwargs["fused"] = True
    return torch.optim.AdamW(groups, **kwargs)