"""Exact PyTorch Diagonal KDN recurrences that carry state across calls.

The Triton kernels cover stateless fixed-length training. These functions
take and return ``(memory, precision)`` so prefill, decode, and padded batches
continue a sequence exactly; the sequential form is also the test oracle.

Per head, with one covariance coordinate ``p`` per key channel::

    p_hat_t = alpha_t**2 * p_{t-1} + omega_t
    kappa_t = p_hat_t * k_t / (r_t + sum_i p_hat_ti * k_ti**2)
    1 / p_t = 1 / p_hat_t + info_scale * k_t**2 / r_t
    S_t     = (I - kappa_t k_t^T) diag(alpha_t) S_{t-1} + kappa_t v_t^T
    o_t     = scale * S_t^T q_t

Shapes: ``q, k, log_alpha, omega`` are ``[B, T, H, K]``, ``v`` is
``[B, T, H, V]``, ``r`` is ``[B, T, H]``, ``memory`` is ``[B, H, K, V]`` and
``precision`` (``1 / p``) is ``[B, H, K]``. All math runs in float32.
"""

from __future__ import annotations

import math

import torch
import torch.nn.functional as F

OMEGA_MIN = 0.0
R_MIN = 0.01
INITIAL_PRECISION = 1.0
INITIAL_PRECISION_FLOOR = 0.1
LOG_MEMORY_ALPHA_FLOOR = math.log(1e-6)
NORM_EPS = 1e-12
DEFAULT_CHUNK_SIZE = 64


def initial_precision_logit() -> float:
    """Raw parameter value whose floored softplus equals ``INITIAL_PRECISION``."""

    return math.log(math.expm1(INITIAL_PRECISION - INITIAL_PRECISION_FLOOR))


def _work(tensor: torch.Tensor) -> torch.Tensor:
    return tensor.to(torch.promote_types(tensor.dtype, torch.float32))


def kalman_noise(
    retention_raw: torch.Tensor,
    omega_raw: torch.Tensor,
    r_raw: torch.Tensor,
    A_log: torch.Tensor,
    dt_bias: torch.Tensor,
    initial_precision_param: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Return ``(log_alpha, omega, r, initial_precision)`` like the Triton frontend."""

    heads, key_dim = retention_raw.shape[-2:]
    rate = _work(A_log).exp().view(heads, 1)
    log_alpha = -rate * F.softplus(
        _work(retention_raw) + _work(dt_bias).view(heads, key_dim)
    )
    omega = OMEGA_MIN + F.softplus(_work(omega_raw))
    r = R_MIN + F.softplus(_work(r_raw))
    initial_precision = INITIAL_PRECISION_FLOOR + F.softplus(
        _work(initial_precision_param)
    )
    return log_alpha, omega, r, initial_precision


def normalize_query_key(
    q: torch.Tensor, k: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    return (
        F.normalize(_work(q), dim=-1, eps=NORM_EPS),
        F.normalize(_work(k), dim=-1, eps=NORM_EPS),
    )


def initial_state(
    initial_precision: torch.Tensor,
    *,
    batch_size: int,
    key_dim: int,
    value_dim: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Zero memory and the learned per-head prior precision for new sequences."""

    heads = initial_precision.shape[0]
    prior = _work(initial_precision)
    memory = prior.new_zeros((batch_size, heads, key_dim, value_dim))
    precision = prior.view(1, heads, 1).expand(batch_size, heads, key_dim).contiguous()
    return memory, precision


def mask_padding(
    token_mask: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    log_alpha: torch.Tensor,
    omega: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Turn padded positions into identity steps that leave the state untouched."""

    keep = token_mask[:, :, None, None]
    return (
        torch.where(keep, k, 0.0),
        torch.where(keep, v, 0.0),
        torch.where(keep, log_alpha, 0.0),
        torch.where(keep, omega, 0.0),
    )


def diag_kdn_step(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    log_alpha: torch.Tensor,
    omega: torch.Tensor,
    r: torch.Tensor,
    memory: torch.Tensor,
    precision: torch.Tensor,
    *,
    scale: float,
    info_scale: float,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Advance one token: ``q, k, log_alpha, omega`` are ``[B, H, K]``."""

    k_squared = k.square()
    predicted = (2.0 * log_alpha).exp() / precision + omega
    noise = r[..., None]
    gain = predicted * k / (noise + (predicted * k_squared).sum(-1, keepdim=True))
    precision = predicted.reciprocal() + info_scale * k_squared / noise
    decayed = memory * log_alpha.clamp_min(LOG_MEMORY_ALPHA_FLOOR).exp()[..., None]
    innovation = v - torch.einsum("bhk,bhkv->bhv", k, decayed)
    memory = decayed + gain[..., None] * innovation[..., None, :]
    output = scale * torch.einsum("bhk,bhkv->bhv", q, memory)
    return output, memory, precision


def diag_kdn_recurrent(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    log_alpha: torch.Tensor,
    omega: torch.Tensor,
    r: torch.Tensor,
    memory: torch.Tensor,
    precision: torch.Tensor,
    *,
    scale: float,
    info_scale: float,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Token-by-token reference; returns ``(output [B,T,H,V], memory, precision)``."""

    outputs = []
    for token in range(q.shape[1]):
        output, memory, precision = diag_kdn_step(
            q[:, token],
            k[:, token],
            v[:, token],
            log_alpha[:, token],
            omega[:, token],
            r[:, token],
            memory,
            precision,
            scale=scale,
            info_scale=info_scale,
        )
        outputs.append(output)
    return torch.stack(outputs, dim=1), memory, precision


def _compose(
    later: tuple[torch.Tensor, ...], earlier: tuple[torch.Tensor, ...]
) -> tuple[torch.Tensor, ...]:
    """Product ``later @ earlier`` of 2x2 maps stored as (m00, m01, m10, m11)."""

    a00, a01, a10, a11 = later
    b00, b01, b10, b11 = earlier
    product = (
        a00 * b00 + a01 * b10,
        a00 * b01 + a01 * b11,
        a10 * b00 + a11 * b10,
        a10 * b01 + a11 * b11,
    )
    # Mobius maps are invariant to scaling, so renormalize to stay in range.
    magnitude = torch.stack(product).abs().amax(dim=0).clamp_min(
        torch.finfo(product[0].dtype).tiny
    )
    return tuple(entry / magnitude for entry in product)


def _inclusive_scan(maps: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]:
    """Hillis-Steele prefix products ``M_t @ ... @ M_1`` along dim 1."""

    length = maps[0].shape[1]
    offset = 1
    while offset < length:
        combined = _compose(
            tuple(entry[:, offset:] for entry in maps),
            tuple(entry[:, :-offset] for entry in maps),
        )
        maps = tuple(
            torch.cat((entry[:, :offset], update), dim=1)
            for entry, update in zip(maps, combined)
        )
        offset *= 2
    return maps


def kalman_gain(
    k: torch.Tensor,
    log_alpha: torch.Tensor,
    omega: torch.Tensor,
    r: torch.Tensor,
    precision: torch.Tensor,
    *,
    info_scale: float,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Vector gains for every token via a parallel Mobius scan, plus final precision."""

    k_squared = k.square()
    noise = r[..., None]
    alpha_squared = (2.0 * log_alpha).exp()
    information = info_scale * k_squared / noise
    maps = _inclusive_scan(
        (
            alpha_squared,
            omega,
            information * alpha_squared,
            1.0 + information * omega,
        )
    )
    start = precision.reciprocal()[:, None]
    covariance = (maps[0] * start + maps[1]) / (maps[2] * start + maps[3])
    previous = torch.cat((start, covariance[:, :-1]), dim=1)
    predicted = alpha_squared * previous + omega
    gain = predicted * k / (noise + (predicted * k_squared).sum(-1, keepdim=True))
    return gain, covariance[:, -1].reciprocal()


def chunked_memory(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    log_alpha: torch.Tensor,
    gain: torch.Tensor,
    memory: torch.Tensor,
    *,
    scale: float,
    chunk_size: int = DEFAULT_CHUNK_SIZE,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Two-key delta-rule memory in WY form, one triangular solve per chunk."""

    if chunk_size <= 0:
        raise ValueError("chunk_size must be positive")
    log_alpha = log_alpha.clamp_min(LOG_MEMORY_ALPHA_FLOOR)
    length = q.shape[1]
    outputs = []
    for start in range(0, length, chunk_size):
        stop = min(start + chunk_size, length)
        size = stop - start
        q_chunk, k_chunk, v_chunk, gain_chunk = (
            tensor[:, start:stop].transpose(1, 2) for tensor in (q, k, v, gain)
        )
        cumulative = log_alpha[:, start:stop].transpose(1, 2).cumsum(dim=2)
        causal = torch.ones(size, size, dtype=torch.bool, device=q.device).tril()
        differences = cumulative[:, :, :, None, :] - cumulative[:, :, None, :, :]
        weights = (
            torch.where(causal[:, :, None], differences, -math.inf).exp()
            * gain_chunk[:, :, None, :, :]
        )
        interaction = torch.einsum("bhtk,bhtsk->bhts", k_chunk, weights)
        readout = torch.einsum("bhtk,bhtsk->bhts", q_chunk, weights)
        decay = cumulative.exp()
        corrected = torch.linalg.solve_triangular(
            interaction,
            v_chunk - (decay * k_chunk) @ memory,
            upper=False,
            unitriangular=True,
        )
        outputs.append(
            (scale * ((q_chunk * decay) @ memory + readout @ corrected)).transpose(1, 2)
        )
        to_end = (cumulative[:, :, -1:] - cumulative).exp() * gain_chunk
        memory = memory * decay[:, :, -1, :, None] + to_end.transpose(-1, -2) @ corrected
    return torch.cat(outputs, dim=1), memory


def diag_kdn_chunked(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    log_alpha: torch.Tensor,
    omega: torch.Tensor,
    r: torch.Tensor,
    memory: torch.Tensor,
    precision: torch.Tensor,
    *,
    scale: float,
    info_scale: float,
    chunk_size: int = DEFAULT_CHUNK_SIZE,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Parallel form of :func:`diag_kdn_recurrent` with the same signature."""

    gain, precision = kalman_gain(
        k, log_alpha, omega, r, precision, info_scale=info_scale
    )
    output, memory = chunked_memory(
        q, k, v, log_alpha, gain, memory, scale=scale, chunk_size=chunk_size
    )
    return output, memory, precision