"""Triton kernels for FastDepth's fused frame body (fused=True).

_add_ln       x' = x + a [+ c[b]] with the same rounding order as torch (each add rounded to the depth dtype), stored;
              y = LayerNorm(x') (fp32 statistics, affine), stored.  Replaces 2-3 elementwise kernels + a LayerNorm.
_topk_sample  per row: logits/temp -> top-k mask -> exponential-race sample (== argmax(p/E), E ~ Exp(1)), i.e. the same
              distribution as the torch path (topk + masked_fill + softmax + exponential_ + div + argmax), one kernel.
              RNG = Philox (tl.rand) keyed by a device-side seed that the captured graph increments every replay.
"""
import torch
import triton
import triton.language as tl


@triton.jit
def _add_ln(X, A, C, XO, Y, W, Bb, eps, T,
            sxb, sxt, sab, sat, scb, sob, sot, syb, syt,
            N: tl.constexpr, HAS_C: tl.constexpr):
    row = tl.program_id(0)
    b = row // T
    t = row % T
    c = tl.arange(0, N)
    dt = XO.dtype.element_ty
    x = (tl.load(X + b * sxb + t * sxt + c).to(tl.float32) + tl.load(A + b * sab + t * sat + c).to(tl.float32)).to(dt)
    if HAS_C:
        x = (x.to(tl.float32) + tl.load(C + b * scb + c).to(tl.float32)).to(dt)
    tl.store(XO + b * sob + t * sot + c, x)
    xf = x.to(tl.float32)
    mean = tl.sum(xf, 0) / N
    d = xf - mean
    var = tl.sum(d * d, 0) / N
    y = d / tl.sqrt(var + eps) * tl.load(W + c).to(tl.float32) + tl.load(Bb + c).to(tl.float32)
    tl.store(Y + b * syb + t * syt + c, y.to(Y.dtype.element_ty))


def add_ln(x, a, ln, c=None):
    """x, a: [B, T, D] (any strides, last dim contiguous); c: [B, 1, D] or None; ln: nn.LayerNorm(D).
    Returns (x + a [+ c], LayerNorm(...)) as new contiguous tensors."""
    B, T, D = x.shape
    xo = torch.empty(B, T, D, device=x.device, dtype=x.dtype)
    y = torch.empty_like(xo)
    cc = c if c is not None else x
    _add_ln[(B * T,)](x, a, cc, xo, y, ln.weight, ln.bias, ln.eps, T,
                      x.stride(0), x.stride(1), a.stride(0), a.stride(1), cc.stride(0), xo.stride(0), xo.stride(1),
                      y.stride(0), y.stride(1), N=D, HAS_C=c is not None)
    return xo, y


@triton.jit
def _topk_sample(L, Out, Seed, Ctr, level, inv_temp, s_l, V: tl.constexpr, K: tl.constexpr, GREEDY: tl.constexpr,
                 PER_ROW: tl.constexpr):
    row = tl.program_id(0)
    c = tl.arange(0, V)
    x = tl.load(L + row * s_l + c).to(tl.float32) * inv_temp
    xs = tl.sort(x, descending=True)
    kth = tl.sum(tl.where(c == K - 1, xs, 0.0), 0)
    xm = tl.where(x >= kth, x, -float("inf"))
    if GREEDY:
        score = xm
    else:
        if PER_ROW:   # per-request stream: seed of this row's request, counter = frames generated by it
            seed = tl.load(Seed + row)
            ctr = tl.load(Ctr + row)
            u = tl.rand(seed, (ctr * 32 + level) * V + c)
        else:
            seed = tl.load(Seed)
            u = tl.rand(seed, (row * 64 + level) * V + c)
        u = tl.maximum(u, 1e-12)
        score = xm - tl.log(-tl.log(u))
    tl.store(Out + row, tl.argmax(score, 0).to(tl.int64))


def topk_sample(logits, seed_t, level, temp, top_k, greedy=False, ctr_t=None):
    """seed_t [1] (engine-wide stream) or, with ctr_t, per-row seeds [B] and counters [B] (reproducible per request)."""
    B, V = logits.shape
    out = torch.empty(B, device=logits.device, dtype=torch.long)
    per_row = ctr_t is not None
    _topk_sample[(B,)](logits, out, seed_t, ctr_t if per_row else seed_t, level, 1.0 / temp, logits.stride(0), V=V, K=top_k,
                       GREEDY=greedy, PER_ROW=per_row)
    return out


@triton.jit
def _row_exponential(Out, Seed, Ctr, level, s_o, V: tl.constexpr, BLOCK: tl.constexpr):
    """Out[row, :V] = Exp(1) noise from the row's own Philox stream (seed, counter): per-request reproducible sampling."""
    row = tl.program_id(0)
    seed = tl.load(Seed + row)
    ctr = tl.load(Ctr + row)
    for start in range(0, V, BLOCK):
        c = start + tl.arange(0, BLOCK)
        u = tl.rand(seed, (ctr * 32 + level) * 4096 + c)
        u = tl.maximum(u, 1e-12)
        tl.store(Out + row * s_o + c, -tl.log(u), mask=c < V)


def row_exponential(B, V, seeds, ctrs, level=0):
    out = torch.empty(B, V, device=seeds.device, dtype=torch.float32)
    _row_exponential[(B,)](out, seeds, ctrs, level, out.stride(0), V=V, BLOCK=1024)
    return out