using Distributions
using Base.Test

# Core testing procedure

function test_mixture(g::UnivariateMixture, n::Int, ns::Int)
    X = zeros(n)
    for i = 1:n
        X[i] = rand(g)
    end

    K = ncomponents(g)
    pr = probs(g)
    @assert length(pr) == K

    # mean
    mu = 0.0
    for k = 1:K
        mu += pr[k] * mean(component(g, k))
    end
    @test mean(g) ≈ mu

    # evaluation of cdf
    cf = zeros(n)
    for k = 1:K
        c_k = component(g, k)
        for i = 1:n
            cf[i] += pr[k] * cdf(c_k, X[i])
        end
    end

    for i = 1:n
        @test cdf(g, X[i]) ≈ cf[i]
    end
    @test cdf.(g, X) ≈ cf

    # evaluation
    P0 = zeros(n, K)
    LP0 = zeros(n, K)
    for k = 1:K
        c_k = component(g, k)
        for i = 1:n
            x_i = X[i]
            P0[i,k] = pdf(c_k, x_i)
            LP0[i,k] = logpdf(c_k, x_i)
        end
    end

    mix_p0 = P0 * pr
    mix_lp0 = log.(mix_p0)

    for i = 1:n
        @test pdf(g, X[i])                  ≈ mix_p0[i]
        @test logpdf(g, X[i])               ≈ mix_lp0[i]
        @test componentwise_pdf(g, X[i])    ≈ vec(P0[i,:])
        @test componentwise_logpdf(g, X[i]) ≈ vec(LP0[i,:])
    end

    @test pdf.(g, X)                  ≈ mix_p0
    @test logpdf.(g, X)               ≈ mix_lp0
    @test componentwise_pdf(g, X)    ≈ P0
    @test componentwise_logpdf(g, X) ≈ LP0

    # sampling
    Xs = rand(g, ns)
    @test isa(Xs, Vector{Float64})
    @test length(Xs) == ns
    @test isapprox(mean(Xs), mean(g), atol=0.01)
end

function test_mixture(g::MultivariateMixture, n::Int, ns::Int)
    X = zeros(length(g), n)
    for i = 1:n
        X[:,i] = rand(g)
    end

    K = ncomponents(g)
    pr = probs(g)
    @assert length(pr) == K

    # mean
    mu = 0.0
    for k = 1:K
        mu += pr[k] * mean(component(g, k))
    end
    @test mean(g) ≈ mu

    # evaluation
    P0 = zeros(n, K)
    LP0 = zeros(n, K)
    for k = 1:K
        c_k = component(g, k)
        for i = 1:n
            x_i = X[:,i]
            P0[i,k] = pdf(c_k, x_i)
            LP0[i,k] = logpdf(c_k, x_i)
        end
    end

    mix_p0 = P0 * pr
    mix_lp0 = log.(mix_p0)

    for i = 1:n
        x_i = X[:,i]
        @test pdf(g, x_i)                  ≈ mix_p0[i]
        @test logpdf(g, x_i)               ≈ mix_lp0[i]
        @test componentwise_pdf(g, x_i)    ≈ vec(P0[i,:])
        @test componentwise_logpdf(g, x_i) ≈ vec(LP0[i,:])
    end

    @show g
    @show size(X)
    @show size(mix_p0)
    @test pdf(g, X)                  ≈ mix_p0
    @test logpdf(g, X)               ≈ mix_lp0
    @test componentwise_pdf(g, X)    ≈ P0
    @test componentwise_logpdf(g, X) ≈ LP0

    # sampling
    Xs = rand(g, ns)
    @test isa(Xs, Matrix{Float64})
    @test size(Xs) == (length(g), ns)
    @test isapprox(vec(mean(Xs, 2)), mean(g), atol=0.1)
    @test isapprox(cov(Xs, 2)      , cov(g) , atol=0.1)
end

function test_params(g::AbstractMixtureModel)
    C = eltype(g.components)
    pars = params(g)
    mm = MixtureModel(C, pars...)
    @test g.prior == mm.prior
    @test g.components == mm.components
end

function test_params(g::UnivariateGMM)
    pars = params(g)
    mm = UnivariateGMM(pars...)
    @test g == mm
end

# Tests

println("    testing UnivariateMixture")

g_u = MixtureModel(Normal, [(0.0, 1.0), (2.0, 1.0), (-4.0, 1.5)], [0.2, 0.5, 0.3])
@test isa(g_u, MixtureModel{Univariate, Continuous, Normal})
@test ncomponents(g_u) == 3
test_mixture(g_u, 1000, 10^6)
test_params(g_u)
@test minimum(g_u) == -Inf
@test maximum(g_u) == Inf

g_u = MixtureModel([TriangularDist(-1,2,0),TriangularDist(-.5,3,1),TriangularDist(-2,0,-1)])
@test minimum(g_u) == -2.0
@test maximum(g_u) == 3.0
@test insupport(g_u, 2.5) == true
@test insupport(g_u, 3.5) == false

g_u = UnivariateGMM([0.0, 2.0, -4.0], [1.0, 1.2, 1.5], Categorical([0.2, 0.5, 0.3]))
@test isa(g_u, UnivariateGMM)
@test ncomponents(g_u) == 3
test_mixture(g_u, 1000, 10^6)
test_params(g_u)
@test minimum(g_u) == -Inf
@test maximum(g_u) == Inf

println("    testing MultivariateMixture")
g_m = MixtureModel(
    IsoNormal[ MvNormal([0.0, 0.0], 1.0),
               MvNormal([0.2, 1.0], 1.0),
               MvNormal([-0.5, -3.0], 1.6) ],
    [0.2, 0.5, 0.3])
@test isa(g_m, MixtureModel{Multivariate, Continuous, IsoNormal})
@test length(components(g_m)) == 3
@test length(g_m) == 2
@test insupport(g_m, [0.0, 0.0]) == true
test_mixture(g_m, 1000, 10^6)
test_params(g_m)
