#===== KernelAbstractions =====#

t4 = time()
using KernelAbstractions

using Tracker

GRAD = :Tracker
_gradient(x...) = Tracker.gradient(x...)

@testset "KernelAbstractions + gradients" begin
    A = (rand(3,4));
    B = (rand(4,5));
    @tullio C[i,k] := A[i,j] * B[j,k]  threads=false  # verbose=2
    @test C ≈ A * B

    @tullio threads=false # else KernelAbstractions CPU kernels not used
    include("gradients.jl")
    @tullio threads=true

    for sy in Tullio.SYMBOLS
        @test !isdefined(@__MODULE__, sy)
    end
end

using CUDA

if is_buildkite
    # If we are on Buildkite, we should assert that we have a CUDA GPU available
    @test CUDA.has_cuda_gpu()
end

if CUDA.has_cuda_gpu()
    @info "===== found a GPU, starting CUDA tests ====="
    @testset "===== CUDA tests on GPU =====" begin
        include("cuda.jl")
    end
end

@info @sprintf("KernelAbstractions tests took %.1f seconds", time()-t4)

@tullio cuda=false


module GenerativeAD

using DrWatson
using StatsBase
using Flux
using Statistics
using ImageTransformations
using MLDataPattern

include("datasets/Datasets.jl")
using .Datasets: load_data

include("evaluation/Evaluation.jl")

include("experiments.jl")
include("bayesian_opt.jl")
include("models/Models.jl")

end #module


const print_levels = Dict(
    :silent => 0,
    :minimal => 1,
    :partial => 2,
    :verbose => 3,
)

    # # Tabu times
    # get!(s, :tabu_time, length_vars(s) ÷ 2) # 10?
    # get!(s, :local_tabu, setting(s, :tabu_time) ÷ 2)
    # get!(s, :δ_tabu, setting(s, :tabu_time) - setting(s, :local_tabu))# 20-30

"""
    Options()
# Arguments:
- `dynamic::Bool`: is the model dynamic?
- `iteration::Union{Int, Float64}`: limit on the number of iterations
- `print_level::Symbol`: verbosity to choose among `:silent`, `:minimal`, `:partial`, `:verbose`
- `solutions::Int`: number of solutions to return
- `specialize::Bool`: should the types of the model be specialized or not. Usually yes for static problems. For dynamic in depends if the user intend to introduce new types. The specialized model is about 10% faster.
- `tabu_time::Int`: DESCRIPTION
- `tabu_local::Int`: DESCRIPTION
- `tabu_delta::Float64`: DESCRIPTION
- `threads::Int`: Number of threads to use
- `time_limit::Float64`: time limit in seconds
- `function Options(; dynamic = false, iteration = 10000, print_level = :minimal, solutions = 1, specialize = !dynamic, tabu_time = 0, tabu_local = 0, tabu_delta = 0.0, threads = typemax(0), time_limit = Inf)

```julia
# Setting options in JuMP syntax: print_level, time_limit, iteration
model = Model(CBLS.Optimizer)
set_optimizer_attribute(model, "iteration", 100)
set_optimizer_attribute(model, "print_level", :verbose)
set_time_limit_sec(model, 5.0)
```
"""
mutable struct Options
    dynamic::Bool
    info_path::String
    iteration::Union{Int,Float64}
    print_level::Symbol
    solutions::Int
    specialize::Bool
    tabu_time::Int
    tabu_local::Int
    tabu_delta::Float64
    threads::Int
    time_limit::Float64 # seconds

    function Options(;
        dynamic=false,
        info_path="",
        iteration=10000,
        print_level=:minimal,
        solutions=1,
        specialize=!dynamic,
        tabu_time=0,
        tabu_local=0,
        tabu_delta=0.0,
        threads=typemax(0),
        time_limit= 60, # seconds
    )
        ds_str = "The model types are specialized to the starting domains, constraints," *
        " and objectives types. Dynamic elements that add a new type will raise an error!"
        dynamic && specialize && @warn ds_str

        notds_str = "The solver types are not specialized in a static model context," *
        " which is sub-optimal."
        !dynamic && !specialize && @info notds_str

        itertime_str = "Both iteration and time limits are disabled. " *
        "Optimization runs will run infinitely."
        iteration == Inf && time_limit == Inf && @warn itertime_str

        new(
            dynamic,
            info_path,
            iteration,
            print_level,
            solutions,
            specialize,
            tabu_time,
            tabu_local,
            tabu_delta,
            threads,
            time_limit,
        )
    end
end

"""
    _verbose(settings, str)
Temporary logging function. #TODO: use better log instead (LoggingExtra.jl)
"""
function _verbose(options, str)
    pl = options.print_level
    print_levels[pl] ≥ 3 && (@info str)
end

"""
    _dynamic(options) = begin

DOCSTRING
"""
_dynamic(options) = options.dynamic

"""
    _dynamic!(options, dynamic) = begin

DOCSTRING
"""
_dynamic!(options, dynamic) = options.dynamic = dynamic

"""
    _info_path(options, path)

DOCSTRING
"""
_info_path(options) = options.info_path

"""
    _info_path!(options, iterations) = begin

DOCSTRING
"""
_info_path!(options, path) = options.info_path = path

"""
    _iteration(options) = begin

DOCSTRING
"""
_iteration(options) = options.iteration

"""
    _iteration!(options, iterations) = begin

DOCSTRING
"""
_iteration!(options, iterations) = options.iteration = iterations

"""
    _print_level(options) = begin

DOCSTRING
"""
_print_level(options) = options.print_level

"""
    _print_level!(options, level) = begin

DOCSTRING
"""
_print_level!(options, level) = options.print_level = level

"""
    _solutions(options) = begin

DOCSTRING
"""
_solutions(options) = options.solutions

"""
    _solutions!(options, solutions) = begin

DOCSTRING
"""
_solutions!(options, solutions) = options.solutions = solutions

"""
    _specialize(options) = begin

DOCSTRING
"""
_specialize(options) = options.specialize

"""
    _specialize!(options, specialize) = begin

DOCSTRING
"""
_specialize!(options, specialize) = options.specialize = specialize

"""
    _tabu_time(options) = begin

DOCSTRING
"""
_tabu_time(options) = options.tabu_time

"""
    _tabu_time!(options, time) = begin

DOCSTRING
"""
_tabu_time!(options, time) = options.tabu_time = time

"""
    _tabu_local(options) = begin

DOCSTRING
"""
_tabu_local(options) = options.tabu_local

"""
    _tabu_local!(options, time) = begin

DOCSTRING
"""
_tabu_local!(options, time) = options.tabu_local = time

"""
    _tabu_delta(options) = begin

DOCSTRING
"""
_tabu_delta(options) = options.tabu_delta

"""
    _tabu_delta!(options, time) = begin

DOCSTRING
"""
_tabu_delta!(options, time) = options.tabu_delta = time

"""
    _threads(options) = begin

DOCSTRING
"""
_threads(options) = options.threads

"""
    _threads!(options, threads) = begin

DOCSTRING
"""
_threads!(options, threads) = options.threads = threads

"""
    _time_limit(options) = begin

DOCSTRING
"""
_time_limit(options) = options.time_limit

"""
    _time_limit!(options, time::Time) = begin

DOCSTRING
"""
_time_limit!(options, time) = options.time_limit = time

function set_option!(options, name, value)
    eval(Symbol("_" * name * "!"))(options, value)
end

function get_option(options, name)
    eval(Symbol("_" * name))(options)
end


# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
# ==============================================================================
#
#   Tests related to ECI to ECI transformations using satellite state vector.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# References
# ==============================================================================
#
#   [1] Vallado, D. A (2013). Fundamentals of Astrodynamics and Applications.
#       Microcosm Press, Hawthorn, CA, USA.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

# Get the current EOP Data.
#
# TODO: The EOP data obtained from IERS website does not match the values in the
# examples in [1]. However, this should be enough, because 1) the individual
# functions at the low level are tested using the same values of [1], and 2) the
# difference is smaller than 30 cm.

eop_iau1980  = read_iers_eop("./eop_IAU1980.txt",  Val(:IAU1980))
eop_iau2000a = read_iers_eop("./eop_IAU2000A.txt", Val(:IAU2000A))

# File: ./src/transformations/sv_eci_to_eci.jl
# ============================================

################################################################################
#                                  IAU-76/FK5
################################################################################

# Functions: sv_eci_to_eci
# ------------------------

# The rotations functions were already heavily tested on `eci_to_eci.jl` file.
# Hence, here we will do only some minor testing involving the following
# transformations:
#
#   GCRF <=> J2000 (FK5)
#   MOD  <=> TOD   (FK5)
#   GCRF <=> CIRS  (IAU2006)

# GCRF <=> J2000
# ==============

################################################################################
#                                 Test Results
################################################################################
#
# Scenario 01
# ===========
#
# Example 3-15: Performing IAU-76/FK5 reduction.
#
# According to this example and Table 3-6, using:
#
#   UTC    = April 6, 2004, 07:51:28.386009
#   r_gcrf = 5102.50895790   i + 6123.01140070   j + 6378.13692820   k [km]
#   v_gcrf =   -4.7432201570 i +    0.7905364970 j +    5.5337557270 k [km/s]
#
# one gets:
#
#   r_j2000 = 5102.50960000   i + 6123.01152000   j + 6378.13630000   k [km]
#   v_j2000 =   -4.7432196000 i +    0.7905366000 j +    5.5337561900 k [km/s]
#
################################################################################

@testset "Function sv_eci_to_eci GCRF <=> J2000" begin
    JD_UTC = date_to_jd(2004, 4, 6, 7, 51, 28.386009)

    ## GCRF => J2000
    ## =============

    r_gcrf   = [5102.50895790; 6123.01140070; 6378.13692820]
    v_gcrf   = [-4.7432201570; 0.7905364970; 5.5337557270]
    sv_gcrf  = orbsv(JD_UTC, r_gcrf, v_gcrf)
    sv_j2000 = sv_eci_to_eci(sv_gcrf, GCRF(), J2000(), JD_UTC, eop_iau1980)

    @test sv_j2000.t === JD_UTC

    @test sv_j2000.r[1] ≈ +5102.50960000 atol=1e-4
    @test sv_j2000.r[2] ≈ +6123.01152000 atol=1e-4
    @test sv_j2000.r[3] ≈ +6378.13630000 atol=1e-4

    @test sv_j2000.v[1] ≈ -4.7432196000  atol=1e-7
    @test sv_j2000.v[2] ≈ +0.7905366000  atol=1e-7
    @test sv_j2000.v[3] ≈ +5.5337561900  atol=1e-7

    ## J2000 => GCRF
    ## =============

    r_j2000  = [5102.50960000; 6123.01152000; 6378.13630000]
    v_j2000  = [-4.7432196000; 0.7905366000; 5.5337561900]
    sv_j2000 = orbsv(JD_UTC, r_j2000, v_j2000)
    sv_gcrf  = sv_eci_to_eci(sv_j2000, J2000(), GCRF(), JD_UTC, eop_iau1980)

    @test sv_gcrf.t === JD_UTC

    @test sv_gcrf.r[1] ≈ +5102.50895790 atol=1e-4
    @test sv_gcrf.r[2] ≈ +6123.01140070 atol=1e-4
    @test sv_gcrf.r[3] ≈ +6378.13692820 atol=1e-4

    @test sv_gcrf.v[1] ≈ -4.7432201570  atol=1e-7
    @test sv_gcrf.v[2] ≈ +0.7905364970  atol=1e-7
    @test sv_gcrf.v[3] ≈ +5.5337557270  atol=1e-7

end

# MOD <=> TOD
# ============

################################################################################
#                                 Test Results
################################################################################
#
# Scenario 01
# ===========
#
# Example 3-15: Performing IAU-76/FK5 reduction.
#
# According to this example and Table 3-6, using:
#
#   UTC    = April 6, 2004, 07:51:28.386009
#   r_mod = 5094.02837450   i + 6127.87081640   j + 6380.24851640   k [km]
#   v_mod =   -4.7462630520 i +    0.7860140450 j +    5.5317905620 k [km/s]
#
# one gets:
#
#   r_tod = 5094.51620300   i + 6127.36527840   j + 6380.34453270   k [km]
#   v_tod =   -4.7460883850 i +    0.7860783240 j +    5.5319312880 k [km/s]
#
################################################################################

@testset "Function sv_eci_to_eci MOD <=> TOD" begin
    JD_UTC = date_to_jd(2004, 4, 6, 7, 51, 28.386009)

    ## MOD => TOD
    ## ===========

    r_mod  = [5094.02837450; 6127.87081640; 6380.24851640]
    v_mod  = [-4.7462630520; 0.7860140450; 5.5317905620]
    sv_mod = orbsv(JD_UTC, r_mod, v_mod)
    sv_tod = sv_eci_to_eci(sv_mod, MOD(), JD_UTC, TOD(), JD_UTC, eop_iau1980)

    @test sv_tod.t === JD_UTC

    @test sv_tod.r[1] ≈ +5094.51620300 atol=1e-4
    @test sv_tod.r[2] ≈ +6127.36527840 atol=1e-4
    @test sv_tod.r[3] ≈ +6380.34453270 atol=1e-4

    @test sv_tod.v[1] ≈ -4.7460883850  atol=1e-7
    @test sv_tod.v[2] ≈ +0.7860783240  atol=1e-7
    @test sv_tod.v[3] ≈ +5.5319312880  atol=1e-7

    ## TOD => MOD
    ## ===========

    r_tod  = [5094.51620300; 6127.36527840; 6380.34453270]
    v_tod  = [-4.7460883850; 0.7860783240; 5.5319312880]
    sv_tod = orbsv(JD_UTC, r_tod, v_tod)
    sv_mod = sv_eci_to_eci(sv_tod, TOD(), JD_UTC, MOD(), JD_UTC, eop_iau1980)

    @test sv_mod.t === JD_UTC

    @test sv_mod.r[1] ≈ +5094.02837450 atol=1e-4
    @test sv_mod.r[2] ≈ +6127.87081640 atol=1e-4
    @test sv_mod.r[3] ≈ +6380.24851640 atol=1e-4

    @test sv_mod.v[1] ≈ -4.7462630520  atol=1e-7
    @test sv_mod.v[2] ≈ +0.7860140450  atol=1e-7
    @test sv_mod.v[3] ≈ +5.5317905620  atol=1e-7

end

################################################################################
#                                IAU-2006/2010
################################################################################

## GCRF <=> CIRS
## =============

################################################################################
#                                 Test Results
################################################################################
#
# Scenario 01
# ===========
#
# Example 3-14: Performing an IAU-2000 reduction [1, p. 220]
#
# According to this example and Table 3-6, using:
#
# NOTE: It seems that the results in the Example 3-14 is computed **without**
# the `dX` and `dY` corrections, whereas in the Table 3-6 they are computed
# **with** the corrections.
#
#   UTC    = April 6, 2004, 07:51:28.386009
#   r_cirs = -5100.01840470   i + 6122.78636480   j + 6380.34453270   k [km]
#   v_cirs =    -4.7453803300 i -    0.7903414530 j +    5.5319312880 k [km/s]
#
# one gets the following (this is the result in Table 3-6):
#
#   r_gcrf = 5102.50895290  i + 6123.01139910 j + 6378.13693380 k [km]
#   v_gcrf =  -4.7432201610 i + 0.7905364950  j + 5.5337557240  k [km/s]
#
################################################################################

@testset "Function sv_eci_to_eci GCRF <=> CIRS" begin
    JD_UTC = date_to_jd(2004, 4, 6, 7, 51, 28.386009)

    ## GCRF => CIRS
    ## ===========

    r_gcrf  = [5102.50895290; 6123.01139910; 6378.13693380]
    v_gcrf  = [-4.7432201610; 0.7905364950; 5.5337557240]
    sv_gcrf = orbsv(JD_UTC, r_gcrf, v_gcrf)
    sv_cirs = sv_eci_to_eci(sv_gcrf, GCRF(), CIRS(), JD_UTC, eop_iau2000a)

    @test sv_cirs.t === JD_UTC

    @test sv_cirs.r[1] ≈ +5100.01840470 atol=1e-4
    @test sv_cirs.r[2] ≈ +6122.78636480 atol=1e-4
    @test sv_cirs.r[3] ≈ +6380.34453270 atol=1e-4

    @test sv_cirs.v[1] ≈ -4.7453803300  atol=1e-7
    @test sv_cirs.v[2] ≈ +0.7903414530  atol=1e-7
    @test sv_cirs.v[3] ≈ +5.5319312880  atol=1e-7

    ## CIRS => GCRF
    ## ===========

    r_cirs  = [+5100.01840470; +6122.78636480; +6380.34453270]
    v_cirs  = [-4.7453803300; +0.7903414530; +5.5319312880]
    sv_cirs = orbsv(JD_UTC, r_cirs, v_cirs)
    sv_gcrf = sv_eci_to_eci(sv_cirs, CIRS(), GCRF(), JD_UTC, eop_iau2000a)

    @test sv_gcrf.t === JD_UTC

    @test sv_gcrf.r[1] ≈ +5102.50895290 atol=1e-4
    @test sv_gcrf.r[2] ≈ +6123.01139910 atol=1e-4
    @test sv_gcrf.r[3] ≈ +6378.13693380 atol=1e-4

    @test sv_gcrf.v[1] ≈ -4.7432201610  atol=1e-7
    @test sv_gcrf.v[2] ≈ +0.7905364950  atol=1e-7
    @test sv_gcrf.v[3] ≈ +5.5337557240  atol=1e-7

end


## Gather data back from RDA files
R"""
load(file = "generation_of_data_for_econ_analysis/waning_imm_scenario_1/kenya_inc_deaths_no_vac_w_vac.rda")
"""
kenya_inc_deaths_sc_1 = @rget kenya_inc_deaths
kenya_inc_deaths_sc_1 = sum(kenya_inc_deaths_sc_1, dims = 2)[:, 1, :]
kenya_inc_deaths = 0
R"""
kenya_inc_deaths = 0 #Free memory
"""
R"""
load(file = "generation_of_data_for_econ_analysis/waning_imm_scenario_2/kenya_inc_deaths_30_perc_w_vac.rda")
"""
kenya_inc_deaths_sc_2 = @rget kenya_inc_deaths
kenya_inc_deaths_sc_2 = sum(kenya_inc_deaths_sc_2, dims = 2)[:, 1, :]
kenya_inc_deaths = 0
R"""
kenya_inc_deaths = 0 #Free memory
"""
R"""
load(file = "generation_of_data_for_econ_analysis/waning_imm_scenario_3/kenya_inc_deaths_50_perc_w_vac.rda")
"""
kenya_inc_deaths_sc_3 = @rget kenya_inc_deaths
kenya_inc_deaths_sc_3 = sum(kenya_inc_deaths_sc_3, dims = 2)[:, 1, :]
kenya_inc_deaths = 0
R"""
kenya_inc_deaths = 0 #Free memory
"""
R"""
load(file = "generation_of_data_for_econ_analysis/waning_imm_scenario_4/kenya_inc_deaths_70_perc_w_vac.rda")
"""
kenya_inc_deaths_sc_4 = @rget kenya_inc_deaths
kenya_inc_deaths_sc_4 = sum(kenya_inc_deaths_sc_4, dims = 2)[:, 1, :]
kenya_inc_deaths = 0
R"""
kenya_inc_deaths = 0 #Free memory
"""
R"""
load(file = "generation_of_data_for_econ_analysis/waning_imm_scenario_5/kenya_inc_deaths_30_perc_rapid_w_vac.rda")
"""
kenya_inc_deaths_sc_5 = @rget kenya_inc_deaths
kenya_inc_deaths_sc_5 = sum(kenya_inc_deaths_sc_5, dims = 2)[:, 1, :]
kenya_inc_deaths = 0
R"""
kenya_inc_deaths = 0 #Free memory
"""
R"""
load(file = "generation_of_data_for_econ_analysis/waning_imm_scenario_6/kenya_inc_deaths_50_perc_rapid_w_vac.rda")
"""
kenya_inc_deaths_sc_6 = @rget kenya_inc_deaths
kenya_inc_deaths_sc_6 = sum(kenya_inc_deaths_sc_6, dims = 2)[:, 1, :]
kenya_inc_deaths = 0
R"""
kenya_inc_deaths = 0 #Free memory
"""
R"""
load(file = "generation_of_data_for_econ_analysis/waning_imm_scenario_7/kenya_inc_deaths_70_perc_rapid_w_vac.rda")
"""
kenya_inc_deaths_sc_7 = @rget kenya_inc_deaths
kenya_inc_deaths_sc_7 = sum(kenya_inc_deaths_sc_7, dims = 2)[:, 1, :]
kenya_inc_deaths = 0
R"""
kenya_inc_deaths = 0 #Free memory
"""


export AngularMeasurements, AngularSpectra, parent, angles, angleidx, closest_angle, between_indices
import Base.parent
# export maximumamp, minimumamp, normalizeamp, normalizeamp!
# export normalizeampangle, normalizeampangle!
# export normalizeampfreq, normalizeampfreq!
# export relativeto, relativeto!, smooth, interpolate, mirror, mirror!

struct AngularMeasurements{T,P<:AbstractSampleArray{T},A<:AbstractVector{<:Real}} <: AbstractSampleArray{T}
    parent::P
    angles::A

    function AngularMeasurements(parent, angles)
        nchannels(parent) == length(angles) || throw(ArgumentError("number of channels ($(nchannels(a))) != number of angles ($(length(angles)))"))
        s = deepcopy(parent) # because we sort it below
        a = normdeg.(Base.vect(angles...)) # normalize degrees to (-180°, 180°]
        length(a) == length(unique(a)) || throw(ArgumentError("non-unique angles!"))
        idxs = sortperm(a)
        new{eltype(parent), typeof(parent), typeof(angles)}(s[:, idxs], collect(a[idxs]))
    end
end

const AngularSpectra{T,P<:AbstractSpectrumArray{T},A<:AbstractVector{<:Real}} = AngularMeasurements{T,P,A}

function Interpolations.interpolate(am::AngularMeasurements, template::T) where {U, T <: AbstractSampleArray{U}}
    AngularMeasurements(Interpolations.interpolate(parent(am), template), angles(am))
end

Base.parent(am::AngularMeasurements) = am.parent
angles(am::AngularMeasurements) = am.angles

SampleArrays.rate(am::AngularMeasurements) = rate(parent(am))
SampleArrays.data(am::AngularMeasurements) = data(parent(am))
SampleArrays.data_no0(as::AngularSpectra) = data_no0(parent(as))

SampleArrays.domain(am::AngularMeasurements) = domain(parent(am))
SampleArrays.domain_no0(as::AngularSpectra) = domain_no0(parent(as))

SampleArrays.names(am::AngularMeasurements) = names(parent(am))
SampleArrays.names!(am::AngularMeasurements) = names!(parent(am))
SampleArrays.nchannels(am::AngularMeasurements) = nchannels(parent(am))
SampleArrays.nframes(am::AngularMeasurements) = nframes(parent(am))

SampleArrays.ntimeframes(as::AngularSpectra) = ntimeframes(parent(as))

function Base.show(io::IO, t::MIME"text/plain", am::AngularMeasurements)
    as = angles(am)
    println(io, "AngularMeasurements: $(minimum(as))°..$(maximum(as))°:")
    show(io, t, parent(am))
end

function angleidx(angles::AbstractVector{<:Real}, α; func=(==))
    # find first angle index (and the actual angle), for which func is true
    for (i, a) in enumerate(angles)
        func(a, α) && return (idx=i, α=a)
    end
    error("angle ($(α)) not found!")
end

angleidx(am::AngularMeasurements, α) = angleidx(angles(am), α)

function closest_angle(angles::AbstractVector{<:Real}, α)
    closeness = abs.(angles .- α)
    _, idx = findmin(closeness) # index of angle spectra closest to the angle α
    (idx=idx, α=angles[idx])
end

closest_angle(am::AngularMeasurements, α) = closest_angle(angles(am), α)

function between_indices(am::AngularMeasurements, α)
    α = normdeg(α)
    nc = nchannels(am)
    if nc == 1
        return 1, 1
    end
    as = [am.angles..., am.angles[1]]
    for i in 1:nc
        j = i % nc + 1
        d1, d2 = deltadeg(am.angles[i], α), deltadeg(α, am.angles[j])
        if d1 >= 0 && d2 >= 0
            if d1 ≈ 0
                return i, i
            elseif d2 ≈ 0
                return j, j
            end
            return i, j
        end
    end
end

function Base.getindex(am::AngularMeasurements, I::R, J::S) where {R <: SampleArrays.FrameIndex, S <: SampleArrays.ChannelIndex}
    idxs = SampleArrays.tochannelidx(parent(am), J)
    p = parent(am)[I, idxs]
    AngularMeasurements(p, angles(am)[idxs])
end
Base.getindex(am::AngularMeasurements, I::R) where {R <: SampleArrays.FrameIndex} = am[I, :]

Base.getindex(am::AngularMeasurements, α::AbstractFloat) = am[:, angleidx(am, α).idx]
function Base.getindex(am::AngularMeasurements, α::AbstractFloat, β::AbstractFloat)
    α <= β || error("α > β!")
    am[:, last(between_indices(am, α)):first(between_indices(am, β))]
end

function Base.similar(am::AngularMeasurements, t::Type, dims::Dims)
    AngularMeasurements(similar(parent(am), t, dims), angles(am))
end

# maximumamp(ds::DirectivitySpectra) = maximum(abs.(ds.spectra))
# minimumamp(ds::DirectivitySpectra) = minimum(abs.(ds.spectra))

# normalizeamp(ds::DirectivitySpectra) = normalizeamp!(deepcopy(ds))

# function normalizeamp!(ds::DirectivitySpectra)
#     maxamp = maximumamp(ds)
#     maxamp == 0 && return ds
#     ds.spectra ./= maxamp
#     return ds
# end

# normalizeampangle(ds::DirectivitySpectra) = normalizeampfreq!(deepcopy(ds))

# function normalizeampangle!(ds::DirectivitySpectra)
#     for α in 1:length(ds)
#         maxamp = maximum(abs.(ds.spectra[:, α]))
#         maxamp == 0 && continue
#         ds.spectra[:, α] ./= maxamp
#     end
#     return ds
# end

# normalizeampfreq(ds::DirectivitySpectra) = normalizeampfreq!(deepcopy(ds))

# function normalizeampfreq!(ds::DirectivitySpectra)
#     for fidx in 1:nfreqs(ds)
#         maxamp = maximum(abs.(data(spectra(ds))[fidx, :]))
#         maxamp == 0 && continue
#         data(spectra(ds))[fidx, :] ./= maxamp
#     end
#     return ds
# end

function Interpolations.interpolate(am::AngularMeasurements, α)
    i, j = between_indices(am, α)
    if i == j
        return am.parent[:, i]
    end
    β, γ = angles(am)[[i, j]]
    b, g = Float32(deltadeg(β, α)), Float32(deltadeg(α, γ)) # TODO better way to keep Float32 of parent(am)!
    @assert b >= 0 && g >= 0
    s = b + g
    @assert s > 0
    return mean([2g/s, 2b/s]' .* parent(am)[:, [i, j]])
end

# relativeto(ds::DirectivitySpectra, to) = relativeto!(deepcopy(ds), to)

# function relativeto!(ds::DirectivitySpectra, to::AbstractSpectrumArray)
#     @assert nchannels(to) == 1
#     to = deepcopy(to.data)
#     any(to .== 0.0) && @warn("relativeto!: zero in denominator!")
#     ds.spectra.data ./= to
#     return ds
# end

# relativeto!(ds::DirectivitySpectra, α::Real) = relativeto!(ds, ds[Float64(α)])

# """
#     mirror(ds::DirectivitySpectra)

# Mirror all spectra, i.e., append copies of spectra with negative angles (if not already in the collection).
# """
# function mirror(ds::DirectivitySpectra)
#     newspectra = []
#     newangles = deepcopy(ds.angles)
#     for i in 1:nchannels(ds)
#         α = normdeg(-newangles[i])
#         if α ∉ newangles
#             push!(newangles, α)
#             push!(newspectra, ds[i])
#         end
#     end
#     newspectra = hcat(ds.spectra, hcat(newspectra...))
#     idxs = sortperm(newangles)
#     newspectra .= newspectra[:, idxs]
#     newangles .= newangles[idxs]
#     return DirectivitySpectra(newspectra, newangles)
# end

# smooth(ds::DirectivitySpectra{T}, octave::Real = 1//3) where {T <: RFFTSpectrumArray} =
#     DirectivitySpectra(smooth(ds.spectra, octave), angles(ds))

# mean(ds::DirectivitySpectra) = mean(ds.spectra)


@doc """
Calculate a causal graph that is consistent with the data:

* `G` : graph of potential connections
* `df` : data frame of observed values for nodes in `G`
""" ->
function pc_algo(G::GenericGraph, df::DataFrame)
end


module Examples

using ..AlphaZero

include("../games/tictactoe/main.jl")
export Tictactoe

include("../games/toetactic/main.jl")
export Toetactic

include("../games/connect-four/main.jl")
export ConnectFour

include("../games/grid-world/main.jl")
export GridWorld

include("../games/mancala/main.jl")
export Mancala

include("../games/gomoku/main.jl")
export Gomoku

const games = Dict(
    "grid-world" => GridWorld.GameSpec(),
    "tictactoe" => Tictactoe.GameSpec(),
    "toetactic" => Toetactic.GameSpec(),
    "connect-four" => ConnectFour.GameSpec(),
    "mancala" => Mancala.GameSpec(),
    "gomoku" => Gomoku.GameSpec(),
)
# "ospiel_ttt" => OSpielTictactoe.GameSpec()
# ospiel_ttt is added from openspiel_example.jl when OpenSpiel.jl is imported

const experiments = Dict(
    "grid-world" => GridWorld.Training.experiment,
    "tictactoe" => Tictactoe.Training.experiment,
    "toetactic" => Toetactic.Training.experiment,
    "connect-four" => ConnectFour.Training.experiment,
    "mancala" => Mancala.Training.experiment,
    "gomoku" => Gomoku.Training.experiment,
)
# "ospiel_ttt" => OSpielTictactoe.Training.experiment

end


struct NcFiles
    u::Union{NcFile,Nothing}        # zonal velocity
    v::Union{NcFile,Nothing}        # meridional velocity
    η::Union{NcFile,Nothing}        # sea surface height anomaly
    sst::Union{NcFile,Nothing}      # tracer / sea surface temperature
    q::Union{NcFile,Nothing}        # potentival vorticity
    ζ::Union{NcFile,Nothing}        # relative vorticity
    du::Union{NcFile,Nothing}       # tendency of u [m^2/s^2]
    dv::Union{NcFile,Nothing}       # tendency of v [m^2/s^2]
    dη::Union{NcFile,Nothing}       # tendency of η [m^2/s]
    iout::Array{Integer,1}          # output index, stored in array for mutability
end

"""Generator function for "empty" NcFiles struct."""
NcFiles(x::Nothing) = NcFiles(x,x,x,x,x,x,x,x,x,[0])

"""Generator function for NcFiles struct, creating the underlying netCDF files."""
function NcFiles(feedback::Feedback,S::ModelSetup)

    if S.parameters.output

        @unpack output_vars,compression_level = S.parameters
        P = S.parameters
        @unpack x_u,x_v,x_T,x_q = S.grid
        @unpack y_u,y_v,y_T,y_q = S.grid

        @unpack run_id,runpath = feedback

        ncu = if "u" in output_vars nc_create(x_u,y_u,"u",runpath,"m/s","zonal velocity",P) else nothing end
        ncv = if "v" in output_vars nc_create(x_v,y_v,"v",runpath,"m/s","meridional velocity",P) else nothing end
        ncη = if "η" in output_vars nc_create(x_T,y_T,"eta",runpath,"m","sea surface height",P) else nothing end
        ncsst = if "sst" in output_vars nc_create(x_T,y_T,"sst",runpath,"1","sea surface temperature",P) else nothing end
        ncq = if "q" in output_vars nc_create(x_q,y_q,"q",runpath,"1/(ms)","potential vorticity",P) else nothing end
        ncζ = if "ζ" in output_vars nc_create(x_q,y_q,"relvort",runpath,"1","relative vorticity",P) else nothing end
        ncdu = if "du" in output_vars nc_create(x_u,y_u,"du",runpath,"m^2/s^2","zonal velocity tendency",P) else nothing end
        ncdv = if "dv" in output_vars nc_create(x_v,y_v,"dv",runpath,"m^2/s^2","meridional velocity tendency",P) else nothing end
        ncdη = if "dη" in output_vars nc_create(x_T,y_T,"deta",runpath,"m^2/s","sea surface height tendency",P) else nothing end

        for nc in (ncu,ncv,ncη,ncsst,ncq,ncζ,ncdu,ncdv,ncdη)
            if nc != nothing
                NetCDF.putatt(nc,"t",Dict("units"=>"s","long_name"=>"time"))
                NetCDF.putatt(nc,"x",Dict("units"=>"m","long_name"=>"zonal coordinate"))
                NetCDF.putatt(nc,"y",Dict("units"=>"m","long_name"=>"meridional coordinate"))
            end
        end

        return NcFiles(ncu,ncv,ncη,ncsst,ncq,ncζ,ncdu,ncdv,ncdη,[0])
    else
        return NcFiles(nothing)
    end
end

"""Creates a netCDF file based on grid vectors x,y the variable name, its path, unit and long_name."""
function nc_create( x::Array{T,1},
                    y::Array{T,1},
                    name::String,
                    path::String,
                    unit::String,
                    long_name::String,
                    P::Parameter) where {T<:Real}

    @unpack compression_level = P

    xdim = NcDim("x",length(x),values=x)
    ydim = NcDim("y",length(y),values=y)
    tdim = NcDim("t",0,unlimited=true)

    var = NcVar(name,[xdim,ydim,tdim],t=Float32,compress=compression_level)
    tvar = NcVar("t",tdim,t=Int32)

    nc = NetCDF.create(joinpath(path,name*".nc"),[var,tvar],mode=NC_NETCDF4)
    # add missing_value although irrelevant for ncview compatibility
    NetCDF.putatt(nc,name,Dict("units"=>unit,"long_name"=>long_name,"missing_value"=>-999999f0))
    return nc
end

function output_nc!(i::Int,
                    ncs::NcFiles,
                    Prog::PrognosticVars,
                    Diag::DiagnosticVars,
                    S::ModelSetup)

    @unpack nout = S.grid
    @unpack output = S.parameters

    if i % nout == 0 && output

        iout = ncs.iout[1]+1        # unpack and increase
        ncs.iout[1] = iout          # pack

        @unpack halo,haloη,halosstx,halossty = S.grid
        @unpack f_q,ep,dtint = S.grid
        @unpack scale,scale_inv,scale_sst = S.constants

        # CUT OFF HALOS
        # As output is before copyto!(u,u0), take u0,v0,η0
        # Tendencies calculate from the last time step, du = u_n+1-u_n etc
        # WRITING THE VARIABLES
        if ncs.u != nothing
            @views u = Float32.(scale_inv*Diag.RungeKutta.u0[halo+1:end-halo,halo+1:end-halo])
            NetCDF.putvar(ncs.u,"u",u,start=[1,1,iout],count=[-1,-1,1])
        end
        if ncs.v != nothing
            @views v = Float32.(scale_inv*Diag.RungeKutta.v0[halo+1:end-halo,halo+1:end-halo])
            NetCDF.putvar(ncs.v,"v",v,start=[1,1,iout],count=[-1,-1,1])
        end
        if ncs.η != nothing
            @views η = Float32.(Diag.RungeKutta.η0[haloη+1:end-haloη,haloη+1:end-haloη])
            NetCDF.putvar(ncs.η,"eta",η,start=[1,1,iout],count=[-1,-1,1])
        end
        if ncs.sst != nothing
            @views sst = Float32.(Prog.sst[halosstx+1:end-halosstx,halossty+1:end-halossty]/scale_sst)
            NetCDF.putvar(ncs.sst,"sst",sst,start=[1,1,iout],count=[-1,-1,1])
        end
        if ncs.q != nothing
            @views q = Float32.(scale_inv*Diag.Vorticity.q[haloη+1:end-haloη,haloη+1:end-haloη])
            NetCDF.putvar(ncs.q,"q",q,start=[1,1,iout],count=[-1,-1,1])
        end
        if ncs.ζ != nothing
            @unpack dvdx,dudy = Diag.Vorticity
            @unpack f_q = S.grid
            @views ζ = Float32.((dvdx[2:end-1,2:end-1]-dudy[2+ep:end-1,2:end-1])./abs.(f_q))
            NetCDF.putvar(ncs.ζ,"relvort",ζ,start=[1,1,iout],count=[-1,-1,1])
        end
        if ncs.du != nothing
            @views u = Float32.(scale_inv*Diag.RungeKutta.u0[halo+1:end-halo,halo+1:end-halo])
            @views du = u-Float32.(scale_inv*Prog.u[halo+1:end-halo,halo+1:end-halo])
            NetCDF.putvar(ncs.du,"du",du,start=[1,1,iout],count=[-1,-1,1])
        end
        if ncs.dv != nothing
            @views v = Float32.(scale_inv*Diag.RungeKutta.v0[halo+1:end-halo,halo+1:end-halo])
            @views dv = v-Float32.(scale_inv*Prog.v[halo+1:end-halo,halo+1:end-halo])
            NetCDF.putvar(ncs.dv,"dv",dv,start=[1,1,iout],count=[-1,-1,1])
        end
        if ncs.dη != nothing
            @views η = Float32.(Diag.RungeKutta.η0[haloη+1:end-haloη,haloη+1:end-haloη])
            @views dη = η-Float32.(Prog.η[haloη+1:end-haloη,haloη+1:end-haloη])
            NetCDF.putvar(ncs.dη,"deta",dη,start=[1,1,iout],count=[-1,-1,1])
        end

        # WRITING THE TIME
        for nc in (ncs.u,ncs.v,ncs.η,ncs.sst,ncs.q,ncs.ζ,ncs.du,ncs.dv,ncs.dη)
            if nc != nothing
                NetCDF.putvar(nc,"t",Int64[i*dtint],start=[iout])
                NetCDF.sync(nc)         # sync to view netcdf while model is still running
            end
        end
    end
end

"""Closes netCDF and progress.txt files."""
function output_close!(ncs::NcFiles,feedback::Feedback,S::ModelSetup)

    @unpack output = S.parameters

    if output
        println("All data stored.")
        write(feedback.progress_txt,"All data stored.")
        close(feedback.progress_txt)
    end
end

"""Finds the first gap in a list of integers."""
function gap(a::Array{Int,1})
    try
        return minimum([i for i in minimum(a):maximum(a) if ~(i in a)])
    catch
        return maximum(a)+1
    end
end

"""Checks output folders to determine a 4-digit run id number."""
function get_run_id_path(S::ModelSetup)

    @unpack output,outpath,get_id_mode = S.parameters

    if output
        runlist = filter(x->startswith(x,"run"),readdir(outpath))
        existing_runs = [parse(Int,id[4:end]) for id in runlist]
        if length(existing_runs) == 0           # if no runfolder exists yet
            runpath = joinpath(outpath,"run0000")
            mkdir(runpath)
            return 0,runpath
        else                                    # create next folder
            if get_id_mode == "fill"  # find the smallest gap in runfolders
                run_id = gap(existing_runs)
                runpath = joinpath(outpath,"run"*@sprintf("%04d",run_id))
                mkdir(runpath)

            elseif get_id_mode == "specific" # specify the run_id as input argument
                @unpack run_id = S.parameters
                runpath = joinpath(outpath,"run"*@sprintf("%04d",run_id))
                try # create folder if not existent
                    mkdir(runpath)
                catch # else rm folder and create new one
                    rm(runpath,recursive=true)
                    mkdir(runpath)
                end

            elseif get_id_mode == "continue" # find largest folder and count one up
                run_id = maximum(existing_runs)+1
                runpath = joinpath(outpath,"run"*@sprintf("%04d",run_id))
                mkdir(runpath)
            else
                throw(error("Order '$get_id_mode' is not valid for get_run_id_path(), choose continue or fill."))
            end
            return run_id,runpath
        end
    else
        return 0,"no runpath"
    end
end


"""
Describes a hydropower unit.

Hydropower is both a purpose and a possible dam output. The optimisation process could ignore the purpose side of
hydropower, and only consider it as a possible (constant) output; but optimising this purpose needs the dedicated
penstock. As a consequence, the
"""
immutable HydropowerUnit
  id::Symbol
  max_discharge::Float64
  max_power::Float64
  efficiency::Float64

  # Private constructor, without units nor defaut values.
  function HydropowerUnit(id::Symbol, max_discharge::Float64, max_power::Float64, efficiency::Float64)
    if efficiency < 0. || efficiency > 1.
      error("Wrong efficiency: " * string(efficiency) * " not between 0.0 and 1.0")
    end

    new(id, max_discharge, max_power, efficiency)
  end
end

# Public constructor, with units and keyword arguments.
HydropowerUnit(;id::Symbol=:hydropower,
                max_discharge::Union{Discharge, Float64}=error("Missing maximum discharge"),
                max_power::Union{Power, Float64}=error("Missing maximum power not specified"),
                efficiency::Float64=error("Missing efficiency")
          ) = HydropowerUnit(id, _from_unitful(max_discharge), _from_unitful(max_power), efficiency)
HydropowerUnit(max_discharge::Union{Discharge, Float64}, max_power::Union{Power, Float64}, efficiency::Float64) =
  HydropowerUnit(max_discharge=max_discharge, max_power=max_power, efficiency=efficiency)

copy(hp::HydropowerUnit; id::Symbol=hp.id,
     max_discharge::Union{Discharge, Float64}=hp.max_discharge, max_power::Union{Power, Float64}=hp.max_power,
     efficiency::Float64=hp.efficiency) =
  # Always call the internal constructor!
  Reservoir(id, _from_unitful(max_discharge), _from_unitful(max_power), efficiency)

getId(hp::HydropowerUnit) = hp.id
getPurposeId(hp::HydropowerUnit) = symbol(getId(hp), "Purpose")
getDamOutputId(hp::HydropowerUnit) = symbol(getId(hp), "DamOutput")
getMaximumDischarge(hp::HydropowerUnit) = hp.max_discharge
getMaximumPower(hp::HydropowerUnit) = hp.max_power
getEfficiency(hp::HydropowerUnit) = hp.efficiency

id(hp::HydropowerUnit) = getId(hp)
purposeId(hp::HydropowerUnit) = getPurposeId(hp)
damOutputId(hp::HydropowerUnit) = getDamOutputId(hp)
maxDischarge(hp::HydropowerUnit) = getMaximumDischarge(hp)
maxPower(hp::HydropowerUnit) = getMaximumPower(hp)
efficiency(hp::HydropowerUnit) = getEfficiency(hp)


const DEFAULT_PRECISION = Ref{Int}(256)

"""
    precision(<:Union{Arf, Arb, Acb, arf_struct, arb_struct, acb_struct})
    precision(<:Ptr{<:Union{arf_struct, arb_struct, acb_struct}})
    precision(x::Union{arf_struct, arb_struct, acb_struct})
    precision(x::Ptr{<:Union{arf_struct, arb_struct, acb_struct}})
Get the default precision (in bits) currently used for `Arblib` arithmetic.
"""
Base.precision(::Type{<:Union{ArbTypes,ArbStructTypes}}) = DEFAULT_PRECISION[]
Base.precision(::Type{<:Ptr{<:ArbStructTypes}}) = DEFAULT_PRECISION[]

Base.precision(x::ArbStructTypes) = DEFAULT_PRECISION[]
Base.precision(x::Ptr{<:ArbStructTypes}) = DEFAULT_PRECISION[]
Base.precision(x::ArbTypes) = x.prec
# MagLike <: ArbTypes
Base.precision(x::MagLike) = DEFAULT_PRECISION[]
# disambiguation
Base.precision(::MagOrRef) = DEFAULT_PRECISION[]

@inline _precision(x::ArbTypes) = precision(x)
@inline _precision(x::BigFloat) = precision(x)
@inline _precision(z::Complex) = max(_precision(real(z)), _precision(imag(z)))
@inline _precision(
    (a, b)::Tuple{S,T},
) where {
    S<:Union{ArbTypes,BigFloat,Complex{<:Union{ArbTypes,BigFloat}}},
    T<:Union{ArbTypes,BigFloat,Complex{<:Union{ArbTypes,BigFloat}}},
} = max(_precision(a), _precision(b))
@inline _precision(
    (a, b)::Tuple{S,T},
) where {S<:Union{ArbTypes,BigFloat,Complex{<:Union{ArbTypes,BigFloat}}},T} = _precision(a)
@inline _precision(
    (a, b)::Tuple{S,T},
) where {S,T<:Union{ArbTypes,BigFloat,Complex{<:Union{ArbTypes,BigFloat}}}} = _precision(b)
@inline _precision((a, b)::Tuple{S,T}) where {S,T} = max(_precision(a), _precision(b))
@inline _precision(@nospecialize _) = DEFAULT_PRECISION[]

"""
    setprecision(::Type{<:Union{Arf, Arb, Acb}}, precision::Int)
Set the precision (in bits) to be used for `Arblib` arithmetic.
!!! warning
    This function is not thread-safe. It will affect code running on all threads, but
    its behavior is undefined if called concurrently with computations that use the
    setting.
"""
function Base.setprecision(::Type{<:ArbTypes}, precision::Integer)
    if precision < 2
        throw(DomainError(precision, "`precision` cannot be less than 2."))
    end
    DEFAULT_PRECISION[] = precision
    return precision
end

Base.setprecision(x::T, prec::Integer) where {T<:Union{Arf,ArfRef,Arb,ArbRef,Acb,AcbRef}} =
    T(x, prec = prec)
Base.setprecision(v::T, prec::Integer) where {T<:Union{ArbVector,ArbRefVector}} =
    T(v.arb_vec, prec)
Base.setprecision(v::T, prec::Integer) where {T<:Union{AcbVector,AcbRefVector}} =
    T(v.acb_vec, prec)
Base.setprecision(A::T, prec::Integer) where {T<:Union{ArbMatrix,ArbRefMatrix}} =
    T(A.arb_mat, prec)
Base.setprecision(A::T, prec::Integer) where {T<:Union{AcbMatrix,AcbRefMatrix}} =
    T(A.acb_mat, prec)

Base.setprecision(poly::ArbPoly, prec::Integer) = ArbPoly(poly.arb_poly, prec = prec)
Base.setprecision(series::ArbSeries, prec::Integer) =
    ArbSeries(series.arb_poly, degree = degree(series), prec = prec)
Base.setprecision(poly::AcbPoly, prec::Integer) = AcbPoly(poly.acb_poly, prec = prec)
Base.setprecision(series::AcbSeries, prec::Integer) =
    AcbSeries(series.acb_poly, degree = degree(series), prec = prec)


using ShaderAbstractions: Buffer, Sampler, VertexArray

# Mesh
mesh(Sphere(Point3f0(0), 1f0)) |> display
mesh(Sphere(Point3f0(0), 1f0), color=:red, ambient=Vec3f0(0.9))

tocolor(x) = RGBf0(x...)
positions = Observable(decompose(Point3f0, Sphere(Point3f0(0), 1f0)))
triangles = Observable(decompose(GLTriangleFace, Sphere(Point3f0(0), 1f0)))
uv = Observable(GeometryBasics.decompose_uv(Sphere(Point3f0(0), 1f0)))
xyz_vertex_color = Observable(tocolor.(positions[]))
texture = Observable(rand(RGBAf0, 10, 10))

pos_buff = Buffer(positions)
triangles_buff = Buffer(triangles)
vert_color_buff = Buffer(xyz_vertex_color)
uv_buff = Buffer(uv)
texture_buff = Sampler(texture)
texsampler = Makie.sampler(:viridis, rand(length(positions)))

coords = VertexArray(pos_buff, triangles_buff, color=vert_color_buff)
mesh = GeometryBasics.Mesh(coords)
GeometryBasics.coordinates(mesh);

using ShaderAbstractions: data

posmeta = Buffer(meta(data(pos_buff); color = data(vert_color_buff)))

program = disp.renderlist[1][3].vertexarray.program

struct OGLContext <: ShaderAbstractions.AbstractContext end

instance = ShaderAbstractions.VertexArray(posmeta, triangles_buff)

ShaderAbstractions.type_string(OGLContext(), Vec2f0)

p = ShaderAbstractions.Program(
    OGLContext(),
    read(loadshader("mesh.vert"), String),
    read(loadshader("mesh.frag"), String),
    instance;
)

println(p.vertex_source)

uniforms = Dict{Symbol, Any}(
    :texturecoordinates => Vec2f0(0),
    :image => nothing
)
rshader = GLMakie.GLAbstraction.gl_convert(shader, uniforms)

vbo = GLMakie.GLAbstraction.GLVertexArray(program, posmeta, triangles_buff)

m = GeometryBasics.Mesh(posmeta, triangles_buff)
disp = display(Makie.mesh(m, show_axis=false));

mesh_normals = GeometryBasics.normals(positions, triangles)
coords = meta(positions, color=xyz_vertex_color, normals=mesh_normals)
vertexcolor_mesh = GeometryBasics.Mesh(coords, triangles)
scren = mesh(vertexcolor_mesh, show_axis=false) |> display

function getter_function(io::IO, ::Fragment, sampler::Sampler, name::Symbol)
    index_t = type_string(context, sampler.values)
    sampler_t = type_string(context, sampler.colors)

    println(io, """
    in $(value_t) fragment_$(name)_index;
    uniform $(sampler_t) $(name)_texture;

    vec4 get_$(name)(){
        return texture($(name)_texture, fragment_$(name)_index);
    }
    """)
end

function getter_function(io::IO, ::Vertex, sampler::Sampler, name::Symbol)
    index_t = type_string(context, sampler.values)
    println(io, """
    in $(index_t) $(name)_index;
    out $(index_t) fragment_$(name)_index;

    vec4 get_$(name)(){
        fragment_uv = uv;
        // color gets calculated in fragment!
        return vec4(0);
    }
    """)
end

function getter_function(io::IO, ::Fragment, ::AbstractVector{T}, name) where T
    t_str = type_string(context, T)
    println(io, """
    in $(t_str) fragment_$(name);
    $(t_str) get_$(name)(){
        return fragment_$(name);
    }
    """)
end

function getter_function(io::IO, ::Vertex, ::AbstractVector{T}, name) where T
    t_str = type_string(context, T)
    println(io, """
    in $(t_str) $(name);
    out $(t_str) fragment_$(name);

    $(t_str) get_$(name)(){
        fragment_$(name) = $(name);
        return $(name);
    }
    """)
end

texsampler = Makie.sampler(rand(RGBf0, 4, 4), uv)
coords = meta(positions, color=texsampler, normals=mesh_normals)
texture_mesh = GeometryBasics.Mesh(coords, triangles)

scren = mesh(texture_mesh, show_axis=false) |> display

texsampler = Makie.sampler(:viridis, rand(length(positions)))
coords = meta(positions, color=texsampler, normals=mesh_normals)
texture_mesh = GeometryBasics.Mesh(coords, triangles)

scren = mesh(texture_mesh, show_axis=false) |> display

glsl"""

out vec2 frag_uv;
out vec3 frag_normal;
out vec3 frag_position;
out vec4 frag_color;
out vec3 frag_lightdir;

uniform mat4 projectionMatrix;
uniform mat4 viewMatrix;
uniform mat4 modelMatrix;

vec3 tovec3(vec2 v){return vec3(v, 0.0);}
vec3 tovec3(vec3 v){return v;}

vec4 tovec4(vec3 v){return vec4(v, 1.0);}
vec4 tovec4(vec4 v){return v;}

void main(){
    // get_* gets the global inputs (uniform, sampler, vertex array)
    // those functions will get inserted by the shader creation pipeline
    vec3 vertex_position = tovec3(get_position());
    frag_normal = $(normals);
    frag_uv = get_texturecoordinates();

    vec3 lightpos = vec3(20);
    vec4 position_world = modelMatrix * vec4(vertex_position, 1);
    frag_lightdir = normalize(lightpos - position_world.xyz);
    // direction to camera
    frag_position = -position_world.xyz;
    frag_uv = vec2(1.0 - frag_uv.y, frag_uv.x);
    frag_color = tovec4(get_color());
    // screen space coordinates of the vertex
    gl_Position = projectionMatrix * viewMatrix * position_world;
}
"""


function get_bgd(obj::GalileoOrbitalElements)
    return jcall(obj, "getBGD", jdouble, ())
end

function get_bgde1e5a(obj::GalileoOrbitalElements)
    return jcall(obj, "getBGDE1E5a", jdouble, ())
end

function get_bgde5b_e1(obj::GalileoOrbitalElements)
    return jcall(obj, "getBGDE5bE1", jdouble, ())
end

function get_iod_nav(obj::GalileoOrbitalElements)
    return jcall(obj, "getIODNav", jint, ())
end


@testset "TupleTabular" begin
    @testset "TupleSeries" begin
        @test @inferred(Series(l"a"=>1)) isa Series{<:Tuple{Vararg{Pair}}}
        @test @inferred(Series(l"a"=>1, l"b"=>true)) isa Series{<:Tuple{Vararg{Pair}}}
        s = Series(l"a"=>1, l"b"=>true)
        @test @inferred(s[l"a"]) === 1
        @test @inferred(s[l"b"]) === true
        @test_throws IndexError s[l"c"]
        @test @inferred(s[(l"a", l"b")]) === s
        @test_throws IndexError s[(l"a",l"c")]
        @test @inferred(s[:]) === s
    end

    @testset "TupleTable" begin
        # TupleTable with inner DictSeries
        @test @inferred(Table(l"a" => Dict(:x=>1,:y=>2,:z=>3), l"b" => Dict(:x=>4.0,:y=>5.0,:z=>6.0))) isa Table{<:Tuple{Vararg{Pair}}}
        t = Table(l"a" => Dict(:x=>1,:y=>2,:z=>3), l"b" => Dict(:x=>4.0,:y=>5.0,:z=>6.0))
        @test @inferred(indices(t)) == (keys(get(t)[1].second), (l"a", l"b"))
        @test @inferred(t[:y, l"a"]) === 2
        @test @inferred(t[:y, (l"a",l"b")]) === Series(l"a" => 2, l"b" => 5.0)
        @test_throws IndexError t[:y, l"c"]
        @test @inferred(t[:y, :]) === Series(l"a" => 2, l"b" => 5.0)
        @test @inferred(t[:, l"a"]) == Series(Dict(:x => 1, :y => 2, :z => 3))
        @test @inferred(t[:, (l"a",l"b")]) == t
        @test @inferred(t[:, :]) == t

        @test_broken (t[:y, l"a"] = 5; t[:y, l"a"] === 5)

        # TupleTable with inner ArraySeries
        @test @inferred(Table(l"a" => [1,2,3], l"b" => [4.0,5.0,6.0])) isa Table{<:Tuple{Vararg{Pair}}}
        t2 = Table(l"a" => [1,2,3], l"b" => [4.0,5.0,6.0])
        @test @inferred(indices(t2)) === (Base.OneTo(3), (l"a", l"b"))
        @test @inferred(t2[2, l"a"]) === 2
        @test_throws IndexError t[2, l"c"]
        @test @inferred(t2[2, :]) == Series(l"a" => 2, l"b" => 5)
        @test @inferred(t2[:, l"a"]) == Series([1,2,3])
        @test @inferred(t2[:, :]) == t2

        @test_broken (t2[2, l"a"] = 5; t2[2, l"a"] === 5)

        # TupleTable with inner TupleSeries
        @test @inferred(Table(l"a" => (l"x"=>true,l"y"=>2,l"z"=>3), l"b" => (l"x"=>4.0,l"y"=>5.0f0,l"z"=>6.0))) isa Table{<:Tuple{Vararg{Pair}}}
        t3 = Table(l"a" => (l"x"=>true,l"y"=>2,l"z"=>3), l"b" => (l"x"=>4.0,l"y"=>5.0f0,l"z"=>6.0))
        @test @inferred(indices(t3)) === ((l"x", l"y", l"z"), (l"a", l"b"))
        @test @inferred(t3[l"y", l"a"]) === 2
        @test_throws IndexError t[l"y", l"c"]
        if VERSION <= v"0.6.0-rc3.0"
            # Inference works at REPL but not in test
            @test t3[l"y", :] === Series(l"a" => 2, l"b" => 5.0f0)
            @test_broken @inferred(t3[l"y", :]) === Series(l"a" => 2, l"b" => 5.0f0)
        else
            @test @inferred(t3[l"y", :]) === Series(l"a" => 2, l"b" => 5.0f0)
        end
        @test @inferred(t3[:, l"a"]) === Series(l"x" => true, l"y" => 2, l"z" => 3)
        @test @inferred(t3[:, :]) === t3

        # TupleTable with inner StructSeries
        @test @inferred(Table(l"a" => 1+2im, l"b" => 3.0+4.0im)) isa Table{<:Tuple{Vararg{Pair}}}
        t4 = Table(l"a" => 1+2im, l"b" => 3.0+4.0im)
        @test @inferred(indices(t4)) === ((l"re", l"im"), (l"a", l"b"))
        @test @inferred(t4[l"im", l"a"]) === 2
        @test_throws IndexError t[l"im", l"c"]
        @test @inferred(t4[l"im", :]) === Series(l"a" => 2, l"b" => 4.0) # Inference works at REPL but not in test
        @test @inferred(t4[:, l"a"]) == Series(1+2im)

        if VERSION <= v"0.6.0-rc3.0"
            @test @inferred(t4[:, :]) == t4
        else
            # TODO - check inference of this on Julia v0.7, which is uninferred at REPL
            # (at least it's consistent...)
            @test t4[:, :] == t4
            @test_broken @inferred(t4[:, :]) == t4
        end
    end
end


using DataFrames
using StatsBase
using CSV
using ProgressMeter
using OnlineStats
using HDF5

"""
    run(model, agglog, bactlog, foodlog; kwargs...)

Runs the given `model`. Aggregate data is logged to file at `agglog` (csv). Data of individual
bacteria is logged to `bactlog` (csv), and the food grid is written to `foodlog` (hdf5).

Keyword arguments
-----------------

- `nsteps = 2000` : Number of steps to run the simulation for
- `log_period = 1` : Periodicity of logging bacterium and food data. Use negative value
  to disable
- `log_bact = true` : Choose whether to log bacteria data
- `log_food = true` : Choose whether to log food map
- `chunk_size = 1000` : How many rows should be written to csv files at a time
- `float_precision = 2` : How many decimals of precision should there be in the csv files
"""
function run!(
    model::ABM,
    agglog::String,
    bactlog::String,
    foodlog::String;
    nsteps::Int = 2000,
    log_period::Int = 1,
    log_bact::Bool = true,
    log_food::Bool = true,
    chunk_size::Int = 1000,
    float_precision::Int = 2,
)
    agg_data = DataFrame(
        step = Int[],
        species = Int[],
        nbact = Int[],
        μ_energy = Float64[],
        σ_energy = Float64[],
        μ_sens = Float64[],
        σ_sens = Float64[],
        μ_repr = Float64[],
        σ_repr = Float64[],
        μ_speed = Float64[],
        σ_speed = Float64[],
    )
    bactlogdf = DataFrame(
        step = Int[],
        age = Int[],
        species = Int[],
        energy = Float64[],
        sens = Float64[],
        repr = Float64[],
        speed = Float64[],
    )

    agglog = open(agglog, "w+")
    log_bact && (bactlog = open(bactlog, "w+"))

    if log_food
        isfile(foodlog) && rm(foodlog)
        foodf = h5open(foodlog, "w")
        foodf["log"] = zeros((floor(Int, nsteps / log_period) + 1, size(model.food)...))
        fdata = HDF5.readmmap(foodf["log"])
    end
    appendbact = false
    appendagg = false

    p = Progress(nsteps; barglyphs = BarGlyphs("[=> ]"), color = :blue)
    for i = 0:nsteps
        nagents(model) == 0 && break
        count = [0 for _ = 1:model.nspecies]
        energy_means = [Mean() for _ = 1:model.nspecies]
        energy_vars = [Variance() for _ = 1:model.nspecies]
        sens_means = [Mean() for _ = 1:model.nspecies]
        sens_vars = [Variance() for _ = 1:model.nspecies]
        repr_means = [Mean() for _ = 1:model.nspecies]
        repr_vars = [Variance() for _ = 1:model.nspecies]
        speed_means = [Mean() for _ = 1:model.nspecies]
        speed_vars = [Variance() for _ = 1:model.nspecies]

        for a in allagents(model)
            count[a.species] += 1
            fit!(energy_means[a.species], a.energy)
            fit!(energy_vars[a.species], a.energy)
            fit!(sens_means[a.species], a.sensory_radius)
            fit!(sens_vars[a.species], a.sensory_radius)
            fit!(repr_means[a.species], a.reproduction_threshold)
            fit!(repr_vars[a.species], a.reproduction_threshold)
            fit!(speed_means[a.species], a.speed)
            fit!(speed_vars[a.species], a.speed)
        end
        for sp = 1:model.nspecies
            push!(
                agg_data,
                (
                    i,
                    sp,
                    count[sp],
                    round(value(energy_means[sp]); digits = float_precision),
                    round(√value(energy_vars[sp]); digits = float_precision),
                    round(value(sens_means[sp]); digits = float_precision),
                    round(√value(sens_vars[sp]); digits = float_precision),
                    round(value(repr_means[sp]); digits = float_precision),
                    round(√value(repr_vars[sp]); digits = float_precision),
                    round(value(speed_means[sp]); digits = float_precision),
                    round(√value(speed_vars[sp]); digits = float_precision),
                ),
            )
        end
        if size(agg_data, 1) >= chunk_size
            CSV.write(agglog, agg_data; append = appendagg)
            appendagg || (appendagg = true)
            empty!(agg_data)
        end

        if log_period > 0 && i % log_period == 0
            if log_bact
                for a in allagents(model)
                    push!(
                        bactlogdf,
                        (
                            i,
                            a.age,
                            a.species,
                            round(a.energy; digits = float_precision),
                            round(a.sensory_radius; digits = float_precision),
                            round(a.reproduction_threshold; digits = float_precision),
                            round(a.speed; digits = float_precision),
                        ),
                    )
                end
                if size(bactlogdf, 1) >= chunk_size
                    CSV.write(bactlog, bactlogdf; append = appendbact)
                    appendbact || (appendbact = true)
                    empty!(bactlogdf)
                end
            end

            log_food && (fdata[floor(Int, i / log_period)+1, :, :] = model.food)
        end
        Agents.step!(model, agent_step!, food_step!)
        ProgressMeter.next!(p; showvalues = [(:steps, i), (:population, nagents(model))])
    end

    size(agg_data, 1) > 0 && CSV.write(agglog, agg_data; append = appendagg)
    size(bactlogdf, 1) > 0 && CSV.write(bactlog, bactlogdf; append = appendbact)

    if log_food
        foodf["cap"] = model.food_data.food_cap
        close(foodf)
    end
    nothing
end

"""
    runconfig(config, dirname, foodpath; kwargs...)

Create a model with given `config`, run it using `foodmap` as the initial food distribution,
and log data to `dirname`.

Keyword arguments
-----------------
- See `run!`
- `saveconfig = false` whether the given config file should also be saved to `dirname` as BSON
"""
function runconfig(
    config,
    dirname::String,
    foodpath::String;
    nsteps::Int = 2000,
    log_period::Int = 10,
    log_bact::Bool = true,
    log_food::Bool = true,
    chunk_size::Int = 1000,
    float_precision::Int = 2,
    saveconfig::Bool = false,
)
    haskey(config, :initial_bacterium) &&
        config[:initial_bacterium] isa Vector &&
        (config[:initial_bacterium] = Vector{Bacterium}(config[:initial_bacterium]))
    mkpath(dirname)
    close(open(joinpath(dirname, "$(configname(config))"), "w"))

    food = image_to_foodmap(foodpath)
    model = initialize_model(food .* config[:food_data].food_cap; config...)

    run!(
        model,
        joinpath(dirname, "agg.csv"),
        joinpath(dirname, "bact.csv"),
        joinpath(dirname, "food.hdf5");
        nsteps,
        log_period,
        log_bact,
        log_food,
        chunk_size,
        float_precision,
    )
    saveconfig && bson(joinpath(dirname, "config.bson"), config)
    return model
end

"""
    configname(config)

Generate a highly abbreviated name for a config
"""
function configname(config::Dict)
    name = ""
    if haskey(config, :n_bacteria)
        if config[:n_bacteria] isa Vector
            name *= "nbx$(length(config[:n_bacteria]))"
        else
            name *= "nb$(config[:n_bacteria])"
        end
    end

    if haskey(config, :initial_bacterium)
        name *= "ib("
        if config[:initial_bacterium] isa Vector
            name *= "e"
            for b in config[:initial_bacterium]
                name *= "$(b.energy)_"
            end
            name *= "s"
            for b in config[:initial_bacterium]
                name *= "$(b.speed)_"
            end
            name *= "sr"
            for b in config[:initial_bacterium]
                name *= "$(b.sensory_radius)_"
            end
            name *= "r"
            for b in config[:initial_bacterium]
                name *= "$(b.reproduction_threshold)_"
            end
        else
            b = config[:initial_bacterium]
            name *= "e$(b.energy)s$(b.speed)sr$(b.sensory_radius)r$(b.reproduction_threshold)"
        end
        name *= ")"
    end

    if haskey(config, :food_data)
        f = config[:food_data]
        name *= "fd(c$(f.food_cap)r$(f.regen_rate))"
    end

    haskey(config, :eat_rate_factor) && (name *= "erf$(config[:eat_rate_factor])")
    haskey(config, :lifetime) && (name *= "l$(config[:lifetime])")
end


@testset "Renewable data misspecification" begin
    # See https://discourse.julialang.org/t/how-to-use-test-warn/15557/5 about testing for warning throwing
    warn_message = "The data doesn't include devices of type RenewableDispatch, consider changing the device models"
    model = DeviceModel(RenewableDispatch, RenewableFullDispatch)
    c_sys5 = build_system("c_sys5")
    c_sys14 = build_system("c_sys14")
    op_problem = OperationsProblem(TestOpProblem, DCPPowerModel, c_sys5)
    @test_logs (:warn, warn_message) construct_device!(op_problem, :Renewable, model)
    op_problem = OperationsProblem(TestOpProblem, DCPPowerModel, c_sys14)
    @test_logs (:warn, warn_message) construct_device!(op_problem, :Renewable, model)
end

@testset "Renewable DCPLossLess FullDispatch" begin
    model = DeviceModel(RenewableDispatch, RenewableFullDispatch)
    c_sys5_re = build_system("c_sys5_re")

    #5 Bus testing case
    op_problem = OperationsProblem(TestOpProblem, DCPPowerModel, c_sys5_re)
    construct_device!(op_problem, :Renewable, model)
    moi_tests(op_problem, false, 72, 0, 72, 0, 0, false)

    psi_checkobjfun_test(op_problem, GAEVF)

    # Using Parameters Testing
    op_problem =
        OperationsProblem(TestOpProblem, DCPPowerModel, c_sys5_re; use_parameters = true)
    construct_device!(op_problem, :Renewable, model)
    moi_tests(op_problem, true, 72, 0, 72, 0, 0, false)
    psi_checkobjfun_test(op_problem, GAEVF)

    # No Forecast - No Parameters Testing
    op_problem = OperationsProblem(
        TestOpProblem,
        DCPPowerModel,
        c_sys5_re;
        use_forecast_data = false,
    )
    construct_device!(op_problem, :Renewable, model)
    moi_tests(op_problem, false, 3, 0, 3, 3, 0, false)
    psi_checkobjfun_test(op_problem, GAEVF)

    # No Forecast Testing
    op_problem = OperationsProblem(
        TestOpProblem,
        DCPPowerModel,
        c_sys5_re;
        use_parameters = true,
        use_forecast_data = false,
    )
    construct_device!(op_problem, :Renewable, model)
    moi_tests(op_problem, true, 3, 0, 3, 0, 0, false)
    psi_checkobjfun_test(op_problem, GAEVF)
end

@testset "Renewable ACPPower Full Dispatch" begin
    model = DeviceModel(RenewableDispatch, RenewableFullDispatch)
    c_sys5_re = build_system("c_sys5_re")
    for p in [true, false]
        op_problem =
            OperationsProblem(TestOpProblem, ACPPowerModel, c_sys5_re; use_parameters = p)
        construct_device!(op_problem, :Renewable, model)
        if p
            moi_tests(op_problem, p, 144, 0, 144, 72, 0, false)
            psi_checkobjfun_test(op_problem, GAEVF)
        else
            moi_tests(op_problem, p, 144, 0, 144, 72, 0, false)

            psi_checkobjfun_test(op_problem, GAEVF)
        end
    end
    # No Forecast Test
    for p in [true, false]
        op_problem = OperationsProblem(
            TestOpProblem,
            ACPPowerModel,
            c_sys5_re;
            use_forecast_data = false,
            use_parameters = p,
        )
        construct_device!(op_problem, :Renewable, model)
        if p
            moi_tests(op_problem, p, 6, 0, 6, 3, 0, false)
            psi_checkobjfun_test(op_problem, GAEVF)
        else
            moi_tests(op_problem, p, 6, 0, 6, 6, 0, false)

            psi_checkobjfun_test(op_problem, GAEVF)
        end
    end
end

@testset "Renewable DCPLossLess ConstantPowerFactor" begin
    model = DeviceModel(RenewableDispatch, RenewableConstantPowerFactor)
    c_sys5_re = build_system("c_sys5_re")

    #5 Bus testing case
    op_problem = OperationsProblem(TestOpProblem, DCPPowerModel, c_sys5_re)
    construct_device!(op_problem, :Renewable, model)
    moi_tests(op_problem, false, 72, 0, 72, 0, 0, false)

    psi_checkobjfun_test(op_problem, GAEVF)

    # Using Parameters Testing
    op_problem =
        OperationsProblem(TestOpProblem, DCPPowerModel, c_sys5_re; use_parameters = true)
    construct_device!(op_problem, :Renewable, model)
    moi_tests(op_problem, true, 72, 0, 72, 0, 0, false)
    psi_checkobjfun_test(op_problem, GAEVF)

    # No Forecast - No Parameters Testing
    op_problem = OperationsProblem(
        TestOpProblem,
        DCPPowerModel,
        c_sys5_re;
        use_forecast_data = false,
    )
    construct_device!(op_problem, :Renewable, model)
    moi_tests(op_problem, false, 3, 0, 3, 3, 0, false)
    psi_checkobjfun_test(op_problem, GAEVF)

    # No Forecast Testing
    op_problem = OperationsProblem(
        TestOpProblem,
        DCPPowerModel,
        c_sys5_re;
        use_parameters = true,
        use_forecast_data = false,
    )
    construct_device!(op_problem, :Renewable, model)
    moi_tests(op_problem, true, 3, 0, 3, 0, 0, false)
    psi_checkobjfun_test(op_problem, GAEVF)
end

@testset "Renewable ACPPower ConstantPowerFactor" begin
    model = DeviceModel(RenewableDispatch, RenewableConstantPowerFactor)
    c_sys5_re = build_system("c_sys5_re")
    for p in [true, false]
        op_problem =
            OperationsProblem(TestOpProblem, ACPPowerModel, c_sys5_re; use_parameters = p)
        construct_device!(op_problem, :Renewable, model)
        if p
            moi_tests(op_problem, p, 144, 0, 72, 0, 72, false)
            psi_checkobjfun_test(op_problem, GAEVF)
        else
            moi_tests(op_problem, p, 144, 0, 72, 0, 72, false)

            psi_checkobjfun_test(op_problem, GAEVF)
        end
    end
    # No Forecast Test
    for p in [true, false]
        op_problem = OperationsProblem(
            TestOpProblem,
            ACPPowerModel,
            c_sys5_re;
            use_forecast_data = false,
            use_parameters = p,
        )
        construct_device!(op_problem, :Renewable, model)
        if p
            moi_tests(op_problem, p, 6, 0, 3, 0, 3, false)
            psi_checkobjfun_test(op_problem, GAEVF)
        else
            moi_tests(op_problem, p, 6, 0, 3, 3, 3, false)

            psi_checkobjfun_test(op_problem, GAEVF)
        end
    end
end

@testset "Renewable DCPLossLess FixedOutput" begin
    model = DeviceModel(RenewableDispatch, FixedOutput)
    c_sys5_re = build_system("c_sys5_re")
    for p in [true, false]
        op_problem =
            OperationsProblem(TestOpProblem, DCPPowerModel, c_sys5_re; use_parameters = p)
        construct_device!(op_problem, :Renewable, model)
        if p
            moi_tests(op_problem, p, 0, 0, 0, 0, 0, false)
            psi_checkobjfun_test(op_problem, GAEVF)
        else
            moi_tests(op_problem, p, 0, 0, 0, 0, 0, false)
            psi_checkobjfun_test(op_problem, GAEVF)
        end
    end
end

@testset "Renewable ACPPowerModel FixedOutput" begin
    model = DeviceModel(RenewableDispatch, FixedOutput)
    c_sys5_re = build_system("c_sys5_re")
    for p in [true, false]
        op_problem =
            OperationsProblem(TestOpProblem, ACPPowerModel, c_sys5_re; use_parameters = p)
        construct_device!(op_problem, :Renewable, model)
        if p
            moi_tests(op_problem, p, 0, 0, 0, 0, 0, false)
            psi_checkobjfun_test(op_problem, GAEVF)
        else
            moi_tests(op_problem, p, 0, 0, 0, 0, 0, false)
            psi_checkobjfun_test(op_problem, GAEVF)
        end
    end
end


function FieldLnsCoefficients(arg0::jint, arg1::jint, arg2::Vector{Vector{RealFieldElement}}, arg3::TreeMap, arg4::RealFieldElement, arg5::Field)
    return FieldLnsCoefficients((jint, jint, Vector{Vector{RealFieldElement}}, TreeMap, RealFieldElement, Field), arg0, arg1, arg2, arg3, arg4, arg5)
end

function get_lns(obj::FieldLnsCoefficients, arg0::jint, arg1::jint)
    return jcall(obj, "getLns", RealFieldElement, (jint, jint), arg0, arg1)
end

function getd_lnsd_gamma(obj::FieldLnsCoefficients, arg0::jint, arg1::jint)
    return jcall(obj, "getdLnsdGamma", RealFieldElement, (jint, jint), arg0, arg1)
end


module CrossSections
using Unitful
const UF = Unitful
using Unitful: MeV, NoUnits, cm, m, J, s, kg
using Random
using Distributions

include("constants.jl")
include("compton.jl")

end # module


using JuliaChem

using Documenter, JuliaChem.JCInput, JuliaChem.JCBasis, JuliaChem.JCRHF

makedocs(
    authors="David Poole",
    sitename="JuliChem Documentation",
    modules = [JCInput,
               JCBasis,
               JCRHF],
    #format = LaTeX()
)


"""
    MOIError{F <: Function} <: MOI.AbstractVectorSet

DOCSTRING

# Arguments:
- `f::F`: DESCRIPTION
- `dimension::Int`: DESCRIPTION
- `MOIError(f, dim = 0) = begin
        #= none:5 =#
        new{typeof(f)}(f, dim)
    end`: DESCRIPTION
"""
struct MOIError{F <: Function} <: MOI.AbstractVectorSet
    f::F
    dimension::Int

    MOIError(f, dim = 0) = new{typeof(f)}(f, dim)
end

"""
    MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOIError}) = begin

DOCSTRING

# Arguments:
- ``: DESCRIPTION
- ``: DESCRIPTION
- ``: DESCRIPTION
"""
function MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOIError{F}}
) where {F <: Function}
    return true
end

"""
    MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, set::MOIError)

DOCSTRING

# Arguments:
- `optimizer`: DESCRIPTION
- `vars`: DESCRIPTION
- `set`: DESCRIPTION
"""
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, set::MOIError{F}) where {F <: Function}
    cidx = constraint!(optimizer, set.f, map(x -> x.value, vars.variables))
    return CI{VOV, MOIError{F}}(cidx)
end

"""
    Base.copy(set::MOIError) = begin

DOCSTRING
"""
Base.copy(set::MOIError) = MOIError(deepcopy(set.f), set.dimension)

"""
    Error{F <: Function} <: JuMP.AbstractVectorSet

The solver will compute a straightforward error function based on the `concept`. To run the solver efficiently, it is possible to provide an *error function* `err` instead of `concept`. `err` must return a nonnegative real number.

```julia
@constraint(model, X in Error(err))
```
"""
struct Error{F <: Function} <: JuMP.AbstractVectorSet
    f::F
end

# @autodoc
JuMP.moi_set(set::Error{F}, dim::Int) where {F <: Function} = MOIError(set.f, dim)

"""
    MOIPredicate{F <: Function} <: MOI.AbstractVectorSet

DOCSTRING

# Arguments:
- `f::F`: DESCRIPTION
- `dimension::Int`: DESCRIPTION
- `MOIPredicate(f, dim = 0) = begin
        #= none:5 =#
        new{typeof(f)}(f, dim)
    end`: DESCRIPTION
"""
struct MOIPredicate{F <: Function} <: MOI.AbstractVectorSet
    f::F
    dimension::Int

    MOIPredicate(f, dim = 0) = new{typeof(f)}(f, dim)
end
function MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOIPredicate{F}}
) where {F <: Function}
    return true
end
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, set::MOIPredicate{F}) where {F <: Function}
    err = x -> convert(Float64, !set.f(x))
    cidx = constraint!(optimizer, err, map(x -> x.value, vars.variables))
    return CI{VOV, MOIPredicate{F}}(cidx)
end

Base.copy(set::MOIPredicate) = MOIPredicate(deepcopy(set.f), copy(set.dimension))

"""
    Predicate{F <: Function} <: JuMP.AbstractVectorSet

Assuming `X` is a (collection of) variables, `concept` a boolean function over `X`, and that a `model` is defined. In `JuMP` syntax we can create a constraint based on `concept` as follows.

```julia
@constraint(model, X in Predicate(concept))
```
"""
struct Predicate{F <: Function} <: JuMP.AbstractVectorSet
    f::F
end
JuMP.moi_set(set::Predicate, dim::Int) = MOIPredicate(set.f, dim)

"""
    MOIAllDifferent <: MOI.AbstractVectorSet

DOCSTRING
"""
struct MOIAllDifferent <: MOI.AbstractVectorSet
    dimension::Int

    MOIAllDifferent(dim = 0) = new(dim)
end
MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOIAllDifferent}) = true
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, ::MOIAllDifferent)
    max_dom_size = max_domains_size(optimizer, map(x -> x.value, vars.variables))
    e = (x; param=nothing, dom_size=max_dom_size) -> error_f(
        usual_constraints[:all_different])(x; param=param, dom_size=dom_size
    )
    cidx = constraint!(optimizer, e, map(x -> x.value, vars.variables))
    return CI{VOV, MOIAllDifferent}(cidx)
end
Base.copy(set::MOIAllDifferent) = MOIAllDifferent(copy(set.dimension))

"""
Global constraint ensuring that all the values of a given configuration are unique.

```julia
@constraint(model, X in AllDifferent())
```
"""
struct AllDifferent <: JuMP.AbstractVectorSet end
JuMP.moi_set(::AllDifferent, dim::Int) = MOIAllDifferent(dim)

"""
    MOIAllEqual <: MOI.AbstractVectorSet

DOCSTRING
"""
struct MOIAllEqual <: MOI.AbstractVectorSet
    dimension::Int

    MOIAllEqual(dim = 0) = new(dim)
end
MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOIAllEqual}) = true
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, ::MOIAllEqual)
    max_dom_size = max_domains_size(optimizer, map(x -> x.value, vars.variables))
    e = (x; param=nothing, dom_size=max_dom_size) -> error_f(
        usual_constraints[:all_equal])(x; param=param, dom_size=dom_size
    )
    cidx = constraint!(optimizer, e, map(x -> x.value, vars.variables))
    return CI{VOV, MOIAllEqual}(cidx)
end

Base.copy(set::MOIAllEqual) = MOIAllEqual(copy(set.dimension))

"""
Global constraint ensuring that all the values of `X` are all equal.

```julia
@constraint(model, X in AllEqual())
```
"""
struct AllEqual <: JuMP.AbstractVectorSet end
JuMP.moi_set(::AllEqual, dim::Int) = MOIAllEqual(dim)

"""
    MOIEq <: MOI.AbstractVectorSet

DOCSTRING
"""
struct MOIEq <: MOI.AbstractVectorSet
    dimension::Int

    MOIEq(dim = 0) = new(dim)
end
MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOIEq}) = true
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, ::MOIEq)
    max_dom_size = max_domains_size(optimizer, map(x -> x.value, vars.variables))
    e = (x; param=nothing, dom_size=max_dom_size) -> error_f(
        usual_constraints[:eq])(x; param=param, dom_size=dom_size
    )
    cidx = constraint!(optimizer, e, map(x -> x.value, vars.variables))
    return CI{VOV, MOIEq}(cidx)
end

Base.copy(set::MOIEq) = MOIEq(copy(set.dimension))

"""
Equality between two variables.

```julia
@constraint(model, X in Eq())
```
"""
struct Eq <: JuMP.AbstractVectorSet end
JuMP.moi_set(::Eq, dim::Int) = MOIEq(dim)

"""
    MOIAlwaysTrue <: MOI.AbstractVectorSet

DOCSTRING
"""
struct MOIAlwaysTrue <: MOI.AbstractVectorSet
    dimension::Int

    MOIAlwaysTrue(dim = 0) = new(dim)
end
MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOIAlwaysTrue}) = true
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, ::MOIAlwaysTrue)
    max_dom_size = max_domains_size(optimizer, map(x -> x.value, vars.variables))
    e = (x; param=nothing, dom_size=max_dom_size) -> error_f(
        usual_constraints[:always_true])(x; param=param, dom_size=dom_size
    )
    cidx = constraint!(optimizer, e, map(x -> x.value, vars.variables))
    return CI{VOV, MOIAlwaysTrue}(cidx)
end

Base.copy(set::MOIAlwaysTrue) = MOIAlwaysTrue(copy(set.dimension))

"""
Always return `true`. Mainly used for testing purpose.

```julia
@constraint(model, X in AlwaysTrue())
```
"""
struct AlwaysTrue <: JuMP.AbstractVectorSet end
JuMP.moi_set(::AlwaysTrue, dim::Int) = MOIAlwaysTrue(dim)

"""
    MOIOrdered <: MOI.AbstractVectorSet

DOCSTRING
"""
struct MOIOrdered <: MOI.AbstractVectorSet
    dimension::Int

    MOIOrdered(dim = 0) = new(dim)
end
MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOIOrdered}) = true
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, ::MOIOrdered)
    max_dom_size = max_domains_size(optimizer, map(x -> x.value, vars.variables))
    e = (x; param=nothing, dom_size=max_dom_size) -> error_f(
        usual_constraints[:ordered])(x; param=param, dom_size=dom_size
    )
    cidx = constraint!(optimizer, e, map(x -> x.value, vars.variables))
    return CI{VOV, MOIOrdered}(cidx)
end

Base.copy(set::MOIOrdered) = MOIOrdered(copy(set.dimension))

"""
Global constraint ensuring that all the values of `x` are ordered.

```julia
@constraint(model, X in Ordered())
```
"""
struct Ordered <: JuMP.AbstractVectorSet end
JuMP.moi_set(::Ordered, dim::Int) = MOIOrdered(dim)

"""
    MOIDistDifferent <: MOI.AbstractVectorSet

DOCSTRING
"""
struct MOIDistDifferent <: MOI.AbstractVectorSet
    dimension::Int

    MOIDistDifferent(dim = 4) = new(dim)
end
MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOIDistDifferent}) = true
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, ::MOIDistDifferent)
    max_dom_size = max_domains_size(optimizer, map(x -> x.value, vars.variables))
    e = (x; param=nothing, dom_size=max_dom_size) -> error_f(
        usual_constraints[:dist_different])(x; param=param, dom_size=dom_size
    )
    cidx = constraint!(optimizer, e, map(x -> x.value, vars.variables))
    return CI{VOV, MOIDistDifferent}(cidx)
end
Base.copy(set::MOIDistDifferent) = MOIDistDifferent(copy(set.dimension))

"""
Local constraint ensuring that, given a vector `X` of size 4, `|X[1] - X[2]| ≠ |X[3] - X[4]|)`.

```julia
@constraint(model, X in DistDifferent())
```
"""
struct DistDifferent <: JuMP.AbstractVectorSet end
JuMP.moi_set(::DistDifferent, dim::Int) = MOIDistDifferent(dim)

"""
    MOIAllEqualParam{T <: Number} <: MOI.AbstractVectorSet

DOCSTRING

# Arguments:
- `param::T`: DESCRIPTION
- `dimension::Int`: DESCRIPTION
- `MOIAllEqualParam(param, dim = 0) = begin
        #= none:5 =#
        new{typeof(param)}(param, dim)
    end`: DESCRIPTION
"""
struct MOIAllEqualParam{T <: Number} <: MOI.AbstractVectorSet
    param::T
    dimension::Int

    MOIAllEqualParam(param, dim = 0) = new{typeof(param)}(param, dim)
end
function MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOIAllEqualParam{T}}
) where {T <: Number}
    return true
end
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, set::MOIAllEqualParam{T}
) where T <: Number
    max_dom_size = max_domains_size(optimizer, map(x -> x.value, vars.variables))
    e = (x; param=set.param, dom_size=max_dom_size) -> error_f(
        usual_constraints[:all_equal_param])(x; param=param, dom_size=dom_size
    )
    cidx = constraint!(optimizer, e, map(x -> x.value, vars.variables))
    return CI{VOV, MOIAllEqualParam{T}}(cidx)
end

Base.copy(set::MOIAllEqualParam) = MOIAllEqualParam(copy(set.param), copy(set.dimension))

"""
Global constraint ensuring that all the values of `X` are all equal to a given parameter `param`.

```julia
@constraint(model, X in AllEqualParam(param))
```
"""
struct AllEqualParam{T <: Number} <: JuMP.AbstractVectorSet
    param::T
end
JuMP.moi_set(set::AllEqualParam, dim::Int) = MOIAllEqualParam(set.param, dim)

"""
    MOISumEqualParam{T <: Number} <: MOI.AbstractVectorSet

DOCSTRING

# Arguments:
- `param::T`: DESCRIPTION
- `dimension::Int`: DESCRIPTION
- `MOISumEqualParam(param, dim = 0) = begin
        #= none:5 =#
        new{typeof(param)}(param, dim)
    end`: DESCRIPTION
"""
struct MOISumEqualParam{T <: Number} <: MOI.AbstractVectorSet
    param::T
    dimension::Int

    MOISumEqualParam(param, dim = 0) = new{typeof(param)}(param, dim)
end
function MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOISumEqualParam{T}}
) where {T <: Number}
    return true
end
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, set::MOISumEqualParam{T}
) where {T <: Number}
    max_dom_size = max_domains_size(optimizer, map(x -> x.value, vars.variables))
    e = (x; param=set.param, dom_size=max_dom_size) -> error_f(
        usual_constraints[:sum_equal_param])(x; param=param, dom_size=dom_size
    )
    cidx = constraint!(optimizer, e, map(x -> x.value, vars.variables))
    return CI{VOV, MOISumEqualParam{T}}(cidx)
end

Base.copy(set::MOISumEqualParam) = MOISumEqualParam(copy(set.param),
copy(set.dimension))

"""
Global constraint ensuring that the sum of the values of `X` is equal to a given parameter `param`.

```julia
@constraint(model, X in SumEqualParam(param))
```
"""
struct SumEqualParam{T <: Number} <: JuMP.AbstractVectorSet
    param::T
end
JuMP.moi_set(set::SumEqualParam, dim::Int) = MOISumEqualParam(set.param, dim)

"""
    MOILessThanParam{T <: Number} <: MOI.AbstractVectorSet

DOCSTRING

# Arguments:
- `param::T`: DESCRIPTION
- `dimension::Int`: DESCRIPTION
- `MOILessThanParam(param, dim = 0) = begin
        #= none:5 =#
        new{typeof(param)}(param, dim)
    end`: DESCRIPTION
"""
struct MOILessThanParam{T <: Number} <: MOI.AbstractVectorSet
    param::T
    dimension::Int

    MOILessThanParam(param, dim = 0) = new{typeof(param)}(param, dim)
end
function MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOILessThanParam{T}}
) where {T <: Number}
    return true
end
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, set::MOILessThanParam{T}
) where {T <: Number}
    max_dom_size = max_domains_size(optimizer, map(x -> x.value, vars.variables))
    e = (x; param=set.param, dom_size=max_dom_size) -> error_f(
        usual_constraints[:less_than_param])(x; param=param, dom_size=dom_size
    )
    cidx = constraint!(optimizer, e, map(x -> x.value, vars.variables))
    return CI{VOV, MOILessThanParam{T}}(cidx)
end

Base.copy(set::MOILessThanParam) = MOILessThanParam(copy(set.param),
copy(set.dimension))

"""
Constraint ensuring that the value of `x` is less than a given parameter `param`.

```julia
@constraint(model, x in LessThanParam(param))
```
"""
struct LessThanParam{T <: Number} <: JuMP.AbstractVectorSet
    param::T
end
JuMP.moi_set(set::LessThanParam, dim::Int) = MOILessThanParam(set.param, dim)

"""
    MOIMinusEqualParam{T <: Number} <: MOI.AbstractVectorSet

DOCSTRING

# Arguments:
- `param::T`: DESCRIPTION
- `dimension::Int`: DESCRIPTION
- `MOIMinusEqualParam(param, dim = 0) = begin
        #= none:5 =#
        new{typeof(param)}(param, dim)
    end`: DESCRIPTION
"""
struct MOIMinusEqualParam{T <: Number} <: MOI.AbstractVectorSet
    param::T
    dimension::Int

    MOIMinusEqualParam(param, dim = 0) = new{typeof(param)}(param, dim)
end
function MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOIMinusEqualParam{T}}
) where {T <: Number}
    return true
end
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, set::MOIMinusEqualParam{T}
) where {T <: Number}
    max_dom_size = max_domains_size(optimizer, map(x -> x.value, vars.variables))
    e = (x; param=set.param, dom_size=max_dom_size) -> error_f(
        usual_constraints[:minus_equal_param])(x; param=param, dom_size=dom_size
    )
    cidx = constraint!(optimizer, e, map(x -> x.value, vars.variables))
    return CI{VOV, MOIMinusEqualParam{T}}(cidx)
end

Base.copy(set::MOIMinusEqualParam) = MOIMinusEqualParam(copy(set.param),
copy(set.dimension))

"""
Constraint ensuring that the value of `x` is less than a given parameter `param`.

```julia
@constraint(model, x in MinusEqualParam(param))
```
"""
struct MinusEqualParam{T <: Number} <: JuMP.AbstractVectorSet
    param::T
end
JuMP.moi_set(set::MinusEqualParam, dim::Int) = MOIMinusEqualParam(set.param, dim)

"""
    MOISequentialTasks <: MOI.AbstractVectorSet

DOCSTRING
"""
struct MOISequentialTasks <: MOI.AbstractVectorSet
    dimension::Int

    MOISequentialTasks(dim = 4) = new(dim)
end
MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOISequentialTasks}) = true
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, ::MOISequentialTasks)
    max_dom_size = max_domains_size(optimizer, map(x -> x.value, vars.variables))
    e = (x; param=nothing, dom_size=max_dom_size) -> error_f(
        usual_constraints[:sequential_tasks])(x; param=param, dom_size=dom_size
    )
    cidx = constraint!(optimizer, e, map(x -> x.value, vars.variables))
    return CI{VOV, MOISequentialTasks}(cidx)
end
Base.copy(set::MOISequentialTasks) = MOISequentialTasks(copy(set.dimension))

"""
Local constraint ensuring that, given a vector `X` of size 4, `|X[1] - X[2]| ≠ |X[3] - X[4]|)`.

```julia
@constraint(model, X in SequentialTasks())
```
"""
struct SequentialTasks <: JuMP.AbstractVectorSet end
JuMP.moi_set(::SequentialTasks, dim::Int) = MOISequentialTasks(dim)


#=
Dash port of Shiny wordcloud example:

https://shiny.rstudio.com/gallery/word-cloud.html

Note: some functionality has been changed from the original
example.
=#

using Dash, DashBootstrapComponents, DashCoreComponents, DashHtmlComponents
using HTTP, WordCloud, Memoize, LRUCache, Base64

BASE_URL = "https://cdn.opensource.faculty.ai/wordcloud"

DOCUMENT_URLS = Dict(
    "midsummer" => "$BASE_URL/a-midsummer-nights-dream.txt",
    "venice" => "$BASE_URL/the-merchant-of-venice.txt",
    "randj" => "$BASE_URL/romeo-and-juliet.txt",
)

@memoize LRU{Any,Tuple{Any,Any}}(maxsize = 3) function load_word_frequencies(book)
    url = DOCUMENT_URLS[book]
    return HTTP.request("GET", url).body |> String |> processtext
end;

app = dash(external_stylesheets = [dbc_themes.BOOTSTRAP]);

dropdown = dcc_dropdown(
    id = "book",
    options = [
        Dict("label" => "A Midsummer Night's Dream", "value" => "midsummer"),
        Dict("label" => "The Merchant of Venice", "value" => "venice"),
        Dict("label" => "Romeo and Juliet", "value" => "randj"),
    ],
    value = "midsummer",
);

word_density_slider = html_div(
    dcc_slider(
        id = "word-density-slider",
        min = 0.1,
        max = 0.5,
        step = 0.05,
        value = 0.25,
        marks = Dict(i => string(convert(Int, i * 100)) * "%" for i = 0.1:0.1:0.5),
    ),
    className = "p-3 mb-2",
);

word_angles_slider = html_div(
    dcc_rangeslider(
        id = "word-angles-slider",
        min = -90,
        max = 90,
        step = 0.1,
        value = [-45, 45],
        marks = Dict(i => string(i) for i = -90:45:90),
    ),
    className = "p-3 mb-2",
);

max_vocab_slider = html_div(
    dcc_slider(
        id = "max-vocab-slider",
        min = 1,
        max = 300,
        step = 1,
        value = 100,
        marks = merge(Dict(1 => "1"), Dict(i => string(i) for i = 30:30:301)),
    ),
    className = "p-3",
);

controls = dbc_card(
    [
        dbc_formgroup([dbc_label("Choose a book:"), dropdown]),
        dbc_formgroup([dbc_label("Layout density:"), word_density_slider]),
        dbc_formgroup([dbc_label("Word angles:"), word_angles_slider]),
        dbc_formgroup([dbc_label("Maximum number of words:"), max_vocab_slider]),
    ],
    className = "mb-3",
    body = true,
);

spinner = dbc_spinner(
    color = "info",
    spinner_style = Dict("width" => "25rem", "height" => "25rem"),
);

app.layout = dbc_container(
    [
        html_h1("Word Cloud"),
        html_hr(),
        dbc_row(
            [
                dbc_col(controls, md = 5),
                dbc_col(
                    dbc_spinner(
                        id = "loading",
                        color = "info",
                        spinner_style = Dict("width" => "3rem", "height" => "3rem"),
                        children = dbc_card([dbc_cardimg(id = "wordcloud")]),
                    ),
                    md = 7,
                    className = "text-center",
                ),
            ],
            align = "center",
        ),
    ],
    fluid = true,
);

callback!(
    app,
    Output("wordcloud", "src"),
    Input("book", "value"),
    Input("word-density-slider", "value"),
    Input("word-angles-slider", "value"),
    Input("max-vocab-slider", "value"),
) do book, word_density, word_angles, max_vocab

    # Restrict the number of results
    word, freq = load_word_frequencies(book)
    frequencies = (word[1:max_vocab], freq[1:max_vocab])

    # Extract the required angle range
    angles = [i for i = word_angles[1]:10:word_angles[2]]

    # Create the wordcloud
    wc = wordcloud(
        frequencies,
        mask = shape(box, 1000, 500, 10),
        colors = WordCloud.randomscheme(),
        angles = angles, # WordCloud.randomangles(),
        density = word_density,
    )
    generate!(wc)

    # Encode the resulting image
    result = base64encode(String(copy(paintsvg(wc).bufferdata)))
    return "data:image/svg+xml;base64,$result"

end;

run_server(app, "0.0.0.0", 8888, debug = true);


## -------------------------------------------------------------------
function load_model(modelkey::String; uncompress = true)
    if modelkey == "raw_model"
        model = _load_raw_model()
    else
        models = lprocdat(iJR904, "base_models.bson"; verbose = false);
        model = models[modelkey]
    end
    return uncompress ? ChU.uncompressed_model(model) : model
end

function load_model(modelkey::String, exp::Int; uncompress = true)
    models = lprocdat(iJR904, "base_models.bson"; verbose = false);
    model = models[modelkey][exp]
    return uncompress ? ChU.uncompressed_model(model) : model
end

function _load_raw_model()
    src_file = rawdir(iJR904, "iJR904.mat")
    mat_model = MAT.matread(src_file)["model"]
    return ChU.MetNet(mat_model; reshape = true)
end


"""
    PowerNetwork <: AbstractPowerSystem

This structure contains constant parameters that define the topology and
physics of the power network.

The object `PowerNetwork` uses its own contiguous indexing for the buses.
The indexing is independent from those specified in the Matpower or the
PSSE input file. However, a correspondence between the two indexing
(Input indexing *to* `PowerNetwork` indexing) is stored inside the
attribute `bus_to_indexes`.

## Note
The object `PowerNetwork` is created in the host memory.
Use a `AbstractFormulation` object to move data to the target device.

"""
struct PowerNetwork <: AbstractPowerSystem
    vbus::Vector{Complex{Float64}}
    # Admittance matrix
    Ybus::SparseArrays.SparseMatrixCSC{Complex{Float64},Int64}
    # Lines
    lines::Branches{Complex{Float64}}
    # Data
    buses::Array{Float64, 2}
    branches::Array{Float64, 2}
    generators::Array{Float64, 2}
    costs::Union{Array{Float64, 2}, Nothing}
    baseMVA::Float64

    nbus::Int64
    ngen::Int64

    bustype::Vector{Int64}
    bus_to_indexes::Dict{Int, Int}
    ref::Vector{Int64}
    pv::Vector{Int64}
    pq::Vector{Int64}

    sbus::Vector{Complex{Float64}}
    sload::Vector{Complex{Float64}}

    function PowerNetwork(data::Dict{String, Array}; remove_lines=Int[])
        # Parsed data indexes
        BUS_I, BUS_TYPE, PD, QD, GS, BS, BUS_AREA, VM, VA, BASE_KV, ZONE, VMAX, VMIN,
        LAM_P, LAM_Q, MU_VMAX, MU_VMIN = IndexSet.idx_bus()

        # retrive required data
        bus = data["bus"]
        gen = data["gen"]
        lines = data["branch"]
        SBASE = data["baseMVA"][1]
        costs = Base.get(data, "cost", nothing)

        bus_id_to_indexes = get_bus_id_to_indexes(data["bus"])

        # Remove specified lines
        lines = get_active_branches(lines, remove_lines)

        # size of the system
        nbus = size(bus, 1)
        ngen = size(gen, 1)

        # obtain V0 from raw data
        vbus = zeros(Complex{Float64}, nbus)
        for i in 1:nbus
            vbus[i] = bus[i, VM]*exp(1im * pi/180 * bus[i, VA])
        end

        # form Y matrix
        topology = makeYbus(bus, lines, SBASE, bus_id_to_indexes)

        branches = Branches{Complex{Float64}}(
            topology.yff, topology.yft, topology.ytf, topology.ytt,
            topology.from_buses, topology.to_buses,
        )

        # bus type indexing
        ref, pv, pq, bustype = bustypeindex(bus, gen, bus_id_to_indexes)

        sbus, sload = assembleSbus(gen, bus, SBASE, bus_id_to_indexes)
        Ybus = topology.ybus

        new(vbus, Ybus, branches, bus, lines, gen, costs, SBASE, nbus, ngen, bustype, bus_id_to_indexes, ref, pv, pq, sbus, sload)
    end
end

function PowerNetwork(datafile::String)
    data = import_dataset(datafile)
    return PowerNetwork(data)
end

# Getters
## Network attributes
get(pf::PowerNetwork, ::NumberOfBuses) = pf.nbus
get(pf::PowerNetwork, ::NumberOfLines) = size(pf.branches, 1)
get(pf::PowerNetwork, ::NumberOfGenerators) = pf.ngen
get(pf::PowerNetwork, ::NumberOfPVBuses) = length(pf.pv)
get(pf::PowerNetwork, ::NumberOfPQBuses) = length(pf.pq)
get(pf::PowerNetwork, ::NumberOfSlackBuses) = length(pf.ref)

## Loads
get(pf::PowerNetwork, ::ActiveLoad) = real.(pf.sload)
get(pf::PowerNetwork, ::ReactiveLoad) = imag.(pf.sload)
function get(pf::PowerNetwork, ::ActivePower)
    GEN_BUS, PG, QG, _ = IndexSet.idx_gen()
    return pf.generators[:, PG] ./ pf.baseMVA
end
function get(pf::PowerNetwork, ::ReactivePower)
    GEN_BUS, PG, QG, _ = IndexSet.idx_gen()
    return pf.generators[:, QG] ./ pf.baseMVA
end

## Indexing
function get(pf::PowerNetwork, ::GeneratorIndexes)
    GEN_BUS = IndexSet.idx_gen()[1]
    gens = pf.generators
    ngens = size(gens, 1)
    # Create array on host memory
    indexing = zeros(Int, ngens)
    # Here, we keep the same ordering as specified in Matpower.
    for i in 1:ngens
        indexing[i] = pf.bus_to_indexes[gens[i, GEN_BUS]]
    end
    return indexing
end
get(pf::PowerNetwork, ::PVIndexes) = pf.pv
get(pf::PowerNetwork, ::PQIndexes) = pf.pq
get(pf::PowerNetwork, ::SlackIndexes) = pf.ref

# Pretty printing
function Base.show(io::IO, pf::PowerNetwork)
    println("Power Network characteristics:")
    @printf("\tBuses: %d. Slack: %d. PV: %d. PQ: %d\n", pf.nbus, length(pf.ref),
            length(pf.pv), length(pf.pq))
    println("\tGenerators: ", pf.ngen, ".")
    # Print system status
    @printf("\t==============================================\n")
    @printf("\tBUS \t TYPE \t VMAG \t VANG \t P \t Q\n")
    @printf("\t==============================================\n")

    for i=1:pf.nbus
        type = pf.bustype[i]
        vmag = abs(pf.vbus[i])
        vang = angle(pf.vbus[i])*(180.0/pi)
        pinj = real(pf.sbus[i])
        qinj = imag(pf.sbus[i])
        @printf("\t%d \t  %d \t %1.3f\t%3.2f\t%3.3f\t%3.3f\n", i,
                type, vmag, vang, pinj, qinj)
    end
end

# Some utils function

function bounds(pf::PowerNetwork, ::Buses, ::VoltageMagnitude)
    BUS_I, BUS_TYPE, PD, QD, GS, BS, BUS_AREA, VM, VA, BASE_KV, ZONE, VMAX, VMIN,
    LAM_P, LAM_Q, MU_VMAX, MU_VMIN = IndexSet.idx_bus()

    bus = pf.buses
    v_min = convert.(Float64, bus[:, VMIN])
    v_max = convert.(Float64, bus[:, VMAX])
    return v_min, v_max
end

function bounds(pf::PowerNetwork, ::Generators, ::ActivePower)
    GEN_BUS, PG, QG, QMAX, QMIN, VG, MBASE, GEN_STATUS, PMAX, PMIN, PC1, PC2, QC1MIN,
    QC2MIN, QC2MAX, RAMP_AGC, RAMP_10, RAMP_30, RAMP_Q, APF, MU_PMAG, MU_PMIN, MU_QMAX,
    MU_QMIN = IndexSet.idx_gen()

    gens = pf.generators
    baseMVA = pf.baseMVA

    p_min = convert.(Float64, gens[:, PMIN] / baseMVA)
    p_max = convert.(Float64, gens[:, PMAX] / baseMVA)
    return p_min, p_max
end

function bounds(pf::PowerNetwork, ::Generators, ::ReactivePower)
    GEN_BUS, PG, QG, QMAX, QMIN, VG, MBASE, GEN_STATUS, PMAX, PMIN, PC1, PC2, QC1MIN,
    QC2MIN, QC2MAX, RAMP_AGC, RAMP_10, RAMP_30, RAMP_Q, APF, MU_PMAG, MU_PMIN, MU_QMAX,
    MU_QMIN = IndexSet.idx_gen()

    gens = pf.generators
    baseMVA = pf.baseMVA

    q_min = convert.(Float64, gens[:, QMIN] / baseMVA)
    q_max = convert.(Float64, gens[:, QMAX] / baseMVA)
    return q_min, q_max
end

function bounds(pf::PowerNetwork, ::Lines, ::ActivePower)
    RATE_A = IndexSet.idx_branch()[6]
    n_lines = get(pf, NumberOfLines())
    # Flow min is always set equal to 0
    flow_min = zeros(n_lines)
    flow_max = (pf.branches[:, RATE_A] ./ pf.baseMVA).^2
    # According to the spec, if RATE_A is equal to 0, then the flow
    # is unconstrained.
    unlimited = findall(isequal(0), flow_max)
    flow_max[unlimited] .= Inf
    return (flow_min, flow_max)
end

"""
    get_costs_coefficients(pf::PowerNetwork)

Return coefficients for costs function.

"""
function get_costs_coefficients(pf::PowerNetwork)
    # indexes
    BUS_I, BUS_TYPE, PD, QD, GS, BS, BUS_AREA, VM, VA, BASE_KV, ZONE, VMAX, VMIN,
    LAM_P, LAM_Q, MU_VMAX, MU_VMIN = IndexSet.idx_bus()
    GEN_BUS, PG, QG, QMAX, QMIN, VG, MBASE, GEN_STATUS, PMAX, PMIN, PC1, PC2, QC1MIN,
    QC2MIN, QC2MAX, RAMP_AGC, RAMP_10, RAMP_30, RAMP_Q, APF, MU_PMAG, MU_PMIN, MU_QMAX,
    MU_QMIN = IndexSet.idx_gen()
    MODEL, STARTUP, SHUTDOWN, NCOST, COST = IndexSet.idx_cost()

    ref = pf.ref
    pv = pf.pv
    pq = pf.pq
    b2i = pf.bus_to_indexes

    # Matpower assumes gens are ordered. Generator in row i has its cost on row i
    # of the cost table.
    gens = pf.generators
    baseMVA = pf.baseMVA
    bus = pf.buses
    ngens = size(gens)[1]
    nbus = size(bus)[1]

    # initialize cost
    # store coefficients in a Float64 array, with 4 columns:
    # - 1st column: bus type
    # - 2nd column: constant coefficient c0
    # - 3rd column: coefficient c1
    # - 4th column: coefficient c2
    coefficients = zeros(ngens, 4)
    # If cost is not specified, we return the array coefficients as is
    if isnothing(pf.costs)
        @warn("PowerSystem: cost is not specified in PowerNetwork dataset")
        return coefficients
    end

    cost_data = pf.costs
    # iterate generators and check if pv or ref.
    for i = 1:ngens
        # only 2nd degree polynomial implemented for now.
        @assert cost_data[i, MODEL] == 2
        @assert cost_data[i, NCOST] == 3
        genbus = b2i[gens[i, GEN_BUS]]
        bustype = bus[genbus, BUS_TYPE]

        # polynomial coefficients
        c0 = cost_data[i, COST+2]
        c1 = cost_data[i, COST+1] * baseMVA
        c2 = cost_data[i, COST] * baseMVA^2
        coefficients[i, :] .= (bustype, c0, c1, c2)
    end
    return coefficients
end