#####################
# comparing metrics
#####################

# ELBO plot for ErgFlow
# this function doesn't support choosing nBurns----always set to be 0
function ELBO_plot(
    o::MixFlow.HamFlow;
    elbo_svi,
    μ,
    D,
    eps,
    Ns,
    elbo_size=1000,
    fig_dir="figure/",
    fig_name="elbo_tune.png",
    res_dir="result/",
    res_name="elbo_dat.jld",
    kwargs...,
)
    # # generate folder 
    if !isdir(fig_dir)
        mkdir(fig_dir)
    end
    Els = zeros(size(eps, 1), size(Ns, 1) + 1)
    Els[:, 1] .= elbo_svi
    for i in 1:size(eps, 1)
        a = MixFlow.HF_params(eps[i], μ, D)
        Els[i, 2:end] .= MixFlow.ELBO_sweep(
            o, a, MixFlow.ref_coord, MixFlow.inv_ref_coord, Ns; elbo_size=elbo_size
        )
        println("ϵ = $(eps[i]), n_mcmc = $(Ns) done")
    end
    JLD.save(joinpath(res_dir, res_name), "elbos", Els, "eps", eps, "Ns", Ns)
    Labels = Matrix{String}(undef, 1, size(eps, 1))
    Labels[1, :] .= ["ϵ=$e" for e in eps]
    p = plot(
        reduce(vcat, [[0], Ns]),
        Els';
        lw=3,
        labels=Labels,
        legend=:outertopright,
        ylabel="ELBO",
        xlabel="#Refreshment",
        kwargs...,
    )
    return savefig(p, joinpath(fig_dir, fig_name))
end