"""Projection-Regularized Gradient optimizer (small-model exact-matrix MVP).

This implementation follows the supplied formula literally. It is intended for
small models; for large models, replace _metric with a diagonal or Woodbury
implementation rather than materializing p x p matrices.
"""
from collections import deque
import torch


class ProjectionRegularizedSGD(torch.optim.Optimizer):
    """SGD using a regularized recent-gradient subspace metric.

    Args:
        params: model parameters
        lr: step size
        history_size: number of recent flattened gradients
        beta: covariance EMA coefficient (the MVP uses the current window)
        ridge: lambda in S + lambda I
        rho: ridge in the m-by-m projected Gram system
        alpha: residual identity strength, in [0, 1]
        max_exact_parameters: safety limit for the exact p-by-p implementation
    """
    def __init__(self, params, lr=1e-2, history_size=4, beta=0.0,
                 ridge=1e-2, rho=0.1, alpha=0.1, max_exact_parameters=4096):
        if lr < 0 or history_size < 1 or not (0 <= beta < 1):
            raise ValueError("invalid lr/history_size/beta")
        if ridge <= 0 or rho <= 0 or not (0 <= alpha <= 1):
            raise ValueError("require ridge,rho > 0 and alpha in [0,1]")
        super().__init__(params, dict(lr=lr))
        self.history_size, self.beta = history_size, beta
        self.ridge, self.rho, self.alpha = ridge, rho, alpha
        self.max_exact_parameters = max_exact_parameters
        self._history = deque(maxlen=history_size)
        self._S = None

    @torch.no_grad()
    def _flat_grad(self):
        chunks = []
        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None:
                    chunks.append(torch.zeros_like(p).reshape(-1))
                else:
                    chunks.append(p.grad.reshape(-1))
        return torch.cat(chunks)

    @torch.no_grad()
    def _metric(self, g):
        G = torch.stack(list(self._history), dim=1)
        p = g.numel()
        if p > self.max_exact_parameters:
            raise RuntimeError("exact projection is limited to small models")
        eye = torch.eye(p, device=g.device, dtype=g.dtype)
        cov = (G @ G.T) / G.shape[1]
        if self._S is None:
            S = cov + self.ridge * eye
        else:
            S = self.beta * self._S + (1.0 - self.beta) * cov + self.ridge * eye
        self._S = S.detach()
        # eigh gives a symmetric, numerically stable inverse square root.
        eigval, eigvec = torch.linalg.eigh(S)
        invsqrt = (eigvec * eigval.clamp_min(1e-12).rsqrt()) @ eigvec.T
        if self.alpha == 1.0:
            return invsqrt @ invsqrt
        SinvG = torch.linalg.solve(S, G)
        K = G.T @ SinvG + self.rho * torch.eye(
            G.shape[1], device=g.device, dtype=g.dtype)
        P = invsqrt @ G @ torch.linalg.solve(K, G.T) @ invsqrt
        return invsqrt @ ((1.0 - self.alpha) * P + self.alpha * eye) @ invsqrt

    @torch.no_grad()
    def step(self, closure=None):
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()
        g = self._flat_grad()
        self._history.append(g.detach().clone())
        d = self._metric(g) @ g
        pos = 0
        for group in self.param_groups:
            for p in group["params"]:
                n = p.numel()
                if p.grad is not None:
                    p.add_(d[pos:pos+n].view_as(p), alpha=-group["lr"])
                pos += n
        return loss