module KMedians

open System.Collections.Generic
open System.Drawing

open Emgu.CV
open Emgu.CV.Structure

type Result = {
    fg: Image<Gray, byte>
    median_bg: float
    median_fg: float
    d_fg: Image<Gray, float32> } // Distances to median_fg.

let kmedians (img: Image<Gray, float32>) (fgFactor: float) : Result =
    let nbIteration = 3
    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 mutable median_bg = (!max).[0] - ((!max).[0] - (!min).[0]) / 4.0
    let mutable median_fg = (!min).[0] + ((!max).[0] - (!min).[0]) / 4.0
    use mutable d_bg = new Image<Gray, float32>(img.Size)
    let mutable d_fg = new Image<Gray, float32>(img.Size)
    let mutable fg = new Image<Gray, byte>(img.Size)

    for i in 1 .. nbIteration do
        CvInvoke.Pow(img - median_bg, 2.0, d_bg)
        CvInvoke.Pow(img - median_fg, 2.0, d_fg)
        fg <- (d_fg * fgFactor).Cmp(d_bg, CvEnum.CmpType.LessThan)

        median_fg <- MathNet.Numerics.Statistics.Statistics.Median(seq {
            for i in 0 .. h - 1 do
                for j in 0 .. w - 1 do
                    if fg.Data.[i, j, 0] > 0uy then yield img.Data.[i, j, 0] |> float })

        median_bg <- MathNet.Numerics.Statistics.Statistics.Median(seq {
            for i in 0 .. h - 1 do
                for j in 0 .. w - 1 do
                    if fg.Data.[i, j, 0] = 0uy then yield img.Data.[i, j, 0] |> float })

    CvInvoke.Sqrt(d_fg, d_fg)

    { fg = fg; median_bg = median_bg; median_fg = median_fg; d_fg = d_fg }