const univariate_path = "CovidAlertVaccinationModel/plots/univariate/" 
const bivariate_path = "CovidAlertVaccinationModel/plots/univariate/" 
function univarate_test(variable, variable_range,samples; progmeter = nothing)
    default_parameters = get_app_parameters() 
    parameter_range_list = [merge(default_parameters,NamedTuple{(variable,)}((value,))) for value in variable_range]
    solve_fn(p) = mean_solve(samples, p,DebugRecorder;progmeter)[1]

    univariate_outlist = ThreadsX.map(solve_fn, parameter_range_list)
    
    p = plot_model(variable,parameter_range_list,univariate_outlist,default_parameters.infection_introduction_day,default_parameters.immunization_begin_day)
    return p
end

if !ispath(univariate_path)
    mkdir(univariate_path)
end
function univariate_simulations()
    len = 5
    samples = 10
    univarate_test_list = (
        # (:I_0_fraction, range(0.0, 0.05; length = len)), 
        # (:base_transmission_probability, range(0.0002, 0.002; length = len)),
        # (:recovery_rate, range(0.1, 0.5; length = len)),
        # (:immunization_loss_prob, range(0.00, 0.05; length = len)),
        # (:π_base, range(-4.5, -3.5;  length = len)),
        (:η, range(0.0, 0.005; length = len)),
        # (:κ, range(0.5, 1.5; length = len)),
        # (:ω, range(0.0, 0.01; length = len)),
        (:ω_en, range(0.0, 0.001; length = len)),
        # (:γ, range(0.0, 0.5; length = len)),
        # (:ξ, range(1, 15; length = len)),
        (:notification_parameter, range(0.000, 0.002; length = len)),
        (:app_user_fraction, range(0.05, 0.8; length = len)),
        (:notification_threshold, (1:len)),
        # (:immunization_delay, [7,10,14,20]),
    )


    numsim = sum(map(t -> length(t[2]), univarate_test_list))
    display(numsim)
    progmeter = Progress(numsim*samples)
    plt_list = ThreadsX.map(univarate_test_list) do ur
        out = univarate_test(ur...,samples;progmeter)
        return out
    end
    for ((varname,_),p) in zip(univarate_test_list,plt_list)
        savefig(p,"$univariate_path/$varname.pdf")
    end
end

using AxisKeys 
function multivariate_simulations()
    len = 10
    samples = 25
    app_simulations = (
        (:η, range(0.0, 0.005; length = len)),
        (:ω_en, range(0.0, 0.05; length = len)),
        # (:notification_threshold, (1:len)),
    )
    run_multivariate_sims(app_simulations,samples)


    # for ((varname,_),p) in zip(univarate_test_list,plt_list)
    #     savefig(p,"$univariate_path/$varname.pdf")
    # end
end
using ProgressMeter
function run_multivariate_sims(sims,samples)
    varnames, sim_ranges = zip(sims...)

    without_app_future = @spawn mean_solve(samples,get_parameters(),HeatmapRecorder)

    default_parameters = get_app_parameters() 
    simvars = Iterators.product(sim_ranges...)
    progmeter = Progress(length(simvars)*samples)
    app_output = ThreadsX.map(simvars) do vars
        vars_with_names = NamedTuple{varnames}(vars)
        parameters = merge(default_parameters,vars_with_names)
        out,_ = mean_solve(samples, parameters,HeatmapRecorder;progmeter) 
  
        return out
    end
    display(length(simvars))
    fname = join(string.(varnames),"_")
    keyed_output = KeyedArray(app_output;NamedTuple{varnames}(sim_ranges)...)
    path = joinpath(PACKAGE_FOLDER,"abm_output","$fname.dat")
    serialize(path,(fetch(without_app_future)[1],keyed_output))
    return fname
end
using ColorSchemes
using LaTeXStrings
function plot_parameter_plane(input_fname)
    path = joinpath(PACKAGE_FOLDER,"abm_output","$input_fname.dat")
    output_no_app, output = deserialize(path)
    var_ranges = axiskeys(output)
    vars = (L"\eta",L"\omega_{en}")



    mean_final_size(p) = mean(reduce(merge!,p.final_size_by_age))
    base_outcome = mean_final_size(output_no_app)
    final_size_map = map(x-> (mean_final_size(x) - base_outcome),output)
    
    mean_weighted_degree_change(p) = mean(p.avg_weighted_degree_of_vaccinators)-mean(p.avg_weighted_degree)
    weighted_degree_map = map(mean_weighted_degree_change,output)
    cs = cgrad([:blue,:orange])

    datamaps = [weighted_degree_map, final_size_map]
    fnames = ["wdg_change.pdf", "final_size_change.pdf"]
    titles = [
        "Average w. deg. of vaccinators minus average w. deg.",
        "Change in average final size of infection outbreak",
    ]
    for (fname,title,datamap) in zip(fnames,titles,datamaps)
        p = heatmap(var_ranges[1],var_ranges[2],datamap; title, xlabel = vars[1], ylabel = vars[2], seriescolor=cs, size = (800,600))
        savefig(p,joinpath(PACKAGE_FOLDER,"plots","app_heatmaps","$fname"))
    end
end