using InstantiateFromURL
# optionally add arguments to force installation: instantiate = true, precompile = true
github_project("QuantEcon/quantecon-notebooks-julia", version = "0.8.0")

using LinearAlgebra, Statistics, Random, SparseArrays

using OrdinaryDiffEq, StochasticDiffEq
using Parameters, Plots

function F(x, p, t)
    s, i, r, d, R₀, δ = x
    @unpack γ, R̄₀, η, σ, ξ, θ, δ_bar = p

    return [-γ*R₀*s*i;        # ds/dt
            γ*R₀*s*i - γ*i;   # di/dt
            (1-δ)*γ*i;        # dr/dt
            δ*γ*i;            # dd/dt
            η*(R̄₀(t, p) - R₀);# dR₀/dt
            θ*(δ_bar - δ);    # dδ/dt
            ]
end

function G(x, p, t)
    s, i, r, d, R₀, δ = x
    @unpack γ, R̄₀, η, σ, ξ, θ, δ_bar = p

    return [0; 0; 0; 0; σ*sqrt(R₀); ξ*sqrt(δ * (1-δ))]
end

p_gen = @with_kw (T = 550.0, γ = 1.0 / 18, η = 1.0 / 20,
                  R₀_n = 1.6, R̄₀ = (t, p) -> p.R₀_n, δ_bar = 0.01,
                  σ = 0.03, ξ = 0.004, θ = 0.2, N = 3.3E8)
p =  p_gen()  # use all defaults
i_0 = 25000 / p.N
r_0 = 0.0
d_0 = 0.0
s_0 = 1.0 - i_0 - r_0 - d_0
R̄₀_0 = 0.5  # starting in lockdown
δ_0 = p.δ_bar
x_0 = [s_0, i_0, r_0, d_0, R̄₀_0, δ_0]

prob = SDEProblem(F, G, x_0, (0, p.T), p)

sol_1 = solve(prob, SOSRI());
@show length(sol_1.t);

sol_2 = solve(prob, SOSRI())
plot(sol_1, vars=[2], title = "Number of Infections", label = "Trajectory 1",
     lm = 2, xlabel = "t", ylabel = "i(t)")
plot!(sol_2, vars=[2], label = "Trajectory 2", lm = 2, ylabel = "i(t)")

plot_1 = plot(sol_1, vars=[4], title = "Cumulative Death Proportion", label = "Trajectory 1",
              lw = 2, xlabel = "t", ylabel = "d(t)", legend = :topleft)
plot!(plot_1, sol_2, vars=[4], label = "Trajectory 2", lw = 2)
plot_2 = plot(sol_1, vars=[3], title = "Cumulative Recovered Proportion", label = "Trajectory 1",
              lw = 2, xlabel = "t", ylabel = "d(t)", legend = :topleft)
plot!(plot_2, sol_2, vars=[3], label = "Trajectory 2", lw = 2)
plot_3 = plot(sol_1, vars=[5], title = "R_0 transition from lockdown", label = "Trajectory 1",
              lw = 2, xlabel = "t", ylabel = "R_0(t)")
plot!(plot_3, sol_2, vars=[5], label = "Trajectory 2", lw = 2)
plot_4 = plot(sol_1, vars=[6], title = "Mortality Rate", label = "Trajectory 1",
              lw = 2, xlabel = "t", ylabel = "delta(t)", ylim = (0.006, 0.014))
plot!(plot_4, sol_2, vars=[6], label = "Trajectory 2", lw = 2)
plot(plot_1, plot_2, plot_3, plot_4, size = (900, 600))

ensembleprob = EnsembleProblem(prob)
sol = solve(ensembleprob, SOSRI(), EnsembleSerial(), trajectories = 10)
plot(sol, vars = [2], title = "Infection Simulations", ylabel = "i(t)", xlabel = "t", lm = 2)

trajectories = 100  # choose larger for smoother quantiles
sol = solve(ensembleprob, SOSRI(), EnsembleThreads(), trajectories = trajectories)
summ = EnsembleSummary(sol) # defaults to saving 0.05, 0.95 quantiles
plot(summ, idxs = (2,), title = "Quantiles of Infections Ensemble", ylabel = "i(t)",
     xlabel = "t", labels = "Middle 95% Quantile", legend = :topright)

sol = solve(ensembleprob, SOSRI(), EnsembleThreads(), trajectories = trajectories)
summ = EnsembleSummary(sol) # defaults to saving 0.05, 0.95 quantiles
summ2 = EnsembleSummary(sol, quantiles = (0.25, 0.75))

plot(summ, idxs = (2,4,5,6),
    title = ["Proportion Infected" "Proportion Dead" "R_0" "delta"],
    ylabel = ["i(t)" "d(t)" "R_0(t)" "delta(t)"], xlabel = "t",
    legend = [:topleft :topleft :bottomright :bottomright],
    labels = "Middle 95% Quantile", layout = (2, 2), size = (900, 600))
plot!(summ2, idxs = (2,4,5,6),
    labels = "Middle 50% Quantile", legend =  [:topleft :topleft :bottomright :bottomright])

function generate_η_experiment(η; p_gen = p_gen, trajectories = 100,
                              saveat = 1.0, x_0 = x_0, T = 120.0)
    p = p_gen(η = η, ξ = 0.0)
    ensembleprob = EnsembleProblem(SDEProblem(F, G, x_0, (0, T), p))
    sol = solve(ensembleprob, SOSRI(), EnsembleThreads(),
                trajectories = trajectories, saveat = saveat)
    return EnsembleSummary(sol)
end

# Evaluate two different lockdown scenarios
η_1 = 1/50
η_2 = 1/20
summ_1 = generate_η_experiment(η_1)
summ_2 = generate_η_experiment(η_2)
plot(summ_1, idxs = (4,5),
    title = ["Proportion Dead" "R_0"],
    ylabel = ["d(t)" "R_0(t)"], xlabel = "t",
    legend = [:topleft :bottomright],
    labels = "Middle 95% Quantile, eta = $η_1",
    layout = (2, 1), size = (900, 900), fillalpha = 0.5)
plot!(summ_2, idxs = (4,5),
    legend = [:topleft :bottomright],
    labels = "Middle 95% Quantile, eta = $η_2", size = (900, 900), fillalpha = 0.5)

R₀_L = 0.5  # lockdown
η_experiment = 1.0/10
σ_experiment = 0.04
R̄₀_lift_early(t, p) = t < 30.0 ? R₀_L : 2.0
R̄₀_lift_late(t, p) = t < 120.0 ? R₀_L : 2.0

p_early = p_gen(R̄₀ = R̄₀_lift_early, η = η_experiment, σ = σ_experiment)
p_late = p_gen(R̄₀ = R̄₀_lift_late, η = η_experiment, σ = σ_experiment)


# initial conditions
i_0 = 100000 / p_early.N
r_0 = 0.0
d_0 = 0.0
s_0 = 1.0 - i_0 - r_0 - d_0
δ_0 = p_early.δ_bar

x_0 = [s_0, i_0, r_0, d_0, R₀_L, δ_0]  # start in lockdown
prob_early = SDEProblem(F, G, x_0, (0, p_early.T), p_early)
prob_late = SDEProblem(F, G, x_0, (0, p_late.T), p_late)

sol_early = solve(prob_early, SOSRI())
sol_late = solve(prob_late, SOSRI())
plot(sol_early, vars = [5, 1,2,4],
    title = ["R_0" "Susceptible" "Infected" "Dead"],
    layout = (2, 2), size = (900, 600),
    ylabel = ["R_0(t)" "s(t)" "i(t)" "d(t)"], xlabel = "t",
        legend = [:bottomright :topright :topright :topleft],
        label = ["Early" "Early" "Early" "Early"])
plot!(sol_late, vars = [5, 1,2,4],
        legend = [:bottomright :topright :topright :topleft],
        label = ["Late" "Late" "Late" "Late"])

trajectories = 400
saveat = 1.0
ensemble_sol_early = solve(EnsembleProblem(prob_early), SOSRI(),
                           EnsembleThreads(), trajectories = trajectories, saveat = saveat)
ensemble_sol_late = solve(EnsembleProblem(prob_late), SOSRI(),
                          EnsembleThreads(), trajectories = trajectories, saveat = saveat)
summ_early = EnsembleSummary(ensemble_sol_early)
summ_late = EnsembleSummary(ensemble_sol_late)

plot(summ_early, idxs = (5, 1, 2, 4),
    title = ["R_0" "Susceptible" "Infected" "Dead"], layout = (2, 2), size = (900, 600),
    ylabel = ["R_0(t)" "s(t)" "i(t)" "d(t)"], xlabel = "t",
    legend = [:bottomright :topright :topright :topleft],
    label = ["Early" "Early" "Early" "Early"])
plot!(summ_late, idxs = (5, 1,2,4),
    legend = [:bottomright :topright :topright :topleft],
    label = ["Late" "Late" "Late" "Late"])

N = p_early.N
t_1 = 350
t_2 = p_early.T  # i.e. the last element
bins_1 = range(0.000, 0.009, length = 30)
bins_2 = 30  # number rather than grid.

hist_1 = histogram([ensemble_sol_early.u[i](t_1)[4] for i in 1:trajectories],
                   fillalpha = 0.5, normalize = :probability,
                   legend = :topleft, bins = bins_1,
                   label = "Early", title = "Death Proportion at t = $t_1")
histogram!(hist_1, [ensemble_sol_late.u[i](t_1)[4] for i in 1:trajectories],
           label = "Late", fillalpha = 0.5, normalize = :probability, bins = bins_1)
hist_2 = histogram([ensemble_sol_early.u[i][4, end] for i in 1:trajectories],
                   fillalpha = 0.5, normalize = :probability, bins = bins_2,
                   label = "Early", title = "Death Proportion at t = $t_2")
histogram!(hist_2, [ensemble_sol_late.u[i][4, end] for i in 1:trajectories],
           label = "Late", fillalpha = 0.5, normalize = :probability, bins = bins_2)
plot(hist_1, hist_2, size = (600,600), layout = (2, 1))

function F_reinfect(x, p, t)
    s, i, r, d, R₀, δ = x
    @unpack γ, R̄₀, η, σ, ξ, θ, δ_bar, ν = p

    return [-γ*R₀*s*i + ν*r;  # ds/dt
            γ*R₀*s*i - γ*i;   # di/dt
            (1-δ)*γ*i - ν*r   # dr/dt
            δ*γ*i;            # dd/dt
            η*(R̄₀(t, p) - R₀);# dR₀/dt
            θ*(δ_bar - δ);    # dδ/dt
            ]
end

p_re_gen = @with_kw ( T = 550.0, γ = 1.0 / 18, η = 1.0 / 20,
                R₀_n = 1.6, R̄₀ = (t, p) -> p.R₀_n,
                δ_bar = 0.01, σ = 0.03, ξ = 0.004, θ = 0.2, N = 3.3E8, ν = 1/360)

p_re_early = p_re_gen(R̄₀ = R̄₀_lift_early, η = η_experiment, σ = σ_experiment)
p_re_late = p_re_gen(R̄₀ = R̄₀_lift_late, η = η_experiment, σ = σ_experiment)

trajectories = 400
saveat = 1.0
prob_re_early = SDEProblem(F_reinfect, G, x_0, (0, p_re_early.T), p_re_early)
prob_re_late = SDEProblem(F_reinfect, G, x_0, (0, p_re_late.T), p_re_late)
ensemble_sol_re_early = solve(EnsembleProblem(prob_re_early), SOSRI(), EnsembleThreads(),
                              trajectories = trajectories, saveat = saveat)
ensemble_sol_re_late = solve(EnsembleProblem(prob_re_late), SOSRI(), EnsembleThreads(),
                              trajectories = trajectories, saveat = saveat)
summ_re_early = EnsembleSummary(ensemble_sol_re_early)
summ_re_late = EnsembleSummary(ensemble_sol_re_late)

plot(summ_late, idxs = (1, 2, 3, 4),
    title = ["Susceptible" "Infected" "Recovered" "Dead"],
    layout = (2, 2), size = (900, 600),
    ylabel = ["s(t)" "i(t)" "r(t)" "d(t)"], xlabel = "t",
    legend = :topleft,
    label = ["s(t)" "i(t)" "r(t)" "d(t)"])
plot!(summ_re_late, idxs =  (1, 2, 3, 4),
    legend = :topleft,
    label = ["s(t); nu > 0" "i(t); nu > 0" "r(t); nu > 0" "d(t); nu > 0"])

bins_re_1 = range(0.003, 0.010, length = 50)
bins_re_2 = range(0.0085, 0.0102, length = 50)
hist_re_1 = histogram([ensemble_sol_re_early.u[i](t_1)[4] for i in 1:trajectories],
                    fillalpha = 0.5, normalize = :probability,
                    legend = :topleft, bins = bins_re_1,
                    label = "Early", title = "Death Proportion at t = $t_1")
histogram!(hist_re_1, [ensemble_sol_re_late.u[i](t_1)[4] for i in 1:trajectories],
        label = "Late", fillalpha = 0.5, normalize = :probability, bins = bins_re_1)
hist_re_2 = histogram([ensemble_sol_re_early.u[i][4, end] for i in 1:trajectories],
                    fillalpha = 0.5, normalize = :probability, bins = bins_re_2,
                    label = "Early", title = "Death Proportion at t = $t_2")
histogram!(hist_re_2, [ensemble_sol_re_late.u[i][4, end] for i in 1:trajectories],
        label = "Late", fillalpha = 0.5, normalize = :probability,
        bins =  bins = bins_re_2)
plot(hist_re_1, hist_re_2, size = (600,600), layout = (2, 1))