"""Alternative point estimators for axial categorical distributions."""

import torch


def estimates(logits, names):
    probability = logits.float().softmax(-2)
    height = logits.shape[-2]
    grid = torch.arange(height, device=logits.device, dtype=torch.float32)[
        None, None, :, None
    ]
    result = {}
    if "mean" in names:
        result["mean"] = (probability * grid).sum(-2)
    if "median" in names:
        cumulative = probability.cumsum(-2)
        index = (cumulative < 0.5).sum(-2).clamp_max(height - 1)
        mass = probability.gather(-2, index.unsqueeze(-2)).squeeze(-2)
        before = cumulative.gather(-2, (index - 1).clamp_min(0).unsqueeze(-2)).squeeze(
            -2
        )
        before = torch.where(index == 0, torch.zeros_like(before), before)
        result["median"] = (
            index.float() - 0.5 + (0.5 - before) / mass.clamp_min(1e-12)
        ).clamp(0, height - 1)
    if "quadratic_mode" in names:
        index = logits.argmax(-2)
        score = logits.float()
        center = score.gather(-2, index.unsqueeze(-2)).squeeze(-2)
        left = score.gather(-2, (index - 1).clamp_min(0).unsqueeze(-2)).squeeze(-2)
        right = score.gather(
            -2, (index + 1).clamp_max(height - 1).unsqueeze(-2)
        ).squeeze(-2)
        curvature = left - 2 * center + right
        offset = torch.where(
            (curvature < -1e-8) & (index > 0) & (index < height - 1),
            0.5 * (left - right) / curvature.clamp_max(-1e-8),
            torch.zeros_like(center),
        ).clamp(-0.5, 0.5)
        result["quadratic_mode"] = (index.float() + offset).clamp(0, height - 1)
    if set(result) != set(names):
        raise ValueError("Unknown posterior estimator")
    return result


def positions_from_logits(logits, config, device=None):
    by_device = config.get("inference", {}).get("posterior_by_device")
    if by_device is not None:
        if device is None or device not in by_device:
            raise ValueError("Configured head estimators require the device")
        keys = config["device_keys"][device]
        mapping = by_device[device]
        if set(mapping) != set(keys) or logits.shape[1] != len(keys):
            raise ValueError("Configured head estimators differ from output identities")
        names = [mapping[key] for key in keys]
        values = estimates(logits, sorted(set(names)))
        return torch.stack(
            [values[name][:, index] for index, name in enumerate(names)], dim=1
        )
    name = config.get("inference", {}).get("posterior", "mean")
    if name not in ["mean", "median", "quadratic_mode"]:
        raise ValueError("Unsupported configured posterior")
    return estimates(logits, [name])[name]