from __future__ import annotations
from typing import Dict, Optional, Tuple, Any
import torch
import torch.nn.functional as F

from . import FLUXSPIKEY as triton_backend


def attend(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    *,
    backend: str = "dense",
    state: Optional[Dict[str, torch.Tensor]] = None,
    neuromorph: Optional[Dict[str, Any]] = None,
) -> Tuple[torch.Tensor, Dict[str, Any]]:
    if backend == "dense":
        return _attend_dense(q, k, v)
    elif backend == "spike":
        neuromorph_cfg = neuromorph or {}
        if triton_backend.has_triton():
            return triton_backend.attend_spike_triton(
                q, k, v,
                tau=neuromorph_cfg.get("tau", 0.15),
                leak=neuromorph_cfg.get("leak", 0.0),
                event_cap_pct=neuromorph_cfg.get("event_cap_pct", 0.45),
                state=state,
            )
        return _attend_spike(q, k, v, state=state, neuromorph=neuromorph_cfg)
    else:
        raise ValueError(f"Unknown backend: {backend}")


def _attend_dense(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
) -> Tuple[torch.Tensor, Dict[str, Any]]:
    scale = q.shape[-1] ** -0.5
    scores = torch.matmul(q, k.transpose(-2, -1)) * scale
    attn_weights = F.softmax(scores, dim=-1)
    output = torch.matmul(attn_weights, v)
    return output, {"backend": "dense", "attn_weights": attn_weights}


def _attend_spike(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    *,
    state: Optional[Dict[str, torch.Tensor]] = None,
    neuromorph: Optional[Dict[str, Any]] = None,
) -> Tuple[torch.Tensor, Dict[str, Any]]:
    cfg = neuromorph or {}
    tau = cfg.get("tau", 0.15)
    leak = cfg.get("leak", 0.0)
    event_cap_pct = cfg.get("event_cap_pct", 0.45)

    batch, heads, seq_len, head_dim = q.shape

    scale = head_dim ** -0.5
    z = torch.matmul(q, k.transpose(-2, -1)) * scale

    z_max_per_row = z.max(dim=-1, keepdim=True)[0]
    M_t = (z >= tau * z_max_per_row).float()

    max_events_per_row = max(1, int(seq_len * event_cap_pct))
    if max_events_per_row < seq_len:
        _, topk_indices = torch.topk(z, k=max_events_per_row, dim=-1, sorted=False)
        M_t_capped = torch.zeros_like(z)
        M_t_capped.scatter_(-1, topk_indices, 1.0)
        M_t = M_t * M_t_capped

    if state is not None and "S_prev" in state and state["S_prev"].shape == M_t.shape:
        S_t = leak * state["S_prev"] + M_t
    else:
        S_t = M_t

    S_t_sum = S_t.sum(dim=-1, keepdim=True)
    S_t_sum = torch.where(S_t_sum > 0, S_t_sum, torch.ones_like(S_t_sum))
    attn_weights = (S_t / S_t_sum * M_t).to(v.dtype)
    output = torch.matmul(attn_weights, v)

    kept = M_t.sum().item()
    total = batch * heads * seq_len * seq_len
    flops_saved = 100.0 * (1.0 - kept / max(1, total))

    return output, {
        "S_prev": S_t.detach(),
        "M_t": M_t,
        "telemetry": {
            "backend": "spike",
            "kept_events_pct": 100.0 * kept / max(1, total),
            "flops_saved_pct": flops_saved,
            "tau": tau,
            "leak": leak,
            "event_cap_pct": event_cap_pct,
        },
    }