module MatchingEllipses

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

open Types
open Utils


// Do not take in account matching score below this when two ellipses are matched.
let matchingScoreThreshold1 = 0.6

// All ellipsee with a score below this is removed.
let matchingScoreThreshold2 = 2.

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

    member this.Ellipse = e

    member this.MatchingScore = matchingScore

    member this.AddMatchingScore(score: float) =
        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 (radiusMin: float) =
    let ellipses = List<EllipseScoreFlaggedKd>()

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

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

    // Process all ellipses and return ellipses ordered from the best score to the worst.
    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 = radiusMin
            for e in ellipses do
                e.Processed <- true
                let areaE = e.Ellipse.Area
                let window = { KdTree.minX = e.Ellipse.Cx - windowSize / 2.0
                               KdTree.maxX = e.Ellipse.Cx + windowSize / 2.0
                               KdTree.minY = e.Ellipse.Cy - windowSize / 2.0
                               KdTree.maxY = e.Ellipse.Cy + windowSize / 2.0 }
                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 (commonArea, _, _) ->
                            let matchingScore = 2.0 * commonArea / (areaE + areaOther)
                            if matchingScore >= matchingScoreThreshold1
                            then
                                other.AddMatchingScore(matchingScore)
                                e.AddMatchingScore(matchingScore)
                        | _ -> ()

            // 3) Sort ellipses by their score.
            ellipses.Sort(fun e1 e2 -> e2.MatchingScore.CompareTo(e1.MatchingScore))

            // 4) Remove ellipses wich have a low score.
            let i = ellipses.BinarySearch(EllipseScoreFlaggedKd(matchingScoreThreshold2, Ellipse(0.0, 0.0, 0.0, 0.0, 0.0)),
                                          { new IComparer<EllipseScoreFlaggedKd> with
                                                member this.Compare(e1, e2) = e2.MatchingScore.CompareTo(e1.MatchingScore) }) |> abs
            let nbToRemove = ellipses.Count - i
            if nbToRemove > 0
            then
                for j in i .. nbToRemove - 1 do
                    ellipses.[j].Removed <- true
                ellipses.RemoveRange(i, nbToRemove)

            // 5) Remove ellipses whose center is into an ellipse with a better score
            for e in ellipses do
                if not e.Removed
                then
                    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 && other.MatchingScore < e.MatchingScore
                        then
                            if e.Ellipse.Scale(0.8).Contains other.Ellipse.Cx other.Ellipse.Cy
                            then
                                other.Removed <- true
            ellipses.RemoveAll(fun e -> e.Removed) |> ignore

            List.ofSeq ellipses |> List.map (fun e -> e.Ellipse)