using HTTP, Plots, DifferentialEquations;

# Signal derivation over days
# x: signal to derivate
# DT: delta time (unit: day)
#   DT=1 => daily
#   DT=7 => weekly
# Return the derived signal
function derivative(x, DT = 1)                
    M = size(x, 1);
    x0 = x[1:M-DT];
    x1 = x[1+DT:M];
    return (x1 .- x0) ./ DT;
end

# Given r, get the a and b coefs of: a * log(accumul)^r + b
# where: cumulated is the total of daily cases (infected, death ... )
# Return a and b coefs and the initial value for the solver.
function model(cumulated, r)
    # Logarithm of accumulated infected people
    log_cumul = log.(cumulated)
    # Linear regression: y = a x + b
    x, y = log_cumul[1:end-1] .^ r, derivative(log_cumul)
    b, a = [ones(size(x, 1)) x] \ y
    return b, a, log_cumul[1]
end

# ODE
# f: function to solve
# u0: initial value
# tend: end of simulation time
function simulation(f::Function, u0, tend)
    tspan = (0.0, tend)
    prob = ODEProblem(f, u0, tspan)
    sol = solve(prob, Tsit5(), reltol=1e-12, abstol=1e-12)
    derivative(exp.(sol(1:tend)))
end

# Daily-infected prediction
# country: the desired country.
# range: days for selecting the desired wave used for creating the model.
# recovered: people recovered from the previous wave
#   (a good estimated value is cumulative[end] from the previous wave)
# r: a * log(accumul)^r + b (default: 1)
# days: estimation up to desired days (default: 100 days)
# wave: the nth wave used for the title of the plot
# position: position of the legend (default: top right)
function prediction(country, range; recovered=0.0, r=1.0, days=100, wave, position=:topright)
    cumul_for_model = cumulative[country][range]
    cumul_for_prediction = cumulative[country][first(range):last(range) + days]
    S = size(cumul_for_model, 1)
    b, a, u0 = model(cumul_for_model .- recovered, r)
    f(u, p, t) = (u .^ r) .* a .+ b
    sim = simulation(f, u0, S + days)
    
    plot(derivative(cumul_for_prediction), color = :green, label="Observations for prediction")
    plot!(derivative(cumul_for_model),
        title = country * " " * wave * " wave",
        xlabel = "Days",
        ylabel = "Daily Infected",
        label = "Observations for model",
        color = :blue,
        legend = position)
    plot!(sim, color = :red, label="Prediction")
end

# Load a CSV file and convert it as a dictionary structure where
# the dictionary keys are unique identifiers for country (ie name) and where
# the dictionary data are arrays of accumulated infected people over days.
#
# file: CSV file
# k: the CSV column id holding the key (ie country names)
# v: the CSV column id holding the accumulated infected people (total cases)
# Return the dictionary String => Vector{Float}
function load(file::AbstractString; k::Int, v::Int)
    dic = Dict{String, Array{Float64,1}}()
    open(file, "r") do data_file
        # Ignore the CSV header
        readline(data_file)
        # Ignore first data
        for row in 1:30
            readline(data_file)
        end
        # For each line of the CSV file ...
        rows = readlines(data_file)
        for row in rows
            # ... Extract columns
            entries = split(strip(row), ",")
            
            # Dictionary key: country identifier
            key = entries[k]
 
            # Dictionary values: convert string to float. Beware
            # some data are missing: in this case we copy the previous
            # non dummy data and when this not possible we use 0.0.
            val = 0.0
            if entries[v] != ""
                val = parse(Float64, entries[v])
            elseif haskey(dic, key)
                val = dic[key][end]
            end
            
            # Create a new if not existing, then append the new data
            if !haskey(dic, key)
                dic[key] = [val]
            else
                append!(dic[key], val)
            end
        end
    end
    dic
end

csv_file = HTTP.download("https://covid.ourworldindata.org/data/owid-covid-data.csv", update_period=Inf)

# Extract cumulative of daily cases by country
cumulative = load(csv_file, k=3, v=5)

# Fix France data with less noisy values. Data get from Fr governement 
cumulative["France"] = vcat(13, 18, 38, 57, 100, 130, 191, 212, 285, 423, 613, 949, 1126, 1412, 1784, 2281, 2876, 3661, 4500, 5423, 6633, 7730,
9134, 10995, 12612, 14459, 16689, 19856, 22302, 25233, 29155, 32964, 37575, 40174, 44450, 52128, 56989, 59105,
64338, 68605, 70478, 74390, 78167, 82048, 86334, 90676, 93790, 95403, 98076, 103573, 106206, 108847, 109252,
111821, 112606, 114657, 117324, 119151, 120804, 122577, 124114, 124575, 125770, 126835, 128442, 129581,
130185, 130979, 131287, 131863, 132967, 137150, 137779, 138421, 138854, 139063, 139519, 140227, 140734,
141356, 141919, 142291, 142411, 142903, 143427, 143845, 144163, 144566, 144806, 144921, 145279, 145555,
145746, 149071, 149668, 151496, 151753, 152091, 152225, 152325, 152444, 153055, 153634, 153977, 154188,
154591, 155136, 155561, 156287, 156813, 157220, 157372, 157716, 158174, 158641, 159452, 160093, 160277,
160750, 161267, 161348, 161648, 162936, 163454, 163980, 164260, 164801, 165719, 166378, 166960, 167711,
168159, 168335, 168810, 169473, 170094, 170752, 171504, 172089, 172377, 172888, 173304, 173838, 174674,
175639, 176404, 176754, 177338, 178336, 179398, 180528, 181528, 182528, 183079, 183804, 185196, 186573,
187919, 189547, 190739, 191295, 192334, 194029, 195633, 197921, 200105, 201990, 202775, 204172, 206696,
209365, 212211, 215521, 218536, 219029, 221267, 225043, 229814, 234400, 238002, 242899, 244854, 248158,
253587, 259698, 267077, 272530, 277943, 281025, 286007, 293024, 300182, 309156, 317706, 324777, 328980,
335524, 344101, 353944, 363350, 373911, 381094, 387252, 395104, 404888, 415481, 428696, 442194, 452763,
458061, 468069, 481141, 497237, 513034, 527546, 538569, 542639, 550690, 563535, 577505, 589653, 606625,
619190, 624274, 634763, 653509, 671638, 691977, 718873, 734974, 743479, 756472, 779063, 809684, 834770,
867197, 897034, 910277, 930475, 957421, 999043, 1041075, 1086497, 1138507, 1165278, 1198695, 1235132,
1282769, 1331984, 1367625, 1413915, 1466433, 1502763, 1543321, 1601367, 1661853, 1748705, 1787324,
1807479, 1829659, 1865538, 1898710, 1922504, 1954599, 1981827, 1991233, 2036755, 2065138, 2086288,
2109170, 2127051, 2140208, 2144660, 2153815, 2170097, 2183660, 2196119, 2208699, 2218483, 2222488,
2230571, 2244635, 2257331, 2268552, 2281475, 2292497, 2295908, 2309621, 2324216, 2337966, 2351372,
2365319, 2376852, 2379915, 2391447, 2409062, 2427316, 2442990, 2460555, 2473354, 2479151, 2490946);

prediction("France", 1:40, days=100, r=2.0, wave="1st")

prediction("France", 150:220, recovered=0*190000, r=0.5, days=32, wave="2nd", position=:topleft)

prediction("India", 150:200, r=8.5, days=150, wave="1st", position=:topright)

prediction("Italy", 60:80, r=5, wave="1st")

prediction("Italy", 40:70, r=3.2, days=150, wave="2nd")

prediction("United Kingdom", 60:80, r=2.1, days=150, wave="1st")

prediction("United States", 151:185, recovered=1500000, r=9.0, wave="3th")

prediction("Japan", 50:110, r=15, wave="1st")

prediction("Japan", 160:200, recovered=17000, r=4, wave="2nd")

prediction("Russia", 65:95, r=0.95, days=150, wave="1st")

prediction("Israel", 60:90, r=7.0, days=80, wave="1st")

prediction("Brazil", 90:190, r=1, days=150, wave="1st")