"""Finite-memory-bank Fisher-floor correction for Gaussian DSM."""
import torch


def estimate_floor(x, alpha, sigma, bank):
    """Estimate F(x)=alpha^2/sigma^4 tr Cov(Y|X=x) from a clean bank."""
    if x.ndim != 2 or bank.ndim != 2 or x.shape[1] != bank.shape[1]:
        raise ValueError("x and bank must be [N,D] and [K,D]")
    a = torch.as_tensor(alpha, device=x.device, dtype=x.dtype)
    s = torch.as_tensor(sigma, device=x.device, dtype=x.dtype)
    if a.ndim == 0:
        a = a.expand(x.shape[0])
    if s.ndim == 0:
        s = s.expand(x.shape[0])
    a, s = a.reshape(-1, 1), s.reshape(-1, 1)
    if a.shape[0] != x.shape[0] or torch.any(s <= 0):
        raise ValueError("schedule shapes invalid or sigma nonpositive")
    with torch.no_grad():
        logits = -((x[:, None, :] - a[:, None, :] * bank[None, :, :]) ** 2).sum(-1)
        logits = logits / (2.0 * s[:, None, :] ** 2)
        probs = torch.softmax(logits, dim=1)
        mean = probs @ bank
        variance_trace = (probs * ((bank[None, :, :] - mean[:, None, :]) ** 2).sum(-1)).sum(1)
        floor = (a[:, 0] ** 2 / s[:, 0] ** 4) * variance_trace
    return floor


def corrected_dsm_loss(score, x, y, alpha, sigma, bank, weight=None):
    """Return corrected mean loss plus raw/floor diagnostics."""
    a = torch.as_tensor(alpha, device=x.device, dtype=x.dtype)
    s = torch.as_tensor(sigma, device=x.device, dtype=x.dtype)
    if a.ndim == 0: a = a.expand(x.shape[0])
    if s.ndim == 0: s = s.expand(x.shape[0])
    target = (a[:, None] * y - x) / s[:, None] ** 2
    raw = ((score - target) ** 2).sum(-1)
    floor = estimate_floor(x, a, s, bank)
    if weight is None:
        weight = torch.ones_like(raw)
    loss = (weight * (raw - floor)).mean()
    return loss, {"raw": raw.mean().detach(), "floor": floor.mean().detach(),
                  "corrected": (raw - floor).mean().detach()}