"""
    UnaryOperation{X, Y, Z, O, A, I, G} <: AbstractOperation{X, Y, Z, G}

An abstract representation of a unary operation on an `AbstractField`; or a function
`f(x)` with on argument acting on `x::AbstractField`.
"""
struct UnaryOperation{X, Y, Z, O, A, I, G} <: AbstractOperation{X, Y, Z, G}
      op :: O
     arg :: A
       ▶ :: I
    grid :: G

    """
        UnaryOperation{X, Y, Z}(op, arg, ▶, grid)

    Returns an abstract `UnaryOperation` representing the action of `op` on `arg`,
    and subsequent interpolation by `▶` on `grid`.
    """
    function UnaryOperation{X, Y, Z}(op, arg, ▶, grid) where {X, Y, Z}
        return new{X, Y, Z, typeof(op), typeof(arg), typeof(▶), typeof(grid)}(op, arg, ▶, grid)
    end
end

"""Create a unary operation for `operator` acting on `arg` which interpolates the
result from `Larg` to `L`."""
function _unary_operation(L, operator, arg, Larg, grid)
    ▶ = interpolation_operator(Larg, L)
    return UnaryOperation{L[1], L[2], L[3]}(operator, data(arg), ▶, grid)
end

@inline Base.getindex(υ::UnaryOperation, i, j, k) = υ.▶(i, j, k, υ.grid, υ.op, υ.arg)

"""
    @unary op1 op2 op3...

Turn each unary function in the list `(op1, op2, op3...)`
into a unary operator on `Oceananigans.Fields` for use in `AbstractOperations`.

Note: a unary function is a function with one argument: for example, `sin(x)` is a unary function.

Also note: a unary function in `Base` must be imported to be extended: use `import Base: op; @unary op`.

Example
=======

julia> square_it(x) = x^2
square_it (generic function with 1 method)

julia> @unary square_it
7-element Array{Any,1}:
 :sqrt
 :sin
 :cos
 :exp
 :tanh
 :-
 :square_it

julia> c = Field(Cell, Cell, Cell, CPU(), RegularCartesianGrid((1, 1, 16), (1, 1, 1)));

julia> square_it(c)
UnaryOperation at (Cell, Cell, Cell)
├── grid: RegularCartesianGrid{Float64,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}
│   ├── size: (1, 1, 16)
│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [0.0, -1.0]
└── tree:

square_it at (Cell, Cell, Cell) via identity
└── OffsetArrays.OffsetArray{Float64,3,Array{Float64,3}}
"""
macro unary(ops...)
    expr = Expr(:block)

    for op in ops
        define_unary_operator = quote
            import Oceananigans: AbstractGrid
            import Oceananigans.Fields: AbstractLocatedField

            local location = Oceananigans.Fields.location

            @inline $op(i, j, k, grid::AbstractGrid, a) = @inbounds $op(a[i, j, k])
            @inline $op(i, j, k, grid::AbstractGrid, a::Number) = $op(a)

            """
                $($op)(Lop::Tuple, a::AbstractLocatedField)

            Returns an abstract representation of the operator `$($op)` acting on the Oceananigans `Field`
            `a`, and subsequently interpolated to the location indicated by `Lop`.
            """
            function $op(Lop::Tuple, a::AbstractLocatedField)
                L = location(a)
                return Oceananigans.AbstractOperations._unary_operation(Lop, $op, a, L, a.grid)
            end

            $op(a::AbstractLocatedField) = $op(location(a), a)

            push!(Oceananigans.AbstractOperations.operators, Symbol($op))
            push!(Oceananigans.AbstractOperations.unary_operators, Symbol($op))
        end

        push!(expr.args, :($(esc(define_unary_operator))))
    end

    return expr
end

const unary_operators = Set()

"Adapt `UnaryOperation` to work on the GPU via CUDAnative and CUDAdrv."
Adapt.adapt_structure(to, unary::UnaryOperation{X, Y, Z}) where {X, Y, Z} =
    UnaryOperation{X, Y, Z}(Adapt.adapt(to, unary.op), Adapt.adapt(to, unary.arg),
                            Adapt.adapt(to, unary.▶), unary.grid)


# complement = -1*
function cmpl(x::CachedResult)
    return CachedResult(-1*x.value, -1*x.signint, x.index)
end

# union = min
# Base.min(x::CachedResult, y::CachedResult) = ifelse(isless(x.value, y.value), x, y)
uni(x::CachedResult, y::CachedResult) = ifelse(isless(x.value, y.value), x, y)

# intersection = max
# Base.max(x::CachedResult, y::CachedResult) = ifelse(isless(y.value, x.value), x, y)
inters(x::CachedResult, y::CachedResult) = ifelse(isless(y.value, x.value), x, y)

# subtraction f,g = f union complement(g)
subtr(x::CachedResult, y::CachedResult) = intersection(x, complement(y))

const opDictCode = Dict(:complement=>cmpl,
                :intersection=>inters,
                :subtraction=>subtr,
                :union=>uni)

function evalf(surface::CachedSurface, cachedvalues, ind)
    val = cachedvalues[ind][surface.index]
    return CachedResult(val, 1, surface.index)
end

function leaf2code(l::CachedCSGNode, _cvals, _ind)
    ex = Expr(Symbol("="), l.sym, Expr(:call, evalf, l.data, _cvals, _ind))
    return ex
end

function node2code(l::CachedCSGNode)
    arg1 = l.children[1].sym
    op = l.data
    if op === :complement
        ex = Expr(Symbol("="), l.sym, Expr(:call, opDictCode[:complement], arg1))
        return ex
    else
        arg2 = l.children[2].sym
        ex = Expr(Symbol("="), l.sym, Expr(:call, opDictCode[op], arg1, arg2))
        return ex
    end
end

function tree2code(tree::CachedCSGNode)
    _cachedv = gensym()
    _inds = gensym()
    exprs = Array{Expr,1}(undef,0)
    for l in Set{CachedCSGNode}(Leaves(tree))
        push!(exprs, leaf2code(l, _cachedv, _inds))
    end
    for n in PostOrderDFS(tree)
        isempty(n.children) && continue
        push!(exprs, node2code(n))
    end
    return CodeWrap(exprs, [_cachedv, _inds])
end

struct CodeWrap
    Core::Array{Expr,1}
    Params::Array{Symbol,1}
end

function code2func(cwrap)
    expr = Expr(:function,
        Expr(:tuple,
        cwrap.Params[1],
        cwrap.Params[2]),
        Expr(:block,cwrap.Core...))
    return eval(expr)
end

tree2func(tree::CachedCSGNode) = code2func(tree2code(tree))


__precompile__(true)
module mdl

using JuliaParser.Diagnostics
using JuliaParser.Lexer
using JuliaParser.Lexer: @tok, TokenStream, eof, eof_token, make_token, peekchar
using JuliaParser.Tokens: ¬, √
using Render
using Render: ParseState, modifycoord!, withem!, withpm!

export mdl_execute, mdl_parser

COMMENT = "//"
peekcomment(ts::TokenStream) = Render.peekcomment(ts, COMMENT)

function next_token{T}(ts::TokenStream{T})
    ts.ateof && return eof_token(ts)
    tmp = Lexer.skipws(ts, false)
    eof(ts) && return eof_token(ts)
    ts.isspace = tmp

    while ! eof(ts)
        peekcomment(ts) && skipcomment(ts, COMMENT)
        c = peekchar(ts)
        if eof(c)
            ts.ateof = true
            return eof_token(T)
        elseif c === ' ' || c === '\t'
            Lexer.skip(ts, 1)
            continue
        elseif c == '\r'
            Lexer.skip(ts, 1)
            peekchar(ts) === '\n' || throw(Lexer.here(ts),
                                           "'\\r' not followed by '\\n' is invalid")
            continue
        elseif Lexer.isnewline(c)
            return @tok Lexer.readchar(ts)
        elseif isdigit(c)
            return Lexer.read_number(ts, false, false)
        elseif c === '-'
            Lexer.skip(ts, 1)
            d = peekchar(ts)
            return isdigit(d) ? Lexer.read_number(ts, d === '.', true) : @tok :(-)
        elseif c === '.'
            Lexer.skip(ts, 1)
            return isdigit(peekchar(ts)) ? Lexer.read_number(ts, true, false) : @tok :(.)
        elseif Lexer.is_identifier_start_char(c)
            return @tok Lexer.accum_julia_symbol(ts, c)
        else
            @assert Lexer.readchar(ts) === c
            throw(Diagnostics.diag(Lexer.here(ts), "invalid character \"$c\""))
        end
    end

    ts.ateof = true
    return eof_token(ts)
end

# simple substitution
const F64 = Float64

# MDL commands
mdlpush(ps::ParseState) = push!(ps.coords, ps.coords[end])
mdlpop(ps::ParseState)  = length(ps.coords) > 1 && pop!(ps.coords)

knobworks(ps, k, thing) = k === nothing ? thing : thing * ps.symtab[k][ps.cframe]

mdlmove(ps::ParseState, x::F64, y::F64, z::F64, knob = nothing) = begin
    modifycoord!(ps, mktranslate(knobworks(ps, knob, [x, y, z])))
end

mdlscale(ps::ParseState, x::F64, y::F64, z::F64, knob = nothing) = begin
    modifycoord!(ps, mkscale(knobworks(ps, knob, [x, y, z])))
end

mdlrotate(ps::ParseState, dir::Symbol, d::F64, knob = nothing) = begin
    modifycoord!(ps, mkrotate(knobworks(ps, knob, d), dir))
end

# optional not at end
mdlsphere(ps::ParseState, cons, x::F64, y::F64, z::F64, r::F64, coord = nothing) = begin
    withpm!(ps, addsphere!, [x, y, z], r, SPHSTEPS)
end
mdlsphere(ps::ParseState, x::F64, y::F64, z::F64, r::F64, coord = nothing) = mdlsphere(ps, nothing, x, y, z, r, coord)

mdltorus(ps::ParseState, cons, x::F64, y::F64, z::F64, r1::F64, r2::F64, coord = nothing) = begin
    withpm!(ps, addtorus!, [x, y, z], r1, r2, TORUSTEPS)
end
mdltorus(ps::ParseState, x::F64, y::F64, z::F64, r1::F64, r2::F64, coord = nothing) = mdltorus(ps, nothing, x, y, z, r1, r2, coord)

mdlbox(ps::ParseState, cons, x1::F64, y1::F64, z1::F64, x2::F64, y2::F64, z2::F64, coord = nothing) = begin
    withpm!(ps, addbox!, [x1, y1, z1], x2, y2, z2)
end
mdlbox(ps::ParseState, x1::F64, y1::F64, z1::F64, x2::F64, y2::F64, z2::F64, coord = nothing) = mdlbox(ps, nothing, x1, y1, z1, x2, y2, z2, coord)

# there should be a cleaner way to do this
mdlline(ps::ParseState, cons, x1::F64, y1::F64, z1::F64, coord1::Symbol, x2::F64, y2::F64, z2::F64, coord2 = nothing) = begin
    withem!(ps, addedge!, [x1, y1, z1], [x2, y2, z2])
end
mdlline(ps::ParseState, x1::F64, y1::F64, z1::F64, coord1::Symbol, x2::F64, y2::F64, z2::F64, coord2 = nothing) = mdlline(ps, nothing, x1, y1, z1, coord1, x2, y2, z2, coord2)
mdlline(ps::ParseState, cons::Symbol, x1::F64, y1::F64, z1::F64, x2::F64, y2::F64, z2::F64, coord2 = nothing) = mdlline(ps, cons, x1, y1, z1, nothing, x2, y2, z2, coord2)
mdlline(ps::ParseState, x1::F64, y1::F64, z1::F64, x2::F64, y2::F64, z2::F64, coord2 = nothing) = mdlline(ps, nothing, x1, y1, z1, nothing, x2, y2, z2, coord2)

mdlbasename(ps::ParseState, bn...) = ps.basename = string(bn...)
mdlframes(ps::ParseState, f::F64) = ps.nframes = Int(f)

mdlvary(ps::ParseState, knob::Symbol, sf, ef, sv::F64, ev::F64) = begin
    sf = Int(sf)
    ef = Int(ef)
    ps.symtab[knob] = get(ps.symtab, knob, Render.Knob(sf, ef, sv, ev, ps.nframes))
    ps.symtab[knob].frdata[sf:ef - 1] = [sv + (ev - sv) / (ef - sf) * i for i in 1:(ef - sf)]
end

mdlsave(ps::ParseState, fn...) = Render.display(ps, `convert - $(string(fn...))`)

mdldisplay(ps::ParseState) = begin
    if ps.nframes > 1
        run(`animate -delay $(Render.ANIM_DELAY) "$tmp/$(ps.basename)*"`)
    else
        Render.display(ps, `display`)
    end
end

function mdl_parser(ts::TokenStream, ps::ParseState)
    Render.skipws_and_comments(ts, COMMENT)
    eof(ts) && return

    args = Vector{Any}()
    t = Render.require_token(ts, next_token)
    command = nothing

    try
        command = getfield(mdl, Symbol(:mdl, ¬t))
    catch UndefVarError
        error("$(¬t) is not a valid MDL command")
    end

    while ! eof(ts)
        x = next_token(ts)
        push!(args, ¬x isa Number ? float(¬x) : ¬x)
    end

    push!(ps.commands, Render.GCommand(command, args))
end

function mdl_execute(ps::ParseState)
    # preprocess "frames", "basename", "vary"
    i = 1
    while i <= length(ps.commands)
        c = ps.commands[i]
        if c.func === mdlframes || c.func === mdlbasename
            c.func(ps, c.args...)
            deleteat!(ps.commands, i)
        else
            i += 1
        end
    end

    i = 1
    while i <= length(ps.commands)
        c = ps.commands[i]
        if c.func === mdlvary
            mdlvary(ps, c.args...)
            deleteat!(ps.commands, i)
        else
            i += 1
        end
    end

    mktempdir() do tmp
        reference = ps
        @sync @parallel for current in 1:reference.nframes
            if reference.nframes > 1
                println("Processing frame ", current)
                ps = deepcopy(reference)
            end

            ps.cframe = current
            for c in ps.commands
                c.func(ps, c.args...)
            end

            ps.nframes > 1 && mdlsave(ps, joinpath(tmp, string(ps.basename, lpad(current, 4, 0), ".png")))
        end
    end
end
end


@testset "Function testing" begin

@testset ":Options" begin
    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        DefaultTestSolver(;traverse_strategy=:DBFS,obj_epsilon=0.5)...
        )
    )

    v = [10,20,12,23,42]
    w = [12,45,12,22,21]
    @variable(m, x[1:5], Bin)

    @objective(m, Max, dot(v,x))

    @NLconstraint(m, sum(w[i]*x[i]^2 for i=1:5) <= 45)

    MOIU.attach_optimizer(m)
    bm = JuMP.backend(m)
    options = bm.optimizer.model.optimizer.options

    nd_options = Juniper.get_non_default_options(options)
    @test nd_options[:obj_epsilon] == 0.5
    @test nd_options[:traverse_strategy] == :DBFS
    @test length(nd_options[:log_levels]) == 0
end

function option_not_available_t()
    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        DefaultTestSolver(;traverse_strategy=:DBS,obj_epsilon=0.5)...
        )
    )
end

function option_not_available_b()
    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        DefaultTestSolver(;branch_strategy=:Pseudo,obj_epsilon=0.5)...
        )
    )
end

function option_not_available()
    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        DefaultTestSolver(;branch=:Pseudo,obj_epsilon=0.5)...
        )
    )
end

function option_no_mip_solver()
    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        DefaultTestSolver(;branch=:Pseudo,obj_epsilon=0.5,feasibility_pump=true)...
        )
    )

    v = [10,20,12,23,42]
    w = [12,45,12,22,21]
    @variable(m, x[1:5], Bin)

    @objective(m, Max, dot(v,x))

    @NLconstraint(m, sum(w[i]*x[i]^2 for i=1:5) <= 45)
    # there it is detected that no mip solver is present
    optimize!(m)

    bm = JuMP.backend(m)
    return bm.optimizer.model.optimizer.options
end

@testset "Silent/TimeLimitSec" begin
    optimizer = Juniper.Optimizer(;nl_solver=optimizer_with_attributes(Ipopt.Optimizer))
    @test MOI.supports(optimizer, MOI.Silent()) === true
    @test MOI.supports(optimizer, MOI.TimeLimitSec()) === true
    MOI.set(optimizer, MOI.RawOptimizerAttribute("mip_gap"), 1.0)
    # parameter doesn't exist
    @test_logs (:error, r"doesn't exist") MOI.set(optimizer, MOI.RawOptimizerAttribute("mip_gap_1"), 1.0)
    # wrong parameter format
    @test_logs (:error, r"different type") MOI.set(optimizer, MOI.RawOptimizerAttribute("mip_gap"), "abc")
    MOI.set(optimizer, MOI.Silent(), true)
    MOI.set(optimizer, MOI.TimeLimitSec(), nothing)
    @test MOI.get(optimizer, MOI.Silent()) === true
    @test MOI.get(optimizer, MOI.RawOptimizerAttribute("mip_gap")) == 1.0
    @test_logs (:error, r"doesn't exist") MOI.get(optimizer, MOI.RawOptimizerAttribute("mip_gap_1"))
    @test isinf(MOI.get(optimizer, MOI.TimeLimitSec()))
    MOI.set(optimizer, MOI.TimeLimitSec(), 12.0)
    @test MOI.get(optimizer, MOI.TimeLimitSec()) == 12.0
end

@testset ":Option not available" begin
    @test_logs (:error, r"not supported") option_not_available_t()
    @test_logs (:error, r"not supported") option_not_available_b()
    opts = option_no_mip_solver()
    @test opts.feasibility_pump == false
    @test !isa(try option_not_available() catch ex ex end, Exception)
end

@testset "Info/Table" begin
    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        DefaultTestSolver(;branch_strategy=:StrongPseudoCost,processors=2, strong_restart=true)...
        )
    )

    v = [10,20,12,23,42]
    w = [12,45,12,22,21]
    @variable(m, x[1:5], Bin)

    @objective(m, Max, dot(v,x))

    @NLconstraint(m, sum(w[i]*x[i]^2 for i=1:5) <= 45)

    MOIU.attach_optimizer(m)
    bm = JuMP.backend(m)
    JuMP.optimize!(m)
    innermodel = internalmodel(m)

    options = innermodel.options
    jp = innermodel
    println("typeof(m): ", typeof(m))
    println("typeof(jp): ", typeof(jp))

    @test !isa(try Juniper.print_info(jp) catch ex ex end, Exception)
    @test !isa(try Juniper.print_options(jp;all=true) catch ex ex end, Exception)
    @test !isa(try Juniper.print_options(jp;all=false) catch ex ex end, Exception)

    fields, field_chars = Juniper.get_table_config(options)
    ln = Juniper.get_table_header_line(fields, field_chars)
    @test length(ln) == sum(field_chars)

    start_time = time()
    tree = Juniper.init(start_time,jp)
    tree.best_bound = 100
    node = Juniper.new_default_node(1,1,zeros(Int64,5),ones(Int64,5),zeros(Int64,5))
    step_obj = Juniper.new_default_step_obj(jp,node)
    step_obj.counter = 1
    tab_ln, tab_arr = Juniper.get_table_line(2,tree,node,step_obj,start_time,fields,field_chars;last_arr=[])
    @test length(fields) == length(field_chars)
    for i in 1:length(fields)
        fc = field_chars[i]
        f = string(fields[i])
        @test fc >= length(f)
    end
    # Test with incumbent
    tree.incumbent = Juniper.Incumbent(42,[0,0,0,0,1],false)
    tree.limit = :Time
    tab_ln, tab_arr = Juniper.get_table_line(2,tree,node,step_obj,start_time,fields,field_chars;last_arr=[])
    @test length(fields) == length(field_chars)
    for i in 1:length(fields)
        fc = field_chars[i]
        f = string(fields[i])
        @test fc >= length(f)
    end
    @test length(fields) == length(tab_arr)
    @test length(tab_ln) == sum(field_chars)

    # test for diff
    tab_arr_new = copy(tab_arr)
    tab_arr_new[1] = 100
    @test Juniper.is_table_diff(fields, tab_arr, tab_arr_new) == true
    @test Juniper.is_table_diff(fields, tab_arr, tab_arr) == false

    # test if :Time not exists
    i = 1
    for f in fields
        if f == :Time
            deleteat!(fields,i)
            deleteat!(field_chars,i)
        end
        i+= 1
    end
    tab_ln, tab_arr = Juniper.get_table_line(2,tree,node,step_obj,start_time,fields,field_chars;last_arr=[])
    tab_arr_new = copy(tab_arr)
    tab_arr_new[1] = 100
    @test Juniper.is_table_diff(field_chars, tab_arr, tab_arr_new) == true
    @test Juniper.is_table_diff(field_chars, tab_arr, tab_arr) == false

    # produce table line with high counter for nrestarts
    step_obj.counter = 100
    tab_ln, tab_arr = Juniper.get_table_line(2,tree,node,step_obj,start_time,fields,field_chars;last_arr=[])
    idx_restarts = findfirst(fields .== :Restarts)
    @test tab_arr[idx_restarts] == "-"

    # normal gain gap
    step_obj.gain_gap = 0.05
    tab_ln, tab_arr = Juniper.get_table_line(2,tree,node,step_obj,start_time,fields,field_chars;last_arr=[])
    idx_gain_gap = findfirst(fields .== :GainGap)
    @test tab_arr[idx_gain_gap] == "5.0%" || tab_arr[idx_gain_gap] == "5.00%"

    # Inf gain gap
    step_obj.gain_gap = Inf
    tab_ln, tab_arr = Juniper.get_table_line(2,tree,node,step_obj,start_time,fields,field_chars;last_arr=[])
    idx_gain_gap = findfirst(fields .== :GainGap)
    @test tab_arr[idx_gain_gap] == "∞"

    # Huge gain gap
    step_obj.gain_gap = 123456789000
    tab_ln, tab_arr = Juniper.get_table_line(2,tree,node,step_obj,start_time,fields,field_chars;last_arr=[])
    idx_gain_gap = findfirst(fields .== :GainGap)
    @test tab_arr[idx_gain_gap] == ">>"

    # very large gap
    tree.best_bound = 1000000
    tree.incumbent.objval = 1
    tab_ln, tab_arr = Juniper.get_table_line(2,tree,node,step_obj,start_time,fields,field_chars;last_arr=[])
    idx_gap = findfirst(fields .== :Gap)
    @test tab_arr[idx_gap] == ">>"

    # way too long
    fields, field_chars = Juniper.get_table_config(options)
    start_time = time()-123456789
    tab_ln, tab_arr = Juniper.get_table_line(2,tree,node,step_obj,start_time,fields,field_chars;last_arr=[])
    idx_time = findfirst(fields .== :Time)
    @test tab_arr[idx_time] == "t.l."
end

@testset "FP: Table config" begin
    mip_obj,nlp_obj,t, fields, field_chars = 5,2,2, [:MIPobj,:NLPobj,:Time], [20,20,10]
    ln, arr = Juniper.get_fp_table(mip_obj,nlp_obj,t, fields, field_chars, 1e-6)
    @test length(ln) == 50
    @test arr == ["5.0","2.0","2.0"]
end

@testset "Random restarts" begin
    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        DefaultTestSolver(;branch_strategy=:StrongPseudoCost,processors=2, strong_restart=true)...
        )
    )

    v = [10,20,12,23,42]
    w = [12,45,12,22,21]
    @variable(m, x[1:5], Int)
    JuMP.set_lower_bound(x[1], 0)
    JuMP.set_lower_bound(x[2], 0.5)
    JuMP.set_upper_bound(x[3], 1)
    JuMP.set_upper_bound(x[2], 1)

    @objective(m, Max, dot(v,x))

    @NLconstraint(m, sum(w[i]*x[i]^2 for i=1:5) <= 45)

    MOIU.attach_optimizer(m)
    bm = JuMP.backend(m)
    JuMP.optimize!(m)
    model = internalmodel(m)

    cont_restart = Juniper.generate_random_restart(model)
    @test length(cont_restart) == 5
    @test 0 <= cont_restart[1] <= 20
    @test 0.5 <= cont_restart[2] <= 1
    @test -19 <= cont_restart[3] <= 1
    @test -10 <= cont_restart[4] <= 10
    @test -10 <= cont_restart[5] <= 10

    disc_restart = Juniper.generate_random_restart(model; cont=false)
    @test length(disc_restart) == 5
    for i=1:5
        @test isapprox(round(disc_restart[i])-disc_restart[i],0,atol=1e-6)
    end

    # test a different seed

    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        DefaultTestSolver(;branch_strategy=:StrongPseudoCost,processors=2, strong_restart=true, seed=92301)...
        )
    )

    v = [10,20,12,23,42]
    w = [12,45,12,22,21]
    @variable(m, x[1:5], Int)
    JuMP.set_lower_bound(x[1], 0)
    JuMP.set_lower_bound(x[2], 0.5)
    JuMP.set_upper_bound(x[3], 1)
    JuMP.set_upper_bound(x[2], 1)

    @objective(m, Max, dot(v,x))

    @NLconstraint(m, sum(w[i]*x[i]^2 for i=1:5) <= 45)

    MOIU.attach_optimizer(m)
    bm = JuMP.backend(m)
    JuMP.optimize!(m)
    model = internalmodel(m)

    cont_restart2 = Juniper.generate_random_restart(model)
    # all random numbers should be different then before
    for i=1:5
        @test cont_restart[i] != cont_restart2[i]
    end

    # test a different seed

    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        DefaultTestSolver(;branch_strategy=:StrongPseudoCost,processors=2, strong_restart=true, seed=92301)...
        )
    )

    v = [10,20,12,23,42]
    w = [12,45,12,22,21]
    @variable(m, x[1:5], Int)
    JuMP.set_lower_bound(x[1], 0)
    JuMP.set_lower_bound(x[2], 0.5)
    JuMP.set_upper_bound(x[3], 1)
    JuMP.set_upper_bound(x[2], 1)

    @objective(m, Max, dot(v,x))

    @NLconstraint(m, sum(w[i]*x[i]^2 for i=1:5) <= 45)

    MOIU.attach_optimizer(m)
    bm = JuMP.backend(m)
    JuMP.optimize!(m)
    model = internalmodel(m)

    cont_restart3 = Juniper.generate_random_restart(model)
    # all random numbers should be the same as before
    for i=1:5
        @test cont_restart2[i] == cont_restart3[i]
    end
end

end


function by_single_type(t::Symbol)
    function single_type(model::ABM)
        ids = collect(allids(model))
        filter!(id -> model[id].type === t, ids)
        return ids
    end
    return single_type
end

function nearby_t(t::Symbol, id, model)
    ids = collect(nearby_ids(model[id].pos, model, model.properties.infection_distance))
    filter!(id -> model[id].type === t, ids)
    return ids
end

function random_without_bacterium(model::ABM{<:DiscreteSpace}, cutoff=0.998)
    function has_no_bacterium(pos, model)
        ids = ids_in_position(pos, model)
        return !any(id -> model[id].type === :bacterium, ids)
    end

    function no_bacterium_positions(model::ABM{<:DiscreteSpace})
        return Iterators.filter(i -> has_no_bacterium(i, model), positions(model))
    end

    if clamp(nagents(model) / prod(size(model.space.s)), 0.0, 1.0) < cutoff
        while true
            pos = random_position(model)
            has_no_bacterium(pos, model) && return pos
        end
    else
        no_bacterium = no_bacterium_positions(model)
        isempty(no_bacterium) && return ()
        return rand(model.rng, collect(no_bacterium))
    end
end

function random_without_bacterium(model::ABM{<:DiscreteSpace}, positions::Vector{Tuple{Int,Int}})
    function has_no_bacterium(pos, model)
        ids = ids_in_position(pos, model)
        return !any(id -> model[id].type === :bacterium, ids)
    end

    function no_bacterium_positions(model::ABM{<:DiscreteSpace}, positions)
        return Iterators.filter(i -> has_no_bacterium(i, model), positions)
    end

    no_bacterium = no_bacterium_positions(model, positions)
    isempty(no_bacterium) && return ()
    return rand(model.rng, collect(no_bacterium))
end

function add_bacterium_single!(
    agent::A,
    model::ABM{<:DiscreteSpace,A};
    cutoff=0.998
) where {A<:AbstractAgent}
    position = random_without_bacterium(model, cutoff)
    isempty(position) && return
    agent.pos = position
    add_agent_pos!(agent, model) >
    return agent
end

function move_bacterium_single!(
    agent::A,
    model::ABM{<:DiscreteSpace,A},
    positions
) where {A<:AbstractAgent}
    position = random_without_bacterium(model, positions)
    isempty(position) && return
    move_agent!(agent, position, model)
    return agent
end


module EMLpfr
	using DanaTypes
	using DotPlusInheritance
	using Reexport
	@reexport using ...streams.EMLstreams
	import EMLtypes.length
	include("pfr/pfr.jl")
end


module NonlinearBandits

using Flux:
    relu, Chain, Dense, gpu, @epochs, train!, ADAM, DataLoader, throttle, params, cpu
using Distributions: Uniform, Normal, loggamma, InverseGamma, MvNormal
using GLMNet: glmnet
using LinearAlgebra: Hermitian, diagm, logdet
using Random: randperm
using Suppressor: @suppress

include("utils.jl")
include("BayesLM.jl")
include("BayesPM.jl")
include("PartitionedBayesPM.jl")
include("bandits.jl")
include("contexts.jl")
include("rewards.jl")
include("metrics.jl")
include("drivers.jl")
include("RandomPolicy.jl")
include("PolynomialThompsonSampling.jl")
include("NeuralLinear.jl")
include("NearestNeighboursUCB.jl")

"""
    fit!(model, X::AbstractMatrix, y::AbstractMatrix)

Update the parameters of `model`.

# Arguments

- `X::AbstractMatrix`: A matrix with observations stored as columns.
- `y::AbstractMatrix`: A matrix with 1 row of response variables.
"""
function fit!(model, X::AbstractMatrix, y::AbstractMatrix)
    throw(ErrorException("must implement fit! for $(typeof(model))"))
end

"""
    shape_scale(model::AbstractBayesianLM)

Return the shape/scale of model.
"""
function shape_scale(model::AbstractBayesianLM)
    throw(ErrorException("must implement shape_scale for $(typeof(model))"))
end

"""
    update!(pol::AbstractPolicy, X::AbstractMatrix, a::AbstractVector{<:Int},
            r::AbstractMatrix)

Update `pol` with a batch of data.
"""
function update!(
    pol::AbstractPolicy, X::AbstractMatrix, a::AbstractVector{<:Int}, r::AbstractMatrix
)
    throw(ErrorException("must implement update! for $(typeof(pol))"))
end

export AbstractBayesianLM,
    AbstractContextSampler,
    AbstractDriver,
    AbstractMetric,
    AbstractPolicy,
    AbstractRewardSampler,
    append_data!,
    arm_data,
    BanditDataset,
    BayesLM,
    BayesPM,
    expand,
    fit!,
    FunctionalRegret,
    GaussianRewards,
    Index,
    lasso_selection,
    LatentDriver,
    locate,
    nearest_neighbours,
    NearestNeighboursUCB,
    NeuralEncoder,
    NeuralLinear,
    Partition,
    PartitionedBayesPM,
    PolynomialThompsonSampling,
    RandomPolicy,
    run!,
    shape_scale,
    split!,
    StandardDriver,
    std,
    tpbasis,
    UniformContexts,
    update!

end


#=
The Vectors defined here are necessary to
ClapeyronParam
=#
struct Compressed4DMatrix{T,V<:AbstractVector{T}}
    values::V
    outer_indices::Vector{Tuple{Int,Int}} #index of components
    inner_indices::Vector{Tuple{Int,Int}} #index of sites
    outer_size::Tuple{Int,Int} #size of component matrix
    inner_size::Tuple{Int,Int} #size of sites matrices
end

function Base.show(io::IO,mime::MIME"text/plain",m::Compressed4DMatrix{T}) where T
    n = length(m.values)
    println(io,typeof(m)," with ",n," entr",(n == 1 ? "y:" : "ies:"))
    for (idx,(i,j),(a,b)) in indices(m)
        if idx != 1
        println(io)
        end
        print(io," ",(i,a)," >=< ",(j,b),": ",m.values[idx])
    end
end

function Base.show(io::IO,m::Compressed4DMatrix{T}) where T
    print(io,typeof(m))
    print(io,m.values)
end
function Compressed4DMatrix{T}() where T
    return Compressed4DMatrix(T[],Tuple{Int,Int}[],Tuple{Int,Int}[],(0,0),(0,0))
end

const MatrixofMatrices{T} = AbstractMatrix{<:AbstractMatrix{T}} where T

function Compressed4DMatrix(x::MatrixofMatrices{T}) where T
    outer_size = size(x)
    is1, is2 = 0, 0
    os1, os2 = outer_size
    for i in eachindex(x)
        _is1,_is2 = size(x[i])
        is1,is2 = max(_is1,is1),max(_is2,is2)
    end
    inner_size = (is1,is2)

    values = T[]
    inner_indices = Tuple{Int,Int}[]
    outer_indices = Tuple{Int,Int}[]

    if iszero(os1) & iszero(os2)
        return Compressed4DMatrix(values,outer_indices,inner_indices,outer_size,inner_size)
    end

    #self association

    for i in 1:os1
        for j in i:os2 #there can be self association
            xi = x[i,j]
            if iszero(length(xi))
                continue
            end
            a1,a2 = size(xi)
            for a in 1:a1
                start = ifelse(i == j,a,1)
                for b in start:a2 #this includes (i,i)(a,a) (sCKSAFT)
                    if !_iszero(xi[a,b])
                        push!(values,xi[a,b])
                        push!(outer_indices,(i,j))
                        push!(inner_indices,(a,b))
                    end
                end
            end
        end
    end
    return Compressed4DMatrix{T,Vector{T}}(values,outer_indices,inner_indices,outer_size,inner_size)
end

function Base.getindex(m::Compressed4DMatrix,i::Int,j::Int)
    i,j = minmax(i,j)
    @inbounds begin
    idx = searchsorted(m.outer_indices,(i,j))
    return AssocView(view(m.values,idx),view(m.inner_indices,idx),m.inner_size)
    end
end

#Base.eltype(m::Compressed4DMatrix{T}) where T = T

function Base.setindex!(m::Compressed4DMatrix,val,i::Int)
    @inbounds begin
        m.values[i] = val
    end
end

struct AssocView{T,V,I} <: AbstractMatrix{T}
    values::V
    indices::I
    size::Tuple{Int,Int}
end

function AssocView(V::Vi,I::Ii,s) where Ii where Vi <:AbstractVector{<:T} where T
    return AssocView{T,Vi,Ii}(V,I,s)
end

Base.size(m::AssocView) = m.size
Base.eltype(m::AssocView{T}) where T = T

function Base.getindex(m::AssocView{T},i::Int,j::Int) where T
    @inbounds begin
        idxs = searchsorted(m.indices,(i,j))
        if iszero(length(idxs))
            idxs = searchsorted(m.indices,(j,i))
            iszero(length(idxs)) &&  return zero(T)
        end
        idx = first(idxs)
        return m.values[idx]
    end
end

function zero_assoc(m::Compressed4DMatrix,::Type{T} = Float64) where T <:Number
    newvalues = zeros(T,length(m.values))
    return Compressed4DMatrix(newvalues,m.outer_indices,m.inner_indices,m.outer_size,m.inner_size)
end

function zero_assoc(m::Matrix{<:Matrix},a::Type{T} = Float64) where T <:Number
    return zero_assoc(Compressed4DMatrix(m),a)
end

#=
function zero_assoc(m::Compressed4DMatrix,x::T = 0.0) where T <:Number
    newvalues = fill(x,length(m.values))
    return Compressed4DMatrix(newvalues,m.outer_indices,m.inner_indices,m.outer_size,m.inner_size)
end
=#
function indices(x::Compressed4DMatrix)
    return zip(1:length(x.values),x.outer_indices,x.inner_indices)
end

function indices(x::PackedVofV)
    return x.p
end

struct SparsePackedMofV{E,P<:PackedVofV}<:SparseArrays.AbstractSparseMatrixCSC{E,Int}
    storage::P
    idx::SparseMatrixCSC{Int,Int}
end

function SparsePackedMofV(storage,idx)
    E = eltype(storage)
    P = typeof(storage)
    return SparsePackedMofV{E,P}(storage,idx)
end

_findnz(x::SparseMatrixCSC) = SparseArrays.findnz(x)
function _findnz(x::AbstractMatrix{<:AbstractVector})
    idxs = findall(z->!iszero(length(z)),x)
    i = first.(idxs)
    j = last.(idxs)
    vals = x[idxs]
    return i,j,vals
end

function SparsePackedMofV(m_of_v::AbstractMatrix{<:AbstractVector})
    i,j,unpack_sparse_vals = _findnz(m_of_v)
    #unpack_sparse_vals = m_of_v[sparse_idx]
    pack_sparse_vals = PackedVectorsOfVectors.pack(unpack_sparse_vals)
    len_linear = length(unpack_sparse_vals)
    idx_linear = collect(1:len_linear)
    m,n = size(m_of_v)
    idx = sparse(i,j,idx_linear,m,n)
    E = eltype(pack_sparse_vals)
    P = typeof(pack_sparse_vals)
    return SparsePackedMofV{E,P}(pack_sparse_vals,idx)
end

function Base.getindex(x::SparsePackedMofV,i::Int,j::Int)
    _idx = x.idx[i,j]
    iszero(_idx) && (return view(x.storage.v,1:0))

    return x.storage[_idx]
end
@inline Base.size(x::SparsePackedMofV) = size(x.idx)
@inline SparseArrays.nnz(x::SparsePackedMofV) = length(x.storage.p)
@inline function SparseArrays.findnz(x::SparsePackedMofV)
    i,j,_ = SparseArrays.findnz(x.idx)
    return i,j,x.storage
end
@inline SparseArrays.nonzeros(x::SparsePackedMofV) = x.storage
@inline SparseArrays.rowvals(x::SparsePackedMofV) = SparseArrays.rowvals(x.idx)
@inline SparseArrays.nzrange(A::SparsePackedMofV, col::Integer) = SparseArrays.nzrange(A.idx,col)
@inline SparseArrays.getcolptr(A::SparsePackedMofV) = SparseArrays.getcolptr(A.idx)
export SparsePackedMofV

function Base.show(io::IO,::MIME"text/plain",A::SparsePackedMofV)
    m,n = size(A)
    println(io,"$(m)×$(n) Sparse Packed Matrix of Vectors of eltype $(eltype(A.storage)) with $(length(A.storage)) non empty values:")
    vals = A.storage
    rows = rowvals(A)
    for j in 1:n
        for ii ∈ nzrange(A, j)
            i = rows[ii]
            val= vals[ii]
            println(io,"  ($i,$j) => $val")
        end
    end
end


const GAEVF = JuMP.GenericAffExpr{Float64, VariableRef}
const GQEVF = JuMP.GenericQuadExpr{Float64, VariableRef}

function moi_tests(op_model::OperationModel,
                   params::Bool,
                   vars::Int64,
                   interval::Int64,
                   lessthan::Int64,
                   greaterthan::Int64,
                   equalto::Int64,
                   binary::Bool)

    JuMPmodel = op_model.canonical.JuMPmodel
    @test (:params in keys(JuMPmodel.ext)) == params
    @test JuMP.num_variables(JuMPmodel) == vars
    @test JuMP.num_constraints(JuMPmodel, GAEVF, MOI.Interval{Float64}) == interval
    @test JuMP.num_constraints(JuMPmodel, GAEVF, MOI.LessThan{Float64}) == lessthan
    @test JuMP.num_constraints(JuMPmodel, GAEVF, MOI.GreaterThan{Float64}) == greaterthan
    @test JuMP.num_constraints(JuMPmodel, GAEVF, MOI.EqualTo{Float64}) == equalto
    @test ((JuMP.VariableRef, MOI.ZeroOne) in JuMP.list_of_constraint_types(JuMPmodel)) == binary

    return

end

function psi_constraint_test(op_model::OperationModel, constraint_names::Vector{Symbol})

    for con in constraint_names
        @test !isnothing(get(op_model.canonical.constraints, con, nothing))
    end

    return

end

function psi_checkbinvar_test(op_model::OperationModel, bin_variable_names::Vector{Symbol})

    for variable in bin_variable_names
        for v in op_model.canonical.variables[variable]
            @test JuMP.is_binary(v)
        end
    end

    return

end

function psi_checkobjfun_test(op_model::OperationModel, exp_type)

    @test JuMP.objective_function_type(op_model.canonical.JuMPmodel) == exp_type

    return

end

function moi_lbvalue_test(op_model::OperationModel, con_name::Symbol, value::Number)

    for con in op_model.canonical.constraints[con_name]
        @test JuMP.constraint_object(con).set.lower == value
    end

    return

end

function psi_checksolve_test(op_model::OperationModel, status)
    JuMP.optimize!(op_model.canonical.JuMPmodel)
    @test termination_status(op_model.canonical.JuMPmodel) in status
end

function psi_checksolve_test(op_model::OperationModel, status, expected_result, tol = 0.0)
    res = solve_op_model!(op_model)
    @test termination_status(op_model.canonical.JuMPmodel) in status
    @test isapprox(res.total_cost[:OBJECTIVE_FUNCTION], expected_result, atol = tol)
end


######################################################################
# backend.jl
#
# Currently uses PyCall and PyNEC as the backend
######################################################################

using PyCall
@pyimport PyNEC as pn

function geo_card(geo::PyObject, w::Wire, tag_id::Int)
    geo[:wire](tag_id, w.n_segments, w.a..., w.b..., w.rad, w.rdel, w.rrad)
end

function ex_card(context::PyObject, vs::VoltageSource, s)

    # loop over wires to find matching segment
    tag_id = 0
    for (sei, se) in enumerate(s)
        se == vs.w && (tag_id = sei)
    end

    context[:ex_card](0,            # 0 for voltage source
                      tag_id,       # I2 - tag number of source segment
                      vs.m,         # I3 - m
                      0,            # I4 -
                      real(vs.v),   # F1 - Real part of voltage in volts
                      imag(vs.v),   # F2 - Imaginary part of voltage
                      0,            # F3 -
                      0,            # F4 -
                      0,            # F5 -
                      0             # F6 -
                     )
end

function ex_card(context::PyObject, lw::LinearWave, s)
    context[:ex_card](1,            # I1 - 1 for linear wave
                      lw.n_thetas,  # I2 - number of theta angles
                      lw.n_phis,    # I3 - number of phi angles
                      0,            # I4 - not sure what this is
                      lw.theta,     # F1
                      lw.phi,       # F2
                      lw.eta,       # F3
                      lw.d_theta,   # F4
                      lw.d_phi,     # F5
                      0             # F6 - min/maj axis
                     )
end

function RadiationPattern(s, ex::Excitation, fmhz::Real;
                          thetas::RealOrRange=0:180,
                          phis::RealOrRange=0:360
                         )
    # create context
    context = pn.nec_context()

    # put structure in
    geo = context[:get_geometry]()
    for (sei, se) in enumerate(s)
        geo_card(geo, se, sei)
    end
    context[:geometry_complete](0)

    # ground
    # TODO: obviously, don't just do it this way
    # nullifies previous parameters and sets free-space condition
    context[:gn_card](-1, 0, 0, 0, 0, 0, 0, 0)

    # excitation
    ex_card(context, ex, s)

    # frequency
    #context[:fr_card](0, 2, 2400, 100.0e6)
    context[:fr_card](0, 1, fmhz, 0)

    # rp_card
    context[:rp_card](0,                    # I1 - 0 = normal mode
                      length(thetas),       # I2 - num theta values
                      length(phis),         # I3 - num phi values
                      0,                # I4-X - output format
                      0,                # I4-N - normalization
                      0,                # I4-D - powergain or directive gain
                      0,                # I4-A - Averaging?
                      istart(thetas),       # F1 - initial theta angle
                      istart(phis),         # F2 - initial phi angle
                      step(thetas),         # F3 - theta increment
                      step(phis),           # F4 - phi increment
                      0.0,              # F5 - Radial distance
                      0.0               # F6 - Gain normalization factor
                     )

    # get radiation pattern
    rp = context[:get_radiation_pattern](0)
    return RadiationPattern(rp)
end

function RadiationPattern(rp::PyObject)
    gains_db = rp[:get_gain]()
    gains = 10.0 .^ (gains_db / 10.0)
    thetas = rp[:get_theta_angles]()
    phis = rp[:get_phi_angles]()
    return RadiationPattern(gains_db, gains, thetas, phis)
end


"""
The FADC supertype corresponds to the FADC component
    of the electronics chain in the DAQ setup.
"""
abstract type FADC end

"""
GenericFADC is a dummy FADC model that performs sampling
    of the input waveform based on the sampling interval <Δt>
"""
@with_kw struct GenericFADC <: FADC
    "Sampling interval of FADC, inverse of sampling rate"
    Δt::typeof(1.0*ns_unit) = uconvert(u"ns", inv(250u"MHz"))

    # to add: number of bits
end

"""
FADC with Flashcam algorithm (ToDo)
"""
struct Flashcam <: FADC end

"""
FADC with Struck algorithm (ToDo)
"""
struct Struck <: FADC end

"""
    GenericFADC(sim_conf)

PropDict -> GenericFADC

Construct a GenericFADC instance based on simulation
    configuration <sim_conf>
"""
function GenericFADC(sim_conf::PropDict)
    GenericFADC(
        Δt = haskey(sim_conf.setup.fadc, :sampling_rate) ? uconvert(u"ns", inv(T(sim_conf.setup.fadc.sampling_rate)u"MHz")) : T(sim_conf.setup.fadc.sampling_interval)u"ns",
    )
end

"""
    FADC(sim_conf)

PropDict -> <FADC>

Construct an FADC supertype instance based on simulation
    configuration <sim_conf>.
Type of <FADC> depends on the type specified in <sim_conf>
    (e.g. generic, Flashcam, Struck)
"""
function FADC(sim_conf::PropDict)
    if sim_conf.setup.fadc.type == "generic"
        GenericFADC(sim_conf)
    else
        @info "FADC type $(sim_conf.setup.fadc.type) not implemented!\n
        Available types: generic\n
        Planned in the future: Flashcam, Struck"
    end
end

# -------------------------------------------------------------------

"""
    simulate(wf, fadc)

RDWaveform, GenericFADC -> RDWaveform

Simulate effects of FADC module <fadc> on the waveform <wf>
    (sampling and digitization)
"""
function simulate(wf::RDWaveform, fadc::GenericFADC)
    # resample -> currently simply pick every Nth sample, more elaborate simulation later
    wf_sampled = wf.value[begin : Int(fadc.Δt/step(wf.time)) : end]
    t_sampled = range(T(0)u"ns", step = fadc.Δt, length = length(wf_sampled))
    wf_daq = RDWaveform(t_sampled, wf_sampled)

    # digitize
    wf_daq = RDWaveform(wf_daq.time, UInt16.(round.(wf_daq.value, digits = 0)))

    wf_daq
end


export
SumProductRule,
@sumProductRule

"""
A non-specific sum-product update
"""
abstract type SumProductRule{factor_type} <: MessageUpdateRule end

"""
`internalSumProductSchedule()` generates a sum-product message passing schedule
for the inner graph of a `CompositeFactor`. This schedule produces the sum-product
message out of the specified `outbound_interface`.
"""
function internalSumProductSchedule(cnode::CompositeFactor,
                                    outbound_interface::Interface,
                                    inferred_outbound_types::Dict{Interface, <:Type})

    # Collect types of messages towards the CompositeFactor
    msg_types = Dict{Interface, Type}()
    for (idx, terminal) in enumerate(cnode.terminals)
        (cnode.interfaces[idx] === outbound_interface) && continue # don't need incoming msg on outbound interface
        msg_types[terminal.interfaces[1]] = inferred_outbound_types[cnode.interfaces[idx].partner]
    end

    terminal_interfaces = [term.interfaces[1] for term in cnode.terminals]
    target_terminal = cnode.interface2terminal[outbound_interface]
    internal_outbound_interface = target_terminal.interfaces[1].partner

    # Generate a feasible summary propagation schedule
    schedule = summaryPropagationSchedule(  Variable[],
                                            limit_set=edges(cnode.inner_graph),
                                            target_sites=[internal_outbound_interface],
                                            breaker_sites=terminal_interfaces)

    # Assign the sum-product update rule to each of the schedule entries
    for entry in schedule
        entry.message_update_rule = SumProductRule{typeof(entry.interface.node)}
    end

    # Sanity check
    if schedule[end].interface !== internal_outbound_interface
        error("The last schedule entry should correspond to the message going out of the CompositeFactor")
    end

    # Type inference for internal messages
    inferUpdateRules!(schedule, inferred_outbound_types=msg_types)

    return schedule
end

function inferUpdateRule!(entry::ScheduleEntry,
                          rule_type::Type{T},
                          inferred_outbound_types::Dict{Interface, <:Type}
                         ) where T<:SumProductRule
    # Collect inbound types
    inbound_types = collectInboundTypes(entry, rule_type, inferred_outbound_types)

    # Find applicable rule(s)
    applicable_rules = Type[]
    for rule in leaftypes(entry.message_update_rule)
        if isApplicable(rule, inbound_types)
            push!(applicable_rules, rule)
        end
    end

    # Select and set applicable rule
    if isempty(applicable_rules)
        if isa(entry.interface.node, CompositeFactor)
            # No 'shortcut rule' available for CompositeFactor.
            # Try to fall back to msg passing on the internal graph.
            entry.internal_schedule = internalSumProductSchedule(entry.interface.node, entry.interface, inferred_outbound_types)
            entry.message_update_rule = entry.internal_schedule[end].message_update_rule
        else
            error("No applicable $(rule_type) update for $(typeof(entry.interface.node)) node with inbound types: $(join(inbound_types, ", "))")
        end
    elseif length(applicable_rules) > 1
        error("Multiple applicable $(rule_type) updates for $(typeof(entry.interface.node)) node with inbound types: $(join(inbound_types, ", ")): $(join(applicable_rules, ", "))")
    else
        entry.message_update_rule = first(applicable_rules)
    end

    return entry
end

function collectInboundTypes(entry::ScheduleEntry,
                             ::Type{T},
                             inferred_outbound_types::Dict{Interface, <:Type}
                            ) where T<:SumProductRule
    inbound_message_types = Type[]
    for node_interface in entry.interface.node.interfaces
        if node_interface === entry.interface
            push!(inbound_message_types, Nothing)
        elseif isDeltaConstrained(node_interface.partner)
            push!(inbound_message_types, Message{PointMass})
        else
            push!(inbound_message_types, inferred_outbound_types[node_interface.partner])
        end
    end

    return inbound_message_types
end

"""
`@sumProductRule` registers a sum-product update rule by defining the rule type
and the corresponding methods for the `outboundType` and `isApplicable`
functions. If no name (type) for the new rule is passed, a unique name (type)
will be generated. Returns the rule type.
"""
macro sumProductRule(fields...)
    # Init required fields in macro scope
    node_type = :unknown
    outbound_type = :unknown
    inbound_types = :unknown
    name = :auto # Triggers automatic naming unless overwritten

    # Loop over fields because order is unknown
    for arg in fields

        (arg.args[1] == :(=>)) || error("Invalid call to @sumProductRule")

        if arg.args[2].value == :node_type
            node_type = arg.args[3]
        elseif arg.args[2].value == :outbound_type
            outbound_type = arg.args[3]
            (outbound_type.head == :curly && outbound_type.args[1] == :Message) || error("Outbound type for SumProductRule should be a Message")
        elseif arg.args[2].value == :inbound_types
            inbound_types = arg.args[3]
            (inbound_types.head == :tuple) || error("Inbound types should be passed as Tuple")
        elseif arg.args[2].value == :name
            name = arg.args[3]
        else
            error("Unrecognized field $(arg.args[2].value) in call to @sumProductRule")
        end
    end

    # Assign unique name if not set already
    if name == :auto
        # Added hash ensures that the rule name is unique
        msg_types_hash = string(hash(vcat([outbound_type], inbound_types)))[1:6]
        name = Symbol("SP$(node_type)$(msg_types_hash)")
    end

    # Build validators for isApplicable
    input_type_validators = Expr[]

    push!(input_type_validators, :(length(input_types) == $(length(inbound_types.args))))
    for (i, i_type) in enumerate(inbound_types.args)
        if i_type != :Nothing
            # Only validate inbounds required for message update
            push!(input_type_validators, :(ForneyLab.matches(input_types[$i], $i_type)))
        end
    end

    expr = quote
        struct $name <: SumProductRule{$node_type} end
        ForneyLab.outboundType(::Type{$name}) = $outbound_type
        ForneyLab.isApplicable(::Type{$name}, input_types::Vector{<:Type}) = begin
            $(reduce((current, item) -> :($current && $item), input_type_validators, init = :true))
        end
    end

    return esc(expr)
end

"""
Collect and construct SP update code for each inbound.
"""
collectInbounds(entry::ScheduleEntry, ::Type{T}) where T<:SumProductRule = collectSumProductNodeInbounds(entry.interface.node, entry)

"""
Construct the inbound code that computes the message for `entry`. Allows for
overloading and for a user the define custom node-specific inbounds collection.
Returns a vector with inbounds that correspond with required interfaces.
"""
function collectSumProductNodeInbounds(::FactorNode, entry::ScheduleEntry)
    interface_to_schedule_entry = current_inference_algorithm.interface_to_schedule_entry

    inbounds = Any[]
    for node_interface in entry.interface.node.interfaces
        inbound_interface = ultimatePartner(node_interface)
        if node_interface === entry.interface
            # Ignore inbound message on outbound interface
            push!(inbounds, nothing)
        elseif isClamped(inbound_interface)
            # Hard-code outbound message of constant node in schedule
            push!(inbounds, assembleClamp!(inbound_interface.node, Message))
        else
            # Collect message from previous result
            push!(inbounds, interface_to_schedule_entry[inbound_interface])
        end
    end

    return inbounds
end


using QMSimFiles
using Test

snpmap1 = "lm_mrk_001.txt"
qtlmap1 = "lm_qtl_001.txt"
qtleff1 = "effect_qtl_001.txt"
snpdata1 = "p1_mrk_001.txt"
snpdata2 = "p1_mrk_002.txt"   # saved as /snp_code
qtldata1 = "p1_qtl_001.txt"

@testset "get_number_of_chromosomes" begin
   nchr = QMSimFiles.get_number_of_chromosomes(snpmap1)
   @test nchr == 3
   nchr = QMSimFiles.get_number_of_chromosomes(qtlmap1)
   @test nchr == 3
end

@testset "get_number_of_loci" begin
   nchr = QMSimFiles.get_number_of_chromosomes(snpmap1)
   nloci = QMSimFiles.get_number_of_loci(snpmap1,nchr)
   @test all( nloci == [4,3,4] )
   nchr = QMSimFiles.get_number_of_chromosomes(qtlmap1)
   nloci = QMSimFiles.get_number_of_loci(qtlmap1,nchr)
   @test all( nloci == [3,4,2] )
end

@testset "get_number_of_QTL_allele" begin
   maxna,maxAllele = QMSimFiles.get_number_of_QTL_allele(qtleff1,3)
   @test maxna == 4
   @test all( maxAllele .== [4,4,3] )
end

@testset "load_SNP_QTL_maps!" begin
   maxna,maxAllele = QMSimFiles.get_number_of_QTL_allele(qtleff1,3)
   maps = Vector{QMSimFiles.QMSimChromosomeMap}(undef,3)
   maps[1] = QMSimFiles.QMSimChromosomeMap(7,4,3,zeros(Int,7),zeros(Float64,7),maxAllele[1],zeros(Int,maxAllele[1]),zeros(Float64,maxAllele[1],3))
   maps[2] = QMSimFiles.QMSimChromosomeMap(7,3,4,zeros(Int,7),zeros(Float64,7),maxAllele[2],zeros(Int,maxAllele[2]),zeros(Float64,maxAllele[2],4))
   maps[3] = QMSimFiles.QMSimChromosomeMap(6,4,2,zeros(Int,6),zeros(Float64,6),maxAllele[3],zeros(Int,maxAllele[3]),zeros(Float64,maxAllele[3],2))
   QMSimFiles.load_SNP_QTL_maps!(snpmap1,qtlmap1,maps)
   @test all( maps[1].seqQTL .== [0,0,0,0,1,2,3] )
   @test all( maps[2].seqQTL .== [0,0,0,4,5,6,7] )
   @test all( maps[3].seqQTL .== [0,0,0,0,8,9] )
   @test all( maps[1].pos .≈ [5.99360,46.98708,71.46866,78.57678, 23.76397,28.15399,89.81130] )
   @test all( maps[2].pos .≈ [8.92501,13.56385,43.67896, 29.88287,38.45590,66.06151,82.55315] )
   @test all( maps[3].pos .≈ [25.60614,86.19237,91.47393,98.69095, 14.34842,24.60904] )
   omaps = deepcopy(maps)

   QMSimFiles.sort_by_position!(maps[1].pos, maps[1].seqQTL)
   QMSimFiles.sort_by_position!(maps[2].pos, maps[2].seqQTL)
   QMSimFiles.sort_by_position!(maps[3].pos, maps[3].seqQTL)
   perm1 = [1,5,6,2,3,4,7]
   perm2 = [1,2,4,5,3,6,7]
   perm3 = [5,6,1,2,3,4]
   @test all( maps[1].seqQTL .== omaps[1].seqQTL[perm1] )
   @test all( maps[2].seqQTL .== omaps[2].seqQTL[perm2] )
   @test all( maps[3].seqQTL .== omaps[3].seqQTL[perm3] )
   @test all( maps[1].pos .≈ omaps[1].pos[perm1] )
   @test all( maps[2].pos .≈ omaps[2].pos[perm2] )
   @test all( maps[3].pos .≈ omaps[3].pos[perm3] )
end

@testset "load_QTL_effect! and read_maps" begin
   maxna,maxAllele = QMSimFiles.get_number_of_QTL_allele(qtleff1,3)
   maps = Vector{QMSimFiles.QMSimChromosomeMap}(undef,3)
   maps[1] = QMSimFiles.QMSimChromosomeMap(7,4,3,zeros(Int,7),zeros(Float64,7),maxAllele[1],zeros(Int,3),zeros(Float64,maxAllele[1],3))
   maps[2] = QMSimFiles.QMSimChromosomeMap(7,3,4,zeros(Int,7),zeros(Float64,7),maxAllele[2],zeros(Int,4),zeros(Float64,maxAllele[2],4))
   maps[3] = QMSimFiles.QMSimChromosomeMap(6,4,2,zeros(Int,6),zeros(Float64,6),maxAllele[3],zeros(Int,2),zeros(Float64,maxAllele[3],2))
   QMSimFiles.load_SNP_QTL_maps!(snpmap1,qtlmap1,maps)
   QMSimFiles.sort_by_position!(maps[1].pos, maps[1].seqQTL)
   QMSimFiles.sort_by_position!(maps[2].pos, maps[2].seqQTL)
   QMSimFiles.sort_by_position!(maps[3].pos, maps[3].seqQTL)

   QMSimFiles.load_QTL_effect!(qtleff1,maps)
   @test all( maps[1].naQTL .== [3,4,4] )
   @test all( maps[2].naQTL .== [3,4,2,2] )
   @test all( maps[3].naQTL .== [3,3] )
   @test all( maps[1].effQTL[1:3,1] .≈ [-0.177111,-0.223061,0.071867] )
   @test all( maps[1].effQTL[1:4,2] .≈ [0.034120,-0.053426,-0.032978,0.022941] )
   @test all( maps[1].effQTL[1:4,3] .≈ [-0.010957,-0.001999,0.033815,-0.010451] )
   @test all( maps[2].effQTL[1:3,1] .≈ [0.054960,-0.038202,-0.063630] )
   @test all( maps[2].effQTL[1:4,2] .≈ [0.130199,0.137865,-0.455851,0.296964] )
   @test all( maps[2].effQTL[1:2,3] .≈ [0.001240,-0.000023] )
   @test all( maps[2].effQTL[1:2,4] .≈ [-0.071812,0.063439] )
   @test all( maps[3].effQTL[1:3,1] .≈ [-0.020500,0.004283,-0.075465] )
   @test all( maps[3].effQTL[1:3,2] .≈ [0.004419,-0.010919,0.021111] )

   refmap = read_maps(snpmap1,qtlmap1,qtleff1)
   @test refmap.nchr == 3
   @test refmap.totalSNP == 11
   @test refmap.totalQTL == 9
   for i in 1:3
      @test refmap.chr[i].nLoci == maps[i].nLoci
      @test refmap.chr[i].nSNP == maps[i].nSNP
      @test refmap.chr[i].nQTL == maps[i].nQTL
      @test all( refmap.chr[i].seqQTL .== maps[i].seqQTL )
      @test all( refmap.chr[i].pos .≈ maps[i].pos )
      @test refmap.chr[i].maxAllele .== maps[i].maxAllele
      @test all( refmap.chr[i].naQTL .== maps[i].naQTL )
      @test all( refmap.chr[i].effQTL .≈ maps[i].effQTL )
   end
end

@testset "generate_chromosome_set" begin
   gmap = read_maps(snpmap1,qtlmap1,qtleff1)
   chromosome_set = QMSimFiles.generate_chromosome_set(gmap)
   @test length(chromosome_set)==gmap.nchr
end

@testset "parse SNP/QTL data" begin
   gmap = read_maps(snpmap1,qtlmap1,qtleff1)
   @test QMSimFiles.is_snpcode_file(snpdata1,gmap.totalSNP)==false
   @test QMSimFiles.is_snpcode_file(snpdata2,gmap.totalSNP)==true

   # for SNP (standard)
   stext = "2       2 1 2 2 1 1 1 1 1 1 1 1 1 1 1 2 1 1 1 2 1 1"
   snp = zeros(Int8,gmap.totalSNP*2)
   QMSimFiles.text_to_code!(stext,gmap.totalSNP,snp,7)
   @test all( snp .== [2, 1, 2, 2, 1, 1, 1, 1,   1, 1, 1, 1, 1, 1,   1, 2, 1, 1, 1, 2, 1, 1] )

   chromosome_set1 = QMSimFiles.generate_chromosome_set(gmap)
   QMSimFiles.convert_snpdata_to_haplotype!(snp,gmap,chromosome_set1)
   @test all( chromosome_set1[1].gp[findall(x->gmap.chr[1].seqQTL[x]==0,1:gmap.chr[1].nLoci)] .== [1,1,0,0] )
   @test all( chromosome_set1[1].gm[findall(x->gmap.chr[1].seqQTL[x]==0,1:gmap.chr[1].nLoci)] .== [0,1,0,0] )
   @test all( chromosome_set1[2].gp[findall(x->gmap.chr[2].seqQTL[x]==0,1:gmap.chr[2].nLoci)] .== [0,0,0] )
   @test all( chromosome_set1[2].gm[findall(x->gmap.chr[2].seqQTL[x]==0,1:gmap.chr[2].nLoci)] .== [0,0,0] )
   @test all( chromosome_set1[3].gp[findall(x->gmap.chr[3].seqQTL[x]==0,1:gmap.chr[3].nLoci)] .== [0,0,0,0] )
   @test all( chromosome_set1[3].gm[findall(x->gmap.chr[3].seqQTL[x]==0,1:gmap.chr[3].nLoci)] .== [1,0,1,0] )

   # for SNP (snp_code)
   ctext = "2      42000003030"
   snp_code = zeros(Int8,gmap.totalSNP)
   QMSimFiles.text_to_snpcode!(ctext,gmap.totalSNP,snp_code,7)
   @test all( snp_code .== [4,2,0,0, 0,0,0, 3,0,3,0] )

   chromosome_set2 = QMSimFiles.generate_chromosome_set(gmap)
   QMSimFiles.convert_snpcode_to_haplotype!(snp_code,gmap,chromosome_set2)
   @test all( chromosome_set2[1].gp .== chromosome_set1[1].gp )
   @test all( chromosome_set2[1].gm .== chromosome_set1[1].gm )
   @test all( chromosome_set2[2].gp .== chromosome_set1[2].gp )
   @test all( chromosome_set2[2].gm .== chromosome_set1[2].gm )
   @test all( chromosome_set2[3].gp .== chromosome_set1[3].gp )
   @test all( chromosome_set2[3].gm .== chromosome_set1[3].gm )

   # for QTL
   qtext = "2       2 2 2 4 4 1 1 1 2 1 2 2 2 2 2 1 1 2"
   qtl = zeros(Int8,gmap.totalQTL*2)
   QMSimFiles.text_to_code!(qtext,gmap.totalQTL,qtl,7)
   @test all( qtl .== [2, 2, 2, 4, 4, 1,   1, 1, 2, 1, 2, 2, 2, 2,   2, 1, 1, 2] )

   QMSimFiles.convert_qtldata_to_haplotype!(qtl,gmap,chromosome_set2)
   @test all( chromosome_set2[1].gp[findall(x->gmap.chr[1].seqQTL[x]>0,1:gmap.chr[1].nLoci)] .== [2,2,4] )
   @test all( chromosome_set2[1].gm[findall(x->gmap.chr[1].seqQTL[x]>0,1:gmap.chr[1].nLoci)] .== [2,4,1] )
   @test all( chromosome_set2[2].gp[findall(x->gmap.chr[2].seqQTL[x]>0,1:gmap.chr[2].nLoci)] .== [1,2,2,2] )
   @test all( chromosome_set2[2].gm[findall(x->gmap.chr[2].seqQTL[x]>0,1:gmap.chr[2].nLoci)] .== [1,1,2,2] )
   @test all( chromosome_set2[3].gp[findall(x->gmap.chr[3].seqQTL[x]>0,1:gmap.chr[3].nLoci)] .== [2,1] )
   @test all( chromosome_set2[3].gm[findall(x->gmap.chr[3].seqQTL[x]>0,1:gmap.chr[3].nLoci)] .== [1,2] )

   # general function
   g1 = read_genotypes(snpdata1,qtldata1,gmap)
   g2 = read_genotypes(snpdata2,qtldata1,gmap)
   for i=1:3
      @test all( g1[2].chr[i].gp .== chromosome_set2[i].gp )
      @test all( g1[2].chr[i].gm .== chromosome_set2[i].gm )
      @test all( g2[2].chr[i].gp .== chromosome_set2[i].gp )
      @test all( g2[2].chr[i].gm .== chromosome_set2[i].gm )
   end
   @test isapprox(g1[1].tbv, -0.412531, atol=1e-5)
   @test isapprox(g1[2].tbv, -0.015914, atol=1e-5)
   @test isapprox(g1[3].tbv, -0.137409, atol=1e-5)
   @test isapprox(g1[4].tbv, -0.242572, atol=1e-5)
   for i=1:4
      @test g1[i].tbv ≈ g2[i].tbv
   end
end

@testset "number_of_crosses" begin
   g = read_qmsim_data(snpmap1,qtlmap1,qtleff1,snpdata2,qtldata1)
   chr = g.map.chr[1]
   refpos = [5.9936, 23.76397, 28.15399, 46.98708, 71.46866, 78.57678, 89.8113]
   n = QMSimFiles.number_of_crosses(chr.pos[end])
   @test (0 <= n) && (n <= 100)
end

@testset "location_of_crossover" begin
   g = read_qmsim_data(snpmap1,qtlmap1,qtleff1,snpdata2,qtldata1)
   chr = g.map.chr[1]
   ntests = 10
   for i=1:ntests
      chrlen = chr.pos[end]
      ncross = QMSimFiles.number_of_crosses(chrlen) + 1
      (ncross,gameteid,loccross) = QMSimFiles.location_of_crossover(chrlen,chr.pos,ncross)
      if ncross>0
         @test issorted(loccross) && (loccross[1]>=0) && (loccross[end]<=chrlen) && length(loccross)==ncross
         @test all( map(x -> ((gameteid[x][1] in [1,2]) && gameteid[x][2] in [3,4]), 1:ncross ) ) && length(gameteid)==ncross
      end
   end
   # adjustment to remove double cross etc
   gameteid = [(1,3),(1,3),(1,4),(1,4),(2,3),(2,3),(2,4),(2,4)]
   loccross = [1.0,  2.0,  3.0,  4.0,  5.0,  6.0,  7.0,  8.0]
   (ncross2,gameteid2,loccross2) = QMSimFiles.adjusted_location_of_crossover(gameteid,loccross)
   @test ncross2==0

   gameteid = [(1,3),(1,3),(1,4),(1,4),(1,4),(1,4),(2,4),(2,4)]
   loccross = [1.0,  2.0,  3.0,  4.0,  5.0,  6.0,  7.0,  8.0]
   (ncross2,gameteid2,loccross2) = QMSimFiles.adjusted_location_of_crossover(gameteid,loccross)
   @test ncross2==0

   gameteid = [(2,3),(1,3),(1,4),(1,3),(2,3),(1,3),(2,4),(1,4),(2,4)]
   loccross = [1.0,  2.0,  3.0,  4.0,  5.0,  6.0,  7.0,  8.0,  9.0]
   (ncross2,gameteid2,loccross2) = QMSimFiles.adjusted_location_of_crossover(gameteid,loccross)
   @test ncross2==1

   gameteid = [(2,3),(1,3),(1,4),(1,3),(2,3),(1,3),(2,4),(1,4),(2,4),(2,3),(1,4)]
   loccross = [1.0,  2.0,  3.0,  4.0,  5.0,  6.0,  7.0,  8.0,  9.0,  10.0, 11.0]
   (ncross2,gameteid2,loccross2) = QMSimFiles.adjusted_location_of_crossover(gameteid,loccross)
   @test ncross2==3
end

@testset "generate_gamete_with_crossover" begin
   g = read_qmsim_data(snpmap1,qtlmap1,qtleff1,snpdata2,qtldata1)
   chr = g.map.chr[1]
   chrlen = chr.pos[end]
   refpos = [5.9936, 23.76397, 28.15399, 46.98708, 71.46866, 78.57678, 89.8113]

   # refgp = [1, 3, 3, 0, 0, 0, 4]
   # refgm = [1, 3, 4, 0, 0, 1, 4]
   gp = g.individual[4].chr[1].gp
   gm = g.individual[4].chr[1].gm
   ncross = 1
   gameteid = [(1,3)]
   loccross = [3]
   ga = QMSimFiles.generate_gamete_with_crossover(gameteid,loccross,1,gp,gm)
   @test all( ga .== [1,3,4,0,0,1,4] )

   ncross = 2
   gameteid = [(1,3),(1,4)]
   loccross = [3,6]
   ga = QMSimFiles.generate_gamete_with_crossover(gameteid,loccross,1,gp,gm)
   @test all( ga .== [1,3,4,0,0,1,4] )
end

@testset "general mating" begin
   g = read_qmsim_data(snpmap1,qtlmap1,qtleff1,snpdata2,qtldata1)
   sumtbv = 0.0
   nanim = 1000
   for i=1:nanim
      animal = mating(g.map, g.individual[1], g.individual[2])
      sumtbv = sumtbv + animal.tbv
   end
   pa = (g.individual[1].tbv + g.individual[2].tbv)/2
   @test isapprox(sumtbv/nanim,pa, atol=1e-1)
end

@testset "equal sign" begin
   g = read_qmsim_data(snpmap1,qtlmap1,qtleff1,snpdata2,qtldata1)
   @test g.map==g.map
   @test g.individual[1]==g.individual[1]
   @test !(g.individual[1]==g.individual[2])
end

@testset "pack/unpack genotypes" begin
   g = read_qmsim_data(snpmap1,qtlmap1,qtleff1,snpdata2,qtldata1)
   for k=1:length(g.individual)
      # pack
      ref_packed = zeros(Int8,0)
      for i=1:g.map.nchr
         ref_packed = [ref_packed; g.individual[k].chr[i].gp]
         ref_packed = [ref_packed; g.individual[k].chr[i].gm]
      end
      ind_packed = QMSimFiles.pack_genome(g.map,g.individual[k])
      @test all( ind_packed .== ref_packed )
      packed = QMSimFiles.pack_genome(g.map,g.individual[k].chr)
      @test all( packed .== ref_packed )

      # unpack
      chromosome_set = QMSimFiles.unpack_genome(g.map,packed)
      @test all( chromosome_set == g.individual[k].chr )
   end
end

@testset "save/read genotypes" begin
   g = read_qmsim_data(snpmap1,qtlmap1,qtleff1,snpdata2,qtldata1)

   #hdf5file = "runtests.h5"
   hdf5file = tempname()
   @info "temporary HDF5 file: $(hdf5file)"
   save_qmsim_data_hdf5(g,hdf5file)

   # map
   saved_map = read_qmsim_map_hdf5(hdf5file)
   @test saved_map == g.map

   # particular individual
   for i=1:length(g.individual)
      individual = QMSimFiles.read_qmsim_individual_hdf5(g.map,hdf5file,i)
      @test individual == g.individual[i]
   end

   # all genotypes
   genotypes = QMSimFiles.read_qmsim_all_genotypes_hdf5(g.map,hdf5file)
   @test g == QMSimPopulationGenome(g.map,genotypes)

   # whole data
   @test g == read_qmsim_data_hdf5(hdf5file)

   # added individual
   add_qmsim_individual_hdf5(g.map, hdf5file, g.individual[1])
   added = QMSimFiles.read_qmsim_individual_hdf5(g.map,hdf5file,5)
   @test added == g.individual[1]

   # size of saved genotypes
   @test get_qmsim_genotype_size(hdf5file) == (5, g.map.totalSNP + g.map.totalQTL)

   # empty hdf5 only with map
   hdf5file2 = tempname()
   @info "temporary HDF5 file: $(hdf5file2)"
   create_qmsim_map_hdf5(g.map, hdf5file2)

   saved_map = read_qmsim_map_hdf5(hdf5file2)
   @test saved_map == g.map

   add_qmsim_individual_hdf5(g.map, hdf5file2, g.individual[1])
   added = QMSimFiles.read_qmsim_individual_hdf5(g.map,hdf5file2,1)
   @test added == g.individual[1]
end

@testset "export to blupf90" begin
   g = read_qmsim_data(snpmap1,qtlmap1,qtleff1,snpdata2,qtldata1)
   @test String( QMSimFiles.pack_markrs(g.map,g.individual[2]) .+ UInt8(48)) == "12000001010"
   @test String( QMSimFiles.pack_markrs(g.map,g.individual[3]) .+ UInt8(48)) == "22220001120"

   snpfile = tempname()
   @info "temporary text file: $(snpfile)"
   if isfile(snpfile); rm(snpfile,force=true); end
   export_qmsim_individual_blupf90(g.map, 2, g.individual[2], snpfile)
   export_qmsim_individual_blupf90(g.map, 3, g.individual[3], snpfile)

   @test countlines(snpfile)==2
   open(snpfile,"r") do io
      # individual 2
      line = readline(io)
      items = split(line)
      @test items[1]=="2" && items[2]=="12000001010"
      # individual 3
      line = readline(io)
      items = split(line)
      @test items[1]=="3" && items[2]=="22220001120"
   end
end


"""
    loss functions
"""
function discriminator_loss(d_real::AbstractArray, d_g_fake_adv::AbstractArray)
    real_loss = mse(1f0, mean(d_real; dims = (1, 2)))
    fake_loss = mse(0f0, mean(d_g_fake_adv; dims = (1, 2)))
    return real_loss + fake_loss
end

generator_adv_loss(d_g_fake_adv::AbstractArray) = mse(1f0, mean(d_g_fake_adv; dims = (1, 2)))

generator_rec_loss(real_img::AbstractArray, g_fake_rec::AbstractArray) = mse(real_img, g_fake_rec)

"""
    update discriminator
"""
function update_discriminator!(opt, dscr, real_img, g_fake_adv)
    @eval Flux.istraining() = true
    ps = params(dscr)
    grad = gradient(ps) do
        discriminator_loss(dscr(real_img), dscr(g_fake_adv))
    end
    update!(opt, ps, grad)
    @eval Flux.istraining() = false
end

"""
    update generator
"""
# function update_generator_rec!(opt, gen, real_img, prev_rec, noise_rec, alpha)
#     @eval Flux.istraining() = true
#     ps = params(gen)
#     grad = gradient(ps) do
#         g_fake_rec = gen(prev_rec, noise_rec)
#         alpha * generator_rec_loss(real_img, g_fake_rec)
#     end
#     update!(opt, ps, grad)
#     @eval Flux.istraining() = false
# end

# function update_generator_adv!(opt, dscr, gen, prev_adv, noise_adv)
#     @eval Flux.istraining() = true
#     ps = params(gen)
#     grad = gradient(ps) do
#         d_g_fake_adv = dscr(gen(prev_adv, noise_adv))
#         generator_adv_loss(d_g_fake_adv)
#     end
#     update!(opt, ps, grad)
#     @eval Flux.istraining() = false
# end

function update_generator!(opt, dscr, gen, real_img, prev_rec, prev_adv, noise_rec, noise_adv, alpha)
    @eval Flux.istraining() = true
    ps = params(gen)
    grad = gradient(ps) do
        g_fake_rec = gen(prev_rec, noise_rec)
        d_g_fake_adv = dscr(gen(prev_adv, noise_adv))
        generator_adv_loss(d_g_fake_adv) + alpha * generator_rec_loss(real_img, g_fake_rec)
    end
    update!(opt, ps, grad)
    @eval Flux.istraining() = false
end

"""
    train
"""
function train_epoch!(opt_dscr, opt_gen, st, loop_dscr, loop_gen,
        dscr, genp, prev_rec, noise_rec, real_img, amplifiers, alpha)
    # training
    # prev_rec is only used to infer array type

    # discriminator
    foreach(1:loop_dscr) do _
        noise_adv = build_noise_pyramid(prev_rec, genp.noise_shapes[1:st], amplifiers)
        g_fake_adv = genp(noise_adv, st, false)
        # add noise to real
        # real_noise = 0.5f0 * amplifiers[st] * randn!(similar(real_img))
        # update_discriminator!(opt_dscr, dscr, real_img + real_noise, g_fake_adv)

        update_discriminator!(opt_dscr, dscr, real_img, g_fake_adv)
    end

    # generator
    foreach(1:loop_gen) do _
        # adv
        noise_adv = build_noise_pyramid(prev_rec, genp.noise_shapes[1:st], amplifiers)
        prev_adv = genp(noise_adv, st - 1, true)
        # update_generator_adv!(opt_gen, dscr, genp.chains[st], prev_adv, last(noise_adv))
        # rec
        # update_generator_rec!(opt_gen, genp.chains[st], real_img, prev_rec, noise_rec, alpha)
        update_generator!(opt_gen, dscr, genp.chains[st], real_img, prev_rec, prev_adv, noise_rec, last(noise_adv), alpha)
    end

    noise_adv = build_noise_pyramid(prev_rec, genp.noise_shapes[1:st], amplifiers)
    g_fake_adv = genp(noise_adv, st, false)
    loss_dscr = discriminator_loss(dscr(real_img), dscr(g_fake_adv))
    d_g_fake_adv = dscr(g_fake_adv)
    loss_gen_adv = generator_adv_loss(d_g_fake_adv)
    g_fake_rec = genp.chains[st](prev_rec, noise_rec)
    loss_gen_rec = generator_rec_loss(real_img, g_fake_rec)

    return loss_dscr, loss_gen_adv, loss_gen_rec
end

function estimate_noise_amplifier(prev_rec::AbstractArray{Float32,4}, real_img::AbstractArray{Float32,4},
        pad::Integer, amplifier_init::Float32)
    prev_rec_crop = @view prev_rec[1 + pad:end - pad, 1 + pad:end - pad, :, :]
    rmse = sqrt(mse(real_img, prev_rec_crop))
    return rmse * amplifier_init
end

function train!(dscrp::DiscriminatorPyramid, genp::GeneratorPyramid,
        real_img_p::Vector{T}, hp::HyperParams) where {T <: AbstractArray{Float32,4}}
    stages = Base.length(genp.image_shapes)
    generate_dirs(stages)
    save_scaled_reals(real_img_p)

    amplifiers = Float32[]

    # fixed noise for rec
    fixed_noise_rec = build_rec_pyramid(first(real_img_p), genp.noise_shapes, 1f0)
    fixed_noise_adv = similar(fixed_noise_rec)

    for st in 1:stages
        @info "Step $(st)"
        # reset optimizer
        opt_dscr = ADAM(hp.lr_dscr, (0.5, 0.999))
        opt_gen = ADAM(hp.lr_gen, (0.5, 0.999))

        # calculate noise amplifier
        prev_rec = genp(fixed_noise_rec, st - 1, true) # padded
        amp = st == 1 ? 1f0 : estimate_noise_amplifier(prev_rec, real_img_p[st], genp.pad, hp.amplifier_init)
        push!(amplifiers, amp)
        # add noise for adv
        fixed_noise_adv[st] = amp * randn_like(prev_rec, expand_dim(genp.noise_shapes[st]...))

        save_noise_amplifiers(st, amp)
        @info "Noise amplifier = $(amp)"

        @time for ep in 1:hp.max_epoch
            # reduce learnint rate
            if ep == hp.reduce_lr_epoch
                @info "Reduce learning rate"
                opt_dscr.eta /= 10
                opt_gen.eta /= 10
            end

            loss_dscr, loss_gen_adv, loss_gen_rec =
                train_epoch!(opt_dscr, opt_gen, st, hp.loop_dscr, hp.loop_gen,
                    dscrp.chains[st], genp, prev_rec, fixed_noise_rec[st], real_img_p[st], amplifiers, hp.alpha)

            # save image/loss
            if ep == 1 || ep % hp.save_image_every_epoch == 0 || ep == hp.max_epoch
                save_generated_images(genp, fixed_noise_rec, fixed_noise_adv, st, ep)
            end

            if ep == 1 || ep % hp.save_loss_every_epoch == 0 || ep == hp.max_epoch
                @info "Epoch $(ep)" loss_dscr loss_gen_adv loss_gen_rec
                save_training_loss(st, ep, loss_dscr, loss_gen_adv, loss_gen_rec)
            end
        end

        save_model_params(dscrp, genp, st)
    end

    return amplifiers, first(fixed_noise_rec)
end


"""
Data structure for particle filters
- effectiveSampleSize(pc :: ParticleContainer)
- normalise!(pc::ParticleContainer)
- consume(pc::ParticleContainer): return incremental likelihood
"""

typealias Particle Trace

type ParticleContainer{T<:Particle}
  model :: Function
  num_particles :: Int
  vals  :: Array{T,1}
  logWs :: Array{Float64,1}  # Log weights (Trace) or incremental likelihoods (ParticleContainer)
  logE  :: Float64           # Log model evidence
  conditional :: Union{Void,Conditional} # storing parameters, helpful for implementing rejuvenation steps
  n_consume :: Int # helpful for rejuvenation steps, e.g. in SMC2
  ParticleContainer(m::Function,n::Int) = new(m,n,Array{Particle,1}(),Array{Float64,1}(),0.0,nothing,0)
end

call{T}(::Type{ParticleContainer{T}}, m) = ParticleContainer{T}(m, 0)

Base.collect(pc :: ParticleContainer) = pc.vals # prev: Dict, now: Array
Base.length(pc :: ParticleContainer)  = length(pc.vals)
Base.similar(pc :: ParticleContainer) = ParticleContainer(pc.model, 0)
# pc[i] returns the i'th particle
Base.getindex(pc :: ParticleContainer, i :: Real) = pc.vals[i]

# registers a new x-particle in the container
function Base.push!(pc :: ParticleContainer, p :: Particle)
  pc.num_particles += 1
  push!(pc.vals, p)
  push!(pc.logWs, 0)
  pc
end
Base.push!(pc :: ParticleContainer) = Base.push!(pc, eltype(pc.vals)(pc.model))
function Base.push!(pc :: ParticleContainer, n :: Int)
  vals = Array{eltype(pc.vals), 1}(n)
  logWs = Array{eltype(pc.logWs), 1}(n)
  for i=1:n
    vals[i]  = eltype(pc.vals)(pc.model)
    logWs[i] = 0
  end
  append!(pc.vals, vals)
  append!(pc.logWs, logWs)
  pc.num_particles += n
  pc
end

# clears the container but keep params, logweight etc.
function Base.empty!(pc :: ParticleContainer)
  pc.num_particles = 0
  pc.vals  = Array{Particle,1}()
  pc.logWs = Array{Float64,1}()
  pc
end

# clones a theta-particle
function Base.copy(pc :: ParticleContainer)
  particles = collect(pc)
  newpc     = similar(pc)
  for p in particles
    newp = forkc(p)
    push!(newpc, newp)
  end
  newpc.logE        = pc.logE
  newpc.logWs       = deepcopy(pc.logWs)
  newpc.conditional = deepcopy(pc.conditional)
  newpc.n_consume   = pc.n_consume
  newpc
end

# run particle filter for one step, return incremental likelihood
function Base.consume(pc :: ParticleContainer)
  @assert pc.num_particles == length(pc)
  # normalisation factor: 1/N
  _, z1      = weights(pc)
  n = length(pc.vals)

  particles = collect(pc)
  num_done = 0
  for i=1:n
    p = pc.vals[i]
    score = consume(p)
    if isa(score, Real)
      increase_logweight(pc, i, Float64(score))
    elseif score == Val{:done}
      num_done += 1
    else
      "[consume]: error in running particle filter."
    end
  end

  if num_done == length(pc)
    res = Val{:done}
  elseif num_done != 0
    error("[consume]: mis-aligned execution traces, num_particles= $(n), num_done=$(num_done).")
  else
    # update incremental likelihoods
    _, z2      = weights(pc)
    res = increase_logevidence(pc, z2 - z1)
    pc.n_consume += 1
    # res = increase_loglikelihood(pc, z2 - z1)
  end

  res
end

function weights(pc :: ParticleContainer)
  @assert pc.num_particles == length(pc)
  logWs = pc.logWs
  Ws = exp(logWs-maximum(logWs))
  logZ = log(sum(Ws)) + maximum(logWs)
  Ws = Ws ./ sum(Ws)
  return Ws, logZ
end

function effectiveSampleSize(pc :: ParticleContainer)
  Ws, _ = weights(pc)
  ess = sum(Ws) ^ 2 / sum(Ws .^ 2)
end

increase_logweight(pc :: ParticleContainer, t :: Int, logw :: Float64) =
  (pc.logWs[t]  += logw)

increase_logevidence(pc :: ParticleContainer, logw :: Float64) =
  (pc.logE += logw)

function resample!( pc :: ParticleContainer,
                   randcat :: Function = Turing.resampleSystematic,
                   ref :: Union{Particle, Void} = nothing)
  n1, particles = pc.num_particles, collect(pc)
  @assert n1 == length(particles)

  # resample
  Ws, _ = weights(pc)
  n2    = isa(ref, Void) ? n1 : n1-1
  indx  = randcat(Ws, n2)

  # fork particles
  empty!(pc)
  num_children = zeros(Int,n1)
  map(i->num_children[i]+=1, indx)
  for i = 1:n1
    p = particles[i] == ref ? Traces.forkc(particles[i]) : particles[i]
    num_children[i] > 0 && push!(pc, p)
    for k=1:num_children[i]-1
      newp = Traces.forkc(p)
      push!(pc, newp)
    end
  end

  if isa(ref, Particle)
    # Insert the retained particle. This is based on the replaying trick for efficiency
    #  reasons. If we implement PG using task copying, we need to store Nx * T particles!
    push!(pc, ref)
  end

  pc
end


@rule typeof(-)(:in1, Marginalisation) (m_out::Any, m_in2::Any, meta::Any) = begin
    @call_rule typeof(+)(:out, Marginalisation) (m_in1 = m_out, m_in2 = m_in2, meta = meta)
end


import UnicodePlots
import UnicodePlots: BrailleCanvas, nrows
import UnicodePlots: lineplot!, scatterplot!

function render_unicode(cur_size, cur_data, cur_layout, cur_config)

  work_layout = deepcopy(cur_layout)

  cur_plot, has_legend, is_linear_x, is_linear_y, max_label_length = (
    _configure_repl_layout(cur_size, cur_data, work_layout)
  )

  if isempty(cur_data)
    return cur_plot
  end

  # repl doesn't allow legend placement
  @assert isempty(work_layout["legend"])
  delete!(work_layout, "legend")

  # no annotations in repl
  delete!(work_layout, "annotations")

  work_config = deepcopy(cur_config)
  delete!(work_config, "responsive")
  @assert isempty(work_config)

  for (cur_index, cur_datum) in enumerate(cur_data)
    @assert cur_datum["mode"] in ["lines", "markers"]

    work_datum = deepcopy(cur_datum)
    singular_mode = cur_datum["mode"][1:end-1]

    if cur_datum["mode"] == "lines"
      plotfunc = lineplot!
    else
      plotfunc = scatterplot!
    end

    cur_kwargs = Dict()
    cur_kwargs[:color] = Symbol(cur_datum[singular_mode]["color"])

    if haskey(cur_datum, "showlegend")
      @assert !cur_datum["showlegend"]
      @assert !haskey(cur_datum, "name")
      delete!(work_datum, "showlegend")
    elseif haskey(cur_datum, "name")
      cur_label = cur_datum["name"]
      if length(cur_label) > max_label_length
        cur_label = cur_label[1:(max_label_length-1)] * "…"
      end

      if has_legend
        cur_kwargs[:name] = cur_label
      end
      delete!(work_datum, "name")
    end

    cur_x = cur_datum["x"]
    cur_y = cur_datum["y"]

    is_linear_x || ( cur_x = log10.(float.(cur_x)) )
    is_linear_y || ( cur_y = log10.(float.(cur_y)) )

    plotfunc(
      cur_plot, cur_x, cur_y; cur_kwargs...
    )

    delete!(work_datum[singular_mode], "color")
    delete!(work_datum[singular_mode], "width")
    delete!(work_datum[singular_mode], "size")

    @assert isempty(work_datum[singular_mode])
    delete!(work_datum, singular_mode)

    delete!(work_datum, "mode")
    delete!(work_datum, "x")
    delete!(work_datum, "y")

    isempty(work_datum) || @show work_datum
    @assert isempty(work_datum)
  end

  @assert isempty(work_layout["shapes"])
  delete!(work_layout, "shapes")
  @assert isempty(work_layout)

  for (k,v) in work_layout
    @show (k,v)
  end

  cur_plot
end

function _configure_repl_layout(cur_size, cur_data, work_layout)
  size_x, size_y = cur_size

  has_legend = work_layout["showlegend"]
  delete!(work_layout, "showlegend")

  if isempty(cur_data)
    max_label_length = 0
  else
    max_label_length = maximum(
      length.(map(
        cur_datum -> haskey(cur_datum, "name") ? cur_datum["name"] : "", cur_data
      ))
    )

    max_label_length = min(16, max_label_length)

    if has_legend
      size_x -= max_label_length
    end
  end

  _get_lims(x) = collect(extrema(Base.Iterators.flatten(
    map(cur_datum -> cur_datum[x], cur_data)
  )))

  function _get_extrema(x)
    if haskey(work_layout["$(x)axis"], "range")
      @assert work_layout["$(x)axis"]["autorange"] == false
      cur_extrema = work_layout["$(x)axis"]["range"]
      delete!(work_layout["$(x)axis"], "range")
    else
      @assert work_layout["$(x)axis"]["autorange"] == true
      if isempty(cur_data)
        cur_extrema = (1, 10)
      else
        cur_extrema = _get_lims(x)
      end
    end

    delete!(work_layout["$(x)axis"], "autorange")
    return cur_extrema
  end

  if isempty(cur_data)
    min_x, max_x = _get_extrema("x")
    min_y, max_y = _get_extrema("y")
  else
    min_x, max_x = _get_extrema("x")
    min_y, max_y = _get_extrema("y")
  end

  function _is_linear(x)
    cur_is_linear = work_layout["$(x)axis"]["type"] == "linear"
    if !cur_is_linear
      @assert work_layout["$(x)axis"]["type"] == "log"
    end
    delete!(work_layout["$(x)axis"], "type")

    cur_is_linear
  end

  is_linear_x = _is_linear("x")
  is_linear_y = _is_linear("y")

  if is_linear_x
    origin_x = min_x
    width = max_x - min_x
  else
    origin_x = log10(min_x)
    width = log10(max_x) - log10(min_x)
  end

  if is_linear_y
    origin_y = min_y
    height = max_y - min_y
  else
    origin_y = log10(min_y)
    height = log10(max_y) - log10(min_y)
  end

  # ignore desired format for repl
  delete!(work_layout["xaxis"], "exponentformat")
  delete!(work_layout["yaxis"], "exponentformat")

  if haskey(work_layout["xaxis"], "title")
    cur_xlabel = work_layout["xaxis"]["title"]
    delete!(work_layout["xaxis"], "title")
  else
    cur_xlabel = ""
  end

  if haskey(work_layout["yaxis"], "title")
    cur_ylabel = work_layout["yaxis"]["title"]
    delete!(work_layout["yaxis"], "title")
  else
    cur_ylabel = ""
  end

  # unicode plots has no concept of ticks
  delete!(work_layout["xaxis"], "tickvals")
  delete!(work_layout["yaxis"], "tickvals")

  @assert isempty(work_layout["xaxis"])
  @assert isempty(work_layout["yaxis"])

  delete!(work_layout, "xaxis")
  delete!(work_layout, "yaxis")

  cur_canvas = BrailleCanvas(
    size_x, size_y, width=width, height=height,
    origin_x=origin_x, origin_y=origin_y
  )

  if is_linear_x
    if is_linear_y
      border = :solid
    else
      border = :ascii
    end
  else
    if is_linear_y
      border = :dashed
    else
      border = :dotted
    end
  end

  if !is_linear_x
    if strip(cur_xlabel) == ""
      cur_xlabel = "log"
    else
      cur_xlabel = "log " * cur_xlabel
    end
  end

  if haskey(work_layout, "title")
    cur_title = work_layout["title"]
    delete!(work_layout, "title")
  else
    cur_title = ""
  end

  cur_plot = UnicodePlots.Plot(
    cur_canvas; title = cur_title,
    xlabel = cur_xlabel, border=border
  )

  function _annotate!(varargs...; color=:light_black)
    UnicodePlots.annotate!(cur_plot, varargs..., color=color)
  end

  ylabel_index = Int(floor(nrows(cur_canvas)/2))

  if strip(cur_ylabel) == ""
    if !is_linear_y
      _annotate!(:l, ylabel_index, "log", color=:gray)
    end
  else
    _annotate!(:l, ylabel_index, cur_ylabel, color=:gray)

    if !is_linear_y
      log_label = "log"
      cur_diff = length(cur_ylabel) - length(log_label)

      log_label = lpad(log_label, length(log_label) + Int(floor(cur_diff/2)))
      log_label = rpad(log_label, length(cur_ylabel))

      _annotate!(:l, ylabel_index-1, log_label, color=:gray)
    end
  end

  function _clean_string(cur_value)
    rounded_value = round(cur_value, sigdigits=3)

    if ( abs(rounded_value) % 1 ) < ( 1e3 * eps() )
      string_value = string(Int(rounded_value))
    elseif 0.1 <= abs(rounded_value) < 1e4
      string_value = string(rounded_value)
    else
      tmp_coeff, tmp_exponent = split(@sprintf("%.2e", rounded_value), "e")
      tmp_coeff = rstrip(rstrip(tmp_coeff, '0'), '.')

      string_value = tmp_coeff * "e" * tmp_exponent
    end

    string_value
  end

  _annotate!(:l, nrows(cur_canvas), _clean_string(min_y))
  _annotate!(:l, 1, _clean_string(max_y))
  _annotate!(:bl, _clean_string(min_x))
  _annotate!(:br, _clean_string(max_x))

  return (cur_plot, has_legend, is_linear_x, is_linear_y, max_label_length)

end


import Base.Dict

export Atoms

"""
`JData`: datatype for storing any data

some data which needs to be updated if the configuration (positions only!) has
changed too much.
"""
mutable struct JData{T}
   max_change::T     # how much X may change before recomputing
   accum_change::T   # how much has it changed already
   data::Any
end

struct LinearConstraint{T}
   C::Matrix{T}
   b::Vector{T}
   desc::String   # description of the constraint
end

mutable struct DofManager{T}
   variablecell::Bool
   xfree::Vector{Int}
   lincons::Vector{LinearConstraint{T}}
   # ----
   X0::Vector{JVec{T}}
   F0::JMat{T}
end

"""
`Atoms{T <: AbstractFloat} <: AbstractAtoms`

The main type storing information about atomic configurations. This is
not normally constructed directly by calling `Atoms`, but using one of the
utility functions such as `bulk`.

## Constructors

TODO

## Convenience functions

* `bulk`

# Getters and Setters

* `positions`, `set_positions!`
* `momenta`, `set_momenta!`
* `masses`, `set_masses!`
* `numbers`, `set_numbers!`
* `cell`, `set_cell!`
* `pbc`, `set_pbc!`
* `calculator`, `set_calculator!`
* `get_data`, `set_data!`
"""
mutable struct Atoms{T <: AbstractFloat} <: AbstractAtoms{T}
   X::Vector{JVec{T}}              # positions
   P::Vector{JVec{T}}              # momenta (or velocities?)
   M::Vector{T}                    # masses
   Z::Vector{AtomicNumber}         # atomic numbers
   cell::JMat{T}                   # cell
   pbc::JVec{Bool}                 # boundary condition
   calc::Union{Nothing, AbstractCalculator}
   dofmgr::DofManager
   data::Dict{Any,JData{T}}
end

Atoms(X::Vector{JVec{T}}, P::Vector{JVec{T}}, M::Vector{T},
      Z::Vector{AtomicNumber}, cell::JMat{T}, pbc::JVec{Bool},
      calc::Union{Nothing, AbstractCalculator} = nothing) where {T} =
   Atoms(X, P, M, Z, cell, pbc, calc,
         DofManager(length(X), T), Dict{Any,JData{T}}())

fltype(::Atoms{T}) where {T} = T

# derived properties
length(at::Atoms) = length(at.X)

getindex(at::Atoms, i::Integer) = at.X[i]

function setindex!(at::Atoms{T}, val, i::Integer) where {T}
   u = norm(at.X[i] - val)
   at.X[i] = convert(JVec{T}, val)
   update_data!(at, u)
   return val
end

# ------------------------ access to struct fields ----------------------

# ----- getters

positions(at::Atoms) = copy(at.X)
momenta(at::Atoms) = copy(at.P)
masses(at::Atoms) = copy(at.M)
numbers(at::Atoms) = copy(at.Z)
atomic_numbers(at::Atoms) = copy(at.Z)
cell(at::Atoms) = at.cell
pbc(at::Atoms) = at.pbc
calculator(at::Atoms) = at.calc

chemical_symbols(at::Atoms) = Chemistry.chemical_symbol.(at.Z)

# ----- setters

function set_positions!(at::Atoms{T}, X::AbstractVector{JVec{T}})  where T
   update_data!(at, dist(at, X))
   at.X .= X
   return at
end
function set_momenta!(at::Atoms{T}, P::AbstractVector{JVec{T}})  where T
   at.P .= P
   return at
end
function set_masses!(at::Atoms{T}, M::AbstractVector{T})  where T
   update_data!(at, Inf)
   at.M .= M
   return at
end
function set_numbers!(at::Atoms, Z::AbstractVector{<:Integer})
   update_data!(at, Inf)
   at.Z .= AtomicNumber.(Z)
   return at
end
function set_cell!(at::Atoms{T}, C::AbstractMatrix) where T
   update_data!(at, Inf)
   at.cell = JMat{T}(C)
   return at
end
function set_pbc!(at::Atoms, p::Union{AbstractVector, Tuple})
   q = JVec{Bool}(p...)
   if at.pbc != q
      update_data!(at, Inf)
      at.pbc = q
   end
   return at
end
function set_calculator!(at::Atoms, calc::Union{AbstractCalculator, Nothing})
   at.calc = calc
   return at
end

# --------------- equality tests -------------------

import Base.==
==(at1::Atoms{T}, at2::Atoms{T}) where T = (
   isapprox(at1, at2, tol = 0.0)  &&
          (at1.data == at2.data)  &&
          (at1.calc == at2.calc)  &&
          (at1.dofmgr == at2.dofmgr) )

import Base.isapprox
function isapprox(at1::Atoms{T}, at2::Atoms{T}; tol = sqrt(eps(T))) where {T}
   if length(at1) != length(at2)
      return false
   end
   if tol > 0
      ndigits = floor(Int, abs(log10(tol)))
      X1 = [round.(x; digits=ndigits) for x in at1.X]
      X2 = [round.(x; digits=ndigits) for x in at2.X]
   else
      X1, X2 = at1.X, at2.X
   end
   p1 = sortperm(X1)
   p2 = sortperm(X2)
   return (maxdist(at1.X[p1], at2.X[p2]) <= tol) &&
          (maxdist(at1.P[p1], at2.P[p2]) <= tol)  &&
          (norm(at1.M[p1] - at2.M[p2], Inf) <= tol) &&
          (at1.Z[p1] == at2.Z[p2]) &&
          (norm(at1.cell - at2.cell, Inf) <= tol) &&
          (at1.pbc == at2.pbc)
end

# --------------- some aliases for easier / direct access --------------
#                 and some access not covered above

function neighbourlist(at::Atoms{T}, rcut::AbstractFloat;
            recompute=false, key="nlist:default", storelist=true,
            int_type = Int, kwargs...) where {T}
   # we are allowed to use a stored list
   if !recompute
      if has_data(at, key)
         nlist = get_data(at, key)::PairList{T, int_type}
         if cutoff(nlist) >= rcut
            return nlist
         end
      end
   end
   # we are either forces to recompute, or no nlist exists or is not
   # suitable for the request.
   nlist = PairList(positions(at), rcut, cell(at), pbc(at);
                    int_type=int_type, kwargs...)::PairList{T, int_type}
   if storelist
      set_data!(at, key, nlist, 0.0)
   end
   return nlist
end

# function neighbourlist(at::Atoms{T}, rcut::AbstractFloat;
#             recompute=false, key="nlist:default", kwargs...) where {T}
#    return PairList(positions(at), rcut, cell(at), pbc(at); kwargs...)
# end

neighbourlist(at::Atoms) = neighbourlist(at, cutoff(at))

# --------------- data Dict handling ------------------

"""
`has_data(at::Atoms, key)`: checks whether `at` is storing a datum with key `key`; see also `set_data!`
"""
has_data(at::Atoms, key) = haskey(at.data, key)

"""
`get_data(at::Atoms, key)`: retrieves a datum;
see also `set_data!`. Unlike `get_positions`, etc, this returns the original
datum, not a copy (unless it is immutable).
"""
get_data(at::Atoms, key) = (at.data[key]).data

"""
`set_data!(a::Atoms, key, value, max_change=Inf)`:

Stores `value` under an arbitrary key `key` inside `at`. The entry is paired
with an updatemeasure `chg`. Whenever atom positions or the cell change by a
distance `d`, the counter is incremented `chg += d`. When `chg > max_change`,
the entry is deleted from the dictionary.

The two most common uses are with `max_change = Inf` (default) - in this case
the entry will never be deleted; and `max_change = 0.0` - in this case the entry
will be deleted whenever atom positions or the cell change.

The change counter is used, e.g., for neighbourlists with buffer, and for
preconditioners so that the list need not be recomputed each time the
configuration changes.
"""
function set_data!(at::Atoms, key, value, max_change=Inf)
   at.data[key] = JData(max_change, 0.0, value)
   return at
end

"""
`function set_transient!(a::Atoms, key, value, max_change=0.0)`: this is an
alias for `set_data!` but with default `max_change = 0.0` instead of `max_change =
Inf`.
"""
set_transient!(at::Atoms, key, value, max_change=0.0) =
      set_data!(at, key, value, max_change)

"""
`update_data!(at::Atoms, r::Real)`:
increment the change counter in all stored data and delete the
expired ones.
"""
function update_data!(at::Atoms, r::Real)
   for (key, t) in at.data
      t.accum_change += r
      if t.accum_change >= t.max_change
         delete!(at.data, key)
      end
   end
   return at
end

# ------------------------ FIO  ----------------------

JuLIP.FIO.write_dict(at::Atoms) =
   Dict( "__id__" => "JuLIP_Atoms",
         "X"      => at.X,
         "P"      => at.P,
         "M"      => at.M,
         "Z"      => Int.(at.Z),
         "cell"   => at.cell,
         "pbc"    => at.pbc,
         "calc"   => nothing,
         "data"   => nothing )

Atoms(D::Dict) = Atoms(D["X"], D["P"], D["M"], AtomicNumber.(D["Z"]),
                       D["cell"], D["pbc"])
          # calc = D["calc"],
          # cons = D["cons"],
          # data = D["data"] )

JuLIP.FIO.read_dict(::Val{:JuLIP_Atoms}, D::Dict) = Atoms(D)

import ExtXYZ

# conversion rules from Atoms.data entries to ExtXYZ standard types
_write_convert(value) = nothing # default rule: do not write
_write_convert(value::AbstractVector) = value
_write_convert(value::AbstractVector{JVec{T}}) where T = value |> mat |> Matrix
_write_convert(value::JVec{T}) where T = value |> Array
_write_convert(value::JMat{T}) where T = value |> Matrix
_write_convert(value::Union{Number,String}) = value

# (value, natoms) -> destination ∈ (:info, :arrays, missing)
_dest(::Any, natoms::Int) = missing

# this is a little fragile, since it assumes all vectors of length `natoms` contain per-atom data
# this is problematic for e.g. `natoms=3`.
_dest(value::AbstractVector, natoms::Int) = length(value) == natoms ? :arrays : :info
_dest(::JVec{T}, natoms::Int) where T = :info
_dest(::JMat{T}, natoms::Int) where T = :info
_dest(::Union{Number,String}, natoms::Int) = :info

function _atoms_to_extxyz_dict(atoms::Atoms{T}) where {T}
   dict = write_dict(atoms)
   natoms = length(atoms)
   dict["N_atoms"] = natoms
   dict["cell"] = _write_convert(dict["cell"])
   dict["pbc"] = _write_convert(dict["pbc"])

   # try to figure out whether data entries are per-atom or per-config
   info = Dict{String}{Any}()
   arrays = Dict{String}{Any}()
   for key in keys(atoms.data)
      value = get_data(atoms, key)
      dest, value = _dest(value, natoms), _write_convert(value)
      if value === nothing || ismissing(dest)
         continue
      end
      key = string(key) # symbol -> string, if needed
      if dest == :arrays
         push!(arrays, key => value)
      else
         push!(info, key => value)
      end
   end

   # move Atoms fields into `arrays` sub-dict
   # first X and Z, which should always be initialised
   arrays["pos"] = _write_convert(pop!(dict, "X"))
   arrays["species"] = _write_convert(string.(chemical_symbol.(atoms.Z)))
   arrays["Z"] = _write_convert(pop!(dict, "Z"))

   # now M and P, which might not be set
   for (source, dest) in zip(["M", "P"], ["masses", "momenta"])
      if source ∈ keys(dict) && dict[source] !== nothing && length(dict[source]) == natoms
         arrays[dest] = _write_convert(pop!(dict, source))
      end
   end

   # tidy up top-level dict entries
   dict["info"] = info
   dict["arrays"] = arrays
   delete!(dict, "data") # remove `nothing` entry
   delete!(dict, "calc") # remove `nothing` entry
   delete!(dict, "__id__")

   return dict
end

# define conversions to be applied when reading from ExtXYZ files to JuLIP Atoms
_read_convert(value, ::Int) = value
_read_convert(value::AbstractMatrix, natoms::Int) = size(value, 2) == natoms ? vecs(value) : value

function _extxyz_dict_to_atoms(dict)
   dict["__id__"] = "JuLIP_Atoms"

   natoms = Int(pop!(dict, "N_atoms"))
   info = pop!(dict, "info")
   arrays = pop!(dict, "arrays")

   "pos" in keys(arrays) || error("arrays dictionary missing 'pos' entry containing positions")
   dict["X"] = _read_convert(pop!(arrays, "pos"), natoms)
   @assert length(dict["X"]) == natoms

   # atomic numbers and symbols
   dict["Z"] = "Z" in keys(arrays) ? pop!(arrays, "Z") : nothing
   if "species" in keys(arrays)
      species = _read_convert(pop!(arrays, "species"), natoms)
      Zsp = atomic_number.([Symbol(sp) for sp in species])
      if dict["Z"] !== nothing
         all(dict["Z"] .== Zsp) || error("inconsistent 'Z' and 'species' properties")
      else
         dict["Z"] = Zsp
      end
   end
   dict["Z"] === nothing && error("atomic numbers not defined - either 'Z' or 'species' must be present")

   # mass - lookup from atomic number if not present
   if "masses" in keys(arrays)
      dict["M"] = _read_convert(pop!(arrays, "masses"), natoms)
   else
      dict["M"] = [atomic_mass(z) for z in AtomicNumber.(dict["Z"])]
   end

   # momenta / velocities
   if "momenta" in keys(arrays)
      dict["P"] = _read_convert(pop!(arrays, "momenta"), natoms)
   else
      dict["P"] = _read_convert(zeros((3, nat)), natoms)
   end

   atoms = read_dict(dict)

   # everything else goes in data
   for (label, D) in zip((:info, :arrays), (info, arrays))
      for (key, value) in D
            value = _read_convert(value, natoms)
            set_data!(atoms, key, value)
      end
   end
   return atoms
end

"""
   read_extxyz(file, atoms[, range])

Read from an extended XYZ file using ExtXYZ.jl.

`file` can be a string filename, open IO stream or `FILE*` pointer.
`range` can be absent, an integer frame number, integer range, or array of integers.
Default behaviour is to read all frames in file.

Returns a vector of `Atoms` structs
"""
function read_extxyz(file, args...; kwargs...)
   results = Atoms[]
   for dict in ExtXYZ.read_frames(file, args...; kwargs...)
      push!(results, _extxyz_dict_to_atoms(dict))
   end
   return results
end

"""
   write_extxyz(file, atoms)

Write atoms to file in extended XYZ format using `ExtXYZ.jl`.

`file` can be a string filename, open IO stream or `FILE*` pointer.
`atoms` can be either a single `JuLIP.Atoms` or a vector of `JuLIP.Atoms`.
"""
write_extxyz(file, atoms::Atoms{T}) where T = ExtXYZ.write_frame(file, _atoms_to_extxyz_dict(atoms))

write_extxyz(file, atoms::Vector{Atoms{T}}) where T = ExtXYZ.write_frames(file, _atoms_to_extxyz_dict.(atoms))

export read_extxyz, write_extxyz


using MechanicalSketch
import MechanicalSketch: background, background, sethue, O, EM, finish, empty_figure
import MechanicalSketch: lumin, get_current_lumin, get_current_RGB, color_with_lumin
import MechanicalSketch: sethue, fontface, textpath, textextents, PT, fontsize, gsave, grestore
import MechanicalSketch: newpath, translate, m, text, fillpreserve, clip, rect, clipreset
import MechanicalSketch: setline, PALETTE, color_from_palette

let
BACKCOLOR = color_with_lumin(PALETTE[2], 70);
function restart()
    empty_figure(filename = joinpath(@__DIR__, "test_10.png"))
    background(BACKCOLOR)
    sethue(PALETTE[2])
end

# Electric power symbol, implemented in 'power.jl' `showpower`

# parameters
power = 3.0kW
maxpower = 5kW |> typeof(power)
fractionpower = power / maxpower
rounded = round(typeof(power), power, digits = 1)
fs = 4EM
symbsize = 1.5fs

restart()

begin
    # Prepare color selection
    luminback = lumin(BACKCOLOR)
    luminfront = get_current_lumin()
    deltalumin = luminback - luminfront
    avglumin = luminfront + 0.5deltalumin
    curcol = get_current_RGB()
    avgcol = color_with_lumin(curcol, luminfront + 0.5 * deltalumin)
    contrastcol = color_with_lumin(curcol, luminfront + 2 * deltalumin);
end;

gsave()
fontsize(symbsize)
fontface("DejaVu Sans")
str = "⚡"
WI, HE = textextents(str)[3:4]
p = O + (-11m, 6m)
begin
    # Outline symbol
    sethue(avgcol)
    fontface("DejaVu Sans Bold")
    fontsize(1.0989010989010988 * symbsize)
    text(str, p - fs.*(0.03, -0.04))
    # Shadow text
    fontface("DejaVu Sans")
    fontsize(fs)
    text(string(rounded), p + (1.5WI, 0) - fs.*(0.03, -0.04))
    # Inside symbol
    sethue(color_with_lumin(color_from_palette("red"), luminback))
    fontsize(symbsize)
    text(str, p)
    # Front text
    fontsize(fs)
    text(string(rounded), p + (1.5WI, 0))
end

# Make a mask
begin
    clipreset()
    newpath()
    translate(p)
    fontsize(symbsize)
    textpath(str)
    fillpreserve()
    clip()
    translate(-p)
end

sethue(curcol)
rect(p, WI , -fractionpower * HE, :fill)
sethue(contrastcol)
setline(PT)
rect(p, WI , -fractionpower * HE, :stroke)
clipreset()
grestore()
finish()
nothing
end