"""Graph building blocks for exporting DETR-style detectors to the playground's WebGPU runtime.

The playground's current `webgpu_forward.wasm` registers 63 operators and links no CPU kernels,
so a program must be one delegate made only of those operators, with every elementwise tensor
at rank <= 4. These helpers rewrite the pieces of RT-DETR / D-FINE that fall outside that set
using only: add/sub/mul/div, relu, sigmoid, where.self, gt/ge/le/lt.Scalar, sum.dim,
gather (fp32 self + int32 index), _to_copy (fp32 <-> int32), view/permute/expand/slice/cat.

Each helper is exact (selection by comparison, integer arithmetic below 2^24) or an
algebraic identity of the original op; see the docstrings.
"""

import torch


def clamp(x, lo=None, hi=None):
    """Exact clamp written with where + compare-against-0.

    The current runtime's aten.clamp returns zeros for integer-valued bounds (the serializer
    stores them as Int), so clamp is never emitted. Bounds are lifted [1] constants.
    """
    if lo is not None:
        x = torch.where((x - lo) < 0, torch.tensor([float(lo)]), x)
    if hi is not None:
        x = torch.where((x - hi) > 0, torch.tensor([float(hi)]), x)
    return x


def where_max(a, b):
    """Elementwise max(a, b), exact: picks b where b - a > 0 (a - b is never 0 for a != b)."""
    return torch.where((b - a) > 0, b, a)


def where_min(a, b):
    return torch.where((b - a) < 0, b, a)


def tree_reduce(x, pick):
    """Reduce the last dim with a pairwise `pick` (where_max / where_min); keeps the dim (size 1)."""
    while x.shape[-1] > 1:
        n = x.shape[-1]
        h = n // 2
        m = pick(x[..., :h], x[..., h : 2 * h])
        x = torch.cat([m, x[..., 2 * h :]], dim=-1) if n % 2 else m
    return x


class TopKSet(torch.nn.Module):
    """Indices of the k largest entries of a [N] score vector, in ascending index order.

    torch.topk is not available on the runtime. The detectors' decoders are permutation
    equivariant over queries (self-attention with box-derived position embeddings, per-query
    deformable cross-attention), so only the selected *set* matters, not its order.

    1. Bisection for the k-th largest value t: `iters` steps of count(s >= mid) >= k.
       After enough steps lo is the k-th largest value (unless another score lies within
       (hi - lo) / 2^iters of it), so {s >= lo} is exactly the top-k set (ties aside).
    2. Stream compaction: inclusive prefix count of the selection mask (two-level block cumsum
       done with a triangular mask and sum, exact for integers), one-hot match of the
       positions 0..k-1 with relu(1 - d^2), and index = sum(onehot * arange) (exact, < 2^24).
    Returns a float tensor [k] of integer-valued indices.
    """

    def __init__(self, n, k, iters=48, block=None):
        super().__init__()
        self.n, self.k, self.iters = n, k, iters
        if block is None:
            block = next(b for b in range(int(n**0.5), 0, -1) if n % b == 0)
        self.block = block
        nb = n // block
        self.nb = nb
        self.register_buffer(
            "tri_in", torch.tril(torch.ones(block, block)), persistent=False
        )
        self.register_buffer(
            "tri_blk", torch.tril(torch.ones(nb, nb), diagonal=-1), persistent=False
        )
        self.register_buffer("one", torch.ones(1), persistent=False)
        self.register_buffer("zero", torch.zeros(1), persistent=False)
        self.register_buffer(
            "slots", torch.arange(k, dtype=torch.float32).view(k, 1), persistent=False
        )
        self.register_buffer(
            "arange", torch.arange(n, dtype=torch.float32).view(1, n), persistent=False
        )

    def forward(self, s):  # s: [1, N]
        lo = tree_reduce(s, where_min) - 1.0  # [1, 1]
        hi = tree_reduce(s, where_max) + 1.0
        for _ in range(self.iters):
            mid = (lo + hi) * 0.5
            cnt = torch.where((s - mid) >= 0, self.one, self.zero).sum(
                dim=-1, keepdim=True
            )
            ok = (cnt - self.k) >= 0
            lo = torch.where(ok, mid, lo)
            hi = torch.where(ok, hi, mid)
        sel = torch.where((s - lo) >= 0, self.one, self.zero)  # [1, N]
        blocks = sel.view(self.nb, 1, self.block)
        within = (blocks * self.tri_in).sum(
            dim=-1
        )  # [nb, block] inclusive count in block
        totals = blocks.sum(dim=-1).view(1, self.nb)  # [1, nb]
        offsets = (
            (totals * self.tri_blk).sum(dim=-1).view(self.nb, 1)
        )  # exclusive block prefix
        pos = (
            sel * (within + offsets).view(1, self.n)
        ) - 1.0  # selected -> 0.., others -> -1
        d = pos - self.slots  # [k, N]
        onehot = torch.relu(1.0 - d * d)
        return (onehot * self.arange).sum(dim=-1)  # [k]


def sigmoid_add_logit(delta, r, eps=1e-5):
    """sigmoid(delta + inverse_sigmoid(r)) without log/exp.

    With x = clamp(r, 0, 1), x1 = clamp(x, eps), x2 = clamp(1 - x, eps) (as in
    inverse_sigmoid), sigmoid(delta + log(x1 / x2)) = x1 s / (x1 s + x2 (1 - s)) with
    s = sigmoid(delta) and 1 - s = sigmoid(-delta).
    """
    x = clamp(r, 0.0, 1.0)
    x1 = clamp(x, eps)
    x2 = clamp(1.0 - x, eps)
    a = x1 * torch.sigmoid(delta)
    return a / (a + x2 * torch.sigmoid(-delta))


class DeformSampler(torch.nn.Module):
    """Multi-scale deformable attention core with bilinear grid_sample written as gathers.

    Matches F.grid_sample(mode="bilinear", padding_mode="zeros", align_corners=False) on grid
    g = 2 * loc - 1, pixel x = ((g + 1) * W - 1) / 2. floor() comes from an fp32 -> int32 cast of a
    clamped, non-negative value; out-of-range corners get weight 0 through an integer mask
    clamp(x0 + 1, 0, 1) * clamp(W - x0, 0, 1). Index clamping to [0, n-1] with a per-point n uses
    min(i - (n - 1), 0) + (n - 1), exact on small integers (clamp.Tensor is not available).
    All levels are concatenated along the spatial axis so one gather serves every
    (level, point, corner).

    points_per_level: number of sampling points per level, e.g. [6, 6] or [4, 4, 4].
    shapes: [(H, W), ...] per level.
    """

    def __init__(self, shapes, points_per_level, heads, head_dim):
        super().__init__()
        self.heads, self.head_dim = heads, head_dim
        self.P = sum(points_per_level)
        self.L = sum(h * w for h, w in shapes)
        Ws, Hs, starts = [], [], []
        start = 0
        for (h, w), p in zip(shapes, points_per_level):
            Ws += [float(w)] * p
            Hs += [float(h)] * p
            starts += [float(start)] * p
            start += h * w
        self.wmax, self.hmax = max(Ws), max(Hs)
        # Shaped [P, 1]: coordinates keep a trailing size-1 dim (slice, not select; the current
        # runtime rejects a select index that the serializer stored as a Double).
        self.register_buffer("W", torch.tensor(Ws).view(-1, 1), persistent=False)
        self.register_buffer("H", torch.tensor(Hs).view(-1, 1), persistent=False)
        self.register_buffer("start", torch.tensor(starts).view(-1, 1), persistent=False)

    def forward(self, value, loc, attn):
        """value [L, heads, head_dim]; loc [Q, heads, P, 2] in [0, 1]; attn [Q, heads, P].

        Returns [Q, heads * head_dim].
        """
        Q = loc.shape[0]
        nh, hd, P = self.heads, self.head_dim, self.P
        loc = loc.permute(1, 0, 2, 3)  # [heads, Q, P, 2]
        W, H = self.W, self.H
        # Same arithmetic as grid_sampler's unnormalize on grid = 2 * loc - 1.
        x = ((loc[..., 0:1] * 2.0 - 1.0 + 1.0) * W - 1.0) / 2.0  # [heads, Q, P, 1]
        y = ((loc[..., 1:2] * 2.0 - 1.0 + 1.0) * H - 1.0) / 2.0
        # trunc == floor once x + 1 >= 0. A point left of -1 or right of the level's W has
        # both corners outside, so clamping to [-1, max W] cannot change the result.
        x = clamp(x, -1.0, self.wmax)
        y = clamp(y, -1.0, self.hmax)
        x0 = (x + 1.0).to(torch.int32).to(torch.float32) - 1.0  # floor: x + 1 >= 0
        y0 = (y + 1.0).to(torch.int32).to(torch.float32) - 1.0
        fx = x - x0
        fy = y - y0
        x1 = x0 + 1.0
        y1 = y0 + 1.0

        def inside(i, n):
            return clamp(i + 1.0, 0.0, 1.0) * clamp(n - i, 0.0, 1.0)

        def clampi(i, n):
            return clamp(clamp(i, 0.0) - (n - 1.0), None, 0.0) + (n - 1.0)

        mx0, mx1, my0, my1 = inside(x0, W), inside(x1, W), inside(y0, H), inside(y1, H)
        cx0, cx1 = clampi(x0, W), clampi(x1, W)
        ry0 = clampi(y0, H) * W + self.start
        ry1 = clampi(y1, H) * W + self.start
        a = attn.permute(1, 0, 2).reshape(nh, Q, P, 1)
        w00 = (1.0 - fx) * (1.0 - fy) * mx0 * my0 * a
        w01 = fx * (1.0 - fy) * mx1 * my0 * a
        w10 = (1.0 - fx) * fy * mx0 * my1 * a
        w11 = fx * fy * mx1 * my1 * a
        idx = torch.cat([ry0 + cx0, ry0 + cx1, ry1 + cx0, ry1 + cx1], dim=-1)  # [heads, Q, P, 4]
        wts = torch.cat([w00, w01, w10, w11], dim=-1)
        K = Q * P * 4
        idx = idx.reshape(nh, 1, K).expand(nh, hd, K).to(torch.int64)
        v = value.permute(1, 2, 0)  # [heads, head_dim, L]
        g = torch.gather(v, 2, idx)  # [heads, head_dim, K]
        out = (
            (g * wts.reshape(nh, 1, K)).reshape(nh, hd, Q, P * 4).sum(dim=-1)
        )  # [heads, hd, Q]
        return out.permute(2, 0, 1).reshape(Q, nh * hd)


def attention(q_proj, k_proj, v_proj, o_proj, hidden, pos, heads):
    """Self-attention with position embeddings on queries/keys (DETR style). hidden [1, N, C]."""
    n, c = hidden.shape[1], hidden.shape[2]
    hd = c // heads
    qk_in = hidden + pos
    q = q_proj(qk_in).view(n, heads, hd).permute(1, 0, 2)
    k = k_proj(qk_in).view(n, heads, hd).permute(1, 2, 0)
    v = v_proj(hidden).view(n, heads, hd).permute(1, 0, 2)
    w = torch.softmax(torch.bmm(q, k) * (hd**-0.5), dim=-1)
    o = torch.bmm(w, v).permute(1, 0, 2).reshape(1, n, c)
    return o_proj(o)


class SiLU(torch.nn.Module):
    def forward(self, x):
        return x * torch.sigmoid(x)


class FoldedBatchNorm(torch.nn.Module):
    """x * scale + shift with the frozen-BN scale/shift precomputed exactly as the module does."""

    def __init__(self, weight, bias, mean, var, eps):
        super().__init__()
        scale = weight * (var + eps).rsqrt()
        self.register_buffer("scale", scale.reshape(1, -1, 1, 1))
        self.register_buffer("shift", (bias - mean * scale).reshape(1, -1, 1, 1))

    def forward(self, x):
        return x * self.scale + self.shift


def fold_frozen_bn_into_conv(root, bn_class):
    """For every module with `.convolution` + frozen-BN `.normalization`, fold the BN into the conv.

    Used where a conv output is too large for the runtime's 1D elementwise dispatch limit
    (~4.2M elements): folded, the following ReLU fuses into et_vk.conv_with_clamp.
    """
    n = 0
    for m in root.modules():
        conv, bn = getattr(m, "convolution", None), getattr(m, "normalization", None)
        if isinstance(conv, torch.nn.Conv2d) and isinstance(bn, bn_class):
            scale = bn.weight * (bn.running_var + 1e-5).rsqrt()
            shift = bn.bias - bn.running_mean * scale
            with torch.no_grad():
                conv.weight.mul_(scale.view(-1, 1, 1, 1))
                bias = shift if conv.bias is None else conv.bias * scale + shift
                conv.bias = torch.nn.Parameter(bias.clone())
            m.normalization = torch.nn.Identity()
            n += 1
    return n


def relabel_golden(out_dir, note):
    """_common.export_vision labels its reference case as the internal goldens; here it is upstream eager."""
    import json
    from pathlib import Path

    p = Path(out_dir) / "validation.json"
    v = json.loads(p.read_text())
    if "golden" in v["cases"]:
        v["cases"]["golden"]["golden_note"] = note
    p.write_text(json.dumps(v, indent=2) + "\n")


def replace_modules(root, predicate, factory):
    n = 0
    for parent in list(root.modules()):
        for name, child in list(parent.named_children()):
            if predicate(child):
                setattr(parent, name, factory(child))
                n += 1
    return n


def yolo_layout_buf(boxes_cxcywh, scores, pix_scale):
    """Like yolo_layout, with the [w, h, w, h] scale passed in as a registered buffer."""
    return torch.cat([boxes_cxcywh * pix_scale, scores], dim=-1).permute(0, 2, 1)


def yolo_layout(boxes_cxcywh, scores, size):
    """[1, Q, 4] normalized boxes + [1, Q, C] probabilities -> [1, 4 + C, Q] in input pixels.

    The layout the playground's YOLO decoder (src/vision/detect.js) already reads.
    """
    w, h = size
    scale = torch.tensor([w, h, w, h], dtype=torch.float32)
    return torch.cat([boxes_cxcywh * scale, scores], dim=-1).permute(0, 2, 1)