# Copyright (c) Together
# Apache 2.0 - Author: Michael Poli
# Adapted for the minimal Evo1 HF port.

"""Inference-time caches for Evo1 blocks.

Evo1 has two block types with different caching needs:

  * `mha` blocks   -> InferenceParams (standard KV cache)
  * `hyena` blocks -> RecurrentInferenceParams (FIR window + IIR modal state)

Per-block dataclasses are wrapped in an HF ``Cache`` subclass (``Evo1Cache``)
so ``model.generate()`` can drive autoregressive decoding without the user
having to instantiate the two caches by hand, and so HF generation helpers
can introspect cache state (``get_seq_length``, ``get_max_cache_shape``).
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional

from torch import Tensor
from transformers.cache_utils import Cache


@dataclass
class InferenceParams:
    """KV-cache parameters for the attention blocks (mha branch)."""

    max_seqlen: int
    max_batch_size: int
    seqlen_offset: int = 0
    batch_size_offset: int = 0
    key_value_memory_dict: dict = field(default_factory=dict)
    lengths_per_sample: Optional[Tensor] = None

    def reset(self, max_seqlen, max_batch_size):
        self.max_seqlen = max_seqlen
        self.max_batch_size = max_batch_size
        self.seqlen_offset = 0
        if self.lengths_per_sample is not None:
            self.lengths_per_sample.zero_()


@dataclass
class RecurrentInferenceParams:
    """SSM-cache parameters for the Hyena blocks (hyena branch)."""

    fir_filter_length: int = 3
    state_dim: int = 16
    seqlen_offset: int = 0
    fir_state_dict: dict = field(default_factory=dict)
    state_dict: dict = field(default_factory=dict)

    def reset(self):
        self.fir_filter_length = 3
        self.state_dim = 16
        self.seqlen_offset = 0


class Evo1Cache(Cache):
    """HF-compatible wrapper around the per-block inference params.

    Internally holds two dataclasses keyed by block type. Exposes
    ``seqlen_offset`` so HF generation helpers can read the current decoded
    length, and implements ``get_seq_length()`` / ``get_max_cache_shape()``
    per the transformers ``Cache`` interface.

    The model internals (``StripedHyena.stateful_forward``) look up caches
    via ``cache["mha"]`` and ``cache["hyena"]``; ``__getitem__`` is delegated
    to attribute access so the original dict-keyed API keeps working.
    """

    is_compileable = False

    def __init__(
        self,
        max_seqlen: int,
        max_batch_size: int,
        short_filter_length: int = 3,
        state_size: int = 8,
    ):
        # transformers >= 4.55 Cache.__init__ requires either ``layers`` or
        # ``layer_class_to_replicate``. We don't use HF's per-layer cache
        # model (our two block-type-specific caches handle storage), so we
        # pass an empty layers list.
        super().__init__(layers=[])
        self.mha = InferenceParams(
            max_seqlen=max_seqlen,
            max_batch_size=max_batch_size,
        )
        self.hyena = RecurrentInferenceParams(
            fir_filter_length=short_filter_length,
            state_dim=state_size,
        )

    # --- HF Cache interface ------------------------------------------------
    @property
    def seqlen_offset(self) -> int:
        return self.mha.seqlen_offset

    def get_seq_length(self, layer_idx: int = 0) -> int:
        return self.mha.seqlen_offset

    def get_max_cache_shape(self) -> int:
        return self.mha.max_seqlen

    def get_max_length(self) -> int:
        # deprecated alias kept for older transformers versions
        return self.mha.max_seqlen

    # --- our convenience helpers ------------------------------------------
    def advance(self, n: int = 1) -> None:
        self.mha.seqlen_offset += n
        self.hyena.seqlen_offset += n

    def set_offset(self, offset: int) -> None:
        self.mha.seqlen_offset = offset
        self.hyena.seqlen_offset = offset

    def reset(self) -> None:
        self.mha.reset(self.mha.max_seqlen, self.mha.max_batch_size)
        self.hyena.reset()

    # --- dict-like access so existing call sites keep working --------------
    def __getitem__(self, name: str):
        return getattr(self, name)

    def by_block_name(self, name: str):
        return getattr(self, name)