"""MLA RoPE case: fused (activation Triton) vs vanilla (PyTorch native).

MLA head dims are fixed by motif3 spec (H_q=80, H_kv=16, D_nope=128,
D_rope=64, D_v=128); the benchmark's only sweep axes are (bs, sl).
The framework's ``dim`` axis is a dummy here — pass 0 in configs.
"""

import torch
from torch import nn

import activation

from common.diff_engine import DiffCase

# ---- MLA shapes (motif3_seq) -----------------------------------------------
H_Q, H_KV = 80, 16
D_NOPE, D_ROPE, D_V = 128, 64, 128
D_QK = D_NOPE + D_ROPE  # 192


# ---- reference (PyTorch native) --------------------------------------------
def _precompute_freqs_cis(dim, end, theta=10000.0):
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: dim // 2].float() / dim))
    t = torch.arange(end, dtype=torch.float32)
    freqs = torch.outer(t, freqs)
    return torch.polar(torch.ones_like(freqs), freqs)


def _apply_rotary_emb_single(x, freqs_cis):
    x_ = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
    freqs_cis = freqs_cis[: x_.shape[1]].view(1, x_.shape[1], 1, x_.shape[3])
    out = torch.view_as_real(x_ * freqs_cis).flatten(3)
    return out.type_as(x)


def _reorder(qk, rope_dim):
    B, S = qk.shape[0], qk.shape[1]
    qk = qk.view(B, S, -1, rope_dim // 2, 2).transpose(3, 4)
    return qk.reshape(B, S, -1, rope_dim)


def _vanilla(q, kv_latent, k_pe, freqs_cis):
    q_nope, q_pe = torch.split(q, [D_NOPE, D_ROPE], dim=-1)
    q_pe = _reorder(_apply_rotary_emb_single(q_pe, freqs_cis), D_ROPE)
    q_total = torch.cat([q_nope, q_pe], dim=-1)
    k_pe_roped = _reorder(_apply_rotary_emb_single(k_pe.unsqueeze(2), freqs_cis), D_ROPE)
    k_nope, v = torch.split(kv_latent, [D_NOPE, D_V], dim=-1)
    k_full = torch.cat([k_nope, k_pe_roped.expand(-1, -1, H_KV, -1)], dim=-1)
    return q_total, k_full, v


def _fused(q, kv_latent, k_pe, freqs_cis):
    q_total = activation.fused_q_rope_inplace(q, freqs_cis, D_NOPE, D_ROPE)
    # k_pe RoPE stays PyTorch native (head-shared, too small for custom kernel)
    k_pe_roped = _reorder(_apply_rotary_emb_single(k_pe.unsqueeze(2), freqs_cis), D_ROPE)
    k_full, v = activation.fused_kv_split_rope_cat(kv_latent, k_pe_roped, D_NOPE, D_V, D_ROPE)
    return q_total, k_full, v


class _VanillaModule(nn.Module):
    def forward(self, q, kv_latent, k_pe, freqs_cis):
        return _vanilla(q, kv_latent, k_pe, freqs_cis)


class _FusedModule(nn.Module):
    def forward(self, q, kv_latent, k_pe, freqs_cis):
        return _fused(q, kv_latent, k_pe, freqs_cis)


class MLARoPE(DiffCase):
    # Framework calls build_inputs(bs, sl, dim, dtype, eps) — dim unused.
    def build_inputs(self, bs, sl, hidden, dtype, eps):
        return {
            "q":         (torch.randn(bs, sl, H_Q, D_QK, dtype=dtype) * 0.5).requires_grad_(True),
            "kv_latent": (torch.randn(bs, sl, H_KV, D_NOPE + D_V, dtype=dtype) * 0.5).requires_grad_(True),
            "k_pe":      (torch.randn(bs, sl, D_ROPE, dtype=dtype) * 0.5).requires_grad_(True),
            "freqs_cis": _precompute_freqs_cis(D_ROPE, sl),
        }

    def make_naive(self, I):
        return _VanillaModule()

    def make_cuda(self, I):
        return _FusedModule()

    def forward(self, obj, I):
        # fused_q_rope_inplace needs non-leaf q; wrap both paths for fairness
        q_in = I["q"] * 1.0
        kv_in = I["kv_latent"] * 1.0
        kpe_in = I["k_pe"] * 1.0
        return obj(q_in, kv_in, kpe_in, I["freqs_cis"])

    def grad_inputs(self, I):
        return [I["q"], I["kv_latent"], I["k_pe"]]


CASE = MLARoPE()