from __future__ import annotations
import torch
import torch.nn.functional as F


def banded_softmax_attention(
    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, w: int
) -> torch.Tensor:
    T = q.shape[-2]
    i = torch.arange(T, device=q.device)[:, None]
    j = torch.arange(T, device=q.device)[None, :]
    mask = (j <= i) & (j > i - w)
    return F.scaled_dot_product_attention(q, k, v, attn_mask=mask)


def swa_fused_forward(mixer, x: torch.Tensor) -> torch.Tensor:
    B, T, C = x.shape
    H, w = (mixer.n_head, mixer.swa_fused_window)
    D = C // H
    q = mixer.swa_fused_q(x).reshape(B, T, H, D).transpose(1, 2)
    k = mixer.swa_fused_k(x).reshape(B, T, H, D).transpose(1, 2)
    v = mixer.swa_fused_v(x).reshape(B, T, H, D).transpose(1, 2)
    o = banded_softmax_attention(q, k, v, w)
    o = o.transpose(1, 2).reshape(B, T, C)
    return x + mixer.swa_fused_o(o)


class CPUSlidingWindow:
    def __init__(self, mixer):
        self.mixer = mixer
        self.H = mixer.n_head
        self.C = mixer.swa_fused_q.out_features
        self.D = self.C // self.H
        self.w = mixer.swa_fused_window
        self.scale = self.D ** (-0.5)

    def init_state(self, batch_size: int = 1, device=None):
        if device is None:
            device = self.mixer.swa_fused_q.weight.device
        W1 = self.w - 1
        return {
            "k": torch.zeros(batch_size, W1, self.H, self.D, device=device),
            "v": torch.zeros(batch_size, W1, self.H, self.D, device=device),
            "n": 0,
        }

    def _project(self, x):
        B, L, _ = x.shape
        q = self.mixer.swa_fused_q(x).reshape(B, L, self.H, self.D)
        k = self.mixer.swa_fused_k(x).reshape(B, L, self.H, self.D)
        v = self.mixer.swa_fused_v(x).reshape(B, L, self.H, self.D)
        return (q, k, v)

    def forward_chunk(self, x: torch.Tensor, state):
        B, L, C = x.shape
        w = self.w
        if state is None:
            state = self.init_state(B, device=x.device)
        n = state["n"]
        W1 = w - 1
        ck = state["k"][:, W1 - n :, :, :] if n > 0 else state["k"][:, :0]
        cv = state["v"][:, W1 - n :, :, :] if n > 0 else state["v"][:, :0]
        q, k, v = self._project(x)
        k_all = torch.cat([ck, k], dim=1)
        v_all = torch.cat([cv, v], dim=1)
        Tk = n + L
        iq = torch.arange(L, device=x.device)[:, None] + n
        jk = torch.arange(Tk, device=x.device)[None, :]
        mask = (jk <= iq) & (jk > iq - w)
        qh = q.transpose(1, 2)
        kh = k_all.transpose(1, 2)
        vh = v_all.transpose(1, 2)
        o = F.scaled_dot_product_attention(qh, kh, vh, attn_mask=mask)
        o = o.transpose(1, 2).reshape(B, L, C)
        out = x + self.mixer.swa_fused_o(o)
        keep = min(W1, Tk)
        new_k = state["k"].clone()
        new_v = state["v"].clone()
        if keep > 0:
            new_k[:, W1 - keep :, :, :] = k_all[:, Tk - keep :, :, :]
            new_v[:, W1 - keep :, :, :] = v_all[:, Tk - keep :, :, :]
        new_state = {"k": new_k, "v": new_v, "n": keep}
        return (out, new_state)

    def step(self, x: torch.Tensor, state):
        if x.dim() == 2:
            x = x.unsqueeze(1)
        o, state = self.forward_chunk(x, state)
        return (o.squeeze(1), state)