"""Mirrors ``tilelang_ops/asc/compress.py``
(compress_and_update_state_prefill_tl / compress_and_update_state_decode_tl).

Softmax-weighted compression of ``COMPRESS_RATIO`` (1 or 2) consecutive (score, latent) fp32 rows into one
bf16 row, plus the rolling ``state_cache[slot, pos % state_cache_size, {score, latent}, :]`` update.

Numerics follow the Ascend kernel: ``m = max_r score``, ``e_r = exp(score_r - m)``, ``sum = sum_r e_r``,
``acc = sum_r latent_r * e_r`` (fp32, row order), ``out = bf16(acc * (1 / sum))``.

Order: prefill reads the previous chunk's tail from ``state_cache`` *before* storing the new tokens;
decode stores every token first and then reads the compression windows (as the two Ascend sub-kernels).
All index math is static-shape tensor ops (no host syncs).
"""

from collections.abc import Callable

import torch


def _compress_rows(score_rows: torch.Tensor, latent_rows: torch.Tensor) -> torch.Tensor:
    """score_rows / latent_rows: [N, R, D] fp32 -> [N, D] bf16 (Ascend operation order)."""
    from deepinfer.neuron_ops import nki_attention

    if nki_attention.supported_pool(score_rows):  # vllm-neuron#71's softmax_pool kernel (issue #11)
        return nki_attention.softmax_pool(score_rows, latent_rows)
    ratio = score_rows.shape[1]
    tile_max = score_rows[:, 0]
    for row in range(1, ratio):
        tile_max = torch.maximum(tile_max, score_rows[:, row])
    tile_sum = torch.zeros_like(tile_max)
    tile_acc = torch.zeros_like(tile_max)
    for row in range(ratio):
        e = torch.exp(score_rows[:, row] - tile_max)
        tile_sum = tile_sum + e
        tile_acc = tile_acc + latent_rows[:, row] * e
    inv_sum = 1.0 / tile_sum
    return (tile_acc * inv_sum).to(torch.bfloat16)


def compress_and_update_state_prefill_tl(
    DIM: int,
    COMPRESS_RATIO: int,
    HAS_INDEX: bool,
) -> Callable[..., None]:
    assert COMPRESS_RATIO in (1, 2)
    ratio = COMPRESS_RATIO

    def _compress_and_update_state_prefill(
        score: torch.Tensor,
        latent: torch.Tensor,
        state_cache: torch.Tensor,
        state_block_idx: torch.Tensor,
        compressed_batch_mapping: torch.Tensor,
        cu_seqlens_q: torch.Tensor,
        cu_seqlens_k: torch.Tensor,
        cu_compressed_seqlens_q: torch.Tensor,
        mirror_cp_reverse_index: torch.Tensor | None,
        kv_compressed: torch.Tensor,
    ) -> None:
        assert score.shape[-1] == DIM and state_cache.shape[-1] == DIM
        assert (mirror_cp_reverse_index is not None) == HAS_INDEX
        device = score.device
        state_size = state_cache.shape[1]
        num_seqs = state_block_idx.shape[0]
        total_c = compressed_batch_mapping.shape[0]
        total_q = mirror_cp_reverse_index.shape[0] if HAS_INDEX else score.shape[0]
        num_rows = score.shape[0]

        cu_q = cu_seqlens_q.to(torch.int64)
        cu_k = cu_seqlens_k.to(torch.int64)
        cu_c = cu_compressed_seqlens_q.to(torch.int64)
        qs_all, qe_all = cu_q[:-1], cu_q[1:]
        qlen_all = qe_all - qs_all
        klen_all = cu_k[1:] - cu_k[:-1]
        qstart_all = klen_all - qlen_all
        state_all = state_block_idx.to(torch.int64)
        reverse = mirror_cp_reverse_index.to(torch.int64) if HAS_INDEX else None

        def to_row(orig_idx: torch.Tensor) -> torch.Tensor:
            orig_idx = orig_idx.clamp(0, max(total_q - 1, 0))
            row = reverse.index_select(0, orig_idx.reshape(-1)).view(orig_idx.shape) if HAS_INDEX else orig_idx
            return row.clamp(0, max(num_rows - 1, 0))

        # ---- phase 1: compression (reads old state before it is overwritten) ----
        if total_c > 0:
            seq = compressed_batch_mapping.to(torch.int64).clamp(0, num_seqs - 1)
            qs, qstart, klen = qs_all[seq], qstart_all[seq], klen_all[seq]
            sidx = state_all[seq]
            total_new_blocks = klen // ratio - qstart // ratio
            block_id = torch.arange(total_c, dtype=torch.int64, device=device) - cu_c[seq]
            valid = (block_id >= 0) & (block_id < total_new_blocks)
            block_start = (qstart // ratio + block_id) * ratio
            pos = block_start[:, None] + torch.arange(ratio, dtype=torch.int64, device=device)[None, :]  # [C, R]
            from_cache = pos < qstart[:, None]
            cache_pos = torch.remainder(pos, state_size)  # pos >= 0 for valid blocks (truncmod == remainder)
            cache_slot = sidx[:, None].expand_as(pos)
            cached = state_cache[cache_slot, cache_pos]  # [C, R, 2, D]
            rows = to_row(pos - qstart[:, None] + qs[:, None])
            s_rows = torch.where(from_cache[..., None], cached[:, :, 0], score[rows].float())
            l_rows = torch.where(from_cache[..., None], cached[:, :, 1], latent[rows].float())
            result = _compress_rows(s_rows, l_rows)
            kv_compressed.copy_(torch.where(valid[:, None], result, torch.zeros_like(result)))

        # ---- phase 2: store the last min(qlen, state_size) tokens of every sequence ----
        if total_q == 0 or num_seqs == 0:
            return
        tok = torch.arange(total_q, dtype=torch.int64, device=device)
        seq_t = torch.searchsorted(qe_all, tok, right=True)
        tok_valid = seq_t < num_seqs
        seq_t = seq_t.clamp_max(num_seqs - 1)
        rel = tok - qs_all[seq_t]
        qlen_t = qlen_all[seq_t]
        # Tokens older than the last ``state_size`` ones are redirected to the stored token occupying the
        # same slot, so duplicate writes carry identical values (no masking / host sync needed).
        first_stored = (qlen_t - state_size).clamp_min(0)
        lag = (first_stored - rel).clamp_min(0)
        rel_src = rel + (lag + state_size - 1) // state_size * state_size
        # Rows past cu_seqlens_q[-1] (if any) duplicate the write of the globally last token.
        anchor = (qe_all[-1] - 1).clamp_min(0)
        anchor_seq = torch.searchsorted(qe_all, anchor[None], right=True).clamp_max(num_seqs - 1)[0]
        src_tok = torch.where(tok_valid, qs_all[seq_t] + rel_src, anchor)
        seq_w = torch.where(tok_valid, seq_t, anchor_seq)
        rel_w = src_tok - qs_all[seq_w]
        any_token = qe_all[-1] > 0
        slot_w = state_all[seq_w]
        pos_w = torch.remainder(qstart_all[seq_w] + rel_w, state_size)
        rows = to_row(src_tok)
        new_vals = torch.stack((score[rows].float(), latent[rows].float()), dim=1)  # [T, 2, D]
        old_vals = state_cache[slot_w, pos_w]
        vals = torch.where(any_token, new_vals, old_vals)
        state_cache.index_put_((slot_w, pos_w), vals.to(state_cache.dtype))

    return _compress_and_update_state_prefill


def compress_and_update_state_decode_tl(
    DIM: int,
    COMPRESS_RATIO: int,
) -> Callable[..., None]:
    assert COMPRESS_RATIO in (1, 2)
    ratio = COMPRESS_RATIO

    def _compress_and_update_state_decode(
        score: torch.Tensor,
        latent: torch.Tensor,
        state_cache: torch.Tensor,
        state_block_idx: torch.Tensor,
        context_lens: torch.Tensor,
        kv_compressed: torch.Tensor,
    ) -> None:
        assert score.shape[-1] == DIM and state_cache.shape[-1] == DIM
        total_q = score.shape[0]
        if total_q == 0:
            return
        device = score.device
        state_size = state_cache.shape[1]
        sidx = state_block_idx.to(torch.int64)
        klen = context_lens.to(torch.int64)
        abs_pos = klen - 1

        # Store all tokens before computing (windows may include tokens of this step).
        cache_pos = torch.remainder(abs_pos, state_size)
        vals = torch.stack((score.float(), latent.float()), dim=1)  # [Q, 2, D]
        state_cache.index_put_((sidx, cache_pos), vals.to(state_cache.dtype))

        is_end = torch.remainder(abs_pos, ratio) == ratio - 1
        block_start = abs_pos.div(ratio, rounding_mode="floor") * ratio
        pos = block_start[:, None] + torch.arange(ratio, dtype=torch.int64, device=device)[None, :]
        win = state_cache[sidx[:, None].expand_as(pos), torch.remainder(pos, state_size)]  # [Q, R, 2, D]
        result = _compress_rows(win[:, :, 0], win[:, :, 1])
        kv_compressed.copy_(torch.where(is_end[:, None], result, torch.zeros_like(result)))

    return _compress_and_update_state_decode