"""Selective-token forward for the block-causal Wan DiT.

MotionCache (token granularity) and FlowCache (frame-group granularity) both need
to run the transformer stack on a *subset* of the current chunk's tokens while the
remaining tokens keep the activations cached from the previous computed step.

For a KV-cached causal model that means, per self-attention layer:
  * compute q/k/v only for the selected rows,
  * scatter the fresh k/v into the slots the current chunk owns in the KV cache,
    leaving the unselected slots holding the previous step's k/v,
  * attend the selected queries against the *full* cache (history + current chunk).

This mirrors ``forward_selective`` in
``MotionCache4SkyReels-V2/skyreels_v2_infer/modules/transformer_motion.py``,
adapted from SkyReels' persistent ``k_cache/v_cache`` buffers to Self-Forcing's
rolling ``kv_cache[block]`` dict.
"""

import math

import torch

from wan.modules.attention import attention


def build_causal_freqs(grid_sizes, freqs, start_frame):
    """Return the per-token RoPE multipliers [seq_len, 1, c] for one chunk.

    Same construction as ``causal_rope_apply`` in ``wan/modules/causal_model.py``,
    factored out so it can be indexed by a token subset and reused across layers.
    """
    c = freqs.size(1)
    parts = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
    f, h, w = grid_sizes[0].tolist()
    seq_len = f * h * w
    freqs_i = torch.cat(
        [
            parts[0][start_frame:start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1),
            parts[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
            parts[2][:w].view(1, 1, w, -1).expand(f, h, w, -1),
        ],
        dim=-1,
    ).reshape(seq_len, 1, -1)
    return freqs_i


def rope_apply_rows(x, freqs_rows):
    """Apply precomputed RoPE multipliers to x [B, S, n, d] given freqs [S, 1, d/2]."""
    b, s, n, _ = x.shape
    x_c = torch.view_as_complex(x.to(torch.float64).reshape(b, s, n, -1, 2))
    out = torch.view_as_real(x_c * freqs_rows.unsqueeze(0)).flatten(3)
    return out.type_as(x)


def self_attn_selective(attn, x_sel, sel_idx, freqs_rows, kv_cache,
                        current_start, block_seq_len):
    """Self-attention over a token subset, refreshing only those KV slots.

    Args:
        attn: a ``CausalWanSelfAttention`` module.
        x_sel: [B, n_sel, C] modulated hidden states for the selected tokens.
        sel_idx: LongTensor [n_sel], token offsets inside the current chunk.
        freqs_rows: [block_seq_len, 1, d/2] RoPE multipliers for the whole chunk.
        kv_cache: this layer's ``{"k", "v", "global_end_index", "local_end_index"}``.
        current_start: token offset of the chunk in global sequence coordinates.
        block_seq_len: number of tokens the chunk owns (all frames of the chunk).
    """
    b, n_sel = x_sel.shape[:2]
    n, d = attn.num_heads, attn.head_dim

    q = attn.norm_q(attn.q(x_sel)).view(b, n_sel, n, d)
    k = attn.norm_k(attn.k(x_sel)).view(b, n_sel, n, d)
    v = attn.v(x_sel).view(b, n_sel, n, d)

    rows = freqs_rows[sel_idx]
    roped_q = rope_apply_rows(q, rows)
    roped_k = rope_apply_rows(k, rows).type_as(v)

    # Index arithmetic must use the *full* chunk length: the chunk still occupies
    # the same KV slots, we are only refreshing some of them.
    current_end = current_start + block_seq_len
    local_end_index = (kv_cache["local_end_index"].item() + current_end
                       - kv_cache["global_end_index"].item())
    local_start_index = local_end_index - block_seq_len

    slots = local_start_index + sel_idx
    kv_cache["k"].index_copy_(1, slots, roped_k)
    kv_cache["v"].index_copy_(1, slots, v)

    lo = max(0, local_end_index - attn.max_attention_size)
    out = attention(roped_q, kv_cache["k"][:, lo:local_end_index],
                    kv_cache["v"][:, lo:local_end_index])

    kv_cache["global_end_index"].fill_(current_end)
    kv_cache["local_end_index"].fill_(local_end_index)

    return attn.o(out.flatten(2))


def block_forward_selective(block, x_sel, sel_idx, frame_of_sel, e0, freqs_rows,
                            context, kv_cache, crossattn_cache,
                            current_start, block_seq_len):
    """One ``CausalWanAttentionBlock`` evaluated on a token subset.

    ``frame_of_sel`` maps each selected token to its frame index inside the chunk,
    so the per-frame AdaLN modulation can be gathered for the selected rows only.
    """
    # e0: [B, F, 6, C] -> six tensors of [B, F, 1, C]; gather per selected token.
    e = (block.modulation.unsqueeze(1) + e0).chunk(6, dim=2)
    e = [t.squeeze(2)[:, frame_of_sel] for t in e]  # each [B, n_sel, C]

    y = self_attn_selective(
        block.self_attn,
        block.norm1(x_sel) * (1 + e[1]) + e[0],
        sel_idx, freqs_rows, kv_cache, current_start, block_seq_len,
    )
    x_sel = x_sel + y * e[2]

    x_sel = x_sel + block.cross_attn(block.norm3(x_sel), context, None,
                                     crossattn_cache=crossattn_cache)
    y = block.ffn(block.norm2(x_sel) * (1 + e[4]) + e[3])
    x_sel = x_sel + y * e[5]
    return x_sel


def run_blocks_selective(model, x, sel_idx, e0, kwargs, kv_cache, crossattn_cache,
                         current_start, cache_start, grid_sizes):
    """Run all transformer blocks on ``x[:, sel_idx]`` and return the updated rows.

    Returns a tensor [B, n_sel, C] holding the new hidden states for the selected
    tokens (the caller merges them into the cached full-length tensor).
    """
    f, h, w = grid_sizes[0].tolist()
    frame_seqlen = h * w
    block_seq_len = f * frame_seqlen
    frame_of_sel = torch.div(sel_idx, frame_seqlen, rounding_mode="floor")

    freqs_rows = build_causal_freqs(grid_sizes, kwargs["freqs"],
                                    current_start // frame_seqlen)

    x_sel = x[:, sel_idx]
    for i, block in enumerate(model.blocks):
        x_sel = block_forward_selective(
            block, x_sel, sel_idx, frame_of_sel, e0, freqs_rows,
            kwargs["context"], kv_cache[i], crossattn_cache[i],
            current_start, block_seq_len,
        )
    return x_sel