ENV["LINES"] = 100

@time using Random: seed!
@time using LinearAlgebra
@time using StatsBase
@time using QuadGK

@time using StatsFuns
logmeanexp(x) = logsumexp(x) - log(length(x))

@time using StatsPlots
@time pyplot()
relax(t=0.2) = (backend() == Plots.PyPlotBackend() && PyPlot.clf(); sleep(t))

rd(x, d=3) = round(x; digits=d)
mrd(t, d=3) = map(x->rd(x, d), t)

@time using Turing
#turnprogress(false);

"""
`get_posterior(chain::AbstractChains, vars::vars)` extracts the posterior sample of `vars` from `chain`.

Example:

```julia
chain = psample(model, NUTS(), 2000, 3)
posterior_sample = get_posterior_sample(chain, [:a, :b, :s])
```
"""
function get_posterior_sample(
        chain::Turing.Inference.AbstractMCMC.AbstractChains,
        vars::AbstractVector=get_vars(chain)
    )
    val = chain[vars].value
    W = vcat((val[:,:,m] for m in 1:size(val, 3))...)
    Array{typeof(W[end,end]), 2}(W)
end

"""
`get_vars(chain::AbstractChains)` returns the array of the names of the parameters in `chain`.
"""
function get_vars(chain::Turing.Inference.AbstractMCMC.AbstractChains)
    vars = chain.name_map[:parameters]
end

"""
`get_vars(model::AbstractModel)` returns the array of the names of the parameters in `model`.
"""
function get_vars(model::Turing.Inference.AbstractMCMC.AbstractModel)
    vi = DynamicPPL.VarInfo(linreg)
    vars = collect(keys(DynamicPPL.tonamedtuple(vi)))
end

function get_posterior_sample(
        chain::Turing.Inference.AbstractMCMC.AbstractChains,
        model::Turing.Inference.AbstractMCMC.AbstractModel
    )
    get_posterior_sample(chain, get_vars(model))
end

display(@doc get_posterior_sample)
display(@doc get_vars)

"""
`logpdf_array(logpdf_func, data, posterior_sample)` returns the array `lp` definef by

```julia
lp[k, l] = logpdf_func(data[k, :], posterior_sample[l, :])
```
"""
function logpdf_array(
        logpdf_func,
        data::AbstractArray,
        posterior_sample::AbstractArray
    )
    n = size(data, 1)
    L = size(posterior_sample, 1)
    lp = Array{eltype(posterior_sample), 2}(undef, n, L)
    for k in 1:n, l in 1:L
        @views lp[k, l] = logpdf_func(data[k, :], posterior_sample[l, :])
    end
    lp
end

"""
`logpdf_array(logpdf_func, data, chain::AbstractChains, vars=get_vars(chain))` returns the array `lp` definef by

```julia
posterior_sample = get_posterior_sample(chain, vars)
lp[k, l] = logpdf_func(data[k, :], posterior_sample[l, :])
```
"""
function logpdf_array(
        logpdf_func,
        data::AbstractArray,
        chain::Turing.Inference.AbstractMCMC.AbstractChains,
        vars::AbstractVector=get_vars(chain)
    )
    posterior_sample = get_posterior_sample(chain, vars)
    logpdf_array(logpdf_func, data, posterior_sample)
end

"""
`training_error(lp)` returns the training error calculated from the logpdf array `lp`.
"""
function training_error(lp::AbstractArray{R, 2}) where {R}
    n, L = size(lp)
    -2sum(logmeanexp(@view lp[k, :]) for k in 1:n)
end

"""
`functional_variance(lp)` returns the functional variance calculated from the logpdf array `lp`.
"""
function functional_variance(lp::AbstractArray{R, 2}) where {R}
    n, L = size(lp)
    2sum(var(@view lp[k,:]) for k in 1:n)
end

"""
`waic(lp)` returns the WAIC, the training error, and the functional variance calculated from the logpdf array `lp`.
"""
function waic(lp::AbstractArray{R, 2}) where {R}
    T = training_error(lp)
    V = functional_variance(lp)
    WAIC = T + V
    (WAIC=WAIC, T=T, V=V)
end

"""
`waic(logpdf_func, data, posterior_sample)` returns the WAIC, the training error, and the functional variance calculated from the arguments.
"""
function waic(
        logpdf_func,
        data::AbstractArray,
        posterior_sample::AbstractArray)
    lp = logpdf_array(logpdf_func, data, posterior_sample)
    waic(lp)
end

"""
`waic(logpdf_func, data, chain, vars=get_vars(chain))` returns the WAIC, the training error, and the functional variance calculated from the arguments.

Example:

```julia
@model ols(X, Y) = begin
    a ~ Flat()
    b ~ Flat()
    σ ~ FlatPos(0.0)
    for k in eachindex(X)
        Y[k] ~ Normal(a + b*X[k], σ)
    end
end

n = 20
X = 3rand(n)
Y = X .+ 1 .+ 0.3randn(n)
chain_ols = psample(ols(X, Y), NUTS(), 2000, 3)

function logpdf_ols(x, w)
    X, Y = x
    a, b, σ = w
    logpdf(Normal(a + b*X, σ), Y)
end
data_ols = hcat(X, Y)
vars_ols = [:a, :b, :σ]

@show waic(logpdf_ols, data_ols, chain_ols, vars_ols)
@show chain_ols
plot(chain_ols)
```
"""
function waic(
        logpdf_func,
        data::AbstractArray,
        chain::Turing.Inference.AbstractMCMC.AbstractChains,
        vars::AbstractVector=get_vars(chain)
    )
    posterior_sample = get_posterior_sample(chain, vars)
    waic(logpdf_func, data, posterior_sample)
end

display(@doc logpdf_array)
display(@doc training_error)
display(@doc functional_variance)
display(@doc waic)

"""
`loocv(lp)` returns the LOOCV from the logpdf array `lp`.
"""
function loocv(lp::AbstractArray{R, 2}) where {R}
    n, L = size(lp)
    2sum(logmeanexp(- @view lp[k, :]) for k in 1:n)
end

"""
`loocv(logpdf_func, data, posterior_sample)` returns the LOOCV calculated from the arguments.
"""
function loocv(
        logpdf_func,
        data::AbstractArray,
        posterior_sample::AbstractArray)
    lp = logpdf_array(logpdf_func, data, posterior_sample)
    loocv(lp)
end

"""
`loocv(logpdf_func, data, chain, vars=get_vars(chain))` returns the LOOCV calculated from the arguments.
"""
function loocv(
        logpdf_func,
        data::AbstractArray,
        chain::Turing.Inference.AbstractMCMC.AbstractChains,
        vars::AbstractVector=get_vars(chain)
    )
    posterior_sample = get_posterior_sample(chain, vars)
    loocv(logpdf_func, data, posterior_sample)
end

@doc loocv

"""
`make_pred_pdf(logpdf_func, posterior_sample)` makes the pdf of the predictive distribution from the arguments.
"""
function make_pred_pdf(logpdf_func, posterior_sample::AbstractArray)
    L = size(posterior_sample, 1)
    pred_pdf(x) = mean(exp(logpdf_func(x, @view(posterior_sample[l, :]))) for l in 1:L)
    pred_pdf
end

"""
`make_pred_pdf(logpdf_func, chain::AbstractChains, vars::AbstractVector)` makes the pdf of the predictive distribution from the arguments.
"""
function make_pred_pdf(
        logpdf_func,
        chain::Turing.Inference.AbstractMCMC.AbstractChains,
        vars::AbstractVector=get_vars(chain)
    )
    posterior_sample = get_posterior_sample(chain, vars)
    make_pred_pdf(logpdf_func, posterior_sample)
end

@doc make_pred_pdf

seed!(4649)
posterior_sample_test = [5, 0] .+ 1.2randn(2, 100)
logpdf_test(x, w) = logpdf(Gamma(exp(w[1]), exp(w[2])), x[1])
pred_pdf_test = make_pred_pdf(logpdf_test, posterior_sample_test)
@show pred_pdf_test([5.0])
@show pred_pdf_test(5.0)

plot(size=(400, 250))
plot!(pred_pdf_test, 0, 50; label="pred_pdf_test")

seed!(20200323);

@model ols(X, Y) = begin
    a ~ Flat()
    b ~ Flat()
    σ ~ FlatPos(0.0)
    for k in eachindex(X)
        Y[k] ~ Normal(a + b*X[k], σ)
    end
end

n = 20
X = 3rand(n)
Y = X .+ 1 .+ 0.3randn(n)
chain_ols = psample(ols(X, Y), NUTS(), 2000, 3)

function logpdf_ols(x, w)
    X, Y = x
    a, b, σ = w
    logpdf(Normal(a + b*X, σ), Y)
end
data_ols = hcat(X, Y)
vars_ols = [:a, :b, :σ]
waic_ols = waic(logpdf_ols, data_ols, chain_ols, vars_ols)
loocv_ols = loocv(logpdf_ols, data_ols, chain_ols, vars_ols)

@show waic_ols
@show loocv_ols
@show chain_ols
plot(chain_ols; lw=0.5)

posterior_sample_ols = get_posterior_sample(chain_ols, vars_ols)
a, b, σ = mapslices(median, posterior_sample_ols, dims=1)
@show a, b, σ
lp = logpdf_array(logpdf_ols, data_ols, posterior_sample_ols)
@show T = training_error(lp)
@show V = functional_variance(lp)
@show WAIC = T + V
@show LOOCV = loocv(lp)
relax()
display(waic(lp))

X, Y = data_ols[:,1], data_ols[:,2]
Phi = hcat(ones(length(X)), X)
a, b = Phi\Y
σ² = mean((Y - Phi*[a,b]).^2)
@show a, b, √σ²
T = -2sum(logpdf(Normal(Phi[k,:]'*[a,b], √σ²), Y[k]) for k in 1:length(X))
nparams = 3
@show T
@show 2nparams
@show AIC = T + 2nparams
relax()

(AIC=AIC, T=T, V=2nparams) |> display

plot(size=(400, 250))
scatter!(X, Y; label="sample")
plot!(x -> x + 1, -0.2, 3.2; label="x+1")
plot!(x -> x + 1 + 2*0.3, -0.2, 3.2; label="±2σ", color=:black, ls=:dot)
plot!(x -> x + 1 - 2*0.3, -0.2, 3.2; label="", color=:black, ls=:dot)

pred_pdf_ols = make_pred_pdf(logpdf_ols, posterior_sample_ols)

x = range(-0.5, 3.5; length=100)
y = range(0.5, 4.5; length=100)
@time z = ((x, y) -> pred_pdf_ols((x, y))).(x', y)
P1 = plot(size=(400, 400), legend=false)
plot!(xtick=-10:10, ytick=-10:10)
plot!(xlim=extrema(x), ylim=extrema(y))
heatmap!(x, y, z)
scatter!(X, Y; label="sample", color=:cyan)

x = range(-7, 10; length=100)
y = range(-6, 11; length=100)
@time z = ((x, y) -> pred_pdf_ols((x, y))).(x', y)
P2 = plot(size=(400, 400), legend=false)
plot!(xtick=-10:2:10, ytick=-10:2:10)
plot!(xlim=extrema(x), ylim=extrema(y))
heatmap!(x, y, z)
scatter!(X, Y; label="sample", color=:cyan)

plot(P1, P2; size=(800, 400), layout=(1, 2))

normal_dist(a, b) = Normal(a, exp(b))
logpdf_normal(x, w) = logpdf(normal_dist(w[1], w[2]), x[1])
@doc Normal

gamma_dist(a, b) = Gamma(exp(a), exp(b))
logpdf_gamma(x, w) = logpdf(gamma_dist(w[1], w[2]), x[1])
@doc Gamma

inversegamma_dist(a, b) = InverseGamma(exp(a), exp(b))
logpdf_inversegamma(x, w) = logpdf(inversegamma_dist(w[1], w[2]), x[1])
@doc InverseGamma

lognormal_dist(a, b) = LogNormal(a, exp(b))
logpdf_lognormal(x, w) = logpdf(lognormal_dist(w[1], w[2]), x[1])
@doc LogNormal

weibull_dist(a, b) = Weibull(exp(a), exp(b))
logpdf_weibull(x, w) = logpdf(weibull_dist(w[1], w[2]), x[1])
@doc Weibull

@show dist_true = weibull_dist(1.5, 1/10)

seed!(4649)
n = 1000
X = rand(dist_true, n)
plot(size=(500, 350), legendfontsize=10)
histogram!(X; norm=true, alpha=0.3, label="sample")
plot!(dist_true, 0, 2; label="dist_true")

@model testmodel(X, dist_func) = begin
    a ~ Flat()
    b ~ Flat()
    X .~ dist_func(a, b)
end

chain_normal = psample(testmodel(X, normal_dist), NUTS(), 2000, 3)
waic_normal = waic(logpdf_normal, X, chain_normal)
@show waic_normal
loocv_normal = loocv(logpdf_normal, X, chain_normal)
@show loocv_normal
println()
@show chain_normal
plot(chain_normal; lw=0.5)

chain_gamma = psample(testmodel(X, gamma_dist), NUTS(), 2000, 3)
waic_gamma = waic(logpdf_gamma, X, chain_gamma)
@show waic_gamma
loocv_gamma = loocv(logpdf_gamma, X, chain_gamma)
@show loocv_gamma
println()
@show chain_gamma
plot(chain_gamma; lw=0.5)

chain_inversegamma = psample(testmodel(X, inversegamma_dist), NUTS(), 2000, 3)
waic_inversegamma = waic(logpdf_inversegamma, X, chain_inversegamma)
@show waic_inversegamma
loocv_inversegamma = loocv(logpdf_inversegamma, X, chain_inversegamma)
@show loocv_inversegamma
println()
@show chain_inversegamma
plot(chain_inversegamma; lw=0.5)

chain_lognormal = psample(testmodel(X, lognormal_dist), NUTS(), 2000, 3)
waic_lognormal = waic(logpdf_lognormal, X, chain_lognormal)
@show waic_lognormal
loocv_lognormal = loocv(logpdf_lognormal, X, chain_lognormal)
@show loocv_lognormal
println()
@show chain_lognormal
plot(chain_lognormal; lw=0.5)

chain_weibull = psample(testmodel(X, weibull_dist), NUTS(), 2000, 3)
waic_weibull = waic(logpdf_weibull, X, chain_weibull)
@show waic_weibull
loocv_weibull = loocv(logpdf_weibull, X, chain_weibull)
@show loocv_weibull
println()
@show chain_weibull
plot(chain_weibull; lw=0.5)

@show waic_normal |> mrd
@show waic_gamma |> mrd
@show waic_inversegamma |> mrd
@show waic_lognormal |> mrd
@show waic_weibull |> mrd
println()
@show loocv_normal |> mrd
@show loocv_gamma |> mrd
@show loocv_inversegamma |> mrd
@show loocv_lognormal |> mrd
@show loocv_weibull |> mrd
;

pred_normal = make_pred_pdf(logpdf_normal, chain_normal)
pred_gamma = make_pred_pdf(logpdf_gamma, chain_gamma)
pred_inversegamma = make_pred_pdf(logpdf_inversegamma, chain_inversegamma)
pred_lognormal = make_pred_pdf(logpdf_lognormal, chain_lognormal)
pred_weibull = make_pred_pdf(logpdf_weibull, chain_weibull)

plot(legendfontsize=10)
histogram!(X; norm=true, alpha=0.2, label="sample")
plot!(x -> pred_weibull([x]), 0, 2.5; label="weibull", ls=:dash, lw=1.5)
plot!(x -> pred_normal([x]), 0, 2.5; label="normal", ls=:dashdot, lw=1.2)
plot!(x -> pred_gamma([x]), 0, 2.5; label="gamma", ls=:dot, lw=1.2)
plot!(x -> pred_inversegamma([x]), 0, 2.5; label="inversegamma", ls=:dot, lw=1.2)
plot!(x -> pred_lognormal([x]), 0, 2.5; label="lognormal", ls=:dot, lw=1.2)
plot!(dist_true, 0, 2.5; label="dist_true", color=:blue, alpha=0.7, lw=0.8)

tStartExposure = [-2, -18, -18, 10, 3, 8, 12, 13, -18, -18, 8, 10, -11, -18, -18, -18, 12, -18, -18, -18, -18, -18, -18, 11, 11, -18, -18, -18, -18, -18, 12, 12, 15, 15, -18, -18, -18, -18, 19, -18, -18, -18, -18, -18, -18, -18, 18, -18, -18, 6, -18, -18, -18, 9, -18, 11, -18, -18, -18, -18, -18, -18, -18, -18, -18, -18, 19, -18, -18, -18, -18, -18, -18, -18, 13, 13, -18, -18, -18, -18, -18, -18, 13, -18, -18, -18, -18, -18]

tEndExposure = [3, 12, 3, 11, 4, 16, 16, 16, 15, 15, 16, 11, 9, 2, 12, 17, 15, 17, 18, 13, 16, 11, 18, 14, 18, 20, 12, 10, 17, 15, 15, 14, 17, 20, 18, 13, 23, 23, 22, 20, 21, 21, 21, 15, 21, 17, 20, 18, 20, 7, 20, 20, 20, 20, 18, 22, 22, 18, 18, 18, 9, 20, 18, 22, 23, 19, 19, 22, 22, 13, 22, 22, 23, 23, 17, 17, 22, 18, 18, 22, 18, 20, 15, 20, 19, 23, 22, 22]

tSymptomOnset = [3, 15, 4, 14, 9, 16, 16, 16, 16, 16, 16, 14, 10, 10, 14, 20, 19, 19, 20, 20, 17, 13, 19, 19, 20, 21, 18, 18, 18, 16, 20, 16, 20, 20, 19, 17, 24, 24, 23, 21, 23, 23, 23, 21, 22, 24, 24, 19, 21, 14, 23, 23, 21, 20, 21, 22, 23, 19, 23, 21, 13, 22, 24, 25, 25, 25, 25, 24, 24, 21, 23, 23, 24, 25, 18, 18, 23, 22, 22, 24, 24, 25, 22, 25, 25, 24, 25, 25]

@show d = length(tStartExposure)

hcat(tStartExposure, tEndExposure, tSymptomOnset)'

data_ip = hcat(tStartExposure', tEndExposure', tSymptomOnset')
size(data_ip)

@model incubation_period(
    dist_func,
    tStartExposure,
    tEndExposure,
    tSymptomOnset,
    ::Type{VT}=Vector{Float64}
) where {VT} = begin
    n = length(tStartExposure)
    uE = VT(undef, n)
    a ~ Flat()
    b ~ Flat()
    uE .~ Uniform()
    for k in 1:n
        tE = tStartExposure[k] + uE[k]*(tEndExposure[k] - tStartExposure[k])
        tSymptomOnset[k] ~ LocationScale(tE, 1.0, dist_func(a, b))
    end
end

model_ip(dist_func) = incubation_period(
    dist_func,
    tStartExposure,
    tEndExposure,
    tSymptomOnset
)

function logpdf_ip(x, w; dist_func=dist_func, d=length(tStartExposure))
    a = w[1]
    b = w[2]
    uE = @view w[3:end]
    sum(logpdf(dist_func(a, b), x[2d+k] - (x[k] + uE[k]*(x[d+k]-x[k]))) for k in 1:d)
end

L = 1000 + 1000
nchains = 3

@time chain_ip_gamma = psample(model_ip(gamma_dist), NUTS(), L, nchains);

#chain_ip_gamma
chain_ip_gamma[["a"; "b"; ["uE[$i]" for i in 1:10:d]]]

#@time plot(chain_ip_gamma[2:2:end]; lw=0.5, alpha=0.8)
@time plot(chain_ip_gamma[["a"; "b"; ["uE[$i]" for i in 1:10:d]]][2:2:end]; lw=0.5, alpha=0.8)

dist_func = gamma_dist
waic_ip_gamma = waic(logpdf_ip, data_ip, chain_ip_gamma)
loocv_ip_gamma = loocv(logpdf_ip, data_ip, chain_ip_gamma)
@show waic_ip_gamma
@show loocv_ip_gamma;

@time chain_ip_lognormal = psample(model_ip(lognormal_dist), NUTS(), L, nchains);

#chain_ip_lognormal
chain_ip_lognormal[["a"; "b"; ["uE[$i]" for i in 1:10:d]]]

#@time plot(chain_ip_lognormal[2:2:end]; lw=0.5, alpha=0.8)
@time plot(chain_ip_lognormal[["a"; "b"; ["uE[$i]" for i in 1:10:d]]][2:2:end]; lw=0.5, alpha=0.8)

dist_func = lognormal_dist
waic_ip_lognormal = waic(logpdf_ip, data_ip, chain_ip_lognormal)
loocv_ip_lognormal = loocv(logpdf_ip, data_ip, chain_ip_lognormal)
@show waic_ip_lognormal
@show loocv_ip_lognormal;

@time chain_ip_weibull = psample(model_ip(weibull_dist), NUTS(), L, nchains);

#chain_ip_weibull
chain_ip_weibull[["a"; "b"; ["uE[$i]" for i in 1:10:d]]]

#@time plot(chain_ip_weibull[2:2:end]; lw=0.5, alpha=0.8)
@time plot(chain_ip_weibull[["a"; "b"; ["uE[$i]" for i in 1:10:d]]][2:2:end]; lw=0.5, alpha=0.8)

dist_func = weibull_dist
waic_ip_weibull = waic(logpdf_ip, data_ip, chain_ip_weibull)
loocv_ip_weibull = loocv(logpdf_ip, data_ip, chain_ip_weibull)
@show waic_ip_weibull
@show loocv_ip_weibull;

@show waic_ip_gamma |> mrd
@show waic_ip_lognormal |> mrd
@show waic_ip_weibull |> mrd
println()
@show loocv_ip_gamma |> mrd
@show loocv_ip_lognormal |> mrd
@show loocv_ip_weibull |> mrd
;

data_ipr = hcat(tStartExposure, tEndExposure, tSymptomOnset)

function logpdf_ipr(x, w; dist_func=dist_func)
    log(quadgk(t -> pdf(dist_func(w[1], w[2]), x[3] - (x[1] + t*(x[2] - x[1]))), 0.0, 1.0)[1])
end

vars_ipr = [:a, :b]
@show vars_ipr;

dist_func = gamma_dist
@time lp_ipr_gamma = logpdf_array(logpdf_ipr, data_ipr, chain_ip_gamma, vars_ipr)
waic_ipr_gamma = waic(lp_ipr_gamma)
loocv_ipr_gamma = loocv(lp_ipr_gamma)
@show waic_ipr_gamma
@show loocv_ipr_gamma;

dist_func = lognormal_dist
@time lp_ipr_lognormal = logpdf_array(logpdf_ipr, data_ipr, chain_ip_lognormal, vars_ipr)
waic_ipr_lognormal = waic(lp_ipr_lognormal)
loocv_ipr_lognormal = loocv(lp_ipr_lognormal)
@show waic_ipr_lognormal
@show loocv_ipr_lognormal;

dist_func = weibull_dist
@time lp_ipr_weibull = logpdf_array(logpdf_ipr, data_ipr, chain_ip_weibull, vars_ipr)
waic_ipr_weibull = waic(lp_ipr_weibull)
loocv_ipr_weibull = loocv(lp_ipr_weibull)
@show waic_ipr_weibull
@show loocv_ipr_weibull;

@show waic_ipr_gamma |> mrd
@show waic_ipr_lognormal |> mrd
@show waic_ipr_weibull |> mrd
println()
@show loocv_ipr_gamma |> mrd
@show loocv_ipr_lognormal |> mrd
@show loocv_ipr_weibull |> mrd
;

w_gamma = get_posterior_sample(chain_ip_gamma, vars_ipr)
function pdf_ipr_gamma(t)
    mean(pdf(gamma_dist(w_gamma[l,1], w_gamma[l,2]), t) for l in 1:size(w_gamma, 1))
end

w_lognormal = get_posterior_sample(chain_ip_lognormal, vars_ipr)
function pdf_ipr_lognormal(t)
    mean(pdf(lognormal_dist(w_lognormal[l,1], w_lognormal[l,2]), t) for l in 1:size(w_lognormal, 1))
end

w_weibull = get_posterior_sample(chain_ip_weibull, vars_ipr)
function pdf_ipr_weibull(t)
    mean(pdf(weibull_dist(w_weibull[l,1], w_weibull[l,2]), t) for l in 1:size(w_weibull, 1))
end

t = range(0, 25, length=400)
plot(size=(500, 300), legendfontsize=10)
plot!(xtick=0:25)
plot!(t, pdf_ipr_gamma.(t); label="gamma", ls=:solid)
plot!(t, pdf_ipr_lognormal.(t); label="lognormal", ls=:dash)
plot!(t, pdf_ipr_weibull.(t); label="weibull", ls=:dashdot)

data_ip2 = hcat(tStartExposure, tEndExposure, tSymptomOnset)

function logpdf_ip2(x, w; dist_func=dist_func)
    z = if x[1] == x[2]
        pdf(dist_func(w[1], w[2]), x[3] - x[1])
    else
        (cdf(dist_func(w[1], w[2]), x[3] - x[2]) - cdf(dist_func(w[1], w[2]), x[3] - x[1]))/(-x[2] + x[1])
    end
    log(z)
end

vars_ip2 = [:a, :b]
@show vars_ip2;

k, l = 67, 1234

println((k = k, data_ip2_k_1 = data_ip2[k,1], data_ip2_k_2 = data_ip2[k,2]))

dist_func = gamma_dist
ps_gamma = get_posterior_sample(chain_ip_gamma, vars_ip2)
@show logpdf_ipr(data_ip2[k,:], ps_gamma[l,:])
@show logpdf_ip2(data_ip2[k,:], ps_gamma[l,:])

dist_func = lognormal_dist
ps_gamma = get_posterior_sample(chain_ip_lognormal, vars_ip2)
@show logpdf_ipr(data_ip2[k,:], ps_gamma[l,:])
@show logpdf_ip2(data_ip2[k,:], ps_gamma[l,:])

dist_func = weibull_dist
ps_gamma = get_posterior_sample(chain_ip_weibull, vars_ip2)
@show logpdf_ipr(data_ip2[k,:], ps_gamma[l,:])
@show logpdf_ip2(data_ip2[k,:], ps_gamma[l,:])
;

k, l = 10, 1234

println((k = k, data_ip2_k_1 = data_ip2[k,1], data_ip2_k_2 = data_ip2[k,2]))

dist_func = gamma_dist
ps_gamma = get_posterior_sample(chain_ip_gamma, vars_ip2)
@show logpdf_ipr(data_ip2[k,:], ps_gamma[l,:])
@show logpdf_ip2(data_ip2[k,:], ps_gamma[l,:])

dist_func = lognormal_dist
ps_lognormal = get_posterior_sample(chain_ip_lognormal, vars_ip2)
@show logpdf_ipr(data_ip2[k,:], ps_lognormal[l,:])
@show logpdf_ip2(data_ip2[k,:], ps_lognormal[l,:])

dist_func = weibull_dist
ps_weibull = get_posterior_sample(chain_ip_weibull, vars_ip2)
@show logpdf_ipr(data_ip2[k,:], ps_weibull[l,:])
@show logpdf_ip2(data_ip2[k,:], ps_weibull[l,:])
;

dist_func = gamma_dist
@time lp_ip2_gamma = logpdf_array(logpdf_ip2, data_ip2, chain_ip_gamma, vars_ip2)
waic_ip2_gamma = waic(lp_ip2_gamma)
loocv_ip2_gamma = loocv(lp_ip2_gamma)
@show waic_ip2_gamma
@show loocv_ip2_gamma;

dist_func = lognormal_dist
@time lp_ip2_lognormal = logpdf_array(logpdf_ip2, data_ip2, chain_ip_lognormal, vars_ip2)
waic_ip2_lognormal = waic(lp_ip2_lognormal)
loocv_ip2_lognormal = loocv(lp_ip2_lognormal)
@show waic_ip2_lognormal
@show loocv_ip2_lognormal;

dist_func = weibull_dist
@time lp_ip2_weibull = logpdf_array(logpdf_ip2, data_ip2, chain_ip_weibull, vars_ip2)
waic_ip2_weibull = waic(lp_ip2_weibull)
loocv_ip2_weibull = loocv(lp_ip2_weibull)
@show waic_ip2_weibull
@show loocv_ip2_weibull;

@show waic_ip2_gamma |> mrd
@show waic_ip2_lognormal |> mrd
@show waic_ip2_weibull |> mrd
println()
@show loocv_ip2_gamma |> mrd
@show loocv_ip2_lognormal |> mrd
@show loocv_ip2_weibull |> mrd
;

@show waic_ipr_gamma |> mrd
@show waic_ipr_lognormal |> mrd
@show waic_ipr_weibull |> mrd
println()
@show loocv_ipr_gamma |> mrd
@show loocv_ipr_lognormal |> mrd
@show loocv_ipr_weibull |> mrd
;