//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Name:        kd_tree.sch																					//
// Description: Implementation of multidimensional kd-trees with option to search for closest points		//
//				in given radius. Example use in comments below.												//
// Written by:  Tymon																						//
// Date:  		2/2/2016																					//
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

USING "net_include.sch"


/* ***********************************************************************************************************
 SAMPLE USE:
**************************************************************************************************************

_T_KD_TREE_GET_OBJECT_SORTABLE_VALUE and _T_KD_TREE_GET_SEARCH_OBJECT_VALUE are used to retrieve values of
elements we store in KD tree. By passing axis argument they serve as a "generic" implementation 
of n dimensional vectors.

KD_TREE is a stuct holding all the nodes in memory. It has a pointer to function to retrieve value of each 
element in each dimension.

To create the KD_TREE we need to implement _T_KD_TREE_GET_OBJECT_SORTABLE_VALUE and create a list of numeric 
indices that serves as an input for tree creation:

INT kdtreeIndices[6] // This is our list of indices

// Implementation of _T_KD_TREE_GET_OBJECT_SORTABLE_VALUE
FUNC FLOAT GET_NUM_LIST_VALUE(INT iElementID, INT iAxis = 0)
	VECTOR list[6]
	list[0] = <<2, 3, 0>>
	list[1] = <<5, 4, 0>>
	list[2] = <<9, 6, 0>>
	list[3] = <<4, 7, 0>>
	list[4] = <<8, 1, 0>>
	list[5] = <<7, 2, 0>>
	
	IF iElementID >= 0 AND iElementID < 6 AND iAxis >= 0 and iAxis < 2
		IF iAxis = 0
			RETURN list[iElementID].x
		ELSE
			RETURN list[iElementID].y
		ENDIF
	ENDIF
	
	RETURN 0.0
ENDFUNC

// Implementation of _T_KD_TREE_GET_SEARCH_OBJECT_VALUE - used to describe point we're quering
FUNC FLOAT GET_SEARCH_VALUE(INT iAxis = 0)
	IF iAxis = 0
		RETURN 3.0
	ENDIF
	
	RETURN 6.0
ENDFUNC

PROC DO_KDTREE_TEST()
	kdtreeIndices[0] = 0
	kdtreeIndices[1] = 1
	kdtreeIndices[2] = 2
	kdtreeIndices[3] = 3
	kdtreeIndices[4] = 4
	kdtreeIndices[5] = 5
	
	KD_TREE kdtree = CONSTRUCT_KD_TREE(&GET_NUM_LIST_VALUE, 2, kdtreeIndices, 6)
	
	#IF IS_DEBUG_BUILD
	DEBUG_PRINT_KD_TREE(kdtree)
	#ENDIF
	
	INT iClosestPoints[10]
	INT iClosestPointsCount
	
	FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS(kdtree, kdtree.nodePool[0], &GET_SEARCH_VALUE, 3.0, iClosestPoints, iClosestPointsCount)
	
	#IF IS_DEBUG_BUILD
		INT i
		PRINTLN("[KDTREE] DO_KDTREE_TEST - got search results, there's ", iClosestPointsCount, " of them: ")
		NET_PRINT("[KDTREE] DO_KDTREE_TEST - ")
		REPEAT iClosestPointsCount i
			NET_PRINT_INT(iClosestPoints[i]) NET_PRINT(" - ")
			DEBUG_PRINT_KD_TREE_SINGLE_OBJECT(kdtree, iClosestPoints[i])
			IF i < iClosestPointsCount - 1
				NET_PRINT(", ")
			ENDIF
		ENDREPEAT
	#ENDIF
ENDPROC

*********************************************************************************************************** */

#IF IS_DEBUG_BUILD
PROC DEBUG_PRINT_KD_TREE_INDICES(INT &indices[], INT iFrom, INT iTo)
	NET_PRINT("[KDTREE] DEBUG_PRINT_KD_TREE_INDICES - Printing from ")
	NET_PRINT_INT(iFrom)
	NET_PRINT(" to ")
	NET_PRINT_INT(iTo)
	NET_PRINT(": ")
	INT i
	FOR i = iFrom TO iTo
		NET_PRINT_INT(indices[i])
		IF i < iTo
			NET_PRINT(", ")
		ENDIF
	ENDFOR
	
	NET_NL()
ENDPROC

PROC DEBUG_PRINT_KD_TREE_SINGLE_OBJECT(KD_TREE &kdtree, INT iElementID)
	INT i
	NET_PRINT("(")
	REPEAT kdtree.iDimension i
		NET_PRINT_FLOAT(CALL kdtree.getValue(iElementID, i))
		IF i < kdtree.iDimension - 1
			NET_PRINT(", ")
		ENDIF
	ENDREPEAT
	NET_PRINT(")")
ENDPROC
#ENDIF

PROC _SWAP_IN_KD_TREE(INT &in[], INT a, INT b)
    INT iTmp = in[a]
    in[a] = in[b]
    in[b] = iTmp
ENDPROC

PROC _INSERT_SORT_KD_TREE(KD_TREE &kdtree, INT &indices[], INT iFrom, INT iTo, INT iAxis = 0)
	INT i, j
	FOR i = iFrom + 1 TO iTo
		j = i
		WHILE j > iFrom AND (CALL kdtree.getValue(indices[j-1], iAxis)) > (CALL kdtree.getValue(indices[j], iAxis))
			_SWAP_IN_KD_TREE(indices, j, j-1)
			j = j-1
		ENDWHILE
	ENDFOR
	
	#IF IS_DEBUG_BUILD
		PRINTLN("[KDTREE] _INSERT_SORT_KD_TREE - Sorted from ", iFrom, " to ", iTo, ". Sorted on axis ", iAxis, ". Indices after sort: ")
		DEBUG_PRINT_KD_TREE_INDICES(indices, iFrom, iTo)
		NET_PRINT("[KDTREE] _INSERT_SORT_KD_TREE - Sorted objects: ")
		FOR i = iFrom TO iTo
			DEBUG_PRINT_KD_TREE_SINGLE_OBJECT(kdtree, indices[i])
			IF i < iTo
				NET_PRINT(", ")
			ENDIF
		ENDFOR
		NET_NL()
	#ENDIF
ENDPROC

PROC _KD_TREE_QUICK_SORT(KD_TREE &kdtree, INT &indices[], INT iFrom, INT iTo, INT iAxis = 0)
	IF iFrom < iTo
		
		INT l = iFrom + 1
		INT r = iTo
		INT p = indices[iFrom]
		FLOAT pValue = (CALL kdtree.getValue(p, iAxis))
		
		WHILE l < r
			IF (CALL kdtree.getValue(indices[l], iAxis)) <= pValue
				l = l + 1
			ELIF (CALL kdtree.getValue(indices[r], iAxis)) >= pValue
				r = r - 1
			ELSE
				_SWAP_IN_KD_TREE(indices, l, r)
			ENDIF
		ENDWHILE
		
		IF (CALL kdtree.getValue(indices[l], iAxis)) < pValue
			_SWAP_IN_KD_TREE(indices, l, iFrom)
			l = l - 1
		ELSE
			l = l - 1
			_SWAP_IN_KD_TREE(indices, l, iFrom)
		ENDIF
		
		_KD_TREE_QUICK_SORT(kdtree, indices, iFrom, l, iAxis)
		_KD_TREE_QUICK_SORT(kdtree, indices, r, iTo, iAxis)
		
		#IF IS_DEBUG_BUILD
			NET_PRINT("[KDTREE] _KD_TREE_QUICK_SORT - Sorted from ") NET_PRINT_INT(iFrom) NET_PRINT(" to ") NET_PRINT_INT(iTo) NET_PRINT(". Sorted on axis ") NET_PRINT_INT(iAxis) NET_PRINT( ". Indices after sort: ") NET_NL()
			DEBUG_PRINT_KD_TREE_INDICES(indices, iFrom, iTo)
			NET_PRINT("[KDTREE] _KD_TREE_QUICK_SORT - Sorted objects: ")
			INT i
			FOR i = iFrom TO iTo
				DEBUG_PRINT_KD_TREE_SINGLE_OBJECT(kdtree, indices[i])
				IF i < iTo
					NET_PRINT(", ")
				ENDIF
			ENDFOR
			NET_NL()
		#ENDIF
	ENDIF
ENDPROC


FUNC INT _BUILD_KD_TREE_NODE(KD_TREE &kdtree, INT &indices[], INT iFrom, INT iTo, INT depth = 0)
	
	INT iLength = iTo - iFrom + 1
	
	IF iLength <= 0
		RETURN -1
	ENDIF
	
	PRINTLN("[KDTREE] _BUILD_KD_TREE_NODE - Building from ", iFrom, " to ", iTo, " at depth ", depth, ", length is ", iLength)
	
	#IF IS_DEBUG_BUILD
		DEBUG_PRINT_KD_TREE_INDICES(indices, iFrom, iTo)
	#ENDIF
	
	INT iAxis = depth % kdtree.iDimension
	
	_KD_TREE_QUICK_SORT(kdtree, indices, iFrom, iTo, iAxis)
	//_INSERT_SORT_KD_TREE(kdtree, iNewIndices, iFrom, iTo, iAxis)
	INT iPivot = iFrom + FLOOR(TO_FLOAT(iLength) * 0.5)
	
	PRINTLN("[KDTREE] _BUILD_KD_TREE_NODE - Choosing pivot at ", iPivot, " which is ", indices[iPivot])
	
	INT iNodeID = kdtree.iNodesCount
	kdtree.iNodesCount += 1
	
	kdtree.nodePool[iNodeID].iElementID = indices[iPivot]
	
	PRINTLN("[KDTREE][INSERT] _BUILD_KD_TREE_NODE - Inserted into pool at ", iNodeID, " element id ", indices[iPivot])
	
	kdtree.nodePool[iNodeID].iLeftNodePointer = _BUILD_KD_TREE_NODE(kdtree, indices, iFrom, iPivot - 1, depth + 1)
	kdtree.nodePool[iNodeID].iRightNodePointer = _BUILD_KD_TREE_NODE(kdtree, indices, iPivot + 1, iTo, depth + 1)
	
	RETURN iNodeID
ENDFUNC

PROC CONSTRUCT_KD_TREE(KD_TREE &out, _T_KD_TREE_GET_OBJECT_SORTABLE_VALUE getValue, INT iDimension, INT &indices[], INT iLength)

	PRINTLN("[KDTREE] CONSTRUCT_KD_TREE - OK let's do it, dimensions: ", iDimension, " and number of points: ", iLength)
	#IF IS_DEBUG_BUILD
		DEBUG_PRINT_KD_TREE_INDICES(indices, 0, iLength)
	#ENDIF

	out.iDimension = iDimension
	out.getValue = getValue
	
	out.iNodesCount = 0
	_BUILD_KD_TREE_NODE(out, indices, 0, iLength)
ENDPROC

FUNC FLOAT _GET_SQ_DIST_OF_KD_TREE_ELEMENT_AND_SEARCH_OBJECT(KD_TREE &kdtree, KD_TREE_NODE &currentNode, _T_KD_TREE_GET_SEARCH_OBJECT_VALUE getSearchObjectValue)
	INT i
	FLOAT fDistSq, fTmp
	REPEAT kdtree.iDimension i
		fTmp = (CALL kdtree.getValue(currentNode.iElementID, i)) - (CALL getSearchObjectValue(i))
		fDistSq += (fTmp * fTmp)
	ENDREPEAT
	
	RETURN fDistSq
ENDFUNC

PROC FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS(KD_TREE &kdtree, KD_TREE_NODE &currentNode, _T_KD_TREE_GET_SEARCH_OBJECT_VALUE getSearchObjectValue, FLOAT fRadius, INT &out[], INT &iOutCount, INT depth = 0)
	
	BOOL bGoingLeft
	INT iAxis = depth % kdtree.iDimension
	FLOAT fSearchDistSq = fRadius * fRadius
	
	#IF IS_DEBUG_BUILD
		IF g_SimpleInteriorData.bVerboseKDTree
			PRINTLN("[KDTREE] FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS - We're at node with element ID ", currentNode.iElementID, " at the axis ", iAxis)
		ENDIF
	#ENDIF
	
	FLOAT fSearchValue = CALL getSearchObjectValue(iAxis)
	FLOAT fTreeValue = CALL kdtree.getValue(currentNode.iElementID, iAxis)
	
	#IF IS_DEBUG_BUILD
		IF g_SimpleInteriorData.bVerboseKDTree
			PRINTLN("[KDTREE] FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS - searchValue is ", fSearchValue, " while treeValue is ", fTreeValue)
		ENDIF
	#ENDIF
	
	IF fSearchValue > fTreeValue
		bGoingLeft = FALSE
		
		// Go right
		IF currentNode.iRightNodePointer > -1
			#IF IS_DEBUG_BUILD
				IF g_SimpleInteriorData.bVerboseKDTree
					PRINTLN("[KDTREE] FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS - Going right (", fSearchValue, " > ", fTreeValue, ")")
				ENDIF
			#ENDIF
			FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS(kdtree, kdtree.nodePool[currentNode.iRightNodePointer], getSearchObjectValue, fRadius, out, iOutCount, depth + 1)
		ELSE
			#IF IS_DEBUG_BUILD
				IF g_SimpleInteriorData.bVerboseKDTree
					PRINTLN("[KDTREE] FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS - No right to go.")
				ENDIF
			#ENDIF
		ENDIF
	ELSE
		bGoingLeft = TRUE
		
		// Go left
		IF currentNode.iLeftNodePointer > -1
			#IF IS_DEBUG_BUILD
				IF g_SimpleInteriorData.bVerboseKDTree
					PRINTLN("[KDTREE] FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS - Going left (", fSearchValue, " <= ", fTreeValue, ")")
				ENDIF
			#ENDIF
			FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS(kdtree, kdtree.nodePool[currentNode.iLeftNodePointer], getSearchObjectValue, fRadius, out, iOutCount, depth + 1)
		ELSE
			#IF IS_DEBUG_BUILD
				IF g_SimpleInteriorData.bVerboseKDTree
					PRINTLN("[KDTREE] FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS - No left to go.")
				ENDIF
			#ENDIF
		ENDIF
	ENDIF
	
	// Check if there might be any other qualifying points on the other side of the splitting plane
	IF ABSF( (CALL getSearchObjectValue(iAxis)) - (CALL kdtree.getValue(currentNode.iElementID, iAxis)) ) <= fRadius
		IF bGoingLeft
			IF currentNode.iRightNodePointer > -1
				#IF IS_DEBUG_BUILD
					IF g_SimpleInteriorData.bVerboseKDTree
						PRINTLN("[KDTREE] FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS - We are at ", currentNode.iElementID ,". There might be some more points on the right, going right.")
					ENDIF
				#ENDIF
				FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS(kdtree, kdtree.nodePool[currentNode.iRightNodePointer], getSearchObjectValue, fRadius, out, iOutCount, depth + 1)
			ENDIF
		ELSE
			IF currentNode.iLeftNodePointer > -1
				#IF IS_DEBUG_BUILD
					IF g_SimpleInteriorData.bVerboseKDTree
						PRINTLN("[KDTREE] FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS - We are at ", currentNode.iElementID ,". There might be some more points on the left, going left.")
					ENDIF
				#ENDIF
				FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS(kdtree, kdtree.nodePool[currentNode.iLeftNodePointer], getSearchObjectValue, fRadius, out, iOutCount, depth + 1)
			ENDIF
		ENDIF
	ENDIF
	
	FLOAT fDistSq = _GET_SQ_DIST_OF_KD_TREE_ELEMENT_AND_SEARCH_OBJECT(kdtree, currentNode, getSearchObjectValue)
	
	#IF IS_DEBUG_BUILD
		IF g_SimpleInteriorData.bVerboseKDTree
			PRINTLN("[KDTREE] FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS - fDistSq is ", fDistSq, " and fSearchDistSq is ", fSearchDistSq)
		ENDIF
	#ENDIF
	
	IF fDistSq <= fSearchDistSq
		IF iOutCount < COUNT_OF(out)
			out[iOutCount] = currentNode.iElementID
			iOutCount += 1
			#IF IS_DEBUG_BUILD
				IF g_SimpleInteriorData.bVerboseKDTree
					PRINTLN("[KDTREE] FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS - Added current node (", currentNode.iElementID ,") to results, results count is ", iOutCount)
				ENDIF
			#ENDIF
		ELSE
			#IF IS_DEBUG_BUILD
				IF g_SimpleInteriorData.bVerboseKDTree
					PRINTLN("[KDTREE] FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS - Reached max number of results!")
				ENDIF
			#ENDIF
		ENDIF
	ENDIF
	
ENDPROC

#IF IS_DEBUG_BUILD
/// PURPOSE:
///    Prints one node like: ((x, y), (node), (node))
/// PARAMS:
///    node - 
///    kdtree - 
///    iSpaces - 
PROC _DEBUG_PRINT_KD_NODE(KD_TREE_NODE &node, KD_TREE &kdtree, INT iSpaces = 0)
	INT i
	
	REPEAT (iSpaces * 2) i
		//NET_PRINT(" ")
	ENDREPEAT
	
	NET_PRINT("(")

		DEBUG_PRINT_KD_TREE_SINGLE_OBJECT(kdtree, node.iElementID)
		NET_PRINT(", ")
		
		IF node.iLeftNodePointer > -1
			_DEBUG_PRINT_KD_NODE(kdtree.nodePool[node.iLeftNodePointer], kdtree)
		ELSE
			NET_PRINT("None")
		ENDIF
		
		NET_PRINT(", ")
		
		IF node.iRightNodePointer > -1
			_DEBUG_PRINT_KD_NODE(kdtree.nodePool[node.iRightNodePointer], kdtree)
		ELSE
			NET_PRINT("None")
		ENDIF
		
	NET_PRINT(")")
ENDPROC

PROC DEBUG_PRINT_KD_TREE(KD_TREE &kdtree)
	NET_PRINT("[KDTREE] DEBUG_PRINT_KD_TREE - Node pool size: ") NET_PRINT_INT(kdtree.iNodesCount)
	NET_PRINT("[KDTREE] DEBUG_PRINT_KD_TREE - ")
	IF kdtree.iNodesCount > 0
		_DEBUG_PRINT_KD_NODE(kdtree.nodePool[0], kdtree)
	ELSE
		NET_PRINT("EMPTY")
	ENDIF
	
	NET_NL()
ENDPROC

PROC DEBUG_DRAW_KD_TREE(KD_TREE &kdtree)
	INT i
	VECTOR vMaxCoord = <<-9999.9, -9999.9, 0.0>>
	VECTOR vMinCoord = <<9999.9, 9999.9, 0.0>>
	VECTOR vCurrent
	
	REPEAT kdtree.iNodesCount i
		vCurrent.X = CALL kdtree.getValue(kdtree.nodePool[i].iElementID, 0)
		vCurrent.Y = CALL kdtree.getValue(kdtree.nodePool[i].iElementID, 1)
		
		IF vCurrent.X > vMaxCoord.X
			vMaxCoord.X = vCurrent.X
		ENDIF
		
		IF vCurrent.Y > vMaxCoord.Y
			vMaxCoord.Y = vCurrent.Y
		ENDIF
		
		IF vCurrent.X < vMinCoord.X
			vMinCoord.X = vCurrent.X
		ENDIF
		
		IF vCurrent.Y < vMinCoord.Y
			vMinCoord.Y = vCurrent.Y
		ENDIF
	ENDREPEAT
	
	FLOAT fNormalisedX, fNormalisedY
	FLOAT fXScale = 1.0 / GET_ASPECT_RATIO(FALSE)
	
	REPEAT kdtree.iNodesCount i
		vCurrent.X = CALL kdtree.getValue(kdtree.nodePool[i].iElementID, 0)
		vCurrent.Y = CALL kdtree.getValue(kdtree.nodePool[i].iElementID, 1)
		
		fNormalisedX = (vCurrent.X - vMinCoord.X) / ABSF(vMaxCoord.X - vMinCoord.X)
		fNormalisedY = (1.0 - (vCurrent.Y - vMinCoord.Y) / ABSF(vMaxCoord.Y - vMinCoord.Y))
		
		//PRINTLN("[KDTREE] DEBUG_DRAW_KD_TREE - fNormalisedX: ", fNormalisedX, ", fNormalisedY: ", fNormalisedY, ", vMinCoord: ", vMinCoord, ", vMaxCoord: ", vMaxCoord, ", vCurrent: ", vCurrent)
		
		DRAW_RECT( (fNormalisedX * 0.8 + 0.1) * fXScale, (fNormalisedY * 0.8 + 0.1), 0.01 * fXScale, 0.01, 255, 0, 0, 255)
	ENDREPEAT
	
	// Draw player position
	vCurrent = GET_PLAYER_COORDS(PLAYER_ID())
	
	fNormalisedX = (vCurrent.X - vMinCoord.X) / ABSF(vMaxCoord.X - vMinCoord.X)
	fNormalisedY = (1.0 - (vCurrent.Y - vMinCoord.Y) / ABSF(vMaxCoord.Y - vMinCoord.Y))
	
	FLOAT fPlayerX = (fNormalisedX * 0.8 + 0.1) * fXScale
	FLOAT fPlayerY = (fNormalisedY * 0.8 + 0.1)
	
	DRAW_RECT(fPlayerX, fPlayerY , 0.01 * fXScale, 0.01, 0, 255, 0, 255)
	
	FLOAT fPlayerHeading = GET_ENTITY_HEADING(PLAYER_PED_ID())
	
	DRAW_LINE_2D(fPlayerX, fPlayerY, fPlayerX + (SIN(fPlayerHeading) * 0.8 + 0.1) * fXScale, fPlayerY + (COS(fPlayerHeading) * 0.8 + 0.1), 0.002, 0, 255, 0, 255)
ENDPROC
#ENDIF