package scientific.stats.discrete

import std.math.*
import std.unittest.*
import std.unittest.testmacro.*

import scientific.numbers.*
import scientific.stats.utils.*
import scientific.stats.random.*

/*
 * Log of Probability density function
 */
public func hypergeomLogPDF(x: Int64, M: Int64, n: Int64, N: Int64): Float64 {
    let l = max(0, N + n - M)
    let r = min(n, N)
    if (x < l || x > r) {
        throw IllegalArgumentException("hypergeomLogPDF: parameters out of bound.")
    }

    let res = hypergeomPDF(x, M, n, N)
    if (res < 0.000001) {
        throw IllegalArgumentException("hypergeomLogCDF: return-value too small.")
    }
    return log(res)
}


/*
 * Probability density function
 */
public func hypergeomPDF(x: Int64, M: Int64, n: Int64, N: Int64): Float64 {
    let l = max(0, N + n - M)
    let r = min(n, N)
    if (x < l || x > r) {
        throw IllegalArgumentException("hypergeomLogPDF: parameters out of bound.")
    }

    let t1 = comb(n, x)
    let t2 = comb(M - n, N - x)
    let t3 = comb(M, N)

    let t = t1 * t2
    return Float64(t) / Float64(t3) 
}


/*
 * Cumulative probability density function
 */
public func hypergeomCDF(x: Int64, M: Int64, n: Int64, N: Int64): Float64 {
    let l = max(0, N + n - M)
    let r = min(n, N)
    if (x < l || x > r) {
        throw IllegalArgumentException("hypergeomCDF: parameters out of bound.")
    }

    var res = 0.0
    for (i in l..(x + 1)) {
        res += hypergeomPDF(i, M, n, N)
    }
    return res 
}


/*
 * Log of cumulative probability density function
 */
public func hypergeomLogCDF(x: Int64, M: Int64, n: Int64, N: Int64): Float64 {
    let l = max(0, N + n - M)
    let r = min(n, N)
    if (x < l || x > r) {
        throw IllegalArgumentException("hypergeomCDF: parameters out of bound.")
    }

    let res = hypergeomCDF(x, M, n, N)
    if (res < 0.000001) {
        throw IllegalArgumentException("hypergeomLogCDF: return-value too small.")
    }

    return log(res)
}

@Test
public class TestHypergeom {
    @TestCase
    func testHypergeom(): Unit {
        @Assert(approxEqual(hypergeomPDF(2, 20, 7, 12), 0.047678018575851404, atol:1e-13))
        @Assert(approxEqual(hypergeomCDF(2, 20, 7, 12), 0.05211558307533541,  atol:1e-8))
    }
}