using Pkg
using Plots
using DataFrames
pyplot()

### functions for simulations
function calc_new_state(y::AbstractArray,
    a::Dict, 
    b::Dict, 
    l::Dict;
    from::Float64=0.01)
    if length(y)==0
        y = [from]
    end
    y[end]+(1-y[end])*(a["walk"]*(1-exp(-b["walk"]*l["walk"]*a["walk"]*y[end]))+a["TTC"]*(1-exp(-b["TTC"]*l["TTC"]*a["TTC"]*y[end])))
end

function simul(s::Dict,
               b::Dict, 
               l::Dict;
               to_perc::Bool=true,
               from::Float64=0.01,
               to::Union{Float64, Int}=0.99)

               # initial calculations
                e1 = (exp(-(s["walk"])/min(s["walk"], s["TTC"])))
                e2 = (exp(-(s["TTC"])/min(s["walk"], s["TTC"])))
                a = Dict("walk" => e1/sum([e1,e2]), "TTC" => e2/sum([e1,e2]))
                # simulation loop
                y = []
                append!(y, from)
                if to_perc
                    while y[end] <= to
                        append!(y, calc_new_state(y,a,b,l))
                    end
                else
                    for i in 1:to
                        append!(y, calc_new_state(y,a,b,l))
                    end
                end
                y
end

#####################################
### scenario 1:
s = Dict("walk" => 15, "TTC" => 10) # time of travel
l = Dict("walk" => 0.01, "TTC" => 0.02) # contagion
b = Dict("walk" => 5, "TTC" => 5) # interaction frequency

x1 = simul(s,b,l)
length(x1) - 1

### scenario 2:
s = Dict("walk" => 15, "TTC" => 10) # time of travel
l = Dict("walk" => 0.01, "TTC" => 0.02) # contagion
b = Dict("walk" => 3, "TTC" => 5) # interaction frequency

x2 = simul(s,b,l)
length(x2) - 1

# PLOT1 
Plots.plot(1:length(x1), x1, linecolor="purple", lw = 2, legend=false, grid=false)
ylabel!("% infected")
xlabel!("days")
savefig("plot1.eps")


# PLOT2 
Plots.plot(1:length(x2), x2, linecolor="purple", lw = 2, legend=false, grid=false)
ylabel!("% infected")
xlabel!("days")
savefig("plot2.eps")


###############################

#PLOT3
Plots.plot(1:length(x1), x1, linecolor="purple", label="Scenario 1", lw=2, grid = false, legend=:topleft)
Plots.plot!(1:length(x2), x2, linecolor="black", label="Scenario 2", lw=2 #linestyle = :dot, )
)
ylabel!("% infected")
xlabel!("days")
savefig("plot3.eps")

################################

# PLOT 4
s = Dict("walk" => 15, "TTC" => 10) # time of travel
l = Dict("walk" => 0.01, "TTC" => 0.02) # contagion
b = Dict("walk" => 3, "TTC" => 5) # interaction frequency

results_full = Dict()
results_short = DataFrame()
for i in 1:50
    s["TTC"] = i
    results_full[i] = simul(s,b,l)
    append!(results_short, DataFrame(s_ttc = [i], days = [length(results_full[i])-1]))
end

results_short # ok

Plots.plot(results_short[!,:s_ttc], results_short[!,:days], linecolor="purple", lw = 2, legend=false, grid=false)
ylabel!("days")
xlabel!("time when taking TTC")
savefig("plot4.eps")

#############################

# PLOT 5
s = Dict("walk" => 15, "TTC" => 10) # time of travel
l = Dict("walk" => 0.01, "TTC" => 0.02) # contagion
b = Dict("walk" => 3, "TTC" => 5) # interaction frequency

results_full2 = Dict()
results_short2 = DataFrame()
for i in 1:50
    s["TTC"] = i
    results_full2[i] = simul(s,b,l, to_perc=false, to=101)
    append!(results_short2, DataFrame(s_ttc = [i], infected = results_full2[i][101]))
end

results_short2 # ok

Plots.plot(results_short2[!,:s_ttc], results_short2[!,:infected], linecolor="purple", lw = 2, legend=false, grid=false)
ylabel!("% infected after 100 days")
xlabel!("time when taking TTC")
savefig("plot5.eps")



