package scientific.stats.continuous

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

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

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

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

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

    let temp = -(y - k) * (y - k) / (2.0 * y * k * k)
    return  temp - 0.5 * (log(2.0 * Float64.getPI()) + 3.0 * log(y)) - log(scale)
}

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

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

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

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

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

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

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

    let res = invgaussLogCDF(x, k, loc: loc, scale: scale)
    return exp(res)
}

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

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

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

    let t = 1.0 / sqrt(y)
    let a = log(normalCDF(t * (y / k - 1.0)))
    let b = 2.0 / k + log(normalCDF(-t * (y / k + 1.0)))
    let temp = 1.0 + exp(b - a)
    return a + log(temp)
}

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

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

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

    let temp = invgaussVar(k, loc: loc, scale: scale)

    if (temp < 0.000001) {
        throw IllegalArgumentException("invgaussStd: return-value too small.")
    }

    return sqrt(temp)
}

@Test
public class TestInvgauss {
    @TestCase
    func testInvgauss(): Unit {
        @Assert(approxEqual(invgaussLogPDF(2.0, 2.0, loc: 1.0, scale: 2.0), -1.1348649429247,    atol:1e-13))
        @Assert(approxEqual(invgaussLogCDF(2.0, 2.0, loc: 1.0, scale: 2.0), -1.3894522486353145, atol:1e-13)) 
        @Assert(approxEqual(invgaussMean(2.0, loc: 1.0, scale: 2.0),         5.0,                atol:1e-13)) 
        @Assert(approxEqual(invgaussVar(2.0, loc: 1.0, scale: 2.0),          32.0,               atol:1e-13)) 
    }
}