module ParasitemiaCore.KdTree

open System

type I2DCoords =
    abstract X : float32
    abstract Y : float32

// Compare 'e1' and 'e2' by X.
let cmpX (e1 : I2DCoords) (e2 : I2DCoords) : int =
    match e1.X.CompareTo e2.X with
    | 0 -> match e1.Y.CompareTo e2.Y with
           | 0 -> e1.GetHashCode().CompareTo (e2.GetHashCode ())
           | v -> v
    | v -> v

// Compare 'e1' and 'e2' by Y.
let cmpY (e1 : I2DCoords) (e2 : I2DCoords) : int =
    match e1.Y.CompareTo e2.Y with
    | 0 -> match e1.X.CompareTo e2.X with
           | 0 -> e1.GetHashCode().CompareTo (e2.GetHashCode ())
           | v -> v
    | v -> v

type Region = { minX : float32; maxX : float32; minY : float32; maxY : float32 } with
    member this.Contains px py : bool =
        px >= this.minX && px <= this.maxX &&
        py >= this.minY && py <= this.maxY

    member this.IsSub otherRegion : bool =
        this.minX >= otherRegion.minX && this.maxX <= otherRegion.maxX &&
        this.minY >= otherRegion.minY && this.maxY <= otherRegion.maxY

    member this.Intersects otherRegion : bool =
        this.minX < otherRegion.maxX && this.maxX >= otherRegion.minX &&
        this.minY < otherRegion.maxY && this.maxY >= otherRegion.minY

type Tree<'a when 'a :> I2DCoords> =
    | Node of float32 * Tree<'a> * Tree<'a>
    | Leaf of 'a

    static member BuildTree (l : 'a list) : Tree<'a> =
        let xSorted = List.toArray l
        let ySorted = List.toArray l
        Array.sortInPlaceWith cmpX xSorted
        Array.sortInPlaceWith cmpY ySorted

        let rec buildTreeFromSortedArray (pXSorted : 'a[]) (pYSorted : 'a[]) (depth : int) : Tree<'a> =
            if pXSorted.Length = 1 then
                Leaf pXSorted.[0]
            else
                if depth % 2 = 1 then // 'depth' is odd -> vertical splitting else horizontal splitting.
                    let leftX, rightX = Array.splitAt ((pXSorted.Length + 1) / 2) pXSorted
                    let splitElement = Array.last leftX
                    let leftY, rightY = Array.partition (fun (e : 'a) -> cmpX e splitElement <= 0) pYSorted // FIXME: Maybe this operation can be optimized.
                    Node (splitElement.X, buildTreeFromSortedArray leftX leftY (depth + 1), buildTreeFromSortedArray rightX rightY (depth + 1))
                else
                    let downY, upY = Array.splitAt ((pYSorted.Length + 1) / 2) pYSorted
                    let splitElement = Array.last downY
                    let downX, upX = Array.partition (fun (e : 'a) -> cmpY e splitElement <= 0) pXSorted // FIXME: Maybe this operation can be optimized.
                    Node (splitElement.Y, buildTreeFromSortedArray downX downY (depth + 1), buildTreeFromSortedArray upX upY (depth + 1))

        buildTreeFromSortedArray xSorted ySorted 1

    member this.Search (searchRegion : Region) : 'a list =
        let rec valuesFrom (tree : Tree<'a>) (acc : 'a list) : 'a list =
            match tree with
            | Node (_, left, right) -> (valuesFrom right (valuesFrom left acc))
            | Leaf v -> v :: acc

        let rec searchWithRegion (tree : Tree<'a>) (currentRegion : Region) (depth : int) : 'a list =
            match tree with
            | Leaf v -> if searchRegion.Contains v.X v.Y then [v] else []
            | Node (splitValue, part1, part2) ->
                let valuesInRegion (region : Region) (treeRegion : Tree<'a>) =
                    if region.IsSub searchRegion then
                        valuesFrom treeRegion []
                    elif region.Intersects searchRegion then
                        searchWithRegion treeRegion region (depth + 1)
                    else
                        []

                if depth % 2 = 1 then // Vertical splitting.
                    let leftRegion = { currentRegion with maxX = splitValue }
                    let rightRegion = { currentRegion with minX = splitValue }
                    (valuesInRegion leftRegion part1) @ (valuesInRegion rightRegion part2)
                else // Horizontal splitting.
                    let downRegion = { currentRegion with maxY = splitValue }
                    let upRegion = { currentRegion with minY = splitValue }
                    (valuesInRegion downRegion part1) @ (valuesInRegion upRegion part2)

        searchWithRegion this { minX = Single.MinValue; maxX = Single.MaxValue; minY = Single.MinValue; maxY = Single.MaxValue } 1