"""Independent SOAP: Adam in Shampoo eigen-coordinates (Appendix C.2).

Uses orthogonal iteration (QR) to track eigenvectors every two steps. First
moments live in physical parameter coordinates; second moments are permuted
when eigenvectors are re-ordered. Vectors/scalars use unrotated Adam.
"""
import torch


class SOAP(torch.optim.Optimizer):
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
                 shampoo_beta=0.999, precondition_frequency=2):
        super().__init__(params, dict(lr=lr, betas=betas, eps=eps,
                                     shampoo_beta=shampoo_beta,
                                     precondition_frequency=precondition_frequency))

    @staticmethod
    def _basis(cov):
        # Initialization only. The DCU float32 eigensolver can fail for repeated
        # eigenvalues, especially the zero gradients of alpha=0 residual blocks.
        # Scaling and CPU float64 preserve the symmetric eigenproblem exactly.
        work = cov.detach().to(device='cpu', dtype=torch.float64)
        scale = float(work.abs().max())
        if scale == 0.0:
            return torch.eye(cov.shape[0], device=cov.device, dtype=cov.dtype)
        work = (work + work.T) * (0.5 / scale)
        _, q = torch.linalg.eigh(work)
        return q.flip(-1).to(device=cov.device, dtype=cov.dtype)

    @staticmethod
    def _iterate(cov, q):
        # Descending Rayleigh quotient prevents arbitrary eigenvector swaps.
        scores = (q * (cov @ q)).sum(0)
        order = torch.argsort(scores, descending=True)
        q = q[:, order]
        scale = cov.abs().amax().clamp_min(1e-30)
        q, _ = torch.linalg.qr((cov / scale) @ q)
        return q, order

    @torch.no_grad()
    def step(self, closure=None):
        result = None
        if closure is not None:
            with torch.enable_grad():
                result = closure()
        for group in self.param_groups:
            b1, b2 = group['betas']
            sb, freq = group['shampoo_beta'], group['precondition_frequency']
            for p in group['params']:
                if p.grad is None:
                    continue
                g = p.grad
                s = self.state[p]
                if not s:
                    s['step'] = 0
                    s['m'] = torch.zeros_like(p)
                    s['v'] = torch.zeros_like(p)
                    if p.ndim == 2:
                        s['left'] = (1.0 - sb) * (g @ g.T)
                        s['right'] = (1.0 - sb) * (g.T @ g)
                        s['ql'] = self._basis(s['left'])
                        s['qr'] = self._basis(s['right'])
                        # SOAP initializes the preconditioner on the first gradient.
                        # The first scheduled LR is zero, so no training is skipped.
                    s['initialized_only'] = True
                    continue
                s['step'] += 1
                step = s['step']
                s['m'].lerp_(g, 1.0 - b1)
                if p.ndim == 2:
                    gl = s['ql'].T @ g @ s['qr']
                    ml = s['ql'].T @ s['m'] @ s['qr']
                else:
                    gl, ml = g, s['m']
                s['v'].mul_(b2).addcmul_(gl, gl, value=1.0 - b2)
                denom = (s['v'] / (1.0 - b2 ** step)).sqrt().add_(group['eps'])
                update = (ml / (1.0 - b1 ** step)) / denom
                if p.ndim == 2:
                    update = s['ql'] @ update @ s['qr'].T
                p.add_(update, alpha=-group['lr'])
                if p.ndim == 2:
                    s['left'].mul_(sb).addmm_(g, g.T, beta=1.0, alpha=1.0 - sb)
                    s['right'].mul_(sb).addmm_(g.T, g, beta=1.0, alpha=1.0 - sb)
                    if step % freq == 0:
                        s['ql'], il = self._iterate(s['left'], s['ql'])
                        s['qr'], ir = self._iterate(s['right'], s['qr'])
                        s['v'] = s['v'][il][:, ir].contiguous()
        return result