include(joinpath(@__DIR__, "run_nf.jl"))

if !isdir(joinpath(@__DIR__, "csvs/"))
    mkpath(joinpath(@__DIR__, "csvs/"))
end

n_sample_eval = 1024

for name in ["SparseRegression", "TReg", "Brownian", "LGCP"]
    for flowtype in ["real_nvp", "neural_spline_flow"]
        for nlayer in [3, 5]
            for niters in [50000]
                for bs in [32]
                    for lr in [1e-3, 1e-4, 1e-5]
                        for seed in 1:5

                            # run simulation
                            df = try
                                run_norm_flow(
                                    seed, name, flowtype, nlayer, lr;
                                    batchsize=bs, niters=niters, show_progress=true,
                                    nsample_eval=n_sample_eval,
                                    save_jld=true
                                )

                            catch

                                DataFrame(
                                    time=NaN,
                                    elbo=NaN,
                                    logZ=NaN,
                                    ess=NaN,
                                )

                            end

                            CSV.write(
                                joinpath(@__DIR__, "csvs/normflow_$(name)_$(flowtype)_$(nlayer)_$(lr)_$(seed).csv"),
                                df,
                            )
                        end
                    end
                end
            end
        end
    end
end
