package scientific.stats.multivariate

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

import scientific.utils.*
import scientific.numbers.*
import scientific.linear.*
import scientific.stats.random.*
import scientific.stats.normal.randn

/*
 *  Log of Probability density function.
 */
public func multiNormalLogPDF(x: Vector<Float64>, mean: Vector<Float64>, cov: Matrix<Float64>): Float64 {
    if (x.size() != mean.size()) {
        throw IllegalArgumentException("multiNormalLogPDF: parameters dimensions incompatible.")
    }
    let k = Float64(mean.size())
    let cov_inv = inv(cov)
    let cov_det = det(cov)

    let temp = x - mean
    let c = 2.0 * Float64.getPI()
    return - 0.5 * temp.dot(cov_inv * temp) - 0.5 * k * log(c) - 0.5 * log(cov_det)
}

/*
 *  Probability density function.
 */
public func multiNormalPDF(x: Vector<Float64>, mean: Vector<Float64>, cov: Matrix<Float64>): Float64 {
    if (x.size() != mean.size()) {
        throw IllegalArgumentException("multiNormalPDF: parameters dimensions incompatible.")
    }
    
    return exp(multiNormalLogPDF(x, mean, cov))
}

/*
 * sample
 */
public func multiNormalSample(r: Random, mean: Vector<Float64>, cov: Matrix<Float64>): Vector<Float64> {
    let n = mean.size()

    let cov_chol = cholesky(cov)
    let x = randn(r, n, 0.0, 1.0)
    let res = cov_chol * x + mean
    return res
}


@Test
public class TestMultiNormal {
    @TestCase
    func testMultiNormal(): Unit {
        let x = vector([1.0, 1.0])
        let mean = vector([1.0, 2.0])
        let cov = matrix([[1.0, 1.0], [1.0, 2.0]])
        let res = multiNormalPDF(x, mean, cov)
        @Assert(approxEqual(res, 0.09653235263005393, atol:1e-6))
    }
}