import JSON
import Combinatorics

dirs = readdir("input")

function combine_modnames(modnames, div)
    output = ""
    for name in modnames
        output = "$output$div$name"
    end
    output[(length(div)+1):end]
end

struct ModMeta
    abbrev::String
    name::String
    priority::Int
    unit_list::Set{String}
    comm_list::Set{String}
end

mods = Vector{ModMeta}()

for dir_end in dirs
    dir = joinpath("input", dir_end)
    if !isdir(dir)
        continue
    end

    metafile = joinpath(dir, "meta.json")
    if !isfile(metafile)
        println("meta.json not found in ", dir)
        continue
    end
    meta = JSON.parsefile(metafile)
    req_fields = ["name", "priority"]
    for field in req_fields
        if !(field in keys(meta))
            println("meta.json is missing \"$field\" field")
        end
    end
    unit_list_path = joinpath(dir, "unit_list.json")
    comm_list_path = joinpath(dir, "commander_list.json")

    abbrev = dir_end
    name = meta["name"]
    priority = meta["priority"]
    unit_list = Set{String}(
        if isfile(unit_list_path)
            JSON.parsefile(unit_list_path)["units"]
        else
            []
        end
    )
    comm_list = Set{String}(
        if isfile(comm_list_path)
            JSON.parsefile(comm_list_path)["commanders"]
        else
            []
        end
    )

    push!(mods, ModMeta(
        abbrev,
        name,
        priority,
        unit_list,
        comm_list
    ))
end

for combo in Combinatorics.powerset(mods, 2)
    abbrevs = getproperty.(combo, :abbrev)

    filename = "compatch_$(combine_modnames(abbrevs, "+"))"
    folder = joinpath("output", filename)
    mkpath(folder)

    modinfo = JSON.parsefile("modinfo_template.json")
    modinfo["display_name"] = "DAEDALUS ComPatch - $(combine_modnames(abbrevs, " + "))"
    modinfo["identifier"] = "com.pa.daedalus.compatch_$(combine_modnames(abbrevs, "-"))"
    modinfo["description"] = "Compatibility Patch for the following unit-adding mods: $(abbrevs...)"
    modinfo["priority"] = minimum(getproperty.(combo, :priority)) - 1
    write(joinpath(folder, "modinfo.json"), JSON.json(modinfo))

    unit_lists = getproperty.(combo, :unit_list)
    if sum([!isempty(i) for i in unit_lists]) > 1
        new_unit_list = Dict("units" => collect(union(unit_lists...)))
        unit_path = joinpath(folder, "pa/units")
        mkpath(unit_path)
        write(joinpath(unit_path, "unit_list.json"), JSON.json(new_unit_list))
    end

    comm_lists = getproperty.(combo, :comm_list)
    if sum([!isempty(i) for i in comm_lists]) > 1
        new_comm_list = Dict("commanders" => collect(union(comm_lists...)))
        comm_path = joinpath(folder, "pa/units/commanders")
        mkpath(comm_path)
        write(joinpath(comm_path, "unit_list.json"), JSON.json(new_comm_list))
    end
end


using InvariantCausal
using DelimitedFiles
using Test

include(joinpath(@__DIR__, "test_search.jl"))
include(joinpath(@__DIR__, "test_regression.jl"))

X = readdlm(joinpath(@__DIR__, "X1.dat"))
env = Vector{Int}(X[:,1])
X = X[:,2:end]
S = 1:size(X,2)

@time @testset "causal search" begin
    result = map(i -> causalSearch(X, X[:, i], env, setdiff(S,i) , α=0.01), S)
    @test result[2].S == [5]
    @test result[3].S == [5]
    @test result[7].S == [4] || result[7].S == []
    for i in [1, 4, 5, 6]
        @test result[i].model_reject == true
    end
end

@time @testset "causal search with selection_only" begin
    result = map(i -> causalSearch(X, X[:, i], env, setdiff(S,i), α=0.01, selection_only=true), S)
    @test result[2].S == [5]
    @test result[3].S == [5]
    @test result[7].S == [4] || result[7].S == []
    for i in [1, 4, 5, 6]
        @test result[i].model_reject == true
    end
end

@time @testset "causal search with limited # of true causes" begin
    result = map(i -> causalSearch(X, X[:, i], env, setdiff(S,i), α=0.01, max_num_true_causes=3), S)
    @test result[2].S == [5]
    @test result[3].S == [5]
    @test result[7].S == [4] || result[7].S == []
    for i in [1, 4, 5, 6]
        @test result[i].model_reject == true
    end
end

@time @testset "causal search with subsampling" begin
    result = map(i -> causalSearch(X, X[:, i], env, setdiff(S,i), α=0.01, n_max_for_exact=90), S)
    @test result[2].S == [5]
    @test result[3].S == [5] || result[3].S == []
    @test result[7].S == [4] || result[7].S == []
    for i in [1, 4, 5, 6]
        @test result[i].model_reject == true
    end
end

function generate_setting(setting_configs; n_environments=2)
    n_int = rand(setting_configs["n_int"])
    n_obs = rand(setting_configs["n_obs"])
    p = rand(setting_configs["p"])
    k = rand(setting_configs["k"])
    lb_obs = rand(setting_configs["lb_obs"])
    ub_obs = lb_obs + rand(setting_configs["ub_lb_delta_obs"])
    _var_1 = rand(setting_configs["err_var_min"])
    _var_2 = rand(setting_configs["err_var_min"])
    err_var_min = min(_var_1, _var_2)
    err_var_max = max(_var_1, _var_2)
    noise_multiplier_min = rand(setting_configs["noise_multiplier_min"])
    if rand() < setting_configs["prob_fixed_noise_multiplier"]
        noise_multiplier_max = noise_multiplier_min
    else
        noise_multiplier_max = noise_multiplier_min + rand(setting_configs["noise_multiplier_max_min_delta"])
    end
    _l = rand(setting_configs["lb_int"])
    _u = rand(setting_configs["ub_int"])
    lb_int = min(_l, _u)
    ub_int = max(_l, _u)
    if rand() < setting_configs["prob_int_single"]
        n_int_variables = 1
        frac_int_variables = 1 / p
    else
        frac_int_variables = rand(setting_configs["frac_int_variables"])
        n_int_variables = floor(Int64, p * frac_int_variables)
    end
    prob_coefficient_unchanged = setting_configs["prob_coefficient_unchanged"]
    target = rand(1:p)
    SEMs = Vector{GaussianSEM}(undef, n_environments)
    env = repeat([1], inner=n_obs)
    intervened_variables = Vector{Vector{Int64}}(undef, n_environments)
    intervened_variables[1] = []
    SEMs[1] = random_gaussian_SEM(p, k,
              lb=lb_obs, ub=ub_obs, var_min=err_var_min, var_max=err_var_max)
    true_β = SEMs[1].B[target, setdiff(1:p, target)]
    true_causes = (1:(p-1))[true_β .!= 0]
    for j in 2:n_environments
        SEMs[j], intervened_variables[j] = random_noise_intervened_SEM(SEMs[1], p_intervened=n_int_variables,
                noise_multiplier_min=noise_multiplier_min, noise_multiplier_max=noise_multiplier_max,
                avoid=[target], prob_coeff_unchanged=prob_coefficient_unchanged,
                lb=lb_int, ub=ub_int)
        env = vcat(env, repeat([j], inner=n_int))
    end
    return Dict(
        "n_environments" => n_environments,
        "env" => env,
        "SEMs" => SEMs,
        "intervened_variables" => intervened_variables,
        "target" => target,
        "true_causes" => true_causes,
        "true_β" => true_β,
        "n_int" => n_int,
        "n_obs" => n_obs,
        "p" => p,
        "k" => k,
        "lb_obs" => lb_obs,
        "ub_obs" => ub_obs,
        "err_var_min" => err_var_min,
        "err_var_max" => err_var_max,
        "noise_multiplier_min" => noise_multiplier_min,
        "noise_multiplier_max" => noise_multiplier_max,
        "lb_int" => lb_int,
        "ub_int" => ub_int,
        "prob_coefficient_unchanged" => prob_coefficient_unchanged,
        "n_int_variables" => n_int_variables,
        "frac_int_variables" => frac_int_variables
    )
end

@time @testset "random instances" begin
    setting_configs = Dict(
    "n_int" => [100, 200, 300, 400, 500],
    "n_obs" => [100, 200, 300, 400, 500],
    "p" =>  collect(5:9),
    "k" => collect(1:4),
    "lb_obs" => collect(0.1:0.1:2),
    "ub_lb_delta_obs" => collect(0.1:0.1:1),
    "err_var_min" => collect(0.1:0.1:2),
    "err_var_max" => collect(0.1:0.1:2),
    "noise_multiplier_min" => collect(0.1:0.1:4),
    "noise_multiplier_max_min_delta" => collect(0.1:0.1:2),
    "lb_int" => collect(0.1:0.1:2),
    "ub_int" => collect(0.1:0.1:2),
    "prob_fixed_noise_multiplier" => 1/3,
    "prob_coefficient_unchanged" => 2/3,
    "prob_int_single" => 1/6,
    "frac_int_variables" => 1. ./ collect(1.1:0.1:3)
    )
    for trial in 1:10
        for n_environments in [2, 3, 5]
            setting = generate_setting(setting_configs, n_environments=n_environments)
            env = setting["env"]
            p = setting["p"]
            target = setting["target"]
            true_β = setting["true_β"]
            true_causes = setting["true_causes"]
            intervened_variables = setting["intervened_variables"]
            SEM_obs = setting["SEMs"][1]
            X = simulate(SEM_obs, setting["n_obs"])
            for j in 2:n_environments
                _X = simulate(setting["SEMs"][j], setting["n_int"])
                X = vcat(X, _X)
            end
            y = X[:, target]
            X = X[:, setdiff(1:p, target)]
            result = causalSearch(X, y, env, verbose=false)
            println("S = $(result.S), truth = $true_causes")
            @test issubset(result.S, true_causes)
            for k in result.S
                @test result.confint[k, 1] < true_β[k] < result.confint[k, 2]
            end
        end
    end
end

@time @testset "random instances, many variables with screening" begin
    setting_configs = Dict(
    "n_int" => [100, 200, 300, 400, 500],
    "n_obs" => [100, 200, 300, 400, 500],
    "p" =>  collect(20:50),
    "k" => collect(1:3),
    "lb_obs" => collect(0.1:0.1:2),
    "ub_lb_delta_obs" => collect(0.1:0.1:1),
    "err_var_min" => collect(0.1:0.1:2),
    "err_var_max" => collect(0.1:0.1:2),
    "noise_multiplier_min" => collect(0.1:0.1:4),
    "noise_multiplier_max_min_delta" => collect(0.1:0.1:2),
    "lb_int" => collect(0.1:0.1:2),
    "ub_int" => collect(0.1:0.1:2),
    "prob_fixed_noise_multiplier" => 1/3,
    "prob_coefficient_unchanged" => 2/3,
    "prob_int_single" => 1/6,
    "frac_int_variables" => 1. ./ collect(1.1:0.1:3)
    )
    for trial in 1:10
        for n_environments in [2, 3, 5]
            setting = generate_setting(setting_configs, n_environments=n_environments)
            env = setting["env"]
            p = setting["p"]
            target = setting["target"]
            true_β = setting["true_β"]
            true_causes = setting["true_causes"]
            intervened_variables = setting["intervened_variables"]
            SEM_obs = setting["SEMs"][1]
            X = simulate(SEM_obs, setting["n_obs"])
            for j in 2:n_environments
                _X = simulate(setting["SEMs"][j], setting["n_int"])
                X = vcat(X, _X)
            end
            y = X[:, target]
            X = X[:, setdiff(1:p, target)]
            result = causalSearch(X, y, env, verbose=false, p_max=8)
            # performance only guaranteed if lasso selected all the causes
            println("S = $(result.S), truth = $true_causes")
            if issubset(true_causes, result.variables_considered)
                @test issubset(result.S, true_causes)
            end
        end
    end
end


using IRTools.All
using IRTools: Pipe
using Base: tail

ntail(x, n) = n <= 0 ? x : xcall(:tail, ntail(x, n-1))

function instrument!(pr, v, st)
  ex = st.expr
  if isexpr(ex, :new)
    st = stmt(st, expr = xcall(Zygote, :__new__, ex.args...))
    pr[v] = st
  elseif isexpr(ex, :splatnew)
    st = stmt(st, expr = xcall(Zygote, :__splatnew__, ex.args...))
    pr[v] = st
  end
  return st
end

function dual(ir)
  args = copy(arguments(ir))
  dx = argument!(ir, at = 1)
  Δs = Dict()
  for bl in blocks(ir)[2:end], arg in copy(arguments(bl))
    Δs[arg] = argument!(bl, insert = false)
  end
  pr = Pipe(ir)
  partial(x::Variable) = Δs[x]
  partial(x) = push!(pr, xcall(Forward, :zerolike, x))
  partial(v, x::Variable) = Δs[x]
  partial(v, x) = insert!(pr, v, xcall(Forward, :zerolike, x))
  for (i, x) in enumerate(args)
    if i == length(args) && ir.meta.method.isva
      Δs[x] = push!(pr, ntail(dx, i-1))
    else
      Δs[x] = push!(pr, xcall(:getindex, dx, i))
    end
  end
  branches(pr) do br
    args = arguments(br)
    if isreturn(br)
      args[1] = push!(pr, xcall(:tuple, args[1], partial(args[1])))
    else
      for arg in copy(args)
        push!(args, partial(arg))
      end
    end
    br
  end
  for (v, st) in pr
    st = instrument!(pr, v, st)
    if isexpr(st.expr, :meta, :inbounds, :loopinfo)
      Δs[v] = nothing
    elseif isexpr(st.expr, :boundscheck) ||
           (isexpr(st.expr, :call) && st.expr.args[1] == GlobalRef(Base, :not_int)) ||
           (isexpr(st.expr, :call) && st.expr.args[1] == GlobalRef(Core, :(===))) ||
           (isexpr(st.expr, :call) && st.expr.args[1] == GlobalRef(Main, :(===)))
      Δs[v] = false
    elseif isexpr(st.expr, :call)
      dargs = insert!(pr, v, xcall(:tuple, partial.((v,), st.expr.args)...))
      result = insert!(pr, v, stmt(st, expr = xcall(Forward, :_pushforward, dargs, st.expr.args...)))
      pr[v] = xcall(:getindex, result, 1)
      Δs[v] = push!(pr, xcall(:getindex, result, 2))
    elseif !isexpr(st.expr)
      Δs[v] = push!(pr, xcall(Forward, :zerolike, v))
    else
      error("Unsupported $(st.expr.head) expression")
    end
  end
  ir = finish(pr)
  return ir
end

@dynamo function _pushforward(_, x...)
  ir = IR(x...)
  ir === nothing && return :(error("non-differentiable function $(args[2])"))
  ir = Zygote.instrument(ir)
  ir.meta.code.inlineable = true
  return dual(ir)
end


mutable struct Config{T <: AbstractFloat}
    folds::Integer
    epochs::Integer
    bs::Integer
    lr::T
end

show(io::IO,c::Config) = print(io, "A trainer config.\n(epochs $(c.epochs), batch size $(c.bs), learn rate $(c.lr)).\nThere are $(c.folds) folds.")

mutable struct Status
    epoch::Integer
    loss::Dict{String,Array{Float64}}
end

function reset_model!(model;init=Flux.glorot_uniform)
    Flux.loadparams!(model, map(x -> init(size(x)), Flux.params(model)))
    model
end

struct Trainer
    model
    loss::Function
    data::Dataset
    folds::Int
    epochs::Int
    bs::Int
end

Trainer(model,loss,data;folds=5,epochs=100,batch_size=16) = Trainer(model, loss, data, folds, epochs, batch_size)

function training(trainer::Trainer, opt;log_dir::String="./result",use_gpu=true)
    RD = trainer.data.associations
    knowns = findall(RD .== 1)
    unknowns = findall(RD .== 0)
    shuffle!(knowns)
    shuffle!(unknowns)
    selected = eachindex(knowns)
    cv = CrossValidation(trainer.folds, selected)
    ps = params(trainer.model)
    plot_train = plot()
    plot_validate = plot()
    isdir("./result") || mkdir("./result")
    writedlm("./result/RD.csv", RD, ',')
    writedlm("./result/drugs.csv", trainer.data.drugs, ',')
    writedlm("./result/diseases.csv", trainer.data.diseases, ',')
    for f in 1:trainer.folds
        @info "fold $f"
        rd = copy(RD)
        loss_train = []
        loss_val = []
        train, test = cv(f)
        rd[knowns[test]] .= false
        mask = falses(size(RD))
        mask[knowns[test]] .= true
        useful = map(>=(1), sum(mask, dims=2) |> vec)
        mask = trues(size(RD))
        mask[knowns[train]] .= false
        mask[unknowns[train]] .= false
        reset_model!(trainer.model)
        samples = Samples(trainer.data, rd, RD, findall(!, mask);use_gpu=use_gpu,mode=:both)
        @info "training..."
        for e in 1:trainer.epochs
            loss_epoch = []
            val_epoch = []
            change_indices!(samples, findall(!, mask))
            loader = Flux.Data.DataLoader(samples, batchsize=trainer.bs)
            for (i, (x1, x2), y) in loader
                local loss
                gs = gradient(ps) do
                    ŷ = trainer.model(x1, x2)
                    loss = trainer.loss(ŷ, y)
                end
                Flux.update!(opt, ps, gs)
                push!(loss_epoch, loss)
            end
            push!(loss_train, mean(loss_epoch))
            validate = map(>=(4), sum(mask, dims=2) |> vec)
            change_indices!(samples, findall(mask .& validate))
            loader = Flux.Data.DataLoader(samples, batchsize=512)
            for (i, (x1, x2), y) in loader
                ŷ = trainer.model(x1, x2)
                loss = trainer.loss(ŷ, y)
                push!(val_epoch, loss)
            end
            push!(loss_val, mean(val_epoch))
        end
        plot!(plot_train, loss_train, label="fold $f")
        plot!(plot_validate, loss_val, label="fold $f")
        plot(plot_train,plot_validate,layout=(2,1))|>display
        @info "testing..."
        @info "$(sum(useful)) drugs to test."
        mask = mask .& useful
        @info "total $(sum(mask)) drug-disease pairs to test."
        change_indices!(samples, findall(mask))
        loader = Flux.Data.DataLoader(samples, batchsize=1024)
        result = zeros(size(RD))
        for (i, x, y) in loader
            ŷ = trainer.model(x...)
            ȳ = cpu(ŷ)
            result[i] .= softmax(ȳ)[2,:]
        end
        h = heatmap(result[useful,:])
        yticks!(h, [1:length(useful);], trainer.data.drugs[useful])
        display(h)
        isdir("./result/fold_$f") || mkdir("./result/fold_$f")
        writedlm("./result/fold_$f/result.csv", result, ',')
        writedlm("./result/fold_$f/mask.csv", mask, ',')
    end
end

export Config,Status,reset_model!,Trainer,training


# utility functions needed for discerning topology from geometry

using Combinatorics

using dataset_pb2

export bohr_to_angstroms, angstroms_to_bohr, distance_between_atoms, bonded, distances
export add_atom!, add_bond!
export smu_atype_to_atomic_number, number_to_smu_btype, smu_btype_to_number, smu_atom_type_to_max_con
export atomic_number_to_smu_atom_type
export add_atom_by_atomic_number!, add_atom_by_type!
export distances, bonded
export canonical_bond_topology!, same_bond_topology
export is_single_fragment

function bohr_to_angstroms(distance)
  return distance * 0.529177249
end

function angstroms_to_bohr(distance)
  return distance / 0.529177249
end

function distance_between_atoms(geom::dataset_pb2.Geometry, a1, a2)::Float32
  return bohr_to_angstroms(sqrt(
    (geom.atom_positions[a1].x - geom.atom_positions[a2].x) *
    (geom.atom_positions[a1].x - geom.atom_positions[a2].x) +
    (geom.atom_positions[a1].y - geom.atom_positions[a2].y) *
    (geom.atom_positions[a1].y - geom.atom_positions[a2].y) +
    (geom.atom_positions[a1].z - geom.atom_positions[a2].z) *
    (geom.atom_positions[a1].z - geom.atom_positions[a2].z)
  ))
end

_smu_atype_to_atomic_number = Dict{Int32, Int32}(
  BondTopology_AtomType.ATOM_H => 1,
  BondTopology_AtomType.ATOM_C => 6,
  BondTopology_AtomType.ATOM_N => 7,
  BondTopology_AtomType.ATOM_NPOS => 7,
  BondTopology_AtomType.ATOM_O => 8,
  BondTopology_AtomType.ATOM_ONEG => 8,
  BondTopology_AtomType.ATOM_F => 9
)

function smu_atype_to_atomic_number(smu_type::T)::Int32 where {T<:Integer}
  return _smu_atype_to_atomic_number[smu_type]
end

"""Given an atomic number and charge, return the BondTopology_AtomType
Args:
  atomic_number:
  charge:
Returns:
  BondTopology_AtomType
"""
function atomic_number_to_smu_atom_type(atomic_number::Integer, charge::Integer=0)::Int32
  if atomic_number == 6
    return BondTopology_AtomType.ATOM_C
  elseif atomic_number == 7
    if charge == 0
      return BondTopology_AtomType.ATOM_N
    else
      return BondTopology_AtomType.ATOM_NPOS
    end
  elseif atomic_number == 8
    if charge == 0
      return BondTopology_AtomType.ATOM_O
    else
      return BondTopology_AtomType.ATOM_ONEG
    end
  elseif atomic_number == 9
    return BondTopology_AtomType.ATOM_F
  elseif atomic_number == 1
    return BondTopology_AtomType.ATOM_H
  else
    @error("Unrecognized atomic number $(atomic_number)")
  end
end

_number_to_smu_btype = Dict{Int32, Int32}(
  0 => BondTopology_BondType.BOND_UNDEFINED,
  1 => BondTopology_BondType.BOND_SINGLE,
  2 => BondTopology_BondType.BOND_DOUBLE,
  3 => BondTopology_BondType.BOND_TRIPLE
)

function number_to_smu_btype(n::T)::Int32 where {T<:Integer}
  return _number_to_smu_btype[n]
end

_smu_btype_to_number = Dict{Int32, Int32}(
  BondTopology_BondType.BOND_UNDEFINED => 0,
  BondTopology_BondType.BOND_SINGLE => 1,
  BondTopology_BondType.BOND_DOUBLE => 2,
  BondTopology_BondType.BOND_TRIPLE => 3
)

_smu_atom_type_to_max_con = Dict{Int32, Int32}(
  BondTopology_AtomType.ATOM_H => 1,
  BondTopology_AtomType.ATOM_C => 4,
  BondTopology_AtomType.ATOM_N => 3,
  BondTopology_AtomType.ATOM_NPOS => 3,
  BondTopology_AtomType.ATOM_O => 2,
  BondTopology_AtomType.ATOM_ONEG => 1,
  BondTopology_AtomType.ATOM_F => 1
)

# Return the maximum number of bonds for `smu_atype`.
# Unfortunately the type declaration on the next line does not work.
# function smu_atom_type_to_max_con(smu_atype::BondTopology_AtomType)::Int32

function smu_atom_type_to_max_con(smu_atype::T)::Int32 where {T <:Integer}
  return _smu_atom_type_to_max_con[smu_atype]
end

function smu_btype_to_number(smu_btype::T)::Int32 where {T<:Integer}
  return _smu_btype_to_number[smu_btype]
end

function add_bond!(a1::T, a2::T, btype, bond_topology::dataset_pb2.BondTopology) where {T<:Integer}
  smu_btype = number_to_smu_btype(btype)

  new_bond = BondTopology_Bond(atom_a=a1, atom_b=a2, bond_type=smu_btype)

  if !hasproperty(bond_topology, :bonds)
    setproperty!(bond_topology, :bonds, [new_bond])
  else
    push!(bond_topology.bonds, new_bond)
  end
end

function add_atom_by_atomic_number!(bond_topology::dataset_pb2.BondTopology, atomic_number::T, charge::T=0) where {T<:Integer}
  smu_atype = atomic_number_to_smu_atom_type(atomic_number, charge)
  add_atom_by_type!(bond_topology, smu_atype)
end

function add_atom_by_type!(bond_topology::dataset_pb2.BondTopology, smu_atype::T) where {T <:Integer}
  if !hasproperty(bond_topology, :atoms)
    setproperty!(bond_topology, :atoms, [smu_atype])
  else
    push!(bond_topology.atoms, smu_atype)
  end
end

"""Return an int array of the bonded atoms in `bond_topology`.
The returned array is square[natoms,natoms]. If there is no bond
between the atoms, the value is zero. If the atoms are bonded,
the entry will be the bond type.
Args:
  bond_topology:
Returns:
  Vector of BondType's
"""
function bonded(bond_topology::dataset_pb2.BondTopology)::Array{Int32, 2}
  natoms = length(bond_topology.atoms)
  connected = zeros(Int32, (natoms, natoms))  # to be returned
  # Note need to convert to 1 indexing for Julia
  for bond in bond_topology.bonds
    a1 = bond.atom_a + 1
    a2 = bond.atom_b + 1
    connected[a1, a2] = connected[a2, a1] = bond.bond_type
  end
  return connected
end

"""Given a connection matrix `bonded`, return a vector of
vectors, one for each atom, containing the indices of the
bonded atoms.
Args:
  bonded:
Returns:
"""
function connections(bonded::Array{Int32, 2})::Vector{Vector{Int32}}
  natoms = size(bonded, 1)
  result = Vector{Vector{Int32}}(undef, (natoms))
  for i in 1:natoms
    nbrs = Vector{Int32}()
    for j in 1:natoms
      if bonded[i,j] > 0
        push!(nbrs, j)
      end
    end
    result[i] = nbrs
  end
  return result
end

"""Return a float array of the interatomic distances in `geometry`.
Args:
  geometry:
Returns:
  a numpy array of distances
"""
function distances(geometry)::Array{Float32, 2}
  natoms = length(geometry.atom_positions)
  result = zeros(Float32, (natoms, natoms))
  for atoms in combinations(1:natoms, 2)
    i = atoms[1]
    j = atoms[2]
    result[i, j] = result[j, i] = distance_between_atoms(geometry, i, j)
  end
  return result
end

"""Transform the bonds attribute of `bond_topology` to a canonical form.

Args:
  bond_topology:
Returns:
  BondTopology
"""
function canonical_bond_topology!(bond_topology)
  if length(bond_topology.bonds) < 2
    return
  end

  for bond in bond_topology.bonds
    if bond.atom_a > bond.atom_b
      bond.atom_a, bond.atom_b = bond.atom_b, bond.atom_a
    end
  end

  sort!(bond_topology.bonds, by = b -> (b.atom_a, b.atom_b))
end

Base.:(==)(b1::BondTopology_Bond, b2::BondTopology_Bond)::Bool = b1.atom_a == b2.atom_a &&
                b1.atom_b == b2.atom_b && b1.bond_type == b2.bond_type

"""Return True if bt1 == bt2.
Note that there is no attempt to canonialise the protos.
Args:
Returns:
"""
function same_bond_topology(bt1, bt2)::Bool
  length(bt2.atoms) == length(bt2.atoms) || return false
  length(bt2.bonds) == length(bt2.bonds) || return false

  # https://stackoverflow.com/questions/47564825/check-if-all-the-elements-of-a-julia-array-are-equal
  # all(y->y==x[1], x);

  atoms1 = bt1.atoms
  if ! all(atoms1->atoms1==bt2.atoms[1], bt2.atoms)
    return false
  end

  # This should work, but raises an exception.
  # bonds1 = bt1.bonds
  # bonds2 = bt2.bonds
  # return all(bonds1->bonds1==bonds2[1], bonds2)

  for i in 1:length(bt1.bonds)
    bt1.bonds[i] != bt2.bonds[i] && return false
  end
  return true
end

"""Recusrively visit nodes in the graph defined by `nbrs`.
Args:
  nbrs:
  atom:
  visited:
Returns:
  The number of nodes visited - including `atom`.
"""
function visit(nbrs::Vector, atom::T, visited::Vector{Bool})::Int32 where {T<:Integer}
  visited[atom] = true
  result = 1    # To be returned.
  for nbr in nbrs[atom]
    if visited[nbr] > 0
      continue
    end
    result += visit(nbrs, nbr, visited)
  end

  return result
end

"""Return True if `bond_topology` is a single fragment.
Args:
  bond_topology:
Returns:
  True if `bond_topology` is a single fragment.
"""
function is_single_fragment(bond_topology::BondTopology)::Bool
  natoms = length(bond_topology.atoms)
  nbonds = length(bond_topology.bonds)
  # Some special cases are easy.
  natoms == 1 && return true
  if natoms == 2 && nbonds == 1
    return true
  end
  if natoms == 3 && nbonds == 2
    return true
  end
  if natoms == nbonds && natoms <= 4
    return true
  end

  attached = connections(bonded(bond_topology))
  # Any atom with zero neighbors means a detached atom.
  any(n->length(n) == 0, attached) && return false

  visited = zeros(Bool, natoms)
  # Mark anything with a single connection as visited.
  # Record the index of an atom that has multiple connections.
  a_multiply_connected_atom = -1
  for i in 1:natoms
    if bond_topology.atoms[i] == BondTopology_AtomType.ATOM_H
      visited[i] = 1
      continue
    end

    if length(attached[i]) > 1
      a_multiply_connected_atom = i
      continue
    end

    # A singly connected heavy atom. Mark visited if not part of a two atom fragment.
    if length(attached[attached[i][1]]) > 1
      visited[i] = 1
    end
  end

  # Not sure this can happen.
  a_multiply_connected_atom < 0 && return false

  number_visited = sum(visited) + visit(attached, a_multiply_connected_atom, visited)
  return number_visited == natoms
end


@testset "BatchQueryStrategies" begin
    BQS = OneClassActiveLearning.QueryStrategies.BatchQueryStrategies

    dummy_data, dummy_labels = load_data(TEST_DATA_FILE)

    classifier = SVDD.SVDDneg(dummy_data, dummy_labels)
    init_strategy = SimpleCombinedStrategy(FixedGammaStrategy(MLKernels.GaussianKernel(1)), FixedCStrategy(0.5))
    initialize!(classifier, init_strategy)

    empty_params = Dict{Symbol, Any}()
    sequential_strategy = Dict{Symbol, Any}(
        :type => :RandomPQs,
        :param => Dict{Symbol, Any}()
    )
    batch_size = 5

    function test_batch_qs(qs_type, params::Dict{Symbol, Any}, classifier, data, labels, candidate_indices)
        @test_throws Union{UndefKeywordError, MethodError, ArgumentError} initialize_qs(qs_type, classifier, data, empty_params)
        qs = initialize_qs(qs_type, classifier, data, Dict(params))
        for candidates in candidate_indices
            batch_indices = select_batch(qs, data, labels, candidates)
            @test length(batch_indices) == min(length(candidates), batch_size)
            @test issubset(batch_indices, Set(candidates))
        end
    end

    @testset "batch_qs" begin
        labels = labelmap(fill(:U, 10))
        candidate_indices = [collect(inds) for inds in [1:3, 1:10]]

        @testset "RandomBatchQs" begin
            qs_type = RandomBatchQs
            params = Dict{Symbol, Any}(
                :k => batch_size
            )
            test_batch_qs(qs_type, params, classifier, dummy_data, labels, candidate_indices)
        end

        @testset "TopK" begin
            qs_type = TopKBatchQs
            params = Dict{Symbol, Any}(
                :k => batch_size,
                :SequentialStrategy => sequential_strategy
            )
            test_batch_qs(qs_type, params, classifier, dummy_data, labels, candidate_indices)
        end

        @testset "GappedTopKBatchQs" begin
            qs_type = GappedTopKBatchQs
            params = Dict{Symbol, Any}(
                :k => batch_size,
                :m => 0,
                :SequentialStrategy => sequential_strategy
            )
            @test_throws ArgumentError initialize_qs(qs_type, SVDD.RandomOCClassifier(dummy_data), dummy_data, Dict(params..., :m => -1))
            test_batch_qs(qs_type, params, classifier, dummy_data, labels, candidate_indices)
        end

        @testset "ClusterBatchQs" begin
            qs_type = ClusterBatchQs
            params = Dict{Symbol, Any}(
                :k => batch_size
            )
            @test_throws UndefKeywordError initialize_qs(qs_type, SVDD.RandomOCClassifier(dummy_data), dummy_data, empty_params)
            test_batch_qs(qs_type, params, classifier, dummy_data, labels, candidate_indices)
        end

        @testset "ClusterTopKBatchQs" begin
            qs_type = ClusterTopKBatchQs
            params = Dict{Symbol, Any}(
                :k => batch_size,
                :SequentialStrategy => sequential_strategy
            )
            qs = initialize_qs(qs_type, SVDD.RandomOCClassifier(dummy_data), dummy_data, params)
            test_batch_qs(qs_type, params, classifier, dummy_data, labels, candidate_indices)
        end

        @testset "EnsembleBatchQs" begin
            qs_type = EnsembleBatchQs
            params = Dict{Symbol, Any}(
                :k => batch_size,
                :SequentialStrategy => sequential_strategy,
                :solver => TEST_SOLVER
            )
            qs = initialize_qs(EnsembleBatchQs, classifier, dummy_data, params)
            labels_ensemble = labelmap(fill(:U, TEST_DATA_NUM_OBSERVATIONS))
            for candidates in candidate_indices
                try
                    batch_indices = select_batch(qs, dummy_data, labels_ensemble, candidates)
                    @test length(batch_indices) == min(length(candidates), batch_size)
                    @test issubset(batch_indices, Set(candidates))
                catch e
                    # This strategy depends on the solver finding a solution
                    # for the submodels.
                    # The standard Ipopt solver does not always find one.
                    # If no solution is found, this strategy simply does not
                    # work and the use of another solver is recommended.
                    # In that case we skip these tests.
                    if e isa ErrorException
                        println(e)
                        println("Solver could not find a solution, skipping test.")
                        continue
                    else
                        throw(e)
                    end
                end
            end
        end

        @testset "weighted sum objective" begin
            qs_types = [IterativeBatchQs, EnumerativeBatchQs]
            params = [Dict{Symbol, Any}(
                :k => batch_size,
                :SequentialStrategy => sequential_strategy,
                :representativeness => rep,
                :diversity => div,
                :λ_inf => 1,
                :λ_rep => 1,
                :λ_div => 1
            ) for rep in [:KDE] for div in [:AngleDiversity, :EuclideanDistance]]
            invalid_rep_params = Dict{Symbol, Any}(
                :k => batch_size,
                :SequentialStrategy => sequential_strategy,
                :representativeness => :INVALID,
                :diversity => :AngleDiversity,
                :λ_inf => 1,
                :λ_rep => 1,
                :λ_div => 1)
            invalid_div_params = Dict{Symbol, Any}(
                :k => batch_size,
                :SequentialStrategy => sequential_strategy,
                :representativeness => :KDE,
                :diversity => :INVALID,
                :λ_inf => 1,
                :λ_rep => 1,
                :λ_div => 1)
            for qs_type in qs_types
                @testset "$(qs_type)" begin
                    for invalid_params in [invalid_rep_params, invalid_div_params]
                        @test_throws ArgumentError initialize_qs(qs_type, classifier, dummy_data, invalid_params)
                    end
                    for filled_params in params
                        @testset "$(filled_params)" begin
                            test_batch_qs(qs_type, filled_params, classifier, dummy_data, labels, candidate_indices)
                        end
                    end
                end
            end
        end

        @testset "hierarchical" begin
            qs_types = [FilterHierarchicalBatchQs, EnumFilterHierarchicalBatchQs]
            params = Dict{Symbol, Any}(
                :k => batch_size,
                :SequentialStrategy => sequential_strategy,
                :representativeness => :KDE,
                :diversity => :EuclideanDistance
            )
            invalid_rep_params = Dict{Symbol, Any}(
                :k => batch_size,
                :SequentialStrategy => sequential_strategy,
                :representativeness => :INVALID,
                :diversity => :AngleDiversity)
            invalid_div_params = Dict{Symbol, Any}(
                :k => batch_size,
                :SequentialStrategy => sequential_strategy,
                :representativeness => :KDE,
                :diversity => :INVALID)
            for qs_type in qs_types
                @testset "$(qs_type)" begin
                    for invalid_params in [invalid_rep_params, invalid_div_params]
                        @test_throws ArgumentError initialize_qs(qs_type, classifier, dummy_data, invalid_params)
                    end
                    test_batch_qs(qs_type, params, classifier, dummy_data, labels, candidate_indices)
                end
            end
        end

        @testset "FilterSimilarBatchQs" begin
            qs_type = FilterSimilarBatchQs
            params = [Dict{Symbol, Any}(
                :k => batch_size,
                :SequentialStrategy => sequential_strategy,
                :diversity => div
            ) for div in [:AngleDiversity, :EuclideanDistance]]
            @test_throws UndefVarError initialize_qs(qs_type, classifier, dummy_data, invalid_div_params)
            for filled_params in params
                @testset "$(filled_params)" begin
                    test_batch_qs(qs_type, filled_params, classifier, dummy_data, labels, candidate_indices)
                end
            end
        end
    end

    @testset "get_query_objects" begin
        qs = RandomBatchQs(k=4)
        @testset "a" begin
            data = rand(2, 6)
            pools = fill(:U, 6)
            pools[5] = :Lin
            indices = [1, 3, 5, 7, 9, 11]
            history = [[7]]
            @test_throws ArgumentError OneClassActiveLearning.get_query_objects(qs, data, fill(:Lin, 5), indices, history)
            @test OneClassActiveLearning.get_query_objects(qs, data, pools, indices, history) == [1, 3, 5, 11]
        end

        @testset "b" begin
            data = rand(2, 5)
            pools = fill(:U, 5)
            pools[5] = :Lin
            indices = [1, 2, 4, 7, 9]
            history = [[7]]
            @test_throws ArgumentError OneClassActiveLearning.get_query_objects(qs, data, fill(:Lin, 5), indices, history)
            @test OneClassActiveLearning.get_query_objects(qs, data, pools, indices, history) == [1, 2, 4]
        end
    end
end


include("../src/Rand.jl")
include("../src/Padding.jl")
include("../src/Iteration.jl")
include("../src/BlockCiphers.jl")
include("../src/CipherModes.jl")
include("testing.jl")

importall Padding
importall CipherModes
importall BlockCiphers

println("testing encryption modes:")

for key_len = [128, 192, 256]
    for plain = {zeros(Uint8, 20), zeros(Uint8, 0), zeros(Uint8, 1)}
        for key = [AES(zeros(Uint8, int32(key_len/8))), DES(zeros(Uint8, 8)), DES3(zeros(Uint8, 8 * 3)), DES3ede(zeros(Uint8, 8 * 3))]
            print("testing ", typeof(key), " ")

            base_len = length(plain)
            padded_len = block_size(key) * int64(ceil(length(plain) / block_size(key)))
            if(length(plain) % block_size(key) == 0)
                padded_len += block_size(key)
            end
            iv_len = block_size(key)

            out = ecb_encrypt(Pkcs7, key, plain)
            @assert_equal(padded_len + 0, length(out))
            @assert_not_equal(plain, out)
            @assert_equal(plain, ecb_decrypt(Pkcs7, key, out))

            out = cbc_encrypt(Pkcs7, key, plain)
            @assert_equal(padded_len + iv_len, length(out))
            @assert_not_equal(plain, out)
            @assert_equal(plain, cbc_decrypt(Pkcs7, key, out))

            out = pcbc_encrypt(Pkcs7, key, plain)
            @assert_equal(padded_len + iv_len, length(out))
            @assert_not_equal(plain, out)
            @assert_equal(plain, pcbc_decrypt(Pkcs7, key, out))

            out = ctr_encrypt(key, 31337, plain)
            @assert_equal(base_len, length(out)) # no padding :D
            if(length(plain) != 0)
                @assert_not_equal(plain, out)
            end
            @assert_equal(plain, ctr_decrypt(key, 31337, out))

            out = cfb_encrypt(key, plain)
            @assert_equal(base_len + iv_len, length(out)) # no padding + iv though :D
            @assert_not_equal(plain, out)
            @assert_equal(plain, cfb_decrypt(key, out))

            out = ofb_encrypt(key, plain)
            @assert_equal(base_len + iv_len, length(out)) # no padding + iv though :D
            @assert_not_equal(plain, out)
            @assert_equal(plain, ofb_decrypt(key, out))
            println()
        end
    end
end

println()


@testset "Reddit_full" begin
    data  = Reddit(full=true)
    @test length(data) == 1
    g = data[1]
    @test g.num_nodes == 232965
    @test g.num_edges == 114615892
    @test size(g.node_data.features) == (602, g.num_nodes)
    @test size(g.node_data.labels) == (g.num_nodes,)
    @test size(data.metadata["split"]["train"]) == (153431,)
    @test size(data.metadata["split"]["val"]) == (23831,)
    @test size(data.metadata["split"]["test"]) == (55703,)
    s, t = g.edge_index
    @test length(s) == length(t) == g.num_edges
    @test minimum(s) == minimum(t) == 1
    @test maximum(s) == maximum(t) == g.num_nodes
    @test sum(length.(values(data.metadata["split"]))) == g.num_nodes
end

@testset "Reddit_subset" begin
    data  = Reddit(full=false)
    @test length(data) == 1
    g = data[1]
    @test g.num_nodes == 231443
    @test g.num_edges == 23213838
    @test size(g.node_data.features) == (602, g.num_nodes)
    @test size(g.node_data.labels) == (g.num_nodes,)
    @test size(data.metadata["split"]["train"]) == (152410,)
    @test size(data.metadata["split"]["val"]) == (23699,)
    @test size(data.metadata["split"]["test"]) == (55334,)
    s, t = g.edge_index
    @test length(s) == length(t) == g.num_edges
    @test minimum(s) == minimum(t) == 1
    @test maximum(s) == maximum(t) == g.num_nodes
    @test sum(length.(values(data.metadata["split"]))) == g.num_nodes
end

@testset "TUDataset - PROTEINS" begin
    data  = TUDataset("PROTEINS")

    @test data.num_nodes == 43471
    @test data.num_edges == 162088
    @test data.num_graphs == 1113

    @test data.num_nodes == sum(g->g.num_nodes, data.graphs)
    @test data.num_edges == sum(g->g.num_edges, data.graphs)
    @test data.num_edges == sum(g->length(g.edge_index[1]), data.graphs)
    @test data.num_edges == sum(g->length(g.edge_index[2]), data.graphs)
    @test data.num_graphs == length(data) == length(data.graphs)

    i = rand(1:length(data))
    di = data[i]
    @test di isa NamedTuple
    g, targets = di.graphs, di.targets
    @test targets == 1
    @test g isa Graph
    @test all(1 .<= g.edge_index[1] .<= g.num_nodes)
    @test all(1 .<= g.edge_index[2] .<= g.num_nodes)

    # graph data
    @test size(data.graph_data.targets) == (data.num_graphs,)

    # node data
    @test size(g.node_data.features) == (1, g.num_nodes)
    @test size(g.node_data.targets) == (g.num_nodes,)

    # edge data
    @test g.edge_data === nothing
end

@testset "TUDataset - QM9" begin
    data  = TUDataset("QM9")

    @test data.num_nodes == 2333625
    @test data.num_edges == 4823498
    @test data.num_graphs == 129433

    @test data.num_nodes == sum(g->g.num_nodes, data.graphs)
    @test data.num_edges == sum(g->g.num_edges, data.graphs)
    @test data.num_edges == sum(g->length(g.edge_index[1]), data.graphs)
    @test data.num_edges == sum(g->length(g.edge_index[2]), data.graphs)
    @test data.num_graphs == length(data) == length(data.graphs)

    i = rand(1:length(data))
    g, features = data[i]
    @test g isa Graph
    @test all(1 .<= g.edge_index[1] .<= g.num_nodes)
    @test all(1 .<= g.edge_index[2] .<= g.num_nodes)

    # graph data
    @test size(data.graph_data.features) == (19, data.num_graphs)

    # node data
    @test size(g.node_data.features) == (16, g.num_nodes)

    # edge data
    @test size(g.edge_data.features) == (4, g.num_edges)
end


#ДАНО: Робот - Робот - в произвольной клетке ограниченного прямоугольного поля с перегородками
#РЕЗУЛЬТАТ: Робот - в исходном положении, и клетки поля промакированы так: нижний ряд - полностью,
#следующий - весь, за исключением одной последней клетки на Востоке, следующий - за исключением двух последних клеток
#на Востоке, и т.д.
include("library.jl")
function putmarker!(robot::CoordsRobot)
    x, y = get_xy(get_coords(robot))
    if x - y <= 0
        putmarker!(get_robot(robot))
    end
end
function lestniza_with_obstacles(r::Robot)
    back_path = BackPath(r)
    R = CoordsRobot(r)
    marks!(R)
    BackPath(R)
    back!(R, back_path)
end


struct CenterCropCat
    trms
end

function CenterCropCat(t, vol)
    t = vol ? (t, t, t) : (t, t)
    CenterCropCat((t..., 0, 0))
end

function (cc::CenterCropCat)(x1, x2) # x2 : input
    lo = cc.trms .+ 1
    up = size(x2) .- cc.trms
    chcat(x2[UnitRange.(lo, up)...], x1)
end

"""
    unet(; inchannels, nclasses = 1, volume = false, base = 64,
    batchnorm = false, padding = false, upsample = :convt, nlevels = 4)

Build a [U-Net](https://arxiv.org/abs/1505.04597v1) to process images with
`inchannels` channels (e.g. `inchannels` = 3 for RGB images). Default argument
values correspond to the original paper implementation with unpadded
convolutions.

* `nclasses`: number of pixel classes. If `nclasses` = 2 the final layer have
  one channel with a sigmoid activation function, otherwise the final layer have
  `nclasses` channels followed by a softmax function.
* `volume`: set to `true` to process tri-dimensional datas
* `basewidth`: base number of channels, the number of filters/channels is
  multiplied by two at each U-Net level. Notice that `basewidth` should be a
  power of two number, other values could lead to unfunctionnal U-Net.
* `batchnorm`: if `true` add a `BatchNorm` layer after each convolution
* `padding`: if `true` convolutions are padded
* `upsample`: either `:convt` (`ConvTranspose` layer), `:nearest` or `:bilinear`
  (`Upsample layer`)
* `nlevels`: number of level or depth of the U-Net
"""
function unet(; inchannels,
                nclasses = 2,
                volume = false,
                basewidth = 64,
                # expansion = 2,
                batchnorm = false,
                padding = false,
                upsample = :convt,
                nlevels = 4)
    if volume && upsample == :bilinear
        throw(ArgumentError(
            "Bilinear upsampling can not be used with tri-dimensionnal datas."))
    end

    # Pack parameters in a named tuple for convenience
    pars = (k₁ = volume ? (1, 1, 1) : (1, 1),
            k₂ = volume ? (2, 2, 2) : (2, 2),
            k₃ = volume ? (3, 3, 3) : (3, 3),
            nclasses = nclasses,
            inchannels = inchannels,
            basewidth = basewidth,
            batchnorm = batchnorm,
            padding = padding,
            upsample = upsample)

    # Build encoders and decoders for each level
    enc, dec, con = [], [], []
    for l ∈ 1:nlevels
        # e, d = level_enc_dec(l, pars)
        # push!(enc, e)
        # push!(dec, d)
        push!.((enc, dec), level_enc_dec(l, pars))
        if padding
            push!(con, chcat)
        else
            push!(con, CenterCropCat(trim(l, nlevels), volume))
        end
    end

    # Build bridge (i.e. bottom part)
    ec, dc = level_channels(nlevels + 1, pars)
    ic, mc, _ = ec
    _, _, oc = dc
    dw, up = samplers(oc, pars)
    bdg  = [dw, double_conv(ic, mc, oc, pars)..., up]

    uchain(encoders = enc, decoders = dec, bridge = bdg, connection = con)
end

# Build encoder and decoder blocks for a given level
function level_enc_dec(lvl, pars)
    # number of channels (input, middle, output)
    ec, dc = level_channels(lvl, pars)
    enc = []
    dec = []
    if lvl == 1
        enc = double_conv(ec, pars)
        # build classifier layers
        k₁ = pars[:k₁]
        nc = pars[:nclasses]
        if nc > 2
            d = length(k₁)
            cl = [Conv(k₁, dc[3] => nc), x -> softmax(x, dims = d)]
        else
            cl = [Conv(k₁, dc[3] => 1, sigmoid)]
        end
        dec = [double_conv(dc, pars)..., cl...]
    else
        dw, up = samplers(dc[3], pars)
        enc = [dw, double_conv(ec, pars)...]
        dec = [double_conv(dc, pars)..., up]
    end
    enc, dec
end

# Return two convolutional layers, with batchnom if needed
function double_conv(chs, pars)
    ic, mc, oc = chs
    padding, k₃, batchnorm = pars[:padding], pars[:k₃], pars[:batchnorm]
    p = padding ? 1 : 0
    cv1(chs) = [Conv(k₃, chs, pad = p, bias = false), BatchNorm(chs[2], relu)]
    cv2(chs) = [Conv(k₃, chs, relu, pad = p)]
    cnv(chs) =  batchnorm ? cv1(chs) : cv2(chs)
    [cnv(ic => mc)..., cnv(mc => oc)...]
end

double_conv(ic, mc, oc, pars) = double_conv((ic, mc, oc), pars)

# Compute encoder/decoder number of channels at a given level (lvl)
# return two tuples one for encoder and one for decoder
# formula :
#  ice = expansion^(lvl - 2) * basewidth
#  mce = expansion^(lvl - 1) * basewidth
function level_channels(lvl, pars)
    inchannels, basewidth = pars[:inchannels], pars[:basewidth]
    upsample = pars[:upsample]
    # encoder channels: input, middle, ouptput = (ice, mce, oce)
    ice = (lvl == 1) ? inchannels : 2^(lvl - 2) * basewidth
    mce = oce = 2^(lvl - 1) * basewidth
    # decoder channels: input, middle, ouptput = (icd, mcd, ocd)
    icd, mcd = 2 * oce, oce
    if upsample == :convt
        ocd = mcd
    else
        ocd = (lvl == 1) ? mcd : mcd ÷ 2
    end
    (ice, mce, oce), (icd, mcd, ocd)
end

# Downsampler, Upsampler
function samplers(chs, pars)
    k₂, upsample = pars[:k₂], pars[:upsample]
    dw = MaxPool(k₂)
    if upsample == :convt
        up =  ConvTranspose(k₂, chs => (chs ÷ 2), stride = 2)
    else
        up = Upsample(upsample, scale = k₂)
    end
    dw, up
end


function validate_constraints(
        source::UVT1,
        target::UVT2,
        resampling::ConstrainedResampling{N}) where {
            UVT1<:Union{AbstractUncertainValueDataset, Vector{<:AbstractUncertainValue}},
            UVT2<:Union{AbstractUncertainValueDataset, Vector{<:AbstractUncertainValue}},
            N}
    if N != 2
        msg = "Both `source` and `target` are uncertain, so `resampling` must contain exactly two separate sets of sampling constraints"
        throw(ArgumentError(msg))
    end
end

function validate_constraints(
        source::Vector{<:Real},
        target::UVT,
        resampling::ConstrainedResampling{N}) where {
            UVT<:Union{AbstractUncertainValueDataset, Vector{<:AbstractUncertainValue}}, N}

    if N != 1
        msg = "Only `target` is uncertain, but $N sets of constraints were provided. Only one should be given."
        throw(ArgumentError(msg))
    end
end

function validate_constraints(
        source::UVT,
        target::Vector{<:Real},
        resampling::ConstrainedResampling{N}) where {
            UVT<:Union{AbstractUncertainValueDataset, Vector{<:AbstractUncertainValue}}, N}

    if N != 1
        msg = "Only `source` is uncertain, but $N sets of constraints were provided. Only one should be given."
        throw(ArgumentError(msg))
    end
end

function validate_constraints(
        source::Vector{<:Real},
        target::Vector{<:Real},
        resampling::ConstrainedResampling{N}) where {
            UVT<:AbstractUncertainValue, N}
    @warn "Neither `source` nor `target` are uncertain. Not constraining them."
end


function hum_dewpoint end
function hum_massfrac end
function hum_molfrac end
function rel_hum end
function hum_wetbulb end
function hum_ratio end
function dry_volume end
function dry_enthalpy end
function dry_entropy end


module MLJTuringInterface

export TuringModel

import MLJModelInterface

using MLJModelInterface:
    Probabilistic,
    Table,
    Continuous,
    metadata_model,
    metadata_pkg
using Statistics: mean
using Tables
using Turing:
    MCMCThreads,
    DynamicPPL,
    chainscat,
    group,
    predict,
    replacenames,
    sample

const MMI = MLJModelInterface
const PKG = "MLJTuringInterface"

include("interface.jl")
include("metadata.jl")

end # module


export MultiThreadEnv

using Base.Threads: @spawn

"""
    MultiThreadEnv(envs::Vector{<:AbstractEnv})

Wrap multiple instances of the same environment type into one environment.
Each environment will run in parallel by leveraging `Threads.@spawn`.
So remember to set the environment variable `JULIA_NUM_THREADS`!
"""
struct MultiThreadEnv{E,S,R,AS,SS,L} <: AbstractEnv
    envs::Vector{E}
    states::S
    rewards::R
    terminals::BitArray{1}
    action_space::AS
    state_space::SS
    legal_action_space_mask::L
end

function Base.show(io::IO, t::MIME"text/markdown", env::MultiThreadEnv)
    print(io, "MultiThreadEnv($(length(env)) x $(nameof(env[1])))")
end

"""
    MultiThreadEnv(f, n::Int)

`f` is a lambda function which creates an `AbstractEnv` by calling `f()`.
"""
MultiThreadEnv(f, n::Int) = MultiThreadEnv([f() for _ in 1:n])

function MultiThreadEnv(envs::Vector{<:AbstractEnv})
    n = length(envs)
    S = state_space(envs[1])
    s = state(envs[1])
    if S isa Space
        S_batch = similar(S, size(S)..., n)
        s_batch = similar(s, size(s)..., n)
        for j in 1:n
            Sₙ = state_space(envs[j])
            sₙ = state(envs[j])
            for i in CartesianIndices(size(S))
                S_batch[i, j] = Sₙ[i]
                s_batch[i, j] = sₙ[i]
            end
        end
    else
        S_batch = Space(state_space.(envs))
        s_batch = state.(envs)
    end

    A = action_space(envs[1])
    if A isa Space
        A_batch = similar(A, size(A)..., n)
        for j in 1:n
            Aⱼ = action_space(envs[j])
            for i in CartesianIndices(size(A))
                A_batch[i, j] = Aⱼ[i]
            end
        end
    else
        A_batch = Space(action_space.(envs))
    end

    r_batch = reward.(envs)
    t_batch = is_terminated.(envs)
    if ActionStyle(envs[1]) === FULL_ACTION_SET
        m_batch = BitArray(undef, size(A)..., n)
        for j in 1:n
            L = legal_action_space_mask(envs[j])
            for i in CartesianIndices(size(A))
                m_batch[i, j] = L[i]
            end
        end
    else
        m_batch = nothing
    end
    MultiThreadEnv(envs, s_batch, r_batch, t_batch, A_batch, S_batch, m_batch)
end

MacroTools.@forward MultiThreadEnv.envs Base.getindex, Base.length, Base.iterate

function (env::MultiThreadEnv)(actions)
    N = ndims(actions)
    @sync for i in 1:length(env)
        @spawn begin
            if N == 1
                env[i](actions[i])
            else
                env[i](selectdim(actions, N, i))
            end
        end
    end
end

function RLBase.reset!(env::MultiThreadEnv; is_force = false)
    if is_force
        for i in 1:length(env)
            reset!(env[i])
        end
    else
        @sync for i in 1:length(env)
            if is_terminated(env[i])
                @spawn begin
                    reset!(env[i])
                end
            end
        end
    end
end

const MULTI_THREAD_ENV_CACHE = IdDict{AbstractEnv,Dict{Symbol,Array}}()

function RLBase.state(env::MultiThreadEnv)
    N = ndims(env.states)
    @sync for i in 1:length(env)
        @spawn selectdim(env.states, N, i) .= state(env[i])
    end
    env.states
end

function RLBase.reward(env::MultiThreadEnv)
    env.rewards .= reward.(env.envs)
    env.rewards
end

function RLBase.is_terminated(env::MultiThreadEnv)
    env.terminals .= is_terminated.(env.envs)
    env.terminals
end

function RLBase.legal_action_space_mask(env::MultiThreadEnv)
    N = ndims(env.states)
    @sync for i in 1:length(env)
        @spawn selectdim(env.legal_action_space_mask, N, i) .=
            legal_action_space_mask(env[i])
    end
    env.legal_action_space_mask
end

RLBase.action_space(env::MultiThreadEnv) = env.action_space
RLBase.state_space(env::MultiThreadEnv) = env.state_space
RLBase.legal_action_space(env::MultiThreadEnv) = Space(legal_action_space.(env.envs))
# RLBase.current_player(env::MultiThreadEnv) = current_player.(env.envs)

for f in RLBase.ENV_API
    if endswith(String(f), "Style")
        @eval RLBase.$f(x::MultiThreadEnv) = $f(x[1])
    end
end

#####
# Patches
#####

(env::MultiThreadEnv)(action::EnrichedAction) = env(action.action)

function (π::QBasedPolicy)(env::MultiThreadEnv, ::MinimalActionSet, A)
    [A[i][a] for (i, a) in enumerate(π.explorer(π.learner(env)))]
end

function (π::QBasedPolicy)(env::MultiThreadEnv, ::FullActionSet, A)
    [
        A[i][a] for
        (i, a) in enumerate(π.explorer(π.learner(env), legal_action_space_mask(env)))
    ]
end

function (π::QBasedPolicy)(
    env::MultiThreadEnv,
    ::MinimalActionSet,
    ::Space{<:Vector{<:Base.OneTo{<:Integer}}},
)
    π.explorer(π.learner(env))
end

function (π::QBasedPolicy)(
    env::MultiThreadEnv,
    ::FullActionSet,
    ::Space{<:Vector{<:Base.OneTo{<:Integer}}},
)
    π.explorer(π.learner(env), legal_action_space_mask(env))
end


using Gurobi

using Compat
using Compat.Test, Compat.SparseArrays, Compat.Random

@testset "C API" begin
    include("c_wrapper.jl")
end

const mpb_tests = [
    "lp_01a",
    "lp_01b",
    "lp_02",
    "lp_03",
    "lp_04",
    "mip_01",
    "qp_01",
    "qp_02",
    "qcqp_01",
    "mathprog",
    "test_grb_attrs",
    "env",
    "range_constraints",
    "test_get_strarray",
    "large_coefficients",
    "multiobj",
    "test_read"
]

@testset "MathProgBase Tests" begin
    for file in mpb_tests
        evalfile(joinpath("MathProgBase", "$(file).jl"))
    end
end

@testset "MathOptInterface Tests" begin
    include("MOIWrapper.jl")
end


using Colors, FixedPointNumbers
using Test
using ColorTypes: eltype_default, parametric3

struct C3{T} <: Color{T,3}
    c1::T; c2::T; c3::T;
end
struct C4{T} <: Color{T,4}
    c1::T; c2::T; c3::T; c4::T;
end
function ColorTypes._convert(::Type{Cdest}, ::Type{Odest}, ::Type{Osrc}, c) where {Cdest<:AbstractRGB,Odest,Osrc<:C4}
    Cdest(c.c1 + c.c4, c.c2 + c.c4, c.c3 + c.c4)
end
function ColorTypes._convert(::Type{Cdest}, ::Type{Odest}, ::Type{Osrc}, c) where {Cdest<:C4,Odest,Osrc<:AbstractRGB}
    r, g, b = red(c), green(c), blue(c)
    c4 = min(r, g, b)
    Cdest(r - c4, g - c4, b - c4, c4)
end

@testset "Conversion" begin
    r8(x) = reinterpret(N0f8, x)

    # srgb_compand / invert_srgb_compand
    @test Colors.srgb_compand(0.5) ≈ 0.7353569830524494 atol=eps()
    @test Colors.invert_srgb_compand(0.7353569830524494) ≈ 0.5 atol=eps()
    @test Colors.invert_srgb_compand(0.735357f0) ≈ 0.5f0 atol=eps(Float32)

    fractional_types = (RGB, BGR, XRGB, RGBX)  # types that support Fractional

    redF64 = RGB{Float64}(1,0,0)
    redF32 = RGB{Float32}(1,0,0)
    red24 = reinterpret(RGB24, 0x00ff0000)
    red32 = reinterpret(ARGB32, 0xffff0000)
    @testset "type check: RGB, RGB{$T}" for T in (Float64, Float32, N0f8)
        c = RGB(one(T), zero(T), zero(T))
        @test eltype(c) == T
        @test convert(RGB{Float64}, c) === redF64
        @test convert(RGB{T}, redF64) == c
    end
    @test RGB(1,0,0) == redF64
    @test RGB(UInt8(1), 0, 0) == redF64
    @test RGB(UInt8(1), UInt8(0), UInt8(0)) == redF64 # != colorant"#010000"
    @test convert(RGB, red24) == redF64

    @testset "type check: RGB24-->$Cto" for Cto in parametric3
        @test typeof(convert(Cto, red24)) == Cto{eltype_default(Cto)}
        @test typeof(convert(Cto{Float64}, red24)) == Cto{Float64}
    end

    @testset "type check: C{Float}-->C{Float}" for Cfrom in parametric3, Tfrom in (Float32, Float64)
        c = convert(Cfrom{Tfrom}, redF64)
        @test typeof(c) == Cfrom{Tfrom}
        @testset "$Cfrom{$Tfrom}-->$Cto" for Cto in parametric3
            c1 = convert(Cto, c)
            @test eltype(c1) == Tfrom
        end
        @testset "$Cfrom{$Tfrom}-->$Cto{$Tto}" for Cto in parametric3, Tto in (Float32, Float64)
            c2 = convert(Cto{Tto}, c)
            @test typeof(c2) == Cto{Tto}
        end
    end

    normed_types = (N0f8, N6f10, N4f12, N2f14, N0f16)
    # Test conversion from Normed types
    @testset "type check: C{Normed}-->C{Float}" for Cfrom in fractional_types, Tfrom in normed_types
        c = convert(Cfrom{Tfrom}, redF64)
        @test typeof(c) == Cfrom{Tfrom}
        @testset "$Cfrom{$Tfrom}-->$Cto" for Cto in parametric3
            eltype_default(Cto) <: FixedPoint && continue
            c1 = convert(Cto, c)
            @test eltype(c1) == eltype_default(Cto)
        end
        @testset "$Cfrom{$Tfrom}-->$Cto{$Tto}" for Cto in parametric3, Tto in (Float32, Float64)
            c2 = convert(Cto{Tto}, c)
            @test typeof(c2) == Cto{Tto}
        end
    end

    # Test conversion to Normed types
    @testset "type check: C{Float}-->C{Normed}" for Cfrom in parametric3, Tfrom in (Float32, Float64)
        c = convert(Cfrom{Tfrom}, redF64)
        @test typeof(c) == Cfrom{Tfrom}
        @testset "$Cfrom{$Tfrom}-->$Cto{$Tto}" for Cto in fractional_types, Tto in normed_types
            c2 = convert(Cto{Tto}, c)
            @test typeof(c2) == Cto{Tto}
        end
    end

    ac = RGBA(redF64)

    @test convert(RGB, ac) == RGB(1,0,0)
    @test convert(RGB{N0f8}, ac) == RGB{N0f8}(1,0,0)
    @test convert(RGBA{N0f8}, ac) == RGBA{N0f8}(1,0,0,1)
    @test convert(HSVA, ac) == HSVA{Float64}(convert(HSV, redF64), 1.0)
    @test convert(HSVA{Float32}, ac) == HSVA{Float32}(convert(HSV{Float32}, redF64), 1.0f0)
    @test convert(RGBA, redF64) == ac

    @test convert(ARGB32, ac) == reinterpret(ARGB32, 0xffff0000)
    @test convert(ARGB32, RGB(1,0,0)) == reinterpret(ARGB32, 0xffff0000)
    ac32 = convert(ARGB32, RGB(1,0,0), 0.5)
    @test ac32 == reinterpret(ARGB32, 0x80ff0000)
    @test convert(ARGB32, ac32) === ac32
    @test convert(ARGB32, BGRA(1,0,0,0.5)) == reinterpret(ARGB32, 0x80ff0000)
    @test reinterpret(UInt32, convert(ARGB32, ac)) == 0xffff0000
    @test convert(RGB24, RGB(r8(0xff),r8(0x00),r8(0x00))) == reinterpret(RGB24, 0x00ff0000)
    @test reinterpret(UInt32, convert(RGB24, RGB(r8(0xff),r8(0x00),r8(0x00)))) == 0x00ff0000
    redhsv = convert(HSV, redF64)
    @test convert(RGB24, red24) === red24
    @test convert(RGB24, redhsv) == reinterpret(RGB24, 0x00ff0000)
    @test_throws ArgumentError convert(RGB24,  RGB(0, 1.1, 0))
    @test_throws ArgumentError convert(ARGB32, RGBA(0, 1.1, 0, 0.8))
    @test_throws ArgumentError convert(ARGB32, RGBA(0, 0.8, 0, 1.1))

    @test convert(RGB{N0f8}, red24) == RGB{N0f8}(1,0,0)
    @test convert(RGBA{N0f8}, red32) == RGBA{N0f8}(1,0,0,1)
    @test convert(HSVA{Float64}, red32) == HSVA{Float64}(0, 1, 1, 1)

    # whitepoint conversions
    @test isa(convert(XYZ, convert(Lab, redF64), Colors.WP_DEFAULT), XYZ{Float64})
    @test isa(convert(XYZ{Float32}, convert(Lab, redF64), Colors.WP_DEFAULT), XYZ{Float32})
    @test isa(convert(XYZ, convert(Luv, redF64), Colors.WP_DEFAULT), XYZ{Float64})
    @test isa(convert(XYZ{Float32}, convert(Luv, redF64), Colors.WP_DEFAULT), XYZ{Float32})
    @test isa(convert(Lab, convert(XYZ, redF64), Colors.WP_DEFAULT), Lab{Float64})
    @test isa(convert(Lab{Float32}, convert(XYZ, redF64), Colors.WP_DEFAULT), Lab{Float32})
    @test isa(convert(Luv, convert(XYZ, redF64), Colors.WP_DEFAULT), Luv{Float64})
    @test isa(convert(Luv{Float32}, convert(XYZ, redF64), Colors.WP_DEFAULT), Luv{Float32})

    #59
    @test Colors.xyz_to_uv(XYZ{Float64}(0.0, 0.0, 0.0)) === (0.0, 0.0)
    @test Colors.xyz_to_uv(XYZ{Float64}(0.0, 1.0, 0.0)) === (0.0, 0.6)
    @test Colors.xyz_to_uv(XYZ{Float64}(0.0, 1.0, 1.0)) === (0.0, 0.5)
    @test Colors.xyz_to_uv(XYZ{Float64}(1.0, 0.0, 1.0)) === (1.0, 0.0)
    @test Colors.xyz_to_uv(XYZ{Float64}(1.0, 0.0, 0.0)) === (4.0, 0.0)

    # ColorTypes.jl issue #40
    @test convert(HSL, RGB{N0f8}(0.678, 0.847, 0.902)) ≈ HSL{Float32}(194.73685f0,0.5327105f0,0.7901961f0) atol=100eps(Float32)

    # YIQ
    @test convert(YIQ, RGB(1,0,0)) == YIQ{Float32}(0.299, 0.595716, 0.211456)
    @test convert(YIQ, RGB(0,1,0)) == YIQ{Float32}(0.587, -0.274453, -0.522591)
    @test convert(YIQ, RGB(0,0,1)) == YIQ{Float32}(0.114, -0.321263, 0.311135)
    @test convert(RGB, YIQ(1.0,0.0,0.0)) == RGB(1,1,1)
    v = 0.5957
    @test convert(RGB, YIQ(0.0,1.0,0.0)) == RGB(0.9563*v,0,0)
    v = -0.5226
    @test convert(RGB, YIQ(0.0,0.0,-1.0)) == RGB(0,-0.6474*v,0)

    # Gray
    @test convert(Gray{N0f8}, Gray{N0f8}(0.1)) == Gray{N0f8}(0.1)
    @test convert(Gray{N0f8}, Gray(0.1))     == Gray{N0f8}(0.1)
    @test convert(Gray{N0f8}, Gray24(0.1))   == Gray{N0f8}(0.1)
    @test convert(Gray24, Gray{N0f8}(0.1))   == Gray24(0.1)

    @test convert(RGB{N0f8}, Gray{N0f8}(0.1)) == RGB{N0f8}(0.1,0.1,0.1)
    @test convert(RGB{N0f8}, Gray24(0.1))   == RGB{N0f8}(0.1,0.1,0.1)

    grayN0f16 = Gray{N0f16}(0.8)
    @test convert(RGB, grayN0f16) == RGB{N0f16}(0.8,0.8,0.8)
    @test convert(RGB{Float32}, grayN0f16) == RGB{Float32}(0.8,0.8,0.8)

    @testset "$C-->Gray" for C in parametric3
        c = convert(C, RGB(1,1,1))
        g1 = convert(Gray, c)
        @test isa(g1, Gray)
        @test gray(g1) ≈ 1 atol=0.01
        g2 = convert(Gray{Float64}, c)
        @test typeof(g2) == Gray{Float64}
        @test gray(g2) ≈ 1 atol=0.01
    end
    # Issue #377
    @test convert(Gray, RGB24(1,0,0)) === convert(Gray, RGB(1,0,0)) === Gray{N0f8}(0.298)
    @test convert(Gray24, RGB(1,0,0)) === Gray24(0.298)
    # Check for roundoff error
    for N in (N0f8, N0f16, N0f32)
        @test convert(Gray{N}, RGB{N}(1,1,1)) === Gray{N}(1)
    end
    @test gray(convert(Gray{N0f64}, RGB{N0f64}(1,1,1))) ≈ 1.0

    # Images issue #382
    @test convert(Gray, RGBA(1,1,1,1)) == Gray(N0f8(1))

    # https://github.com/timholy/Images.jl/pull/445#issuecomment-189866806
    @test convert(Gray, RGB{N0f8}(0.145,0.145,0.145)) == Gray{N0f8}(0.145)

    # More AbstractRGB
    r4 = RGBX(1,0,0)
    @test convert(RGB, r4) == RGB(1,0,0)
    @test convert(RGB{N0f8}, r4) == RGB{N0f8}(1,0,0)
    @test convert(RGBX{N0f8}, r4) == RGBX{N0f8}(1,0,0)
    @test convert(RGBX{Float32}, r4) == RGBX{Float32}(1,0,0)
    @test convert(BGR{Float32}, r4) == BGR{Float32}(1,0,0)

    # Issue #257
    c = RGB{Float16}(0.9473,0.962,0.9766)
    convert(HSI, c)
    hsi = convert(HSI, c)
    @test hsi.i > 0.96 && hsi.h ≈ 210

    # {HSV, HSL, HSI} --> RGB (issue #379)
    @testset "HSx --> RGB" begin
        @test convert(RGB, HSV{Float32}( 780, 1, 1)) === RGB{Float32}(1,1,0)
        @test convert(RGB, HSV{Float32}(   0, 1, 1)) === RGB{Float32}(1,0,0)
        @test convert(RGB, HSV{Float32}(-780, 1, 1)) === RGB{Float32}(1,0,1)
        @test convert(RGB, HSV{Float32}(30,  2, .5)) === RGB{Float32}(.5,.25,0)
        @test convert(RGB, HSV{Float32}(30, .5, -1)) === RGB{Float32}(0,0,0)
        @test convert(RGB{Float64}, HSV{BigFloat}(-360120, 2, 1)) === RGB{Float64}(0,0,1)

        @test convert(RGB, HSL{Float32}( 780, 1, .5)) === RGB{Float32}(1,1,0)
        @test convert(RGB, HSL{Float32}(   0, 1, .5)) === RGB{Float32}(1,0,0)
        @test convert(RGB, HSL{Float32}(-780, 1, .5)) === RGB{Float32}(1,0,1)
        @test convert(RGB, HSL{Float32}(30,  2, .25)) === RGB{Float32}(.5,.25,0)
        @test convert(RGB, HSL{Float32}(30, .5,  -1)) === RGB{Float32}(0,0,0)
        @test convert(RGB{Float64}, HSL{BigFloat}(-360120, 2, .5)) === RGB{Float64}(0,0,1)

        @test convert(RGB, HSI{Float32}( 780, .5, .5)) ≈ RGB{Float32}(.625,.625,.25)
        @test convert(RGB, HSI{Float32}(   0, .5, .5)) ≈ RGB{Float32}(1,.25,.25)
        @test convert(RGB, HSI{Float32}(-780, .5, .5)) ≈ RGB{Float32}(.625,.25,.625)
        @test convert(RGB, HSI{Float32}(30,  2, .25)) ≈ RGB{Float32}(.5,.25,0)
        @test convert(RGB, HSI{Float32}(30, .5,  -1)) ≈ RGB{Float32}(0,0,0)
        @test convert(RGB{Float64}, HSI{BigFloat}(-360120, .5, .5)) ≈ RGB{Float64}(.25,.25,1)
    end

    @testset "custom types" begin
        # issue #465
        @test_throws ErrorException convert(RGB,   C3{Float32}(1, 2, 3))
        @test_throws ErrorException convert(RGB24, C3{Float64}(1, 2, 3))
        @test_throws ErrorException convert(Lab,   C3{Float16}(1, 2, 3))

        @test @inferred(convert(RGB,   C4{Float32}(0.1, 0.2, 0.3, 0.4))) ≈ RGB{Float32}(0.5, 0.6, 0.7)
        @test @inferred(convert(RGB24, C4{Float64}(0.1, 0.2, 0.3, 0.4))) ≈ RGB24(0.5, 0.6, 0.7)
        @test @inferred(convert(ALab,  C4{Float16}(0.1, 0.2, 0.3, 0.4))) ≈ ALab{Float16}(62.12, -2.86, -16.25, 1.0)

        @test @inferred(convert(C4, RGB{Float32}(0.5, 0.6, 0.7))) ≈ C4{Float32}(0.0, 0.1, 0.2, 0.5)
        @test @inferred(convert(C4{N0f8}, ARGB32(0.5, 0.6, 0.7))) ≈ C4{N0f16}(0.0, 0.1, 0.2, 0.5)
        @test @inferred(convert(C4, Lab{Float16}(62.12, -2.86, -16.25))) ≈ C4{Float16}(0.0, 0.1, 0.2, 0.5)
    end

    # Test accuracy of conversion
    include("test_conversions.jl")

    # Since `colordiff`(e.g. `DE_2000`) involves a color space conversion, it is
    # not suitable for evaluating the conversion itself. On the other hand,
    # since the tolerance varies from component to component, a homogeneous
    # error evaluation function (e.g. a simple sum of differences) is also not
    # appropriate. Therefore, a series of `diffnorm`, which returns the
    # normalized Euclidean distance, is defined as follows. They are just for
    # testing purposes as the cyclicity of hue is ignored.
    sqd(a, b, s=1.0) = ((float(a) - float(b))/s)^2
    function diffnorm(a::T, b::T) where {T<:Color3} # RGB,XYZ,xyY,LMS
        sqrt(sqd(comp1(a), comp1(b)) + sqd(comp2(a), comp2(b)) + sqd(comp3(a),comp3(b)))/sqrt(3)
    end
    function diffnorm(a::T, b::T) where {T<:Union{HSV,HSL,HSI}}
        sqrt(sqd(a.h, b.h, 360) + sqd(a.s, b.s) + sqd(comp3(a), comp3(b)))/sqrt(3)
    end
    function diffnorm(a::T, b::T) where {T<:Union{Lab,Luv}}
        sqrt(sqd(a.l, b.l, 100) + sqd(comp2(a), comp2(b), 200) + sqd(comp3(a), comp3(b), 200))/sqrt(3)
    end
    function diffnorm(a::T, b::T) where {T<:Union{LCHab,LCHuv}}
        sqrt(sqd(a.l, b.l, 100) + sqd(a.c, b.c, 100) + sqd(a.h, b.h, 360))/sqrt(3)
    end
    function diffnorm(a::T, b::T) where {T<:Union{DIN99,DIN99d,DIN99o}} # csconv has no DIN99 case
        sqrt(sqd(a.l, b.l, 100) + sqd(a.a, b.a, 100) + sqd(a.b, b.b, 100))/sqrt(3)
    end
    function diffnorm(a::T, b::T) where {T<:YIQ}
        sqrt(sqd(a.y, b.y) + sqd(a.i, b.i, 1.2) + sqd(a.q, b.q, 1.2))/sqrt(3)
    end
    function diffnorm(a::T, b::T) where {T<:YCbCr}
        sqrt(sqd(a.y, b.y, 219) + sqd(a.cb, b.cb, 224) + sqd(a.cr, b.cr, 224))/sqrt(3)
    end

    for C in ColorTypes.parametric3
        y = convert(C, RGB(1.0,1.0,0.0))
        b = convert(C, RGB(0.1,0.1,0.2))
        diffnorm(y, b) < 0.5 && @warn("`diffnorm` for $C may be broken")
    end

    function convcompare(from, to, tol; showfailure::Bool=false)
        errmax = 0.0
        for i = 1:length(from)
            t = to[i]
            f = convert(typeof(t), from[i])
            diff = diffnorm(t, f)
            if showfailure && diff>tol
                original = from[i]
                @show original f t diff
            end
            errmax = max(errmax, diff)
        end
        errmax > tol && @warn("Error on conversions from $(eltype(from)) to $(eltype(to)), relative error = $errmax")
        errmax <= tol
    end

    @testset "accuracy test: $Cfrom-->$Cto" for Cto in keys(csconv), Cfrom in keys(csconv)
        f = csconv[Cfrom]
        t = csconv[Cto]
        @test convcompare(f, t, 2e-3, showfailure=false)
    end
end


@inline getbody(mechanism::Mechanism, id::Integer) = mechanism.bodies[id]
@inline getbody(mechanism::Mechanism, id::Nothing) = mechanism.origin
function getbody(mechanism::Mechanism, name::String)
    if mechanism.origin.name == name
        return mechanism.origin
    else
        for body in mechanism.bodies
            if body.name == name
                return body
            end
        end
    end
    return
end
@inline geteqconstraint(mechanism::Mechanism, id::Integer) = mechanism.eqconstraints[id]
function geteqconstraint(mechanism::Mechanism, name::String)
    for eqc in mechanism.eqconstraints
        if eqc.name == name
            return eqc
        end
    end
    return
end
@inline getineqconstraint(mechanism::Mechanism, id::Integer) = mechanism.ineqconstraints[id]
function getineqconstraint(mechanism::Mechanism, name::String)
    for ineqc in mechanism.eqconstraints
        if ineqc.name == name
            return ineqc
        end
    end
    return
end

getcomponent(mechanism::Mechanism, id::Nothing) = mechanism.origin
function getcomponent(mechanism::Mechanism, id::Integer)
    if haskey(mechanism.bodies, id)
        return getbody(mechanism, id)
    elseif haskey(mechanism.eqconstraints, id)
        return geteqconstraint(mechanism, id)
    elseif haskey(mechanism.ineqconstraints, id)
        return getineqconstraint(mechanism, id)
    else
        return
    end
end

function getcomponent(mechanism::Mechanism, name::String)
    component = getbody(mechanism,name)
    if component === nothing
        component = geteqconstraint(mechanism,name)
    end
    if component === nothing
        component = getineqconstraint(mechanism,name)
    end
    return component
end

@inline function normf(mechanism::Mechanism, body::Body)
    f = dynamics(mechanism, body)
    return dot(f, f)
end

@inline function normf(mechanism::Mechanism, eqc::EqualityConstraint)
    f = g(mechanism, eqc)
    return dot(f, f)
end

@inline function normf(mechanism::Mechanism, ineqc::InequalityConstraint)
    f = gs(mechanism, ineqc)
    d = h(ineqc)
    return dot(f, f) + dot(d, d)
end

@inline function normfμ(mechanism::Mechanism, ineqc::InequalityConstraint)
    f = gs(mechanism, ineqc)
    d = hμ(ineqc, mechanism.μ)
    return dot(f, f) + dot(d, d)
end

@inline function normf(mechanism::Mechanism)
    mechanism.normf = 0

    foreachactive(addNormf!, mechanism.bodies, mechanism)
    foreachactive(addNormf!, mechanism.eqconstraints, mechanism)
    foreachactive(addNormf!, mechanism.ineqconstraints, mechanism)

    return sqrt(mechanism.normf)
end

@inline function normf(mechanism::LinearMechanism)
    mechanism.normf = norm([fdynamics(mechanism); fconstraints(mechanism)])

    return sqrt(mechanism.normf)
end

@inline function meritf(mechanism::Mechanism)
    mechanism.normf = 0

    foreachactive(addNormf!, mechanism.bodies, mechanism)
    foreachactive(addNormf!, mechanism.eqconstraints, mechanism)
    foreachactive(addNormfμ!, mechanism.ineqconstraints, mechanism)

    return sqrt(mechanism.normf)
end

@inline function normΔs(mechanism::Mechanism)
    mechanism.normΔs = 0

    foreachactive(addNormΔs!, mechanism.bodies, mechanism)
    foreachactive(addNormΔs!, mechanism.eqconstraints, mechanism)
    foreachactive(addNormΔs!, mechanism.ineqconstraints, mechanism)

    return sqrt(mechanism.normΔs)
end

@inline function addNormf!(component::Component, mechanism::Mechanism)
    mechanism.normf += normf(mechanism, component)
    return
end

@inline function addNormfμ!(ineqc::InequalityConstraint, mechanism::Mechanism)
    mechanism.normf += normfμ(mechanism, ineqc)
    return
end

@inline function addNormΔs!(component::Component, mechanism::Mechanism)
    mechanism.normΔs += normΔs(component)
    return
end

@inline function discretizestate!(mechanism::Mechanism)
    foreach(discretizestate!, mechanism.bodies, mechanism.Δt)
    return
end

@inline function activate!(mechanism::Mechanism, id::Integer)
    component = getcomponent(mechanism, id)
    activate!(component)
    activate!(mechanism.graph,id)
    return
end

@inline function deactivate!(mechanism::Mechanism, id::Integer)
    component = getcomponent(mechanism, id)
    deactivate!(component)
    deactivate!(mechanism.graph,id)
    return
end

@inline function fdynamics(mechanism::LinearMechanism)
    A = mechanism.A
    Bu = mechanism.Bu
    Bλ = mechanism.Bλ
    z0 = mechanism.z
    z1 = mechanism.zsol[2]
    λ1 = mechanism.λsol[2]
    u = mechanism.u

    return A*z0 + Bu*u + Bλ*λ1 - z1
end

@inline function fconstraints(mechanism::LinearMechanism)
    G = mechanism.G
    z1 = mechanism.zsol[2]

    return G*z1
end

function disassemble(mechanism::Mechanism{T}) where T
    origin = mechanism.origin
    bodies = mechanism.bodies.values
    eqconstraints = mechanism.eqconstraints.values
    ineqconstraints = mechanism.ineqconstraints.values

    # Flip component ids
    for body in bodies
        body.id *= -1
    end
    for eqc in eqconstraints
        eqc.id *= -1
        if eqc.parentid !== nothing
            eqc.parentid *= -1
        end
        eqc.childids *= -1
    end
    for ineqc in ineqconstraints
        ineqc.id *= -1
        if ineqc.parentid !== nothing
            ineqc.parentid *= -1
        end
        ineqc.childids *= -1
    end

    # Set CURRENTID
    global CURRENTID = -1
    for body in bodies
        if body.id <= CURRENTID
            CURRENTID = body.id-1
        end
    end
    for eqc in eqconstraints
        if eqc.id <= CURRENTID
            CURRENTID = eqc.id-1
        end
    end
    for ineqc in ineqconstraints
        if ineqc.id <= CURRENTID
            CURRENTID = ineqc.id-1
        end
    end

    # Set origin to next id
    oldoid = origin.id
    origin.id = getGlobalID()
    for eqc in eqconstraints
        if eqc.parentid === nothing
            eqc.parentid = origin.id
        end
    end
    for ineqc in ineqconstraints
        if ineqc.parentid === nothing
            ineqc.parentid = origin.id
        end
    end

    return origin, bodies, eqconstraints, ineqconstraints
end


export DataSetofSets,
    DataFull,
    DataRoyaltyProduce,
    DataDrillOnly,
    split_thetas,
    merge_thetas

# -------------------------------------------------
# Data structure
# -------------------------------------------------

function check_theta(d::AbstractDataStructure, theta)
    _nparm(d) == length(theta) || throw(DimensionMismatch("_nparm(d) = $(_nparm(d)) != length(theta) = $(length(theta))"))
    all(isfinite.(theta)) || throw(error("theta = $theta not finite"))
    return true
end

"check coef restrictions fit in parameter vector"
function check_coef_restr_in_parm_vec(cf, d::AbstractDataSet)
    n = _nparm(d)
    issubset(Set(cf), OneTo(n)) || throw(error("coefs $cf not in 1:$n"))
    allunique(cf) || throw(error("coefs $cf not unique"))
    return true
end

struct DataSetofSets{
    D<:Union{AbstractDataDrill,EmptyDataSet},  R<:Union{DataRoyalty,EmptyDataSet}, P<:Union{DataProduce,EmptyDataSet},
    } <: AbstractDataSetofSets
    data::Tuple{D,R,P}
    coef_link_pdxn::Vector{Int}
    coef_link_drill::Vector{Int}
    function DataSetofSets(d::D, r::R, p::P, cfp::V, cfd::V) where {D,R,P,V<:Vector{Int}}
        drp = (d, r, p)
        lengths = length.(drp)
        issubset(lengths, (0, maximum(lengths),) ) || throw(DimensionMismatch("datasets must same or 0 lengths. lengths = $lengths"))
        length(cfp) == length(cfd) || throw(DimensionMismatch())
        check_coef_restr_in_parm_vec(cfp, p)
        check_coef_restr_in_parm_vec(cfd, d)
        any(lengths .> 0 ) || throw(error("one dataset must have nonzero length"))
        return new{D,R,P}(drp, cfp, cfd)
    end
end

# some versions
const DataFull = DataSetofSets{<:AbstractDataDrill, <:DataRoyalty, <:DataProduce}
const DataRoyaltyProduce = DataSetofSets{EmptyDataSet, <:DataRoyalty, <:DataProduce}
const DataDrillOnly = DataSetofSets{<:AbstractDataDrill, EmptyDataSet, EmptyDataSet}

# If no coef_links provided
DataSetofSets(d,r,p) = DataSetofSets(d,r,p,zeros(Int,0),zeros(Int,0))

function DataSetofSets(d,r,p, cfp::Vector, cfd::Vector)
    DataSetofSets(d,r,p, map(f -> f(p), cfp), map(f -> f(d), cfd))
end

function DataSetofSets(d,r,p, ct::Tuple{Vector,Vector})
    DataSetofSets(d,r,p,ct...)
end

function DataSetofSets(d,r,p,cfl::Vector{<:Tuple})
    return DataSetofSets(d,r,p, first.(cfl), last.(cfl))
end

function SimulationDraws(data::DataSetofSets, M)
    n = num_i(data)
    k = _nparm(drill(data))
    return SimulationDraws(M,n,k)
end

function DataSetofSets(dsos::DataSetofSets, ddrill::DataDrill)
    return DataSetofSets(ddrill, royalty(dsos), produce(dsos), coef_link_pdxn(dsos), coef_link_drill(dsos))
end

# -------------------------------------------------
# methods
# -------------------------------------------------

data(d::DataSetofSets) = d.data
coef_link_drill(d::DataSetofSets) = d.coef_link_drill
coef_link_pdxn(d::DataSetofSets) = d.coef_link_pdxn
coef_links(d::DataSetofSets) = zip(coef_link_pdxn(d), coef_link_drill(d))
length(d::DataSetofSets) = 3
num_i(d::DataSetofSets) = maximum(length.(data(d)))

iterate(d::DataSetofSets, state...) = iterate(data(d), state...)
Broadcast.broadcastable(d::DataSetofSets) = data(d)

drill(  d::DataSetofSets) = d.data[1] # drill
royalty(d::DataSetofSets) = d.data[2] # royalty
produce(d::DataSetofSets) = d.data[3] # produce

_nparm(d::DataFull)      = sum(_nparm.(d)) - length(coef_links(d)) - 1
_nparm(d::DataSetofSets) = sum(_nparm.(d)) - length(coef_links(d))

# @deprecate DataDrill(  d::DataSetofSets) drill(d)
# @deprecate DataRoyalty(d::DataSetofSets) royalty(d)
# @deprecate DataProduce(d::DataSetofSets) produce(d)

# -------------------------------------------------
# Indexing into parameter vector
# -------------------------------------------------

# if dataset is empty
idx_drill(  ::EmptyDataSet) = 1:0
idx_royalty(::EmptyDataSet) = 1:0
idx_produce(::EmptyDataSet) = 1:0

# standard indexing, no overlap
idx_drill(  data::DataSetofSets) = idx_drill(drill(data))
idx_royalty(data::DataSetofSets) = last(idx_drill(data)) .+ idx_royalty(royalty(data))
idx_produce(data::DataSetofSets) = last(idx_royalty(data)) .+ idx_produce(produce(data))

# theta_ρ(data::DataSetofSets{EmptyDataSet,<:DataRoyalty}, theta) = theta[1]
theta_ρ(data::DataSetofSets{<:AbstractDataDrill}, theta) = theta[_nparm(drill(data))]
theta_ρ(data::DataSetofSets{<:EmptyDataSet}, theta) = theta[1]

function merge_thetas(thetas::NTuple{3,AbstractVector}, data::DataFull)
    length.(thetas) == _nparm.(data) ||
        throw(DimensionMismatch("len thetas = $(length.(thetas)) but nparm = $(_nparm.(data))"))
    thet_d, thet_r, thet_p = thetas

    data_p = produce(data)
    idx_p = OneTo(_nparm(data_p))
    idx_p_drops = coef_link_pdxn(data)

    thet_p_short = [thet_p[i] for i in idx_p if i ∉ idx_p_drops]

    return vcat(thet_d, thet_r[2:end], thet_p_short)
end

merge_thetas(thetas, data) = vcat(thetas...)

@deprecate theta_linked(thetas::NTuple{3,AbstractVector}, data::DataFull) merge_thetas(thetas,data)

# full datasets
function idx_royalty(data::DataFull)
    kd = _nparm(drill(data))
    return (kd-1) .+ idx_royalty(royalty(data))
end

function idx_produce(data::DataFull)
    dd, dr, dp = drill(data), royalty(data), produce(data)
    kd = _nparm(dd)
    kr = _nparm(dr)
    idx = collect((kd+kr-1) .+ idx_produce(dp))
    for (idxp, idxd) in coef_links(data)
        idx[idxp] = idxd
        idx[idxp+1:end] .-= 1
    end
    return idx
end

function theta_indexes(data::DataSetofSets)
    d = idx_drill(  data)
    r = idx_royalty(data)
    p = idx_produce(data)
    return d, r, p
end

# to ensure that we get a copy... NOT a view
theta_produce(d::DataFull, theta) = theta[idx_produce(d)]

function split_thetas(data::DataSetofSets, theta::AbstractVector)
    d = theta_drill(  data, theta)
    r = theta_royalty(data, theta)
    p = theta_produce(data, theta)
    return d, r, p
end

@deprecate thetas(data::DataSetofSets, theta::AbstractVector) split_thetas(data,theta)
@deprecate split_theta(data, theta) split_thetas(data, theta)

function coefnames(data::DataSetofSets)
    nms = coefnames.(data)
    return merge_thetas(nms, data)
end


"""
    Origin{T} <: Node{T}

    Global reference frame

    id: always 0
    name: always :origin
    state: State; defaults to zero values
    shape: Shape; defaults to EmptyShape
"""
mutable struct Origin{T} <: Node{T}
    id::Int64
    name::Symbol
    state::State{T}
    shape::Shape{T}

    function Origin{T}(;
        name::Symbol=:origin,
        state=State{T}(),
        shape::Shape=EmptyShape()) where T
        new{T}(getGlobalID(), name, state, shape)
    end

    function Origin(body::Body{T}) where T
        new{T}(body.id, body.name, body.state, body.shape)
    end
end


############################   MAIN FUNCTIONS ###############################

# Main function that takes care of translating the model into source code
# of target language
function generate_code!(ode_model::dt.OdeModel; language = "julia",
                        unit_analysis = true, name = "autogenerated_model",
                        file = "autogenerated_model", jacobian = false,
                        sensitivities = false)

  # Create dictionaries containing the default values of the inputs (at the reference scale)
  # and the conversion coefficients for the units. This prepares the inputs for the
  # code generation functions
  coef_observed = OrderedDict{ASCIIString, Float64}()
  for i in ode_model.NamesObserved
    coef_observed[i] = ode_model.Equations[i].Units.f
  end
  named_states = OrderedDict{ASCIIString, Float64}()
  coef_states = OrderedDict{ASCIIString, Float64}()
  for (key,val) in ode_model.States
    named_states[key] = val.Value * val.Units.f
    coef_states[key] = val.Units.f
  end
  named_parameters = OrderedDict{ASCIIString, Float64}()
  coef_parameters = OrderedDict{ASCIIString, Float64}()
  for (key,val) in ode_model.Parameters
    named_parameters[key] = val.Value * val.Units.f
    coef_parameters[key] = val.Units.f
  end
  forcings = OrderedDict{ASCIIString,Any}()
  coef_forcings = OrderedDict{ASCIIString,Float64}()
  c = 0
  for (key,value) in ode_model.Forcings
      c += 1
      forcings[key] = (float(value.Time), float(value.Value)*value.Units.f)
      coef_forcings[key] = value.Units.f
  end

  # Write a file with the default inputs for the model in JSON format
  write_inputs!(named_states,coef_states, named_parameters, coef_parameters,
               forcings, coef_forcings, ode_model.NamesObserved,coef_observed,
               file)

  # Sort the equations of the model
  sorted_model = OdeSorted(ode_model)

  # Check the physical dimensions of the equations in the model (if demanded)
  unit_analysis &&  check_units(sorted_model)

  # Generate the rhs function (includes observed variables)
  # If DynamicType = DAE generate semiexplicit representation of dae model
  if (sorted_model.DynamicType == "ODE")
    model_function = create_derivatives(language, sorted_model, name)
    if(language == "cpp")
      observed_function = ""
    else
      observed_function = create_observed(language, sorted_model, name)
    end
  else
    error("Only DynamicType = \"ode\" is currently supported")
  end

  # Write down the code!
  write_code!(language, ode_model.DynamicType, model_function, observed_function, name, file)

  nothing

end

############################ WRITE TIME DERIVATIVES ################################

# This function writes the source for calculating the time derivatives and observed variables
# and it dispatches the correct function depending on the language argument
function create_derivatives(language, sorted_model, name)
  if (language == "julia")
    return create_derivatives_julia(sorted_model, name)
  elseif (language == "cpp")
    return create_derivatives_rcpp(language, sorted_model, name)
  else
    error("Currently only \"julia\" and \"cpp\" are supported")
  end
end

############################ WRITE INDEPENDENT OBSERVER ################################

# This function writes the source for calculating the time derivatives and observed variables
# and it dispatches the correct function depending on the language argument
function create_observed(language, sorted_model, name)
  if (language == "julia")
    return create_observed_julia(sorted_model, name)
  else
    error("Currently only \"julia\" and \"cpp\" are supported")
  end
end

############################### WRITE CODE FILES ###################################

# Write a JSON file with all the inputs of the model.
# This should be readable in every language
function write_inputs!(named_states,coef_states, named_parameters, coef_parameters,
                       forcings, coef_forcings, observed,coef_observed,
                       file)
  f = open("$(file)_inputs.json","w")
  inputs = {"states" => named_states,
            "coef_states" => coef_states,
            "parameters" => named_parameters,
            "coef_parameters" => coef_parameters,
            "forcings" => forcings,
            "coef_forcings" => coef_forcings,
            "observed" => observed,
            "coef_observed" => coef_observed}
  JSON.print(f, inputs)
  close(f)
end

# Write a source file with all the functions associated to the model in the target language
function write_code!(language, dynamic_type, model_function, observed_function, name, file)

  if (language == "julia")
    return write_code_julia!(dynamic_type, model_function, observed_function, name, file)
  elseif (language == "cpp")
    return write_code_rcpp!(dynamic_type, model_function, name, file)
  end

end