module Test06

import Automa
import Automa.RegExp: @re_str
using Test

@testset "Test06" begin
    re = Automa.RegExp

    foo = re.cat("foo")
    foos = re.rep(re.cat(foo, re" *"))
    foo.actions[:exit]  = [:foo]
    actions = Dict(:foo => :(push!(ret, state.p:p-1); @escape))

    machine = Automa.compile(foos)

    @eval mutable struct MachineState
        p::Int
        cs::Int
        function MachineState()
            $(Automa.generate_init_code(Automa.CodeGenContext(), machine))
            return new(p, cs)
        end
    end

    for generator in (:table, :inline, :goto), checkbounds in (true, false), clean in (true, false)
        ctx = Automa.CodeGenContext(generator=generator, checkbounds=checkbounds, clean=clean)
        run! = @eval function (state, data)
            ret = []
            p = state.p
            cs = state.cs
            p_end = p_eof = lastindex(data)
            $(Automa.generate_exec_code(ctx, machine, actions))
            state.p = p
            state.cs = cs
            return ret
        end
        state = MachineState()
        data = b"foo foofoo   foo"
        @test run!(state, data) == [1:3]
        @test run!(state, data) == [5:7]
        @test run!(state, data) == [9:10]
        @test run!(state, data) == [12:16]
        @test run!(state, data) == []
        @test run!(state, data) == []
    end
end

end


"""
    mutable struct Reaction
        id::String
        metabolites::Dict{String,Float64}
        lb::Float64
        ub::Float64
        grr::Maybe{GeneAssociation}
        subsystem::Maybe{String}
        notes::Notes
        annotations::Annotations
        objective_coefficient::Float64
    end

A structure for representing a single reaction in a [`StandardModel`](@ref).
"""
mutable struct Reaction
    id::String
    metabolites::Dict{String,Float64}
    lb::Float64
    ub::Float64
    grr::Maybe{GeneAssociation}
    subsystem::Maybe{String}
    notes::Notes
    annotations::Annotations
    objective_coefficient::Float64
end

"""
    Reaction(
        id = "";
        metabolites = Dict{String,Float64}(),
        lb = -_constants.default_reaction_bound,
        ub = _constants.default_reaction_bound,
        grr = nothing,
        subsystem = nothing,
        notes = Notes(),
        annotations = Annotations(),
        objective_coefficient = 0.0,
    )

A constructor for Reaction that only takes a reaction `id` and
assigns default/uninformative values to all the fields that are not
explicitely assigned.
"""
function Reaction(
    id = "";
    metabolites = Dict{String,Float64}(),
    lb = -_constants.default_reaction_bound,
    ub = _constants.default_reaction_bound,
    grr = nothing,
    subsystem = nothing,
    notes = Notes(),
    annotations = Annotations(),
    objective_coefficient = 0.0,
)
    mets = Dict(k => float(v) for (k, v) in metabolites)
    return Reaction(
        id,
        mets,
        lb,
        ub,
        grr,
        subsystem,
        notes,
        annotations,
        objective_coefficient,
    )
end

"""
    Reaction(
        id::String,
        metabolites::Dict{String,Union{Int, Float64}},
        dir = :bidirectional;
        default_bound = _constants.default_reaction_bound,
    )

Convenience constructor for `Reaction`. The reaction equation is specified using
`metabolites`, which is a dictionary mapping metabolite ids to stoichiometric
coefficients. The direcion of the reaction is set through `dir` which can take
`:bidirectional`, `:forward`, and `:reverse` as values. Finally, the
`default_bound` is the value taken to mean infinity in the context of constraint
based models, often this is set to a very high flux value like 1000.
"""
function Reaction(
    id::String,
    metabolites,
    dir = :bidirectional;
    default_bound = _constants.default_reaction_bound,
)
    if dir == :forward
        lb = 0.0
        ub = default_bound
    elseif dir == :reverse
        lb = -default_bound
        ub = 0.0
    else
        lb = -default_bound
        ub = default_bound
    end
    Reaction(id; metabolites = metabolites, lb = lb, ub = ub)
end


module Enzyme

export Forward, Reverse
export Const, Active, Duplicated, DuplicatedNoNeed, BatchDuplicated, BatchDuplicatedNoNeed
export autodiff, jacobian, gradient, gradient!
export markType, batch_size, onehot, chunkedonehot

using LinearAlgebra

# Independent code, must be loaded before "compiler.jl"
include("pmap.jl")

using Adapt

"""
    abstract type Annotation{T}

Abstract type for [`autodiff`](@ref) function argument wrappers like
[`Const`](@ref), [`Active`](@ref) and [`Duplicated`](@ref).
"""
abstract type Annotation{T} end
Base.eltype(::Type{<:Annotation{T}}) where T = T

"""
    Const(x)

Mark a function argument `x` of [`autodiff`](@ref) as constant,
Enzyme will not auto-differentiate in respect `Const` arguments.
"""
struct Const{T} <: Annotation{T}
    val::T
end
Adapt.adapt_structure(to, x::Const) = Const(adapt(to, x.val))

# To deal with Const(Int) and prevent it to go to `Const{DataType}(T)`
Const(::Type{T}) where T = Const{Type{T}}(T)

"""
    Active(x)

Mark a function argument `x` of [`autodiff`](@ref) as active,
Enzyme will auto-differentiate in respect `Active` arguments.

!!! note

    Enzyme gradients with respect to integer values are zero.
    [`Active`](@ref) will automatically convert plain integers to floating
    point values, but cannot do so for integer values in tuples and structs.
"""
struct Active{T} <: Annotation{T}
    val::T
end
Adapt.adapt_structure(to, x::Active) = Active(adapt(to, x.val))

Active(i::Integer) = Active(float(i))

"""
    Duplicated(x, ∂f_∂x)

Mark a function argument `x` of [`autodiff`](@ref) as duplicated, Enzyme will
auto-differentiate in respect to such arguments, with `dx` acting as an
accumulator for gradients (so ``\\partial f / \\partial x`` will be *added to*)
`∂f_∂x`.
"""
struct Duplicated{T} <: Annotation{T}
    val::T
    dval::T
end
Adapt.adapt_structure(to, x::Duplicated) = Duplicated(adapt(to, x.val), adapt(to, x.dval))

"""
    DuplicatedNoNeed(x, ∂f_∂x)

Like [`Duplicated`](@ref), except also specifies that Enzyme may avoid computing
the original result and only compute the derivative values.
"""
struct DuplicatedNoNeed{T} <: Annotation{T}
    val::T
    dval::T
end
Adapt.adapt_structure(to, x::DuplicatedNoNeed) = DuplicatedNoNeed(adapt(to, x.val), adapt(to, x.dval))

"""
    BatchDuplicated(x, ∂f_∂xs)

Like [`Duplicated`](@ref), except contains several shadows to compute derivatives
for all at once. Argument `∂f_∂xs` should be a tuple of the several values of type `x`.
"""
struct BatchDuplicated{T,N} <: Annotation{T}
    val::T
    dval::NTuple{N,T}
end
Adapt.adapt_structure(to, x::BatchDuplicated) = BatchDuplicated(adapt(to, x.val), adapt(to, x.dval))

"""
    BatchDuplicatedNoNeed(x, ∂f_∂xs)

Like [`DuplicatedNoNeed`](@ref), except contains several shadows to compute derivatives
for all at once. Argument `∂f_∂xs` should be a tuple of the several values of type `x`.
"""
struct BatchDuplicatedNoNeed{T,N} <: Annotation{T}
    val::T
    dval::NTuple{N,T}
end
batch_size(::BatchDuplicated{T,N}) where {T,N} = N
batch_size(::BatchDuplicatedNoNeed{T,N}) where {T,N} = N
Adapt.adapt_structure(to, x::BatchDuplicatedNoNeed) = BatchDuplicatedNoNeed(adapt(to, x.val), adapt(to, x.dval))

"""
    abstract type Mode

Abstract type for what differentiation mode will be used.
"""
abstract type Mode end

"""
    struct Reverse <: Mode

Reverse mode differentiation
"""
struct ReverseMode <: Mode
end
const Reverse = ReverseMode()
guess_activity(::Type{T}, ::ReverseMode) where T = guess_activity(T)

"""
    struct Forward <: Mode

Forward mode differentiation
"""
struct ForwardMode <: Mode
end
const Forward = ForwardMode()
guess_activity(::Type{T}, ::ForwardMode) where T = guess_activity(T, API.DEM_ForwardMode)

import LLVM
include("api.jl")

@inline function guess_activity(::Type{T}, Mode::API.CDerivativeMode=API.DEM_ReverseModeCombined) where {T}
    return Const{T}
end
@inline function guess_activity(::Type{T}, Mode::API.CDerivativeMode=API.DEM_ReverseModeCombined) where {T<:AbstractFloat}
    if Mode == API.DEM_ForwardMode
        return DuplicatedNoNeed{T}
    else
        return Active{T}
    end
end
@inline function guess_activity(::Type{T}, Mode::API.CDerivativeMode=API.DEM_ReverseModeCombined) where {T<:Complex{<:AbstractFloat}}
    if Mode == API.DEM_ForwardMode
        return DuplicatedNoNeed{T}
    else
        return Active{T}
    end
end

@inline function guess_activity(::Type{T}, Mode::API.CDerivativeMode=API.DEM_ReverseModeCombined) where {T<:AbstractArray}
    if Mode == API.DEM_ForwardMode
        return DuplicatedNoNeed{T}
    else
        return Duplicated{T}
    end
end

include("logic.jl")
include("typeanalysis.jl")
include("typetree.jl")
include("utils.jl")
include("compiler.jl")

import .Compiler: CompilationException

# @inline annotate() = ()
# @inline annotate(arg::A, args::Vararg{Any, N}) where {A<:Annotation, N} = (arg, annotate(args...)...)
# @inline annotate(arg, args::Vararg{Any, N}) where N = (Const(arg), annotate(args...)...)

@inline function annotate(args::Vararg{Any, N}) where N
    ntuple(Val(N)) do i
        Base.@_inline_meta
        arg = @inbounds args[i]
        if arg isa Annotation
            return arg
        else
            return Const(arg)
        end
    end
end

@inline function same_or_one(args::Vararg{Any, N}) where N
    current = -1
    for arg in args
        if arg isa BatchDuplicated
            if current == -1
                current = batch_size(arg)
            else
                @assert current == batch_size(arg)
            end
        elseif arg isa BatchDuplicatedNoNeed
            if current == -1
                current = batch_size(arg)
            else
                @assert current == batch_size(arg)
            end
        end
    end

    if current == -1
        current = 1
    end

    return current
end

"""
    autodiff(::ReverseMode, f, Activity, args...)

Auto-differentiate function `f` at arguments `args` using reverse mode.

Limitations:

* `f` may only return a `Real` (of a built-in/primitive type) or `nothing`,
  not an array, struct, `BigFloat`, etc. To handle vector-valued return
  types, use a mutating `f!` that returns `nothing` and stores it's return
  value in one of the arguments, which must be wrapped in a
  [`Duplicated`](@ref).

`args` may be numbers, arrays, structs of numbers, structs of arrays and so
on. Enzyme will only differentiate in respect to arguments that are wrapped
in an [`Active`](@ref) (for immutable arguments like primitive types and
structs thereof) or [`Duplicated`](@ref) (for mutable arguments like arrays,
`Ref`s and structs thereof). Non-annotated arguments will automatically be
treated as [`Const`](@ref).

`Activity` is the Activity of the return value, it may be `Const` or `Active`.

Example:

```jldoctest
a = 4.2
b = [2.2, 3.3]; ∂f_∂b = zero(b)
c = 55; d = 9

f(a, b, c, d) = a * √(b[1]^2 + b[2]^2) + c^2 * d^2
∂f_∂a, ∂f_∂d = autodiff(Reverse, f, Active, Active(a), Duplicated(b, ∂f_∂b), c, Active(d))

# output

(3.966106403010388, 54450.0)
```

here, `autodiff` returns a tuple
``(\\partial f/\\partial a, \\partial f/\\partial d)``,
while ``\\partial f/\\partial b`` will be *added to* `∂f_∂b` (but not returned).
`c` will be treated as `Const(c)`.

!!! note

    Enzyme gradients with respect to integer values are zero.
    [`Active`](@ref) will automatically convert plain integers to floating
    point values, but cannot do so for integer values in tuples and structs.
"""
@inline function autodiff(::ReverseMode, f::F, ::Type{A}, args...) where {F, A<:Annotation}
    args′  = annotate(args...)
    tt′    = Tuple{map(Core.Typeof, args′)...}
    width = Val(same_or_one(args...))
    if A <: Active
        tt    = Tuple{map(T->eltype(Core.Typeof(T)), args′)...}
        rt = Core.Compiler.return_type(f, tt)
        if !allocatedinline(rt)
            forward, adjoint = Enzyme.Compiler.thunk(f, #=df=#nothing, Duplicated{rt}, tt′, #=Split=# Val(API.DEM_ReverseModeGradient), width, #=ModifiedBetween=#Val(false))
            res = forward(args′...)
            @assert length(res) == 2
            tape = res[1]
            if res[2] isa Base.RefValue
                res[2][] += one(eltype(typeof(res[2])))
            else
                res[2] += one(eltype(typeof(res[2])))
            end
            return adjoint(args′..., tape)
        end
    elseif A <: Duplicated || A<: DuplicatedNoNeed || A <: BatchDuplicated || A<: BatchDuplicatedNoNeed
        throw(ErrorException("Duplicated Returns not yet handled"))
    end
    thunk = Enzyme.Compiler.thunk(f, #=df=#nothing, A, tt′, #=Split=# Val(API.DEM_ReverseModeCombined), width)
    if A <: Active
        tt    = Tuple{map(T->eltype(Core.Typeof(T)), args′)...}
        rt = eltype(Compiler.return_type(thunk))
        args′ = (args′..., one(rt))
    end
    thunk(args′...)
end

@inline function autodiff(::ReverseMode, dupf::Duplicated{F}, ::Type{A}, args...) where {F, A<:Annotation}
    args′  = annotate(args...)
    tt′    = Tuple{map(Core.Typeof, args′)...}
    width = Val(same_or_one(args...))
    thunk = Enzyme.Compiler.thunk(#=f=#dupf.val, #=df=#dupf.dval, A, tt′, #=Split=# Val(API.DEM_ReverseModeCombined), width)
    if A <: Active
        rt = eltype(Compiler.return_type(thunk))
        args′ = (args′..., one(rt))
    end
    thunk(args′...)
end

"""
    autodiff(mode::Mode, f, args...)
    autodiff(f, mode::Mode, args...)

Like [`autodiff`](@ref) but will try to guess the activity of the return value.
"""
@inline function autodiff(mode::Mode, f::F, args...) where {F}
    args′ = annotate(args...)
    tt′   = Tuple{map(Core.Typeof, args′)...}
    tt    = Tuple{map(T->eltype(Core.Typeof(T)), args′)...}
    rt    = Core.Compiler.return_type(f, tt)
    A     = guess_activity(rt, mode)
    autodiff(mode, f, A, args′...)
end

@inline function autodiff(mode::Mode, dupf::Duplicated{F}, args...) where {F}
    args′ = annotate(args...)
    tt′   = Tuple{map(Core.Typeof, args′)...}
    tt    = Tuple{map(T->eltype(Core.Typeof(T)), args′)...}
    rt    = Core.Compiler.return_type(dupf.val, tt)
    A     = guess_activity(rt, mode)
    autodiff(mode, dupf, A, args′...)
end

"""
    autodiff(::ForwardMode, f, Activity, args...)

Auto-differentiate function `f` at arguments `args` using forward mode.

`args` may be numbers, arrays, structs of numbers, structs of arrays and so
on. Enzyme will only differentiate in respect to arguments that are wrapped
in a [`Duplicated`](@ref) or similar argument. Non-annotated arguments will
automatically be treated as [`Const`](@ref). Unlike reverse mode in
[`autodiff`](@ref), [`Active`](@ref) arguments are not allowed here, since
all

`Activity` is the Activity of the return value, it may be:
* `Const` if the return is not to be differentiated with respect to
* `Duplicated`, if the return is being differentiated with respect to and
  both the original value and the derivative return are desired
* `DuplicatedNoNeed`, if the return is being differentiated with respect to
  and only the derivative return is desired.

Example returning both original return and derivative:

```jldoctest
a = 4.2
b = [2.2, 3.3]; ∂f_∂b = zero(b)
c = 55; d = 9

f(x) = x*x
res, ∂f_∂x = autodiff(Forward, f, Duplicated, Duplicated(3.14, 1.0))

# output

(9.8596, 6.28)
```

Example returning just the derivative:

```jldoctest
a = 4.2
b = [2.2, 3.3]; ∂f_∂b = zero(b)
c = 55; d = 9

f(x) = x*x
∂f_∂x = autodiff(Forward, f, DuplicatedNoNeed, Duplicated(3.14, 1.0))

# output

(6.28,)
```
"""
@inline function autodiff(::ForwardMode, f::F, ::Type{A}, args...) where {F, A<:Annotation}
    args′  = annotate(args...)
    tt′    = Tuple{map(Core.Typeof, args′)...}
    width = Val(same_or_one(args...))
    if A <: Active
        throw(ErrorException("Active Returns not allowed in forward mode"))
    end
    ReturnPrimal = Val(A <: Duplicated || A <: BatchDuplicated)
    thunk = Enzyme.Compiler.thunk(f, #=df=#nothing, A, tt′, #=Mode=# Val(API.DEM_ForwardMode), width,
                                     #=ModifiedBetween=#Val(false), ReturnPrimal)
    thunk(args′...)
end

@inline function autodiff(::ForwardMode, dupf::Duplicated{F}, ::Type{A}, args...) where {F, A<:Annotation}
    args′  = annotate(args...)
    tt′    = Tuple{map(Core.Typeof, args′)...}
    width = Val(same_or_one(args...))
    if A <: Active
        throw(ErrorException("Active Returns not allowed in forward mode"))
    end
    ReturnPrimal = Val(A <: Duplicated || A <: BatchDuplicated)
    thunk = Enzyme.Compiler.thunk(#=f=#dupf.val, #=df=#dupf.dval, A, tt′, #=Mode=# Val(API.DEM_ForwardMode), width, #=ModifiedBetween=#Val(false), ReturnPrimal)
    thunk(args′...)
end

# Compat
@deprecate fwddiff(f::F, ::Type{A}, args...) where {F, A<:Annotation} autodiff(Forward, f, A, args...)
@deprecate fwddiff(f::F, args...) where {F} autodiff(Forward, f, args...)

# Compat
@inline autodiff(f::F, ::Type{A}, args...) where {F, A<:Annotation} = autodiff(Reverse, f, A, args...)
@inline autodiff(f::F, args...) where {F} = autodiff(Reverse, f, args...)

# F as first arg for `do` syntax
@inline autodiff(dupf::Duplicated{F}, mode::Mode, ::Type{A}, args...) where {F,A<:Annotation} = autodiff(mode, dupf, A, args...)
@inline autodiff(f::F, mode::Mode, ::Type{A}, args...) where {F,A<:Annotation} = autodiff(mode, f, A, args...)
@inline autodiff(dupf::Duplicated{F}, mode::Mode, args...) where {F} = autodiff(mode, dupf, args...)
@inline autodiff(f::F, mode::Mode, args...) where {F} = autodiff(mode, f, args...)

"""
    autodiff_deferred(f, Activity, args...)

Same as [`autodiff`](@ref) but uses deferred compilation to support usage in GPU
code, as well as high-order differentiation.
"""
@inline function autodiff_deferred(f::F, ::Type{A}, args...) where {F, A<:Annotation}
    args′ = annotate(args...)
    tt′   = Tuple{map(Core.Typeof, args′)...}
    width = Val(same_or_one(args...))
    if A isa UnionAll
        tt = Tuple{map(T->eltype(Core.Typeof(T)), args′)...}
        rt = Core.Compiler.return_type(f, tt)
        rt = A{rt}
    else
        @assert A isa DataType
        rt = A
    end

    if eltype(rt) == Union{}
        error("Return type inferred to be Union{}. Giving up.")
    end

    ReturnPrimal = Val(false)
    ptr   = Compiler.deferred_codegen(f, Val(tt′), Val(rt), #=dupClosure=#Val(false), Val(API.DEM_ReverseModeCombined), width, #=ModifiedBetween=#Val(false), ReturnPrimal)
    thunk = Compiler.CombinedAdjointThunk{F, rt, tt′, typeof(width), Nothing, ReturnPrimal}(f, ptr, #=df=#nothing)
    if rt <: Active
        args′ = (args′..., one(eltype(rt)))
    elseif A <: Duplicated || A<: DuplicatedNoNeed || A <: BatchDuplicated || A<: BatchDuplicatedNoNeed
        throw(ErrorException("Duplicated Returns not yet handled"))
    end
    thunk(args′...)
end

"""
    autodiff_deferred(f, args...)

Like [`autodiff_deferred`](@ref) but will try to guess the activity of the return value.
"""
@inline function autodiff_deferred(f::F, args...) where {F}
    args′ = annotate(args...)
    tt′   = Tuple{map(Core.Typeof, args′)...}
    tt    = Tuple{map(T->eltype(Core.Typeof(T)), args′)...}
    rt    = Core.Compiler.return_type(f, tt)
    rt    = guess_activity(rt)
    autodiff_deferred(f, rt, args′...)
end

"""
    fwddiff_deferred(f, Activity, args...)

Same as `autodiff(::ForwardMode, ...)` but uses deferred compilation to support usage in GPU
code, as well as high-order differentiation.
"""
@inline function fwddiff_deferred(f::F, ::Type{A}, args...) where {F, A<:Annotation}
    args′ = annotate(args...)
    tt′   = Tuple{map(Core.Typeof, args′)...}
    width = Val(same_or_one(args...))
    if A isa UnionAll
        tt = Tuple{map(T->eltype(Core.Typeof(T)), args′)...}
        rt = Core.Compiler.return_type(f, tt)
        rt = A{rt}
    else
        @assert A isa DataType
        rt = A
    end

    if eltype(rt) == Union{}
        error("Return type inferred to be Union{}. Giving up.")
    end

    if A <: Active
        throw(ErrorException("Active Returns not allowed in forward mode"))
    end

    ReturnPrimal = Val(A <: Duplicated || A <: BatchDuplicated)
    ptr   = Compiler.deferred_codegen(f, Val(tt′), Val(rt), #=dupClosure=#Val(false), Val(API.DEM_ForwardMode), width, #=ModifiedBetween=#Val(false), ReturnPrimal)
    thunk = Compiler.ForwardModeThunk{F, rt, tt′, typeof(width), Nothing, ReturnPrimal}(f, ptr, #=df=#nothing)
    thunk(args′...)
end

"""
    fwddiff_deferred(f, args...)

Like [`fwddiff_deferred`](@ref) but will try to guess the activity of the return value.
"""
@inline function fwddiff_deferred(f::F, args...) where {F}
    args′ = annotate(args...)
    tt    = Tuple{map(T->eltype(Core.Typeof(T)), args′)...}
    rt    = Core.Compiler.return_type(f, tt)
    rt    = guess_activity(rt, API.DEM_ForwardMode)
    fwddiff_deferred(f, rt, args′...)
end

# White lie, should be `Core.LLVMPtr{Cvoid, 0}` but that's not supported by ccallable
Base.@ccallable function __enzyme_float(x::Ptr{Cvoid})::Cvoid
    return nothing
end

Base.@ccallable function __enzyme_double(x::Ptr{Cvoid})::Cvoid
    return nothing
end

@inline function markType(::Type{T}, ptr::Ptr{Cvoid}) where T
    markType(Base.unsafe_convert(Ptr{T}, ptr))
end

@inline function markType(data::Array{T}) where T
    GC.@preserve data markType(pointer(data))
end

# TODO(WM): We record the type of a single index here, we could give it a range
@inline function markType(data::SubArray)
    GC.@preserve data markType(pointer(data))
end

@inline function markType(data::Ptr{Float32})
    Base.llvmcall(("declare void @__enzyme_float(i8* nocapture) nounwind define void @c(i64 %q) nounwind alwaysinline { %p = inttoptr i64 %q to i8* call void @__enzyme_float(i8* %p) ret void }", "c"), Cvoid, Tuple{Ptr{Float32}}, data)
    nothing
end

@inline function markType(data::Ptr{Float64})
    Base.llvmcall(("declare void @__enzyme_double(i8* nocapture) nounwind define void @c(i64 %q) nounwind alwaysinline { %p = inttoptr i64 %q to i8* call void @__enzyme_double(i8* %p) ret void }", "c"), Cvoid, Tuple{Ptr{Float64}}, data)
    nothing
end

@inline function onehot(x, start=1, endl=length(x))
    ntuple(Val(endl-start+1)) do i
        Base.@_inline_meta
        res = similar(x)
        for idx in 1:length(x)
            @inbounds res[idx] = (i + start - 1== idx) ? 1.0 : 0.0
        end
        return res
    end
end

@inline function onehot(x::NTuple{N, T}, start=1, endl=N) where {T, N}
    ntuple(Val(endl-start+1)) do i
        Base.@_inline_meta
        ntuple(N) do idx
            Base.@_inline_meta
            return (i + start - 1 == idx) ? 1.0 : 0.0
        end
    end
end

"""
    gradient(::ReverseMode, f, x)

Compute the gradient of an array-input function `f` using reverse mode.
This will allocate and return new array with the gradient result.

Example:

```jldoctest
f(x) = x[1]*x[2]

grad = gradient(Reverse, f, [2.0, 3.0])

# output

2-element Vector{Float64}:
 3.0
 2.0
```
"""
@inline function gradient(::ReverseMode, f, x)
    dx = zero(x)
    autodiff(Reverse, f, Duplicated(x, dx))
    dx
end

"""
    gradient!(::ReverseMode, dx, f, x)

Compute the gradient of an array-input function `f` using reverse mode,
storing the derivative result in an existing array `dx`.

Example:

```jldoctest
f(x) = x[1]*x[2]

dx = [0.0, 0.0]
gradient!(Reverse, dx, f, [2.0, 3.0])

# output

2-element Vector{Float64}:
 3.0
 2.0
```
"""
@inline function gradient!(::ReverseMode, dx, f, x)
    dx .= 0
    autodiff(Reverse, f, Duplicated(x, dx))
    dx
end

"""
    gradient(::ForwardMode, f, x; shadow=onehot(x))

Compute the gradient of an array-input function `f` using forward mode. The
optional keyword argument `shadow` is a vector of one-hot vectors of type `x`
which are used to forward-propagate into the return. For performance reasons,
this should be computed once, outside the call to `gradient`, rather than
within this call.

Example:

```jldoctest
f(x) = x[1]*x[2]

grad = gradient(Forward, f, [2.0, 3.0])

# output

(3.0, 2.0)
```
"""
@inline function gradient(::ForwardMode, f, x; shadow=onehot(x))
    only(autodiff(Forward, f, BatchDuplicatedNoNeed, BatchDuplicated(x, shadow)))
end

@inline function chunkedonehot(x, ::Val{chunk}) where chunk
    sz = length(x)
    num = ((sz + chunk - 1) ÷ chunk)
    ntuple(Val(num)) do i
        Base.@_inline_meta
        onehot(x, (i-1)*chunk+1, i == num ? sz : (i*chunk) )
    end
end

@inline tupleconcat(x) = x
@inline tupleconcat(x, y) = (x..., y...)
@inline tupleconcat(x, y, z...) = (x..., tupleconcat(y, z...)...)

"""
    gradient(::ForwardMode, f, x, ::Val{chunk}; shadow=onehot(x))

Compute the gradient of an array-input function `f` using vector forward mode.
Like [`gradient`](@ref), except it uses a chunk size of `chunk` to compute
`chunk` derivatives in a single call.

Example:

```jldoctest
f(x) = x[1]*x[2]

grad = gradient(Forward, f, [2.0, 3.0], Val(2))

# output

(3.0, 2.0)
```
"""
@inline function gradient(::ForwardMode, f, x, ::Val{chunk}; shadow=chunkedonehot(x, Val(chunk))) where chunk
    tmp = ntuple(length(shadow)) do i
        autodiff(Forward, f, BatchDuplicatedNoNeed, BatchDuplicated(x, shadow[i]))[1]
    end
    tupleconcat(tmp...)
end

@inline function gradient(::ForwardMode, f, x, ::Val{1}; shadow=onehot(x))
    ntuple(length(shadow)) do i
        autodiff(Forward, f, DuplicatedNoNeed, Duplicated(x, shadow[i]))[1]
    end
end

"""
    jacobian(::ForwardMode, f, x; shadow=onehot(x))
    jacobian(::ForwardMode, f, x, ::Val{chunk}; shadow=onehot(x))

Compute the jacobian of an array-input function `f` using (potentially vector)
forward mode. This is a simple rename of the [`gradient`](@ref) function,
and all relevant arguments apply here.

Example:

```jldoctest
f(x) = [x[1]*x[2], x[2]]

grad = jacobian(Forward, f, [2.0, 3.0])

# output

2×2 Matrix{Float64}:
 3.0  2.0
 0.0  1.0
```
"""
@inline function jacobian(::ForwardMode, args...; kwargs...)
    cols = gradient(Forward, args...; kwargs...)
    reduce(hcat, cols)
end

"""
    jacobian(::ReverseMode, f, x, ::Val{num_outs}, ::Val{chunk})

Compute the jacobian of an array-input function `f` using (potentially vector)
reverse mode. The `chunk` argument denotes the chunk size to use and `num_outs`
denotes the number of outputs `f` will return in an array.

Example:

```jldoctest
f(x) = [x[1]*x[2], x[2]]

grad = jacobian(Reverse, f, [2.0, 3.0], Val(2))

# output

2×2 Matrix{Float64}:
 3.0  2.0
 0.0  1.0
```
"""
@inline function jacobian(::ReverseMode, f, x, n_outs::Val{n_out_val}, ::Val{chunk}) where {chunk, n_out_val}
    num = ((n_out_val + chunk - 1) ÷ chunk)

    tt′   = Tuple{BatchDuplicated{Core.Typeof(x), chunk}}
    tt    = Tuple{Core.Typeof(x)}
    rt = Core.Compiler.return_type(f, tt)
    primal, adjoint = Enzyme.Compiler.thunk(f, #=df=#nothing, BatchDuplicatedNoNeed{rt}, tt′, #=Split=# Val(API.DEM_ReverseModeGradient), #=width=#Val(chunk), #=ModifiedBetween=#Val(false))

    if num * chunk == n_out_val
        last_size = chunk
        primal2, adjoint2 = primal, adjoint
    else
        last_size = n_out_val - (num-1)*chunk
        tt′ = Tuple{BatchDuplicated{Core.Typeof(x), last_size}}
        primal2, adjoint2 = Enzyme.Compiler.thunk(f, #=df=#nothing, BatchDuplicatedNoNeed{rt}, tt′, #=Split=# Val(API.DEM_ReverseModeGradient), #=width=#Val(last_size), #=ModifiedBetween=#Val(false))
    end

    tmp = ntuple(num) do i
        Base.@_inline_meta
        dx = ntuple(i == num ? last_size : chunk) do idx
            Base.@_inline_meta
            zero(x)
        end
        res = (i == num ? primal2 : primal)(BatchDuplicated(x, dx))
        tape = res[1]
        j = 0
        for shadow in res[2]
            j += 1
            @inbounds shadow[(i-1)*chunk+j] += one(eltype(typeof(shadow)))
        end
        (i == num ? adjoint2 : adjoint)(BatchDuplicated(x, dx), tape)
        return dx
    end
    rows = tupleconcat(tmp...)
    mapreduce(LinearAlgebra.adjoint, vcat, rows)
end

@inline function jacobian(::ReverseMode, f, x, n_outs::Val{n_out_val}, ::Val{1} = Val(1)) where {n_out_val}
    tt′   = Tuple{Duplicated{Core.Typeof(x)}}
    tt    = Tuple{Core.Typeof(x)}
    rt = Core.Compiler.return_type(f, tt)
    primal, adjoint = Enzyme.Compiler.thunk(f, #=df=#nothing, DuplicatedNoNeed{rt}, tt′, #=Split=# Val(API.DEM_ReverseModeGradient), #=width=#Val(1), #=ModifiedBetween=#Val(false))
    rows = ntuple(n_outs) do i
        Base.@_inline_meta
        dx = zero(x)
        res = primal(Duplicated(x, dx))
        tape = res[1]
        @inbounds res[2][i] += one(eltype(typeof(res[2])))
        adjoint(Duplicated(x, dx), tape)
        return dx
    end
    mapreduce(LinearAlgebra.adjoint, vcat, rows)
end

end # module


##########
##########
# Functions to perform perturbations on the food webs
##########
##########
# AUTHOR: Cole B. Brookson
# DATE OF CREATION: 2021-08-09
##########
##########

# set up =======================================================================
using DrWatson
@quickactivate "trait-based-rewiring"
# logger set up
using Logging: global_logger
using TerminalLoggers: TerminalLogger
global_logger(TerminalLogger())
using BenchmarkTools

include(srcdir("01_eqn_component_functions.jl"))
include(srcdir("00_species_list_creation.jl"))
include(srcdir("03_initial_sims_functions.jl"))
include(srcdir("04_perturbation_sims_functions.jl"))

# set all positional values
organism_id_pos = 1
organism_type_pos = 2
body_size_pos = 3
lower_prey_limit_pos = 4
upper_prey_limit_pos = 5
habitat_midpoint_pos = 6
habitat_lower_limit_pos = 7
habitat_upper_limit_pos = 8
nutritional_value_pos = 9
biomass_pos = 10
network_id_pos = 11
alive_pos = 12

# pulse perturbations ==========================================================

# little line to set the network id (this happens from the shell line)
network_id = parse(Int64, ARGS[1])

species_list = readdlm(datadir("./intermediate-files/species-lists-pre-perturbation/species_list_0001.csv"), ',', Any)
network = create_interaction_matrix(species_list)

# set pulse amounts
c = 210
x = 100
f = 1
z = c/(x*(1/f))
# set vector callback
extinction_callback =
    VectorContinuousCallback(extinction_threshold,
                            extinction_affect!,
                            size(species_list, 1),
                            save_positions = ntuple(x->true,
                                                    size(species_list, 1)))
steady_callback = TerminateSteadyState(1e-8,
                                        1e-6,
                                        DiffEqCallbacks.allDerivPass,
                                        min_t = 1205)

# set callback set
"""
    affect_inc!(integrator)

Affect for the condition to work on. Increase the mortality by making the amount
(the fourth element in p) equal to z.
"""
function affect_inc!(integrator, z)
  integrator.p[4] = z
end

"""
    affect_dec!(integrator)

Affect for the condition to work on. Decrease the fourth element in p,
returning it to 1.0.
"""
function affect_dec!(integrator)
  integrator.p[4] = nothing
end

save_positions1 = (true, true)
cb_inc = DiscreteCallback(inc_mort_cond, affect_inc!,
                            save_positions=save_positions1)
save_positions2 = (false, true)
cb_dec = DiscreteCallback(dec_mort_cond, affect_dec!,
                            save_positions=save_positions2)
callback_set = CallbackSet(extinction_callback,
                            #steady_callback, #keep this out until fixed
                            cb_inc,
                            cb_dec)

#  allocate x0
u0 = Float64[]
for i in 1:size(species_list, 1)
    append!(u0, species_list[i, biomass_pos])
end

# set perturbation tspan
tspan = (1000.0, 1100.0)

# set constant stop values
const stop1 = [1010.0]
const stop2 = [1080.0]
const tstop = [1010.;1080.]

# loop through random selections of species to perturb
#@info{"starting loop..."}

# set times to save from the simulation
save_times = [collect(range(1000.0, 1300.0, step = 0.1));
                collect(range(1250, 4800, step = sample(200:400, 1)[1]));
                collect(range(4800.0, 5000.0, step = 0.1))]

for j in 1:100

    # pick set of random species to perturb
    all = Array(1:size(species_list, 1))
    perturb_sp = rand(Array(1:size(species_list, 1)),
                        Int(round(size(species_list, 1)*(1/2))))
    non_perturb = setdiff(Set(all), Set(perturb_sp))

    # set p
    amount = nothing
    p = [species_list, network, 0.3, amount, perturb_sp, non_perturb]

    prob_pulse = ODEProblem(biomass_perturb_random,
                            u0,
                            tspan,
                            p,
                            tstops=tstop)
    @benchmark sol_prob_pulse = solve(prob_pulse,
                                        OwrenZen5(),
                                        callback = callback_set,
                                        reltol = 1e-4,
                                        # save every part for the perturbation and
                                        # the first 150 time steps after the
                                        # perturbation, but then just grab a few
                                        # from there to the end
                                        saveat = save_times,
                                        progress = true,
                                        progress_steps = 1)
    # save output of the solver
    sol_name = string("sol-pulse-net-", lpad(network_id,4,"0"), "-sim-", lpad(j, 4, "0"), ".csv")
    writedlm(datadir("./simulation-outpus/pulse-sims/output/$name"), sol_prob_pulse, ",")
    #plot(sol_prob_pulse)
    u_matrix = zeros(length(sol_prob_pulse.u), length(sol_prob_pulse.u[1]))
    for i in 1:size(u_matrix, 1)
        u_matrix[i,:] = sol_prob_pulse.u[i]
    end
    output = hcat(sol_prob_pulse.t, u_matrix)
    # write out the biomass outputs
    output_name = string("pulse-net-", lpad(network_id,4,"0"), "-sim-", lpad(j, 4, "0"), ".csv")
    writedlm(datadir("./simulation-outpus/pulse-sims/output/$name"), output_name, ",")
end


#Example command line to run this program on multiple processors:
# julia -p 4 "run_evolve_goals.jl"

@everywhere include("EvolveGoals.jl")

# Example runs.  Uncomment one of these lines
#run_evolve_goals("prun",2,3,10,4)
#run_evolve_goals("rrun",4,15,200,3,4)
#run_evolve_goals("erun",3,14,50,6)
run_evolve_goals("erun",3,14,50,6,4)
#run_evolve_goals("erun",3,14,50,20)

# In run_evolve_goals("frun",3,14,50,6),
#   numinputs = 3
#   random seed = 14
#   genome size = 50
#   number of repetitions per goal = 6
# Evolves all goals of length 3 with 6 repetitions per goal

# In run_evolve_goals("frun",4,14,200,3,4),
#   numinputs = 4
#   random seed = 14
#   genome size = 200
#   number of repetitions per goal = 3
#   number of random goals = 4

# An alternative that does not involve running this file
# julia -p 2  -L EvolveGoals.jl -e "run_evolve_goals(\"prun\",3,16,20,3,4)"


for SpatialVector in (:Twist, :SpatialAcceleration, :Momentum, :Wrench)
    @eval begin
        # Basics
        Base.eltype(::Type{$SpatialVector{T}}) where {T} = T
        angular(v::$SpatialVector) = v.angular
        linear(v::$SpatialVector) = v.linear

        # Construct/convert given another spatial vector
        $SpatialVector{T}(v::$SpatialVector{T}) where {T} = v
        Base.convert(::Type{V}, v::$SpatialVector) where {V<:$SpatialVector} = V(v)

        # Construct/convert to SVector
        StaticArrays.SVector{6, T}(v::$SpatialVector) where {T} = convert(SVector{6, T}, [angular(v); linear(v)])
        StaticArrays.SVector{6}(v::$SpatialVector{X}) where {X} = SVector{6, X}(v)
        StaticArrays.SVector(v::$SpatialVector) = SVector{6}(v)
        StaticArrays.SArray(v::$SpatialVector) = SVector(v)
        Base.convert(::Type{SA}, v::$SpatialVector) where {SA <: SArray} = SA(v)

        function Base.Array(v::$SpatialVector)
            Base.depwarn("Array(v) has been deprecated. Please use SVector(v) instead.", :Array)
            SVector(v)
        end
    end
end

for SpatialMatrix in (:GeometricJacobian, :MomentumMatrix, :WrenchMatrix)
    @eval begin
        # Basics
        Base.eltype(::Type{$SpatialMatrix{A}}) where {T, A<:AbstractMatrix{T}} = T
        Base.size(m::$SpatialMatrix) = (6, size(angular(m), 2))
        Base.size(m::$SpatialMatrix, d) = size(m)[d]
        Base.transpose(m::$SpatialMatrix) = Transpose(m)
        angular(m::$SpatialMatrix) = m.angular
        linear(m::$SpatialMatrix) = m.linear

        # Construct/convert given another spatial matrix
        $SpatialMatrix(m::$SpatialMatrix{A}) where {A} = SpatialMatrix{A}(m)
        Base.convert(::Type{$SpatialMatrix{A}}, m::$SpatialMatrix) where {A} = $SpatialMatrix{A}(m)
        Base.convert(::Type{$SpatialMatrix{A}}, m::$SpatialMatrix{A}) where {A} = m

        # Construct/convert to Matrix
        (::Type{A})(m::$SpatialMatrix) where {A<:Array} = convert(A, [angular(m); linear(m)])
        Base.convert(::Type{A}, m::$SpatialMatrix) where {A<:Array} = A(m)
    end
end


module PWscf

using Test

using Setfield
using StructArrays: StructArray

using QuantumESPRESSOBase
using QuantumESPRESSOBase.Inputs.PWscf

@testset "Constructing `AtomicSpecies`" begin
    # Data from https://github.com/QEF/q-e/blob/7be27df/PW/examples/gatefield/run_example#L128.
    x = AtomicSpecies("S", 32.066, "S.pz-n-rrkjus_psl.0.1.UPF")
    @test_throws AssertionError @set x.atom = "sulfur"
    @test_throws InexactError @set x.mass = 1im
    @test x == AtomicSpecies('S', 32.066, "S.pz-n-rrkjus_psl.0.1.UPF")
    @test AtomicSpecies(
        AtomicPosition('S', [0.500000000, 0.288675130, 1.974192764]),
        32.066,
        "S.pz-n-rrkjus_psl.0.1.UPF",
    ) == x
end

@testset "Test constructing `AtomicSpeciesCard` from `StructArray`s" begin
    a = ["Al", "As"]
    m = [24590.7655930491, 68285.4024548272]
    pp = ["Al.pbe-n-kjpaw_psl.1.0.0.UPF", "As.pbe-n-kjpaw_psl.1.0.0.UPF"]
    card = AtomicSpeciesCard(StructArray{AtomicSpecies}((a, m, pp)))
    @test card.data == [
        AtomicSpecies("Al", 24590.7655930491, "Al.pbe-n-kjpaw_psl.1.0.0.UPF"),
        AtomicSpecies("As", 68285.4024548272, "As.pbe-n-kjpaw_psl.1.0.0.UPF"),
    ]
    push!(card.data, AtomicSpecies("Si", 25591.1924913552, "Si.pbe-n-kjpaw_psl.1.0.0.UPF"))
    @test card.data == [
        AtomicSpecies("Al", 24590.7655930491, "Al.pbe-n-kjpaw_psl.1.0.0.UPF"),
        AtomicSpecies("As", 68285.4024548272, "As.pbe-n-kjpaw_psl.1.0.0.UPF"),
        AtomicSpecies("Si", 25591.1924913552, "Si.pbe-n-kjpaw_psl.1.0.0.UPF"),
    ]
end

@testset "Constructing `AtomicPosition`" begin
    # Data from https://github.com/QEF/q-e/blob/7be27df/PW/examples/gatefield/run_example#L129-L132.
    x = AtomicPosition("S", [0.500000000, 0.288675130, 1.974192764])
    @test_throws AssertionError @set x.atom = "sulfur"
    @test_throws InexactError @set x.pos = [1im, 2im, 3im]
    @test_throws ErrorException x.posi = [0.1, 0.2, 0.3]  # Given a wrong field name
    @test x.if_pos == [1, 1, 1]
    @test x == AtomicPosition('S', [0.500000000, 0.288675130, 1.974192764])
    @test AtomicPosition(
        AtomicSpecies('S', 32.066, "S.pz-n-rrkjus_psl.0.1.UPF"),
        [0.500000000, 0.288675130, 1.974192764],
    ) == x
end

@testset "Test constructing `AtomicPositionsCard` from `StructArray`s" begin
    # Data from https://github.com/QEF/q-e/blob/7be27df/PW/examples/gatefield/run_example#L129-L132.
    a = ["S", "Mo", "S"]
    pos = [
        [0.500000000, 0.288675130, 1.974192764],
        [0.000000000, 0.577350270, 2.462038339],
        [0.000000000, -0.577350270, 2.950837559],
    ]
    card = AtomicPositionsCard(
        StructArray{AtomicPosition}((a, pos, [[1, 1, 1], [1, 1, 1], [1, 1, 1]])),
        "alat",
    )
    @test card.data == [
        AtomicPosition("S", [0.500000000, 0.288675130, 1.974192764]),
        AtomicPosition("Mo", [0.000000000, 0.577350270, 2.462038339]),
        AtomicPosition("S", [0.000000000, -0.577350270, 2.950837559]),
    ]
end

@testset "Constructing `CellParametersCard`" begin
    #Data from https://gitlab.com/QEF/q-e/blob/master/NEB/examples/neb1.in
    option = "bohr"
    data = [
        12 0 0
        0 5 0
        0 0 5
    ]
    card = CellParametersCard(data, option)
    @test_throws AssertionError @set card.option = "ala"
    @test_throws AssertionError @set card.option = "crystal" # Allowed options are alat, angstrom, bohr
    @test_throws DimensionMismatch @set card.data = [ # Matrix size should be (3, 3)
        1 2
        3 4
    ]
    @test_throws DimensionMismatch @set card.data = [
        1 2 3 4
        5 6 7 8
        4 3 2 1
        8 7 6 5
    ]
    @test CellParametersCard(data).option == "alat" # default option is alat
end

@testset "Constructing `AtomicForce`" begin
    x = AtomicForce("H", [1, 2, 3])
    @test_throws DimensionMismatch @set x.force = [1, 2]
    @test_throws DimensionMismatch @set x.force = [1, 2, 3, 4]
end

@testset "Test constructing `AtomicForce` from `StructArray`" begin
    a = ["H", "O", "H"]
    f = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
    card = AtomicForcesCard(StructArray{AtomicForce}((a, f)))
    @test card.data == [
        AtomicForce("H", [1, 2, 3]),
        AtomicForce("O", [4, 5, 6]),
        AtomicForce("H", [7, 8, 9]),
    ]
end

@testset "Test constructing a `PWInput`: silicon" begin
    # This example is from https://github.com/QEF/q-e/blob/master/PW/examples/example01/run_example.
    for diago in ("david", "cg", "ppcg")
        control = ControlNamelist(
            tstress = true,
            tprnfor = true,
            outdir = raw"$TMP_DIR/",
            prefix = "silicon",
            pseudo_dir = raw"$PSEUDO_DIR/",
        )
        system =
            SystemNamelist(ibrav = 2, celldm = [10.2], nat = 2, ntyp = 1, ecutwfc = 18.0)
        electrons = ElectronsNamelist(conv_thr = 1.0e-8, diagonalization = "$diago")
        atomic_species = AtomicSpeciesCard([AtomicSpecies("Si", 28.086, "Si.pz-vbc.UPF")])
        atomic_positions = AtomicPositionsCard([
            AtomicPosition("Si", [0.0, 0.0, 0.0]),
            AtomicPosition("Si", [0.25, 0.25, 0.25]),
        ])
        k_points = SpecialPointsCard(
            [
                0.125 0.125 0.125 1.0
                0.125 0.125 0.375 3.0
                0.125 0.125 0.625 3.0
                0.125 0.125 0.875 3.0
                0.125 0.375 0.375 3.0
                0.125 0.375 0.625 6.0
                0.125 0.375 0.875 6.0
                0.125 0.625 0.625 3.0
                0.375 0.375 0.375 1.0
                0.375 0.375 0.625 3.0
            ],
        )
        object = PWInput(
            control = control,
            system = system,
            electrons = electrons,
            atomic_species = atomic_species,
            atomic_positions = atomic_positions,
            k_points = k_points,
            cell_parameters = nothing,
        )
    end
end

@testset "Test constructing a `PWInput`: silicon bands" begin
    # This example is from https://github.com/QEF/q-e/blob/master/PW/examples/example01/run_example.
    for diago in ("david", "cg", "ppcg")
        control = ControlNamelist(
            calculation = "bands",
            pseudo_dir = raw"$PSEUDO_DIR/",
            outdir = raw"$TMP_DIR/",
            prefix = "silicon",
        )
        system = SystemNamelist(
            ibrav = 2,
            celldm = [10.2],
            nat = 2,
            ntyp = 1,
            ecutwfc = 18.0,
            nbnd = 8,
        )
        electrons = ElectronsNamelist(diagonalization = "$diago")
        atomic_species = AtomicSpeciesCard([AtomicSpecies("Si", 28.086, "Si.pz-vbc.UPF")])
        atomic_positions = AtomicPositionsCard([
            AtomicPosition("Si", [0.0, 0.0, 0.0]),
            AtomicPosition("Si", [0.25, 0.25, 0.25]),
        ])
        k_points = SpecialPointsCard(
            [
                0.0 0.0 0.0 1.0
                0.0 0.0 0.1 1.0
                0.0 0.0 0.2 1.0
                0.0 0.0 0.3 1.0
                0.0 0.0 0.4 1.0
                0.0 0.0 0.5 1.0
                0.0 0.0 0.6 1.0
                0.0 0.0 0.7 1.0
                0.0 0.0 0.8 1.0
                0.0 0.0 0.9 1.0
                0.0 0.0 1.0 1.0
                0.0 0.0 0.0 1.0
                0.0 0.1 0.1 1.0
                0.0 0.2 0.2 1.0
                0.0 0.3 0.3 1.0
                0.0 0.4 0.4 1.0
                0.0 0.5 0.5 1.0
                0.0 0.6 0.6 1.0
                0.0 0.7 0.7 1.0
                0.0 0.8 0.8 1.0
                0.0 0.9 0.9 1.0
                0.0 1.0 1.0 1.0
                0.0 0.0 0.0 1.0
                0.1 0.1 0.1 1.0
                0.2 0.2 0.2 1.0
                0.3 0.3 0.3 1.0
                0.4 0.4 0.4 1.0
                0.5 0.5 0.5 1.0
            ],
        )
        object = PWInput(
            control = control,
            system = system,
            electrons = electrons,
            atomic_species = atomic_species,
            atomic_positions = atomic_positions,
            k_points = k_points,
            cell_parameters = nothing,
        )
    end
end

@testset "Test constructing a `PWInput`: aluminium" begin
    # This example is from https://github.com/QEF/q-e/blob/master/PW/examples/example01/run_example.
    for diago in ("david", "cg", "ppcg")
        control = ControlNamelist(
            calculation = "scf",
            restart_mode = "from_scratch",
            pseudo_dir = raw"$PSEUDO_DIR/",
            outdir = raw"$TMP_DIR/",
            prefix = "al",
            tprnfor = true,
            tstress = true,
        )
        system = SystemNamelist(
            ibrav = 2,
            celldm = [7.50],
            nat = 1,
            ntyp = 1,
            ecutwfc = 15.0,
            occupations = "smearing",
            smearing = "marzari-vanderbilt",
            degauss = 0.05,
        )
        electrons = ElectronsNamelist(diagonalization = "$diago", mixing_beta = 0.7)
        atomic_species = AtomicSpeciesCard([AtomicSpecies("Al", 26.98, "Al.pz-vbc.UPF")])
        atomic_positions = AtomicPositionsCard([AtomicPosition("Al", [0.0, 0.0, 0.0])])
        k_points = SpecialPointsCard(
            [
                0.0625000 0.0625000 0.0625000 1.00
                0.0625000 0.0625000 0.1875000 3.00
                0.0625000 0.0625000 0.3125000 3.00
                0.0625000 0.0625000 0.4375000 3.00
                0.0625000 0.0625000 0.5625000 3.00
                0.0625000 0.0625000 0.6875000 3.00
                0.0625000 0.0625000 0.8125000 3.00
                0.0625000 0.0625000 0.9375000 3.00
                0.0625000 0.1875000 0.1875000 3.00
                0.0625000 0.1875000 0.3125000 6.00
                0.0625000 0.1875000 0.4375000 6.00
                0.0625000 0.1875000 0.5625000 6.00
                0.0625000 0.1875000 0.6875000 6.00
                0.0625000 0.1875000 0.8125000 6.00
                0.0625000 0.1875000 0.9375000 6.00
                0.0625000 0.3125000 0.3125000 3.00
                0.0625000 0.3125000 0.4375000 6.00
                0.0625000 0.3125000 0.5625000 6.00
                0.0625000 0.3125000 0.6875000 6.00
                0.0625000 0.3125000 0.8125000 6.00
                0.0625000 0.3125000 0.9375000 6.00
                0.0625000 0.4375000 0.4375000 3.00
                0.0625000 0.4375000 0.5625000 6.00
                0.0625000 0.4375000 0.6875000 6.00
                0.0625000 0.4375000 0.8125000 6.00
                0.0625000 0.4375000 0.9375000 6.00
                0.0625000 0.5625000 0.5625000 3.00
                0.0625000 0.5625000 0.6875000 6.00
                0.0625000 0.5625000 0.8125000 6.00
                0.0625000 0.6875000 0.6875000 3.00
                0.0625000 0.6875000 0.8125000 6.00
                0.0625000 0.8125000 0.8125000 3.00
                0.1875000 0.1875000 0.1875000 1.00
                0.1875000 0.1875000 0.3125000 3.00
                0.1875000 0.1875000 0.4375000 3.00
                0.1875000 0.1875000 0.5625000 3.00
                0.1875000 0.1875000 0.6875000 3.00
                0.1875000 0.1875000 0.8125000 3.00
                0.1875000 0.3125000 0.3125000 3.00
                0.1875000 0.3125000 0.4375000 6.00
                0.1875000 0.3125000 0.5625000 6.00
                0.1875000 0.3125000 0.6875000 6.00
                0.1875000 0.3125000 0.8125000 6.00
                0.1875000 0.4375000 0.4375000 3.00
                0.1875000 0.4375000 0.5625000 6.00
                0.1875000 0.4375000 0.6875000 6.00
                0.1875000 0.4375000 0.8125000 6.00
                0.1875000 0.5625000 0.5625000 3.00
                0.1875000 0.5625000 0.6875000 6.00
                0.1875000 0.6875000 0.6875000 3.00
                0.3125000 0.3125000 0.3125000 1.00
                0.3125000 0.3125000 0.4375000 3.00
                0.3125000 0.3125000 0.5625000 3.00
                0.3125000 0.3125000 0.6875000 3.00
                0.3125000 0.4375000 0.4375000 3.00
                0.3125000 0.4375000 0.5625000 6.00
                0.3125000 0.4375000 0.6875000 6.00
                0.3125000 0.5625000 0.5625000 3.00
                0.4375000 0.4375000 0.4375000 1.00
                0.4375000 0.4375000 0.5625000 3.00
            ],
        )
        object = PWInput(
            control = control,
            system = system,
            electrons = electrons,
            atomic_species = atomic_species,
            atomic_positions = atomic_positions,
            k_points = k_points,
            cell_parameters = nothing,
        )
    end
end

end


const UniFinArr = UnivariateFiniteArray

Base.size(u::UniFinArr, args...) =
    size(first(values(u.prob_given_ref)), args...)

function Base.getindex(u::UniFinArr{<:Any,<:Any,R,P,N},
                       i::Integer...) where {R,P,N}
    prob_given_ref = LittleDict{R,P}()
    for ref in keys(u.prob_given_ref)
        prob_given_ref[ref] = getindex(u.prob_given_ref[ref], i...)
    end
    return UnivariateFinite(u.scitype, u.decoder, prob_given_ref)
end

function Base.getindex(u::UniFinArr{<:Any,<:Any,R,P,N},
                       I...) where {R,P,N}
    prob_given_ref = LittleDict{R,Array{P,N}}()
    for ref in keys(u.prob_given_ref)
        prob_given_ref[ref] = getindex(u.prob_given_ref[ref], I...)
    end
    return UnivariateFiniteArray(u.scitype, u.decoder, prob_given_ref)
end

function Base.setindex!(u::UniFinArr{S,V,R,P,N},
                        v::UnivariateFinite{S,V,R,P},
                        i::Integer...) where {S,V,R,P,N}
    for ref in keys(u.prob_given_ref)
       setindex!(u.prob_given_ref[ref], v.prob_given_ref[ref], i...)
    end
    return u
end

_err_incompatible_levels() = throw(DomainError(
    "Cannot concatenate `UnivariateFiniteArray`s with "*
    "different categorical levels (classes), "*
    "or whose levels, when ordered, are not  "*
    "consistently ordered. "))

# terminology:

# "classes"  - full pool of `CategoricalElement`s, even "unseen" ones (those
#             missing from support)
# "levels"   - same thing but in raw form (eg, `Symbol`s) aka "labels"
# "suppport" - those classes with a corresponding probability (the ones
#              named at time of construction of the `UnivariateFiniteArray`)

function Base.cat(us::UnivariateFiniteArray{S,V,R,P,N}...;
                      dims::Integer) where {S,V,R,P,N}

    isempty(us) && return []

    # build combined raw_support and check compatibility of levels:
    u1 = first(us)
    ordered = isordered(classes(u1))
    support_with_duplicates = Dist.support(u1)
    _classes = classes(u1)
    for i in 2:length(us)
        isordered(us[i]) == ordered || _err_incompatible_levels()
        if ordered
            classes(us[i]) ==
                _classes|| _err_incompatible_levels()
        else
            Set(classes(us[i])) ==
                Set(_classes) || _err_incompatible_levels()
        end
        support_with_duplicates =
            vcat(support_with_duplicates, Dist.support(us[i]))
    end
    _support = unique(support_with_duplicates) # no-longer categorical!

    # build the combined `prob_given_class` dictionary:
    pairs = (class => cat((pdf.(u, class) for u in us)..., dims=dims)
             for class in _support)
    prob_given_class = Dict(pairs)

    return UnivariateFinite(prob_given_class, pool=_classes)
end

Base.vcat(us::UnivariateFiniteArray...) = cat(us..., dims=1)
Base.hcat(us::UnivariateFiniteArray...) = cat(us..., dims=2)

## CONVENIENCE METHOD pdf(array_of_univariate_finite, labels)

# next bit is not specific to `UnivariateFiniteArray` but is for any
# abstract array with eltype `UnivariateFinite`.

# this is type piracy that has been adopted only after much
# agonizing over alternatives. Note that pdf.(u, labels) must
# necessarily have a different meaning (and only makes sense if u and
# labels have the same length or labels is a scalar)

function Distributions.pdf(
    u::AbstractArray{UnivariateFinite{S,V,R,P},N},
    C::AbstractVector{<:Union{V, CategoricalValue{V,R}}}) where {S,V,R,P,N}

    ret = Array{P,N+1}(undef, size(u)..., length(C))
    for i in eachindex(C)
        ret[fill(:,N)...,i] = broadcast(pdf, u, C[i])
    end
    return ret
end

## PERFORMANT BROADCASTING OF pdf

# u - a UnivariateFiniteArray
# cv - a CategoricalValue
# v - a vector of CategoricalArrays

# pdf.(u, cv)
function Base.Broadcast.broadcasted(
    ::typeof(pdf),
    u::UniFinArr{S,V,R,P,N},
    cv::CategoricalValue) where {S,V,R,P,N}

    cv in classes(u) || _err_missing_class(cv)

    return get(u.prob_given_ref, int(cv), zeros(P, size(u)))
end

# pdf.(u, v)
function Base.Broadcast.broadcasted(
    ::typeof(pdf),
    u::UniFinArr{S,V,R,P,N},
    v::AbstractArray{<:CategoricalValue{V,R},N}) where {S,V,R,P,N}

    length(u) == length(v) ||throw(DimensionMismatch(
        "Arrays could not be broadcast to a common size; "*
        "got a dimension with lengths $(length(u)) and $(length(v))"))
    for cv in v
        cv in classes(u) || _err_missing_class(c)
    end

    # will use linear indexing:
    v_flat = ((v[i], i) for i in 1:length(v))

    getter((cv, i), dtype) = _getindex(get(u.prob_given_ref, int(cv), nothing), i, dtype)

    ret_flat = getter.(v_flat, P)
    return reshape(ret_flat, size(u))
end

#dummy function
# returns `x[i]` for `Array` inputs `x`
# For non-Array inputs returns `zero(dtype)`
#This avoids using an if statement
_getindex(x::Array,i, dtype)=x[i]
_getindex(::Nothing, i, dtype) = zero(dtype)

# pdf.(u, raw) where raw is scalar or vec
function Base.Broadcast.broadcasted(
    ::typeof(pdf),
    u::UnivariateFiniteArray{S,V,R,P,N},
    raw::Union{V,AbstractArray{V,N}}) where {S,V,R,P,N}

    cat = transform(classes(u), raw)
    return broadcast(pdf, u, cat)
end

## PERFORMANT BROADCASTING OF mode:

function Base.Broadcast.broadcasted(::typeof(mode),
                                    u::UniFinArr{S,V,R,P,N}) where {S,V,R,P,N}
    dic = u.prob_given_ref

    # using linear indexing:
    mode_flat = map(1:length(u)) do i
        max_prob = maximum(dic[ref][i] for ref in keys(dic))
        m = zero(R)
        for ref in keys(dic)
            if dic[ref][i] == max_prob
                m = ref
                break
            end
        end
        return u.decoder(m)
    end
    return reshape(mode_flat, size(u))
end


const cufftNumber = Union{cufftDoubleReal,cufftReal,cufftDoubleComplex,cufftComplex}
const cufftReals = Union{cufftDoubleReal,cufftReal}
const cufftComplexes = Union{cufftDoubleComplex,cufftComplex}
const cufftDouble = Union{cufftDoubleReal,cufftDoubleComplex}
const cufftSingle = Union{cufftReal,cufftComplex}
const cufftTypeDouble = Union{Type{cufftDoubleReal},Type{cufftDoubleComplex}}
const cufftTypeSingle = Union{Type{cufftReal},Type{cufftComplex}}

cufftfloat(x) = _cufftfloat(float(x))
_cufftfloat(::Type{T}) where {T<:cufftReals} = T
_cufftfloat(::Type{Float16}) = Float32
_cufftfloat(::Type{Complex{T}}) where {T} = Complex{_cufftfloat(T)}
_cufftfloat(::Type{T}) where {T} = error("type $T not supported")
_cufftfloat(x::T) where {T} = _cufftfloat(T)(x)

complexfloat(x::DenseCuArray{Complex{<:cufftReals}}) = x
realfloat(x::DenseCuArray{<:cufftReals}) = x

complexfloat(x::DenseCuArray{T}) where {T<:Complex} = copy1(typeof(cufftfloat(zero(T))), x)
complexfloat(x::DenseCuArray{T}) where {T<:Real} = copy1(typeof(complex(cufftfloat(zero(T)))), x)

realfloat(x::DenseCuArray{T}) where {T<:Real} = copy1(typeof(cufftfloat(zero(T))), x)

function copy1(::Type{T}, x) where T
    y = CuArray{T}(undef, map(length, axes(x)))
    #copy!(y, x)
    y .= broadcast(xi->convert(T,xi),x)
end


function RandomPercentile()
    return RandomPercentile(())
end

function RandomPercentile(arg0::RandomGenerator)
    return RandomPercentile((RandomGenerator,), arg0)
end

function RandomPercentile(arg0::RandomPercentile)
    return RandomPercentile((RandomPercentile,), arg0)
end

function RandomPercentile(arg0::jdouble)
    return RandomPercentile((jdouble,), arg0)
end

function RandomPercentile(arg0::jdouble, arg1::RandomGenerator)
    return RandomPercentile((jdouble, RandomGenerator), arg0, arg1)
end

function aggregate(obj::RandomPercentile, arg0::RandomPercentile)
    return jcall(obj, "aggregate", void, (RandomPercentile,), arg0)
end

function clear(obj::RandomPercentile)
    return jcall(obj, "clear", void, ())
end

function copy(obj::RandomPercentile)
    return jcall(obj, "copy", RandomPercentile, ())
end

function evaluate(obj::RandomPercentile, arg0::Vector{jdouble}, arg1::jint, arg2::jint)
    return jcall(obj, "evaluate", jdouble, (Vector{jdouble}, jint, jint), arg0, arg1, arg2)
end

function evaluate(obj::RandomPercentile, arg0::jdouble, arg1::Vector{jdouble})
    return jcall(obj, "evaluate", jdouble, (jdouble, Vector{jdouble}), arg0, arg1)
end

function evaluate(obj::RandomPercentile, arg0::jdouble, arg1::Vector{jdouble}, arg2::jint, arg3::jint)
    return jcall(obj, "evaluate", jdouble, (jdouble, Vector{jdouble}, jint, jint), arg0, arg1, arg2, arg3)
end

function get_aggregate_n(obj::RandomPercentile, arg0::Collection)
    return jcall(obj, "getAggregateN", jdouble, (Collection,), arg0)
end

function get_aggregate_quantile_rank(obj::RandomPercentile, arg0::jdouble, arg1::Collection)
    return jcall(obj, "getAggregateQuantileRank", jdouble, (jdouble, Collection), arg0, arg1)
end

function get_aggregate_rank(obj::RandomPercentile, arg0::jdouble, arg1::Collection)
    return jcall(obj, "getAggregateRank", jdouble, (jdouble, Collection), arg0, arg1)
end

function get_n(obj::RandomPercentile)
    return jcall(obj, "getN", jlong, ())
end

function get_quantile_rank(obj::RandomPercentile, arg0::jdouble)
    return jcall(obj, "getQuantileRank", jdouble, (jdouble,), arg0)
end

function get_rank(obj::RandomPercentile, arg0::jdouble)
    return jcall(obj, "getRank", jdouble, (jdouble,), arg0)
end

function get_result(obj::RandomPercentile)
    return jcall(obj, "getResult", jdouble, ())
end

function get_result(obj::RandomPercentile, arg0::jdouble)
    return jcall(obj, "getResult", jdouble, (jdouble,), arg0)
end

function increment(obj::RandomPercentile, arg0::jdouble)
    return jcall(obj, "increment", void, (jdouble,), arg0)
end

function max_values_retained(::Type{RandomPercentile}, arg0::jdouble)
    return jcall(RandomPercentile, "maxValuesRetained", jlong, (jdouble,), arg0)
end

function reduce(obj::RandomPercentile, arg0::jdouble, arg1::Collection)
    return jcall(obj, "reduce", jdouble, (jdouble, Collection), arg0, arg1)
end


# Populate a SeisData structure with random channels
"""
    randSeisData()

Generate 8 to 24 channels of random seismic data as a SeisData object.

    randSeisData(N)

Generate `N` channels of random seismic data as a SeisData object.

### Keywords
| KW      | Default | Type      | Meaning                                 |
|-------- |:--------|:----------|:-------------------------------         |
| s       | 0.6     | Float64   | chance each channel is seismic data     |
| c       | 0.2     | Float64   | chance each channel is irregular (fs=0) |
| nx      | 0       | Int64     | number of samples in each channel       |
| fs_min  | 0.0     | Float64   | regular channels have fs ≥ fs_min       |
| a0      | false   | Bool      | recalculate :resp constants?            |

### Behavior
* If `nx` ≤ 0, the number of samples is determined randomly for each channel
* `100*s` is the minimum percentage of channels with guaranteed seismic data
* `100*c` is the maximum percentage of channels with irregularly-sampled data
* `s` takes precedence over `c`; they are not renormalized. Thus, for example:
  + `randSeisData(12, c=1.0, s=1.0)` works like `randSeisData(12, c=0.0, s=1.0)`
  + `randSeisData(c=1.0)` works like `randSeisData(s=0.0)`
* For `N` channels, `s*N` is rounded up; `c*N` is rounded down. Thus:
  + `randSeisData(10, c=0.28)` and `randSeisData(10, c=0.2)` are equivalent
  + `randSeisData(10, s=0.28)` and `randSeisData(10, s=0.3)` are equivalent
* By default, `:resp.a0 = 1.0` for all channels; `a0=true` calls `resp_a0!(S)`
  + `:resp.p` and `:resp.z` are random; `a0=true` isn't guaranteed to work

See also: `randSeisChannel`, `resp_a0!`
"""
function randSeisData(N::Int64; c::Float64=0.2, s::Float64=0.6, nx::Integer=0, fs_min::Float64=0.0, a0::Bool=false)
  S = SeisData(N)

  # Evaluate probabilities, with s taking precedence
  if s != 0.6
    c = max(0.0, 1.0-s)
  end
  if c != 0.2
    s = max(0.0, 1.0-c)
  end

  # determine number of channels of each
  n_seis = max(min(ceil(Int, s*S.n), S.n), 0)
  n_irr = max(min(floor(Int, c*S.n), S.n-n_seis), 0)

  # fill in and shuffle data_spec
  data_spec = zeros(UInt8, S.n)
  data_spec[1:n_seis] .= 0x01
  data_spec[n_seis+1:n_seis+n_irr] .= 0x02
  shuffle!(data_spec)

  # populate all channels according to data_spec
  for i = 1:S.n
    if data_spec[i] == 0x01
      S[i] = randSeisChannel(s=true, nx=nx, fs_min=fs_min)
    elseif data_spec[i] == 0x02
      S[i] = randSeisChannel(c=true, nx=nx, fs_min=fs_min)
    else
      S[i] = randSeisChannel(nx=nx, fs_min=fs_min)
    end
  end
  a0 && resp_a0!(S)
  return S
end
randSeisData(; c::Float64=0.2, s::Float64=0.6, nx::Int64=0, fs_min::Float64=0.0, a0::Bool=false) = randSeisData(rand(8:24), c=c, s=s, nx=nx, fs_min=fs_min, a0=a0)


using RomeoApp
using Test

@testset "ROMEO function tests" begin

niread = RomeoApp.niread
savenii = RomeoApp.savenii

p = joinpath("data", "small")
phasefile_me = joinpath(p, "Phase.nii")
phasefile_me_nan = joinpath(p, "phase_with_nan.nii")
magfile_me = joinpath(p, "Mag.nii")
tmpdir = mktempdir()
phasefile_1eco = joinpath(tmpdir, "Phase.nii")
magfile_1eco = joinpath(tmpdir, "Mag.nii")
phasefile_1arreco = joinpath(tmpdir, "Phase.nii")
magfile_1arreco = joinpath(tmpdir, "Mag.nii")
maskfile = joinpath(tmpdir, "Mask.nii")
savenii(niread(magfile_me)[:,:,:,1] |> I -> I .> RomeoApp.MriResearchTools.median(I), maskfile)
savenii(niread(phasefile_me)[:,:,:,1], phasefile_1eco)
savenii(niread(magfile_me)[:,:,:,1], magfile_1eco)
savenii(niread(phasefile_me)[:,:,:,[1]], phasefile_1arreco)
savenii(niread(magfile_me)[:,:,:,[1]], magfile_1arreco)

phasefile_me_5D = joinpath(tmpdir, "phase_multi_channel.nii")
magfile_5D = joinpath(tmpdir, "mag_multi_channel.nii")
savenii(repeat(niread(phasefile_me),1,1,1,1,2), phasefile_me_5D)
savenii(repeat(niread(magfile_me),1,1,1,1,2), magfile_5D)

function test_romeo(args)
    folder = tempname()
    args = [args..., "-o", folder]
    try
        msg = unwrapping_main(args)
        @test msg == 0
        @test isfile(joinpath(folder, "unwrapped.nii"))
    catch e
        println(args)
        println(sprint(showerror, e, catch_backtrace()))
        @test "test failed" == "with error" # signal a failed test
    end
end

configurations_se(pf, mf) = vcat(configurations_se.([[pf], [pf, "-m", mf]])...)
configurations_se(pm) = [
    [pm...],
    [pm..., "-g"],
    [pm..., "-N"],
    [pm..., "-i"],
    [pm..., "-q"],
    [pm..., "-Q"],
    [pm..., "-m", magfile_me, "-u"],
    [pm..., "-w", "romeo"],
    [pm..., "-w", "bestpath"],
    [pm..., "-w", "1010"],
    [pm..., "-w", "101011"],
    [pm..., "--threshold", "4"],
    [pm..., "-s", "50"],
    [pm..., "-s", "50", "--merge-regions"],
    [pm..., "-s", "50", "--merge-regions", "--correct-regions"],
    [pm..., "--wrap-addition", "0.1"],
    [pm..., "-k", "robustmask"],
    [pm..., "-k", "nomask"],
    [pm..., "-k", "qualitymask"],
    [pm..., "-k", "qualitymask", "0.1"],
    [pm..., "-k", maskfile],
]
configurations_me(phasefile_me, magfile_me) = vcat(configurations_me.([[phasefile_me], [phasefile_me, "-m", magfile_me]])...)
configurations_me(pm) = [
    [pm..., "-e", "1:2", "-t", "[2,4]"], # giving two echo times for two echoes used out of three
    [pm..., "-e", "[1,3]", "-t", "[2,4,6]"], # giving three echo times for two echoes used out of three
    [pm..., "-e", "[1", "3]", "-t", "[2,4,6]"],
    [pm..., "-t", "[2,4,6]"],
    [pm..., "-t", "2:2:6"],
    [pm..., "-t", "[2.1,4.2,6.3]"],
    [pm..., "-t", "epi"], # shorthand for "ones(<num-echoes>)"
    [pm..., "-t", "epi", "5.3"], # shorthand for "5.3*ones(<num-echoes>)"
    [pm..., "-B", "-t", "[2,4,6]"],
    [pm..., "-B", "-t", "[2" ,"4", "6]"], # when written like [2 4 6] in command line
    [pm..., "--temporal-uncertain-unwrapping", "-t", "[2,4,6]"],
    [pm..., "--template", "1", "-t", "[2,4,6]"],
    [pm..., "--template", "3", "-t", "[2,4,6]"],
    [pm..., "--phase-offset-correction", "-t", "[2,4,6]"],
    [pm..., "--phase-offset-correction", "bipolar", "-t", "[2,4,6]"],
    [pm..., "--phase-offset-correction", "-t", "[2,4,6]", "--phase-offset-smoothing-sigma-mm", "[5,8,4]"],
    [pm..., "--phase-offset-correction", "-t", "[2,4,6]", "--write-phase-offsets"],
]

files = [(phasefile_1eco, magfile_1eco), (phasefile_1arreco, magfile_1arreco), (phasefile_1eco, magfile_1arreco), (phasefile_1arreco, magfile_1eco)]
for (pf, mf) in files, args in configurations_se(pf, mf)
    test_romeo(args)
end
for args in configurations_me(phasefile_me, magfile_me)
    test_romeo(args)
end
for args in configurations_se(["-p", phasefile_me, "-m", magfile_me, "-t", "[2,4,6]"])
    test_romeo(args)
end
for args in configurations_me(phasefile_me_5D, magfile_5D)[end-2:end] # test the last 3 configurations_me lines for coil combination
    test_romeo(args)
end
files_se = [(phasefile_1eco, magfile_1eco), (phasefile_1arreco, magfile_1arreco)]
for (pf, mf) in files_se
    b_args = ["-B", "-t", "3.06"]
    test_romeo(["-p", pf, b_args...])
    test_romeo(["-p", pf, "-m", mf, b_args...])
end

test_romeo([phasefile_me_nan, "-t", "[2,4]", "-k", "nomask"])

## Test error and warning messages
m = "multi-echo data is used, but no echo times are given. Please specify the echo times using the -t option."
@test_throws ErrorException(m) unwrapping_main(["-p", phasefile_me, "-o", tmpdir, "-v"])

m = "masking option '0.8' is undefined (Maybe '-k qualitymask 0.8' was meant?)"
@test_throws ErrorException(m) unwrapping_main(["-p", phasefile_1eco, "-o", tmpdir, "-v", "-k", "0.8"])

m = "masking option 'blub' is undefined"
@test_throws ErrorException(m) unwrapping_main(["-p", phasefile_1eco, "-o", tmpdir, "-v", "-k", "blub"])

m = "Phase offset determination requires all echo times!"
@test_throws ErrorException(m) unwrapping_main(["-p", phasefile_me_5D, "-o", tmpdir, "-v", "-t", "[1,2]", "-e", "[1,2]", "--phase-offset-correction"])

m = "echoes=[1,5]: specified echo out of range! Number of echoes is 3"
@test_throws ErrorException(m) unwrapping_main(["-p", phasefile_me, "-o", tmpdir, "-v", "-t", "[1,2,3]", "-e", "[1,5]"])

m = "echoes=[1,5} wrongly formatted!"
@test_throws ErrorException(m) unwrapping_main(["-p", phasefile_me, "-o", tmpdir, "-v", "-t", "[1,2,3]", "-e", "[1,5}"])

m = "Number of chosen echoes is 2 (3 in .nii data), but 5 TEs were specified!"
@test_throws ErrorException(m) unwrapping_main(["-p", phasefile_me, "-o", tmpdir, "-v", "-t", "[1,2,3,4,5]", "-e", "[1,2]"])

m = "size of magnitude and phase does not match!"
@test_throws ErrorException(m) unwrapping_main(["-p", phasefile_me, "-o", tmpdir, "-v", "-t", "[1,2,3]", "-m", magfile_1eco])

m = "robustmask was chosen but no magnitude is available. No mask is used!"
@test_logs (:warn, m) match_mode=:any unwrapping_main(["-p", phasefile_1eco, "-o", tmpdir])

@test_logs unwrapping_main(["-p", phasefile_1eco, "-o", tmpdir, "-m", magfile_1eco]) # test that no warning appears

## test no-rescale
readphase = RomeoApp.readphase
phasefile_me_uw = joinpath(tempname(), "unwrapped.nii")
phasefile_me_uw_wrong = joinpath(tempname(), "wrong_unwrapped.nii")
phasefile_me_uw_again = joinpath(tempname(), "again_unwrapped.nii")
unwrapping_main([phasefile_me, "-o", phasefile_me_uw, "-t", "[2,4,6]"])
unwrapping_main([phasefile_me_uw, "-o", phasefile_me_uw_wrong, "-t", "[2,4,6]"])
unwrapping_main([phasefile_me_uw, "-o", phasefile_me_uw_again, "-t", "[2,4,6]", "--no-rescale"])

@test readphase(phasefile_me_uw_again; rescale=false).raw == readphase(phasefile_me_uw; rescale=false).raw
@test readphase(phasefile_me_uw_wrong; rescale=false).raw != readphase(phasefile_me_uw; rescale=false).raw

## test mcpc3ds output files
testpath = joinpath(tmpdir, "test5d")
unwrapping_main([phasefile_me_5D, "-o", testpath, "-m", magfile_5D, "-t", "[2,4,6]"])
@test isfile(joinpath(testpath, "combined_mag.nii"))
@test isfile(joinpath(testpath, "combined_phase.nii"))

testpath = joinpath(tmpdir, "test4d")
unwrapping_main([phasefile_me, "-o", testpath, "-m", magfile_me, "-t", "[2,4,6]", "--phase-offset-correction"])
@test isfile(joinpath(testpath, "corrected_phase.nii"))

## test quality map
unwrapping_main([phasefile_me, "-m", magfile_me, "-o", tmpdir, "-t", "[2,4,6]", "-qQ"])
fns = joinpath.(tmpdir, ["quality.nii", ("quality_$i.nii" for i in 1:4)...])
for i in 1:length(fns), j in i+1:length(fns)
    @test niread(fns[i]).raw != niread(fns[j]).raw
end

end

## print version to verify
println()
unwrapping_main(["--version"])


module Voronoi
    export Point, Segment, Vertex, Edge, get_edges

    @static if is_windows()
        const library_path =  joinpath(dirname(@__FILE__), "voronoi.dll")
    end

    @static if is_unix()
        const library_path = joinpath(dirname(@__FILE__), "voronoi.so")
    end

    type Point
        X::Int32
        Y::Int32
    end

    type Segment
      p0::Point
      p1::Point
    end

    type Vertex
        X::Float64
        Y::Float64
    end

    type Edge
        start_index::UInt64
        end_index::UInt64

        is_primary::Bool

        site1::UInt64
        site2::UInt64
    end

    function Base.show(io::IO, point::Point)
      print(io, "[$(point.X),$(point.Y)]")
    end

    function Base.show(io::IO, vertex::Vertex)
      print(io, "[$(vertex.X),$(vertex.Y)]")
    end

    function Base.show(io::IO, segment::Segment)
      print(io, "($(segment.p0),$(segment.p1))")
    end

    function appendVertex!(outputArray::Ptr{Void}, vertex::Vertex)
        ourArray = unsafe_pointer_to_objref(outputArray)::Vector{Vertex}

        push!(ourArray, vertex)
    end

    function appendEdge!(outputArray::Ptr{Void}, edge::Edge)
        ourArray = unsafe_pointer_to_objref(outputArray)::Vector{Edge}

        push!(ourArray, Edge(edge.start_index, edge.end_index, edge.is_primary,
                              edge.site1, edge.site2))
    end

    function get_edges(points::Vector{Point}, segments::Vector{Segment})
      vertices = Vector{Vertex}()
      edges = Vector{Edge}()

      ccall((:get_edges, library_path), Void, (Ptr{Ptr{Cint}}, Csize_t,
                                               Ptr{Ptr{Cint}}, Csize_t,
                                               Any, Ptr{Void},
                                               Any, Ptr{Void}),
            [[p.X, p.Y] for p in points],
            length(points),
            [[s.p0.X, s.p0.Y, s.p1.X, s.p1.Y] for s in segments],
            length(segments),
            vertices,
            cfunction(appendVertex!, Any, (Ptr{Void}, Vertex)),
            edges,
            cfunction(appendEdge!, Any, (Ptr{Void}, Edge)))

      vertices, edges
    end
end


using Statistics

function format(x)
    if abs(x) < 999
        return @sprintf("%12.7f", x)
    else
        return @sprintf("%12.5e", x)
    end
end

import Base.write
"""

```
write(R::ComplexMixtures.Result, filename::String, solute::Selection, solvent::Selection)
```

Function to write the final results to output files as simple tables
that are human-readable and easy to analyze with other software

If the solute and solvent selections are provides, pass on the atom
names.

"""
write(R::Result, filename::String, solute::Selection, solvent::Selection) =
    write(R, filename, solute_names = solute.names, solvent_names = solvent.names)

"""

```
write(R::ComplexMixtures.Result, filename::String;
      solute_names::Vector{String} = ["nothing"],
      solvent_names::Vector{String} = ["nothing"])
```

Optional passing of atom names.

"""
function write(
    R::Result,
    filename::String;
    solute_names::Vector{String} = ["nothing"],
    solvent_names::Vector{String} = ["nothing"],
)

    # Names of output files containing atomic contibutions
    atom_contrib_solvent =
        FileOperations.remove_extension(filename) *
        "-ATOM_CONTRIB_SOLVENT." *
        FileOperations.file_extension(filename)
    atom_contrib_solute =
        FileOperations.remove_extension(filename) *
        "-ATOM_CONTRIB_SOLUTE." *
        FileOperations.file_extension(filename)

    #
    # GMD computed with minimum distance
    #

    # Conversion factor for volumes (as KB integrals), from A^3 to cm^3/mol
    mole = 6.022140857e23
    convert = mole / 1.e24

    output = open(filename, "w")
    println(output, @sprintf("#"))
    println(output, @sprintf("# Output of ComplexMixtures - MDDF"))
    println(output, @sprintf("# Trajectory files and weights:"))
    for i = 1:length(R.files)
        println(output, "#  $(R.files[i]) - w = $(R.weights[i])")
    end
    println(output, @sprintf("#"))
    println(
        output,
        @sprintf(
            "# Density of solvent in simulation box (sites/A^3): %15.8f",
            R.density.solvent
        )
    )
    println(
        output,
        @sprintf(
            "# Density of solvent in bulk (estimated) (sites/A^3): %15.8f",
            R.density.solvent_bulk
        )
    )
    println(
        output,
        @sprintf(
            "# Molar volume of solvent in simulation (cc/mol): %15.8f",
            convert / R.density.solvent
        )
    )
    println(
        output,
        @sprintf(
            "# Molar volume of solvent in bulk (estimated) (cc/mol): %15.8f",
            convert / R.density.solvent_bulk
        )
    )
    println(output, @sprintf("#"))
    println(output, @sprintf("# Number of atoms solute: %i9", size(R.solute_atom, 2)))
    println(
        output,
        @sprintf("# Number of atoms of the solvent: %i9", size(R.solvent_atom, 2))
    )
    println(output, @sprintf("#"))

    if R.options.usecutoff
        ibulk = setbin(R.options.dbulk, R.options.binstep)
    else
        ibulk = round(Int, R.nbins - 1 / R.options.binstep)
    end
    bulkerror = Statistics.mean(R.mddf[ibulk:R.nbins])
    sdbulkerror = Statistics.std(R.mddf[ibulk:R.nbins])
    println(output, "#")
    println(
        output,
        @sprintf(
            "# Average and standard deviation of bulk-gmd: %12.5f +/- %12.5f",
            bulkerror,
            sdbulkerror
        )
    )
    println(output, "#")

    # Output table

    println(output, "# COLUMNS CORRESPOND TO: ")
    println(output, "#       1  Minimum distance to solute (dmin)")
    println(
        output,
        "#       2  GMD distribution (md count normalized by md count of random-solute distribution)",
    )
    println(
        output,
        "#       3  Kirwood-Buff integral (cc/mol) computed [(1/bulkdensity)*(col(6)-col(7))].",
    )
    println(output, "#       4  Minimum distance site count for each dmin.")
    println(
        output,
        "#       5  Minimum distance site count for each dmin for random solute distribution.",
    )
    println(
        output,
        "#       6  Cumulative number of molecules within dmin in the simulation.",
    )
    println(
        output,
        "#       7  Cumulative number of molecules within dmin for random solute distribution.",
    )
    println(output, "#       8  Volume of the shell of distance dmin and width binstep.")
    println(output, "#")
    println(
        output,
        "#   1-DISTANCE         2-GMD      3-KB INT    4-MD COUNT  5-COUNT RAND      6-SUM MD    7-SUM RAND   8-SHELL VOL",
    )

    for i = 1:R.nbins
        line = "  " * format(R.d[i])                                 #  1-DISTANCE
        line = line * "  " * format(R.mddf[i])                         #  2-GMD
        line = line * "  " * format(R.kb[i])                           #  3-KB INT
        line = line * "  " * format(R.md_count[i])                     #  4-MD COUNT
        line = line * "  " * format(R.md_count_random[i])              #  5-COUNT RAND
        line = line * "  " * format(R.sum_md_count[i])                 #  6-SUM MD
        line = line * "  " * format(R.sum_md_count_random[i])          #  7-SUM RAND
        line = line * "  " * format(R.volume.shell[i])                 #  8-SHELL VOL
        println(output, line)
    end
    close(output)

    # Writting gmd per atom contributions for the solvent

    output = open(atom_contrib_solvent, "w")
    println(output, "#")
    println(output, "# Solvent atomic contributions to total MDDF. ")
    println(output, "#")
    println(output, "# Trajectory files: ", R.files)
    println(output, "#")
    println(output, "# Atoms: ")
    for i = 1:size(R.solvent_atom, 2)
        if solvent_names[1] == "nothing"
            println(output, @sprintf("# %9i", i))
        else
            println(output, @sprintf("# %9i %5s", i, solvent_names[i]))
        end
    end
    println(output, "#")
    string = "#   DISTANCE     GMD TOTAL"
    for i = 1:size(R.solvent_atom, 2)
        string = string * @sprintf("  %12i", i)
    end
    println(output, string)
    for i = 1:R.nbins
        string = format(R.d[i])
        string = string * "  " * format(R.mddf[i])
        for j = 1:size(R.solvent_atom, 2)
            string = string * "  " * format(R.solvent_atom[i, j])
        end
        println(output, string)
    end
    close(output)

    # Writting gmd per atom contributions for the solute

    output = open(atom_contrib_solute, "w")
    println(output, "#")
    println(output, "# Solute atomic contributions to total MDDF. ")
    println(output, "#")
    println(output, "# Trajectory files: ", R.files)
    println(output, "#")
    println(output, "# Atoms:")
    for i = 1:size(R.solute_atom, 2)
        if solute_names[1] == "nothing"
            println(output, @sprintf("# %9i", i))
        else
            println(output, @sprintf("# %9i %5s", i, solute_names[i]))
        end
    end
    println(output, "#")
    string = "#   DISTANCE      GMD TOTAL"
    for i = 1:size(R.solute_atom, 2)
        string = string * @sprintf("  %12i", i)
    end
    println(output, string)
    for i = 1:R.nbins
        string = format(R.d[i]) * "  " * format(R.mddf[i])
        for j = 1:size(R.solute_atom, 2)
            string = string * "  " * format(R.solute_atom[i, j])
        end
        println(output, string)
    end
    close(output)

    # Write final messages with names of output files and their content

    println()
    println(" OUTPUT FILES: ")
    println()
    println(" Wrote solvent atomic GMD contributions to file: ", atom_contrib_solvent)
    println(" Wrote solute atomic GMD contributions to file: ", atom_contrib_solute)
    println()
    println(" Wrote main output file: ", filename)
    println()

    return nothing
end


@recipe function f(ta::T) where {T<:TimeArray}
    seriestype --> :path

    st = plotattributes[:seriestype]
    if st ∈ [:candlestick, :heikinashi]
        Candlestick(ta)
    else
        label --> reshape(string.(colnames(ta)), 1, :)
        timestamp(ta), values(ta)
    end
end

# FIXME: refine Candlestick as a subtype of AbstractTimeSeries
#        it requires the Base.view supports. (#419)
"""
    Candlestick(ta::TimeArray)

# Argument

- There are four required columns from `ta::TimeArray`: `open`, `high`, `low` and
  `close`. The column names is case-insensitive.

# Examples

```julia-repl
julia> using MarketData

julia> TimeSeries.Candlestick(ohlcv)
```
"""
mutable struct Candlestick{D <: TimeType}
    time::Vector{D}
    open::AbstractVector
    high::AbstractVector
    low::AbstractVector
    close::AbstractVector
end

Candlestick(ta::TimeArray) = Candlestick(extract_ohlc(ta)...)

function extract_ohlc(ta::TimeArray)
    C     = ["open", "high", "low", "close"]
    cols  = colnames(ta)
    cols′ = lowercase.(string.(colnames(ta)))
    V     = map(C) do x
        i = findfirst(isequal(x), cols′)
        i ≡ nothing && throw(ArgumentError("the TimeArray did not have column `$x`"))
        values(ta[cols[i]])
    end
    (timestamp(ta), V...)
end

function HeikinAshi!(cs::Candlestick) # some values here are made too high!
    cs.close[1] = (cs.open[1] + cs.low[1] + cs.close[1] + cs.high[1]) / 4
    cs.open[1] = (cs.open[1] + cs.close[1]) / 2
    cs.high[1] = cs.high[1]
    cs.low[1] = cs.low[1]

    for i in 2:length(cs.open)
        cs.close[i] = (cs.open[i] + cs.low[i] + cs.close[i] + cs.high[i]) / 4
        cs.open[i] = (cs.open[i-1] + cs.close[i-1]) / 2
        cs.high[i] = maximum([cs.high[i], cs.open[i], cs.close[i]])
        cs.low[i] = minimum([cs.low[i], cs.open[i], cs.close[i]])
    end
end

@recipe function f(cs::Candlestick)
    st = get(plotattributes, :seriestype, :candlestick)
    st == :heikinashi && HeikinAshi!(cs)

    legend    --> false
    linewidth --> 0.5
    grid      --> true
    xrotation --> 90

    bw = get(plotattributes, :bar_width, 0.9)

    # allow passing alternative colors as a vector
    cols = get(plotattributes, :seriescolor, [:green, :red])
    if !(cols isa Vector{Symbol}) || length(cols) != 2
        throw(ArgumentError(":seriescolor should be a Vector{Symbol} of two elements."))
    end

    xseg, yseg = Float64[], Float64[]
    xcenter, ycenter = Float64[], Float64[]

    idx  = 1:length(cs.time)
    tick = 1

    colors = Symbol[]
    hovers = String[]
    margin = (1 - bw) / 2 * tick
    for (i, o, h, l, c) ∈ zip(idx, cs.open, cs.high, cs.low, cs.close)
        x₁, x₂ = i - tick + margin, i - margin
        x̄ = i - tick / 2
        append!(xseg, [
            x̄,
            x̄,
            x₁,
            x₁,
            x̄,
            x̄,
            x̄,
            x₂,
            x₂,
            x̄,
            NaN
        ])

        y₁, y₂ = min(o, c), max(o, c)
        append!(yseg, [
            l,
            y₁,
            y₁,
            y₂,
            y₂,
            h,
            y₂,
            y₂,
            y₁,
            y₁,
            NaN
        ])

        push!(colors, ifelse(o ≤ c, cols[1], cols[2]))

        push!(xcenter, x̄)
        push!(ycenter, (o + c) / 2)
        push!(hovers, "Open: $o
High: $h
Low: $l
Close: $c")
    end

    xticks = get(plotattributes, :xticks, 1)
    if xticks isa Integer
        if xticks != 1
            xticks := (idx .- 0.5, map(x -> (x % xticks) != 1 ? "" : string(cs.time[x]) ,idx))
        else
            xticks := (idx .- 0.5, string.(cs.time))
        end
    end

    # tweak the hover function for Plotly backend
    extract_plot_kwargs = get(plotattributes, :extract_plot_kwargs, Dict{Symbol,Any}())
    get!(extract_plot_kwargs, :hovermode, :x)
    get!(extract_plot_kwargs, :hoverdistance, 5)

    @series begin
        seriestype  := :shape
        seriescolor := colors
        linecolor  --> colors
        xseg, yseg
    end

    @series begin
        seriestype        := :scatter
        seriescolor       := colors
        markersize        := 1
        markerstrokewidth := 0
        hover             --> hovers
        xcenter, ycenter
    end
end


type PlaceOccupied <: Exception end

type Position
     x::Int
     y::Int
end

abstract Zwierze

type NULL <: Zwierze end

type Drapieznik <: Zwierze
     name::String
     attack::Float16
     position::Position
end

type Ofiara <: Zwierze
     name::String
     defense::Float16
     position::Position
end

type Swiat
     width::Int
     height::Int
     mapa
end

function initWorld(world::Swiat)
	 for i = 1:n
	     for j = 1:n
	         oneWorld.mapa[i,j] = NULL()
	     end
	 end
end

function printWorld(world::Swiat)
	 println("Swiat: D - drapieznik, O - ofiara")
	 for i = 1:world.width
	     for j = 1:world.height
	         if (world.mapa[i,j] == NULL())
		     print("N")
		 elseif (typeof(world.mapa[i,j]) == Ofiara)
		     print("O")
		 elseif (typeof(world.mapa[i,j]) == Drapieznik)
		     print("D")
		 end
		 print(" ")
	     end
	     println()
	 end
end

function addAnimal(animal::Zwierze, world::Swiat)
	 if world.mapa[animal.position.x, animal.position.y] != NULL()
	     throw(PlaceOccupied)
	 else
	     world.mapa[animal.position.x, animal.position.y] = animal
	 end
end

function seek(animal::Zwierze, world::Swiat)
	 for i = 1:world.width
	     for j = 1:world.height
	     	 if world.mapa[i,j] == NULL()
		    return Position(i,j)
		 end
	     end
	 end

	 return Position(0,0)
end

function odleglosc{T1 <: Zwierze, T2 <: Zwierze}(a1::T1, a2::T2)
	 return abs(a1.position.x - a2.position.x) + abs(a1.position.y - a2.position.y)
end

function interakcja(o::Ofiara, d::Drapieznik, world::Swiat)
  local newPositionVictim = seek(o,world)
  local newPositionPredator = o.position

  world.mapa[d.position.x, d.position.y] = NULL()
  world.mapa[o.position.x, o.position.y] = d
  world.mapa[newPositionVictim.x, newPositionVictim.y] = o

  d.position = newPositionPredator
  o.position = newPositionVictim

end

function interakcja(d::Drapieznik, o::Ofiara, world::Swiat)
  local newPositionVictim = o.position

  world.mapa[d.position.x,d.position.y] = NULL()
  world.mapa[o.position.x,o.position.y] = d

  d.position = newPositionVictim
  o = NULL()
end

function interakcja(d1::Drapieznik, d2::Drapieznik, world::Swiat)
	 "Wrrrr"
end

function interakcja(o1::Ofiara, o2::Ofiara, world::Swiat)
	 "Beeee"
end

#run section

n = 3

oneWorld = Swiat(n,n,Array(Zwierze,n,n))

initWorld(oneWorld)

predator = Drapieznik("lew",1.32,Position(1,1))
predator2 = Drapieznik("nosorożec",1.22,Position(1,2))

victim = Ofiara("antylopa",0.59,Position(2,3))
victim2 = Ofiara("zebra",1.09,Position(3,3))

addAnimal(predator,oneWorld)
addAnimal(predator2,oneWorld)
#addAnimal(predator2,oneWorld)

addAnimal(victim,oneWorld)
addAnimal(victim2,oneWorld)

printWorld(oneWorld)

println("Odleglosc w metryce taksowkowej miedzy (1,1) a (1,2): ",odleglosc(predator,victim))

println(interakcja(predator,predator2,oneWorld))
println(interakcja(victim, victim2, oneWorld))

println("Interakcja ofiara-drapieznik")
interakcja(victim,predator,oneWorld)
printWorld(oneWorld)

println("Interakcja drapieznik-ofiara")
interakcja(predator,victim,oneWorld)
printWorld(oneWorld)


_isfulldummy(x::AbstractTerm) = false
function _isfulldummy(x::CategoricalTerm)
    return isa(x.contrasts, StatsModels.ContrastsMatrix{StatsModels.FullDummyCoding})
end

function _fulldummycheck(outer::InteractionTerm)
    all(_isfulldummy, outer.terms[1:end-1]) ||
        throw(ArgumentError("Outer interactions in a nesting must consist only " *
                            " of categorical terms with FullDummyCoding, got $outer"))
    return nothing
end

"""
    group / term

Generate predictors for `term` within each level of `group`.  Implemented as
`group + fulldummy(group) & term`.
"""
function Base.:(/)(outer::CategoricalTerm, inner::AbstractTerm)
    return outer + fulldummy(outer) & inner
end

function Base.:(/)(outer::CategoricalTerm, inner::TermTuple)
    fd = fulldummy(outer)
    return mapfoldl(x -> fd & x, +, inner; init=outer)
end

function Base.:(/)(outer::TermTuple, inner::Union{AbstractTerm, TermTuple})
    return outer[1:end-1] + last(outer) / inner
end

function Base.:(/)(outer::InteractionTerm, inner::AbstractTerm)
    # we should only get here via expansion where the interaction term,
    # but who knows what devious things users will try
    _fulldummycheck(outer)
    return outer + outer & inner
end

function Base.:(/)(outer::InteractionTerm, inner::TermTuple)
    # we should only get here via expansion where the interaction term,
    # but who knows what devious things users will try
    _fulldummycheck(outer)
    return mapfoldl(x -> outer & x, +, inner; init=outer)
end

function Base.:(/)(outer::AbstractTerm, inner::AbstractTerm)
    throw(ArgumentError("Nesting terms requires categorical grouping term, got $outer / $inner " *
                        "Manually specify $outer as `CategoricalTerm` in hints/contrasts"))
end

function StatsModels.apply_schema(
    t::FunctionTerm{typeof(/)},
    sch::StatsModels.FullRank,
    Mod::Type{<:RegressionModel},
)
    length(t.args_parsed) == 2 ||
        throw(ArgumentError("malformed nesting term: $t (Exactly two arguments required)"))

    any(x -> isa(x, ConstantTerm), t.args_parsed) && return t

    args = apply_schema.(t.args_parsed, Ref(sch), Mod)

    return first(args) / last(args)
end


const re_format = Dict([('a', "[A-Z]"), ('c', "[A-Z0-9]"), ('n', "[0-9]")])

const iban_global_re = r"^([A-Z]{2})([0-9]{2})([A-Z0-9]{10,34})$"

const iban_country_length = Dict([
                            ("AL", 28),  # Albania
                            ("AD", 24),  # Andorra
                            ("AE", 23),  # United Arab Emirates
                            ("AT", 20),  # Austria
                            ("AZ", 28),  # Azerbaijan
                            ("BA", 20),  # Bosnia and Herzegovina
                            ("BE", 16),  # Belgium
                            ("BG", 22),  # Bulgaria
                            ("BH", 22),  # Bahrain
                            ("BR", 29),  # Brazil
                            ("CH", 21),  # Switzerland
                            ("CR", 21),  # Costa Rica
                            ("CY", 28),  # Cyprus
                            ("CZ", 24),  # Czech Republic
                            ("DE", 22),  # Germany
                            ("DK", 18),  # Denmark
                            ("DO", 28),  # Dominican Republic
                            ("EE", 20),  # Estonia
                            ("ES", 24),  # Spain
                            ("FI", 18),  # Finland
                            ("FO", 18),  # Faroe Islands
                            ("FR", 27),  # France
                            ("GB", 22),  # United Kingdom
                            ("GE", 22),  # Georgia
                            ("GI", 23),  # Gibraltar
                            ("GL", 18),  # Greenland
                            ("GR", 27),  # Greece
                            ("GT", 28),  # Guatemala
                            ("HR", 21),  # Croatia
                            ("HU", 28),  # Hungary
                            ("IE", 22),  # Ireland
                            ("IL", 23),  # Israel
                            ("IS", 26),  # Iceland
                            ("IT", 27),  # Italy
                            ("JO", 30),  # Jordan
                            ("KZ", 20),  # Kazakhstan
                            ("KW", 30),  # Kuwait
                            ("LB", 28),  # Lebanon
                            ("LI", 21),  # Liechtenstein
                            ("LT", 20),  # Lithuania
                            ("LU", 20),  # Luxembourg
                            ("LV", 21),  # Latvia
                            ("MC", 27),  # Monaco
                            ("MD", 24),  # Moldova
                            ("ME", 22),  # Montenegro
                            ("MK", 19),  # Macedonia
                            ("MT", 31),  # Malta
                            ("MR", 27),  # Mauritania
                            ("MU", 30),  # Mauritius
                            ("NL", 18),  # Netherlands
                            ("NO", 15),  # Norway
                            ("PS", 29),  # Palestine
                            ("PK", 24),  # Pakistan
                            ("PL", 28),  # Poland
                            ("PT", 25),  # Portugal
                            ("QA", 29),  # Qatar
                            ("RO", 24),  # Romania
                            ("RS", 22),  # Serbia
                            ("SA", 24),  # Saudi Arabia
                            ("SE", 24),  # Sweden
                            ("SI", 19),  # Slovenia
                            ("SK", 24),  # Slovakia
                            ("SM", 27),  # San Marino
                            ("TN", 24),  # Tunisia
                            ("TR", 26),  # Turkey
                            ("VG", 24)  # British Virgin Islands
                             ])

const bban_format = Dict([
                          ("GB", ((4, 'a'), (14, 'n'))),
                          ("DE", ((18, 'n'), ))
                          ])


include("bench.jl")
include("visualize.jl")
include("place.jl")

import Random
using HTTP

# FIXME the documentation of HTTP is not very good. There are lots of
# undocumented APIs, like HTTP.payload

# test this: accept a file, append hello to the file, and respond back

function web_server()
    @info "Listening on localhost:8082 .."
    # FIXME 8081 seems to be used by jupyterhub, I'm thus using 8082
#     HTTP.Sockets.localhost
    HTTP.serve("0.0.0.0", 8082) do request::HTTP.Request
        @show request
        @show request.method

        payload = HTTP.payload(request)
        # CAUTION the String constructor method will clear the payload!
        # String(payload)

        @info "parsing payload .."
        jstr = String(payload)
        jobj = JSON.parse(jstr)
        xs, ys, as, ws, hs, Es, mask, diearea, params = decode_place_spec(jobj)
        @info "running placement .."

        # set seed. This should make the placement deterministic given the same
        # hyper-parameters
        Random.seed!(1234)

        solxs, solys = place(xs, ys, ws, hs, Es, mask, diearea,
                             nsteps=params["place-nsteps"],
                             nbins=params["place-nbins"])
        solxs, solys, solas, conflicts = simulated_annealing_legalization(
            solxs, solys, as, ws, hs, mask, diearea,
            ncycles=params["sa-ncycles"],
            nsteps=params["sa-nsteps"],
            stepsize=params["sa-stepsize"],
            theta_stepsize=params["sa-theta-stepsize"])

        # FIXME run iterations

        # @info "visualizing .."
        # R = Region(xs, ys, ws, hs, 300)
        # TODO for each incoming request, plot the process and save as log
        # visualize(xs, ys, ws, hs, R)
        # visualize(solxs, solys, ws, hs, R)

        @info "sending results back .."
        # I'll only send back cells, or cell locations
        # res_payload = encode(jobj, solxs, solys)
        #
        # the new payload just encode the xs and ys
        #
        # transform the result to corner based, because racket pict system convention
        #
        # UPDATE I'm using center at every locations, to be consistent for fixed locations
        # solxs = solxs .- ws ./ 2
        # solys = solys .- hs ./ 2
        res_payload = Dict("xs"=>solxs, "ys"=>solys,
                           "as"=>solas,
                           # DEBUG
                           "ws"=>ws,
                           "hs"=>hs,
                           # should I add diearea here?
                           "diearea"=>diearea,
                           # for visualization and debug
                           "conflicts"=>conflicts) |> JSON.json

        # TODO I also want to send back visualizations UPDATE probably just send
        # back the coordinates and optionally meta data during the process. The
        # rendering can be done on the client side.

        # parse this directly as json, and save data to internal data structure
        # and start placement and routing.
        try
            # return HTTP.Response("Hello")
            return HTTP.Response(res_payload)
        catch e
            return HTTP.Response(404, "Error: $e")
        end
    end
end

function test()
    # FIXME relative path might not work
    str = open("place-spec.json") do io
        read(io, String)
    end;
    jobj = JSON.parse(str);
    xs, ys, as, ws, hs, Es, mask, diearea, params = decode_place_spec(jobj);

    solxs, solys = place(xs, ys, ws, hs, Es, mask, diearea,
                         nsteps=params["place-nsteps"],
                         nbins=params["place-nbins"])

    solxs2, solys2, solas = simulated_annealing_legalization(
        solxs, solys, as, ws, hs, mask, diearea,
        # vis=true,
        ncycles=params["sa-ncycles"],
        nsteps=params["sa-nsteps"],
        stepsize=params["sa-stepsize"],
        theta_stepsize=params["sa-theta-stepsize"])

    R = Region(xs, ys, ws, hs, diearea, 300)
    visualize(solxs, solys, ws, hs, R)
    visualize(solxs2, solys2, ws, hs, R)
    # visualize(solxs, solys, ws, hs, R)
end