parameter_optimization.jl

using KissABC

const vaccination_data = [0.0,0.043,0.385,0.424,0.115,0.03,0.005] #by month starting in august
const ymo_vac = [0.255,0.278,0.602]
const ymo_attack_rate = [10.376,5.636,7.2]./100


function solve_w_parameters(default_p, p_names, new_p_list)

    new_params = merge(default_p, NamedTuple{p_names}(ntuple(i -> new_p_list[i],length(p_names))))
    out = DebugRecorder(0,default_p.sim_length)
    model = abm(new_params,out)
    return out, model
end

function fit_all_parameters(p_tuple)
    p_names = (:ω,:β_y,:β_m,:β_o,:π_base_y,:π_base_m,:π_base_o,:α_y,:α_m,:α_o)
    priors = Factored(
        Uniform(0.0,0.01),
        Uniform(0.0,0.2),
        Uniform(0.0,0.1),
        Uniform(0.0,1.0), 
        Uniform(-5.0,3.0),
        Uniform(-5.0,3.0),
        Uniform(-5.0,3.0),
        TriangularDist(23.0, 44.0, 34.0),
        TriangularDist(22.0, 54.0, 40.0),
        TriangularDist(9.0, 59.0, 39.0)   
    )
    #simulation begins in july
    #60 days for opinion dynamics to stabilize, then immunization begins in september,
    #infection introduced at beginning of december
    sim_length = 300
    p_tuple_adjust = merge(p_tuple,
        (
            sim_length = sim_length,
            I_0_fraction = 0.005,
            immunization_begin_day = 60, 
            infection_introduction_day = 180,
            immunizing = true,
        )
    )
    function cost(p)
        output,model = solve_w_parameters(p_tuple_adjust, p_names,p)
        target_ymo_vac = ymo_vac .* sum(vaccination_data[1:end]) .* length.(model.index_vectors)
        ymo_vaccination_ts = output.daily_immunized_by_age

        normalize_by_pop(v) = v./length.(model.index_vectors)
        total_postinf_vaccination = sum.(eachrow(ymo_vaccination_ts[:,180:end]))

        final_size = sum.(eachrow(output.daily_unvac_cases_by_age))

        target_final_size = ymo_attack_rate .* length.(model.index_vectors)
        target_ymo_vac = ymo_vac .* sum(vaccination_data[1:4]) .* length.(model.index_vectors)

        ymo_vaccination_ts = output.daily_immunized_by_age
        total_preinfection_vaccination = sum.(eachrow(ymo_vaccination_ts))
    
        # return sum((normalize_by_pop(total_preinfection_vaccination .- target_ymo_vac)).^2)  
        # + sum(normalize_by_pop((total_postinf_vaccination .- target_ymo_vac)).^2)  
        # + sum(normalize_by_pop((final_size .- target_final_size)).^2)
        display(final_size)
        display( sum.(eachrow(ymo_vaccination_ts[:,1:end])))
        display( length.(model.index_vectors))
        return sum(normalize_by_pop((final_size .- target_final_size)).^2)
    end
    display(cost((0.0,0.0005,0.0005,1.0,-3.0,-3.0,-0.2,0.25,0.25,0.0)))
    # out = smc(priors,cost; verbose = true, nparticles = 1000, parallel = true)#this one has better NaN handling
    return NamedTuple{p_names}(ntuple(i -> out.P[i].particles,length(p_names)))
end
function fit_distribution_parameters(p_tuple)
    p_names = (:ω,:β_y,:π_base_y,:π_base_m,:π_base_o,:α_y,:α_m,:α_o,:O_distribution_shift)
    priors = Factored(
        Uniform(0.0,0.01),
        Uniform(0.0,0.5),  
        Uniform(-5.0,3.0),
        Uniform(-5.0,3.0),
        Uniform(-5.0,3.0),
        TriangularDist(23.0, 44.0, 34.0),
        TriangularDist(22.0, 54.0, 40.0),
        TriangularDist(9.0, 59.0, 39.0),   
        Uniform(0.0,0.2),   
    )
    #simulation begins in july
    #60 days for opinion dynamics to stabilize, then immunization begins in september,
    #infection introduced at beginning of december
    sim_length = 300
    p_tuple_adjust = merge(p_tuple,
        (
            sim_length = sim_length,
            I_0_fraction = 0.005,
            immunization_begin_day = 60, 
            infection_introduction_day = 180,
            immunizing = true,
        )
    )
    function cost(p)
        p_mergedbeta = merge(p_tuple_adjust, (β_m = p[2],β_o = p[2], )) #assume one beta
        output,model = solve_w_parameters(p_mergedbeta, p_names,p)
        target_ymo_vac = ymo_vac .* sum(vaccination_data[1:end]) .* length.(model.index_vectors)
        ymo_vaccination_ts = output.daily_immunized_by_age

        normalize_by_pop(v) = v./length.(model.index_vectors)
        total_postinf_vaccination = sum.(eachrow(ymo_vaccination_ts[:,180:end]))

        final_size = sum.(eachrow(output.daily_unvac_cases_by_age))

        target_final_size = ymo_attack_rate .* length.(model.index_vectors)
        target_ymo_vac = ymo_vac .* sum(vaccination_data[1:4]) .* length.(model.index_vectors)

        ymo_vaccination_ts = output.daily_immunized_by_age
        total_preinfection_vaccination = sum.(eachrow(ymo_vaccination_ts))
    
        return sum((normalize_by_pop(total_preinfection_vaccination .- target_ymo_vac)).^2)  
        + sum(normalize_by_pop((total_postinf_vaccination .- target_ymo_vac)).^2)  
        + sum(normalize_by_pop((final_size .- target_final_size)).^2)
        # display(final_size)
        # display( sum.(eachrow(ymo_vaccination_ts[:,1:end])))
        # display( length.(model.index_vectors))
        # return sum(normalize_by_pop((final_size .- target_final_size)).^2)
    end
    # display(cost((0.0,0.0005,-3.0,-3.0,-0.2,0.25,0.25,0.0,0.1)))
    #
    out = smc(priors,cost; verbose = true, nparticles = 300, parallel = true)
    return NamedTuple{p_names}(ntuple(i -> out.P[i].particles,length(p_names)))
end
function fit_parameters(default_parameters)
    # pre_inf_behaviour_parameters_path =joinpath(PACKAGE_FOLDER,"abm_parameter_fits","pre_inf_behaviour_parameters.dat")
    # post_inf_behaviour_parameters_path = joinpath(PACKAGE_FOLDER,"abm_parameter_fits","post_inf_behaviour_parameters.dat")
    fit_all_parameters_path = joinpath(PACKAGE_FOLDER,"abm_parameter_fits","fit_all_parameters.dat")
    fit_dist_parameters_path = joinpath(PACKAGE_FOLDER,"abm_parameter_fits","fit_dist_parameters.dat")


    output = fit_distribution_parameters(default_parameters)
    # serialize(fit_all_parameters_path,output)

    
    # fitted_parameter_tuple = deserialize(fit_all_parameters_path)
    # fitted_sol,avg_populations = plot_fitting_posteriors("post_inf_fitting",fitted_parameter_tuple,default_parameters)

    # final_vac = sum.(eachrow(mean.(fitted_sol.daily_immunized_by_age)))
    # final_size = sum.(eachrow(mean.(fitted_sol.daily_unvac_cases_by_age)))
    # display(final_vac./avg_populations)
    # display(final_size./avg_populations)
    return fitted_sol
end 

function plot_fitting_posteriors(fname,particles_tuple,parameters)
    # p_adjust = merge(parameters,(β_y = p.β_y*0.1, β_m = p.β_m*0.1 ,β_o = p.β_o)) 
    stat_recorder = DebugRecorder(Variance(), parameters.sim_length)
    output_recorder = DebugRecorder(parameters.sim_length)
    avg_populations = [0.0,0.0,0.0]
    names = keys(particles_tuple)
    samples = collect(zip(particles_tuple...))
    for p_set in sample(samples,10)
        p_set_named = NamedTuple{names}(p_set)
        sol = abm(merge(parameters,p_set_named),output_recorder)
        avg_populations .+= length.(sol.index_vectors)
        fit!(stat_recorder,output_recorder)
    end
    avg_populations ./= 10
    p = plot_model(nothing,[nothing],[stat_recorder],parameters.infection_introduction_day,parameters.immunization_begin_day)
    savefig(p, "$fname.pdf")
    
    plts = [plot() for i in 1:length(particles_tuple)]
    for (plt,(k,v)) in zip(plts,pairs(particles_tuple))
        hist = StatsBase.fit(Histogram,v; nbins = 30)
        plot!(plt,hist;legend = false,xlabel = k)
    end
    p = plot(plts...; size = (1400,800),bottommargin = 5Plots.mm)
    savefig(p,"$(fname)_posteriors.pdf")
    return stat_recorder,avg_populations
end