"""Diagonal Kalman Delta Network token mixer (arXiv:2609.07816)."""

from __future__ import annotations

import math
from typing import Any

import torch
from einops import rearrange
from fla.layers.utils import get_layer_cache, update_layer_cache
from fla.modules import FusedRMSNormGated, ShortConvolution
from torch import nn

from tiny_gdn.kdn import recurrence
from tiny_gdn.kdn.kernels import (
    _diag_kdn_frontend_precumsum_paired,
    _diag_kdn_gain,
    _diag_kdn_normalize,
    chunk_kalman,
)

SUPPORTED_HEAD_DIMS = (64, 128)


class DiagonalKalmanDeltaNet(nn.Module):
    """Delta-rule memory whose write strength is a per-channel Kalman gain.

    Each head keeps a diagonal posterior covariance over its key channels:
    channels that have absorbed little evidence take large corrections and
    well-determined ones small corrections, replacing KDA's scalar ``beta``.
    Stateless CUDA BF16 calls run the fused Triton kernels. Cached prefill,
    decode, padding masks, and other dtypes use the exact PyTorch recurrences.

    The cache holds ``recurrent_state = (memory [B,H,K,K], precision [B,H,K])``
    in float32 plus the three short-convolution states.
    """

    def __init__(
        self,
        hidden_size: int,
        head_dim: int,
        num_heads: int,
        *,
        conv_size: int = 4,
        norm_eps: float = 1e-5,
        layer_idx: int | None = None,
    ) -> None:
        super().__init__()
        if head_dim not in SUPPORTED_HEAD_DIMS:
            raise ValueError(
                f"Diagonal KDN head_dim must be one of {SUPPORTED_HEAD_DIMS}, got {head_dim}"
            )
        if num_heads <= 0 or hidden_size <= 0:
            raise ValueError("hidden_size and num_heads must be positive")
        if conv_size <= 0:
            raise ValueError("conv_size must be positive")
        self.hidden_size = hidden_size
        self.head_dim = head_dim
        self.num_heads = num_heads
        self.key_dim = num_heads * head_dim
        self.layer_idx = layer_idx
        self.scale = head_dim**-0.5
        self.info_scale = float(head_dim)

        self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
        self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
        self.v_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
        self.q_conv1d, self.k_conv1d, self.v_conv1d = (
            ShortConvolution(
                hidden_size=self.key_dim,
                kernel_size=conv_size,
                bias=False,
                activation="silu",
            )
            for _ in range(3)
        )

        self.f_proj = nn.Sequential(
            nn.Linear(hidden_size, head_dim, bias=False),
            nn.Linear(head_dim, self.key_dim, bias=False),
        )
        self.A_log = nn.Parameter(
            torch.log(torch.empty(num_heads, dtype=torch.float32).uniform_(1, 16))
        )
        self.A_log._no_weight_decay = True
        dt = torch.exp(
            torch.rand(self.key_dim, dtype=torch.float32)
            * (math.log(0.1) - math.log(0.001))
            + math.log(0.001)
        ).clamp(min=1e-4)
        self.dt_bias = nn.Parameter(dt + torch.log(-torch.expm1(-dt)))
        self.dt_bias._no_weight_decay = True

        self.omega_proj = nn.Sequential(
            nn.Linear(hidden_size, head_dim, bias=False),
            nn.Linear(head_dim, self.key_dim, bias=True),
        )
        self.r_proj = nn.Linear(hidden_size, num_heads, bias=True)
        self.initial_precision_param = nn.Parameter(
            torch.full(
                (num_heads,),
                recurrence.initial_precision_logit(),
                dtype=torch.float32,
            )
        )
        self.initial_precision_param._no_weight_decay = True

        self.g_proj = nn.Sequential(
            nn.Linear(hidden_size, head_dim, bias=False),
            nn.Linear(head_dim, self.key_dim, bias=True),
        )
        self.o_norm = FusedRMSNormGated(head_dim, activation="sigmoid", eps=norm_eps)
        self.o_proj = nn.Linear(self.key_dim, hidden_size, bias=False)
        self.reset_noise_parameters()

    def reset_noise_parameters(self) -> None:
        """Isotropic process noise and a neutral observation-noise bias at init.

        Blanket initializers overwrite these, so callers re-run this afterwards.
        """

        nn.init.zeros_(self.omega_proj[1].weight)
        nn.init.zeros_(self.omega_proj[1].bias)
        nn.init.zeros_(self.r_proj.bias)

    def _fused_supported(self, hidden_states: torch.Tensor) -> bool:
        return hidden_states.is_cuda and hidden_states.dtype == torch.bfloat16

    def _split_heads(self, tensor: torch.Tensor) -> torch.Tensor:
        return rearrange(tensor, "... (h d) -> ... h d", d=self.head_dim).contiguous()

    def _fused(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        retention_raw: torch.Tensor,
        omega_raw: torch.Tensor,
        r_raw: torch.Tensor,
    ) -> torch.Tensor:
        q_memory, k_gain, k_memory = _diag_kdn_normalize(q, k)
        g_cumsum, alpha, omega, r, initial_precision = (
            _diag_kdn_frontend_precumsum_paired(
                retention_raw,
                omega_raw,
                r_raw,
                self.A_log.float().contiguous(),
                self.dt_bias.float().view(self.num_heads, self.head_dim).contiguous(),
                self.initial_precision_param.float().contiguous(),
            )
        )
        kappa = _diag_kdn_gain(
            k_gain, alpha, omega, r, initial_precision, omega_raw, r_raw
        )
        return chunk_kalman(q_memory, k_memory, kappa, v, g_cumsum, self.scale)

    def _stateful(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        retention_raw: torch.Tensor,
        omega_raw: torch.Tensor,
        r_raw: torch.Tensor,
        state: tuple[torch.Tensor, torch.Tensor] | None,
        token_mask: torch.Tensor | None,
    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        log_alpha, omega, r, initial_precision = recurrence.kalman_noise(
            retention_raw,
            omega_raw,
            r_raw,
            self.A_log,
            self.dt_bias,
            self.initial_precision_param,
        )
        q, k = recurrence.normalize_query_key(q, k)
        v = v.to(q.dtype)
        if token_mask is not None:
            k, v, log_alpha, omega = recurrence.mask_padding(
                token_mask, k, v, log_alpha, omega
            )
        if state is None:
            memory, precision = recurrence.initial_state(
                initial_precision,
                batch_size=q.shape[0],
                key_dim=self.head_dim,
                value_dim=self.head_dim,
            )
        else:
            memory, precision = state
        if q.shape[1] == 1:
            output, memory, precision = recurrence.diag_kdn_step(
                q[:, 0],
                k[:, 0],
                v[:, 0],
                log_alpha[:, 0],
                omega[:, 0],
                r[:, 0],
                memory,
                precision,
                scale=self.scale,
                info_scale=self.info_scale,
            )
            output = output[:, None]
        else:
            output, memory, precision = recurrence.diag_kdn_chunked(
                q,
                k,
                v,
                log_alpha,
                omega,
                r,
                memory,
                precision,
                scale=self.scale,
                info_scale=self.info_scale,
            )
        return output, (memory, precision)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        past_key_values: Any | None = None,
        use_cache: bool = False,
        output_attentions: bool = False,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, None, Any | None]:
        if hidden_states.ndim != 3 or hidden_states.shape[-1] != self.hidden_size:
            raise ValueError(
                f"hidden_states must have shape [B, T, {self.hidden_size}]"
            )
        if attention_mask is not None and attention_mask.ndim != 2:
            raise ValueError("attention_mask must be a [batch, sequence] padding mask")
        if output_attentions:
            raise ValueError("Diagonal KDN has no attention weights to return")
        length = hidden_states.shape[1]
        last_state = get_layer_cache(self, past_key_values)
        token_mask = (
            None
            if attention_mask is None
            else attention_mask[:, -length:].to(device=hidden_states.device, dtype=torch.bool)
        )
        conv_mask = (
            None if token_mask is None else token_mask.to(hidden_states.dtype)
        )
        conv_q, conv_k, conv_v = (
            (None, None, None) if last_state is None else last_state["conv_state"]
        )
        q, conv_q = self.q_conv1d(
            x=self.q_proj(hidden_states),
            mask=conv_mask,
            cache=conv_q,
            output_final_state=use_cache,
        )
        k, conv_k = self.k_conv1d(
            x=self.k_proj(hidden_states),
            mask=conv_mask,
            cache=conv_k,
            output_final_state=use_cache,
        )
        v, conv_v = self.v_conv1d(
            x=self.v_proj(hidden_states),
            mask=conv_mask,
            cache=conv_v,
            output_final_state=use_cache,
        )
        q, k, v = (self._split_heads(tensor) for tensor in (q, k, v))
        retention_raw = self._split_heads(self.f_proj(hidden_states))
        omega_raw = self._split_heads(self.omega_proj(hidden_states))
        r_raw = self.r_proj(hidden_states).contiguous()

        recurrent_state = None if last_state is None else last_state["recurrent_state"]
        stateless = recurrent_state is None and token_mask is None and not use_cache
        if stateless and self._fused_supported(hidden_states):
            output = self._fused(q, k, v, retention_raw, omega_raw, r_raw)
        else:
            output, recurrent_state = self._stateful(
                q, k, v, retention_raw, omega_raw, r_raw, recurrent_state, token_mask
            )
        update_layer_cache(
            self,
            past_key_values,
            recurrent_state=recurrent_state,
            conv_state=(conv_q, conv_k, conv_v),
            offset=length,
        )
        output = self.o_norm(
            output.to(v.dtype),
            self._split_heads(self.g_proj(hidden_states)),
        )
        output = self.o_proj(rearrange(output, "b t h d -> b t (h d)"))
        return output, None, past_key_values