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 tLogPDF(x: Float64, df: Int64, loc!: Float64 = 0.0, scale!: Float64 = 1.0): Float64 {
    let y = (x - loc) / scale

    if (df <= 0) {
        throw IllegalArgumentException("tLogPDF: freedom parameter df should be positive.")
    }

    if (scale == 0.0) {
        throw IllegalArgumentException("tLogPDF: scale cannot be 0.")
    }
    
    let d = Float64(df)

    return gammaLog((d + 1.0) * 0.5) - gammaLog(d * 0.5) - 0.5 * (log(Float64.getPI()) + log(d)) - 0.5 * (d + 1.0) * log(1.0 + y * y / d) - log(scale)
}  

/*
 * Probability density function
 */
public func tPDF(x: Float64, df: Int64, loc!: Float64 = 0.0, scale!: Float64 = 1.0): Float64 {
    if (scale == 0.0) {
        throw IllegalArgumentException("tPDF: scale cannot be 0.")
    }
    if (df <= 0) {
        throw IllegalArgumentException("tPDF: freedom parameter df should be positive.")
    }

    return exp(tLogPDF(x, df, loc: loc, scale: scale))
}

public func logBeta(a: Float64, b: Float64): Float64 {
    return gammaLog(a) + gammaLog(b) - gammaLog(a + b)
}

public func regularizedIncompleteBeta(a: Float64, b: Float64, x: Float64,
                                      epsilon!: Float64 = 1.0e-15,
                                      maxIterations!: Int64 = 200): Float64 {
    let tiny: Float64 = 1.0e-30

    // Computes the regularized incomplete beta function I_x(a, b)
    if (x < 0.0 || x > 1.0) {
        throw IllegalArgumentException("x must be in the interval [0, 1]")
    }
    if (a <= 0.0 || b <= 0.0) {
        throw IllegalArgumentException("a and b must be positive")
    }

    if (x == 0.0) {
        return 0.0
    }
    if (x == 1.0) {
        return 1.0
    }

    // The continued fraction converges faster for smaller x
    // A symmetry relation I_x(a,b) = 1 - I_{1-x}(b,a) is used
    // The switch point is chosen to optimize convergence
    if (x > (a + 1.0) / (a + b + 2.0)) {
        return 1.0 - regularizedIncompleteBeta(b, a, 1.0 - x)
    }

    // The factor in front of the continued fraction
    // log(prefix) = a*log(x) + b*log(1-x) - log(a) - log(B(a,b))
    let log_prefix = a * log(x) + b * log(1.0 - x) - log(a) - logBeta(a, b)
    let prefix = exp(log_prefix)
    
    // Re-implementation of the continued fraction part using a more standard Lentz method
    // for f = 1/(1+d1/(1+d2/...))
    
    let qab = a + b
    let qap = a + 1.0
    let qam = a - 1.0
    
    var c = 1.0
    var d = 1.0 - qab * x / qap
    if (abs(d) < tiny) {
        d = tiny
    }
    d = 1.0 / d
    var h = d

    for (m in 1..(maxIterations + 1)) {
        let fm = Float64(m)
        let m2 = 2.0 * fm
        // Odd term
        let num_odd = fm * (b - fm) * x / ((qam + m2) * (a + m2))
        d = 1.0 + num_odd * d
        if (abs(d) < tiny) {
            d = tiny
        }
        c = 1.0 + num_odd / c
        if (abs(c) < tiny) {
            c = tiny
        }
        d = 1.0 / d
        h = h * d * c
        
        // Even term
        let num_even = -(a + fm) * (qab + fm) * x / ((a + m2) * (qap + m2))
        d = 1.0 + num_even * d
        if (abs(d) < tiny) {
            d = tiny
        }
        c = 1.0 + num_even / c
        if (abs(c) < tiny) {
            c = tiny
        }
        d = 1.0 / d
        let h_new = h * d * c

        if (abs(h_new - h) < epsilon) {
             return prefix * h_new
        }
        h = h_new
    }

    return prefix * h
}

public func tCDF(t: Float64, df: Int64, loc!: Float64 = 0.0, scale!: Float64 = 1.0): Float64 {
    // Computes the cumulative distribution function (CDF) for the
    // Student's t-distribution
    let y = (t - loc) / scale
    if (df <= 0) {
        throw IllegalArgumentException("Degrees of freedom must be positive")
    }

    // Parameters for the regularized incomplete beta function
    let a = Float64(df) / 2.0
    let b = 0.5
    let x = Float64(df) / (Float64(df) + y * y)

    // The incomplete beta function is defined for x in [0, 1]
    // which is always true for our definition of x
    let inc_beta_val = regularizedIncompleteBeta(a, b, x)

    if (y > 0.0) {
        return 1.0 - 0.5 * inc_beta_val
    } else { // t <= 0
        return 0.5 * inc_beta_val
    }
}

/*
 * Mean of the distribution.
 */
public func tMean(df: Int64, loc!: Float64 = 0.0, scale!: Float64 = 1.0): Float64 {
    if (df < 2 || scale == 0.0) {
        return Float64.Inf    // return infinite
    }

    return loc
}   

/*
 * Var of the distribution.
 */
public func tVar(df: Int64, loc!: Float64 = 0.0, scale!: Float64 = 1.0): Float64 {
    if (df <= 2 || scale == 0.0) {
        return Float64.Inf    // return infinite
    }

    let d = Float64(df)
    return d / (d - 2.0) * scale * scale
}   

@Test
public class TestT {
    @TestCase
    func testT(): Unit {
        @Assert(approxEqual(tLogPDF(1.0, 2, loc: 1.0, scale: 2.0), -1.732867951399863,  atol:1e-10))
        @Assert(approxEqual(tPDF(1.0, 2, loc: 1.0, scale: 2.0),    0.1767766952966369,  atol:1e-10))
        @Assert(approxEqual(tCDF(1.0, 2, loc: 1.0, scale: 2.0),    0.5,                 atol:1e-10))
        @Assert(approxEqual(tCDF(1.5, 2, loc: 1.0, scale: 2.0),    0.5870388279778489,  atol:1e-10))

        @Assert(approxEqual(tCDF(-2.236068, 5), 0.037793408))
        @Assert(approxEqual(tCDF(2.236068, 5),  0.962206591))
        @Assert(approxEqual(tCDF(-2.48869, 3),  0.044293270))
        @Assert(approxEqual(tCDF(2.48869, 3),   0.955706729))

        @Assert(approxEqual(tMean(4, loc: 3.0, scale: 2.0), 3.0, atol:1e-10))
        @Assert(approxEqual(tVar(4, loc: 3.0, scale: 2.0),  8.0, atol:1e-10))
    }
}