# Copyright (c) Together / Apache 2.0.
#
# Minimal multi-head attention block for the Evo1 HF port.
#
# Replaces flash_attn.modules.mha.MHA with a small, dependency-light
# implementation that:
#   - keeps the same parameter names (Wqkv, out_proj, rotary_emb.inv_freq)
#     so existing checkpoints load directly,
#   - supports attn_implementation in {"eager", "sdpa", "flash_attention_2"},
#   - returns attention weights when output_attentions=True (eager path),
#   - falls back to eager when output_attentions=True for sdpa/flash backends
#     (per the standard HuggingFace dispatch convention),
#   - keeps a one-method KV cache compatible with the existing
#     InferenceParams dataclass for autoregressive generation.
#
# Math is causal, single-stream (no cross-attention), no ALiBi, no sliding
# window. Evo1 only ever exercised the qkv-packed self-attention path.

from __future__ import annotations

import math

import torch
import torch.nn as nn
import torch.nn.functional as F

from .rotary import RotaryEmbedding


def _flash_attn_required():
    try:
        from flash_attn import flash_attn_func, flash_attn_varlen_func  # noqa: F401
        from flash_attn.bert_padding import pad_input, unpad_input  # noqa: F401
    except ImportError as exc:  # pragma: no cover - optional dep
        raise ImportError(
            "attn_implementation='flash_attention_2' requires the flash-attn "
            "package. Install with `pip install flash-attn --no-build-isolation`."
        ) from exc


def _update_kv_cache(kv: torch.Tensor, inference_params, layer_idx: int) -> torch.Tensor:
    """Append `kv` to inference_params.key_value_memory_dict[layer_idx].

    kv: (B, S, 2, H_kv, D) where S is the new-token chunk length (may be 1).
    Returns the cumulative kv up to the current sequence position.
    """
    num_heads, head_dim = kv.shape[-2:]
    if layer_idx not in inference_params.key_value_memory_dict:
        kv_cache = torch.empty(
            inference_params.max_batch_size,
            inference_params.max_seqlen,
            2,
            num_heads,
            head_dim,
            dtype=kv.dtype,
            device=kv.device,
        )
        inference_params.key_value_memory_dict[layer_idx] = kv_cache
    else:
        kv_cache = inference_params.key_value_memory_dict[layer_idx]
    batch_start = inference_params.batch_size_offset
    batch_end = batch_start + kv.shape[0]
    sequence_start = inference_params.seqlen_offset
    sequence_end = sequence_start + kv.shape[1]
    kv_cache[batch_start:batch_end, sequence_start:sequence_end, ...] = kv
    return kv_cache[batch_start:batch_end, :sequence_end, ...]


class MHA(nn.Module):
    """Multi-head self-attention with backend-dispatch.

    Constructor signature is a strict subset of flash_attn.modules.mha.MHA so
    that the existing AttentionBlock instantiation site is left untouched.
    Unsupported kwargs (cross_attn, dwconv, alibi, window_size, ...) are
    accepted and ignored or hard-asserted: Evo1 never exercises them.
    """

    def __init__(
        self,
        embed_dim: int,
        num_heads: int,
        num_heads_kv: int | None = None,
        cross_attn: bool = False,
        qkv_proj_bias: bool = True,
        out_proj_bias: bool = True,
        dropout: float = 0.0,
        softmax_scale: float | None = None,
        causal: bool = False,
        layer_idx: int | None = None,
        rotary_emb_dim: int = 0,
        rotary_emb_base: float = 10000.0,
        rotary_emb_scale_base: float | None = None,
        rotary_emb_interleaved: bool = False,
        use_flash_attn: bool = False,  # legacy kwarg, kept for ctor compatibility
        attn_implementation: str = "eager",
        device=None,
        dtype=None,
    ) -> None:
        super().__init__()
        if cross_attn:
            raise NotImplementedError("Cross-attention is not supported in this minimal MHA.")

        factory_kwargs = {"device": device, "dtype": dtype}
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.num_heads_kv = num_heads_kv if num_heads_kv is not None else num_heads
        if self.embed_dim % num_heads != 0:
            raise ValueError("embed_dim must be divisible by num_heads")
        if self.num_heads % self.num_heads_kv != 0:
            raise ValueError("num_heads must be divisible by num_heads_kv")
        self.head_dim = self.embed_dim // num_heads
        self.causal = causal
        self.softmax_scale = softmax_scale
        self.layer_idx = layer_idx
        self.rotary_emb_dim = rotary_emb_dim
        self.attn_implementation = attn_implementation
        self.dropout_p = dropout

        if self.rotary_emb_dim > 0:
            self.rotary_emb = RotaryEmbedding(
                self.rotary_emb_dim,
                base=rotary_emb_base,
                interleaved=rotary_emb_interleaved,
                scale_base=rotary_emb_scale_base,
                device=device,
            )

        qkv_dim = self.head_dim * (self.num_heads + 2 * self.num_heads_kv)
        self.Wqkv = nn.Linear(embed_dim, qkv_dim, bias=qkv_proj_bias, **factory_kwargs)
        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=out_proj_bias, **factory_kwargs)

    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None):
        dtype = self.out_proj.weight.dtype if dtype is None else dtype
        device = self.out_proj.weight.device
        return torch.empty(
            batch_size, max_seqlen, 2, self.num_heads_kv, self.head_dim,
            dtype=dtype, device=device,
        )

    def _project_qkv(self, x: torch.Tensor) -> torch.Tensor:
        """Compute Wqkv(x) and reshape to (B, T, 3, H, D) when MHA, or
        return (q, kv) tuple-like layout when GQA. Returns the packed qkv
        tensor in either case (kv heads broadcast for SDPA/flash later).

        For Evo1 we have num_heads_kv == num_heads (proj_groups=1), so the
        common-case packed layout is fine; we keep a GQA branch for future
        flexibility but assert MHA at construction time.
        """
        qkv = self.Wqkv(x)
        if self.num_heads_kv == self.num_heads:
            return qkv.view(*qkv.shape[:-1], 3, self.num_heads, self.head_dim)
        # GQA path (unused by Evo1):
        q = qkv[..., : self.num_heads * self.head_dim]
        kv = qkv[..., self.num_heads * self.head_dim:]
        q = q.view(*q.shape[:-1], self.num_heads, self.head_dim)
        kv = kv.view(*kv.shape[:-1], 2, self.num_heads_kv, self.head_dim)
        return q, kv  # type: ignore[return-value]

    # ------------------------------------------------------------------ eager
    def _forward_eager(
        self,
        qkv: torch.Tensor,
        output_attentions: bool,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        # qkv: (B, T, 3, H, D). Match flash_attn / sdpa numerical behaviour by
        # doing the attention math in fp32 internally (q*scale, QK^T matmul,
        # softmax, attn @ V). Without this, the bf16 matmuls accumulate
        # ~1e-2 absolute error per attention block and diverge meaningfully
        # from flash_attn (which always accumulates in fp32 inside its CUDA
        # kernel). Output is cast back to the original dtype for the residual
        # add.
        orig_dtype = qkv.dtype
        q, k, v = qkv.unbind(dim=2)
        q = q.permute(0, 2, 1, 3).float()  # (B, H, T, D), fp32
        k = k.permute(0, 2, 1, 3).float()
        v = v.permute(0, 2, 1, 3).float()
        scale = self.softmax_scale if self.softmax_scale is not None else 1.0 / math.sqrt(self.head_dim)

        scores = torch.matmul(q, k.transpose(-2, -1)) * scale
        if self.causal:
            T = q.shape[-2]
            mask = torch.triu(
                torch.ones(T, T, device=scores.device, dtype=torch.bool), diagonal=1
            )
            scores = scores.masked_fill(mask, float("-inf"))
        attn = F.softmax(scores, dim=-1)
        if self.training and self.dropout_p > 0:
            attn = F.dropout(attn, p=self.dropout_p)
        out = torch.matmul(attn, v).permute(0, 2, 1, 3)  # (B, T, H, D), fp32
        out = out.to(orig_dtype)
        return out, (attn.to(orig_dtype) if output_attentions else None)

    # -------------------------------------------------------------------- sdpa
    def _forward_sdpa(self, qkv: torch.Tensor) -> torch.Tensor:
        q, k, v = qkv.unbind(dim=2)
        q = q.permute(0, 2, 1, 3)  # (B, H, T, D)
        k = k.permute(0, 2, 1, 3)
        v = v.permute(0, 2, 1, 3)
        scale = self.softmax_scale if self.softmax_scale is not None else None
        out = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=None,
            dropout_p=self.dropout_p if self.training else 0.0,
            is_causal=self.causal,
            scale=scale,
        )
        return out.permute(0, 2, 1, 3)  # (B, T, H, D)

    # -------------------------------------------------------- flash_attention_2
    def _forward_flash(self, qkv: torch.Tensor) -> torch.Tensor:
        _flash_attn_required()
        from flash_attn import flash_attn_qkvpacked_func
        # flash_attn expects (B, T, 3, H, D) in fp16/bf16 already; Evo1 attn
        # blocks already cast to bf16 in __init__.
        out = flash_attn_qkvpacked_func(
            qkv,
            dropout_p=self.dropout_p if self.training else 0.0,
            softmax_scale=self.softmax_scale,
            causal=self.causal,
        )
        return out  # (B, T, H, D)

    # ----------------------------------------------------------- KV-cache path
    def _forward_with_cache(
        self,
        qkv: torch.Tensor,
        inference_params,
    ) -> torch.Tensor:
        # qkv: (B, T, 3, H, D). Apply rotary at the current offset, append kv
        # to cache, attend over the cumulative kv. For correctness we use SDPA
        # which has stable behaviour at all sequence lengths.
        if self.rotary_emb_dim > 0:
            qkv = self.rotary_emb(
                qkv,
                seqlen_offset=inference_params.seqlen_offset,
                max_seqlen=inference_params.max_seqlen,
            )
        q, k, v = qkv.unbind(dim=2)
        kv = torch.stack((k, v), dim=2)  # (B, T, 2, H, D)
        kv = _update_kv_cache(kv, inference_params, self.layer_idx)
        k_full, v_full = kv.unbind(dim=2)  # (B, S_total, H, D)

        q = q.permute(0, 2, 1, 3)
        k_full = k_full.permute(0, 2, 1, 3)
        v_full = v_full.permute(0, 2, 1, 3)
        scale = self.softmax_scale if self.softmax_scale is not None else None
        is_causal = self.causal and q.shape[-2] == k_full.shape[-2]
        out = F.scaled_dot_product_attention(
            q, k_full, v_full, is_causal=is_causal, scale=scale,
        )
        return out.permute(0, 2, 1, 3)  # (B, T, H, D)

    # ---------------------------------------------------------------- forward
    def forward(
        self,
        x: torch.Tensor,
        inference_params=None,
        output_attentions: bool = False,
        **_unused,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """Returns (out, attn_weights_or_None) where out is (B, T, embed_dim)."""
        if self.num_heads_kv != self.num_heads:
            raise NotImplementedError("GQA is not exercised by Evo1; please file an issue if needed.")

        qkv = self._project_qkv(x)  # (B, T, 3, H, D)

        if inference_params is not None:
            out_btd = self._forward_with_cache(qkv, inference_params)
            attn_weights = None
        else:
            if self.rotary_emb_dim > 0:
                qkv = self.rotary_emb(qkv, seqlen_offset=0, max_seqlen=qkv.shape[1])

            backend = self.attn_implementation
            if output_attentions and backend != "eager":
                # Standard HF behaviour: silently fall back to eager so we can
                # actually compute and return the attention matrix.
                backend = "eager"

            if backend == "eager":
                out_btd, attn_weights = self._forward_eager(qkv, output_attentions=output_attentions)
            elif backend == "sdpa":
                out_btd = self._forward_sdpa(qkv)
                attn_weights = None
            elif backend == "flash_attention_2":
                out_btd = self._forward_flash(qkv)
                attn_weights = None
            else:
                raise ValueError(f"Unknown attn_implementation: {backend!r}")

        # (B, T, H, D) -> (B, T, embed_dim)
        B, T, H, D = out_btd.shape
        out_flat = out_btd.reshape(B, T, H * D)
        return self.out_proj(out_flat), attn_weights