"""Superpoint (over-segmentation) extraction.

Faithful Python/Numba port of ``segment_mesh`` from
https://github.com/kenomo/segmentator (itself the ScanNet ``Segmentator``
Felzenszwalb-Huttenlocher graph segmentation).  AQ3D's ScanNet preprocessing
calls it with ``kThresh=0.01, segMinVerts=20``; the network consumes those
superpoints directly, so reproducing this exactly matters far more than it
looks.

The original is a small C++/PyTorch extension that has to be built with CMake
against libtorch, which is not practical inside a Space image, hence the port.
"""

import numpy as np
from numba import njit


@njit(cache=True, nogil=True)
def _segment_graph(ea, eb, w, order, num_vertices, c, seg_min_verts):
    parent = np.arange(num_vertices, dtype=np.int64)
    rank = np.zeros(num_vertices, dtype=np.int64)
    size = np.ones(num_vertices, dtype=np.int64)
    threshold = np.full(num_vertices, c, dtype=np.float32)

    # --- pass 1: Felzenszwalb-Huttenlocher merging over ascending weights ---
    for i in range(order.shape[0]):
        e = order[i]
        # find(a)
        x = ea[e]
        y = x
        while y != parent[y]:
            y = parent[y]
        parent[x] = y
        a = y
        # find(b)
        x = eb[e]
        y = x
        while y != parent[y]:
            y = parent[y]
        parent[x] = y
        b = y

        if a != b:
            ww = w[e]
            if ww <= threshold[a] and ww <= threshold[b]:
                # join(a, b)
                if rank[a] > rank[b]:
                    parent[b] = a
                    size[a] += size[b]
                    root = a
                else:
                    parent[a] = b
                    size[b] += size[a]
                    if rank[a] == rank[b]:
                        rank[b] += 1
                    root = b
                threshold[root] = ww + c / size[root]

    # --- pass 2: absorb segments smaller than seg_min_verts ---
    for i in range(order.shape[0]):
        e = order[i]
        x = ea[e]
        y = x
        while y != parent[y]:
            y = parent[y]
        parent[x] = y
        a = y
        x = eb[e]
        y = x
        while y != parent[y]:
            y = parent[y]
        parent[x] = y
        b = y

        if a != b and (size[a] < seg_min_verts or size[b] < seg_min_verts):
            if rank[a] > rank[b]:
                parent[b] = a
                size[a] += size[b]
            else:
                parent[a] = b
                size[b] += size[a]
                if rank[a] == rank[b]:
                    rank[b] += 1

    out = np.empty(num_vertices, dtype=np.int64)
    for q in range(num_vertices):
        y = q
        while y != parent[y]:
            y = parent[y]
        parent[q] = y
        out[q] = y
    return out


def unit_face_normals(vertices: np.ndarray, faces: np.ndarray) -> np.ndarray:
    v = vertices.astype(np.float64)
    n = np.cross(v[faces[:, 1]] - v[faces[:, 0]], v[faces[:, 2]] - v[faces[:, 0]])
    length = np.linalg.norm(n, axis=1, keepdims=True)
    return np.divide(n, length, out=np.zeros_like(n), where=length > 1e-20)


def segmentator_vertex_normals(vertices: np.ndarray, faces: np.ndarray) -> np.ndarray:
    """Mean of the incident *unit* face normals (the C++ code's running lerp)."""
    nf = unit_face_normals(vertices, faces)
    idx = faces.reshape(-1)
    vals = np.repeat(nf, 3, axis=0)
    nv = np.zeros((vertices.shape[0], 3), dtype=np.float64)
    for a in range(3):
        nv[:, a] = np.bincount(idx, weights=vals[:, a], minlength=vertices.shape[0])
    counts = np.bincount(idx, minlength=vertices.shape[0])
    nv /= np.maximum(counts, 1)[:, None]
    return nv


def segment_mesh(vertices: np.ndarray, faces: np.ndarray, k_thresh: float = 0.01,
                 seg_min_verts: int = 20) -> np.ndarray:
    """Return a superpoint id per vertex, relabelled to ``0..S-1``."""
    vertices = np.ascontiguousarray(vertices, dtype=np.float32)
    faces = np.ascontiguousarray(faces, dtype=np.int64)
    num_vertices = vertices.shape[0]

    normals = segmentator_vertex_normals(vertices, faces)

    # edges (i1,i2), (i1,i3), (i3,i2) per face -- exactly as in the C++ source
    ea = np.empty(faces.shape[0] * 3, dtype=np.int64)
    eb = np.empty(faces.shape[0] * 3, dtype=np.int64)
    ea[0::3], eb[0::3] = faces[:, 0], faces[:, 1]
    ea[1::3], eb[1::3] = faces[:, 0], faces[:, 2]
    ea[2::3], eb[2::3] = faces[:, 2], faces[:, 1]

    p = vertices.astype(np.float64)
    d = p[eb] - p[ea]
    dd = np.linalg.norm(d, axis=1, keepdims=True)
    d = np.divide(d, dd, out=np.zeros_like(d), where=dd > 1e-20)
    n1, n2 = normals[ea], normals[eb]
    dot = (n1 * n2).sum(1)
    dot2 = (n2 * d).sum(1)
    w = 1.0 - dot
    w = np.where(dot2 > 0, w * w, w).astype(np.float32)

    order = np.argsort(w, kind="stable").astype(np.int64)
    comps = _segment_graph(ea, eb, w, order, num_vertices,
                           np.float32(k_thresh), np.int64(seg_min_verts))

    # remap component ids to 0..S-1 (upstream does the same)
    _, remapped = np.unique(comps, return_inverse=True)
    return remapped.astype(np.int64).reshape(-1)


def warmup() -> None:
    """Trigger Numba JIT compilation once, at import time."""
    v = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0]], dtype=np.float32)
    f = np.array([[0, 1, 2], [1, 3, 2]], dtype=np.int64)
    segment_mesh(v, f)