# This file defines different sets of functions that are used for comparisons of some of our
# opacity calculations. Each set of functions is organized into it's own submodule.
#
# Tests are explicitly NOT defined in this file so that functions defined in this file can be used
# in scripts outside of this framework (to create plots in order to assess disagreement)

# ================================================================================================

# The first set of functions are defined to facillitate comparisons against the values shown in
# several panels of Figure 8.5 from Gray's 2005 edition of "The Observation and Analysis of Stellar
# Photospheres".

# For these panels, data was extracted using the WebPlotDigitizer tool. We've recorded values with
# greater precision than the values should actually have. The values are primarily meant to be used
# to make sure that our implementation give results to the correct order of magniutde and
# qualitatively has the expected wavelength dependences.

module Gray_opac_compare
using Korg, HDF5
using Korg: Interval, closed_interval, contained
using Interpolations: LinearInterpolation, Throw

# Load Gray05 data for a given panel. This returns a tuple holding two dictionaries:
# 1. The first holds the panel properties (i.e. the temperature and partial electron pressure that
#    was assumed for the calculation)
# 2. The second holds a 2-tuple for each source of opacity source. The first tuple entry specifies
#    wavelengths (in Ångstroms) and the second entry holds the opacity contribution at those
#    wavelengths.
function load_panel_data(panel_name, h5fname = joinpath(@__DIR__, "data/gray05_fig8.5.h5"))
    attribute_dict = Dict{String,Float64}()
    data_dict = Dict{String,Tuple{Array{Float64,1},Array{Float64,1}}}()

    h5open(h5fname, "r") do file
        # at the highest level, the file is divided into one group per panel
        panel_group = file[panel_name]

        # load the attributes from the panel group. This nominally holds the temperature and
        # partial electron pressure for which the values are computed.
        for key in keys(attributes(panel_group))
            attribute_dict[key] = read(attributes(panel_group)[key])
        end

        # In a given panel_group there is one subgroup per opacity source.
        for opacity_subgroup_name in keys(panel_group)
            opacity_subgroup = panel_group[opacity_subgroup_name]
            if !(length(keys(opacity_subgroup)) == 2 &&
                 haskey(opacity_subgroup, "lambda") &&
                 haskey(opacity_subgroup, "data"))
                msg = string("Each opacity subgroup is expected to just have a \"lambda\" ",
                             "dataset and a \"data\" dataset")
                throw(ErrorException(msg))
            end
            tmp = (read(opacity_subgroup["lambda"]), read(opacity_subgroup["data"]))
            if length(tmp[1]) != length(tmp[2])
                throw(ErrorException("Datasets within a given subgroup should have a fixed length"))
            end
            data_dict[opacity_subgroup_name] = tmp
        end
    end
    (attribute_dict, data_dict)
end

# To facillitate comparisons against opacity from H₂⁺ and He⁻, we need to have reasonable estimates
# of the ratio of free electrons to protons. In the future, we might want to hardcode the
# calculated values so that changes to the atomic data doesn't break all of these tests.

# from table D.2. This interpolates the log of the partition function.
# I updated cases were the table records 0.301 with log10(2.0)
function _gray05_H_I_partition_func(T)
    log_2 = log10(2.0)
    θ_vals = [  0.2,   0.4,   0.6,   0.8,   1.0,   1.2,   1.4,   1.6,   1.8,   2.0]
    table  = [0.368, 0.303, log_2, log_2, log_2, log_2, log_2, log_2, log_2, log_2]
    interpolator = LinearInterpolation(θ_vals, table, extrapolation_bc=Throw())
    10^interpolator(5040.0 / T)
end

# compute the ratio of the nₑ to n_H (the number density of all H I and H II)
# this function could be reused in the future for testing solutions to coupled Saha equations
function free_electrons_per_Hydrogen_particle(nₑ, T, abundances = Korg.solar_abundances,
                                              ionization_energies = Korg.ionization_energies,
                                              partition_funcs = Korg.partition_funcs)
    out = 0.0
    for element in 1:Korg.Natoms
        wII, wIII = Korg.saha_ion_weights(T, nₑ, element, Korg.ionization_energies,
                                          Korg.partition_funcs)

        nₑ_per_ndens_species = wII/(1 + wII + wIII) + 2wIII/(1 + wII + wIII)
        abundance = 10.0^(abundances[element]-12.0)
        out += abundance * nₑ_per_ndens_species
    end
    out
end

# At a given Temperature and number density of electrons, generate semi-realistic values for:
#    - number densities of H I
#    - number density of all hydgrogen particles (H I and H II)
#    - the mass density
#
# This is mostly used for comparisons against panels of figure 8.5 from Gray (2005). For most
# curves, we can simply assert that there is a fixed ionization fraction because we divide out the
# dependence on all of these terms during the calculation
function _semi_realisitic_dens(ne::F, fion::F= 0.02, HydrogenMassFrac::F= 0.76) where F <: Real
    # In the interest of semi-realistic numbers, Let's assume all free electrons come from Hydrogen
    nH_I = ne/fion
    nH = ne + nH_I
    ρ = nH * 1.67e-24/HydrogenMassFrac
    (nH_I, nH, ρ)
end

# Gray05 uses some confusing terminology and notation.
# - he defines the "atomic absorption coefficient" as the area per absorber and represent it with
#   the variable α. For reference, Rybicki & Lightman call this same quantity the cross-section and
#   represent it with the variable σ_ν.
# - he defines the "continuous absorption coefficient per neutral hydrogen atom" and denotes it
#   with κ. Rybicki & Lightman don't explicitly define this quantity, but it's compatible with
#   their nomenclature. R&L would express this quantity as α_ν/nH_I (in their notation). This is
#   NOT a cross-section in the general case (e.g. for continuum absorption of He, Gray still
#   defines the "continuous absorption coefficient per neutral hydrogen atom" as κ = α_ν/nH_I).
# - he uses κ_ν to denote the mass opacity (just like Rybicki & Lightman)
#
# Now we actually define the functions that compute the opacities in the form comparable with
# Gray05. These functons return the "continuous absorption coefficient per neutral hydrogen atom"
# divided by Pₑ, the partial electron pressure.
#
# For the sake of clarity, we label Rybicki & Lightman's α_ν = κ_ν*ρ as the "linear absorption
# coefficient" within these functions (note: throughout the rest of the codebase, we interchangably
# use the simpler term, "absorption coefficient", to refer to this same quantity)

# Combined H I bound-free and free-free opacity
function HI_coefficient(λ, T, Pₑ, H_I_ion_energy = 13.598)
    # λ should have units of Ångstroms
    # Pₑ is the partial pressure of the electrons in dyne/cm²
    nₑ =  Pₑ/(Korg.kboltz_cgs * T)
    ν = (Korg.c_cgs*1e8)/λ

    bf_coef = begin
        H_I_partition_val = 2.0 # implicitly in the implementation provided by Gray (2005)
        nH_I = nₑ * 100.0 # this is totally arbitrary
        bf_linear_absorption_coef = Korg.ContinuumAbsorption.H_I_bf([ν], T, nH_I/H_I_partition_val,
                                                                    H_I_ion_energy)[1]
        bf_linear_absorption_coef / (Pₑ * nH_I)
    end

    ff_coef = begin
        nH_total = nₑ * 100.0 # this is totally arbitrary
        ρ = nH_total * 1.67e-24/0.76 # this is totally arbitrary

        #Assume U_I(T) = 2.0 and U_II(T) = 1.0, as in Gray (2005)
        Us = Dict([Korg.species"H_I"=>(T -> 2.0), Korg.species"H_II"=>(T -> 1.0)])
        χs = [(Korg.RydbergH_eV, -1.0, -1.0)]
        wII, wIII = Korg.saha_ion_weights(T, nₑ, 1, χs, Us)

        nH_I = nH_total / (1 + wII)
        nH_II = nH_total * wII/(1 + wII)

        # compute the linear absorption coefficient  = dτ/ds = opacity*ρ
        ff_linear_absorption_coef = Korg.ContinuumAbsorption.H_I_ff([ν], T, nH_II, nₑ)[1]
        ff_linear_absorption_coef / (Pₑ * nH_I)
    end

    bf_coef + ff_coef
end

function Hminus_bf_coefficient(λ, T, Pₑ, ion_energy_H⁻ = 0.7552)
    # λ should have units of Ångstroms
    # Pₑ is the partial pressure of the electrons in dyne/cm²
    ne =  Pₑ/(Korg.kboltz_cgs * T)

    # the values of ndens, ρ and nH_I shouldn't actually matter since we will
    # just divide out all dependence on these variables.
    nH_I, nH, ρ = _semi_realisitic_dens(ne)

    partition_func = 2.0 # may want to include the temperature dependence of the partition function
    ν = (Korg.c_cgs*1e8)/λ

    linear_absorb_coef = Korg.ContinuumAbsorption.Hminus_bf([ν], T, nH_I/partition_func, ne,
                                                            ion_energy_H⁻)[1]
    linear_absorb_coef / (Pₑ * nH_I)
end

function Hminus_ff_coefficient(λ, T, Pₑ)
    # λ should have units of Ångstroms
    # Pₑ is the partial pressure of the electrons in dyne/cm²
    ne =  Pₑ/(Korg.kboltz_cgs * T)

    # the values of ndens, ρ and nH_I shouldn't actually matter since we will
    # just divide out all dependence on these variables.
    nH_I, nH, ρ = _semi_realisitic_dens(ne)

    partition_func = 2.0 # may want to include the temperature dependence of the partition function
    ν = (Korg.c_cgs*1e8)/λ

    linear_absorb_coef = Korg.ContinuumAbsorption.Hminus_ff([ν], T, nH_I/partition_func, ne)[1]
    linear_absorb_coef / (Pₑ * nH_I)
end

# computes the combine H₂⁺ free-free and bound-free absorption in units of cm^2 per H atom (not a
# typo)
# There seems to be an error in our function to compute H₂⁺
function H2plus_coefficient(λ, T, Pₑ)
    # λ should have units of Ångstroms
    # Pₑ is the partial pressure of the electrons in dyne/cm²
    ne =  Pₑ/(Korg.kboltz_cgs * T)

    # in the future, we may want to use the abundances, ionization_energies and partition functions
    # used in Gray (2005) since we are comparing to his plots
    ne_div_nH =
        free_electrons_per_Hydrogen_particle(ne, T)
    nH = ne/ne_div_nH

    wII, wIII = Korg.saha_ion_weights(T, ne, 0x01, Korg.ionization_energies,
                                         Korg.partition_funcs)
    nH_I = nH / (1 + wII)
    nH_II = nH * wII / (1 + wII)

    # set the partition function to 2.0 in nH_I_div_partition so that n(H I, n=1) = n(H I) for this
    # calculation, which is an assumption that Gray (2005) implicitly uses
    nH_I_div_partition = nH_I/2.0
    ρ = 1.0 # arbitrary value because we divide it out after

    ν = (Korg.c_cgs*1e8)/λ
    linear_absorb_coef = Korg.ContinuumAbsorption.H2plus_bf_and_ff([ν], T, nH_I, nH_II)[1]
    linear_absorb_coef / (Pₑ * nH_I)
end

# compute He⁻ free-free absorption in units of cm^2 per H atom (not a typo)
function Heminus_ff_coefficient(λ, T, Pₑ)
    # λ should have units of Ångstroms
    # Pₑ is the partial pressure of the electrons in dyne/cm²
    ne =  Pₑ/(Korg.kboltz_cgs * T)

    # the values of ndens, ρ and nH_I shouldn't actually matter since we will
    # just divide out all dependence on these variables.

    # In the interest of semi-realistic numbers, Let's assume all free
    # electrons come from Hydrogen
    nH_I, nH, ρ = _semi_realisitic_dens(ne)

    # we need to compute the number density of He I. Since this function exists to facillitate
    # comparisons against plots from Gray (2005), we should employ the abundances from
    # their book. They record that there is 8.51e-2 He particles per H particle in table 16.3
    nHe = nH*8.51e-2
    # Gray accounts for only neutral and singly ionized helium in his calculation.
    # We use the ionization energy tabulated in table D.1 (χ(He I) = 24.587).
    # Table D.2 of Gray (2005) records the partion function of He II to be
    # as 10^0.301, the closest value at the table's precision to 2.0
    UI, UII = (1.0, 10^-0.301)
    # Because this neglects He III, we don't use Korg.saha_ion_weights.
    wII = (2.0/ne * (UII/UI) * Korg.translational_U(Korg.electron_mass_cgs, T)
           * exp(-24.587/(Korg.kboltz_eV*T)))
    nHe_I = nHe * 1/(1+wII)

    ν = (Korg.c_cgs*1e8)/λ
    linear_absorb_coef = Korg.ContinuumAbsorption.Heminus_ff([ν], T, nHe_I/UI, ne)[1]
    linear_absorb_coef / (Pₑ * nH_I)
end

# There appears to be some errors in the H₂⁺ opacities, skipping them for now
const Gray05_opacity_form_funcs =
    Dict("H"          => (HI_coefficient,
                          Interval(0, Inf), "H I bf and ff"),
         "Hminus_bf"  => (Hminus_bf_coefficient,
                          closed_interval(2250.0, 15000.0),  "H⁻ bound-free"),
         "Hminus_ff"  => (Hminus_ff_coefficient,
                          closed_interval(2604.0, 113918.0), "H⁻ free-free"),
         "Heminus_ff" => (Heminus_ff_coefficient,
                          closed_interval(5063.0, 151878.0), "He⁻ free-free"),
         "H2plus"     => (H2plus_coefficient,
                          closed_interval(3847.0, 25000.0),   "H₂⁺ ff and bf"),
         )

# the absolute tolerances for values the opacity contributions from
#    - H⁻ bound-free
#    - H⁻ free-free
#    - He⁻ free-free
# There is a much larger discrepancy between the summed H I bf and ff opacity at large λ. But
# unlike for H₂⁺, this could simply be a consequence of the different approximations that are used
const Gray05_atols = Dict("a" => 0.05, "b" => 0.02, "c" => 0.04)

# this function is defined to make tests easier
function Gray05_comparison_vals(panel, opacity_func_name)
    panel_prop, panel_dict = load_panel_data(panel)
    orig_λ_vals, ref_data = panel_dict[opacity_func_name]
    temperature = panel_prop["temperature"]
    Pₑ = 10^panel_prop["log10(electron_pressure)"]

    func, bounds = Gray05_opacity_form_funcs[opacity_func_name][1:2]

    w = contained.(orig_λ_vals, Ref(bounds))
    calculated_vals = func.(orig_λ_vals[w], temperature, Pₑ)
    (calculated_vals*1e26, ref_data[w])
end

end # end the definition of the Gray_opac_compare module

# ================================================================================================

# Below, we define a set of functions that are used to compare our implementation of bound-free
# opacity calculations using data from the Opacity Project against calculations of the same
# opacities using the hydrogenic aproximation. This nominally means that we're comparing
# opacities from H I and He II
#
# In the future, it might also be nice to include comparisons for multi-electron data against
# experimental data.
module OP_compare
using Korg

function calc_hydrogenic_bf_absorption_coef(λ_vals,  T, ndens_species, species_name;
                                            use_OP_data = false)
    @assert species_name in ["H_I", "He_II"]

    if use_OP_data
        cross_sec_file = if species_name == "H_I"
            joinpath(@__DIR__, "data/TOPbase_cross_section_H_I.txt")
        else
            joinpath(@__DIR__, "data/TOPbase_cross_section_He_II.txt")
        end
        Korg.ContinuumAbsorption.absorption_coef_bf_TOPBase(λ_vals.*1e-8, [T], [ndens_species],
                                                            Korg.Species(species_name);
                                                            extrapolation_bc = 0.0,
                                                            cross_sec_file = cross_sec_file)[:, 1]
    else
        ν_vals = (Korg.c_cgs*1e8)./λ_vals # Hz
        ndens_div_partition = ndens_species/Korg.partition_funcs[Korg.Species(species_name)](T)
        κ_dflt_approach = if species_name == "H_I"
            H_I_ion_energy = 13.598
            Korg.ContinuumAbsorption.H_I_bf(ν_vals, T, ndens_div_partition, H_I_ion_energy)
        else
            He_II_ion_energy = 54.418
            Korg.ContinuumAbsorption.He_II_bf(ν_vals, T, ndens_div_partition, He_II_ion_energy)
        end
    end
end

# the following functions could all be refactored so that they are more concise... (they could
# honestly just be constants)

function _dflt_λ_vals(species_name)
    first, last = (species_name == "H_I") ? (80,80000) : (25, 40000)
    map((λ) -> Float32(λ), first:1.0:last)
end

function _λ_comp_intervals(species_name)
    # because the energy state data used by the Opacity Project is somewhat inaccurate (note the
    # disagreement gets worse at higher energy levels), we can only make meaningful comparisons
    # over specific wavelength intervals.

    if species_name == "H_I"
        [80.0       911.0;
         1000.0     3645.0;
         4000.0     8203.0;
         9300.0     14579.0;
         18000.0    22789.0]
    else
        [25.0       10.0^2.35;
         10.0^2.37  10.0^2.95;
         10.0^2.99  10.0^3.3;
         10.0^3.37  10.0^3.55;
         10.0^3.65  10.0^3.75]
    end
end

_hydrogenic_rtol(species_name) = (species_name == "H_I") ? 0.11 : 0.06
end # end the definition of OP_compare module


#various utilities to help with differentiation.

d(::Function) = throw(MethodError(d, (Function)))
dxasy(::Function) = throw(MethodError(dbyx, (Function)))


#abstrcattype.jl

# Metida.jl

abstract type MetidaModel <: StatisticalModel end

abstract type AbstractCovarianceStructure end

abstract type AbstractCovmatMethod end

abstract type AbstractCovarianceType end

abstract type AbstractLMMDataBlocks end

# Data

abstract type AbstractData end

abstract type AbstractDataSet{AbstractData} end

# All have field: id::Dict
abstract type AbstractIdData <: AbstractData end

# All have field:  result::Dict
abstract type AbstractResultData <: AbstractData end

function Base.getindex(a::T, s::Symbol) where T <: AbstractResultData
    return a.result[s]
end

# All have field: result::Dict, data::AbstractIdData
abstract type  AbstractIDResult{AbstractIdData} <: AbstractResultData end

# MetidaFreq.jl

# MetidaNCA.jl
# All have fields: time::Vector, obs::Vector, id::Dict
abstract type AbstractSubject <: AbstractIdData end

# All have field: result::Dict, data::AbstractSubject
#abstract type AbstractSubjectResult{AbstractSubject} <: AbstractResultData end
abstract type AbstractSubjectResult{AbstractSubject} <: AbstractIDResult{AbstractSubject} end

# Descriptive


# Indexing and dimensions

export workitemIdx, workgroupIdx
export threadIdx, blockIdx

@generated function _index(::Val{fname}, ::Val{name}, ::Val{range}) where {fname, name, range}
    T_int32 = LLVM.Int32Type(JuliaContext())

    # create function
    llvm_f, _ = create_function(T_int32)
    mod = LLVM.parent(llvm_f)

    # generate IR
    Builder(JuliaContext()) do builder
        entry = BasicBlock(llvm_f, "entry", JuliaContext())
        position!(builder, entry)

        # call the indexing intrinsic
        intr_typ = LLVM.FunctionType(T_int32)
        intr = LLVM.Function(mod, "llvm.amdgcn.$fname.id.$name", intr_typ)
        idx = call!(builder, intr)

        # attach range metadata
        range_metadata = MDNode([ConstantInt(Int32(range.start), JuliaContext()),
                                 ConstantInt(Int32(range.stop), JuliaContext())],
                                JuliaContext())
        metadata(idx)[LLVM.MD_range] = range_metadata

        ret!(builder, idx)
    end

    call_function(llvm_f, UInt32)
end

# TODO: look these up for the current agent/queue
const max_block_size = (x=1024, y=1024, z=1024)
const max_grid_size  = (x=2147483647, y=65535, z=65535)

for dim in (:x, :y, :z)
    # Workitem index
    fname = Symbol("workitem")
    fn = Symbol("workitemIdx_$dim")
    intr = Symbol("$dim")
    @eval @inline $fn() = Int(_index($(Val(fname)), $(Val(intr)), $(Val(0:max_block_size[dim])))) + 1

    # Workgroup index
    fname = Symbol("workgroup")
    fn = Symbol("workgroupIdx_$dim")
    intr = Symbol("$dim")
    @eval @inline $fn() = Int(_index($(Val(fname)), $(Val(intr)), $(Val(0:max_block_size[dim])))) + 1
end

@inline workitemIdx() = (x=workitemIdx_x(), y=workitemIdx_y(), z=workitemIdx_z())
@inline workgroupIdx() = (x=workgroupIdx_x(), y=workgroupIdx_y(), z=workgroupIdx_z())

# For compat with CUDAnative et. al

@inline threadIdx() = (x=workitemIdx_x(), y=workitemIdx_y(), z=workitemIdx_z())
@inline blockIdx() = (x=workgroupIdx_x(), y=workgroupIdx_y(), z=workgroupIdx_z())


module TestECOS

using Test
import ECOS
import MathOptInterface

const MOI = MathOptInterface

function runtests()
    for name in names(@__MODULE__; all = true)
        if startswith("$(name)", "test_")
            @testset "$(name)" begin
                getfield(@__MODULE__, name)()
            end
        end
    end
    return
end

function test_runtests()
    model = MOI.Utilities.CachingOptimizer(
        MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
        MOI.instantiate(ECOS.Optimizer; with_bridge_type = Float64),
    )
    @test model.optimizer.model.model_cache isa
          MOI.Utilities.UniversalFallback{ECOS.OptimizerCache}
    MOI.set(model, MOI.Silent(), true)
    exclude = String[
        # ZerosBridge does not support ConstraintDual. These are tested below in
        # test_runtests_ZerosBridge
        "test_conic_RotatedSecondOrderCone_INFEASIBLE_2",
        "test_conic_linear_VectorOfVariables_2",
        "test_linear_integration",
        "test_quadratic_constraint_GreaterThan",
        "test_quadratic_constraint_LessThan",
    ]
    if Sys.WORD_SIZE == 32
        # These tests fail on x86 Linux, returning ITERATION_LIMIT instead of
        # proving {primal,dual}_INFEASIBLE.
        push!(exclude, "test_conic_linear_INFEASIBLE")
        push!(exclude, "test_solve_TerminationStatus_DUAL_INFEASIBLE")
    end
    MOI.Test.runtests(
        model,
        MOI.Test.Config(
            atol = 1e-3,
            rtol = 1e-3,
            exclude = Any[
                MOI.ConstraintBasisStatus,
                MOI.VariableBasisStatus,
                MOI.ObjectiveBound,
            ],
        ),
        exclude = exclude,
    )
    return
end

function test_runtests_ZerosBridge()
    optimizer = MOI.instantiate(ECOS.Optimizer; with_bridge_type = Float64)
    MOI.Bridges.remove_bridge(
        optimizer,
        MOI.Bridges.Variable.ZerosBridge{Float64},
    )
    model = MOI.Utilities.CachingOptimizer(
        MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
        optimizer,
    )
    MOI.Test.runtests(
        model,
        MOI.Test.Config(
            atol = 1e-3,
            rtol = 1e-3,
            exclude = Any[
                MOI.ConstraintBasisStatus,
                MOI.VariableBasisStatus,
                MOI.ObjectiveBound,
            ],
        ),
        include = String[
            # ZerosBridge does not support ConstraintDual
            "test_conic_RotatedSecondOrderCone_INFEASIBLE_2",
            "test_conic_linear_VectorOfVariables_2",
            "test_linear_integration",
            "test_quadratic_constraint_GreaterThan",
            "test_quadratic_constraint_LessThan",
        ],
    )
    return
end

function test_RawOptimizerAttribute()
    model = ECOS.Optimizer()
    MOI.set(model, MOI.RawOptimizerAttribute("abstol"), 1e-5)
    @test MOI.get(model, MOI.RawOptimizerAttribute("abstol")) ≈ 1e-5
    MOI.set(model, MOI.RawOptimizerAttribute("abstol"), 2e-5)
    @test MOI.get(model, MOI.RawOptimizerAttribute("abstol")) == 2e-5
    return
end

function test_iteration_limit()
    v = [5.0, 3.0, 1.0]
    w = [2.0, 1.5, 0.3]
    solver = MOI.OptimizerWithAttributes(ECOS.Optimizer, MOI.Silent() => true)
    model = MOI.instantiate(solver, with_bridge_type = Float64)
    maxit = 1
    MOI.set(model, MOI.RawOptimizerAttribute("maxit"), maxit)
    MOI.set(model, MOI.Silent(), true)
    x = MOI.add_variables(model, 3)
    MOI.add_constraint.(model, x, MOI.Interval(0.0, 1.0))
    MOI.add_constraint(
        model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(w, x), 0.0),
        MOI.LessThan(3.0),
    )
    MOI.set(
        model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(v, x), 0.0),
    )
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.ITERATION_LIMIT
    @test MOI.get(model, MOI.BarrierIterations()) == maxit
    return
end

function test_empty_problem()
    model = MOI.Utilities.Model{Float64}()
    ecos = ECOS.Optimizer()
    MOI.optimize!(ecos, model)
    @test MOI.get(ecos, MOI.TerminationStatus()) == MOI.INVALID_MODEL
    @test MOI.get(ecos, MOI.PrimalStatus()) == MOI.NO_SOLUTION
    @test MOI.get(ecos, MOI.DualStatus()) == MOI.NO_SOLUTION
    return
end

function test_conic_no_variables()
    model = MOI.Utilities.Model{Float64}()
    ecos = ECOS.Optimizer()
    f = MOI.VectorAffineFunction(
        MOI.VectorAffineTerm{Float64}[],
        [1.0, 0.5, 0.5],
    )
    MOI.add_constraint(model, f, MOI.SecondOrderCone(3))
    MOI.optimize!(ecos, model)
    @test MOI.get(ecos, MOI.TerminationStatus()) == MOI.INVALID_MODEL
    @test MOI.get(ecos, MOI.PrimalStatus()) == MOI.NO_SOLUTION
    @test MOI.get(ecos, MOI.DualStatus()) == MOI.NO_SOLUTION
    return
end

end  # module

TestECOS.runtests()


export NEAT_populate!, NEAT_evaluate!, GA_NEAT_populate!, NEAT_tournament, NEAT_random_top

function NEAT_tournament(pop::Array{NEATIndiv}, t_size::Int)
    if length(pop) == 1
        return pop[1]
    end
    inds = shuffle!(collect(1:length(pop)))
    sort(pop[inds[1:t_size]])[end]
end

function NEAT_random_top(pop::Array{NEATIndiv}, threshold::Float64)
    sort!(pop, by = x -> x.fitness, rev = true)
    max_index = Integer(ceil(threshold * length(pop)))
    pop[rand(1:max_index)]
end

function NEAT_populate!(e::Evolution, selection::Function)
    # Create offsprings and add to the population
    new_pop::Array{NEATIndiv} = []
    if e.cfg["elites"]>0
        sort!(e.population, by = x -> x.fitness, rev = true)
        for i in 1:e.cfg["elites"]
            push!(new_pop, e.population[i])
        end
    end
    # println("\n-------\nTotal fitness: ", e.cfg["total_fitness"], " Pop: ", length(e.population))
    for s in values(e.cfg["Species"])
        # println("\nFitness: ", s.fitness_val, " Members: ", length(s.members))
        if e.cfg["total_fitness"] > 0
            nb_offspring = Integer(round(
                e.cfg["n_population"] * s.fitness_val / e.cfg["total_fitness"],
            ))
        else
            nb_offspring = length(s.members) #Same species size
        end
        # println("Children: ", nb_offspring)
        reproduction!(s, selection, nb_offspring, e.cfg)
        append!(new_pop, s.members)
    end

    e.population = new_pop
    if e.cfg["verbose"]
        println(
            "\n",
            e.gen,
            " - Pop: ",
            length(e.population),
            "  Species: ",
            length(e.cfg["Species"]),
            " Inno: ",
            e.cfg["innovation_max"],
        )
        println("Best fitness: ", maximum(getfield.(e.population, :fitness))[1])
    end
end

function NEAT_evaluate!(e::Evolution, fitness::Function)
    # Build
    build!.(e.population)
    # Assign species
    if length(e.cfg["Species"]) == 0
        s = Species(e.cfg)
    end
    for i in e.population
        find_species!(i, e.cfg)
    end

    # update distance threshold to keep species number stable
    update_threshold!(e.cfg)

    # Update species dict
    for i in keys(e.cfg["Species"])
        if length(e.cfg["Species"][i].members) == 0
            pop!(e.cfg["Species"], i)
        end
    end

    # Compute fitness
    e.cfg["total_fitness"] = 0
    if e.cfg["use_max_fitness"]
        for s in values(e.cfg["Species"])
            e.cfg["total_fitness"] += compute_fitness_max!(s, fitness)
        end
    else
        for s in values(e.cfg["Species"])
            e.cfg["total_fitness"] += compute_fitness_mean!(s, fitness)
        end
    end
end

function update_threshold!(cfg::Dict)
    if length(cfg["Species"]) == 0
        println("NO SPECIES")
        return nothing
    end
    species_pop = []
    for s in values(cfg["Species"])
        push!(species_pop, length(s.members))
    end
    if cfg["verbose"]
        if length(species_pop) > 1
            sort!(species_pop)
            println(
                "Indiv/species | Mean: ",
                sum(species_pop) / length(species_pop),
                " | Med: ",
                species_pop[Integer(floor(length(species_pop) / 2))],
            )
        else
            println("Only 1 species: ", sum(species_pop), " elements")
        end
    end

    nb_species = length(cfg["Species"])
    nb_excess = nb_species - cfg["target_species_number"]
    dist_mod = cfg["dist_mod"] * abs(nb_excess) / cfg["target_species_number"]
    if nb_excess < 0
        cfg["dist_threshold"] -= dist_mod
        if cfg["verbose"]
            println(
                "Dist Threshold - ",
                dist_mod,
                " -> ",
                cfg["dist_threshold"],
            )
        end
    elseif nb_excess > 0
        cfg["dist_threshold"] += dist_mod
        if cfg["verbose"]
            println(
                "Dist Threshold + ",
                dist_mod,
                " -> ",
                cfg["dist_threshold"],
            )
        end
    end

    if cfg["dist_threshold"] < cfg["dist_mod"]
        cfg["dist_threshold"] = cfg["dist_mod"]
    end
end

function GA_NEAT_populate!(e; args...)
    Cambrian.ga_populate(
        e,
        mutation = NeuroEvolution.mutate,
        crossover = NeuroEvolution.crossover,
        selection = x ->
            tournament_selection(x, e.cfg["tournament_size"]),
    )
end


using NeidSolarScripts    # For continuum, solar rotaiton, differential rotation,...
 using RvSpectMLBase      # For spectra data structures
 using EchelleInstruments # For reading spectra
 using JLD2, FileIO       # For outputs
 using ArgParse           # For command line arguments

 # using CSV, DataFrames, Query, Dates, SHA
 #=
 println("# Parsing arguments...")
    function parse_commandline()
      s = ArgParseSettings( description = "Calculate continuum normalization from NEID L1 FITS files.")
      add_arg_group!(s, "Files", :argg_files)
      @add_arg_table! s begin
          "manifest"
              help = "Manifest file (CVS) containing FITS files to analyze.\nExpects columns named Filename, bjd, and target."
              arg_type = String
              default = "manifest.csv"
              #required = true
          "output"
              help = "Filename for output CCFs (jld2)"
              arg_type = String
              default = "daily_ccfs.jld2"
          "--overwrite"
             help = "Specify it's ok to overwrite the output file."
             action = :store_true
       end
       add_arg_group!(s, "Normalization parameters", :argg_norm_param)
       @add_arg_table! s begin
         "--orders_to_use"
            help = "First and last order _index_ to compute CCFs for"
            nargs = 2
            arg_type = Int64
            #default = [ min_order(NEID2D()), max_order(NEID2D()) ]
            #default = [ first(orders_to_use_default(NEID2D())), last(orders_to_use_default(NEID2D())) ]
       end
       add_arg_group!(s, "Filter manifest for ", :argg_filter_param)
       @add_arg_table! s begin
          "--target"
             help = "Target field"
             arg_type = String
             default = "Sun"
          "--datestr"
             help = "Filenames that contain date string."
             arg_type = String
         "--max_solar_hour_angle"
            help = "Maximum absolute value of solar hour angle."
            arg_type = Float64
            default = 6.0
          "--max_airmass"
             help = "Maximum airmass."
             arg_type = Float64
             #default = 10.0
          "--min_snr_factor"
             help = "Minimum SNR relative to max_snr."
             arg_type = Float64
             #default = 0.0
          "--max_snr"
             help = "Specify max_snr manually."
             arg_type = Float64
          "--start_time"
             help = "Specify daily start time for CCFs to be processed (HH MM)"
             nargs = 2
             arg_type = Int64
             #default = [ min_order(NEID2D()), max_order(NEID2D()) ]
             default = [0, 0]
         "--stop_time"
            help = "Specify daily stop time for CCFs to be processed (HH MM)"
            nargs = 2
            arg_type = Int64
            #default = [ min_order(NEID2D()), max_order(NEID2D()) ]
            default = [23, 59]
          "--max_num_spectra"
             help = "Maximum number of spectra to process."
             arg_type = Int
             default = 250
       end

      return parse_args(s)
  end
  args = parse_commandline()
  =#

sed_path = "/home/eford/Code/RvSpectMLEcoSystem/NeidSolarScripts/data"
   sed = Continuum.read_master_sed_neid(path=sed_path)

fits_path = "/mnt/data_simons/NEID/DRPv0.7-fixedflatfielding2/"
   fits_fn = "neidL1_20210104T182500.fits"
   continuum_afs_path = joinpath(fits_path,"output","continuum")

   continuum_afs_fn = "neidL1_20210104T182500_continuum=afs.jld2"
   continuum_ras_path = joinpath(fits_path,"output","continuum")
   continuum_ras_fn = "neidL1_20210104T182500_continuum=ras1.jld2"
   write_output = false
   spec = NEID.read_data(joinpath(fits_path,fits_fn))

min_order_for_continuum = min_order(NEID2D())
  max_order_for_continuum = max_order(NEID2D())
  orders_to_use = min_order_for_continuum:max_order_for_continuum
  (f_norm, var_norm) = Continuum.normalize_by_sed(spec.λ,spec.flux,spec.var, sed; orders_to_use=orders_to_use)

using Plots
order_idx = 110
 plt = scatter(spec.λ[:,order_idx],f_norm[:,order_idx],ms=1.5,legend=:none)
 #plot!(plt, lambda,poly, color=:blue, lw=2)
 #plt = scatter(lambda,f_obs./(sed_view.*coeff.(lambda)),ms=1.5,legend=:none)
 #plt
#savefig("fig.png")
 ylims!(0.,2.0)
true

 #f_norm = spec.flux ./sed
 #xlabel!("λ")
 #ylabel!("Flux/SED")
 #savefig("fig.png")

get_pixel_range_for_continuum = GetPixelRangeFunctor(NEID2D())

#f_norm = spec.flux
@time (anchors, continuum_output, f_filtered) = Continuum.calc_continuum(spec.λ,f_norm,var_norm,stretch_factor=5.0, ν=1.0,
               merging_threshold=0.25, orders_to_use=orders_to_use, get_pixel_range=get_pixel_range_for_continuum,
               verbose=false)

order_idx = 90
  lambda_obs = view(spec.λ,:,order_idx)
  f_obs = view(f_norm,:,order_idx)
  var_obs = view(var_norm,:,order_idx)
  (anchors, continuum, f_filtered_internal) = Continuum.calc_continuum(lambda_obs,f_obs,var_obs,
       stretch_factor=5.0, ν=1.0, merging_threshold=0.25,  min_R_factor = 100.0, smoothing_half_width = 40, verbose=false)
  f_filtered = f_filtered_internal
  scatter(lambda_obs,f_obs, color=:black, ms=2, legend=:none)
  plot!(lambda_obs,f_filtered,color=:black)
  plot!(lambda_obs,continuum,color=:cyan)
  scatter!(lambda_obs[anchors], f_filtered[anchors], color=:blue, ms=4)

if write_output
  jldopen(continuum_ras_fn, "w") do file
    file["order_idx_with_continuum"] = orders_to_use
    file["continuum_anchors_idx"] = anchors
    file["continuum"] = continuum_output
    file["flux_filtered"] = f_filtered
  end
end

using Plots
order_idx = 50
 println("# Order Index: ", order_idx)
 lambda_obs = view(spec.λ,:,order_idx)
 f_obs = view(f_norm,:,order_idx)
 var_obs = view(var_norm,:,order_idx)
 (anchors, continuum, f_filtered_internal) = Continuum.calc_continuum(lambda_obs,f_obs,var_obs,
      stretch_factor=10.0, ν=1.0, merging_threshold=0.25, smoothing_half_width = 40, min_R_factor = 200.0, verbose=false)
  f_filtered = f_filtered_internal
  println(" size(anchors) = ", size(anchors), " size(continuum) = ", size(continuum), " size(f_filtered) = ", size(f_filtered) )

  #plt = scatter(lambda_obs,f_obs, color=:black, ms=2, legend=:none)
  plot!(plt,lambda_obs,f_filtered,color=:black)
  plot!(plt,lambda_obs,continuum,color=:cyan)
  #=
  if length(anchors) >=1
    scatter!(plt,lambda_obs[anchors], f_filtered[anchors], color=:blue, ms=4)
  end
  plt
  =#
  #=
  (anchors2, continuum2, f_filtered_internal2) = Continuum.calc_continuum(lambda,f_norm[:,order_idx],
       stretch_factor=5.0, ν=1.0, merging_threshold=0.25, smoothing_half_width = 40, verbose=false)
  println("# num_anchors1 = ", length(anchors), "  num_anchors2 = ", length(anchors2))
  plot!(lambda,continuum2,color=:magenta)
  scatter!(lambda[anchors2], f_filtered_internal2[anchors2], color=:red, ms=4)
  =#

xmin = first(lambda)-0.05; xmax = xmin+20;
 xlims!(xmin,xmax)
 idx_pix = searchsortedfirst(lambda,xmin):searchsortedlast(lambda,xmax)
 ylims!(NaNMath.extrema(f_obs[idx_pix]))

xmax = last(lambda)+0.05; xmin = xmax-20;
  xlims!(xmin,xmax)
  idx_pix = searchsortedfirst(lambda,xmin):searchsortedlast(lambda,xmax)
  ylims!(NaNMath.extrema(f_obs[idx_pix]))

  #xmin = 4570; xmax = 4580;
  xmin = 4945; xmax = 4965;
    xlims!(xmin,xmax)
    idx_pix = searchsortedfirst(lambda,xmin):searchsortedlast(lambda,xmax)
    ylims!(NaNMath.extrema(f_obs[idx_pix]))

xlims!(10840,10860)
xlims!(5880,5915)
  #xlims!(xmin,xmax)
xlims!(5350,5385)
xlims!(5385,5400)

xlims!(4950,4960)
true

  xlims!(5875,5905)
  #ylims!(0.95,1.15)
  #xmin = minimum(lambda_plt)
  #xmax = minimum(lambda_plt)
  #xlims!(xmin,xmax)
  #ylims!(extrema(f_filtered[xmin.<=lambda.<=xmax]))

  lambda = spec.λ[:,order_idx]
  f_obs = spec.flux[:,order_idx]
  findall(isnan.(f_obs))

 (anchors, continuum, f_filtered_internal) = Continuum.calc_continuum(lambda,f_obs,
       stretch_factor=5.0, ν=1.0, merging_threshold=0.25, verbose=true)

@time continuum = Continuum.calc_continuum(lambda,f_obs,stretch_factor=5.0, ν=1.0, merging_threshold=0.25, verbose=false)

#using Plots
continuum = zeros(size(spec.flux))
  for order_idx in 4:118
    #order_idx =30
    #possible_pix = get_pixel_range(get_inst(spec),order_idx)
    possible_pix = min_pixel_in_order(get_inst(spec)):max_pixel_in_order(get_inst(spec))
    bad_pix = bad_col_ranges(get_inst(spec),order_idx)
    pix_rng = EchelleInstruments.calc_complement_index_ranges(possible_pix,bad_pix)
    pix = mapreduce(p->collect(p),vcat,pix_rng)
    lambda = spec.λ[pix,order_idx]
    f_obs = spec.flux[pix,order_idx]
    (anchors, continuum_order, f_filtered) = Continuum.calc_continuum(lambda,f_obs,λout=spec.λ[possible_pix,order_idx],
         stretch_factor=10.0, merging_threshold=1.0, verbose=true)
    continuum[possible_pix,order_idx] .= continuum_order
    lambda_plt = spec.λ[possible_pix,order_idx]
    plot(lambda_plt,continuum_order)
    scatter!(lambda[anchors], f_filtered[anchors], color=:blue, ms=4)
  end

plot(spec.λ[:,60],continuum[:,60])
true

plot(lambda, spec.flux[:,order_idx]./continuum, legend=:none)
 xlabel!("λ")
 ylabel!("Flux (continuum-normalized)")
 ylims!(0.9,1.05)
 #ylims!(extrema( f_obs./continuum))
 #xlims!(minimum(lambda),5340)
true
=#

#=
  #plot!(lambda,f_clean,ms=1.5, color=:green)
  #plot!(lambda,f_clip_threshold,color=:blue)
  #ylims!(0.9,1.2)
  #plot!(lambda,f_median_filtered.-f_clean)

 #plot!(lambda,rollingpin_radius./rollingpin_R_max)
 #plot(lambda,f_obs, legend=:none)
 #plot(lambda,f_obs, legend=:none)
 #plot!(lambda,f_interp)
 anchor_vals = f_median_filtered[anch]
 #scatter!(lambda[anch],anchor_vals, color=:blue)
 scatter!(lambda[anch],anchor_vals, color=:red)
 scatter!(lambda[anch[anch_mask]],anchor_vals[anch_mask], color=:blue)
 #interp_cubic = CubicSpline(f_median_filtered[anch],anchor_vals )
 #plot!(lambda,interp_cubic.(lambda), color=:green)
 ylims!(0.9,1.2)
 #xlims!(5006.5,5008)

  scatter!(lambda[anchors_merged],f_median_filtered[anchors_merged], color=:magenta)
 plot!(lambda,calc_continuum_from_anchors(lambda,f_median_filtered,anchors_merged),color=:red)

 #=
 plot!(λ_med_filtered,f_median_filtered)
 plot!(λ_med_filtered,f_median_filtered.+iqr)
 plot!(λ_med_filtered,f_median_filtered.-iqr)
 plot!(λ_med_filtered,f_interp, ms=2.0)
 scatter!(λ_med_filtered[anch],f_median_filtered[anch], ms=4.0)
 =#
 #ylims!(0.9,1.03)

#interp_bsp = BSplineApprox(lambda, f_obs,3,4,:ArcLen,:Average)
=#


module Images

import Base.Order: Ordering, ForwardOrdering, ReverseOrdering
import Base.Graphics: width, height
importall Base

using Color

# We need a couple of extra features not present in Color
immutable RGB8 <: ColorValue
    r::Uint8
    g::Uint8
    b::Uint8

    function RGB8(r::Real, g::Real, b::Real)
        new(r, g, b)
    end

    RGB8() = RGB8(0,0,0)
end

typemin(::Type{RGB}) = RGB(0,0,0)
typemax(::Type{RGB}) = RGB(1,1,1)

if VERSION.minor < 3
    using Cartesian
else
    using Base.Cartesian
    include("compatibility/forcartesian.jl")
end

# if isdefined(module_parent(Images), :Grid)
#     import ..Grid.restrict
# end

include("core.jl")
include("iterator.jl")
const have_imagemagick = include("ioformats/libmagickwand.jl")
@osx_only include("ioformats/OSXnative.jl")
include("io.jl")
include("scaling.jl")
include("labeledarrays.jl")
include("algorithms.jl")
include("connected.jl")
include("edge.jl")

__init__() = LibMagick.init()

function precompile()
    for T in (Uint8, Uint16, Int, Float32, Float64)
        Tdiv = typeof(one(T)/2)
        for N = 2:3
            precompile(restrict!, (Array{Tdiv, N}, Array{T,N}, Int))
            precompile(imfilter, (Array{T,N}, Array{Float64,N}))
            precompile(imfilter, (Array{T,N}, Array{Float64,2}))
            precompile(imfilter, (Array{T,N}, Array{Float32,N}))
            precompile(imfilter, (Array{T,N}, Array{Float32,2}))
        end
    end
    for T in (Float32, Float64)
        for N = 2:3
            precompile(_imfilter_gaussian!, (Array{T,N}, Vector{Float64}))
            precompile(_imfilter_gaussian!, (Array{T,N}, Vector{Int}))
            precompile(imfilter_gaussian_no_nans!, (Array{T,N}, Vector{Float64}))
            precompile(imfilter_gaussian_no_nans!, (Array{T,N}, Vector{Int}))
            precompile(fft, (Array{T,N},))
        end
    end
    for T in (Complex{Float32}, Complex{Float64})
        for N = 2:3
            precompile(ifft, (Array{T,N},))
        end
    end
end

export # types
    AbstractImage,
    AbstractImageDirect,
    AbstractImageIndexed,
    Image,
    ImageCmap,
    BitShift,
    ClipMin,
    ClipMax,
    ClipMinMax,
    LabeledArray,
    Overlay,
    OverlayImage,
    ScaleAutoMinMax,
    ScaleInfo,
    ScaleMinMax,
    ScaleNone,
    ScaleSigned,
    SliceData,

    # constants
    palette_fire,
    palette_gray32,
    palette_gray64,
    palette_rainbow,

    # core functions
    assert2d,
    assert_scalar_color,
    assert_timedim_last,
    assert_xfirst,
    assert_yfirst,
    colordim,
    colorspace,
    coords,
    coords_spatial,
    data,
    dimindex,
    dimindexes,
    getindexim,
    grayim,
    colorim,
    height,
    isdirect,
    isxfirst,
    isyfirst,
    limits,
    maxabsfinite,
    maxfinite,
    minfinite,
    ncolorelem,
    nimages,
    pixelspacing,
    properties,
    refim,
    rerange!,
    reslice!,
    restrict,
    sdims,
    size_spatial,
    share,
    sliceim,
    spacedirections,
    spatialorder,
    spatialpermutation,
    spatialproperties,
    storageorder,
    subim,
    timedim,
    width,
    widthheight,

    # io functions
    add_image_file_format,
    imread,
    imwrite,
    loadformat,

    # iterator functions
    first_index,
    iterate_spatial,
    parent,

    # color-related functions
    alphaval,
    blueval,
    greenval,
    redval,
    red,
    green,
    blue,
    hsi2rgb,
    imadjustintensity,
    indexedcolor,
    lut,
    ntsc2rgb,
    rgb2gray,
    rgb2hsi,
    rgb2ntsc,
    rgb2ycbcr,
    uint32color,
    uint32color!,
    ycbcr2rgb,

    # Scaling of intensity
    climdefault,
    float32sc,
    float64sc,
    sc,
    scale,
    scaleinfo,
    scaleminmax,
    scalesigned,
    uint8sc,
    uint16sc,
    uint32sc,

    # algorithms
    ando3,
    ando4,
    ando5,
    backdiffx,
    backdiffy,
    dilate,
    erode,
    opening,
    closing,
    forwarddiffx,
    forwarddiffy,
    gaussian2d,
    imaverage,
    imcomplement,
    imdog,
    imedge,
    imfilter,
    imfilter_fft,
    imfilter_gaussian,
    imfilter_gaussian!,
    imfilter_LoG,
    imgaussiannoise,
    imgradients,
    imlaplacian,
    imlineardiffusion,
    imlog,
    imROF,
    imstretch,
#     imthresh,
    label_components,
    label_components!,
    magnitude,
    magnitude_phase,
    ncc,
    orientation,
    padarray,
    phase,
    prewitt,
    sad,
    sadn,
    sobel,
    ssd,
    ssdn,
    thin_edges,
    thin_edges_subpix,
    thin_edges_nonmaxsup,
    thin_edges_nonmaxsup_subpix,

    # phantoms
    shepp_logan

@deprecate cairoRGB     uint32color!
@deprecate refim        getindexim
@deprecate scaledefault climdefault

if VERSION < v"0.3-"
  __init__()
end

end


"""

```
mddf_linkedcells_parallel(trajectory::Trajectory, options::Options,samples::Samples, RNG, mddf_compute!)
```

Computes the MDDF using linked cells, in parallel.

"""
function mddf_linkedcells_parallel(
    trajectory::Trajectory,
    options::Options,
    samples::Samples,
    RNG,
    mddf_compute!,
)

    # Simplify code by assigning some shortened names
    solute = trajectory.solute
    solvent = trajectory.solvent
    x_solute = trajectory.x_solute
    x_solvent = trajectory.x_solvent

    # Number of threads
    if options.nthreads < 0
        nspawn = Threads.nthreads() - 1
    else
        nspawn = options.nthreads - 1
    end
    if nspawn == 0
        error(" Parallel version must be executed only with more than 1 thread. ")
    end

    # Initializing the structure that carries the result per thread
    R0 = Result(trajectory, options)
    R = [Result(trajectory, options, irefatom = R0.irefatom) for i = 1:nspawn]

    # Safe passing of frame counter to the threads
    tframe = zeros(Int, nspawn)

    # Create data structures required for multithreading
    framedata = Vector{FrameData}(undef, nspawn)
    for ispawn = 1:nspawn
        framedata[ispawn] = FrameData(deepcopy(trajectory), R[1])
    end

    # Tasks
    t = Vector{Task}(undef, nspawn)
    free = ones(Bool, nspawn)

    # Skip initial frames if desired
    iframe = 0 # Counter for all frames of the input file
    while iframe < options.firstframe - 1
        nextframe!(trajectory)
        iframe += 1
    end

    # Print some information about this run
    options.silent || title(R[1], solute, solvent, nspawn)

    ndone = 0
    iframe_read = 0 # Counter for the frames that are actually being considered
    if !options.silent
        progress = Progress(R[1].nframes_read, 1)
    end
    while ndone < R[1].nframes_read

        # Launch for each free thread the computation of a frame
        while iframe_read < R[1].nframes_read && count(free) > 0
            ifree = findfirst(x -> x == true, free)
            # Skip frames that are not going to be considered
            while (iframe + 1) % options.stride != 0
                nextframe!(trajectory)
                iframe += 1
            end
            # Actually read the data of the frame that will be considered
            nextframe!(trajectory)
            iframe += 1
            iframe_read += 1
            tframe[ifree] = iframe
            @. framedata[ifree].trajectory.x_solute = trajectory.x_solute
            @. framedata[ifree].trajectory.x_solvent = trajectory.x_solvent
            @. framedata[ifree].trajectory.sides = trajectory.sides
            # Spawn the calculations for this frame
            t[ifree] = ThreadPools.@tspawnat ifree + 1 mddf_compute!(
                tframe[ifree],
                framedata[ifree],
                options,
                RNG,
                R[ifree],
            )
            free[ifree] = false
        end

        # Wait a little bit before checking
        sleep(options.sleep)

        # Check thread status
        for ispawn = 1:nspawn
            if !free[ispawn]
                if istaskfailed(t[ispawn])
                    error(
                        " Computation of MDDF failed in thread: $ispawn",
                        fetch(t[ispawn]),
                    )
                end
                if istaskdone(t[ispawn])
                    ndone += 1
                    free[ispawn] = true
                    options.silent || next!(progress)
                    if options.GC &&
                       (Sys.free_memory() / Sys.total_memory() < options.GC_threshold)
                        GC.gc() # why we need this anyway??? There should not be so much garbage.
                    end
                end
            end
        end

    end
    closetraj(trajectory)

    # Sum up the results of all threads into the data of thread one (R1<-R1+R2)
    for ispawn = 2:nspawn
        sum!(R[1], R[ispawn])
    end

    # Setup the final data structure with final values averaged over the number of frames,
    # sampling, etc, and computes final distributions and integrals
    finalresults!(R[1], options, trajectory, samples)
    options.silent || println(bars)

    return R[1]

end


# Implicit solvent models
# Based on the OpenMM source code

export
    AbstractGBSA,
    ImplicitSolventOBC,
    ImplicitSolventGBN2,
    born_radii_and_grad

"""
Generalized Born (GB) implicit solvent models augmented with the
hydrophobic solvent accessible surface area (SA) term.
Custom GBSA methods should sub-type this type.
"""
abstract type AbstractGBSA end

# Default solvent dielectric is 78.5 for consistency with AMBER
# Elsewhere it is 78.3
const gb_solvent_dielectric = 78.5
const gb_solute_dielectric = 1.0

const obc_offset = 0.009u"nm"
const gbn2_offset = 0.0195141u"nm"

const gb_probe_radius = 0.14u"nm"
const gb_sa_factor = 28.3919551u"kJ * mol^-1 * nm^-2"

const gbn2_neck_scale = 0.826836
const gbn2_neck_cut = 0.68u"nm"

const mbondi2_element_to_radius = Dict(
    "N"     => 0.155u"nm",
    "O"     => 0.15u"nm" ,
    "F"     => 0.15u"nm" ,
    "Si"    => 0.21u"nm" ,
    "P"     => 0.185u"nm",
    "S"     => 0.18u"nm" ,
    "Cl"    => 0.17u"nm" ,
    "C"     => 0.17u"nm" ,
    "H"     => 0.12u"nm" ,
    "H_N"   => 0.13u"nm" ,
    "H_ARG" => 0.117u"nm",
    "O_CAR" => 0.14u"nm" ,
    "-"     => 0.15u"nm" ,
)

const obc_element_to_screen = Dict(
    "H" => 0.85,
    "C" => 0.72,
    "N" => 0.79,
    "O" => 0.85,
    "F" => 0.88,
    "P" => 0.86,
    "S" => 0.96,
    "-" => 0.80,
)

const gbn2_element_to_screen = Dict(
    "H" => 1.425952,
    "C" => 1.058554,
    "N" => 0.733599,
    "O" => 1.061039,
    "F" => 0.5,
    "P" => 0.5,
    "S" => -0.703469,
    "-" => 0.5,
)

const gbn2_element_to_screen_nucleic = Dict(
    "H" => 1.696538,
    "C" => 1.268902,
    "N" => 1.4259728,
    "O" => 0.1840098,
    "F" => 0.5,
    "P" => 1.5450597,
    "S" => 0.05,
    "-" => 0.5,
)

const gbn2_atom_params = Dict(
    "H_α" => 0.788440, "H_β" => 0.798699, "H_γ" => 0.437334,
    "D_α" => 0.788440, "D_β" => 0.798699, "D_γ" => 0.437334,
    "C_α" => 0.733756, "C_β" => 0.506378, "C_γ" => 0.205844,
    "N_α" => 0.503364, "N_β" => 0.316828, "N_γ" => 0.192915,
    "O_α" => 0.867814, "O_β" => 0.876635, "O_γ" => 0.387882,
    "S_α" => 0.867814, "S_β" => 0.876635, "S_γ" => 0.387882,
    "-_α" => 1.0     , "-_β" => 0.8     , "-_γ" => 4.851   ,
)

const gbn2_atom_params_nucleic = Dict(
    "H_α" => 0.537050, "H_β" => 0.362861, "H_γ" => 0.116704 ,
    "D_α" => 0.537050, "D_β" => 0.362861, "D_γ" => 0.116704 ,
    "C_α" => 0.331670, "C_β" => 0.196842, "C_γ" => 0.093422 ,
    "N_α" => 0.686311, "N_β" => 0.463189, "N_γ" => 0.138722 ,
    "O_α" => 0.606344, "O_β" => 0.463006, "O_γ" => 0.142262 ,
    "S_α" => 0.606344, "S_β" => 0.463006, "S_γ" => 0.142262 ,
    "P_α" => 0.418365, "P_β" => 0.290054, "P_γ" => 0.1064245,
    "-_α" => 1.0     , "-_β" => 0.8     , "-_γ" => 4.851    ,
)

const gbn2_data_d0 = [
    2.26685, 2.32548, 2.38397, 2.44235, 2.50057, 2.55867, 2.61663, 2.67444,
    2.73212, 2.78965, 2.84705, 2.9043, 2.96141, 3.0184, 3.07524, 3.13196,
    3.18854, 3.24498, 3.30132, 3.35752, 3.4136,
    2.31191, 2.37017, 2.4283, 2.48632, 2.5442, 2.60197, 2.65961, 2.71711,
    2.77449, 2.83175, 2.88887, 2.94586, 3.00273, 3.05948, 3.1161, 3.1726,
    3.22897, 3.28522, 3.34136, 3.39738, 3.45072,
    2.35759, 2.41549, 2.47329, 2.53097, 2.58854, 2.646, 2.70333, 2.76056,
    2.81766, 2.87465, 2.93152, 2.98827, 3.0449, 3.10142, 3.15782, 3.21411,
    3.27028, 3.32634, 3.3823, 3.43813, 3.49387,
    2.4038, 2.46138, 2.51885, 2.57623, 2.63351, 2.69067, 2.74773, 2.80469,
    2.86152, 2.91826, 2.97489, 3.0314, 3.08781, 3.1441, 3.20031, 3.25638,
    3.31237, 3.36825, 3.42402, 3.4797, 3.53527,
    2.45045, 2.50773, 2.56492, 2.62201, 2.679, 2.7359, 2.7927, 2.8494, 2.90599,
    2.9625, 3.0189, 3.07518, 3.13138, 3.18748, 3.24347, 3.29937, 3.35515,
    3.41085, 3.46646, 3.52196, 3.57738,
    2.4975, 2.5545, 2.61143, 2.66825, 2.72499, 2.78163, 2.83818, 2.89464,
    2.95101, 3.00729, 3.06346, 3.11954, 3.17554, 3.23143, 3.28723, 3.34294,
    3.39856, 3.45409, 3.50952, 3.56488, 3.62014,
    2.54489, 2.60164, 2.6583, 2.71488, 2.77134, 2.8278, 2.88412, 2.94034,
    2.9965, 3.05256, 3.10853, 3.16442, 3.22021, 3.27592, 3.33154, 3.38707,
    3.44253, 3.49789, 3.55316, 3.60836, 3.66348,
    2.59259, 2.6491, 2.70553, 2.76188, 2.81815, 2.87434, 2.93044, 2.98646,
    3.04241, 3.09827, 3.15404, 3.20974, 3.26536, 3.32089, 3.37633, 3.4317,
    3.48699, 3.54219, 3.59731, 3.65237, 3.70734,
    2.64054, 2.69684, 2.75305, 2.80918, 2.86523, 2.92122, 2.97712, 3.03295,
    3.0887, 3.14437, 3.19996, 3.25548, 3.31091, 3.36627, 3.42156, 3.47677,
    3.5319, 3.58695, 3.64193, 3.69684, 3.75167,
    2.68873, 2.74482, 2.80083, 2.85676, 2.91262, 2.96841, 3.02412, 3.07976,
    3.13533, 3.19082, 3.24623, 3.30157, 3.35685, 3.41205, 3.46718, 3.52223,
    3.57721, 3.63213, 3.68696, 3.74174, 3.79644,
    2.73713, 2.79302, 2.84884, 2.90459, 2.96027, 3.01587, 3.0714, 3.12686,
    3.18225, 3.23757, 3.29282, 3.34801, 3.40313, 3.45815, 3.51315, 3.56805,
    3.6229, 3.67767, 3.73237, 3.78701, 3.84159,
    2.78572, 2.84143, 2.89707, 2.95264, 3.00813, 3.06356, 3.11892, 3.17422,
    3.22946, 3.28462, 3.33971, 3.39474, 3.44971, 3.5046, 3.55944, 3.61421,
    3.66891, 3.72356, 3.77814, 3.83264, 3.8871,
    2.83446, 2.89, 2.94547, 3.00088, 3.05621, 3.11147, 3.16669, 3.22183,
    3.27689, 3.33191, 3.38685, 3.44174, 3.49656, 3.55132, 3.60602, 3.66066,
    3.71523, 3.76975, 3.82421, 3.8786, 3.93293,
    2.88335, 2.93873, 2.99404, 3.04929, 3.10447, 3.15959, 3.21464, 3.26963,
    3.32456, 3.37943, 3.43424, 3.48898, 3.54366, 3.5983, 3.65287, 3.70737,
    3.76183, 3.81622, 3.87056, 3.92484, 3.97905,
    2.93234, 2.9876, 3.04277, 3.09786, 3.15291, 3.20787, 3.26278, 3.31764,
    3.37242, 3.42716, 3.48184, 3.53662, 3.591, 3.64551, 3.69995, 3.75435,
    3.80867, 3.86295, 3.91718, 3.97134, 4.02545,
    2.98151, 3.0366, 3.09163, 3.14659, 3.20149, 3.25632, 3.3111, 3.36581,
    3.42047, 3.47507, 3.52963, 3.58411, 3.63855, 3.69293, 3.74725, 3.80153,
    3.85575, 3.90991, 3.96403, 4.01809, 4.07211,
    3.03074, 3.08571, 3.14061, 3.19543, 3.25021, 3.30491, 3.35956, 3.41415,
    3.46869, 3.52317, 3.57759, 3.63196, 3.68628, 3.74054, 3.79476, 3.84893,
    3.90303, 3.95709, 4.01111, 4.06506, 4.11897,
    3.08008, 3.13492, 3.1897, 3.2444, 3.29905, 3.35363, 3.40815, 3.46263,
    3.51704, 3.57141, 3.62572, 3.67998, 3.73418, 3.78834, 3.84244, 3.8965,
    3.95051, 4.00447, 4.05837, 4.11224, 4.16605,
    3.12949, 3.18422, 3.23888, 3.29347, 3.348, 3.40247, 3.45688, 3.51124,
    3.56554, 3.6198, 3.674, 3.72815, 3.78225, 3.83629, 3.8903, 3.94425,
    3.99816, 4.05203, 4.10583, 4.15961, 4.21333,
    3.17899, 3.23361, 3.28815, 3.34264, 3.39706, 3.45142, 3.50571, 3.55997,
    3.61416, 3.66831, 3.72241, 3.77645, 3.83046, 3.8844, 3.93831, 3.99216,
    4.04598, 4.09974, 4.15347, 4.20715, 4.26078,
    3.22855, 3.28307, 3.33751, 3.39188, 3.4462, 3.50046, 3.55466, 3.6088,
    3.6629, 3.71694, 3.77095, 3.82489, 3.8788, 3.93265, 3.98646, 4.04022,
    4.09395, 4.14762, 4.20126, 4.25485, 4.3084,
]u"nm" ./ 10

const gbn2_data_m0 = [
    0.0381511, 0.0338587, 0.0301776, 0.027003, 0.0242506, 0.0218529,
    0.0197547, 0.0179109, 0.0162844, 0.0148442, 0.0135647, 0.0124243,
    0.0114047, 0.0104906, 0.00966876, 0.008928, 0.0082587, 0.00765255,
    0.00710237, 0.00660196, 0.00614589,
    0.0396198, 0.0351837, 0.0313767, 0.0280911, 0.0252409, 0.0227563,
    0.0205808, 0.0186681, 0.0169799, 0.0154843, 0.014155, 0.0129696,
    0.0119094, 0.0109584, 0.0101031, 0.00933189, 0.0086348, 0.00800326,
    0.00742986, 0.00690814, 0.00643255,
    0.041048, 0.0364738, 0.0325456, 0.0291532, 0.0262084, 0.0236399,
    0.0213897, 0.0194102, 0.0176622, 0.0161129, 0.0147351, 0.0135059,
    0.0124061, 0.0114192, 0.0105312, 0.00973027, 0.00900602, 0.00834965,
    0.0077535, 0.00721091, 0.00671609,
    0.0424365, 0.0377295, 0.0336846, 0.0301893, 0.0271533, 0.0245038,
    0.0221813, 0.0201371, 0.018331, 0.0167295, 0.0153047, 0.014033,
    0.0128946, 0.0118727, 0.0109529, 0.0101229, 0.00937212, 0.00869147,
    0.00807306, 0.00751003, 0.00699641,
    0.0437861, 0.0389516, 0.0347944, 0.0311998, 0.0280758, 0.0253479,
    0.0229555, 0.0208487, 0.0189864, 0.0173343, 0.0158637, 0.0145507,
    0.0133748, 0.0123188, 0.0113679, 0.0105096, 0.0097329, 0.00902853,
    0.00838835, 0.00780533, 0.0072733,
    0.0450979, 0.0401406, 0.0358753, 0.0321851, 0.0289761, 0.0261726,
    0.0237125, 0.0215451, 0.0196282, 0.017927, 0.0164121, 0.0150588,
    0.0138465, 0.0127573, 0.0117761, 0.0108902, 0.0100882, 0.00936068,
    0.00869923, 0.00809665, 0.00754661,
    0.0463729, 0.0412976, 0.0369281, 0.0331456, 0.0298547, 0.026978,
    0.0244525, 0.0222264, 0.0202567, 0.0185078, 0.0169498, 0.0155575,
    0.0143096, 0.0131881, 0.0121775, 0.0112646, 0.010438, 0.00968781,
    0.00900559, 0.00838388, 0.00781622,
    0.0476123, 0.0424233, 0.0379534, 0.034082, 0.0307118, 0.0277645,
    0.0251757, 0.0228927, 0.0208718, 0.0190767, 0.0174768, 0.0160466,
    0.0147642, 0.0136112, 0.0125719, 0.0116328, 0.0107821, 0.0100099,
    0.00930735, 0.00866695, 0.00808206,
    0.0488171, 0.0435186, 0.038952, 0.0349947, 0.0315481, 0.0285324,
    0.0258824, 0.0235443, 0.0214738, 0.0196339, 0.0179934, 0.0165262,
    0.0152103, 0.0140267, 0.0129595, 0.0119947, 0.0111206, 0.0103268,
    0.00960445, 0.00894579, 0.00834405,
    0.0499883, 0.0445845, 0.0399246, 0.0358844, 0.032364, 0.0292822,
    0.0265729, 0.0241815, 0.0220629, 0.0201794, 0.0184994, 0.0169964,
    0.0156479, 0.0144345, 0.0133401, 0.0123504, 0.0114534, 0.0106386,
    0.00989687, 0.00922037, 0.00860216,
    0.0511272, 0.0456219, 0.040872, 0.0367518, 0.0331599, 0.0300142,
    0.0272475, 0.0248045, 0.0226392, 0.0207135, 0.0189952, 0.0174574,
    0.0160771, 0.0148348, 0.0137138, 0.0126998, 0.0117805, 0.0109452,
    0.0101846, 0.00949067, 0.00885636,
    0.0522348, 0.0466315, 0.0417948, 0.0375973, 0.0339365, 0.030729,
    0.0279067, 0.0254136, 0.023203, 0.0212363, 0.0194809, 0.0179092,
    0.016498, 0.0152275, 0.0140807, 0.013043, 0.012102, 0.0112466,
    0.0104676, 0.00975668, 0.00910664,
    0.0533123, 0.0476145, 0.042694, 0.0384218, 0.0346942, 0.0314268,
    0.0285507, 0.026009, 0.0237547, 0.0217482, 0.0199566, 0.018352,
    0.0169108, 0.0156128, 0.0144408, 0.0133801, 0.0124179, 0.011543,
    0.010746, 0.0100184, 0.00935302,
    0.0543606, 0.0485716, 0.04357, 0.0392257, 0.0354335, 0.0321082,
    0.02918, 0.0265913, 0.0242943, 0.0222492, 0.0204225, 0.0187859,
    0.0173155, 0.0159908, 0.0147943, 0.0137111, 0.0127282, 0.0118343,
    0.0110197, 0.0102759, 0.00959549,
    0.0553807, 0.0495037, 0.0444239, 0.0400097, 0.0361551, 0.0327736,
    0.0297949, 0.0271605, 0.0248222, 0.0227396, 0.0208788, 0.0192111,
    0.0177122, 0.0163615, 0.0151413, 0.0140361, 0.013033, 0.0121206,
    0.0112888, 0.0105292, 0.00983409,
    0.0563738, 0.0504116, 0.0452562, 0.0407745, 0.0368593, 0.0334235,
    0.0303958, 0.0277171, 0.0253387, 0.0232197, 0.0213257, 0.0196277,
    0.0181013, 0.0167252, 0.0154817, 0.0143552, 0.0133325, 0.0124019,
    0.0115534, 0.0107783, 0.0100688,
    0.0573406, 0.0512963, 0.0460676, 0.0415206, 0.0375468, 0.0340583,
    0.030983, 0.0282614, 0.0258441, 0.0236896, 0.0217634, 0.020036,
    0.0184826, 0.017082, 0.0158158, 0.0146685, 0.0136266, 0.0126783,
    0.0118135, 0.0110232, 0.0102998,
    0.0582822, 0.0521584, 0.0468589, 0.0422486, 0.038218, 0.0346784,
    0.0315571, 0.0287938, 0.0263386, 0.0241497, 0.0221922, 0.0204362,
    0.0188566, 0.0174319, 0.0161437, 0.0149761, 0.0139154, 0.0129499,
    0.0120691, 0.0112641, 0.0105269,
    0.0591994, 0.0529987, 0.0476307, 0.042959, 0.0388734, 0.0352843,
    0.0321182, 0.0293144, 0.0268225, 0.0246002, 0.0226121, 0.0208283,
    0.0192232, 0.0177751, 0.0164654, 0.015278, 0.0141991, 0.0132167,
    0.0123204, 0.0115009, 0.0107504,
    0.0600932, 0.053818, 0.0483836, 0.0436525, 0.0395136, 0.0358764,
    0.0326669, 0.0298237, 0.0272961, 0.0250413, 0.0230236, 0.0212126,
    0.0195826, 0.0181118, 0.0167811, 0.0155744, 0.0144778, 0.0134789,
    0.0125673, 0.0117338, 0.0109702,
    0.0609642, 0.0546169, 0.0491183, 0.0443295, 0.0401388, 0.036455,
    0.0332033, 0.030322, 0.0277596, 0.0254732, 0.0234266, 0.0215892,
    0.0199351, 0.018442, 0.0170909, 0.0158654, 0.0147514, 0.0137365,
    0.0128101, 0.0119627, 0.0111863,
]u"nm^-1" .* 10

# This is force field dependent
is_carboxylate_O(at_data) = at_data.atom_type == "O2"

function atoms_bonded_to_N(atoms_data, bonds)
    bonded_to_N = falses(length(atoms_data))
    for (i, j) in zip(bonds.is, bonds.js)
        if atoms_data[i].element == "N"
            bonded_to_N[j] = true
        end
        if atoms_data[j].element == "N"
            bonded_to_N[i] = true
        end
    end
    return bonded_to_N
end

function mbondi2_radii(atoms_data, bonds; use_mbondi3=false,
                        element_to_radius=mbondi2_element_to_radius)
    bonded_to_N = atoms_bonded_to_N(atoms_data, bonds)
    return map(atoms_data, bonded_to_N) do at_data, at_bonded_to_N
        if use_mbondi3 && at_data.res_name == "ARG" &&
                (startswith(at_data.atom_name, "HH") || startswith(at_data.atom_name, "HE"))
            radius = element_to_radius["H_ARG"]
        elseif use_mbondi3 && is_carboxylate_O(at_data)
            radius = element_to_radius["O_CAR"]
        elseif at_data.element in ("H", "D")
            radius = at_bonded_to_N ? element_to_radius["H_N"] : element_to_radius["H"]
        else
            radius = dict_get(element_to_radius, at_data.element, element_to_radius["-"])
        end
        return radius
    end
end

function mbondi3_radii(atoms_data, bonds; element_to_radius=mbondi2_element_to_radius)
    return mbondi2_radii(atoms_data, bonds; use_mbondi3=true, element_to_radius=element_to_radius)
end

# We use a full atom pairwise table rather than looking up a value with the atom radius
# This works better with broadcasting
function lookup_table(full_table::AbstractArray{T}, radii) where T
    n_atoms = length(radii)
    table_positions = [(r - 0.1u"nm") * 200 for r in radii]
    # These zero-based indexes are converted to one-based when looking up the full table
    index_1, index_2 = zeros(Int, n_atoms), zeros(Int, n_atoms)
    weight_1, weight_2 = zeros(n_atoms), zeros(n_atoms)
    for (i, p) in enumerate(table_positions)
        if p <= 0.0u"nm"
            weight_1[i] = 1.0
        elseif p >= 20.0u"nm"
            index_1[i] = 20
            weight_1[i] = 1.0
        else
            ps = ustrip(u"nm", p)
            index_1[i] = Int(floor(ps))
            index_2[i] = index_1[i] + 1
            weight_1[i] = index_2[i] - ps
            weight_2[i] = 1.0 - weight_1[i]
        end
    end
    table = zeros(T, n_atoms, n_atoms)
    for i in 1:n_atoms
        for j in 1:n_atoms
            table[j, i] = weight_1[i] * weight_1[j] * full_table[index_1[i] * 21 + index_1[j] + 1] +
                          weight_1[i] * weight_2[j] * full_table[index_1[i] * 21 + index_2[j] + 1] +
                          weight_2[i] * weight_1[j] * full_table[index_2[i] * 21 + index_1[j] + 1] +
                          weight_2[i] * weight_2[j] * full_table[index_2[i] * 21 + index_2[j] + 1]
        end
    end
    return table
end

function lookup_table(full_table::AbstractArray{T}, radii::AbstractArray{<:AbstractFloat}) where T
    return lookup_table(full_table, radii * u"nm")
end

"""
    ImplicitSolventOBC(atoms, atoms_data, bonds)

Onufriev-Bashford-Case GBSA model.
Should be used along with a Coulomb or CoulombReactionField interaction.
The keyword argument `use_OBC2` determines whether to use parameter set
I (`false`, the default) or II (`true`).
"""
struct ImplicitSolventOBC{T, D, V, K, S, F, I} <: AbstractGBSA
    offset_radii::V
    scaled_offset_radii::V
    solvent_dielectric::T
    solute_dielectric::T
    kappa::K
    offset::D
    cutoff::D
    use_ACE::Bool
    α::T
    β::T
    γ::T
    probe_radius::D
    sa_factor::S
    factor_solute::F
    factor_solvent::F
    is::I
    js::I
end

function ImplicitSolventOBC(atoms::AbstractArray{Atom{T, M, D, E}},
                            atoms_data,
                            bonds;
                            solvent_dielectric=gb_solvent_dielectric,
                            solute_dielectric=gb_solute_dielectric,
                            kappa=0.0u"nm^-1",
                            offset=obc_offset,
                            cutoff=0.0u"nm",
                            probe_radius=gb_probe_radius,
                            sa_factor=gb_sa_factor,
                            use_ACE=true,
                            use_OBC2=false,
                            element_to_radius=mbondi2_element_to_radius,
                            element_to_screen=obc_element_to_screen) where {T, M, D, E}
    units = dimension(D) == u"𝐋"
    radii = mbondi2_radii(atoms_data, bonds; element_to_radius=element_to_radius)

    if units
        offset_radii = T.(radii .- offset)
    else
        offset_radii = ustrip.(T.(radii .- offset))
    end
    scaled_offset_radii = map(atoms_data, offset_radii) do at_data, offset_radius
        screen = dict_get(element_to_screen, at_data.element, element_to_screen["-"])
        return T(screen) * offset_radius
    end

    if use_OBC2
        # GBOBCII parameters
        α, β, γ = T(1.0), T(0.8), T(4.85)
    else
        # GBOBCI parameters
        α, β, γ = T(0.8), T(0.0), T(2.909125)
    end

    n_atoms = length(atoms)
    inds_i = hcat([collect(1:n_atoms) for i in 1:n_atoms]...)
    inds_j = permutedims(inds_i, (2, 1))

    coulomb_const = units ? coulombconst : ustrip(coulombconst)
    if !iszero(solute_dielectric)
        factor_solute = -T(coulomb_const) / T(solute_dielectric)
    else
        factor_solute = zero(T(coulomb_const))
    end
    if !iszero(solvent_dielectric)
        factor_solvent = T(coulomb_const) / T(solvent_dielectric)
    else
        factor_solvent = zero(T(coulomb_const))
    end

    if isa(atoms, CuArray)
        or = CuArray(offset_radii) # cu converts to Float32
        sor = CuArray(scaled_offset_radii)
        is, js = cu(inds_i), cu(inds_j)
    else
        or = offset_radii
        sor = scaled_offset_radii
        is, js = inds_i, inds_j
    end

    if units
        return ImplicitSolventOBC{T, D, typeof(or), typeof(T(kappa)), typeof(T(sa_factor)),
                        typeof(factor_solute), typeof(is)}(
                    or, sor, solvent_dielectric, solute_dielectric, T(kappa), offset,
                    cutoff, use_ACE, α, β, γ, probe_radius, T(sa_factor), factor_solute,
                    factor_solvent, is, js)
    else
        return ImplicitSolventOBC{T, T, typeof(or), typeof(T(ustrip(kappa))), T, T, typeof(is)}(
                    or, sor, solvent_dielectric, solute_dielectric, T(ustrip(kappa)),
                    ustrip(offset), ustrip(cutoff), use_ACE, α, β, γ, ustrip(probe_radius),
                    ustrip(sa_factor), factor_solute, factor_solvent, is, js)
    end
end

"""
    ImplicitSolventGBN2(atoms, atoms_data, bonds)

GBn2 solvation model.
Should be used along with a Coulomb or CoulombReactionField interaction.
"""
struct ImplicitSolventGBN2{T, D, VT, VD, K, S, F, I, TD, TM} <: AbstractGBSA
    offset_radii::VD
    scaled_offset_radii::VD
    solvent_dielectric::T
    solute_dielectric::T
    kappa::K
    offset::D
    cutoff::D
    use_ACE::Bool
    αs::VT
    βs::VT
    γs::VT
    probe_radius::D
    sa_factor::S
    factor_solute::F
    factor_solvent::F
    is::I
    js::I
    d0s::TD
    m0s::TM
    neck_scale::T
    neck_cut::D
end

function ImplicitSolventGBN2(atoms::AbstractArray{Atom{T, M, D, E}},
                                atoms_data,
                                bonds;
                                solvent_dielectric=gb_solvent_dielectric,
                                solute_dielectric=gb_solute_dielectric,
                                kappa=0.0u"nm^-1",
                                offset=gbn2_offset,
                                cutoff=0.0u"nm",
                                probe_radius=gb_probe_radius,
                                sa_factor=gb_sa_factor,
                                use_ACE=true,
                                neck_scale=gbn2_neck_scale,
                                neck_cut=gbn2_neck_cut,
                                element_to_radius=mbondi2_element_to_radius,
                                element_to_screen=gbn2_element_to_screen,
                                element_to_screen_nucleic=gbn2_element_to_screen_nucleic,
                                atom_params=gbn2_atom_params,
                                atom_params_nucleic=gbn2_atom_params_nucleic,
                                data_d0=gbn2_data_d0,
                                data_m0=gbn2_data_m0) where {T, M, D, E}
    units = dimension(D) == u"𝐋"
    radii = mbondi3_radii(atoms_data, bonds; element_to_radius=element_to_radius)
    nucleic_acid_residues = ("A", "C", "G", "U", "DA", "DC", "DG", "DT")

    if units
        offset_radii = T.(radii .- offset)
    else
        offset_radii = T.(ustrip.(radii) .- ustrip(offset))
    end
    scaled_offset_radii = map(atoms_data, offset_radii) do at_data, offset_radius
        if at_data.res_name in nucleic_acid_residues
            screen = dict_get(element_to_screen_nucleic, at_data.element, element_to_screen_nucleic["-"])
        else
            screen = dict_get(element_to_screen, at_data.element, element_to_screen["-"])
        end
        return T(screen) * offset_radius
    end

    αs_cpu = map(atoms_data) do at_data
        if at_data.res_name in nucleic_acid_residues
            α = dict_get(atom_params_nucleic, at_data.element * "_α", atom_params_nucleic["-_α"])
        else
            α = dict_get(atom_params, at_data.element * "_α", atom_params["-_α"])
        end
        return T(α)
    end
    βs_cpu = map(atoms_data) do at_data
        if at_data.res_name in nucleic_acid_residues
            β = dict_get(atom_params_nucleic, at_data.element * "_β", atom_params_nucleic["-_β"])
        else
            β = dict_get(atom_params, at_data.element * "_β", atom_params["-_β"])
        end
        return T(β)
    end
    γs_cpu = map(atoms_data) do at_data
        if at_data.res_name in nucleic_acid_residues
            γ = dict_get(atom_params_nucleic, at_data.element * "_γ", atom_params_nucleic["-_γ"])
        else
            γ = dict_get(atom_params, at_data.element * "_γ", atom_params["-_γ"])
        end
        return T(γ)
    end

    n_atoms = length(atoms)
    inds_i = hcat([collect(1:n_atoms) for i in 1:n_atoms]...)
    inds_j = permutedims(inds_i, (2, 1))

    table_d0_units = T.(lookup_table(data_d0, radii))
    table_m0_units = T.(lookup_table(data_m0, radii))
    if units
        table_d0 = table_d0_units
        table_m0 = table_m0_units
    else
        table_d0 = ustrip.(table_d0_units)
        table_m0 = ustrip.(table_m0_units)
    end

    coulomb_const = units ? coulombconst : ustrip(coulombconst)
    if !iszero(solute_dielectric)
        factor_solute = -T(coulomb_const) / T(solute_dielectric)
    else
        factor_solute = zero(T(coulomb_const))
    end
    if !iszero(solvent_dielectric)
        factor_solvent = T(coulomb_const) / T(solvent_dielectric)
    else
        factor_solvent = zero(T(coulomb_const))
    end

    if isa(atoms, CuArray)
        or = CuArray(offset_radii) # cu converts to Float32
        sor = CuArray(scaled_offset_radii)
        is, js = cu(inds_i), cu(inds_j)
        d0s, m0s = CuArray(table_d0), CuArray(table_m0)
        αs, βs, γs = CuArray(αs_cpu), CuArray(βs_cpu), CuArray(γs_cpu)
    else
        or = offset_radii
        sor = scaled_offset_radii
        is, js = inds_i, inds_j
        d0s, m0s = table_d0, table_m0
        αs, βs, γs = αs_cpu, βs_cpu, γs_cpu
    end

    if units
        return ImplicitSolventGBN2{T, D, typeof(αs), typeof(or), typeof(T(kappa)), typeof(T(sa_factor)),
                        typeof(factor_solute), typeof(is), typeof(d0s), typeof(m0s)}(
                    or, sor, solvent_dielectric, solute_dielectric, T(kappa), offset, cutoff,
                    use_ACE, αs, βs, γs, probe_radius, T(sa_factor), factor_solute,
                    factor_solvent, is, js, d0s, m0s, neck_scale, neck_cut)
    else
        return ImplicitSolventGBN2{T, T, typeof(αs), typeof(or), typeof(T(ustrip(kappa))), T, T,
                        typeof(is), typeof(d0s), typeof(m0s)}(
                    or, sor, solvent_dielectric, solute_dielectric, T(ustrip(kappa)), ustrip(offset),
                    ustrip(cutoff), use_ACE, αs, βs, γs, ustrip(probe_radius), ustrip(sa_factor),
                    factor_solute, factor_solvent, is, js, d0s, m0s, neck_scale, ustrip(neck_cut))
    end
end

function born_radii_loop_OBC(coord_i, coord_j, ori, srj, cutoff, box_size)
    I = zero(coord_i[1] / unit(cutoff)^2)
    r = norm(vector(coord_i, coord_j, box_size))
    if iszero(r) || (!iszero(cutoff) && r > cutoff)
        return I
    end
    U = r + srj
    if ori < U
        D_ij = abs(r - srj)
        L = max(ori, D_ij)
        I += (1/L - 1/U + (r - (srj^2)/r)*(1/(U^2) - 1/(L^2))/4 + log(L/U)/(2*r)) / 2
        if ori < (srj - r)
            I += 2 * (1/ori - 1/L)
        end
    end
    return I
end

"""
    born_radii_and_grad(inter, coords, box_size)

Calculate Born radii and gradients of Born radii and surface area overlap
with respect to atomic distance.
Custom GBSA methods should implement this function.
"""
function born_radii_and_grad(inter::ImplicitSolventOBC, coords, box_size)
    coords_i = @view coords[inter.is]
    coords_j = @view coords[inter.js]
    oris = @view inter.offset_radii[inter.is]
    srjs = @view inter.scaled_offset_radii[inter.js]
    loop_res = born_radii_loop_OBC.(coords_i, coords_j, oris, srjs, inter.cutoff,
                                    (box_size,))
    Is = dropdims(sum(loop_res; dims=2); dims=2)
    I_grads = zero(loop_res) ./ unit(inter.cutoff)

    ori = inter.offset_radii
    radii = ori .+ inter.offset
    ψs = Is .* ori
    ψs2 = ψs .^ 2
    α, β, γ = inter.α, inter.β, inter.γ
    tanh_sums = tanh.(α .* ψs .- β .* ψs2 .+ γ .* ψs2 .* ψs)
    Bs = 1 ./ (1 ./ ori .- tanh_sums ./ radii)
    grad_terms = ori .* (α .- 2 .* β .* ψs .+ 3 .* γ .* ψs2)
    B_grads = (1 .- tanh_sums .^ 2) .* grad_terms ./ radii

    return Bs, B_grads, I_grads
end

struct BornRadiiGBN2LoopResult{I, G}
    I::I
    I_grad::G
end

get_I(r::BornRadiiGBN2LoopResult) = r.I
get_I_grad(r::BornRadiiGBN2LoopResult) = r.I_grad

function born_radii_loop_GBN2(coord_i::SVector{D, T}, coord_j, ori, orj, srj, cutoff, offset,
                                neck_scale, neck_cut, d0, m0, box_size) where {D, T}
    I = zero(coord_i[1] / unit(cutoff)^2)
    I_grad = zero(coord_i[1] / unit(cutoff)^3)
    r = norm(vector(coord_i, coord_j, box_size))
    if iszero(r) || (!iszero(cutoff) && r > cutoff)
        return BornRadiiGBN2LoopResult(I, I_grad)
    end
    U = r + srj
    if ori < U
        D_ij = abs(r - srj)
        L = max(ori, D_ij)
        I += (1/L - 1/U + (r - (srj^2)/r)*(1/(U^2) - 1/(L^2))/4 + log(L/U)/(2*r)) / 2
        if ori < (srj - r)
            I += 2 * (1/ori - 1/L)
        end
    end
    radius_i = ori + offset
    radius_j = orj + offset
    if r < (radius_i + radius_j + neck_cut)
        if dimension(T) == u"𝐋"
            r_d0_strip = 10 * ustrip(u"nm", r - d0) # The integral uses Å
        else
            r_d0_strip = 10 * (r - d0)
        end
        denom = 1 + r_d0_strip^2 + 3 * r_d0_strip^6 / 10
        I += neck_scale * m0 / denom
        numer = 2 * r_d0_strip + 9 * r_d0_strip^5 / 5
        I_grad -= 10 * neck_scale * m0 * numer / (denom^2 * unit(cutoff))
    end
    return BornRadiiGBN2LoopResult(I, I_grad)
end

function born_radii_and_grad(inter::ImplicitSolventGBN2, coords, box_size)
    coords_i = @view coords[inter.is]
    coords_j = @view coords[inter.js]
    oris = @view inter.offset_radii[inter.is]
    orjs = @view inter.offset_radii[inter.js]
    srjs = @view inter.scaled_offset_radii[inter.js]
    loop_res = born_radii_loop_GBN2.(coords_i, coords_j, oris, orjs, srjs, inter.cutoff,
                    inter.offset, inter.neck_scale, inter.neck_cut, inter.d0s,
                    inter.m0s, (box_size,))
    Is = dropdims(sum(get_I.(loop_res); dims=2); dims=2)
    I_grads = get_I_grad.(loop_res)

    ori = inter.offset_radii
    radii = ori .+ inter.offset
    ψs = Is .* ori
    ψs2 = ψs .^ 2
    αs, βs, γs = inter.αs, inter.βs, inter.γs
    tanh_sums = tanh.(αs .* ψs .- βs .* ψs2 .+ γs .* ψs2 .* ψs)
    Bs = 1 ./ (1 ./ ori .- tanh_sums ./ radii)
    grad_terms = ori .* (αs .- 2 .* βs .* ψs .+ 3 .* γs .* ψs2)
    B_grads = (1 .- tanh_sums .^ 2) .* grad_terms ./ radii

    return Bs, B_grads, I_grads
end

# Store the results of the ij broadcasts during force calculation
struct ForceLoopResult1{T, V}
    bi::T
    bj::T
    fi::V
    fj::V
end

get_bi(r::ForceLoopResult1) = r.bi
get_bj(r::ForceLoopResult1) = r.bj

struct ForceLoopResult2{V}
    fi::V
    fj::V
end

get_fi(r::Union{ForceLoopResult1, ForceLoopResult2}) = r.fi
get_fj(r::Union{ForceLoopResult1, ForceLoopResult2}) = r.fj

function gb_force_loop_1(coord_i, coord_j, i, j, charge_i, charge_j, Bi, Bj, cutoff,
                            factor_solute, factor_solvent, kappa, box_size)
    if j < i
        zero_force = zero(factor_solute ./ coord_i .^ 2)
        return ForceLoopResult1(zero_force[1], zero_force[1], zero_force, zero_force)
    end
    dr = vector(coord_i, coord_j, box_size)
    r2 = sum(abs2, dr)
    if !iszero(cutoff) && r2 > cutoff^2
        zero_force = zero(factor_solute ./ coord_i .^ 2)
        return ForceLoopResult1(zero_force[1], zero_force[1], zero_force, zero_force)
    end
    alpha2_ij = Bi * Bj
    D = r2 / (4 * alpha2_ij)
    exp_term = exp(-D)
    denominator2 = r2 + alpha2_ij * exp_term
    denominator = sqrt(denominator2)
    if iszero(kappa)
        pre_factor = factor_solute + factor_solvent
    else
        pre_factor = factor_solute + exp(-kappa * denominator) * factor_solvent +
                        kappa * denominator * exp(-kappa * denominator) * factor_solvent
    end
    Gpol = (pre_factor * charge_i * charge_j) / denominator
    dGpol_dr = -Gpol * (1 - exp_term/4) / denominator2
    dGpol_dalpha2_ij = -Gpol * exp_term * (1 + D) / (2 * denominator2)
    change_born_force_i = dGpol_dalpha2_ij * Bj
    if i != j
        change_born_force_j = dGpol_dalpha2_ij * Bi
        fdr = dr * dGpol_dr
        change_fs_i =  fdr
        change_fs_j = -fdr
        return ForceLoopResult1(change_born_force_i, change_born_force_j,
                                change_fs_i, change_fs_j)
    else
        zero_force = zero(factor_solute ./ coord_i .^ 2)
        return ForceLoopResult1(change_born_force_i, zero_force[1], zero_force, zero_force)
    end
end

function gb_force_loop_2(coord_i, coord_j, bi, ig, ori, srj, cutoff, box_size)
    dr = vector(coord_i, coord_j, box_size)
    r = norm(dr)
    if iszero(r) || (!iszero(cutoff) && r > cutoff)
        zero_force = zero(bi ./ coord_i .^ 2)
        return ForceLoopResult2(zero_force, zero_force)
    end
    rsrj = r + srj
    if ori < rsrj
        D = abs(r - srj)
        L = inv(max(ori, D))
        U = inv(rsrj)
        rinv = inv(r)
        r2inv = rinv^2
        t3 = (1 + (srj^2)*r2inv)*(L^2 - U^2)/8 + log(U/L)*r2inv/4
        de = bi * (t3 - ig) * rinv
        fdr = dr * de
        return ForceLoopResult2(-fdr, fdr)
    else
        zero_force = zero(bi ./ coord_i .^ 2)
        return ForceLoopResult2(zero_force, zero_force)
    end
end

function forces(inter::AbstractGBSA, sys, neighbors=nothing)
    coords, atoms, box_size = sys.coords, sys.atoms, sys.box_size
    Bs, B_grads, I_grads = born_radii_and_grad(inter, coords, box_size)

    if inter.use_ACE
        radii = inter.offset_radii .+ inter.offset
        sa_terms = inter.sa_factor .* (radii .+ inter.probe_radius) .^ 2 .* (radii ./ Bs) .^ 6
        born_forces = (-6 .* sa_terms ./ Bs) .* (Bs .> zero(inter.offset))
    else
        born_forces = zeros(typeof(inter.sa_factor * inter.cutoff), length(sys))
    end

    coords_i = @view coords[inter.is]
    coords_j = @view coords[inter.js]
    charges = charge.(atoms)
    charges_i = @view charges[inter.is]
    charges_j = @view charges[inter.js]
    Bsi = @view Bs[inter.is]
    Bsj = @view Bs[inter.js]
    loop_res_1 = gb_force_loop_1.(coords_i, coords_j, inter.is, inter.js, charges_i, charges_j,
                                    Bsi, Bsj, inter.cutoff, inter.factor_solute,
                                    inter.factor_solvent, inter.kappa, (box_size,))
    born_forces = born_forces .+ dropdims(sum(get_bi.(loop_res_1); dims=2); dims=2)
    born_forces = born_forces .+ dropdims(sum(get_bj.(loop_res_1); dims=1); dims=1)
    fs = dropdims(sum(get_fi.(loop_res_1); dims=2); dims=2) .+ dropdims(sum(get_fj.(loop_res_1); dims=1); dims=1)

    born_forces_2 = born_forces .* (Bs .^ 2) .* B_grads

    bis = @view born_forces_2[inter.is]
    oris = @view inter.offset_radii[inter.is]
    srjs = @view inter.scaled_offset_radii[inter.js]
    loop_res_2 = gb_force_loop_2.(coords_i, coords_j, bis, I_grads, oris, srjs,
                                    inter.cutoff, (box_size,))

    return fs .+ dropdims(sum(get_fi.(loop_res_2); dims=2); dims=2) .+ dropdims(sum(get_fj.(loop_res_2); dims=1); dims=1)
end

function gb_energy_loop(coord_i, coord_j, i, j, charge_i, charge_j, Bi, Bj, ori, cutoff,
                        factor_solute, factor_solvent, kappa, offset, probe_radius,
                        sa_factor, use_ACE, box_size)
    if i == j
        if iszero(kappa)
            pre_factor = factor_solute + factor_solvent
        else
            pre_factor = factor_solute + exp(-kappa * Bi) * factor_solvent
        end
        E = pre_factor * (charge_i^2) / (2*Bi)
        if use_ACE && (Bi > zero(offset))
            radius_i = ori + offset
            E += sa_factor * (radius_i + probe_radius)^2 * (radius_i / Bi)^6
        end
        return E
    elseif j > i
        r2 = sum(abs2, vector(coord_i, coord_j, box_size))
        if !iszero(cutoff) && r2 > cutoff^2
            return zero(factor_solute / offset)
        end
        f = sqrt(r2 + Bi*Bj*exp(-r2/(4*Bi*Bj)))
        if iszero(cutoff)
            f_cutoff = 1/f
        else
            f_cutoff = (1/f - 1/cutoff)
        end
        if iszero(kappa)
            pre_factor = factor_solute + factor_solvent
        else
            pre_factor = factor_solute + exp(-kappa * f) * factor_solvent
        end
        return pre_factor * charge_i * charge_j * f_cutoff
    else
        return zero(factor_solute / offset)
    end
end

function potential_energy(inter::AbstractGBSA, sys, neighbors=nothing)
    coords, atoms, box_size = sys.coords, sys.atoms, sys.box_size
    Bs, B_grads, I_grads = born_radii_and_grad(inter, coords, box_size)

    coords_i = @view coords[inter.is]
    coords_j = @view coords[inter.js]
    charges = charge.(atoms)
    charges_i = @view charges[inter.is]
    charges_j = @view charges[inter.js]
    Bsi = @view Bs[inter.is]
    Bsj = @view Bs[inter.js]
    oris = @view inter.offset_radii[inter.is]
    return sum(gb_energy_loop.(coords_i, coords_j, inter.is, inter.js, charges_i, charges_j,
                                Bsi, Bsj, oris, inter.cutoff, inter.factor_solute,
                                inter.factor_solvent, inter.kappa, inter.offset,
                                inter.probe_radius, inter.sa_factor, inter.use_ACE,
                                (box_size,)))
end


#=== elastic ===#
#=== write MSB ===#
function wMSB(
    model::Union{elmod2d},
    sou::Array{MTsource},
    sourcenumber::Int64,
    ParbinPath::Union{String, Void})

    nwf = model.nwf
    pkf = model.medium.pkf
    iflag = model.medium.iflag
    medium = model.medium
    ext = model.medium.ext

    if ParbinPath != nothing
        if ParbinPath[end-2:end] != "bin"
            error("Check output file extensions. It has to be in 'bin' format.")
        else
            fid = open(ParbinPath,"a+")
            write(fid,iflag)
            write(fid,ext)
            write(fid,nwf.BDntpp)
            write(fid,medium.pvel)
            write(fid,medium.svel)
            write(fid,medium.rho)
            write(fid,Float64(medium.dx))
            write(fid,Float64(medium.dz))
            write(fid,Float64(medium.dt))
            write(fid,medium.nT)
            write(fid,Float64(pkf))

            write(fid,sourcenumber)
            for isn = 1:sourcenumber
                write(fid,sou[i].nloc)
            end
            for isn = 1:sourcenumber
                write(fid,sou[i].waveform)
            end
            for isn = 1:sourcenumber
                write(fid,sou[i].not)
            end
            for isn = 1:sourcenumber
                write(fid,sou[i].mt)
            end
            close(fid)
        end
    end
end

#=== acoustic ===#
#=== write MSB ===#
function wMSB(
    model::Union{acmod2d},
    sou::Array{source},
    sourcenumber::Int64,
    ParbinPath::Union{String, Void})

    nwf = model.nwf
    pkf = model.medium.pkf
    iflag = model.medium.iflag
    medium = model.medium
    ext = model.medium.ext

    if ParbinPath != nothing
        if ParbinPath[end-2:end] != "bin"
            error("Check output file extensions. It has to be in 'bin' format.")
        else
            fid = open(ParbinPath,"a+")
            write(fid,iflag)
            write(fid,ext)
            write(fid,nwf.BDntpp)
            write(fid,medium.pvel)
            write(fid,medium.rho)
            write(fid,Float64(medium.dx))
            write(fid,Float64(medium.dz))
            write(fid,Float64(medium.dt))
            write(fid,medium.nT)
            write(fid,Float64(pkf))

            write(fid,sourcenumber)
            for i = 1:sourcenumber
                write(fid,sou[i].nloc)
            end
            for i = 1:sourcenumber
                write(fid,sou[i].waveform)
            end
            for i = 1:sourcenumber
                write(fid,sou[i].not)
            end
            for i = 1:sourcenumber
                if sou[i].tp == "expl"
                    write(fid,1)
                elseif sou[i].tp == "sfx"
                    write(fid,2)
                elseif sou[i].tp == "sfz"
                    write(fid,4)
                elseif sou[i].tp == "scx"
                    write(fid,5)
                elseif sou[i].tp == "scz"
                    write(fid,7)
                end
            end
            close(fid)
        end
    end
end


# %% imports

using PyCall
gym = pyimport("gym")

include("../../Modules/Gridcells/StateSpikeGaussGrid.jl")
include("../../Modules/Networks/net_arch.jl")
include("../../Modules/Gridcells/GridGenerator.jl")

# %% sim functions

function count_spikes(arch::Array{Int64}, spiked::Array{Int64}, layer::Int64)
    count_layer = zeros(arch[layer])
    interval = [sum(arch[1:layer-1])+1, sum(arch[1:layer])]
    spikes_in = filter(x-> interval[1] <= x <= interval[2], spiked)
    spikes_in .-= sum(arch[1:layer-1])
    count_layer[spikes_in] .+= 1
    return count_layer
end

function left_right(count_spikes::Array{Float64}, n_left::Int64, n_right::Int64)
    sl = sum(count_spikes[1:n_left])
    sr = sum(count_spikes[n_left+1:end])
    (sl > sr) ? action = 0 : action = 1
    return action, sl, sr
end

# %% fitness function play 100 episodes and check les contraintes a chaque épisode

function play_episode(net::Network, n_steps::Int64 = 200)
    env = gym.make("CartPole-v1")
    obs = env.reset()
    win_size = 10 # ms
    grid_cells = grid_gen()
    spike_hist = Array{Int64}[]
    score_hist = [[],[]]
    decision  = 0
    rng_action = Random.MersenneTwister(0)
    @inbounds for j in 1:n_steps
        count_tot = zeros(net.arch[end])
        spikes = input_spikes(obs, grid_cells)
        influence = 2
        @inbounds for msec in 1:win_size
            (msec==10) ? train=true : train=false
            spiked = step!(net, spikes, train)
            if msec >= length(net.arch)-1
                count_tot .+= influence .* count_spikes(net.arch, spiked, length(net.arch))
                influence *= 0.95
                push!(spike_hist,spiked)
            end
        end
        n_left = div(net.arch[end],2)
        n_right = net.arch[end] - n_left
        action, sl, sr = left_right(count_tot, n_left, n_right)
        if sl == sr
            action = rand(rng_action, 0:1)
            decision += 0.001
        else
            decision += 0.01
        end
        push!(score_hist[1], sl)
        push!(score_hist[2], sr)
        obs_new, _, _, _ = env.step(action)
        obs = obs_new
    end
    env.close()
    env = nothing
    Base.GC.gc()
    return score_hist, spike_hist, decision
end

# CONSTRAINTS
function constraints(score::Array{Array{Any,1},1}, spikes::Array{Array{Int64}}, net::Network)
    mw_min = 0.0 #FILL
    mw_max = 1.0#FILL
    s_min = 1000#FILL
    s_max = 15000#FILL
    m_min = -5.0#FILL
    m_max = 5.0 #FILL
    # constraint #1 weight average
    weights = net.s.e
    for layer in 1:length(net.arch)-1
        mw = mean(weights[layer])
        if (mw < mw_min) | (mw > mw_max)
            println("-----> weights mean out of bounds")
            return true
        end
    end
    # constraint #2 layer activity
    for layer in 1:length(net.arch)
        count_layer = zeros(net.arch[layer])
        for i in 1:length(spikes)
            count_layer .+= count_spikes(net.arch, spikes[i], layer)
        end
        spikes_layer = sum(count_layer)
        if (spikes_layer < s_min) | (spikes_layer > s_max)
            println("-----> spike activity out of bounds")
            return true
        end
    end
    # constraint #3 mean left right
    mlr = mean(score[1] .- score[2])
    if (mlr > m_max) | (mlr < m_min)
        println("-----> left right mean out of bounds")
        return true
    end
    return false
end

# plays 100 episodes and returns the fitness ( and tells wether to drop batch )
function fitness(episodes::Int64, net::Network)
    #score_hist = [[0.,0.],[0.,0.]]
    drop = false
    fit = 0
    for episode in 1:episodes
        score, spikes, decision = play_episode(net)
        drop = constraints(score, spikes, net)
        #push!(score_hist[1],score[1]...)
        #push!(score_hist[2],score[2]...)
        fit += decision
        if drop
            break
        end

    end
    #fit = std(score_hist[1] .- score_hist[2]) + pond
    return fit
end

# %% objective function ( fitness + contraintes ) ( create new network and play 100 episodes with fitness function )

function rescale(genes::Array{Float64}, tunekeys::Array{String})
    cfg = YAML.load(open("/Users/titou/Documents/PFE/RL-STDP/Julia/CartPole/CMAES/cfg_cmaes.yml"))
    N = length(genes)
    scaled_genes = zeros(N)
    for idx in 1:N
        down = cfg[tunekeys[idx]*"_d"]
        up = cfg[tunekeys[idx]*"_u"]
        scaled_genes[idx] = ((genes[idx]+1)*0.5)*(up-down) + down
    end
    return scaled_genes
end

function descale(params::Array{Float64}, tunekeys::Array{String})
    cfg = YAML.load(open("/Users/titou/Documents/PFE/RL-STDP/Julia/CartPole/CMAES/cfg_cmaes.yml"))
    N = length(tunekeys)
    descaled_genes = zeros(N)
    for idx in 1:N
        down = cfg[tunekeys[idx]*"_d"]
        up = cfg[tunekeys[idx]*"_u"]
        descaled_genes[idx] = 2*(params[idx]-down)/(up-down) - 1
    end
    return descaled_genes
end

function generate(genes::Array{Float64})
    tunekeys = ["vthresh", "ae", "b", "c", "de","imin", "imax",
                "m", "std", "apost", "apre"]
    params = Dict(tunekeys .=> rescale(genes,tunekeys))
    keys = ["tda", "wmax" ,"theta", "ttheta", "ai", "di", "wei", "wie"]
    values = [0.995, 1, 0.0, 0.998, 0.1, 2.0, 0.0, 0.0]
    for idx in 1:length(keys)
        params[keys[idx]] = values[idx]
    end
    return params
end

function objective(genes::Array{Float64})
    params = generate(genes)
    n = Network([120,60,80],params)
    fit = fitness(50,n)
    @show fit
    -fit
end

# %% tests
#
# objective([-50.0,0.02,0.2,-65.0,8.0,0.5,0.1,0.0,0.998,1.0,1.5])


import Statistics
using LinearAlgebra
using Cassette
using Cassette: Tagged, tag, untag, istagged, metadata, hasmetadata,
    enabletagging, @overdub, overdub, canrecurse, similarcontext, fallback
using JuliaInterpreter
using Espresso
using CUDAapi

include("utils.jl")
include("helpers.jl")
include("devices.jl")
include("tape.jl")
include("tapeutils.jl")
include("trace/trace.jl")
include("diffrules/diffrules.jl")
include("grad.jl")
include("compile.jl")
include("update.jl")
include("transform.jl")

# function __init__()
#     @require CuArrays="3a865a2d-5b23-5a0f-bc46-62713ec82fae" include("cuda.jl")
# end

const BEST_AVAILABLE_DEVICE = Ref{AbstractDevice}(CPU())

if has_cuda()
    try
        using CuArrays
        using CUDAnative

        BEST_AVAILABLE_DEVICE[] = GPU(0)

        include("cuda.jl")
    catch ex
        # something is wrong with the user's set-up (or there's a bug in CuArrays)
        @warn "CUDA is installed, but CuArrays.jl fails to load" exception=(ex,catch_backtrace())

    end
end

best_available_device() = BEST_AVAILABLE_DEVICE[]


function read_nodal_segy(fname::String,
  nn::String,
  s::TimeSpec,
  t::TimeSpec,
  chans::ChanSpec,
  memmap::Bool)

  # Preprocessing
  (d0, d1) = parsetimewin(s, t)
  t0 = DateTime(d0).instant.periods.value*1000 - dtconst
  f = memmap ? IOBuffer(Mmap.mmap(fname)) : open(fname, "r")
  trace_fh = Array{Int16, 1}(undef, 3)
  shorts  = getfield(BUF, :int16_buf)
  fhd = Dict{String,Any}()
  # ww = Array{String, 1}(undef, 0)

  # ==========================================================================
  # Read file header
  fhd["filehdr"]  = fastread(f, 3200)
  fhd["jobid"]    = bswap(fastread(f, Int32))
  fhd["lineid"]   = bswap(fastread(f, Int32))
  fhd["reelid"]   = bswap(fastread(f, Int32))
  fast_readbytes!(f, BUF.buf, 48)
  fillx_i16_be!(shorts, BUF.buf, 24, 0)
  fastskip(f, 240)
  for i in 25:27
    shorts[i] = read(f, Int16)
  end
  fastskip(f, 94)

  # check endianness; can be inconsistent; 0x0400 is a kludge
  # (stands for SEG Y rev 10.24, current is rev 2.0)
  # if (unsigned(shorts[25]) > 0x0400) || (shorts[26] < zero(Int16)) || (shorts[26] > one(Int16)) || (shorts[27] < zero(Int16))
  #   shorts[25:27] .= bswap.(shorts[25:27])
  #   push!(ww, "Inconsistent file header endianness")
  # end
  # fastskip(f, 94)

  # Process
  nh = max(zero(Int16), getindex(shorts, 27))
  fhd["exthdr"] = Array{String,1}(undef, nh)
  [fhd["exthdr"][i] = fastread(f, 3200) for i in 1:nh]
  for (j,i) in enumerate(String[ "ntr", "naux", "filedt", "origdt", "filenx",
                                 "orignx", "fmt", "cdpfold", "trasort", "vsum",
                                 "swst", "swen0", "swlen", "swtyp", "tapnum",
                                 "swtapst", "swtapen", "taptyp", "corrtra", "bgainrec",
                                 "amprec", "msys", "zupdn", "vibpol", "segyver",
                                 "isfixed", "n_exthdr" ])
    fhd[i] = shorts[j]
  end

  nt = getindex(shorts, 1)
  trace_fh[1] = getindex(shorts,3)
  trace_fh[2] = getindex(shorts,5)
  trace_fh[3] = getindex(shorts,7)

  # ==========================================================================
  # initialize NodalData container; set variables
  if isempty(chans)
    chans = 1:nt
  end
  fs = sμ / Float64(trace_fh[1])
  data = Array{Float32, 2}(undef, trace_fh[2], nt)
  S = NodalData(data, fhd, chans, t0)
  net = nn * "."
  cha = string("..O", getbandcode(fs), "0")

  # ==========================================================================
  # Read traces
  j = 0

  for i = 1:nt
    C = do_trace(f, false, true, 0x00, true, trace_fh)
    if i in chans
      j += 1
      S.id[j]       = string(net, lpad(i, 5, '0'), cha)
      S.name[j]     = string(C.misc["rec_no"], "_", i)
      S.fs[j]       = C.fs
      S.gain[j]     = C.gain
      S.misc[j]     = C.misc
      S.t[j]        = C.t
      S.data[:, j] .= C.x
    end
  end

  # TO DO: actually use s, t here

  # ==========================================================================
  # Cleanup
  close(f)
  resize!(BUF.buf, 65535)
  fill!(S.fs, fs)
  fill!(S.src, realpath(fname))
  fill!(S.units, "m/m")

  # Output warnings to STDOUT
  # if !isempty(ww)
  #   for i = 1:length(ww)
  #     @warn(ww[i])
  #   end
  # end
  # S.info["Warnings"] = ww
  return S
end


module VsRecorderBase

using Reexport: @reexport

import Dates, TOML

import DataStructures
import YAML
@reexport using DataStructures: Queue, enqueue!, dequeue!,
    OrderedDict
using FileIO: Stream, @format_str

@reexport using LinearAlgebra, Statistics
@reexport using VideoIO, Images

import StringDistances
using Tesseract, ImageFiltering
@reexport using MacroTools: @forward

export Missable, Nullable, @missable, @nullable,
    Yes,
    SimpleTypeWrapper, @type_wrapper,
    ∞, ±
include("utils/misc.jl")

export AbstractVsStrategy, AbstractVsSource, AbstractVsScene, AbstractVsStream,
    VsStream, VsFrame,
    VsContextData,
    VsConfig,
    VsContext
export image, time
export vs_setup, vs_init!, vs_parse_frame!, vs_update!, vs_result, vs_tryparse_scene
include("types.jl")

export BoundedBinaryHeap
include("utils/data_structures.jl")

export find_closest_n, find_closest,
    data_search_n, data_search
include("utils/data_search.jl")

export remove_spaces, parse_int
include("utils/string_utils.jl")

export to_gray_image, blur,
    color_distance, floodfill, floodfill!,
    draw_outline, draw_outline!,
    cycled_translate,
    blend_color,
    bounding_rect, shrink,
    prepare_text_for_ocr
include("utils/image_transform.jl")

export Rect, subimage, topleft, is_gray
include("utils/image_data.jl")

export image_distance, block,
    template_match_all, template_match,
    table_search
export SpriteSheet
include("utils/template_match.jl")

export open_video, open_camera,
    read_frame, load_image,
    enum_to_string, enum_from_string,
    serialize, deserialize
include("io.jl")
include("yaml.jl")

export create_ocr_instance,
    download_ocr_language,
    default_ocr_instance,
    ocr, ocr_tsv,
    is_cjk,
    init_multiple_ocr!, ocr_multiple_lang
include("ocr.jl")

export DefaultStrategy
include("strategies/default.jl")
using .DefaultStrategyModule: DefaultStrategy

include("main.jl")
include("api.jl")

function __init__()
    mkpath(OCR_DATA_DIR)
    nothing
end

end