using BilevelJuMP
using Test, MathOptInterface, JuMP, Dualization
# using MathOptFormat

const MOI  = MathOptInterface
const MOIU = MathOptInterface.Utilities
const MOIB = MathOptInterface.Bridges
const MOIT = MathOptInterface.Test

# TODO
# add JUMPExtension test

struct Config
    atol::Float64
    rtol::Float64
    bound_hint::Bool
    start_value::Bool
    function Config(;
        atol = 1e-6,
        rtol = 1e-6,
        bound_hint = false,
        start_value = false,
        )
        new(atol, rtol, bound_hint, start_value)
    end
end

config = Config()
CONFIG_1 = Config(atol = 1e-1, rtol = 1e-2)
CONFIG_2 = Config(atol = 1e-2, rtol = 1e-2)
CONFIG_3 = Config(atol = 1e-3, rtol = 1e-3)
CONFIG_3_start = Config(atol = 1e-3, rtol = 1e-3, start_value = true)
CONFIG_3_hint = Config(atol = 1e-3, rtol = 1e-3, bound_hint = true)
CONFIG_4 = Config(atol = 1e-4, rtol = 1e-4)
CONFIG_5 = Config(atol = 1e-5, rtol = 1e-5)

OptModeType = NamedTuple{(:opt, :mode),Tuple{Any,Any}}

solvers = OptModeType[]
solvers_cached = OptModeType[]
solvers_sos = OptModeType[]
solvers_unit = OptModeType[]
solvers_indicator = OptModeType[]
solvers_quad = OptModeType[]
solvers_bin_exp = OptModeType[]
solvers_sos_quad = OptModeType[]
solvers_nlp = OptModeType[]
solvers_nlp_sd = OptModeType[]
solvers_nlp_sd_e = OptModeType[]
solvers_nlp_sd_i = OptModeType[]
solvers_nlp_lowtol = OptModeType[]
solvers_sos_quad_bin = OptModeType[]
solvers_fa_quad_bin = OptModeType[]
solvers_fa_quad_bin_mixed = OptModeType[]
solvers_fa = OptModeType[]
solvers_fa2 = OptModeType[]
solvers_complements = OptModeType[]

include("solvers/ipopt.jl")
include("solvers/scip.jl")
# DONE
# include("solvers/cbc.jl")
# include("solvers/gurobi.jl")
# include("solvers/xpress.jl")
# include("solvers/knitro.jl")
# include("solvers/gams.jl")
# include("solvers/couenne.jl")
# include("solvers/bonmin.jl")
# include("solvers/baron.jl")
# TODO
# include("solvers/alpine.jl") # require NLP from JuMP
# include("solvers/pavito.jl") # require NLP from JuMP
# include("solvers/juniper.jl") # require NLP from JuMP
# include("solvers/path.jl") # require LCP method

include("moi.jl")
include("jump.jl")
include("jump_unit.jl")
include("mibs.jl")
include("jump_nlp.jl")

@testset "BilevelJuMP tests" begin

@testset "nlp" begin
    jump_nlp_01(IPO_OPT, mode = BilevelJuMP.ProductMode(1e-8))
    jump_nlp_02(IPO_OPT, mode = BilevelJuMP.ProductMode(1e-8))
    jump_nlp_03(IPO_OPT, mode = BilevelJuMP.ProductMode(1e-8))
    jump_nlp_04(IPO_OPT, mode = BilevelJuMP.ProductMode(1e-8))
end

@testset "Unit" begin
    jump_display()
    jump_objective()
    jump_bounds()
    jump_attributes()
    mixed_mode_unit()
    jump_constraints()
    jump_variables()
    variables_unit()
    jump_no_cb()
    constraint_unit()
    for solver in solvers_unit
        invalid_lower_objective(solver.opt, solver.mode)
        jump_display_solver(solver.opt, solver.mode)
        invalid_optimizer(solver.opt, solver.mode)
        jump_objective_solver(solver.opt, solver.mode)
        jump_attributes_solver(solver.opt, solver.mode)
    end
end

@testset "Simple LP" begin
    for solver in solvers_cached
        moi_01(solver.opt)
        moi_02(solver.opt, solver.mode)
        moi_03(solver.opt, solver.mode)
    end
end

@testset "Simple BLP JuMP" begin
    for solver in solvers_nlp
        jump_01(solver.opt, solver.mode, CONFIG_3)
        jump_01vec(solver.opt, solver.mode, CONFIG_3)
        jump_02(solver.opt, solver.mode) # numerical isntability in ipopt
        jump_03(solver.opt, solver.mode, CONFIG_3_start)
        jump_03_vec(solver.opt, solver.mode, CONFIG_3_start)
        jump_04(solver.opt, solver.mode, CONFIG_3_start)
        jump_05(solver.opt, solver.mode)
        jump_3SAT(solver.opt, solver.mode, CONFIG_3)
        jump_06(solver.opt, solver.mode)
        jump_06_sv(solver.opt, solver.mode)
        jump_07(solver.opt, solver.mode, CONFIG_2)
        jump_08(solver.opt, solver.mode, CONFIG_3_start)
        jump_09a(solver.opt, solver.mode)
        jump_09b(solver.opt, solver.mode)
        jump_11a(solver.opt, solver.mode)
        jump_11b(solver.opt, solver.mode)
        jump_12(solver.opt, solver.mode)
        jump_14(solver.opt, solver.mode)
    end
    for solver in solvers_sos
        jump_01_mixed(solver.opt)
    end
    for solver in solvers_sos
        jump_01(solver.opt, solver.mode)#
        jump_02(solver.opt, solver.mode)#
        jump_03(solver.opt, solver.mode)
        jump_04(solver.opt, solver.mode)
        jump_05(solver.opt, solver.mode)#
        jump_3SAT(solver.opt, solver.mode)
        jump_06(solver.opt, solver.mode)#
        jump_06_sv(solver.opt, solver.mode)
        jump_07(solver.opt, solver.mode)#
        jump_08(solver.opt, solver.mode)#
        jump_09a(solver.opt, solver.mode) # fail on cbc positive SOS
        jump_09b(solver.opt, solver.mode)
        jump_11a(solver.opt, solver.mode)#
        jump_11b(solver.opt, solver.mode)
        jump_12(solver.opt, solver.mode)#
        jump_14(solver.opt, solver.mode)
        #
        jump_16(solver.opt, solver.mode)
    end
    for solver in solvers_bin_exp
        jump_01(solver.opt, solver.mode, CONFIG_3_hint)
        jump_01vec(solver.opt, solver.mode, CONFIG_3_hint)
        jump_02(solver.opt, solver.mode, CONFIG_3_hint)
        jump_03(solver.opt, solver.mode, CONFIG_3_hint)
    end
    for solver in solvers_fa
        jump_01(solver.opt, solver.mode, CONFIG_3_hint)
        jump_01vec(solver.opt, solver.mode, CONFIG_3_hint)
        jump_02(solver.opt, solver.mode, CONFIG_3_hint)
        jump_03(solver.opt, solver.mode, CONFIG_3_hint)
    end
    for solver in solvers_fa2
        jump_01(solver.opt, solver.mode, CONFIG_3)
        jump_01vec(solver.opt, solver.mode, CONFIG_3)
        jump_02(solver.opt, solver.mode, CONFIG_3)
        jump_03(solver.opt, solver.mode, CONFIG_3)
    end
    for solver in solvers_indicator
        jump_01(solver.opt, solver.mode)
        jump_01vec(solver.opt, solver.mode)
        jump_02(solver.opt, solver.mode) # fail cbc - pass xpress 8.9
        jump_03(solver.opt, solver.mode)
    end
    for solver in solvers_complements # requires KNITRO
        # jump_01(solver.opt, solver.mode, CONFIG_3)
        # jump_01vec(solver.opt, solver.mode, CONFIG_3)
        # jump_02(solver.opt, solver.mode) # numerical isntability in ipopt
        jump_03(solver.opt, solver.mode, CONFIG_3_start)
        jump_03_vec(solver.opt, solver.mode, CONFIG_3_start)
        jump_04(solver.opt, solver.mode, CONFIG_3_start)
        # jump_05(solver.opt, solver.mode)
        jump_3SAT(solver.opt, solver.mode)
        jump_06(solver.opt, solver.mode)
        jump_06_sv(solver.opt, solver.mode)
        # jump_07(solver.opt, solver.mode, CONFIG_2)
        jump_08(solver.opt, solver.mode, CONFIG_3_start)
        jump_09a(solver.opt, solver.mode)
        # jump_09b(solver.opt, solver.mode)
        jump_11a(solver.opt, solver.mode)
        jump_11b(solver.opt, solver.mode)
        # jump_12(solver.opt, solver.mode)
        # jump_14(solver.opt, solver.mode)
    end
end

@testset "JuMP quad" begin
    for solver in solvers_quad
        jump_quad_01_a(solver.opt, solver.mode)
        jump_quad_01_b(solver.opt, solver.mode)
        jump_quad_01_c(solver.opt, solver.mode)
        jump_13_quad(solver.opt, solver.mode)
    end
end

@testset "Princeton Handbook Linear" begin
    for solver in vcat(solvers_nlp, solvers_nlp_sd)
        jump_HTP_lin01(solver.opt, solver.mode, CONFIG_3_start)
        # jump_HTP_lin02(solver.opt, solver.mode, CONFIG_4)
        jump_HTP_lin03(solver.opt, solver.mode)
        jump_HTP_lin03_vec(solver.opt, solver.mode)
        jump_HTP_lin04(solver.opt, solver.mode)
        jump_HTP_lin05(solver.opt, solver.mode) # broken on cbc linux on julia 1.0 and 1.2 but not 1.1 see: https://travis-ci.org/joaquimg/BilevelJuMP.jl/builds/619335351
        jump_HTP_lin06(solver.opt, solver.mode)
        jump_HTP_lin07(solver.opt, solver.mode, CONFIG_2)
        jump_HTP_lin09(solver.opt, solver.mode)
        # jump_HTP_lin10(solver.opt, solver.mode)
    end
    for solver in solvers_nlp
        jump_HTP_lin02(solver.opt, solver.mode)
        jump_HTP_lin10(solver.opt, solver.mode)
    end
    for solver in solvers_nlp_sd
        jump_HTP_lin02(solver.opt, solver.mode, CONFIG_2)
    end
    for solver in solvers_nlp_sd_i
        # jump_HTP_lin08(solver.opt, solver.mode, CONFIG_1)
        jump_HTP_lin10(solver.opt, solver.mode, CONFIG_4)
    end
    for solver in solvers_nlp_sd_e
        # TODO add dual start
        # jump_HTP_lin08(solver.opt, solver.mode, CONFIG_4)
        # jump_HTP_lin10(solver.opt, solver.mode)
    end
    for solver in solvers_sos
        jump_HTP_lin01(solver.opt, solver.mode)
        jump_HTP_lin02(solver.opt, solver.mode)
        jump_HTP_lin03(solver.opt, solver.mode) # failing cbc
        jump_HTP_lin04(solver.opt, solver.mode)
        jump_HTP_lin05(solver.opt, solver.mode) # broken on cbc linux on julia 1.0 and 1.2 but not 1.1 see: https://travis-ci.org/joaquimg/BilevelJuMP.jl/builds/619335351
        jump_HTP_lin06(solver.opt, solver.mode)
        jump_HTP_lin07(solver.opt, solver.mode)
        jump_HTP_lin08(solver.opt, solver.mode)
        jump_HTP_lin09(solver.opt, solver.mode)
        jump_HTP_lin10(solver.opt, solver.mode)
    end
    for solver in solvers_complements # requires KNITRO
        jump_HTP_lin01(solver.opt, solver.mode, CONFIG_3_start)
        jump_HTP_lin02(solver.opt, solver.mode, CONFIG_4)
        # jump_HTP_lin03(solver.opt, solver.mode)
        # jump_HTP_lin03_vec(solver.opt, solver.mode)
        # jump_HTP_lin04(solver.opt, solver.mode)
        jump_HTP_lin05(solver.opt, solver.mode) # broken on cbc linux on julia 1.0 and 1.2 but not 1.1 see: https://travis-ci.org/joaquimg/BilevelJuMP.jl/builds/619335351
        jump_HTP_lin06(solver.opt, solver.mode)
        # jump_HTP_lin07(solver.opt, solver.mode, CONFIG_2)
        # jump_HTP_lin09(solver.opt, solver.mode)
        jump_HTP_lin10(solver.opt, solver.mode)
    end
end

@testset "Princeton Handbook Quadratic" begin
    for solver in solvers_nlp
        jump_HTP_quad01(solver.opt, solver.mode)
        jump_HTP_quad02(solver.opt, solver.mode)
        jump_HTP_quad04(solver.opt, solver.mode, CONFIG_3)
        jump_HTP_quad05(solver.opt, solver.mode)
        jump_HTP_quad06(solver.opt, solver.mode, CONFIG_3)
        # jump_HTP_quad06b(solver.opt, solver.mode) # TODO weird results
        jump_HTP_quad07(solver.opt, solver.mode)
        jump_HTP_quad08(solver.opt, solver.mode) # not PSD
    end
    for solver in solvers_nlp
        jump_HTP_quad03(solver.opt, solver.mode)
        jump_HTP_quad09(solver.opt, solver.mode)
    end
    for solver in solvers_sos_quad
        jump_HTP_quad01(solver.opt, solver.mode, CONFIG_5)
        jump_HTP_quad02(solver.opt, solver.mode)
        jump_HTP_quad04(solver.opt, solver.mode)
        jump_HTP_quad05(solver.opt, solver.mode)
        jump_HTP_quad06(solver.opt, solver.mode)
        jump_HTP_quad06b(solver.opt, solver.mode)
        jump_HTP_quad07(solver.opt, solver.mode)
        jump_HTP_quad08(solver.opt, solver.mode) # not PSD
    end
    for solver in solvers_sos
        jump_HTP_quad03(solver.opt, solver.mode)
        jump_HTP_quad09(solver.opt, solver.mode)
    end
end

@testset "Ferris educational" begin
    for solver in solvers
        jump_jointc1(solver.opt, solver.mode)
        jump_jointc2(solver.opt, solver.mode)
    end
end

@testset "Bard Efficient Point Example" begin
    for solver in solvers_sos
        jump_EffPointAlgo(solver.opt, solver.mode)
    end
end

@testset "SemiRandom" begin
    for solver in solvers_sos
        jump_SemiRand(solver.opt, solver.mode)
    end
end

@testset "DNE" begin
    for solver in solvers
        jump_DTMP_01(solver.opt, solver.mode)
    end
end

@testset "DNE - modifications" begin
    for solver in solvers
        jump_DTMP_01_mod1(solver.opt, solver.mode)
        jump_DTMP_01_mod2_error(solver.opt, solver.mode)
        jump_DTMP_01_mod3vec(solver.opt, solver.mode)
        jump_DTMP_01_mod3(solver.opt, solver.mode)
    end
end

@testset "equilibrium" begin
    for solver in solvers_nlp
        jump_conejo2016(solver.opt, solver.mode)
        jump_fanzeres2017(solver.opt, solver.mode)
        # jump_eq_price(solver.opt, solver.mode)
    end
    for solver in solvers_nlp_sd
        jump_conejo2016(solver.opt, solver.mode)
        # jump_fanzeres2017(solver.opt, solver.mode)
        # jump_eq_price(solver.opt, solver.mode)
    end
    for solver in solvers_sos_quad_bin
        jump_conejo2016(solver.opt, solver.mode, config, bounds = true) # fail travis on cbc
        # jump_fanzeres2017(solver.opt, solver.mode)
        jump_eq_price(solver.opt, solver.mode) # fail travis on cbc
    end
    for solver in solvers_fa_quad_bin
        jump_conejo2016(solver.opt, solver.mode, config, bounds = true)
        # jump_fanzeres2017(solver.opt, solver.mode)
        jump_eq_price(solver.opt, solver.mode)
    end
end

@testset "Fruits" begin
    for solver in solvers_fa2
        jump_fruits(solver.opt, solver.mode, CONFIG_4, 0.05)
        jump_fruits(solver.opt, solver.mode, CONFIG_4, 0.09)
    end
end

@testset "Bilevel Conic JuMP NLP" begin
    for solver in solvers_nlp_lowtol
        jump_conic01(solver.opt, solver.mode)
        jump_conic02(solver.opt, solver.mode)
        jump_conic03(solver.opt, solver.mode)
        jump_conic04(solver.opt, solver.mode)
    end
end

@testset "Bilevel Conic JuMP MIP" begin
    for solver in solvers_fa_quad_bin_mixed
        # @time jump_conic01(solver.opt, solver.mode, config, bounds = true)
        @time jump_conic02(solver.opt, solver.mode, config, bounds = true)
        @time jump_conic03(solver.opt, solver.mode, config, bounds = true)
        @time jump_conic04(solver.opt, solver.mode, config, bounds = true)
    end
end

end


@testset "Test Joins" begin
    y = rand(10)
    z = rand(10)

    t  = table((x=1:10,   y=y), pkey=:x)
    t2 = table((x=1:2:20, z=z), pkey=:x)

    @testset "how = :inner" begin
        t_inner = table((x = 1:2:9, y = y[1:2:9], z = z[1:5]), pkey = :x)
        @test isequal(join(t, t2; how=:inner), t_inner)
    end
    @testset "how = :left" begin
        # Missing
        z_left = Union{Float64,Missing}[missing for i in 1:10]
        z_left[1:2:9] = z[1:5]
        t_left = table((x = 1:10, y = y, z = z_left))
        @test isequal(join(t, t2; how=:left), t_left)

        # DataValue
        z_left2 = [DataValue{Float64}() for i in 1:10]
        z_left2[1:2:9] = z[1:5]
        t_left2 = table((x=1:10, y = y, z = z_left2))
        @test isequal(join(t, t2, how=:left, missingtype=DataValue), t_left2)
    end
    @testset "how = :outer" begin
        # Missing
        x_outer = union(1:10, 1:2:20)
        y_outer = vcat(y, fill(missing, 5))
        z_left = Union{Float64,Missing}[missing for i in 1:10]
        z_left[1:2:9] = z[1:5]
        z_outer = vcat(z_left, z[6:10])
        t_outer = table((x=x_outer, y=y_outer, z=z_outer); pkey=:x)
        @test isequal(join(t, t2; how=:outer), t_outer)

        # DataValue
        y_outer2 = vcat(y, fill(DataValue{Float64}(), 5))
        z_left2 = [DataValue{Float64}() for i in 1:10]
        z_left2[1:2:9] = z[1:5]
        z_outer2 = vcat(z_left2, z[6:10])
        t_outer2 = table((x=x_outer, y=y_outer2, z=z_outer2); pkey=:x)
        @test isequal(join(t, t2; how=:outer, missingtype=DataValue), t_outer2)
    end
    @testset "how = :anti" begin
        t_anti = table((x=2:2:10, y=y[2:2:10]), pkey=:x)
        @test isequal(join(t, t2; how=:anti), t_anti)
    end
end


# ******************************************************************************************
# Notices:
#
# Copyright © 2022 United States Government as represented by the Administrator of the
# National Aeronautics and Space Administration.  All Rights Reserved.
#
# Disclaimers
#
# No Warranty: THE SUBJECT SOFTWARE IS PROVIDED "AS IS" WITHOUT ANY WARRANTY OF ANY KIND,
# EITHER EXPRESSED, IMPLIED, OR STATUTORY, INCLUDING, BUT NOT LIMITED TO, ANY WARRANTY
# THAT THE SUBJECT SOFTWARE WILL CONFORM TO SPECIFICATIONS, ANY IMPLIED WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, OR FREEDOM FROM INFRINGEMENT, ANY
# WARRANTY THAT THE SUBJECT SOFTWARE WILL BE ERROR FREE, OR ANY WARRANTY THAT DOCUMENTATION,
# IF PROVIDED, WILL CONFORM TO THE SUBJECT SOFTWARE. THIS AGREEMENT DOES NOT, IN ANY MANNER,
# CONSTITUTE AN ENDORSEMENT BY GOVERNMENT AGENCY OR ANY PRIOR RECIPIENT OF ANY RESULTS,
# RESULTING DESIGNS, HARDWARE, SOFTWARE PRODUCTS OR ANY OTHER APPLICATIONS RESULTING FROM
# USE OF THE SUBJECT SOFTWARE.  FURTHER, GOVERNMENT AGENCY DISCLAIMS ALL WARRANTIES AND
# LIABILITIES REGARDING THIRD-PARTY SOFTWARE, IF PRESENT IN THE ORIGINAL SOFTWARE, AND
# DISTRIBUTES IT "AS IS."
#
# Waiver and Indemnity:  RECIPIENT AGREES TO WAIVE ANY AND ALL CLAIMS AGAINST THE UNITED
# STATES GOVERNMENT, ITS CONTRACTORS AND SUBCONTRACTORS, AS WELL AS ANY PRIOR RECIPIENT.
# IF RECIPIENT'S USE OF THE SUBJECT SOFTWARE RESULTS IN ANY LIABILITIES, DEMANDS, DAMAGES,
# EXPENSES OR LOSSES ARISING FROM SUCH USE, INCLUDING ANY DAMAGES FROM PRODUCTS BASED ON,
# OR RESULTING FROM, RECIPIENT'S USE OF THE SUBJECT SOFTWARE, RECIPIENT SHALL INDEMNIFY AND
# HOLD HARMLESS THE UNITED STATES GOVERNMENT, ITS CONTRACTORS AND SUBCONTRACTORS, AS WELL
# AS ANY PRIOR RECIPIENT, TO THE EXTENT PERMITTED BY LAW.  RECIPIENT'S SOLE REMEDY FOR ANY
# SUCH MATTER SHALL BE THE IMMEDIATE, UNILATERAL TERMINATION OF THIS AGREEMENT.
# ******************************************************************************************

"""
Convert cell to `Shape` type.
"""
shape(cell::Cell) = Shape(Tuple.(collect(vertices(cell).data))[[1,2,4,3]])

"""
Return x- and y-coordinates of cell vertices.
"""
function vs(cell::Cell)
    v = hcat(collect(vertices(cell))...)
    vx = v[1,[1,2,4,3,1]]
    vy = v[2,[1,2,4,3,1]]
    return vx, vy
end

FILL_OPTIONS = [:none, :proof, :pid]

"""
    visualize!([p], root::Cell; fill::Symbol=:none, tol::Float64=0.01)

Visualize edges and/or shapes of cells. Optional first argument is a pre-existing plot.
"""
function visualize!(p, root::Cell; fill::Symbol=:none, tol::Float64=0.01)
    # plot parameters
    @assert fill in FILL_OPTIONS
    format = (alpha=1.0, lw=0.0, label=:none)

    # gray background
    fill == :none || plot!(p, shape(root); color=:gray, format...)

    # iterate through leaves
    pr = Progress(num_leaves(root), 1)
    for leaf in allleaves(root)
        w = leaf.boundary.widths
        if minimum(w) >= tol
            if fill == :none
                plot!(p, vs(leaf)..., color=:black, lw=0.5, label=:none)
            elseif fill == :proof
                if leaf.data.proven
                    color = leaf.data.member ? :green : :red
                    plot!(p, shape(leaf); color=color, format...)
                end
            elseif fill == :pid
                plot!(p, shape(leaf); color=leaf.data.pid, format...)
            end
        end
        next!(pr)
    end
    return p
end

visualize(p, root::Cell; kwargs...) = visualize!(deepcopy(p), root; kwargs...)
visualize(root::Cell; kwargs...) = visualize!(plot(), root; kwargs...)

"""
    visualize(nnet::Network, limits::NTuple{2, Vector{Float64}})

Visualize network cross-section.
"""
function visualize(nnet::Network, limits::NTuple{2, Vector{Float64}})
    if size(nnet.layers[1].weights, 2) != 2
        error("Number of free variables must be 2.")
    end

    n = 100
    f = (nnet, x, y) -> compute_output(nnet, [x, y])[]
    xlims = getindex.(limits, 1)
    ylims = getindex.(limits, 2)
    xs = collect(range(xlims..., length=n))
    ys = collect(range(ylims..., length=n))
    zs = f.(Ref(nnet), xs, ys')
    p = heatmap(xs, ys, zs'; aspect_ratio=:equal, xlims=xlims, ylims=ylims)
    return p
end


# Run conventional GA (or infinite alleles) algorithm to optimize digital circuits.
export run_inf_alleles, inf_alleles, inf_alleles_to_tuple

using DataFrames
using CSV

mutable struct inf_alleles_result_type
  numinputs::Int64
  numoutputs::Int64
  numints::Int64
  levsback::Int64
  gl::Vector{Goal}
  popsize::Int64
  max_pop_gens::Int64
  tourn_size::Int64
  func_evals::Int64
  fitness::Float64
  gen_finished::Int64
end

# This is the top-level function to run multiple simulations in parallel using pmap.
# Parameters of type IntRange can be Int64s or Bools or AbstractRanges{Int64} or AbstractRanges{Bool}.
function run_inf_alleles( nreps::Int64,
    numinputs::IntRange, numoutputs::IntRange, numinteriors::IntRange, ngoals_rng::IntRange, levelsback::IntRange,
    popsize_rng::IntRange, max_pop_gens_rng::IntRange, tournsize::Int64, csvfile::String )
  nodearity = 2
  inf_alleles_param_list = inf_alleles_result_type[]
  df = DataFrame()
  df.numinputs = Int64[]
  df.numoutputs = Int64[]
  df.numints = Int64[]
  df.levsback = Int64[]
  #df.gl = Vector{Goal}[]
  df.ngoals = Int64[]
  df.popsize = Int64[]
  df.max_pop_gens = Int64[]
  df.tournsize = Int64[]
  df.func_evals=Int64[]
  df.fitness = Float64[]
  df.gen_finished = Int64[]
  iter = 1   # iteration
  for num_inputs = numinputs
    for num_outputs = numoutputs
      funcs = default_funcs(num_inputs)
      for ngoals = ngoals_rng
        gl = randgoallist(ngoals,num_inputs,num_outputs)
        for num_interiors = numinteriors
          for levsback = levelsback
            for max_pop_gens = max_pop_gens_rng
              for popsize = popsize_rng
                for i = 1:nreps
                  p = Parameters( num_inputs, num_outputs, num_interiors, levsback )
                  #println("run inf_alleles: iter: ",iter)
                  ia = inf_alleles_result( p, gl, popsize=popsize,  max_pop_gens=max_pop_gens, tourn_size=tournsize )
                  push!(inf_alleles_param_list,ia)
                  iter += 1
                end
              end
            end
          end
        end
      end
    end
  end
  inf_alleles_result_list = pmap( x->inf_alleles(x), inf_alleles_param_list )
  #inf_alleles_result_list = map( x->inf_alleles(x), inf_alleles_param_list )
  for ia in inf_alleles_result_list
    new_row = inf_alleles_to_tuple(ia)
    Base.push!( df, new_row )
  end
  println()
  open( csvfile, "w" ) do f
    println(f,"MyInt: ",Main.CGP.MyInt)
    println(f,"funcs: ", Main.CGP.default_funcs(numinputs[end]))
    println(f,"nodearity: ",nodearity)
    CSV.write( f, df, append=true, writeheader=true )
  end
  df
end

function inf_alleles( ia::inf_alleles_result_type )
  funcs = default_funcs( ia.numinputs )
  context = construct_context(ia.numinputs)
  p = Parameters( numinputs=ia.numinputs, numoutputs=ia.numoutputs, numinteriors=ia.numints, numlevelsback=ia.levsback )
  pop = [ random_chromosome(p,funcs) for _ = 1:ia.popsize ]
  @assert ia.numinputs == p.numinputs
  for g = 1:ia.max_pop_gens
    fit_vector = zeros(Float64,ia.popsize)
    for i = 1:ia.popsize
      new_c = deepcopy(pop[i])
      mutate_chromosome!( new_c, funcs )
      output = execute_chromosome(new_c,context)
      ( new_c.fitness, matched_goals, matched_goals_list ) = goals_matched_hamming( output, ia.gl, p.numinputs)
      pop[i] = new_c
      fit_vector[i] = pop[i].fitness
    end
    #println("g:",g,"  fitv: ",fit_vector)
    maxfit = findmax( fit_vector )
    if Int(trunc(maxfit[1])) == p.numoutputs
      #println("optimum at generation ",g," = ",pop[maxfit[2]].fitness)
      ia.fitness = maxfit[1]
      ia.gen_finished = g
      ia.func_evals = g*ia.popsize
      print(".")
      #return (maxfit[1], g)
      return ia
    elseif g == ia.max_pop_gens  # return best fitness chromosome
      #println("Best fitness found: ",maxfit," after ",max_pop_gens," generations" )
      ia.fitness = maxfit[1]
      ia.gen_finished = ia.max_pop_gens
      ia.func_evals = ia.max_pop_gens*ia.popsize
      print(".")
      #return (maxfit[1], max_pop_gens)
      return ia
    end
    if ia.tourn_size == 0
      propsel!( pop, fit_vector )
    else
      tournsel!( pop, fit_vector, ia.tourn_size )
    end
  end
end

function inf_alleles_result( p::Parameters, gl::Vector{Goal} ;
    popsize::Int64=3, max_pop_gens::Int64=10, tourn_size::Int64=0 )
  inf_alleles_result_type( p.numinputs, p.numoutputs, p.numinteriors, p.numlevelsback, gl, popsize, max_pop_gens,
        tourn_size, 0, 0.0, 0 )
end

function inf_alleles_to_tuple( ia::inf_alleles_result_type )
  (
    ia.numinputs,
    ia.numoutputs,
    ia.numints,
    ia.levsback,
    length(ia.gl),
    ia.popsize,
    ia.max_pop_gens,
    ia.tourn_size,
    ia.func_evals,
    ia.fitness,
    ia.gen_finished
  )
end


"""
TestProcessingCore.jl defines the TestProcessingCore types and methods

------------------------------------------------------------------------------
COPYRIGHT/LICENSE. This file is part of the XYZ package. It is subject to
the license terms in the LICENSE file found in the top-level directory of
this distribution. No part of the XYZ package, including this file, may be
copied, modified, propagated, or distributed except according to the terms
contained in the LICENSE file.
------------------------------------------------------------------------------
"""
# --- Imports

using Boid: Boid
using Boid: AbstractProcessingCore

# --- Type definitions

mutable struct TestProcessingCore <: AbstractProcessingCore
    # Data
    field_1::Float64
    field_2::Float64

    # Default constructor
    TestProcessingCore() = new(0.1, 1)
end

# --- Method definitions

function Boid.process_data!(processing_core::TestProcessingCore, data::Vector)
    processing_core.field_1 = data[1]
    processing_core.field_2 = data[2]
    return processing_core.field_1 + processing_core.field_2
end


"""
    coef(MLM; isXSum, isZSum)

Extracts coefficients from Mlm object, accounting for the possibility of the
object containing back-estimated interactions

# Arguments

- MLM = Mlm object

# Keyword arguments

- isXSum = boolean flag indicating if a left-out `X` sum contrast
  interaction was back-estimated AND should not be included in the extracted
  coefficients. Defaults to `false`.
- isZSum = boolean flag indicating if a left-out `Z` sum contrast
  interaction was back-estimated AND should not be included in the extracted
  coefficients. Defaults to `false`.

# Value

2d array of floats

# Some notes

If the left-out `X` and/or `Z` sum contrast interactions were back-estimated,
but the user wishes to include them in the extracted coefficients anyway, set
`isXSum` and/or `isZSum` to `false`.

"""
function coef(MLM::Mlm; isXSum::Bool=false, isZSum::Bool=false)

    return MLM.B[1:(end-isXSum), 1:(end-isZSum)]
end

"""
    predict(MLM, newPredictors; isXSum, isZSum)

Calculates new predictions based on Mlm object, accounting for the possibility
of the object containing back-estimated interactions

# Arguments

- MLM = Mlm object
- newPredictors = Predictors object. Defaults to the `data.predictors` field
  in the Mlm object used to fit the model.

# Keyword arguments

- isXSum = boolean flag indicating whether a left-out `X` sum contrast
  interaction was back-estimated. Defaults to `false`.
- isZSum = boolean flag indicating whether a left-out `Z` sum contrast
  interaction was back-estimated. Defaults to `false`.

# Value

Response object

"""
function predict(MLM::Mlm, newPredictors::Predictors=MLM.data.predictors;
                 isXSum::Bool=false, isZSum::Bool=false)

  	# Include X and Z intercepts in new data if necessary
  	if MLM.data.predictors.isXIntercept==true &&
       newPredictors.isXIntercept==false
    	newPredictors.X = add_intercept(newPredictors.X)
    	newPredictors.isXIntercept = true
    	println("Adding X intercept to newPredictors.")
  	end
  	if MLM.data.predictors.isZIntercept==true &&
       newPredictors.isZIntercept==false
    	newPredictors.Z = add_intercept(newPredictors.Z)
    	newPredictors.isZIntercept = true
    	println("Adding Z intercept to newPredictors.")
  	end

  	# Remove X and Z intercepts in new data if necessary
  	if MLM.data.predictors.isXIntercept==false &&
       newPredictors.isXIntercept==true
    	newPredictors.X = remove_intercept(newPredictors.X)
    	newPredictors.isXIntercept = false
    	println("Removing X intercept from newPredictors.")
  	end
  	if MLM.data.predictors.isZIntercept==false &&
       newPredictors.isZIntercept==true
    	newPredictors.Z = remove_intercept(newPredictors.Z)
    	newPredictors.isZIntercept = false
    	println("Removing Z intercept from newPredictors.")
  	end

    # Calculate predictions
    return Response(calc_preds(newPredictors.X, newPredictors.Z,
                               coef(MLM, isXSum=isXSum, isZSum=isZSum)))
end

"""
    fitted(MLM; isXSum, isZSum)

Calculate fitted values of an Mlm object, accounting for the possibility of
the object containing back-estimated interactions

# Arguments

- MLM = Mlm object

# Keyword arguments

- isXSum = boolean flag indicating whether a left-out `X` sum contrast
  interaction was back-estimated. Defaults to `false`.
- isZSum = boolean flag indicating whether a left-out `Z` sum contrast
  interaction was back-estimated. Defaults to `false`.

# Value

Response object

"""
function fitted(MLM::Mlm; isXSum::Bool=false, isZSum::Bool=false)

    # Call the predict function with default newPredictors
    return predict(MLM, isXSum=isXSum, isZSum=isZSum)
end

"""
    resid(MLM, newData; isXSum, isZSum)

Calculates residuals of an Mlm object, accounting for the possibility of the
object containing back-estimated interactions

# Arguments

- MLM = Mlm object
- newData = RawData object. Defaults to the `data` field in the Mlm object
  used to fit the model.

# Keyword arguments

- isXSum = boolean flag indicating whether a left-out `X` sum contrast
  interaction was back-estimated. Defaults to `false`.
- isZSum = boolean flag indicating whether a left-out `Z` sum contrast
  interaction was back-estimated. Defaults to `false`.

# Value

2d array of floats

"""
function resid(MLM::Mlm, newData::RawData=MLM.data;
               isXSum::Bool=false, isZSum::Bool=false)

    # Include X and Z intercepts in new data if necessary
    if MLM.data.predictors.isXIntercept==true &&
       newData.predictors.isXIntercept==false
        newData.predictors.X = add_intercept(newData.predictors.X)
        newData.predictors.isXIntercept = true
        newData.p = newData.p + 1
        println("Adding X intercept to newData.")
    end
    if MLM.data.predictors.isZIntercept==true &&
       newData.predictors.isZIntercept==false
        newData.predictors.Z = add_intercept(newData.predictors.Z)
        newData.predictors.isZIntercept = true
        newData.q = newData.q + 1
        println("Adding Z intercept to newData.")
    end

    # Remove X and Z intercepts in new data if necessary
    if MLM.data.predictors.isXIntercept==false &&
       newData.predictors.isXIntercept==true
        newData.predictors.X = remove_intercept(newData.predictors.X)
        newData.predictors.isXIntercept = false
        newData.p = newData.p - 1
        println("Removing X intercept from newData.")
    end
    if MLM.data.predictors.isZIntercept==false &&
       newData.predictors.isZIntercept==true
        newData.predictors.Z = remove_intercept(newData.predictors.Z)
        newData.predictors.isZIntercept = false
        newData.q = newData.q - 1
        println("Removing Z intercept from newData.")
    end

	  # Calculate residuals
    return calc_resid(get_X(newData), get_Y(newData), get_Z(newData),
                      coef(MLM, isXSum=isXSum, isZSum=isZSum))
end


function plot_lines!(vis::Visualizer, model::Flamingo, q::AbstractVector;
		r=0.005, offset=0.05, size=10, name::Symbol=:Flamingo, col::Bool=true)
	r_contact = r*8/7
	p_shift = [0.0, 0.0, r_contact]
	orange_mat, blue_mat, black_mat = get_line_material(size)

	# Point Traj
	torso_point = Vector{Point{3,Float64}}()
	t1_point = Vector{Point{3,Float64}}()
	t2_point = Vector{Point{3,Float64}}()
	for qi in q
		k_torso = kinematics_1(model, qi, body = :torso, mode = :com)
		p_torso = [k_torso[1], -offset, k_torso[2]] + p_shift

		k_toe_1 = kinematics_3(model, qi, body = :foot_1, mode = :toe)
		p_toe_1 = [k_toe_1[1], -offset, k_toe_1[2]] + p_shift

		k_toe_2 = kinematics_3(model, qi, body = :foot_2, mode = :toe)
		p_toe_2 = [k_toe_2[1], -offset, k_toe_2[2]] + p_shift

		push!(torso_point, Point(p_torso...))
		push!(t1_point, Point(p_toe_1...))
		push!(t2_point, Point(p_toe_2...))
	end

	# Set lines
	if col
		setobject!(vis[name][:lines][:torso], MeshCat.Line(torso_point, orange_mat))
		setobject!(vis[name][:lines][:toe1], MeshCat.Line(t1_point, blue_mat))
		setobject!(vis[name][:lines][:toe2], MeshCat.Line(t2_point, blue_mat))
	else
		setobject!(vis[name][:lines][:torso], MeshCat.Line(torso_point, black_mat))
		setobject!(vis[name][:lines][:toe1], MeshCat.Line(t1_point, black_mat))
		setobject!(vis[name][:lines][:toe2], MeshCat.Line(t2_point, black_mat))
	end
	return nothing
end

function build_robot!(vis::Visualizer, model::Flamingo; name::Symbol=:Flamingo, r=0.005, α=1.0)
	r = convert(Float32, r)
	r_contact = convert(Float32, r*8/7)
	body_mat = MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, α))
	contact_mat = MeshPhongMaterial(color = RGBA(1.0, 165.0 / 255.0, 0.0, α))

	default_background!(vis)

	torso = Cylinder(Point3f0(0.0), Point3f0(0.0, 0.0, model.l_torso), r)
	setobject!(vis[name][:robot]["torso"], torso, body_mat)

	thigh_1 = Cylinder(Point3f0(0.0), Point3f0(0.0, 0.0, model.l_thigh1), r)
	setobject!(vis[name][:robot]["thigh1"], thigh_1, body_mat)

	calf_1 = Cylinder(Point3f0(0.0), Point3f0(0.0, 0.0, model.l_calf1), r)
	setobject!(vis[name][:robot]["calf1"], calf_1, body_mat)

	foot_1 = Cylinder(Point3f0(0.0), Point3f0(0.0, 0.0, model.l_foot1 + model.d_foot1), r)
	setobject!(vis[name][:robot]["foot1"], foot_1, body_mat)

	thigh_2 = Cylinder(Point3f0(0.0), Point3f0(0.0, 0.0, model.l_thigh2), r)
	setobject!(vis[name][:robot]["thigh2"], thigh_2, body_mat)

	calf_2 = Cylinder(Point3f0(0.0), Point3f0(0.0, 0.0, model.l_calf2), r)
	setobject!(vis[name][:robot]["calf2"], calf_2, body_mat)

	foot_2 = Cylinder(Point3f0(0.0), Point3f0(0.0, 0.0, model.l_foot2 + model.d_foot2), r)
	setobject!(vis[name][:robot]["foot2"], foot_2, body_mat)

	setobject!(vis[name][:robot]["heel1"], Sphere(Point3f0(0.0), r_contact), contact_mat)
	setobject!(vis[name][:robot]["heel2"], Sphere(Point3f0(0.0), r_contact), contact_mat)
	setobject!(vis[name][:robot]["toe1"], Sphere(Point3f0(0.0), r_contact), contact_mat)
	setobject!(vis[name][:robot]["toe2"], Sphere(Point3f0(0.0), r_contact), contact_mat)
	setobject!(vis[name][:robot]["knee1"], Sphere(Point3f0(0.0), r), body_mat)
	setobject!(vis[name][:robot]["knee2"], Sphere(Point3f0(0.0), r), body_mat)
	setobject!(vis[name][:robot]["hip"], Sphere(Point3f0(0.0), r), body_mat)
	setobject!(vis[name][:robot]["torso_top"], Sphere(Point3f0(0.0), r), body_mat)

	return nothing
end

function set_robot!(vis::Visualizer, model::Flamingo, q::AbstractVector;
		name::Symbol=:Flamingo, r=0.005)

	r = convert(Float32, r)
	r_contact = convert(Float32, r*8/7)
	p_shift = [0.0; 0.0; r_contact]

	p = [q[1]; 0.0; q[2]] + p_shift

	k_torso = kinematics_1(model, q, body = :torso, mode = :ee)
	p_torso = [k_torso[1], 0.0, k_torso[2]] + p_shift

	k_thigh_1 = kinematics_1(model, q, body = :thigh_1, mode = :ee)
	p_thigh_1 = [k_thigh_1[1], 0.0, k_thigh_1[2]] + p_shift

	k_calf_1 = kinematics_2(model, q, body = :calf_1, mode = :ee)
	p_calf_1 = [k_calf_1[1], 0.0, k_calf_1[2]] + p_shift

	k_thigh_2 = kinematics_1(model, q, body = :thigh_2, mode = :ee)
	p_thigh_2 = [k_thigh_2[1], 0.0, k_thigh_2[2]] + p_shift

	k_calf_2 = kinematics_2(model, q, body = :calf_2, mode = :ee)
	p_calf_2 = [k_calf_2[1], 0.0, k_calf_2[2]] + p_shift

	k_toe_1 = kinematics_3(model, q, body = :foot_1, mode = :toe)
	p_toe_1 = [k_toe_1[1], 0.0, k_toe_1[2]] + p_shift

	k_heel_1 = kinematics_3(model, q, body = :foot_1, mode = :heel)
	p_heel_1 = [k_heel_1[1], 0.0, k_heel_1[2]] + p_shift

	k_toe_2 = kinematics_3(model, q, body = :foot_2, mode = :toe)
	p_toe_2 = [k_toe_2[1], 0.0, k_toe_2[2]] + p_shift

	k_heel_2 = kinematics_3(model, q, body = :foot_2, mode = :heel)
	p_heel_2 = [k_heel_2[1], 0.0, k_heel_2[2]] + p_shift

	settransform!(vis[name][:robot]["thigh1"], cable_transform(p, p_thigh_1))
	settransform!(vis[name][:robot]["calf1"], cable_transform(p_thigh_1, p_calf_1))
	settransform!(vis[name][:robot]["foot1"], cable_transform(p_toe_1, p_heel_1))

	settransform!(vis[name][:robot]["thigh2"], cable_transform(p, p_thigh_2))
	settransform!(vis[name][:robot]["calf2"], cable_transform(p_thigh_2, p_calf_2))
	settransform!(vis[name][:robot]["foot2"], cable_transform(p_toe_2, p_heel_2))

	settransform!(vis[name][:robot]["torso"], cable_transform(p_torso,p))
	settransform!(vis[name][:robot]["heel1"], Translation(p_heel_1))
	settransform!(vis[name][:robot]["heel2"], Translation(p_heel_2))
	settransform!(vis[name][:robot]["toe1"], Translation(p_toe_1))
	settransform!(vis[name][:robot]["toe2"], Translation(p_toe_2))
	settransform!(vis[name][:robot]["knee1"], Translation(p_thigh_1))
	settransform!(vis[name][:robot]["knee2"], Translation(p_thigh_2))
	settransform!(vis[name][:robot]["hip"], Translation(p))
	settransform!(vis[name][:robot]["torso_top"], Translation(p_torso))

	return nothing
end

function contact_point(model::Flamingo, q::AbstractVector)
	k_toe_1 = kinematics_3(model, q, body = :foot_1, mode = :toe)
	p_toe_1 = [k_toe_1[1], 0.0, k_toe_1[2]]

	k_heel_1 = kinematics_3(model, q, body = :foot_1, mode = :heel)
	p_heel_1 = [k_heel_1[1], 0.0, k_heel_1[2]]

	k_toe_2 = kinematics_3(model, q, body = :foot_2, mode = :toe)
	p_toe_2 = [k_toe_2[1], 0.0, k_toe_2[2]]

	k_heel_2 = kinematics_3(model, q, body = :foot_2, mode = :heel)
	p_heel_2 = [k_heel_2[1], 0.0, k_heel_2[2]]

	pc = [p_toe_1, p_heel_1, p_toe_2, p_heel_2]
	return pc
end

function build_meshrobot!(vis::Visualizer, model::Flamingo; name::Symbol=:Flamingo, α=1.0)
	default_background!(vis)

	# urdf = joinpath(@__DIR__, "mesh", "flamingo.urdf")
	# mechanism = MeshCatMechanisms.parse_urdf(urdf)
	# mvis = MechanismVisualizer(mechanism, URDFVisuals(urdf, package_path=[@__DIR__]), vis[name])
	urdf = joinpath(@__DIR__, "mesh", "flamingo.urdf")
	package_path = @__DIR__
	build_meshrobot!(vis, model, urdf, package_path; name=name, α=α)
end

function convert_config(model::Flamingo, q::AbstractVector)
    _q = zeros(7)
    _q[1] = q[3]
    _q[2] = q[4] - _q[1]
    _q[4] = (q[5] - _q[2] - _q[1])
    _q[6] = q[8] - _q[4] - _q[2] - _q[1] - 0.5 * π

    _q[3] = q[6] - _q[1]
    _q[5] = q[7] - _q[3] - _q[1]
    _q[7] = q[9] - _q[5] - _q[3] - _q[1] - 0.5 * π

    return _q, compose(Translation(q[1], 0.0, q[2] + 0.02), LinearMap(RotZ(π)))
end

function flamingo_animation!(vis, sim, surf; lines = true)
	x_mean = 1.5 #0.5 * (sim.traj.q[1][1] + sim.traj.q[end][1])
	shift_traj = deepcopy(sim.traj)
	sim.traj.q[end][1]
	for t = 1:length(shift_traj.q)
		shift_traj.q[t][1] -= x_mean
	end

	# surf_shift = x -> x[3] - sim.model.env.surf(x[1] + x_mean)[1]
	# plot_surface!(vis, surf_shift, xlims=[-4.5, 4.5], ylims = [-0.5, 0.5],
	# 	col = (0.9, 0.9, 0.9), α = 1.0)
	x_range = range(-3.0, stop = 3.0, length = 100)
	surface_points = [Point(x, 0.0, surf(x + x_mean)) for x in x_range]
	surface_mat = LineBasicMaterial(color=RGBA(0.0, 0.0, 0.0, 1.0), linewidth=5)
	setobject!(vis[:lines][:surface], MeshCat.Line(surface_points, surface_mat))

	# lines && plot_lines!(vis, model, shift_traj.q[1:end], offset=-0.25, size = 5)

	settransform!(vis["/Cameras/default"],
			compose(Translation(0.0, -95.0, -1.0), LinearMap(RotY(0.0 * π) * RotZ(-π / 2.0))))
	setprop!(vis["/Cameras/default/rotated/<object>"], "zoom", 50)

	anim = visualize_meshrobot!(vis, sim.model, shift_traj, sample=10)

	return anim, shift_traj
end

function flamingo_ghost!(vis, sim, surf; lines = true)
	x_mean = 1.5 #0.5 * (sim.traj.q[1][1] + sim.traj.q[end][1])
	shift_traj = deepcopy(sim.traj)
	sim.traj.q[end][1]
	for t = 1:length(shift_traj.q)
		shift_traj.q[t][1] -= x_mean
	end

	# surf_shift = x -> x[3] - sim.model.env.surf(x[1] + x_mean)[1]
	# plot_surface!(vis, surf_shift, xlims=[-4.5, 4.5], ylims = [-0.5, 0.5],
	# 	col = (0.9, 0.9, 0.9), α = 1.0)
	x_range = range(-3.0, stop = 3.0, length = 100)
	surface_points = [Point(x, 0.0, surf(x + x_mean)) for x in x_range]
	surface_mat = LineBasicMaterial(color=RGBA(0.0, 0.0, 0.0, 1.0), linewidth=5)
	setobject!(vis[:lines][:surface], MeshCat.Line(surface_points, surface_mat))

	# lines && plot_lines!(vis, model, shift_traj.q[1:end], offset=-0.25, size = 5)

	settransform!(vis["/Cameras/default"],
			compose(Translation(0.0, -95.0, -1.0), LinearMap(RotY(0.0 * π) * RotZ(-π / 2.0))))
	setprop!(vis["/Cameras/default/rotated/<object>"], "zoom", 50)

	t = 1
	mvis1 = build_meshrobot!(vis, sim.model; name=:Flamingo1, α = 0.20)
	set_meshrobot!(vis, mvis1, sim.model, shift_traj.q[t], name = :Flamingo1)

	t = 1175
	mvis2 = build_meshrobot!(vis, sim.model; name=:Flamingo2, α = 0.40)
	set_meshrobot!(vis, mvis2, sim.model, shift_traj.q[t], name = :Flamingo2)

	t = 2565
	mvis3 = build_meshrobot!(vis, sim.model; name=:Flamingo3, α = 0.60)
	set_meshrobot!(vis, mvis3, sim.model, shift_traj.q[t], name = :Flamingo3)

	t = 3750
	mvis4 = build_meshrobot!(vis, sim.model; name=:Flamingo4, α = 0.80)
	set_meshrobot!(vis, mvis4, sim.model, shift_traj.q[t], name = :Flamingo4)

	t = shift_traj.H
	mvis5 = build_meshrobot!(vis, sim.model; name=:Flamingo5, α = 1.0)
	set_meshrobot!(vis, mvis5, sim.model, shift_traj.q[t], name = :Flamingo5)
end


module Bandits
import Base: >=, <=, +, /

export
    Bandit, StaticBandit, staticbandit, pull, regret,
    Policy, GreedyPolicy, EpsilonGreedyPolicy, ExploreThenExploit, UCB1, ThompsonSampling, choose,
    Agent, BasicAgent, BetaBernoulliAgent, NormalAgent, randomSample, choose, observe,
    Action, BanditStats, simulate, aggregate_simulate

include("utils.jl")
include("actions.jl")
include("bandit.jl")
include("policy.jl")
include("agent.jl")
include("simulate.jl")

end


# define the special PCNone type for use with direct solvers

"""
  Preconditioner type for direct solve.
  Do not use with PetscLinearOperator.

  **Public Fields**

   * none
"""
mutable struct PCNone <: AbstractPC
  is_shared::Bool
end

"""
  Outer constructor

  **Inputs**

   * mesh
   * sbp
   * eqn
   * opts
"""
function PCNone(mesh::AbstractMesh, sbp::AbstractOperator,
                eqn::AbstractSolutionData, opts::Dict)

  is_shared = false
  return PCNone(is_shared)
end

function free(pc::PCNone)

  return nothing
end

function calcPC(pc::PCNone, mesh::AbstractMesh, sbp::AbstractOperator,
                eqn::AbstractSolutionData, opts::Dict, ctx_residual, t)

  return nothing
end

function applyPC(pc::PCNone, mesh::AbstractMesh, sbp::AbstractOperator,
                 eqn::AbstractSolutionData, opts::Dict, t, b::AbstractVector,
                 x::AbstractVector)

  return nothing
end

function applyPCTranspose(pc::PCNone, mesh::AbstractMesh, sbp::AbstractOperator,
                 eqn::AbstractSolutionData, opts::Dict, t, b::AbstractVector,
                 x::AbstractVector)

  return nothing
end

function getBasePC(pc::PCNone)

  # this is the bottom of the recursion tree
  return pc
end


using Test
import KenyaCoV
include("test_functions.jl");
@testset "Simulation tests" begin
    @test test_no_infecteds()
    @test stays_nonnegative()
end


@testset "Webscale" begin
    seed!(42)
    nInd = 20
    kernel = SqExponentialKernel()
    alg = Webscale(nInd)
    @test repr(alg) == "Webscale (m = $nInd)"
    test_Zalg(alg; kernel=kernel)
end


using MonteCarloMeasurements: Particles

function test_namedtuple_data(
    idata,
    group,
    names,
    nchains,
    ndraws;
    library = "MyLib",
    coords = Dict(),
    dims = Dict(),
)
    @test idata isa InferenceData
    @test group in ArviZ.groupnames(idata)
    ds = getproperty(idata, group)
    sizes = dimsizes(ds)
    @test length(sizes) == 2 + length(coords)
    vars = vardict(ds)
    for name in string.(names)
        @test name in keys(vars)
        dim = get(dims, name, [])
        s = (x -> length(get(coords, x, []))).(dim)
        @test size(vars[name]) == (nchains, ndraws, s...)
    end
    @test "inference_library" in keys(attributes(ds))
    @test attributes(ds)["inference_library"] == library
    return nothing
end

@testset "from_namedtuple" begin
    rng = MersenneTwister(42)

    nchains, ndraws = 4, 10
    sizes = (x = (), y = (2,), z = (3, 4))
    dims = Dict("y" => ["yx"], "z" => ["zx", "zy"])
    coords = Dict("yx" => ["y1", "y2"], "zx" => 1:3, "zy" => 1:4)

    nts = [
        "NamedTuple" =>
            (; (k => randn(rng, nchains, ndraws, v...) for (k, v) in pairs(sizes))...),
        "Vector{NamedTuple}" => [
            (; (k => randn(rng, ndraws, v...) for (k, v) in pairs(sizes))...)
            for _ in 1:nchains
        ],
        "Matrix{NamedTuple}" => [
            (; (k => randn(rng, v...) for (k, v) in pairs(sizes))...)
            for _ in 1:nchains, _ in 1:ndraws
        ],
        "Vector{Vector{NamedTuple}}" => [
            [(; (k => randn(rng, v...) for (k, v) in pairs(sizes))...) for _ in 1:ndraws] for _ in 1:nchains
        ],
        "Vector{NamedTuple} particles" => [
            (; (k => Particles(randn(rng, ndraws, v...)) for (k, v) in pairs(sizes))...) for _ in 1:nchains
        ],
    ]

    @testset "posterior::$(type)" for (type, nt) in nts
        idata1 = from_namedtuple(nt; dims = dims, coords = coords, library = "MyLib")
        test_namedtuple_data(
            idata1,
            :posterior,
            keys(sizes),
            nchains,
            ndraws;
            library = "MyLib",
            coords = coords,
            dims = dims,
        )

        idata2 =
            convert_to_inference_data(nt; dims = dims, coords = coords, library = "MyLib")
        test_namedtuple_data(
            idata2,
            :posterior,
            keys(sizes),
            nchains,
            ndraws;
            library = "MyLib",
            coords = coords,
            dims = dims,
        )
    end

    @testset "$(group)" for group in [
        :posterior_predictive,
        :sample_stats,
        :predictions,
        :log_likelihood,
    ]
        @testset "::$(type)" for (type, nt) in nts
            idata1 = convert_to_inference_data(
                nt;
                (group => nt,)...,
                dims = dims,
                coords = coords,
                library = "MyLib",
            )
            test_namedtuple_data(
                idata1,
                :posterior,
                keys(sizes),
                nchains,
                ndraws;
                library = "MyLib",
                coords = coords,
                dims = dims,
            )
            test_namedtuple_data(
                idata1,
                group,
                keys(sizes),
                nchains,
                ndraws;
                library = "MyLib",
                coords = coords,
                dims = dims,
            )

            idata2 = convert_to_inference_data(
                nt;
                (group => keys(sizes),)...,
                dims = dims,
                coords = coords,
                library = "MyLib",
            )
            test_namedtuple_data(
                idata2,
                group,
                keys(sizes),
                nchains,
                ndraws;
                library = "MyLib",
                coords = coords,
                dims = dims,
            )
        end
    end

    @testset "$(group)" for group in [:prior_predictive, :sample_stats_prior]
        @testset "::$(type)" for (type, nt) in nts
            idata1 = from_namedtuple(;
                prior = nt,
                (group => nt,)...,
                dims = dims,
                coords = coords,
                library = "MyLib",
            )
            test_namedtuple_data(
                idata1,
                :prior,
                keys(sizes),
                nchains,
                ndraws;
                library = "MyLib",
                coords = coords,
                dims = dims,
            )
            test_namedtuple_data(
                idata1,
                group,
                keys(sizes),
                nchains,
                ndraws;
                library = "MyLib",
                coords = coords,
                dims = dims,
            )

            idata2 = from_namedtuple(;
                prior = nt,
                (group => keys(sizes),)...,
                dims = dims,
                coords = coords,
                library = "MyLib",
            )
            test_namedtuple_data(
                idata2,
                group,
                keys(sizes),
                nchains,
                ndraws;
                library = "MyLib",
                coords = coords,
                dims = dims,
            )
        end
    end

    @testset "$(group)" for group in
                            [:observed_data, :constant_data, :predictions_constant_data]
        _, nt = nts[1]
        idata = from_namedtuple(
            nt;
            (group => Dict("w" => [1.0, 2.0]),)...,
            dims = Dict("w" => ["wx"]),
            coords = Dict("wx" => 1:2),
            library = "MyLib",
        )
        @test idata isa InferenceData
        @test group in ArviZ.groupnames(idata)
        ds = getproperty(idata, group)
        sizes = dimsizes(ds)
        @test length(sizes) == 1
        @test "w" in keys(vardict(ds))
        @test "inference_library" in keys(attributes(ds))
        @test attributes(ds)["inference_library"] == "MyLib"
    end
end


struct Labeled{T}
    obs::T
    label::Union{Missing,Int}
end

Base.broadcastable(x::Labeled) = Ref(x)

Labeled(x) = Labeled(x, missing)

putatives(p::InfiniteParticle, x::Labeled) =
    x.label === missing ? putatives(p, x.obs) : (fit(p, x.obs, x.label), )


function available(interpreter::Interpreter, domain::Domain, state::State)
    # Ground all action definitions with arguments
    actions = Term[]
    for act in values(get_actions(domain))
        act_name = get_name(act)
        act_vars, act_types = get_argvars(act), get_argtypes(act)
        # Directly check precondition if action has no parameters
        if isempty(act_vars) && satisfy(domain, state, get_precond(act))
            push!(actions, Const(act_name))
            continue
        end
        # Include type conditions when necessary for correctness
        typecond = (@julog($ty(:v)) for (v, ty) in zip(act_vars, act_types))
        p = get_precond(act)
        if has_func(p, domain) || has_derived(p, domain) || has_quantifier(p)
            conds = prepend!(flatten_conjs(p), typecond)
        elseif get_requirements(domain)[:typing]
            conds = append!(flatten_conjs(p), typecond)
        else
            conds = flatten_conjs(p)
        end
        # Find all substitutions that satisfy preconditions
        subst = satisfiers(interpreter, domain, state, conds)
        if isempty(subst) continue end
        for s in subst
            args = [s[v] for v in act_vars if v in keys(s)]
            if any(!is_ground(a) for a in args) continue end
            term = isempty(args) ? Const(act_name) : Compound(act_name, args)
            push!(actions, term)
        end
    end
    return actions
end

function available(interpreter::Interpreter,
                   domain::Domain, state::State, action::Action, args)
    if any(!is_ground(a) for a in args)
       error("Not all arguments are ground.")
    end
    act_vars, act_types = get_argvars(action), get_argtypes(action)
    subst = Subst(var => val for (var, val) in zip(act_vars, args))
    # Construct type conditions of the form "type(val)"
    typecond = (all(ty == :object for ty in action.types) ? Term[] :
               [@julog($ty(:v)) for (v, ty) in zip(args, act_types)])
    # Check whether preconditions hold
    precond = substitute(get_precond(action), subst)
    conds = has_func(precond, domain) || has_quantifier(precond) ?
        [typecond; precond] : [precond; typecond]
    return satisfy(interpreter, domain, state, conds)
end


# Code taken from FluxML/Optimisers.jl

const Param{T<:Number} = Union{AbstractArray{T},T}

_apply(opt, x, x̄, state) = apply(opt, x, x̄, state)
_apply(opt, x, x̄, ::Nothing) = apply(opt, x, x̄)

# Immutable updates
update!(opt, x, xb, state=nothing) = _update!(opt, x, xb, state)
update(opt, x, xb, state=nothing)  = _update(opt, x, xb, state)

function _update(opt, x::Param, x̄::Param, state = nothing)
  Δ, state = _apply(opt, x, x̄, state)
  return x .- Δ, state
end

# Mutable updates

# Figure out if we can do in-place
inplace(x, y) = false
inplace(x, y::Nothing) = true
inplace(x::AbstractArray, x̄::AbstractArray) = true
inplace(x, x̄::NamedTuple) = all(inplace(getfield(x, f), getfield(x̄, f)) for f in fieldnames(typeof(x̄)))

function _update!(opt, x::AbstractArray{<:Number}, x̄::AbstractArray, state = nothing)
  Δ, state = _apply(opt, x, x̄, state)
  #println("fr = ", x)
  #println("to = ", Δ)
  x .-= Δ
  #println("after = ", x)
  return state
end

function _update!(opt, x::Tuple, x̄::Tuple, state=nothing)
  for f in propertynames(x̄)
    f̄ = getfield(x̄, f)
#  println(f,"-->",f̄)
    f̄ === nothing || _update!(opt, getfield(x, f), f̄, state)
  end
end

function _update!(opt, x, x̄::NamedTuple, state=nothing)
  for f in propertynames(x̄)
    f == :tuple_all_weights && continue
    f̄ = getproperty(x̄, f)
#    println(f,"-->",f̄)
    # println(f,"-->", typeof(getfield(x, f)))

    f̄ === nothing || _update!(opt, getfield(x, f), f̄, state)
  end
end

function _update!(opt, x, x̄, state=nothing)
  for f in propertynames(x̄)
    f̄ = getproperty(x̄, f)
#    println(f,"-->",f̄)
    # println(f,"-->", typeof(getfield(x, f)))

    f̄ === nothing || _update!(opt, getfield(x, f), f̄, state)
  end
end

# Package Integration

using Requires

@init @require Colors="5ae59095-9a9b-59fe-a467-6f913c188581" begin
  function _update(opt, x::Colors.RGB{T}, x̄::NamedTuple) where T
    Colors.RGB{T}(clamp(_update(opt, x.r, x̄.r)[1], 0, 1),
                  clamp(_update(opt, x.g, x̄.g)[1], 0, 1),
                  clamp(_update(opt, x.b, x̄.b)[1], 0, 1)), nothing
  end
end


include(joinpath("src", "axelrodnet.jl"))

Axelrodnet.run_experiment(
    experiment_name="test_run_grid",
    agentcount=100,
    n_iter=5_000_000,
    nettopology="grid",
    networkprops=Dict("grid_height" => 10),
    stubborncount=0,
    rndseed=0,
    n_replicates=10,
    export_every_n=1_000_000,
    exportdata=true,
    exportmode="default",
    keep_rawnet=false
)

Axelrodnet.run_experiment(
    experiment_name="test_run_grid",
    agentcount=100,
    n_iter=1_000_000,
    nettopology="grid",
    networkprops=Dict("grid_heit" => 10),
    stubborncount=1001,
    rndseed=0,
    n_replicates=10,
    export_every_n=100_000,
    exportdata=true,
    exportmode="aggregated",
    keep_rawnet=true
)

Axelrodnet.run_experiment(
    experiment_name="test_run_watts_strogatz",
    agentcount=100,
    n_iter=1_000_000,
    nettopology="watts_strogatz",
    networkprops=Dict("k" => 101, "beta" => 1.2),
    stubborncount=1001,
    rndseed=0,
    n_replicates=10,
    export_every_n=1_000_000,
    exportdata=true,
    exportmode="default",
    keep_rawnet=true
)

Axelrodnet.run_experiment(
    experiment_name="test_run_barabasi_albert",
    agentcount=100,
    n_iter=1_000_000,
    nettopology="barabasi_albert",
    networkprops=Dict("m0" => 200),
    stubborncount=1001,
    rndseed=0,
    n_replicates=10,
    export_every_n=100_000,
    exportdata=true,
    exportmode="aggregated",
    keep_rawnet=true
)


# ####################### Parameters #############################################
# # Abstract type
# import Base.NamedTuple as Parameter

# # Operators
# ⊕(p::Parameter, q::Parameter) = merge(p, q)
# ≅(p::Tuple, q::Tuple) = length(p) == length(q) && all([size(pp) == size(qq) for (pp, qq) in zip(p, q)])
# ≅(x::NamedTuple{N,T}, y::NamedTuple{N2,T2}) where {N,T,N2,T2} =
#   all( [key1 == key2 for (key1, key2) in zip(keys(x), keys(y))] ) && (values(x) ≅ values(y))

# +(p::Tuple, q::Tuple) = (length(p) == length(q)) ? [v1 + v2 for (v1, v2) in zip(values(p), values(q))] : AssertionError("p and q must have the same length and sizes.")
# +(p::Parameter, q::Parameter) = ( p ≅ q ) ? Parameter{Tuple(keys(p))}( values(p) + values(q) ) : AssertionError("p and q must have the same keys and sizes.")

# -(p::Tuple, q::Tuple) = (length(p) == length(q)) ? [v1 - v2 for (v1, v2) in zip(values(p), values(q))] : AssertionError("p and q must have the same length and sizes.")
# -(p::Parameter, q::Parameter) = ( p ≅ q ) ? Parameter{Tuple(keys(p))}( values(p) - values(q) ) : AssertionError("p and q must have the same keys and sizes.")

# *(a::Real, p::Parameter) = Parameter{Tuple(keys(p))}( a .* values(p) )

# ## Convert parameter (named tuple) to vector for use in fitting
# function parameter_to_vec( p :: Parameter )
#     l = length(p)
#     v = Vector{Float64}()
#     for (k, vals) in zip(keys(p), p)
#         if typeof(vals) <: AbstractFloat
#             append!(v, vals)
#         else
#             append!(v, vec(vals))
#         end
#     end
#     return v
# end

# ## Convert vector back to parameter for use after fitting
# function vec_to_parameter!(p::Parameter, v::Vector{Float64})
#     i = 1
#     j = 1
#     t = Vector{Any}(undef, length(p))
#     for val in p
#         l = length(val)
#         if l == 1
#             t[j] =  v[i]
#             i += 1
#         else
#             s = size(val)
#             t[j] = reshape( v[i:(i+l-1)], s)
#             i += l
#         end
#         j += 1
#     end
#     p = Parameter{keys(p)}(t)
#     return p
# end


@testset "Remove reactions" begin
    lp = CoreModel(
        [1.0 1 1 0; 1 1 1 0; 1 1 1 0; 0 0 0 1],
        zeros(4),
        zeros(4),
        zeros(4),
        zeros(4),
        ["r1"; "r2"; "r3"; "r4"],
        ["m1"; "m2"; "m3"; "m4"],
    )
end

@testset "Find exchange reactions and metabolites" begin
    cp = test_LP()
    @test isempty(find_exchange_reactions(cp))
    @test isempty(find_exchange_metabolites(cp))

    cp = test_simpleLP()
    @test isempty(find_exchange_reactions(cp))
    @test isempty(find_exchange_metabolites(cp))

    cp = CoreModel(
        [-1.0 -1 -2; 0 -1 0; 0 0 0],
        zeros(3),
        ones(3),
        ones(3),
        ones(3),
        ["EX_m1"; "r2"; "r3"],
        ["m1"; "m2"; "m3"],
    )
    @test find_exchange_reactions(cp) == [1]

    cp = CoreModel(
        [-1.0 0 0; 0 0 -1; 0 -1 0],
        zeros(3),
        ones(3),
        ones(3),
        ones(3),
        ["EX_m1"; "Exch_m3"; "Ex_m2"],
        ["m1"; "m2"; "m3"],
    )
    @test find_exchange_reactions(cp) == [1; 2; 3]
    @test find_exchange_metabolites(cp) == [1; 3; 2]
    @test find_exchange_reactions(cp, exc_prefs = ["Exch_"]) == [2]
    @test find_exchange_metabolites(cp, exc_prefs = ["Exch_"]) == [3]

    cp = load_model(CoreModel, joinpath("data", "toyModel1.mat"))
    @test find_exchange_reactions(cp) == [4; 5; 6]
    @test find_exchange_metabolites(cp) == [4; 5; 6]
    @test find_exchange_reactions(cp, exclude_biomass = true) == [4; 5]
    @test find_exchange_metabolites(cp, exclude_biomass = true) == [4; 5]
    @test find_exchange_reactions(cp, exclude_biomass = true, biomass_str = "biom") ==
          [4; 5]
    @test find_exchange_metabolites(cp, exclude_biomass = true, biomass_str = "biom") ==
          [4; 5]
end

@testset "Change bounds" begin
    cp = test_LP()
    change_bounds!(cp, [3; 1], xl = [-10.0; -20], xu = [10.0; 20])
    @test cp.xl == [-20; 1; -10]
    @test cp.xu == [20; 1; 10]
    change_bounds!(
        cp,
        ["gibberish1"; "r3"; "r1"; "gibberish2"],
        xl = [0; -30.0; -40; 0],
        xu = [0; 30.0; 40; 0],
    )
    @test cp.xl == [-40; 1; -30]
    @test cp.xu == [40; 1; 30]
    change_bounds!(cp, ["r1"; "r3"], xl = [-50.0; -60])
    @test cp.xl == [-50; 1; -60]
end

@testset "Verify consistency" begin
    cp = test_LP()
    @test size(cp.S) == (4, 3)
    (new_reactions, new_mets) = verify_consistency(
        cp,
        reshape(cp.S[:, end], :, 1),
        [1.0, 2.0, 3.0, 4.0],
        [2.0],
        [-1.0],
        [1.0],
        ["r4"],
        ["m1", "m2", "m3", "m6"],
        [1],
        [4],
    )
    @test new_reactions == [1]
    @test new_mets == [4]

    (new_reactions, new_mets) = verify_consistency(
        cp,
        reshape(cp.S[:, end], :, 1),
        [1.0, 2.0, 3.0, 4.0],
        [2.0],
        [-1.0],
        [1.0],
        ["r1"],
        ["m1", "m2", "m3", "m6"],
        [],
        [4],
    )
    @test new_reactions == []
    @test new_mets == [4]
end

@testset "Add reactions (checking existence and consistency)" begin
    cp = test_LP()
    @test size(cp.S) == (4, 3)
    (new_cp, new_reactions, new_mets) = add_reactions(
        cp,
        cp.S[:, end],
        [1.0, 2.0, 3.0, 4.0],
        2.0,
        -1.0,
        1.0,
        check_consistency = true,
    )
    @test n_reactions(cp) + 1 == n_reactions(new_cp)

    (new_cp, new_reactions, new_mets) = add_reactions(
        cp,
        cp.S[:, end],
        [1.0, 2.0, 3.0, 4.0],
        2.0,
        -1.0,
        1.0,
        "r1",
        ["m1", "m2", "m3", "m6"],
        check_consistency = true,
    )
    @test n_reactions(cp) == n_reactions(new_cp)
    @test n_metabolites(cp) + 1 == n_metabolites(new_cp)
end

@testset "Add reactions" begin
    cp = test_LP()
    @test size(cp.S) == (4, 3)
    cp = add_reactions(cp, 2.0 * ones(4), 3 .* ones(4), 2.0, -1.0, 1.0)
    @test size(cp.S) == (8, 4)
    cp = add_reactions(cp, 2.0 * ones(4, 1), 3 .* ones(4), 2 .* ones(1), -ones(1), ones(1))
    @test size(cp.S) == (12, 5)
    cp = add_reactions(
        cp,
        2.0 * ones(4, 10),
        3 .* ones(4),
        2 .* ones(10),
        -ones(10),
        ones(10),
    )
    @test size(cp.S) == (16, 15)

    cp = test_sparseLP()
    @test size(cp.S) == (4000, 3000)
    cp = add_reactions(cp, 2.0 * sprand(4000, 0.5), 3 .* sprand(4000, 0.5), 2.0, -1.0, 1.0)
    @test size(cp.S) == (8000, 3001)
    cp = add_reactions(
        cp,
        2.0 * sprand(4000, 1, 0.5),
        3 .* sprand(4000, 0.5),
        2 .* sprand(1, 0.5),
        -sprand(1, 0.5),
        sprand(1, 0.5),
    )
    @test size(cp.S) == (12000, 3002)
    cp = add_reactions(
        cp,
        2.0 * sprand(4000, 1000, 0.5),
        3 .* sprand(4000, 0.5),
        2 .* sprand(1000, 0.5),
        -sprand(1000, 0.5),
        sprand(1000, 0.5),
    )
    @test size(cp.S) == (16000, 4002)

    cp = test_sparseLP()
    @test size(cp.S) == (4000, 3000)
    cp = add_reactions(cp, 2.0 * ones(4000), 3 .* ones(4000), 2.0, -1.0, 1.0)
    @test size(cp.S) == (8000, 3001)
    cp = add_reactions(
        cp,
        2.0 * ones(4000, 1),
        3 .* ones(4000),
        2 .* ones(1),
        -ones(1),
        ones(1),
    )
    @test size(cp.S) == (12000, 3002)
    cp = add_reactions(
        cp,
        2.0 * ones(4000, 1000),
        3 .* ones(4000),
        2 .* ones(1000),
        -ones(1000),
        ones(1000),
    )
    @test size(cp.S) == (16000, 4002)

    # proper subset of existing metabolites
    cp = test_LP()
    new_cp = add_reactions(cp, [-1.0], zeros(1), 1.0, 0.0, 1.0, "r4", ["m1"])
    @test n_reactions(cp) + 1 == n_reactions(new_cp)

    @test_throws DimensionMismatch add_reactions(
        cp,
        2.0 * ones(4000, 1),
        3 .* ones(4000),
        2 .* ones(2),
        -ones(1),
        ones(1),
    )
end

@testset "Remove reactions" begin
    cp = test_LP()
    @test size(cp.S) == (4, 3)
    cp = remove_reactions(cp, 2)
    @test size(cp.S) == (0, 2)
    cp = remove_reactions(cp, [2, 1])
    @test size(cp.S) == (0, 0)

    cp = test_LP()
    @test size(cp.S) == (4, 3)
    cp = remove_reactions(cp, "r0")
    @test size(cp.S) == (4, 3)
    cp = remove_reactions(cp, "r1")
    @test size(cp.S) == (0, 2)
    cp = remove_reactions(cp, ["r2"])
    @test size(cp.S) == (0, 1)
end


const Tms = typeof(1.0f0*ms)

"""
    CellsTrial(t, dFoF)

Store cellular responses for a single trial. `t` is the list of frame times (in units of time, see Unitful.jl) and
`dFoF` is a matrix with one row for each time in `t`, and one column per cell.

# Examples

```jldoctest demo; setup=:(using EcoTrialStructure)
julia> dFoF = [ 0.1 0.8;
               -0.1 0.7;
                0.2 0.6;
                0.1 0.5;
                0.0 0.4];

julia> ct = CellsTrial((100:100:500) * ms, dFoF)
2 cells with 5 timepoints
```

It's straightforward to extract the dFoF values over a particular time interval:

```jldoctest demo
julia> tframes, df = ct[175ms..310ms, :];  # frames range nearest to given start..stop times

julia> tframes
200.0f0 ms..300.0f0 ms

julia> df
2×2 Matrix{Float64}:
 -0.1  0.7
  0.2  0.6
```
This extracted the `dFoF` values recorded between 200ms and 300ms, inclusive.

In other cases, you might want a specific number of frames, starting at a particular time:

```jldoctest demo
julia> tframes, df = ct[FrameSeq(175ms, 3), :];  # start nearest to 175ms, and grab 3 frames' worth

julia> tframes
200.0f0 ms..400.0f0 ms

julia> df
3×2 OffsetArray(::Matrix{Float64}, 0:2, 1:2) with eltype Float64 with indices 0:2×1:2:
 -0.1  0.7
  0.2  0.6
  0.1  0.5
```
"""
struct CellsTrial{T<:AbstractFloat}
    t::Vector{Tms}
    dFoF::Matrix{T}

    function CellsTrial{T}(t, dFoF) where T
        issorted(t) || throw(ArgumentError("frame times must be increasing"))
        axes(dFoF, 1) == axes(t, 1) || throw(DimensionMismatch("frame times and dFoF must match, got sizes $(length(t)) and $(size(dFoF))"))
        return new{T}(t, dFoF)
    end
end
CellsTrial(t, dFoF::AbstractMatrix{T}) where T<:AbstractFloat = CellsTrial{T}(t, dFoF)
CellsTrial(t, dFoF::AbstractMatrix) = CellsTrial{Float64}(t, dFoF)

Base.:(==)(a::CellsTrial, b::CellsTrial) = a.t == b.t && a.dFoF == b.dFoF
Base.isequal(a::CellsTrial, b::CellsTrial) = isequal(a.t, b.t) && isequal(a.dFoF, b.dFoF)

const hash_seed_CellsTrial = Int === Int64 ? 0x80450cff5158f34c : 0xbcdbf5e0
Base.hash(ct::CellsTrial, h::UInt) = hash(hash_seed_CellsTrial, hash(ct.t, hash(ct.dFoF, h)))

Base.show(io::IO, ct::CellsTrial) = print(io, size(ct.dFoF, 2), " cells with ", length(ct.t), " timepoints")

# indexing (primarily, selecting frames based on time)

Base.getindex(ct::CellsTrial, ti::Union{Colon,AbstractVector{<:Integer}}, ci) = (ct.t[ti], ct.dFoF[ti, ci])

function Base.getindex(ct::CellsTrial, ti::AbstractInterval, ci)
    ibegin, iend = idxof(ct.t, minimum(ti)), idxof(ct.t, maximum(ti))
    return ct.t[ibegin] .. ct.t[iend], ct.dFoF[ibegin:iend, ci]
end

"""
    FrameSeq(tstart, nframes)
    FrameSeq(tstart, offsetrange)
    FrameSeq(eventfield::Symbol, nframes)

A sequence of `nframes` frames starting at the nearest timepoint to `tstart`.
Optionally specify a range of frames offset from `tstart`.

```jldoctest; setup=:(using EcoTrialStructure)
julia> fs = FrameSeq(100ms, 5)
FrameSeq(100.0f0 ms, 0:4)

julia> fs = FrameSeq(100ms, -2:4)
FrameSeq(100.0f0 ms, -2:4)
```

See [`CellsTrial`](@ref) for an example using this in indexing.

Alternatively, this can be constructed specifying a particular fieldname of [`EventTiming`](@ref),
in which case the concrete timing can be deferred until a later time based on a specific
trial:

```jldoctest; setup=:(using EcoTrialStructure)
julia> fs = FrameSeq(:go, 5)
FrameSeq(:go, 0:4)

julia> et = EventTiming(0ms, 100ms, 400ms, 450ms, 837ms, 1.2s)
EventTiming(trial_start=0.0f0 ms, offer_on=100.0f0 ms, offer_off=400.0f0 ms, go=450.0f0 ms, choice=837.0f0 ms, trial_end=1200.0f0 ms)

julia> fs(et)
FrameSeq(450.0f0 ms, 0:4)
```
"""
struct FrameSeq
    start::Union{Tms,Symbol}
    idx::UnitRange{Int}
end
FrameSeq(start::Unitful.Quantity, len::Integer) = FrameSeq(Tms(start), 0:len-1)
FrameSeq(start::Symbol, len::Integer) = FrameSeq(start, 0:len-1)
FrameSeq(start::Unitful.Quantity, idx::AbstractUnitRange) = FrameSeq(Tms(start), idx)

function Base.getindex(ct::CellsTrial, ti::FrameSeq, ci)
    isdeferred(ti) && throw(ArgumentError("indexing requires a concrete `FrameSeq`, use `fs(et::EventTiming)`"))
    irange = idxsof(ct.t, ti)
    return ct.t[first(irange)] .. ct.t[last(irange)], ct.dFoF[irange, ci]
end

Base.length(fs::FrameSeq) = length(fs.idx)
Base.axes(fs::FrameSeq) = (Base.IdentityUnitRange(fs.idx),)
Base.axes(fs::FrameSeq, d) = axes(fs)[d]

function Base.checkbounds(::Type{Bool}, ct::CellsTrial, ti::FrameSeq, ci)
    r1, r2 = _idxof(ct.t, ti.start)
    checkbounds(Bool, ct.t, r1) && checkbounds(Bool, ct.t, r2) || return false
    istart = _idxof(ct.t, ti.start, r1, r2)
    return checkbounds(Bool, ct.dFoF, istart .+ ti.idx, ci)
end

"""
    OfferType(nA, nB, leftA::Bool)

Encode the offers (`nA` and `nB` are the number of drops of A and B, respectively),
and whether A was on the left.

If you have a list of `OfferType`s, you can `sort` them. High-A trials will come before high-B trials,
and left before right.
"""
struct OfferType
    nA::Int8
    nB::Int8
    leftA::Bool
end

# This constructor makes it possible to copy/paste the output of the `show` method below
OfferType(; nA, nB, leftA) = OfferType(nA, nB, leftA)

OfferType(tt::OfferType) = tt

Base.show(io::IO, tt::OfferType) = print(io, "OfferType(nA=", tt.nA, ", nB=", tt.nB, ", leftA=", tt.leftA, ")")

function Base.isless(a::OfferType, b::OfferType)
    ratio(tt::OfferType) = tt.nB/tt.nA

    # Ordering puts high-A trials before high-B trials
    rA, rB = ratio(a), ratio(b)
    isless(rA, rB) && return true
    isless(rB, rA) && return false
    a.nA > b.nA && return true
    a.nA < b.nA && return false
    a.nB > b.nB && return false
    a.nB < b.nB && return true
    return a.leftA > b.leftA   # if all else are equal, put left before right
end

"""
    TrialResult(nA, nB, leftA::Bool, choseA::Union{Bool,Missing})
    TrialResult(tt::OfferType, choseA::Union{Bool,Missing})

Encode the offer configuration (see [`OfferType`](@ref)), and whether the animal chose A, B, or failed to make a choice
(`choseA = true | false | missing`, respectively).
"""
struct TrialResult
    tt::OfferType
    choseA::Union{Bool,Missing}   # missing if the animal didn't lick
end
TrialResult(nA, nB, leftA, choseA) = TrialResult(OfferType(nA, nB, leftA), choseA)
TrialResult(; nA, nB, leftA, choseA) = TrialResult(nA, nB, leftA, choseA)
TrialResult(tr::TrialResult) = tr

function Base.getproperty(tr::TrialResult, name::Symbol)
    name === :nA && return getfield(tr, :tt).nA
    name === :nB && return getfield(tr, :tt).nB
    name === :leftA && return getfield(tr, :tt).leftA
    name === :tt && return getfield(tr, :tt)
    return getfield(tr, name)
end

Base.show(io::IO, tr::TrialResult) = print(io, "TrialResult(nA=", tr.nA, ", nB=", tr.nB, ", leftA=", tr.leftA, ", choseA=", tr.choseA, ")")

OfferType(tr::TrialResult) = tr.tt

function Base.isless(a::TrialResult, b::TrialResult)
    tta, ttb = OfferType(a), OfferType(b)
    isless(tta, ttb) && return true
    isless(ttb, tta) && return false
    # This implements `true` > `missing` > `false` logic
    rank(choice) = isa(choice, Bool) ? Float64(choice) : 0.5
    return isless(rank(b.choseA), rank(a.choseA))
end

"""
    EventTiming(trial_start, offer_on, offer_off, go, choice, trial_end)

Specify the timing of events during a trial. All times should be in physical units (`s` or `ms`).

# Examples

```jldoctest; setup=:(using EcoTrialStructure)
julia> EventTiming(0ms, 100ms, 400ms, 450ms, 837ms, 1.2s)
EventTiming(trial_start=0.0f0 ms, offer_on=100.0f0 ms, offer_off=400.0f0 ms, go=450.0f0 ms, choice=837.0f0 ms, trial_end=1200.0f0 ms)
```
"""
struct EventTiming
    trial_start::Tms
    offer_on::Tms
    offer_off::Tms
    go::Tms
    choice::Tms
    trial_end::Tms
end
EventTiming(; trial_start, offer_on, offer_off, go, choice, trial_end) =
    EventTiming(trial_start, offer_on, offer_off, go, choice, trial_end)

Base.show(io::IO, et::EventTiming) = print(io, "EventTiming(trial_start=", et.trial_start, ", offer_on=", et.offer_on, ", offer_off=", et.offer_off, ", go=", et.go, ", choice=", et.choice, ", trial_end=", et.trial_end, ")")

function (fs::FrameSeq)(et::EventTiming)
    start = fs.start
    return isa(start, Symbol) ? FrameSeq(getfield(et, fs.start), fs.idx) : fs
end


"Type signature for a PDDL fluent (i.e. predicate or function)."
struct Signature{N}
    name::Symbol
    type::Symbol
    args::NTuple{N, Var}
    argtypes::NTuple{N, Symbol}
end

Signature(name, type, args, argtypes) =
    Signature{length(args)}(name, type, Tuple(args), Tuple(argtypes))

Signature(term::Term, argtypes, type) =
    Signature(term.name, type, term.args, argtypes)

"Returns the arity (i.e. number of arguments) of a fluent signature."
arity(::Signature{N}) where {N} = N

function Base.convert(::Type{Term}, sig::Signature{0})
    return Const(sig.name)
end

function Base.convert(::Type{Term}, sig::Signature)
    return Compound(sig.name, collect(sig.args))
end

function Base.show(io::IO, sig::Signature)
    args = ["?$(lowercase(string(a.name))) - $t"
            for (a, t) in zip(sig.args, sig.argtypes)]
    print(io, "($(join([sig.name; args], " "))) - $(sig.type)")
end


# YOLO.jl
# loss.jl
# SOME DEFAULT SETTINGS
# ANCHORS = [1.07709888,  1.78171903, 2.71054693,  5.12469308,  10.47181473,
#                 10.09646365, 5.48531347,  8.11011331]
# LABELS = ["aeroplane","bicycle","bird","boat","bottle",
#                 "bus","car","cat","chair","cow",
#                 "diningtable","dog","horse","motorbike","person",
#                 "pottedplant","sheep","sofa","train","tvmonitor"]
# BATCH_SIZE        = 1 # default was 500
# IMAGE_H, IMAGE_W  = 416, 416
# GRID_H,  GRID_W   = 13 , 13
# TRUE_BOX_BUFFER   = 50
# BOX               = Int(length(ANCHORS)/2)
# CLASS             = length(LABELS)
# LAMBDA_NO_OBJECT = 1.0
# LAMBDA_OBJECT    = 5.0
# LAMBDA_COORD     = 1.0
# LAMBDA_CLASS     = 1.0
# GRIDSIZE         = (GRID_H,  GRID_W) #usually
#
# # VOC folder paths
# train_image_folder = "C:/Framework/Julia/YOLO_loss/VOCdevkit/VOC2012/JPEGImages/"
# #train_image_folder = "C:/Framework/Julia/YOLO_loss/VOCdevkit/VOC2007/JPEGImages/"
# train_annot_folder = "C:/Framework/Julia/YOLO_loss/VOCdevkit/VOC2012/Annotations/"
# #train_annot_folder = "C:/Framework/Julia/YOLO_loss/VOCdevkit/VOC2007/Annotations/"
#
# # Just a nice dict
# generator_config = Dict(
#     "IMAGE_H"         => IMAGE_H,
#     "IMAGE_W"         => IMAGE_W,
#     "GRID_H"          => GRID_H,
#     "GRID_W"          => GRID_W,
#     "LABELS"          => LABELS,
#     "ANCHORS"         => ANCHORS,
#     "BATCH_SIZE"      => BATCH_SIZE,
#     "TRUE_BOX_BUFFER" => TRUE_BOX_BUFFER,
# )
#
# # Stuff used for basic testing
# b_batch = zeros(Float32,1,1,1,1,50,4)# this is just a placeholder for true_boxes dimensions
# true_boxes = zeros(size(b_batch))
# y_batch = zeros(Float32,1, 13, 13, 4, 25)
# sz   = BATCH_SIZE*GRID_W*GRID_H*BOX*(4 + 1 + CLASS)
# #y_pred = np.random.normal(size=size,scale = 10/(GRID_H*GRID_W))
# y_pred = randn(Float32,sz)/33
# y_pred = reshape(y_pred,BATCH_SIZE,GRID_H,GRID_W,BOX,4 + 1 + CLASS)
# y_true = y_pred
# custom_loss_core(y_pred,y_pred,true_boxes,GRIDSIZE,
#         BATCH_SIZE,ANCHORS,LAMBDA_COORD,LAMBDA_CLASS,LAMBDA_NO_OBJECT, LAMBDA_OBJECT)
#
# Basic helper functions
"""
```
sigmoid(x) = 1 / (1 + exp(-x))
```
"""
sigmoid(x) = 1 / (1 + exp(-x))

"""
```
sparse_softmax_cross_entropy_with_logits based on:
https://stackoverflow.com/questions/37775596/where-is-the-origin-coding-of-sparse-softmax-cross-entropy-with-logits-function
```
"""
function sparse_softmax_cross_entropy_with_logits(true_box_class, pred_box_class)
    x = -1.0 .* log.(softmax(pred_box_class)) # tensorflow gives different results
    # for softmax applied on pred_box_class(why???)
    return x[true_box_class]
end

"""
```
expand_dims(x::Array,axis::Int)
Add a new 1-sized dim at axis value, using python style -ve indexing for positions from end of dims
i.e.
size(expand_dims(rand(10,24,12),2))
(10, 1, 24, 12)
size(expand_dims(rand(10,24,12),2))
(10, 24, 1, 12)
size(expand_dims(rand(10,24,12),-1))
(10, 24, 12, 1)
size(expand_dims(rand(10,24,12),-2))
(10, 24, 1, 12)
```
"""
function expand_dims(x::Array,axis::Int64)
    dims = collect(size(x))
    if abs(axis) > length(dims)
        error(string("`axis` value (",abs(axis),") is larger than dimensions of `x` (",length(dims),")"))
    end
    if axis > 0
        splice!(dims, axis, [1,dims[axis]])
        elseif axis < 0
        splice!(dims, length(dims)+axis+1, [dims[length(dims)+axis+1],1])
    else
        error("Axis cannot be equal to 0")
    end

    reshape(x, dims...)
end

# Implementation based on https://fairyonice.github.io/Part_4_Object_Detection_with_Yolo_using_VOC_2012_data_loss.html
"""
```
get_cell_grid(gridsize::Tuple{Int,Int},nbatchsize::Int,nboxes::Int)
Helper function to assure that the bounding box x and y are in the grid cell scale
```
"""
function get_cell_grid(gridsize::Tuple{Int,Int},batchsize::Int,nboxes::Int)
    # gridsize = (13,13) #usually
    # batchsize = 1 # adjust based on RAM size
    # nboxes = 4 #(in this case) length(ANCHORS)/2 in python
    cell_grid = Array{Float32}(undef,batchsize,gridsize[1],gridsize[2],nboxes,2)
    for batch in 1:batchsize
        for gh in 1:gridsize[1]
            for gw in 1:gridsize[2]
                for box in 1:nboxes
                    cell_grid[batch,gh,gw,box,1] =  gw - 1
                    cell_grid[batch,gh,gw,box,2] =  gh - 1
                end
            end
        end
    end
    return cell_grid
end

"""
```
Adjust prediction
== input ==
y_pred : takes any real values
tensor of shape = (N batch, NGrid h, NGrid w, NAnchor, 4 + 1 + N class)
ANCHORS : list containing width and height specializaiton of anchor box
== output ==
pred_box_xy : shape = (N batch, N grid x, N grid y, N anchor, 2), contianing [center_y, center_x] rangining [0,0]x[grid_H-1,grid_W-1]
pred_box_xy[irow,igrid_h,igrid_w,ianchor,0] =  center_x
pred_box_xy[irow,igrid_h,igrid_w,ianchor,1] =  center_1
calculation process:
sigmoid(y_pred[...,:2]) : takes values between 0 and 1
sigmoid(y_pred[...,:2]) + cell_grid : takes values between 0 and grid_W - 1 for x coordinate
takes values between 0 and grid_H - 1 for y coordinate
pred_Box_wh : shape = (N batch, N grid h, N grid w, N anchor, 2),
containing width and height, rangining [0,0]x[grid_H-1,grid_W-1]
pred_box_conf : shape = (N batch, N grid h, N grid w, N anchor, 1), containing confidence to range between 0 and 1
pred_box_class : shape = (N batch, N grid h, N grid w, N anchor, N class), containing
```
"""
function adjust_scale_prediction(y_pred::Array{Float32,5}, cell_grid::Array{Float32,5}, anchors::Vector{Float64})
    # anchors is an array with width,height,width,height....
    # size(y_pred) = (1,13,13,4,25)
    nboxes = Int(length(anchors)/2)

    ### adjust x and y
    # the bounding box bx and by are rescaled to range between 0 and 1 for given gird.
    # Since there are nboxes x nboxes grids, we rescale each bx and by to range between 0 to nboxes + 1
    pred_box_xy = sigmoid.(y_pred[:,:,:,:,1:2]) + cell_grid # bx, by
    #size(pred_box_xy) = (1,13,13,4,2)

    ### adjust w and h
    # exp to make width and height positive
    # rescale each grid to make some anchor "good" at representing certain shape of bounding box
    anchor_aux = permutedims(reshape(anchors,1,1,1,2,nboxes),[1,2,3,5,4]) # Julia uses column major
    pred_box_wh = exp.(y_pred[:,:,:,:,3:4]) .* anchor_aux
    #size(pred_box_wh) = (1,13,13,4,2)

    ### adjust confidence
    pred_box_conf = sigmoid.(y_pred[:,:,:,:,5])# prob bb
    #size(pred_box_conf) = (1,13,13,4)

    ### adjust class probabilities
    pred_box_class = y_pred[:,:,:,:,6:end] # prC1, prC2, ..., prC20
    #size(pred_box_class) = (1,13,13,4,20)

    return pred_box_xy,pred_box_wh,pred_box_conf,pred_box_class
end

function extract_ground_truth(y_true::Array{Float32,5})
    # size(y_true) = (1,13,13,4,25)
    true_box_xy = y_true[:,:,:,:,1:2]
    # size(true_box_xy) = (1,13,13,4,2)
    true_box_wh = y_true[:,:,:,:,3:4]
    # size(true_box_wh) = (1,13,13,4,2)
    true_box_conf = y_true[:,:,:,:,5]
    # size(true_box_conf) = (1,13,13,4)
    true_box_class = dropdims(argmax(y_true[:,:,:,:,6:end], dims=5);dims=5)#?????
    # size(true_box_class) = (1,13,13,4)??
    return true_box_xy, true_box_wh, true_box_conf, true_box_class
end

"""
```
coord_mask:      np.array of shape (Nbatch, Ngrid h, N grid w, N anchor, 1)
lambda_{coord} L_{i,j}^{obj}
```
"""
function calc_loss_xywh(true_box_conf,LAMBDA_COORD,true_box_xy, pred_box_xy,true_box_wh,pred_box_wh)
    # lambda_{coord} L_{i,j}^{obj}
    # np.array of shape (Nbatch, Ngrid h, N grid w, N anchor, 1)
    coord_mask   = expand_dims(true_box_conf, -1) .* LAMBDA_COORD
    # size(coord_mask) = (1,13,13,4,1)
    nb_coord_box = sum(Float32.(coord_mask .> 0.0))
    #size(nb_coord_box) = 1
    loss_xy      = sum((true_box_xy.-pred_box_xy).^2 .* coord_mask) / (nb_coord_box + 1e-6) / 2.0
    loss_wh      = sum((true_box_wh.-pred_box_wh).^2 .* coord_mask) / (nb_coord_box + 1e-6) / 2.0
    return loss_xy + loss_wh, coord_mask
end

"""
```
== input ==
true_box_conf  : tensor of shape (N batch, N grid h, N grid w, N anchor)
true_box_class : tensor of shape (N batch, N grid h, N grid w, N anchor), containing class index
pred_box_class : tensor of shape (N batch, N grid h, N grid w, N anchor, N class)
LAMBDA_CLASS    : 1.0
== output ==
class_mask
if object exists in this (grid_cell, anchor) pair and the class object receive nonzero weight
class_mask[iframe,igridy,igridx,ianchor] = 1
else:
0
```
"""
function calc_loss_class(true_box_conf,LAMBDA_CLASS, true_box_class,pred_box_class)
    class_mask   = true_box_conf  .* LAMBDA_CLASS
    # size(class_mask) = (1,13,13,4)
    nb_class_box = sum(Float32.(class_mask .> 0.0))
    loss_class   = sparse_softmax_cross_entropy_with_logits(true_box_class, pred_box_class)#????tied to true_box_class argmax issue
    loss_class   = sum(loss_class .* class_mask) / (nb_class_box + 1e-6)
    return loss_class
end

"""
== INPUT ==
true_box_xy,pred_box_xy, true_box_wh and pred_box_wh must have the same shape length

p1 : pred_mins = (px1,py1)
p2 : pred_maxs = (px2,py2)
t1 : true_mins = (tx1,ty1)
t2 : true_maxs = (tx2,ty2)
             p1______________________
             |      t1___________   |
             |       |           |  |
             |_______|___________|__|p2
                     |           |rmax
                     |___________|
                                  t2
intersect_mins : rmin = t1  = (tx1,ty1)
intersect_maxs : rmax = (rmaxx,rmaxy)
intersect_wh   : (rmaxx - tx1, rmaxy - ty1)
"""
function get_intersect_area(true_box_xy, true_box_wh, pred_box_xy, pred_box_wh)

    true_box_wh_half = true_box_wh ./ 2.0
    true_mins    = true_box_xy .- true_box_wh_half
    true_maxes   = true_box_xy .+ true_box_wh_half

    pred_box_wh_half = pred_box_wh ./ 2.0
    pred_mins    = pred_box_xy .- pred_box_wh_half
    pred_maxes   = pred_box_xy .+ pred_box_wh_half

    intersect_mins  = max.(pred_mins,  true_mins)
    intersect_maxes = min.(pred_maxes, true_maxes)
    intersect_wh    = max.(intersect_maxes .- intersect_mins, 0.0)
    if ndims(true_box_xy) == 5
        intersect_areas = intersect_wh[:,:,:,:, 1] .* intersect_wh[:,:,:,:, 2]
        true_areas = true_box_wh[:,:,:,:, 1] .* true_box_wh[:,:,:,:, 2]
        pred_areas = pred_box_wh[:,:,:,:, 1] .* pred_box_wh[:,:,:,:, 2]
    elseif ndims(true_box_xy) == 6
        intersect_areas = intersect_wh[:,:,:,:,:, 1] .* intersect_wh[:,:,:,:,:, 2]
        true_areas = true_box_wh[:,:,:,:,:, 1] .* true_box_wh[:,:,:,:,:, 2]
        pred_areas = pred_box_wh[:,:,:,:,:, 1] .* pred_box_wh[:,:,:,:,:, 2]
    end
    union_areas = pred_areas .+ true_areas .- intersect_areas
    iou_scores  = intersect_areas ./ union_areas
    return iou_scores
end

"""
== input ==

true_box_conf : tensor of shape (N batch, N grid h, N grid w, N anchor )
true_box_xy   : tensor of shape (N batch, N grid h, N grid w, N anchor , 2)
true_box_wh   : tensor of shape (N batch, N grid h, N grid w, N anchor , 2)
pred_box_xy   : tensor of shape (N batch, N grid h, N grid w, N anchor , 2)
pred_box_wh   : tensor of shape (N batch, N grid h, N grid w, N anchor , 2)

== output ==

true_box_conf : tensor of shape (N batch, N grid h, N grid w, N anchor)

true_box_conf value depends on the predicted values
true_box_conf = IOU_{true,pred} if objecte exist in this anchor else 0
"""
function calc_IOU_pred_true_assigned(true_box_conf,true_box_xy, true_box_wh,pred_box_xy,  pred_box_wh)
    iou_scores = get_intersect_area(true_box_xy,true_box_wh,pred_box_xy,pred_box_wh)
    true_box_conf_IOU = iou_scores .* true_box_conf
    return true_box_conf_IOU
end

"""
```
== input ==
pred_box_xy : tensor of shape (N batch, N grid h, N grid w, N anchor, 2)
pred_box_wh : tensor of shape (N batch, N grid h, N grid w, N anchor, 2)
true_boxes  : tensor of shape (N batch, N grid h, N grid w, N anchor, 2)
== output ==
best_ious
for each iframe,
best_ious[iframe,igridy,igridx,ianchor] contains
the IOU of the object that is most likely included (or best fitted)
within the bounded box recorded in (grid_cell, anchor) pair
NOTE: a same object may be contained in multiple (grid_cell, anchor) pair
from best_ious, you cannot tell how may actual objects are captured as the "best" object
```
"""
function calc_IOU_pred_true_best(pred_box_xy,pred_box_wh,true_boxes)
    true_box_xy = true_boxes[:,:,:,:,:, 1:2]           # (N batch, 1, 1, 1, TRUE_BOX_BUFFER, 2)
    true_box_wh = true_boxes[:,:,:,:,:, 3:4]           # (N batch, 1, 1, 1, TRUE_BOX_BUFFER, 2)

    pred_box_xy = expand_dims(pred_box_xy, 5) # (N batch, N grid_h, N grid_w, N anchor, 1, 2)
    pred_box_wh = expand_dims(pred_box_wh, 5) # (N batch, N grid_h, N grid_w, N anchor, 1, 2)

    iou_scores  =  get_intersect_area(true_box_xy,true_box_wh,pred_box_xy,pred_box_wh) # (N batch, N grid_h, N grid_w, N anchor, 50)

    best_ious = dropdims(maximum(iou_scores, dims=5);dims=5 ) # (N batch, N grid_h, N grid_w, N anchor)
    return best_ious
end

"""
```
== input ==
best_ious           : tensor of shape (Nbatch, N grid h, N grid w, N anchor)
true_box_conf       : tensor of shape (Nbatch, N grid h, N grid w, N anchor)
true_box_conf_IOU   : tensor of shape (Nbatch, N grid h, N grid w, N anchor)
LAMBDA_NO_OBJECT    : 1.0
LAMBDA_OBJECT       : 5.0
== output ==
conf_mask : tensor of shape (Nbatch, N grid h, N grid w, N anchor)
conf_mask[iframe, igridy, igridx, ianchor] = 0
when there is no object assigned in (grid cell, anchor) pair and the region seems useless i.e.
y_true[iframe,igridx,igridy,4] = 0 "and" the predicted region has no object that has IoU > 0.6
conf_mask[iframe, igridy, igridx, ianchor] =  NO_OBJECT_SCALE
when there is no object assigned in (grid cell, anchor) pair but region seems to include some object
y_true[iframe,igridx,igridy,4] = 0 "and" the predicted region has some object that has IoU > 0.6
conf_mask[iframe, igridy, igridx, ianchor] =  OBJECT_SCALE
when there is an object in (grid cell, anchor) pair
```
"""
function get_conf_mask(best_ious, true_box_conf, true_box_conf_IOU,LAMBDA_NO_OBJECT, LAMBDA_OBJECT)
    conf_mask = Float32.(best_ious .< 0.6) .* (1 .- true_box_conf) .* LAMBDA_NO_OBJECT
    # penalize the confidence of the boxes, which are reponsible for corresponding ground truth box
    conf_mask = conf_mask .+ true_box_conf_IOU .* LAMBDA_OBJECT
    return conf_mask
end

"""
```
== input ==
conf_mask         : tensor of shape (Nbatch, N grid h, N grid w, N anchor)
true_box_conf_IOU : tensor of shape (Nbatch, N grid h, N grid w, N anchor)
pred_box_conf     : tensor of shape (Nbatch, N grid h, N grid w, N anchor)
```
"""
function calc_loss_conf(conf_mask,true_box_conf_IOU, pred_box_conf)
    # the number of (grid cell, anchor) pair that has an assigned object or
    # that has no assigned object but some objects may be in bounding box.
    # N conf
    nb_conf_box  = sum(Float32.(conf_mask .> 0.0))
    loss_conf    = sum((true_box_conf_IOU .- pred_box_conf).^2 .* conf_mask)  / (nb_conf_box  + 1e-6) / 2.0
    return loss_conf
end

"""
```
y_true : (N batch, N grid h, N grid w, N anchor, 4 + 1 + N classes)
y_true[irow, i_gridh, i_gridw, i_anchor, :4] = center_x, center_y, w, h
center_x : The x coordinate center of the bounding box.
Rescaled to range between 0 and N gird  w (e.g., ranging between [0,13)
center_y : The y coordinate center of the bounding box.
Rescaled to range between 0 and N gird  h (e.g., ranging between [0,13)
w        : The width of the bounding box.
Rescaled to range between 0 and N gird  w (e.g., ranging between [0,13)
h        : The height of the bounding box.
Rescaled to range between 0 and N gird  h (e.g., ranging between [0,13)
y_true[irow, i_gridh, i_gridw, i_anchor, 4] = ground truth confidence
ground truth confidence is 1 if object exists in this (anchor box, gird cell) pair
y_true[irow, i_gridh, i_gridw, i_anchor, 5 + iclass] = 1 if the object is in category <iclass> else 0
=====================================================
tensor that connect to the YOLO model's hack input
=====================================================
true_boxes
=========================================
training parameters specification example
=========================================
gridsize           = (13,13)
batchsize          = 34
anchors = np.array([1.07709888,  1.78171903,  # anchor box 1, width , height
2.71054693,  5.12469308,  # anchor box 2, width,  height
10.47181473, 10.09646365,  # anchor box 3, width,  height
5.48531347,  8.11011331]) # anchor box 4, width,  height
λno_object = 1.0
λobject    = 5.0
λcoord     = 1.0
λclass     = 1.0
```
"""
function custom_loss_core(y_true,y_pred,true_boxes,GRIDSIZE::Tuple{Int,Int},
        BATCH_SIZE::Int,ANCHORS,LAMBDA_COORD,LAMBDA_CLASS,LAMBDA_NO_OBJECT, LAMBDA_OBJECT)
    nboxes = Int64(length(ANCHORS)/2)
    # Step 1: Adjust prediction output
    cell_grid   = get_cell_grid(GRIDSIZE,BATCH_SIZE,nboxes)
    pred_box_xy, pred_box_wh, pred_box_conf, pred_box_class = adjust_scale_prediction(y_pred,cell_grid,ANCHORS)
    # Step 2: Extract ground truth output
    true_box_xy, true_box_wh, true_box_conf, true_box_class = extract_ground_truth(y_true::Array{Float32,5})
    # Step 3: Calculate loss for the bounding box parameters
    loss_xywh, coord_mask = calc_loss_xywh(true_box_conf,LAMBDA_COORD,true_box_xy, pred_box_xy,true_box_wh,pred_box_wh)
    # Step 4: Calculate loss for the class probabilities
    loss_class  = calc_loss_class(true_box_conf,LAMBDA_CLASS,true_box_class,pred_box_class)
    # Step 5: For each (grid cell, anchor) pair,
    #         calculate the IoU between predicted and ground truth bounding box
    true_box_conf_IOU =
        calc_IOU_pred_true_assigned(true_box_conf,true_box_xy, true_box_wh,pred_box_xy,  pred_box_wh)
    # Step 6: For each predicted bounded box from (grid cell, anchor box),
    #         calculate the best IOU, regardless of the ground truth anchor box that each object gets assigned.
    best_ious = calc_IOU_pred_true_best(pred_box_xy,pred_box_wh,true_boxes)
    # Step 7: For each grid cell, calculate the L_{i,j}^{noobj}
    conf_mask = get_conf_mask(best_ious, true_box_conf, true_box_conf_IOU,LAMBDA_NO_OBJECT, LAMBDA_OBJECT)
    # Step 8: Calculate loss for the confidence
    loss_conf = calc_loss_conf(conf_mask,true_box_conf_IOU, pred_box_conf)

    loss = loss_xywh + loss_conf + loss_class
    return loss
end