using Printf, DataFrames, CSV, LinearAlgebra, Statistics, Plots, Plots.PlotMeasures, Dates
# pyplot()

# abbreviate the name of a state to 2 capital letters
function abbreviate_state(state)
    d = Dict(
        "Alabama" => "AL",
        "Alaska" => "AK",
        "Arizona" => "AZ",
        "Arkansas" => "AR",
        "California" => "CA",
        "Colorado" => "CO",
        "Connecticut" => "CT",
        "Delaware" => "DE",
        "District of Columbia" => "DC",
        "Florida" => "FL",
        "Georgia" => "GA",
        "Hawaii" => "HI",
        "Idaho" => "ID",
        "Illinois" => "IL",
        "Indiana" => "IN",
        "Iowa" => "IA",
        "Kansas" => "KS",
        "Kentucky" => "KY",
        "Louisiana" => "LA",
        "Maine" => "ME",
        "Maryland" => "MD",
        "Massachusetts" => "MA",
        "Michigan" => "MI",
        "Minnesota" => "MN",
        "Mississippi" => "MS",
        "Missouri" => "MO",
        "Montana" => "MT",
        "Nebraska" => "NE",
        "Nevada" => "NV",
        "New Hampshire" => "NH",
        "New Jersey" => "NJ",
        "New Mexico" => "NM",
        "New York" => "NY",
        "North Carolina" => "NC",
        "North Dakota" => "ND",
        "Ohio" => "OH",
        "Oklahoma" => "OK",
        "Oregon" => "OR",
        "Pennsylvania" => "PA",
        "Puerto Rico" => "PR",
        "Rhode Island" => "RI",
        "South Carolina" => "SC",
        "South Dakota" => "SD",
        "Tennessee" => "TN",
        "Texas" => "TX",
        "Utah" => "UT",
        "Vermont" => "VT",
        "Virginia" => "VA",
        "Washington" => "WA",
        "West Virginia" => "WV",
        "Wisconsin" => "WI",
        "Wyoming" => "WY"
    )
    get(d, state, state)
end

# load population data for a county
function load_pop_county(county::String, state::String)
    df = CSV.read("county_pop.csv", DataFrame)
    if county == "New York City"
        pop = df[(df.county.=="New York").&(df.state.=="New York"), :population][1] +
              df[(df.county.=="Kings").&(df.state.=="New York"), :population][1] +
              df[(df.county.=="Queens").&(df.state.=="New York"), :population][1] +
              df[(df.county.=="Bronx").&(df.state.=="New York"), :population][1] +
              df[(df.county.=="Richmond").&(df.state.=="New York"), :population][1]
    else
        pop = df[(df.county.==county).&(df.state.==state), :population][1]
    end
    pop
end

# load population data for a state
function load_pop_state(state::String)
    df = CSV.read("state_pop.csv", DataFrame)
    df[df.state.==state, :population][1]
end

# load population data for the nation
function load_pop_nation()
    df = CSV.read("state_pop.csv", DataFrame)
    sum(df[:, :population])
end

# load case and death data for a single county from a dataframe
function load_cd_county(df::DataFrame, county::String, state::String)
    Matrix(df[(df.county.==county).&(df.state.==state), [:cases, :deaths]])
end

# load case and death data for all states
function load_cd_states()
    df = CSV.read("covid-19-data/us-states.csv", DataFrame)
    date = df.date[end]
    states = unique(df.state)
    n_states = length(states) # not 50, includes territories, etc
    cd_data = Array{Any,1}(undef, n_states)
    for i = 1:n_states
        cd_data[i] = Matrix(df[df.state.==states[i], [:cases, :deaths]])
    end
    states, cd_data, date
end

# load case and death data for the nation
function load_cd_nation()
    df = CSV.read("covid-19-data/us.csv", DataFrame)
    cd_data = Matrix(df[:, [:cases, :deaths]])
end

# apply bidirectional first-order IIR to smooth data
function smooth_iir_1(y, alpha = 0.3)
    n = length(y)
    # first order IIR forward
    z0 = copy(y)
    z = y[1]
    for i = 1:n
        z = alpha * y[i] + (1 - alpha) * z
        z0[i] = z
    end
    # first order IIR backward
    z1 = copy(y)
    z = y[end]
    for i = n:-1:1
        z = alpha * y[i] + (1 - alpha) * z
        z1[i] = z
    end
    (z0 + z1) * 0.5
end

# estimate the differences of a time series and fill in gaps
function est_diff(y)
    n = length(y)
    d = diff(y)
    # remove negative dips
    d[d.<0.0] .= 0.0
    # redistribute spikes after zeros
    i0 = 1
    i = 2
    while i < n
        # X0 = hcat(ones(2 * m + 1), collect(-m:m))
        # if zero is found
        if d[i] == 0
            # find first non-zero
            while d[i] == 0.0 && i < n - 1
                i += 1
            end
            if d[i] == 0.0 && i == n - 1
                break
            end
            # redistribute d[i] to make a line from d[i0] to d[i]
            m = i - i0
            @assert(m > 0)
            t = (m * (m + 1.0)) / 2.0
            a = (d[i] - m * d[i0]) / t
            if a * m + d[i0] >= 0
                X = hcat(ones(m), collect(1:m))
                d[i0+1:i] = X * [d[i0], a]
            end
        end
        i0 = i
        i += 1
    end
    # remove huge spikes
    b = [-3, -2, -1, 1, 2, 3]
    for i = 4:n-4
        m = mean(d[b.+i])
        if d[i] > 5 * max(1, m)
            d[i] = m
        end
    end
    d
end

# sum of derivatives of sigmoids (generalized logistic function)
function sigmoid_deriv_sum(a, t)
    @assert(length(a) % 4 == 0)
    b = reshape(a, 4, :)
    m = size(b, 2)
    y = zeros(length(t))
    for i = 1:m
        e = exp.((-5.0 * b[2, i]) * (t .- b[3, i]))
        y .+= (4.0 * b[1, i]) * e ./ ((1.0 .+ e) .^ (1.0 + b[4, i]))
    end
    y
end

# Jacobian of the weighted residual in a sum of sigmoids fit
function sigmoid_deriv_sum_jac(a, t, w)
    @assert(length(a) % 4 == 0)
    b = reshape(a, 4, :)
    m = size(b, 2)
    n = length(t)
    jac = zeros(n, 4, m)
    u = (4.0 / sqrt(n)) * w
    for i = 1:m
        e = exp.(-5.0 * b[2, i] * (t .- b[3, i]))
        d = (1 .+ e)
        g = b[1, i] * u .* e
        h = g .* d .^ (-(2.0 + b[4, i])) .* (1.0 .- b[4, i] * e)
        jac[:, 1, i] = u .* e ./ (d .^ (1.0 + b[4, i]))
        jac[:, 2, i] = h .* (-5.0 * (t .- b[3, i]))
        jac[:, 3, i] = h * (5.0 * b[2, i])
        jac[:, 4, i] = g .* d .^ (-(1.0 + b[4, i])) .* (-log.(d))
    end
    reshape(jac, n, :)
end

# test for Jacobian of the residual in a sum of sigmoids fit
function test_sigmoid_deriv_sum_jac()
    n = 100
    t = collect(range(0, stop = 1, length = n))
    w = ones(n)
    a = [1.5, 0.6, 0.8, 1, 2, 0.3, 1, 1.1]
    y0 = rand(n)
    # direct computation of Jacobian of residual
    jac0 = sigmoid_deriv_sum_jac(a, t, w)
    # finite difference approximation
    jac1 = zeros(n, 8)
    eps = 1e-6
    for i = 1:8
        a1 = copy(a)
        a1[i] -= eps
        res0 = w .* (sigmoid_deriv_sum(a1, t) - y0) / sqrt(n)
        a1[i] += 2 * eps
        res1 = w .* (sigmoid_deriv_sum(a1, t) - y0) / sqrt(n)
        jac1[:, i] = (res1 - res0) / (2 * eps)
    end
    println(maximum(abs.(jac0 - jac1)))
end

# fit a sigmoid sum using the Gauss-Newton method to optimize weighted least squares
function fit_sigmoid_deriv_sum_wls_gn(x_init, t, w, y0_raw, iter_max, debug::Bool = false)
    # initialize return value
    x_opt = copy(x_init)

    # normalize data to fit
    n = length(x_init)
    y0_max = maximum(y0_raw)
    y0 = y0_raw / y0_max

    # Levenberg-Marquardt parameters
    eps_min = 1e-6
    eps_max = 1e+2
    eps_inc = 3.0
    eps_dec = 0.5

    # set up initial values
    eps = 1.0
    eps_it = 0
    x1 = copy(x_init)
    res = w .* (sigmoid_deriv_sum(x1, t) - y0) / sqrt(n)
    jac = sigmoid_deriv_sum_jac(x1, t, w)
    err1 = norm(res)
    err_opt = err1

    # regularization coefficients
    x_init[1:4:end] .= 0.0
    c_reg = 3e-2 * ones(n)
    c_reg[3:4:end] .= 0.0

    # Gauss-Newton iterations
    for it = 1:iter_max

        # update parameters
        x0 = copy(x1)
        x1 -= (jac' * jac + eps * I) \ (jac' * res) * 0.5 - c_reg .* (x1 - x_init)
        x1[1:4:end] = max.(min.(x1[1:4:end], 100.0), 0.01)
        x1[2:4:end] = max.(min.(x1[2:4:end], 10.0), 0.1)
        x1[3:4:end] = max.(min.(x1[3:4:end], 1.1), -0.1)
        x1[4:4:end] = max.(min.(x1[4:4:end], 10.0), 0.1)

        # update weighted residual and its Jacobian
        err0 = err1
        res = w .* (sigmoid_deriv_sum(x1, t) - y0) / sqrt(n)
        jac = sigmoid_deriv_sum_jac(x1, t, w)
        err1 = norm(res)

        # lock in parameters if the error goes down
        if err1 < err_opt
            err_opt = err1
            x_opt = copy(x1)
        end

        # update epsilon based on the error going up or down
        if err1 <= err0 + 1e-2
            eps *= eps_dec
            if eps < eps_min
                eps = eps_min
            end
            if eps == eps_min
                if eps_it == 3
                    break
                else
                    eps_it += 1
                end
            end
        else # backtrack
            if eps == eps_max
                if eps_it == 3
                    break
                else
                    eps_it += 1
                end
            else
                eps_it = 0
                err1 = err0
                x1 = copy(x0)
                eps *= eps_inc
                if eps > eps_max
                    eps = eps_max
                end
            end
        end

        # possibly output debug information
        if debug
            @printf("Iter: %2d,    Eps: %.6f,    Err: %.6f,    Best Err: %.6f\n",
                it, eps, err1, err_opt)
        end
    end
    if debug
        println()
    end

    x_opt[1:4:end] *= y0_max
    x_opt, err_opt / length(y0)
end

# fit sum of sigmoid functions (generalized logistics) using iteratively re-weighted least squares
function fit_sigmoid_deriv_sum(y0, n_ext)
    n = length(y0)
    t0 = collect(range(0, stop = 1, length = n))
    w0 = ones(n)
    w0[end-5:end] .= [0.8, 0.6, 0.4, 0.2, 0, 0]
    opt_a1 = []
    opt_b1 = Inf
    for n_sig = 1:10
        a0 = ones(4, n_sig)
        a0[1, :] .= 0.5
        for i_sig = 1:n_sig
            a0[3, i_sig] = (i_sig - 1.0) / n_sig
        end
        a1, err = fit_sigmoid_deriv_sum_wls_gn(a0[:], t0, w0, y0, 300)
        b1 = n_sig / (n / 60) + log(err) # bayesian information criterion
        if b1 < opt_b1
            opt_b1 = b1
            opt_a1 = a1
        end
    end
    t1 = [t0; 1.0 .+ (t0[2] - t0[1]) * collect(1:n_ext)]
    y1 = sigmoid_deriv_sum(opt_a1, t1)
    y1_max = sigmoid_deriv_sum(opt_a1, 10.0)
    y1, y1_max
end

function filter_first_of_month(t0)
    # find unique first days of each month
    t1 = unique(Dates.firstdayofmonth.(t0))
    if t1[1] < t0[1]
        t1 = t1[2:end]
    end
    t1
end

# hack to deal with twinx() undocumented/buggy behavior for dates: convert to int64
function dates_int64_xticks(t0)
    x = collect(0:length(t0)-1)
    x_offset = Dates.value(t0[1])
    tt = filter_first_of_month(t0)
    xt = zeros(length(tt))
    j = 1
    for i = 1:length(t0)
        if tt[j] == t0[i]
            xt[j] = x[i]
            j += 1
            if j > length(tt)
                break
            end
        end
    end
    xt, x, x_offset
end

# plot daily cases/deaths per day
function plot_dcd(cd, title_str, date, pop, set_ymax, n_days_ext, n_days_plot)
    # cutoff for fit
    cd = cd[1:end-3, :]
    n_days_plot = min(size(cd, 1), n_days_plot)
    cd = cd[end-(n_days_plot-1):end, :]
    @assert(size(cd, 1) == n_days_plot)
    # normalize by population (per 100,000)
    c0 = cd[:, 1] / pop * 1e5
    d0 = cd[:, 2] / pop * 1e5
    # compute new cases
    dc0 = est_diff(c0)
    dd0 = est_diff(d0)
    # remove trailing zeros
    n_days = length(dc0)
    n_days_zero = 0
    while dc0[end-n_days_zero] == 0.0
        n_days_zero += 1
        # exclude case when cases go smoothly to zero
        if 0.0 < dc0[end-n_days_zero] && dc0[end-n_days_zero] < 1.0
            n_days_zero = 0
            break
        end
    end
    n_days_plot -= n_days_zero
    dc0 = dc0[1:end-n_days_zero]
    dd0 = dd0[1:end-n_days_zero]
    t0 = collect(date-Day(n_days - 1):Day(1):date-Day(n_days_zero))
    t1 = collect(date-Day(n_days - 1):Day(1):date-Day(n_days_zero)+Day(n_days_ext))
    # smooth data
    dc0 = smooth_iir_1(dc0)
    dd0 = smooth_iir_1(dd0, 0.15)
    # fit sigmoid derivative to new cases
    dc1, dc1_max = fit_sigmoid_deriv_sum(dc0, n_days_ext)
    dd1, dd1_max = fit_sigmoid_deriv_sum(dd0, n_days_ext)
    # current cases
    curr_cases = round(dc0[end], digits = 1)
    title_str = title_str * "  [$curr_cases]"
    # plot cases
    if set_ymax
        c_ymax = 150.0
        d_ymax = 2.0
    else
        c_ymax = max(1.05 * max(maximum(dc0), maximum(dc1[1:n_days_plot])), 10.0)
        d_ymax = max(2.00 * max(maximum(dd0), maximum(dd1[1:n_days_plot])), 1.0)
    end
    xt, x, x_offset = dates_int64_xticks(t0)
    p = plot(x, dc0, w = 2, color = :black, xticks = false, ylabel = "Daily Cases (per 100k)", yguidefont = font(8),
        title = title_str, titlefontsize = 10, legend = false, ylims = (0, c_ymax))
    tx = twinx()
    plot!(tx, dd0, w = 2, color = :red, grid = :off, xticks = false, legend = false,
        ylabel = "Daily Deaths (per 100k)", yguidefont = font(8, :red), ylims = (0, d_ymax))
    xt, x, x_offset = dates_int64_xticks(t1)
    plot!(x, dc1, color = :grey, xticks = xt, xtickfontsize = 8, legend = false, ylims = (0, c_ymax))
    plot!(tx, dd1, color = :orange, grid = :off, xticks = false, legend = false, ylims = (0, d_ymax))
    plot!(xformatter = x -> Dates.monthabbr(Date(Dates.UTD(x + x_offset))))
    p
end

# plot cases/deaths over time and new cases versus total cases
function plot_county_list(df, county_list, plot_national = true)
    set_ymax = false
    n_days_ext = 5
    n_days_plot = 150
    n_cnty = length(county_list)
    n_plot = n_cnty + plot_national
    date = df[end, :date]
    p = Array{Any,1}(undef, n_plot)
    for i = 1:n_cnty
        county, state = county_list[i]
        cd = load_cd_county(df, county, state)
        pop = load_pop_county(county, state)
        state_abbr = abbreviate_state(state)
        if county == "New York City"
            title_str = "New York City, NY"
        else
            title_str = "$county County, $state_abbr"
        end
        p[i] = plot_dcd(cd, title_str, date, pop, set_ymax, n_days_ext, n_days_plot)
    end
    if plot_national
        cd = load_cd_nation()
        pop = load_pop_nation()
        p[end] = plot_dcd(cd, "United States", date, pop, set_ymax, n_days_ext, n_days_plot)
    end
    plot(p..., layout = (n_plot, 1), size = (600, 250 * n_plot), right_margin = 18mm, left_margin = 8mm)
end

# load data
print("Loading Covid-19 county cases/deaths data...")
t = @elapsed df = CSV.read("covid-19-data/us-counties.csv", DataFrame)
@printf("finished in %.2f sec.\n", t)

# plot data
# display(plot_county_list(df, [("Winnebago", "Illinois"),
#                               ("Boone", "Illinois"),
#                               ("DeKalb", "Illinois")]))

display(plot_county_list(df, [("Santa Clara", "California"),
    ("Santa Barbara", "California"),
    ("Los Angeles", "California"),
    ("New York City", "New York")]))

