using Distributions, ForwardDiff, Random, Plots
using ProgressMeter, LinearAlgebra, LogExpFunctions
using Base.Threads: @threads
using JLD

d = 2

function logp(z)
    g1 = -0.5 * z[1]^2 / 0.15^2 - 0.5 * (z[2] - 2)^2 / 1^2 - log(2π) - log(0.15)
    g2 = -0.5 * (z[1] + 2.0)^2 / 1^2 - 0.5 * z[2]^2 / 0.15^2 - log(2π) - log(0.15)
    g3 = -0.5 * z[2]^2 / 0.15^2 - 0.5 * (z[1] - 2)^2 / 1^2 - log(2π) - log(0.15)
    g4 = -0.5 * (z[2] + 2.0)^2 / 1^2 - 0.5 * z[1]^2 / 0.15^2 - log(2π) - log(0.15)
    return LogExpFunctions.logsumexp([g1, g2, g3, g4]) - log(4)
end

∇logp(x) = ForwardDiff.gradient(logp, x)
function logq(x, μ, D)
    return -0.5 * d * log(2π) - sum(log, abs.(D)) - 0.5 * sum(abs2, (x .- μ) ./ (D .+ 1e-8))
end

if !isdir("figure")
    mkdir("figure")
end
if !isdir("result")
    mkdir("result")
end