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


class CPULandmark:
    def __init__(self, mixer):
        self.mixer = mixer
        self.H = mixer.n_head
        self.D = mixer.d_head
        self.C = self.H * self.D
        self.chunk = mixer.chunk
        self.max_land = mixer.max_land

    def init_state(self, batch_size: int = 1, device=None, c: int | None = None):
        if device is None:
            device = self.mixer.q_proj.weight.device
        return {
            "k": torch.zeros(batch_size, self.max_land, self.H, self.D, device=device),
            "v": torch.zeros(batch_size, self.max_land, self.H, self.D, device=device),
            "n_land": torch.zeros(batch_size, dtype=torch.long, device=device),
            "acc_sum": torch.zeros(batch_size, self.C, device=device),
            "acc_cnt": torch.zeros(batch_size, dtype=torch.long, device=device),
            "c": int(c) if c else self.chunk,
        }

    def _sink_kv(self, batch_size: int, device):
        sink = self.mixer.sink
        sk = self.mixer.k_proj(sink).reshape(1, 1, self.H, self.D)
        sv = self.mixer.v_proj(sink).reshape(1, 1, self.H, self.D)
        exp = (batch_size, 1, self.H, self.D)
        return (sk.to(device).expand(*exp), sv.to(device).expand(*exp))

    def step(self, x: torch.Tensor, state):
        if x.dim() == 3:
            x = x.squeeze(1)
        B = x.shape[0]
        if state is None:
            state = self.init_state(B, device=x.device)
        c = state["c"]
        n = state["n_land"]
        Lmax = self.max_land
        sk, sv = self._sink_kv(B, x.device)
        keys = torch.cat([sk, state["k"]], dim=1)
        vals = torch.cat([sv, state["v"]], dim=1)
        idx = torch.arange(Lmax, device=x.device)[None, :]
        valid = torch.cat(
            [torch.ones(B, 1, dtype=torch.bool, device=x.device), idx < n[:, None]],
            dim=1,
        )
        q = self.mixer.q_proj(x).reshape(B, 1, self.H, self.D).transpose(1, 2)
        o = F.scaled_dot_product_attention(
            q,
            keys.transpose(1, 2),
            vals.transpose(1, 2),
            attn_mask=valid[:, None, None, :],
        )
        o = self.mixer.o_proj(o.transpose(1, 2).reshape(B, self.C))
        acc_sum = state["acc_sum"] + x
        acc_cnt = state["acc_cnt"] + 1
        k, v = (state["k"].clone(), state["v"].clone())
        fin = acc_cnt == c
        if bool(fin.any()):
            mean = acc_sum / c
            lk = self.mixer.k_proj(mean).reshape(B, self.H, self.D)
            lv = self.mixer.v_proj(mean).reshape(B, self.H, self.D)
            rows = torch.arange(B, device=x.device)
            full = fin & (n >= Lmax)
            if bool(full.any()):
                k[full] = torch.cat([k[full][:, 1:], k[full][:, :1]], dim=1)
                v[full] = torch.cat([v[full][:, 1:], v[full][:, :1]], dim=1)
            write_pos = n.clamp(max=Lmax - 1)
            fr, wp = (rows[fin], write_pos[fin])
            k[fr, wp] = lk[fr]
            v[fr, wp] = lv[fr]
            n = (n + fin.long()).clamp(max=Lmax)
            acc_sum = torch.where(fin[:, None], torch.zeros_like(acc_sum), acc_sum)
            acc_cnt = torch.where(fin, torch.zeros_like(acc_cnt), acc_cnt)
        return (
            o,
            {
                "k": k,
                "v": v,
                "n_land": n,
                "acc_sum": acc_sum,
                "acc_cnt": acc_cnt,
                "c": c,
            },
        )

    def forward_chunk(self, x: torch.Tensor, state):
        B, L, _ = x.shape
        if state is None:
            state = self.init_state(B, device=x.device)
        outs = []
        for i in range(L):
            o, state = self.step(x[:, i], state)
            outs.append(o)
        return (torch.stack(outs, dim=1), state)