import torch
import torch.nn.functional as F
from .utils import l2norm

try:
    from . import _C
    _HAS_METAL = True
except ImportError:
    _HAS_METAL = False


def _metal_recurrent(q, k, v, g, beta, state, scale):
    """Metal recurrent loop. Uses fused kernel for T>1, step kernel for T=1.
    q, k: (B, H, T, K) float16
    v: (B, H, T, V) float16
    g: (B, H, T) float32
    beta: (B, H, T) float16
    state: (B, H, K, V) float32 — mutated in-place
    Returns: (out, state) where out is (B, H, T, V) float16
    """
    T = q.shape[2]
    if T == 1:
        step_out = _C.delta_rule_step(
            q[:, :, 0], k[:, :, 0], v[:, :, 0],
            g[:, :, 0], beta[:, :, 0], state, scale)
        return step_out.unsqueeze(2), state
    out = _C.delta_rule_fused(q, k, v, g, beta, state, scale)
    return out, state


def fused_recurrent_gated_delta_rule(
    q, k, v, g, beta, gk=None, gv=None,
    scale=None,
    initial_state=None,
    output_final_state=False,
    use_qk_l2norm_in_kernel=False,
    cu_seqlens=None,
):
    """MPS-compatible recurrent gated delta rule."""
    initial_dtype = q.dtype
    if use_qk_l2norm_in_kernel:
        q = l2norm(q, dim=-1)
        k = l2norm(k, dim=-1)

    B, T, H, K = q.shape
    V = v.shape[-1]
    if scale is None:
        scale = K ** -0.5

    # State is always float32 for precision
    state = (
        torch.zeros(B, H, K, V, device=q.device, dtype=torch.float32)
        if initial_state is None
        else initial_state.clone().to(dtype=torch.float32, device=q.device)
    )

    if _HAS_METAL and q.is_mps:
        # Transpose to (B, H, T, K/V)
        q_h = q.transpose(1, 2).contiguous().to(torch.float16)
        k_h = k.transpose(1, 2).contiguous().to(torch.float16)
        v_h = v.transpose(1, 2).contiguous().to(torch.float16)
        g_f = g.transpose(1, 2).contiguous().to(torch.float32)
        beta_h = beta.transpose(1, 2).contiguous().to(torch.float16)

        out_h, state = _metal_recurrent(q_h, k_h, v_h, g_f, beta_h, state, scale)
        out = out_h.to(torch.float32).transpose(1, 2).contiguous().to(initial_dtype)
    else:
        # Transpose to (B, H, T, K/V)
        q, k, v, beta, g = [
            x.transpose(1, 2).contiguous().to(torch.float32) for x in (q, k, v, beta, g)
        ]
        q = q * scale
        out = torch.zeros(B, H, T, V, device=q.device, dtype=q.dtype)
        for i in range(T):
            q_t = q[:, :, i]
            k_t = k[:, :, i]
            v_t = v[:, :, i]
            g_t = g[:, :, i].exp().unsqueeze(-1).unsqueeze(-1)
            beta_t = beta[:, :, i].unsqueeze(-1)
            state = state * g_t
            kv_mem = (state * k_t.unsqueeze(-1)).sum(dim=-2)
            delta = (v_t - kv_mem) * beta_t
            state = state + k_t.unsqueeze(-1) * delta.unsqueeze(-2)
            out[:, :, i] = (state * q_t.unsqueeze(-1)).sum(dim=-2)
        out = out.transpose(1, 2).contiguous().to(initial_dtype)

    last_state = state if output_final_state else None
    return out, last_state


def chunk_gated_delta_rule(
    q, k, v, g, beta,
    scale=None,
    initial_state=None,
    output_final_state=False,
    use_qk_l2norm_in_kernel=False,
    cu_seqlens=None,
    chunk_size=64,
):
    """MPS-compatible chunked gated delta rule (prefill)."""
    initial_dtype = q.dtype
    if use_qk_l2norm_in_kernel:
        q = l2norm(q, dim=-1)
        k = l2norm(k, dim=-1)
    q, k, v, beta, g = [
        x.transpose(1, 2).contiguous().to(torch.float32) for x in (q, k, v, beta, g)
    ]

    B, H, T, K = k.shape
    V = v.shape[-1]
    pad_size = (chunk_size - T % chunk_size) % chunk_size
    q = F.pad(q, (0, 0, 0, pad_size))
    k = F.pad(k, (0, 0, 0, pad_size))
    v = F.pad(v, (0, 0, 0, pad_size))
    beta = F.pad(beta, (0, pad_size))
    g = F.pad(g, (0, pad_size))
    total_len = T + pad_size
    if scale is None:
        scale = K ** -0.5
    q = q * scale

    v_beta = v * beta.unsqueeze(-1)
    k_beta = k * beta.unsqueeze(-1)
    q, k, v, k_beta, v_beta = [
        x.reshape(B, H, -1, chunk_size, x.shape[-1]) for x in (q, k, v, k_beta, v_beta)
    ]
    g = g.reshape(B, H, -1, chunk_size)
    mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0)

    g = g.cumsum(dim=-1)
    decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
    attn = -((k_beta @ k.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0)
    for i in range(1, chunk_size):
        row = attn[..., i, :i].clone()
        sub = attn[..., :i, :i].clone()
        attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)
    attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
    v2 = attn @ v_beta
    k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))

    state = (
        torch.zeros(B, H, K, V, device=q.device, dtype=q.dtype)
        if initial_state is None
        else initial_state.to(q)
    )
    out = torch.zeros_like(v2)
    mask2 = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1)

    num_chunks = total_len // chunk_size
    for i in range(num_chunks):
        q_i, k_i, v_i = q[:, :, i], k[:, :, i], v2[:, :, i]
        attn_intra = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask2, 0)
        v_prime = k_cumdecay[:, :, i] @ state
        v_new = v_i - v_prime
        attn_inter = (q_i * g[:, :, i, :, None].exp()) @ state
        out[:, :, i] = attn_inter + attn_intra @ v_new
        state = (
            state * g[:, :, i, -1, None, None].exp()
            + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new
        )

    last_state = state if output_final_state else None
    out = out.reshape(B, H, -1, out.shape[-1])[:, :, :T]
    out = out.transpose(1, 2).contiguous().to(initial_dtype)
    return out, last_state
