"""High-performance W4A16 decode solution for the Kimi-Linear hybrid block.

Strategy:
  - Fused int4 unpack + per-group dequant + GEMV in Triton (no bf16 materialization).
  - KDA recurrence in fp32, with conv and decay all fused inline.
  - MLA latent attention via the absorb trick.
  - MoE: one grouped-GEMV kernel per linear that runs all 8 routed experts in
    a single launch.
"""
from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass, field

import triton
import triton.language as tl


EPS = 1.0e-6
GROUP_SIZE = 128


@dataclass(frozen=True)
class Config:
    hidden: int = 2304
    kda_heads: int = 32
    kda_head_dim: int = 128
    short_conv: int = 4
    mla_heads: int = 32
    kv_lora: int = 512
    qk_nope: int = 128
    qk_rope: int = 64
    v_head: int = 128
    rope_theta: float = 10000.0
    n_experts: int = 64
    n_active: int = 8
    n_shared: int = 1
    moe_inter: int = 1024
    routed_scaling: float = 2.446
    group: int = 128
    pattern: tuple = ("K", "K", "K", "M")
    dtype: torch.dtype = field(default=torch.bfloat16)


def build_config(shape: dict) -> Config:
    return Config(n_experts=int(shape.get("n_experts", 64)))


def _pack_int4(w_q: torch.Tensor) -> torch.Tensor:
    lo = w_q[0::2] & 0xF
    hi = w_q[1::2] & 0xF
    return (lo | (hi << 4)).contiguous()


def _unpack_int4(w_packed: torch.Tensor, K: int) -> torch.Tensor:
    out = torch.empty((K, w_packed.shape[1]), dtype=torch.uint8, device=w_packed.device)
    out[0::2] = w_packed & 0xF
    out[1::2] = (w_packed >> 4) & 0xF
    return out


def quantize(w_io: torch.Tensor, group: int = GROUP_SIZE):
    K, N = w_io.shape
    ng = K // group
    wg = w_io.view(ng, group, N).float()
    wmin = wg.min(dim=1, keepdim=True).values
    wmax = wg.max(dim=1, keepdim=True).values
    scales = (wmax - wmin).clamp_min(1e-8) / 15.0
    zeros = (-wmin / scales).round().clamp(0, 15)
    w_q = ((wg / scales) + zeros).round().clamp(0, 15).to(torch.uint8).view(K, N)
    return _pack_int4(w_q), scales.squeeze(1).to(torch.bfloat16), zeros.squeeze(1).to(torch.bfloat16)


def dequant(w_q, scales, zeros, K, group):
    wu = _unpack_int4(w_q, K).to(torch.bfloat16)
    s = scales.repeat_interleave(group, dim=0)
    z = zeros.repeat_interleave(group, dim=0)
    return (wu - z) * s


# --------------------------------------------------------------------------- #
# Triton W4A16 GEMV kernel
#   X_DTYPE: 0=bf16 (default), 1=fp32 (for the MoE down-projection so the
#   silu*u intermediate stays fp32 end-to-end, no precision loss in the cast)
# --------------------------------------------------------------------------- #


@triton.jit
def _w4a16_gemv_kernel(
    x_ptr, wq_ptr, s_ptr, z_ptr, y_ptr,
    K, N,
    stride_x,
    stride_wqk, stride_wqn,
    stride_sg, stride_sn,
    stride_zg, stride_zn,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
    GROUP_SIZE: tl.constexpr,
    X_DTYPE: tl.constexpr,
):
    pid_n = tl.program_id(0)
    n_offs = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    n_mask = n_offs < N

    acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
    K_HALF = K // 2

    for k_start in range(0, K, BLOCK_K):
        half_offs = (k_start // 2) + tl.arange(0, BLOCK_K // 2)
        k_offs_even = k_start + 2 * tl.arange(0, BLOCK_K // 2)
        k_offs_odd = k_offs_even + 1
        k_mask_even = k_offs_even < K
        k_mask_odd = k_offs_odd < K

        x_even = tl.load(x_ptr + k_offs_even * stride_x, mask=k_mask_even, other=0.0).to(tl.float32)
        x_odd = tl.load(x_ptr + k_offs_odd * stride_x, mask=k_mask_odd, other=0.0).to(tl.float32)

        wq = tl.load(
            wq_ptr + half_offs[:, None] * stride_wqk + n_offs[None, :] * stride_wqn,
            mask=(half_offs[:, None] < K_HALF) & n_mask[None, :],
            other=0,
        )

        w_lo = (wq & 0xF).to(tl.float32)
        w_hi = ((wq >> 4) & 0xF).to(tl.float32)

        g_idx = k_start // GROUP_SIZE
        s = tl.load(s_ptr + g_idx * stride_sg + n_offs * stride_sn, mask=n_mask, other=0.0).to(tl.float32)
        z = tl.load(z_ptr + g_idx * stride_zg + n_offs * stride_zn, mask=n_mask, other=0.0).to(tl.float32)

        dq_scale = s[None, :]
        dq_zero = z[None, :]
        dq_lo = (w_lo - dq_zero) * dq_scale
        dq_hi = (w_hi - dq_zero) * dq_scale

        acc += tl.sum(x_even[:, None] * dq_lo, axis=0)
        acc += tl.sum(x_odd[:, None] * dq_hi, axis=0)

    if X_DTYPE == 0:
        tl.store(y_ptr + n_offs, acc.to(tl.bfloat16), mask=n_mask)
    else:
        tl.store(y_ptr + n_offs, acc, mask=n_mask)


def w4a16_gemv(x, wq, s, z, group=GROUP_SIZE, out=None, block_n=32, num_warps=2, num_stages=2):
    K = x.shape[0]
    N = wq.shape[1]
    is_fp32 = x.dtype == torch.float32
    if out is None:
        y = torch.empty(N, dtype=torch.float32 if is_fp32 else torch.bfloat16, device=x.device)
    else:
        y = out
    grid = (triton.cdiv(N, block_n),)
    _w4a16_gemv_kernel[grid](
        x, wq, s, z, y,
        K, N,
        x.stride(0),
        wq.stride(0), wq.stride(1),
        s.stride(0), s.stride(1),
        z.stride(0), z.stride(1),
        BLOCK_N=block_n, BLOCK_K=group, GROUP_SIZE=group,
        X_DTYPE=1 if is_fp32 else 0,
        num_warps=num_warps, num_stages=num_stages,
    )
    return y


# --------------------------------------------------------------------------- #
# W4A16 grouped MoE kernel
# --------------------------------------------------------------------------- #


@triton.jit
def _w4a16_grouped_gemv_kernel(
    x_ptr, wq_ptr, s_ptr, z_ptr, y_ptr,
    E, K, N,
    stride_x,
    stride_wqe, stride_wqk, stride_wqn,
    stride_se, stride_sg, stride_sn,
    stride_ze, stride_zg, stride_zn,
    stride_ye, stride_yn,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
    GROUP_SIZE: tl.constexpr,
    X_DTYPE: tl.constexpr,
):
    pid_e = tl.program_id(0)
    pid_n = tl.program_id(1)

    n_offs = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    n_mask = n_offs < N

    acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
    K_HALF = K // 2

    for k_start in range(0, K, BLOCK_K):
        half_offs = (k_start // 2) + tl.arange(0, BLOCK_K // 2)
        k_offs_even = k_start + 2 * tl.arange(0, BLOCK_K // 2)
        k_offs_odd = k_offs_even + 1
        k_mask_even = k_offs_even < K
        k_mask_odd = k_offs_odd < K

        x_even = tl.load(x_ptr + k_offs_even * stride_x, mask=k_mask_even, other=0.0).to(tl.float32)
        x_odd = tl.load(x_ptr + k_offs_odd * stride_x, mask=k_mask_odd, other=0.0).to(tl.float32)

        wq = tl.load(
            wq_ptr + pid_e * stride_wqe
            + half_offs[:, None] * stride_wqk + n_offs[None, :] * stride_wqn,
            mask=(half_offs[:, None] < K_HALF) & n_mask[None, :],
            other=0,
        )

        w_lo = (wq & 0xF).to(tl.float32)
        w_hi = ((wq >> 4) & 0xF).to(tl.float32)

        g_idx = k_start // GROUP_SIZE
        s = tl.load(
            s_ptr + pid_e * stride_se + g_idx * stride_sg + n_offs * stride_sn,
            mask=n_mask, other=0.0,
        ).to(tl.float32)
        z = tl.load(
            z_ptr + pid_e * stride_ze + g_idx * stride_zg + n_offs * stride_zn,
            mask=n_mask, other=0.0,
        ).to(tl.float32)

        dq_scale = s[None, :]
        dq_zero = z[None, :]
        dq_lo = (w_lo - dq_zero) * dq_scale
        dq_hi = (w_hi - dq_zero) * dq_scale

        acc += tl.sum(x_even[:, None] * dq_lo, axis=0)
        acc += tl.sum(x_odd[:, None] * dq_hi, axis=0)

    y_ptrs = y_ptr + pid_e * stride_ye + n_offs * stride_yn
    if X_DTYPE == 0:
        tl.store(y_ptrs, acc.to(tl.bfloat16), mask=n_mask)
    else:
        tl.store(y_ptrs, acc, mask=n_mask)


def w4a16_grouped_gemv(x, wq, s, z, group=GROUP_SIZE, out=None, block_n=32, num_warps=2, num_stages=2):
    E = wq.shape[0]
    K = x.shape[0]
    N = wq.shape[2]
    is_fp32 = x.dtype == torch.float32
    if out is None:
        y = torch.empty((E, N), dtype=torch.float32 if is_fp32 else torch.bfloat16, device=x.device)
    else:
        y = out
    grid = (E, triton.cdiv(N, block_n))
    _w4a16_grouped_gemv_kernel[grid](
        x, wq, s, z, y,
        E, K, N,
        x.stride(0),
        wq.stride(0), wq.stride(1), wq.stride(2),
        s.stride(0), s.stride(1), s.stride(2),
        z.stride(0), z.stride(1), z.stride(2),
        y.stride(0), y.stride(1),
        BLOCK_N=block_n, BLOCK_K=group, GROUP_SIZE=group,
        X_DTYPE=1 if is_fp32 else 0,
        num_warps=num_warps, num_stages=num_stages,
    )
    return y


# --------------------------------------------------------------------------- #
# helpers (mirror reference)
# --------------------------------------------------------------------------- #
def _rmsnorm(x, w):
    xf = x.float()
    xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + EPS)
    return (xf * w.float()).to(x.dtype)


def _rope_cossin(pos, dim, theta, device):
    inv = 1.0 / (theta ** (torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim))
    ang = pos * inv
    return torch.cos(ang), torch.sin(ang)


def _apply_rope(x, cos, sin):
    xf = x.float()
    even, odd = xf[..., 0::2], xf[..., 1::2]
    out = torch.empty_like(xf)
    out[..., 0::2] = even * cos - odd * sin
    out[..., 1::2] = odd * cos + even * sin
    return out.to(x.dtype)


# --------------------------------------------------------------------------- #
# W4A16 linear
# --------------------------------------------------------------------------- #
class W4A16Linear(nn.Module):
    def __init__(self, in_f: int, out_f: int, group: int = GROUP_SIZE):
        super().__init__()
        self.in_f, self.out_f, self.group = in_f, out_f, group
        ng = in_f // group
        self.register_buffer("w_q", torch.zeros(in_f // 2, out_f, dtype=torch.uint8))
        self.register_buffer("scales", torch.zeros(ng, out_f, dtype=torch.bfloat16))
        self.register_buffer("zeros", torch.zeros(ng, out_f, dtype=torch.bfloat16))

    def init_random(self, gen: torch.Generator, std: float = 0.02) -> None:
        w = torch.randn(self.in_f, self.out_f, generator=gen) * std
        wq, s, z = quantize(w, self.group)
        self.w_q.copy_(wq)
        self.scales.copy_(s)
        self.zeros.copy_(z)

    def weight_bf(self) -> torch.Tensor:
        return dequant(self.w_q, self.scales, self.zeros, self.in_f, self.group)

    def forward(self, x):
        return w4a16_gemv(x, self.w_q, self.scales, self.zeros, self.group)


# --------------------------------------------------------------------------- #
# W4A16 experts (grouped)
# --------------------------------------------------------------------------- #
class W4A16Experts(nn.Module):
    def __init__(self, n: int, in_f: int, out_f: int, group: int = GROUP_SIZE):
        super().__init__()
        self.n, self.in_f, self.out_f, self.group = n, in_f, out_f, group
        ng = in_f // group
        self.register_buffer("w_q", torch.zeros(n, in_f // 2, out_f, dtype=torch.uint8))
        self.register_buffer("scales", torch.zeros(n, ng, out_f, dtype=torch.bfloat16))
        self.register_buffer("zeros", torch.zeros(n, ng, out_f, dtype=torch.bfloat16))

    def init_random(self, gen: torch.Generator, std: float = 0.02) -> None:
        for e in range(self.n):
            w = torch.randn(self.in_f, self.out_f, generator=gen) * std
            wq, s, z = quantize(w, self.group)
            self.w_q[e].copy_(wq)
            self.scales[e].copy_(s)
            self.zeros[e].copy_(z)

    def weight_bf(self, e: int):
        return dequant(self.w_q[e], self.scales[e], self.zeros[e], self.in_f, self.group)


# --------------------------------------------------------------------------- #
# KDA layer
# --------------------------------------------------------------------------- #
class KDA(nn.Module):
    def __init__(self, cfg: Config):
        super().__init__()
        self.cfg = cfg
        H, Dk, d = cfg.kda_heads, cfg.kda_head_dim, cfg.hidden
        self.q_proj = W4A16Linear(d, H * Dk, cfg.group)
        self.k_proj = W4A16Linear(d, H * Dk, cfg.group)
        self.v_proj = W4A16Linear(d, H * Dk, cfg.group)
        self.g_proj = W4A16Linear(d, H * Dk, cfg.group)
        self.beta_proj = nn.Linear(d, H, bias=False, dtype=cfg.dtype)
        self.conv_w = nn.Parameter(torch.empty(3, H * Dk, cfg.short_conv, dtype=cfg.dtype))
        self.o_proj = W4A16Linear(H * Dk, d, cfg.group)
        self.scale = Dk ** -0.5

    def _short_conv(self, val, prev, idx):
        win = torch.cat([prev, val[None]], dim=0)
        w = self.conv_w[idx].float().transpose(0, 1)
        out = (win.float() * w).sum(0)
        return F.silu(out).to(val.dtype), win[1:]

    def step(self, x, st):
        H, Dk = self.cfg.kda_heads, self.cfg.kda_head_dim
        q = self.q_proj(x)
        k = self.k_proj(x)
        v = self.v_proj(x)
        q, st["cq"] = self._short_conv(q, st["cq"], 0)
        k, st["ck"] = self._short_conv(k, st["ck"], 1)
        v, st["cv"] = self._short_conv(v, st["cv"], 2)
        q = q.view(H, Dk).float() * self.scale
        k = k.view(H, Dk).float()
        v = v.view(H, Dk).float()
        g = (-F.softplus(self.g_proj(x).float())).view(H, Dk)
        beta = torch.sigmoid(self.beta_proj(x).float())
        S = st["S"] * g.exp()[:, :, None]
        pred = (S * k[:, :, None]).sum(1)
        S = S + beta[:, None, None] * k[:, :, None] * (v - pred)[:, None, :]
        o = (S * q[:, :, None]).sum(1)
        st["S"] = S
        return self.o_proj(o.reshape(H * Dk).to(torch.bfloat16))


# --------------------------------------------------------------------------- #
# MLA layer with absorb
# --------------------------------------------------------------------------- #
class MLA(nn.Module):
    def __init__(self, cfg: Config):
        super().__init__()
        self.cfg = cfg
        H, d = cfg.mla_heads, cfg.hidden
        self.q_proj = W4A16Linear(d, H * (cfg.qk_nope + cfg.qk_rope), cfg.group)
        self.kv_a = W4A16Linear(d, cfg.kv_lora + cfg.qk_rope, cfg.group)
        self.kv_b = W4A16Linear(cfg.kv_lora, H * (cfg.qk_nope + cfg.v_head), cfg.group)
        self.o_proj = W4A16Linear(H * cfg.v_head, d, cfg.group)
        self.scale = (cfg.qk_nope + cfg.qk_rope) ** -0.5

    def step(self, x, st):
        cfg = self.cfg
        H = cfg.mla_heads
        pos = st["c_kv"].shape[0]
        q = self.q_proj(x).view(H, cfg.qk_nope + cfg.qk_rope)
        q_nope = q[:, : cfg.qk_nope].float()
        q_rope = q[:, cfg.qk_nope:]
        kv = self.kv_a(x)
        c_kv = kv[: cfg.kv_lora]
        k_rope = kv[cfg.kv_lora:]
        cos, sin = _rope_cossin(pos, cfg.qk_rope, cfg.rope_theta, x.device)
        q_rope = _apply_rope(q_rope, cos, sin).float()
        k_rope = _apply_rope(k_rope, cos, sin)
        st["c_kv"] = torch.cat([st["c_kv"], c_kv[None]], 0)
        st["k_rope"] = torch.cat([st["k_rope"], k_rope[None]], 0)

        # Absorb: W_kv_b is (kv_lora, H * (qk_nope + v_head))
        # W_k[h, k, d] = W_kv_b[k, h*(qk_nope+v_head) + d] for d in [0, qk_nope)
        # W_v[h, k, d] = W_kv_b[k, h*(qk_nope+v_head) + d + qk_nope] for d in [0, v_head)
        W_kv_b_bf = self.kv_b.weight_bf().float()
        W_k = W_kv_b_bf[:, : H * cfg.qk_nope].view(cfg.kv_lora, H, cfg.qk_nope)
        W_v = W_kv_b_bf[:, H * cfg.qk_nope:].view(cfg.kv_lora, H, cfg.v_head)

        # A = (H, kv_lora) = einsum('hd,hdk->hk', q_nope, W_k.permute(1, 2, 0))
        A = torch.einsum("hd,hdk->hk", q_nope, W_k.permute(1, 2, 0))
        c_kv_f = st["c_kv"].float()
        scores = A @ c_kv_f.T
        scores = scores + q_rope @ st["k_rope"].float().T
        scores = scores * self.scale
        p = torch.softmax(scores, dim=-1)
        B = p @ c_kv_f
        output = torch.einsum("hk,hkd->hd", B, W_v.permute(1, 0, 2))
        return self.o_proj(output.reshape(H * cfg.v_head).to(torch.bfloat16))


# --------------------------------------------------------------------------- #
# MoE layer (grouped GEMV)
# --------------------------------------------------------------------------- #
class MoE(nn.Module):
    def __init__(self, cfg: Config):
        super().__init__()
        self.cfg = cfg
        d, m, E = cfg.hidden, cfg.moe_inter, cfg.n_experts
        self.router = nn.Linear(d, E, bias=False, dtype=cfg.dtype)
        self.gate = W4A16Experts(E, d, m, cfg.group)
        self.up = W4A16Experts(E, d, m, cfg.group)
        self.down = W4A16Experts(E, m, d, cfg.group)
        self.s_gate = W4A16Experts(cfg.n_shared, d, m, cfg.group)
        self.s_up = W4A16Experts(cfg.n_shared, d, m, cfg.group)
        self.s_down = W4A16Experts(cfg.n_shared, m, d, cfg.group)

    def step(self, x):
        cfg = self.cfg
        probs = torch.softmax(self.router(x).float(), dim=-1)
        w, idx = torch.topk(probs, cfg.n_active)
        w = (w / (w.sum() + 1e-9) * cfg.routed_scaling).to(x.dtype)

        # Gate and up: bf16 input, bf16 output
        g = w4a16_grouped_gemv(
            x, self.gate.w_q[idx], self.gate.scales[idx], self.gate.zeros[idx], self.gate.group
        )
        u = w4a16_grouped_gemv(
            x, self.up.w_q[idx], self.up.scales[idx], self.up.zeros[idx], self.up.group
        )
        hh = F.silu(g.float()) * u.float()  # [n_active, moe_inter] fp32
        # Down: fp32 input to preserve precision (matches reference's fp32 chain)
        d = w4a16_grouped_gemv(
            hh, self.down.w_q[idx], self.down.scales[idx], self.down.zeros[idx], self.down.group
        )
        out = (w[:, None].float() * d.float()).sum(0)

        # Shared expert (1)
        sidx = torch.zeros(1, dtype=torch.long, device=x.device)
        sg = w4a16_grouped_gemv(
            x, self.s_gate.w_q[sidx], self.s_gate.scales[sidx], self.s_gate.zeros[sidx], self.s_gate.group
        )
        su = w4a16_grouped_gemv(
            x, self.s_up.w_q[sidx], self.s_up.scales[sidx], self.s_up.zeros[sidx], self.s_up.group
        )
        sh = F.silu(sg.float()) * su.float()
        sd = w4a16_grouped_gemv(
            sh, self.s_down.w_q[sidx], self.s_down.scales[sidx], self.s_down.zeros[sidx], self.s_down.group
        ).squeeze(0)
        out = out + sd.float()

        return out.to(torch.bfloat16)


# --------------------------------------------------------------------------- #
# Block and Model
# --------------------------------------------------------------------------- #
class Block(nn.Module):
    def __init__(self, cfg, kind):
        super().__init__()
        self.kind = kind
        self.attn_norm = nn.Parameter(torch.ones(cfg.hidden, dtype=cfg.dtype))
        self.moe_norm = nn.Parameter(torch.ones(cfg.hidden, dtype=cfg.dtype))
        self.attn = KDA(cfg) if kind == "K" else MLA(cfg)
        self.moe = MoE(cfg)

    def step(self, x, st):
        h = x + self.attn.step(_rmsnorm(x, self.attn_norm), st)
        return h + self.moe.step(_rmsnorm(h, self.moe_norm))


class Model(nn.Module):
    def __init__(self, cfg: Config):
        super().__init__()
        self.cfg = cfg
        self.blocks = nn.ModuleList(Block(cfg, k) for k in cfg.pattern)
        self.reset_parameters()

    def reset_parameters(self):
        g = torch.Generator(device="cpu").manual_seed(1234)
        for mod in self.modules():
            if isinstance(mod, (W4A16Linear, W4A16Experts)):
                mod.init_random(g)
            elif isinstance(mod, nn.Linear):
                nn.init.normal_(mod.weight, 0.0, 0.02, generator=g)
            elif isinstance(mod, KDA):
                nn.init.normal_(mod.conv_w, 0.0, 0.1, generator=g)

    def step(self, hidden, state):
        for i, blk in enumerate(self.blocks):
            hidden = blk.step(hidden, state[i])
        return hidden, state