# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
# CPU adaptation Copyright (c) 2026 malaiwah.
# KDA recurrence/gate and convolution contracts adapted from FLA at
# 9d981ffef3b361ba931102b633ae2a9fd91ca6c3. See LICENSE-FLA.
# This module is MIT licensed. No GPU imports, compilation, network or subprocesses.
"""Explicit Torch reference backend, not the upstream fused serving kernels."""
import torch
from torch import nn
from torch.nn import functional as F


def prepare_lens_from_mask(mask):
    return mask.sum(-1, dtype=torch.int32)


def prepare_cu_seqlens_from_mask(mask):
    return F.pad(prepare_lens_from_mask(mask).cumsum(0, dtype=torch.int32), (1, 0))


def tensor_cache(function):
    # Optimization-only upstream memoization is deliberately absent on CPU.
    return function


def rearrange(x, pattern, **axes):
    if pattern == 'b s ... -> (b s) ...':
        return x.flatten(0, 1)
    if pattern == '... (h d) -> ... h d':
        return x.unflatten(-1, (x.shape[-1] // axes['d'], axes['d']))
    if pattern == 'b t h d -> b t (h d)':
        return x.flatten(-2)
    raise ValueError('Unsupported reference reshape: ' + pattern)


def _segments(x, cu_seqlens):
    if cu_seqlens is None:
        return [(index, 0, x.shape[1]) for index in range(x.shape[0])]
    if x.shape[0] != 1:
        raise ValueError('Packed input must have batch size one')
    bounds = cu_seqlens.tolist()
    if bounds[0] != 0 or bounds[-1] != x.shape[1] or any(a > b for a, b in zip(bounds, bounds[1:])):
        raise ValueError('Invalid packed cumulative lengths')
    return [(0, a, b) for a, b in zip(bounds, bounds[1:])]


class ShortConvolution(nn.Conv1d):
    def __init__(self, hidden_size, kernel_size, activation='silu', bias=False):
        super().__init__(hidden_size, hidden_size, kernel_size, groups=hidden_size, bias=bias)
        if activation not in ('silu', 'swish', None):
            raise ValueError(activation)
        self.activation = activation

    def forward(self, x, cache=None, output_final_state=False, cu_seqlens=None):
        width = self.kernel_size[0]
        output = torch.empty_like(x)
        final = []
        for index, (batch, start, end) in enumerate(_segments(x, cu_seqlens)):
            state = x.new_zeros(self.in_channels, width) if cache is None else cache[index].clone()
            for offset in range(start, end):
                state = torch.cat((state[:, 1:], x[batch, offset, :, None]), -1)
                value = (state.float() * self.weight[:, 0].float()).sum(-1)
                if self.bias is not None:
                    value = value + self.bias.float()
                if self.activation is not None:
                    value = F.silu(value)
                output[batch, offset] = value.to(x.dtype)
            final.append(state)
        final = torch.stack(final)
        if cache is not None:
            cache.copy_(final)
        return output, final if output_final_state else None


class FusedRMSNormGated(nn.Module):
    def __init__(self, hidden_size, eps=1e-5, activation='sigmoid'):
        super().__init__()
        if activation != 'sigmoid':
            raise ValueError('This Kimi reference requires sigmoid output gating')
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, x, gate):
        normalized = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + self.eps)
        return (normalized * self.weight.float() * gate.float().sigmoid()).to(x.dtype)


def _reference_kda(q, k, v, g, beta, A_log=None, dt_bias=None, initial_state=None,
                   output_final_state=False, use_qk_l2norm_in_kernel=False,
                   use_gate_in_kernel=False, use_beta_sigmoid_in_kernel=False,
                   lower_bound=None, transpose_state_layout=False, cu_seqlens=None,
                   safe_gate=False):
    dtype = v.dtype
    if q.shape[2] != v.shape[2]:
        raise ValueError('Kimi K3 reference requires matching query/value heads')
    q, k, v, g, beta = [value.float() for value in (q, k, v, g, beta)]
    if use_qk_l2norm_in_kernel:
        # FLA l2norm uses rsqrt(sum(x*x) + 1e-6), not max(norm, eps).
        q = q * torch.rsqrt(q.square().sum(-1, keepdim=True) + 1e-6)
        k = k * torch.rsqrt(k.square().sum(-1, keepdim=True) + 1e-6)
    if use_gate_in_kernel:
        if dt_bias is not None:
            g = g + dt_bias.float().view(g.shape[-2:])
        if lower_bound is None:
            g = -A_log.float().exp().view(-1, 1) * F.softplus(g)
        else:
            g = lower_bound * (A_log.float().exp().view(-1, 1) * g).sigmoid()
    if use_beta_sigmoid_in_kernel:
        beta = beta.sigmoid()
    q = q * q.shape[-1] ** -0.5
    output = torch.empty_like(v)
    states = []
    for index, (batch, start, end) in enumerate(_segments(q, cu_seqlens)):
        state = q.new_zeros(q.shape[2], q.shape[3], v.shape[3])
        if initial_state is not None:
            state = initial_state[index].float().clone()
            if transpose_state_layout:
                state = state.transpose(-1, -2)
        for offset in range(start, end):
            key = k[batch, offset]
            state = state * g[batch, offset].exp().unsqueeze(-1)
            delta = v[batch, offset] - (key.unsqueeze(-1) * state).sum(-2)
            state = state + (beta[batch, offset, :, None] * key).unsqueeze(-1) * delta.unsqueeze(-2)
            output[batch, offset] = (q[batch, offset].unsqueeze(-1) * state).sum(-2)
        states.append(state.transpose(-1, -2).contiguous() if transpose_state_layout else state)
    return output.to(dtype), torch.stack(states) if output_final_state else None


def chunk_kda(**kwargs):
    """Evaluate chunk-mode API by its exact sequential recurrence on CPU."""
    return _reference_kda(**kwargs)


def fused_recurrent_kda(**kwargs):
    """Evaluate recurrent-mode API with the same reference state convention."""
    return _reference_kda(**kwargs)