ENV["PYTHON"]=""; import Pkg; Pkg.activate(@__DIR__); Pkg.instantiate()
using GpABC, OrdinaryDiffEq, Distributions, Distances, Plots
pyplot()

times = [0.0, 0.6, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0]
data = [[ 20.     ,  10.     ,   0.     ],
        [  0.12313,  13.16813,   9.42344],
        [  0.12102,   7.17251,  11.18957],
        [  0.09898,   2.36466,  10.0365 ],
        [  0.37887,   0.92019,   6.87117],
        [  1.00661,   0.61958,   4.44955],
        [  1.20135,   0.17449,   3.01271],
        [  1.46433,   0.28039,   1.76431],
        [  1.37789,   0.0985 ,   1.28868],
        [  1.57073,   0.03343,   0.81813],
        [  1.4647 ,   0.28544,   0.52111],
        [  1.24719,   0.10138,   0.22746],
        [  1.56065,   0.21671,   0.19627]]

data = hcat(data...); 

# Need initial conditions for each model since they have different numbers of states
ic1 = [20.0, 10.0, 0.0]
ic2 = [20.0, 0.0, 10.0, 0.0]
ics = [ic1, ic2]

# Define a simulator function for each model. Each function returns the solution of the model
# as a 2D array for some parameter values. The appropriate ODE is solved at the same time 
# points as the observed data
function simulator1(params::Array{Float64,1})
    
    # p = (alpha, gamma, d, v)
    # x = (S, I, R)
    function model1(dx, x, p, t)
        dx[1] = p[1] - p[2]*x[1]*x[2] - p[3]*x[1] # dS/dt = alpha - gamma*S*I - d*S
        dx[2] = p[3]*x[1]*x[2] - p[4]*x[2] - p[3]*x[2] # dI/dt = gamma*S*I - v*I - d*I
        dx[3] = p[4]*x[2] - p[3]*x[3] # dR/dt = v*I - d*R
    end
    
    sol = solve(ODEProblem(model1, ics[1], (times[1], times[end]), params),
                RK4(), saveat=times, force_dtmin=true)
    hcat(sol.u...)
end

function simulator2(params::Array{Float64,1})

    # p = (alpha, gamma, d, v, delta)
    # x = (S, L, I, R)
    function model2(dx, x, p, t)
        dx[1] = p[1] - p[2]*x[1]*x[3] - p[3]*x[1] # dS/dt = alpha - gamma*S*I - d*S
        dx[2] = p[2]*x[1]*x[3] - p[5]*x[2] - p[3]*x[2] # dL/dt = gamma*S*I - delta*L - d*L
        dx[3] = p[5]*x[2] - p[4]*x[3] - p[3]*x[3] # dI/dt = delta*L - v*I - d*I
        dx[4] = p[4]*x[3] - p[3]*x[4] # dR/dt = v*I - d*R
    end
    
    # Model2 contains the species L, which is not measured - we remove it from the returned ODE solution
    # so that it can be compared to the reference data "data", which only contains S, I and R
    sol = solve(ODEProblem(model2, ics[2], (times[1], times[end]), params),
                RK4(), saveat=times, force_dtmin=true)
    hcat(sol.u...)[[1,3,4],:]
end;

#
# Priors and initial conditions - these are model-specfic as each model can 
# have different numbers of parameters/species
#
priors1 = [Uniform(0.0, 5.0) for i in 1:4]
priors2 = vcat([Uniform(0.0, 5.0) for i in 1:4], Uniform(0.0, 10.0))
priors3 = vcat([Uniform(0.0, 5.0) for i in 1:4], Uniform(0.0, 10.0))

threshold_schedule = [20, 15, 10, 5, 3, 2.5, 2, 1.7, 1.5]
summary_statistic = "keep_all";

n_particles = 200

ms_sim_result = SimulatedModelSelection(data,
    [simulator1, simulator2],
    [priors1, priors2],
    threshold_schedule,
    n_particles);

plot(ms_sim_result, title="Model selection result - simulation")

ms_emu_result = EmulatedModelSelection(data,
    [simulator1, simulator2],
    [priors1, priors2],
    threshold_schedule,
    n_particles,
    200;
    summary_statistic = "keep_all",
    distance_function=Distances.euclidean);

plot(ms_emu_result, title="Model selection result - emulation")