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.*

/*
 * Probability density function for binom distribution
 */
public func binomPDF(k: Int64, n: Int64, p: Float64): Float64 {
    if (k >= 0 && n >= k) {
        return Float64(comb(n, k)) * pow(p, Float64(k)) * pow(1.0 - p, Float64(n - k))
    } else {
        return 0.0; // Input not in valid range
    }
}

/*
 * Cumulative density function for binom distribution
 */
public func binomCDF(x: Int64, n: Int64, p: Float64): Float64 {
    if (x < 0) {
        return 0.0
    }
    if (x > n) {
        return 1.0
    }

    var ans: Float64 = 0.0
    for (i in 0..(x + 1)) {
        ans += binomPDF(i, n, p)
    }
    return ans
}

/**
 * Sample the binom distribution n times, return the sum
 */
public func binomSample(m: Random, p: Float64, n: Int64) {
    var ans = 0
    for (i in 0..n) {
        let u = m.nextFloat64()
        if (u < p) {
            ans++
        }
    }
    return ans
}

/**
 * Sample the binom distribution once
 */
public func binomSample(m: Random, p: Float64) {
    return binomSample(m, p, 1)
}

@Test
public class TestBinom {
    @TestCase
    func testBinomPDF(): Unit {
        @Assert(approxEqual(binomPDF(1, 10, 0.2), 0.268435, atol:1e-6))
        @Assert(approxEqual(binomPDF(2, 10, 0.2), 0.301990, atol:1e-6))
        @Assert(approxEqual(binomPDF(3, 10, 0.2), 0.201327, atol:1e-6))
        @Assert(approxEqual(binomPDF(4, 10, 0.2), 0.088080, atol:1e-6))
        @Assert(approxEqual(binomPDF(5, 10, 0.2), 0.026424, atol:1e-6))
    }

    @TestCase
    func testBinomCDF(): Unit {
        @Assert(approxEqual(binomCDF(1, 10, 0.2), 0.3758096384, atol:1e-13))
        @Assert(approxEqual(binomCDF(2, 10, 0.2), 0.6777995264, atol:1e-13))
        @Assert(approxEqual(binomCDF(3, 10, 0.2), 0.8791261184, atol:1e-13))
        @Assert(approxEqual(binomCDF(4, 10, 0.2), 0.9672065024, atol:1e-13))
        @Assert(approxEqual(binomCDF(5, 10, 0.2), 0.9936306176, atol:1e-13))
    }
}