"""Muon optimizer (Newton-Schulz orthogonalized momentum) + hybrid helpers.
Convention: Muon on 2D hidden matrices, AdamW on everything else (emb/pos/LN/bias)."""
import torch


def newton_schulz(G, steps=5, eps=1e-7):
    """approximate polar factor of G via the quintic NS iteration (Keller Jordan coefficients)."""
    a, b, c = 3.4445, -4.7750, 2.0315
    X = G / (G.norm() + eps)
    transposed = X.size(0) > X.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
    return X.T if transposed else X


class Muon(torch.optim.Optimizer):
    def __init__(self, params, lr=0.02, momentum=0.95, ns_steps=5, nesterov=True, wd=0.0):
        super().__init__(params, dict(lr=lr, momentum=momentum, ns_steps=ns_steps, nesterov=nesterov, wd=wd))

    @torch.no_grad()
    def step(self, closure=None):
        for group in self.param_groups:
            for p in group['params']:
                if p.grad is None: continue
                g = p.grad
                st = self.state[p]
                if 'mom' not in st: st['mom'] = torch.zeros_like(g)
                buf = st['mom']
                buf.mul_(group['momentum']).add_(g)
                u = g.add(buf, alpha=group['momentum']) if group['nesterov'] else buf
                if u.ndim == 2:
                    u = newton_schulz(u, group['ns_steps'])
                    u = u * max(1.0, u.size(0) / u.size(1)) ** 0.5   # rms-matched scaling
                if group['wd'] > 0: p.mul_(1 - group['lr'] * group['wd'])   # decoupled decay (Moonshot: required for scale)
                p.add_(u, alpha=-group['lr'])


class MultiOpt:
    """duck-typed bundle of optimizers (step/zero_grad/state_dict API-compatible)."""
    def __init__(self, opts): self.optimizers = opts
    def step(self):
        for o in self.optimizers: o.step()
    def zero_grad(self, set_to_none=True):
        for o in self.optimizers: o.zero_grad(set_to_none=set_to_none)
    def state_dict(self):
        return [o.state_dict() for o in self.optimizers]
    def load_state_dict(self, sds):
        for o, sd in zip(self.optimizers, sds): o.load_state_dict(sd)


class MultiSched:
    def __init__(self, scheds): self.scheds = scheds
    def step(self):
        for s in self.scheds: s.step()


def build_hybrid(blocks, other_params, lr_adamw, lr_muon, warmup, total_steps=0, lr_min_ratio=0.1,
                 muon_mom=0.95, adam_b1=0.9, head_param=None, head_lr_mult=1.0, muon_wd=0.0):
    """Muon(2D block matrices) + AdamW(everything else). Scheds: linear warmup, then cosine decay to
    lr_min_ratio*peak if total_steps>0 (long runs), else constant after warmup (legacy).
    muon_mom/adam_b1: momentum knobs (late-SNR noise-averaging arms, 2026-07-13)."""
    import math as _m
    mats = [p for p in blocks.parameters() if p.ndim == 2]
    mat_ids = {id(p) for p in mats}
    rest = [p for p in other_params if id(p) not in mat_ids]
    om = Muon(mats, lr=lr_muon, momentum=muon_mom, wd=muon_wd)
    if head_param is not None and head_lr_mult != 1.0:
        hid = id(head_param)
        groups = [{'params': [p for p in rest if id(p) != hid], 'lr': lr_adamw},
                  {'params': [p for p in rest if id(p) == hid], 'lr': lr_adamw * head_lr_mult}]
        oa = torch.optim.AdamW(groups, lr=lr_adamw, weight_decay=1e-4, betas=(adam_b1, 0.999))
    else:
        oa = torch.optim.AdamW(rest, lr=lr_adamw, weight_decay=1e-4, betas=(adam_b1, 0.999))
    if total_steps > 0:
        def fn(s):
            if s < warmup: return (s + 1) / max(warmup, 1)
            p = min(1.0, (s - warmup) / max(1, total_steps - warmup))
            return lr_min_ratio + 0.5 * (1 - lr_min_ratio) * (1 + _m.cos(_m.pi * p))
    else:
        fn = lambda s: min(1.0, (s + 1) / max(warmup, 1))
    scheds = [torch.optim.lr_scheduler.LambdaLR(om, fn), torch.optim.lr_scheduler.LambdaLR(oa, fn)]
    return MultiOpt([om, oa]), MultiSched(scheds)


class Lion(torch.optim.Optimizer):
    """Lion (Chen et al. 2023): sign of the beta1-mixed momentum; decoupled wd."""
    def __init__(self, params, lr=3e-4, betas=(0.9, 0.99), wd=0.0):
        super().__init__(params, dict(lr=lr, betas=betas, wd=wd))

    @torch.no_grad()
    def step(self, closure=None):
        for g_ in self.param_groups:
            b1, b2 = g_['betas']
            for p in g_['params']:
                if p.grad is None: continue
                st = self.state[p]
                if 'm' not in st: st['m'] = torch.zeros_like(p)
                m = st['m']
                u = (b1 * m + (1 - b1) * p.grad).sign_()
                if g_['wd'] > 0: p.mul_(1 - g_['lr'] * g_['wd'])
                p.add_(u, alpha=-g_['lr'])
                m.mul_(b2).add_(p.grad, alpha=1 - b2)


class OLion(torch.optim.Optimizer):
    """OLion (arXiv:2602.01105): Lion-style momentum -> Newton-Schulz -> entrywise sign,
    RMS alignment gamma (||sign||_F = sqrt(numel) exactly, so D = gamma * sign(Q)), decoupled wd.
    2D params only; caller routes 1D elsewhere. betas default to Lion's (0.9, 0.99) — the paper's
    defaults were not in the pages we read; flagged as an assumption in the battery notes."""
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.99), gamma=0.2, ns_steps=5, wd=0.0):
        super().__init__(params, dict(lr=lr, betas=betas, gamma=gamma, ns_steps=ns_steps, wd=wd))

    @torch.no_grad()
    def step(self, closure=None):
        for g_ in self.param_groups:
            b1, b2 = g_['betas']
            for p in g_['params']:
                if p.grad is None: continue
                st = self.state[p]
                if 'm' not in st: st['m'] = torch.zeros_like(p)
                m = st['m']
                m.mul_(b2).add_(p.grad, alpha=1 - b2)          # slow momentum
                gt = (1 - b1) * p.grad + b1 * m                # Nesterov mix
                q = newton_schulz(gt, g_['ns_steps'])
                if g_['wd'] > 0: p.mul_(1 - g_['lr'] * g_['wd'])
                p.add_(q.sign_(), alpha=-g_['lr'] * g_['gamma'])


class Adafactor2D(torch.optim.Optimizer):
    """Minimal Shazeer-Stern factored Adam for 2D params: row/col second-moment statistics,
    RMS-1 update clipping, no relative-step magic (external lr + schedule). The analog-native
    Adam per the 07-11 BoM audit (row/col stats = AGC channels)."""
    def __init__(self, params, lr=1e-3, beta2=0.999, eps=1e-30, clip=1.0, wd=0.0):
        super().__init__(params, dict(lr=lr, beta2=beta2, eps=eps, clip=clip, wd=wd))

    @torch.no_grad()
    def step(self, closure=None):
        for g_ in self.param_groups:
            for p in g_['params']:
                if p.grad is None: continue
                g = p.grad
                st = self.state[p]
                if 'r' not in st:
                    st['r'] = torch.zeros(g.shape[0], device=g.device)
                    st['c'] = torch.zeros(g.shape[1], device=g.device)
                r, c = st['r'], st['c']
                g2 = g.float().pow(2) + g_['eps']
                r.mul_(g_['beta2']).add_(g2.mean(1), alpha=1 - g_['beta2'])
                c.mul_(g_['beta2']).add_(g2.mean(0), alpha=1 - g_['beta2'])
                v = r[:, None] * c[None, :] / max(float(r.mean()), g_['eps'])
                u = g / v.sqrt().to(g.dtype)
                rms = float(u.pow(2).mean().sqrt())
                if rms > g_['clip']: u = u * (g_['clip'] / rms)
                if g_['wd'] > 0: p.mul_(1 - g_['lr'] * g_['wd'])
                p.add_(u, alpha=-g_['lr'])


class PSGDQuadWrap(torch.optim.Optimizer):
    """Xi-Lin Li's KronWhiten (dQ='QUAD') as a proper torch Optimizer so LambdaLR accepts it.
    Our EP flow pre-populates p.grad; the closure hands PSGD a synthetic scalar whose autograd
    gradient equals the stored p.grad (loss = sum <p, p.grad_detached>). Screening-tier: the
    inner preconditioner state is NOT checkpointed (base state_dict covers param_groups only)."""
    def __init__(self, params, lr=1e-3, momentum=0.95):
        params = list(params)
        super().__init__(params, dict(lr=lr))
        from psgd_vendor import KronWhiten
        self._flat = [p for g_ in self.param_groups for p in g_['params']]
        self.inner = KronWhiten(self._flat, preconditioner_init_scale=1.0,
                                lr_params=lr, lr_preconditioner=0.1, momentum=momentum,
                                whiten_grad=True, dQ="QUAD")

    @torch.no_grad()
    def step(self, closure=None):
        self.inner.lr_params = self.param_groups[0]['lr']
        flat = self._flat
        def _closure():
            return sum((p * p.grad.detach()).sum() for p in flat if p.grad is not None)
        self.inner.step(_closure)


def build_alt(opt_name, blocks, other_params, lr, warmup, total_steps=0, lr_min_ratio=0.1,
              lr_matrix=None, wd=0.0):
    """Screening-tier builder for the optimizer price list: OPT on block matrices + AdamW on the
    rest (same split as build_hybrid so arms differ only in the matrix rule). sgdm applies SGD to
    everything (no split) — the fully-local baseline."""
    import math as _m
    mats = [p for p in blocks.parameters() if p.ndim == 2]
    mat_ids = {id(p) for p in mats}
    rest = [p for p in other_params if id(p) not in mat_ids]
    lm = lr_matrix if lr_matrix is not None else lr
    if opt_name == 'sgdm':
        om = torch.optim.SGD(mats + rest, lr=lm, momentum=0.95, nesterov=True)
        opts = [om]
    else:
        om = {'lion': lambda: Lion(mats, lr=lm, wd=wd),
              'olion': lambda: OLion(mats, lr=lm, wd=wd),
              'adafactor': lambda: Adafactor2D(mats, lr=lm, wd=wd),
              'signline': lambda: SignLine(mats, lr=lm, wd=wd),
              'conslion': lambda: ConsensusLion(mats, lr=lm, wd=wd),
              'ditherlion': lambda: DitherLion(mats, lr=lm, wd=wd),
              'cautlion': lambda: CautiousLion(mats, lr=lm, wd=wd),
              'olionns': lambda: OLion(mats, lr=lm, wd=wd, ns_steps=0),
              'olionk1': lambda: OLion(mats, lr=lm, wd=wd, ns_steps=1),
              'olionk2': lambda: OLion(mats, lr=lm, wd=wd, ns_steps=2),
              'olionk3': lambda: OLion(mats, lr=lm, wd=wd, ns_steps=3),
              'psgdquad': lambda: PSGDQuadWrap(mats, lr=lm)}[opt_name]()
        oa = torch.optim.AdamW(rest, lr=lr, weight_decay=1e-4)
        opts = [om, oa]
    if total_steps > 0:
        def fn(s):
            if s < warmup: return (s + 1) / max(warmup, 1)
            pr = min(1.0, (s - warmup) / max(1, total_steps - warmup))
            return lr_min_ratio + 0.5 * (1 - lr_min_ratio) * (1 + _m.cos(_m.pi * pr))
    else:
        fn = lambda s: min(1.0, (s + 1) / max(warmup, 1))
    scheds = [torch.optim.lr_scheduler.LambdaLR(o, fn) for o in opts]
    return MultiOpt(opts), MultiSched(scheds)


class SignLine(torch.optim.Optimizer):
    """sign(momentum) with per-row/per-column pulse amplitudes from leaky RMS line statistics.
    Insight: positive row/col scaling INSIDE a sign is a no-op; applied OUTSIDE as amplitudes
    (a_i b_j sign(m_ij)) it is the analog-native form of factored adaptivity: per-line DACs set
    drive amplitude, comparators give the sign. Amplitudes normalized to unit mean so lr keeps
    its scale; alpha in [0,1] interpolates flat-Lion (0) -> full line-adaptive (1)."""
    def __init__(self, params, lr=3e-4, betas=(0.9, 0.99), beta_line=0.999, alpha=1.0, wd=0.0):
        super().__init__(params, dict(lr=lr, betas=betas, bl=beta_line, alpha=alpha, wd=wd))

    @torch.no_grad()
    def step(self, closure=None):
        for g_ in self.param_groups:
            b1, b2 = g_['betas']
            for p in g_['params']:
                if p.grad is None: continue
                st = self.state[p]
                if 'm' not in st:
                    st['m'] = torch.zeros_like(p)
                    st['r'] = torch.ones(p.shape[0], device=p.device)
                    st['c'] = torch.ones(p.shape[1], device=p.device)
                m, r, c = st['m'], st['r'], st['c']
                u = b1 * m + (1 - b1) * p.grad
                r.mul_(g_['bl']).add_(u.float().pow(2).mean(1), alpha=1 - g_['bl'])
                c.mul_(g_['bl']).add_(u.float().pow(2).mean(0), alpha=1 - g_['bl'])
                a = (r / r.mean()).clamp_min(1e-12).pow(-0.25 * g_['alpha'])
                b = (c / c.mean()).clamp_min(1e-12).pow(-0.25 * g_['alpha'])
                amp = torch.outer(a, b).to(p.dtype)
                amp = amp / amp.mean()
                if g_['wd'] > 0: p.mul_(1 - g_['lr'] * g_['wd'])
                p.add_(u.sign_() * amp, alpha=-g_['lr'])
                m.mul_(b2).add_(p.grad, alpha=1 - b2)


class ConsensusLion(torch.optim.Optimizer):
    """two-timescale sign consensus: update only where sign(fast momentum) == sign(slow momentum).
    Two leaky integrators with different leaks + one comparator + coincidence gate — the entire
    optimizer is capacitors and logic. Where they disagree, write nothing (noise veto)."""
    def __init__(self, params, lr=3e-4, beta_fast=0.9, beta_slow=0.99, wd=0.0):
        super().__init__(params, dict(lr=lr, bf=beta_fast, bs=beta_slow, wd=wd))

    @torch.no_grad()
    def step(self, closure=None):
        for g_ in self.param_groups:
            for p in g_['params']:
                if p.grad is None: continue
                st = self.state[p]
                if 'mf' not in st:
                    st['mf'] = torch.zeros_like(p); st['ms'] = torch.zeros_like(p)
                mf, ms = st['mf'], st['ms']
                mf.mul_(g_['bf']).add_(p.grad, alpha=1 - g_['bf'])
                ms.mul_(g_['bs']).add_(p.grad, alpha=1 - g_['bs'])
                sf, ss = mf.sign(), ms.sign()
                u = sf * (sf == ss)
                if g_['wd'] > 0: p.mul_(1 - g_['lr'] * g_['wd'])
                p.add_(u, alpha=-g_['lr'])


class CautiousLion(torch.optim.Optimizer):
    """C-Lion (Liang et al., arXiv:2411.16085, ICLR'26): Lion masked where the update sign
    disagrees with the CURRENT gradient sign. The mandatory prior-art baseline for ConsensusLion;
    the difference under test is instantaneous-gradient gating (this) vs filtered two-EMA gating."""
    def __init__(self, params, lr=3e-4, betas=(0.9, 0.99), wd=0.0):
        super().__init__(params, dict(lr=lr, betas=betas, wd=wd))

    @torch.no_grad()
    def step(self, closure=None):
        for g_ in self.param_groups:
            b1, b2 = g_['betas']
            for p in g_['params']:
                if p.grad is None: continue
                st = self.state[p]
                if 'm' not in st: st['m'] = torch.zeros_like(p)
                m = st['m']
                u = (b1 * m + (1 - b1) * p.grad).sign_()
                u = u * ((u * p.grad) > 0)
                if g_['wd'] > 0: p.mul_(1 - g_['lr'] * g_['wd'])
                p.add_(u, alpha=-g_['lr'])
                m.mul_(b2).add_(p.grad, alpha=1 - b2)


class DitherLion(torch.optim.Optimizer):
    """Lion with dithered sign: sign(m + tau*noise*rms(m)). Free substrate noise turns the hard
    sign into an unbiased soft-sign in expectation, letting small entries carry proportional
    information across steps. tau=0 recovers Lion."""
    def __init__(self, params, lr=3e-4, betas=(0.9, 0.99), tau=0.5, wd=0.0):
        super().__init__(params, dict(lr=lr, betas=betas, tau=tau, wd=wd))

    @torch.no_grad()
    def step(self, closure=None):
        for g_ in self.param_groups:
            b1, b2 = g_['betas']
            for p in g_['params']:
                if p.grad is None: continue
                st = self.state[p]
                if 'm' not in st: st['m'] = torch.zeros_like(p)
                m = st['m']
                u = b1 * m + (1 - b1) * p.grad
                d = torch.randn_like(u) * (g_['tau'] * u.float().pow(2).mean().sqrt().to(u.dtype))
                if g_['wd'] > 0: p.mul_(1 - g_['lr'] * g_['wd'])
                p.add_((u + d).sign_(), alpha=-g_['lr'])
                m.mul_(b2).add_(p.grad, alpha=1 - b2)