"""Utilities for KV smashing."""

from collections.abc import Iterable
from functools import reduce
from typing import Any

import torch
from transformers.cache_utils import DynamicCache
from transformers.tokenization_utils_base import BatchEncoding

TokenizedCacheIterleaving = Iterable[BatchEncoding | DynamicCache]
LegacyCache = Any


def legacy_cache_smash(a: LegacyCache, b: LegacyCache) -> LegacyCache:
    """Concatenates two LegacyCache Ks and Vs along the time axis."""
    legacy_merged = tuple(
        (torch.cat([a[i][0], b[i][0]], dim=2), torch.cat([a[i][1], b[i][1]], dim=2))
        for i in range(len(a))
    )
    return legacy_merged


def merge_dynamic_caches(caches: Iterable[DynamicCache]) -> DynamicCache:
    """Merges two DynamicCache Ks and Vs along the time axis."""
    legacies = [c.to_legacy_cache() for c in caches]  # type: ignore
    assert len(legacies) >= 1
    rv = DynamicCache.from_legacy_cache(reduce(legacy_cache_smash, legacies))  # type: ignore
    return rv  # type: ignore


def tokens_to_legacy_cache(
    model, device: str, tokens_or_cache: BatchEncoding | DynamicCache
) -> Iterable[LegacyCache]:
    """Prefills and returns Ks and Vs as a LegacyCache."""
    if type(tokens_or_cache) is DynamicCache:
        return tokens_or_cache.to_legacy_cache()  # type: ignore
    else:
        tokens = tokens_or_cache
        dc = DynamicCache()
        with torch.no_grad():
            dc = model(
                tokens["input_ids"].to(device),  # type: ignore
                attention_mask=tokens["attention_mask"].to(device),  # type: ignore
                past_key_values=dc,
            ).past_key_values
        return dc.to_legacy_cache()
