"""
TITAN kernel dispatch layer.

Every primitive has:
  * a FAST path (chosen by measurement — see reports/01_pmlf_audit.md)
  * a REFERENCE path (simple, obviously-correct)
Toggle with TITAN_KERNELS=ref|fast (env) or kernels.set_mode().

Measured decisions (Apple M5 Max, MLX 0.32.0, bf16):
  rmsnorm   : mx.fast.rms_norm  0.276ms  vs PMLF fuse 0.670ms  -> NATIVE (2.43x)
  swiglu    : PMLF fuse.swiglu  2.212ms  vs native 2.110ms     -> PMLF (tied, modular)
  attention : mx.fast.sdpa      1.364ms  vs manual 2.420ms     -> NATIVE (PMLF has none)
  rope      : mx.fast.rope      0.278ms  vs manual 0.507ms     -> NATIVE (PMLF has none)
  packing   : PMLF pack         100% vs 26.5% efficiency        -> PMLF (3.77x FLOP win)
"""
import os, sys
import mlx.core as mx
import mlx.nn as nn

_PMLF_PATH = '/Users/jonah/dataset_prep/kernels'
if _PMLF_PATH not in sys.path:
    sys.path.insert(0, _PMLF_PATH)

# ------------------------------------------------------------------ PMLF load
PMLF_AVAILABLE = False
PMLF_VERSION = 'unavailable'
try:
    from pmlf.fuse import swiglu as _pmlf_swiglu
    from pmlf.pack import plan_packing as _pmlf_plan_packing
    PMLF_AVAILABLE = True
    PMLF_VERSION = 'local-2026-08-06+sha:3597f8a2'
except Exception as _e:                                    # pragma: no cover
    _pmlf_swiglu = None
    _pmlf_plan_packing = None

MODE = os.environ.get('TITAN_KERNELS', 'fast')


def set_mode(m):
    global MODE
    assert m in ('fast', 'ref')
    MODE = m


def used_kernels():
    return {
        'mode': MODE,
        'pmlf_available': PMLF_AVAILABLE,
        'pmlf_version': PMLF_VERSION,
        'rmsnorm': 'mx.fast.rms_norm' if MODE == 'fast' else 'reference',
        'swiglu': ('pmlf.fuse.swiglu' if (MODE == 'fast' and PMLF_AVAILABLE)
                   else 'reference'),
        'attention': 'mx.fast.scaled_dot_product_attention' if MODE == 'fast' else 'reference',
        'rope': 'mx.fast.rope' if MODE == 'fast' else 'reference',
        'packing': 'pmlf.pack.plan_packing (FFD)' if PMLF_AVAILABLE else 'reference FFD',
    }


# ------------------------------------------------------------------- RMSNorm
def rmsnorm(x, weight, eps=1e-6):
    if MODE == 'fast':
        return mx.fast.rms_norm(x, weight, eps)
    v = mx.mean(x.astype(mx.float32) ** 2, axis=-1, keepdims=True)
    return (x.astype(mx.float32) * mx.rsqrt(v + eps)).astype(x.dtype) * weight


def rmsnorm_ref(x, weight, eps=1e-6):
    v = mx.mean(x.astype(mx.float32) ** 2, axis=-1, keepdims=True)
    return (x.astype(mx.float32) * mx.rsqrt(v + eps)).astype(x.dtype) * weight


# -------------------------------------------------------------------- SwiGLU
def swiglu(x, w_gate, w_up, w_down):
    if MODE == 'fast' and PMLF_AVAILABLE:
        return _pmlf_swiglu(x, w_gate, w_up, w_down)
    g = x @ w_gate
    return ((g * mx.sigmoid(g)) * (x @ w_up)) @ w_down


def swiglu_ref(x, w_gate, w_up, w_down):
    g = x @ w_gate
    return ((g * mx.sigmoid(g)) * (x @ w_up)) @ w_down


# ---------------------------------------------------------------------- RoPE
def rope(x, dims, base=10000.0, offset=0):
    """x: (B, H, T, Dh). offset may be an int or an int array (packed restarts)."""
    if MODE == 'fast':
        return mx.fast.rope(x, dims, traditional=False, base=base, scale=1.0,
                            offset=offset)
    return rope_ref(x, dims, base=base, offset=offset)


def rope_ref(x, dims, base=10000.0, offset=0):
    import numpy as np
    T, d = x.shape[-2], dims
    inv = mx.array((1.0 / (base ** (np.arange(0, d, 2) / d))).astype(np.float32))
    pos = mx.arange(T, dtype=mx.float32) + (offset if isinstance(offset, int) else 0)
    ang = pos[:, None] * inv[None, :]
    cos, sin = mx.cos(ang).astype(x.dtype), mx.sin(ang).astype(x.dtype)
    x1, x2 = x[..., 0::2], x[..., 1::2]
    o1 = x1 * cos - x2 * sin
    o2 = x1 * sin + x2 * cos
    return mx.stack([o1, o2], axis=-1).reshape(x.shape)


# ----------------------------------------------------------------- attention
def sdpa(q, k, v, scale, mask=None):
    """q,k,v: (B, H, T, Dh). mask: None | 'causal' | bool array."""
    if MODE == 'fast':
        return mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask=mask)
    return sdpa_ref(q, k, v, scale, mask)


def sdpa_ref(q, k, v, scale, mask=None):
    s = (q * scale) @ mx.swapaxes(k, -1, -2)
    if isinstance(mask, str) and mask == 'causal':
        T = q.shape[-2]
        import numpy as np
        cm = mx.array(np.tril(np.ones((T, T), dtype=bool)))
        s = mx.where(cm, s, mx.array(-1e9, dtype=s.dtype))
    elif mask is not None and not isinstance(mask, str):
        s = mx.where(mask, s, mx.array(-1e9, dtype=s.dtype))
    return mx.softmax(s.astype(mx.float32), axis=-1).astype(q.dtype) @ v


# ----------------------------------------------------------- GQA head repeat
def repeat_kv(x, n_rep):
    """(B, n_kv, T, Dh) -> (B, n_kv*n_rep, T, Dh)"""
    if n_rep == 1:
        return x
    B, H, T, D = x.shape
    return mx.repeat(x, n_rep, axis=1)


# --------------------------------------------------------------- chunked CE
def cross_entropy_chunked(hidden, w_out, targets, mask=None, chunks=4,
                          label_smoothing=0.0, z_loss=0.0):
    """Memory-bounded CE. hidden (N,d) bf16, w_out (d,V), targets (N,) int32.
    mask (N,) bool/float: 1 where the token counts. Returns (loss, n_tokens)."""
    N = hidden.shape[0]
    cs = max(1, (N + chunks - 1) // chunks)
    tot = mx.zeros((), dtype=mx.float32)
    cnt = mx.zeros((), dtype=mx.float32)
    for i in range(0, N, cs):
        h = hidden[i:i + cs]
        y = targets[i:i + cs]
        lg = (h @ w_out).astype(mx.float32)
        lse = mx.logsumexp(lg, axis=-1)
        tgt = mx.take_along_axis(lg, y[:, None], axis=-1).squeeze(-1)
        nll = lse - tgt
        if label_smoothing > 0.0:
            V = lg.shape[-1]
            smooth = lse - mx.mean(lg, axis=-1)
            nll = (1 - label_smoothing) * nll + label_smoothing * smooth
        if z_loss > 0.0:
            nll = nll + z_loss * (lse ** 2)
        if mask is not None:
            m = mask[i:i + cs].astype(mx.float32)
            tot = tot + mx.sum(nll * m)
            cnt = cnt + mx.sum(m)
        else:
            tot = tot + mx.sum(nll)
            cnt = cnt + nll.shape[0]
    return tot / mx.maximum(cnt, 1.0), cnt