@testset "varinfo.jl" begin
    @testset "TypedVarInfo" begin
        @model gdemo(x, y) = begin
            s ~ InverseGamma(2, 3)
            m ~ truncated(Normal(0.0, sqrt(s)), 0.0, 2.0)
            x ~ Normal(m, sqrt(s))
            y ~ Normal(m, sqrt(s))
        end
        model = gdemo(1.0, 2.0)

        vi = VarInfo()
        model(vi, SampleFromUniform())
        tvi = TypedVarInfo(vi)

        meta = vi.metadata
        for f in fieldnames(typeof(tvi.metadata))
            fmeta = getfield(tvi.metadata, f)
            for vn in fmeta.vns
                @test tvi[vn] == vi[vn]
                ind = meta.idcs[vn]
                tind = fmeta.idcs[vn]
                @test meta.dists[ind] == fmeta.dists[tind]
                @test meta.orders[ind] == fmeta.orders[tind]
                @test meta.gids[ind] == fmeta.gids[tind]
                for flag in keys(meta.flags)
                    @test meta.flags[flag][ind] == fmeta.flags[flag][tind]
                end
                range = meta.ranges[ind]
                trange = fmeta.ranges[tind]
                @test all(meta.vals[range] .== fmeta.vals[trange])
            end
        end
    end
    @testset "Base" begin
        # Test Base functions:
        #   string, Symbol, ==, hash, in, keys, haskey, isempty, push!, empty!,
        #   getindex, setindex!, getproperty, setproperty!
        csym = gensym()
        vn1 = @varname x[1][2]
        @test string(vn1) == "x[1][2]"
        @test Symbol(vn1) == Symbol("x[1][2]")

        vn2 = @varname x[1][2]
        @test vn2 == vn1
        @test hash(vn2) == hash(vn1)
        @test inspace(vn1, (:x,))

        function test_base!(vi)
            empty!(vi)
            @test getlogp(vi) == 0
            @test get_num_produce(vi) == 0

            vn = @varname x
            dist = Normal(0, 1)
            r = rand(dist)
            gid = Selector()

            @test isempty(vi)
            @test ~haskey(vi, vn)
            @test !(vn in keys(vi))
            push!(vi, vn, r, dist, gid)
            @test ~isempty(vi)
            @test haskey(vi, vn)
            @test vn in keys(vi)

            @test length(vi[vn]) == 1
            @test length(vi[SampleFromPrior()]) == 1

            @test vi[vn] == r
            @test vi[SampleFromPrior()][1] == r
            vi[vn] = [2 * r]
            @test vi[vn] == 2 * r
            @test vi[SampleFromPrior()][1] == 2 * r
            vi[SampleFromPrior()] = [3 * r]
            @test vi[vn] == 3 * r
            @test vi[SampleFromPrior()][1] == 3 * r

            empty!(vi)
            @test isempty(vi)
            push!(vi, vn, r, dist, gid)

            function test_inspace()
                space = (:x, :y, @varname(z[1]), @varname(M[1:10, :]))

                @test inspace(@varname(x), space)
                @test inspace(@varname(y), space)
                @test inspace(@varname(x[1]), space)
                @test inspace(@varname(z[1][1]), space)
                @test inspace(@varname(z[1][:]), space)
                @test inspace(@varname(z[1][2:3:10]), space)
                @test inspace(@varname(M[[2, 3], 1]), space)
                @test inspace(@varname(M[:, 1:4]), space)
                @test inspace(@varname(M[1, [2, 4, 6]]), space)
                @test !inspace(@varname(z[2]), space)
                @test !inspace(@varname(z), space)
            end
            return test_inspace()
        end
        vi = VarInfo()
        test_base!(vi)
        test_base!(empty!(TypedVarInfo(vi)))
    end
    @testset "flags" begin
        # Test flag setting:
        #    is_flagged, set_flag!, unset_flag!
        function test_varinfo!(vi)
            vn_x = @varname x
            dist = Normal(0, 1)
            r = rand(dist)
            gid = Selector()

            push!(vi, vn_x, r, dist, gid)

            # del is set by default
            @test !is_flagged(vi, vn_x, "del")

            set_flag!(vi, vn_x, "del")
            @test is_flagged(vi, vn_x, "del")

            unset_flag!(vi, vn_x, "del")
            @test !is_flagged(vi, vn_x, "del")
        end
        vi = VarInfo()
        test_varinfo!(vi)
        test_varinfo!(empty!(TypedVarInfo(vi)))
    end
    @testset "setgid!" begin
        vi = VarInfo()
        meta = vi.metadata
        vn = @varname x
        dist = Normal(0, 1)
        r = rand(dist)
        gid1 = Selector()
        gid2 = Selector(2, :HMC)

        push!(vi, vn, r, dist, gid1)
        @test meta.gids[meta.idcs[vn]] == Set([gid1])
        setgid!(vi, gid2, vn)
        @test meta.gids[meta.idcs[vn]] == Set([gid1, gid2])

        vi = empty!(TypedVarInfo(vi))
        meta = vi.metadata
        push!(vi, vn, r, dist, gid1)
        @test meta.x.gids[meta.x.idcs[vn]] == Set([gid1])
        setgid!(vi, gid2, vn)
        @test meta.x.gids[meta.x.idcs[vn]] == Set([gid1, gid2])
    end
    @testset "setval! & setval_and_resample!" begin
        @model function testmodel(x)
            n = length(x)
            s ~ truncated(Normal(), 0, Inf)
            m ~ MvNormal(zeros(n), I)
            return x ~ MvNormal(m, s^2 * I)
        end

        @model function testmodel_univariate(x, ::Type{TV}=Vector{Float64}) where {TV}
            n = length(x)
            s ~ truncated(Normal(), 0, Inf)

            m = TV(undef, n)
            for i in eachindex(m)
                m[i] ~ Normal()
            end

            for i in eachindex(x)
                x[i] ~ Normal(m[i], s)
            end
        end

        x = randn(5)
        model_mv = testmodel(x)
        model_uv = testmodel_univariate(x)

        for model in [model_uv, model_mv]
            m_vns = model == model_uv ? [@varname(m[i]) for i in 1:5] : @varname(m)
            s_vns = @varname(s)

            vi_typed = VarInfo(model)
            vi_untyped = VarInfo()
            model(vi_untyped, SampleFromPrior())

            for vi in [vi_untyped, vi_typed]
                vicopy = deepcopy(vi)

                ### `setval` ###
                DynamicPPL.setval!(vicopy, (m=zeros(5),))
                # Setting `m` fails for univariate due to limitations of `setval!`
                # and `setval_and_resample!`. See docstring of `setval!` for more info.
                if model == model_uv
                    @test_broken vicopy[m_vns] == zeros(5)
                else
                    @test vicopy[m_vns] == zeros(5)
                end
                @test vicopy[s_vns] == vi[s_vns]

                # Ordering is NOT preserved => fails for multivariate model.
                DynamicPPL.setval!(
                    vicopy, (; (Symbol("m[$i]") => i for i in (1, 3, 5, 4, 2))...)
                )
                if model == model_uv
                    @test vicopy[m_vns] == 1:5
                else
                    @test vicopy[m_vns] == [1, 3, 5, 4, 2]
                end
                @test vicopy[s_vns] == vi[s_vns]

                DynamicPPL.setval!(
                    vicopy, (; (Symbol("m[$i]") => i for i in (1, 2, 3, 4, 5))...)
                )
                DynamicPPL.setval!(vicopy, (s=42,))
                @test vicopy[m_vns] == 1:5
                @test vicopy[s_vns] == 42

                ### `setval_and_resample!` ###
                if model == model_mv && vi == vi_untyped
                    # Trying to re-run model with `MvNormal` on `vi_untyped` will call
                    # `MvNormal(μ::Vector{Real}, Σ)` which causes `StackOverflowError`
                    # so we skip this particular case.
                    continue
                end

                vicopy = deepcopy(vi)
                DynamicPPL.setval_and_resample!(vicopy, (m=zeros(5),))
                model(vicopy)
                # Setting `m` fails for univariate due to limitations of `subsumes(::String, ::String)`
                if model == model_uv
                    @test_broken vicopy[m_vns] == zeros(5)
                else
                    @test vicopy[m_vns] == zeros(5)
                end
                @test vicopy[s_vns] != vi[s_vns]

                # Ordering is NOT preserved.
                DynamicPPL.setval_and_resample!(
                    vicopy, (; (Symbol("m[$i]") => i for i in (1, 3, 5, 4, 2))...)
                )
                model(vicopy)
                if model == model_uv
                    @test vicopy[m_vns] == 1:5
                else
                    @test vicopy[m_vns] == [1, 3, 5, 4, 2]
                end
                @test vicopy[s_vns] != vi[s_vns]

                # Correct ordering.
                DynamicPPL.setval_and_resample!(
                    vicopy, (; (Symbol("m[$i]") => i for i in (1, 2, 3, 4, 5))...)
                )
                model(vicopy)
                @test vicopy[m_vns] == 1:5
                @test vicopy[s_vns] != vi[s_vns]

                DynamicPPL.setval_and_resample!(vicopy, (s=42,))
                model(vicopy)
                @test vicopy[m_vns] != 1:5
                @test vicopy[s_vns] == 42
            end
        end

        # https://github.com/TuringLang/DynamicPPL.jl/issues/250
        @model function demo()
            return x ~ filldist(MvNormal([1, 100], I), 2)
        end

        vi = VarInfo(demo())
        vals_prev = vi.metadata.x.vals
        ks = [@varname(x[1, 1]), @varname(x[2, 1]), @varname(x[1, 2]), @varname(x[2, 2])]
        DynamicPPL.setval!(vi, vi.metadata.x.vals, ks)
        @test vals_prev == vi.metadata.x.vals

        DynamicPPL.setval_and_resample!(vi, vi.metadata.x.vals, ks)
        @test vals_prev == vi.metadata.x.vals
    end
end


#### Multiperiod DC Optimal Power Shutoff ####

# This file provides a basic implentation of the multiperiod DC Optimal Power Shutoff
# problem for experimentation on algorithm and problem variations.

# This file can be run by calling `include("mp-dc-ops.jl")` from the Julia REPL

# Developed by Noah Rhodes(@noahrhodes) and Line Roald (@lroald)

###############################################################################
# 0. Initialization
###############################################################################
using Pkg; Pkg.activate(".")
Pkg.instantiate()
Pkg.precompile()

# Load Julia Packages
#--------------------
using PowerModels
using PowerModelsWildfire
using Gurobi
using JuMP
using CSV
using DataFrames

# Load System Data
# ----------------

# load the data file
data = PowerModels.parse_file(joinpath(@__DIR__, "data", "case5_risk_sys1.m"))
load_data = CSV.read(joinpath(@__DIR__, "data", "case5_load.csv"),DataFrame)
risk_data = CSV.read(joinpath(@__DIR__, "data", "case5_risk.csv"),DataFrame)

# Modify risk weighting
# data["risk_weight"] = 0.2

# Add zeros to turn linear objective functions into quadratic ones
# so that additional parameter checks are not required
PowerModels.standardize_cost_terms!(data, order=2)

# Adds reasonable rate_a values to branches without them
PowerModels.calc_thermal_limits!(data)

# Create multiperiod data model
mp_data = PowerModels.replicate(data, size(load_data,1))

for row in eachrow(load_data)
    for col in names(row[2:end])
        col
        nwid = string(row[1])
        comp_type, comp_id = split(col,"_")
        load = row[col]

        mp_data["nw"][nwid][comp_type][comp_id]["pd"] = load
    end
end

for row in eachrow(risk_data)
    for col in names(row[2:end])
        col
        nwid = string(row[1])
        comp_type, comp_id = split(col,"_")
        risk = row[col]

        mp_data["nw"][nwid][comp_type][comp_id]["power_risk"] = risk
    end
end

# use build_ref to filter out inactive components
ref = PowerModels.build_ref(mp_data)
ref_add_on_off_va_bounds!(ref, mp_data)
ref = ref[:it][:pm]

# Note: ref contains all the relevant system parameters needed to build the OPS model
# When we introduce constraints and variable bounds below, we use the parameters in ref.

###############################################################################
# 1. Building the Optimal Power Shutoff Model
###############################################################################

# Initialize a JuMP Optimization Model
#-------------------------------------
# model = Model(GLPK.Optimizer)
model = Model(Gurobi.Optimizer)

# Add Optimization and State Variables
# ------------------------------------

# Add voltage angles va for each bus
@variable(model, va[t in keys(ref[:nw]),i in keys(ref[:nw][1][:bus])])
# note: [i in keys(ref[:nw][1][:bus])] adds one `va` variable for each bus in the network

# Add active power generation variable pg for each generator
@variable(model, pg[t in keys(ref[:nw]),i in keys(ref[:nw][1][:gen])])

# Add power flow variables p to represent the active power flow for each branch
@variable(model,  p[t in keys(ref[:nw]), (l,i,j) in ref[:nw][1][:arcs_from]])

# Build JuMP expressions for the value of p[(l,i,j)] and p[(l,j,i)] on the branches
p_expr = Dict([((t,(l,i,j)), 1.0*p[t,(l,i,j)]) for (l,i,j) in ref[:nw][1][:arcs_from], t in keys(ref[:nw])])
p_expr = merge(p_expr, Dict([((t,(l,j,i)), -1.0*p[t,(l,i,j)]) for (l,i,j) in ref[:nw][1][:arcs_from], t in keys(ref[:nw])]))
# note: this is used to make the definition of nodal power balance simpler

# Add binary variables to represent that state of all components
@variable(model, z_bus[t in keys(ref[:nw]), i in keys(ref[:nw][1][:bus])], Bin)
@variable(model, z_gen[t in keys(ref[:nw]), i in keys(ref[:nw][1][:gen])], Bin)
@variable(model, z_branch[t in keys(ref[:nw]), (l,i,j) in ref[:nw][1][:arcs_from]], Bin)

# Add continous load shed variable
@variable(model, 0.0 <= x_load[t in keys(ref[:nw]), l in keys(ref[:nw][1][:load])] <= 1.0)
@variable(model, 0.0 <= x_shunt[t in keys(ref[:nw]), s in keys(ref[:nw][1][:shunt])] <= 1.0)

# Add Objective Function
# ----------------------
# Maximize power delivery while minimizing wildfire risk

if haskey(ref[:nw][1], :risk_weight)
    alpha = ref[:nw][1][:risk_weight]
else
    @warn "network data should specify risk_weight, using 0.5 as a default"
    alpha = 0.5
end

for comp_type in [:gen, :load, :bus, :branch]
    for (nwid,nw) in ref[:nw]
        for (id,comp) in  nw[comp_type]
            if ~haskey(comp, "power_risk")
                @warn "$(comp_type) $(id) does not have a power_risk value, using 0.1 as a default"
                comp["power_risk"] = 0.1
            end
        end
    end
end

JuMP.@objective(model, Max,
    sum(
        (1-alpha)*(
                sum(x_load[t,i]*load["pd"] for (i,load) in nw[:load])
        )
        - alpha*(
            sum(z_gen[t,i]*gen["power_risk"] for (i,gen) in nw[:gen])
            + sum(z_bus[t,i]*bus["power_risk"] for (i,bus) in nw[:bus])
            + sum(z_branch[t,(l,i,j)]*nw[:branch][l]["power_risk"] for (l,i,j) in nw[:arcs_from])
            + sum(x_load[t,i]*load["power_risk"] for (i,load) in nw[:load])
        )
        for (t,nw) in ref[:nw]
    )
)

# Add Constraints
# ---------------

#  constraints for each period of OPS
for (t,nw) in ref[:nw]

    # Fix the voltage angle to zero at the reference bus
    for (i,bus) in nw[:ref_buses]
        @constraint(model, va[t,i] == 0)
    end

    # Nodal power balance constraints
    for (i,bus) in nw[:bus]
        # Build a list of the loads and shunt elements connected to the bus i
        bus_loads = [(l,nw[:load][l]) for l in nw[:bus_loads][i]]
        bus_shunts = [(s,nw[:shunt][s]) for s in nw[:bus_shunts][i]]

        # Active power balance at node i
        @constraint(model,
            sum(p_expr[(t,a)] for a in nw[:bus_arcs][i]) ==                    # sum of active power flow on lines from bus i +
            sum(pg[t,g] for g in nw[:bus_gens][i]) -                         # sum of active power generation at bus i -
            sum(x_load[t,l]*load["pd"] for (l,load) in bus_loads) -           # sum of active load * load shed  at bus i -
            sum(x_shunt[t,s]*shunt["gs"] for (s,shunt) in bus_shunts)*1.0^2   # sum of active shunt element injections * shed at bus i
        )
    end

    # Branch power flow physics and limit constraints
    for (i,branch) in nw[:branch]
        # Build the from variable id of the i-th branch, which is a tuple given by (branch id, from bus, to bus)
        f_idx = (i, branch["f_bus"], branch["t_bus"])
        p_fr = p[t,f_idx]                     # p_fr is a reference to the optimization variable p[f_idx]
        z_br = z_branch[t,f_idx]
        va_fr = va[t,branch["f_bus"]]         # va_fr is a reference to the optimization variable va on the from side of the branch
        va_to = va[t,branch["t_bus"]]         # va_fr is a reference to the optimization variable va on the to side of the branch
        # Compute the branch parameters and transformer ratios from the data
        g, b = PowerModels.calc_branch_y(branch)

        # Voltage angle difference limit
        JuMP.@constraint(model, va_fr - va_to <= branch["angmax"]*z_br + nw[:off_angmax]*(1-z_br))
        JuMP.@constraint(model, va_fr - va_to >= branch["angmin"]*z_br + nw[:off_angmin]*(1-z_br))

        # DC Power Flow Constraint
        if b <= 0
            JuMP.@constraint(model, p_fr <= -b*(va_fr - va_to + nw[:off_angmax]*(1-z_br)) )
            JuMP.@constraint(model, p_fr >= -b*(va_fr - va_to + nw[:off_angmin]*(1-z_br)) )
        else # account for bound reversal when b is positive
            JuMP.@constraint(model, p_fr >= -b*(va_fr - va_to + nw[:off_angmax]*(1-z_br)) )
            JuMP.@constraint(model, p_fr <= -b*(va_fr - va_to + nw[:off_angmin]*(1-z_br)) )
        end

        # Thermal limit
        JuMP.@constraint(model, p_fr <=  branch["rate_a"]*z_br)
        JuMP.@constraint(model, p_fr >= -branch["rate_a"]*z_br)

        # Connectivity constraint
        JuMP.@constraint(model, z_br <= z_bus[t,branch["f_bus"]])
        JuMP.@constraint(model, z_br <= z_bus[t,branch["t_bus"]])
    end

    # Generator constraints
    for (i,gen) in nw[:gen]
        # Power limit
        JuMP.@constraint(model, pg[t,i] <= gen["pmax"]*z_gen[t,i])
        JuMP.@constraint(model, pg[t,i] >= gen["pmin"]*z_gen[t,i])

        # Connectivity constraint
        JuMP.@constraint(model, z_gen[t,i] <= z_bus[t,gen["gen_bus"]])
    end

    # Load constraints
    for (i, load) in nw[:load]
        # Connectivity constraints
        JuMP.@constraint(model, x_load[t,i] <= z_bus[t,load["load_bus"]])
    end
end

# Inter-period constraints
for t in keys(ref[:nw])
    if t != 1
        # if component is disabled, it must stay disabled for future periods
        for (i,branch) in ref[:nw][t][:branch]
            f_idx = (i, branch["f_bus"], branch["t_bus"])
            JuMP.@constraint(model, z_branch[t-1,f_idx] >= z_branch[t,f_idx])
        end

        for (i,bus) in ref[:nw][t][:bus]
            JuMP.@constraint(model, z_bus[t-1,i] >= z_bus[t,i])
        end
        for (i,gen) in ref[:nw][t][:gen]
            JuMP.@constraint(model, z_gen[t-1,i] >= z_gen[t,i])
        end
    end
end

###############################################################################
# 3. Solve the Optimal Power Flow Model and Review the Results
###############################################################################

# Solve the optimization problem
optimize!(model)

# Check that the solver terminated without an error
println("The solver termination status is $(termination_status(model))")

# Check the value of the objective function
println("The objective value is $(objective_value(model))")
total_load = sum(sum(value(x_load[t,i])*load["pd"] for (i,load) in nw[:load]) for (t,nw) in ref[:nw])
println("The load served is $total_load ")
total_risk = sum(
    sum(value(z_branch[t,(l,i,j)])*nw[:branch][l]["power_risk"] for (l,i,j) in nw[:arcs_from]) +
    sum(value(z_gen[t,g])*gen["power_risk"] for (g,gen) in nw[:gen]) +
    sum(value(z_bus[t,i])*bus["power_risk"] for (i,bus) in nw[:bus]) +
    sum(value(x_load[t,i])*load["power_risk"] for (i,load) in nw[:load])
    for (t,nw) in ref[:nw]
)
println("The system risk is $total_risk")

# Get inactive devices
for t in sort(collect(keys(ref[:nw])))
    nw = ref[:nw][t]
    for l in sort(collect(keys(nw[:branch])))
        f_idx = (l, nw[:branch][l]["f_bus"], nw[:branch][l]["t_bus"])
        value(z_branch[t,f_idx])==0 ? println("Branch $l is disabled in period $t.") : false
    end
    for g in sort(collect(keys(nw[:gen])))
        value(z_gen[t,g])==0 ? println("Generator $g is disabled in period $t.") : false
    end
    for i in sort(collect(keys(nw[:bus])))
        value(z_bus[t,i])==0 ? println("Bus $i is disabled in period $t.") : false
    end
    for i in sort(collect(keys(nw[:load])))
        value(x_load[t,i])!=1 ? println("Load $i is serving $(value(x_load[t,i])*100)% of demand in period $t.") : false
    end
end


#= extract_pcd_points.jl
extract the radar-points in pcd file(ascii) =#

using XLSX;

prj_path = dirname(dirname(@__FILE__))
data_path = prj_path * raw"\data\data_no_anno\radar_ascii";

xls_file_path = prj_path * raw"\data\radar_process\radar_points.xlsx";
XLSX.openxlsx(xls_file_path, mode="w") do xf
    for file_id in 1:38
        file_path = data_path * "\\$file_id.pcd"
        f = open(file_path, "r");
        n = countlines(f)

        # skip the first 12 lines
        seekstart(f)
        for i = 1:12
            readline(f)
        end

        data = Array{Float64}(undef, n - 12, 18)
        for i = 13:n
            row_id = i - 12;
            str_list = split(readline(f), " ")

            x = parse(Float64, str_list[1]);
            data[row_id, 1] = x;

            y = parse(Float64, str_list[2]);
            data[row_id, 2] = y;

            z = parse(Float64, str_list[3]);
            data[row_id, 3] = z;

            dyn_prop = parse(Float64, str_list[4]);
            data[row_id, 4] = dyn_prop;

            id = parse(Float64, str_list[5]);
            data[row_id, 5] = id;

            rcs = parse(Float64, str_list[6]);
            data[row_id, 6] = rcs;

            vx = parse(Float64, str_list[7]);
            data[row_id, 7] = vx;

            vy = parse(Float64, str_list[8]);
            data[row_id, 8] = vy;

            vx_comp = parse(Float64, str_list[9]);
            data[row_id, 9] = vx_comp;

            vy_comp = parse(Float64, str_list[10]);
            data[row_id, 10] = vy_comp;

            is_quality_valid = parse(Float64, str_list[11]);
            data[row_id, 11] = is_quality_valid;

            ambig_state = parse(Float64, str_list[12]);
            data[row_id, 12] = ambig_state;

            x_rms = parse(Float64, str_list[13]);
            data[row_id, 13] = x_rms;

            y_rms = parse(Float64, str_list[14]);
            data[row_id, 14] = y_rms;

            invalid_state = parse(Float64, str_list[15]);
            data[row_id, 15] = invalid_state;

            pdh0 = parse(Float64, str_list[16]);
            data[row_id, 16] = pdh0;

            vx_rms = parse(Float64, str_list[17]);
            data[row_id, 17] = vx_rms;

            vy_rms = parse(Float64, str_list[18]);
            data[row_id, 18] = vy_rms;
        end

        if file_id == 1
            sheet = xf[file_id];
            XLSX.rename!(sheet, "$file_id.pcd");
        else
            XLSX.addsheet!(xf, "$file_id.pcd")
        end
        sheet = xf["$file_id.pcd"]

        sheet["A1", dim=2] = [
            "x";
            "y ";
            "z ";
            "dyn_prop ";
            "id ";
            "rcs ";
            "vx ";
            "vy ";
            "vx_comp";
            "vy_comp ";
            "is_quality_valid";
            "ambig_state";
            "x_rms";
            "y_rms ";
            "invalid_state";
            "pdh0";
            "vx_rms";
            "vy_rms";
        ];
        sheet["A2:R$(n - 11)"] = data;
    end
end


"""
Main code for running the branching process. This invokes other Julia scripts in order
to run:
    - read_in_TP.jl:
        Code for reading the TP csv files saved from the fitting and
        forecasting parts of the model.
    - helper_functions.jl:
        Code which is commonly reused in the model. This involves non-specialised functions
        for sampling onset and infection times, and for sampling from a NegativeBinomial
        and Binomial distribution using asymptotic limiting properties.
    - forecast_types.jl
        Code for some user defined types. This lets us group variables of importance. There
        is a struct for the dynamical constants, properties of the data and a
        simple mapping for moving between types and indices.
    - simulation_consistency.jl
        Code for checking consistency of simulations against the data. The code in this file
        is also responsible for injecting cases into the simulation if we are below some
        threshold.
    - assumptions.jl
        One of the most important files. Keeps track of all assumptions in a function that
        is used to initialise the system constants.

The key logic of this simulation is featured in the function simulate_branching_process().
This is written in such a way that it is relatively straightforward to follow the logic of
the processes and functions are called in a systematic way. simulate_branching_process()
will return reshaped arrays that are the good simulated realisations based off
the consistency checks.
"""

using Distributed
using Distributions
using Random
using Statistics
using LinearAlgebra
using ProgressBars
using StaticArrays
using SharedArrays
using ProgressMeter

include("helper_functions.jl")
include("forecast_types.jl")
include("simulation_consistency.jl")

##### Methods
function initialise_population!(
    forecast::Forecast,
    D0,
)
    """
    Fill the infection array for each simulation, inititialising the states of
    all the simulations.
    """
    Z = forecast.sim_realisation.Z
    Z_historical = forecast.sim_realisation.Z_historical
    D = forecast.sim_realisation.D
    U = forecast.sim_realisation.U

    # extract prob of being symptomatic and prob of symptoms given detection
    # NOTE: this currently assumes we are starting in the Delta wave (which is true as
    # of 31/1/2022) but we will need to update this if we shift the starting date into
    # an only Omicron period.
    T_end = forecast.sim_features.T_end
    individual_type_map = forecast.individual_type_map

    # assume detection probabilities for initialisation is whatever it was at day 0
    p_detect_given_symp = forecast.sim_constants.p_detect_given_symp[map_day_to_index_p(0)]
    p_symp = forecast.sim_constants.p_symp[map_day_to_index_p(0)]

    # infer some initial undetected symptomatic
    num_S_undetected = 0

    if D0.S == 0
        num_S_undetected = rand(NegativeBinomial(1, p_detect_given_symp))
    else
        num_S_undetected = rand(NegativeBinomial(D0.S, p_detect_given_symp))
    end

    # infer some initial undetected asymptomatic
    total_S = D0.S + num_S_undetected
    num_A_undetected = 0

    if total_S == 0
        num_A_undetected = rand(NegativeBinomial(1, p_symp))
    else
        num_A_undetected = rand(NegativeBinomial(total_S, p_symp))
    end

    # total initial cases include the undetected
    total_initial_cases = sum(D0)
    total_initial_cases += num_S_undetected + num_A_undetected

    start_day = 0

    # reset arrays
    D .= 0
    U .= 0
    Z .= 0
    Z_historical .= 0

    if total_initial_cases > 0
        inf_times = MVector{total_initial_cases}(
            zeros(Int, total_initial_cases)
        )

        # sample negative infection times as these individuals are assumed to occur
        # before the simulation
        for i in 1:total_initial_cases
            inf_times[i] = -sample_onset_time(false)
        end

        # the starting day is the earliest infection time deemed to contribute to the sim.
        start_day = minimum(inf_times)
        idx = 1

        for _ in 1:D0.S
            Z[map_day_to_index_Z(inf_times[idx]), individual_type_map.S] += 1
            idx += 1
            D[1, individual_type_map.S] += 1
        end

        for _ in 1:D0.A
            Z[map_day_to_index_Z(inf_times[idx]), individual_type_map.A] += 1
            idx += 1
            D[1, individual_type_map.A] += 1
        end

        for _ in 1:num_S_undetected
            Z[map_day_to_index_Z(inf_times[idx]), individual_type_map.S] += 1
            idx += 1
            U[1, individual_type_map.S] += 1
        end

        for _ in 1:num_A_undetected
            Z[map_day_to_index_Z(inf_times[idx]), individual_type_map.A] += 1
            idx += 1
            U[1, individual_type_map.A] += 1
        end
    end

    return start_day

end

function import_cases_model!(
	forecast::Forecast,
    import_cases,
	forecast_start_date,
)
	"""
	A model to handle international imports in the forecasting of cases. This is run
    prior to the standard forecasting model and will pre-fill D, U and Z for each
    simulation.
	We assume that cases arise from
	D[t] ~ NegBin(a[t], 1/(b+1))
	where
	a[t] = alpha + f(I[t])
	b = beta + 1
	are the parameters. The function f() is an exponential weighted moving average such
    that
	f(I[t]) = phi I[t] + (1-phi) f(I[t-1])
	and alpha, beta are hyperparameters fixed at 0.5 and 0.2, respectively. We
    take b = β + 1 as we assume a fixed period of 1 day for estimating the posterior.
	"""

    Z = forecast.sim_realisation.Z
    D = forecast.sim_realisation.D
    U = forecast.sim_realisation.U

    T_observed = forecast.sim_features.T_observed
    T_end = forecast.sim_features.T_end

	# number of days after t = 0 (2020/03/01)
	days_into_forecast = Day(Date(forecast_start_date) - Date("2020-03-01")).value

    # grab the correct parameters for the particular dominant strain
    p_detect_import = 0.0

	# the number of days the simulation is run for
	forecast_days_plus_padding = -10:T_end
	# posterior a, b for sampling number of imports on day t
	a_post = MVector{length(forecast_days_plus_padding)}(
        zeros(Float64, length(forecast_days_plus_padding))
    )

	# current ema is taken 11 days before the forecast start date
	current_ema = zero(Float64)

    # imports show up in the data before
    # import_infections = import_cases[5:end]
    inc_period_shift = 5
    current_ema += import_cases[inc_period_shift + days_into_forecast - 11]

	for (i, t) in enumerate(forecast_days_plus_padding)
		# if within observation period, update the posterior. Otherwise use the last
		# available measurement
		if t <= T_observed - inc_period_shift
			count_on_day = import_cases[inc_period_shift + days_into_forecast + t]
			current_ema = forecast.sim_constants.ϕ * count_on_day +
                (1 - forecast.sim_constants.ϕ) * current_ema
			a_post[i] = forecast.sim_constants.prior_alpha + current_ema
		else
			a_post[i] = forecast.sim_constants.prior_alpha + current_ema
		end
	end

    # assuming a fixed period of 1 day (so this is the same for the whole period)
    b_post = forecast.sim_constants.prior_beta + 1

    for (i, t) in enumerate(forecast_days_plus_padding)
        # sample detected imports
        D_I = rand(NegativeBinomial(a_post[i], b_post / (1 + b_post)))
        # see whether they're observed
        unobserved_I = D_I == 0 ? 1 : D_I
        # based on the switch to omicron (in terms of infections so -5 days)
        if map_day_to_index_p(t) <= length(forecast.sim_constants.p_detect_import)
            p_detect_import = forecast.sim_constants.p_detect_import[map_day_to_index_p(t)]
        else
            p_detect_import = forecast.sim_constants.p_detect_import[end]
        end
        U_I = rand(NegativeBinomial(unobserved_I, p_detect_import))

        # day of infection is time t which is why the infection time in
        if D_I + U_I > 0
            assign_to_arrays_and_times!(
                forecast,
                t,
                is_omicron = false,
                num_I_detected = D_I,
                num_I_undetected = U_I,
            )
        end
    end

	return nothing

end

# function calculate_sus_dep_factor(
#     susceptible_depletion_sim,
#     Z_historical,
#     day,
#     N,
# )

#     ψ = 0.0

#     for (i, τ) in enumerate((30, 60, 90, 120))
#         # guard clauses for when we do not have enough observations to consider the longer term
#         # effects (in other words, all the infections are relatively recent)
#         map_day_to_index_Z(day) < 30 && continue

#         # Get starting and ending index. Start must be at least the beginning of the array and end
#         # must be at least the forecast length.
#         τ_start = max(map_day_to_index_Z(day) - (τ - 1), 1)
#         τ_end = min(map_day_to_index_Z(day) - (τ - 30), size(Z_historical, 1))

#         # sum infections over the time period of interest
#         infections = sum(@view Z_historical[τ_start:τ_end])

#         # we index from 4 - 1 as the largest reduction occurs as a result of recent infections
#         prop_inf = min(1.0, infections / N)
#         ψ += susceptible_depletion_sim[5 - i] * prop_inf
#     end

#     scale_factor = max(0.0, 1.0 - ψ)

#     return scale_factor

# end

function calculate_sus_dep_factor(
    susceptible_depletion_sim,
    Z_historical,
    day,
    N,
)

    ψ = 0.0

    τ_horizons_start = (0, 14, 28, 42, 72, 102)
    τ_horizons_end = (0, 13, 27, 41, 71, 101, 131)

    for i in 1:length(τ_horizons_start)
        # guard clause to ensure we have a reasonable number of observations to calculate the
        # contribution to ψ from period i
        map_day_to_index_Z(day) <= τ_horizons_start[i] && continue

        # Get starting and ending index. Start must be at least the beginning of the array and end
        # must be at least the forecast length.
        τ_start = max(map_day_to_index_Z(day) - τ_horizons_end[i + 1], 1)
        τ_end = min(map_day_to_index_Z(day) - τ_horizons_start[i], size(Z_historical, 1))

        # sum infections over the time period of interest
        infections = sum(@view Z_historical[τ_start:τ_end])

        # we index from 4 - 1 as the largest reduction occurs as a result of recent infections
        prop_inf = min(1.0, infections / N)
        ψ += susceptible_depletion_sim[7 - i] * prop_inf
    end

    scale_factor = max(0.0, 1.0 - ψ)

    return scale_factor

end

# function calculate_sus_dep_factor(
#     susceptible_depletion_sim,
#     Z_historical,
#     day,
#     N,
# )

#     ψ = 0.0

#     τ_end = map_day_to_index_Z(day)

#     infections = sum(@view Z_historical[begin:τ_end])
#     ψ = susceptible_depletion_sim * infections / N

#     scale_factor = max(0.0, 1.0 - ψ)

#     return scale_factor

# end

function sigmoid(t, prop_pars)
    (m₀, m₁, r, τ) = prop_pars
    res = m₀ + (m₁ - m₀) / (1 + exp(-r * (t - τ)))

    return res

end

function sample_offspring!(
    forecast::Forecast,
    day,
    TP,
    prop_pars,
    susceptible_depletion_sim,
    local_cases;
    adjust_TP = false,
)
    """
    Sample offspring for all adults at time t. This uses the fact that the sum of
    N NB(s, p) is NB(N*s, p) and hence we can sample ALL offspring for a particular type
    of adults.
    """

    Z = forecast.sim_realisation.Z
    Z_historical = forecast.sim_realisation.Z_historical
    D = forecast.sim_realisation.D
    U = forecast.sim_realisation.U
    omicron_start_day = forecast.sim_features.omicron_start_day
    omicron_only_day = forecast.sim_features.omicron_only_day
    T_observed = forecast.sim_features.T_observed
    state = forecast.sim_features.state

    # extract the CAR related quantities
    p_detect = forecast.sim_constants.p_detect
    p_symp = forecast.sim_constants.p_symp[map_day_to_index_Z(day)]
    p_detect_given_symp = forecast.sim_constants.p_detect_given_symp[map_day_to_index_Z(day)]
    p_detect_given_asymp = forecast.sim_constants.p_detect_given_asymp[map_day_to_index_Z(day)]

    # extract fixed simulation values
    N = forecast.sim_features.N

    # pointer to the current day's infections
    Z_tmp = @view Z[map_day_to_index_Z(day), :]

    total_parents = sum(Z_tmp)

    if total_parents > 0
        (S_parents, A_parents, I_parents) = Z_tmp

        # sum up the number local infections on the current day
        Z_historical[map_day_to_index_Z(day)] += S_parents + A_parents
        # Z_tmp gets emptied at the end of sampling so we won't be double counting
        # zero out infections on day t so we can restart sim and keep the progress
        Z_tmp .= 0

        S_parents_omicron = 0
        A_parents_omicron = 0
        I_parents_omicron = 0

        if day >= omicron_start_day && day < omicron_only_day
            # prior to 2022-02-01, there is some probability of non-Omicron (Delta) cases
            # number of days into the omicron wave
            t_omicron = day - omicron_start_day
            prop_omicron = sigmoid(t_omicron, prop_pars)
            # sample a random number of Omicron parents
            S_parents_omicron = rand(Binomial(S_parents, prop_omicron))
            A_parents_omicron = rand(Binomial(A_parents, prop_omicron))
            I_parents_omicron = rand(Binomial(I_parents, prop_omicron))
        elseif day >= omicron_only_day
            # after 2022-02-01 everything is just Omicron
            S_parents_omicron = S_parents
            A_parents_omicron = A_parents
            I_parents_omicron = I_parents
        end

        # put the number of parents in a tuple to iterate over
        num_parents = (
            S_parents - S_parents_omicron,
            A_parents - A_parents_omicron,
            I_parents - I_parents_omicron,
            S_parents_omicron,
            A_parents_omicron,
            I_parents_omicron,
        )

        # adjust the TP/Reff by susceptible depletion based off the current simulation state
        scale_factor = calculate_sus_dep_factor(
            susceptible_depletion_sim,
            Z_historical,
            day,
            N,
        )

        for (i, n_parents) in enumerate(num_parents)
            # a guard clause for skipping the current iteration
            n_parents == 0 && continue

            # reset counters
            TP_parent = 0.0
            k = 0.0
            # relative infectiousness of asymptomatic
            γ = 0.5

            # essentially a (gross) switch statement
            if i == 1
                α_s = 1 / (p_symp + γ * (1 - p_symp))
                TP_parent = TP.TP_local_delta_parent * α_s
                k = 0.15
            elseif i == 2
                α_s = 1 / (p_symp + γ * (1 - p_symp))
                α_a = γ * α_s
                TP_parent = TP.TP_local_delta_parent * α_a
                k = 0.15
            elseif i == 3
                TP_parent = TP.TP_import_delta_parent
                k = 0.15
            elseif i == 4
                α_s = 1 / (p_symp + γ * (1 - p_symp))
                α_a = γ * α_s
                TP_parent = TP.TP_local_omicron_parent * α_s
                k = 0.5
            elseif i == 5
                α_s = 1 / (p_symp + γ * (1 - p_symp))
                α_a = γ * α_s
                TP_parent = TP.TP_local_omicron_parent * α_a
                k = 0.6
            elseif i == 6
                TP_parent = TP.TP_import_omicron_parent
                k = 0.6
            end

            # scale the TP by depletion of susceptibles
            # println("Day: ", day)
            # println("Scale factor: ", scale_factor)
            TP_parent = max(0, TP_parent * scale_factor)

            TP_parent == 0 && continue

            num_offspring = sample_negative_binomial_limit(TP_parent * n_parents, k * n_parents)

            if num_offspring > 0
                is_omicron = false

                if i in 1:3
                    is_omicron = false
                elseif i in 4:6
                    is_omicron = true
                end

                num_S_offspring = sample_binomial_limit(num_offspring, p_symp)
                num_A_offspring = num_offspring - num_S_offspring

                num_S_offspring_detected = sample_binomial_limit(
                    num_S_offspring, p_detect_given_symp
                )
                num_A_offspring_detected = sample_binomial_limit(
                    num_A_offspring, p_detect_given_asymp
                )

                num_S_offspring_undetected = num_S_offspring - num_S_offspring_detected
                num_A_offspring_undetected = num_A_offspring - num_A_offspring_detected

                assign_to_arrays_and_times!(
                    forecast,
                    day,
                    is_omicron = is_omicron,
                    num_S_detected = num_S_offspring_detected,
                    num_S_undetected = num_S_offspring_undetected,
                    num_A_detected = num_A_offspring_detected,
                    num_A_undetected = num_A_offspring_undetected,
                )
            end
        end
    end

    return nothing

end

function assign_to_arrays_and_times!(
    forecast::Forecast,
    day;
    is_omicron = false,
    num_S_detected = 0,
    num_S_undetected = 0,
    num_A_detected = 0,
    num_A_undetected = 0,
    num_I_detected = 0,
    num_I_undetected = 0,
)
    """
    Increment the arrays D and U based on sampled offsrping counts. The number of
    each offspring are keyword parameters which are default 0. This allows us to call
    this function from multiple spots.
    """

    Z = forecast.sim_realisation.Z
    D = forecast.sim_realisation.D
    U = forecast.sim_realisation.U
    individual_type_map = forecast.individual_type_map
    T_end = forecast.sim_features.T_end

    # sampling detected cases
    for _ in 1:num_S_detected
        (inf_time, onset_time) = sample_times(day, is_omicron)

        if inf_time < T_end
            Z[map_day_to_index_Z(inf_time), individual_type_map.S] += 1
        end

        if onset_time < T_end
            D[map_day_to_index_UD(onset_time), individual_type_map.S] += 1
        end
    end

    for _ in 1:num_A_detected
        (inf_time, onset_time) = sample_times(day, is_omicron)

        if inf_time < T_end
            Z[map_day_to_index_Z(inf_time), individual_type_map.A] += 1
        end

        if onset_time < T_end
            D[map_day_to_index_UD(onset_time), individual_type_map.A] += 1
        end
    end

    # sampling undetected cases
    for _ in 1:num_S_undetected
        (inf_time, onset_time) = sample_times(day, is_omicron)

        if inf_time < T_end
            Z[map_day_to_index_Z(inf_time), individual_type_map.S] += 1
        end

        if onset_time < T_end
            U[map_day_to_index_UD(onset_time), individual_type_map.S] += 1
        end
    end

    for _ in 1:num_A_undetected
        (inf_time, onset_time) = sample_times(day, is_omicron)

        if inf_time < T_end
            Z[map_day_to_index_Z(inf_time), individual_type_map.A] += 1
        end

        if onset_time < T_end
            U[map_day_to_index_UD(onset_time), individual_type_map.A] += 1
        end
    end

    # Sampling imports. These are treated slightly differently as the infection
    # time is inferred through the import model during initialisation
    for _ in 1:num_I_detected
        inf_time = day
        if inf_time < T_end
            Z[map_day_to_index_Z(inf_time), individual_type_map.I] += 1
        end

        onset_time = inf_time + sample_onset_time(is_omicron)
        if onset_time < T_end
            # time+1 as first index is day 0
            D[map_day_to_index_UD(onset_time), individual_type_map.I] += 1
        end
    end

    for _ in 1:num_I_undetected
        inf_time = day
        if inf_time < T_end
            Z[map_day_to_index_Z(inf_time), individual_type_map.I] += 1
        end

        onset_time = inf_time + sample_onset_time(is_omicron)
        if onset_time < T_end
            # time+1 as first index is day 0
            U[map_day_to_index_UD(onset_time), individual_type_map.I] += 1
        end
    end

    return nothing

end

function clean_up_sim_results(
    Z_historical_results,
    D_results,
    U_results,
    good_sims,
)
    """
    A simple function that takes in the simulation output, cleans it up and spits out the
    processed format. We also account for the effect of susceptible depletion in the TP
    here so that the output includes the adjusted TP.
    """
    # return (
    #     Z_historical_results,
    #     D_results,
    #     U_results,
    # )
    return (
        Z_historical_results[:, 1:good_sims - 1],
        D_results[:, :, 1:good_sims - 1],
        U_results[:, :, 1:good_sims - 1],
    )
end

function simulate_branching_process(
    D0,
    N,
    nsims,
    local_cases,
    import_cases,
    forecast_start_date,
    date,
    state;
    adjust_TP = false,
)
    """
    Simulate branching process nsims times conditional on the cumulative observed cases,
    with initial state X0 = (S, A, I) (a named tuple). T is the duration of the simulation
    and N is the population size which is currently unused.
    """
    # read in the TPs for the state and strains
    (TP_indices, TP_local_delta, TP_import_delta) = get_single_state_TP(
        forecast_start_date,
        date,
        state,
        strain = "Delta",
        adjust_TP = adjust_TP,
    )

    (_, TP_local_omicron, TP_import_omicron) = get_single_state_TP(
        forecast_start_date,
        date,
        state,
        strain = "Omicron",
        adjust_TP = adjust_TP,
    )

    susceptible_depletion = read_in_susceptible_depletion(date)

    omicron_prop_pars = read_in_prop_omicron(date, state)

    max_restarts = 10
    good_sims = ones(Bool, nsims)

    # calculate the upper and lower limits for checking sims against data
    sim_features = get_simulation_limits(
        local_cases,
        forecast_start_date,
        TP_indices,
        N,
        state,
    )

    # the range of days for the simulation -1 as the "days", day ↦ day + 1
    day_range = -30:sim_features.T_end - 1

    earliest_start_date = Dates.Date(forecast_start_date) - Dates.Day(30)
    forecast_end_date = Dates.Date(forecast_start_date) + Dates.Day(sim_features.T_end - 1)

    # get simulation constants
    (sim_constants, individual_type_map) = set_simulation_constants(
        earliest_start_date,
        forecast_end_date,
        state,
    )

    # initialise state arrays (T_end + 1 as we include 0 day)
    sim_realisation = Realisation(sim_features.T_end)
    # always targeting 2000 sims
    sim_results = Results(sim_features.T_end, 2000)

    # put everything into a forecast object
    forecast = Forecast(
        sim_features,
        sim_realisation,
        sim_results,
        sim_constants,
        individual_type_map,
    )

    good_TPs_inds = zeros(Int, 2000)

    # Applying a reduction in NSW TP based off observations of the model fit. This issue is
    # at the beginning of the third wave so doesn't really influence the results for
    # Omicron. I DON'T LIKE THIS BUT FOR NOW IT'S CORRECTING SOME ISSUE IN THE TP/REFF
    # affecting NSW only.
    if state == "NSW" || state == "VIC"
        TP_ind = findfirst(40 == ind for ind in TP_indices)
        TP_local_delta[1:TP_ind, :] *= 0.5
        TP_local_omicron[1:TP_ind, :] *= 0.5
    end

    # number of TP samples saved
    num_TP = size(TP_local_delta, 2)

    # counts in each window are the same size as the limit arrays
    max_cases = forecast.sim_features.max_cases
    case_counts = zero(similar(max_cases))

    # a counter for the good sims, will terminate if good_sims > 2000
    good_sims = 1

    Z_results = forecast.sim_results.Z
    D_results = forecast.sim_results.D
    U_results = forecast.sim_results.U
    Z_historical_results = forecast.sim_results.Z_historical

    @showprogress for sim in 1:nsims
        # sample the TP/susceptible_depletion for this sim
        TP_ind = sim % num_TP == 0 ? num_TP : sim % num_TP
        susceptible_depletion_sim = susceptible_depletion[TP_ind, :]
        # susceptible_depletion_sim = susceptible_depletion[TP_ind]

        # reset boolean flags
        bad_sim = false
        injected_cases = false

        start_day = initialise_population!(
            forecast,
            D0,
        )

        import_cases_model!(
            forecast,
            import_cases,
            forecast_start_date,
        )

        # initialise counter to first day, when earliest infection occured based on
        # initialisation
        day = start_day

        # find index of the first TP to use based on the indices representing time
        # from forecast origin
        TP_start_day_ind = findfirst(start_day == d for d in TP_indices)
        TP_day_ind = TP_start_day_ind

        last_injection = 0
        n_restarts = 0

        while day <= day_range[end]
            # get TP's for the current day under the two strains and put in a named tuple
            TP_day = (
                TP_local_delta_parent = TP_local_delta[TP_day_ind, TP_ind],
                TP_import_delta_parent = TP_import_delta[TP_day_ind, TP_ind],
                TP_local_omicron_parent = TP_local_omicron[TP_day_ind, TP_ind],
                TP_import_omicron_parent = TP_import_omicron[TP_day_ind, TP_ind],
            )
            # get the sampled parameters on a given day
            prop_pars_day = (
                m0 = omicron_prop_pars["m0"][TP_ind],
                m1 = omicron_prop_pars["m1"][TP_ind],
                # m1 = 1.0,
                r = omicron_prop_pars["r"][TP_ind],
                tau = omicron_prop_pars["tau"][TP_ind],
            )

            sample_offspring!(
                forecast,
                day,
                TP_day,
                prop_pars_day,
                susceptible_depletion_sim,
                local_cases;
                adjust_TP = adjust_TP
            )

            reinitialise_allowed = day >= 3 &&
                day < forecast.sim_features.T_observed - 3 &&
                n_restarts < max_restarts &&
                day - last_injection >= 3
            # reinitialise_allowed = day >= 7 &&
            #     day < forecast.sim_features.T_observed - 7 &&
            #     n_restarts < max_restarts

            (bad_sim, injected_cases) = check_sim!(
                forecast,
                case_counts,
                local_cases,
                reinitialise_allowed;
                day = day,
                sim = sim,
            )

            if !injected_cases
                day += 1
                TP_day_ind += 1
            else
                n_restarts += 1
                if n_restarts > max_restarts
                    bad_sim = true
                end
                last_injection = day
                day = start_day
                TP_day_ind = TP_start_day_ind
                injected_cases = false
            end

            bad_sim && break

        end

        if !bad_sim
            Z_historical_results[:, good_sims] .= forecast.sim_realisation.Z_historical
            D_results[:, :, good_sims] .= forecast.sim_realisation.D
            U_results[:, :, good_sims] .= forecast.sim_realisation.U
            good_sims += 1
        end

        if good_sims > 2000
            break
        end

    end

    println("##########")
    println(state, " finished with ", good_sims - 1, " good simulations.")
    println("##########")

    results = clean_up_sim_results(
        Z_historical_results,
        D_results,
        U_results,
        good_sims,
    )

    return results

end


abstract type AbstractChargeDriftModel{T <: SSDFloat} end
abstract type AbstractTemperatureModel{T <: SSDFloat} end

"""
    get_electron_drift_field(ef::Array{SVector{3, T},3}, chargedriftmodel::AbstractChargeDriftModel)::Array{SVector{3,T},3} where {T <: SSDFloat}

Applies the charge drift model onto the electric field vectors. The field vectors have to be in cartesian coordinates.
"""
function get_electron_drift_field(ef::Array{SVector{3, T},3}, chargedriftmodel::AbstractChargeDriftModel;
            use_nthreads::Int = Base.Threads.nthreads())::Array{SVector{3,T},3} where {T <: SSDFloat}
    df = Array{SVector{3,T}, 3}(undef, size(ef))
    @onthreads 1:use_nthreads for i3 in workpart(axes(df, 3), 1:use_nthreads, Base.Threads.threadid())
        for i2 in axes(df, 2)
            for i1 in axes(df, 1)
                @inbounds df[i1, i2, i3] = getVe(ef[i1, i2, i3], chargedriftmodel)
            end
        end
    end
    return df
end

"""
    get_hole_drift_field(ef::Array{SVector{3, T},3}, chargedriftmodel::AbstractChargeDriftModel)::Array{SVector{3,T},3} where {T <: SSDFloat}

Applies the charge drift model onto the hole field vectors. The field vectors have to be in cartesian coordinates.
"""
function get_hole_drift_field(ef::Array{SVector{3,T},3}, chargedriftmodel::AbstractChargeDriftModel;
            use_nthreads::Int = Base.Threads.nthreads())::Array{SVector{3,T},3} where {T <: SSDFloat}
    df = Array{SVector{3,T}, 3}(undef, size(ef))
    @onthreads 1:use_nthreads for i3 in workpart(axes(df, 3), 1:use_nthreads, Base.Threads.threadid())
        for i2 in axes(df, 2)
            for i1 in axes(df, 1)
                @inbounds df[i1, i2, i3] = getVh(ef[i1, i2, i3], chargedriftmodel)
            end
        end
    end
    return df
end

include("Vacuum/Vacuum.jl")
include("ADL/ADL.jl")


module ProcessLibrariesIO

export
dependencies_by_month,
monthly_adjacency_matrices,
get_depversions_and_adj_mats,
make_depversions_and_adj_mats

using Serialization
using JuliaDB

import Dates
import LightGraphs

"""Add Project_Node and Dependency_Project_Node columns to `deps`

Node IDs are natural numbers between 1 and {number of unique IDs}
"""
function _IDs_to_nodes(deps)
    project_ids = vcat(select(deps, :Project_ID), select(deps, :Dependency_Project_ID)) |> unique
    node_numbers = Dict(project_ids[i] => i for i in 1:length(project_ids))

    deps = setcol(deps, :Project_Node, :Project_ID => pid->node_numbers[pid])
    setcol(deps, :Dependency_Project_Node, :Dependency_Project_ID => pid->node_numbers[pid])
end

"""All dependencies of the latest version of each project in each month"""
function dependencies_by_month(versions, dependencies)
    parsed_dates = select(versions, :Published_Timestamp) .|>
        x -> Dates.DateTime(x, Dates.dateformat"y-m-d H:M:S \U\T\C")
    versions = setcol(versions, :Timestamp, parsed_dates)
    versions = setcol(versions, :Month, :Timestamp => ts -> floor(ts, Dates.Month))

    # Latest ID for each project for each month:
    # Sort by timestamp, then number
    latest_version_each_month = groupby(
        (Version_ID=x->sort(x)[end].ID,),
        versions, (:Project_ID, :Month),
        select=(:Timestamp, :Number, :ID))

    # Inner join to get only the deps that match a selected versionID
    # We lose packages here if they don't have any dependents and aren't depended on by anything.
    lonely_packages = setdiff(
        versions.columns.Project_ID,
        dependencies.columns.Project_ID,
        dependencies.columns.Dependency_Project_ID)

    if length(lonely_packages) > 0
        @warn "$(length(lonely_packages)) unconnected packages omitted"
    end

    depversions = join(
        select(latest_version_each_month, (:Version_ID, :Month)),
        dependencies, lkey=:Version_ID, rkey=:Version_ID)

    # Just need to rekey all project_ids to natural numbers and reindex
    depversions |> _IDs_to_nodes |> t -> reindex(t, :Project_Node)
end

"Number of projects in the table"
function number_of_nodes(depversion)
    select(depversion, (:Project_Node, :Dependency_Project_Node)) |>
        maximum ∘ maximum
end

"DiGraph from dependency table"
function graph_from_table(depversion,
                          nnodes = number_of_nodes(depversion))
    g = LightGraphs.SimpleDiGraph(nnodes)
    for row in depversion
        # Create edge from dependency to dependent
        LightGraphs.add_edge!(g, row.Dependency_Project_Node, row.Project_Node)
    end
    return g
end

"""
Make an adjacency matrix from each month.

The nth column contains the dependents of package n. The nth row contains the dependees.
"""
function latest_at_month(deps, month)
    deps = filter(row->row.Month <= month, deps)
    # This allocates memory. Could compute row indexes if we need to.
    groupby(
        (x = grp -> filter(row -> row.Month == maximum(grp.Month), grp),),
        deps, :Project_Node) |> flatten
end

"One project adjacency matrix for each month in deps"
function monthly_adjacency_matrices(deps)
    months = deps.columns.Month
    months = minimum(months):Dates.Month(1):maximum(months)
    nnodes = number_of_nodes(deps)

    # This is slow - do it in parallel

    # optimize grouping in latest_at_month
    if colnames(deps)[1] !== :Project_Node
        deps = reindex(deps, :Project_Node)
    end

    # Drop columns we don't need to avoid allocating excess memory
    deps = select(deps, (:Project_Node, :Month, :Dependency_Project_Node))

    # pre-allocate
    result = Array{Any}(undef, length(months))

    Threads.@threads for i in 1:length(months)
        # Can't use enumerate with @threads
        month = months[i]
        result[i] = latest_at_month(deps, month) |>
            t -> graph_from_table(t, nnodes) |>
            LightGraphs.adjacency_matrix
    end
    result
end

function tensor(adj_mats)
    # Guessing that float float calculations are faster, but haven't checked.
    X = Array{Float64}(undef, (size(adj_mats[1])..., length(adj_mats)))
    [X[:,:,i]=f for (i,f) in adj_mats|>enumerate]
    X
end

# These should obviously be a Make/Drake/DAG thing, but I haven't bothered to make that.

function filter_to_platform(table, platform)
    filter(row -> row.Platform == platform, table)
end

function platpath(platform)
    path = "../data/processed/pacman/$platform/"
    run(`mkdir -p $path`)
    path
end

function make_depversions_and_adj_mats(platform)
    versions = loadtable("../data/sample-1.4/$(platform)_versions-1.4.0-2018-12-22.csv")
    dependencies = loadtable("../data/sample-1.4/$(platform)_dependencies-1.4.0-2018-12-22.csv", type_detect_rows=4000)

    versions = filter_to_platform(versions, platform)
    dependences = filter_to_platform(dependencies, platform)

    depversions = ProcessLibrariesIO.dependencies_by_month(versions, dependencies)
    adj_mats = ProcessLibrariesIO.monthly_adjacency_matrices(depversions);

    ppath = platpath(platform)
    serialize(ppath * "meta.jls", depversions)
    serialize(ppath * "adj-mats.jls", adj_mats)

    depversions, adj_mats
end

function get_meta(platform)
    ppath = platpath(platform)
    if isfile(ppath * "meta.jls")
        meta = deserialize(ppath * "meta.jls")
    else
        meta = make_depversions_and_adj_mats(platform)[1]
    end
    meta
end

function get_adj_mats(platform)
    ppath = platpath(platform)
    if isfile(ppath * "adj-mats.jls")
        adj_mats = deserialize(ppath * "adj-mats.jls")
    else
        adj_mats = make_depversions_and_adj_mats(platform)[2]
    end
    adj_mats
end

"Check ../data/processed and make only if missing"
function get_depversions_and_adj_mats(platform)
    ppath = platpath(platform)
    if isfile(ppath * "meta.jls")
        adj_mats = deserialize(ppath * "adj-mats.jls")
        meta = deserialize(ppath * "meta.jls")
    else
        meta, adj_mats = make_depversions_and_adj_mats(platform)
    end
    meta, adj_mats
end

PLATFORMS = [
   "Elm",
   "Cargo",
   "Pypi",
   "CRAN",
   "Maven",
   "NPM",
  ]

function _make_all()
    # Printing causes segfaults with @threads.
    #= oldlogger = Logging.global_logger() =#
    #= Logging.global_logger(Logging.NullLogger()) =#
    for p in PLATFORMS
        get_depversions_and_adj_mats(p)
        @info "$p done!"
    end
    #= Logging.global_logger(oldlogger) =#
end

end


#=
    Used from https://github.com/lanl-ansi/Juniper.jl
=#

#=
    Used from https://github.com/lanl-ansi/Alpine.jl
=#
function expr_dereferencing!(expr, m)
    for i in 2:length(expr.args)
        if isa(expr.args[i], Union{Float64,Int64,JuMP.VariableRef})
            k = 0
        elseif expr.args[i].head == :ref
            @assert isa(expr.args[i].args[2], MOI.VariableIndex)
            expr.args[i] = JuMP.VariableRef(m, expr.args[i].args[2])
        elseif expr.args[i].head == :call
            expr_dereferencing!(expr.args[i], m)
        else
            error("expr_dereferencing :: Unexpected term in expression tree.")
        end
    end
end

function expr_dereferencing(expr, m)
    c_expr = copy(expr)
    expr_dereferencing!(c_expr, m)
    return c_expr
end

function get_type_dict(obj)
    T = typeof(obj)
    type_dict = Dict{Symbol,Type}()
    for (name, typ) in zip(fieldnames(T), T.types)
        type_dict[name] = typ
    end
    return type_dict
end

"""
    is_global_status(state::MOI.TerminationStatusCode)
Returns true if either ALMOST_OPTIMAL, OPTIMAL or INFEASIBLE and false otherwise
"""
function is_global_status(state::MOI.TerminationStatusCode)
    return state == MOI.ALMOST_OPTIMAL || state == MOI.OPTIMAL || state == MOI.INFEASIBLE
end

"""
    only_almost_solved(state::MOI.TerminationStatusCode)
Returns true if either ALMOST_OPTIMAL or ALMOST_LOCALLY_SOLVED
"""
function only_almost_solved(state::MOI.TerminationStatusCode)
    return state == MOI.ALMOST_OPTIMAL || state == MOI.ALMOST_LOCALLY_SOLVED
end

"""
    state_is_optimal(state::MOI.TerminationStatusCode; allow_almost=false)
Returns true if either optimal or locally solved. If allow_almost then check for `ALMOST_LOCALLY_SOLVED` and `ALMOST_OPTIMAL`
"""
function state_is_optimal(state::MOI.TerminationStatusCode; allow_almost=false)
    return state == MOI.OPTIMAL || state == MOI.LOCALLY_SOLVED ||
            (allow_almost && state == MOI.ALMOST_LOCALLY_SOLVED) || (allow_almost && state == MOI.ALMOST_OPTIMAL)
end

"""
    state_is_infeasible(state::MOI.TerminationStatusCode)
Returns true if either infeasible or locally infeasible
"""
function state_is_infeasible(state::MOI.TerminationStatusCode)
    return state == MOI.INFEASIBLE || state == MOI.LOCALLY_INFEASIBLE
end

"""
    add_obj_constraint(jp::HybridProblem, rhs::Float64)
Add a constraint for the objective based on whether the objective is linear/quadratic or non linear.
If the objective sense is :MIN than add objective <= rhs else objective >= rhs
"""
function add_obj_constraint(jp::HybridProblem, rhs::Float64)
    if jp.has_nl_objective
        obj_expr = MOI.objective_expr(jp.nlp_evaluator)
        if jp.obj_sense == :Min
            obj_constr = Expr(:call, :<=, obj_expr, rhs)
        else
            obj_constr = Expr(:call, :>=, obj_expr, rhs)
        end
        HybridMINLPSolver.expr_dereferencing!(obj_constr, jp.model)
        JuMP.add_NL_constraint(jp.model, obj_constr)
    else # linear or quadratic
        backend = JuMP.backend(jp.model);
        if isa(jp.objective, MOI.SingleVariable)
            if jp.obj_sense == :Min
                JuMP.set_upper_bound(jp.x[jp.objective.variable.value], rhs)
            else
                JuMP.set_lower_bound(jp.x[jp.objective.variable.value], rhs)
            end
        else
            if jp.obj_sense == :Min
                MOI.add_constraint(backend, jp.objective, MOI.LessThan(rhs))
            else
                MOI.add_constraint(backend, jp.objective, MOI.GreaterThan(rhs))
            end
        end
    end
end

"""
    evaluate_objective(optimizer::MOI.AbstractOptimizer, jp::HybridProblem, xs::Vector{Float64})
If no objective exists => return 0
Evaluate the objective whether it is non linear, linear or quadratic
"""
function evaluate_objective(optimizer::MOI.AbstractOptimizer, jp::HybridProblem, xs::Vector{Float64})
    if optimizer.nlp_data.has_objective
        return MOI.eval_objective(optimizer.nlp_data.evaluator, xs)
    elseif optimizer.objective !== nothing
        return MOIU.eval_variables(vi -> xs[vi.value], optimizer.objective)
    else
        return 0.0
    end
end

function set_time_limit!(optimizer, time_limit::Union{Nothing,Float64})
    old_time_limit = Inf
    if MOI.supports(optimizer, MOI.TimeLimitSec())
        old_time_limit = MOI.get(optimizer, MOI.TimeLimitSec())
        MOI.set(optimizer, MOI.TimeLimitSec(), time_limit)
    end
    return old_time_limit
end

"""
    optimize_get_status_backend(model::JuMP.Model; solver=nothing)
Run optimize! and get the status and the backend
"""
function optimize_get_status_backend(model::JuMP.Model; solver=nothing)
    if solver === nothing
        JuMP.optimize!(model)
    else
        JuMP.set_optimizer(model, solver)
        JuMP.optimize!(model)
    end
    backend = JuMP.backend(model)
    status = MOI.get(backend, MOI.TerminationStatus())
    return status, backend
end

function register_functions!(model, registered_functions)
    if registered_functions !== nothing
        for reg_f in registered_functions
            if reg_f.gradf === nothing
                JuMP.register(model, reg_f.s, reg_f.dimension, reg_f.f; autodiff=reg_f.autodiff)
            elseif reg_f.grad2f === nothing
                JuMP.register(model, reg_f.s, reg_f.dimension, reg_f.f, reg_f.gradf)
            else
                JuMP.register(model, reg_f.s, reg_f.dimension, reg_f.f, reg_f.gradf, reg_f.grad2f)
            end
        end
    end
end


SingleView{T} = SubArray{T,0,Array{T,1},Tuple{Int64},true} where T

abstract type AbstractTuneable <: Number end

struct Tuneable{T<:Number} <: AbstractTuneable
    val::T
    Tuneable(val::T) where T = new{T}(val)
end

struct InitializedTuneable{T<:Number} <: AbstractTuneable
    view::SingleView{T}
end

initialize!(vect, untuneable) = untuneable
function initialize!(vect::Vector{T}, tuneable::Tuneable) where T
    push!(vect, tuneable.val)
    return InitializedTuneable{T}(@view vect[end])
end
function initialize!(vect::Vector, array::AbstractArray{T}) where T
    new_array = T[]
    for elem in array
        push!(new_array, initialize!(vect, elem))
    end
    return new_array
end
function initialize!(vect::Vector, tup::NamedTuple)
    array = Pair[]
    for (key, val) in zip(keys(tup), values(tup))
        push!(array, key => initialize!(vect, val))
    end
    return (;array...)
end

struct TuneableModel{M, A}
    model::M
    array::Vector{<:A}
end
function TuneableModel{A}(model) where {A}
    arr = A[]
    for pname in propertynames(model)
        prop = getproperty(model, pname)
        initialize!(arr, prop)
    end
end


abstract type Component end

mutable struct Bus
    contents::UInt8
end
value(c::Bus) = c.contents
set_value(c::Bus, v::UInt8) = c.contents = v

mutable struct Register <: Component
    value::UInt8
    bus::Bus
end
value(c::Register) = c.value
set_value(c::Register, v::UInt8) = c.value = v

mutable struct ProgramCounter <: Component
    value::UInt8
    bus::Bus
end
value(c::ProgramCounter) = c.value
set_value(c::ProgramCounter, v::UInt8) = c.value = v
enable(c::ProgramCounter) = c.value += 1

mutable struct ALU <: Component
    in1::Register
    in2::Register
    bus::Bus
end
value(c::ALU) = value(c.in1) + value(c.in2)

mutable struct RAM <: Component
    contents::Vector{UInt8}
    mar::Register
    bus::Bus

    RAM(mar::Register, bus::Bus) = new(zeros(UInt8, 16), mar, bus)
end
value(c::RAM) = c.contents[value(c.mar) << 4 >> 4 + 1]
set_value(c::RAM, v::UInt8) = c.contents[value(c.mar) << 4 >> 4 + 1] = v

output(c::Component) = c.bus.contents = value(c)
input(c::Component) = set_value(c, c.bus.contents)


# A helper function to let our `method` keyword correctly understands `Degree` inputs.
@inline wrap_BSpline(itp::Interpolations.InterpolationType) = itp
@inline wrap_BSpline(degree::Interpolations.Degree) = BSpline(degree)

# The default values used by extrapolation for off-domain points
const FillType = Union{Number,Colorant,Flat,Periodic,Reflect}
const FloatLike{T<:AbstractFloat} = Union{T,AbstractGray{T}}
const FloatColorant{T<:AbstractFloat} = Colorant{T}
@inline _default_fillvalue(::Type{T}) where {T<:FloatLike} = convert(T, NaN)
@inline _default_fillvalue(::Type{T}) where {T<:FloatColorant} = _nan(T)
@inline _default_fillvalue(::Type{T}) where {T} = zero(T)

@inline _make_compatible(T, fill) = fill
@inline _make_compatible(::Type{T}, fill::Number) where {T} = T(fill)

Interpolations.tweight(A::AbstractArray{C}) where C<:Colorant{T} where T = T

const MethodType = Union{Interpolations.Degree, Interpolations.InterpolationType}

function box_extrapolation(
        parent::AbstractArray;
        fillvalue::FillType = _default_fillvalue(eltype(parent)),
        method=Linear(),
        kwargs...)
    T = typeof(zero(Interpolations.tweight(parent)) * zero(eltype(parent)))
    itp = maybe_skip_prefilter(T, parent, method)
    extrapolate(itp, _make_compatible(T, fillvalue))
end
box_extrapolation(etp::AbstractExtrapolation) = etp
box_extrapolation(itp::AbstractInterpolation{T}; fillvalue=_default_fillvalue(T)) where T =
    extrapolate(itp, _make_compatible(T, fillvalue))

@inline function maybe_skip_prefilter(::Type{T}, A::AbstractArray, degree::D) where {T, D<:Union{Linear, Constant}}
    axs = axes(A)
    return Interpolations.BSplineInterpolation{T,ndims(A),typeof(A),BSpline{D},typeof(axs)}(A, axs, BSpline(degree))
end
@inline function maybe_skip_prefilter(::Type{T}, A::AbstractArray, method::BSpline{D}) where {T, D<:Union{Linear, Constant}}
    maybe_skip_prefilter(T, A, construct_interpolation_type(D))
end

@inline function maybe_skip_prefilter(::Type{T}, A::AbstractArray, method::MethodType) where T
    return interpolate(A, wrap_BSpline(method))
end

# This is type-piracy, but necessary if we want Interpolations to be
# independent of OffsetArrays.
function AxisAlgorithms.A_ldiv_B_md!(dest::OffsetArray, F, src::OffsetArray, dim::Integer, b::AbstractVector)
    indsdim = axes(parent(src), dim)
    indsF = axes(F)[2]
    if indsF == indsdim
        AxisAlgorithms.A_ldiv_B_md!(parent(dest), F, parent(src), dim, b)
        return dest
    end
    throw(DimensionMismatch("indices $(axes(parent(src))) do not match $(axes(F))"))
end


using Test
using GradientBoost
using GradientBoost.GB
using GradientBoost.LossFunctions

struct DummyGradientBoost <: GBAlgorithm; end

struct StubGradientBoost <: GBAlgorithm
  loss_function::LossFunction
  sampling_rate::Float64
  learning_rate::Float64
  num_iterations::Int
end

function GradientBoost.GB.build_base_func(
  gb::StubGradientBoost,
  instances,
  labels,
  prev_func_pred,
  psuedo)

  function pred(instances)
    num_instances = size(instances, 1)
    predictions = Array{Float64}(undef, num_instances)
    for i = 1:num_instances
      predictions[i] = sum(instances[i,:])
    end
    predictions
  end

  model_const = 0.5
  return (instances) -> model_const .* pred(instances)
end

sgb_instances = [
  2 2;
  2 4
]
sgb_labels = [
  1.0;
  3.0
]

@testset "Gradient Boost" begin
  @testset "not implemented functions throw an error" begin
    dgb = DummyGradientBoost()
    emp_mat = fill(1,1,1)
    emp_vec = Array[]
    @test_throws ErrorException build_base_func(
      dgb,
      emp_mat,
      emp_vec,
      emp_vec,
      emp_vec
    )
  end

  @testset "stochastic_gradient_boost works" begin
    # Sanity check
    sgb = StubGradientBoost(LeastSquares(), 1.0, 0.5, 1)
    model = stochastic_gradient_boost(sgb, sgb_instances, sgb_labels)
    @test 1 == 1
  end

  @testset "fit returns model" begin
    sgb = StubGradientBoost(LeastSquares(), 1.0, 0.5, 1)
    model = stochastic_gradient_boost(sgb, sgb_instances, sgb_labels)
    @test model isa GBModel
  end

  @testset "predict works" begin
    expected = [
      3.0
      3.5
    ]
    sgb = StubGradientBoost(LeastSquares(), 1.0, 0.5, 1)
    model = stochastic_gradient_boost(sgb, sgb_instances, sgb_labels)
    predictions = predict(model, sgb_instances)
    @test predictions == expected
  end

  @testset "create_sample_indices works" begin
    instances = [1:5 6:10]
    labels = [1:5]

    sgb = StubGradientBoost(LeastSquares(), 1, 1, 1)
    indices = create_sample_indices(sgb, instances, labels)
    @test length(indices) == 5
    @test length(unique(indices)) == 5
    @test minimum(indices) >= 1
    @test maximum(indices) <= 5

    sgb = StubGradientBoost(LeastSquares(), 0.5, 1, 1)
    indices = create_sample_indices(sgb, instances, labels)
    @test length(indices) == 3
    @test length(unique(indices)) == 3
    @test minimum(indices) >= 1
    @test maximum(indices) <= 5
  end
end


module SlaterKoster

using DataFrames
using LinearAlgebra

##################################################
# Constants
##################################################

# Atomic mass unit in atomic units (`amu / m_e`)
const amu = 1_822.888_486_192
# Hartree in units of electronvolts (`eV / hartree`)
const electronvolt = 27.211_386_02
# Bohr radius in units of Ångström (`Å / a0`)
const a0 = 0.529_177_210_67

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

abstract type PrimitiveTable{T} end

struct HomoNuclearTable{T} <: PrimitiveTable{T}
    element_symbol::Symbol
    grid_dist::T
    n_grid_points::Int
    energy_table::DataFrame
    mass_table::DataFrame
    integral_table::DataFrame
end

struct HeteroNuclearTable{T} <: PrimitiveTable{T}
    first_element_symbol::Symbol
    second_element_symbol::Symbol
    grid_dist::T
    n_grid_points::Int
    mass_table::DataFrame
    integral_table::DataFrame
end

function Base.:(==)(t1::T, t2::T) where {T<:PrimitiveTable}
    fields = fieldnames(T)
    return all(field -> getfield(t1, field) == getfield(t2, field), fields)
end

struct SlaterKosterTable{T}
    data::Dict{Tuple{Symbol,Symbol},PrimitiveTable{T}}
end

Base.getindex(skt::SlaterKosterTable, x) = getindex(skt.data, x)

function Base.show(io::IO, skt::SlaterKosterTable)
    print(io, "SlaterKosterTable with entries for elements: ")
    print(io, join(elementsymbols(skt), ", ", " and "))
    print(io, ":\n")
    print(io, "    ")
    print(io, join(("$(first(key)) => $(last(key))" for key in keys(skt.data) |> collect |> sort!), "\n    "))
end

elementsymbols(skt::HomoNuclearTable) = [skt.element_symbol]
elementsymbols(skt::HeteroNuclearTable) = [skt.first_element_symbol, skt.second_element_symbol]
elementsymbols(skt::SlaterKosterTable) = reduce(vcat, elementsymbols.(values(skt.data))) |> unique! |> sort!

include("load.jl")
include("util.jl")

include("hamiltonian.jl")
include("overlap.jl")
include("mass.jl")

end # module


# ------------ Trace typing system ------------ #

# Utilities.
keys(::Type{NamedTuple{K, V}}) where {K, V} = K
keys(::NamedTuple{K, V}) where {K, V} = K
value_types(::Type{NamedTuple{K, V}}) where {K, V} = V
NamedTuple{K, V}() where {K, V} = NamedTuple{K}(map(V.parameters) do p
                                                    p()
                                                end)

function Base.:(<<)(nt1::Type{NamedTuple{K1, V1}}, nt2::Type{NamedTuple{K2, V2}}) where {K1, K2, V1, V2}
    for k in K1
        k in K2 || return false
    end
    for (v, k) in zip(V1.parameters, V2.parameters)
        v << k || return false
    end
    true
end

# ------------ Support types ------------ #

abstract type SupportType end

abstract type BaseLebesgue <: SupportType end
pretty(::Type{BaseLebesgue}) = :lebesgue

abstract type BaseCounting <: SupportType end
pretty(::Type{BaseCounting}) = :counting

# Base types.
struct Reals{N} <: BaseLebesgue end
Base.:(<<)(::Type{Reals{N}}, ::Type{Reals{N}}) where N = true
Base.:(<<)(::Type{Reals}, a) = false

struct RealInterval <: BaseLebesgue
    a
    b
end

struct DiscreteInterval{A, B} <: BaseCounting end

struct PositiveReals{N} <: BaseLebesgue end
Base.:(<<)(::Type{PositiveReals{N}}, ::Type{PositiveReals{N}}) where N = true
Base.:(<<)(::Type{PositiveReals}, a) = false

struct Integers{N} <: BaseCounting end
Base.:(<<)(::Type{Integers{N}}, ::Type{Reals{N}}) where N = true
Base.:(<<)(::Type{Integers{N}}, ::Type{Integers{N}}) where N = true
Base.:(<<)(::Type{Integers}, a) = false

struct Discrete{N} <: BaseCounting end

# Derived types.
struct List{N} <: SupportType
    tt::N
end
List{N}() where N = List{N}(N())

struct Sum{T <: Tuple} <: SupportType
    tt::T
end
Sum{T}() where T <: Tuple = Sum{T}(map(T) do a
                                       a()
                                   end)

# Inferred record type is just a NamedTuple.
const TraceType = NamedTuple

# ------------ Tracer ------------ #

# Type inference.
function infer_support_types(ctx, fn, arg_types...)
    ir = lower_to_ir(fn, arg_types...)
    dh = detect_dynamic_addresses(ir)
    !isempty(dh) && begin
        display(dh)
        println()
        return missing
    end
    try
        tr = ir |> prepare_ir! |> ir -> infer!(ctx, ir)
        tr
    catch e
        @info "Failed to trace $(fn).\nCause: $e\n\nReverting to untyped IR."
        return missing
    end
end

# Extract trace type from inferred IR.
function trace_type(tr)
    ks = Set{Symbol}([])
    types = []
    for (v, st) in tr
        st.expr isa Expr || continue
        st.expr.head == :call || continue
        st.expr.args[1] == trace || continue
        st.expr.args[2] isa QuoteNode || continue
        push!(ks, st.expr.args[2].value)
        if !(st.type isa Type)
            push!(types, st.type)
        elseif st.type isa Type && (st.type <: SupportType || st.type <: Missing)
            push!(types, st.type())
        elseif _lift(st.type) <: NamedTuple
            new = NamedTuple{keys(st.type)}(map(value_types(st.type).parameters) do p
                                                p()
                                            end)
            push!(types, new)
        else
            push!(types, missing)
        end
    end
    NamedTuple{tuple(ks...)}(types)
end


module Setters

using Test

using Setfield: set
using Unitful
using UnitfulAtomic

using QuantumESPRESSOBase.Inputs.PWscf
using QuantumESPRESSOBase.Setters

@testset "Test `VerbositySetter` on `ControlNamelist`" begin
    control = ControlNamelist()
    @test (  # Default values
        control.verbosity,
        control.wf_collect,
        control.tstress,
        control.tprnfor,
        control.disk_io,
    ) == ("low", true, false, false, "low")
    control = set(control, VerbositySetter(:high))
    @test (
        control.verbosity,
        control.wf_collect,
        control.tstress,
        control.tprnfor,
        control.disk_io,
    ) == ("high", true, true, true, "high")
    control = set(control, VerbositySetter(:low))
    @test (
        control.verbosity,
        control.wf_collect,
        control.tstress,
        control.tprnfor,
        control.disk_io,
    ) == ("low", false, false, false, "low")
end # testset

@testset "Test `FiniteTemperatureSetter` on `SystemNamelist`" begin
    system = SystemNamelist()
    @test (system.occupations, system.degauss, system.smearing) ==
          ("fixed", 0.0, "gaussian")
    system = set(system, FiniteTemperatureSetter(0.0019000869380733452))
    @test (system.occupations, system.degauss, system.smearing) ==
          ("smearing", 0.0019000869380733452, "fermi-dirac")
    system = set(system, FiniteTemperatureSetter(3e8 * u"μK"))  # 300K
    @test (system.occupations, system.degauss, system.smearing) ==
          ("smearing", 0.0019000869380733452, "fermi-dirac")
    system = set(system, FiniteTemperatureSetter(2.5851999786e-8 * u"MeV"))  # 300K
    @test (system.occupations, system.degauss, system.smearing) ==
          ("smearing", 0.001900086938034304, "fermi-dirac")
    system = set(system, FiniteTemperatureSetter(0.00095004346903668 * u"hartree"))  # 300K
    @test (system.occupations, system.degauss, system.smearing) ==
          ("smearing", 0.00190008693807336, "fermi-dirac")
    system = set(system, FiniteTemperatureSetter(6.250985736 * u"THz"))  # 300K
    @test (system.occupations, system.degauss, system.smearing) ==
          ("smearing", 0.0019000869377698247, "fermi-dirac")
    system = set(system, FiniteTemperatureSetter(20851.044012 * u"m^-1"))  # 300K
    @test (system.occupations, system.degauss, system.smearing) ==
          ("smearing", 0.001900086937837879, "fermi-dirac")
    system = set(system, FiniteTemperatureSetter(208.51044012 * u"cm^-1"))  # 300K
    @test (system.occupations, system.degauss, system.smearing) ==
          ("smearing", 0.001900086937837879, "fermi-dirac")
    system = set(system, FiniteTemperatureSetter(4.141947e-24 * u"kJ"))  # 300K
    @test (system.occupations, system.degauss, system.smearing) ==
          ("smearing", 0.0019000869380663146, "fermi-dirac")
    system = set(system, FiniteTemperatureSetter(4.6085375610000005e-35 * u"g"))  # 300K
    @test (system.occupations, system.degauss, system.smearing) ==
          ("smearing", 0.0019000869377760491, "fermi-dirac")
    system = set(system, FiniteTemperatureSetter(4.6085375610000005e-32 * u"mg"))  # 300K
    @test (system.occupations, system.degauss, system.smearing) ==
          ("smearing", 0.001900086937776049, "fermi-dirac")
end # testset

end


solvers_const = Dict(
    :ARCLDLt_abs => (HessDense, PDataLDLt, solve_modelARCDiagAbs, ()),
    :ARCLDLt_HO_Sham => (
        HessDense,
        PDataLDLt,
        (v, w, x, y, z) -> solve_modelARCDiag_HO(
            v,
            w,
            x,
            y,
            z,
            ho_correction = :Shamanskii,
            λfact = 100.0,
        ),
        (),
    ),
    :ARCLDLt_HO_vs_Nwt => (
        HessDense,
        PDataLDLt,
        (v, w, x, y, z) -> solve_modelARCDiag_HO_vs_Nwt(
            v,
            w,
            x,
            y,
            z,
            λfact = 100.0,
            nwt_res_fact = 0.8,
        ),
        (),
    ),
    :ARCLDLt => (HessDense, PDataLDLt, solve_modelARCDiag, ()),
    # :ARCMA57_abs => (HessSparse, PDataMA57, solve_modelARCDiagAbs, ()),
    # :ARCMA57_Sham_vs_Nwt_λ => (HessSparse, PDataMA57, (v, w, x, y, z) -> solve_modelARCDiag_HO_vs_Nwt(v, w, x, y, z, λfact = 1.0, nwt_res_fact = 0.25, ho_correction = :Shamanskii_MA57), ()),
    # :ARCMA57_Sham_vs_Nwt => (HessSparse, PDataMA57, solve_modelARCDiag, ()),
    # :ARCMA57_Sham_λ => (HessSparse, PDataMA57, (v, w, x, y, z) -> solve_modelARCDiag_HO(v, w, x, y, z, ho_correction = :Shamanskii_MA57, λfact = 1.0), ()),
    # :ARCMA57 => (HessSparse, PDataMA57, solve_modelARCDiag, ()),
    # :ARCMA97_abs => (HessDense, PDataMA97, solve_modelARCDiagAbs, ()),
    # :ARCMA97 => (HessDense, PDataMA97, solve_modelARCDiag, ()),
    :ARCqKdense => (
        HessDense,
        PDataKARC,
        solve_modelKARC,
        [(:shifts => 10.0 .^ (collect(-20.0:1.0:20.0)))],
    ),
    :ARCqKOp => (
        HessOp,
        PDataKARC,
        solve_modelKARC,
        [:shifts => 10.0 .^ (collect(-20.0:1.0:20.0))],
    ),
    :ARCqKsparse => (
        HessSparse,
        PDataKARC,
        solve_modelKARC,
        [:shifts => 10.0 .^ (collect(-20.0:1.0:20.0))],
    ),
    :ARCSpectral_abs => (HessDense, PDataSpectral, solve_modelARCDiagAbs, ()),
    :ARCSpectral => (HessDense, PDataSpectral, solve_modelARCDiag, ()),
    :ST_TRdense => (HessDense, PDataST, solve_modelST_TR, ()),
    :ST_TROp => (HessOp, PDataST, solve_modelST_TR, ()),
    :ST_TRsparse => (HessSparse, PDataST, solve_modelST_TR, ()),
    :TRKdense => (HessDense, PDataKTR, solve_modelKTR, ()),
    :TRKOp => (HessOp, PDataKTR, solve_modelKTR, ()),
    :TRKsparse => (HessSparse, PDataKTR, solve_modelKTR, ()),
    :TRLDLt_abs => (HessDense, PDataLDLt, solve_modelTRDiagAbs, ()),
    :TRLDLt_HO_vs_Nwt_λ => (
        HessDense,
        PDataLDLt,
        (v, w, x, y, z) -> solve_modelTRDiag_HO_vs_Nwt_λ(
            v,
            w,
            x,
            y,
            z;
            ho_correction = :Shamanskii,
            nwt_res_fact = 0.25,
            λfact = 1.0,
        ),
        (),
    ),
    :TRLDLt => (HessDense, PDataLDLt, solve_modelTRDiag, ()),
    # :TRMA57_abs => (HessSparse, PDataMA57, solve_modelTRDiagAbs, ()),
    # :TRMA57_Sham_vs_Nwt_λ => (HessSparse, PDataMA57, (v, w, x, y, z) -> solve_modelTRDiag_HO_vs_Nwt_λ(v, w, x, y, z; ho_correction = :Shamanskii_MA57, nwt_res_fact = 0.25, λfact = 1.0), ()),
    # :TRMA57 => (HessSparse, PDataMA57, solve_modelTRDiag, ()),
    # :TRMA97_abs => (HessDense, PDataMA97, solve_modelTRDiagAbs, ()),
    # :TRMA97 => (HessSparse, PDataMA97, solve_modelTRDiag, ()),
    :TRSpectral_abs => (HessDense, PDataSpectral, solve_modelTRDiagAbs, ()),
    :TRSpectral => (HessDense, PDataSpectral, solve_modelTRDiag, ()),
)