"""Late-time fractional-order optimizer MVP.

The fractional direction is a normalized piecewise-constant Caputo-memory
average of recent gradients. Normalization keeps the learning-rate meaning
comparable to Adam/SGD while retaining the intended power-law kernel shape.
"""
from collections import deque
import math
import torch


class LateTimeOrderEstimator:
    def __init__(self, alpha0=0.5, amin=0.2, amax=0.95, beta=0.8,
                 warmup=40, lag=20, check_every=5, monotone_checks=2,
                 eps=1e-12):
        self.alpha = float(alpha0)
        self.amin, self.amax = amin, amax
        self.beta, self.warmup, self.lag = beta, warmup, lag
        self.check_every, self.monotone_checks, self.eps = check_every, monotone_checks, eps
        self.signals = []
        self.ratios = deque(maxlen=monotone_checks)

    def update(self, step, loss):
        # A positive loss is the scalar observation M. EMA suppresses minibatch noise.
        loss = max(float(loss), self.eps)
        if self.signals:
            self.signals.append(self.beta * self.signals[-1] + (1-self.beta) * loss)
        else:
            self.signals.append(loss)
        if step < self.warmup or step % self.check_every:
            return self.alpha
        old_i = max(0, len(self.signals)-1-self.lag)
        if old_i == len(self.signals)-1:
            return self.alpha
        ratio = self.signals[-1] / max(self.signals[old_i], self.eps)
        self.ratios.append(ratio)
        # The late-time hypothesis requires decreasing positive observations.
        if not (0.0 < ratio < 1.0) or len(self.ratios) >= 2 and any(
                self.ratios[i] >= self.ratios[i-1] for i in range(1, len(self.ratios))):
            return self.alpha
        rho = (step + 1) / float(step + 1 - self.lag)
        raw = -math.log(ratio) / math.log(rho)
        raw = min(self.amax, max(self.amin, raw))
        self.alpha = self.beta * self.alpha + (1-self.beta) * raw
        return self.alpha


class FractionalMemory(torch.optim.Optimizer):
    def __init__(self, params, lr=0.05, alpha=0.5, history=32, weight_decay=0.0,
                 adaptive=False, estimator=None):
        defaults = dict(lr=lr, alpha=alpha, history=history, weight_decay=weight_decay)
        super().__init__(params, defaults)
        self.adaptive = adaptive
        self.estimator = estimator or LateTimeOrderEstimator(alpha0=alpha)
        self.step_count = 0

    @torch.no_grad()
    def step(self, closure=None, loss_value=None):
        if closure is not None:
            with torch.enable_grad():
                loss = closure()
            loss_value = float(loss)
        self.step_count += 1
        alpha = self.estimator.update(self.step_count, loss_value) if self.adaptive and loss_value is not None else None
        for group in self.param_groups:
            a = alpha if alpha is not None else group['alpha']
            a = min(0.999, max(0.05, float(a)))
            H = group['history']
            # normalized interval weights: w_j=(j+1)^a-j^a, recent gradient j=0 largest
            weights = torch.tensor([(j+1)**a-j**a for j in range(H)], device=group['params'][0].device)
            weights = weights / weights.sum()
            for p in group['params']:
                if p.grad is None:
                    continue
                g = p.grad.detach().clone()
                state = self.state[p]
                hist = state.setdefault('grad_history', deque(maxlen=H))
                hist.appendleft(g)
                n = len(hist)
                v = sum(weights[j] * hist[j] for j in range(n))
                if group['weight_decay']:
                    p.mul_(1 - group['lr'] * group['weight_decay'])
                p.add_(v, alpha=-group['lr'])
        return None


class AdamW(torch.optim.AdamW):
    """Named wrapper so the experiment has an explicit standard baseline."""
    pass