"""Block-causal flex-attention mask over a packed span.

Rows come in contiguous runs of one (role, chunk). The BlockMask's full / partial 128x128 tiles are
derived from those few runs (never a token-level [S,S] mask; the builder follows SolarWM's
backends/minimax_h3/torch_mask.py, Apache-2.0), and mask_mod decides single tokens from a handful of
captured scalars, so the compiled kernel does not specialise on the caption length.

A span is packed [text | Picture 1 | DUV chunks 0..5 | audio | video chunks first .. first+n-1]. Roles:
STATIC (text, Picture 1, audio), DUV (the coarse proxy, one run per chunk) and VIDEO. Nothing sees a
proxy chunk later than its own (real-time control: the future proxy does not exist at inference).
"""
from __future__ import annotations

import math
from dataclasses import dataclass

import torch
from torch.nn.attention.flex_attention import BlockMask

from .layout import CHUNK_SIZES, FRAME_ROWS, NUM_CHUNKS, SPAN_CHUNKS

STATIC, DUV, VIDEO = 0, 1, 2
TILE = 128
assert len(set(CHUNK_SIZES)) == 1
CHUNK_ROWS = CHUNK_SIZES[0] * FRAME_ROWS


def df_allows(qr, qc, kr, kc, static_sees_first: bool) -> bool:
    """Diffusion forcing with a causal proxy.
        static    : static rows, and -- static_sees_first, i.e. the span starts the stream -- proxy and video
                    chunk 0 (with no history, chunk 0 needs condition rows that see it, as in the released
                    model; with a history, conditions that see only themselves do better)
        proxy p   : static rows and proxy chunks <= p
        video c   : static rows, proxy chunks <= c and video chunks c-5..c"""
    if qr == STATIC:
        return kr == STATIC or (static_sees_first and kc == 0 and kr in (DUV, VIDEO))
    if qr == DUV:
        return kr == STATIC or (kr == DUV and kc <= qc)
    return kr == STATIC or (kr == DUV and kc <= qc) or (kr == VIDEO and 0 <= qc - kc < SPAN_CHUNKS)


def _tiles(runs, allows, S, device):
    allowed = {(r, c): [(a, b) for a, b, r2, c2 in runs if allows(r, c, r2, c2)] for _, _, r, c in runs}
    nb = math.ceil(S / TILE)
    ks = torch.arange(nb) * TILE
    ke = torch.clamp(ks + TILE, max=S)
    full = torch.zeros(nb, nb, dtype=torch.bool)
    anyt = torch.zeros(nb, nb, dtype=torch.bool)
    for qb in range(nb):
        q0, q1 = qb * TILE, min(qb * TILE + TILE, S)
        rf, ra = torch.ones(nb, dtype=torch.bool), torch.zeros(nb, dtype=torch.bool)
        for a, b, r, c in runs:
            if a < q1 and q0 < b:
                sf, sa = torch.zeros(nb, dtype=torch.bool), torch.zeros(nb, dtype=torch.bool)
                for a2, b2 in allowed[(r, c)]:
                    sf |= (a2 <= ks) & (ke <= b2)
                    sa |= (ks < b2) & (a2 < ke)
                rf &= sf
                ra |= sa
        full[qb], anyt[qb] = rf, ra
    mb = 1 << (nb - 1).bit_length()

    def meta(m):
        o = torch.zeros(mb, mb, dtype=torch.int32)
        o[:nb, :nb] = m
        return (o.sum(-1).to(torch.int32)[None, None].to(device),
                torch.argsort(o, dim=-1, descending=True, stable=True).to(torch.int32)[None, None].to(device))

    pc, pi = meta(anyt & ~full)
    fc, fi = meta(full)
    return pc, pi, fc, fi


def _scalars(device, *values):
    return [torch.tensor(v, dtype=torch.long, device=device) for v in values]


def span_runs(head_rows: int, duv_rows: int, audio_rows: int, n_chunks: int, first_chunk: int = 0):
    """(runs, total rows) of [head | DUV x6 | audio | video chunks first..]; runs are (start, end, role, chunk)."""
    assert duv_rows % NUM_CHUNKS == 0, duv_rows
    dr = duv_rows // NUM_CHUNKS
    runs = [(0, head_rows, STATIC, -1)]
    runs += [(head_rows + p * dr, head_rows + (p + 1) * dr, DUV, p) for p in range(NUM_CHUNKS)]
    v0 = head_rows + duv_rows + audio_rows
    runs.append((head_rows + duv_rows, v0, STATIC, -1))
    runs += [(v0 + j * CHUNK_ROWS, v0 + (j + 1) * CHUNK_ROWS, VIDEO, first_chunk + j) for j in range(n_chunks)]
    return [r for r in runs if r[1] > r[0]], v0 + n_chunks * CHUNK_ROWS


def df_mask(head_rows: int, duv_rows: int, audio_rows: int, n_chunks: int, device, first_chunk: int = 0,
            static_sees_first: bool = True) -> BlockMask:
    """Mask of a packed span [text | Picture 1 | DUV chunks 0..5 | audio | video chunks first ..]."""
    runs, S = span_runs(head_rows, duv_rows, audio_rows, n_chunks, first_chunk)
    pc, pi, fc, fi = _tiles(runs, lambda *a: df_allows(*a, static_sees_first), S, device)
    d0, d1, v0, se, dr, cr, wb, first, ssf = _scalars(
        device, head_rows, head_rows + duv_rows, head_rows + duv_rows + audio_rows, S, duv_rows // NUM_CHUNKS,
        CHUNK_ROWS, SPAN_CHUNKS, first_chunk, int(static_sees_first))

    def classify(i):
        duv = (i >= d0) & (i < d1)
        vid = (i >= v0) & (i < se)
        chunk = torch.where(duv, torch.div(i - d0, dr, rounding_mode="floor"),
                            torch.div(i - v0, cr, rounding_mode="floor") + first)
        return ~duv & ~vid, duv, vid, chunk

    def mask_mod(_b, _h, qi, ki):
        q_st, q_du, q_vi, qc = classify(qi)
        k_st, k_du, k_vi, kc = classify(ki)
        first_rows = (ssf == 1) & (kc == 0) & (k_du | k_vi)
        return ((q_st & (k_st | first_rows)) | (q_du & (k_st | (k_du & (kc <= qc))))
                | (q_vi & (k_st | (k_du & (kc <= qc)) | (k_vi & (qc - kc >= 0) & (qc - kc < wb)))))

    return BlockMask.from_kv_blocks(pc, pi, fc, fi, BLOCK_SIZE=TILE, mask_mod=mask_mod, seq_lengths=(S, S))


@dataclass(frozen=True)
class DenseDF:
    """The df_mask of a training span (video chunks 0..n-1, n <= SPAN_CHUNKS) for dense FlashAttention.

    Under the rule every query run's keys are one contiguous prefix of one of two key orders:
        main  [static | DUV 0 | video 0 | DUV 1 | video 1 | ...]   static rows, video chunk c
        proxy [static | DUV 0 | DUV 1 | ...]                        proxy chunk p
    (static = text, Picture 1 and audio), so the masked attention is a few unmasked prefix calls with the
    same result as the BlockMask, without its per-tile mask evaluation (rtcwm.model.attention)."""
    head_rows: int
    duv_rows: int
    audio_rows: int
    n_chunks: int
    static_sees_first: bool = True

    def __post_init__(self):
        assert 0 < self.n_chunks <= SPAN_CHUNKS and self.duv_rows % NUM_CHUNKS == 0, self