package scientific.stats.continuous

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

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

/*
 * Log of Probability density function
 */
public func chiLogPDF(x: Float64, k: Float64, loc!: Float64 = 0.0, scale!: Float64 = 1.0): Float64 {
    let y = (x - loc) / scale

    if (y < 0.0) {
        throw IllegalArgumentException("chiLogPDF: input value x out of bound.")
    }

    if (k <= 0.0) {
        throw IllegalArgumentException("chiLogPDF: shape parameter out of bound.")
    }

    return - (0.5 * k - 1.0) * log(2.0) - gammaLog(k * 0.5) + (k - 1.0) * log(y) - 0.5 * y * y - log(scale)
}

/*
 * Probability density function
 */
public func chiPDF(x: Float64, k: Float64, loc!: Float64 = 0.0, scale!: Float64 = 1.0): Float64 {
    let y = (x - loc) / scale

    if (y < 0.0) {
        throw IllegalArgumentException("chiPDF: input value x out of bound.")
    }

    if (k <= 0.0) {
        throw IllegalArgumentException("chiPDF: shape parameter out of bound.")
    }

    return exp(chiLogPDF(x, k, loc: loc, scale: scale))
}

/*
 * Sample
 */
public func chiSample(r: Random, k: Float64, loc!: Float64 = 0.0, scale!: Float64 = 1.0): Float64 {
    if (k <= 0.0) {
        throw IllegalArgumentException("chiSample: shape parameter out of bound.")
    }

    return sqrt(gamSample(r, k)) * scale + loc
}

/*
 * Mean
 */
public func chiMean(k: Float64, loc!: Float64 = 0.0, scale!: Float64 = 1.0): Float64 {
    if (k <= 0.0) {
        throw IllegalArgumentException("chiMean: shape parameter out of bound.")
    }

    let temp = sqrt(2.0) * exp(gammaLog(0.5 * k + 0.5) - gammaLog(0.5 * k))
    return temp * scale + loc
}


/*
 * Var
 */
public func chiVar(k: Float64): Float64 {
    if (k <= 0.0) {
        throw IllegalArgumentException("chiVar: shape parameter out of bound.")
    }

    let temp = chiMean(k)
    return k - temp * temp
}

@Test
public class TestChi {
    @TestCase
    func testChiLogPDF(): Unit {
        @Assert(approxEqual(chiLogPDF(2.0, 2.0, loc: 1.0, scale: 2.0), -1.5112943611198906, atol:1e-13))
    }

    @TestCase
    func testChiMean(): Unit {
        @Assert(approxEqual(chiMean(2.0, loc: 1.0, scale: 2.0), 3.5066282746310007, atol:1e-13))  
    }

    @TestCase
    func testChiVar(): Unit {
        @Assert(approxEqual(chiVar(2.0), 0.4292036732051032, atol:1e-13))
    }
}