module Inference

using ..Core, ..Utilities
using DynamicPPL: Metadata, _tail, VarInfo, TypedVarInfo,
    islinked, invlink!, getlogp, tonamedtuple, VarName, getsym, vectorize,
    settrans!, _getvns, getdist, split_var_str, CACHERESET, AbstractSampler,
    Model, runmodel!, Sampler, SampleFromPrior, SampleFromUniform,
    Selector, AbstractSamplerState, DefaultContext, PriorContext,
    LikelihoodContext, MiniBatchContext, set_flag!, unset_flag!, NamedDist, NoDist
using Distributions, Libtask, Bijectors
using DistributionsAD: VectorOfMultivariate
using LinearAlgebra
using ..Turing: PROGRESS, Turing
using StatsFuns: logsumexp
using Random: GLOBAL_RNG, AbstractRNG, randexp
using DynamicPPL
using AbstractMCMC: AbstractModel, AbstractSampler
using Bijectors: _debug
using MCMCChains: Chains

import AbstractMCMC
import AdvancedHMC; const AHMC = AdvancedHMC
import AdvancedMH; const AMH = AdvancedMH
import ..Core: getchunksize, getADbackend
import DynamicPPL: tilde, dot_tilde, getspace, get_matching_type,
    VarName, _getranges, _getindex, getval, _getvns
import EllipticalSliceSampling
import Random

export  InferenceAlgorithm,
        Hamiltonian,
        GibbsComponent,
        StaticHamiltonian,
        AdaptiveHamiltonian,
        SampleFromUniform,
        SampleFromPrior,
        MH,
        ESS,
        Gibbs,      # classic sampling
        HMC,
        SGLD,
        SGHMC,
        HMCDA,
        NUTS,       # Hamiltonian-like sampling
        DynamicNUTS,
        IS,
        SMC,
        CSMC,
        PG,
        assume,
        dot_assume,
        observe,
        dot_observe,
        resume

#######################
# Sampler abstraction #
#######################
abstract type AbstractAdapter end
abstract type InferenceAlgorithm end
abstract type ParticleInference <: InferenceAlgorithm end
abstract type Hamiltonian{AD} <: InferenceAlgorithm end
abstract type StaticHamiltonian{AD} <: Hamiltonian{AD} end
abstract type AdaptiveHamiltonian{AD} <: Hamiltonian{AD} end

getchunksize(::Type{<:Hamiltonian{AD}}) where AD = getchunksize(AD)
getADbackend(::Hamiltonian{AD}) where AD = AD()

"""
    mh_accept(logp_current::Real, logp_proposal::Real, log_proposal_ratio::Real)

Decide if a proposal ``x'`` with log probability ``\\log p(x') = logp_proposal`` and
log proposal ratio ``\\log k(x', x) - \\log k(x, x') = log_proposal_ratio`` in a
Metropolis-Hastings algorithm with Markov kernel ``k(x_t, x_{t+1})`` and current state
``x`` with log probability ``\\log p(x) = logp_current`` is accepted by evaluating the
Metropolis-Hastings acceptance criterion
```math
\\log U \\leq \\log p(x') - \\log p(x) + \\log k(x', x) - \\log k(x, x')
```
for a uniform random number ``U \\in [0, 1)``.
"""
function mh_accept(logp_current::Real, logp_proposal::Real, log_proposal_ratio::Real)
    # replacing log(rand()) with -randexp() yields test errors
    return log(rand()) + logp_current ≤ logp_proposal + log_proposal_ratio
end

######################
# Default Transition #
######################

struct Transition{T, F<:AbstractFloat}
    θ  :: T
    lp :: F
end

function Transition(spl::Sampler, nt::NamedTuple=NamedTuple())
    theta = merge(tonamedtuple(spl.state.vi), nt)
    lp = getlogp(spl.state.vi)
    return Transition{typeof(theta), typeof(lp)}(theta, lp)
end

function additional_parameters(::Type{<:Transition})
    return [:lp]
end

##########################################
# Internal variable names for MCMCChains #
##########################################

const TURING_INTERNAL_VARS = (internals = [
    "elapsed",
    "eval_num",
    "lf_eps",
    "lp",
    "weight",
    "le",
    "acceptance_rate",
    "hamiltonian_energy",
    "hamiltonian_energy_error",
    "max_hamiltonian_energy_error",
    "is_accept",
    "log_density",
    "n_steps",
    "numerical_error",
    "step_size",
    "nom_step_size",
    "tree_depth",
    "is_adapt",
],)

#########################################
# Default definitions for the interface #
#########################################

function AbstractMCMC.sample(
    model::AbstractModel,
    alg::InferenceAlgorithm,
    N::Integer;
    kwargs...
)
    return AbstractMCMC.sample(Random.GLOBAL_RNG, model, alg, N; kwargs...)
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    alg::InferenceAlgorithm,
    N::Integer;
    chain_type=Chains,
    resume_from=nothing,
    progress=PROGRESS[],
    kwargs...
)
    if resume_from === nothing
        return AbstractMCMC.sample(rng, model, Sampler(alg, model), N;
                                   chain_type=chain_type, progress=progress, kwargs...)
    else
        return resume(resume_from, N; chain_type=chain_type, progress=progress, kwargs...)
    end
end

function AbstractMCMC.psample(
    model::AbstractModel,
    alg::InferenceAlgorithm,
    N::Integer,
    n_chains::Integer;
    kwargs...
)
    return AbstractMCMC.psample(Random.GLOBAL_RNG, model, alg, N, n_chains; kwargs...)
end

function AbstractMCMC.psample(
    rng::AbstractRNG,
    model::AbstractModel,
    alg::InferenceAlgorithm,
    N::Integer,
    n_chains::Integer;
    chain_type=Chains,
    progress=PROGRESS[],
    kwargs...
)
    return AbstractMCMC.psample(rng, model, Sampler(alg, model), N, n_chains;
                                chain_type=chain_type, progress=progress, kwargs...)
end

function AbstractMCMC.sample_init!(
    ::AbstractRNG,
    model::Model,
    spl::Sampler,
    N::Integer;
    kwargs...
)
    # Resume the sampler.
    set_resume!(spl; kwargs...)

    # Set the parameters to a starting value.
    initialize_parameters!(spl; kwargs...)
end

function AbstractMCMC.sample_end!(
    ::AbstractRNG,
    ::Model,
    ::Sampler,
    ::Integer,
    ::Vector;
    kwargs...
)
    # Silence the default API function.
end

function initialize_parameters!(
    spl::Sampler;
    init_theta::Union{Nothing,Vector}=nothing,
    verbose::Bool=false,
    kwargs...
)
    # Get `init_theta`
    if init_theta !== nothing
        verbose && @info "Using passed-in initial variable values" init_theta
        # Convert individual numbers to length 1 vector; `ismissing(v)` is needed as `size(missing)` is undefined`
        init_theta = [ismissing(v) || size(v) == () ? [v] : v for v in init_theta]
        # Flatten `init_theta`
        init_theta_flat = foldl(vcat, map(vec, init_theta))
        # Create a mask to inidicate which values are not missing
        theta_mask = map(x -> !ismissing(x), init_theta_flat)
        # Get all values
        theta = spl.state.vi[spl]
        @assert length(theta) == length(init_theta_flat) "Provided initial value doesn't match the dimension of the model"
        # Update those which are provided (i.e. not missing)
        theta[theta_mask] .= init_theta_flat[theta_mask]
        # Update in `vi`
        spl.state.vi[spl] = theta
    end
end

##########################
# Chain making utilities #
##########################

function _params_to_array(ts::Vector, spl::Sampler)
    names_set = Set{String}()
    # Extract the parameter names and values from each transition.
    dicts = map(ts) do t
        nms, vs = flatten_namedtuple(t.θ)
        for nm in nms
            push!(names_set, nm)
        end
        # Convert the names and values to a single dictionary.
        return Dict(nms[j] => vs[j] for j in 1:length(vs))
    end
    names = collect(names_set)
    vals = [get(dicts[i], key, missing) for i in eachindex(dicts),
        (j, key) in enumerate(names)]

    return names, vals
end

function flatten_namedtuple(nt::NamedTuple)
    names_vals = mapreduce(vcat, keys(nt)) do k
        v = nt[k]
        if length(v) == 1
            return [(string(k), v)]
        else
            return mapreduce(vcat, zip(v[1], v[2])) do (vnval, vn)
                return collect(FlattenIterator(vn, vnval))
            end
        end
    end
    return [vn[1] for vn in names_vals], [vn[2] for vn in names_vals]
end

function get_transition_extras(ts::Vector)
    # Get the extra field names from the sampler state type.
    # This handles things like :lp or :weight.
    extra_params = additional_parameters(eltype(ts))

    # Get the values of the extra parameters.
    local extra_names
    all_vals = []

    # Iterate through each transition.
    for t in ts
        extra_names = String[]
        vals = []

        # Iterate through each of the additional field names
        # in the struct.
        for p in extra_params
            # Check whether the field contains a NamedTuple,
            # in which case we need to iterate through each
            # key/value pair.
            prop = getproperty(t, p)
            if prop isa NamedTuple
                for (k, v) in pairs(prop)
                    push!(extra_names, string(k))
                    push!(vals, v)
                end
            else
                push!(extra_names, string(p))
                push!(vals, prop)
            end
        end
        push!(all_vals, vals)
    end

    # Convert the vector-of-vectors to a matrix.
    valmat = [all_vals[i][j] for i in 1:length(ts), j in 1:length(all_vals[1])]

    return extra_names, valmat
end

# Default Chains constructor.
function AbstractMCMC.bundle_samples(
    rng::AbstractRNG,
    model::Model,
    spl::Sampler,
    N::Integer,
    ts::Vector,
    chain_type::Type{Chains};
    discard_adapt::Bool=true,
    save_state=false,
    kwargs...
)
    # Check if we have adaptation samples.
    if discard_adapt && :n_adapts in fieldnames(typeof(spl.alg))
        ts = ts[(spl.alg.n_adapts+1):end]
    end

    # Convert transitions to array format.
    # Also retrieve the variable names.
    nms, vals = _params_to_array(ts, spl)

    # Get the values of the extra parameters in each Transition struct.
    extra_params, extra_values = get_transition_extras(ts)

    # Extract names & construct param array.
    nms = [nms; extra_params]
    parray = hcat(vals, extra_values)

    # If the state field has average_logevidence or final_logevidence, grab that.
    le = missing
    if :average_logevidence in fieldnames(typeof(spl.state))
        le = getproperty(spl.state, :average_logevidence)
    elseif :final_logevidence in fieldnames(typeof(spl.state))
        le = getproperty(spl.state, :final_logevidence)
    end

    # Check whether to invlink! the varinfo
    if islinked(spl.state.vi, spl)
        invlink!(spl.state.vi, spl)
    end

    # Set up the info tuple.
    if save_state
        info = (range = rng, model = model, spl = spl, vi = spl.state.vi)
    else
        info = NamedTuple()
    end

    # Chain construction.
    return Chains(
        parray,
        string.(nms),
        deepcopy(TURING_INTERNAL_VARS);
        evidence=le,
        info=info,
        sorted=true
    )
end

function AbstractMCMC.bundle_samples(
    rng::AbstractRNG,
    model::Model,
    spl::Sampler,
    N::Integer,
    ts::Vector,
    chain_type::Type{Vector{NamedTuple}};
    discard_adapt::Bool=true,
    save_state=false,
    kwargs...
)
    nts = Vector{NamedTuple}(undef, N)

    for (i,t) in enumerate(ts)
        k = collect(keys(t.θ))
        vs = []
        for v in values(t.θ)
            push!(vs, v[1])
        end

        push!(k, :lp)

        nts[i] = NamedTuple{tuple(k...)}(tuple(vs..., t.lp))
    end

    return map(identity, nts)
end

function save(c::Chains, spl::Sampler, model, vi, samples)
    nt = NamedTuple{(:spl, :model, :vi, :samples)}((spl, model, deepcopy(vi), samples))
    return setinfo(c, merge(nt, c.info))
end

function resume(c::Chains, n_iter::Int; chain_type=Chains, progress=PROGRESS[], kwargs...)
    @assert !isempty(c.info) "[Turing] cannot resume from a chain without state info"

    # Sample a new chain.
    newchain = AbstractMCMC.sample(
        c.info[:range],
        c.info[:model],
        c.info[:spl],
        n_iter;
        resume_from=c,
        reuse_spl_n=n_iter,
        chain_type=Chains,
        progress=progress,
        kwargs...
    )

    # Stick the new samples at the end of the old chain.
    return vcat(c, newchain)
end

function set_resume!(
    s::Sampler;
    resume_from::Union{Chains, Nothing}=nothing,
    kwargs...
)
    # If we're resuming, grab the sampler info.
    if resume_from !== nothing
        s = resume_from.info[:spl]
    end
end

#########################
# Default sampler state #
#########################

"""
A blank `AbstractSamplerState` that contains only `VarInfo` information.
"""
mutable struct SamplerState{VIType<:VarInfo} <: AbstractSamplerState
    vi :: VIType
end

#######################################
# Concrete algorithm implementations. #
#######################################

include("ess.jl")
include("hmc.jl")
include("mh.jl")
include("is.jl")
include("AdvancedSMC.jl")
include("gibbs.jl")
include("../contrib/inference/sghmc.jl")

################
# Typing tools #
################

for alg in (:SMC, :PG, :MH, :IS, :ESS, :Gibbs)
    @eval getspace(::$alg{space}) where {space} = space
end
for alg in (:HMC, :HMCDA, :NUTS, :SGLD, :SGHMC)
    @eval getspace(::$alg{<:Any, space}) where {space} = space
end

floatof(::Type{T}) where {T <: Real} = typeof(one(T)/one(T))
floatof(::Type) = Real # fallback if type inference failed

function get_matching_type(
    spl::AbstractSampler,
    vi::VarInfo,
    ::Type{T},
) where {T}
    return T
end
function get_matching_type(
    spl::AbstractSampler,
    vi::VarInfo,
    ::Type{<:AbstractFloat},
)
    return floatof(eltype(vi, spl))
end
function get_matching_type(
    spl::Sampler{<:Hamiltonian},
    vi::VarInfo,
    ::Type{<:Union{Missing, AbstractFloat}},
)
    return Union{Missing, floatof(eltype(vi, spl))}
end
function get_matching_type(
    spl::Sampler{<:Hamiltonian},
    vi::VarInfo,
    ::Type{<:AbstractFloat},
)
    return floatof(eltype(vi, spl))
end
function get_matching_type(
    spl::Sampler{<:Hamiltonian},
    vi::VarInfo,
    ::Type{TV},
) where {T, N, TV <: Array{T, N}}
    return Array{get_matching_type(spl, vi, T), N}
end
function get_matching_type(
    spl::Sampler{<:Union{PG, SMC}},
    vi::VarInfo,
    ::Type{TV},
) where {T, N, TV <: Array{T, N}}
    return TArray{T, N}
end

##############
# Utilities  #
##############

getspace(spl::Sampler) = getspace(spl.alg)
function ambiguity_error_msg()
    return "Ambiguous `lhs .~ rhs` or `@. lhs ~ rhs` syntax. The broadcasting can either be
    column-wise following the convention of Distributions.jl or element-wise following
    Julia's general broadcasting semantics. Please make sure that the element type of `lhs`
    is not a supertype of the support type of `AbstractVector` to eliminate ambiguity."
end

end # module


module SARA
using JSON
using Base.Threads
using LinearAlgebra
using Interpolations
using ForwardDiff

using CovarianceFunctions
const Kernel = CovarianceFunctions
using GaussianDistributions
using GaussianDistributions: input_transformation

export innerloop, inner, uncertainty_sampling, evaluate, TemperatureProfile,
        next_stripe, get_relevant_T, iSARA, oSARA, xrd_to_global, stripe_to_global,
        stripe_uncertainty_sampling

include("util.jl")
include("estimate_noise.jl")
include("temperature_profile.jl")
include("acquisition.jl") # acquisition functions
include("inner.jl") # inner and outer loops for optical spectroscopy
include("stripe_to_global.jl")
include("outer.jl")

end # SARA


# The following was inspired by AtariAlgos.jl
# Modifications by Dennis Wilson @d9w

using Colors
using ImageCore
using ImageTransformations
export
    Game,
    close!,
    draw,
    get_inputs,
    get_rgb,
    get_float,
    get_int,
    get_bool,
    get_string

rom_directory(x::String) = joinpath(dirname(@__FILE__), "..", "deps", "rom_files", "$x.bin")

struct Game
    ale::ALEPtr
    width::Int
    height::Int
    actions::Array{Int32}
end

function Game(romfile::String, seed::Int64)
    ale = ALE_new()
    setInt(ale, "random_seed", Cint(seed))
    loadROM(ale, string(rom_directory(romfile)))
    w = getScreenWidth(ale)
    h = getScreenHeight(ale)
    actions = getMinimalActionSet(ale)
    Game(ale, w, h, actions)
end

function close!(game::Game)
    ALE_del(game.ale)
end

function draw(game::Game)
    rawscreen = Array{Cuchar}(undef, game.width * game.height * 3)
    getScreenRGB(game.ale, rawscreen)
    colorview(RGB, Float64.(reshape(rawscreen/256.,
                                    (3, game.width, game.height))))';
end

function get_inputs(game::Game)
    screen = getScreen(game.ale)/(0xff*1.0)
    screen = reshape(screen, (game.width, game.height))'
    # imresize(screen, (42, 32))/256.
    screen
end

function get_rgb(game::Game)
    rawscreen = Array{Cuchar}(undef, game.width * game.height * 3)
    getScreenRGB(game.ale, rawscreen)
    rgb = Float64.(reshape(rawscreen/256.,(3, game.width, game.height)));
    [rgb[i,:,:]' for i in 1:3]
end

# Read access to A.L.E. configuration:
get_float(game::Game, key::String) = getFloat(game.ale, key)
get_int(game::Game, key::String) = getInt(game.ale, key)
get_bool(game::Game, key::String) = getBool(game.ale, key)
get_string(game::Game, key::String) = getString(game.ale, key)


using Test
using ThreeBodyTB
using Suppressor

#basic loading

#name of test directory
TESTDIR=ThreeBodyTB.TESTDIR

@testset "basic crystal dftout tests" begin
    types=["Li"]
    pos=zeros((1,3))
    A=[ [10.0 0 0]; [0 10.0 0]; [ 0 0 10.0]]

    c=makecrys(A, pos, types, units="Bohr")

    @test c.A == A
    @test c.coords == pos
    @test c.types == types

    c2=makecrys("$TESTDIR/POSCAR_testing")
    @test c2.types == ["Cl", "N","N","Sr", "Sr"]

    c3=makecrys("$TESTDIR/qe_testing.in")
    @test c3.types == ["Si", "C"]

    forces=ones((1,3))
    energy=0.0
    energy_smear = 0.0
    d=makedftout(c, energy, energy_smear, forces, zeros(3,3))
    @test energy == d.energy
    @test forces == d.forces

    d2=makedftout(A,pos,types, energy, energy_smear, forces, zeros(3,3))
    @test d.forces == d2.forces

    #supposed to throw an error
    @suppress begin
        forces2=zeros((2,3))
        @test_throws ErrorException("Error forces") makedftout(c, energy,energy_smear, forces2, zeros(3,3))
    end

    dft_out = ThreeBodyTB.QE.loadXML("$TESTDIR/dimer.save/")

    convert_ryd_ha = 2.0

    @test convert_ryd_ha*-1.428445472333963e1 ≈ dft_out.energy
    @test convert_ryd_ha*-6.425031283667987e-2 ≈ dft_out.bandstruct.efermi
    @test convert_ryd_ha*-1.832485909355843e0 ≈ dft_out.bandstruct.eigs[1,1]

end

@testset "advanced crystal tests" begin
    @suppress begin
        types=["Li"]
        pos=zeros((1,3))
        A=[ [10.0 0 0]; [0 10.0 0]; [ 0 0 10.0]]

        c=makecrys(A, pos, types, units="Bohr")

#        println(c)
        c2 = c * [2,2,2]

        ThreeBodyTB.CrystalMod.generate_random_distortion(c2,0.01,0.01)
        ThreeBodyTB.CrystalMod.write_poscar(c2,"$TESTDIR/POSCAR_tmp")
        rm("$TESTDIR/POSCAR_tmp")

        ThreeBodyTB.CrystalMod.write_efs(c, 0.0, zeros(1,3), zeros(3,3), "$TESTDIR/qe.out")
        rm("$TESTDIR/qe.out")

        @test 1 == 1

    end
end

if false

    #run QE . You actually need quantum espresso working to run this test, not appropriate for CI.
    @testset "run QE" begin
        types=["Li"]
        pos=zeros((1,3))
        A=[ [6.0 0 0]; [0 6.0 0]; [ 0 0 6.0]]

        c=makecrys(A, pos, types, units="Bohr")

        d = ThreeBodyTB.DFT.runSCF(c, nprocs=4)

        @test -14.35143737 ≈ d.energy

    end
end


############################
# Types & Constructors
############################
mutable struct FixedLengthArray{T<:Real} <: Solution
    sol::Array{T,1};
    status::Status;
end

"""
    constructSolution(solutionType, args)

Constructs a container of type `solutionType` for the solution. Note that the arguments of the constructor are provided in `args`.

# Example
```jldoctest
julia> typeVariables = Int; numVariables = 5;  numConstraints = 3;
julia> status = OREnvironment.constructStatus(numConstraints);
julia> args = (typeVariables, numVariables, status);
julia> s = OREnvironment.constructSolution(:FixedLengthArray, args)
```
"""
function constructSolution(solutionType::Symbol, args)
    return constructSolution(Val(solutionType), args...);
end

function constructSolution(::Val{:FixedLengthArray},
                           typeVariables::DataType,
                           numVariables::Int,
                           status::Status)
    return FixedLengthArray(zeros(typeVariables, numVariables), status);
end

"""
    create_empty_solution(problem, T)

Constructs an empty solution of type `:FixedLengthArray` with variables of type `T` with the information contained in `problem` (num of constraints, of variables).
"""
function create_empty_solution(problem::Problem, typeVariables::DataType)
	local numConstraints::Int = get_number_of_constraints(problem);
	local numVars::Int 		  = get_number_of_variables(problem);
	local status::Status	  = constructStatus(numConstraints);
	return constructSolution(:FixedLengthArray, (typeVariables, numVars, status));
end

############################
# General methods for Solution
############################
"""
    is_feasible(s)

Returns the feasibility value of solution `s`.
# Example
```jldoctest
julia> typeVariables = Int; numVariables = 5;  numConstraints = 3;
julia> status = OREnvironment.constructStatus(numConstraints);
julia> args = (typeVariables, numVariables, status);
julia> s = OREnvironment.constructSolution(:FixedLengthArray, args)
julia> OREnvironment.is_feasible(s)
false
```
"""
is_feasible(s::Solution)::Bool = is_feasible(s.status);

"""
    set_feasible!(s, val)

Sets the feasibility value of solution `s` to `val`.

# Example
```jldoctest
julia> typeVariables = Int; numVariables = 5;  numConstraints = 3;
julia> status = OREnvironment.constructStatus(numConstraints);
julia> args = (typeVariables, numVariables, status);
julia> s = OREnvironment.constructSolution(:FixedLengthArray, args)
julia> OREnvironment.set_feasible!(s, true)
julia> OREnvironment.is_feasible(s)
true
```
"""
set_feasible!(s::Solution, value::Bool) = set_feasible!(s.status, value);

"""
    is_optimal(s)

Returns the optimality value of solution `s`.
# Example
```jldoctest
julia> typeVariables = Int; numVariables = 5;  numConstraints = 3;
julia> status = OREnvironment.constructStatus(numConstraints);
julia> args = (typeVariables, numVariables, status);
julia> s = OREnvironment.constructSolution(:FixedLengthArray, args)
julia> OREnvironment.is_optimal(s)
false
```
"""
is_optimal(s::Solution)::Bool = is_optimal(s.status);

"""
    set_optimal!(s, val)

Sets the optimality value of solution `s` to `val`.

# Example
```jldoctest
julia> typeVariables = Int; numVariables = 5;  numConstraints = 3;
julia> status = OREnvironment.constructStatus(numConstraints);
julia> args = (typeVariables, numVariables, status);
julia> s = OREnvironment.constructSolution(:FixedLengthArray, args)
julia> OREnvironment.set_optimal!(s, true)
julia> OREnvironment.is_optimal(s)
true
```
"""
set_optimal!(s::Solution, value::Bool) = set_optimal!(s.status, value);

"""
    get_objfunction(s)

Returns the objective function value of solution `s`.
# Example
```jldoctest
julia> typeVariables = Int; numVariables = 5;  numConstraints = 3;
julia> status = OREnvironment.constructStatus(numConstraints);
julia> args = (typeVariables, numVariables, status);
julia> s = OREnvironment.constructSolution(:FixedLengthArray, args)
julia> OREnvironment.get_objfunction(s)
0.0
```
"""
get_objfunction(s::Solution)::Float64 = get_objfunction(s.status);

"""
    set_objfunction!(s, val)

Sets the objective function value of solution `s` to `val`.

# Example
```jldoctest
julia> typeVariables = Int; numVariables = 5;  numConstraints = 3;
julia> status = OREnvironment.constructStatus(numConstraints);
julia> args = (typeVariables, numVariables, status);
julia> s = OREnvironment.constructSolution(:FixedLengthArray, args)
julia> OREnvironment.set_objfunction!(s, 15.66)
julia> OREnvironment.get_objfunction(s)
15.66
```
"""
set_objfunction!(s::Solution, value::Float64) = set_objfunction!(s.status, value);

"""
    get_constraint_consumption(s, idx)

Returns the consumption of the `idx` constraint in solution `s`.

**Note:** It returns the stored value, This method does not compute the consumption based on the current solution.

# Example
```jldoctest
julia> typeVariables = Int; numVariables = 5;  numConstraints = 3;
julia> status = OREnvironment.constructStatus(numConstraints);
julia> args = (typeVariables, numVariables, status);
julia> s = OREnvironment.constructSolution(:FixedLengthArray, args)
julia> OREnvironment.get_constraint_consumption(s, 2)
0.0
```
"""
@inline function get_constraint_consumption(s::Solution, idxConstraint::Int)::Float64
    return get_constraint_consumption(s.status, idxConstraint);
end

"""
    set_constraint_consumption!(s, idx, val)

Sets the consumption of the `idx` constraint in solution `s` to value `val`.

# Example
```jldoctest
julia> typeVariables = Int; numVariables = 5;  numConstraints = 3;
julia> status = OREnvironment.constructStatus(numConstraints);
julia> args = (typeVariables, numVariables, status);
julia> s = OREnvironment.constructSolution(:FixedLengthArray, args)
julia> OREnvironment.set_constraint_consumption(s, 2, 3.4);
julia> OREnvironment.get_constraint_consumption(s, 2)
3.4
```
"""
@inline function set_constraint_consumption!(s::Solution, idxConstraint::Int, value::Float64)
    set_constraint_consumption!(s.status, idxConstraint, value);
end

"""
    is_first_solution_better(s1, s2, objSense, feasibilityRequired)

Returns if solution `s1` is better than solution `s2`. For the comparaison is required the objective function sense and if feasibility is a must in the comparaison.

# Example
```jldoctest
julia> typeVariables = Int; numVariables = 5;  numConstraints = 3;
julia> status = OREnvironment.constructStatus(numConstraints);
julia> args = (typeVariables, numVariables, status);
julia> s1 = OREnvironment.constructSolution(:FixedLengthArray, args)
julia> OREnvironment.set_feasible!(s1, true);
julia> OREnvironment.set_objfunction!(s1, 19.8);
julia> s2 = OREnvironment.constructSolution(:FixedLengthArray, args)
julia> OREnvironment.set_feasible!(s2, false);
julia> OREnvironment.set_objfunction!(s2, 100.0);
julia> OREnvironment.is_first_solution_better(s1, s2, :max, true)
julia> true
julia> OREnvironment.is_first_solution_better(s1, s2, :max, false)
julia> false
julia> OREnvironment.set_feasible!(s2, true);
julia> OREnvironment.is_first_solution_better(s1, s2, :max, true)
julia> false
julia> OREnvironment.is_first_solution_better(s1, s2, :min, true)
julia> true
```
"""
function is_first_solution_better(s1::Solution,
                                  s2::Solution,
                                  objSense::Symbol,
                                  feasibilityRequired::Bool)::Bool
    return is_first_status_better(s1.status, s2.status, objSense, feasibilityRequired);
end

############################
# Specific methods for Solution
############################

"""
    get_solution(s, var)

Returns the solution value of variable/index `var`.

# Example
```jldoctest
julia> typeVariables = Int; numVariables = 5;  numConstraints = 3;
julia> status = OREnvironment.constructStatus(numConstraints);
julia> args = (typeVariables, numVariables, status);
julia> s = OREnvironment.constructSolution(:FixedLengthArray, args)
julia> OREnvironment.get_solution(s, 2)
0
```
"""
@inline function get_solution(s::FixedLengthArray{T}, variable::Int) where {T<:Real}
    return @inbounds s.sol[variable];
end

"""
    add_solution!(s, var, val)

Sets the solution value of variable/index `var` to `val`

# Example
```jldoctest
julia> typeVariables = Int; numVariables = 5;  numConstraints = 3;
julia> status = OREnvironment.constructStatus(numConstraints);
julia> args = (typeVariables, numVariables, status);
julia> s = OREnvironment.constructSolution(:FixedLengthArray, args)
julia> OREnvironment.add_solution!(s, 2, 4)
julia> OREnvironment.get_solution(s, 2)
4
```
"""
@inline function add_solution!(s::FixedLengthArray{T}, variable::Int, value::T) where {T<:Real}
    @inbounds s.sol[variable] = value;
end

"""
    remove_solution!(s, var)

Sets the solution value of variable/index `var` to 0.

# Example
```jldoctest
julia> typeVariables = Int; numVariables = 5;  numConstraints = 3;
julia> status = OREnvironment.constructStatus(numConstraints);
julia> args = (typeVariables, numVariables, status);
julia> s = OREnvironment.constructSolution(:FixedLengthArray, args)
julia> OREnvironment.add_solution!(s, 2, 4)
julia> OREnvironment.get_solution(s, 2)
4
julia> OREnvironment.remove_solution!(s, 2)
julia> OREnvironment.get_solution(s, 2)
0
```
"""
@inline function remove_solution!(s::FixedLengthArray{T}, variable::Int) where {T<:Real}
    @inbounds s.sol[variable] = zero(T);
end

"""
    remove_all_solution!(s)

Sets the solution value of all the variables to 0.

# Example
```jldoctest
julia> typeVariables = Int; numVariables = 5;  numConstraints = 3;
julia> status = OREnvironment.constructStatus(numConstraints);
julia> args = (typeVariables, numVariables, status);
julia> s = OREnvironment.constructSolution(:FixedLengthArray, args)
julia> OREnvironment.add_solution!(s, 1, 12)
julia> OREnvironment.get_solution(s, 1)
12
julia> OREnvironment.add_solution!(s, 2, 4)
julia> OREnvironment.get_solution(s, 2)
4
julia> OREnvironment.remove_all_solution!(s)
julia> OREnvironment.get_solution(s, 1)
0
julia> OREnvironment.get_solution(s, 2)
0
```
"""
function remove_all_solutions!(s::FixedLengthArray{T}) where {T<:Real}
    s.sol .= zero(T);
    return nothing;
end

"""
    copy_first_solution_to_second(s1, s2)

Copies all the information of the first solution (`s1`) into the second one (`s2`).

See file *./test/Solution.jl* for an example.
"""
@inline function copy_first_solution_to_second!(s1::FixedLengthArray{T},
                                                s2::FixedLengthArray{T}) where {T<:Real}
    update_status!(s2.status, s1.status);
    copy!(s2.sol, s1.sol);
    return nothing;
end

############################
# General methods for Solution when dealing with constraints
############################
"""
    update_constraint_consumption!(s, vconstr)

Given solution `s` and the constraints of the problem `vconstr`, this method computes the left-hand side of the constraints (constraint consumption) and saves the values in memory.

See file *./test/Solution.jl* for an example.
"""
function update_constraint_consumption!(s::Solution,
                                        constraints::Array{<:Constraint,1})
    local N::Int = length(constraints);
    @inbounds for i in 1:N
        local lhs::Float64 = compute_lhs(constraints[i], s);
        set_constraint_consumption!(s, i, lhs);
    end
    return nothing;
end

"""
    update_constraint_consumption_and_feasibility!(s, vconstr)

Given solution `s` and the constraints of the problem `vconstr`, this method computes the left-hand side of the constraints (constraint consumption), saves the values in memory and update the feasibility status according to the results obtained.

See file *./test/Solution.jl* for an example.
"""
function update_constraint_consumption_and_feasibility!(s::Solution,
                                                        constraints::Array{<:Constraint,1})
    local feasible::Bool = true;
    local N::Int = length(constraints);
    @inbounds for i in 1:N
        local lhs::Float64 = compute_lhs(constraints[i], s);
        set_constraint_consumption!(s, i, lhs);
        local isfeasible::Bool = is_feasible(constraints[i], lhs);
        if !isfeasible feasible = false; end
    end
    set_feasible!(s, feasible);
    return nothing;
end

"""
    update_constraint_consumption!(s, vconstr, var, Δ, idxConstraints)

Updates constraint consuption of solution `s` for the constraints of the problem (`vconstr`) when variable `var` has changed its value by `Δ`. `idxConstraints` is the vector of constraint indexes where variable `var` appears.

Note that this method is more efficient than updating the constraint consumption by recomputing again all the left-hand sides of the constraints: Most of the constraints does not change when modifying a variable, and those which change, only do it in one variable! Thus, if only a few changes have occurred in the vector of solutions since the last update, this incremental method should be preferred.

See file *./test/Solution.jl* for an example.
"""
function update_constraint_consumption!(s::Solution,
                                        constraints::Array{<:Constraint,1},
                                        variable::Int,
                                        Δvariable::Real,
                                        idxConstraints::Array{Int,1})
    @inbounds for i in idxConstraints
        local currentLHS::Float64 = get_constraint_consumption(s, i);
        local lhs::Float64 = compute_lhs_after_increment(variable, Δvariable,
                                                     currentLHS, constraints[i]);
        set_constraint_consumption!(s, i, lhs);
    end
    return nothing;
end

"""
    update_constraint_consumption_and_feasibility!(s, vconstr, var, Δ, idxConstraints)

Updates constraint consuption of solution `s`, and its feasibility status, for the constraints of the problem (`vconstr`) when variable `var` has changed its value by `Δ`. `idxConstraints` is the vector of constraint indexes where variable `var` appears.

Note that this method is more efficient than updating the constraint consumption by recomputing again all the left-hand sides of the constraints: Most of the constraints does not change when modifying a variable, and those which change, only do it in one variable! Thus, if only a few changes have occurred in the vector of solutions since the last update, this incremental method should be preferred.

See file *./test/Solution.jl* for an example.
"""
function update_constraint_consumption_and_feasibility!(s::Solution,
                                                        constraints::Array{<:Constraint,1},
                                                        variable::Int,
                                                        Δvariable::Real,
                                                        idxConstraints::Array{Int,1})
    local feasible::Bool = true;
    @inbounds for i in idxConstraints
        local currentLHS::Float64 = get_constraint_consumption(s, i);
        local lhs::Float64 = compute_lhs_after_increment(variable, Δvariable,
                                                    currentLHS, constraints[i]);
        set_constraint_consumption!(s, i, lhs);
        local isfeasible::Bool = is_feasible(constraints[i], lhs);
        if !isfeasible feasible = false; end
    end

    if should_global_feasibility_be_checked(is_feasible(s), feasible)
        feasible = is_current_consumption_feasible(s, constraints);
    end
    set_feasible!(s, feasible);
    return nothing;
end

function should_global_feasibility_be_checked(isSolutionFeasible::Bool,
                                              isIncrementFeasible::Bool)::Bool
    # if the solution was not already feasible but the current movement is,
    # then feasibility must be set checking all the constraints.
    return !isSolutionFeasible && isIncrementFeasible;
end

############################
# Specific methods for Solution
############################
"""
    add_solution_and_update_status!(s, var, val, p)

Sets the solution value of variable/index `var` to `val` and updates the status of the solution (including the consumption of the constraints in problem `p`).

See file *./test/Solution.jl* for an example.
"""
function add_solution_and_update_status!(s::FixedLengthArray{T},
                                         variable::Int,
                                         value::T,
                                         p::Problem) where {T<:Real}
    @inbounds local Δ::T  = value - s.sol[variable];
    local newObj::Float64 = get_objfunction(s) + Δ*get_cost(p, variable);
    @inbounds s.sol[variable] = value;
    update_constraint_consumption_and_feasibility!(s, get_constraints(p),
                            variable, Δ, get_constraints_of_variable(p, variable));
    set_objfunction!(s, newObj);
    return nothing;
end

"""
    remove_solution_and_update_status!(s, var, p)

Sets the solution value of variable/index `var` to 0 and updates the status of the solution (including the consumption of the constraints in problem `p`).

See file *./test/Solution.jl* for an example.
"""
function remove_solution_and_update_status!(s::FixedLengthArray{T},
                                            variable::Int,
                                            p::Problem) where {T<:Real}
    add_solution_and_update_status!(s, variable, zero(T), p);
    return nothing;
end

"""
    remove_all_solution_and_update_status!(s, p)

Sets the solution value of all the variables to 0 and updates the status of the solution (including the consumption of the constraints in problem `p`).

See file *./test/Solution.jl* for an example.
"""
function remove_all_solutions_and_update_status!(s::FixedLengthArray{T},
                                                 p::Problem) where {T<:Real}
    s.sol .= zero(T);
    set_objfunction!(s, 0.0);
    s.status.constraintLhsConsumption .= 0.0;
    local feasible::Bool = is_current_consumption_feasible(s, get_constraints(p));
    set_feasible!(s, feasible);
    return nothing;
end


module KmerAnalysis

export
    CountMode,
    Canonical,
    CANONICAL,
    NonCanonical,
    NONCANONICAL,

    # MerCount
    MerCount,
    mer,
    freq,
    merge,
    merge_into!,
    collapse_into_counts,
    collapse_into_counts!,
    collect_mers,
    collect_mers!,

    # Counters
    serial_mem,
    count!,
    dist_mem,

    # K-mer frequency spectra
    KmerFrequencySpectra,
    spectra,
    spectra!,
    find_next_peak,
    find_spectra_peaks,

    # Indexed k-mer counts
    IndexedCounts,
    add_count!,
    project_count

using BioSequences, ReadDatastores, Distributed

abstract type CountMode end

struct Canonical    <: CountMode end
struct NonCanonical <: CountMode end

@inline (::Canonical)(x) = canonical(x)
@inline (::NonCanonical)(x) = fwmer(x)

const CANONICAL = Canonical()
const NONCANONICAL = NonCanonical()

abstract type AbstractKmerCounter{M<:AbstractMer} <: Function end

include("MerCount.jl")
include("counters/serial_mem.jl")
include("counters/dist_mem.jl")
include("IndexedCounts.jl")
include("KmerFrequencySpectra.jl")
include("kci.jl")

end # module


"""
    Controls(mitigate, remove, geoeng, adapt)

Create data structure for climate controls.

# Examples
```jldoctest
a = zeros(4);
C = Controls(a, a, a, a);
C.geoeng

# output
4-element Array{Float64,1}:
 0.0
 0.0
 0.0
 0.0
```

See also: [`ClimateModel`](@ref)
"""
mutable struct Controls
    mitigate::Array{Float64,1}
    remove::Array{Float64,1}
    geoeng::Array{Float64,1}
    adapt::Array{Float64,1}
end


module Coord

using Gadfly
using Compose
using DataArrays
import Gadfly.Maybe
import Iterators: cycle

export cartesian

# Cartesian coordinates with position given by the x and y (and similar)
# aesthetics.
#
# Args:
#   xvars: Aesthetics to consider when choosing x bounds.
#   yvars: Aesthetics to consider when choosing y bounds.
#   xmin, xmax, ymin, ymax:
#     Force a particular x or y bound, rather than trying to choose it
#     from the data.
#
#
immutable Cartesian <: Gadfly.CoordinateElement
    xvars::Vector{Symbol}
    yvars::Vector{Symbol}
    xmin
    xmax
    ymin
    ymax
    xflip
    yflip

    function Cartesian(
            xvars=[:x, :xviewmin, :xviewmax, :xmin, :xmax],
            yvars=[:y, :yviewmin, :yviewmax, :ymin, :ymax, :middle,
                   :lower_hinge, :upper_hinge, :lower_fence, :upper_fence, :outliers];
            xflip::Bool=false, yflip::Bool=false,
            xmin=nothing, xmax=nothing, ymin=nothing, ymax=nothing)
        new(xvars, yvars, xmin, xmax, ymin, ymax, xflip, yflip)
    end
end

const cartesian = Cartesian

# Produce a canvas with suitable cartesian coordinates.
#
# Args:
#   coord: cartesian coordinate instance.
#
# Returns:
#   A compose Canvas.
#
function apply_coordinate(coord::Cartesian, aess::Gadfly.Aesthetics...)
    xmin = nothing
    xmax = nothing
    for var in coord.xvars
        for aes in aess
            if getfield(aes, var) === nothing
                continue
            end

            vals = getfield(aes, var)
            if !isa(vals, AbstractArray)
                vals = {vals}
            end

            for val in vals
                if !Gadfly.isconcrete(val)
                    continue
                end

                if xmin === nothing || val < xmin
                    xmin = val
                end

                if xmax === nothing || val > xmax
                    xmax = val
                end
            end
        end
    end

    ymin = nothing
    ymax = nothing
    for var in coord.yvars
        for aes in aess
            if getfield(aes, var) === nothing
                continue
            end

            # Outliers is an odd aesthetic that needs special treatment.
            if var == :outliers
                for vals in aes.outliers
                    for val in vals
                        if !Gadfly.isconcrete(val)
                            continue
                        end

                        if ymin === nothing || val < ymin
                            ymin = val
                        end

                        if ymax === nothing || val > ymax
                            ymax = val
                        end
                    end
                end

                continue
            end

            vals = getfield(aes, var)
            if !isa(vals, AbstractArray)
                vals = {vals}
            end

            for val in vals
                if !Gadfly.isconcrete(val)
                    continue
                end

                if ymin === nothing || val < ymin
                    ymin = val
                end

                if ymax === nothing || val > ymax
                    ymax = val
                end
            end
        end
    end

    xmin = coord.xmin === nothing ? xmin : coord.xmin
    xmax = coord.xmax === nothing ? xmax : coord.xmax
    ymin = coord.ymin === nothing ? ymin : coord.ymin
    ymax = coord.ymax === nothing ? ymax : coord.ymax

    if xmin === nothing || !isfinite(xmin)
        xmin = 0.0
        xmax = 1.0
    end

    if ymin === nothing || !isfinite(ymin)
        ymin = 0.0
        ymax = 1.0
    end

    # A bit of kludge. We need to extend a bit to be able to fit bars when
    # using discrete scales. TODO: Think more carefully about this. Is there a
    # way for the geometry to let the coordinates know that a little extra room
    # is needed to draw everything?
    if all([aes.x === nothing || typeof(aes.x) <: PooledDataArray for aes in aess])
        xmin -= 0.5
        xmax += 0.5
        xpadding = 0.0
    else
        xpadding = 0.03 * (xmax - xmin)
    end

    if all([aes.y === nothing || typeof(aes.y) <: PooledDataArray for aes in aess])
        ymin -= 0.5
        ymax += 0.5
        ypadding = 0.0
    else
        ypadding = 0.03 * (ymax - ymin)
    end

    width  = xmax - xmin + 2.0 * xpadding
    height = ymax - ymin + 2.0 * ypadding

    compose(
        canvas(unit_box=UnitBox(
            coord.xflip ? xmax + xpadding : xmin - xpadding,
            coord.yflip ? ymin - ypadding : ymax + ypadding,
            coord.xflip ? -width : width,
            coord.yflip ? height : -height)),
        svgclass("plotpanel"))
end

end # module Coord


@testmodule DataStreamTest begin
  using MLLibrary

  function test_DataSetCount()

    path = joinpath(Pkg.dir("MLLibrary"),"resources","datasets")
    @test MLLibrary.DataParser.DataStream.getDataSetCount(joinpath(path,"abalone.data")) ==4177
    @test MLLibrary.DataParser.DataStream.getDataSetCount(joinpath(path,"3percent-miscategorization.csv")) ==10000
    @test MLLibrary.DataParser.DataStream.getDataSetCount(joinpath(path,"10percent-miscatergorization.csv")) ==10000
    @test MLLibrary.DataParser.DataStream.getDataSetCount(joinpath(path,"Seperable.csv")) ==10000

  end

  function test_getFeatureCount()
    path = joinpath(Pkg.dir("MLLibrary"),"resources","datasets")
    @test MLLibrary.DataParser.DataStream.getFeatureCount(joinpath(path,"abalone.data")) ==9
    @test MLLibrary.DataParser.DataStream.getFeatureCount(joinpath(path,"3percent-miscategorization.csv")) ==10
    @test MLLibrary.DataParser.DataStream.getFeatureCount(joinpath(path,"10percent-miscatergorization.csv")) ==10
    @test MLLibrary.DataParser.DataStream.getFeatureCount(joinpath(path,"Seperable.csv")) ==10
  end

end


println("Compile libraries on main processor..."); flush(stdout)
include("imports.jl")
println("Finished loading libraries."); flush(stdout)

println("ARGS: ", ARGS)

marker = Symbol(ARGS[1])
println("marker: ", marker); flush(stdout)

K = parse(Int, ARGS[2])
println("K: ", K)

stm = parse(Bool, ARGS[3])
println("skewtmix: ", stm); flush(stdout)

sims = dict_list(Dict(:K => K, :marker => marker, :skewtmix => stm))

# Run simulations in parallel.
println("Running analyses ..."); flush(stdout)
@time run(sims[1])

# Post process
println("Post process ..."); flush(stdout)
@time postprocess(sims[1])

# Print number of small clusters
@time print_number_of_small_clusters(sims[1], p=0.02)


# This file was generated, do not modify it. # hide
fp = fitted_params(mach_uni)
@show round.(fp.coefs, sigdigits=3)
@show round(fp.intercept, sigdigits=3)


using BenchmarkTools
using JLD

# Define a parent BenchmarkGroup to contain our suite
const suite = BenchmarkGroup()

# Add some child groups to our benchmark suite.
suite["utf8"] = BenchmarkGroup(["string", "unicode"])
suite["trigonometry"] = BenchmarkGroup(["math", "triangles"])

# This string will be the same every time because we're seeding the RNG
teststr = UTF8String(join(rand(MersenneTwister(1), 'a':'d', 10^4)))

# Add some benchmarks to the "utf8" group
suite["utf8"]["replace"] = @benchmarkable replace($teststr, "a", "b")
suite["utf8"]["join"] = @benchmarkable join($teststr, $teststr)

# Add some benchmarks to the "trigonometry" group
for f in (sin, cos, tan)
    for x in (0.0, pi)
        suite["trigonometry"][string(f), x] = @benchmarkable $(f)($x)
    end
end

# Load the suite's cached parameters as part of including the file. This is much
# faster and more reliable than re-tuning `suite` every time the file is included
paramspath = joinpath(Pkg.dir("BenchmarkTools"), "benchmark", "params.jld")
#tune!(suite); JLD.save(paramspath, "suite", params(suite));
loadparams!(suite, JLD.load(paramspath, "suite"), :evals, :samples);


using Dates
using TimeZones

struct DateWrapper
    days::Integer
end

Base.:(==)(x::DateWrapper, y::DateWrapper) = x.days == y.days

function DateWrapper(y, m, d)
    days = (Date(y,m,d)-Date(1970,1,1)).value
    return DateWrapper(days)
end

# timezone_offset here is in seconds
struct TimeWrapper
    nanoseconds::Integer
    timezone_offset::Union{Integer, Nothing}
end

Base.:(==)(x::TimeWrapper, y::TimeWrapper) = x.nanoseconds == y.nanoseconds && x.timezone_offset == y.timezone_offset

# if o is an integer then it is timezone_offset in mins
function TimeWrapper(h, m, s, n=0; timezone=nothing)
    nanoseconds = n + s*1e+9 + m*6e+10 + h*3.6e+12
    return TimeWrapper(nanoseconds, if isa(timezone, Number) Integer(floor(timezone*60)) else timezone end)
end

TimeWrapper(nanoseconds) = TimeWrapper(nanoseconds, nothing)

struct DateTimeWrapper
    seconds::Integer
    nanoseconds::Integer
    timezone_offset::Union{Integer, Nothing}
    timezone_name::Union{String, Nothing}
end

Base.:(==)(x::DateTimeWrapper, y::DateTimeWrapper) = x.seconds == y.seconds && x.nanoseconds == y.nanoseconds && x.timezone_offset == y.timezone_offset && x.timezone_name == y.timezone_name

DateTimeWrapper(seconds, nanoseconds) = DateTimeWrapper(seconds, nanoseconds, nothing, nothing)
DateTimeWrapper(seconds, nanoseconds, timezone::Integer) = DateTimeWrapper(seconds, nanoseconds, timezone, nothing)
DateTimeWrapper(seconds, nanoseconds, timezone::String) = DateTimeWrapper(seconds, nanoseconds, nothing, timezone)

function DateTimeWrapper(y,m,d,h,i,s=0,n=0;timezone=nothing)
    millisecond = (DateTime(y,m,d,h,i,s)-DateTime(1970,1,1)).value
    seconds = Integer(floor(millisecond/1000))
    nanoseconds = n

    if timezone == nothing
        return DateTimeWrapper(seconds, nanoseconds, nothing, nothing)
    elseif isa(timezone, String)
        return DateTimeWrapper(seconds, nanoseconds, nothing, timezone)
    elseif isa(timezone, Number)
        return DateTimeWrapper(seconds, nanoseconds, Integer(floor(timezone*60)), nothing)
    else
        throw(ErrorException("Not a valid DateTime"))
    end
end

struct DurationWrapper
    months::Integer
    days::Integer
    seconds::Integer
    nanoseconds::Integer
end

Base.:(==)(x::DurationWrapper, y::DurationWrapper) = x.months == y.months && x.days == y.days && x.seconds == y.seconds && x.nanoseconds == y.nanoseconds

function DurationWrapper(years, months, days, hours, minutes, seconds=0, nanoseconds=0)
    m = months + years * 12
    d = days
    s = seconds + hours * 3600 + minutes * 60
    n = nanoseconds
    return DurationWrapper(m, d, s, n)
end

DurationWrapper(;years=0, months=0, days=0, hours=0, minutes=0, seconds=0, nanoseconds=0) = DurationWrapper(years, months, days, hours, minutes, seconds, nanoseconds)

function hydrate_date(days; use_julia_dates=true)
    if use_julia_dates
        return Date(1970,1,1) + Dates.Day(days)
    else
        return DateWrapper(days)
    end
end

function dehydrate_date(value::Date)
    return Structure(UInt8('D'), Any[(value-Date(1970,1,1)).value])
end

function dehydrate_date(value::DateWrapper)
    return Structure(UInt8('D'), Any[value.days])
end

# Julia Times do not support time zones
function hydrate_time(nanoseconds::Integer, timezone::Union{Integer, Nothing}=nothing; use_julia_dates=true)
    if use_julia_dates
        return Time(Dates.Nanosecond(nanoseconds))
    else
        return TimeWrapper(nanoseconds, timezone)
    end
end

function dehydrate_time(value::Time)
    return Structure(UInt8('t'), Any[value.instant.value])
end

function dehydrate_time(value::TimeWrapper)
    if value.timezone_offset != nothing
        return Structure(UInt8('T'), Any[value.nanoseconds, value.timezone_offset])
    else
        return Structure(UInt8('t'), Any[value.nanoseconds])
    end
end

# Julia DateTimes do not support nanosecond accuracy time
function hydrate_datetime(seconds::Integer, nanoseconds::Integer, timezone::Union{String, Integer, Nothing}=nothing; use_julia_dates=true)
    if use_julia_dates
        date_time =  DateTime(1970,1,1) + Dates.Second(seconds) + Dates.Nanosecond(nanoseconds)
        if timezone == nothing
            return date_time
        elseif isa(timezone, String)
            return ZonedDateTime(date_time, TimeZone(timezone))
        else
            return ZonedDateTime(date_time, FixedTimeZone("Time Zone", timezone))
        end
    else
        if timezone == nothing
            return DateTimeWrapper(seconds, nanoseconds, nothing, nothing)
        elseif isa(timezone, String)
            return DateTimeWrapper(seconds, nanoseconds, nothing, timezone)
        else
            return DateTimeWrapper(seconds, nanoseconds, timezone, nothing)
        end
    end
end

function dehydrate_datetime(value::DateTime)
    millisecond = (value-DateTime(1970,1,1)).value # miliseconds
    seconds = Integer(floor(millisecond/1000.0))
    nanoseconds = (millisecond % 1000) * 1000000
    return Structure(UInt8('d'), Any[seconds, nanoseconds])
end

function dehydrate_datetime(value::ZonedDateTime)
    d = TimeZones.localtime(value)
    millisecond = (d-DateTime(1970,1,1)).value # miliseconds
    seconds = Integer(floor(millisecond/1000.0))
    nanoseconds = (millisecond % 1000) * 1000000
    if isa(value.timezone, FixedTimeZone)
        return Structure(UInt8('F'), Any[seconds, nanoseconds, value.timezone.offset.std.value])
    else
        return Structure(UInt8('f'), Any[seconds, nanoseconds, string(value.timezone.name)])
    end
end

function dehydrate_datetime(value::DateTimeWrapper)
    if value.timezone_offset != nothing
        return Structure(UInt8('F'), Any[value.seconds, value.nanoseconds, value.timezone_offset])
    elseif value.timezone_name != nothing
        return Structure(UInt8('f'), Any[value.seconds, value.nanoseconds, value.timezone_name])
    else
        return Structure(UInt8('d'), Any[value.seconds, value.nanoseconds])
    end

end

function hydrate_duration(months, days, seconds, nanoseconds; use_julia_dates=true)
    if use_julia_dates
        return Month(months) + Day(days) + Second(seconds) + Nanosecond(nanoseconds)
    else
        return DurationWrapper(months, days, seconds, nanoseconds)
    end
end

function tons(c::Dates.CompoundPeriod)
    s::Float64 = 0.0
    for p in c.periods
        v = Dates.tons(p)
        s += v
    end
    return s
end

function dehydrate_duration(value::Dates.CompoundPeriod)
    value = canonicalize(value)

    months = 0
    periods = Dates.Period[]

    for p in value.periods
        if isa(p, Dates.Month)
            months += p.value
        elseif isa(p, Dates.Year)
            months += p.value * 12
        else
            push!(periods, p)
        end
    end

    value = Dates.CompoundPeriod(periods)
    v = tons(value)

    nanoseconds = v % 1000000000
    v = (v - nanoseconds) / 1000000000
    seconds = v % 86400
    v = (v - seconds) / 86400
    days = v

    return Structure(UInt8('E'), Any[Integer(months), Integer(days), Integer(seconds), Integer(nanoseconds)])
end

dehydrate_duration(value::Period) = dehydrate_duration(Dates.CompoundPeriod(value))

function dehydrate_duration(value::DurationWrapper)
    return Structure(UInt8('E'), Any[value.months, value.days, value.seconds, value.nanoseconds])
end

dehydrate(value::DateWrapper) = dehydrate_date(value)
dehydrate(value::TimeWrapper) = dehydrate_time(value)
dehydrate(value::DateTimeWrapper) = dehydrate_datetime(value)
dehydrate(value::DurationWrapper) = dehydrate_duration(value)
dehydrate(value::Dates.Date) = dehydrate_date(value)
dehydrate(value::Dates.Time) = dehydrate_time(value)
dehydrate(value::Dates.DateTime) = dehydrate_datetime(value)
dehydrate(value::ZonedDateTime) = dehydrate_datetime(value)
dehydrate(value::Dates.CompoundPeriod) = dehydrate_duration(value)
dehydrate(value::Period) = dehydrate_date(value)

function temporal_hydration_functions()
    return Dict(
        UInt8('D') => (psh, values)->hydrate_date(values..., use_julia_dates=psh.use_julia_dates),
        UInt8('T') => (psh, values)->hydrate_time(values..., use_julia_dates=psh.use_julia_dates),
        UInt8('t') => (psh, values)->hydrate_time(values..., use_julia_dates=psh.use_julia_dates),
        UInt8('F') => (psh, values)->hydrate_datetime(values..., use_julia_dates=psh.use_julia_dates),
        UInt8('f') => (psh, values)->hydrate_datetime(values..., use_julia_dates=psh.use_julia_dates),
        UInt8('d') => (psh, values)->hydrate_datetime(values..., use_julia_dates=psh.use_julia_dates),
        UInt8('E') => (psh, values)->hydrate_duration(values..., use_julia_dates=psh.use_julia_dates)
    )
end


abstract type AbstractWindFarmModel end

"""
WindFarm(windfarm, windresource, windfarmstates)

Struct defining a wind farm

# Arguments
- `turbine_x::Array{Float}(Nturbines)`: contains windturbine x coordinates in the global reference frame
- `turbine_y::Array{Float}(Nturbines)`: contains windturbine y coordinates in the global reference frame
- `turbine_z::Array{Float}(Nturbines)`: contains windturbine base/z coordinates in the global reference frame
- `turbine_definition_ids::Array{Int}(Nturbines)`: contains integers for each wind turbine specifying its definition
- `turbine_definitions::Array{AbstractTurbineDefinition}(Ntypes)`: contains structs defining each wind turbine definition (design) used in the farm
"""
struct WindFarm{AF1,AF2,AF3,AI,AS} <: AbstractWindFarmModel

    # farm design properties
    turbine_x::AF1
    turbine_y::AF2
    turbine_z::AF3
    turbine_definition_ids::AI
    turbine_definitions::AS

end

struct SingleWindFarmState{TI,AF1,AF2,AF3,AF4,AF5,AF6,AF7,AF8,AF9,AI} <: AbstractWindFarmModel

    # farm properties in rotated frame
    id::TI
    turbine_x::AF1
    turbine_y::AF2
    turbine_z::AF3
    turbine_yaw::AF4
    turbine_ct::AF5
    turbine_ai::AF6
    turbine_inflow_velcities::AF7
    turbine_generators_powers::AF8
    turbine_local_ti::AF9
    sorted_turbine_index::AI

end


export visualize

using DifferentiableNAS
using Flux
using CUDA
using LightGraphs
using GraphPlot
using Cairo
using Compose
using BSON
using LaTeXStrings
using PyCall
graphviz = pyimport("graphviz")
include("training_utils.jl")

function visualize(αs, filename::String)
    nodelabels = ["c_{k-2}", "c_{k-1}", "0", "1", "2", "3", "c_{k}"]
    inputindices, _, opnames = discretize(αs, 1, false, 4; plotting = true)
    g = graphviz.Digraph(
          format="png",
          edge_attr=Dict("fontsize"=>"20", "fontname"=>"times"),
          node_attr=Dict("style"=>"filled", "shape"=>"rect", "align"=>"center", "fontsize"=>"20", "height"=>"0.5", "width"=>"0.5", "penwidth"=>"2", "fontname"=>"times"),
          engine="dot")
    g.body = ["rankdir=LR"]

    g.node(nodelabels[1], fillcolor="darkseagreen2")
    g.node(nodelabels[2], fillcolor="darkseagreen2")

    for i in 3:6
      g.node(nodelabels[i], fillcolor="lightblue")
    end

    for i in 1:length(inputindices)
        g.edge(nodelabels[inputindices[i][1]], nodelabels[i+2], label=opnames[i][1], fillcolor="gray")
        g.edge(nodelabels[inputindices[i][2]], nodelabels[i+2], label=opnames[i][2], fillcolor="gray")
    end

    g.node(nodelabels[7], fillcolor="palegoldenrod")
    for i in 3:6
      g.edge(nodelabels[i], nodelabels[7], fillcolor="gray")
    end
    g.render(filename, view=false)
end

function visualize(filename::String)
    nodelabels = ["c_{k-2}", "c_{k-1}", "0", "1", "2", "3", "c_{k}"]
    inputindices, _, opnames = parse_genotype(channels=1, reduce=false, steps=4)
    g = graphviz.Digraph(
          format="png",
          edge_attr=Dict("fontsize"=>"20", "fontname"=>"times"),
          node_attr=Dict("style"=>"filled", "shape"=>"rect", "align"=>"center", "fontsize"=>"20", "height"=>"0.5", "width"=>"0.5", "penwidth"=>"2", "fontname"=>"times"),
          engine="dot")
    g.body = ["rankdir=LR"]

    g.node(nodelabels[1], fillcolor="darkseagreen2")
    g.node(nodelabels[2], fillcolor="darkseagreen2")

    for i in 3:6
      g.node(nodelabels[i], fillcolor="lightblue")
    end

    for i in 1:length(inputindices)
        g.edge(nodelabels[inputindices[i][1]], nodelabels[i+2], label=opnames[i][1], fillcolor="gray")
        g.edge(nodelabels[inputindices[i][2]], nodelabels[i+2], label=opnames[i][2], fillcolor="gray")
    end

    g.node(nodelabels[7], fillcolor="palegoldenrod")
    for i in 3:6
      g.edge(nodelabels[i], nodelabels[7], fillcolor="gray")
    end
    g.render(filename, view=false)
end

function visualize_supercell(filename::String)
    nodelabels = ["c_{k-2}", "c_{k-1}", "0", "1", "2", "3", "c_{k}"]
    g = graphviz.Digraph(
          format="png",
          edge_attr=Dict("fontsize"=>"20", "fontname"=>"times"),
          node_attr=Dict("style"=>"filled", "shape"=>"rect", "align"=>"center", "fontsize"=>"20", "height"=>"0.5", "width"=>"0.5", "penwidth"=>"2", "fontname"=>"times"),
          engine="dot")
    g.body = ["rankdir=LR"]

    g.node(nodelabels[1], fillcolor="darkseagreen2")
    g.node(nodelabels[2], fillcolor="darkseagreen2")

    for i in 3:6
      g.node(nodelabels[i], fillcolor="lightblue")
    end

    connects = vcat([[(j,i) for j = 1:i-1] for i = 3:6]...)

    for connect in connects
        g.edge(nodelabels[connect[1]], nodelabels[connect[2]], fillcolor="black", penwidth="3")
    end

    g.node(nodelabels[7], fillcolor="palegoldenrod")
    for i in 3:6
      g.edge(nodelabels[i], nodelabels[7], fillcolor="gray")
    end
    g.render(filename, view=false)
end


#InspectDR: Secondary annotation (Cairo layer)
#-------------------------------------------------------------------------------

#==Constants
===============================================================================#
const ALIGN_MAP = Dict{Symbol, CAlignment}(
	:tl => ALIGN_TOP | ALIGN_LEFT,
	:tc => ALIGN_TOP | ALIGN_HCENTER,
	:tr => ALIGN_TOP | ALIGN_RIGHT,

	:cl => ALIGN_VCENTER | ALIGN_LEFT,
	:cc => ALIGN_VCENTER | ALIGN_HCENTER,
	:cr => ALIGN_VCENTER | ALIGN_RIGHT,

	:bl => ALIGN_BOTTOM | ALIGN_LEFT,
	:bc => ALIGN_BOTTOM | ALIGN_HCENTER,
	:br => ALIGN_BOTTOM | ALIGN_RIGHT,
)

#==Rendering text annotation
===============================================================================#
function render(ctx::CairoContext, rstrip::RStrip2D, a::TextAnnotation)
	align = get(ALIGN_MAP, a.align, ALIGN_BOTTOM | ALIGN_LEFT)
	angle = deg2rad(a.angle)
	pt = axis2dev(a.pos, rstrip.xf, rstrip.ixf, rstrip.bb)
	render(ctx, a.text, pt, a.font, angle=angle, align=align)
	return
end

#==Rendering markers
===============================================================================#

function render(ctx::CairoContext, rstrip::RStrip2D, mkr::HVMarker)
	graphbb = rstrip.bb #WANTCONST
	if :none == mkr.line.style
		return
	end

	setlinestyle(ctx, LineStyle(mkr.line))
	pt = axis2aloc(mkr.pos, rstrip.ixf)
	pt = apply(rstrip.xf, pt)

	if mkr.vmarker
		drawline(ctx, Point2D(pt.x, graphbb.ymin), Point2D(pt.x, graphbb.ymax))
	end
	if mkr.hmarker
		drawline(ctx, Point2D(graphbb.xmin, pt.y), Point2D(graphbb.xmax, pt.y))
	end

	return
end

#==Rendering polyline for annotation
===============================================================================#
function render(ctx::CairoContext, rstrip::RStrip2D, a::PolylineAnnotation)
	x = a.x; y = a.y
	setlinestyle(ctx, LineStyle(a.line))
	pt = apply(rstrip.xf, Point2D(x[1], y[1]))
	Cairo.move_to(ctx, pt.x, pt.y)
	for i in 2:length(x)
		pt = apply(rstrip.xf, Point2D(x[i], y[i]))
		Cairo.line_to(ctx, pt.x, pt.y)
	end
	if a.closepath
		Cairo.close_path(ctx)
		renderfill(ctx, a.fillcolor)
	end
	Cairo.stroke(ctx)
	return
end

#==Generic rendering algorithm for annotations
===============================================================================#

#Basic render function for PlotAnnotation types:
function render(ctx::CairoContext, rstrip::RStrip2D, a::PlotAnnotation)
	if 0 == a.strip || a.strip == rstrip.istrip
		render(ctx, rstrip, a)
	end
	return
end

function render(ctx::CairoContext, rstrip::RStrip2D, alist::Vector{T}) where T<:PlotAnnotation
	for a in alist
		render(ctx, rstrip, a)
	end
	return
end

#Last line


using StochasticBlockModel
using Test

@testset "OptResults struct" begin

    @test_throws ArgumentError StochasticBlockModel.OptResults(2.0, 1.0, :Optimal, 2.0, 0, 0, 0) # Check LB, UB
    @test_throws ArgumentError StochasticBlockModel.OptResults(2.0, 4.0, :Optimal, 2.0, 0, 0, 0) # Check optimal status
    #@test_throws DomainError StochasticBlockModel.OptResults(2.0, 2.0, :Optimal, 0.0, 0, 0, 0) # Check solvetime domain
    @test_throws DomainError StochasticBlockModel.OptResults(2.0, 2.0, :Optimal, -2.0, 0, 0, 0) # Check solvetime domain
    @test_throws DomainError StochasticBlockModel.OptResults(2.0, 2.0, :Optimal, 2.0, -1, 0, 0) # Check iterations domain
    @test_throws DomainError StochasticBlockModel.OptResults(2.0, 2.0, :Optimal, 2.0, 0, -1, 0) # Check nodecount domain
    @test_throws DomainError StochasticBlockModel.OptResults(2.0, 2.0, :Optimal, 2.0, 0, 0, -1) # Check nodecount domain
end


import Dates

struct HolidayRoundingNBD <: HolidayRoundingConvention end
function apply(::HolidayRoundingNBD, date::Dates.Date, calendar::Calendar)::Dates.Date
    while is_holiday(calendar, date)
        date += Dates.Day(1)
    end
    date
end
Base.show(io::IO, ::HolidayRoundingNBD) = print(io, "NBD")

struct HolidayRoundingPBD <: HolidayRoundingConvention end
function apply(::HolidayRoundingPBD, date::Dates.Date, calendar::Calendar)::Dates.Date
    while is_holiday(calendar, date)
        date -= Dates.Day(1)
    end
    date
end
Base.show(io::IO, ::HolidayRoundingPBD) = print(io, "PBD")

struct HolidayRoundingNBDSM <: HolidayRoundingConvention end
function apply(::HolidayRoundingNBDSM, date::Dates.Date, calendar::Calendar)::Dates.Date
    new_date = date
    while is_holiday(calendar, new_date)
        new_date += Dates.Day(1)
    end

    if Dates.month(new_date) != Dates.month(date)
        new_date = date
        while is_holiday(calendar, new_date)
            new_date -= Dates.Day(1)
        end
    end

    new_date
end
Base.show(io::IO, ::HolidayRoundingNBDSM) = print(io, "NBDSM")

struct HolidayRoundingPBDSM <: HolidayRoundingConvention end
function apply(::HolidayRoundingPBDSM, date::Dates.Date, calendar::Calendar)::Dates.Date
    new_date = date
    while is_holiday(calendar, new_date)
        new_date -= Dates.Day(1)
    end

    if Dates.month(new_date) != Dates.month(date)
        new_date = date
        while is_holiday(calendar, new_date)
            new_date += Dates.Day(1)
        end
    end

    new_date
end
Base.show(io::IO, ::HolidayRoundingPBDSM) = print(io, "PBDSM")

struct HolidayRoundingNR <: HolidayRoundingConvention end
apply(::HolidayRoundingNR, date::Dates.Date, ::Calendar) = date
Base.show(io::IO, ::HolidayRoundingNR) = print(io, "NR")

const HOLIDAY_ROUNDING_MAPPINGS = Dict("NR" => HolidayRoundingNR(), "NBD" => HolidayRoundingNBD(), "PBD" => HolidayRoundingPBD(), "NBDSM" => HolidayRoundingNBDSM(), "PBDSM" => HolidayRoundingPBDSM())


# Conditionals

for op ∈ (:(==), :(!=), :(<), :(<=), :(>), :(>=) )
    rename = VECTOR_SYMBOLS[op]
    erename = Symbol(:e, rename)
    # rename_mask = Symbol(rename, :_mask)
    @eval begin
        # scalar versions handled in floating_point_arithmetic.jl
        # @inline $rename(s1::ScalarTypes, s2::ScalarTypes) = $op(s1,s2)
        @vectordef $rename function Base.$op(v1, v2) where {W,T}
            llvmwrap_bitmask(Val{$(QuoteNode(op))}(), extract_data(v1), extract_data(v2))
        end
        @vectordef $rename function Base.$op(v1, v2::T) where {W,T}
            $rename(extract_data(v1), v2)
        end
        @vectordef $rename function Base.$op(v1::T, v2) where {W,T}
            $rename(extract_data(v1), v2)
        end
        @evectordef $erename function Base.$op(v1, v2) where {W,T <: FloatingTypes}
            llvmwrap_bitmask_notfast(Val{$(QuoteNode(op))}(), extract_data(v1), extract_data(v2))
        end
    end
end
@inline visfinite(s::ScalarTypes) = isfinite(s)
@inline function visfinite(v1::_Vec{_W,T}) where {_W,T<:FloatingTypes}
    W = _W + 1
    U = uint_type(T)
    em = vbroadcast(Vec{W,U}, exponent_mask(T))
    iv = vreinterpret(Vec{W,U}, v1)
    # evnot_equal(vand(iv, em), em)
    vnot_equal(vand(iv, em), em)
end
@inline visfinite(v1::SVec{W}) where {W} = SVec{W}(visfinite(extract_data(v1)))
@inline Base.isfinite(v1::SVec{W}) where {W} = SVec{W}(visfinite(extract_data(v1)))

@inline visinf(s1::ScalarTypes) = isinf(s1)
@inline visinf(v1::_Vec{W,T}) where {W,T<:FloatingTypes} = evisequal(vabs(v1), vbroadcast(_Vec{W,T},typemax(T)))
@inline visinf(v1::SVec{W}) where {W} = SVec{W}(visinf(extract_data(v1)))
@inline Base.isinf(v1::SVec{W}) where {W} = SVec{W}(visinf(extract_data(v1)))

@inline visnan(s1::ScalarTypes) = isnan(s1)
# @inline visnan(v1::_Vec{W,T}) where {W,T<:FloatingTypes} = evnot_equal(v1, v1)
@inline visnan(v1::_Vec{W,T}) where {W,T<:FloatingTypes} = vnot_equal(v1, v1)
@inline visnan(v1::SVec{W}) where {W} = SVec{W}(visnan(extract_data(v1)))
@inline Base.isnan(v1::SVec{W}) where {W} = SVec{W}(visnan(extract_data(v1)))

@inline vissubnormal(s1::ScalarTypes) = issubnormal(s1)
@inline function vissubnormal(v1::_Vec{_W,T}) where {_W,T<:FloatingTypes}
    W = _W + 1
    U = uint_type(T)
    em = vbroadcast(Vec{W,U}, exponent_mask(T))
    sm = vbroadcast(Vec{W,U}, significand_mask(T))
    iv = vreinterpret(Vec{W,U}, v1)
    vand(evisequal(vand(iv, em), vbroadcast(Vec{W,U}, 0)), evnot_equal(vand(iv, sm), vbroadcast(Vec{W,U}, 0)))
end
@inline vissubnormal(v1::SVec{W}) where {W} = SVec{W}(vissubnormal(extract_data(v1)))
@inline Base.issubnormal(v1::SVec{W}) where {W} = SVec{W}(vissubnormal(extract_data(v1)))

@inline signbit(s1::ScalarTypes) = signbit(s1)
@inline function vsignbit(v1::_Vec{W,T}) where {W,T<:FloatingTypes}
    U = uint_type(T)
    sm = vbroadcast(_Vec{W,U}, sign_mask(T))
    iv = vreinterpret(_Vec{W,U}, v1)
    vnot_equal(vand(iv, sm), vbroadcast(_Vec{W,U}, 0))
end
@inline function vsignbit(v1::SVec{W,T}) where {W,T<:FloatingTypes}
    SVec{W}(vsignbit(extract_data(v1)))
end
@inline function Base.signbit(v1::SVec{W,T}) where {W,T<:FloatingTypes}
    SVec{W}(vsignbit(extract_data(v1)))
end

# @inline vifelse(c::Bool, x, y) = c ? x : y
@inline vifelse(c::Bool, x, y) = ifelse(c, x, y)
@generated function vifelse(v1::_Vec{_W,Bool}, v2::_Vec{_W,T}, v3::_Vec{_W,T}) where {_W,T}
    W = _W + 1
    btyp = llvmtype(Bool)
    vbtyp = "<$W x $btyp>"
    abtyp = "[$W x $btyp]"
    typ = llvmtype(T)
    vtyp = "<$W x $typ>"
    atyp = "[$W x $typ]"
    decls = String[]
    instrs = String[]
    push!(instrs, "%cond = trunc $vbtyp %0 to <$W x i1>")
    if T <: FloatingTypes && Base.libllvm_version >= v"9"
        push!(instrs, "%res = select $(fastflags(T)) <$W x i1> %cond, $vtyp %1, $vtyp %2")
    else
        push!(instrs, "%res = select <$W x i1> %cond, $vtyp %1, $vtyp %2")
    end
    push!(instrs, "ret $vtyp %res")
    quote
        $(Expr(:meta, :inline))
        llvmcall(
            $((join(decls, "\n"), join(instrs, "\n"))),
            Vec{$W,$T},
            Tuple{Vec{$W,Bool}, Vec{$W,T}, Vec{$W,$T}},
            v1, v2, v3
        )
    end
end
@inline function vifelse(
    v1::_Vec{W,Bool}, v2::AbstractSIMDVector{W}, v3::AbstractSIMDVector{W}
) where {W}
    vifelse(extract_data(v1), extract_data(v2), extract_data(v3))
end
@inline function vifelse(
    v1::AbstractStructVec{W,Bool}, v2::AbstractStructVec{W,T}, v3::AbstractStructVec{W,T}
) where {W,T}
    SVec(vifelse(extract_data(v1), extract_data(v2), extract_data(v3)))
end
# @inline function vifelse(v1::VecOrProd{W,Bool},
#         v2::VecOrProd{W,T}, v3::VecOrProd{W,T}) where {W,T}
#     vifelse(extract_data(v1), extract_data(v2), extract_data(v3))
# end
# @inline function Base.ifelse(v1::SVec{W,Bool},
#         v2::SVec{W,T}, v3::SVec{W,T}) where {W,T}
#     SVec(vifelse(extract_data(v1), extract_data(v2), extract_data(v3)))
# end

@generated function vifelse(mask::U, v2::_Vec{_W,T}, v3::_Vec{_W,T}) where {_W,T,U<:Unsigned}
    W = _W + 1
    @assert 8sizeof(U) >= W
    btyp = llvmtype(Bool)
    vbtyp = "<$W x $btyp>"
    abtyp = "[$W x $btyp]"
    typ = llvmtype(T)
    vtyp = "<$W x $typ>"
    atyp = "[$W x $typ]"
    decls = String[]
    instrs = String[]

    mtyp_input = llvmtype(U)
    mtyp_trunc = "i$W"

    if mtyp_input == mtyp_trunc
        push!(instrs, "%cond = bitcast $mtyp_input %0 to <$W x i1>")
    else
        push!(instrs, "%condtrunc = trunc $mtyp_input %0 to $mtyp_trunc")
        push!(instrs, "%cond = bitcast $mtyp_trunc %condtrunc to <$W x i1>")
    end
    # push!(instrs, "%cond = trunc $vbtyp %0 to <$W x i1>")
    if T <: FloatingTypes && Base.libllvm_version >= v"9"
        push!(instrs, "%res = select $(fastflags(T)) <$W x i1> %cond, $vtyp %1, $vtyp %2")
    else
        push!(instrs, "%res = select <$W x i1> %cond, $vtyp %1, $vtyp %2")
    end
    push!(instrs, "ret $vtyp %res")
    quote
        $(Expr(:meta, :inline))
        llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            Vec{$W,$T},
            Tuple{$U, Vec{$W,$T}, Vec{$W,$T}},
            mask, v2, v3)
    end
end

@vpromote vifelse 3

@inline vifelse(b::Bool, v2::_Vec{W,T}, s::Union{T,Int}) where {W,T} = ifelse(b, extract_data(v2), vbroadcast(_Vec{W,T}, s))
@inline vifelse(b::Bool, v2::AbstractStructVec{W,T}, s::Union{T,Int}) where {W,T} = SVec(ifelse(b, extract_data(v2), vbroadcast(Vec{W,T}, s)))
@inline vifelse(b::Bool, s::Union{T,Int}, v2::_Vec{W,T}) where {W,T} = ifelse(b, vbroadcast(_Vec{W,T}, s), v2)
@inline vifelse(b::Bool, s::Union{T,Int}, v2::AbstractStructVec{W,T}) where {W,T} = SVec(ifelse(b, vbroadcast(Vec{W,T}, s), extract_data(v2)))

@inline vifelse(U::Unsigned, v2::AbstractSIMDVector, v3::AbstractSIMDVector) = SVec(vifelse(U, extract_data(v2), extract_data(v3)))
@inline vifelse(U::Unsigned, v2::_Vec{W,T}, s::Union{T,Int}) where {W,T} = vifelse(U, extract_data(v2), vbroadcast(_Vec{W,T}, s))
@inline vifelse(U::Unsigned, v2::AbstractStructVec{W,T}, s::Union{T,Int}) where {W,T} = SVec(vifelse(U, extract_data(v2), vbroadcast(Vec{W,T}, s)))
@inline vifelse(U::Unsigned, s::Union{T,Int}, v2::_Vec{W,T}) where {W,T} = vifelse(U, vbroadcast(_Vec{W,T}, s), v2)
@inline vifelse(U::Unsigned, s::Union{T,Int}, v2::AbstractStructVec{W,T}) where {W,T} = SVec(vifelse(U, vbroadcast(Vec{W,T}, s), extract_data(v2)))
@inline vifelse(m::Mask{W}, v1::Vec{W}, v2::Vec{W}) where {W} = vifelse(m.u, v1, v2)

@inline function vifelse(f::F, m::AbstractMask{W}, vargs::Vararg{<:Any,N}) where {F<:Function,W,N}
    vifelse(m, f(vargs...), @inbounds(vargs[N]))
end

@vectordef visodd function Base.isodd(v) where {W,T<:Integer}
    visequal(vand(v, one(T)), one(T))
end

# @inline vand(U::Unsigned, vb::Vec{W,Bool}) where {W} = U & vecbool_to_unsigned(vb)
# @inline vand(vb::Vec{W,Bool}, U::Unsigned) where {W} = U & vecbool_to_unsigned(vb)


using DDR2import

function get_arefile(arefile::String)
    are = DDR2import.Are.read(arefile)
    sectors = Vector{Vector{Array{Float64, 2}}}()
    for sector in are
        append!(sectors, [[sector[2].points]])
    end
    return sectors
end

function add_sectors!(sectors::String, projparam=Equirectangular();
    linecolor=:grey, linewidth=0.01, linealpha=0.33, kwargs...)
    sectors_array = get_arefile(sectors)
    add_polygon!(sectors_array, projparam; linecolor=linecolor,
    linewidth=linewidth, linealpha=linealpha, kwargs...)
end

add_sectors!(sectors::String, mapsettings::MapSettings; kwargs...) =
add_sectors!(sectors, mapsettings.projection; kwargs...)

function add_freeflightairspaces!(sectors::String, projparam=Equirectangular();
    linecolor=:orange, linewidth=0.01, linealpha=0.2, fillcolor=:orange,
    fillalpha=0.1, fillrange=0, kwargs...)
    sectors_array = get_arefile(sectors)
    add_polygon!(sectors_array, projparam; linecolor=linecolor,
    linewidth=linewidth, linealpha=linealpha, fillcolor=fillcolor,
    fillalpha=fillalpha, fillrange=fillrange, kwargs...)
end

add_freeflightairspaces!(sectors::String, mapsettings::MapSettings; kwargs...) =
add_freeflightairspaces!(sectors, mapsettings.projection; kwargs...)


include("../src/NMRSpectraSimulator.jl")
import .NMRSpectraSimulator

using LinearAlgebra
using FFTW

import PlotlyJS
using Plots; plotly()

import OffsetArrays
import Interpolations

import BSON
import JSON

import Random
Random.seed!(25)

save_root_folder = "/home/roy/MEGAsync/outputs/NMR/SH_explore"
isdir(save_root_folder) || mkpath(save_root_folder)

# molecule_names = Vector{String}(undef, 0)
# #root_folders = Vector{String}(undef, 0)
# #project_names = Vector{String}(undef, 0)
#
# push!(molecule_names, "D-(+)-Glucose")
# #push!(root_folders, "/home/roy/MEGAsync/outputs/NMR/experiments/BMRB-500-100mM")
# #push!(project_names, "D-(+)-Glucose")

#compound_select = 2

##############

# get mapping from molecule names to their spin system info json files.
H_params_path = "/home/roy/Documents/repo/NMRData/input/coupling_info"
dict_compound_to_filename = JSON.parsefile("/home/roy/Documents/repo/NMRData/input/compound_mapping/select_compounds.json")

molecule_names = collect( key for (key,val) in dict_compound_to_filename)

notes = collect( "$(i) $(molecule_names[i])" for i = 1:length(molecule_names))
outfile = joinpath(save_root_folder, "compound_numbering.txt")
open(outfile, "w") do f
  for i in notes
    println(f, i)
  end
end

#molecule_names = ["L-Serine"; "Ethanol"; "L-Isoleucine"; "L-Leucine"]
#molecule_names = ["Creatinine"; "Ethanol"]

# # machine settings taken from the BMRB 20mM 700MHz 1D 1H NOSEY experiment for glucose.
# fs = 14005.602240896402
# SW = 20.0041938620844
# ν_0ppm = 10656.011933076665

# machine and 0ppm estimate from the NRC Jan 2022 experiment.
fs = 9615.38461538462
SW = 16.0196918511501
ν_0ppm = 6753.577042707225
α_0ppm = 307479.48631657596
β_0ppm = 3.0967187803295806
λ_0ppm = 1.9474607841303089

# spin-Hamiltonian-related.
tol_coherence = 1e-2 # resonances are pairs of eigenvalues of the Hamiltonian that have quantum numbers that differ by -1. This is the tolerance away from -1 that is allowed.
α_relative_threshold = 0.05 # resonances with relative amplitude less than this factor compared to the maximum resonance in the spin group will be removed. Set to 0.0 to see every single resonance component.
Δc_partition_radius = 0.17 # determines how many resonances get grouped together. Larger number means less groups and thus more resonances per group.
SH_config_path = "/home/roy/Documents/repo/NMRData/input/SH_configs/select_compounds_SH_configs_reduce.json"

# surrogate-related.
# default values. The alternative is to load from a config file.
Δr_default = 1.0 # the samples used to build the surrogate is taken every `Δr` radian on the frequency axis. Decrease for improved accuracy at the expense of computation resources.
Δκ_λ_default = 0.05 # the samples used to build thes urrogate for κ_λ are taken at this sampling spacing. Decrease for improved accuracy at the expense of computation resources.
Δcs_max_scalar_default = 0.2 # In units of ppm. interpolation border that is added to the lowest and highest resonance frequency component of the mixture being simulated.
κ_λ_lb_default = 0.5 # interpolation lower limit for κ_λ.
κ_λ_ub_default = 2.5 # interpolation upper limit for κ_λ.
λ0 = 3.4
surrogate_config_path = "/home/roy/Documents/repo/NMRData/input/surrogate_configs/select_compounds_surrogate_configs.json"

dummy_SSFID = NMRSpectraSimulator.SpinSysParamsType1(0.0) # level 2 model.

##############

hz2ppmfunc = uu->(uu - ν_0ppm)*SW/fs
ppm2hzfunc = pp->(ν_0ppm + pp*fs/SW)

println("Timing: mag equivalence")
@time MEs = NMRSpectraSimulator.getmageqinfomixture(molecule_names,
    H_params_path,
    dict_compound_to_filename;
    unique_cs_atol = 1e-6)

println("Timing: setupmixtureproxies()")
@time mixture_params = NMRSpectraSimulator.setupmixtureSH(molecule_names,
    H_params_path, dict_compound_to_filename, fs, SW,
    ν_0ppm;
    config_path = SH_config_path,
    MEs = MEs,
    tol_coherence = tol_coherence,
    α_relative_threshold = α_relative_threshold,
    Δc_partition_radius = Δc_partition_radius)

############## The rest are for plotting only.

#As = mixture_params[compound_select:compound_select]
As = mixture_params
dummy_SSFID = NMRSpectraSimulator.SpinSysParamsType1(0.0)

## frequency locations. For plotting.
ΩS_ppm = NMRSpectraSimulator.getPsnospininfo(As, hz2ppmfunc)
ΩS_ppm_sorted = sort(NMRSpectraSimulator.combinevectors(ΩS_ppm))

u_offset = 0.5 # in units ppm.
u_min = ppm2hzfunc(ΩS_ppm_sorted[1] - u_offset)
u_max = ppm2hzfunc(ΩS_ppm_sorted[end] + u_offset)

# This is the frequency range that we shall work with.
P = LinRange(hz2ppmfunc(u_min), hz2ppmfunc(u_max), 50000)
U = ppm2hzfunc.(P)
U_rad = U .* (2*π)

# #select_ind = findfirst(xx->xx=="Ethanol", molecule_names)
# select_ind = findfirst(xx->xx=="L-Isoleucine", molecule_names)
# A = As[select_ind]
# ME = MEs[select_ind]
#
# import Destruct
# tmp = collect( NMRSpectraSimulator.createorderingfromeqinds(ME[i], A.N_spins_sys[i]) for i = 1:length(A.N_spins_sys) )
# κs_β_ordering, κs_β_DOF = Destruct.destruct(tmp)
#
# A.Δc_m_compound
# A.part_inds_compound
# c_bar = A.Δc_bar
# c0 = A.Δc_m_compound
#
# # c = NMRSpectraSimulator.reduceΔc(A.Δc_m_compound, ME, A.N_spins_sys)
# # c2 = NMRSpectraSimulator.reduceΔc(A.Δc_m_compound[1], ME[1], A.N_spins_sys[1])
# # sum(c[1]-c2)

#@assert 1==2

include("plot.jl")


#jl
"""
Struct Material
Used for solid Materials in Pipes. So far includes just name, thermal
conductivity, and surface roughness.
"""
struct Material
    name::String
    thermal_conductivity::Function
    roughness::Real
end
struct Pipe{T<:Real}
    length::T
    diameter::T
    radius::T
    flowarea::T
    thickness::T
    outer_diameter::T
    outer_radius::T
    totalarea::T
    roughness::T
    elevation_change::T
    therm_cond::Function
end

const valid_Pipe_diams = [0.125, 0.25, 0.375, 0.5, 0.75, 1.0, 1.25, 1.5, 2.0, 2.5,
                    3.0, 3.5, 4.0, 5.0, 6.0, 8.0, 10.0, 12.0, 16.0]

function Pipe(l, d, t, ϵ, z, kfunc)
    # Assumes a circular Pipe of diameter d
    return Pipe(l,d,d/2,pi*(d/2)^2,t,d+2t,(d+2t)/2,pi*((d+2t)/2)^2,ϵ,z,kfunc)
end

mutable struct PipeProps{T<:Real}
    dᵢ::T
    dₒ::T
    thickness::T
    rᵢ::T
    rₒ::T
    units::String
end

function PipeProps(d, t, units="in")
  di = d - 2t
  ro = d/2
  ri = di/2
  return PipeProps(di, d, t, ri, ro, units)
end

const conversions = Dict(["cm","m"] => 1e-2,
                   ["m","cm"] => 100,
                   ["in","cm"] => 2.54,
                   ["cm","in"] => 1/2.54,
                   ["in","m"] => 2.54e-2,
                   ["m","in"] => 1/2.54e-2,
                   ["m","m"] => 1.0,
                   ["cm","cm"] => 1.0,
                   ["in","in"] => 1.0)

function convert(p::PipeProps, newUnits::String)
    key_ = [p.units, newUnits]
    conversion_factor = conversions[key_]
    p.dᵢ *= conversion_factor
    p.dₒ *= conversion_factor
    p.rᵢ *= conversion_factor
    p.rₒ *= conversion_factor
    p.units = newUnits
end

const schedule40_props = [PipeProps(0.405, 0.068), PipeProps(0.540, 0.088),
                    PipeProps(0.675, 0.091), PipeProps(0.840, 0.109),
                    PipeProps(1.050, 0.113), PipeProps(1.315, 0.133),
                    PipeProps(1.660, 0.140), PipeProps(1.900, 0.145),
                    PipeProps(2.375, 0.154), PipeProps(2.875, 0.203),
                    PipeProps(3.500, 0.216), PipeProps(4.000, 0.226),
                    PipeProps(4.500, 0.237), PipeProps(5.563, 0.258),
                    PipeProps(6.625, 0.280), PipeProps(8.625, 0.322),
                    PipeProps(10.750, 0.365), PipeProps(12.750, 0.406),
                    PipeProps(16.000, 0.500)]
# Note: These outer diameters and wall thicknesses were obtained at this web
# site: http://www.engineeringtoolbox.com/steel-dimensions-d_43.html

for props in schedule40_props
    convert(props, "m")
end

const props40_by_nominal_diameter = Dict(zip(valid_Pipe_diams, schedule40_props))

const schedule_props = Dict("40" => props40_by_nominal_diameter,
                           40 => props40_by_nominal_diameter)

function Pipe(schedule_info::Tuple, mat::Material, l, Δz)
  # schedule_info is a Tuple: the schedule number comes first and the nominal
  # diameter comes second
  props = schedule_props[schedule_info[1]][schedule_info[2]]
  return Pipe(l,props.dᵢ,props.dᵢ/2,pi*(props.rᵢ)^2,props.thickness,props.dₒ,
            props.rₒ,pi*(props.rₒ)^2,mat.roughness,Δz,mat.thermal_conductivity)
end

function Pipe(schedule_info::Tuple, mat::String, l, Δz)
  return Pipe(schedule_info, Pipe_Materials[mat], l, Δz)
end

function colebrook(Re, p::Pipe, f_g)
    if f_g < 0
        return f_g * 1e5
    else
        return 10^(-1 / (2 * √(f_g))) - (p.roughness / (3.7 * p.diameter) + 2.51 / (Re * √(f_g)))
    end
end

function getf(Re, p::Pipe)
    print("Sanity check: Re = ",Re)
    function dummysolve!(f_guess, fvec)
        fvec[1] = colebrook(Re, p, f_guess[1])
    end
    return NLsolve.nlsolve(dummysolve!, [0.001]).zero[1]

end

laminarf(Re, p::Pipe) = 64.0/Re

function swameejain(Re, p::Pipe)
    return 0.25/(log10(p.roughness/(3.7*p.diameter)+ 5.74/(Re^  0.9)))^2
end

function getf2(Re, p::Pipe)
    if Re > 2000
        return swameejain(Re, p)
    else
        return laminarf(Re, p)
    end
end

function stainless_steel_thermal_conductivity(T)
    # Source: http://www.mace.manchester.ac.uk/project/research/structures/strucfire/MaterialInFire/Steel/StainlessSteel/thermalProperties.htm
    return 14.6 + 1.27e-2 *(T-273.15)
end

function copper_thermal_conductivity(T)
    # Source: http://www-ferp.ucsd.edu/LIB/PROPS/PANOS/cu.html
    return 14.6 + 1.27e-2 *(T-273.15)
end

const ss_roughness_by_finish = Dict("2D" => 1E-6,
                              "2B" => 0.5E-6,
                              "2R" => 0.2E-6,
                              "BA" => 0.2E-6,
                              "2BB" => 0.1E-6)
# These surface roughness values are in m and were obtained from the following
# corporate site: http://www.outokumpu.com/en/products-properties/
# more-stainless/stainless-steel-surface-finishes/cold-rolled-finishes/Pages/default.aspx
# Note that the given values are the maximum roughness values in the range on
# the website, except for finish 2BB, which didn't have a roughness value
# so I took the lowest value for finish 2B since the compnay says this:
# "Due to the fact that surface roughness of the finish 2BB is lower than that
# of 2B, some modifications on lubrication during forming might be needed."

const ss_mat = Material("Stainless Steel", stainless_steel_thermal_conductivity,
                  ss_roughness_by_finish["2B"])

const cu_roughness = 0.03e-3
# Copper roughness.
# Source: http://www.pressure-drop.com/Online-Calculator/rauh.html

const cu_mat = Material("Copper", copper_thermal_conductivity, cu_roughness)

const pipematerials = Dict("Copper" => cu_mat, "Stainless Steel"=>ss_mat)


#####
##### 'Tupled closure' implementation: 1-tuple, 2-tuple, and then n-tuple by induction
#####

#####
##### Stress divergences
#####

for stress_div in (:∂ⱼ_τ₁ⱼ, :∂ⱼ_τ₂ⱼ, :∂ⱼ_τ₃ⱼ)
    @eval begin
        @inline $stress_div(i, j, k, grid::AbstractGrid, closures::Tuple{<:Any}, clock, U, Ks, args...) =
                    $stress_div(i, j, k, grid, closures[1], clock, U, Ks[1], args...)

        @inline $stress_div(i, j, k, grid::AbstractGrid, closures::Tuple{<:Any, <:Any}, clock, U, Ks, args...) = (
                    $stress_div(i, j, k, grid, closures[1], clock, U, Ks[1], args...)
                  + $stress_div(i, j, k, grid, closures[2], clock, U, Ks[2], args...))

        @inline $stress_div(i, j, k, grid::AbstractGrid, closures::Tuple{<:Any, <:Any, <:Any}, clock, U, Ks, args...) = (
                    $stress_div(i, j, k, grid, closures[1], clock, U, Ks[1], args...)
                  + $stress_div(i, j, k, grid, closures[2], clock, U, Ks[2], args...)
                  + $stress_div(i, j, k, grid, closures[3], clock, U, Ks[3], args...))

        @inline $stress_div(i, j, k, grid::AbstractGrid, closures::Tuple, clock, U, Ks, args...) = (
                    $stress_div(i, j, k, grid, closures[1:2], clock, U, Ks[1:2], args...)
                  + $stress_div(i, j, k, grid, closures[3:end], clock, U, Ks[3:end], args...))
    end
end

#####
##### Tracer flux divergences
#####

@inline ∇_dot_qᶜ(i, j, k, grid::AbstractGrid, closures::Tuple{<:Any}, c, iᶜ, clock, Ks, args...) =
        ∇_dot_qᶜ(i, j, k, grid, closures[1], c, iᶜ, clock, Ks[1], args...)

@inline ∇_dot_qᶜ(i, j, k, grid::AbstractGrid, closures::Tuple{<:Any, <:Any}, c, iᶜ, clock, Ks, args...) = (
        ∇_dot_qᶜ(i, j, k, grid, closures[1], c, iᶜ, clock, Ks[1], args...)
      + ∇_dot_qᶜ(i, j, k, grid, closures[2], c, iᶜ, clock, Ks[2], args...))

@inline ∇_dot_qᶜ(i, j, k, grid::AbstractGrid, closures::Tuple{<:Any, <:Any, <:Any}, c, iᶜ, clock, Ks, args...) = (
        ∇_dot_qᶜ(i, j, k, grid, closures[1], c, iᶜ, clock, Ks[1], args...)
      + ∇_dot_qᶜ(i, j, k, grid, closures[2], c, iᶜ, clock, Ks[2], args...)
      + ∇_dot_qᶜ(i, j, k, grid, closures[3], c, iᶜ, clock, Ks[3], args...))

@inline ∇_dot_qᶜ(i, j, k, grid::AbstractGrid, closures::Tuple, c, iᶜ, clock, Ks, args...) = (
        ∇_dot_qᶜ(i, j, k, grid, closures[1:2], c, iᶜ, clock, Ks[1:2], args...)
      + ∇_dot_qᶜ(i, j, k, grid, closures[3:end], c, iᶜ, clock, Ks[3:end], args...))

#####
##### Utilities
#####

with_tracers(tracers, closure_tuple::Tuple) =
    Tuple(with_tracers(tracers, closure) for closure in closure_tuple)

function calculate_diffusivities!(diffusivity_fields_tuple, closure_tuple::Tuple, args...)
    for (α, closure) in enumerate(closure_tuple)
        @inbounds diffusivity_fields = diffusivity_fields_tuple[α]
        calculate_diffusivities!(diffusivity_fields, closure, args...)
    end
    return nothing
end

#####
##### Support for VerticallyImplicitTimeDiscretization
#####

const EC = AbstractTurbulenceClosure{<:ExplicitTimeDiscretization}
const VIC = AbstractTurbulenceClosure{<:VerticallyImplicitTimeDiscretization}

# Filter explicitly-discretized closures.
@inline z_diffusivity(clo::Tuple{<:EC},        iᶜ, Ks, args...) = tuple(0)
@inline z_diffusivity(clo::Tuple{<:VIC},       iᶜ, Ks, args...) = tuple(z_diffusivity(clo[1], iᶜ, Ks[1], args...))
@inline z_diffusivity(clo::Tuple{<:VIC, <:EC}, iᶜ, Ks, args...) = tuple(z_diffusivity(clo[1], iᶜ, Ks[1], args...))
@inline z_diffusivity(clo::Tuple{<:EC, <:VIC}, iᶜ, Ks, args...) = tuple(z_diffusivity(clo[2], iᶜ, Ks[2], args...))

@inline z_diffusivity(clo::Tuple{<:VIC, <:VIC}, iᶜ, Ks, args...) = tuple(z_diffusivity(clo[1], iᶜ, Ks[1], args...),
                                                                         z_diffusivity(clo[2], iᶜ, Ks[2], args...))

@inline z_diffusivity(clo::Tuple, iᶜ, Ks, args...) = tuple(z_diffusivity(clo[1:2],   iᶜ, Ks[1:2],   args...)...,
                                                           z_diffusivity(clo[3:end], iᶜ, Ks[3:end], args...)...)

@inline z_viscosity(clo::Tuple{<:EC},         Ks, args...) = tuple(0)
@inline z_viscosity(clo::Tuple{<:VIC},        Ks, args...) = tuple(z_viscosity(clo[1], Ks[1], args...))
@inline z_viscosity(clo::Tuple{<:VIC, <:EC},  Ks, args...) = tuple(z_viscosity(clo[1], Ks[1], args...))
@inline z_viscosity(clo::Tuple{<:EC, <:VIC},  Ks, args...) = tuple(z_viscosity(clo[2], Ks[2], args...))

@inline z_viscosity(clo::Tuple{<:VIC, <:VIC}, Ks, args...) = tuple(z_viscosity(clo[1], Ks[1], args...),
                                                                   z_viscosity(clo[2], Ks[2], args...))

@inline z_viscosity(clo::Tuple, Ks, args...) = tuple(z_viscosity(clo[1:2],   Ks[1:2], args...)...,
                                                     z_viscosity(clo[3:end], Ks[3:end], args...)...)

for coeff in (:νᶜᶜᶜ, :νᶠᶠᶜ, :νᶠᶜᶠ, :νᶜᶠᶠ, :κᶜᶜᶠ, :κᶜᶠᶜ, :κᶠᶜᶜ)
    @eval begin
        @inline $coeff(i, j, k, grid, clock, ν::Tuple{C1})     where C1       = $coeff(i, j, k, grid, clock, ν[1])
        @inline $coeff(i, j, k, grid, clock, ν::Tuple{C1, C2}) where {C1, C2} = $coeff(i, j, k, grid, clock, ν[1])   + $coeff(i, j, k, grid, clock, ν[2])
        @inline $coeff(i, j, k, grid, clock, ν::Tuple)                        = $coeff(i, j, k, grid, clock, ν[1:2]) + $coeff(i, j, k, grid, clock, ν[3:end])
    end
end

const ImplicitClosure = AbstractTurbulenceClosure{TD} where TD <: VerticallyImplicitTimeDiscretization
const ExplicitOrNothing = Union{ExplicitTimeDiscretization, Nothing}

@inline combine_time_discretizations(disc) = disc

@inline combine_time_discretizations(::ExplicitTimeDiscretization, ::VerticallyImplicitTimeDiscretization)           = VerticallyImplicitTimeDiscretization()
@inline combine_time_discretizations(::VerticallyImplicitTimeDiscretization, ::ExplicitTimeDiscretization)           = VerticallyImplicitTimeDiscretization()
@inline combine_time_discretizations(::VerticallyImplicitTimeDiscretization, ::VerticallyImplicitTimeDiscretization) = VerticallyImplicitTimeDiscretization()
@inline combine_time_discretizations(::ExplicitTimeDiscretization, ::ExplicitTimeDiscretization)                     = ExplicitTimeDiscretization()

@inline combine_time_discretizations(disc1, disc2, other_discs...) =
    combine_time_discretizations(combine_time_discretizations(disc1, disc2), other_discs...)

@inline time_discretization(closures::Tuple) = combine_time_discretizations(time_discretization.(closures)...)