"""Inside-outside and collision queries straight from an FCP descriptor (no decode, no mesh).

A point is inside when the signed sum over the parts that span its height of "inside the part's ring at that height"
is positive (a hole part counts -1). The ring at the point's canonical height s:

  mode "poly":  the stored polynomial itself: the ring's 2K Fourier coefficients are the Chebyshev series of the part
                at the local height u (clamped to the outer cell centres, as `_parts.part_u`), then the ring's NT
                vertices x(t_v), y(t_v) = sum_k c_k F_k(t_v), then an even-odd crossing test.
  mode "table": the rings at the part's layer centres are evaluated once (NT vertices each, the decode's
                `_heights.centre_values`); at s the two neighbouring rings are blended vertex by vertex (end rings
                held in the outer half cells), then the crossing test.

Both skip a part whose (ring) box misses the point. Everything is float32 in canonical coordinates (the crossing test
is affine-invariant). Collision of two placed descriptors: the ring vertices of each (the table rings at the layer
centres, optionally strided) are mapped into the other's canonical frame and tested; a pose collides when any is
inside. CPU: numba (parallel); GPU: CuPy raw kernels.
"""
from __future__ import annotations

import sys

import numba as nb
import numpy as np

sys.path.insert(0, "/home/joao/code/fcp_devel/api")
from fcp import _heights, _parts  # noqa: E402
from fcp._slice import L  # noqa: E402

F = np.float32


class FCPQuery:
    def __init__(self, d, nt: int = 128, nb: int | None = None):
        """nt: ring vertices; nb: y-bins per layer interval (default nt / 4: about 5-10 edges per bin at any nt)."""
        nb = max(32, nt // 4) if nb is None else nb
        self.d, self.nt, self.NB = d, nt, nb
        K, M, layers = d.K, d.M, d.layers
        self.K, self.M, self.layers = K, M, layers
        t = (2 * np.pi * np.arange(nt) / nt).astype(F)
        Fm = np.zeros((K, nt), F)
        Fm[0] = 1
        h = np.arange(1, d.N + 1, dtype=F)[:, None]
        Fm[1::2], Fm[2::2] = np.cos(h * t[None]), np.sin(h * t[None])
        self.Fm = Fm
        coef, B = d.coef(), 2 * K
        sign, s0, s1, h0s, h1s, first, C, rings = [], [], [], [], [], [], [], []
        nring = 0
        for i, (sg, h0, h1) in enumerate(d.parts):
            blk = coef[i * B:(i + 1) * B]
            if not np.any(blk):
                continue
            n = h1 - h0 + 1
            a, b = _parts.part_range({"h0": h0, "h1": h1}, layers)
            sign.append(sg), s0.append(a), s1.append(b), h0s.append(h0), h1s.append(h1), first.append(nring)
            C.append(blk.T.astype(F))  # (M + 1, 2K)
            Y = _heights.centre_values(blk.T, n).reshape(n, 2, K)
            rings.append(np.einsum("nck,kv->nvc", Y, Fm).astype(F))  # (n, nt, 2)
            nring += n
        self.P = len(sign)
        self.sign = np.array(sign, np.int32)
        self.s0, self.s1 = np.array(s0, F), np.array(s1, F)
        self.h0, self.h1 = np.array(h0s, np.int32), np.array(h1s, np.int32)
        self.first = np.array(first, np.int32)
        self.C = np.ascontiguousarray(np.stack(C)) if C else np.zeros((0, M + 1, B), F)  # (P, M + 1, 2K)
        self.rings = np.ascontiguousarray(np.concatenate(rings)) if rings else np.zeros((0, nt, 2), F)
        R = self.rings
        self.rbox = np.ascontiguousarray(np.stack([R[:, :, 0].min(1), R[:, :, 0].max(1), R[:, :, 1].min(1), R[:, :, 1].max(1)], 1), F)
        self.pbox = np.ascontiguousarray(np.stack([
            [self.rbox[f:f + n, 0].min(), self.rbox[f:f + n, 1].max(), self.rbox[f:f + n, 2].min(), self.rbox[f:f + n, 3].max()]
            for f, n in zip(self.first, self.h1 - self.h0 + 1)]), F) if self.P else np.zeros((0, 4), F)
        # the poly mode's box is the table box padded (the polynomial between centres can bulge a little)
        pad = 0.02 * max(1e-6, float(np.abs(self.pbox).max())) if self.P else 0
        self.pbox_poly = self.pbox + np.array([-pad, pad, -pad, pad], F)
        c = -L + (np.arange(layers) + 0.5) * 2 * L / layers
        self.zc = c.astype(F)
        self.w = F(2 * L / layers)
        self.slot_first, self.sl_r, self.sl_f, self.bin_off, self.edges = _build_slots(self.rings, self.first, self.h0, self.h1, nb)
        self._gpu = None

    def bin_args(self):
        return (self.rings, self.sign, self.h0, self.h1, self.slot_first, self.sl_r, self.sl_f, self.bin_off, self.edges,
                self.zc[0], self.w, np.int32(self.NB))

    # ------------------------------------------------------------ surface samples (for collision)
    def surface_points(self, layer_stride: int = 1, vertex_stride: int = 1, caps: bool = False):
        """Canonical (n, 3) ring vertices at the layer centres (every layer_stride-th layer of each part, its end
        layers always, every vertex_stride-th vertex). caps: also the part ends' discs (the end ring on the part's
        outer edge plane, and a grid inside it at layer_stride layer heights). Every sample lies in the solid or on
        its surface (a disc inside the solid at a junction too), so a sample inside the other object is a real overlap."""
        from matplotlib.path import Path

        out = []
        h = float(self.w) * layer_stride
        for p in range(self.P):
            n = self.h1[p] - self.h0[p] + 1
            ls = sorted(set(range(0, n, layer_stride)) | {n - 1})
            for l in ls:
                xy = self.rings[self.first[p] + l, ::vertex_stride]
                out.append(np.column_stack([xy, np.full(len(xy), self.zc[self.h0[p] + l], F)]))
            if not caps:
                continue
            for l, z in ((0, self.s0[p]), (n - 1, self.s1[p])):
                ring = self.rings[self.first[p] + l]
                lo, hi = ring.min(0), ring.max(0)
                gx, gy = np.meshgrid(np.arange(lo[0] + h / 2, hi[0], h), np.arange(lo[1] + h / 2, hi[1], h))
                g = np.column_stack([gx.ravel(), gy.ravel()])
                g = g[Path(ring).contains_points(g)] if len(g) else g
                for xy in (ring[::vertex_stride], g):
                    out.append(np.column_stack([xy, np.full(len(xy), z, F)]))
        return np.concatenate(out).astype(F)

    # ------------------------------------------------------------ CPU
    def inside_cpu(self, Pc, mode="table"):
        """Pc: (n, 3) canonical points -> bool (n,)."""
        Pc = np.ascontiguousarray(Pc, F)
        out = np.zeros(len(Pc), np.bool_)
        if mode == "table":
            _inside_bins(Pc, *self.bin_args(), out)
        elif mode == "scan":
            _inside_table(Pc, self.rings, self.rbox, self.pbox, self.sign, self.h0, self.h1, self.first, self.zc[0], self.w,
                          self.layers, out)
        else:
            _inside_poly(Pc, self.C, self.Fm, self.pbox_poly, self.sign, self.s0, self.s1, self.h0, self.h1, out)
        return out

    def inside_world_cpu(self, P, mode="table"):
        return self.inside_cpu(self.d.pose.to_canonical(np.asarray(P, float)), mode)

    # ------------------------------------------------------------ GPU
    def gpu(self):
        if self._gpu is None:
            import cupy as cp

            g = {k: cp.asarray(getattr(self, k)) for k in ("rings", "rbox", "pbox", "pbox_poly", "sign", "h0", "h1", "first", "s0",
                                                           "s1", "C", "Fm", "slot_first", "sl_r", "sl_f",
                                                                                          "bin_off", "edges")}
            self._gpu = g
        return self._gpu

    def inside_gpu(self, Pc, mode="table"):
        """Pc: CuPy (n, 3) float32 canonical points -> CuPy bool (n,)."""
        import cupy as cp

        g = self.gpu()
        n = len(Pc)
        out = cp.zeros(n, cp.bool_)
        if n == 0 or self.P == 0:
            return out
        tb = 128
        if mode == "table":
            _kern("inside_bins")((-(-n // tb),), (tb,), (Pc, np.int32(n), g["rings"], g["sign"], g["h0"], g["h1"], g["slot_first"],
                                                      g["sl_r"], g["sl_f"], g["bin_off"], g["edges"], np.int32(self.P),
                                                      self.zc[0], self.w, np.int32(self.nt), np.int32(self.NB), out))
        elif mode == "scan":
            _kern("inside_table")((-(-n // tb),), (tb,), (Pc, np.int32(n), g["rings"], g["rbox"], g["pbox"], g["sign"], g["h0"], g["h1"],
                                                       g["first"], np.int32(self.P), self.zc[0], self.w, np.int32(self.layers),
                                                       np.int32(self.nt), out))
        else:
            _kern("inside_poly")((-(-n // tb),), (tb,), (Pc, np.int32(n), g["C"], g["Fm"], g["pbox_poly"], g["sign"], g["s0"], g["s1"],
                                                      g["h0"], g["h1"], np.int32(self.P), np.int32(self.M), np.int32(self.K),
                                                      np.int32(self.nt), out),
                                 shared_mem=self.K * self.nt * 4)
        return out


# ===================================================================== numba

@nb.njit(cache=True, inline="always")
def _cross(x, y, ax, ay, bx, by, c):
    if (ay > y) != (by > y):
        if x < ax + (y - ay) * (bx - ax) / (by - ay):
            return 1 - c
    return c


@nb.njit(cache=True, inline="always")
def _ring_inside_lerp(rings, r0, r1, a, x, y, nt):
    c = 0
    b = np.float32(1) - a
    px = b * rings[r0, nt - 1, 0] + a * rings[r1, nt - 1, 0]
    py = b * rings[r0, nt - 1, 1] + a * rings[r1, nt - 1, 1]
    for v in range(nt):
        qx = b * rings[r0, v, 0] + a * rings[r1, v, 0]
        qy = b * rings[r0, v, 1] + a * rings[r1, v, 1]
        c = _cross(x, y, px, py, qx, qy, c)
        px, py = qx, qy
    return c


@nb.njit(cache=True, inline="always")
def _point_table(x, y, z, rings, rbox, pbox, sign, h0, h1, first, z0, w, layers):
    nt = rings.shape[1]
    s = (z - z0) / w  # layer coordinate: centre j at s = j
    acc = 0
    for p in range(len(sign)):
        if s < h0[p] - 0.5 or s > h1[p] + 0.5:
            continue
        if x < pbox[p, 0] or x > pbox[p, 1] or y < pbox[p, 2] or y > pbox[p, 3]:
            continue
        if s <= h0[p]:
            j0, j1, a = h0[p], h0[p], np.float32(0)
        elif s >= h1[p]:
            j0, j1, a = h1[p], h1[p], np.float32(0)
        else:
            j0 = int(np.floor(s))
            j1 = j0 + 1
            a = np.float32(s - j0)
        r0, r1 = first[p] + j0 - h0[p], first[p] + j1 - h0[p]
        if (x < min(rbox[r0, 0], rbox[r1, 0]) or x > max(rbox[r0, 1], rbox[r1, 1]) or
                y < min(rbox[r0, 2], rbox[r1, 2]) or y > max(rbox[r0, 3], rbox[r1, 3])):
            continue
        if _ring_inside_lerp(rings, r0, r1, a, x, y, nt):
            acc += sign[p]
    return acc > 0


@nb.njit(cache=True, parallel=True)
def _inside_table(Pc, rings, rbox, pbox, sign, h0, h1, first, z0, w, layers, out):
    for i in nb.prange(len(Pc)):
        out[i] = _point_table(Pc[i, 0], Pc[i, 1], Pc[i, 2], rings, rbox, pbox, sign, h0, h1, first, z0, w, layers)


@nb.njit(cache=True)
def _build_slots(rings, first, h0, h1, NB):
    """Slots of a part with n rings: 0 = the first ring held (lower half cell), l = between rings l - 1 and l, n = the
    last ring held. Per slot: its rows (r0, r1), its box and NB bins in y, each listing the edges v (vertex v - 1 to
    v) whose blended y range (any blend of the two rings) meets the bin."""
    P, nt = len(first), rings.shape[1]
    S = 0
    for p in range(P):
        S += h1[p] - h0[p] + 2
    slot_first = np.empty(P, np.int32)
    sl_r = np.empty((S, 2), np.int32)
    sl_f = np.empty((S, 6), np.float32)
    lo = np.empty(nt, np.float32)
    hi = np.empty(nt, np.float32)
    cnt = np.zeros(S * NB + 1, np.int32)
    k = 0
    for p in range(P):
        slot_first[p] = k
        n = h1[p] - h0[p] + 1
        for l in range(n + 1):
            r0 = first[p] + max(l - 1, 0)
            r1 = first[p] + min(l, n - 1)
            sl_r[k, 0], sl_r[k, 1] = r0, r1
            k += 1
    # pass 1: boxes and counts; pass 2: fill
    edges_n = 0
    for k in range(S):
        r0, r1 = sl_r[k, 0], sl_r[k, 1]
        x0, x1, y0, y1 = np.float32(3e38), np.float32(-3e38), np.float32(3e38), np.float32(-3e38)
        for v in range(nt):
            for r in (r0, r1):
                x0 = min(x0, rings[r, v, 0]); x1 = max(x1, rings[r, v, 0])
                y0 = min(y0, rings[r, v, 1]); y1 = max(y1, rings[r, v, 1])
        h = (y1 - y0) / np.float32(NB)
        sl_f[k, 0], sl_f[k, 2], sl_f[k, 3], sl_f[k, 4], sl_f[k, 5] = y0, x0, x1, y0, y1
        sl_f[k, 1] = np.float32(1) / h if h > 0 else np.float32(0)
        ih = sl_f[k, 1]  # the stored float32 value in both passes and in the query (the same bins)
        for v in range(nt):
            u = v - 1 if v > 0 else nt - 1
            a = min(min(rings[r0, u, 1], rings[r0, v, 1]), min(rings[r1, u, 1], rings[r1, v, 1]))
            b = max(max(rings[r0, u, 1], rings[r0, v, 1]), max(rings[r1, u, 1], rings[r1, v, 1]))
            b0 = min(max(int((a - y0) * ih), 0), NB - 1)
            b1 = min(max(int((b - y0) * ih), 0), NB - 1)
            for q in range(b0, b1 + 1):
                cnt[k * NB + q + 1] += 1
                edges_n += 1
    off = np.cumsum(cnt).astype(np.int32)
    edges = np.empty(edges_n, np.int32)
    fill = off[:-1].copy()
    for k in range(S):
        r0, r1 = sl_r[k, 0], sl_r[k, 1]
        y0, ih = sl_f[k, 0], sl_f[k, 1]
        for v in range(nt):
            u = v - 1 if v > 0 else nt - 1
            a = min(min(rings[r0, u, 1], rings[r0, v, 1]), min(rings[r1, u, 1], rings[r1, v, 1]))
            b = max(max(rings[r0, u, 1], rings[r0, v, 1]), max(rings[r1, u, 1], rings[r1, v, 1]))
            b0 = min(max(int((a - y0) * ih), 0), NB - 1)
            b1 = min(max(int((b - y0) * ih), 0), NB - 1)
            for q in range(b0, b1 + 1):
                edges[fill[k * NB + q]] = v
                fill[k * NB + q] += 1
    return slot_first, sl_r, sl_f, off, edges


@nb.njit(cache=True, inline="always")
def _point_bins(x, y, z, rings, sign, h0, h1, slot_first, sl_r, sl_f, bin_off, edges, z0, w, NB):
    nt = rings.shape[1]
    s = (z - z0) / w
    acc = 0
    for p in range(len(sign)):
        if s < h0[p] - 0.5 or s > h1[p] + 0.5:
            continue
        if s <= h0[p]:
            l, a = 0, np.float32(0)
        elif s >= h1[p]:
            l, a = h1[p] - h0[p] + 1, np.float32(0)
        else:
            j0 = int(np.floor(s))
            l, a = j0 - h0[p] + 1, np.float32(s - j0)
        k = slot_first[p] + l
        if x < sl_f[k, 2] or x > sl_f[k, 3] or y < sl_f[k, 4] or y > sl_f[k, 5]:
            continue
        q = min(int((y - sl_f[k, 0]) * sl_f[k, 1]), NB - 1)
        r0, r1 = sl_r[k, 0], sl_r[k, 1]
        b = np.float32(1) - a
        c = 0
        for e in range(bin_off[k * NB + q], bin_off[k * NB + q + 1]):
            v = edges[e]
            u = v - 1 if v > 0 else nt - 1
            px = b * rings[r0, u, 0] + a * rings[r1, u, 0]
            py = b * rings[r0, u, 1] + a * rings[r1, u, 1]
            qx = b * rings[r0, v, 0] + a * rings[r1, v, 0]
            qy = b * rings[r0, v, 1] + a * rings[r1, v, 1]
            c = _cross(x, y, px, py, qx, qy, c)
        if c:
            acc += sign[p]
    return acc > 0


@nb.njit(cache=True, parallel=True)
def _inside_bins(Pc, rings, sign, h0, h1, slot_first, sl_r, sl_f, bin_off, edges, z0, w, NB, out):
    for i in nb.prange(len(Pc)):
        out[i] = _point_bins(Pc[i, 0], Pc[i, 1], Pc[i, 2], rings, sign, h0, h1, slot_first, sl_r, sl_f, bin_off, edges, z0, w, NB)


@nb.njit(cache=True, parallel=True)
def _inside_poly(Pc, C, Fm, pbox, sign, s0, s1, h0, h1, out):
    M1, K2 = C.shape[1], C.shape[2]
    K, nt = Fm.shape
    for i in nb.prange(len(Pc)):
        x, y, z = Pc[i, 0], Pc[i, 1], Pc[i, 2]
        T = np.empty(M1, np.float32)
        c = np.empty(K2, np.float32)
        acc = 0
        for p in range(len(sign)):
            if z < s0[p] or z > s1[p]:
                continue
            if x < pbox[p, 0] or x > pbox[p, 1] or y < pbox[p, 2] or y > pbox[p, 3]:
                continue
            e = np.float32(1) - np.float32(1) / np.float32(h1[p] - h0[p] + 1)
            u = np.float32(2) * (z - s0[p]) / (s1[p] - s0[p]) - np.float32(1)
            u = min(max(u, -e), e)
            T[0] = 1
            if M1 > 1:
                T[1] = u
            for j in range(2, M1):
                T[j] = 2 * u * T[j - 1] - T[j - 2]
            for k in range(K2):
                acc_k = np.float32(0)
                for j in range(M1):
                    acc_k += T[j] * C[p, j, k]
                c[k] = acc_k
            cr = 0
            px = np.float32(0)
            py = np.float32(0)
            for k in range(K):
                px += c[k] * Fm[k, nt - 1]
                py += c[K + k] * Fm[k, nt - 1]
            for v in range(nt):
                qx = np.float32(0)
                qy = np.float32(0)
                for k in range(K):
                    qx += c[k] * Fm[k, v]
                    qy += c[K + k] * Fm[k, v]
                cr = _cross(x, y, px, py, qx, qy, cr)
                px, py = qx, qy
            if cr:
                acc += sign[p]
        out[i] = acc > 0


@nb.njit(cache=True, parallel=True)
def collide_cpu_batch(PA, PB, MAB, tAB, MBA, tBA, A_args, B_args, out):
    """Poses k: canonical A -> canonical B by x_B = MAB[k] x_A + tAB[k] (and the inverse). out[k] = any point of A
    inside B or of B inside A. A_args / B_args: the table arrays of each (rings, rbox, pbox, sign, h0, h1, first, z0,
    w, layers)."""
    for k in nb.prange(len(out)):
        hit = False
        for i in range(len(PA)):
            x = MAB[k, 0, 0] * PA[i, 0] + MAB[k, 0, 1] * PA[i, 1] + MAB[k, 0, 2] * PA[i, 2] + tAB[k, 0]
            y = MAB[k, 1, 0] * PA[i, 0] + MAB[k, 1, 1] * PA[i, 1] + MAB[k, 1, 2] * PA[i, 2] + tAB[k, 1]
            z = MAB[k, 2, 0] * PA[i, 0] + MAB[k, 2, 1] * PA[i, 1] + MAB[k, 2, 2] * PA[i, 2] + tAB[k, 2]
            if abs(x) > 1.0 or abs(y) > 1.0 or abs(z) > 1.0:
                continue
            if _point_bins(x, y, z, B_args[0], B_args[1], B_args[2], B_args[3], B_args[4], B_args[5], B_args[6], B_args[7],
                           B_args[8], B_args[9], B_args[10], B_args[11]):
                hit = True
                break
        if not hit:
            for i in range(len(PB)):
                x = MBA[k, 0, 0] * PB[i, 0] + MBA[k, 0, 1] * PB[i, 1] + MBA[k, 0, 2] * PB[i, 2] + tBA[k, 0]
                y = MBA[k, 1, 0] * PB[i, 0] + MBA[k, 1, 1] * PB[i, 1] + MBA[k, 1, 2] * PB[i, 2] + tBA[k, 1]
                z = MBA[k, 2, 0] * PB[i, 0] + MBA[k, 2, 1] * PB[i, 1] + MBA[k, 2, 2] * PB[i, 2] + tBA[k, 2]
                if abs(x) > 1.0 or abs(y) > 1.0 or abs(z) > 1.0:
                    continue
                if _point_bins(x, y, z, A_args[0], A_args[1], A_args[2], A_args[3], A_args[4], A_args[5], A_args[6], A_args[7],
                               A_args[8], A_args[9], A_args[10], A_args[11]):
                    hit = True
                    break
        out[k] = hit


def table_args(q: FCPQuery):
    return q.bin_args()


def relative_maps(qA: FCPQuery, qB: FCPQuery, R, t):
    """B placed by world x = R[k] x_B + t[k] (A at its encoded world frame): the canonical maps A -> B and B -> A,
    float32 (k, 3, 3), (k, 3)."""
    pa, pb = qA.d.pose, qB.d.pose
    SA = pa.rotation * pa.scale  # canonical A -> world: x = ca + SA c
    SB = pb.rotation * pb.scale
    # world of B (the frame B was encoded in) = R^T (x - t); canonical B = SB^-1 (xB - cb)
    SBi, SAi = np.linalg.inv(SB), np.linalg.inv(SA)
    RT = np.transpose(R, (0, 2, 1))
    MAB = SBi @ RT @ SA  # (k, 3, 3)
    tAB = (SBi @ (RT @ (pa.center[None, :, None] - t[:, :, None]))).squeeze(-1) - SBi @ pb.center
    MBA = SAi @ R @ SB
    tBA = (SAi @ ((R @ pb.center[None, :, None]).squeeze(-1) + t - pa.center).T).T
    return tuple(np.ascontiguousarray(a, F) for a in (MAB, tAB, MBA, tBA))


# ===================================================================== CUDA

_SRC = r"""
__device__ __forceinline__ int cross_(float x, float y, float ax, float ay, float bx, float by, int c) {
    if ((ay > y) != (by > y)) { if (x < ax + (y - ay) * (bx - ax) / (by - ay)) return 1 - c; }
    return c;
}

__device__ bool point_table(float x, float y, float z, const float* rings, const float* rbox, const float* pbox,
                            const int* sign, const int* h0, const int* h1, const int* first, int P, float z0, float w, int nt) {
    float s = (z - z0) / w;
    int acc = 0;
    for (int p = 0; p < P; ++p) {
        if (s < h0[p] - 0.5f || s > h1[p] + 0.5f) continue;
        if (x < pbox[4*p] || x > pbox[4*p+1] || y < pbox[4*p+2] || y > pbox[4*p+3]) continue;
        int j0, j1; float a;
        if (s <= h0[p]) { j0 = j1 = h0[p]; a = 0.f; }
        else if (s >= h1[p]) { j0 = j1 = h1[p]; a = 0.f; }
        else { j0 = (int)floorf(s); j1 = j0 + 1; a = s - j0; }
        int r0 = first[p] + j0 - h0[p], r1 = first[p] + j1 - h0[p];
        const float* b0 = rbox + 4*r0; const float* b1 = rbox + 4*r1;
        if (x < fminf(b0[0], b1[0]) || x > fmaxf(b0[1], b1[1]) || y < fminf(b0[2], b1[2]) || y > fmaxf(b0[3], b1[3])) continue;
        const float* R0 = rings + (size_t)r0 * nt * 2; const float* R1 = rings + (size_t)r1 * nt * 2;
        float b = 1.f - a;
        float px = b * R0[2*(nt-1)] + a * R1[2*(nt-1)], py = b * R0[2*(nt-1)+1] + a * R1[2*(nt-1)+1];
        int c = 0;
        for (int v = 0; v < nt; ++v) {
            float qx = b * R0[2*v] + a * R1[2*v], qy = b * R0[2*v+1] + a * R1[2*v+1];
            c = cross_(x, y, px, py, qx, qy, c);
            px = qx; py = qy;
        }
        if (c) acc += sign[p];
    }
    return acc > 0;
}

__device__ bool point_bins(float x, float y, float z, const float* rings, const int* sign, const int* h0, const int* h1,
                           const int* slot_first, const int* sl_r, const float* sl_f, const int* bin_off, const int* edges,
                           int P, float z0, float w, int nt, int NB) {
    float s = (z - z0) / w;
    int acc = 0;
    for (int p = 0; p < P; ++p) {
        if (s < h0[p] - 0.5f || s > h1[p] + 0.5f) continue;
        int l; float a;
        if (s <= h0[p]) { l = 0; a = 0.f; }
        else if (s >= h1[p]) { l = h1[p] - h0[p] + 1; a = 0.f; }
        else { int j0 = (int)floorf(s); l = j0 - h0[p] + 1; a = s - j0; }
        int k = slot_first[p] + l;
        const float* f = sl_f + 6*k;
        if (x < f[2] || x > f[3] || y < f[4] || y > f[5]) continue;
        int q = min((int)((y - f[0]) * f[1]), NB - 1);
        const float2* R0 = (const float2*)rings + (size_t)sl_r[2*k] * nt;
        const float2* R1 = (const float2*)rings + (size_t)sl_r[2*k+1] * nt;
        float b = 1.f - a;
        int c = 0;
        for (int e = bin_off[k*NB + q], e1 = bin_off[k*NB + q + 1]; e < e1; ++e) {
            int v = edges[e], u = v > 0 ? v - 1 : nt - 1;
            float2 A0 = R0[u], A1 = R1[u], B0 = R0[v], B1 = R1[v];
            c = cross_(x, y, b*A0.x + a*A1.x, b*A0.y + a*A1.y, b*B0.x + a*B1.x, b*B0.y + a*B1.y, c);
        }
        if (c) acc += sign[p];
    }
    return acc > 0;
}

extern "C" __global__ void inside_bins(const float* P3, int n, const float* rings, const int* sign, const int* h0, const int* h1,
                                       const int* slot_first, const int* sl_r, const float* sl_f, const int* bin_off,
                                       const int* edges, int P, float z0, float w, int nt, int NB, bool* out) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    out[i] = point_bins(P3[3*i], P3[3*i+1], P3[3*i+2], rings, sign, h0, h1, slot_first, sl_r, sl_f, bin_off, edges, P, z0, w, nt, NB);
}

extern "C" __global__ void inside_table(const float* P3, int n, const float* rings, const float* rbox, const float* pbox,
                                        const int* sign, const int* h0, const int* h1, const int* first, int P,
                                        float z0, float w, int layers, int nt, bool* out) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    out[i] = point_table(P3[3*i], P3[3*i+1], P3[3*i+2], rings, rbox, pbox, sign, h0, h1, first, P, z0, w, nt);
}

extern "C" __global__ void inside_poly(const float* P3, int n, const float* C, const float* Fm, const float* pbox,
                                       const int* sign, const float* s0, const float* s1, const int* h0, const int* h1,
                                       int P, int M, int K, int nt, bool* out) {
    extern __shared__ float sF[];  // the Fourier table (K, nt)
    for (int q = threadIdx.x; q < K * nt; q += blockDim.x) sF[q] = Fm[q];
    __syncthreads();
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    float x = P3[3*i], y = P3[3*i+1], z = P3[3*i+2];
    float T[64]; float c[2*64+2];
    int acc = 0, M1 = M + 1, K2 = 2 * K;
    for (int p = 0; p < P; ++p) {
        if (z < s0[p] || z > s1[p]) continue;
        if (x < pbox[4*p] || x > pbox[4*p+1] || y < pbox[4*p+2] || y > pbox[4*p+3]) continue;
        float e = 1.f - 1.f / (float)(h1[p] - h0[p] + 1);
        float u = 2.f * (z - s0[p]) / (s1[p] - s0[p]) - 1.f;
        u = fminf(fmaxf(u, -e), e);
        T[0] = 1.f; if (M1 > 1) T[1] = u;
        for (int j = 2; j < M1; ++j) T[j] = 2.f * u * T[j-1] - T[j-2];
        const float* Cp = C + (size_t)p * M1 * K2;
        for (int k = 0; k < K2; ++k) { float a = 0.f; for (int j = 0; j < M1; ++j) a += T[j] * Cp[j*K2 + k]; c[k] = a; }
        float px = 0.f, py = 0.f;
        for (int k = 0; k < K; ++k) { px += c[k] * sF[k*nt + nt-1]; py += c[K+k] * sF[k*nt + nt-1]; }
        int cr = 0;
        for (int v = 0; v < nt; ++v) {
            float qx = 0.f, qy = 0.f;
            for (int k = 0; k < K; ++k) { float f = sF[k*nt + v]; qx += c[k] * f; qy += c[K+k] * f; }
            cr = cross_(x, y, px, py, qx, qy, cr);
            px = qx; py = qy;
        }
        if (cr) acc += sign[p];
    }
    out[i] = acc > 0;
}

// one block row per pose: threads stride over the points of A (into B), then of B (into A); a hit sets the flag
#define TBL(X) const float* rings##X, const int* sign##X, const int* h0##X, const int* h1##X, const int* sf##X, const int* sr##X, \
               const float* sfl##X, const int* bo##X, const int* ed##X, int Pc##X, float z0##X, float w##X, int nt##X, int NB##X
#define ARG(X) rings##X, sign##X, h0##X, h1##X, sf##X, sr##X, sfl##X, bo##X, ed##X, Pc##X, z0##X, w##X, nt##X, NB##X
extern "C" __global__ void collide(const float* PA, int nA, const float* PB, int nB, int npose,
                                   const float* MAB, const float* tAB, const float* MBA, const float* tBA,
                                   TBL(A), TBL(B), int* hit) {
    int k = blockIdx.y;
    if (k >= npose) return;
    int stride = gridDim.x * blockDim.x;
    const float* m = MAB + 9*k; const float* tt = tAB + 3*k;
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < nA; i += stride) {
        if (*(volatile int*)(hit + k)) return;
        float a0 = PA[3*i], a1 = PA[3*i+1], a2 = PA[3*i+2];
        float x = m[0]*a0 + m[1]*a1 + m[2]*a2 + tt[0];
        float y = m[3]*a0 + m[4]*a1 + m[5]*a2 + tt[1];
        float z = m[6]*a0 + m[7]*a1 + m[8]*a2 + tt[2];
        if (fabsf(x) > 1.f || fabsf(y) > 1.f || fabsf(z) > 1.f) continue;
        if (point_bins(x, y, z, ARG(B))) { hit[k] = 1; return; }
    }
    m = MBA + 9*k; tt = tBA + 3*k;
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < nB; i += stride) {
        if (*(volatile int*)(hit + k)) return;
        float a0 = PB[3*i], a1 = PB[3*i+1], a2 = PB[3*i+2];
        float x = m[0]*a0 + m[1]*a1 + m[2]*a2 + tt[0];
        float y = m[3]*a0 + m[4]*a1 + m[5]*a2 + tt[1];
        float z = m[6]*a0 + m[7]*a1 + m[8]*a2 + tt[2];
        if (fabsf(x) > 1.f || fabsf(y) > 1.f || fabsf(z) > 1.f) continue;
        if (point_bins(x, y, z, ARG(A))) { hit[k] = 1; return; }
    }
}
"""
_MOD = None


def _kern(name):
    global _MOD
    import cupy as cp

    if _MOD is None:
        _MOD = cp.RawModule(code=_SRC, options=("--use_fast_math",))
    return _MOD.get_function(name)


def _targs(q):
    g = q.gpu()
    return (g["rings"], g["sign"], g["h0"], g["h1"], g["slot_first"], g["sl_r"], g["sl_f"], g["bin_off"], g["edges"], np.int32(q.P),
            q.zc[0], q.w, np.int32(q.nt), np.int32(q.NB))


def collide_gpu(qA: FCPQuery, qB: FCPQuery, PA, PB, MAB, tAB, MBA, tBA, blocks_x: int = 8, tb: int = 128):
    """PA, PB, MAB, ...: CuPy float32 arrays (canonical points of each; the maps of `relative_maps`). -> CuPy int32
    (npose,) hit flags."""
    import cupy as cp

    k = len(MAB)
    hit = cp.zeros(k, cp.int32)
    _kern("collide")((blocks_x, k), (tb,), (PA, np.int32(len(PA)), PB, np.int32(len(PB)), np.int32(k), MAB, tAB, MBA, tBA)
                     + _targs(qA) + _targs(qB) + (hit,))
    return hit