"""
    SMCModel(M!::F1, lG::F2, maxn::Int64, particle::Type, pScratch::Type) where
      {F1<:Function,F2<:Function}
- ```M!``` Mutation function
- ```lG``` Log potential function
- ```maxn``` Maximum n for which the model is well-defined
- ```particle``` Type of a particle
- ```pScratch``` Type of particle scratch space
"""
struct SMCModel{F1<:Function,F2<:Function}
  M!::F1
  lG::F2
  maxn::Int64
  particle::Type
  pScratch::Type
end

const SubVectorF64 = SubArray{Float64, 1, Array{Float64, 1},
  Tuple{UnitRange{Int64}}, true}
const SubVectorI64 = SubArray{Int64, 1, Array{Int64, 1},
  Tuple{UnitRange{Int64}}, true}
const SubVector{Particle} = SubArray{Particle,1,Array{Particle,1},Tuple{UnitRange{Int64}},true}

# the fields of this struct are used by the parallel SMC algorithm
# they are not intended for use by users of the package
mutable struct _SMCInternalParallel{Particle, ParticleScratch}
  Nperthread::Int64
  maxlws::Vector{Float64}
  partialSums::Vector{Float64}
  partialSumSqs::Vector{Float64}
  Ns::Vector{Int64}
  NsPartial::Vector{Int64}
  particleScratches::Vector{ParticleScratch}

  localAs::Vector{SubVectorI64}
  localEves::Vector{SubVectorI64}
  localOldEves::Vector{SubVectorI64}
  localZetas::Vector{SubVector{Particle}}
  localZetaAncs::Vector{SubVector{Particle}}
  localWs::Vector{SubVectorF64}
  localLogWs::Vector{SubVectorF64}
  localScratch1s::Vector{SubVectorF64}

  ## below are only populated if fullOutput = true
  localAllAs::Vector{Vector{SubVectorI64}}
  localAllZetas::Vector{Vector{SubVector{Particle}}}
  localAllEves::Vector{Vector{SubVectorI64}}
  localAllWs::Vector{Vector{SubVectorF64}}
end

function _SMCInternalParallel{Particle, ParticleScratch}(N::Int64, n::Int64,
  nthreads::Int64, fullOutput::Bool) where {Particle, ParticleScratch}
  @assert mod(N, nthreads) == 0 "N must be a multiple of nthreads"
  Nperthread = div(N, nthreads)

  if nthreads > 1
    maxlws::Vector{Float64} = Vector{Float64}(undef, nthreads)
    partialSums::Vector{Float64} = Vector{Float64}(undef, nthreads)
    partialSumSqs::Vector{Float64} = Vector{Float64}(undef, nthreads)
    Ns::Vector{Int64} = Vector{Int64}(undef, nthreads)
    NsPartial::Vector{Int64} = Vector{Int64}(undef, nthreads)

    particleScratches::Vector{ParticleScratch} =
      Vector{ParticleScratch}(undef, nthreads)
    # avoid false sharing
    Threads.@threads for i = 1:nthreads
      @inbounds particleScratches[i] = ParticleScratch()
    end

    localAs::Vector{SubVectorI64} =
      Vector{SubVectorI64}(undef, nthreads)
    localEves::Vector{SubVectorI64} =
      Vector{SubVectorI64}(undef, nthreads)
    localOldEves::Vector{SubVectorI64} =
      Vector{SubVectorI64}(undef, nthreads)
    localZetas::Vector{SubVector{Particle}} =
      Vector{SubVector{Particle}}(undef, nthreads)
    localZetaAncs::Vector{SubVector{Particle}} =
      Vector{SubVector{Particle}}(undef, nthreads)
    localWs::Vector{SubVectorF64} =
      Vector{SubVectorF64}(undef, nthreads)
    localLogWs::Vector{SubVectorF64} =
      Vector{SubVectorF64}(undef, nthreads)
    localScratch1s::Vector{SubVectorF64} =
      Vector{SubVectorF64}(undef, nthreads)
  else
    maxlws = Vector{Float64}(undef, 0)
    partialSums = Vector{Float64}(undef, 0)
    partialSumSqs = Vector{Float64}(undef, 0)
    Ns = Vector{Int64}(undef, 0)
    NsPartial = Vector{Int64}(undef, 0)
    particleScratches = Vector{ParticleScratch}(undef, 0)

    localAs = Vector{SubVectorI64}(undef, 0)
    localEves = Vector{SubVectorI64}(undef, 0)
    localOldEves = Vector{SubVectorI64}(undef, 0)
    localZetas = Vector{SubVector{Particle}}(undef, 0)
    localZetaAncs = Vector{SubVector{Particle}}(undef, 0)
    localWs = Vector{SubVectorF64}(undef, 0)
    localLogWs = Vector{SubVectorF64}(undef, 0)
    localScratch1s = Vector{SubVectorF64}(undef, 0)
  end

  if fullOutput && nthreads > 1
    localAllZetas::Vector{Vector{SubVector{Particle}}} =
      Vector{Vector{SubVector{Particle}}}(undef, n)
    localAllAs::Vector{Vector{SubVectorI64}} =
      Vector{Vector{SubVectorI64}}(undef, n-1)
    localAllEves::Vector{Vector{SubVectorI64}} =
      Vector{Vector{SubVectorI64}}(undef, n)
    localAllWs::Vector{Vector{SubVectorF64}} =
      Vector{Vector{SubVectorF64}}(undef, n)
    for p = 1:n
      localAllZetas[p] = Vector{SubVector{Particle}}(undef, nthreads)
      localAllEves[p] = Vector{SubVectorI64}(undef, nthreads)
      localAllWs[p] = Vector{SubVectorF64}(undef, nthreads)
      p < n && (localAllAs[p] = Vector{SubVectorI64}(undef, nthreads))
    end
  else
    localAllAs = Vector{Vector{SubVectorI64}}(undef, 0)
    localAllZetas = Vector{Vector{SubVector{Particle}}}(undef, 0)
    localAllEves = Vector{Vector{SubVectorI64}}(undef, 0)
    localAllWs = Vector{Vector{SubVectorF64}}(undef, 0)
  end

  return _SMCInternalParallel(Nperthread, maxlws, partialSums, partialSumSqs,
    Ns, NsPartial, particleScratches, localAs, localEves, localOldEves,
    localZetas, localZetaAncs, localWs, localLogWs, localScratch1s, localAllAs,
    localAllZetas, localAllEves, localAllWs)
end

# the fields of this struct are used by the SMC algorithm
# they are not intended for use by users of the package
mutable struct _SMCInternal{Particle, ParticleScratch}
  zetaAncs::Vector{Particle}
  oldEves::Vector{Int64}
  as::Vector{Int64}
  lws::Vector{Float64}
  scratch1::Vector{Float64}
  scratch2::Vector{Float64}
  vecOnes::Vector{Float64}
  nresamples::Int64

  sws::Float64
  mws::Float64
  maxlw::Float64
  particleScratch::ParticleScratch

  parallel::_SMCInternalParallel{Particle, ParticleScratch}
end

## constructor for _SMCInternal
function _SMCInternal{Particle, ParticleScratch}(N::Int64, n::Int64,
  nthreads::Int64, fullOutput::Bool) where {Particle, ParticleScratch}

  lws = Vector{Float64}(undef, N)
  as = Vector{Int64}(undef, N)
  fill!(as, 1)
  scratch1 = Vector{Float64}(undef, N)
  scratch2 = Vector{Float64}(undef, N)
  vecOnes = ones(N)
  nresamples::Int64 = 0
  zetaAncs = Vector{Particle}(undef, N)
  for i in 1:N
    zetaAncs[i] = Particle()
  end
  oldEves = Vector{Int64}(undef, N)

  ## assign user-defined particle scratch space
  @assert ParticleScratch == Nothing || !isbitstype(ParticleScratch)
  particleScratch = ParticleScratch()

  parallel::_SMCInternalParallel{Particle, ParticleScratch} =
    _SMCInternalParallel{Particle, ParticleScratch}(N, n, nthreads, fullOutput)

  return _SMCInternal(zetaAncs, oldEves, as, lws, scratch1, scratch2, vecOnes,
    nresamples, 0.0, 0.0, 0.0, particleScratch, parallel)
end

function _assignThreadViews(internal::_SMCInternal{Particle},
  zetas::Vector{Particle}, eves::Vector{Int64}, ws::Vector{Float64},
  allZetas::Vector{Vector{Particle}}, allWs::Vector{Vector{Float64}},
  allAs::Vector{Vector{Int64}}, allEves::Vector{Vector{Int64}}) where Particle

  zetaAncs::Vector{Particle} = internal.zetaAncs
  as::Vector{Int64} = internal.as
  oldEves::Vector{Int64} = internal.oldEves
  lws::Vector{Float64} = internal.lws
  scratch1::Vector{Float64} = internal.scratch1

  ip::_SMCInternalParallel = internal.parallel
  nthreads::Int64 = length(ip.localWs)
  Nperthread::Int64 = ip.Nperthread
  for i = 1:nthreads
    start = Nperthread * (i - 1) + 1
    finish = start + Nperthread - 1
    ip.localAs[i] = view(as, start:finish)
    ip.localEves[i] = view(eves, start:finish)
    ip.localOldEves[i] = view(oldEves, start:finish)
    ip.localZetas[i] = view(zetas, start:finish)
    ip.localZetaAncs[i] = view(zetaAncs, start:finish)
    ip.localWs[i] = view(ws, start:finish)
    ip.localLogWs[i] = view(lws, start:finish)
    ip.localScratch1s[i] = view(scratch1, start:finish)
  end

  fullOutput::Bool = length(ip.localAllWs) > 0
  if fullOutput && nthreads > 1
    n::Int64 = length(allWs)
    for p = 1:n
      for i = 1:nthreads
        start = Nperthread * (i - 1) + 1
        finish = start + Nperthread - 1
        ip.localAllZetas[p][i] = view(allZetas[p], start:finish)
        ip.localAllEves[p][i] = view(allEves[p], start:finish)
        ip.localAllWs[p][i] = view(allWs[p], start:finish)
        p < n && (ip.localAllAs[p][i] = view(allAs[p], start:finish))
      end
    end
  end
end

## SMC input / output struct. Also contains internal state for the SMC
## implementation
"""
    SMCIO{Particle, ParticleScratch}
Structs of this type should be constructed using the provided constructor.
Important fields:
- ```N::Int64``` Number of particles ``N``
- ```n::Int64``` Number of steps ``n``
- ```nthreads::Int64``` Number of threads
- ```fullOutput::Bool``` Whether particle system history should be recorded
- ```essThreshold::Float64``` Relative ESS Threshold ``\\tau``
- ```zetas::Vector{Particle}``` Time n particles ``\\zeta_n^1, \\ldots, \\zeta_n^N``
- ```eves::Vector{Int64}``` Time n Eve indices ``E_n^1, \\ldots, E_n^N``
- ```ws::Vector{Float64}``` Time n weights ``W_n^1, \\ldots, W_n^N``
- ```logZhats::Vector{Float64}``` ``\\log(\\hat{Z}^N_1), \\ldots, \\log(\\hat{Z}^N_n)``
- ```Vhat1s::Vector{Float64}``` ``\\hat{V}_1^N(1), \\ldots, \\hat{V}_n^N(1)``
- ```esses::Vector{Float64}``` Relative ESS values ``\\mathcal{E}_1^N, \\ldots, \\mathcal{E}_n^N``
- ```resample::Vector{Bool}``` Resampling indicators ``R_1, \\ldots, R_{n-1}``
Populated only if ```fullOutput == true```
- ```allZetas::Vector{Vector{Particle}}``` All the particles
- ```allWs::Vector{Vector{Float64}}``` All the weights
- ```allAs::Vector{Vector{Int64}}``` All the ancestor indices
- ```allEves::Vector{Vector{Int64}}``` All the Eve indices
"""
struct SMCIO{Particle, ParticleScratch}
  N::Int64
  n::Int64
  nthreads::Int64
  zetas::Vector{Particle}
  eves::Vector{Int64}
  ws::Vector{Float64}
  logZhats::Vector{Float64}
  Vhat1s::Vector{Float64}
  esses::Vector{Float64}
  resample::Vector{Bool}
  fullOutput::Bool
  essThreshold::Float64

  internal::_SMCInternal{Particle, ParticleScratch} # for internal use only

  # these are only populated if fullOutput = true
  allZetas::Vector{Vector{Particle}}
  allWs::Vector{Vector{Float64}}
  allAs::Vector{Vector{Int64}}
  allEves::Vector{Vector{Int64}}
end

"""
    SMCIO{Particle, ParticleScratch}(N::Int64, n::Int64, nthreads::Int64,
      fullOutput::Bool, essThreshold::Float64 = 2.0) where
      {Particle, ParticleScratch}
Constructor for ```SMCIO``` structs.
"""
function SMCIO{Particle, ParticleScratch}(N::Int64, n::Int64, nthreads::Int64,
  fullOutput::Bool, essThreshold::Float64 = 2.0) where {Particle,
  ParticleScratch}
  @assert hasmethod(Particle, ()) "Particle() must exist"
  @assert hasmethod(ParticleScratch, ()) "ParticleScratch() must exist"

  zetas = Vector{Particle}(undef, N)
  for i=1:N
    zetas[i] = Particle()
  end

  eves = Vector{Int64}(undef, N)
  ws = Vector{Float64}(undef, N)
  logZhats = Vector{Float64}(undef, n)
  Vhat1s = Vector{Float64}(undef, n)
  esses = Vector{Float64}(undef, n)
  resample = Vector{Bool}(undef, n-1)

  if fullOutput
    allZetas = Vector{Vector{Particle}}(undef, n)
    for i=1:n
      allZetas[i] = Vector{Particle}(undef, N)
      for j=1:N
        allZetas[i][j] = Particle()
      end
    end
    allWs = Vector{Vector{Float64}}(undef, n)
    for i=1:n
      allWs[i] = Vector{Float64}(undef, N)
    end
    allAs = Vector{Vector{Int64}}(undef, n-1)
    for i=1:n-1
      allAs[i] = Vector{Int64}(undef, N)
    end
    allEves = Vector{Vector{Int64}}(undef, n)
    for i=1:n
      allEves[i] = Vector{Int64}(undef, N)
    end
  else
    allZetas = Vector{Vector{Particle}}(undef, 0)
    allWs = Vector{Vector{Float64}}(undef, 0)
    allAs = Vector{Vector{Int64}}(undef, 0)
    allEves = Vector{Vector{Int64}}(undef, 0)
  end

  internal::_SMCInternal = _SMCInternal{Particle, ParticleScratch}(N, n,
    nthreads, fullOutput)

  _assignThreadViews(internal, zetas, eves, ws, allZetas, allWs, allAs, allEves)

  return SMCIO(N, n, nthreads, zetas, eves, ws, logZhats, Vhat1s, esses,
    resample, fullOutput, essThreshold, internal, allZetas, allWs, allAs,
    allEves)
end


#=
Julia sample software pipeline usinn coroutine.
https://docs.julialang.org/en/v1.1/manual/parallel-computing/#Parallel-Computing-1
=#

const loader = Channel{Array{Float64,2}}(5)
const producer = Channel{Array{Float64,2}}(5)

function make_jobs(n)
   while true
       sleep(rand())
       put!(loader, rand(3, 3))
   end
end

function comm()
    println("start")
    while true
        data = take!(loader)
        sleep(rand())
        put!(producer, data)
    end
    println("end")
end

function use_pipeline(n)
    @async make_jobs(n)
    @async comm()
    println("process start")
    for i in 1:n
        @show take!(producer)
    end
end

function sequential(n)
    println("process start")
    for i in 1:n
        #do load something e.g. FileIO
        sleep(rand())
        # do some calculation
        sleep(rand())
        data = rand(3, 3)
        @show data
    end
end

function main()
    n=10
    @time sequential(n)
    @time use_pipeline(n)
end


# device intrinsics for querying the compute SimpleVersion and PTX ISA version

## a GPU-compatible version number

export SimpleVersion, @sv_str

struct SimpleVersion
    major::UInt32
    minor::UInt32

    SimpleVersion(major, minor=0) = new(major, minor)
end

function Base.tryparse(::Type{SimpleVersion}, v::AbstractString)
    parts = split(v, ".")
    1 <= length(parts) <= 2 || return nothing

    int_parts = map(parts) do part
        tryparse(Int, part)
    end
    any(isnothing, int_parts) && return nothing

    SimpleVersion(int_parts...)
end

function Base.parse(::Type{SimpleVersion}, v::AbstractString)
    ver = tryparse(SimpleVersion, v)
    ver === nothing && throw(ArgumentError("invalid SimpleVersion string: '$v'"))
    return ver
end

SimpleVersion(v::AbstractString) = parse(SimpleVersion, v)

@inline function Base.isless(a::SimpleVersion, b::SimpleVersion)
    (a.major < b.major) && return true
    (a.major > b.major) && return false
    (a.minor < b.minor) && return true
    (a.minor > b.minor) && return false
    return false
end

macro sv_str(str)
    SimpleVersion(str)
end

## accessors for the compute SimpleVersion and PTX ISA version

export compute_capability, ptx_isa_version

for var in ["sm_major", "sm_minor", "ptx_major", "ptx_minor"]
    @eval @inline $(Symbol(var))() =
        Base.llvmcall(
            $("""@$var = external global i32
                 define i32 @entry() #0 {
                     %val = load i32, i32* @$var
                     ret i32 %val
                 }
                 attributes #0 = { alwaysinline }
            """, "entry"), UInt32, Tuple{})
end

@device_function @inline compute_capability() = SimpleVersion(sm_major(), sm_minor())
@device_function @inline ptx_isa_version() = SimpleVersion(ptx_major(), ptx_minor())


module PiADS1x15

using PiGPIO

export ADS1015, ADS1115
export set_config, read_conv_result, set_threshold_alert, enable_conv_rdy_alert, is_busy

# I2C Addresses
const addresses = Dict(
  :GND => 0x48, # I2C address if ADDR pin is connected to GND
  :VDD => 0x49, # I2C address if ADDR pin is connected to VCC
  :SDA => 0x4a, # I2C address if ADDR pin is connected to SDA
  :SCL => 0x4b # I2C address if ADDR pin is connected to SCL
)

# Pointer register
const pointers = Dict(
  :CONVERT => 0x00, # Convert register
  :CONFIG => 0x01, # Config register
  :LOWTHRESH => 0x02, # Low threshold register
  :HITHRESH => 0x03 # High threshold register
)

# Config register
const config_os = Dict(
  :MASK_OS => 0x8000,
  :STARTSINGLE => 0x8000, # Write: Start a single-conversion
  :BUSY => 0x0000, # Read: Device is currently performing a conversion
  :NOTBUSY => 0x8000 # Read: Device is in idle mode
)

const config_mux = Dict(
  :MASK_MUX => 0x7000,
  :DIFF_0_1 => 0x0000, # AIN_P = AIN0 and AIN_N = AIN1 (differential, default)
  :DIFF_0_3 => 0x1000, # AIN_P = AIN0 and AIN_N = AIN3 (differential)
  :DIFF_1_3 => 0x2000, # AIN_P = AIN1 and AIN_N = AIN3 (differential)
  :DIFF_2_3 => 0x3000, # AIN_P = AIN2 and AIN_N = AIN3 (differential)
  :SINGLE_0 => 0x4000, # AIN_P = AIN0 and AIN_N = GND (single-ended)
  :SINGLE_1 => 0x5000, # AIN_P = AIN1 and AIN_N = GND (single-ended)
  :SINGLE_2 => 0x6000, # AIN_P = AIN2 and AIN_N = GND (single-ended)
  :SINGLE_3 => 0x7000 # AIN_P = AIN3 and AIN_N = GND (single-ended)
)

const config_pga = Dict(
  :MASK_PGA => 0x0E00,
  :V6_144 => 0x0000, # +-6.144V range
  :V4_096 => 0x0200, # +-4.096V range
  :V2_048 => 0x0400, # +-2.048V range (default)
  :V1_024 => 0x0600, # +-1.024V range
  :V0_512 => 0x0800, # +-0.512V range
  :V0_256 => 0x0A00 # +-0.256V range
)

const config_mode = Dict(
  :MASK_MODE => 0x0100,
  :CONTINOUS => 0x0000, # Continuous-conversion mode
  :SINGLE => 0x0100 # Single-shot mode or power-down state (default)
)

const config_dr = Dict(
  :MASK_DR => 0x00E0,
  :SPS128 => 0x0000, # 128 samples per second
  :SPS250 => 0x0020, # 250 samples per second
  :SPS490 => 0x0040, # 490 samples per second
  :SPS920 => 0x0060, # 920 samples per second
  :SPS1600 => 0x0080, # 1600 samples per second (default)
  :SPS2400 => 0x00A0, # 2400 samples per second
  :SPS3300 => 0x00C0 # 3300 samples per second
)

const config_compmode = Dict(
  :MASK_COMPMODE => 0x0010,
  :TRADITIONAL => 0x0000, # Traditional comparator (default)
  :WINDOW => 0x0010 # Window comparator
)

const config_comppol = Dict(
  :MASK_COMPPOL => 0x0008,
  :ACTIVELOW => 0x0000, # Active low ALERT/RDY pin (default)
  :ACTIVEHI => 0x0008 # Active high ALERT/RDY pin
)

const config_complat = Dict(
  :MASK_COMPLAT => 0x0004,
  :NONLATCH => 0x0000, # Non-latching comparator (default)
  :LATCH => 0x0004 # Latching comparator
)

const config_compque = Dict(
  :MASK_COMPQUE => 0x0003,
  :CONV1 => 0x0000, # Assert ALERT/RDY after one conversions
  :CONV2 => 0x0001, # Assert ALERT/RDY after two conversions
  :CONV4 => 0x0002, # Assert ALERT/RDY after four conversions
  :NONE => 0x0003 # Disable the comparator and put ALERT/RDY to high impedance (default)
)

current_config = nothing # Current config register, needs to be set after startup

abstract type ADS1x15 end

mutable struct ADS1015 <: ADS1x15
  handle          # PiGPIO handle
  i2c_address      # I2C address
  conversion_delay # Conversion delay (s)
  bit_shift        # Bit shift
end

mutable struct ADS1115 <: ADS1x15
  handle          # PiGPIO handle
  i2c_address      # I2C address
  conversion_delay # Conversion delay (s)
  bit_shift        # Bit shift
end

function ADS1015(pi::Pi, i2c_bus::Integer, i2c_address::Symbol=:GND)
  global addresses
  handle = PiGPIO.i2c_open(pi, i2c_bus, addresses[i2c_address])
  return ADS1015(handle, i2c_address, 1//1000, 4)
end

function ADS1115(pi::Pi, i2c_bus::Integer, i2c_address::Symbol=:GND)
  global addresses
  handle = PiGPIO.i2c_open(pi, i2c_bus, addresses[i2c_address])
  return ADS1115(handle, i2c_address, 9//1000, 0)
end

"""
    write_register(pi, ads, register, value)

Write a 16-bits `value` to the specified destination `register`.
"""
function write_register(pi::Pi, ads::ADS1x15, register::Number, value::Number)
  return PiGPIO.i2c_write_word_data(pi, ads.handle, UInt8(register), UInt16(value))
end

"""
  read_register(pi, ads, register)

Read 16-bits from the specified destination `register`.
"""
function read_register(pi::Pi, ads::ADS1x15, register::Number)
  return PiGPIO.i2c_read_word_data(pi, ads.handle, UInt8(register))
end

"""
  set_config(pi, ads [, os; mux, pga, mode, dr, comp_mode, comp_pol, comp_lat, comp_que])

Set the config register of a ADS1x15 device `ads` that is connected to a Raspberry Pi via a PiGPIO.jl object `pi`.
The conversion results can be collected afterwards using [`read_conv_result`](@ref).
"""
function set_config(pi::Pi, ads::ADS1x15;
  os::Symbol=:STARTSINGLE,
  mux::Symbol=:DIFF_0_1,
  pga::Symbol=:V2_048,
  mode::Symbol=:SINGLE,
  dr::Symbol=:SPS1600,
  compmode::Symbol=:TRADITIONAL,
  comppol::Symbol=:ACTIVELOW,
  complat::Symbol=:NONLATCH,
  compque::Symbol=:NONE)

  global pointers
  global config_os
  global config_mux
  global config_pga
  global config_mode
  global config_dr
  global config_compmode
  global config_comppol
  global config_complat
  global config_compque
  global current_config

  config = config_os[os] | config_mux[mux] | config_pga[pga] | config_mode[mode] | config_dr[dr] |
  config_compmode[compmode] | config_comppol[comppol] | config_complat[complat] | config_compque[compque]

  if write_register(pi, ads, pointers[:CONFIG], config) == 0
    current_config = config # Save current config register
    return true
  else
    return false
  end
end

"""
  read_conv_result(pi, ads)

Get a conversion result from an ADS1x15 device connected to a Raspberry Pi via a PiGPIO.jl object `pi`.
The config register needs to be set sepeartely beforehand using [`set_config`](@ref).
"""
function read_conv_result(pi::Pi, ads::ADS1x15)
  global pointers
  global config_mux
  global current_config

  if current_config == nothing
    error("Config register needs to be set before reading values.")
  end

  sleep(ads.conversion_delay)
  result = read_register(pi, ads, pointers[:CONVERT]) >> ads.bit_shift
  if ads.bit_shift != 0 && (current_config & config_mux[:MASK_MUX]) <= config_mux[:DIFF_2_3] && result > 0x07FF
    result |= 0xF000 # Move sign to 16th bit in case of negative differential values
  end
  return result
end

"""
  read_config(pi, ads)

Read the current config register from an ADS1x15 device connected to a Raspberry Pi via a PiGPIO.jl object `pi`.
"""
function read_config(pi::Pi, ads::ADS1x15)
  global pointers
  global config_os
  global config_mux
  global config_pga
  global config_mode
  global config_dr
  global config_compmode
  global config_comppol
  global config_complat
  global config_compque
  global current_config

  res = read_register(pi, ads, pointers[:CONFIG])
  if res != 0
    current_config = res
    if field == :NONE
      return current_config
    elseif field == :OS
      return current_config & config_os[:MASK_OS]
    elseif field == :MUX
      return current_config & config_mux[:MASK_MUX]
    elseif field == :PGA
      return current_config & config_pga[:MASK_PGA]
    elseif field == :MODE
      return current_config & config_mode[:MASK_MODE]
    elseif field == :DR
      return current_config & config_dr[:MASK_DR]
    elseif field == :COMPMODE
      return current_config & config_compmode[:MASK_COMPMODE]
    elseif field == :COMPPOL
      return current_config & config_comppol[:MASK_COMPPOL]
    elseif field == :COMPLAT
      return current_config & config_complat[:MASK_COMPLAT]
    elseif field == :COMPQUE
      return current_config & config_compque[:MASK_COMPQUE]
    else
      error("Unknown field $field.")
    end
  end
end

"""
  set_threshold_ADS1x15(pi, ads, thld_value[, write_low_thld])

Set the low or high (see boolean `write_low_thld`) comperator alert threshold to a specific digital value `thld_value`.
Use function [`enable_conv_rdy_alert`](@ref) in order to enable a alert pin changes on newly available conversion results.
"""
function set_threshold_alert(pi::Pi, ads::ADS1x15, thld_value::Number, write_low_thld::Bool=true)
  global pointers
  register = write_low_thld ? pointers[:LOWTHRESH] : pointers[:HITHRESH]
  return write_register(pi, ads, register, thld_value << ads.bit_shift)
end

"""
  enable_conv_rdy_alert(pi, ads)

Enable alert pin changes on newly available conversion results.
"""
function enable_conv_rdy_alert(pi::Pi, ads::ADS1x15)
  global pointers
  write_register(pi, ads, pointers[:LOWTHRESH], 0x800)
  write_register(pi, ads, pointers[:HITHRESH], 0x7ff)
end

"""
  is_busy(pi, ads)

Check if ADC is currenty performing a measurement.
"""
function is_busy(pi::Pi, ads::ADS1x15)
  global pointers
  global config_os
  return read_register(pi, ads, pointers[:CONFIG]) & config_os[:MASK_OS] == config_os[:BUSY]
end

end #module


import JSON
using DataFrames
import CSV
using Glob
using Statistics

function pmap(f, c)
    out = Vector{typeof(f(c[1]))}(undef, length(c))
    Threads.@threads for i in 1:length(c)
        out[i] = f(c[i])
    end
    out
end

function main()
    df = pmap(glob("$(ARGS[1])/*.json")) do f
        map(readlines(f)) do l
            o = JSON.parse(l)
            config = o["i"]
            input = JSON.parse(config[1][1])
            task = input[1]
            target = config[1][2]
            cost = mean(o["r"][1])
            schedule = JSON.json(config[2][2])
            loop_shape = mapreduce(x -> x[1], *, config[2][2])
            DataFrame(input=JSON.json(input),task=task,target=target,cost=cost,schedule=schedule,loop_shape=loop_shape)
        end |> x -> vcat(x...)
    end |> x -> vcat(x...)

    CSV.write(ARGS[2], df)
end

main()


struct Parameters
	N                        ::Int
	d                        ::Int
	m                        ::Float64
	L                        ::Float64
	T                        ::Float64
	τ                        ::Float64
	ε_LJ                     ::Float64
	σ_LJ                     ::Float64
	k_B                      ::Float64
	seed                     ::Int
	n_iter                   ::Int
	ensemble_type            ::Symbol
	rescale_velocity_interval::Int
	rng                      ::AbstractRNG
	diag_indices             ::Vector{CartesianIndex{3}}

	function Parameters(params_dict)
		new(
			params_dict[:N],
			params_dict[:d],
			params_dict[:m],
			params_dict[:L],
			params_dict[:T],
			params_dict[:τ],
			params_dict[:ε_LJ],
			params_dict[:σ_LJ],
			params_dict[:k_B],
			params_dict[:seed],
			params_dict[:n_iter],
			params_dict[:ensemble_type],
			params_dict[:rescale_velocity_interval],
			MersenneTwister(params_dict[:seed]),
			[CartesianIndex((i, 1, i)) for i in 1:params_dict[:N]]
		)
	end
end

mutable struct State
	x                        ::Matrix{Float64}
	v                        ::Matrix{Float64}
	a                        ::Matrix{Float64}
	a_prev                   ::Matrix{Float64}
	r_nearest_image          ::Array{Float64, 3}
	r_nearest_image_magnitude::Array{Float64, 3}
	E_kin                    ::Float64
	E_pot                    ::Float64

	function State(P::Parameters, params_dict::Dict)
		x = get(params_dict, :x, initialize_random_positions(P))
		v = get(params_dict, :v, initialize_random_velocities(P))
		a                         = zeros((P.N, P.d))
		a_prev                    = zeros((P.N, P.d))
		r_nearest_image           = zeros((P.N, P.d, P.N))
		r_nearest_image_magnitude = zeros((P.N, 1, P.N))
		E_kin                     = 0
		E_pot                     = 0

		new(
			x,
			v,
			a,
			a_prev,
			r_nearest_image,
			r_nearest_image_magnitude,
			E_kin,
			E_pot,
		)
	end
end

mutable struct History
	x    ::Array{Float64, 3}
	v    ::Array{Float64, 3}
	E_kin::Vector{Float64}
	E_pot::Vector{Float64}

	function History(P::Parameters)
		new(
			zeros((P.N, P.d, P.n_iter + 1)),
			zeros((P.N, P.d, P.n_iter + 1)),
			zeros(P.n_iter + 1),
			zeros(P.n_iter + 1),
		)
	end
end

mutable struct MDSimulation
	params  ::Parameters
	state   ::State
	history ::History

	function MDSimulation(params_dict)
		params  = Parameters(params_dict)
		state   = State(params, params_dict)
		history = History(params)

		mdsim = new(params, state, history)
		compute_accelerations!(mdsim)
		compute_energies!(mdsim)
		update_history!(mdsim, 1)

		return mdsim
	end
end


struct OmniscapeFlags
    calc_flow_potential::Bool
    calc_normalized_current::Bool
    compute_flow_potential::Bool #bad name, but need a variable that stores the "or" of the previous two, internal use only
    write_raw_currmap::Bool
    parallelize::Bool
    correct_artifacts::Bool
    source_from_resistance::Bool
    conditional::Bool
    mask_nodata ::Bool
    resistance_is_conductance::Bool
    write_as_tif::Bool
    allow_different_projections::Bool
    reclassify::Bool
    write_reclassified_resistance::Bool
end

struct Target
    x_coord::Int64
    y_coord::Int64
    amps::Float64
end

struct Conditions
    comparison1::String
    comparison2::String
    condition1_lower::Number
    condition1_upper::Number
    condition2_lower::Number
    condition2_upper::Number
end

struct ConditionLayers
    condition1_present::Array{Union{Missing, Number}, 2} # where T <: Number
    condition1_future::Array{Union{Missing, Number}, 2} # where U <: Number
    condition2_present::Array{Union{Missing, Number}, 2} # where V <: Number
    condition2_future::Array{Union{Missing, Number}, 2} # where W <: Number
end


import MathOptInterface
import MathOptInterface: Utilities

const MOI = MathOptInterface
const MOIU = MOI.Utilities

# indices
const VI = MOI.VariableIndex
const CI = MOI.ConstraintIndex

# function aliases
const SV = MOI.SingleVariable
const SAF = MOI.ScalarAffineFunction{Float64}
const SQF = MOI.ScalarQuadraticFunction{Float64}

# set aliases
const Bounds{T} = Union{
    MOI.EqualTo{T},
    MOI.GreaterThan{T},
    MOI.LessThan{T},
    MOI.Interval{T}
}

mutable struct Optimizer <: MOI.AbstractOptimizer
    inner::BaronModel
    nlp_block_data::Union{Nothing, MOI.NLPBlockData}
end

Optimizer(; options...) = Optimizer(BaronModel(; options...), nothing)

MOI.get(model::Optimizer, ::MOI.SolverName) = "BARON"

# empty
function MOI.is_empty(model::Optimizer)
    BARON.is_empty(model.inner) && model.nlp_block_data === nothing
end

function MOI.empty!(model::Optimizer)
    model.inner = BaronModel(; model.inner.options...)
    model.nlp_block_data = nothing
end

# copy
MOIU.supports_default_copy_to(model::Optimizer, copy_names::Bool) = !copy_names
function MOI.copy_to(dest::Optimizer, src::MOI.ModelLike; kws...)
    return MOIU.automatic_copy_to(dest, src; kws...)
end

# optimize
function MOI.optimize!(model::Optimizer)
    write_bar_file(model.inner)
    run(`$baron_exec $(model.inner.problem_file_name)`)
    read_results(model.inner)
end

# TimeLimitSec
MOI.supports(::Optimizer, ::MOI.TimeLimitSec) = true
function MOI.set(model::Optimizer, ::MOI.TimeLimitSec, val::Real)
    model.inner.options[:MaxTime] = Float64(val)
    return
end

# BARON's default time limit is 1000 seconds.
function MOI.get(model::Optimizer, ::MOI.TimeLimitSec)
    return get(model.inner.options, :MaxTime, 1000.0)
end

include(joinpath("moi", "util.jl"))
include(joinpath("moi", "variables.jl"))
include(joinpath("moi", "constraints.jl"))
include(joinpath("moi", "objective.jl"))
include(joinpath("moi", "nlp.jl"))
include(joinpath("moi", "results.jl"))


const AVec = AbstractVector
const Rydberg2eV = 13.605698066

struct ElecHam
   basic_data
   ws_rvecs
   rvec_idx
   hopping
end

#struct LatticeIFC
#   basic_data
#   ws_rvecs
#   rvec_idx
#   forceconst
#end


# See interface docs
@inline inbounds(data::Union{GridData,AbstractSimData}, I::Tuple) = inbounds(data, I...)
@inline inbounds(data::Union{GridData,AbstractSimData}, I...) =
    _inbounds(boundary(data), gridsize(data), I...)

@inline function _inbounds(boundary::BoundaryCondition, size::Tuple, i1, i2)
    a, inbounds_a = _inbounds(boundary, size[1], i1)
    b, inbounds_b = _inbounds(boundary, size[2], i2)
    (a, b), inbounds_a & inbounds_b
end
@inline _inbounds(::Remove, size::Number, i::Number) = i, _isinbounds(size, i)
@inline function _inbounds(::Wrap, size::Number, i::Number)
    if i < oneunit(i)
        size + rem(i, size), true
    elseif i > size
        rem(i, size), true
    else
        i, true
    end
end

# See interface docs
@inline isinbounds(data::Union{GridData,AbstractSimData}, I::Tuple) = isinbounds(data, I...)
@inline isinbounds(data::Union{GridData,AbstractSimData}, I...) = _isinbounds(gridsize(data), I...)

@inline _isinbounds(size::Tuple, I...) = all(map(_isinbounds, size, I))
@inline _isinbounds(size, i) = i >= one(i) && i <= size

# _updateboundary
# Reset or wrap boundary where required. This allows us to ignore
# bounds checks on neighborhoods and still use a wraparound grid.
_updateboundary!(grids::Tuple) = map(_updateboundary!, grids)
function _updateboundary!(g::GridData{S,R}) where {S<:Tuple{Y,X},R} where {Y,X}
    R < 1 && return g
    return _updateboundary!(g, boundary(g))
end
function _updateboundary!(g::GridData{S,R}, ::Remove) where {S<:Tuple{Y,X},R} where {Y,X}
    src = parent(source(g))
    # Left
    @inbounds src[1:Y, 1:R] .= Ref(padval(g))
    # Right
    @inbounds src[1:Y, X+R+1:X+2R] .= Ref(padval(g))
    # Top middle
    @inbounds src[1:R, R+1:X+R] .= Ref(padval(g))
    # Bottom middle
    @inbounds src[Y+R+1:Y+2R, R+1:X+R] .= Ref(padval(g))
    status = sourcestatus(g)
    return g
end
function _updateboundary!(g::GridData{S,R}, ::Wrap) where {S<:Tuple{Y,X},R} where {Y,X}
    src = parent(source(g))
    nrows, ncols = gridsize(g)
    startpadrow = startpadcol = 1:R
    endpadrow = nrows+R+1:nrows+2R
    endpadcol = ncols+R+1:ncols+2R
    startrow = startcol = 1+R:2R
    endrow = nrows+1:nrows+R
    endcol = ncols+1:ncols+R
    rows = 1+R:nrows+R
    cols = 1+R:ncols+R

    # Sides ---
    # Left
    @inbounds copyto!(src, CartesianIndices((rows, startpadcol)),
                      src, CartesianIndices((rows, endcol)))
    # Right
    @inbounds copyto!(src, CartesianIndices((rows, endpadcol)),
                      src, CartesianIndices((rows, startcol)))
    # Top
    @inbounds copyto!(src, CartesianIndices((startpadrow, cols)),
                      src, CartesianIndices((endrow, cols)))
    # Bottom
    @inbounds copyto!(src, CartesianIndices((endpadrow, cols)),
                      src, CartesianIndices((startrow, cols)))

    # Corners ---
    # Top Left
    @inbounds copyto!(src, CartesianIndices((startpadrow, startpadcol)),
                      src, CartesianIndices((endrow, endcol)))
    # Top Right
    @inbounds copyto!(src, CartesianIndices((startpadrow, endpadcol)),
                      src, CartesianIndices((endrow, startcol)))
    # Botom Left
    @inbounds copyto!(src, CartesianIndices((endpadrow, startpadcol)),
                      src, CartesianIndices((startrow, endcol)))
    # Botom Right
    @inbounds copyto!(src, CartesianIndices((endpadrow, endpadcol)),
                      src, CartesianIndices((startrow, startcol)))

    _wrapstatus!(sourcestatus(g))
    return g
end

# _wrapstatus!
# Copies status from opposite sides/corners in Wrap boundary mode
_wrapstatus!(status::Nothing) = nothing
function _wrapstatus!(status::AbstractArray)
    # !!! The end row/column is always empty !!!
    # Its padding for block opt. So we work with end-1 and end-2

    # This could be further optimised by not copying the end-2
    # block column/row when the blocks are aligned at both ends.

    # Sides
    status[1, :] .|= status[end-1, :] .| status[end-2, :]
    status[:, 1] .|= status[:, end-1] .| status[:, end-2]
    status[end-1, :] .|= status[1, :]
    status[:, end-1] .|= status[:, 1]
    status[end-2, :] .|= status[1, :]
    status[:, end-2] .|= status[:, 1]

    # Corners
    status[1, 1] |= status[end-1, end-1] | status[end-2, end-1] |
                    status[end-1, end-2] | status[end-2, end-2]
    status[end-1, 1] |= status[1, end-1] | status[1, end-2]
    status[end-2, 1] |= status[1, end-1] | status[1, end-2]
    status[1, end-1] |= status[end-1, 1] | status[end-2, 1]
    status[1, end-2] |= status[end-1, 1] | status[end-2, 1]
    status[end-1, end-1] |= status[1, 1]
    status[end-2, end-2] |= status[1, 1]
end


function Base.get{S<:VehicleState,D,I,R}(::Feature_PosFt, rec::EntityQueueRecord{S,D,I}, roadway::R, vehicle_index::Int, pastframe::Int=0)
    FeatureValue(rec[pastframe][vehicle_index].state.posF.t)
end
function Base.get{S<:VehicleState,D,I,R}(::Feature_PosFyaw, rec::EntityQueueRecord{S,D,I}, roadway::R, vehicle_index::Int, pastframe::Int=0)
    FeatureValue(rec[pastframe][vehicle_index].state.posF.ϕ)
end
function Base.get{S<:VehicleState,D,I,R}(::Feature_Speed, rec::EntityQueueRecord{S,D,I}, roadway::R, vehicle_index::Int, pastframe::Int=0)
    FeatureValue(rec[pastframe][vehicle_index].state.v)
end
function Base.get{S<:VehicleState,D,I,R}(::Feature_VelFs, rec::EntityQueueRecord{S,D,I}, roadway::R, vehicle_index::Int, pastframe::Int=0)
    veh = rec[pastframe][vehicle_index]
    FeatureValue(veh.state.v*cos(veh.state.posF.ϕ))
end
function Base.get{S<:VehicleState,D,I,R}(::Feature_VelFt, rec::EntityQueueRecord{S,D,I}, roadway::R, vehicle_index::Int, pastframe::Int=0)
    veh = rec[pastframe][vehicle_index]
    FeatureValue(veh.state.v*sin(veh.state.posF.ϕ))
end

generate_feature_functions("TurnRateG", :turnrateG, Float64, "rad/s")
function Base.get{S<:VehicleState,D,I,R}(::Feature_TurnRateG, rec::EntityQueueRecord{S,D,I}, roadway::R, vehicle_index::Int, pastframe::Int=0; frames_back::Int=1)

    id = rec[pastframe][vehicle_index].id

    retval = FeatureValue(0.0, FeatureState.INSUF_HIST)
    pastframe2 = pastframe - frames_back
    if pastframe_inbounds(rec, pastframe2)

        veh_index_curr = vehicle_index
        veh_index_prev = findfirst(rec[pastframe2], id)

        if veh_index_prev != 0
            curr = rec[pastframe][veh_index_curr].state.posG.θ
            past = rec[pastframe2][veh_index_prev].state.posG.θ
            Δt = get_elapsed_time(rec, pastframe2, pastframe)
            retval = FeatureValue(deltaangle(past, curr) / Δt)
        end
    end

    retval
end
generate_feature_functions("TurnRateF", :turnrateF, Float64, "rad/s")
function Base.get{S<:VehicleState,D,I,R}(::Feature_TurnRateF, rec::EntityQueueRecord{S,D,I}, roadway::R, vehicle_index::Int, pastframe::Int=0)
    _get_feature_derivative_backwards(POSFYAW, rec, roadway, vehicle_index, pastframe)
end
generate_feature_functions("AngularRateG", :angrateG, Float64, "rad/s²")
function Base.get{S<:VehicleState,D,I,R}(::Feature_AngularRateG, rec::EntityQueueRecord{S,D,I}, roadway::R, vehicle_index::Int, pastframe::Int=0)
    _get_feature_derivative_backwards(TURNRATEG, rec, roadway, vehicle_index, pastframe)
end
generate_feature_functions("AngularRateF", :angrateF, Float64, "rad/s²")
function Base.get{S<:VehicleState,D,I,R}(::Feature_AngularRateF, rec::EntityQueueRecord{S,D,I}, roadway::R, vehicle_index::Int, pastframe::Int=0)
    _get_feature_derivative_backwards(TURNRATEF, rec, roadway, vehicle_index, pastframe)
end
generate_feature_functions("DesiredAngle", :desang, Float64, "rad")
function Base.get{S<:VehicleState,D,I,R}(::Feature_DesiredAngle, rec::EntityQueueRecord{S,D,I}, roadway::R, vehicle_index::Int, pastframe::Int=0;
    kp_desired_angle::Float64 = 1.0,
    )

    retval = FeatureValue(0.0, FeatureState.INSUF_HIST)
    if pastframe_inbounds(rec, pastframe) && pastframe_inbounds(rec, pastframe-1)

        id = rec[pastframe][vehicle_index].id

        pastϕ = get_state(rec, id, pastframe-1).posF.ϕ
        currϕ = get_state(rec, id, pastframe).posF.ϕ

        Δt = rec.timestep
        expconst = exp(-kp_desired_angle*Δt)
        retval = FeatureValue((currϕ - pastϕ*expconst) / (1.0 - expconst))
    end
    retval
end

generate_feature_functions("MarkerDist_Left", :d_ml, Float64, "m")
Base.get(::Feature_MarkerDist_Left, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0) =
    FeatureValue(get_markerdist_left(rec[pastframe][vehicle_index], roadway))

generate_feature_functions("MarkerDist_Right", :d_mr, Float64, "m")
Base.get(::Feature_MarkerDist_Right, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0) =
    FeatureValue(get_markerdist_right(rec[pastframe][vehicle_index], roadway))

generate_feature_functions("MarkerDist_Left_Left", :d_mll, Float64, "m", can_be_missing=true)
function Base.get(::Feature_MarkerDist_Left_Left, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    #=
    Distance to the left lane marker one lane to the left
    =#

    veh = rec[pastframe][vehicle_index]
    lane = roadway[veh.state.posF.roadind.tag]
    if n_lanes_left(lane, roadway) > 0
        offset = veh.state.posF.t
        lane_left = roadway[LaneTag(lane.tag.segment, lane.tag.lane + 1)]
        FeatureValue(lane.width/2 - offset + lane_left.width)
    else
        FeatureValue(NaN, FeatureState.MISSING) # there is no left lane
    end
end
generate_feature_functions("MarkerDist_Right_Right", :d_mrr, Float64, "m", can_be_missing=true)
function Base.get(::Feature_MarkerDist_Right_Right, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    #=
    Distance to the right lane marker one lane to the right
    =#

    veh = rec[pastframe][vehicle_index]
    lane = roadway[veh.state.posF.roadind.tag]
    if n_lanes_right(lane, roadway) > 0
        offset = veh.state.posF.t
        lane_right = roadway[LaneTag(lane.tag.segment, lane.tag.lane - 1)]
        FeatureValue(lane.width/2 + offset + lane_right.width)
    else
        FeatureValue(NaN, FeatureState.MISSING) # there is no right lane
    end
end
generate_feature_functions("RoadEdgeDist_Left", :d_edgel, Float64, "m")
function Base.get(::Feature_RoadEdgeDist_Left, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    veh = rec[pastframe][vehicle_index]
    offset = veh.state.posF.t
    footpoint = get_footpoint(veh)
    seg = roadway[veh.state.posF.roadind.tag.segment]
    lane = seg.lanes[end]
    roadproj = proj(footpoint, lane, roadway)
    curvept = roadway[RoadIndex(roadproj)]
    lane = roadway[roadproj.tag]
    FeatureValue(lane.width/2 + abs(curvept.pos - footpoint) - offset)
end
generate_feature_functions("RoadEdgeDist_Right", :d_edger, Float64, "m")
function Base.get(::Feature_RoadEdgeDist_Right, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    veh = rec[pastframe][vehicle_index]
    offset = veh.state.posF.t
    footpoint = get_footpoint(veh)
    seg = roadway[veh.state.posF.roadind.tag.segment]
    lane = seg.lanes[1]
    roadproj = proj(footpoint, lane, roadway)
    curvept = roadway[RoadIndex(roadproj)]
    lane = roadway[roadproj.tag]
    FeatureValue(lane.width/2 + abs(curvept.pos - footpoint) + offset)
end
generate_feature_functions("LaneOffsetLeft", :posFtL, Float64, "m", can_be_missing=true)
function Base.get(::Feature_LaneOffsetLeft, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    veh_ego = rec[pastframe][vehicle_index]
    t = veh_ego.state.posF.t
    lane = roadway[veh_ego.state.posF.roadind.tag]
    if n_lanes_left(lane, roadway) > 0
        lane_left = roadway[LaneTag(lane.tag.segment, lane.tag.lane + 1)]
        lane_offset = t - lane.width/2 - lane_left.width/2
        FeatureValue(lane_offset)
    else
        FeatureValue(NaN, FeatureState.MISSING)
    end
end
generate_feature_functions("LaneOffsetRight", :posFtR, Float64, "m", can_be_missing=true)
function Base.get(::Feature_LaneOffsetRight, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    veh_ego = rec[pastframe][vehicle_index]
    t = veh_ego.state.posF.t
    lane = roadway[veh_ego.state.posF.roadind.tag]
    if n_lanes_right(lane, roadway) > 0
        lane_right = roadway[LaneTag(lane.tag.segment, lane.tag.lane - 1)]
        lane_offset = t + lane.width/2 + lane_right.width/2
        FeatureValue(lane_offset)
    else
        FeatureValue(NaN, FeatureState.MISSING)
    end
end
generate_feature_functions("N_Lane_Right", :n_lane_right, Int, "-", lowerbound=0.0)
function Base.get(::Feature_N_Lane_Right, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    nlr = rec[pastframe][vehicle_index].state.posF.roadind.tag.lane - 1
    FeatureValue(convert(Float64, nlr))
end
generate_feature_functions("N_Lane_Left", :n_lane_left, Int, "-", lowerbound=0.0)
function Base.get(::Feature_N_Lane_Left, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    veh = rec[pastframe][vehicle_index]
    seg = roadway[veh.state.posF.roadind.tag.segment]
    nll = length(seg.lanes) - veh.state.posF.roadind.tag.lane
    FeatureValue(convert(Float64, nll))
end
generate_feature_functions("Has_Lane_Right", :has_lane_right, Bool, "-", lowerbound=0.0, upperbound=1.0)
function Base.get(::Feature_Has_Lane_Right, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    val = get(N_LANE_RIGHT, rec, roadway, vehicle_index, pastframe).v > 0.0
    FeatureValue(convert(Float64, val))
end
generate_feature_functions("Has_Lane_Left", :has_lane_left, Bool, "-", lowerbound=0.0, upperbound=1.0)
function Base.get(::Feature_Has_Lane_Left, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    val = get(N_LANE_LEFT, rec, roadway, vehicle_index, pastframe).v > 0.0
    FeatureValue(convert(Float64, val))
end
generate_feature_functions("LaneCurvature", :curvature, Float64, "1/m", can_be_missing=true)
function Base.get(::Feature_LaneCurvature, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    veh = rec[pastframe][vehicle_index]
    curvept = roadway[veh.state.posF.roadind]
    val = curvept.k
    if isnan(val)
        FeatureValue(0.0, FeatureState.MISSING)
    else
        FeatureValue(val)
    end
end

# Dist_Merge
# Dist_Split

generate_feature_functions("TimeToCrossing_Right", :ttcr_mr, Float64, "s", lowerbound=0.0, censor_hi=10.0)
function Base.get(::Feature_TimeToCrossing_Right, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    d_mr = get(MARKERDIST_RIGHT, rec, roadway, vehicle_index, pastframe).v
    velFt = get(VELFT, rec, roadway, vehicle_index, pastframe).v

    if d_mr > 0.0 && velFt < 0.0
        FeatureValue(-d_mr / velFt)
    else
        FeatureValue(Inf)
    end
end
generate_feature_functions("TimeToCrossing_Left", :ttcr_ml, Float64, "s", lowerbound=0.0, censor_hi=10.0)
function Base.get(::Feature_TimeToCrossing_Left, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    d_ml = get(MARKERDIST_RIGHT, rec, roadway, vehicle_index, pastframe).v
    velFt = get(VELFT, rec, roadway, vehicle_index, pastframe).v

    if d_ml > 0.0 && velFt < 0.0
        FeatureValue(-d_ml / velFs)
    else
        FeatureValue(Inf)
    end
end
generate_feature_functions("EstimatedTimeToLaneCrossing", :est_ttcr, Float64, "s", lowerbound=0.0, censor_hi=10.0)
function Base.get(::Feature_EstimatedTimeToLaneCrossing, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    ttcr_left = get(TIMETOCROSSING_LEFT, rec, roadway, vehicle_index, pastframe).v
    ttcr_right = get(TIMETOCROSSING_RIGHT, rec, roadway, vehicle_index, pastframe).v
    FeatureValue(min(ttcr_left, ttcr_right))
end
generate_feature_functions("A_REQ_StayInLane", :a_req_stayinlane, Float64, "m/s²", can_be_missing=true)
function Base.get(::Feature_A_REQ_StayInLane, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    velFt = get(VELFT, rec, roadway, vehicle_index, pastframe).v

    if velFt > 0.0
        d_mr = get(MARKERDIST_RIGHT, rec, roadway, vehicle_index, pastframe).v
        if d_mr > 0.0
            return FeatureValue(0.5velFt*velFt / d_mr)
        else
            return FeatureValue(NaN, FeatureState.MISSING)
        end
    else
        d_ml = get(MARKERDIST_LEFT, rec, roadway, vehicle_index, pastframe)
        if d_ml < 0.0
            return FeatureValue(-0.5velFt*velFt / d_ml)
        else
            return FeatureValue(NaN, FeatureState.MISSING)
        end
    end
end

generate_feature_functions("Time_Consecutive_Brake", :time_consec_brake, Float64, "s", lowerbound=0.0)
function Base.get(::Feature_Time_Consecutive_Brake, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)

    prev_accel = convert(Float64, get(ACC, rec, roadway, vehicle_index, pastframe))
    if prev_accel ≥ 0.0
        FeatureValue(0.0)
    else
        pastframe_orig = pastframe
        id = rec[pastframe][vehicle_index].id
        while pastframe_inbounds(rec, pastframe-1) &&
              get(ACC, rec, roadway, findfirst(rec[pastframe-1], id)) < 0.0

            pastframe -= 1
        end

        FeatureValue(get_elapsed_time(rec, pastframe, pastframe_orig))
    end
end
generate_feature_functions("Time_Consecutive_Accel", :time_consec_accel, Float64, "s", lowerbound=0.0)
function Base.get(::Feature_Time_Consecutive_Accel, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)

    prev_accel = convert(Float64, get(ACC, rec, roadway, vehicle_index, pastframe))
    if prev_accel ≤ 0.0
        FeatureValue(0.0)
    else

        pastframe_orig = pastframe
        id = rec[pastframe][vehicle_index].id
        while pastframe_inbounds(rec, pastframe-1) &&
              get(ACC, rec, roadway, findfirst(rec[pastframe-1], id)) > 0.0

            pastframe -= 1
        end

        FeatureValue(get_elapsed_time(rec, pastframe, pastframe_orig))
    end
end
generate_feature_functions("Time_Consecutive_Throttle", :time_consec_throttle, Float64, "s", lowerbound=0.0)
function Base.get(::Feature_Time_Consecutive_Throttle, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    tc_accel = get(TIME_CONSECUTIVE_ACCEL, rec, roadway, vehicle_index, pastframe).v
    tc_brake = get(TIME_CONSECUTIVE_BRAKE, rec, roadway, vehicle_index, pastframe).v
    FeatureValue(tc_accel ≥ tc_brake ? tc_accel : -tc_brake)
end

#############################################
#
# FRONT
#
#############################################

generate_feature_functions("Dist_Front", :d_front, Float64, "m", lowerbound=0.0, can_be_missing=true)
function Base.get(::Feature_Dist_Front, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0;
    neighborfore::NeighborLongitudinalResult = get_neighbor_fore_along_lane(rec[pastframe], vehicle_index, roadway),
    censor_hi::Float64=100.0,
    )

    if neighborfore.ind == 0
        FeatureValue(100.0, FeatureState.CENSORED_HI)
    else
        scene = rec[pastframe]
        len_ego = scene[vehicle_index].def.length
        len_oth = scene[neighborfore.ind].def.length
        FeatureValue(neighborfore.Δs - len_ego/2 - len_oth/2)
    end
end
generate_feature_functions("Speed_Front", :v_front, Float64, "m/s", can_be_missing=true)
function Base.get(::Feature_Speed_Front, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0;
    neighborfore::NeighborLongitudinalResult = get_neighbor_fore_along_lane(rec[pastframe], vehicle_index, roadway),
    )

    if neighborfore.ind == 0
        FeatureValue(0.0, FeatureState.MISSING)
    else
        FeatureValue(rec[pastframe][neighborfore.ind].state.v)
    end
end
generate_feature_functions("Timegap", :timegap, Float64, "s", can_be_missing=true)
function Base.get(::Feature_Timegap, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0;
    neighborfore::NeighborLongitudinalResult = get_neighbor_fore_along_lane(rec[pastframe], vehicle_index, roadway),
    censor_hi::Float64 = 10.0,
    )

    v = rec[pastframe][vehicle_index].state.v

    if v ≤ 0.0 || neighborfore.ind == 0
        FeatureValue(censor_hi, FeatureState.CENSORED_HI)
    else
        scene = rec[pastframe]
        len_ego = scene[vehicle_index].def.length
        len_oth = scene[neighborfore.ind].def.length
        Δs = neighborfore.Δs - len_ego/2 - len_oth/2

        if Δs > 0.0
            FeatureValue(Δs / v)
        else
            FeatureValue(0.0) # collision!
        end
    end
end

generate_feature_functions("Inv_TTC", :inv_ttc, Float64, "1/s", can_be_missing=true)
function Base.get(::Feature_Inv_TTC, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0;
    neighborfore::NeighborLongitudinalResult = get_neighbor_fore_along_lane(rec[pastframe], vehicle_index, roadway),
    censor_hi::Float64 = 10.0,
    )

    if neighborfore.ind == 0
        FeatureValue(0.0, FeatureState.MISSING)
    else
        scene = rec[pastframe]
        veh_fore = scene[neighborfore.ind]
        veh_rear = scene[vehicle_index]

        len_ego = veh_fore.def.length
        len_oth = veh_rear.def.length
        Δs = neighborfore.Δs - len_ego/2 - len_oth/2

        Δv = veh_fore.state.v - veh_rear.state.v

        if Δs < 0.0 # collision!
            FeatureValue(censor_hi, FeatureState.CENSORED_HI)
        elseif Δv > 0.0 # front car is pulling away
            FeatureValue(0.0)
        else
            f = -Δv/Δs
            if f > censor_hi
                FeatureValue(f, FeatureState.CENSORED_HI)
            else
                FeatureValue(f)
            end
        end
    end
end
generate_feature_functions("TTC", :ttc, Float64, "s", can_be_missing=true)
function Base.get(::Feature_TTC, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0;
    neighborfore::NeighborLongitudinalResult = get_neighbor_fore_along_lane(rec[pastframe], vehicle_index, roadway),
    censor_hi::Float64 = 10.0,
    inv_ttc::FeatureValue = get(INV_TTC, rec, roadway, vehicle_index, pastframe, neighborfore=neighborfore, censor_hi=censor_hi),
    )

    if inv_ttc.i == FeatureState.MISSING
        # if the value is missing then front car not found and set TTC to censor_hi
        return FeatureValue(censor_hi, FeatureState.MISSING)
    else
        @assert is_feature_valid(inv_ttc) || inv_ttc.i == FeatureState.CENSORED_HI
        return FeatureValue(min(1.0 / inv_ttc.v, censor_hi))
    end
end

#############################################
#
# FRONT LEFT
#
#############################################
generate_feature_functions("Dist_Front_Left", :d_front_left, Float64, "m", lowerbound=0.0, can_be_missing=true)
function Base.get(::Feature_Dist_Front_Left, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0;
    neighborfore::NeighborLongitudinalResult = get_neighbor_fore_along_left_lane(rec[pastframe], vehicle_index, roadway),
    )

    get(DIST_FRONT, rec, roadway, vehicle_index, pastframe, neighborfore=neighborfore)
end

#############################################
#
# FRONT RIGHT
#
#############################################

generate_feature_functions("Dist_Front_Right", :d_front_right, Float64, "m", lowerbound=0.0, can_be_missing=true)
function Base.get(::Feature_Dist_Front_Right, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0;
    neighborfore::NeighborLongitudinalResult = get_neighbor_fore_along_right_lane(rec[pastframe], vehicle_index, roadway),
    )

    get(DIST_FRONT, rec, roadway, vehicle_index, pastframe, neighborfore=neighborfore)
end

#############################################
#
# REAR
#
#############################################

generate_feature_functions("Dist_Rear", :d_rear, Float64, "m", lowerbound=0.0, can_be_missing=true)
function Base.get(::Feature_Dist_Rear, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0;
    neighborrear::NeighborLongitudinalResult = get_neighbor_rear_along_lane(rec[pastframe], vehicle_index, roadway),
    censor_hi::Float64=100.0,
    )

    if neighborrear.ind == 0
        FeatureValue(100.0, FeatureState.CENSORED_HI)
    else
        scene = rec[pastframe]
        len_ego = scene[vehicle_index].def.length
        len_oth = scene[neighborrear.ind].def.length
        FeatureValue(neighborrear.Δs - len_ego/2 - len_oth/2)
    end
end
generate_feature_functions("Speed_Rear", :v_rear, Float64, "m/s", can_be_missing=true)
function Base.get(::Feature_Speed_Rear, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0;
    neighborrear::NeighborLongitudinalResult = get_neighbor_rear_along_lane(rec[pastframe], vehicle_index, roadway),
    )

    if neighborrear.ind == 0
        FeatureValue(0.0, FeatureState.MISSING)
    else
        FeatureValue(rec[pastframe][neighborrear.ind].state.v)
    end
end

#############################################
#
# SCENE WISE
#
#############################################

generate_feature_functions("Is_Colliding", :is_colliding, Bool, "-", lowerbound=0.0, upperbound=1.0)
function Base.get{S,D,I,R}(::Feature_Is_Colliding, rec::EntityQueueRecord{S,D,I}, roadway::R, vehicle_index::Int, pastframe::Int=0;
    mem::CPAMemory=CPAMemory(),
    )

    scene = rec[pastframe]
    is_colliding = convert(Float64, get_first_collision(scene, vehicle_index, mem).is_colliding)
    FeatureValue(is_colliding)
end


"""
    TexasHoldem

A no-limit Texas Holdem simulator.

# Terminology
 - `game` a single "game", where players are dealt hands,
   winner(s) are declared once.
 - `state` a point or process in the game, including
   `PreFlop`, `Flop`, `Turn`, `River`.
 - `round` the process of each player deciding which
   actions to take, until no further actions are taking.
"""
module TexasHoldem

using PlayingCards
using PokerHandEvaluator
using Printf

export AbstractGameState, PreFlop, Flop, Turn, River

abstract type AbstractGameState end
struct PreFlop <: AbstractGameState end
struct Flop <: AbstractGameState end
struct Turn <: AbstractGameState end
struct River <: AbstractGameState end

include("player_types.jl")
include("transactions.jl")
include("table.jl")
include("game.jl")
include("player_actions.jl")
include("player_options.jl")
include("config_game.jl")

end # module


function CachedNormalizedSphericalHarmonicsProvider(arg0::NormalizedSphericalHarmonicsProvider, arg1::jdouble, arg2::jint, arg3::jint, arg4::jdouble, arg5::jdouble)
    return CachedNormalizedSphericalHarmonicsProvider((NormalizedSphericalHarmonicsProvider, jdouble, jint, jint, jdouble, jdouble), arg0, arg1, arg2, arg3, arg4, arg5)
end

function get_ae(obj::CachedNormalizedSphericalHarmonicsProvider)
    return jcall(obj, "getAe", jdouble, ())
end

function get_max_degree(obj::CachedNormalizedSphericalHarmonicsProvider)
    return jcall(obj, "getMaxDegree", jint, ())
end

function get_max_order(obj::CachedNormalizedSphericalHarmonicsProvider)
    return jcall(obj, "getMaxOrder", jint, ())
end

function get_mu(obj::CachedNormalizedSphericalHarmonicsProvider)
    return jcall(obj, "getMu", jdouble, ())
end

function get_offset(obj::CachedNormalizedSphericalHarmonicsProvider, arg0::AbsoluteDate)
    return jcall(obj, "getOffset", jdouble, (AbsoluteDate,), arg0)
end

function get_reference_date(obj::CachedNormalizedSphericalHarmonicsProvider)
    return jcall(obj, "getReferenceDate", AbsoluteDate, ())
end

function get_tide_system(obj::CachedNormalizedSphericalHarmonicsProvider)
    return jcall(obj, "getTideSystem", TideSystem, ())
end

function on_date(obj::CachedNormalizedSphericalHarmonicsProvider, arg0::AbsoluteDate)
    return jcall(obj, "onDate", NormalizedSphericalHarmonicsProvider_NormalizedSphericalHarmonics, (AbsoluteDate,), arg0)
end


include(joinpath(@__DIR__, "..", "src/dynamics", "quadruped", "visuals.jl"))
vis = Visualizer()
open(vis)
# const ContactImplicitMPC = Main

s = get_simulation("quadruped", "flat_2D_lc", "flat")
model = s.model
env = s.env

ref_traj_ = deepcopy(ContactImplicitMPC.get_trajectory(s.model, s.env,
    joinpath(module_dir(), "src/dynamics/quadruped/gaits/gait2.jld2"),
    load_type = :split_traj_alt))
ref_traj = deepcopy(ref_traj_)
update_friction_coefficient!(ref_traj, model, env)

# time
H = ref_traj.H
h = ref_traj.h
N_sample = 5
H_mpc = 10
h_sim = h / N_sample
H_sim = 1000 #4000 #3000

# barrier parameter
κ_mpc = 2.0e-4

obj = TrackingObjective(model, env, H_mpc,
    q = [Diagonal(1e-2 * [1.0; 0.02; 0.25; 0.25 * ones(model.dim.q-3)]) for t = 1:H_mpc],
    u = [Diagonal(3e-2 * ones(model.dim.u)) for t = 1:H_mpc],
    γ = [Diagonal(1.0e-100 * ones(model.dim.c)) for t = 1:H_mpc],
    b = [Diagonal(1.0e-100 * ones(model.dim.c * friction_dim(env))) for t = 1:H_mpc])

p = linearized_mpc_policy(ref_traj, s, obj,
    H_mpc = H_mpc,
    N_sample = N_sample,
    κ_mpc = κ_mpc,
	mode = :configuration,
    n_opts = NewtonOptions(
		solver = :lu_solver,
		r_tol = 3e-4,
		max_iter = 5,
		max_time = ref_traj.h, # HARD REAL TIME
		),
    mpc_opts = LinearizedMPCOptions(
        # live_plotting=true,
        # altitude_update = true,
        # altitude_impact_threshold = 0.05,
        # altitude_verbose = true,
        ),
	ip_opts = InteriorPointOptions(
					undercut = 5.0,
					γ_reg = 0.1,
					κ_tol = κ_mpc,
					r_tol = 1.0e-8,
					diff_sol = true,
					solver = :empty_solver,
					max_time = 1000.0,
					),
    )

q1_ref = copy(ref_traj.q[2])
q0_ref = copy(ref_traj.q[1])
q1_sim = SVector{model.dim.q}(q1_ref)
q0_sim = SVector{model.dim.q}(copy(q1_sim - (q1_ref - q0_ref) / N_sample))
@assert norm((q1_sim - q0_sim) / h_sim - (q1_ref - q0_ref) / h) < 1.0e-8

sim = simulator(s, q0_sim, q1_sim, h_sim, H_sim,
    p = p,
	ip_opts = InteriorPointOptions(
		undercut = Inf,
		γ_reg = 0.0,
		# verbose = true,
        r_tol = 1.0e-8,
        κ_tol = 1.0e-8),
    sim_opts = SimulatorOptions(warmstart = true),
    )

@time status = ContactImplicitMPC.simulate!(sim, verbose = true)

################################################################################
# Timing result
################################################################################
process!(sim)
# Time budget
ref_traj.h
# Time used on average
sim.stats.μ_dt
# Speed ratio
H_sim * h_sim / sum(sim.stats.dt)

# @elapsed status = ContactImplicitMPC.simulate!(sim)
# @profiler status = ContactImplicitMPC.simulate!(sim)

plot_lines!(vis, model, sim.traj.q[1:1:end])
plot_surface!(vis, env, ylims=[0.3, -0.05])
anim = visualize_meshrobot!(vis, model, sim.traj, sample=5)
# anim = visualize_robot!(vis, model, sim.traj, anim=anim)
anim = visualize_force!(vis, model, env, sim.traj, anim=anim, h=h_sim, sample = 5)

# Display ghosts
t_ghosts = [1, 1333, 2666]
mvis_ghosts = []
for (i,t) in enumerate(t_ghosts)
    α = i/(length(t_ghosts)+1)
    name = Symbol("ghost$i")
    mvis = build_meshrobot!(vis, model, name=name, α=α)
    push!(mvis_ghosts, mvis)
end

for (i,t) in enumerate(t_ghosts)
    name = Symbol("ghost$i")
    set_meshrobot!(vis, mvis_ghosts[i], model, sim.traj.q[t], name=name)
end

# R = Dict{Symbol,Vector}(
# 	:dyn => [],
# 	:imp => [],
# 	:mdp => [],
# 	:fri => [],
# 	:bimp => [],
# 	:bmdp => [],
# 	:bfri => [],
# 	)
# nz = num_var(model, env)
# dyn = index_dyn(model, env, quat = true)
# imp = index_imp(model, env, quat = true)
# mdp = index_mdp(model, env, quat = true)
# fri = index_fri(model, env, quat = true)
# bimp = index_bimp(model, env, quat = true)
# bmdp = index_bmdp(model, env, quat = true)
# bfri = index_bfri(model, env, quat = true)
# for t = 1:ref_traj.H
#
# 	z = deepcopy(ref_traj.z[t])
# 	θ = deepcopy(ref_traj.θ[t])
# 	r = zeros(nz)
# 	r = residual(model, env, z, θ, 0.0)
# 	s.res.r!(r, z, θ, 1e-6)
# 	push!(R[:dyn], norm(r[dyn]))
# 	push!(R[:imp], norm(r[imp]))
# 	push!(R[:mdp], norm(r[mdp]))
# 	push!(R[:fri], norm(r[fri]))
# 	push!(R[:bimp], norm(r[bimp]))
# 	push!(R[:bmdp], norm(r[bmdp]))
# 	push!(R[:bfri], norm(r[bfri]))
# end
# plt = plot()
# plot!(plt, R[:dyn], label = "dyn")
# plot!(plt, R[:imp], label = "imp")
# plot!(plt, R[:mdp], label = "mdp")
# plot!(plt, R[:fri], label = "fri")
# plot!(plt, R[:bimp], label = "bimp")
# plot!(plt, R[:bmdp], label = "bmdp")
# plot!(plt, R[:bfri], label = "bfri")


export LinearHBModel

# Linear model for 1D IMEX
"""
    LinearHBModel <: BalanceLaw

A `BalanceLaw` for modeling vertical diffusion implicitly.

write out the equations here

# Usage

    model = HydrostaticBoussinesqModel(problem)
    linear = LinearHBModel(model)

"""
struct LinearHBModel{M} <: BalanceLaw
    ocean::M
    function LinearHBModel(ocean::M) where {M}
        return new{M}(ocean)
    end
end

"""
    Copy over state, aux, and diff variables from HBModel
"""
vars_state(lm::LinearHBModel, ::Prognostic, FT) =
    vars_state(lm.ocean, Prognostic(), FT)
vars_state(lm::LinearHBModel, st::Gradient, FT) = vars_state(lm.ocean, st, FT)
vars_state(lm::LinearHBModel, ::GradientFlux, FT) =
    vars_state(lm.ocean, GradientFlux(), FT)
vars_state(lm::LinearHBModel, st::Auxiliary, FT) = vars_state(lm.ocean, st, FT)
vars_state(lm::LinearHBModel, ::UpwardIntegrals, FT) = @vars()

"""
    No integration, hyperbolic flux, or source terms
"""
@inline integrate_aux!(::LinearHBModel, _...) = nothing
@inline flux_first_order!(::LinearHBModel, _...) = nothing
@inline source!(::LinearHBModel, _...) = nothing

"""
    No need to init, initialize by full model
"""
init_state_auxiliary!(
    lm::LinearHBModel,
    state_auxiliary::MPIStateArray,
    grid,
    direction,
) = nothing
init_state_prognostic!(lm::LinearHBModel, Q::Vars, A::Vars, coords, t) = nothing

"""
    compute_gradient_argument!(::LinearHBModel)

copy u and θ to var_gradient
this computation is done pointwise at each nodal point

# arguments:
- `m`: model in this case HBModel
- `G`: array of gradient variables
- `Q`: array of state variables
- `A`: array of aux variables
- `t`: time, not used
"""
@inline function compute_gradient_argument!(
    m::LinearHBModel,
    G::Vars,
    Q::Vars,
    A,
    t,
)
    G.∇u = Q.u
    G.∇θ = Q.θ

    return nothing
end

"""
    compute_gradient_flux!(::LinearHBModel)

copy ν∇u and κ∇θ to var_diffusive
this computation is done pointwise at each nodal point

# arguments:
- `m`: model in this case HBModel
- `D`: array of diffusive variables
- `G`: array of gradient variables
- `Q`: array of state variables
- `A`: array of aux variables
- `t`: time, not used
"""
@inline function compute_gradient_flux!(
    lm::LinearHBModel,
    D::Vars,
    G::Grad,
    Q::Vars,
    A::Vars,
    t,
)
    ν = viscosity_tensor(lm.ocean)
    D.ν∇u = -ν * G.∇u

    κ = diffusivity_tensor(lm.ocean, G.∇θ[3])
    D.κ∇θ = -κ * G.∇θ

    return nothing
end

"""
    flux_second_order!(::HBModel)

calculates the parabolic flux contribution to state variables
this computation is done pointwise at each nodal point

# arguments:
- `m`: model in this case HBModel
- `F`: array of fluxes for each state variable
- `Q`: array of state variables
- `D`: array of diff variables
- `A`: array of aux variables
- `t`: time, not used

# computations
∂ᵗu = -∇⋅(ν∇u)
∂ᵗθ = -∇⋅(κ∇θ)
"""
@inline function flux_second_order!(
    lm::LinearHBModel,
    F::Grad,
    Q::Vars,
    D::Vars,
    HD::Vars,
    A::Vars,
    t::Real,
)
    F.u += D.ν∇u
    F.θ += D.κ∇θ

    return nothing
end

"""
    wavespeed(::LinaerHBModel)

calculates the wavespeed for rusanov flux
"""
function wavespeed(lm::LinearHBModel, n⁻, _...)
    C = abs(SVector(lm.ocean.cʰ, lm.ocean.cʰ, lm.ocean.cᶻ)' * n⁻)
    return C
end

"""
    boundary_state!(nf, ::LinearHBModel, args...)

applies boundary conditions for the hyperbolic fluxes
dispatches to a function in OceanBoundaryConditions.jl based on bytype defined by a problem such as SimpleBoxProblem.jl
"""
@inline function boundary_state!(nf, linear::LinearHBModel, args...)
    ocean = linear.ocean
    boundary_conditions = ocean.problem.boundary_conditions

    return ocean_boundary_state!(nf, boundary_conditions, ocean, args...)
end


"""
A `struct` for individuals that keeps individual-specific variables.
"""
mutable struct Ind{B<:AbstractFloat, C<:AbstractArray, D<:AbstractArray, E<:AbstractArray} <: AbstractAgent
  id::Int  # the individual ID
  pos::Tuple{Int, Int}  # the individuals position
  species::Int  # the species ID the individual belongs to
  biotic_phenotype::Vector{B}
  abiotic_phenotype::Vector{B}
  epistasisMat::C  # epistasis matrix
  pleiotropyMat::D  # pleiotropy matrix
  q::E  # expression array
  age::Int
  sex::Bool
  interaction_history::MArray{S, Int} where S # records the last interaction with all species
  energy::B  # determines whether need to feed (energy=0) or not.
  W::B  # survival probability
  isalive::Bool
end

struct Params{F<:AbstractFloat, I<:Int, N<:AbstractString}
  ngenes::Vector{I}
  nphenotypes::Vector{I}
  growthrates::Vector{F}
  selectionCoeffs::Vector{F}
  ploidy::Vector{I}
  optvals::Vector{Vector{Vector{Matrix{F}}}}
  optinds::Vector{Vector{I}}
  mutProbs::Vector{Vector{DiscreteNonParametric{Bool, Float64, Vector{Bool}, Vector{Float64}}}}
  mutMagnitudes::Vector{Vector{UnivariateDistribution{S} where S<:ValueSupport}}
  N::Vector{Vector{I}}
  E::Vector{Normal{F}}
  generations::I
  nspecies::I
  new_N::Vector{Vector{I}}
  migration_traits::Vector{I}
  vision_radius::Vector{I}
  check_fraction::Vector{F}
  migration_thresholds::Vector{F}
  step::MVector{1, Int64}
  nodes::Matrix{Tuple{I, I}}
  biotic_phenotypes::Vector{Vector{I}}
  abiotic_phenotypes::Vector{Vector{I}}
  max_ages::Vector{I}
  names::Dict{I, N}
  food_sources::Matrix{F}
  interactions::Matrix{F}
  resources::Matrix{I}
  resources_org::Matrix{I}
  recombination::Vector{Poisson{F}}
  initial_energy::Vector{F}
  bottlenecks::Vector{BitMatrix}
  repro_start::Vector{Int}
  repro_end::Vector{Int}
  seed::Int
end

const variance = 1.0

"""
    model_initiation(dd)

Innitializes the model.
"""
function model_initiation(dd)

  if !isnothing(dd[:model]["seed"])
    Random.seed!(dd[:model]["seed"])
  else
    rnd = Int(rand(1:1e6))
    dd[:model]["seed"] = rnd
    Random.seed!(rnd)
  end

  properties, species_arrays = create_properties(dd)
  epistasisMat, pleiotropyMat, expressionArrays = species_arrays

  space = dd[:model]["space"]

  if isnothing(space)
    fspace = GridSpace((1, 1))
  elseif typeof(space) <: Tuple
    fspace = GridSpace(space, periodic=dd[:model]["periodic"], metric=dd[:model]["metric"])
  end

  indtype = EvoDynamics.Ind{typeof(0.1), eltype(epistasisMat), eltype(pleiotropyMat), eltype(expressionArrays)}

  model = ABM(indtype, fspace, properties=properties, scheduler=Schedulers.randomly)

  # create and add agents
  for sp in 1:model.nspecies
    for (pos, n) in enumerate(model.N[sp])
      for ind in 1:n
        abiotic_ph = get_abiotic_phenotype(sp, epistasisMat[sp], pleiotropyMat[sp], expressionArrays[sp], model)
        biotic_ph = get_biotic_phenotype(sp, epistasisMat[sp], pleiotropyMat[sp], expressionArrays[sp], model)
        W = abiotic_fitness(abiotic_ph, sp, pos, model)
        W = adjust_fitness(W, sp, model)
        sex = false
        if model.ploidy[sp] == 2
          sex = rand((true, false))
        end
        interaction_history = MVector{model.nspecies, Int}(fill(0, model.nspecies))
        initial_energy = model.initial_energy[sp]
        add_agent!(model.nodes[pos], model, sp, biotic_ph, abiotic_ph, epistasisMat[sp], pleiotropyMat[sp], expressionArrays[sp], 0, sex, interaction_history, initial_energy, W, true)
      end
    end
  end

  return model
end

nnodes(x::GridSpace) = prod(size(x))
nnodes(x::ABM) = nnodes(x.space)

function create_properties(dd)
  nspecies = length(dd[:species])

  Ed = [Normal(0.0, dd[:species][i]["environmental noise"]) for i in 1:nspecies]
  Mdists = [[DiscreteNonParametric([true, false], [i, 1-i]) for i in dd[:species][arr]["mutation probabilities"]] for arr in 1:nspecies]  # μ (probability of change)
  Ddists = [[Normal(0, dd[:species][ar]["mutation magnitudes"][1]), DiscreteNonParametric([true, false], [dd[:species][ar]["mutation magnitudes"][2], 1-dd[:species][ar]["mutation magnitudes"][2]]), Normal(0, dd[:species][ar]["mutation magnitudes"][3])] for ar in 1:nspecies]  # amount of change in case of mutation

  # make single-element arrays 2D so that linAlg functions will work
  newA = Array{Array{Float64}}(undef, nspecies)
  newQ = Array{Array{Float64}}(undef, nspecies)
  for i in 1:nspecies
    if length(dd[:species][i]["epistasis matrix"]) == 1
      newA[i] = reshape(dd[:species][i]["epistasis matrix"], 1, 1)
      newQ[i] = reshape(dd[:species][i]["expression array"], 1, 1)
    else
      newA[i] = dd[:species][i]["epistasis matrix"]
      newQ[i] = dd[:species][i]["expression array"]
    end
  end

  epistasisMatS = [MArray{Tuple{size(newA[i])...}}(newA[i]) for i in eachindex(newA)]
  pleiotropyMatS = [MArray{Tuple{size(dd[:species][i]["pleiotropy matrix"])...}}(dd[:species][i]["pleiotropy matrix"]) for i in 1:nspecies]
  expressionArraysS = [MArray{Tuple{size(newQ[i])...}}(newQ[i]) for i in eachindex(newQ)]

  ngenes = [dd[:species][i]["number of genes"] for i in 1:nspecies]
  nphenotypes = [dd[:species][i]["number of phenotypes"] for i in 1:nspecies]
  growthrates = [dd[:species][i]["growth rate"] for i in 1:nspecies]
  selectionCoeffs = [dd[:species][i]["selection coefficient"] for i in 1:nspecies]
  ploidy = [dd[:species][i]["ploidy"] for i in 1:nspecies]
  optvals = [dd[:species][i]["optimal phenotype values"] for i in 1:nspecies]
  optinds = [dd[:species][i]["optimal phenotypes"] for i in 1:nspecies]
  Ns = [dd[:species][i]["N"] for i in 1:nspecies]
  migration_traits = [dd[:species][i]["migration phenotype"] for i in 1:nspecies]
  vision_radius = [dd[:species][i]["vision radius"] for i in 1:nspecies]
  check_fraction = [dd[:species][i]["check fraction"] for i in 1:nspecies]
  migration_thresholds = [dd[:species][i]["migration threshold"] for i in 1:nspecies]
  generations = dd[:model]["generations"]
  step = MVector{1, Int}(undef)
  step[1] = 0
  nnodes = Matrix{Tuple{Int, Int}}(undef, dd[:model]["space"]...)
  for col in 1:size(nnodes, 2)
    for row in 1:size(nnodes, 1)
      nnodes[row, col] = (row, col)
    end
  end
  biotic_phenotyps = [dd[:species][i]["biotic phenotypes"] for i in 1:nspecies]
  abiotic_phenotyps = [dd[:species][i]["abiotic phenotypes"] for i in 1:nspecies]
  max_ages = [dd[:species][i]["age"] for i in 1:nspecies]
  names = Dict(i => dd[:species][i]["name"] for i in 1:nspecies)
  recombination = [Poisson(dd[:species][i]["recombination"]) for i in 1:nspecies]
  initial_energy = [AbstractFloat(dd[:species][i]["initial energy"]) for i in 1:nspecies]
  bottlenecks = [dd[:species][i]["bottleneck times"] for i in 1:nspecies]
  repro_start = [dd[:species][i]["reproduction start age"] for i in 1:nspecies]
  repro_end = [dd[:species][i]["reproduction end age"] for i in 1:nspecies]

  properties = Params(ngenes, nphenotypes, growthrates, selectionCoeffs, ploidy, optvals, optinds, Mdists, Ddists, Ns, Ed, generations, nspecies, Ns, migration_traits, vision_radius, check_fraction, migration_thresholds, step, nnodes, biotic_phenotyps, abiotic_phenotyps, max_ages, names, dd[:model]["food sources"], dd[:model]["interactions"], dd[:model]["resources"], deepcopy(dd[:model]["resources"]), recombination, initial_energy, bottlenecks, repro_start, repro_end, dd[:model]["seed"])

  return properties, (epistasisMatS, pleiotropyMatS, expressionArraysS)
end

function return_opt_phenotype(species::Int, generation::Int, site::Int, model::ABM)
  nabiotic = length(model.abiotic_phenotypes[species])
  output = Array{typeof(0.1)}(undef, nabiotic)
  for trait in 1:nabiotic
    output[trait] = model.optvals[species][model.optinds[species][generation+1]][trait][site]
  end
  return output
end

function return_opt_phenotype(species::Int, generation::Int, site::Tuple{Int, Int}, model::ABM)
  nabiotic = length(model.abiotic_phenotypes[species])
  output = Array{typeof(0.1)}(undef, nabiotic)
  for trait in 1:nabiotic
    output[trait] = model.optvals[species][model.optinds[species][generation+1]][trait][site[1],site[2]]
  end
  return output
end

function model_step!(model::ABM)
  model.step[1] += 1
  model.resources .= model.resources_org
end

function agent_step!(agent::Ind, model::ABM)
  # update age
  agent.age += 1
  # abiotic survive
  abiotic_survive!(agent, model)  # the agent first survives then feeds, reproduces, and interacts.
  if !agent.isalive
    return
  end
  # use food
  burn_energy!(agent)
  # consume basic energy if agent can
  consume_food!(agent, model)
  # interact with other species
  interact!(agent, model)
  # Kill the agent if it doesn't have energy
  if agent.isalive && agent.energy < 0
    remove_agent!(agent, model)
    return
  end
  # survive
  survive!(agent, model)
  # migrate
  migrate!(agent, model)
  # reproduction for the haploid
  reproduce!(agent, model)
  # bottleneck
  if agent.isalive && model.bottlenecks[agent.species][coord2vertex(agent.pos, model), model.step[1]]
    if EvoDynamics.bottleneck(agent, model)
      remove_agent!(agent, model)
    end
  end
end

function bottleneck(agent, model)
  # @info "Package bottleneck function"
  return false
end

function remove_agent!(agent, model)
  agent.isalive = false
  kill_agent!(agent, model)
end

"""
    survive!(agent::Ind, model::ABM)

Kills the agent if it has no energy, or is too old.
"""
function survive!(agent::Ind, model::ABM)
  if !agent.isalive
    return
  elseif agent.energy < 0
    remove_agent!(agent, model)
  elseif agent.age ≥ model.max_ages[agent.species]
    remove_agent!(agent, model)
  # elseif rand() > agent.W
  #   remove_agent!(agent, model)
  end
end

"""
Kills the agent by chance given its fitness `W`
"""
function abiotic_survive!(agent::Ind, model::ABM)
  if agent.isalive && rand() > agent.W
    remove_agent!(agent, model)
  end
end

function adjust_fitness!(agent::Ind, model::ABM)
  W = agent.W < 0 ? 0.0 : agent.W
  newW = 1.0 - ( (1.0 - W) * model.selectionCoeffs[agent.species])
  agent.W = newW
end

function adjust_fitness(W, species, model::ABM)
  W2 = W < 0 ? 0.0 : W
  newW = 1.0 - ( (1.0 - W2) * model.selectionCoeffs[species])
  return newW
end

function burn_energy!(agent)
  agent.energy -= 1
end

function consume_food!(agent::Ind, model::ABM)
  environmental_consumption = model.food_sources[agent.species, agent.species]
  if environmental_consumption > 0 && model.resources[agent.pos...] >= environmental_consumption
    model.resources[agent.pos...] -= environmental_consumption
    agent.energy += environmental_consumption
  end
end

"Mutate an agent."
function mutate!(agent::Ind, model::ABM)
  mutated = false
  # mutate gene expression
  if rand(model.mutProbs[agent.species][1])
    agent.q .+= rand(model.mutMagnitudes[agent.species][1], model.ngenes[agent.species] * model.ploidy[agent.species])
    mutated = true
  end
  # mutate pleiotropy matrix
  if rand(model.mutProbs[agent.species][2])
    randnumbers = rand(model.mutMagnitudes[agent.species][2], size(agent.pleiotropyMat))
    agent.pleiotropyMat[randnumbers] .= .!agent.pleiotropyMat[randnumbers]
    mutated = true
  end
  # mutate epistasis matrix
  if rand(model.mutProbs[agent.species][3])
    agent.epistasisMat .+= rand(model.mutMagnitudes[agent.species][3], size(agent.epistasisMat))
    mutated = true
  end
  # update biotic and abiotic phenotypes and W
  if mutated
    update_fitness!(agent, model)
  end
end

coord2vertex(c, model) = c[1] + (size(model.space)[1] * (c[2]-1))

function update_fitness!(agent::Ind, model::ABM)
  abiotic_phenotype = get_abiotic_phenotype(agent, model)
  biotic_phenotype = get_biotic_phenotype(agent, model)
  W = abiotic_fitness(abiotic_phenotype, agent.species, agent.pos, model)
  W = adjust_fitness(W, agent.species, model)
  agent.biotic_phenotype .= biotic_phenotype
  agent.abiotic_phenotype .= abiotic_phenotype
  agent.W = W
end


using Makie, YOLO

#First time only (downloads 5011 images & labels!)
#YOLO.download_dataset("voc2007")

settings = YOLO.pretrained.v2_tiny_voc.load(minibatch_size=1) #Process one image at a time
model = YOLO.v2_tiny.load(settings)
YOLO.loadWeights!(model, settings)

voc = YOLO.datasets.VOC.populate()
vocloaded = YOLO.load(voc, settings, indexes = [rand(1:5011)]) #Load a single random image

#Run the model
res = model(vocloaded.imstack_mat);

#Convert the output into readable predictions
predictions = YOLO.postprocess(res, settings, conf_thresh = 0.3, iou_thresh = 0.3)

scene = YOLO.renderResult(vocloaded.imstack_mat[:,:,:,1], predictions, settings, save_file = "test.png")
display(scene)


function schedule_day(employees, env)
    ok, message = validate_day(employees, env)
    if !ok
        Logging.info( message)
        return false, {}
    end

    ok, schedule = calculate_day(employees, env)
    if !ok
        return false, {}
    end

    # again a boolean "ok" followed by message
    ok, message = validate_day_schedule(employees, env, schedule)
    if !ok
        Logging.info( message)
        return false, {}
    end

    # boolean "ok" followed by data
    return true, schedule
end

function calculate_day(employees, env)
    total_time = length(env["coverage"])
    num_employees = length(employees)

    sum_coverage = sum(env["coverage"])

    if GUROBI
        m = Model(
            solver=GurobiSolver(
                LogToConsole=0,
                TimeLimit=CALCULATION_TIMEOUT,
                OutputFlag=0,
            )
        )
    else
        m = Model(
            solver=CbcSolver(
                threads=Base.CPU_CORES,
            )
        )
    end

    ###########################################################################
    # SHIFTS
    ###########################################################################

    # Build variable
    shift_start_max = Uint8[]
    length_max = Uint8[]
    length_min = Uint8[]
    worker_keys = Any[]
    for e in keys(employees) # e is a string
        push!(shift_start_max, total_time + 1 - employees[e]["min"])
        push!(length_max, employees[e]["max"])
        push!(length_min, employees[e]["min"])
        push!(worker_keys, e)
    end

    # Shift start time
    @defVar(m, 1 <= shift_start[1:num_employees] <= total_time, Int)
    for e in 1:num_employees # e is an int
        setUpper(shift_start[e], shift_start_max[e])
    end
    # Shift length - max & min
    @defVar(m, shift_length[e=1:num_employees], Int)
    for e in 1:num_employees # e is an int
        setLower(shift_length[e], length_min[e])
        setUpper(shift_length[e], length_max[e])

        # The "delivery" constraint:
        # prevent unassigned shifts from starting after a certain time
        if "no_shifts_after" in keys(env)
            @addConstraint(m,
                shift_start[e] <= env["no_shifts_after"],
            )
        end
    end

    ###########################################################################
    # BINARY CONVERSION
    ###########################################################################
    #=
    Set up binary variable for every employee for every week to determine
    whether somebody is working. This makes it easy to sum over the week and
    see how many people are working and whether s comply with availability.
    =#

    # build some helper variables to make my life more sane
    bin_total = total_time * num_employees
    i_to_t = Int[]
    i_to_e = Uint8[] # basically what is the equivalent s index
    i_employee = String[]
    #(could have used a modulus of cycle?)

    # we are stacking cycles together, offset by total_time
    # so all t=1 first, then all t=2, etc
    for t in 1:total_time
        e_index = 0
        for e in keys(employees)
            e_index+=1
            push!(i_to_t, t)
            push!(i_to_e, e_index)
            push!(i_employee, e)
        end
    end

    # This whole section is about converting a shift start time and length to
    # a binary array of whether it has started at time t and whether it is
    # active at time t

    # define variables
    @defVar(m, started[1:bin_total], Bin)
    @defVar(m, active[1:bin_total], Bin)

    # Need to split this into its constituent parts
    # (allows us to determine whether it is equal to zero, meaning
    # that the shift started)
    @defVar(m, 0 <= start_min_now_pos[1:bin_total] <= total_time, Int)
    @defVar(m, 0 <= start_min_now_neg[1:bin_total] <= total_time, Int)
    # Is it positive?
    @defVar(m, start_min_now_helper[1:bin_total], Bin)

    # Constraints
    for i in 1:bin_total
        # instead of SOS constraint on the start_min variables:
        # Use a big M approach
        M = total_time + 1
        @addConstraints m begin
            start_min_now_pos[i] <= M * start_min_now_helper[i]
            start_min_now_neg[i] <= M * (1 - start_min_now_helper[i])
            start_min_now_pos[i] - start_min_now_neg[i] == shift_start[i_to_e[i]] - i_to_t[i]
            1 <= start_min_now_pos[i] + start_min_now_neg[i] + started[i]
        end

        # Create one contiguous chunk of ones
        if i_to_t[i] == 1
            @addConstraint(m,
                started[i] == active[i],
            )
        else # not first time
            @addConstraint(m,
                # basically "t-1" for same shift
                started[i] >= active[i] - active[i-num_employees],
            )
        end
    end

    for e in 1:num_employees
        @addConstraints m begin
            # each shift can only have one start time
            1 == sum{started[i], i=e:num_employees:bin_total}
            # Shift binary array must sum to length variable
            shift_length[e] == sum{active[i], i=e:num_employees:bin_total}
        end
    end

    ###########################################################################
    # AVAILABILITY CONSTRAINTS
    ###########################################################################

    for i in 1:bin_total
        # Employees dont' work when they are unavailable
        if employees[i_employee[i]]["availability"][i_to_t[i]] == 0
            @addConstraint(m, active[i] == 0)
        end
    end

    ###########################################################################
    # ENVIRONMENT CONSTRAINTS
    ###########################################################################

    # Make sure that the organization has the proper number of people working
    for t in 1:total_time
        # No coverage == nobody working
        if env["coverage"][t] > 0
            @addConstraint(m,
                env["coverage"][t] <= sum{active[i], i=1:bin_total; i_to_t[i] == t}
            )
        else
            # nobody needed, so nobody should be working
            @addConstraint(m,
                0 == sum{active[i], i=1:bin_total; i_to_t[i] == t}
            )
        end
    end

    ###########################################################################
    # OPTIMIZATION
    ###########################################################################

    # Set the objective
    @setObjective(m, Min, sum{shift_length[e], e=1:num_employees})

    # optimize
    status = solve(m)
    sumHours = 0
    if status == :Optimal
        hours = getObjectiveValue(m)
        Logging.debug("Day scheduler: scheduled $hours / min $sum_coverage")
        e_index = 0
        schedule = Dict()
        for e in keys(employees)
            e_index += 1
            # NOTE - we round here. If we're debugging - don't round. Might get weird floats.
            start = convert(Int, round(getValue(shift_start[e_index])))
            length = convert(Int, round(getValue(shift_length[e_index])))
            sumHours += length
            schedule[e] = {
                "start" => start,
                "length" => length,
            }
        end
        # TODO - build return function
    else
        Logging.debug( "No optimal solution found")
        gc()
        return false, {}
    end
    # ok
    gc()
    return true, schedule
end

function validate_day(employees, env)
    total_time = length(env["coverage"])
    min_hours = 0
    max_hours = 0
    coverage_total = sum(env["coverage"])

    for e in keys(employees)
        if length(employees[e]["availability"]) != total_time
            return false, string("Incorrect availability length for ", e)
        end

        longest = longest_consecutive(employees[e]["availability"])

        if longest < employees[e]["min"]
            return false, string("employee ", e, " not available long enough")
        end

        if employees[e]["min"] > total_time
            return false, string("employee ", e, " has too long a min length")
        end

        min_hours += employees[e]["min"]
        max_hours += employees[e]["max"]
    end

    Logging.debug("Coverage $coverage_total / Min hours: $min_hours / Max hours: $max_hours")
    if coverage_total > max_hours
        return false, "coverage exceeds max employee hours"
    end

    # check that enough employees are available to work each shift
    for t in 1:total_time
        availability_sum = 0
        for e in keys(employees)
            availability_sum += employees[e]["availability"][t]
        end

        if availability_sum < env["coverage"][t]
            return false, string("not enough employees available at time ", t)
        end
    end

    return true, "42"
end

function validate_day_schedule(employees, env, schedule)
    total_time = length(env["coverage"])

    # helper variable for later
    coverage_check = zeros(Int, total_time)

    for e in keys(employees)
        start = schedule[e]["start"]
        length = schedule[e]["length"]

        # shift start time >= 1
        if start < 1
            return false, string("employee ", e, " start time out of bounds")
        end

        if length <= 0
            return false, string("employee ", e, " shift length invalid")
        end

        # end  time in bounds
        if (start + length - 1) > total_time # fucking non-zero-indices
            return false, string("employee ", e, " length out of bounds")
        end

        for t in start:(start + length - 1)
            # make sure employee is available then
            if employees[e]["availability"][t] == 0
                return false, string("employee ", e, " scheduled when unavailable")
            end
            # also add to coverage tab for checking later
            coverage_check[t] += 1
        end
    end

    for t in 1:total_time
        # make sure we have enough peeps working
        if coverage_check[t] < env["coverage"][t]
            return false, string("not enough people working at time ", t)
        end
    end
    return true, ""
end


struct Match
    rule::AbstractRule
    # the rhs pattern to instantiate
    pat_to_inst::Union{Nothing,Pattern}
    # the substitution
    sub::Sub
    # the id the matched the lhs
    id::EClassId
end

const MatchesBuf = Vector{Match}

function cached_ids(g::EGraph, p::Pattern)# ::Vector{Int64}
    collect(keys(g.classes))
end

# FIXME
function cached_ids(g::EGraph, p::PatTerm)
    # println("pattern $p, $(p.head)")
    # println("all ids")
    # keys(g.classes) |> println
    # println("cached symbols")
    # cached = get(g.symcache, p.head, Set{Int64}())
    # println("symbols where $(p.head) appears")
    # appears = Set{Int64}()
    # for (id, class) ∈ g.classes
    #     for n ∈ class
    #         if n.head == p.head
    #             push!(appears, id)
    #         end
    #     end
    # end
    # # println(appears)
    # if !(cached == appears)
    #     @show cached
    #     @show appears
    # end

    collect(keys(g.classes))
    # cached
    # get(g.symcache, p.head, [])
end

# function cached_ids(g::EGraph, p::PatLiteral)
#     get(g.symcache, p.val, [])
# end

function (r::SymbolicRule)(g::EGraph, id)
    ematch(g, r.ematch_program, id) .|> sub -> Match(r, r.right, sub, id)
end

function (r::DynamicRule)(g::EGraph, id)
    ematch(g, r.ematch_program, id) .|> sub -> Match(r, nothing, sub, id)
end

function (r::BidirRule)(g::EGraph, id)
    vcat(ematch(g, r.ematch_program_l, id) .|> sub -> Match(r, r.right, sub, id),
        ematch(g, r.ematch_program_r, id) .|> sub -> Match(r, r.left, sub, id))
end

function eqsat_search_threaded!(egraph::EGraph, theory::Vector{<:AbstractRule},
        scheduler::AbstractScheduler)::MatchesBuf
    matches = MatchesBuf()
    mlock = ReentrantLock()

    inequalities = filter(Base.Fix2(isa, UnequalRule), theory)
    # never skip contradiction checks
    Threads.@threads for rule ∈ inequalities
        ids = cached_ids(egraph, rule.left)

        lock(mlock) do
            for id in ids
                append!(matches, rule(egraph, id))
            end
        end
    end

    other_rules = filter(theory) do rule
        !(rule isa UnequalRule)
    end
    Threads.@threads for rule ∈ other_rules
        # don't apply banned rules
        if !cansearch(scheduler, rule)
            # println("skipping banned rule $rule")
            continue
        end

        rule_matches = []
        ids = cached_ids(egraph, rule.left)
        for id in ids
            append!(rule_matches, rule(egraph, id))
        end

        can_yield = inform!(scheduler, rule, rule_matches)
        if can_yield
            lock(mlock) do
                append!(matches, rule_matches)
            end
        end
    end
    return matches
end

function eqsat_search!(egraph::EGraph, theory::Vector{<:AbstractRule},
    scheduler::AbstractScheduler)::MatchesBuf
    matches = MatchesBuf()

    inequalities = filter(Base.Fix2(isa, UnequalRule), theory)
    # never skip contradiction checks
    for rule ∈ inequalities
        ids = cached_ids(egraph, rule.left)

        for id in ids
            append!(matches, rule(egraph, id))
        end
    end

    other_rules = filter(theory) do rule
        !(rule isa UnequalRule)
    end
    for rule ∈ other_rules
    # don't apply banned rules
        if !cansearch(scheduler, rule)
        # println("skipping banned rule $rule")
            continue
        end

        rule_matches = Match[]
        ids = cached_ids(egraph, rule.left)
        for id in ids
            append!(rule_matches, rule(egraph, id))
        end
        can_yield = inform!(scheduler, rule, rule_matches)
        if can_yield
            append!(matches, rule_matches)
        end
    end
    return matches
end


export get_opts, get_tol_norm
import Base.Iterators: Pairs

const SPQR_DEFAULTS = pairs((
                             get_details = 0,
                             repeatable = true,
                             k = 1,
                             tol = -1.0,
                             atol = 0.0,
                             rtol = 0.0,
                             tol_norm_type = 2,
                             nsvals_large = 1,
                             nsvals_small = 1,
                             implicit_null_space_basis = true,
                             start_with_A_transpose = false,
                             ssi_min_block = 3,
                             ssi_max_block = 10,
                             ssi_nblock_increment = 5,
                             ssi_min_iters = 3,
                             ssi_max_iters = 100,
                             ssi_convergence_factor = 0.1, # also 0.25 is often safe
                             ssp_min_iters = 4,
                             ssp_max_iters = 10,
                             ssp_convergence_factor = 0.1,
));

function get_opts(opts::Pairs, symbols::Symbol...)
    res = []
    for sym in symbols
        @assert sym in keys(SPQR_DEFAULTS)
        push!(res, get(opts, sym) do; getindex(SPQR_DEFAULTS, sym); end)
    end
    pushfirst!(res, merge_opts(opts, zip_named(symbols, res)))
    length(res) == 1 ? res[1] : Tuple(res)
end

function get_tol_norm(opts::Pairs, A::AbstractMatrix{T}) where T
    opts, tol = get_opts(opts, :tol)
    normest_A = -1.0
    if tol == -1.0
        opts, atol, rtol, tol_norm_type = get_opts(opts, :atol, :rtol, :tol_norm_type)
        if ( atol <= 0.0 && rtol <= 0.0 )
            rtol = max(size(A)...) * eps(real(T))
        end
        if rtol > 0.0
            normest_A = tol_norm_type != 2 ? opnorm(A, tol_norm_type) : normest(A, tol=0.01)
            xtol = rtol * normest_A
        end
        tol = max(atol, xtol)
        opts = merge_opts(opts, tol=tol, rtol=rtol)
    end
    opts, tol, normest_A
end

zip_named(vs, vv) = pairs(NamedTuple{Tuple(vs)}(Tuple(vv)))
kwlist(;kw...) = kw
merge_opts(x::Pairs; kw...) = kwlist(;x..., kw...)
merge_opts(x::Pairs, kw) = kwlist(;x..., kw...)