using .NumericalFluxes: CentralHyperDiffusiveFlux, CentralDivPenalty

struct DGModel{BL, G, NFND, NFD, GNF, AS, DS, HDS, D, DD, MD}
    balancelaw::BL
    grid::G
    numfluxnondiff::NFND
    numfluxdiff::NFD
    gradnumflux::GNF
    auxstate::AS
    diffstate::DS
    hyperdiffstate::HDS
    direction::D
    diffusion_direction::DD
    modeldata::MD
end
function DGModel(
    balancelaw,
    grid,
    numfluxnondiff,
    numfluxdiff,
    gradnumflux;
    auxstate = create_auxstate(balancelaw, grid),
    diffstate = create_diffstate(balancelaw, grid),
    hyperdiffstate = create_hyperdiffstate(balancelaw, grid),
    direction = EveryDirection(),
    diffusion_direction = direction,
    modeldata = nothing,
)
    DGModel(
        balancelaw,
        grid,
        numfluxnondiff,
        numfluxdiff,
        gradnumflux,
        auxstate,
        diffstate,
        hyperdiffstate,
        direction,
        diffusion_direction,
        modeldata,
    )
end

function (dg::DGModel)(dQdt, Q, ::Nothing, t; increment = false)

    bl = dg.balancelaw
    device = typeof(Q.data) <: Array ? CPU() : CUDA()

    grid = dg.grid
    topology = grid.topology

    dim = dimensionality(grid)
    N = polynomialorder(grid)
    Nq = N + 1
    Nqk = dim == 2 ? 1 : Nq
    Nfp = Nq * Nqk
    nrealelem = length(topology.realelems)

    Qvisc = dg.diffstate
    Qhypervisc_grad, Qhypervisc_div = dg.hyperdiffstate
    auxstate = dg.auxstate

    FT = eltype(Q)
    nviscstate = num_diffusive(bl, FT)
    nhyperviscstate = num_hyperdiffusive(bl, FT)

    Np = dofs_per_element(grid)

    workgroups_volume = (Nq, Nq, Nqk)
    ndrange_volume = (nrealelem * Nq, Nq, Nqk)
    workgroups_surface = Nfp
    ndrange_surface = Nfp * nrealelem

    communicate =
        !(isstacked(topology) && typeof(dg.direction) <: VerticalDirection)

    aux_comm = update_aux!(dg, bl, Q, t)
    @assert typeof(aux_comm) == Bool

    if nhyperviscstate > 0
        hypervisc_indexmap = create_hypervisc_indexmap(bl)
    else
        hypervisc_indexmap = nothing
    end

    ########################
    # Gradient Computation #
    ########################
    if communicate
        MPIStateArrays.start_ghost_exchange!(Q)
        if aux_comm
            MPIStateArrays.start_ghost_exchange!(auxstate)
        end
    end

    if nviscstate > 0 || nhyperviscstate > 0

        event = Event(device)
        event = volumeviscterms!(device, workgroups_volume)(
            bl,
            Val(dim),
            Val(N),
            dg.diffusion_direction,
            Q.data,
            Qvisc.data,
            Qhypervisc_grad.data,
            auxstate.data,
            grid.vgeo,
            t,
            grid.D,
            hypervisc_indexmap,
            topology.realelems,
            ndrange = ndrange_volume,
            dependencies = (event,),
        )
        wait(device, event)

        if communicate
            MPIStateArrays.finish_ghost_recv!(Q)
            if aux_comm
                MPIStateArrays.finish_ghost_recv!(auxstate)
            end
        end

        event = Event(device)
        event = faceviscterms!(device, workgroups_surface)(
            bl,
            Val(dim),
            Val(N),
            dg.diffusion_direction,
            dg.gradnumflux,
            Q.data,
            Qvisc.data,
            Qhypervisc_grad.data,
            auxstate.data,
            grid.vgeo,
            grid.sgeo,
            t,
            grid.vmap⁻,
            grid.vmap⁺,
            grid.elemtobndy,
            hypervisc_indexmap,
            topology.realelems;
            ndrange = ndrange_surface,
            dependencies = (event,),
        )
        wait(device, event)

        if communicate
            nviscstate > 0 && MPIStateArrays.start_ghost_exchange!(Qvisc)
            nhyperviscstate > 0 &&
            MPIStateArrays.start_ghost_exchange!(Qhypervisc_grad)
        end

        if nviscstate > 0
            aux_comm = update_aux_diffusive!(dg, bl, Q, t)
            @assert typeof(aux_comm) == Bool
        end

        if aux_comm
            MPIStateArrays.start_ghost_exchange!(auxstate)
        end
    end

    if nhyperviscstate > 0
        #########################
        # Laplacian Computation #
        #########################

        event = Event(device)
        event = volumedivgrad!(device, workgroups_volume)(
            bl,
            Val(dim),
            Val(N),
            dg.diffusion_direction,
            Qhypervisc_grad.data,
            Qhypervisc_div.data,
            grid.vgeo,
            grid.D,
            topology.realelems;
            ndrange = ndrange_volume,
            dependencies = (event,),
        )
        wait(device, event)

        communicate && MPIStateArrays.finish_ghost_recv!(Qhypervisc_grad)

        event = Event(device)
        event = facedivgrad!(device, workgroups_surface)(
            bl,
            Val(dim),
            Val(N),
            dg.diffusion_direction,
            CentralDivPenalty(),
            Qhypervisc_grad.data,
            Qhypervisc_div.data,
            grid.vgeo,
            grid.sgeo,
            grid.vmap⁻,
            grid.vmap⁺,
            grid.elemtobndy,
            topology.realelems;
            ndrange = ndrange_surface,
            dependencies = (event,),
        )
        wait(device, event)

        communicate && MPIStateArrays.start_ghost_exchange!(Qhypervisc_div)

        ####################################
        # Hyperdiffusive terms computation #
        ####################################

        event = Event(device)
        event = volumehyperviscterms!(device, workgroups_volume)(
            bl,
            Val(dim),
            Val(N),
            dg.diffusion_direction,
            Qhypervisc_grad.data,
            Qhypervisc_div.data,
            Q.data,
            auxstate.data,
            grid.vgeo,
            grid.ω,
            grid.D,
            topology.realelems,
            t;
            ndrange = ndrange_volume,
            dependencies = (event,),
        )
        wait(device, event)

        communicate && MPIStateArrays.finish_ghost_recv!(Qhypervisc_div)

        event = Event(device)
        event = facehyperviscterms!(device, workgroups_surface)(
            bl,
            Val(dim),
            Val(N),
            dg.diffusion_direction,
            CentralHyperDiffusiveFlux(),
            Qhypervisc_grad.data,
            Qhypervisc_div.data,
            Q.data,
            auxstate.data,
            grid.vgeo,
            grid.sgeo,
            grid.vmap⁻,
            grid.vmap⁺,
            grid.elemtobndy,
            topology.realelems,
            t;
            ndrange = ndrange_surface,
            dependencies = (event,),
        )
        wait(device, event)

        communicate && MPIStateArrays.start_ghost_exchange!(Qhypervisc_grad)
    end

    ###################
    # RHS Computation #
    ###################
    event = Event(device)
    event = volumerhs!(device, workgroups_volume)(
        bl,
        Val(dim),
        Val(N),
        dg.direction,
        dQdt.data,
        Q.data,
        Qvisc.data,
        Qhypervisc_grad.data,
        auxstate.data,
        grid.vgeo,
        t,
        grid.ω,
        grid.D,
        topology.realelems,
        increment;
        ndrange = ndrange_volume,
        dependencies = (event,),
    )
    wait(device, event)

    if communicate
        if nviscstate > 0 || nhyperviscstate > 0
            if nviscstate > 0
                MPIStateArrays.finish_ghost_recv!(Qvisc)
                if aux_comm
                    MPIStateArrays.finish_ghost_recv!(auxstate)
                end
            end
            nhyperviscstate > 0 &&
            MPIStateArrays.finish_ghost_recv!(Qhypervisc_grad)
        else
            MPIStateArrays.finish_ghost_recv!(Q)
            if aux_comm
                MPIStateArrays.finish_ghost_recv!(auxstate)
            end
        end
    end

    event = Event(device)
    event = facerhs!(device, workgroups_surface)(
        bl,
        Val(dim),
        Val(N),
        dg.direction,
        dg.numfluxnondiff,
        dg.numfluxdiff,
        dQdt.data,
        Q.data,
        Qvisc.data,
        Qhypervisc_grad.data,
        auxstate.data,
        grid.vgeo,
        grid.sgeo,
        t,
        grid.vmap⁻,
        grid.vmap⁺,
        grid.elemtobndy,
        topology.realelems;
        ndrange = ndrange_surface,
        dependencies = (event,),
    )
    wait(device, event)

    # Just to be safe, we wait on the sends we started.
    if communicate
        MPIStateArrays.finish_ghost_send!(Qhypervisc_div)
        MPIStateArrays.finish_ghost_send!(Qvisc)
        MPIStateArrays.finish_ghost_send!(Qhypervisc_grad)
        MPIStateArrays.finish_ghost_send!(Q)
    end
end

function init_ode_state(
    dg::DGModel,
    args...;
    init_on_cpu = false,
    commtag = 888,
)
    device = arraytype(dg.grid) <: Array ? CPU() : CUDA()

    bl = dg.balancelaw
    grid = dg.grid

    state = create_state(bl, grid, commtag)

    topology = grid.topology
    Np = dofs_per_element(grid)

    auxstate = dg.auxstate
    dim = dimensionality(grid)
    N = polynomialorder(grid)
    nrealelem = length(topology.realelems)

    if !init_on_cpu
        event = Event(device)
        event = initstate!(device, Np)(
            bl,
            Val(dim),
            Val(N),
            state.data,
            auxstate.data,
            grid.vgeo,
            topology.realelems,
            args...;
            ndrange = Np * nrealelem,
            dependencies = (event,),
        )
        wait(device, event)
    else
        h_state = similar(state, Array)
        h_auxstate = similar(auxstate, Array)
        h_auxstate .= auxstate
        event = initstate!(CPU(), Np)(
            bl,
            Val(dim),
            Val(N),
            h_state.data,
            h_auxstate.data,
            Array(grid.vgeo),
            topology.realelems,
            args...;
            ndrange = Np * nrealelem,
        )
        wait(event) # XXX: This could be `wait(device, event)` once KA supports that.
        state .= h_state
    end

    MPIStateArrays.start_ghost_exchange!(state)
    MPIStateArrays.finish_ghost_exchange!(state)

    return state
end

# fallback
function update_aux!(dg::DGModel, bl::BalanceLaw, Q::MPIStateArray, t::Real)
    return false
end

function update_aux_diffusive!(
    dg::DGModel,
    bl::BalanceLaw,
    Q::MPIStateArray,
    t::Real,
)
    return false
end

function indefinite_stack_integral!(
    dg::DGModel,
    m::BalanceLaw,
    Q::MPIStateArray,
    auxstate::MPIStateArray,
    t::Real,
)

    device = typeof(Q.data) <: Array ? CPU() : CUDA()

    grid = dg.grid
    topology = grid.topology

    dim = dimensionality(grid)
    N = polynomialorder(grid)
    Nq = N + 1
    Nqk = dim == 2 ? 1 : Nq

    FT = eltype(Q)

    # do integrals
    nelem = length(topology.elems)
    nvertelem = topology.stacksize
    nhorzelem = div(nelem, nvertelem)

    event = Event(device)
    event = knl_indefinite_stack_integral!(device, (Nq, Nqk))(
        m,
        Val(dim),
        Val(N),
        Val(nvertelem),
        Q.data,
        auxstate.data,
        grid.vgeo,
        grid.Imat,
        1:nhorzelem;
        ndrange = (nhorzelem * Nq, Nqk),
        dependencies = (event,),
    )
    wait(device, event)
end

function reverse_indefinite_stack_integral!(
    dg::DGModel,
    m::BalanceLaw,
    Q::MPIStateArray,
    auxstate::MPIStateArray,
    t::Real,
)

    device = typeof(auxstate.data) <: Array ? CPU() : CUDA()

    grid = dg.grid
    topology = grid.topology

    dim = dimensionality(grid)
    N = polynomialorder(grid)
    Nq = N + 1
    Nqk = dim == 2 ? 1 : Nq

    FT = eltype(auxstate)

    # do integrals
    nelem = length(topology.elems)
    nvertelem = topology.stacksize
    nhorzelem = div(nelem, nvertelem)

    event = Event(device)
    event = knl_reverse_indefinite_stack_integral!(device, (Nq, Nqk))(
        m,
        Val(dim),
        Val(N),
        Val(nvertelem),
        Q.data,
        auxstate.data,
        1:nhorzelem;
        ndrange = (nhorzelem * Nq, Nqk),
        dependencies = (event,),
    )
    wait(device, event)
end

function nodal_update_aux!(
    f!,
    dg::DGModel,
    m::BalanceLaw,
    Q::MPIStateArray,
    t::Real;
    diffusive = false,
)
    device = typeof(Q.data) <: Array ? CPU() : CUDA()

    grid = dg.grid
    topology = grid.topology

    dim = dimensionality(grid)
    N = polynomialorder(grid)
    Nq = N + 1
    nrealelem = length(topology.realelems)

    Np = dofs_per_element(grid)

    nodal_update_aux! = knl_nodal_update_aux!(device, Np)
    ### update aux variables
    event = Event(device)
    if diffusive
        event = nodal_update_aux!(
            m,
            Val(dim),
            Val(N),
            f!,
            Q.data,
            dg.auxstate.data,
            dg.diffstate.data,
            t,
            topology.realelems;
            ndrange = Np * nrealelem,
            dependencies = (event,),
        )
    else
        event = nodal_update_aux!(
            m,
            Val(dim),
            Val(N),
            f!,
            Q.data,
            dg.auxstate.data,
            t,
            topology.realelems;
            ndrange = Np * nrealelem,
            dependencies = (event,),
        )
    end
    wait(device, event)
end

"""
    courant(local_courant::Function, dg::DGModel, m::BalanceLaw,
            Q::MPIStateArray, direction=EveryDirection())
Returns the maximum of the evaluation of the function `local_courant`
pointwise throughout the domain.  The function `local_courant` is given an
approximation of the local node distance `Δx`.  The `direction` controls which
reference directions are considered when computing the minimum node distance
`Δx`.
An example `local_courant` function is
    function local_courant(m::AtmosModel, state::Vars, aux::Vars,
                           diffusive::Vars, Δx)
      return Δt * cmax / Δx
    end
where `Δt` is the time step size and `cmax` is the maximum flow speed in the
model.
"""
function courant(
    local_courant::Function,
    dg::DGModel,
    m::BalanceLaw,
    Q::MPIStateArray,
    Δt,
    simtime,
    direction = EveryDirection(),
)
    grid = dg.grid
    topology = grid.topology
    nrealelem = length(topology.realelems)

    if nrealelem > 0
        N = polynomialorder(grid)
        dim = dimensionality(grid)
        Nq = N + 1
        Nqk = dim == 2 ? 1 : Nq
        device = grid.vgeo isa Array ? CPU() : CUDA()
        pointwise_courant = similar(grid.vgeo, Nq^dim, nrealelem)
        event = Event(device)
        event = Grids.knl_min_neighbor_distance!(device, (Nq, Nq, Nqk))(
            Val(N),
            Val(dim),
            direction,
            pointwise_courant,
            grid.vgeo,
            topology.realelems;
            ndrange = (nrealelem * Nq, Nq, Nqk),
            dependencies = (event,),
        )
        event = knl_local_courant!(device, Nq * Nq * Nqk)(
            m,
            Val(dim),
            Val(N),
            pointwise_courant,
            local_courant,
            Q.data,
            dg.auxstate.data,
            dg.diffstate.data,
            topology.realelems,
            Δt,
            simtime,
            direction;
            ndrange = nrealelem * Nq * Nq * Nqk,
            dependencies = (event,),
        )
        wait(device, event)
        rank_courant_max = maximum(pointwise_courant)
    else
        rank_courant_max = typemin(eltype(Q))
    end

    MPI.Allreduce(rank_courant_max, max, topology.mpicomm)
end

function copy_stack_field_down!(
    dg::DGModel,
    m::BalanceLaw,
    auxstate::MPIStateArray,
    fldin,
    fldout,
)

    device = typeof(auxstate.data) <: Array ? CPU() : CUDA()

    grid = dg.grid
    topology = grid.topology

    dim = dimensionality(grid)
    N = polynomialorder(grid)
    Nq = N + 1
    Nqk = dim == 2 ? 1 : Nq

    # do integrals
    nelem = length(topology.elems)
    nvertelem = topology.stacksize
    nhorzelem = div(nelem, nvertelem)

    event = Event(device)
    event = knl_copy_stack_field_down!(device, (Nq, Nqk))(
        Val(dim),
        Val(N),
        Val(nvertelem),
        auxstate.data,
        1:nhorzelem,
        Val(fldin),
        Val(fldout);
        ndrange = (nhorzelem * Nq, Nqk),
        dependencies = (event,),
    )
    wait(device, event)
end

function MPIStateArrays.MPIStateArray(dg::DGModel, commtag = 888)
    bl = dg.balancelaw
    grid = dg.grid

    state = create_state(bl, grid, commtag)

    return state
end

function create_hypervisc_indexmap(bl::BalanceLaw)
    # helper function
    _getvars(v, ::Type) = v
    function _getvars(v::Vars, ::Type{T}) where {T <: NamedTuple}
        fields = getproperty.(Ref(v), fieldnames(T))
        collect(Iterators.Flatten(_getvars.(fields, fieldtypes(T))))
    end

    gradvars = vars_gradient(bl, Int)
    gradlapvars = vars_gradient_laplacian(bl, Int)
    indices = Vars{gradvars}(1:varsize(gradvars))
    SVector{varsize(gradlapvars)}(_getvars(indices, gradlapvars))
end


using Random
using Distributions
using MultiAgentPathFinding
using MultiAgentAllocationTransit
using BenchmarkTools
using Statistics
using JSON
using Logging
using JLD2
using NearestNeighbors
using StaticArrays
global_logger(SimpleLogger(stderr, Logging.Error))

rng = MersenneTwister(2345)

# Script arguments relating to transit files etc.
const city_params_file = ARGS[1]
const stop_coords_file = ARGS[2]
const trips_file = ARGS[3]
const drone_params_file = ARGS[4]
const bb_params_file = ARGS[5]
const out_file_pref = ARGS[6]
const city_travel_time_estimates = ARGS[7]
const N_DEPOTS = parse(Int64, ARGS[8])
const N_AGENTS = parse(Int64, ARGS[9])
const N_TRIALS = parse(Int64, ARGS[10])

const TRANSIT_CAP_RANGE = (3, 5)
const ECBS_WEIGHT = 1.1

## NOTE: EXAMPLE CALL
# julia scripts/benchmark_replanning.jl data/sfmta/sf_params.toml data/sfmta/stop_to_coords.json data/sfmta/trips.json
# data/drone_params.toml data/sfmta/sf_bb_params.toml data/sf_replan_20trials data/sfmta/sf_halton_tt_estimates.jld2 5 10 20

function main()

    # Load the various parameters
    city_params = parse_city_params(city_params_file)
    bb_params = parse_city_params(bb_params_file)
    drone_params = parse_drone_params(drone_params_file)

    lat_dist = Uniform(bb_params.lat_start, bb_params.lat_end)
    lon_dist = Uniform(bb_params.lon_start, bb_params.lon_end)

    # Transit Graph Preprocessing
    tg = load_transit_graph_latlong(stop_coords_file, trips_file, TRANSIT_CAP_RANGE, rng)
    tg, stop_idx_to_trips = transit_graph_preprocessing(tg, MultiAgentAllocationTransit.distance_lat_lon_euclidean, drone_params)

    # Load Halton stuff
    @load city_travel_time_estimates halton_nn_tree city_halton_points travel_time_estimates

    replan_results = Dict("cap_range"=>TRANSIT_CAP_RANGE, "weight"=>ECBS_WEIGHT,
                        "depots"=>N_DEPOTS, "agents"=>N_AGENTS)

    rp_one_times = Float64[]
    rp_one_msps = Float64[]
    rp_all_times = Float64[]
    rp_all_msps = Float64[]

    # No need to ignore first trial now!
    TRIAL = 1
    while TRIAL <= N_TRIALS+1

        N_SITES = 2*N_AGENTS

        println("Trial $(TRIAL)")

        # Generate depots and sites
        depots = [LatLonCoords((lat = rand(rng, lat_dist), lon = rand(rng, lon_dist))) for i = 1:N_DEPOTS]
        sites = [LatLonCoords((lat = rand(rng, lat_dist), lon = rand(rng, lon_dist))) for i = 1:N_SITES]
        depot_sites = vcat(depots, sites)

        # Off transit preprocessing
        otg = OffTransitGraph(depots = depots, sites = sites)
        aug_trips_fws_dists = augmented_trip_meta_graph_fws_dists(tg, MultiAgentAllocationTransit.distance_lat_lon_euclidean,
                                                              length(depots), length(sites),
                                                              vcat(depots, sites),
                                                              drone_params)
        state_graph, depot_sites_to_vtx, trip_to_vtx_range = setup_state_graph(tg, otg)

        # Set up the shadow env to do task allocation with
        env = MAPFTransitEnv(off_transit_graph = otg, transit_graph = tg, state_graph = state_graph,
                             agent_states = AgentState[], depot_sites_to_vtx = depot_sites_to_vtx, trip_to_vtx_range = trip_to_vtx_range,
                             stop_idx_to_trips = stop_idx_to_trips, aug_trips_fws_dists = aug_trips_fws_dists,
                             drone_params = drone_params, dist_fn = MultiAgentAllocationTransit.distance_lat_lon_euclidean,
                             curr_site_points = [], threshold_global_conflicts = 10)

        # run the task allocation, obtain the agent tasks and true number of agents
        cost_fn(i, j) = allocation_cost_wrapper_estimate(env, ECBS_WEIGHT, N_DEPOTS, N_SITES,
                                                        halton_nn_tree, travel_time_estimates, i, j)
        agent_tours = task_allocation(N_DEPOTS, N_SITES, N_AGENTS,
                                      depot_sites, cost_fn)

        println("Task Allocation Done!")

        agent_tasks = get_agent_task_set(agent_tours, N_DEPOTS, N_SITES)
        true_n_agents = length(agent_tasks)

        # Now set n_agent dependent env params
        env.agent_states = [AgentState(task=agt_task) for agt_task in agent_tasks]
        env.curr_site_points = zeros(Int64, true_n_agents)
        env.plan_ref_times = zeros(true_n_agents)

        # Setup ECBS Solver; Run search
        initial_states = Vector{MAPFTransitVertexState}(undef, true_n_agents)
        for i = 1:true_n_agents
            initial_states[i] = env.state_graph.vertices[env.depot_sites_to_vtx[string("d-", agent_tasks[i].origin)]]
        end

        solver = ECBSSolver{MAPFTransitVertexState,MAPFTransitAction,Float64,Makespan,MAPFTransitConflict,MAPFTransitConstraints,MAPFTransitEnv}(env = env, weight = ECBS_WEIGHT)

        solution = search!(solver, initial_states)

        println("Initial solution obtained!")

        # Compute env diagnostics and throw away if invalid
        set_solution_diagnostics!(env, solution)
        n_valid_firstpaths = length(env.valid_path_dists)

        if n_valid_firstpaths > 0.67 * true_n_agents

            # Copy env and solution for replanning
            env_copy = deepcopy(env)
            solution_copy = deepcopy(solution)

            # Actually run the replanning
            did_replan_indiv, el_time_indiv = replan_individual!(env, solution, N_DEPOTS, N_SITES, agent_tours, ECBS_WEIGHT)
            did_replan_collec, el_time_collec, new_soln_collec = replan_collective!(env_copy, solution_copy, N_DEPOTS, N_SITES, agent_tours, ECBS_WEIGHT)

            # Enter only if both valid
            if did_replan_indiv && did_replan_collec

                println("Replanning done successfully!")

                # Now look at times and makespans
                msp_one = maximum([s.cost for s in solution])
                msp_all = maximum([s.cost for s in new_soln_collec])

                push!(rp_one_times, el_time_indiv)
                push!(rp_one_msps, msp_one)

                push!(rp_all_times, el_time_collec)
                push!(rp_all_msps, msp_all)

                # Increment trial
                TRIAL = TRIAL + 1

            else

                println("No replan! Indiv - $(did_replan_indiv); Collec - $(did_replan_collec)")

            end
        end
    end

    replan_results["results"] = Dict("replan_one_times" => rp_one_times,
                                     "replan_one_makespans" => rp_one_msps,
                                     "replan_all_times" => rp_all_times,
                                     "replan_all_makespans" => rp_all_msps)

    out_file = string(out_file_pref, "_", N_DEPOTS, "deps_", N_AGENTS, "_agts.json")
    open(out_file, "w") do f
        JSON.print(f, replan_results, 2)
    end
end

main()


module CSDP

using LinearAlgebra # For Diagonal
using SparseArrays # For SparseMatrixCSC

if VERSION < v"1.3"
    if isfile(joinpath(dirname(@__FILE__),"..","deps","deps.jl"))
        include("../deps/deps.jl")
    else
        error("CSDP not properly installed. Please run Pkg.build(\"CSDP\")")
    end
else
    import CSDP_jll: libcsdp
end

# This is the size of int used by the LAPACK library used by CSDP.
# If libcsdp is patched to use a 64-bit integer LAPACK, this should be replaced by `Clong`.
const CSDP_INT = Cint

export Blockmatrix

include("blockdiag.jl")
include("blockmat.h.jl")
include("blockmat.jl")
include("declarations.h.jl")
include("declarations.jl")
include("debug-mat.jl")
include("options.jl")
include("MOI_wrapper.jl")

end # module


const BOUND_LOWER = UInt8(1)
const BOUND_UPPER = UInt8(2)
const BOUND_EXACT = UInt8(3)

import Base.hash

# Below we override the base hashing method as we already generated a hash!
# Improves performance for larger depth searches where TT calls are numerous.
Base.hash(x::UInt32) = x
Base.hash(x::UInt16) = x

mutable struct TT_Entry
    eval::Int16
    move::Move
    depth::UInt8
    bound::UInt8
    hash16::UInt16
end
const NO_ENTRY = TT_Entry(0, Move(), 0, 0, 0)

mutable struct TT_Table
    table::Dict{UInt32, TT_Entry}
    hashmask::UInt64
end
TT_Table() = TT_Table(16)

#=
Chess engines usually face memory restrictions, as we can't possibly hash and store all positions.
The below function aims to cap the hash memory size, by restricting the keys in the dict.
=#
function TT_Table(size_MB::Int)
    bytes = size_MB << 20

    # work out size of slots in TT
    #slotsize = Base.summarysize(TT_Entry())
    slotsize = 14 # At present...
    maxslots = bytes / slotsize

    size = 1
    while size <= maxslots
        size <<= 1
    end
    # We shift right again as we "overshifted" by one in the while loop.
    size >>= 1
    size -= 1

    TT_Table(Dict{UInt32, TT_Entry}(), UInt64(size))
end

function getTTentry(tt::TT_Table, hash::ZobristHash)
    res = get(tt.table, UInt32(hash.hash & tt.hashmask), NO_ENTRY)
    if (res !== NO_ENTRY) && (res.hash16 === UInt16(hash.hash >> 48))
        return res
    else
        return NO_ENTRY
    end
end

function setTTentry!(tt::TT_Table, hash::ZobristHash, eval::Integer, move::Move, depth::Integer, bound::UInt8)
    hash16 = UInt16(hash.hash >> 48)
    entry = TT_Entry(eval, move, depth, bound, hash16)
    tt.table[UInt32(hash.hash & tt.hashmask)] = entry
end

function ttvalue(tt_entry::TT_Entry, ply::Int)
    if tt_entry.eval >= (MATE - MAX_PLY)
        return tt_entry.eval + ply
    elseif tt_entry.eval <= (MAX_PLY - MATE)
        return tt_entry.eval - ply
    else
        return tt_entry.eval
    end
end

function hashfull(tt::TT_Table)
    full = 0
    for i in 1:3000
        if haskey(tt.table, UInt32(i))
            full += 1
        end
    end
    fld(full,3)
end

function getPKentry(pktable::PawnKingTable, pkhash::ZobristHash)
    key16 = UInt16(pkhash.hash >> 48)
    if (res = get(pktable, key16, PKT_BLANK)) !== PKT_BLANK
        if res.pkhash === pkhash
            return res
        end
    end
    return PKT_BLANK
end

const PKT_BLANK = PKT_Entry(ZobristHash(0), EMPTY, 0)

function storePKentry!(pktable::PawnKingTable, pkhash::ZobristHash, passed::Bitboard, eval::Int)
    key16 = UInt16(pkhash.hash >> 48)
    pktable[key16] = PKT_Entry(pkhash, passed, eval)
end


"""

```
contrib(s::Selection, atom_contributions::Array{Float64}, selection)
```

Extract the contribution of a given atom type selection from the
solute or solvent atomic contributions to the MDDF.

`s` here is the solute or solvent selection (type `ComplexMixtures.Selection`)
`atom_contributions` is the `R.solute_atom` or `R.solvent_atom` arrays of the `Result` structure,
and the last argument is the selection of atoms from the solute to be considered, given as a list
of indexes, list of atom names, vector of `PDBTools.Atom`s, or a `PDBTools.Residue`.

"""
function contrib(s::Selection, atom_contributions::Array{Float64}, indexes::Vector{Int})
  nbins = size(atom_contributions,1)
  c = zeros(nbins)
  for it in indexes
    if it > s.natomspermol
      error("The index list contains atoms with indexes greater than the number of atoms of the molecule.")
    end
    c += atom_contributions[:,it]
  end
  return c
end

#
# If a list of atom names is provided
#
function contrib(s::Selection, atom_contributions::Array{Float64}, names::Vector{String})
  indexes = Vector{Int}(undef,0)
  for name in names
    index = findall(isequal(name), s.names)
    if length(index) == 0
      error(" Atom in input list is not part of solvent (or solute): $name")
    end
    append!(indexes,index)
  end
  return contrib(s, atom_contributions, indexes)
end

#
# If a list of atoms of PDBTools.Atom is provided
#
function contrib(s::Selection,
                 atom_contributions::Array{Float64},
                 atoms::Vector{PDBTools.Atom};
                 warning=true)
  (warning && s.nmols > 1) && warning_nmols_types()
  indexes = [ atom.index for atom in atoms ]
  # Check which types of atoms belong to this selection
  selected_types = which_types(s, indexes, warning=warning)
  return contrib(s, atom_contributions, selected_types)
end

#
# If a residue of type PDBTools.Residue is provided
#
function contrib(s::Selection,
                 atom_contributions::Array{Float64},
                 residue::Residue;
                 warning=true)
  (warning && s.nmols > 1) && warning_nmols_types()
  indexes = collect(residue.range)
  # Check which types of atoms belong to this selection
  selected_types = which_types(s, indexes, warning=warning)
  return contrib(s, atom_contributions, selected_types)
end

function warning_nmols_types()
  println("WARNING: there is more than one molecule in this selection. "*
          "Contributions are summed over all atoms of the same type.")
end


# This file is a part of Julia. License is MIT: https://julialang.org/license

abstract type AbstractTime end

"""
    Period
    Year
    Month
    Week
    Day
    Hour
    Minute
    Second
    Millisecond
    Microsecond
    Nanosecond

`Period` types represent discrete, human representations of time.
"""
abstract type Period     <: AbstractTime end
abstract type DatePeriod <: Period end
abstract type TimePeriod <: Period end

for T in (:Year, :Month, :Week, :Day)
    @eval struct $T <: DatePeriod
        value::Int64
        $T(v::Number) = new(v)
    end
end
for T in (:Hour, :Minute, :Second, :Millisecond, :Microsecond, :Nanosecond)
    @eval struct $T <: TimePeriod
        value::Int64
        $T(v::Number) = new(v)
    end
end

"""
    Year(v)
    Month(v)
    Week(v)
    Day(v)
    Hour(v)
    Minute(v)
    Second(v)
    Millisecond(v)
    Microsecond(v)
    Nanosecond(v)

Construct a `Period` type with the given `v` value. Input must be losslessly convertible
to an [`Int64`](@ref).
"""
Period(v)

"""
    Instant

`Instant` types represent integer-based, machine representations of time as continuous
timelines starting from an epoch.
"""
abstract type Instant <: AbstractTime end

"""
    UTInstant{T}

The `UTInstant` represents a machine timeline based on UT time (1 day = one revolution of
the earth). The `T` is a `Period` parameter that indicates the resolution or precision of
the instant.
"""
struct UTInstant{P<:Period} <: Instant
    periods::P
end

# Convenience default constructors
UTM(x) = UTInstant(Millisecond(x))
UTD(x) = UTInstant(Day(x))

# Calendar types provide rules for interpretating instant
# timelines in human-readable form.
abstract type Calendar <: AbstractTime end

# ISOCalendar implements the ISO 8601 standard (en.wikipedia.org/wiki/ISO_8601)
# Notably based on the proleptic Gregorian calendar
# ISOCalendar provides interpretation rules for UTInstants to civil date and time parts
struct ISOCalendar <: Calendar end

abstract type TimeZone end
struct UTC <: TimeZone end

"""
    TimeType

`TimeType` types wrap `Instant` machine instances to provide human representations of the
machine instant. `Time`, `DateTime` and `Date` are subtypes of `TimeType`.
"""
abstract type TimeType <: AbstractTime end

abstract type AbstractDateTime <: TimeType end

"""
    DateTime

`DateTime` wraps a `UTInstant{Millisecond}` and interprets it according to the proleptic
Gregorian calendar.
"""
struct DateTime <: AbstractDateTime
    instant::UTInstant{Millisecond}
    DateTime(instant::UTInstant{Millisecond}) = new(instant)
end

"""
    Date

`Date` wraps a `UTInstant{Day}` and interprets it according to the proleptic Gregorian calendar.
"""
struct Date <: TimeType
    instant::UTInstant{Day}
    Date(instant::UTInstant{Day}) = new(instant)
end

"""
    Time

`Time` wraps a `Nanosecond` and represents a specific moment in a 24-hour day.
"""
struct Time <: TimeType
    instant::Nanosecond
    Time(instant::Nanosecond) = new(instant)
end

# Convert y,m,d to # of Rata Die days
# Works by shifting the beginning of the year to March 1,
# so a leap day is the very last day of the year
const SHIFTEDMONTHDAYS = (306, 337, 0, 31, 61, 92, 122, 153, 184, 214, 245, 275)
function totaldays(y, m, d)
    # If we're in Jan/Feb, shift the given year back one
    z = m < 3 ? y - 1 : y
    mdays = SHIFTEDMONTHDAYS[m]
    # days + month_days + year_days
    return d + mdays + 365z + fld(z, 4) - fld(z, 100) + fld(z, 400) - 306
end

# If the year is divisible by 4, except for every 100 years, except for every 400 years
isleapyear(y) = ((y % 4 == 0) && (y % 100 != 0)) || (y % 400 == 0)

# Number of days in month
const DAYSINMONTH = (31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31)
daysinmonth(y,m) = DAYSINMONTH[m] + (m == 2 && isleapyear(y))

### UTILITIES ###

# These are necessary because the type constructors for TimeType subtypes can
# throw, and we want to be able to use tryparse without requiring a try/catch.
# This is made easier by providing a helper function that checks arguments, so
# we can validate arguments in tryparse.

"""
    validargs(::Type{<:TimeType}, args...) -> Union{ArgumentError, Nothing}

Determine whether the given arguments consitute valid inputs for the given type.
Returns either an `ArgumentError`, or [`nothing`](@ref) in case of success.
"""
function validargs end

"""
    argerror([msg]) -> Union{ArgumentError, Nothing}

Return an `ArgumentError` object with the given message,
or [`nothing`](@ref) if no message is provided. For use by `validargs`.
"""
argerror(msg::String) = ArgumentError(msg)
argerror() = nothing

### CONSTRUCTORS ###
# Core constructors
"""
    DateTime(y, [m, d, h, mi, s, ms]) -> DateTime

Construct a `DateTime` type by parts. Arguments must be convertible to [`Int64`](@ref).
"""
function DateTime(y::Int64, m::Int64=1, d::Int64=1,
                  h::Int64=0, mi::Int64=0, s::Int64=0, ms::Int64=0)
    err = validargs(DateTime, y, m, d, h, mi, s, ms)
    err === nothing || throw(err)
    rata = ms + 1000 * (s + 60mi + 3600h + 86400 * totaldays(y, m, d))
    return DateTime(UTM(rata))
end

function validargs(::Type{DateTime}, y::Int64, m::Int64, d::Int64,
                   h::Int64, mi::Int64, s::Int64, ms::Int64)
    0 < m < 13 || return argerror("Month: $m out of range (1:12)")
    0 < d < daysinmonth(y, m) + 1 || return argerror("Day: $d out of range (1:$(daysinmonth(y, m)))")
    -1 < h < 24 || (h == 24 && mi==s==ms==0) ||
        return argerror("Hour: $h out of range (0:23)")
    -1 < mi < 60 || return argerror("Minute: $mi out of range (0:59)")
    -1 < s < 60 || return argerror("Second: $s out of range (0:59)")
    -1 < ms < 1000 || return argerror("Millisecond: $ms out of range (0:999)")
    return argerror()
end

"""
    Date(y, [m, d]) -> Date

Construct a `Date` type by parts. Arguments must be convertible to [`Int64`](@ref).
"""
function Date(y::Int64, m::Int64=1, d::Int64=1)
    err = validargs(Date, y, m, d)
    err === nothing || throw(err)
    return Date(UTD(totaldays(y, m, d)))
end

function validargs(::Type{Date}, y::Int64, m::Int64, d::Int64)
    0 < m < 13 || return argerror("Month: $m out of range (1:12)")
    0 < d < daysinmonth(y, m) + 1 || return argerror("Day: $d out of range (1:$(daysinmonth(y, m)))")
    return argerror()
end

"""
    Time(h, [mi, s, ms, us, ns]) -> Time

Construct a `Time` type by parts. Arguments must be convertible to [`Int64`](@ref).
"""
function Time(h::Int64, mi::Int64=0, s::Int64=0, ms::Int64=0, us::Int64=0, ns::Int64=0)
    err = validargs(Time, h, mi, s, ms, us, ns)
    err === nothing || throw(err)
    return Time(Nanosecond(ns + 1000us + 1000000ms + 1000000000s + 60000000000mi + 3600000000000h))
end

function validargs(::Type{Time}, h::Int64, mi::Int64, s::Int64, ms::Int64, us::Int64, ns::Int64)
    -1 < h < 24 || return argerror("Hour: $h out of range (0:23)")
    -1 < mi < 60 || return argerror("Minute: $mi out of range (0:59)")
    -1 < s < 60 || return argerror("Second: $s out of range (0:59)")
    -1 < ms < 1000 || return argerror("Millisecond: $ms out of range (0:999)")
    -1 < us < 1000 || return argerror("Microsecond: $us out of range (0:999)")
    -1 < ns < 1000 || return argerror("Nanosecond: $ns out of range (0:999)")
    return argerror()
end

# Convenience constructors from Periods
function DateTime(y::Year, m::Month=Month(1), d::Day=Day(1),
                  h::Hour=Hour(0), mi::Minute=Minute(0),
                  s::Second=Second(0), ms::Millisecond=Millisecond(0))
    return DateTime(value(y), value(m), value(d),
                        value(h), value(mi), value(s), value(ms))
end

Date(y::Year, m::Month=Month(1), d::Day=Day(1)) = Date(value(y), value(m), value(d))

function Time(h::Hour, mi::Minute=Minute(0), s::Second=Second(0),
              ms::Millisecond=Millisecond(0),
              us::Microsecond=Microsecond(0), ns::Nanosecond=Nanosecond(0))
    return Time(value(h), value(mi), value(s), value(ms), value(us), value(ns))
end

# To allow any order/combination of Periods

"""
    DateTime(periods::Period...) -> DateTime

Construct a `DateTime` type by `Period` type parts. Arguments may be in any order. DateTime
parts not provided will default to the value of `Dates.default(period)`.
"""
function DateTime(period::Period, periods::Period...)
    y = Year(1); m = Month(1); d = Day(1)
    h = Hour(0); mi = Minute(0); s = Second(0); ms = Millisecond(0)
    for p in (period, periods...)
        isa(p, Year) && (y = p::Year)
        isa(p, Month) && (m = p::Month)
        isa(p, Day) && (d = p::Day)
        isa(p, Hour) && (h = p::Hour)
        isa(p, Minute) && (mi = p::Minute)
        isa(p, Second) && (s = p::Second)
        isa(p, Millisecond) && (ms = p::Millisecond)
    end
    return DateTime(y, m, d, h, mi, s, ms)
end

"""
    Date(period::Period...) -> Date

Construct a `Date` type by `Period` type parts. Arguments may be in any order. `Date` parts
not provided will default to the value of `Dates.default(period)`.
"""
function Date(period::Period, periods::Period...)
    y = Year(1); m = Month(1); d = Day(1)
    for p in (period, periods...)
        isa(p, Year) && (y = p::Year)
        isa(p, Month) && (m = p::Month)
        isa(p, Day) && (d = p::Day)
    end
    return Date(y, m, d)
end

"""
    Time(period::TimePeriod...) -> Time

Construct a `Time` type by `Period` type parts. Arguments may be in any order. `Time` parts
not provided will default to the value of `Dates.default(period)`.
"""
function Time(period::TimePeriod, periods::TimePeriod...)
    h = Hour(0); mi = Minute(0); s = Second(0)
    ms = Millisecond(0); us = Microsecond(0); ns = Nanosecond(0)
    for p in (period, periods...)
        isa(p, Hour) && (h = p::Hour)
        isa(p, Minute) && (mi = p::Minute)
        isa(p, Second) && (s = p::Second)
        isa(p, Millisecond) && (ms = p::Millisecond)
        isa(p, Microsecond) && (us = p::Microsecond)
        isa(p, Nanosecond) && (ns = p::Nanosecond)
    end
    return Time(h, mi, s, ms, us, ns)
end

# Fallback constructors
DateTime(y, m=1, d=1, h=0, mi=0, s=0, ms=0) = DateTime(Int64(y), Int64(m), Int64(d), Int64(h), Int64(mi), Int64(s), Int64(ms))
Date(y, m=1, d=1) = Date(Int64(y), Int64(m), Int64(d))
Time(h, mi=0, s=0, ms=0, us=0, ns=0) = Time(Int64(h), Int64(mi), Int64(s), Int64(ms), Int64(us), Int64(ns))

# Traits, Equality
Base.isfinite(::Union{Type{T}, T}) where {T<:TimeType} = true
calendar(dt::DateTime) = ISOCalendar
calendar(dt::Date) = ISOCalendar

"""
    eps(::DateTime) -> Millisecond
    eps(::Date) -> Day
    eps(::Time) -> Nanosecond

Returns `Millisecond(1)` for `DateTime` values, `Day(1)` for `Date` values, and `Nanosecond(1)` for `Time` values.
"""
Base.eps

Base.eps(dt::DateTime) = Millisecond(1)
Base.eps(dt::Date) = Day(1)
Base.eps(t::Time) = Nanosecond(1)

Base.typemax(::Union{DateTime, Type{DateTime}}) = DateTime(146138512, 12, 31, 23, 59, 59)
Base.typemin(::Union{DateTime, Type{DateTime}}) = DateTime(-146138511, 1, 1, 0, 0, 0)
Base.typemax(::Union{Date, Type{Date}}) = Date(252522163911149, 12, 31)
Base.typemin(::Union{Date, Type{Date}}) = Date(-252522163911150, 1, 1)
Base.typemax(::Union{Time, Type{Time}}) = Time(23, 59, 59, 999, 999, 999)
Base.typemin(::Union{Time, Type{Time}}) = Time(0)
# Date-DateTime promotion, isless, ==
Base.promote_rule(::Type{Date}, x::Type{DateTime}) = DateTime
Base.isless(x::T, y::T) where {T<:TimeType} = isless(value(x), value(y))
Base.isless(x::TimeType, y::TimeType) = isless(promote(x, y)...)
(==)(x::T, y::T) where {T<:TimeType} = (==)(value(x), value(y))
function ==(a::Time, b::Time)
    return hour(a) == hour(b) && minute(a) == minute(b) &&
        second(a) == second(b) && millisecond(a) == millisecond(b) &&
        microsecond(a) == microsecond(b) && nanosecond(a) == nanosecond(b)
end
(==)(x::TimeType, y::TimeType) = (===)(promote(x, y)...)
Base.hash(x::Time, h::UInt) =
    hash(hour(x), hash(minute(x), hash(second(x),
        hash(millisecond(x), hash(microsecond(x), hash(nanosecond(x), h))))))

import Base: sleep, Timer, timedwait
sleep(time::Period) = sleep(toms(time) / 1000)
Timer(time::Period; interval::Period = Second(0)) =
    Timer(toms(time) / 1000, interval = toms(interval) / 1000)
timedwait(testcb::Function, time::Period) = timedwait(testcb, toms(time) / 1000)

Base.OrderStyle(::Type{<:AbstractTime}) = Base.Ordered()
Base.ArithmeticStyle(::Type{<:AbstractTime}) = Base.ArithmeticWraps()


import Base: .<, .>, .<=, .>=, .==, .!=
export .<, .>, .<=, .>=, .==, .!=, and, or, maxOf, minOf

import Base: .+, .-, .*, ./, .\, .%
export .+, .-, .*, ./, .\, .%

export addAssign!,subAssign!,divAssign!,mulAssign!,modAssign!

immutable Binary <: AFImpl
    le::Ptr{Void}
    lt::Ptr{Void}
    ge::Ptr{Void}
    gt::Ptr{Void}
    eq::Ptr{Void}
    neq::Ptr{Void}
    and::Ptr{Void}
    or::Ptr{Void}
    add::Ptr{Void}
    sub::Ptr{Void}
    mul::Ptr{Void}
    div::Ptr{Void}
    maxOf::Ptr{Void}
    minOf::Ptr{Void}
    mod::Ptr{Void}

    function Binary(ptr)
        new(
            Libdl.dlsym(ptr, :af_le),
            Libdl.dlsym(ptr, :af_lt),
            Libdl.dlsym(ptr, :af_ge),
            Libdl.dlsym(ptr, :af_gt),
            Libdl.dlsym(ptr, :af_eq),
            Libdl.dlsym(ptr, :af_neq),
            Libdl.dlsym(ptr, :af_and),
            Libdl.dlsym(ptr, :af_or),
            Libdl.dlsym(ptr, :af_add),
            Libdl.dlsym(ptr, :af_sub),
            Libdl.dlsym(ptr, :af_mul),
            Libdl.dlsym(ptr, :af_div),
            Libdl.dlsym(ptr, :af_maxof),
            Libdl.dlsym(ptr, :af_minof),
            Libdl.dlsym(ptr, :af_mod)
        )
    end
end

macro binOp(op, cFunc)
    quote
        function $(esc(op)){B}(lhs::AFArray{B}, rhs::AFArray{B})
            verifyAccess(lhs)
            verifyAccess(rhs)
            af = lhs.af
            result = af.results.ptr
            err = ccall(af.binary.$cFunc,
                Cint, (Ptr{Ptr{Void}}, Ptr{Void}, Ptr{Void}, Bool),
                result, lhs.ptr, rhs.ptr, af.batch)
            assertErr(err)
            AFArray{B}(af, result[])
        end

        function $(esc(op)){B}(lhs::AFArray{B}, rhsConst::Number)
            verifyAccess(lhs)
            af = lhs.af
            result = af.results.ptr
            rhs = constant(af, rhsConst, size(lhs)...)
            try
                err = ccall(af.binary.$cFunc,
                    Cint, (Ptr{Ptr{Void}}, Ptr{Void}, Ptr{Void}, Bool),
                    result, lhs.ptr, rhs.ptr, af.batch)
                assertErr(err)
                AFArray{B}(af, result[])
            finally
                release!(rhs)
            end
        end

        function $(esc(op)){B}(lhsConst::Number, rhs::AFArray{B})
            verifyAccess(rhs)
            af = rhs.af
            result = af.results.ptr
            lhs = constant(af, lhsConst, size(rhs)...)
            try
                err = ccall(af.binary.$cFunc,
                    Cint, (Ptr{Ptr{Void}}, Ptr{Void}, Ptr{Void}, Bool),
                    result, lhs.ptr, rhs.ptr, af.batch)
                assertErr(err)
                AFArray{B}(af, result[])
            finally
                release!(lhs)
            end
        end
    end
end

@binOp(.<, lt)
@binOp(.<=, le)
@binOp(.>, gt)
@binOp(.>=, ge)
@binOp(.==, eq)
@binOp(.!=, neq)
@binOp(and, and)
@binOp(or, or)

@binOp(.+, add)
@binOp(.-, sub)
@binOp(.*, mul)
@binOp(./, div)
@binOp(.%, mod)
@binOp(maxOf, maxOf)
@binOp(minOf, minOf)

function .\{B}(lhs::AFArray{B}, rhs::AFArray{B})
    rhs ./ lhs
end

function .\(lhs::AFArray, rhs::Number)
    rhs ./ lhs
end

function .\(lhs::Number, rhs::AFArray)
    rhs ./ lhs
end

macro binOpAssign(op, cFunc)
    quote
        function $(esc(op)){B}(lhs::AFArray{B}, rhs::AFArray{B})
            verifyAccess(lhs)
            verifyAccess(rhs)
            af = lhs.af
            result = af.results.ptr
            err = ccall(af.binary.$cFunc,
                Cint, (Ptr{Ptr{Void}}, Ptr{Void}, Ptr{Void}, Bool),
                result, lhs.ptr, rhs.ptr, af.batch)
            assertErr(err)
            updatePtr(lhs, result[])
        end

        function $(esc(op)){B}(lhs::AFArray{B}, rhsConst::Number)
            verifyAccess(lhs)
            af = lhs.af
            result = af.results.ptr
            rhs = constant(af, rhsConst, size(lhs)...)
            try
                err = ccall(af.binary.$cFunc,
                    Cint, (Ptr{Ptr{Void}}, Ptr{Void}, Ptr{Void}, Bool),
                    result, lhs.ptr, rhs.ptr, af.batch)
                assertErr(err)
                updatePtr(lhs, result[])
            finally
                release!(rhs)
            end
        end

        function $(esc(op)){B}(lhsConst::Number, rhs::AFArray{B})
            verifyAccess(rhs)
            af = rhs.af
            result = af.results.ptr
            lhs = constant(af, lhsConst, size(rhs)...)
            try
                err = ccall(af.binary.$cFunc,
                    Cint, (Ptr{Ptr{Void}}, Ptr{Void}, Ptr{Void}, Bool),
                    result, lhs.ptr, rhs.ptr, af.batch)
                assertErr(err)
                updatePtr(lhs, result[])
            finally
                release!(lhs)
            end
        end
    end
end

function updatePtr(arr::AFArray, ptr::Ptr{Void})
    if ptr != arr.ptr
        release!(arr.af, arr.ptr)
        arr.ptr = ptr
    end
end

@binOpAssign(addAssign!, add)
@binOpAssign(subAssign!, sub)
@binOpAssign(mulAssign!, mul)
@binOpAssign(divAssign!, div)
@binOpAssign(modAssign!, mod)


using Revise
using Lava
using VulkanCore
using LinearAlgebra
using GeometryTypes
using Lava: @class, @scope, @autodestroy

shaderFolder = String(@__DIR__) * "/../shaders/"

mutable struct CameraData
    view::Mat4f0
    proj::Mat4f0

    CameraData() = new(Mat4f0(1I), Mat4f0(1I))
    CameraData(view::Mat4f0, proj::Mat4f0) = new(view, proj)
end

@class CameraData

function getTransposeData(data::CameraData)
    return [transpose(data.view) transpose(data.proj)]
end

const u8vec4 = Vec{4, UInt8}
struct Vertex
    position::Vec3f0
    color::u8vec4
end

cubeVertices = [
    Vertex(Vec3f0(-1.0, -1.0, -1.0), u8vec4(0, 0, 0, 255)),
    Vertex(Vec3f0(-1.0, -1.0, 1.0), u8vec4(0, 0, 255, 255)),
    Vertex(Vec3f0(-1.0, 1.0, -1.0), u8vec4(0, 255, 0, 255)),
    Vertex(Vec3f0(-1.0, 1.0, 1.0), u8vec4(0, 255, 255, 255)),
    Vertex(Vec3f0(1.0, -1.0, -1.0), u8vec4(255, 0, 0, 255)),
    Vertex(Vec3f0(1.0, -1.0, 1.0), u8vec4(255, 0, 255, 255)),
    Vertex(Vec3f0(1.0, 1.0, -1.0), u8vec4(255, 255, 0, 255)),
    Vertex(Vec3f0(1.0, 1.0, 1.0), u8vec4(255, 255, 255, 255))
]

cubeIndices = Vector{UInt32}(  [0, 2, 3,
                                0, 3, 1,
                                0, 1, 5,
                                0, 5, 4,
                                2, 7, 3,
                                2, 6, 7,
                                1, 7, 5,
                                1, 3, 7,
                                0, 4, 6,
                                0, 6, 2,
                                4, 5, 7,
                                4, 7, 6])

function main()
    # Create a Vulkan instance, tell it we need glfw and the validation as
    # extension features
    @scope begin
    fs = Vector{features.IFeature}()
    glfw = features.create(features.GlfwOutput)
    push!(fs, features.create(features.Validation))
    push!(fs, glfw)
    instance = @autodestroy lava.create(lava.Instance, fs)
    queues = [lava.createGraphics(lava.QueueRequest, "graphics")]#::Vector{QueueRequest}

    # Create a logical device with a single graphics queue named "graphics" on
    # a discrete (non-integrated) GPU. For more complex apps you can request
    # additional queues (e.g. a separate transfer queue)
    device = @autodestroy instance.createDevice(queues,
                                    lava.NthOfTypeStrategy(vk.VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU))
                                   # lava.NthOfTypeStrategy(vk.VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU

    # Create a pipeline layout that transfers the CameraData as push constants
    # and has no descriptors (textures or uniform buffers)
    plLayout = @autodestroy device.createPipelineLayout(CameraData)

    # Test descriptor set layout
    # descCreateInfo = lava.DescriptorSetLayoutCreateInfo()
    # lava.addUniformBuffer(descCreateInfo, vk.VK_SHADER_STAGE_VERTEX_BIT);
    # lava.addCombinedImageSampler(descCreateInfo, vk.VK_SHADER_STAGE_FRAGMENT_BIT);
    # descLayout = lava.createDescriptorSetLayout(device, descCreateInfo)
    # plLayout = lava.createPipelineLayout(device, CameraData, [descLayout])

    pass = @autodestroy device.createRenderPass(lava.createSimpleForward(lava.RenderPassCreateInfo, glfw.format()))

    vertShader = @autodestroy device.createShaderFromFile(shaderFolder * "cube_vert.spv")
    fragShader = @autodestroy device.createShaderFromFile(shaderFolder * "cube_frag.spv")
    stageVert = lava.defaults(lava.PipelineShaderStageCreateInfo, _module = vertShader)
    stageFrag = lava.defaults(lava.PipelineShaderStageCreateInfo, _module = fragShader)
    stages = [stageVert.handleRef()[], stageFrag.handleRef()[]]

    attributes = Vector{vk.VkVertexInputAttributeDescription}()
    bindings = Vector{vk.VkVertexInputBindingDescription}()
    lava.binding(attributes, bindings, UInt32(0), Vertex, :position)
    lava.binding(attributes, bindings, UInt32(0), Vertex, :color)
    vertexInputState = lava.PipelineVertexInputStateCreateInfo(attributes = attributes, bindings = bindings)

    ci = lava.defaults(lava.GraphicsPipelineCreateInfo,
        stages = stages,
        layout = plLayout,
        renderPass = lava.handleRef(pass)[],
        subpass = UInt32(0),
        depthTestEnable = VkExt.VK_TRUE,
        depthWriteEnable = VkExt.VK_TRUE,
        # Due to the flipped y axis in NDC, the triangle winding order is inverted, too
        frontFace = vk.VK_FRONT_FACE_CLOCKWISE,
        inputState = vertexInputState
    )
    pipeline = @autodestroy lava.GraphicsPipeline(ci)

    camera = Camera.GenericCamera()
    camera.setTarget(Vec3f0(0.0, 0.0, 0.0))
    fbos = @autodestroy Vector{lava.Framebuffer}()
    window = @autodestroy glfw.openWindow(UInt32(800), UInt32(600), true)
    imgList = @autodestroy []
    imgViewList = @autodestroy []
    window.buildSwapchainWith(views::Vector{lava.ImageView}->begin
        depthImageCreateInfo = lava.attachment2D(lava.getPhysicalDevice(device), window.mWidth, window.mHeight, vk.VK_FORMAT_D32_SFLOAT)
        depthImage = lava.createImage(depthImageCreateInfo, device)
        depthImage.realizeAttachment()
        depthView = depthImage.createView()
        lava.destroy!(imgList)
        lava.destroy!(imgViewList)
        push!(imgList, depthImage)
        push!(imgViewList, depthView)

        lava.destroy!(fbos)
        for view in views
            push!(fbos, lava.Framebuffer(pass, [depthView, view]))
        end

        camera.setAspectRatio(Float32(features.getWidth(window)) / Float32(features.getHeight(window)))
    end)

    #Upload of the cube mesh
    vab = @autodestroy lava.createBuffer(device, lava.arrayBuffer())
    vab.setDataVRAM(cubeVertices, Vertex)
    eab = @autodestroy lava.createBuffer(device, lava.indexBuffer())
    eab.setDataVRAM(cubeIndices, UInt32)

    camera.setPosition(Vec3f0(2.0, 2.0, 2.0))
    camera.setTarget(Vec3f0(0.0, 0.0, 0.0))
    camera.rotateAroundTarget_GlobalAxes(Float32(0.0), Float32(0.001), Float32(0.0))

    while true
        features.pollEvent()
        lava.graphicsQueue(device).catchUp(Int32(1))
        camera.rotateAroundTarget_GlobalAxes(Float32(0.0), Float32(0.001), Float32(0.0))
        #println(CameraData(Camera.getViewMatrix(camera), Camera.getProjectionMatrix(camera)))
        frame = window.startFrame()
        cmd = lava.graphicsQueue(device).beginCommandBuffer()

        # Vulkan doesn't take care of synchronization for us.
        # We need to tell it to wait with rendering until the image is ready
        # and to wait with presenting the image until the rendering is done.
        # The frame struct provides the necessary semaphores.
        cmd.wait(frame.imageReady())
        cmd.signal(frame.renderingComplete())

        forward = cmd.beginRenderpass(fbos[features.imageIndex(frame) + 1])

        sub = forward.startInlineSubpass()
        sub.bindPipeline(pipeline)
        sub.setViewports([vk.VkViewport(0, 0, features.getWidth(window), features.getHeight(window), 0.0, 1.0)])
        sub.bindVertexBuffers([vab])
        sub.bindIndexBuffer(eab)
        cameraData = CameraData(camera.getViewMatrix(), camera.getProjectionMatrix())
        sub.pushConstantBlock(cameraData.getTransposeData())

        sub.drawIndexed(UInt32(length(cubeIndices)))

        forward.endRenderPass()
        cmd.endCommandBuffer()
        frame.endFrame()
    end
    end
end


#####
##### Profile Asynchronous MCTS to figure out the optimal number of workers
#####

using AlphaZero
using ProgressMeter
using Plots

# It is important not to profile async mcts on tictactoe as the number
# of states in this game is small and so the mcts tree quickly contains
# every state.

# Do NOT run this from the REPL!

ENV["JULIA_CUDA_MEMORY_POOL"] = "split"

include("../games.jl")
GAME = get(ENV, "GAME", "connect-four")
SelectedGame = GAME_MODULE[GAME]
using .SelectedGame: Game, Training

NUM_GAMES = 10
MAX_LOG_NWORKERS = 8 # 2^8 = 256

TIME_FIG = "mcts_speed"
INFERENCE_TIME_RATIO_FIG = "inference_time_ratio"

network = Training.Network{Game}(Training.netparams)

function profile(nworkers, ngames)
  params = MctsParams(Training.self_play.mcts, num_workers=nworkers)
  player = MctsPlayer(network, params)
  GC.gc(true)
  time = @elapsed begin
    @showprogress for i in 1:ngames
      AlphaZero.play_game(player, player)
    end
  end
  return time, MCTS.inference_time_ratio(player.mcts)
end

println("Compile everything...")
profile(1, 1)
profile(2, 1)

results = []
ts = []
itrs = []
for i in 0:MAX_LOG_NWORKERS
  nworkers = 2 ^ i
  println("Profiling MCTS with $nworkers workers...")
  time, inference_time_ratio = profile(nworkers, NUM_GAMES)
  push!(ts, time)
  push!(itrs, inference_time_ratio)
end

xticks = (
  collect(0:MAX_LOG_NWORKERS),
  ["$(2^i)" for i in 0:MAX_LOG_NWORKERS])

plot(0:MAX_LOG_NWORKERS, ts[1] ./ ts,
  title="Async MCTS Speedup",
  ylabel="Speedup",
  xlabel="Number of workers",
  ylims=(0, Inf),
  legend=:none,
  xticks=xticks)

hline!([1])

savefig(TIME_FIG)

plot(0:MAX_LOG_NWORKERS, 100 * itrs,
  title="Percentage of time spent in inference",
  ylims=(0, 100),
  xlabel="Number of workers",
  legend=:none,
  xticks=xticks)

savefig(INFERENCE_TIME_RATIO_FIG)


abstract type pseudoGPUArray <: SolidStateDetectors.GPUArrays.AbstractGPUArray{T, N where N} end
SolidStateDetectors.get_device(::Type{pseudoGPUArray}) = SolidStateDetectors.KernelAbstractions.CPU()

@testset "Simulate example detector: Inverted Coax" begin
    sim = Simulation{T}(SSD_examples[:InvertedCoax])

    calculate_electric_potential!(
        sim,
        device_array_type = pseudoGPUArray,
        convergence_limit = 0, # For this we have to probably add an addition kernel function
        max_n_iterations = 2000, # For now just perform always `max_n_iterations` iterations
        refinement_limits = [0.2, 0.1],
        depletion_handling = true
    )

    W = SolidStateDetectors.calculate_stored_energy(sim)
    @test isapprox(ustrip(W), 2.41e-4, atol = 1e-5)
end

@testset "Simulate example detector: Inverted Coax" begin
    sim = Simulation{T}(SSD_examples[:CGD])

    calculate_electric_potential!(
        sim,
        device_array_type = pseudoGPUArray,
        convergence_limit = 0, # For this we have to probably add an addition kernel function
        max_n_iterations = 2000, # For now just perform always `max_n_iterations` iterations
        refinement_limits = [0.2],
        depletion_handling = true
    )

    W = SolidStateDetectors.calculate_stored_energy(sim)
    @test isapprox(ustrip(W), 4.39e-6, atol = 1e-5)
end


module PulseController

# Bring in other files as submodules.  Every file is exactly one module.
# The dependency tree is as follows:

# Clifford - standalone
# Waveforms - standalone

# Qubits - imports from Clifford and Waveforms

# Seqeuncing - imports from Qubits

# ParticleSwarm - standalone

# Orbit - imports from Sequencing and ParticleSwarm

include("clifford.jl")
include("waveforms.jl")
include("qubits.jl")
include("sequencing.jl")
include("particleswarm.jl")
include("orbit.jl")

using .ORBIT # Carries all other modules
export sendSequence

# For writing out sequences, we will provide all of the pulse names.  Hopefully
# having these in scope doesn't clash with user variables
import .Clifford: Idle, Xpi, Xpi2, X3pi2, Ypi, Ypi2, Y3pi2, Zpi, Zpi2, Z3pi2, CZ
export Idle, Xpi, Xpi2, X3pi2, Ypi, Ypi2, Y3pi2, Zpi, Zpi2, Z3pi2, CZ

# From qubits
export Qubit
export cosInit
export gaussInit
export generalInit

# From waveforms, for setting qubit waveforms manually
export FloatWaveform
export ExactWaveform

end # End module PulseController

using PulseController # Bad style probably


function DSSTPropagatorBuilder(arg0::Orbit, arg1::ODEIntegratorBuilder, arg2::jdouble, arg3::PropagationType, arg4::PropagationType)
    return DSSTPropagatorBuilder((Orbit, ODEIntegratorBuilder, jdouble, PropagationType, PropagationType), arg0, arg1, arg2, arg3, arg4)
end

function DSSTPropagatorBuilder(arg0::Orbit, arg1::ODEIntegratorBuilder, arg2::jdouble, arg3::PropagationType, arg4::PropagationType, arg5::AttitudeProvider)
    return DSSTPropagatorBuilder((Orbit, ODEIntegratorBuilder, jdouble, PropagationType, PropagationType, AttitudeProvider), arg0, arg1, arg2, arg3, arg4, arg5)
end

function add_force_model(obj::DSSTPropagatorBuilder, arg0::DSSTForceModel)
    return jcall(obj, "addForceModel", void, (DSSTForceModel,), arg0)
end

function build_kalman_model(obj::DSSTPropagatorBuilder, arg0::List, arg1::List, arg2::ParameterDriversList)
    return jcall(obj, "buildKalmanModel", DSSTKalmanModel, (List, List, ParameterDriversList), arg0, arg1, arg2)
end

function build_ls_model(obj::DSSTPropagatorBuilder, arg0::Vector{IntegratedPropagatorBuilder}, arg1::List, arg2::ParameterDriversList, arg3::ModelObserver)
    return jcall(obj, "buildLSModel", DSSTBatchLSModel, (Vector{IntegratedPropagatorBuilder}, List, ParameterDriversList, ModelObserver), arg0, arg1, arg2, arg3)
end

function build_propagator(obj::DSSTPropagatorBuilder, arg0::Vector{jdouble})
    return jcall(obj, "buildPropagator", DSSTPropagator, (Vector{jdouble},), arg0)
end

function copy(obj::DSSTPropagatorBuilder)
    return jcall(obj, "copy", DSSTPropagatorBuilder, ())
end

function get_all_force_models(obj::DSSTPropagatorBuilder)
    return jcall(obj, "getAllForceModels", List, ())
end

function get_integrator_builder(obj::DSSTPropagatorBuilder)
    return jcall(obj, "getIntegratorBuilder", ODEIntegratorBuilder, ())
end

function get_mass(obj::DSSTPropagatorBuilder)
    return jcall(obj, "getMass", jdouble, ())
end

function set_mass(obj::DSSTPropagatorBuilder, arg0::jdouble)
    return jcall(obj, "setMass", void, (jdouble,), arg0)
end


mutable struct TestClassifier <: AbstractClassifier
    dummy_loss::Float64
    classes::Vector
end

Lighthouse.classes(c::TestClassifier) = c.classes

function Lighthouse.train!(c::TestClassifier, dummy_batches, logger)
    for (dummy_input_batch, loss_delta) in dummy_batches
        c.dummy_loss += loss_delta
        Lighthouse.log_value!(logger, "train/loss_per_batch", c.dummy_loss)
    end
    return c.dummy_loss
end

const RNG_LOSS = StableRNG(22)

function Lighthouse.loss_and_prediction(c::TestClassifier, dummy_input_batch)
    dummy_soft_label_batch = rand(RNG_LOSS, length(c.classes), size(dummy_input_batch)[end])
    # Fake a softmax
    dummy_soft_label_batch .= dummy_soft_label_batch ./ sum(dummy_soft_label_batch; dims=1)
    return c.dummy_loss, dummy_soft_label_batch
end

@testset "Multi-class learn!(::TestModel, ...)" begin
    mktempdir() do tmpdir
        model = TestClassifier(1000000.0, ["class_$i" for i in 1:5])
        k, n = length(model.classes), 3
        rng = StableRNG(22)
        train_batches = [(rand(rng, 4 * k, n), -rand(rng)) for _ in 1:100]
        test_batches = [((rand(rng, 4 * k, n),), (n * i - n + 1):(n * i)) for i in 1:10]
        possible_vote_labels = collect(0:k)
        votes = [rand(rng, possible_vote_labels) for sample in 1:(n * 10), voter in 1:7]
        votes[:, [1, 2, 3]] .= votes[:, 4] # Voter 1-3 voted identically to voter 4 (force non-zero agreement)
        logger = LearnLogger(joinpath(tmpdir, "logs"), "test_run")
        limit = 5
        let counted = 0
            upon_loss_decrease = Lighthouse.upon(logger,
                                                 "test_set_prediction/mean_loss_per_epoch";
                                                 condition=<, initial=Inf)
            callback = n -> begin
                upon_loss_decrease() do _
                    counted += n
                    @info counted n
                end
            end
            elected = majority.((rng,), eachrow(votes), (1:length(Lighthouse.classes(model)),))
            Lighthouse.learn!(model, logger, () -> train_batches, () -> test_batches, votes, elected;
                              epoch_limit=limit, post_epoch_callback=callback)
            @test counted == sum(1:limit)
        end
        @test length(logger.logged["train/loss_per_batch"]) == length(train_batches) * limit
        for key in ["test_set_prediction/loss_per_batch",
             "test_set_prediction/time_in_seconds_per_batch",
             "test_set_prediction/gc_time_in_seconds_per_batch",
             "test_set_prediction/allocations_per_batch",
             "test_set_prediction/memory_in_mb_per_batch"]
            @test length(logger.logged[key]) == length(test_batches) * limit
        end
        for key in ["test_set_prediction/mean_loss_per_epoch",
             "test_set_evaluation/time_in_seconds_per_epoch",
             "test_set_evaluation/gc_time_in_seconds_per_epoch",
             "test_set_evaluation/allocations_per_epoch",
             "test_set_evaluation/memory_in_mb_per_epoch"]
            @test length(logger.logged[key]) == limit
        end
        @test length(logger.logged["test_set_evaluation/metrics_per_epoch"]) == limit

        # Test multiclass optimal_threshold param invalid
        @test_throws ArgumentError Lighthouse.learn!(model, logger, () -> train_batches,
                                                     () -> test_batches, votes;
                                                     epoch_limit=limit,
                                                     optimal_threshold_class=1)

        # Test `predict!`
        num_samples = sum(b -> size(b[1][1], 2), test_batches)
        predicted_soft = zeros(num_samples, length(model.classes))
        predict!(model, predicted_soft, test_batches, logger; logger_prefix="halloooo")
        @test !all(predicted_soft .== 0)
        @test length(logger.logged["halloooo/mean_loss_per_epoch"]) == 1
        @test length(logger.logged["halloooo/loss_per_batch"]) == length(test_batches)
        @test length(logger.logged["halloooo/time_in_seconds_per_batch"]) ==
              length(test_batches)

        # Test `evaluate!`
        n_examples = 40
        num_voters = 20
        predicted_soft = rand(rng, Float32, n_examples, length(model.classes))
        predicted_hard = map(label -> Lighthouse.onecold(model, label),
                             eachrow(predicted_soft))
        votes = [rand(rng, possible_vote_labels)
                 for sample in 1:n_examples, voter in 1:num_voters]
        votes[:, 3] .= votes[:, 4] # Voter 4 voted identically to voter 3 (force non-zero agreement)
        elected_hard = map(row -> majority(rng, row, 1:length(model.classes)),
                           eachrow(votes))
        evaluate!(predicted_hard, predicted_soft, elected_hard, model.classes, logger;
                  logger_prefix="wheeeeeee", logger_suffix="_for_all_time", votes=votes)
        @test length(logger.logged["wheeeeeee/time_in_seconds_for_all_time"]) == 1
        @test length(logger.logged["wheeeeeee/metrics_for_all_time"]) == 1

        # Test startified eval
        strata = [Set("group $(j % Int(ceil(sqrt(j))))" for j in 1:(i - 1))
                  for i in 1:size(votes, 1)]
        plot, plot_data = evaluation_metrics_plot(predicted_hard, predicted_soft,
                                                  elected_hard, model.classes, 0.0:0.01:1.0;
                                                  votes=votes, strata=strata)
        @test haskey(plot_data, "stratified_kappas")

        # Test plotting
        plot_data = last(logger.logged["test_set_evaluation/metrics_per_epoch"])
        @test isa(plot_data["thresholds"], AbstractVector)

        pr = plot_pr_curves(plot_data["per_class_pr_curves"], plot_data["class_labels"])
        @testplot pr

        roc = plot_roc_curves(plot_data["per_class_roc_curves"],
                              plot_data["per_class_roc_aucs"], plot_data["class_labels"])
        @testplot roc

        # Kappa no IRA
        kappas_no_ira = plot_kappas(vcat(plot_data["multiclass_kappa"], plot_data["per_class_kappas"]),
                        vcat("Multiclass", plot_data["class_labels"]))
        @testplot kappas_no_ira

        # Kappa with IRA
        kappas_ira = plot_kappas(vcat(plot_data["multiclass_kappa"], plot_data["per_class_kappas"]),
                        vcat("Multiclass", plot_data["class_labels"]),
                        vcat(plot_data["multiclass_IRA_kappas"],
                             plot_data["per_class_IRA_kappas"]))
        @testplot kappas_ira

        reliability_calibration = plot_reliability_calibration_curves(plot_data["per_class_reliability_calibration_curves"],
                                                                      plot_data["per_class_reliability_calibration_scores"],
                                                                      plot_data["class_labels"])
        @testplot reliability_calibration

        confusion_row = plot_confusion_matrix(plot_data["confusion_matrix"],
                                              plot_data["class_labels"], :Row)
        @testplot confusion_row

        confusion_col = plot_confusion_matrix(plot_data["confusion_matrix"],
                                              plot_data["class_labels"], :Column)
        @testplot confusion_col

        all_together_2 = evaluation_metrics_plot(plot_data)
        @testplot all_together_2

        #savefig(all_together_2, "/tmp/multiclass.png")
    end
end

@testset "2-class `learn!(::TestModel, ...)`" begin
    mktempdir() do tmpdir
        model = TestClassifier(1000000.0, ["class_$i" for i in 1:2])
        k, n = length(model.classes), 3
        rng = StableRNG(23)
        train_batches = [(rand(rng, 4 * k, n), -rand(rng)) for _ in 1:100]
        test_batches = [((rand(rng, 4 * k, n),), (n * i - n + 1):(n * i)) for i in 1:10]
        possible_vote_labels = collect(0:k)
        votes = [rand(rng, possible_vote_labels) for sample in 1:(n * 10), voter in 1:7]
        votes[:, [1, 2, 3]] .= votes[:, 4] # Voter 1-3 voted identically to voter 4 (force non-zero agreement)
        logger = LearnLogger(joinpath(tmpdir, "logs"), "test_run")
        limit = 5
        let counted = 0
            upon_loss_decrease = Lighthouse.upon(logger,
                                                 "test_set_prediction/mean_loss_per_epoch";
                                                 condition=<, initial=Inf)
            callback = n -> begin
                upon_loss_decrease() do _
                    counted += n
                    @info counted n
                end
            end
            elected = majority.((rng,), eachrow(votes), (1:length(Lighthouse.classes(model)),))
            Lighthouse.learn!(model, logger, () -> train_batches, () -> test_batches, votes, elected;
                              epoch_limit=limit, post_epoch_callback=callback)
            @test counted == sum(1:limit)
        end
        # Binary classification logs some additional metrics
        @test length(logger.logged["test_set_evaluation/spearman_correlation_per_epoch"]) ==
              limit
        plot_data = last(logger.logged["test_set_evaluation/metrics_per_epoch"])
        @test haskey(plot_data, "spearman_correlation")

        # No `optimal_threshold_class` during learning...
        @test !haskey(plot_data, "optimal_threshold")
        @test !haskey(plot_data, "optimal_threshold_class")

        # And now, `optimal_threshold_class` during learning
        elected = majority.((rng,), eachrow(votes), (1:length(Lighthouse.classes(model)),))
        Lighthouse.learn!(model, logger, () -> train_batches, () -> test_batches, votes, elected;
                          epoch_limit=limit, optimal_threshold_class=2,
                          test_set_logger_prefix="validation_set")
        plot_data = last(logger.logged["validation_set_evaluation/metrics_per_epoch"])
        @test haskey(plot_data, "optimal_threshold")
        @test haskey(plot_data, "optimal_threshold_class")
        @test plot_data["optimal_threshold_class"] == 2

        # `optimal_threshold_class` param invalid
        @test_throws ArgumentError Lighthouse.learn!(model, logger, () -> train_batches,
                                                     () -> test_batches, votes;
                                                     epoch_limit=limit,
                                                     optimal_threshold_class=3)

        # Test `evaluate!` for votes, no votes
        n_examples = 40
        num_voters = 20
        predicted_soft = rand(rng, Float32, n_examples, length(model.classes))
        predicted_soft .= predicted_soft ./ sum(predicted_soft; dims=2)
        predicted_hard = map(label -> Lighthouse.onecold(model, label),
                             eachrow(predicted_soft))
        votes = [rand(rng, possible_vote_labels)
                 for sample in 1:n_examples, voter in 1:num_voters]
        votes[:, 3] .= votes[:, 4] # Voter 4 voted identically to voter 3 (force non-zero agreement)
        elected_hard = map(row -> majority(rng, row, 1:length(model.classes)),
                           eachrow(votes))

        evaluate!(predicted_hard, predicted_soft, elected_hard, model.classes, logger;
                  logger_prefix="wheeeeeee", logger_suffix="_for_all_time", votes=nothing)
        plot_data = last(logger.logged["wheeeeeee/metrics_for_all_time"])
        @test !haskey(plot_data, "per_class_IRA_kappas")
        @test !haskey(plot_data, "multiclass_IRA_kappas")

        evaluate!(predicted_hard, predicted_soft, elected_hard, model.classes, logger;
                  logger_prefix="wheeeeeee", logger_suffix="_for_all_time", votes=votes)
        plot_data = last(logger.logged["wheeeeeee/metrics_for_all_time"])
        @test haskey(plot_data, "per_class_IRA_kappas")
        @test haskey(plot_data, "multiclass_IRA_kappas")

        # Test `evaluate` for different optimal_threshold classes
        evaluate!(predicted_hard, predicted_soft, elected_hard, model.classes, logger;
                  logger_prefix="wheeeeeee", logger_suffix="_for_all_time", votes=votes,
                  optimal_threshold_class=1)
        plot_data_1 = last(logger.logged["wheeeeeee/metrics_for_all_time"])
        evaluate!(predicted_hard, predicted_soft, elected_hard, model.classes, logger;
                  logger_prefix="wheeeeeee", logger_suffix="_for_all_time", votes=votes,
                  optimal_threshold_class=2)
        plot_data_2 = last(logger.logged["wheeeeeee/metrics_for_all_time"])

        # The thresholds should not be identical (since they are *inclusive* when applied:
        # values greater than _or equal to_ the threshold are given the class value)
        @test plot_data_1["optimal_threshold"] != plot_data_2["optimal_threshold"]

        # The two threshold options yield different results
        thresh_from_roc = Lighthouse._get_optimal_threshold_from_ROC(plot_data_2["per_class_roc_curves"];
                                                                     thresholds=plot_data_2["thresholds"],
                                                                     class_of_interest_index=plot_data_2["optimal_threshold_class"])
        thresh_from_calibration = Lighthouse._calculate_optimal_threshold_from_discrimination_calibration(predicted_soft,
                                                                                votes;
                                                                                thresholds=plot_data_2["thresholds"],
                                                                                class_of_interest_index=plot_data_2["optimal_threshold_class"]).threshold
        @test !isequal(thresh_from_roc, thresh_from_calibration)
        @test isequal(thresh_from_calibration, plot_data_2["optimal_threshold"])

        # Test that plotting succeeds (no specialization relative to the multi-class tests)
        plot_data = last(logger.logged["validation_set_evaluation/metrics_per_epoch"])
        all_together = evaluation_metrics_plot(plot_data)
        #savefig(all_together, "/tmp/binary.png")
        @testplot all_together
    end
end

@testset "`_calculate_ira_kappas`" begin
    classes = ["roy", "gee", "biv"]
    @test isnothing(Lighthouse._calculate_ira_kappas([1; 1; 1; 1], classes))  # Only one voter...
    @test isnothing(Lighthouse._calculate_ira_kappas([1 0; 1 0; 0 1], classes))  # No observations in common...
end

@testset "Calculate `_spearman_corr`" begin
    # Test failures for...
    # ...nonbinary input
    @test_throws ArgumentError Lighthouse._calculate_spearman_correlation([0.9 0.1], [1],
                                                                          ["oh" "em" "gee"])
    # ...non-softmax input
    @test_throws ArgumentError Lighthouse._calculate_spearman_correlation([0.9 0.9], [1],
                                                                          ["oh" "em"])

    # Test class order doesn't matter
    predicted = [0.1 0.9; 0.3 0.7; 1 0]
    votes = [0 1 1 2; 2 2 0 0; 0 2 2 1]
    predicted_flipped = [0.9 0.1; 0.7 0.3; 0 1]
    votes_flipped = [0 2 2 1; 1 1 0 0; 0 1 1 2]
    @test Lighthouse._calculate_spearman_correlation(predicted, votes, ["oh" "em"]) ==
          Lighthouse._calculate_spearman_correlation(predicted_flipped, votes_flipped,
                                                     ["em" "oh"])

    # Test complete agreement
    predicted_soft = [0.0 1.0; 1.0 0.0; 1.0 0.0]
    votes = [2; 1; 1]
    sp = Lighthouse._calculate_spearman_correlation(predicted_soft, votes, ["oh" "em"])
    @test sp.ρ == 1

    # Test complete disagreement
    votes = [1; 2; 2]
    sp = Lighthouse._calculate_spearman_correlation(predicted_soft, votes, ["oh" "em"])
    @test sp.ρ == -1

    # Test NaN spearman due to unranked input
    votes = [1; 2; 2]
    predicted_soft = [
        0.3 0.7
        0.3 0.7
        0.3 0.7
    ]
    sp = Lighthouse._calculate_spearman_correlation(predicted_soft, votes, ["oh" "em"])
    @test isnan(sp.ρ)

    # Test CI
    sp = Lighthouse._spearman_corr([0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.2],
                                   [0.1, 0.3, 0.2, 0.1, 0.3, 0.1, 0.6])
    @test -1.0 < sp.ci_lower < sp.ρ < sp.ci_upper < 1.0
end

@testset "Discrimination calibration" begin
    # Test single voter has same calibration as "whole" by definition
    votes = [1; 1; 1; 2; 2; 2]
    @test length(Lighthouse._elected_probabilities(votes, 1)) == 6
    @test Lighthouse._elected_probabilities(votes, 1) == [1; 1; 1; 0; 0; 0]

    # Test single voter discrimination calibration
    single_voter_calibration = Lighthouse._calculate_voter_discrimination_calibration(votes;
                                                                                class_of_interest_index=1)
    @test length(single_voter_calibration.mse) == 1

    # Test multi-voter voter discrimination calibration
    votes = [
        0 1 1 1
        1 2 0 0
        2 1 2 2
    ] # Note: voters 3 and 4 have voted identically
    voter_calibration = Lighthouse._calculate_voter_discrimination_calibration(votes;
                                                                               class_of_interest_index=1)
    @test length(voter_calibration.mse) == size(votes, 2)
    @test length(voter_calibration.plot_curve_data) == size(votes, 2)
    @test voter_calibration.mse[1] > voter_calibration.mse[3]
    @test !isequal(voter_calibration.plot_curve_data[1],
                   voter_calibration.plot_curve_data[2])
    @test isequal(voter_calibration.plot_curve_data[3],
                  voter_calibration.plot_curve_data[4])

    # Test predicted voter discrimination calibration
    predicted_soft_labels = [1.0 0.0; 0.0 1.0; 0.0 1.0]
    cal = Lighthouse._calculate_optimal_threshold_from_discrimination_calibration(predicted_soft_labels,
                                                                                votes;
                                                                                thresholds=0.0:0.01:1.0,
                                                                                class_of_interest_index=1)
    @test all(cal.mse .<= voter_calibration.mse)
    cal2 = Lighthouse._calculate_optimal_threshold_from_discrimination_calibration(predicted_soft_labels,
                                                                                votes;
                                                                                thresholds=0.0:0.01:1.0,
                                                                                class_of_interest_index=2)
    @test cal.mse == cal2.mse
    @test cal.plot_curve_data[2] != cal2.plot_curve_data[2]
end

@testset "2-class per_class_confusion_statistics" begin
    predicted_soft_labels = [
        0.51 0.49
        0.49 0.51
        0.1 0.9
        0.9 0.1
        0.0 1.0
    ]
    elected_hard_labels = [1, 2, 2, 2, 1]
    thresholds = [0.25, 0.5, 0.75]
    class_1, class_2 = Lighthouse.per_class_confusion_statistics(predicted_soft_labels,
                                                                 elected_hard_labels,
                                                                 thresholds)
    stats_1, stats_2 = class_1[1], class_2[1] # threshold == 0.25
    @test stats_1.actual_negatives == stats_2.actual_positives == 3
    @test stats_1.actual_positives == stats_2.actual_negatives == 2
    @test stats_1.predicted_positives == 3
    @test stats_2.predicted_positives == 4
    @test stats_1.predicted_negatives == 2
    @test stats_2.predicted_negatives == 1
    @test stats_1.true_positives == 1
    @test stats_2.true_positives == 2
    @test stats_1.true_negatives == 1
    @test stats_2.true_negatives == 0
    @test stats_1.false_positives == 2
    @test stats_2.false_positives == 2
    @test stats_1.false_negatives == 1
    @test stats_2.false_negatives == 1
    @test stats_1.true_positive_rate == 0.5
    @test stats_2.true_positive_rate == 2 / 3
    @test stats_1.true_negative_rate == 1 / 3
    @test stats_2.true_negative_rate == 0.0
    @test stats_1.false_positive_rate == 2 / 3
    @test stats_2.false_positive_rate == 1.0
    @test stats_1.false_negative_rate == 0.5
    @test stats_2.false_negative_rate == 1 / 3
    @test stats_1.precision == 1 / 3
    @test stats_2.precision == 0.5

    stats_1, stats_2 = class_1[2], class_2[2] # threshold == 0.5
    @test stats_1.actual_negatives == stats_2.actual_positives == 3
    @test stats_1.actual_positives == stats_2.actual_negatives == 2
    @test stats_1.predicted_negatives == stats_2.predicted_positives == 3
    @test stats_1.predicted_positives == stats_2.predicted_negatives == 2
    @test stats_1.true_negatives == stats_2.true_positives == 2
    @test stats_1.true_positives == stats_2.true_negatives == 1
    @test stats_1.false_positives == stats_2.false_negatives == 1
    @test stats_1.false_negatives == stats_2.false_positives == 1
    @test stats_1.true_positive_rate == stats_2.true_negative_rate == 0.5
    @test stats_1.true_negative_rate == stats_2.true_positive_rate == 2 / 3
    @test stats_1.false_positive_rate == stats_2.false_negative_rate == 1 / 3
    @test stats_1.false_negative_rate == stats_2.false_positive_rate == 0.5
    @test stats_1.precision == 0.5 && stats_2.precision == 2 / 3

    stats_1, stats_2 = class_1[3], class_2[3] # threshold == 0.75
    @test stats_1.actual_negatives == stats_2.actual_positives == 3
    @test stats_1.actual_positives == stats_2.actual_negatives == 2
    @test stats_1.predicted_positives == 1
    @test stats_2.predicted_positives == 2
    @test stats_1.predicted_negatives == 4
    @test stats_2.predicted_negatives == 3
    @test stats_1.true_positives == 0
    @test stats_2.true_positives == 1
    @test stats_1.true_negatives == 2
    @test stats_2.true_negatives == 1
    @test stats_1.false_positives == 1
    @test stats_2.false_positives == 1
    @test stats_1.false_negatives == 2
    @test stats_2.false_negatives == 2
    @test stats_1.true_positive_rate == 0.0
    @test stats_2.true_positive_rate == 1 / 3
    @test stats_1.true_negative_rate == 2 / 3
    @test stats_2.true_negative_rate == 0.5
    @test stats_1.false_positive_rate == 1 / 3
    @test stats_2.false_positive_rate == 0.5
    @test stats_1.false_negative_rate == 1.0
    @test stats_2.false_negative_rate == 2 / 3
    @test stats_1.precision == 0.0
    @test stats_2.precision == 0.5
end

@testset "3-class per_class_confusion_statistics" begin
    predicted_soft_labels = [
        1/3 1/3 1/3
        0.1 0.7 0.2
        0.25 0.25 0.5
        0.4 0.5 0.1
        0.0 0.0 1.0
        0.2 0.5 0.3
        0.5 0.4 0.1
    ]
    elected_hard_labels = [1, 2, 2, 1, 3, 3, 1]
    # TODO would be more robust to have multiple thresholds, but our naive tests
    # here will have to be refactored to avoid becoming a nightmare if we do that
    thresholds = [0.5]
    class_1, class_2, class_3 = Lighthouse.per_class_confusion_statistics(predicted_soft_labels,
                                                                          elected_hard_labels,
                                                                          thresholds)
    stats_1, stats_2, stats_3 = class_1[], class_2[], class_3[] # threshold == 0.5
    @test stats_1.predicted_positives == 1
    @test stats_2.predicted_positives == 3
    @test stats_3.predicted_positives == 2
    @test stats_1.predicted_negatives == 6
    @test stats_2.predicted_negatives == 4
    @test stats_3.predicted_negatives == 5
    @test stats_1.actual_positives == 3
    @test stats_2.actual_positives == 2
    @test stats_3.actual_positives == 2
    @test stats_1.actual_negatives == 4
    @test stats_2.actual_negatives == 5
    @test stats_3.actual_negatives == 5
    @test stats_1.true_positives == 1
    @test stats_2.true_positives == 1
    @test stats_3.true_positives == 1
    @test stats_1.true_negatives == 4
    @test stats_2.true_negatives == 3
    @test stats_3.true_negatives == 4
    @test stats_1.false_positives == 0
    @test stats_2.false_positives == 2
    @test stats_3.false_positives == 1
    @test stats_1.false_negatives == 2
    @test stats_2.false_negatives == 1
    @test stats_3.false_negatives == 1
    @test stats_1.true_positive_rate == 1 / 3
    @test stats_2.true_positive_rate == 0.5
    @test stats_3.true_positive_rate == 0.5
    @test stats_1.true_negative_rate == 1.0
    @test stats_2.true_negative_rate == 0.6
    @test stats_3.true_negative_rate == 0.8
    @test stats_1.false_positive_rate == 0.0
    @test stats_2.false_positive_rate == 0.4
    @test stats_3.false_positive_rate == 0.2
    @test stats_1.false_negative_rate == 2 / 3
    @test stats_2.false_negative_rate == 0.5
    @test stats_3.false_negative_rate == 0.5
    @test stats_1.precision == 1.0
    @test stats_2.precision == 1 / 3
    @test stats_3.precision == 0.5
end


module FractrixZeta

using MacroTools
using LightGraphs
using Optim
using LinearAlgebra
using GeometryBasics
using FileIO
using Dates
using ColorTypes
using VideoIO
using StaticArrays
using Observables
using CUDA
using StructArrays
using Zygote
using Parameters
import Makie
import AbstractPlotting.MakieLayout
import AbstractPlotting
using MacroTools

const tau = 2pi
const τ = tau
const TAU = τ

include("abstracttypes.jl")
include("conformal_transforms.jl")
include("inputhandler.jl")
include("state.jl")
include("lighthouse.jl")
include("renderer.jl")
include("config.jl")
include("mandelbrot.jl")
include("game.jl")

function main()
    game = init_game()
    run_game(game)
    game
end

end

#=

High level organization:

AbstractGameState abstract type
AbstractPlayerState abstract type
AbstractWorldState abstract type
AbstractConformalTransform abstract type
AbstractLightHouse abstract type
AbstractRenderer abstract type

GameState struct <: AbstractGameState
    Designed to be a full reproduction of the state of the game.
    Supports saving and loading.

    PlayerState <: AbstractPlayerState
        Position
        Rotation
        Scale
        Resources,etc,other relevant state
    end

    history of player states

    WorldState <: AbstractWorldState
        Current Transform
        LightHouse graph
            history of transforms to warp between spaces
        SimTime
    end
end

=#


using Printf
using TimerOutputs
using Oceananigans
using Oceananigans.Utils

include("benchmark_utils.jl")

#####
##### Benchmark setup and parameters
#####

const timer = TimerOutput()

FT = Float64
Nt = 10  # Number of iterations to use for benchmarking time stepping.

         archs = [CPU()]        # Architectures to benchmark on.
@hascuda archs = [CPU(), GPU()] # Benchmark GPU on systems with CUDA-enabled GPUs.

#####
##### Utility functions for generating tracer lists
#####

function active_tracers(n)
    n == 0 && return []
    n == 1 && return [:b]
    n == 2 && return [:T, :S]
    throw(ArgumentError("Can't have more than 2 active tracers!"))
end

passive_tracers(n) = [Symbol("C" * string(n)) for n in 1:n]

tracer_list(na, np) = Tuple(vcat(active_tracers(na), passive_tracers(np)))

""" Number of active tracers to buoyancy """
function na2buoyancy(n)
    n == 0 && return nothing
    n == 1 && return BuoyancyTracer()
    n == 2 && return SeawaterBuoyancy()
    throw(ArgumentError("Can't have more than 2 active tracers!"))
end

#####
##### Run benchmarks
#####

# Each test case specifies (number of active tracers, number of passive tracers)
test_cases = [(0, 0), (0, 1), (0, 2), (1, 0), (2, 0), (2, 3), (2, 5), (2, 10)]

for arch in archs, test_case in test_cases
    N = arch isa CPU ? (32, 32, 32) : (256, 256, 128)
    na, np = test_case
    tracers = tracer_list(na, np)

    grid = RegularCartesianGrid(size=N, extent=(1, 1, 1))
    model = IncompressibleModel(architecture=arch, float_type=FT, grid=grid,
                                buoyancy=na2buoyancy(na), tracers=tracers)

    time_step!(model, 1)  # precompile

    bname =  benchmark_name(N, "$na active + $(lpad(np, 2)) passive", arch, FT)
    @printf("Running benchmark: %s...\n", bname)
    for i in 1:Nt
        @timeit timer bname time_step!(model, 1)
    end
end

#####
##### Print benchmark results
#####

println()
println(oceananigans_versioninfo())
println(versioninfo_with_gpu())
print_timer(timer, title="Tracer benchmarks", sortby=:name)
println()


"""
Calculates the exact boundingbox of a Scene/Plot, without considering any transformation
"""
function raw_boundingbox(x::Atomic)
    bb = data_limits(x)
end
rootparent(x) = rootparent(parent(x))
rootparent(x::Scene) = x

function raw_boundingbox(x::Annotations)
    bb = raw_boundingbox(x.plots)
    inv(modelmatrix(rootparent(x))) * bb
end
function raw_boundingbox(x::Combined)
    raw_boundingbox(x.plots)
end
function boundingbox(x)
    raw_boundingbox(x)
end

function combined_modelmatrix(x)
    m = Mat4f0(I)
    while true
        m = modelmatrix(x) * m
        if parent(x) !== nothing && parent(x) isa Combined
            x = parent(x)
        else
            break
        end
    end
    m
end

function modelmatrix(x)
    t = transformation(x)
    transformationmatrix(t.translation[], t.scale[], t.rotation[])
end

function boundingbox(x::Atomic)
    bb = raw_boundingbox(x)
    combined_modelmatrix(x) * bb
end

boundingbox(scene::Scene) = raw_boundingbox(scene)
function raw_boundingbox(scene::Scene)
    if scene[Axis] !== nothing
        return raw_boundingbox(scene[Axis])
    elseif scene.limits[] !== automatic
        return scene_limits(scene)
    elseif cameracontrols(scene) == EmptyCamera()
        # Empty camera means this is a parent scene that itself doesn't display anything
        return raw_boundingbox(scene.children)
    else
        plots = plots_from_camera(scene)
        children = filter(scene.children) do child
            child.camera == scene.camera
        end
        return raw_boundingbox([plots; children])
    end
end

function raw_boundingbox(plots::Vector)
    isempty(plots) && return FRect3D()
    plot_idx = iterate(plots)
    bb = FRect3D()
    while plot_idx !== nothing
        plot, idx = plot_idx
        plot_idx = iterate(plots, idx)
        # isvisible(plot) || continue
        bb2 = boundingbox(plot)
        isfinite(bb) || (bb = bb2)
        isfinite(bb2) || continue
        bb = union(bb, bb2)
    end
    bb
end

function project_widths(matrix, vec)
    pr = project(matrix, vec)
    zero = project(matrix, zeros(typeof(vec)))
    pr - zero
end

function boundingbox(x::Text, text::String)
    position = to_value(x[:position])
    @get_attribute x (textsize, font, align, rotation)
    bb = boundingbox(text, position, textsize, font, align, rotation, modelmatrix(x))
    pm = inv(transformationmatrix(parent(x))[])
    wh = widths(bb)
    whp = project_widths(pm, wh)
    aoffset = whp .* to_ndim(Vec3f0, align, 0f0)
    return FRect3D(minimum(bb) .- aoffset, whp)
end

boundingbox(x::Text) = boundingbox(x, to_value(x[1]))

function boundingbox(
        text::String, position, textsize;
        font = "default", align = (:left, :bottom), rotation = 0.0
    )
    boundingbox(
        text, position, textsize,
        to_font(font), to_align(align), to_rotation(rotation)
    )

end

# function boundingbox(
#         text::String, position, textsize, fonts,
#         align, rotation, model = Mat4f0(I)
#     )
#     isempty(text) && return FRect3D()
#     pos_per_char = !isa(position, VecTypes)
#
#     start_pos = Vec(pos_per_char ? first(position) : position)
#     start_pos3d = project(model, to_ndim(Vec3f0, start_pos, 0.0))
#     bb = FRect3D(start_pos3d, Vec3f0(0))
#
#     if pos_per_char
#         broadcast_foreach(position, textsize, fonts, collect(text)) do pos, scale, font, char
#             rect, extent = FreeTypeAbstraction.metrics_bb(char, font, scale)
#             bb = union(FRect3D(rect) + to_ndim(Vec3f0, pos, 0.0), bb)
#             @show pos scale
#         end
#     else
#         y_advance = 0.0
#         line_advance = FreeTypeAbstraction.get_extent(fonts, 'x').advance[2]
#         for line in split(text, r"(\r\n|\r|\n)")
#             rectangles = FreeTypeAbstraction.glyph_rects(line, fonts, textsize)
#             bb2d = reduce(union, rectangles)
#             bb2d = bb2d + Vec2f0(0, y_advance)
#             bb = union(bb, FRect3D(bb2d))
#             y_advance += line_advance
#             @show y_advance
#         end
#     end
#     return bb
# end

function boundingbox(
        text::String, position, textsize, font,
        align, rotation, model = Mat4f0(I)
    )
    atlas = get_texture_atlas()
    N = length(text)
    ctext_state = iterate(text)
    ctext_state === nothing && return FRect3D()
    pos_per_char = !isa(position, VecTypes)
    start_pos = Vec(pos_per_char ? first(position) : position)
    start_pos2D = to_ndim(Point2f0, start_pos, 0.0)
    last_pos = Point2f0(0, 0)
    start_pos3d = project(model, to_ndim(Vec3f0, start_pos, 0.0))
    bb = FRect3D(start_pos3d, Vec3f0(0))
    broadcast_foreach(1:N, rotation, font, textsize) do i, rotation, font, scale
        c, text_state = ctext_state
        ctext_state = iterate(text, text_state)
        # TODO fix center + align + rotation
        if c != '\r'
            pos = if pos_per_char
                to_ndim(Vec3f0, position[i], 0.0)
            else
                last_pos = calc_position(last_pos, Point2f0(0, 0), atlas, c, font, scale)
                start_pos3d .+ (rotation * to_ndim(Vec3f0, last_pos, 0.0))
            end
            s = glyph_scale!(atlas, c, font, scale)
            srot = rotation * to_ndim(Vec3f0, s, 0.0)
            bb = GeometryTypes.update(bb, pos)
            bb = GeometryTypes.update(bb, pos .+ srot)
        end
    end
    return bb
end


using Dates
using LibSerialPort
using SnakeIterator

include("utils.jl")

portname = LibSerialPort.get_port_list()[1]
baudrate = 115_200 # Creality Ender 3D uses a non standard baud rate

origin = Point(0, 0, 55, "origin")
plate = Point(125, 152, 29, "plate top left")
waste = Point(65, plate.y, 55, "waste container")

# let 3D printer initialize, it reboots when we setup the serial connection
sp = LibSerialPort.open(portname, baudrate)
sleep(7)

# move head up and calibrate the origin

write(sp, "G1 Z55\n")
write(sp, "G28\n")

# raise head and wait, this is useful for priming the lines
moveto(sp, origin; wait=true)

wells = create_wells(plate, well_spacing = 9.0, well_depth = 11.0)

groups = [
    wells[:, 1:6],
    wells[:, 7:12]
    ]

for group in groups
    println("Fractionating first tube")

    moveto(sp, plate)
    wait_for_key("start flow now, press any key")

    waste_start = now()
    moveto(sp, waste; wait = true)
    waste_stop = now()
    println("$(waste_stop - waste_start)")

    start_pos = first(group)
    start_pos = Point(start_pos.x, start_pos.y, "Start pos")
    moveto(sp, start_pos)
    sleep(2)

    prev_pos = start_pos

    # visit each well in a group in a snaking pattern
    for well in snake(group)
        println(well.name)
        moveto(sp, well; safe_height=plate.z)
        sleep(13)
        prev_pos = well
    end

    moveto(sp, Point(prev_pos.x, prev_pos.y + 9, waste.z*1.0, "outside"))
    moveto(sp, waste; wait=true)
end