import type { Fn , ICopy , IEmpty , Pair } from "@thi.ng/api" ; import { ensureArray } from "@thi.ng/arrays/ensure-array" ; import type { IRegionQuery , ISpatialMap } from "@thi.ng/geom-api" ; import { Heap } from "@thi.ng/heaps/heap" ; import { EPS } from "@thi.ng/math/api" ; import { map } from "@thi.ng/transducers/map" ; import type { ReadonlyVec , Vec } from "@thi.ng/vectors" ; import { distSq } from "@thi.ng/vectors/distsq" ; import { addResults , CMP , into } from "./utils.js" ; type MaybeKdNode < K extends ReadonlyVec , V > = KdNode < K , V > | undefined ; export class KdNode < K extends ReadonlyVec , V > { d : number ; parent ? : KdNode < K , V > ; l ? : KdNode < K , V > ; r ? : KdNode < K , V > ; k : K ; v : V ; constructor ( parent , dim , key , val ) { this . parent = parent ; this . d = dim ; this . k = key ; this . v = val ; } get height ( ) : number { return ( 1 + Math . max ( this . l ? this . l . height : 0 , this . r ? this . r . height : 0 ) ) ; } } /**
 *  https *
 * Partially based on:
 *  https export class KdTreeMap < K extends ReadonlyVec , V > implements ICopy < KdTreeMap < K , V > > , IEmpty < KdTreeMap < K , V > > , IRegionQuery < K , V , number > , ISpatialMap < K , V > { readonly dim : number ; protected root : MaybeKdNode < K , V > ; protected _size : number ; constructor ( dim , pairs ? ) { this . dim = dim ; this . _size = 0 ; this . root = pairs ? this . buildTree ( ensureArray ( pairs ) , 0 ) : undefined ; } * [ Symbol . iterator ] ( ) { let queue = this . root ? [ this . root ] : [ ] ; while ( queue . length ) { const n = queue . pop ( ) ; if ( n ) { yield < Pair < K , V > > [ n . k , n . v ] ; queue . push ( n . r , n . l ) ; } } } * keys ( ) { let queue = this . root ? [ this . root ] : [ ] ; while ( queue . length ) { const n = queue . pop ( ) ; if ( n ) { yield n . k ; queue . push ( n . r , n . l ) ; } } } values ( ) { return map ( ( p ) => p [ 1 ] , this ) ; } get size ( ) { return this . _size ; } get height ( ) { return this . root ? this . root . height : 0 ; } get ratio ( ) { return this . _size ? this . height / Math . log2 ( this . _size ) : 0 ; } copy ( ) { return new KdTreeMap ( this . dim , this ) ; } clear ( ) { delete this . root ; this . _size = 0 ; } empty ( ) { return new KdTreeMap < K , V > ( this . dim ) ; } set ( key , val , eps = EPS ) { eps = Math . max ( 0 , eps ) ; eps *= eps ; const search = ( node , parent ) : MaybeKdNode < K , V > => node ? search ( key [ node . d ] < node . k [ node . d ] ? node . l : node . r , node ) : parent ; let parent ; if ( this . root ) { parent = nearest1 ( key , [ eps , undefined ] , this . dim , this . root ) [ 1 ] ; if ( parent ) { parent . v = val ; return false ; } parent = search ( this . root , undefined ) ! ; const dim = parent . d ; parent [ key [ dim ] < parent . k [ dim ] ? "l" : "r" ] = new KdNode < K , V > ( parent , ( dim + 1 ) % this . dim , key , val ) ; } else { this . root = new KdNode < K , V > ( undefined , 0 , key , val ) ; } this . _size ++ ; return true ; } into ( pairs , eps = EPS ) { return into ( this , pairs , eps ) ; } remove ( key ) { const node = find ( key , this . root , 0 ) ; if ( node ) { remove ( node ) && ( this . root = undefined ) ; this . _size -- ; return true ; } return false ; } has ( key , eps = EPS ) { return ( ! ! this . root && ! ! nearest1 ( key , [ eps * eps , undefined ] , this . dim , this . root ) [ 1 ] ) ; } get ( key , eps = EPS ) { if ( this . root ) { const node = nearest1 ( key , [ eps * eps , undefined ] , this . dim , this . root ) [ 1 ] ; return node ? node . v : undefined ; } } query ( q , maxDist , limit ? , acc ? ) { return this . doSelect ( q , ( x ) => [ x . k , x . v ] , maxDist , limit , acc ) ; } queryKeys ( q , maxDist , limit ? , acc ? ) { return this . doSelect ( q , ( x ) => x . k , maxDist , limit , acc ) ; } queryValues ( q , maxDist , limit ? , acc ? ) { return this . doSelect ( q , ( x ) => x . v , maxDist , limit , acc ) ; } protected doSelect < T > ( q , f , maxDist , maxNum = 1 , acc = [ ] ) { if ( ! this . root ) return [ ] ; maxDist *= maxDist ; if ( maxNum === 1 ) { const sel = nearest1 ( q , [ maxDist , undefined ] , this . dim , this . root ) [ 1 ] ; sel && acc . push ( f ( sel ) ) ; } else { const nodes = new Heap < [ number , MaybeKdNode < K , V > ] > ( [ [ maxDist , undefined ] ] , { compare : CMP , } ) ; nearest ( q , nodes , this . dim , maxNum , this . root ! ) ; return addResults ( f , nodes . values , acc ) ; } return acc ; } protected buildTree ( points , depth , parent ? ) { const n = points . length ; if ( n === 0 ) { return ; } this . _size ++ ; let dim = depth % this . dim ; if ( n === 1 ) { return new KdNode < K , V > ( parent , dim , ... points [ 0 ] ) ; } points . sort ( ( a , b ) => a [ 0 ] [ dim ] - b [ 0 ] [ dim ] ) ; const med = n >>> 1 ; const node = new KdNode < K , V > ( parent , dim , ... points [ med ] ) ; node . l = this . buildTree ( points . slice ( 0 , med ) , depth + 1 , node ) ; node . r = this . buildTree ( points . slice ( med + 1 ) , depth + 1 , node ) ; return node ; } } param p param node param epsSq const find = < K extends ReadonlyVec , V > ( p , node , epsSq ) : KdNode < K , V > | undefined => { if ( ! node ) return ; return distSq ( p , node . k ) <= epsSq ? node : find ( p , p [ node . d ] < node . k [ node . d ] ? node . l : node . r , epsSq ) ; } ; const findMin = < K extends ReadonlyVec , V > ( node , dim ) : MaybeKdNode < K , V > => { if ( ! node ) return ; if ( node . d === dim ) { return node . l ? findMin ( node . l , dim ) : node ; } const q = node . k [ dim ] ; const l = findMin ( node . l , dim ) ; const r = findMin ( node . r , dim ) ; let min = node ; if ( l && l . k [ dim ] < q ) { min = l ; } if ( r && r . k [ dim ] < min . k [ dim ] ) { min = r ; } return min ; } ; param node const remove = < K extends ReadonlyVec , V > ( node ) => { if ( ! node . l && ! node . r ) { if ( ! node . parent ) { return true ; } const parent = node . parent ; const pdim = parent . d ; parent [ node . k [ pdim ] < parent . k [ pdim ] ? "l" : "r" ] = undefined ; return ; } let next ; let nextP ; if ( node . r ) { next = findMin ( node . r , node . d ) ! ; nextP = next . k ; remove ( next ) ; node . k = nextP ; } else { next = findMin ( node . l , node . d ) ! ; nextP = next . k ; remove ( next ) ; node . r = node . l ; node . l = undefined ; node . k = nextP ; } } ; const nearest = < K extends ReadonlyVec , V > ( q , acc , dims , maxNum , node ) => { const p = node . k ; const ndist = distSq ( p , q ) ; if ( ! node . l && ! node . r ) { collect ( acc , maxNum , node , ndist ) ; return ; } const tdist = nodeDist ( node , dims , q , p ) ; let best = bestChild ( node , q ) ; nearest ( q , acc , dims , maxNum , best ! ) ; collect ( acc , maxNum , node , ndist ) ; if ( tdist < acc . values [ 0 ] [ 0 ] ) { best = best === node . l ? node . r : node . l ; best && nearest ( q , acc , dims , maxNum , best ) ; } } ; /**
 * Optimized version of  nearest for single closest point search.
 *
 *  param q param acc param dims param node const nearest1 = < K extends ReadonlyVec , V > ( q , acc , dims , node ) : [ number , MaybeKdNode < K , V > ] => { const p = node . k ; const ndist = distSq ( p , q ) ; if ( ! node . l && ! node . r ) { collect1 ( acc , node , ndist ) ; return acc ; } const tdist = nodeDist ( node , dims , q , p ) ; let best = bestChild ( node , q ) ; nearest1 ( q , acc , dims , best ! ) ; collect1 ( acc , node , ndist ) ; if ( tdist < acc [ 0 ] ) { best = best === node . l ? node . r : node . l ; best && nearest1 ( q , acc , dims , best ) ; } return acc ; } ; const bestChild = < K extends ReadonlyVec , V > ( node , q ) => { const d = node . d ; return ! node . r ? node . l : ! node . l ? node . r : q [ d ] < node . k [ d ] ? node . l : node . r ; } ; const collect = < K extends ReadonlyVec , V > ( acc , maxNum , node , ndist ) => ( ! acc . length || ndist < acc . peek ( ) ! [ 0 ] ) && ( acc . length >= maxNum ? acc . pushPop ( [ ndist , node ] ) : acc . push ( [ ndist , node ] ) ) ; const collect1 = < K extends ReadonlyVec , V > ( acc , node , ndist ) => ndist < acc [ 0 ] && ( ( acc [ 0 ] = ndist ) , ( acc [ 1 ] = node ) ) ; const TMP = [ ] ; const nodeDist = < K extends ReadonlyVec , V > ( node , dims , q , p ) => { for ( let i = dims , d = node . d ; i -- > 0 ; ) { TMP [ i ] = i === d ? q [ i ] : p [ i ] ; } return distSq ( TMP , p ) ; } ;