from typing import Tuple, Union

import torch
from torch import nn, Tensor


class KVCache(nn.Module):
    """
    Standalone ``nn.Module`` containing a kv-cache to cache past key and values during inference.

    Args:
        batch_size (int): batch size model will be run with
        max_seq_len (int): maximum sequence length model will be run with
        num_heads (int): number of heads. We take num_heads instead of num_kv_heads because
            the cache is created after we've expanded the key and value tensors to have the
            same shape as the query tensor. See attention.py for more details
        head_dim (int): per-attention head embedding dimension
        dtype (torch.dtype): dtype for the caches
    """

    def __init__(
        self,
        batch_size: int,
        max_seq_len: int,
        num_heads: int,
        head_dim: int,
        dtype: torch.dtype,
        device: torch.device,
    ) -> None:
        super().__init__()
        cache_shape = (batch_size, num_heads, max_seq_len, head_dim)
        self.register_buffer(
            "k_cache",
            torch.zeros(cache_shape, dtype=dtype, device=device),
            persistent=False,
        )
        self.register_buffer(
            "v_cache",
            torch.zeros(cache_shape, dtype=dtype, device=device),
            persistent=False,
        )
        self.batch_size = batch_size
        self.current_idx = torch.tensor([0], device=device)

    def reset(self, device: Union[str, int]) -> None:
        """Reset the cache to zero."""
        self.k_cache.zero_().to(device)
        self.v_cache.zero_().to(device)

    def update(
        self, input_pos: Tensor, k_val: Tensor, v_val: Tensor
    ) -> Tuple[Tensor, Tensor]:
        """Update KV cache with the new k_val, v_val and return the updated cache.

        Args:
            input_pos (Tensor): Current position tensor with shape [S]
            k_val (Tensor): Current key tensor with shape [B, H, S, D]
            v_val (Tensor): Current value tensor with shape [B, H, S, D]

        Raises:
            ValueError: if ``input_pos`` is longer than the maximum sequence length

        Returns:
            Tuple[Tensor, Tensor]: Updated KV cache with key first
        """
        seq_len = input_pos.shape[-1]
        assert seq_len == k_val.shape[2]
        k_out = self.k_cache
        v_out = self.v_cache
        if seq_len > 1 or self.current_idx == 0:
            # first pass when we prefill the cache
            assert self.current_idx == 0
            k_out[:, :, torch.arange(seq_len, device=input_pos.device)] = k_val
            v_out[:, :, torch.arange(seq_len, device=input_pos.device)] = v_val
            self.current_idx = torch.tensor([seq_len], device=k_val.device)
        else:
            # other passes where the cache is updated by one token
            k_out[:, :, self.current_idx] = k_val
            v_out[:, :, self.current_idx] = v_val
            self.current_idx += 1
        return k_out, v_out