"""数组化 + Numba 的「单层」八叉树编码,等价替换
  MatrixOctree(max_depth=target_depth+1).build + OctreeFaceTokenBuilder(mc=8).encode_level_tokens(target_depth)

前提(per-row 栈已满足):max_depth = target_depth+1 -> 只有一次「叶子(td+1)->节点(td)」聚合,
不需要多层 neighbor 传播;且 max_children=8 时 reduce_node_children_to_max 是 no-op(活跃 octant ≤ 8)。

去掉 OctreeNode 对象 / dict-of-tuple / set-of-tuple:
  * 叶子坐标量化 -> 边表(numpy) -> 节点坐标(leaf>>1)+ octant 局部索引(位运算, 向量化)
  * 按节点聚合 intra 8×8 / 活跃掩码 F / 节点-节点邻接(向量化 scatter / np.unique)
  * 三角形(3-clique)枚举走 Numba njit(CSR 两指针求交), 这是原 get_level_info 的热点
  * inter 8×8 按 (节点,邻居) 分组填充
输出 (raw_t, out_t) 与原实现「逐 slot 解码后网格一致」(顺序可不同;decode 与顺序无关)。
"""
import numpy as np
from numba import njit

from utils.data_utils import discretize

_OCT_W = np.array([1, 2, 4], dtype=np.int64)   # octant 局部索引 = x + 2y + 4z


@njit(cache=False)
def _triangles_csr(indptr, indices):
    """前向邻接 CSR(indices[indptr[u]:indptr[u+1]] = 与 u 相连且 >u 的节点, 升序)。
    枚举 u<v<w 三角形, 返回 [[u,v,w],...]。两指针求 adj[u]∩adj[v]。"""
    out_u = []
    out_v = []
    out_w = []
    n = indptr.shape[0] - 1
    for u in range(n):
        su, eu = indptr[u], indptr[u + 1]
        for pa in range(su, eu):
            v = indices[pa]
            sv, ev = indptr[v], indptr[v + 1]
            # 交集: adj[u] 中 >v 的部分(即 pa+1..eu, 因为升序且后面都 >v) 与 adj[v](sv..ev)
            i = pa + 1
            j = sv
            while i < eu and j < ev:
                wi = indices[i]
                wj = indices[j]
                if wi == wj:
                    out_u.append(u); out_v.append(v); out_w.append(wi)
                    i += 1; j += 1
                elif wi < wj:
                    i += 1
                else:
                    j += 1
    return (np.array(out_u, dtype=np.int64),
            np.array(out_v, dtype=np.int64),
            np.array(out_w, dtype=np.int64))