module ParasitemiaCore.MatchingEllipses

open System
open System.Linq
open System.Collections
open System.Collections.Generic

open Types
open Utils

type private EllipseScoreFlaggedKd (matchingScore: float32, e: Ellipse) =
    let mutable matchingScore = matchingScore

    member this.Ellipse = e

    member this.MatchingScore = matchingScore

    member this.AddMatchingScore (score: float32) =
        matchingScore <- matchingScore + score

    member val Processed = false with get, set
    member val Removed = false with get, set

    interface KdTree.I2DCoords with
        member this.X = this.Ellipse.Cx
        member this.Y = this.Ellipse.Cy

type MatchingEllipses (radius: float32) =
    let ellipses = List<EllipseScoreFlaggedKd>()

    // All ellipses with a score below this are removed.
    let matchingScoreThreshold = 0.4f

    member this.Add (e: Ellipse) =
        ellipses.Add(EllipseScoreFlaggedKd(0.f, e))

    member this.Ellipses : Ellipse list =
        List.ofSeq ellipses |> List.map (fun e -> e.Ellipse)

    member this.PrunedEllipses : Ellipse list =
        if ellipses.Count = 0
        then
            []
        else
            // 1) Create a kd-tree from the ellipses list.
            let tree = KdTree.Tree.BuildTree (List.ofSeq ellipses)

            // 2) Compute the matching score of each ellipses.
            let windowSize = radius / 2.f
            for e in ellipses do
                e.Processed <- true
                let areaE = e.Ellipse.Area
                let window = { KdTree.minX = e.Ellipse.Cx - windowSize / 2.f
                               KdTree.maxX = e.Ellipse.Cx + windowSize / 2.f
                               KdTree.minY = e.Ellipse.Cy - windowSize / 2.f
                               KdTree.maxY = e.Ellipse.Cy + windowSize / 2.f }
                for other in tree.Search window do
                    if not other.Processed
                    then
                        let areaOther = other.Ellipse.Area
                        match EEOver.EEOverlapArea e.Ellipse other.Ellipse with
                        | Some (overlapArea, _, _)
                            // Because of approximation error, see https://github.com/chraibi/EEOver/issues/4
                            when overlapArea - areaE < 1.f && overlapArea - areaOther < 1.f ->
                                let matchingScore = (2.f * overlapArea / (areaE + areaOther)) ** 30.f
                                other.AddMatchingScore(matchingScore)
                                e.AddMatchingScore(matchingScore)
                        | _ -> ()

            // 3) Remove ellipses whose center is near the center of another ellipse with a better score.
            for e in ellipses do
                if e.MatchingScore < matchingScoreThreshold
                then
                    e.Removed <- true
                else
                    let window = { KdTree.minX = e.Ellipse.Cx - e.Ellipse.A
                                   KdTree.maxX = e.Ellipse.Cx + e.Ellipse.A
                                   KdTree.minY = e.Ellipse.Cy - e.Ellipse.A
                                   KdTree.maxY = e.Ellipse.Cy + e.Ellipse.A }
                    for other in tree.Search window do
                        if not other.Removed && e.MatchingScore > other.MatchingScore
                        then
                            // Case where ellipses are too close.
                            if distanceTwoPoints (PointF(e.Ellipse.Cx, e.Ellipse.Cy)) (PointF(other.Ellipse.Cx, other.Ellipse.Cy)) < 0.3f * e.Ellipse.B
                            then
                                other.Removed <- true
                            else
                                // Case where ellipses are overlapped.
                                match EEOver.EEOverlapArea e.Ellipse other.Ellipse with
                                | Some (overlapArea, _, _) when e.Ellipse.Area < 1.1f * overlapArea || other.Ellipse.Area < 1.1f * overlapArea ->
                                    other.Removed <- true
                                | _ ->
                                    ()
            ellipses
            |> List.ofSeq
            |> List.filter (fun e -> not e.Removed)
            |> List.sortWith (fun e1 e2 -> e2.MatchingScore.CompareTo(e1.MatchingScore))
            |> List.map (fun e -> e.Ellipse)