module ParasitemiaCore.KMeans

open System.Collections.Generic
open System.Drawing

open Emgu.CV
open Emgu.CV.Structure

type Result =
    {
        fg : Image<Gray, byte>
        mean_bg : float32
        mean_fg : float32
        d_fg : Image<Gray, float32> // Euclidean distances of the foreground to mean_fg.
    }

let kmeans (img : Image<Gray, float32>) : Result =
    let nbIteration = 4
    let w = img.Width
    let h = img.Height

    let min = ref [| 0.0 |]
    let minLocation = ref <| [| Point () |]
    let max = ref [| 0.0 |]
    let maxLocation = ref <| [| Point () |]
    img.MinMax (min, max, minLocation, maxLocation)

    let minf = float32 (!min).[0]
    let maxf = float32 (!max).[0]

    let mutable mean_bg = maxf - (maxf - minf) / 4.f
    let mutable mean_fg = minf + (maxf - minf) / 4.f
    use mutable d_bg : Image<Gray, float32> = null
    let mutable d_fg : Image<Gray, float32> = null
    let fg = new Image<Gray, byte> (img.Size)

    let imgData = img.Data
    let fgData = fg.Data

    for i = 1 to nbIteration do
        match d_bg with
        | null -> ()
        | _ ->
            d_bg.Dispose ()
            d_fg.Dispose ()

        // EmGu doesn't import the in-place version of 'AbsDiff' so we have to create two images for each iteration.
        d_bg <- img.AbsDiff (Gray (float mean_bg))
        d_fg <- img.AbsDiff (Gray (float mean_fg))

        CvInvoke.Compare (d_fg, d_bg, fg, CvEnum.CmpType.LessThan)

        let mutable bg_total = 0.f
        let mutable bg_nb = 0

        let mutable fg_total = 0.f
        let mutable fg_nb = 0

        for i = 0 to h - 1 do
            for j = 0 to w - 1 do
                if fgData.[i, j, 0] > 0uy then
                    fg_total <- fg_total + imgData.[i, j, 0]
                    fg_nb <- fg_nb + 1
                else
                    bg_total <- bg_total + imgData.[i, j, 0]
                    bg_nb <- bg_nb + 1

        mean_bg <- bg_total / float32 bg_nb
        mean_fg <- fg_total / float32 fg_nb

    { fg = fg; mean_bg = mean_bg; mean_fg = mean_fg; d_fg = d_fg }