"""Empirical Dimuon direction candidate, not TML's implementation.

Recovered using synthetic Qwen3-8B LoRA checkpoints on 2026-10-16.
The direction geometry and lagged-momentum statistics are empirically supported.
Batching FP16 Horner evaluation of the published five-step Polar Express table
gives about 0.15% mean update error across the original 288 rank-32 factors.
Matched warm first/second steps give about 0.00048%/0.067% with recovered
gradients rounded to their independently observed BF16 grid. Working EMA is
FP32 and persisted momentum is BF16. See results/residual-summary.json.
These improvements hold on layers excluded from selecting the arithmetic.
This is not bit-exact; per-factor GEMM gives materially different rounding.
Persistent-state replay additionally requires FP32 multiply-then-add-alpha
before the BF16 write. Fixed-weight mixed-gradient state probes give about
0.088--0.095% mean error on the following default-damping nonzero updates.
"""
import math
import torch

# Published Polar Express table (arXiv:2505.16932), first five iterations.
# This differs slightly from regenerating coefficients with the current
# upstream helper's default cushion. The published table fits the probes better.
_PAPER_COEFFICIENTS = (
    (8.28721201814563, -23.595886519098837, 17.300387312530933),
    (4.107059111542203, -2.9478499167379106, 0.5448431082926601),
    (3.9486908534822946, -2.908902115962949, 0.5518191394370137),
    (3.3184196573706015, -2.4884880243148474, 0.51004894012372),
    (2.300652019954817, -1.6689039845747493, 0.4188073119525673),
)
POLAR_EXPRESS_COEFFICIENTS = tuple(
    (a / 1.01, b / 1.01**3, c / 1.01**5)
    for a, b, c in _PAPER_COEFFICIENTS
)


def gradient_statistics(gradient):
    squared = gradient.float().square()
    return squared.mean(dim=-1, keepdim=True), squared.mean(dim=-2, keepdim=True)


def polar5(matrix, arithmetic='horner_fp16', norm_epsilon=0.0):
    """Accept a matrix or a batch of matrices, normalizing each independently.

    GPU batched GEMM (batch size >= 2) matches the API much more closely than
    per-matrix GEMM. Use candidate_steps to group compatible LoRA factors.
    """
    transposed = matrix.shape[-2] > matrix.shape[-1]
    x = matrix.mT if transposed else matrix
    x = x.to(torch.bfloat16 if arithmetic == 'bf16' else torch.float32)
    # A paper-style additive 1e-7 breaks the observed invariance when damping
    # rescales the input by 1e-4. Matched API probes favor zero/tiny epsilon.
    # Keep each reduction identical to the original individual-factor probe.
    norm = (torch.stack([item.norm() for item in x])[:, None, None]
            if x.ndim == 3 else x.norm())
    denominator = (norm * 1.01 + norm_epsilon).clamp_min(1e-30)
    x = x / denominator
    # The observed update fits FP16 rounding at this point much better than
    # either all-FP32 arithmetic or BF16 rounding. The historical baseline
    # iterates in FP32; the better-fitting Horner form iterates in FP16.
    if arithmetic.startswith('normalized_fp16') or arithmetic == 'horner_fp16':
        x = x.half().float()
    if arithmetic == 'horner_fp16':
        # Same polynomial as a*X + (b*S+c*S^2)@X, but round its Horner
        # evaluation in FP16. This order matches TML's public msign structure:
        # https://github.com/thinking-machines-lab/manifolds/blob/main/src/msign.py
        x = x.half()
        for a, b, c in POLAR_EXPRESS_COEFFICIENTS:
            gram = x @ x.mT
            polynomial = float(c) * gram
            polynomial.diagonal(dim1=-2, dim2=-1).add_(float(b))
            polynomial = polynomial @ gram
            polynomial.diagonal(dim1=-2, dim2=-1).add_(float(a))
            x = polynomial @ x
        return (x.mT if transposed else x).float()
    def mm(a, b):
        if arithmetic == 'bf16acc32':
            operation = torch.bmm if a.ndim == 3 else torch.mm
            return operation(a.bfloat16(), b.bfloat16(), out_dtype=torch.float32)
        return a @ b
    for a, b, c in POLAR_EXPRESS_COEFFICIENTS:
        gram = mm(x, x.mT)
        polynomial = float(b) * gram + float(c) * mm(gram, gram)
        x = float(a) * x + mm(polynomial, x)
    # Fresh rank-32 API updates lie on a commonly scaled FP16 mantissa grid.
    # Keep this optional: the matching kernel arithmetic is not yet recovered.
    if arithmetic.endswith('_output'):
        x = x.half().float()
    return (x.mT if transposed else x).float()


@torch.no_grad()
def direction(momentum, row_stat, col_stat, damping=1e-4,
              arithmetic='horner_fp16', placement='both',
              normalization_epsilon=0.0):
    """Return a direction with Frobenius norm sqrt(min(momentum.shape)).

    At beta1=0, use row/column mean-square statistics of the previous gradient.
    Use ones for a fresh optimizer. Both factors are treated independently.
    """
    m = momentum.float()
    row = row_stat.float() / row_stat.float().amax(dim=(-2,-1),keepdim=True).clamp_min(1e-30)
    col = col_stat.float() / col_stat.float().amax(dim=(-2,-1),keepdim=True).clamp_min(1e-30)
    left = (row + damping).pow(-0.25)
    right = (col + damping).pow(-0.25)
    transformed = left * m * right if placement in ('both', 'inside') else m
    output = polar5(transformed, arithmetic, normalization_epsilon)
    if placement in ('both', 'outside'):
        output = left * output * right
    return output * (math.sqrt(min(m.shape[-2:])) /
                     output.norm(dim=(-2,-1),keepdim=True).clamp_min(1e-30))


@torch.no_grad()
def candidate_step(parameter, gradient, previous_momentum=None, *,
                   learning_rate=0.01, beta1=0.9, damping=1e-4,
                   rank_lr_correction_exponent=0.175,
                   normalization_epsilon=0.0,
                   momentum_storage_dtype=torch.bfloat16,
                   arithmetic='horner_fp16'):
    """Functional reference step; returns (updated_parameter, new_momentum).

    Compute preconditioners BEFORE updating momentum. This ordering is essential.
    Use the FP32 working EMA for this update, then round its persisted copy to
    BF16 for the next step. The storage dtype is an empirical inference; callers
    can select another dtype to compare hypotheses. Pass back the returned
    momentum unchanged on the next call.
    Set arithmetic='normalized_fp16' for the earlier FP32-iteration baseline.
    Accepts (rows, cols) or (batch, rows, cols). Prefer candidate_steps for
    a collection of factors so GPU matrix multiplications run in batches.
    Gradients retain their input precision. Tinker probes support BF16 native
    gradients; round recovered gradients to BF16 to remove inversion noise.
    Global gradient clipping, when requested, belongs before this function.
    Row/column statistics are derived from previous_momentum; this reference
    has no independently accumulated row/column second-moment state.
    """
    if previous_momentum is None:
        previous_momentum = torch.zeros_like(gradient, dtype=torch.float32)
    row, col = gradient_statistics(previous_momentum)
    # Preserve the operation order: round the beta*old product in FP32, then
    # add the scaled gradient with add_'s fused alpha operation. Separately
    # forming (1-beta)*gradient changes BF16 rounding at a small set of entries
    # and produces a substantially worse state on the following update.
    momentum = previous_momentum.float().mul(beta1)
    momentum.add_(gradient.float(), alpha=1 - beta1)
    update = direction(momentum, row, col, damping, arithmetic=arithmetic,
                       normalization_epsilon=normalization_epsilon)
    rank = min(parameter.shape[-2:])
    scale = learning_rate * (rank / 32) ** rank_lr_correction_exponent
    return parameter.float() - scale * update, momentum.to(momentum_storage_dtype)


@torch.no_grad()
def candidate_steps(parameters, gradients, previous_momenta=None, *,
                    batch_size=8, **options):
    """Step a dictionary of LoRA factors using real batches of matching shapes.

    Return (updated_parameters, new_momenta), both dictionaries with the input
    keys. Canonicalize orientation before stacking. A final singleton is padded
    with a copy to retain batched GEMM dispatch, then the padding is discarded.
    Batch sizes 2--32 agree closely in the saved rank-32 controls. Exact backend
    kernels and input gradients are not fully recovered.
    """
    from collections import defaultdict
    if batch_size < 1:
        raise ValueError('batch_size must be positive')
    groups=defaultdict(list)
    for key,parameter in parameters.items():
        transposed=parameter.shape[0]>parameter.shape[1]
        shape=tuple(sorted(parameter.shape))
        groups[(shape,parameter.device,parameter.dtype)].append((key,transposed))
    updated={};momenta={}
    for entries in groups.values():
        for start in range(0,len(entries),batch_size):
            real=entries[start:start+batch_size]
            chunk=real*2 if len(real)==1 and batch_size>1 else real
            def stack(values):
                return torch.stack([values[key].T if transposed else values[key]
                                    for key,transposed in chunk])
            previous=stack(previous_momenta) if previous_momenta is not None else None
            result,momentum=candidate_step(stack(parameters),stack(gradients),previous,**options)
            for i,(key,transposed) in enumerate(real):
                updated[key]=result[i].T if transposed else result[i]
                momenta[key]=momentum[i].T if transposed else momentum[i]
    return updated,momenta