"""First-Spike Laplacian attention for latency-coded tokens."""
import torch
from torch import nn


class FirstSpikeLaplacianAttention(nn.Module):
    """Attention using exp(-||q-k||_1 / sigma), with one bandwidth per head.

    q, k: [batch, heads, query_tokens/key_tokens, channels]
    v:    [batch, heads, key_tokens, value_channels]
    """
    def __init__(self, heads: int, init_sigma: float = 1.0, eps: float = 1e-6):
        super().__init__()
        if init_sigma <= eps:
            raise ValueError("init_sigma must be positive")
        self.eps = eps
        # inverse softplus gives the requested initial positive bandwidth
        theta = torch.log(torch.expm1(torch.tensor(float(init_sigma - eps))))
        self.theta = nn.Parameter(theta.repeat(heads))

    def bandwidth(self):
        return torch.nn.functional.softplus(self.theta) + self.eps

    def forward(self, q, k, v, return_attention=False):
        if q.ndim != 4 or k.ndim != 4 or v.ndim != 4:
            raise ValueError("q, k, v must have shape [B,H,T,C]")
        if q.shape[0] != k.shape[0] or q.shape[1] != k.shape[1]:
            raise ValueError("q and k batch/head dimensions must match")
        if k.shape[:2] != v.shape[:2] or k.shape[2] != v.shape[2]:
            raise ValueError("k and v batch/head/token dimensions must match")
        distances = (q[:, :, :, None, :] - k[:, :, None, :, :]).abs().sum(dim=-1)
        logits = -distances / self.bandwidth()[None, :, None, None]
        attention = torch.softmax(logits, dim=-1)
        output = attention @ v
        return (output, attention, distances) if return_attention else output