// Geo-spatial index over (latitude, longitude) node properties. Points
// embed on the unit sphere (lat/lon -> 3-D Cartesian), so Euclidean chord
// distance is monotonic with great-circle distance and a 3-D k-d tree
// answers radius and k-NN queries exactly -- correct at the poles and
// across the antimeridian. Distances report in kilometres.

package chickpeas

import (
	
	
	

	
	
)

// earthRadiusKM is the mean Earth radius (IUGG).
const earthRadiusKM = 6371.0088

type kdNode struct {
	point       int
	left, right *kdNode
}

// GeoIndex is the immutable index for one (label, latKey, lonKey) field.
type GeoIndex struct {
	nodes []uint32
	lats  []float64
	lons  []float64
	xyz   [][3]float64
	root  *kdNode
}

// BuildGeoIndex builds from (node, latDeg, lonDeg) points; non-finite or
// out-of-range coordinates are skipped.
func ( func( func( uint32, ,  float64) bool)) *GeoIndex {
	 := &GeoIndex{}
	(func( uint32, ,  float64) bool {
		if math.IsNaN() || math.IsInf(, 0) || math.IsNaN() || math.IsInf(, 0) ||
			 < -90 ||  > 90 ||  < -180 ||  > 180 {
			return true
		}
		.nodes = append(.nodes, )
		.lats = append(.lats, )
		.lons = append(.lons, )
		.xyz = append(.xyz, unitVec(, ))
		return true
	})
	 := make([]int, len(.nodes))
	for  := range  {
		[] = 
	}
	.root = buildKD(.xyz, , 0)
	return 
}

// Len is the number of indexed points.
func ( *GeoIndex) () int { return len(.nodes) }

// WithinRadius returns the nodes within km great-circle distance of
// (lat, lon).
func ( *GeoIndex) (, ,  float64) *nodeset.Set {
	if  < 0 || .root == nil {
		return nodeset.New()
	}
	 := unitVec(, )
	 := chordThreshold()
	 := roaring.New()
	.rangeSearch(.root, &, *, 0, )
	return nodeset.FromBitmap()
}

// GeoHit is one KNN result.
type GeoHit struct {
	Node uint32
	KM   float64
}

// KNN returns up to k nearest nodes to (lat, lon), sorted by increasing
// distance, ties by ascending node id.
func ( *GeoIndex) (,  float64,  int) []GeoHit {
	if  <= 0 || .root == nil {
		return nil
	}
	 := unitVec(, )
	 := &neighborHeap{}
	.nn(.root, &, , 0, )
	 := make([]GeoHit, 0, .Len())
	for ,  := range * {
		 = append(, GeoHit{Node: .node, KM: chordToKM(math.Sqrt(.distSq))})
	}
	sort.Slice(, func(,  int) bool {
		if [].KM != [].KM {
			return [].KM < [].KM
		}
		return [].Node < [].Node
	})
	return 
}

// WithinBBox returns the nodes inside the lat/lon rectangle; minLon >
// maxLon treats the box as crossing the antimeridian.
func ( *GeoIndex) (, , ,  float64) *nodeset.Set {
	 := func( float64) bool {
		if  <=  {
			return  >=  &&  <= 
		}
		return  >=  ||  <= 
	}
	 := roaring.New()
	for ,  := range .nodes {
		if .lats[] >=  && .lats[] <=  && (.lons[]) {
			.Add()
		}
	}
	return nodeset.FromBitmap()
}

func ( *GeoIndex) ( *kdNode,  *[3]float64,  float64,  int,  *roaring.Bitmap) {
	if  == nil {
		return
	}
	 := .point
	if distSq(&.xyz[], ) <=  {
		.Add(.nodes[])
	}
	 := [] - .xyz[][]
	 := ( + 1) % 3
	,  := .left, .right
	if  > 0 {
		,  = .right, .left
	}
	.(, , , , )
	if * <=  {
		.(, , , , )
	}
}

func ( *GeoIndex) ( *kdNode,  *[3]float64, ,  int,  *neighborHeap) {
	if  == nil {
		return
	}
	 := .point
	pushBounded(, , neighbor{distSq: distSq(&.xyz[], ), node: .nodes[]})
	 := [] - .xyz[][]
	 := ( + 1) % 3
	,  := .left, .right
	if  > 0 {
		,  = .right, .left
	}
	.(, , , , )
	 := math.Inf(1)
	if .Len() >=  {
		 = (*)[0].distSq
	}
	if * <=  {
		.(, , , , )
	}
}

// neighborHeap is a max-heap by (distance, node id) so its root is the
// current worst of the k best.
type neighbor struct {
	distSq float64
	node   uint32
}

type neighborHeap []neighbor

func ( neighborHeap) () int { return len() }
func ( neighborHeap) (,  int) bool {
	if [].distSq != [].distSq {
		return [].distSq > [].distSq
	}
	return [].node > [].node
}
func ( neighborHeap) (,  int) { [], [] = [], [] }
func ( *neighborHeap) ( any)   { * = append(*, .(neighbor)) }
func ( *neighborHeap) () any {
	 := *
	 := len()
	 := [-1]
	* = [:-1]
	return 
}

func pushBounded( *neighborHeap,  int,  neighbor) {
	heap.Push(, )
	if .Len() >  {
		heap.Pop()
	}
}

// buildKD builds a balanced tree by median split (axis = depth % 3),
// partitioning around the median in O(n) per level.
func buildKD( [][3]float64,  []int,  int) *kdNode {
	if len() == 0 {
		return nil
	}
	 :=  % 3
	 := len() / 2
	selectNth(, , , )
	return &kdNode{
		point: [],
		left:  (, [:], +1),
		right: (, [+1:], +1),
	}
}

// selectNth partitions indices so indices[nth] holds the nth element by
// the axis coordinate (quickselect).
func selectNth( [][3]float64,  []int, ,  int) {
	,  := 0, len()-1
	for  <  {
		 := [[(+)/2]][]
		,  := , 
		for  <=  {
			for [[]][] <  {
				++
			}
			for [[]][] >  {
				--
			}
			if  <=  {
				[], [] = [], []
				++
				--
			}
		}
		if  <=  {
			 = 
		} else if  >=  {
			 = 
		} else {
			return
		}
	}
}

func unitVec(,  float64) [3]float64 {
	,  := *math.Pi/180, *math.Pi/180
	 := math.Cos()
	return [3]float64{ * math.Cos(),  * math.Sin(), math.Sin()}
}

func distSq(,  *[3]float64) float64 {
	, ,  := [0]-[0], [1]-[1], [2]-[2]
	return * + * + *
}

// chordThreshold is the unit-sphere chord length for a great-circle
// distance of km.
func chordThreshold( float64) float64 {
	 := math.Min(/earthRadiusKM, math.Pi)
	return 2 * math.Sin(/2)
}

// chordToKM inverts chordThreshold.
func chordToKM( float64) float64 {
	return 2 * math.Asin(math.Max(-1, math.Min(1, /2))) * earthRadiusKM
}

// HaversineKM is the great-circle distance in kilometres between two
// lat/lon points -- a public utility independent of the index.
func (, , ,  float64) float64 {
	,  := *math.Pi/180, *math.Pi/180
	 := ( - ) * math.Pi / 180
	 := ( - ) * math.Pi / 180
	 := math.Sin( / 2)
	 := math.Sin( / 2)
	 := * + math.Cos()*math.Cos()**
	return 2 * math.Asin(math.Max(-1, math.Min(1, math.Sqrt()))) * earthRadiusKM
}

// ln32 is float32 natural log via float64.
func ln32( float32) float32 {
	return float32(math.Log(float64()))
}