import torch
import torch.nn as nn


class StaticKVCache(nn.Module):
    def __init__(
        self,
        max_batch_size,
        n_heads,
        context_length,
        head_dim,
        device,
        dtype=torch.float32
    ):
        super().__init__()
        self.register_buffer(
            "k_cache",
            torch.zeros(
                (max_batch_size, n_heads, context_length, head_dim),
                dtype=dtype,
                device=device
            )
        )
        self.register_buffer(
            "v_cache",
            torch.zeros(
                (max_batch_size, n_heads, context_length, head_dim),
                dtype=dtype,
                device=device
            )
        )

    def update(self, k_val, v_val, cache_position):
        self.k_cache[:, :, cache_position] = k_val
        self.v_cache[:, :, cache_position] = v_val

        total_len = cache_position[-1].item() + 1
        return (
            self.k_cache[:, :, :total_len],
            self.v_cache[:, :, :total_len]
        )