### Load the required packages ###

using DimensionalData, GeoData, ArchGDAL, NCDatasets, HDF5
using CSV, DataFrames, DataFramesMeta, TableView
using DynamicGrids, Dispersal
using Statistics, Dates, Unitful, Setfield, Flatten, FieldMetadata
using Plots, ColorSchemes, Colors
using DimensionalData: setdims, rebuild, Between

const DG = DynamicGrids

### Override default parameter range bounds ###

# For both interfaces and optimiser

import FieldMetadata: @bounds, bounds, @flattenable, flattenable
using Flatten: flatten

@bounds HumanDispersal begin
    human_exponent  | (0.0, 3.0)
    dist_exponent   | (0.0, 3.0)
    dispersalperpop | (0.0, 1e-8)
    max_dispersers  | (1e1, 1e5)
end
# Constant growth
@bounds ExactLogisticGrowth begin
    intrinsicrate | (0.0, 10.0)
end
# Alee
@bounds AlleeExtinction begin
    minfounders | (2e0, 1e3)
end

# Dont parametrise carrycap
@flattenable ExactLogisticGrowthMap begin
    carrycap | false
end
@flattenable ExactLogisticGrowth begin
    carrycap | false
end
# Don't parametrise λ for main models, it will just be the top bound
# so we need to estimate it. This is reverted below for local-dispersal
# only models.
@flattenable ExponentialKernel begin
    λ | false
end

### Define simulation settings ###

timestep = Month(1)

us_tspan = DateTime(2008, 5):timestep:DateTime(2013, 12)
eu_tspan = DateTime(2008, 5):timestep:DateTime(2016, 12)

aus = Lon(Between(113.3402, 153.9523)), Lat(Between(-43.62234, -10.65125))
us = Lon(Between(-130.332, -66.9)), Lat(Between(18.977, 71.59))
eu = Lon(Between(-13.7, 40.15)), Lat(Between(34.9, 70.30))

aus_incursionpoints = (
    Melbourne=(-37.805896, 144.959527),
    Mildura=(-34.219504, 142.130864),
    Coffs_Harbour=(-30.287245, 153.092991),
    Sydney=(-33.839943, 151.006101),
    Adelaide=(-34.901608, 138.601547),
    Port_Augusta=(-32.466201, 137.813850),
    Devonport=(-41.180545, 146.314887),
    Hobart=(-42.881742, 147.323879),
    Brisbane=(-27.436190, 152.990588),
    Cairns=(-16.937281, 145.747709),
    Perth=(-31.9505, 115.8605),
    Geraldton=(-28.778138, 114.615632)
)

us_incursionpoints = (
    San_Francisco=(36.9453, -122.0695),
)

obs = CSV.File(joinpath(basedir, "data/Oersted_occurrence.csv")) |> DataFrame
eu_obs = @where(obs, :Year .!== missing, :Year .== 2008, :Country .!= "USA")
eu_incursionpoints = collect(zip(eu_obs.Latitude, eu_obs.Longitude))

incursionpointstatekeys = (
    Melbourne = (:Vic),
    Mildura	 = (:Vic),
    Coffs_Harbour =	(:NSW),
    Sydney = (:NSW),
    Adelaide = (:SA),
    Port_Augusta = (:SA),
    Devonport  = (:Tas),
    Hobart = (:Tas),
    Brisbane = (:QLD),
    Cairns = (:QLD),
    Perth =	(:WA),
    Geraldton =	(:WA)
)

### Load numbered states rasters ###

loadstates(path, selectors) = GDALarray(joinpath(basedir, path);
    name=:States, usercrs=EPSG(4326))[Band(1), selectors...] |>
    x->permutedims(x, (Lat, Lon)) |> x -> replace_missing(x, 0) |>
    x->trunc.(Int, x)
aus_states = loadstates(joinpath(basedir, "data/aus_states.tif"), aus)
us_states = loadstates(joinpath(basedir, "data/US_states.tif"), us)
eu_states = loadstates(joinpath(basedir, "data/EU_states.tif"), eu)
us_states = loadstates(joinpath(basedir, "data/US_states.tif"), ())
# aus_states |> plot
# eu_states |> plot
# us_states |> plot

### Define the RuleSets ###

growthratesfilepath = joinpath(basedir, "output/sensitivity/growthrates_standard.ncd")
isfile(growthratesfilepath)
growthrates = NCDarray(growthratesfilepath; crs=crs(aus_states), dimcrs=EPSG(4326)) |> GeoArray |>
    x -> setdims(x, (@set dims(x, Ti).mode.span = Regular(Month(1)))) |>
    x -> setdims(x, convertmode(Projected, dims(x, Lon))) |>
    x -> setdims(x, convertmode(Projected, dims(x, Lat))) |>
    x -> permutedims(x, (Lat, Lon, Ti)) |>
    x -> reverse(x; dims=Lat)

# Make growth rates aux grids, replacing the `missing` values with zero.
# Dispersal.jl `init` can't contain `missing` - it will spread everywhere.

grzm = floattype.(replace_missing(growthrates, 0.0))
aus_growthrates_zeromissing = grzm[aus...]
us_growthrates_zeromissing = grzm[us...]
eu_growthrates_zeromissing = grzm[eu...]
aus_growthrates_zeromissing[Ti(1)] |> plot
us_growthrates_zeromissing[Ti(1)] |> plot
eu_growthrates_zeromissing[Ti(1)] |> plot
aus_aux = (growthrates=aus_growthrates_zeromissing,)
us_aux = (growthrates=us_growthrates_zeromissing,)
eu_aux = (growthrates=eu_growthrates_zeromissing,)

### Define masking layers ###

# These define regions that are not run, specifically, the sea.

aus_boolmask = GeoData.boolmask(aus_states)
aus_missingmask = GeoData.missingmask(aus_states)
aus_missingmask = GeoData.missingmask(aus_states)
us_boolmask = GeoData.boolmask(growthrates[us..., Ti(1)])
us_missingmask = GeoData.missingmask(growthrates[us..., Ti(1)])
eu_boolmask = GeoData.boolmask(eu_states)
eu_missingmask = GeoData.missingmask(eu_states)
# plot(aus_boolmask)
# plot(aus_missingmask)
# plot(us_boolmask)
# plot(us_missingmask)
# plot(eu_boolmask)
# plot(eu_missingmask)

# Growth masks. These are the cells where the mean growth rate is above zero.

aus_growthmask = rebuild(aus_boolmask, (mean(aus_growthrates_zeromissing; dims=Ti)[Ti(1)] .* aus_boolmask) .> 0)
us_growthmask = rebuild(us_boolmask, (mean(us_growthrates_zeromissing; dims=Ti)[Ti(1)] .* us_boolmask) .> 0)
eu_growthmask = rebuild(eu_boolmask, (mean(eu_growthrates_zeromissing; dims=Ti)[Ti(1)] .* eu_boolmask) .> 0)
# aus_growthmask |> plot
# us_growthmask |> plot
# eu_growthmask |> plot

### Load human population

human_pop_path = joinpath(basedir, "data/population_density.tif")
human_pop = GDALarray(human_pop_path; usercrs=EPSG(4326))[Band(1)] |>
            permutedims |> replace_missing
# human_pop |> plot
aus_humanpop = human_pop[aus...]
us_humanpop = human_pop[us...]
eu_humanpop = human_pop[eu...]
# aus_humanpop |> plot
# us_humanpop |> plot
# eu_humanpop |> plot

gdp_path = joinpath(basedir, "data/GDP_USD_per_m2.tif")
gdp = GDALarray(gdp_path; usercrs=EPSG(4326))[Band(1)] .* 1e6 |> permutedims
gdp = (@set gdp.missingval = minimum(gdp)) |> replace_missing
# gdp |> plot
aus_gdp = gdp[aus...]
us_gdp = gdp[us...]
eu_gdp = gdp[eu...]
# aus_gdp |> plot
# us_gdp |> plot
# eu_gdp |> plot

@assert size(eu_humanpop) == size(eu_growthmask[Ti(1)]) == size(eu_states)
@assert size(us_humanpop) == size(us_growthmask[Ti(1)]) == size(us_states)
@assert size(aus_humanpop) == size(aus_growthmask[Ti(1)]) == size(aus_states)

### Define Rules ###

carrycap = floattype(1e9)
growth = ExactLogisticGrowthMap{:population,:population}(
    layerkey=Val(:growthrates),
    carrycap=carrycap,
    timestep=Day(1),
    nsteps=floattype(1.0),
);
growth

### Constant growth ###

constant_growth = ExactLogisticGrowth{:population,:population}(
    intrinsicrate=floattype(0.1),
    timestep=Day(1),
    carrycap=carrycap
)

### Local dispersal ###

λ = floattype(0.0125)
radius = 1

@time hood = DispersalKernel{radius}(
    formulation=ExponentialKernel(λ),
    distancemethod=AreaToArea(30),
    cellsize=floattype(1.0)
)
localdisp = InwardsPopulationDispersal{:population,:population}(hood)
log.(hood.kernel) |> heatmap
savefig(joinpath(basedir, "output/dispersal_kernel.png"))

### Allee effects

allee = AlleeExtinction{:population,:population}(
    minfounders=floattype(22.0)
);

### Human Dispersal ###

scale = 10
hkwargs = (
    read=:population,
    scale=scale,
    human_exponent=floattype(1.075),
    dist_exponent=floattype(1.429),
    dispersalperpop=floattype(8.703e-9),
    max_dispersers=floattype(3.264e4),
    nshortlisted=100,
)

if humanactivity == :pop
    println("human activity from population")
    aus_humandisp = HumanDispersal(; human_pop=parent(aus_humanpop), hkwargs...)
    us_humandisp = HumanDispersal(; human_pop=parent(us_humanpop), hkwargs...)
    eu_humandisp = HumanDispersal(; human_pop=parent(eu_humanpop), hkwargs...)
elseif humanactivity == :gdp
    println("human activity from gdp")
    aus_humandisp = HumanDispersal(; human_pop=parent(aus_gdp), hkwargs...)
    us_humandisp = HumanDispersal(; human_pop=parent(us_gdp), hkwargs...)
    eu_humandisp = HumanDispersal(; human_pop=parent(eu_gdp), hkwargs...)
end

I = (DimensionalData.dims2indices(aus_states, (Lat(Contains(-36.805896)), Lon(Contains(143.959527)))) .-1) .÷ scale .+ 1
melb_dests = Dispersal.populate!(zero(aus_growthrates_zeromissing[Ti(1)]), aus_humandisp, I...)
plot(melb_dests, title="Melbourne destinations")
savefig(joinpath(basedir, "output/human_dispersal_destinations.png"))

### Define initialisation grids

assign_coord!(popgrid, (lat, lon), initsize) = begin
   coord = Lat(Contains(lat)), Lon(Contains(lon))
   popgrid[coord...] = initsize
end
init_popgrid!(popgrid, (lat, lon), initsize) = begin
   popgrid .= floattype(0.0)
    assign_coord!(popgrid, (lat, lon), initsize)
   nothing
end

aus_populationgrid = zero(aus_growthrates_zeromissing[Ti(1)])
init_popgrid!(aus_populationgrid, aus_incursionpoints[:Brisbane], carrycap)
us_populationgrid = zero(us_growthrates_zeromissing[Ti(1)])
init_popgrid!(us_populationgrid, us_incursionpoints[:San_Francisco], carrycap)
eu_populationgrid = zero(eu_growthrates_zeromissing[Ti(1)])
map(ip -> assign_coord!(eu_populationgrid, ip, carrycap), eu_incursionpoints)
# p = plot(eu_boolmask, label="")
# scatter!(p, reverse.(eu_incursionpoints); title="")
# eu_populationgrid |> plot

### Define the combined rulesets ###

# HumanDispersal has location specific precalculated
# data that the optimiser needs to update, making this a little
# messier than it otherwise would be

chain = Chain(localdisp, allee, growth)

# Test rule type stability
output = ArrayOutput((population=aus_populationgrid,);
    mask=collect(aus_boolmask),
    aux=(growthrates=aus_growthrates_zeromissing,),
    tspan=DateTime(2020, 1):timestep:DateTime(2026, 1)
)
DynamicGrids.isinferred(output, localdisp)
DynamicGrids.isinferred(output, allee)
DynamicGrids.isinferred(output, growth)
DynamicGrids.isinferred(output, chain)
DynamicGrids.isinferred(output, aus_humandisp)

ruleset_kwargs = (
    timestep=timestep,
    opt=SparseOpt(),
)
aus_ruleset = Ruleset(
    aus_humandisp, chain;
    ruleset_kwargs...
);
us_ruleset = Ruleset(
    us_humandisp, chain;
    ruleset_kwargs...
);
eu_ruleset = Ruleset(
    eu_humandisp, chain;
    ruleset_kwargs...
);
DynamicGrids.isinferred(output, us_ruleset)

### Output ###

# Define some color processors to use in live simulations.

zerocolor = ARGB32(0.6)
maskcolor = ARGB32(0.3)
textconfig = TextConfig(font="cantarel", bcolor=maskcolor, fcolor=ARGB32(0.8))

greyscale = ColorProcessor(Greyscale(), zerocolor, maskcolor, textconfig)
oranges = ColorProcessor(ColorSchemes.Oranges_3, zerocolor, maskcolor, textconfig)
jet = ColorProcessor(ColorSchemes.jet, zerocolor, maskcolor, textconfig)
viridis = ColorProcessor(ColorSchemes.viridis, zerocolor, maskcolor, textconfig)
inferno = ColorProcessor(ColorSchemes.inferno, zerocolor, maskcolor, textconfig)
magma = ColorProcessor(ColorSchemes.magma, zerocolor, maskcolor, textconfig)


# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

abstract type AbstractFeasibilityAlgorithm end
abstract type AbstractFeasibilityStrategy end

"""
    NoFeasibilityAlgorithm

Empty functor object for running an L-shaped algorithm without dealing with second-stage feasibility.
"""
struct NoFeasibilityAlgorithm <: AbstractFeasibilityAlgorithm end

handle_feasibility(::NoFeasibilityAlgorithm) = false
num_cuts(::NoFeasibilityAlgorithm) = 0
active!(::MOI.ModelLike, ::NoFeasibilityAlgorithm) = nothing
deactive!(::MOI.ModelLike, ::NoFeasibilityAlgorithm) = nothing
restore!(::MOI.ModelLike, ::NoFeasibilityAlgorithm) = nothing

"""
    FeasibilityCutsMaster

Master functor object for using feasibility cuts in an L-shaped algorithm. Create by supplying a [`FeasibilityCuts`](@ref) object through `feasibility_strategy` in `LShaped.Optimizer` or set the [`FeasibilityStrategy`](@ref) attribute.
"""
struct FeasibilityCutsMaster{T <: AbstractFloat} <: AbstractFeasibilityAlgorithm
    cuts::Vector{SparseFeasibilityCut{T}}

    function FeasibilityCutsMaster(::Type{T}) where T <: AbstractFloat
        return new{T}(Vector{SparseFeasibilityCut{T}}())
    end
end

handle_feasibility(::FeasibilityCutsMaster) = true
worker_type(::FeasibilityCutsMaster) = FeasibilityCutsWorker
num_cuts(feasibility::FeasibilityCutsMaster) = length(feasibility.cuts)
active!(::MOI.ModelLike, ::FeasibilityCutsMaster) = nothing
deactive!(::MOI.ModelLike, ::FeasibilityCutsMaster) = nothing
restore!(::MOI.ModelLike, ::FeasibilityCutsMaster) = nothing

"""
    FeasibilityCutsWorker

Worker functor object for using feasibility cuts in an L-shaped algorithm. Create by supplying a [`FeasibilityCuts`](@ref) object through `feasibility_strategy` in `LShaped.Optimizer` or set the [`FeasibilityStrategy`](@ref) attribute.
"""
mutable struct FeasibilityCutsWorker <: AbstractFeasibilityAlgorithm
    objective::MOI.AbstractScalarFunction
    linking_constraints::Vector{MOI.ConstraintIndex}
    feasibility_variables::Vector{MOI.VariableIndex}
    aux_constraint::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64}}

    function FeasibilityCutsWorker(objective::MOI.AbstractScalarFunction,
                                   linking_constraints::Vector{MOI.ConstraintIndex},
                                   feasibility_variables::Vector{MOI.VariableIndex})
        return new(objective, linking_constraints, feasibility_variables, CI{MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64}}(0))
    end
end

handle_feasibility(::FeasibilityCutsWorker) = true
num_cuts(::FeasibilityCutsWorker) = 0

function prepare!(model::MOI.ModelLike, worker::FeasibilityCutsWorker)
    # Set objective to zero
    G = MOI.ScalarAffineFunction{Float64}
    MOI.set(model, MOI.ObjectiveFunction{G}(), zero(MOI.ScalarAffineFunction{Float64}))
    i = 1
    # Create auxiliary feasibility variables
    for ci in worker.linking_constraints
        i = add_auxiliary_variables!(model, worker, ci, i)
    end
    return nothing
end
function prepared(worker::FeasibilityCutsWorker)
    return length(worker.feasibility_variables) > 0
end

function add_auxiliary_variables!(model::MOI.ModelLike,
                                  worker::FeasibilityCutsWorker,
                                  ci::CI{F,S},
                                  idx::Integer) where {F <: MOI.AbstractFunction, S <: MOI.AbstractSet}
    # Nothing to do for most most constraints
    return idx
end

function add_auxiliary_variables!(model::MOI.ModelLike,
                                  worker::FeasibilityCutsWorker,
                                  ci::CI{F,S},
                                  idx::Integer) where {F <: AffineDecisionFunction, S <: MOI.AbstractScalarSet}
    G = MOI.ScalarAffineFunction{Float64}
    obj_sense = MOI.get(model, MOI.ObjectiveSense())
    # Positive feasibility variable
    pos_aux_var = MOI.add_variable(model)
    name = add_subscript(:v⁺, idx)
    MOI.set(model, MOI.VariableName(), pos_aux_var, name)
    push!(worker.feasibility_variables, pos_aux_var)
    # Nonnegativity constraint
    MOI.add_constraint(model, MOI.SingleVariable(pos_aux_var),
                       MOI.GreaterThan{Float64}(0.0))
    # Add to objective
    MOI.modify(model, MOI.ObjectiveFunction{G}(),
               MOI.ScalarCoefficientChange(pos_aux_var, obj_sense == MOI.MAX_SENSE ? -1.0 : 1.0))
    # Add to constraint
    MOI.modify(model, ci, MOI.ScalarCoefficientChange(pos_aux_var, 1.0))
    # Negative feasibility variable
    neg_aux_var = MOI.add_variable(model)
    name = add_subscript(:v⁻, idx)
    MOI.set(model, MOI.VariableName(), neg_aux_var, name)
    push!(worker.feasibility_variables, neg_aux_var)
    # Nonnegativity constraint
    MOI.add_constraint(model, MOI.SingleVariable(neg_aux_var),
                       MOI.GreaterThan{Float64}(0.0))
    # Add to objective
    MOI.modify(model, MOI.ObjectiveFunction{G}(),
               MOI.ScalarCoefficientChange(neg_aux_var, obj_sense == MOI.MAX_SENSE ? -1.0 : 1.0))
    # Add to constraint
    MOI.modify(model, ci, MOI.ScalarCoefficientChange(neg_aux_var, -1.0))
    # Return updated identification index
    return idx + 1
end

function add_auxiliary_variables!(model::MOI.ModelLike,
                                  worker::FeasibilityCutsWorker,
                                  ci::CI{F,S},
                                  idx::Integer) where {F <: VectorAffineDecisionFunction, S <: MOI.AbstractVectorSet}
    G = MOI.ScalarAffineFunction{Float64}
    obj_sense = MOI.get(model, MOI.ObjectiveSense())
    n = MOI.dimension(MOI.get(model, MOI.ConstraintSet(), ci))
    for (i, id) in enumerate(idx:(idx + n - 1))
        # Positive feasibility variable
        pos_aux_var = MOI.add_variable(model)
        name = add_subscript(:v⁺, id)
        MOI.set(model, MOI.VariableName(), pos_aux_var, name)
        push!(worker.feasibility_variables, pos_aux_var)
        # Nonnegativity constraint
        MOI.add_constraint(model, MOI.SingleVariable(pos_aux_var),
                           MOI.GreaterThan{Float64}(0.0))
        # Add to objective
        MOI.modify(model, MOI.ObjectiveFunction{G}(),
                   MOI.ScalarCoefficientChange(pos_aux_var, obj_sense == MOI.MAX_SENSE ? -1.0 : 1.0))
        # Add to constraint
        MOI.modify(model, ci, MOI.MultirowChange(pos_aux_var, [(i, 1.0)]))
    end
    for (i, id) in enumerate(idx:(idx + n - 1))
        # Negative feasibility variable
        neg_aux_var = MOI.add_variable(model)
        name = add_subscript(:v⁻, id)
        MOI.set(model, MOI.VariableName(), neg_aux_var, name)
        push!(worker.feasibility_variables, neg_aux_var)
        # Nonnegativity constraint
        MOI.add_constraint(model, MOI.SingleVariable(neg_aux_var),
                           MOI.GreaterThan{Float64}(0.0))
        # Add to objective
        MOI.modify(model, MOI.ObjectiveFunction{G}(),
                   MOI.ScalarCoefficientChange(neg_aux_var, obj_sense == MOI.MAX_SENSE ? -1.0 : 1.0))
        # Add to constraint
        MOI.modify(model, ci, MOI.MultirowChange(neg_aux_var, [(i, -1.0)]))
    end
    # Return updated identification index
    return idx + n + 1
end

function activate!(model::MOI.ModelLike, worker::FeasibilityCutsWorker)
    # Set objective to zero
    G = MOI.ScalarAffineFunction{Float64}
    MOI.set(model, MOI.ObjectiveFunction{G}(), zero(MOI.ScalarAffineFunction{Float64}))
    obj_sense = MOI.get(model, MOI.ObjectiveSense())
    # Add auxiliary variables to objective and linking constraints
    idx = 0
    for ci in worker.linking_constraints
        dim = MOI.dimension(MOI.get(model, MOI.ConstraintSet(), ci))
        for i in 1:dim
            pos_aux_var = worker.feasibility_variables[idx + 2*(i-1) + 1]
            MOI.modify(model, MOI.ObjectiveFunction{G}(),
                   MOI.ScalarCoefficientChange(pos_aux_var, obj_sense == MOI.MAX_SENSE ? -1.0 : 1.0))
            neg_aux_var = worker.feasibility_variables[idx + 2*(i-1) + 2]
            MOI.modify(model, MOI.ObjectiveFunction{G}(),
                       MOI.ScalarCoefficientChange(neg_aux_var, obj_sense == MOI.MAX_SENSE ? -1.0 : 1.0))
        end
        idx += 2*dim
    end
    if MOI.is_valid(model, worker.aux_constraint)
        MOI.delete(model, worker.aux_constraint)
    end
    return nothing
end

function deactivate!(model::MOI.ModelLike, worker::FeasibilityCutsWorker)
    # Force auxiliary variables to zero
    func_type = MOI.get(model, MOI.ObjectiveFunctionType())
    obj = MOI.get(model, MOI.ObjectiveFunction{func_type}())
    worker.aux_constraint = MOI.add_constraint(model, obj, MOI.EqualTo{Float64}(0.0))
    # Restore objective
    F = typeof(worker.objective)
    MOI.set(model, MOI.ObjectiveFunction{F}(), worker.objective)
    return nothing
end

function restore!(model::MOI.ModelLike, worker::FeasibilityCutsWorker)
    # Remove aux constraint
    if MOI.is_valid(model, worker.aux_constraint)
        MOI.delete(model, worker.aux_constraint)
    end
    worker.aux_constraint = CI{MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64}}(0)
    # Delete any feasibility variables
    if prepared(worker)
        MOI.delete(model, worker.feasibility_variables)
    end
    empty!(worker.feasibility_variables)
    # Restore objective
    F = typeof(worker.objective)
    MOI.set(model, MOI.ObjectiveFunction{F}(), worker.objective)
    return nothing
end

# API
# ------------------------------------------------------------
"""
    IgnoreFeasibility

Factory object for [`NoFeasibilityAlgorithm`](@ref). Passed by default to `feasibility_strategy` in `LShaped.Optimizer`.

"""
struct IgnoreFeasibility <: AbstractFeasibilityStrategy end

function master(::IgnoreFeasibility, ::Type{T}) where T <: AbstractFloat
    return NoFeasibilityAlgorithm()
end

function worker(::IgnoreFeasibility, ::Vector{MOI.ConstraintIndex}, ::MOI.ModelLike)
    return NoFeasibilityAlgorithm()
end
function worker_type(::IgnoreFeasibility)
    return NoFeasibilityAlgorithm
end

"""
    IgnoreFeasibility

Factory object for using feasibility cuts in an L-shaped algorithm.

"""
struct FeasibilityCuts <: AbstractFeasibilityStrategy end

function master(::FeasibilityCuts, ::Type{T}) where T <: AbstractFloat
    return FeasibilityCutsMaster(T)
end

function worker(::FeasibilityCuts, linking_constraints::Vector{MOI.ConstraintIndex}, model::MOI.ModelLike)
    # Cache objective
    func_type = MOI.get(model, MOI.ObjectiveFunctionType())
    obj = MOI.get(model, MOI.ObjectiveFunction{func_type}())
    return FeasibilityCutsWorker(obj, linking_constraints, Vector{MOI.VariableIndex}())
end
function worker_type(::FeasibilityCuts)
    return FeasibilityCutsWorker
end


MOIU.@model(
    DualizableModel,
    (),
    (MOI.EqualTo, MOI.GreaterThan, MOI.LessThan),
    (
        MOI.Reals,
        MOI.Zeros,
        MOI.Nonnegatives,
        MOI.Nonpositives,
        MOI.SecondOrderCone,
        MOI.RotatedSecondOrderCone,
        MOI.ExponentialCone,
        MOI.DualExponentialCone,
        MOI.PositiveSemidefiniteConeTriangle,
    ),
    (MOI.PowerCone, MOI.DualPowerCone),
    (),
    (MOI.ScalarAffineFunction,),
    (MOI.VectorOfVariables,),
    (MOI.VectorAffineFunction,)
)

"""
    PrimalDualMap{T}

Maps information from all structures of the primal to the dual model.

* `constrained_var_idx::Dict{VI,Tuple{CI,Int}}`: maps original primal
constrained variables to their primal original constraints (the special ones
that makes them constrained variables) and their internal index (if vector
constraints, VectorOfVariables-in-Set), 1 otherwise (VariableIndex-in-Set).

* `constrained_var_dual::Dict{CI,CI}`: maps the original primal constraint index
of constrained variables to the dual model's constraint index of the
associated dual constraint.

* `constrained_var_zero::Dict{CI,Unions{SAF,VAF}}`: caches scalar affine
functions or vector affine functions associated with constrained variables
of type `VectorOfVariables`-in-`Zeros` or
`VariableIndex`-in-`EqualTo(zero(T))` as
their duals would be `func`-in-`Reals`, which are "irrelevant" to the model.
This information is cached for completeness of the `DualOptimizer` for
`get`ting `ConstraintDuals`.

* `primal_var_dual_con::Dict{VI,CI}`: maps "free" primal variables to their
associated dual constraints. Free variables as opposed to constrained
variables. Note that Dualization will select automatically which variables
are free and which are constrained.

* `primal_con_dual_var::Dict{CI,Vector{VI}}`: maps primal constraint indices to
vectors of dual variable indices. For scalar constraints those vectors will be
single element vectors.

* `primal_con_dual_con::Dict{CI,CI}`: maps primal constraints to
dual variable constraints (if there is such constraint the dual
dual variable is said to be constrained). If the primal constraint's set
is EqualTo or Zeros, no constraint is added in the dual variable (the
dual variable is said to be free).

* `primal_con_constants::Dict{CI,Vector{T}}`: maps primal constraints to
their respective constants, which might be inside the set.
This map is used in `MOI.get(::DualOptimizer,::MOI.ConstraintPrimal,ci)`
that requires extra information in the case that the scalar set constrains
a constant (`EqualtTo`, `GreaterThan`, `LessThan`).

* `primal_parameter::Dict{VI,VI}`: maps parameters in the primal to parameters
in the dual model.

* `primal_var_dual_quad_slack::Dict{VI,VI}`: maps primal variables
(that appear in quadratic objective terms) to dual "slack" variables.

"""
mutable struct PrimalDualMap{T}
    constrained_var_idx::Dict{VI,Tuple{CI,Int}}
    constrained_var_dual::Dict{CI,CI}
    constrained_var_zero::Dict{
        CI,
        Union{MOI.VectorAffineFunction{T},MOI.ScalarAffineFunction{T}},
    }
    primal_var_dual_con::Dict{VI,CI}
    primal_con_dual_var::Dict{CI,Vector{VI}}
    primal_con_dual_con::Dict{CI,CI}
    primal_con_constants::Dict{CI,Vector{T}}

    primal_parameter::Dict{VI,VI}
    primal_var_dual_quad_slack::Dict{VI,VI}

    function PrimalDualMap{T}() where {T}
        return new(
            Dict{VI,Tuple{CI,Int}}(),
            Dict{CI,CI}(),
            Dict{
                CI,
                Union{MOI.VectorAffineFunction{T},MOI.ScalarAffineFunction{T}},
            }(),
            Dict{VI,CI}(),
            Dict{CI,Vector{VI}}(),
            Dict{CI,CI}(),
            Dict{CI,Vector{T}}(),
            Dict{VI,VI}(),
            Dict{VI,VI}(),
        )
    end
end

function is_empty(primal_dual_map::PrimalDualMap{T}) where {T}
    return isempty(primal_dual_map.constrained_var_idx) &&
           isempty(primal_dual_map.constrained_var_dual) &&
           isempty(primal_dual_map.constrained_var_zero) &&
           isempty(primal_dual_map.primal_var_dual_con) &&
           isempty(primal_dual_map.primal_con_dual_var) &&
           isempty(primal_dual_map.primal_con_dual_con) &&
           isempty(primal_dual_map.primal_con_constants) &&
           isempty(primal_dual_map.primal_parameter) &&
           isempty(primal_dual_map.primal_var_dual_quad_slack)
end

function empty!(primal_dual_map::PrimalDualMap)
    Base.empty!(primal_dual_map.constrained_var_idx)
    Base.empty!(primal_dual_map.constrained_var_dual)
    Base.empty!(primal_dual_map.constrained_var_zero)
    Base.empty!(primal_dual_map.primal_var_dual_con)
    Base.empty!(primal_dual_map.primal_con_dual_var)
    Base.empty!(primal_dual_map.primal_con_dual_con)
    Base.empty!(primal_dual_map.primal_con_constants)
    Base.empty!(primal_dual_map.primal_parameter)
    Base.empty!(primal_dual_map.primal_var_dual_quad_slack)
    return
end

"""
    DualProblem{T,OT<:MOI.ModelLike}

Result of the `dualize` function. Contains the fields:

* `dual_model::OT`: contaninf a optimizer or data structure with the
`MathOptInterface` definition of the resulting dual model.

* `primal_dual_map::PrimalDualMap{T}`: a data structure to hold information of
how the primal and dual model are related in terms of indices (`VariableIndex`
and `ConstraintIndex`) and other data.
"""
struct DualProblem{T,OT<:MOI.ModelLike}
    dual_model::OT #It can be a model from an optimizer or a DualizableModel{T}
    primal_dual_map::PrimalDualMap{T}

    function DualProblem{T}(
        dual_optimizer::OT,
        pdmap::PrimalDualMap{T},
    ) where {T,OT<:MOI.ModelLike}
        return new{T,OT}(dual_optimizer, pdmap)
    end
end

function DualProblem{T}(dual_optimizer::OT) where {T,OT<:MOI.ModelLike}
    return DualProblem{T}(dual_optimizer, PrimalDualMap{T}())
end

function DualProblem(dual_optimizer::OT) where {OT<:MOI.ModelLike}
    return DualProblem{Float64}(dual_optimizer)
end

# Empty DualProblem cosntructor
function DualProblem{T}() where {T}
    return DualProblem{T}(DualizableModel{T}(), PrimalDualMap{T}())
end


struct DfPair
    df1::DataFrame
    df2::DataFrame
end

abstract type AbstractMatch end

struct Match <: AbstractMatch
    i::Int
    j::Int
    dist::Real
end

struct HalfMatch <: AbstractMatch
    ind::Int
    dist::Real
end

isless(x::AbstractMatch, y::AbstractMatch) = x.dist < y.dist;

abstract type StepParameters end

struct InitParams <: StepParameters
    entity_type::String
    left_sample::Int
    right_sample::Int
end

struct ExtendParams <: StepParameters
    rel_pair_ind::Int
    root_side::Int
end

struct MatchParams <: StepParameters
    needed_matches::Int
end

struct IntegrateParams <: StepParameters end

struct Normalizer
    means::Array{Real, 1}
    stds::Array{Real, 1}
end

struct SortedJoiner
    arr::Array{Int, 2}
    sorted_inds::Array{Array{Int, 1}, 1}
    sorted_views::Array{
        SubArray{Int, 1, Array{Int, 2}, Tuple{Array{Int, 1}, Int}, false},
        1,
    }
end

struct RelPair
    sj1::SortedJoiner
    sj2::SortedJoiner
    colmap::Tuple{String, String}
end

struct CorefSystem
    esp_dict::Dict{String, DfPair}
    normalizers::Dict{String, Normalizer}
    relps::Array{RelPair, 1}
end

mutable struct EntityRefSpace
    space::Array{Int, 2}
    height::Int
    width::Int
    entity_types::Array{String, 1}
end

struct SpacePair
    erspace1::EntityRefSpace
    erspace2::EntityRefSpace
end

heights(space_pair::SpacePair) = (space_pair.erspace1.height, space_pair.erspace2.height)
spaces(space_pair::SpacePair) = (space_pair.erspace1, space_pair.erspace2)
innerspaces(space_pair::SpacePair) = (ers.space for ers in spaces(space_pair))

mutable struct CorefProposal
    ind::Int
    latent_vars::Array{Real,1}
end

struct CorefSpace
    latent_match_vars::Dict{String, Array{Array{CorefProposal,1},1}}
    prepspace::Array{HalfMatch, 2}
    free_i::Array{Bool, 1}
    free_j::Array{Bool, 1}
    param_log::Array{StepParameters, 1}
end

mutable struct CorefResolver
    crs::CorefSystem
    crspace::CorefSpace
    space_pairs::Tuple{SpacePair, SpacePair}
    leftspace_main::Bool
end

function valid(proposal::CorefProposal)
    proposal.ind > 0
end

function parse(lvars::Array{Array{CorefProposal,1},1})
    out = Array{Array{Int, 2}, 1}()
    for (fun, lvar_arr) in zip([x -> x, reverse], lvars)
        for (i, propsal) in enumerate(lvar_arr)
            if valid(propsal)
                r = reshape(fun([i, propsal.ind]), (1, 2))
                (r in out) && continue
                push!(out, r)
            end
        end
    end
    vcat(out...)
end

result(crr::CorefResolver) = Dict(name=> parse(lvars) for (name, lvars) in crr.crspace.latent_match_vars)
main(crr::CorefResolver) = crr.space_pairs[crr.leftspace_main ? 1 : 2]
innerspaces(crr::CorefResolver) = (crr.leftspace_main ? x -> x : reverse)(crr.space_pairs)
switch!(crr::CorefResolver) = (crr.leftspace_main = !crr.leftspace_main)


module TwoElectronTDDFT

using LinearAlgebra

include("abstract_types.jl")
include("operator.jl")
include("Potential.jl")
include("Hamiltonian.jl")

end


"""
This file is used to test the simulation. This is just the simulate_states code unwrapped
so will be a little less optimal than when directly run but should enable an easier
observation of correctness.
"""

using Revise
using Distributed
using Plots
using Chain
using ProfileView

include("simulate_states.jl")

# parameters to pass to the main function

file_date = "2022-04-12"

# set seed for consistent plots
rng = Random.Xoshiro(2022)

jurisdiction_assumptions = JurisdictionAssumptions()

# get the latest onset date
latest_start_date = Dates.Date(
    maximum(v for v in values(jurisdiction_assumptions.simulation_start_dates))
)

omicron_start_date = "2021-11-15"
omicron_dominant_date = "2021-12-15"
(dates, local_case_dict, import_case_dict) = read_in_cases(file_date, rng)
last_date_in_data = dates[end]
forecast_end_date = last_date_in_data + Dates.Day(35)
# create vector for dates
onset_dates = latest_start_date:Dates.Day(1):forecast_end_date
# add a small truncation to the simulations as we don't trust the most recent data
truncation_days = 7
# states to simulate
state = "QLD"
nsims = 1000
forecast_start_date = Dates.Date(jurisdiction_assumptions.simulation_start_dates[state])

# named tuple for initial conditions
D0 = jurisdiction_assumptions.initial_conditions[state]
N = jurisdiction_assumptions.pop_sizes[state]

# get the observed cases
local_cases = local_case_dict[state][dates .>= forecast_start_date]
# cutoff the last bit of the local cases
import_cases = import_case_dict[state]
local_cases = Int.(local_cases[begin:end-truncation_days+1])
import_cases = Int.(import_cases)

local_cases_input = CSV.read("data/local_cases_input_2022-04-13.csv", DataFrame)
local_cases_input = filter(:state => ==(state), local_cases_input)

plot(dates, local_case_dict[state])
plot!(local_cases_input.date_onset, local_cases_input.count)
xlims!((Dates.Date("2022-01-01"), Dates.Date("2022-04-15")))

include("simulate_states.jl")

(Z, D, U) = simulate_branching_process(
    D0,
    N,
    nsims,
    local_cases,
    import_cases,
    forecast_start_date,
    file_date,
    state,
    adjust_TP = false,
)

save_simulations(
    D,
    state,
    file_date,
    onset_dates,
    rng,
)

let
    forecast_start_date = Dates.Date(jurisdiction_assumptions.simulation_start_dates[state])
    local_cases = local_case_dict[state][dates .>= forecast_start_date]
    local_cases = local_cases[begin:end]
    D_local = D[:, 1, :] + D[:, 2, :]
    D_local_median = median(D_local, dims = 2)
    D_local_mean = mean(D_local, dims = 2)

    d_start = Dates.Date(forecast_start_date)
    Δd = Dates.Day(1)
    d_end = Dates.Date(forecast_start_date) + Δd * (size(D_local, 1) - 1)
    D_local_dates = d_start:Δd:d_end

    f = plot()
    # plot!(f, dates[dates .>= forecast_start_date], local_cases, legend = false, linealpha = 1, lc = 1)
    # plot!(f, Z_historical[36:end, :], legend = false, lc = 3, linealpha = 0.5)
    plot!(f, D_local_dates, D_local, legend = false, lc = 1, linealpha = 0.5)
    # plot!(f, sims_old[!, "onset date"], sims_old_mat, legend = false, lc = 2, linealpha = 0.5)
    # plot!(f, dates[dates .>= forecast_start_date][begin:end-truncation_days+1], local_cases[begin:end-truncation_days+1], legend = false, lc = :black, lw = 2)
    plot!(f, dates[dates .>= forecast_start_date], local_cases, legend = false, lc = :black, lw = 2)
    # plot!(f, local_cases2.date_onset, local_cases2.count, legend = false, linealpha = 1, lc = 1)
    plot!(f, D_local_dates, D_local_median, legend = false, lc = 2, lw = 2)
    # plot!(f, df_state[!, "onset date"], Matrix(Int.(df_state[!, 3:end-1])))
    vline!(f, [Dates.Date("2021-11-15")], ls = :dash, lc = :black)
    # xlims!(f, 0, length(local_cases) + 35)
    ylims!(f, 0, 40000)
    xlims!(f, Dates.value(Dates.Date("2021-12-15")), Dates.value(Dates.Date("2022-05-20")))
    # ylims!(f, 0, 500)
end

cases_adjusted = CSV.read("results/EpyReff/cases_adjusted_2022-04-12.csv", DataFrame)

cases_adjusted_state = filter(:STATE => ==(state), cases_adjusted)

plot(cases_adjusted_state.local_scaled_14)
plot!(cases_adjusted_state.local_scaled_28)
plot!(cases_adjusted_state.local_scaled_42)
plot!(cases_adjusted_state.local_scaled_72)
plot!(cases_adjusted_state.local_scaled_102)
plot!(cases_adjusted_state.local_scaled_132)

ϕ = read_in_susceptible_depletion(file_date)
ϕ_mean = mean(ϕ, dims = 1)

ψ = (
    ϕ_mean[6] * (cases_adjusted_state.local_scaled_14) / N
    + ϕ_mean[5] * (cases_adjusted_state.local_scaled_28) / N
    + ϕ_mean[4] * (cases_adjusted_state.local_scaled_42) / N
    + ϕ_mean[3] * (cases_adjusted_state.local_scaled_72) / N
    + ϕ_mean[2] * (cases_adjusted_state.local_scaled_102) / N
    + ϕ_mean[1] * (cases_adjusted_state.local_scaled_132) / N
)

scaling_factor = 1 .- ψ

scaling_factor_sim = zeros(size(Z, 1))
prop_inf_full = zeros(size(Z, 1), 6)

Z_historical = mean(Z, dims = 2)

for day in 1:length(Z_historical)

    ψ = 0.0

    τ_horizons_start = [0, 14, 28, 42, 72, 102]
    τ_horizons_end = [1, 14, 28, 42, 72, 102, 132] .- 1

    for i in 1:length(τ_horizons_start)

        # 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(day - τ_horizons_end[i + 1], 1)
        τ_end = min(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)
        ψ += ϕ_mean[7 - i] * prop_inf

        prop_inf_full[day, i] = prop_inf

    end

    scale_factor = max(0.0, 1.0 - ψ)

    scaling_factor_sim[day] = scale_factor

end

d_start = Dates.Date(forecast_start_date) - Dates.Day(30)
Δd = Dates.Day(1)
d_end = Dates.Date(forecast_start_date) + Δd * (size(Z, 1) - 1) - Dates.Day(30)
D_local_dates = d_start:Δd:d_end

plot(cases_adjusted_state.date_inferred, scaling_factor)
plot!(D_local_dates, scaling_factor_sim)
xlims!((Dates.Date("2021-12-01"), D_local_dates[end]))
ylabel!("Scaling factor")

fig = plot(layout = (3, 2), legend = false)
plot!(
    fig,
    subplot = 1,
    cases_adjusted_state.date_inferred,
    cases_adjusted_state.local_scaled_14 ./ N,
    title = "14 days", titlefontsize=8
)
plot!(
    fig,
    subplot = 2,
    cases_adjusted_state.date_inferred,
    cases_adjusted_state.local_scaled_28 ./ N,
    title = "28-14 days", titlefontsize=8
)
plot!(
    fig,
    subplot = 3,
    cases_adjusted_state.date_inferred,
    cases_adjusted_state.local_scaled_42 ./ N,
    title = "42-28 days", titlefontsize=8
)
plot!(
    fig,
    subplot = 4,
    cases_adjusted_state.date_inferred,
    cases_adjusted_state.local_scaled_72 ./ N,
    title = "72-42 days", titlefontsize=8
)
plot!(
    fig,
    subplot = 5,
    cases_adjusted_state.date_inferred,
    cases_adjusted_state.local_scaled_102 ./ N,
    title = "102-72 days", titlefontsize=8)
plot!(
    fig,
    subplot = 6,
    cases_adjusted_state.date_inferred,
    cases_adjusted_state.local_scaled_132 ./ N,
    title = "132-102 days", titlefontsize=8)
plot!(
    fig,
    subplot = 1,
    D_local_dates,
    prop_inf_full[:, 1],
    ls = :dash
)
plot!(
    fig,
    subplot = 2,
    D_local_dates,
    prop_inf_full[:, 2],
    ls = :dash
)
plot!(
    fig,
    subplot = 3,
    D_local_dates,
    prop_inf_full[:, 3],
    ls = :dash
)
plot!(
    fig,
    subplot = 4,
    D_local_dates,
    prop_inf_full[:, 4],
    ls = :dash
)
plot!(
    fig,
    subplot = 5,
    D_local_dates,
    prop_inf_full[:, 5],
    ls = :dash
)
plot!(
    fig,
    subplot = 6,
    D_local_dates,
    prop_inf_full[:, 6],
    ls = :dash
)
xlims!((Dates.Date("2021-12-01"), D_local_dates[end]))
ylims!((0, 0.25))


import Base: setindex!, getindex

include("params.jl")

"""
    LindaEnv

Optimizer environment. Holds parameters' values.
"""
mutable struct LindaEnv

    num_cols_rmp_max::LindaIntParam  # Maximum number of columns in RMP

    num_columns_max::LindaIntParam  # Maximum number of columns before pricing stops

    num_cgiter_max::LindaIntParam  # Maximum number of Column-Generation iterations

    time_limit::LindaFloatParam  # Time limit, in seconds

    tol_reduced_cost::LindaFloatParam  # Numerical tolerance for reduced costs

    verbose::LindaIntParam  # Verbosity level

    function LindaEnv()
        env = new()

        env.num_cols_rmp_max = LindaRealParam(:num_cols_rmp_max,
            typemax(Int64), zero(Int64), typemax(Int64)
        )

        env.num_columns_max = LindaRealParam(:num_columns_max,
            typemax(Int64), zero(Int64), typemax(Int64)
        )

        env.num_cgiter_max = LindaRealParam(:num_cgiter_max,
            typemax(Int64), zero(Int64), typemax(Int64)
        )

        env.time_limit = LindaRealParam(:time_limit, Inf, 0.0, Inf)

        env.tol_reduced_cost = LindaRealParam(:tol_reduced_cost, 10.0^-6, 0.0, Inf)

        env.verbose = LindaRealParam(:verbose, 0, 0, 1 )

        return env
    end

end

# TODO: use `getproperty` and `setproperty!` instead (> Julia 0.7)
getindex(env::LindaEnv, name::Symbol) = get_param_value(Core.getfield(env, name))

getindex(env::LindaEnv, ::Type{Val{:num_cols_rmp_max}}) = env.num_cols_rmp_max.val
getindex(env::LindaEnv, ::Type{Val{:num_columns_max}}) = env.num_columns_max.val
getindex(env::LindaEnv, ::Type{Val{:num_cgiter_max}}) = env.num_cgiter_max.val
getindex(env::LindaEnv, ::Type{Val{:time_limit}}) = env.time_limit.val
getindex(env::LindaEnv, ::Type{Val{:tol_reduced_cost}}) = env.tol_reduced_cost.val
getindex(env::LindaEnv, ::Type{Val{:verbose}}) = env.verbose.val

setindex!(env::LindaEnv, v, name::Symbol) = set_param_value!(Core.getfield(env, name), v)

"""
    setparam!(env; kwargs...)

Set parameters to specified values.
"""
function setparam!(env; kwargs...)
    for (s, v) in kwargs
        set_param_value!(Core.getfield(env, s), v)
    end
end

"""
    reset!(env)

Reset all parameters to their default values.
"""
function reset!(env)
    for s in fieldnames(LindaEnv)
        set_param_default!(Core.getfield(env, s))
    end
    return nothing
end


struct ExactArc{T<:Node,U<:Real,V<:Int}
    uuid                    ::UUID
    origin                  ::T
    destination             ::T
    assignment              ::V
    value                   ::U
end

Base.show(io::IO,arc::ExactArc) = Base.print(
    io,
    "\n   Assignment: ", arc.assignment, "\n",
    "   Arc Value: ", arc.value, "\n"
)

#change hashing of nodes to use uuids
Base.hash(arc::ExactArc) = hash(arc.uuid)
Base.:(==)(x::ExactArc,y::ExactArc) = x.uuid == y.uuid

#do UUIDs automatically
function ExactArc(origin::T, destination::T, assignment::U, value::V)where{T<:Node, U<:Int, V<:Real}
    return ExactArc(uuid4(), origin, destination, assignment, value)
end

"""
    deepCopy(arc::ExactArc)

Deep copy an arc

# Arguments
- `arcs::ExactArc`: the arc
"""
function deepCopy(arc::ExactArc)
    return ExactArc(arc.uuid, arc.origin,arc.destination,arc.assignment,arc.value)
end

"""
    deepCopy(arcs::Vector{ExactArc})

Deep copy a list of arcs

# Arguments
- `arcs::Array{ExactArc}`: a list of arcs
"""
function deepCopy(arcs::Vector{ExactArc})
    newArcs = Vector{ExactArc}()
    for arc in arcs
        push!(newArcs, deepCopy(arc))
    end
    return newArcs
end

"""
    shallowCopy(arc::ExactArc)

Shallow copy an arc

# Arguments
- `arcs::ExactArc`: the arc
"""
function shallowCopy(arc::ExactArc)
    return ExactArc(arc.origin,arc.destination,arc.assignment,arc.value)
end

"""
    shallowCopy(arcs::Vector{ExactArc})

functionally the same as deepCopy a list of arcs, but with different uuids

# Arguments
- `arcs::Array{ExactArc}`: a list of arcs
"""
function shallowCopy(arcs::Vector{ExactArc})
    newArcs = Vector{ExactArc}()
    for arc in arcs
        push!(newArcs, shallowCopy(arc))
    end
    return newArcs
end

"""
    findBestArc(model::SequencingModel, arcs::Vector{ExactArc})

Find the best arc in a list, where best means containing the most optimal path

# Arguments
- `model::SequencingModel`: The Sequencing problem
- `arcs::Vector{ExactArc}`: The list of arcs
"""
function findBestArc(model::SequencingModel, arcs::Vector{ExactArc})
    bestArc = arcs[1]
    bestArcValue = bestArc.origin.lengthToRoot + bestArc.value
    for arc in arcs
        if isBetterSolutionValue(model, arc.origin.lengthToRoot + arc.value, bestArcValue)
            bestArc = arc
            bestArcValue = arc.origin.lengthToRoot + arc.value
        end
    end
    return bestArc, bestArcValue
end

"""
    deleteArc!(arc::ExactArc)
Remove an arc from both relevant nodes

# Arguments
- `arc::ExactArc`: the arc to remove
"""
function deleteArc!(arc::ExactArc)
    filter!(x->x!=arc, arc.origin.arcsOut)
    filter!(x->x!=arc, arc.destination.arcsIn)
end

"""
    deleteArcs!(arcsToRemove::Vector{ExactArc})
Delete a list of arcs

# Arguments
- `arcsToRemove::Vector{ExactArc}`: the arcs to remove
"""
function deleteArcs!(arcsToRemove::Vector{ExactArc})
    toRemove = Vector{ExactArc}()
    for arc in arcsToRemove
        push!(toRemove, arc)
    end
    for arc in toRemove
        deleteArc!(arc)
    end
end


module ConstLab

import Base.show

using NLsolve

export loadcase, driver, MatStatus, MatParameter, create_component_macros

const NON_CONV_MESSAGE = "constitutive iterations did not converge"

type NonConvergenceError <: Exception end
show(io::IO, nce::NonConvergenceError) = print(io, NON_CONV_MESSAGE)

const MAT_NON_CONV_MESSAGE = "material iterations did not converge"

type MaterialNonConvergenceError <: Exception end
show(io::IO, nce::NonConvergenceError) = print(io, MAT_NON_CONV_MESSAGE)

abstract MatStatus
abstract MatParameter
abstract MatModel

include("driver.jl")
include("loadcase.jl")
include("utils.jl")

end # module


"""
    read_drifters(pth,lst;chnk=Inf,rng=(missing,missing))

Read near-surface [drifter data](https://doi.org/10.1002/2016JC011716) from the
[Global Drifter Program](https://doi.org/10.25921/7ntx-z961) into a DataFrame.

Note: need to use NetCDF.jl as NCDatasets.jl errors when TIME = Inf
```
pth,list=drifter_files()
df=read_drifters( pth*lst[end],chnk=1000,rng=(2014.1,2014.2) )
#sort!(df, [:t, :lat])
#CSV.write(pth*"Drifter_hourly_2005_2019.csv", df)
```
"""
function read_drifters(fil::String;chnk=Inf,rng=(missing,missing))
   t=ncread(fil,"TIME")
   t_u=ncgetatt(fil,"TIME","units")
   lo=ncread(fil,"LON")
   la=ncread(fil,"LAT")
   ID=ncread(fil,"ID")

   ii=findall(isfinite.(lo.*la.*t))
   t_ii=t[ii]
   t_ii=timedecode.(t_ii, t_u)
   tmp=dayofyear.(t_ii)+(hour.(t_ii) + minute.(t_ii)/60 ) /24
   t_ii=year.(t_ii)+tmp./daysinyear.(t_ii)
   isfinite(rng[1]) ? jj=findall( (t_ii.>rng[1]).&(t_ii.<=rng[2])) : jj=1:length(ii)

   ii=ii[jj]
   t=t_ii[jj]
   lo=lo[ii]
   la=la[ii]
   ID=ID[ii]

   df = DataFrame(ID=Int[], lon=Float64[], lat=Float64[], t=Float64[])

   !isinf(chnk) ? nn=Int(ceil(length(ii)/chnk)) : nn=1
   for jj=1:nn
      #println([jj nn])
      !isinf(chnk) ? i=(jj-1)*chnk.+(1:chnk) : i=(1:length(ii))
      i=i[findall(i.<length(ii))]
      append!(df,DataFrame(lon=lo[i], lat=la[i], t=t[i], ID=Int.(ID[i])))
   end

   return df
end

"""
    read_drifters( pth, lst )

Read near-surface [hourly drifter data](https://doi.org/10.1002/2016JC011716) from the
[Global Drifter Program](https://doi.org/10.25921/7ntx-z961) into a DataFrame.

Note: need to use NetCDF.jl as NCDatasets.jl errors when TIME = Inf

```
pth,list=drifter_files()
df=read_drifters( pth, lst)
```
"""
function read_drifters( pth, lst )
   df = DataFrame([fill(Int, 1) ; fill(Float64, 3)], [:ID, :lon, :lat, :t])
   for fil in lst
      println(fil)
      append!(df,read_drifters( pth*fil,chnk=10000,rng=(2005.0,2020.0) ))
      #println(size(df))
   end
   return df
end

function drifter_files()
   pth="Drifter_hourly_v013/"
   lst=["driftertrajGPS_1.03.nc","driftertrajWMLE_1.02_block1.nc","driftertrajWMLE_1.02_block2.nc",
      "driftertrajWMLE_1.02_block3.nc","driftertrajWMLE_1.02_block4.nc","driftertrajWMLE_1.02_block5.nc",
      "driftertrajWMLE_1.02_block6.nc","driftertrajWMLE_1.03_block7.nc"]
   return pth,list
end

"""
    read_velocities(γ::gcmgrid,t::Int,pth::String)

Read velocity components `u,v` from files in `pth`for time `t`
"""
function read_velocities(γ::gcmgrid,t::Int,pth::String)
    u=read_nctiles("$pth"*"UVELMASS/UVELMASS","UVELMASS",γ,I=(:,:,:,t))
    v=read_nctiles("$pth"*"VVELMASS/VVELMASS","VVELMASS",γ,I=(:,:,:,t))
    return u,v
end

"""
    read_mds(filRoot::String,x::MeshArray)

Read a gridded variable from 2x2 tile files. This is used
in `example2_setup()` with `flt_example/`
"""
function read_mds(filRoot::String,x::MeshArray)
   prec=eltype(x)
   prec==Float64 ? reclen=8 : reclen=4;

   (n1,n2)=Int64.(x.grid.ioSize ./ 2);
   fil=filRoot*".001.001.data"
   tmp1=stat(fil);
   n3=Int64(tmp1.size/n1/n2/reclen);

   v00=x.grid.write(x)
   for ii=1:2; for jj=1:2;
      fid = open(filRoot*".00$ii.00$jj.data")
      fld = Array{prec,1}(undef,(n1*n2*n3))
      read!(fid,fld)
      fld = hton.(fld)

      n3>1 ? s=(n1,n2,n3) : s=(n1,n2)
      v00[1+(ii-1)*n1:ii*n1,1+(jj-1)*n2:jj*n2,:]=reshape(fld,s)
   end; end;

   return x.grid.read(v00,x)
end

"""
    read_uvetc(k::Int,Γ::Dict,pth::String)

Define `uvetc` given the grid variables `Γ` and a vertical level choice `k`
including velocities obtained from files in `pth`

**deprecated: use set_up_𝑃 and update_𝑃! instead -- only still used in GlobalOceanNotebooks ?**
"""
function read_uvetc(k::Int,Γ::Dict,pth::String)
    𝑃 = dict_to_nt(IndividualDisplacements.NeighborTileIndices_cs(Γ))
    Γ = dict_to_nt( Γ )
    nt=12; msk=(Γ.hFacC[:,k] .> 0.) #select depth

    u=0. ./Γ.DXC; v=0. ./Γ.DYC;
    for t=1:nt
        (U,V)=read_velocities(Γ.XC.grid,t,pth)
        for i=1:size(u,1)
            u[i]=u[i] + U[i,k]
            v[i]=v[i] + V[i,k]
        end
    end
    u=u ./ nt
    v=v ./ nt #time average

    u[findall(isnan.(u))]=0.0; v[findall(isnan.(v))]=0.0 #mask with 0s rather than NaNs
    u=u./Γ.DXC; v=v./Γ.DYC; #normalize to grid units

    (u,v)=exchange(u,v,1) #add 1 point at each edge for u and v
    XC=exchange(Γ.XC) #add 1 lon point at each edge
    YC=exchange(Γ.YC) #add 1 lat point at each edge

    t0=0.0; t1=86400*366*10.0; dt=10*86400.0;
    tmp = (u0=u, u1=u, v0=v, v1=v, t0=t0, t1=t1, dt=dt, msk=msk, XC=XC, YC=YC)

    return merge(𝑃,tmp)
end

"""
    get_ecco_velocity_if_needed()

Download `MITgcm` transport output to `examples/nctiles_climatology` if needed
"""
function get_ecco_velocity_if_needed()
    p=dirname(pathof(OceanStateEstimation))
    lst=joinpath(p,"../examples/nctiles_climatology.csv")
    pth=ECCOclim_path
    !isdir(pth*"UVELMASS") ? get_from_dataverse(lst,"UVELMASS",pth) : nothing
    !isdir(pth*"VVELMASS") ? get_from_dataverse(lst,"VVELMASS",pth) : nothing
    !isdir(pth*"WVELMASS") ? get_from_dataverse(lst,"WVELMASS",pth) : nothing
end

"""
    get_occa_velocity_if_needed()

Download `MITgcm` transport output to `examples/OCCA_climatology` if needed
"""
function get_occa_velocity_if_needed()
    p=dirname(pathof(OceanStateEstimation))
    lst=joinpath(p,"../examples/OCCA_climatology.csv")
    pth=OCCAclim_path
    nams = ("DDuvel.0406clim.nc","DDvvel.0406clim.nc","DDwvel.0406clim.nc","DDtheta.0406clim.nc","DDsalt.0406clim.nc")
    if !isfile("$pth"*"DDuvel.0406clim.nc")
        tmp=joinpath(pth,"tmp/")
        !isdir(tmp) ? mkdir(tmp) : nothing
        [get_from_dataverse(lst,nam,tmp) for nam in nams]
        [mv(joinpath(tmp,nam,nam),joinpath(pth,nam)) for nam in nams]
    end
end


# routines related to the ParallelStencil backend

export initialize_backend!

"""
    start_backend(Backend{BackendParallelStencil}; Scalar=Float64, dim=1)

Starts the `ParallelStencil` backend for the `Scalar`, which loads ParallelStencil (and, if requested, MPI) for the correct number of dimensions
"""
function check_backend(b::Backend{BackendParallelStencil}; dim=1, Scalar=Float64)
    # Check whether MPI is loaded
    if (b.mpi==true && !isdefined(Main, :ImplicitGlobalGrid))
        error("ImplicitGlobalGrid is not loaded; ensure it is loaded first with: using ImplicitGlobalGrid")
    end
    if (b.mpi==true && !isdefined(Main, :MPI))
        error("MPI is not loaded; ensure it is loaded first with: using MPI")
    end

    # Check whether ParallelStencil is loaded
    if !isdefined(Main, :ParallelStencil)
        error("ParallelStencil is not loaded; ensure it is loaded first with: using ParallelStencil")
    end

    # Load GPU or CPU backend
    if ((b.arch==:GPU) || (b.arch==:CUDA))
        # Currently, ParallelStencil, only supports CUDA
        if !isdefined(Main, :CUDA)
            error("CUDA is not loaded; ensure it is loaded first with: using CUDA")
        end
    end

    return nothing
end

function initialize_backend!(grid, b::Backend{BackendParallelStencil}, stencilwidth, opts, fields::NamedTuple)

    N_vec = ones(Int,3);
    D = grid.dim

    Nl,dims = localfromglobalsize(backend.mpi, grid.N, opts)      # local size using MPI (w/out halo)
    N_vec[1:D] = collect(Nl)

    initmpi=false
    if !MPI.Initialized()
        initmpi=true
    end

    # Halo
    overlap  = zeros(Int64,3)
    overlap[1:D] .= 2*stencilwidth
    Nl_halo  = N_vec .+ overlap

    # Periodic?
    period  = zeros(Int64,3)
    period[findall(grid.topology.==Periodic)] .= 1

    IGG_data = init_global_grid(Nl_halo[1], Nl_halo[2], Nl_halo[3],
                                                            init_MPI=initmpi,
                                                            quiet=false,
                                                            overlapx=overlap[1], overlapy=overlap[2], overlapz=overlap[3],
                                                            periodx=period[1], periody=period[2], periodz=period[3];
                                                            opts...);

    IGG_data = ImplicitGlobalGrid.get_global_grid()
    grid.IGG = IGG_data;                    # store

    # Store local grid dimensions
    grid.Nl = (Nl_halo[1:D]...,)            # local with halo
    grid.Ng = (IGG_data.nxyz_g[1:D]...,)    # global + halo

    # Compute local indices of grid in global grid (not including the halo)
    ind_s = IGG_data.coords.*(IGG_data.nxyz .- IGG_data.overlaps)
    ind_e = ind_s[1:D] .+ Nl;
    ind_local = ();
    for i=1:D
        ind_local = (ind_local..., ind_s[i]+1:ind_e[i])
    end
    grid.ind_local = ind_local

    # Update corners accordingly, which provides an easy way to iterate over the local portions of the grid, with/without ghost/halo points
    size = ind_e .- ind_s[1:D];
    grid.corners = ( lower=(CartesianIndex( (ones(Int64,D) .+ stencilwidth)...,)),
                    upper=(CartesianIndex( (zeros(Int64,D) .+ stencilwidth + size)...,)),
                    size=(size...,))
    grid.ghostcorners = ( lower=(CartesianIndex( (ones(Int64,D))...,)),
                        upper=(CartesianIndex( (zeros(Int64,D) .+ (2 .* stencilwidth) + size)...,)),
                        size =(Nl_halo[1:D]...,))

    # Update the 1D coordinate vectors
    local_coordinates!(grid, grid.backend)

    return nothing

end

"""
    initialize_fields!(grid, b::Backend{BackendParallelStencil}, fields::NamedTuple)

This initializes fields to a grid structure, in case the `ParallelStencil` is employed.
The names of the fields are a NamedTuple, which specifies the initial value (constant) value of the arrays, such as `fields=(T=0, P=1, Ci=2)`.

Note: typically, this routine is not called directly but as part of setting up the grid

# Example

```
julia> @init_backend(ParallelStencil, Threads, false, Float64);
julia> initialize_fields!(grid, b::Backend{BackendParallelStencil}, fields::NTuple)
```

"""
function initialize_fields!(grid, b::Backend{BackendParallelStencil, FT}, fields::NamedTuple) where FT

    names = keys(fields)
    for ifield=1:length(fields)

        # Created named tuple for new field. We use @eval here, to avoid
        name = names[ifield]
        value = fields[ifield]

        # Create a string with the expression to be evaluated (as calling @ones directly induces a compilation error)
        if b.arch == :Threads
            new_field = ParallelStencil.ParallelKernel.@ones_threads(grid.Nl)*value
        elseif b.arch == :CUDA
            new_field = ParallelStencil.ParallelKernel.@ones_cuda(grid.Nl)*value
        else
            error("Not yet implemented")
        end

        # Add to Tuple
        grid.fields = add_field(grid.fields,name,new_field)
    end
end

"""
    globalfromlocalsize(mpi, localsize::NTuple, opts::NTuple, stencilwidth::Int, b::Backend{BackendParallelStencil, FT})

Compute global from local size in case we use the ParallelStencil backend
"""
function globalfromlocalsize(size, localsize, opts, stencilwidth, topology, b::Backend{BackendParallelStencil, FT}) where FT

    if !b.mpi & !isempty(localsize)
        size = localsize   # not using MPI/IGG
    elseif b.mpi & !isempty(localsize)
        dim = length(localsize);
        dims = zeros(Int64,3)

        dims[dim+1:end] .= 1;       # fix
        if typeof(opts) <: Dict # process command-line options
            dims[1] = get(opts, :dimx, dims[1]);
            if dim>1;   dims[2] = get(opts, :dimy, dims[2]);    end
            if dim>2;   dims[3] = get(opts, :dimz, dims[3]);    end
        end
        size_vec = ones(Int,3);
        size_vec[1:dim] .= localsize;

        # Periodic boundary
        period  = zeros(Int64,3)
        period[findall(topology.==Periodic)] .= 1

        # Use MPI to compute processor partitioning (dims contain the # of processors in each direction)
        overlap = stencilwidth*2

        # Initialize the global grid
        init_global_grid(size_vec[1],size_vec[2],size_vec[3],
                        dimx=dims[1], dimy=dims[2], dimz=dims[3],
                        overlapx=overlap, overlapy=overlap, overlapz=overlap,
                        periodx=period[1],  periody=period[2], periodz=period[3],
                        init_MPI=false, quiet=true)

        size_global = [nx_g(), ny_g(), nz_g()] .- overlap

        size = (size_global[1:dim]...,)
    else
        size = size
    end

    return size
end


using Pkg
Pkg.activate(".")
using DataFrames

global ULLR_required_columns = [ :upper_left_x, :lower_right_x, :upper_left_y, :lower_right_y ]

"""
    has_required_columns( df::DataFrame, required_columns::Vector )::Bool

Returns `true` if a `DataFrame` instance has all of the `required_columns` present.

"""
function has_required_columns( df::DataFrame, required_columns::Vector )::Bool
    df_names = names( df )
    result = true
    for required_column in required_columns
        if !( required_column in df_names )
            result = false
            break
        end
    end
    return result
end

"""
    highest_vote!( df::DataFrame, class_columns::Vector )::DataFrame

Finds highest vote amongst all specified class columns and creates a new column with that variable name.

"""
function highest_vote!( df::DataFrame, class_columns::Vector )::DataFrame
    df_classes = df[ !, class_columns ]
    df[ !, :Highest_Vote ] = map( row -> argmax( row ), eachrow( df_classes ) )
    df_classes = df[ !, vcat( class_columns, :Highest_Vote ) ]
    df[ !, :Highest_Prediction ] = map( row -> row[row.Highest_Vote], eachrow( df_classes ) )
    return df
end

highest_vote!( dfs::Vector{DataFrame}, class_columns::Vector )::Vector{DataFrame} = [ highest_vote!(df, class_columns) for df in dfs ]

"""
    barycenter_boxes!( df::DataFrame )::DataFrame

Will convert a DataFrame containing the following variables: center_x, center_y, width, height to be formatted in upper left lower right coordinates.

"""
function barycenter_boxes!( df::DataFrame )::DataFrame
    req_cols = [ :center_x, :center_y, :width, :height ]
    @assert( has_required_columns( df, req_cols ), "Input dataframe does not have the required columns: $req_cols. Please check spelling and capitalization of column names." )
    df[ !, [ :width, :height] ] ./= 2
    df[ !, :upper_left_x ] = df.center_x - df.width
    df[ !, :lower_right_x] = df.center_x + df.width
    df[ !, :upper_left_y ] = df.center_y - df.height
    df[ !, :lower_right_y] = df.center_y + df.height
    select!( df, Not( req_cols ) )
    return df
end

"""
    LLUR_boxes!( df::DataFrame )::DataFrame
"""
function LLUR_boxes!( df::DataFrame )::DataFrame
    req_cols = [ :lower_left_x, :upper_right_x, :lower_left_y, :upper_right_y ]
    new_cols = ULLR_required_columns
    @assert( has_required_columns( df, req_cols ), "Input dataframe does not have the required columns: $req_cols. Please check spelling and capitalization of column names." )
    rename!( df, Dict( req_cols .=> new_cols ) )
    return df
end

"""
    area( df::DataFrame )::DataFrame

Calculates the areas of all Boxes in a DataFrame as a new column.

"""
function area!( df::DataFrame )::DataFrame
    @assert( has_required_columns( df, ULLR_required_columns ), "Input dataframe does not have the required columns: $ULLR_required_columns. Please check spelling and capitalization of column names and ensure the format is ULLR." )
    df[ !, :Area ] = ( df.lower_right_x - df.upper_left_x ) .* ( df.lower_right_y - df.upper_left_y )
    return df
end

area!( dfs::Vector{DataFrame} )::Vector{DataFrame} = [ area!(df) for df in dfs ]

"""
    intersection_area( a::DataFrame, b::DataFrame )::DataFrame

Calculates the area of intersection between 2 DataFrames containing ULLR boxes.
This returns a new DataFrame

"""
function intersection_over_union( a::DataFrame, b::DataFrame )::DataFrame
    @assert( has_required_columns( a, ULLR_required_columns ), "Input dataframe does not have the required columns: $ULLR_required_columns. Please check spelling and capitalization of column names and ensure the format is ULLR." )
    @assert( has_required_columns( b, ULLR_required_columns ), "Input dataframe does not have the required columns: $ULLR_required_columns. Please check spelling and capitalization of column names and ensure the format is ULLR." )
    @assert( has_required_columns( a, [:Area] ), "Input dataframe does not have the required columns: Area. Please run `area()` on the individual dataframes." )
    @assert( has_required_columns( b, [:Area] ), "Input dataframe does not have the required columns: Area. Please run `area()` on the individual dataframes." )
    len_a, len_b = first.(size.( [ a, b ] ) )
    total_len = len_a * len_b

    output = DataFrame( (   :ID_A => Vector{Int}( undef, total_len ),
                            :ID_B => Vector{Int}( undef, total_len ),
                            :class_labels => Vector{Symbol}( undef, total_len ),
                            :Intersection_Area => Vector{Float64}( undef, total_len ),
                            :Intersection_Over_Union => Vector{Float64}( undef, total_len ),
                            :Vote_Matches => Vector{Bool}( undef, total_len ),
                            :Highest_Prediction => Vector{Float64}( undef, total_len),
                            :Highest_Vote => Vector{Symbol}( undef, total_len)
                        ) )
    current_row = 1
    for ( row_a_ID, row_a ) in enumerate( eachrow( a ) ), ( row_b_ID, row_b ) in enumerate( eachrow( b ) )
        #current_row = row_a_ID + ( len_a * ( row_b_ID - 1 ) )
        output[ current_row, :ID_A ] = row_a_ID
        output[ current_row, :ID_B ] = row_b_ID
        output[ current_row, :class_labels ] = row_b.class_labels
        output[ current_row, :Vote_Matches ] = row_a.Highest_Vote == row_b.class_labels
        output[ current_row, :Highest_Prediction ] = row_a.Highest_Prediction #row_b[!,row_a.Highest_Vote]
        output[ current_row, :Highest_Vote ] = row_a.Highest_Vote
        output[ current_row, :Intersection_Area ] = max( 0, min( row_a.lower_right_x, row_b.lower_right_x ) - max( row_a.upper_left_x, row_b.upper_left_x ) ) .*
                                                    max( 0, min( row_a.lower_right_y, row_b.lower_right_y ) - max( row_a.upper_left_y, row_b.upper_left_y ) )
        output[ current_row, :Intersection_Over_Union ] = output[ current_row, :Intersection_Area ] / ( row_a.Area + row_b.Area - output[ current_row, :Intersection_Area ])
        current_row += 1
    end
    #remove intermediate calculation column Intersection_Area
    select!(output, Not(:Intersection_Area))
    #Split data by each unique prediction
    #   -> sort by intersection over union (for each prediction all others are constant)
    #   -> grab the highest IoU and toss the rest.
    return combine( map( x -> sort(x, :Intersection_Over_Union, rev = true)[1,:], groupby( output, :ID_A ) ) )
end

mutable struct ObjectDetectionStats
    predictions::Vector{DataFrame}
    groundtruth::Vector{DataFrame}
    IoU_map::Vector{DataFrame}
    classnames::Vector{Symbol}
    len::Int
end

"""
    potential_scores( ods::ObjectDetectionStats; IoU_threshold = 0.5 )

"""
function potential_scores( cur_IoU_df::DataFrame; IoU_threshold = 0.5 )
    cur_IoU_df[ :iou_threshold ] = cur_IoU_df.Intersection_Over_Union .> IoU_threshold
    #is TP or FN?
    cur_IoU_df[ :TP_or_FN ] = cur_IoU_df.Vote_Matches .& cur_IoU_df.iou_threshold
    #Ensure we do not double dip from a GT prediction. Only one prediction can map to a GT here...
    #Group by ground truth class label
    #   -> sort by the highest predicted class score
    #       _-> find first prediced class matches above an IoU threshold
    for gdf in groupby( cur_IoU_df, :class_labels )
        sort!( cur_IoU_df, :Highest_Prediction, rev = true)
        kept_first = false
        for row in eachrow( gdf )
            if row.TP_or_FN
                if kept_first
                    row.TP_or_FN = false
                end
                kept_first = true
            end
        end
    end
    #Bifurcate DF on the possibility of true positives
    TP_or_FN, FP_or_FN = DataFrame.( collect( groupby( cur_IoU_df, :TP_or_FN ) ) )
    return TP_or_FN, FP_or_FN
end

"""
    ObjectDetectionStats(   predictions::Vector{DataFrame},
                            groundtruth::Vector{DataFrame},
                            classnames::Vector{Symbol} )::ObjectDetectionStats

Instantiates an `ObjectDetectionStats` object.

"""
function ObjectDetectionStats(  predictions::Vector{DataFrame},
                                groundtruth::Vector{DataFrame},
                                classnames::Vector{Symbol} )::ObjectDetectionStats
    len = length(predictions)
    @assert len == length(groundtruth) "Prediction and groundtruth Vector's must have the same length."
    area!( predictions )
    highest_vote!( predictions, classnames )
    area!( groundtruth )
    IoU_map = [ intersection_over_union( p, gt ) for (p, gt) in zip( predictions, groundtruth) ]
    return ObjectDetectionStats( predictions, groundtruth, IoU_map, classnames, len )
end

# each prediction maps to every GT
# predictions sorted by their confidence scores (highest to lowest)
# Each prediction can ONLY map to one GT
#  for each fixed prediction -> [GT] the prediction score is static
#  Need to find "most likely hit" via IoU.
#So sort by confidence and IoU
# Now for each group of predictions, remove everything but the top ranked ones
#Now we have the most representative prediction for each prediction to [GT]
#Find TP-or-FN mappings
#Find FP-or-FN mappings

using DataFrames
a = DataFrame(Dict([ :a => [1,1,1,2,2,2,3,3,3], :b => randn(9)]))
x,z,t = DataFrame.( collect( groupby(a, :a) ) )

a

KeyDF = DataFrame( Dict( [    :a               => [0.8, 0.1, 0.1],
                              :b               => [0.1, 0.7, 0.1],
                              :c               => [0.1, 0.2, 0.8],
                              :upper_left_x     => [12.5, 10.0 ,20.0],
                              :lower_right_x    => [17.5, 20.0, 30.0],
                              :upper_left_y     => [7.50, 5.00, 30.0],
                              :lower_right_y    => [12.5, 15.0, 35.0]
                         ] ) )

GT = DataFrame( Dict( [     :class_labels     => [:a, :b, :c],
                                :upper_left_x     => [12.5, 10.0, 16.0],
                                :lower_right_x    => [17.5, 20.0, 60.0],
                                :upper_left_y     => [7.50, 5.00, 16.0],
                                :lower_right_y    => [12.5, 15.0, 20.0]
                      ] ) )

ods = ObjectDetectionStats(  [KeyDF], [GT], [:a,:b,:c] )

score( ods, detection_threshold = 0.5, IoU_threshold = 0.5 )

@show ods.IoU_map[1]

#@benchmark join( KeyDF, KeyDFB, kind = :cross; makeunique = true )
#new_df = intersection_over_union( area( KeyDF ), area( KeyDFB ) )
@show( new_df )

# Pkg.add("DataFrames")
# using DataFrames
#
# b = DataFrame(  (:A => [1,2,1,2,2,2,3,3,3], :B => collect(1:9) )  )
#
# for gdf in groupby( b, :A )
#     gdf.B = map(x -> (x.B % 2) == 1, eachrow(gdf))
#     println(gdf)
# end
#
# b


function SimpsonIntegrator()
    return SimpsonIntegrator(())
end

function SimpsonIntegrator(arg0::jdouble, arg1::jdouble, arg2::jint, arg3::jint)
    return SimpsonIntegrator((jdouble, jdouble, jint, jint), arg0, arg1, arg2, arg3)
end

function SimpsonIntegrator(arg0::jint, arg1::jint)
    return SimpsonIntegrator((jint, jint), arg0, arg1)
end


###########################################
# States

"""
    AbstractState{N, W}
Abstract type representing state of the (index) automaton over an alphabet of `N` letters.

The subtypes of `AbstractState{N, W}` should implement the following methods which
constitute `AbstractState` interface:
 * `name`: external name of the state
 * `isterminal`: predicate to determine if state is terminal (i.e. whether the
    state represents left-hand side of some rewriting rule)
 * `rightrule`: if state is terminal, returns the right-hand side of the rewriting rule.
 Undefined in other cases.
 * `inedges`: container (indexed by letters of `AbstractWord`) of source states of in-edges.
 * `outedges`: container (indexed by letters of `AbstractWord`) of destination states of out-edges.
 * `isfailstate`: predicate to determine if state represents "fail" state.
 * `Base.length`: the length of the signature of the shortest path to the
    state, which starts at initial state (length of "simple path")
"""
abstract type AbstractState{N, W} end

"""
    State{N, W<:AbstractWord} <: AbstractState{N, W}
State of an index automaton over an alphabet of `N` letters.
"""
mutable struct State{N, W<:AbstractWord} <: AbstractState{N, W}
    name::W
    terminal::Bool
    failstate::Bool
    rrule::W
    ined::Vector{State{N, W}}
    outed::NTuple{N, State{N, W}}

    function State{N, W}() where {N, W<:AbstractWord}
        x = new{N, W}()
        x.failstate = true
        return x
    end

    function State(name::W, failstate::State{N, W}) where {N, W}
        s = new{N, W}(
            name,
            false,
            false,
            one(W),
            State{N,W}[],
            ntuple(i -> failstate, N),
        )
        return s
    end
end

name(s::State) = s.name
isterminal(s::State) = s.terminal
rightrule(s::State) = s.rrule
inedges(s::State) = s.ined
outedges(s::State) = s.outed
isfailstate(s::State) = s.failstate
Base.length(s::State) = length(name(s))

"""
    declarerightrule!(s::State, w::AbstractWord)
Decalres given word as a right-hand side rule of a given state (and makes this)
state terminal.
"""
function declarerightrule!(s::State, w::AbstractWord)
    s.terminal = true
    s.rrule = w
end

function Base.show(io::IO, s::State)
    if isfailstate(s)
        print(io, "Fail State")
    elseif isterminal(s)
        print(io, "TState($(name(s)) → $(rightrule(s)))")
    else
        print(io, "NTState($(length(inedges(s))), $(length(outedges(s))))")
    end
end

function Base.show(io::IO, ::MIME"text/plain", s::State)
    if isfailstate(s)
        print(io, "Fail state")
    elseif isterminal(s)
        println(io, "Terminal state $(name(s))")
        println(io, "Right rule: $(rightrule(s))")
        println(io, " Edges entering:")
        for (i, e) in enumerate(inedges(s))
            !isfailstate(e) && println(io, "   ", " - with label $(i) from state $(name(e))")
        end
    else
        println(io, "State $(name(s))")
        println(io, " Edges entering:")
        for (i, fromst) in enumerate(inedges(s))
            !isfailstate(fromst) && println(io, "   ", " - with label $(i) from state $(name(fromst))")
        end
        println(io, " Edges leaving:")
        for (i, tost) in enumerate(outedges(s))
            !isfailstate(tost) && println(io, "   ", " - with label $(i) from state $(name(tost))")
        end
    end
end

###########################################
# Automata

"""
    AbstractAutomaton{N, W}
Abstract type representing (index) automaton.

The subtypes of `AbstractAutomaton` should implement the following methods which
constitute `AbstractAutomaton` interface:
 * `alphabet`: return the alphabet of the automaton
 * `states`: return the list of states in the automaton
 * `initialstate`: return the initial state of the automaton
 * `failstate`: return the unique fail state (representing "no edge from source with a given label" situation)
 * `Base.push!`: append a single state to the automaton
 * `addedge!`: add edge between two states
 * `removeedge!`: remove edge between two states
 * `Base.deleteat!`: delete state at given position (together with edges)
 * `walk`: travel through the automaton according to given path and initial state
"""
abstract type AbstractAutomaton{N, W} end

"""
    Automaton{N, W<:AbstractWord} <: AbstractAutomaton{N, W}
Automaton as a vector of states together with alphabet.
"""
struct Automaton{N, W<:AbstractWord} <: AbstractAutomaton{N, W}
    states::Vector{State{N, W}}
    abt::Alphabet
    failstate::State{N, W}
    stateslengths::Vector{Int}
    _past_states::Vector{State{N, W}}
end

# Default type is set to UInt16
function Automaton(abt::Alphabet, W::Type{<:AbstractWord}=Word{UInt16})
    S = State{length(abt), W}
    failstate = S()
    Automaton{length(abt), W}([State(one(W), failstate)], abt, failstate, [0], S[])
end

alphabet(a::Automaton) = a.abt
states(a::Automaton) = a.states
initialstate(a::Automaton) = first(states(a))
failstate(a::Automaton) = a.failstate
stateslengths(a::Automaton) = a.stateslengths

function Base.isempty(a::Automaton)
    σ = initialstate(a)
    length(states(a)) == 1 && isone(name(σ)) && return true
    return false
end

Base.push!(a::Automaton{N, W}, name::W) where {N, W} =
    (push!(states(a), State(name, failstate(a))); push!(stateslengths(a), length(name)); a)
Base.empty!(a::Automaton{N, W}) where{N, W} =
    (empty!(states(a)); empty!(stateslengths(a)); push!(a, one(W)); a)

"""
    replace(k::NTuple{N, T}, val::T, idx::Integer) where {N, T}
Returns a copy of `NTuple` `k` with value at index `idx` changed to `val`.
"""
Base.@propagate_inbounds function replace(k::NTuple{N, T}, val::T, idx::Integer) where {N, T}
    @boundscheck 1 ≤ idx ≤ N || throw(BoundsError(k, idx))
    return @inbounds ntuple( i -> i==idx ? val : k[i], Val(N))
end

"""
    updateoutedges!(src::State, dst::State, label::Integer)
Updates outedges of the state `src` with state `dst` connected through the edge `label`.
"""
Base.@propagate_inbounds function updateoutedges!(src::State{N}, dst::State{N}, label::Integer) where {N}
    @boundscheck 1 ≤ label ≤ N || throw(BoundsError(outedges(src), label))
    @inbounds src.outed = replace(outedges(src), dst, label)
end

"""
    addinedge!(dst::State, src::State)
Add in-edge to `dst` from state `src` to state `dst` storing it in `inedges` of the state `dst`.
"""
addinedge!(dst::State, src::State) = push!(inedges(dst), src)

"""
    addedge!(a::Automaton, label::Integer, src::Integer, dst::Integer)
Adds the edge with a given `label` directed from `src` state to `dst` state.
"""
Base.@propagate_inbounds function addedge!(a::Automaton{N}, label::Integer, src::Integer, dst::Integer) where {N}
    @boundscheck 1 ≤ label ≤ N || throw(BoundsError(outedges(states(a)[src]), label))
    @boundscheck checkbounds(states(a), src)
    @boundscheck checkbounds(states(a), dst)
    @inbounds srcstate = states(a)[src]
    @inbounds dststate = states(a)[dst]
    @inbounds updateoutedges!(srcstate, dststate, label)
    addinedge!(dststate, srcstate)
end

Base.@propagate_inbounds function addedge!(a::Automaton{N}, label::Integer, src::State, dst::State) where {N}
    @boundscheck 1 ≤ label ≤ N || throw(BoundsError(outedges(src), label))
    @assert src in states(a) "Source state is not in automaton!"
    @assert dst in states(a) "Destination state to is not in automaton"
    @inbounds updateoutedges!(src, dst, label)
    addinedge!(dst, src)
end
# Should we also check if the edge corresponding to a given letter does not exist?

"""
    removeedge!(a::Automaton, label::Integer, from::Integer, to::Integer)
Removes the edge with a given `label` `from` one state `to` another.
"""
Base.@propagate_inbounds function removeedge!(a::Automaton{N}, label::Integer, from::Integer, to::Integer) where {N}
    @boundscheck 1 ≤ label ≤ N || throw(BoundsError(outedges(states(a)[from]), label))
    @boundscheck checkbounds(states(a), from)
    @boundscheck checkbounds(states(a), to)
    @inbounds fromstate = states(a)[from]
    @inbounds tostate = states(a)[to]
    @inbounds updateoutedges!(fromstate, failstate(a), label)
    deleteat!(inedges(tostate), findfirst(isequal(fromstate), inedges(tostate)))
end

Base.@propagate_inbounds function Base.deleteat!(a::Automaton, idx::Integer)
    @boundscheck checkbounds(states(a), idx)
    @boundscheck checkbounds(stateslengths(a), idx)
    @inbounds state = states(a)[idx]
    for σ in outedges(state)
        isfailstate(σ) || deleteat!(inedges(σ), findfirst(isequal(state), inedges(σ)))
    end
    for σ in inedges(state)
        updateoutedges!(σ, failstate(a), findfirst(isequal(state), outedges(σ)))
    end
    @inbounds deleteat!(states(a), idx)
    @inbounds deleteat!(stateslengths(a), idx)
end

"""
    walk(a::AbstractAutomaton, signature::AbstractWord[, state=initialstate(a)])
Walk through index automaton according to the path given by the `signature`, starting from `state`.

Return a tuple `(idx, state)` where `idx` is the length of prefix of signature which was successfully traced and
`state` is the final state.
Note that if `idx ≠ length(signature)` there is no path in the automaton corresponding to the full signature.
"""
function walk(a::AbstractAutomaton, signature::AbstractWord, state=initialstate(a))
    idx = 0
    for (i, k) in enumerate(signature)
        next = outedges(state)[k]
        isfailstate(next) ? break : (idx, state) = (i, next)
    end
    return (idx, state)
end

function Base.show(io::IO, a::AbstractAutomaton)
    println(io, "Automaton with $(length(states(a))) states")
    for (i, state) in enumerate(states(a))
        print(io, " $i. Edges leaving state (")
        print_repr(io, name(state), alphabet(a))
        println(io, "):")

        for (i, tostate) in enumerate(outedges(state))
            if !isfailstate(tostate)
                print(io,
                "   ", " - with label ", alphabet(a)[i], " to state (")
                print_repr(io, name(tostate), alphabet(a))
                println(io, ")")
            end
        end
    end
end

###########################################
# Ad index automata

function makeindexautomaton!(a::Automaton, rws::RewritingSystem)
    empty!(a)
    # Determining simple paths
    for (lhs, rhs) in rules(rws)
        σ = initialstate(a)
        for (i, letter) in enumerate(lhs)
            if isfailstate(outedges(σ)[letter])
                push!(a, lhs[1:i])
                addedge!(a, letter, σ, last(states(a)))
            end
            σ = outedges(σ)[letter]
        end
        terminal = last(states(a))
        declarerightrule!(terminal, rhs)
        # Add loops for terminal state
        for i in 1:length(outedges(terminal))
            addedge!(a, i, terminal, terminal)
        end
    end
    # Determining cross paths
    for (i, state) in enumerate(outedges(initialstate(a)))
        isfailstate(state) && addedge!(a, i, 1, 1)
    end
    i = 1
    indcs = findall(isequal(i), stateslengths(a))
    while !isempty(indcs)
        for idx in indcs
            σ = states(a)[idx]
            τ = last(walk(a, @view name(σ)[2:end]))
            @assert !isfailstate(τ) "Tracing $σ in an automaton failed!"
            for (letter, state) in enumerate(outedges(σ))
                isfailstate(state) && addedge!(a, letter, σ, outedges(τ)[letter])
            end
        end
        i += 1
        resize!(indcs, 0)
        for (k,len) in enumerate(stateslengths(a))
            len == i && push!(indcs, k)
        end
    end
    return a
end

"""
    makeindexautomaton(rws::RewritingSystem)
Creates index automaton corresponding to a given rewriting system.
"""
makeindexautomaton(rws::RewritingSystem) =
    (a = Automaton(alphabet(rws)); makeindexautomaton!(a, rws))

"""
    updateautomaton!(a::Automaton, rws::RewritingSystem)
Updates given automaton so it corresponds to a given rewriting system.
"""
updateautomaton!(a::Automaton, rws::RewritingSystem) = makeindexautomaton!(a, rws)

"""
    rewrite_from_left!(v::AbstractWord, w::AbstractWord, a::Automaton)
Rewrites word `w` from left using index automaton `a` and appends the result
to `v`. For standard rewriting `v` should be empty.
"""
function rewrite_from_left!(v::AbstractWord, w::AbstractWord, a::Automaton)
    tape = a._past_states
    resize!(tape, length(w) + 1)
    tape[1] = initialstate(a)
    initial_length = length(v)
    @debug "Rewriting $w with automaton"
    while !isone(w)
        x = popfirst!(w)
        k = length(v) - initial_length + 1
        @debug "" v x w
        s = tape[k]
        state = outedges(s)[x]

        if isterminal(state)
            lhs, rhs = name(state), rightrule(state)
            @debug "applying rewriting rule $lhs → $rhs"
            w = prepend!(w, rhs)
            resize!(v, length(v) - length(lhs) + 1)
        else
            length(tape) <= k ? push!(tape, state) : tape[k+1] = state
            push!(v, x)
            @debug "state is not terminal, pushed $x, v=$v"
        end
    end
    return v
end