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 trapezoidLogPDF(x: Float64, a: Float64, c: Float64, loc!: Float64 = 0.0, scale!: Float64 = 1.0): Float64 {
    let y = (x - loc) / scale
    var paramCheck = true
    if (0.0 <= a && a <= 1.0 && 0.0 <= c && c <= 1.0 && a <= c) {
        paramCheck = false
    }
    if (paramCheck) {
        throw IllegalArgumentException("trapezoidLogPDF: shape parameter out of bound.")
    }

    let temp = trapezoidPDF(x, a, c, loc:loc, scale: scale)
    if (temp < 0.000001) {
        throw IllegalArgumentException("trapezoidLogPDF: return-value too small.")
    }

    return log(temp)
}

/*
 * Probability density function
 */
public func trapezoidPDF(x: Float64, a: Float64, c: Float64, loc!: Float64 = 0.0, scale!: Float64 = 1.0): Float64 {
    let y = (x - loc) / scale
    var paramCheck = true
    if (0.0 <= a && a <= 1.0 && 0.0 <= c && c <= 1.0 && a <= c) {
        paramCheck = false
    }
    if (paramCheck) {
        throw IllegalArgumentException("trapezoidPDF: shape parameter out of bound.")
    }

    let t = 2.0 / (c - a + 1.0)
    var res = 0.0
    if (y < a) {
        res = t * y / a
    } else if (a <= y && y <= c) {
        res = t
    } else {
        res = t * (1.0 - y) / (1.0 - c)
    }

    return res / scale
}

/*
 * Cumulative probability density function
 */
public func trapezoidCDF(x: Float64, a: Float64, c: Float64, loc!: Float64 = 0.0, scale!: Float64 = 1.0): Float64 {
    let y = (x - loc) / scale
    var paramCheck = true
    if (0.0 <= a && a <= 1.0 && 0.0 <= c && c <= 1.0 && a <= c) {
        paramCheck = false
    }
    if (paramCheck) {
        throw IllegalArgumentException("trapezoidLogCDF: shape parameter out of bound.")
    }

    var res = 0.0
    if (y < a) {
        res = y * y / a /(c - a + 1.0)
    } else if (a <= y && y <= c) {
        res = (a + 2.0 * (y - a)) / (c - a + 1.0)
    } else {
        res = 1.0 - ((1.0 - y) * (1.0 - y) / (c - a + 1.0) / (1.0 - c))
    }
    return res
}

/*
 * Log of Cumulative probability density function
 */
public func trapezoidLogCDF(x: Float64, a: Float64, c: Float64, loc!: Float64 = 0.0, scale!: Float64 = 1.0): Float64 {
    let y = (x - loc) / scale
    var paramCheck = true
    if (0.0 <= a && a <= 1.0 && 0.0 <= c && c <= 1.0 && a <= c) {
        paramCheck = false
    }
    if (paramCheck) {
        throw IllegalArgumentException("trapezoidCDF: shape parameter out of bound.")
    }

    let temp = trapezoidCDF(x, a, c, loc:loc, scale: scale)
    if (temp < 0.000001) {
        throw IllegalArgumentException("trapezoidCDF: return-value too small.")
    }

    return log(temp)
}

/*
 * PPF
 */
public func trapezoidPPF(q: Float64, a: Float64, c: Float64, loc!: Float64 = 0.0, scale!: Float64 = 1.0): Float64 {
    var paramCheck = true
    if (0.0 <= a && a <= 1.0 && 0.0 <= c && c <= 1.0 && a <= c) {
        paramCheck = false
    }
    if (paramCheck) {
        throw IllegalArgumentException("trapezoidPPF: shape parameter out of bound.")
    }
    
    if (q <= 0.0 || q >= 1.0) {
        throw IllegalArgumentException("trapezoidPPF: quantile out of bound.")
    }

    let ta = trapezoidCDF(a, a, c)
    let tc = trapezoidCDF(c, a, c)
    var res = 0.0
    if (q < ta) {
        res = sqrt(q * a * (c - a + 1.0))
    } else if (q <= tc) {
        res = 0.5 * q * (c - a + 1.0) + 0.5 * a
    } else {
        res = 1.0 - sqrt((1.0 - q) * (c - a + 1.0) * (1.0 - c))
    }

    return res * scale + loc
}

@Test
public class TestTrapezoid {
    @TestCase
    func testTrapezoid(): Unit {
        @Assert(approxEqual(trapezoidLogPDF(2.0, 0.2, 0.5, loc: 1.0, scale: 2.0), -0.2623642644674911,  atol:1e-13))
        @Assert(approxEqual(trapezoidLogCDF(2.0, 0.2, 0.5, loc: 1.0, scale: 2.0), -0.48550781578170077, atol:1e-13))
        @Assert(approxEqual(trapezoidPPF(0.2, 0.2, 0.5, loc: 1.0, scale: 2.0),     1.46,                atol:1e-6))
    }
}