################################################################################
#
#	CONCRETE TYPE
#
#   Lattice{S,B,U} <: AbstractLattice{S,B,U}
#   --> S is the site type (<: AbstractSite{LS,D})
#       --> D is the dimension of embedding space
#       --> LS is the label type of sites
#   --> B is the bond type (<: AbstractBond{LB,N})
#       --> N is the dimension of the Bravais lattice that the bond is located in
#       --> LB is the label type of bonds
#   --> U is the unitcell type (<: AbstractUnitcell{SU,BU})
#       --> SU is the site type of the unitcell (<: AbstractSite{LUS,DU})
#           --> DU is the dimension of embedding space (of the UC)
#           --> LUS is the label type of sites (of the UC)
#       --> BU is the bond type of the unitcell (<: AbstractBond{LUB,NU})
#           --> NU is the dimension of the Bravais lattice that the bond is located in (of the UC)
#           --> LUB is the label type of bonds (of the UC)
#
#   FILE CONTAINS
#       - concrete struct definition
#       - interface implementation
#
################################################################################

################################################################################
#
#   CONCRETE STRUCT DEFINITION
#
################################################################################
mutable struct Lattice{S,B,U} <: AbstractLattice{S,B,U}

    # basis vectors of the Bravais lattice
    lattice_vectors	:: Vector{Vector{Float64}}

    # all sites within the lattice
    sites			:: Vector{S}

    # list of bonds
    bonds			:: Vector{B}

    # unitcell
    unitcell        :: U

end

# export the concrete type
export Lattice

################################################################################
#
#	INTERFACING / ACCESSING LATTICES
#	(functions have to be overwritten by concrete types)
#
################################################################################

# default constructor interface
# used for creation of new lattices
function newLattice(
            ::Type{Lattice{S,B,U}},
            lattice_vectors :: Vector{<:Vector{<:Real}},
            sites           :: Vector{S},
            bonds           :: Vector{B},
            unitcell        :: U
        ) :: Lattice{S,B,U} where {D,N,LS,LB,U,S<:AbstractSite{LS,D},B<:AbstractBond{LB,N}}

    # return a newly created object
    Lattice{S,B,U}(lattice_vectors, sites, bonds, unitcell)
end

# accessing a list of lattice vectors
function latticeVectors(
            lattice :: Lattice{S,B,U}
        ) :: Vector{Vector{Float64}} where {D,N,LS,LB,U,S<:AbstractSite{LS,D},B<:AbstractBond{LB,N}}

    # return the list of lattice vectors
    return lattice.lattice_vectors
end
# setting a list of lattice vectors
function latticeVectors!(
            lattice         :: Lattice{S,B,U},
            lattice_vectors :: Vector{<:Vector{<:Real}}
        ) where {D,N,LS,LB,U,S<:AbstractSite{LS,D},B<:AbstractBond{LB,N}}

    # set the list of lattice vectors
    lattice.lattice_vectors = lattice_vectors
end

# accessing a list of sites
function sites(
            lattice :: Lattice{S,B,U}
        ) :: Vector{S} where {D,N,LS,LB,U,S<:AbstractSite{LS,D},B<:AbstractBond{LB,N}}

    # return the list of sites
    return lattice.sites
end
# setting a list of sites
function sites!(
            lattice :: Lattice{S,B,U},
            sites   :: Vector{S}
        ) where {D,N,LS,LB,U,S<:AbstractSite{LS,D},B<:AbstractBond{LB,N}}

    # set the list of sites
    lattice.sites = sites
end

# accessing a list of bonds
function bonds(
            lattice :: Lattice{S,B,U}
        ) :: Vector{B} where {D,N,LS,LB,U,S<:AbstractSite{LS,D},B<:AbstractBond{LB,N}}

    # return the list of bonds
    return lattice.bonds
end
# setting a list of bonds
function bonds!(
            lattice :: Lattice{S,B,U},
            bonds   :: Vector{B}
        ) where {D,N,LS,LB,U,S<:AbstractSite{LS,D},B<:AbstractBond{LB,N}}

    # set the list of bonds
    lattice.bonds = bonds
end

# accessing the unitcell
function unitcell(
            lattice :: Lattice{S,B,U}
        ) :: U where {D,N,LS,LB,U,S<:AbstractSite{LS,D},B<:AbstractBond{LB,N}}

    # return the unitcell
    return lattice.unitcell
end
# setting the unitcell
function unitcell!(
            lattice  :: Lattice{S,B,U},
            unitcell :: U
        ) :: U where {D,N,LS,LB,U,S<:AbstractSite{LS,D},B<:AbstractBond{LB,N}}

    # set the unitcell
    lattice.unitcell = unitcell
end

function saveLattice(
        lt :: L,
        fn :: AbstractString,
        group :: AbstractString = "lattice"
        ;
        append :: Bool = false
    ) where {LS,LB,D,N,S<:AbstractSite{LS,D},B<:AbstractBond{LB,N},SU,BU,U<:AbstractUnitcell{SU,BU},L<:Lattice{S,B,U}}

    # determine the mode based on if one wants to append stuff
    if append
        mode = "r+"
    else
        mode = "w"
    end

    # determine the site and bond group names
    group_sites    = group*"/sites"
    group_bonds    = group*"/bonds"
    group_unitcell = group*"/unitcell"

    # open the file in mode
    h5open(fn, mode) do ltfile
        # create the top level group
        group_lt = g_create(ltfile, group)
        # save into the attributes in which groups the sites and bonds are stored
        attrs(group_lt)["sites"]    = group_sites
        attrs(group_lt)["bonds"]    = group_bonds
        attrs(group_lt)["unitcell"] = group_unitcell
        # save types into the attributes
        attrs(group_lt)["site_type"]     = string(S)
        attrs(group_lt)["bond_type"]     = string(B)
        attrs(group_lt)["unitcell_type"] = string(U)
        # number of bravais lattice dimensions
        attrs(group_lt)["N"] = Int64(N)
        # save further attributes

        # write bravais lattice vectors
        for i in 1:N
            group_lt["a"*string(i)] = latticeVectors(lt)[i]
        end
    end

    # save sites, bonds and unitcell into the file
    saveSites(sites(lt), fn, group_sites, append=true)
    saveBonds(bonds(lt), fn, group_bonds, append=true)
    saveUnitcell(unitcell(lt), fn, group_unitcell, append=true)

    # return nothing
    return nothing
end

function loadLattice(
        ::Type{LT},
        fn :: AbstractString,
        group :: AbstractString = "lattice"
    ) where {SLT<:AbstractSite,BLT<:AbstractBond,ULT<:AbstractUnitcell,LT<:Union{Lattice{SLT,BLT,ULT},Lattice}}

    # read the attribute data
    attr_data = h5readattr(fn, group)

    # determine the bravais lattice dimension
    N = attr_data["N"]

    # determine the site and bond group names
    group_sites    = attr_data["sites"]
    group_bonds    = attr_data["bonds"]
    group_unitcell = attr_data["unitcell"]

    # load bravais lattice vectors
    lattice_vectors = Vector{Float64}[]
    for i in 1:N
        push!(lattice_vectors, h5read(fn, group*"/a"*string(i)))
    end

    # load the sites
    site_list = loadSites(fn, group_sites)
    # load the bonds
    bond_list = loadBonds(fn, group_bonds)
    # load the unitcell
    unitcell_loaded = loadUnitcell(fn, group_unitcell)

    # determine site and bond type
    S = typeof(site_list[1])
    B = typeof(bond_list[1])
    U = typeof(unitcell_loaded)

    # return new unitcell
    return newLattice(
        Lattice{S,B,U},
        lattice_vectors,
        site_list,
        bond_list,
        unitcell_loaded
    )
end

# convinience function for standard type
function loadLattice(
        fn :: AbstractString,
        group :: AbstractString = "lattice"
    )

    return loadLattice(Lattice, fn, group)
end


abstract StochasticProcess{F<:VariateForm, S<:ValueSupport} <: Sampleable{F,S}

typealias UnivariateStochasticProcess{S<:ValueSupport} StochasticProcess{Univariate, S}
typealias MultivariateStochasticProcess{S<:ValueSupport} StochasticProcess{Multivariate, S}

typealias DiscreteStochasticProcess{F<:VariateForm} StochasticProcess{F, Discrete}
typealias ContinuousStochasticProcess{F<:VariateForm} StochasticProcess{F, Continuous}

#typealias DiscreteUnivariateStochasticProcess StochasticProcess{Univariate, Discrete}
typealias ContinuousUnivariateStochasticProcess StochasticProcess{Univariate, Continuous}
#typealias DiscreteMultivariateStochasticProcess StochasticProcess{Multivariate, Discrete}
typealias ContinuousMultivariateStochasticProcess StochasticProcess{Multivariate, Continuous}


# integration with AbstractFFTs.jl

@reexport using AbstractFFTs

import AbstractFFTs: plan_fft, plan_fft!, plan_bfft, plan_bfft!, plan_ifft,
    plan_rfft, plan_brfft, plan_inv, normalization, fft, bfft, ifft, rfft,
    Plan, ScaledPlan

using LinearAlgebra

Base.:(*)(p::Plan{T}, x::DenseCuArray) where {T} = p * copy1(T, x)
Base.:(*)(p::ScaledPlan, x::DenseCuArray) = rmul!(p.p * x, p.scale)

## plan structure

# K is an integer flag for forward/backward
# also used as an alias for r2c/c2r

# inplace is a boolean flag

abstract type CuFFTPlan{T<:cufftNumber, K, inplace} <: Plan{T} end

# for some reason, cufftHandle is an integer and not a pointer...
Base.convert(::Type{cufftHandle}, p::CuFFTPlan) = p.handle

function CUDA.unsafe_free!(plan::CuFFTPlan, stream::CuStream=stream())
    cufftDestroy(plan)
    unsafe_free!(plan.workarea, stream)
end

function unsafe_finalize!(plan::CuFFTPlan)
    context!(plan.ctx; skip_destroyed=true) do
        unsafe_free!(plan, default_stream())
    end
end

mutable struct cCuFFTPlan{T<:cufftNumber,K,inplace,N} <: CuFFTPlan{T,K,inplace}
    handle::cufftHandle
    ctx::CuContext
    stream::CuStream
    workarea::CuVector{Int8}
    sz::NTuple{N,Int} # Julia size of input array
    osz::NTuple{N,Int} # Julia size of output array
    xtype::cufftType
    region::Any
    pinv::ScaledPlan # required by AbstractFFT API

    function cCuFFTPlan{T,K,inplace,N}(handle::cufftHandle, workarea::CuVector{Int8},
                                       X::DenseCuArray{T,N}, sizey::Tuple, region, xtype;
                                       stream::CuStream=stream()) where {T<:cufftNumber,K,inplace,N}
        # maybe enforce consistency of sizey
        p = new(handle, context(), stream, workarea, size(X), sizey, xtype, region)
        finalizer(unsafe_finalize!, p)
        p
    end
end

mutable struct rCuFFTPlan{T<:cufftNumber,K,inplace,N} <: CuFFTPlan{T,K,inplace}
    handle::cufftHandle
    ctx::CuContext
    stream::CuStream
    workarea::CuVector{Int8}
    sz::NTuple{N,Int} # Julia size of input array
    osz::NTuple{N,Int} # Julia size of output array
    xtype::cufftType
    region::Any
    pinv::ScaledPlan # required by AbstractFFT API

    function rCuFFTPlan{T,K,inplace,N}(handle::cufftHandle, workarea::CuVector{Int8},
                                       X::DenseCuArray{T,N}, sizey::Tuple, region, xtype;
                                       stream::CuStream=stream()) where {T<:cufftNumber,K,inplace,N}
        # maybe enforce consistency of sizey
        p = new(handle, context(), stream, workarea, size(X), sizey, xtype, region)
        finalizer(unsafe_finalize!, p)
        p
    end
end

const xtypenames = Dict{cufftType,String}(CUFFT_R2C => "real-to-complex",
                                          CUFFT_C2R => "complex-to-real",
                                          CUFFT_C2C => "complex",
                                          CUFFT_D2Z => "d.p. real-to-complex",
                                          CUFFT_Z2D => "d.p. complex-to-real",
                                          CUFFT_Z2Z => "d.p. complex")

function showfftdims(io, sz, T)
    if isempty(sz)
        print(io,"0-dimensional")
    elseif length(sz) == 1
        print(io, sz[1], "-element")
    else
        print(io, join(sz, "×"))
    end
    print(io, " CuArray of ", T)
end

function Base.show(io::IO, p::CuFFTPlan{T,K,inplace}) where {T,K,inplace}
    print(io, inplace ? "CUFFT in-place " : "CUFFT ",
          xtypenames[p.xtype],
          K == CUFFT_FORWARD ? " forward" : " backward",
          " plan for ")
    showfftdims(io, p.sz, T)
end

Base.size(p::CuFFTPlan) = p.sz

# FFT plans can be user-created on a different task, whose stream might be different from
# the one used in the current task. call this function before every API call that performs
# operations on a stream to ensure the plan is using the correct task-local stream.
@inline function update_stream(plan::CuFFTPlan)
    new_stream = stream()
    if plan.stream != new_stream
        plan.stream = new_stream
        cufftSetStream(plan, new_stream)

        # replace the workarea by one (asynchronously) allocated on the current stream
        new_workarea = similar(plan.workarea)
        cufftSetWorkArea(plan, new_workarea)
        CUDA.unsafe_free!(plan.workarea)
        plan.workarea = new_workarea
    end
    return
end

## plan methods

# Note: we don't implement padded storage dimensions
function create_plan(xtype, xdims, region)
    nrank = length(region)
    sz = [xdims[i] for i in region]
    csz = copy(sz)
    csz[1] = div(sz[1],2) + 1
    batch = prod(xdims) ÷ prod(sz)

    # initialize the plan handle
    handle_ref = Ref{cufftHandle}()
    cufftCreate(handle_ref)
    handle = handle_ref[]

    # take control over the workarea
    cufftSetAutoAllocation(handle, 0)
    cufftSetStream(handle, stream())

    # make the plan
    worksize_ref = Ref{Csize_t}()
    if (nrank == 1) && (batch == 1)
        cufftMakePlan1d(handle, sz[1], xtype, 1, worksize_ref)
    elseif (nrank == 2) && (batch == 1)
        cufftMakePlan2d(handle, sz[2], sz[1], xtype, worksize_ref)
    elseif (nrank == 3) && (batch == 1)
        cufftMakePlan3d(handle, sz[3], sz[2], sz[1], xtype, worksize_ref)
    else
        rsz = (length(sz) > 1) ? rsz = reverse(sz) : sz
        if ((region...,) == ((1:nrank)...,))
            # handle simple case ... simply! (for robustness)
           cufftMakePlanMany(handle, nrank, Cint[rsz...], C_NULL, 1, 1, C_NULL, 1, 1,
                             xtype, batch, worksize_ref)
        else
            if nrank==1 || all(diff(collect(region)) .== 1)
                # _stride: successive elements in innermost dimension
                # _dist: distance between first elements of batches
                if region[1] == 1
                    istride = 1
                    idist = prod(sz)
                    cdist = prod(csz)
                else
                    if region[end] != length(xdims)
                        throw(ArgumentError("batching dims must be sequential"))
                    end
                    istride = prod(xdims[1:region[1]-1])
                    idist = 1
                    cdist = 1
                end
                inembed = Cint[rsz...]
                cnembed = (length(csz) > 1) ? Cint[reverse(csz)...] : Cint[csz[1]]
                ostride = istride
                if xtype == CUFFT_R2C || xtype == CUFFT_D2Z
                    odist = cdist
                    onembed = cnembed
                else
                    odist = idist
                    onembed = inembed
                end
                if xtype == CUFFT_C2R || xtype == CUFFT_Z2D
                    idist = cdist
                    inembed = cnembed
                end
            else
                if any(diff(collect(region)) .< 1)
                    throw(ArgumentError("region must be an increasing sequence"))
                end
                cdims = collect(xdims)
                cdims[region[1]] = div(cdims[region[1]],2)+1

                if region[1] == 1
                    istride = 1
                    ii=1
                    while (ii < nrank) && (region[ii] == region[ii+1]-1)
                        ii += 1
                    end
                    idist = prod(xdims[1:ii])
                    cdist = prod(cdims[1:ii])
                    ngaps = 0
                else
                    istride = prod(xdims[1:region[1]-1])
                    idist = 1
                    cdist = 1
                    ngaps = 1
                end
                nem = ones(Int,nrank)
                cem = ones(Int,nrank)
                id = 1
                for ii=1:nrank-1
                    if region[ii+1] > region[ii]+1
                        ngaps += 1
                    end
                    while id < region[ii+1]
                        nem[ii] *= xdims[id]
                        cem[ii] *= cdims[id]
                        id += 1
                    end
                    @assert nem[ii] >= sz[ii]
                end
                if region[end] < length(xdims)
                    ngaps += 1
                end
                # CUFFT represents batches by a single stride (_dist)
                # so we must verify that region is consistent with this:
                if ngaps > 1
                    throw(ArgumentError("batch regions must be sequential"))
                end

                inembed = Cint[reverse(nem)...]
                cnembed = Cint[reverse(cem)...]
                ostride = istride
                if xtype == CUFFT_R2C || xtype == CUFFT_D2Z
                    odist = cdist
                    onembed = cnembed
                else
                    odist = idist
                    onembed = inembed
                end
                if xtype == CUFFT_C2R || xtype == CUFFT_Z2D
                    idist = cdist
                    inembed = cnembed
                end
            end
            cufftMakePlanMany(handle, nrank, Cint[rsz...],
                              inembed, istride, idist, onembed, ostride, odist,
                              xtype, batch, worksize_ref)
        end
    end

    # assign the workarea
    workarea = CuArray{Int8}(undef, worksize_ref[])
    cufftSetWorkArea(handle, workarea)

    handle, workarea
end

# promote to a complex floating-point type (out-of-place only),
# so implementations only need Complex{Float} methods
for f in (:fft, :bfft, :ifft)
    pf = Symbol("plan_", f)
    @eval begin
        $f(x::DenseCuArray{<:Real}, region=1:ndims(x)) = $f(complexfloat(x), region)
        $pf(x::DenseCuArray{<:Real}, region) = $pf(complexfloat(x), region)
        $f(x::DenseCuArray{<:Complex{<:Union{Integer,Rational}}}, region=1:ndims(x)) = $f(complexfloat(x), region)
        $pf(x::DenseCuArray{<:Complex{<:Union{Integer,Rational}}}, region) = $pf(complexfloat(x), region)
    end
end
rfft(x::DenseCuArray{<:Union{Integer,Rational}}, region=1:ndims(x)) = rfft(realfloat(x), region)
plan_rfft(x::DenseCuArray{<:Real}, region) = plan_rfft(realfloat(x), region)

# region is an iterable subset of dimensions
# spec. an integer, range, tuple, or array

# inplace complex
function plan_fft!(X::DenseCuArray{T,N}, region) where {T<:cufftComplexes,N}
    K = CUFFT_FORWARD
    inplace = true
    xtype = (T == cufftComplex) ? CUFFT_C2C : CUFFT_Z2Z

    handle, workarea = create_plan(xtype, size(X), region)

    cCuFFTPlan{T,K,inplace,N}(handle, workarea, X, size(X), region, xtype)
end

function plan_bfft!(X::DenseCuArray{T,N}, region) where {T<:cufftComplexes,N}
    K = CUFFT_INVERSE
    inplace = true
    xtype =  (T == cufftComplex) ? CUFFT_C2C : CUFFT_Z2Z

    handle, workarea = create_plan(xtype, size(X), region)

    cCuFFTPlan{T,K,inplace,N}(handle, workarea, X, size(X), region, xtype)
end

# out-of-place complex
function plan_fft(X::DenseCuArray{T,N}, region) where {T<:cufftComplexes,N}
    K = CUFFT_FORWARD
    xtype =  (T == cufftComplex) ? CUFFT_C2C : CUFFT_Z2Z
    inplace = false

    handle, workarea = create_plan(xtype, size(X), region)

    cCuFFTPlan{T,K,inplace,N}(handle, workarea, X, size(X), region, xtype)
end

function plan_bfft(X::DenseCuArray{T,N}, region) where {T<:cufftComplexes,N}
    K = CUFFT_INVERSE
    inplace = false
    xtype =  (T == cufftComplex) ? CUFFT_C2C : CUFFT_Z2Z

    handle, workarea = create_plan(xtype, size(X), region)

    cCuFFTPlan{T,K,inplace,N}(handle, workarea, X, size(X), region, xtype)
end

# out-of-place real-to-complex
function plan_rfft(X::DenseCuArray{T,N}, region) where {T<:cufftReals,N}
    K = CUFFT_FORWARD
    inplace = false
    xtype =  (T == cufftReal) ? CUFFT_R2C : CUFFT_D2Z

    handle, workarea = create_plan(xtype, size(X), region)

    ydims = collect(size(X))
    ydims[region[1]] = div(ydims[region[1]],2)+1

    rCuFFTPlan{T,K,inplace,N}(handle, workarea, X, (ydims...,), region, xtype)
end

function plan_brfft(X::DenseCuArray{T,N}, d::Integer, region::Any) where {T<:cufftComplexes,N}
    K = CUFFT_INVERSE
    inplace = false
    xtype =  (T == cufftComplex) ? CUFFT_C2R : CUFFT_Z2D
    ydims = collect(size(X))
    ydims[region[1]] = d

    handle, workarea = create_plan(xtype, (ydims...,), region)

    rCuFFTPlan{T,K,inplace,N}(handle, workarea, X, (ydims...,), region, xtype)
end

# FIXME: plan_inv methods allocate needlessly (to provide type parameters)
# Perhaps use FakeArray types to avoid this.

function plan_inv(p::cCuFFTPlan{T,CUFFT_FORWARD,inplace,N}) where {T,N,inplace}
    X = CuArray{T}(undef, p.sz)
    handle, workarea = create_plan(p.xtype, p.sz, p.region)
    ScaledPlan(cCuFFTPlan{T,CUFFT_INVERSE,inplace,N}(handle, workarea, X, p.sz, p.region,
                                                     p.xtype),
               normalization(X, p.region))
end

function plan_inv(p::cCuFFTPlan{T,CUFFT_INVERSE,inplace,N}) where {T,N,inplace}
    X = CuArray{T}(undef, p.sz)
    handle, workarea = create_plan(p.xtype, p.sz, p.region)
    ScaledPlan(cCuFFTPlan{T,CUFFT_FORWARD,inplace,N}(handle, workarea, X, p.sz, p.region,
                                                     p.xtype),
               normalization(X, p.region))
end

function plan_inv(p::rCuFFTPlan{T,CUFFT_INVERSE,inplace,N}
                  ) where {T<:cufftComplexes,N,inplace}
    X = CuArray{real(T)}(undef, p.osz)
    Y = CuArray{T}(undef, p.sz)
    xtype = p.xtype == CUFFT_C2R ? CUFFT_R2C : CUFFT_D2Z
    handle, workarea = create_plan(xtype, p.osz, p.region)
    ScaledPlan(rCuFFTPlan{real(T),CUFFT_FORWARD,inplace,N}(handle, workarea, X, p.sz, p.region, xtype),
               normalization(X, p.region))
end

function plan_inv(p::rCuFFTPlan{T,CUFFT_FORWARD,inplace,N}
                  ) where {T<:cufftReals,N,inplace}
    X = CuArray{complex(T)}(undef, p.osz)
    Y = CuArray{T}(undef, p.sz)
    xtype = p.xtype == CUFFT_R2C ? CUFFT_C2R : CUFFT_Z2D
    handle, workarea = create_plan(xtype, p.sz, p.region)
    ScaledPlan(rCuFFTPlan{complex(T),CUFFT_INVERSE,inplace,N}(handle, workarea, X, p.sz,
                                                              p.region, xtype),
               normalization(Y, p.region))
end

## plan execution

# NOTE: "in-place complex-to-real FFTs may overwrite arbitrary imaginary input point values
#       [...]. Out-of-place complex-to-real FFT will always overwrite input buffer."
#       see # JuliaGPU/CuArrays.jl#345, NVIDIA/cuFFT#2714055.

function assert_applicable(p::CuFFTPlan{T}, X::DenseCuArray{T}) where {T}
    (size(X) == p.sz) ||
        throw(ArgumentError("CuFFT plan applied to wrong-size input"))
end

function assert_applicable(p::CuFFTPlan{T,K,inplace}, X::DenseCuArray{T},
                           Y::DenseCuArray) where {T,K,inplace}
    assert_applicable(p, X)
    if size(Y) != p.osz
        throw(ArgumentError("CuFFT plan applied to wrong-size output"))
    elseif inplace != (pointer(X) == pointer(Y))
        throw(ArgumentError(string("CuFFT ",
                                   inplace ? "in-place" : "out-of-place",
                                   " plan applied to ",
                                   inplace ? "out-of-place" : "in-place",
                                   " data")))
    end
end

function unsafe_execute!(plan::cCuFFTPlan{cufftComplex,K,<:Any,N},
                         x::DenseCuArray{cufftComplex,N},
                         y::DenseCuArray{cufftComplex,N}) where {K,N}
    @assert plan.xtype == CUFFT_C2C
    update_stream(plan)
    cufftExecC2C(plan, x, y, K)
end

function unsafe_execute!(plan::rCuFFTPlan{cufftComplex,K,true,N},
                         x::DenseCuArray{cufftComplex,N},
                         y::DenseCuArray{cufftReal,N}) where {K,N}
    @assert plan.xtype == CUFFT_C2R
    update_stream(plan)
    cufftExecC2R(plan, x, y)
end
function unsafe_execute!(plan::rCuFFTPlan{cufftComplex,K,false,N},
                         x::DenseCuArray{cufftComplex,N},
                         y::DenseCuArray{cufftReal}) where {K,N}
    @assert plan.xtype == CUFFT_C2R
    x = copy(x)
    update_stream(plan)
    cufftExecC2R(plan, x, y)
    unsafe_free!(x)
end

function unsafe_execute!(plan::rCuFFTPlan{cufftReal,K,<:Any,N},
                         x::DenseCuArray{cufftReal,N},
                         y::DenseCuArray{cufftComplex,N}) where {K,N}
    @assert plan.xtype == CUFFT_R2C
    update_stream(plan)
    cufftExecR2C(plan, x, y)
end

function unsafe_execute!(plan::cCuFFTPlan{cufftDoubleComplex,K,<:Any,N},
                         x::DenseCuArray{cufftDoubleComplex,N},
                         y::DenseCuArray{cufftDoubleComplex}) where {K,N}
    @assert plan.xtype == CUFFT_Z2Z
    update_stream(plan)
    cufftExecZ2Z(plan, x, y, K)
end

function unsafe_execute!(plan::rCuFFTPlan{cufftDoubleComplex,K,true,N},
                         x::DenseCuArray{cufftDoubleComplex,N},
                         y::DenseCuArray{cufftDoubleReal}) where {K,N}
    update_stream(plan)
    @assert plan.xtype == CUFFT_Z2D
    cufftExecZ2D(plan, x, y)
end
function unsafe_execute!(plan::rCuFFTPlan{cufftDoubleComplex,K,false,N},
                         x::DenseCuArray{cufftDoubleComplex,N},
                         y::DenseCuArray{cufftDoubleReal}) where {K,N}
    @assert plan.xtype == CUFFT_Z2D
    x = copy(x)
    update_stream(plan)
    cufftExecZ2D(plan, x, y)
    unsafe_free!(x)
end

function unsafe_execute!(plan::rCuFFTPlan{cufftDoubleReal,K,<:Any,N},
                         x::DenseCuArray{cufftDoubleReal,N},
                         y::DenseCuArray{cufftDoubleComplex,N}) where {K,N}
    @assert plan.xtype == CUFFT_D2Z
    update_stream(plan)
    cufftExecD2Z(plan, x, y)
end

## high-level integrations

function LinearAlgebra.mul!(y::DenseCuArray{Ty}, p::CuFFTPlan{T}, x::DenseCuArray{T}
                           ) where {Ty, T}
    assert_applicable(p,x,y)
    unsafe_execute!(p,x,y)
    return y
end

function Base.:(*)(p::cCuFFTPlan{T,K,true,N}, x::DenseCuArray{T,N}) where {T,K,N}
    assert_applicable(p,x)
    unsafe_execute!(p,x,x)
    x
end

function Base.:(*)(p::rCuFFTPlan{T,CUFFT_FORWARD,false,N}, x::DenseCuArray{T,N}
           ) where {T<:cufftReals,N}
    assert_applicable(p,x)
    @assert p.xtype ∈ [CUFFT_R2C,CUFFT_D2Z]
    y = CuArray{complex(T),N}(undef, p.osz)
    unsafe_execute!(p,x,y)
    y
end

function Base.:(*)(p::rCuFFTPlan{T,CUFFT_INVERSE,false,N}, x::DenseCuArray{T,N}
           ) where {T<:cufftComplexes,N}
    assert_applicable(p,x)
    @assert p.xtype ∈ [CUFFT_C2R,CUFFT_Z2D]
    y = CuArray{real(T),N}(undef, p.osz)
    unsafe_execute!(p,x,y)
    y
end

function Base.:(*)(p::cCuFFTPlan{T,K,false,N}, x::DenseCuArray{T,N}) where {T,K,N}
    assert_applicable(p,x)
    y = CuArray{T,N}(undef, p.osz)
    unsafe_execute!(p,x,y)
    y
end


fetchvcat(J::judiVector) = J
fetchvcat(x) = vcat(x...)

function JUDI.time_modeling(model::Model, srcGeometry, srcData, recGeometry, recData, dm, srcnum::UnitRange{Int64}, op::Char, mode::Int64, options)

    # Broadcast common parameters
    _model = @bcast model
    _dm = isnothing(dm) ? dm : @bcast dm
    opts = make_opt([model, dm, op, srcGeometry])
    iter = make_parts(srcnum)
    # Run on azure
    results = @batchexec pmap(sx -> time_modeling_azure(_model, subsample(srcGeometry, sx), subsample(srcData, sx),
                                                        subsample(recGeometry, sx), subsample(recData, sx), _dm,
                                                        op, mode, subsample(options, sx)), iter) opts
    # Gather results
    if op=='F' || (op=='J' && mode==1)
        temp = fetch(results)
        argout1 = fetchvcat(temp)
    elseif op=='J' && mode==-1
        argout1 = fetchreduce(results; op=+, remote=true)
    else
        error("operation no defined")
    end
    return argout1
end

# FWI
function JUDI.fwi_objective(model::Model, source::judiVector, dObs::judiVector; options=JUDI.Options())
    # fwi_objective function for multiple sources. The function distributes the sources and the input data amongst the available workers.
    # Broadcast common parameters
    _model = @bcast model
    opts = make_opt([model, dObs, source])
    iter = make_parts(1:dObs.nsrc)
    results = @batchexec pmap(j -> fwi_objective_azure(_model, source[j], dObs[j], subsample(options, j)), iter) opts

    # Collect and reduce gradients
    obj, gradient = fetchreduce(results; op=+, remote=true)
    # first value corresponds to function value, the rest to the gradient
    return obj, gradient
end

# lsrtm

function JUDI.lsrtm_objective(model::Model, source::judiVector, dObs::judiVector, dm; options=JUDI.Options(), nlind=false)
    # lsrtm_objective function for multiple sources. The function distributes the sources and the input data amongst the available workers.

    # Broadcast common parameters
    _model = @bcast model
    _dm = isnothing(dm) ? dm : @bcast dm

    opts = make_opt([model, dObs, source, dm])
    iter = make_parts(1:dObs.nsrc)
    results = @batchexec pmap(j -> lsrtm_objective_azure(_model, source[j], dObs[j], _dm, subsample(options, j); nlind=nlind), iter) opts
    # Collect and reduce gradients
    obj, gradient = fetchreduce(results; op=+, remote=false)
    delete_job(results)

    # first value corresponds to function value, the rest to the gradient
    return obj, gradient
end

function JUDI.lsrtm_objective(model::Array{Model,1}, source::Array{judiVector{T,Array{T,2}},1}, dObs::Array{judiVector{T,Array{T,2}},1}, dm::Union{Array{Array{T,1},1}, Array{PhysicalParameter{T},1}}; options=Options(), nlind=false) where T
# lsrtm_objective function for multiple sources and multiple vintages. The function distributes the sources and the input data amongst the available workers.

    println("serial lsrtm_objective on each vintage now")
    ctrl = Array{Any, 1}(undef, length(dObs))
    obj = Array{T, 1}(undef, length(dObs))
    gradient = Array{Array{T,2}, 1}(undef, length(dObs))

    for i = 1:length(dObs)
        _model = model[i]
        _source = source[i]
        _dObs = dObs[i]
        _dm = dm[i]
        _model1 = @bcast _model
        _dm1 = @bcast _dm
        opts = make_opt([_model, dObs, source, _dm])
        ctrl[i] = @batchexec pmap(j -> lsrtm_objective_azure(_model1, _source[j], _dObs[j], _dm1, subsample(options, j); nlind=nlind), 1:dObs[i].nsrc) opts
    end
    for i = 1:length(dObs)
        results = fetch(ctrl[i])
        delete_job(ctrl[i])
        obj[i] = sum([results[j][1] for j =  1:dObs[i].nsrc])
        gradient[i] = sum([results[j][2] for j = 1:dObs[i].nsrc])
    end

    obj1 = sum(obj)

    # first value corresponds to function value, the rest to the gradient
    return obj1, gradient
end

# Make options for unique batch ids
make_opt(a::Array) = AzureClusterlessHPC.Options(;task_name="task_$(objectid(a))")

# Batchexec function
slicec(i::Integer, j::Integer) = i==j ? i : (i:j)
chunk(arr, n) = [arr[slicec(i, min(i + n - 1, length(arr)))] for i in 1:n:length(arr)]
make_parts(iter) = _njpi > 1 ? chunk(iter, _njpi) : iter


function evaluate_tc(obj::IERSConventions, arg0::AbsoluteDate)
    return jcall(obj, "evaluateTC", jdouble, (AbsoluteDate,), arg0)
end

function evaluate_tc(obj::IERSConventions, arg0::AbsoluteDate, arg1::TimeScales)
    return jcall(obj, "evaluateTC", jdouble, (AbsoluteDate, TimeScales), arg0, arg1)
end

function evaluate_tc(obj::IERSConventions, arg0::FieldAbsoluteDate)
    return jcall(obj, "evaluateTC", RealFieldElement, (FieldAbsoluteDate,), arg0)
end

function evaluate_tc(obj::IERSConventions, arg0::FieldAbsoluteDate, arg1::TimeScales)
    return jcall(obj, "evaluateTC", RealFieldElement, (FieldAbsoluteDate, TimeScales), arg0, arg1)
end

function get_earth_orientation_angle_function(obj::IERSConventions, arg0::TimeScale)
    return jcall(obj, "getEarthOrientationAngleFunction", TimeScalarFunction, (TimeScale,), arg0)
end

function get_earth_orientation_angle_function(obj::IERSConventions, arg0::TimeScale, arg1::TimeScale)
    return jcall(obj, "getEarthOrientationAngleFunction", TimeScalarFunction, (TimeScale, TimeScale), arg0, arg1)
end

function get_eop_tidal_correction(obj::IERSConventions)
    return jcall(obj, "getEOPTidalCorrection", TimeVectorFunction, ())
end

function get_eop_tidal_correction(obj::IERSConventions, arg0::TimeScales)
    return jcall(obj, "getEOPTidalCorrection", TimeVectorFunction, (TimeScales,), arg0)
end

function get_gast_function(obj::IERSConventions, arg0::TimeScale, arg1::EOPHistory)
    return jcall(obj, "getGASTFunction", TimeScalarFunction, (TimeScale, EOPHistory), arg0, arg1)
end

function get_gast_function(obj::IERSConventions, arg0::TimeScale, arg1::EOPHistory, arg2::TimeScales)
    return jcall(obj, "getGASTFunction", TimeScalarFunction, (TimeScale, EOPHistory, TimeScales), arg0, arg1, arg2)
end

function get_gmst_function(obj::IERSConventions, arg0::TimeScale)
    return jcall(obj, "getGMSTFunction", TimeScalarFunction, (TimeScale,), arg0)
end

function get_gmst_function(obj::IERSConventions, arg0::TimeScale, arg1::TimeScales)
    return jcall(obj, "getGMSTFunction", TimeScalarFunction, (TimeScale, TimeScales), arg0, arg1)
end

function get_gmst_rate_function(obj::IERSConventions, arg0::TimeScale)
    return jcall(obj, "getGMSTRateFunction", TimeScalarFunction, (TimeScale,), arg0)
end

function get_gmst_rate_function(obj::IERSConventions, arg0::TimeScale, arg1::TimeScales)
    return jcall(obj, "getGMSTRateFunction", TimeScalarFunction, (TimeScale, TimeScales), arg0, arg1)
end

function get_love_numbers(obj::IERSConventions)
    return jcall(obj, "getLoveNumbers", LoveNumbers, ())
end

function get_mean_obliquity_function(obj::IERSConventions)
    return jcall(obj, "getMeanObliquityFunction", TimeScalarFunction, ())
end

function get_mean_obliquity_function(obj::IERSConventions, arg0::TimeScales)
    return jcall(obj, "getMeanObliquityFunction", TimeScalarFunction, (TimeScales,), arg0)
end

function get_nominal_tidal_displacement(obj::IERSConventions)
    return jcall(obj, "getNominalTidalDisplacement", Vector{jdouble}, ())
end

function get_nutation_arguments(obj::IERSConventions, arg0::TimeScale)
    return jcall(obj, "getNutationArguments", FundamentalNutationArguments, (TimeScale,), arg0)
end

function get_nutation_arguments(obj::IERSConventions, arg0::TimeScale, arg1::TimeScales)
    return jcall(obj, "getNutationArguments", FundamentalNutationArguments, (TimeScale, TimeScales), arg0, arg1)
end

function get_nutation_correction_converter(obj::IERSConventions)
    return jcall(obj, "getNutationCorrectionConverter", IERSConventions_NutationCorrectionConverter, ())
end

function get_nutation_correction_converter(obj::IERSConventions, arg0::TimeScales)
    return jcall(obj, "getNutationCorrectionConverter", IERSConventions_NutationCorrectionConverter, (TimeScales,), arg0)
end

function get_nutation_function(obj::IERSConventions)
    return jcall(obj, "getNutationFunction", TimeVectorFunction, ())
end

function get_nutation_function(obj::IERSConventions, arg0::TimeScales)
    return jcall(obj, "getNutationFunction", TimeVectorFunction, (TimeScales,), arg0)
end

function get_nutation_reference_epoch(obj::IERSConventions)
    return jcall(obj, "getNutationReferenceEpoch", AbsoluteDate, ())
end

function get_nutation_reference_epoch(obj::IERSConventions, arg0::TimeScales)
    return jcall(obj, "getNutationReferenceEpoch", AbsoluteDate, (TimeScales,), arg0)
end

function get_ocean_pole_tide(obj::IERSConventions, arg0::EOPHistory)
    return jcall(obj, "getOceanPoleTide", TimeVectorFunction, (EOPHistory,), arg0)
end

function get_permanent_tide(obj::IERSConventions)
    return jcall(obj, "getPermanentTide", jdouble, ())
end

function get_precession_function(obj::IERSConventions)
    return jcall(obj, "getPrecessionFunction", TimeVectorFunction, ())
end

function get_precession_function(obj::IERSConventions, arg0::TimeScales)
    return jcall(obj, "getPrecessionFunction", TimeVectorFunction, (TimeScales,), arg0)
end

function get_solid_pole_tide(obj::IERSConventions, arg0::EOPHistory)
    return jcall(obj, "getSolidPoleTide", TimeVectorFunction, (EOPHistory,), arg0)
end

function get_tidal_displacement_frequency_correction_diurnal(obj::IERSConventions)
    return jcall(obj, "getTidalDisplacementFrequencyCorrectionDiurnal", PoissonSeries_CompiledSeries, ())
end

function get_tidal_displacement_frequency_correction_zonal(obj::IERSConventions)
    return jcall(obj, "getTidalDisplacementFrequencyCorrectionZonal", PoissonSeries_CompiledSeries, ())
end

function get_tide_frequency_dependence_function(obj::IERSConventions, arg0::TimeScale)
    return jcall(obj, "getTideFrequencyDependenceFunction", TimeVectorFunction, (TimeScale,), arg0)
end

function get_tide_frequency_dependence_function(obj::IERSConventions, arg0::TimeScale, arg1::TimeScales)
    return jcall(obj, "getTideFrequencyDependenceFunction", TimeVectorFunction, (TimeScale, TimeScales), arg0, arg1)
end

function get_xy_sp_xy2_function(obj::IERSConventions)
    return jcall(obj, "getXYSpXY2Function", TimeVectorFunction, ())
end

function get_xy_sp_xy2_function(obj::IERSConventions, arg0::TimeScales)
    return jcall(obj, "getXYSpXY2Function", TimeVectorFunction, (TimeScales,), arg0)
end

function value_of(::Type{IERSConventions}, arg0::JString)
    return jcall(IERSConventions, "valueOf", IERSConventions, (JString,), arg0)
end

function values(::Type{IERSConventions})
    return jcall(IERSConventions, "values", Vector{IERSConventions}, ())
end


# XXX things to possibly speed this up
# replace recursion with while loop
# cache simulation results

function create_policy{S}(solver::AbstractPOMCPSolver, pomdp::POMDPs.POMDP{S})
    if isa(solver.node_belief_updater, DefaultReinvigoratorStub)
        node_belief_updater = DeadReinvigorator{S}()
    else
        node_belief_updater = solver.node_belief_updater
    end
    return POMCPPlanner(pomdp, solver, node_belief_updater,
                        convert_estimator(solver.estimate_value, solver, pomdp),
                        Nullable{Any}())
end

function action(policy::POMCPPlanner, belief, a=nothing)
    try
        a = search(policy, make_node(policy, belief), policy.solver.tree_queries)
    catch ex
        a = default_action(policy.solver.default_action, belief, ex)
    end
    return a
end

"""
    solve(solver::POMCPSolver, pomdp::POMDPs.POMDP)

Simply return a properly constructed POMCPPlanner object.
"""
function solve{S}(solver::AbstractPOMCPSolver, pomdp::POMDPs.POMDP{S})
    create_policy(solver, pomdp)
end

solve(solver::AbstractPOMCPSolver, pomdp::POMDPs.POMDP, dummy_policy::Any) = solve(solver, pomdp)

"""
    function search(pomcp::POMCPPlanner, b::BeliefNode, tree_queries)

Search the tree for the next best move.
"""
function search(pomcp::POMCPPlanner, b::BeliefNode, tree_queries)
    # XXX hack
    pomcp._tree_ref = b
    # end hack

    all_terminal = true
    for i in 1:tree_queries
        s = rand(pomcp.solver.rng, b)
        if !POMDPs.isterminal(pomcp.problem, s)
            simulate(pomcp, b, s, 0)
            b.N += 1
            all_terminal = false
        end
    end

    if all_terminal
        throw(AllSamplesTerminal(b.B))
    end

    best_node = first(values(b.children))
    best_V = best_node.V
    @assert !isnan(best_V)
    for node in collect(values(b.children))[2:end]
        if node.V >= best_V
            best_V = node.V
            best_node = node
        end
    end

    return best_node.label
end

"""
    simulate{S}(pomcp::POMCPPlanner, h::BeliefNode, s::S, depth)

Move the simulation forward a single step and update the BeliefNode h accordingly.
"""
function simulate{S,A,O,B}(pomcp::POMCPPlanner{S,A,O,B,POMCPSolver}, h::BeliefNode, s::S, depth)

    sol = pomcp.solver

    if POMDPs.discount(pomcp.problem)^depth < sol.eps ||
            POMDPs.isterminal(pomcp.problem, s) ||
            depth >= sol.max_depth
        return 0
    end
	if isempty(h.children)
        action_space_iter = sparse_actions(pomcp, pomcp.problem, h, sol.num_sparse_actions)
		for a in action_space_iter
			h.children[a] = ActNode(a,
                                    init_N(sol.init_N, pomcp.problem, h, a),
                                    init_V(sol.init_V, pomcp.problem, h, a),
                                    Dict{O,ObsNode{B,A,O}}())
		end

        if depth > 0 # no need for a rollout if this is the root node
            if POMDPs.discount(pomcp.problem) == 1.0
                step_to_eps = typemax(Int)
            else
                steps_to_eps = ceil(Int, log(sol.eps)/log(POMDPs.discount(pomcp.problem))-depth)
            end
            steps = min(sol.max_depth-depth, steps_to_eps)
            return POMDPs.discount(pomcp.problem)^depth * estimate_value(pomcp.solved_estimate, pomcp.problem, s, h, steps)
        else
            return 0.
        end
	end

    best_criterion_val = -Inf
    local best_node
    for node in values(h.children)
        if node.N == 0 && h.N == 1
            criterion_value = node.V
        elseif node.N == 0 && node.V == -Inf
            criterion_value = Inf
        else
            criterion_value = node.V + sol.c*sqrt(log(h.N)/node.N)
        end
        if criterion_value >= best_criterion_val
            best_criterion_val = criterion_value
            best_node = node
        end
    end
    a = best_node.label

    (sp, o, r) = POMDPs.generate_sor(pomcp.problem, s, a, sol.rng)

    if r == Inf
        warn("POMCP: +Inf reward. This is not recommended and may cause future errors.")
    end

    if haskey(best_node.children, o)
        hao = best_node.children[o]
    else
        if isa(pomcp.node_belief_updater, ParticleReinvigorator)
            hao = ObsNode(o, 0, ParticleCollection{S}(), Dict{A,ActNode{A,O,ObsNode{B,A,O}}}())
        else
            new_belief = update(pomcp.node_belief_updater, h.B, a, o) # this relies on h.B not being modified
            hao = ObsNode(o, 0, new_belief, Dict{A,ActNode{A,O,ObsNode{B,A,O}}}())
        end
        best_node.children[o]=hao
    end

    R = r + POMDPs.discount(pomcp.problem)*simulate(pomcp, hao, sp, depth+1)

    if uses_states_from_planner(hao.B)
        push!(hao.B, sp) # insert this state into the particle collection
    end
    hao.N += 1

    best_node.N += 1
    if best_node.V != -Inf
        best_node.V += (R-best_node.V)/best_node.N
    end

    return R
end

function simulate{S,A,O,B}(pomcp::POMCPPlanner{S,A,O,B,POMCPDPWSolver}, h::BeliefNode, s::S, depth)

    sol = pomcp.solver

    if POMDPs.discount(pomcp.problem)^depth < sol.eps ||
            POMDPs.isterminal(pomcp.problem, s) ||
            depth >= sol.max_depth
        return 0
    end

    total_N = reduce(add_N, 0, values(h.children))
    if sol.enable_action_pw
        if length(h.children) <= sol.k_action*total_N^sol.alpha_action
            a = next_action(sol.next_action, pomcp.problem, h.B, h) # XXX should this be a function of s or h.B?
            if !(a in keys(h.children))
                h.children[a] = ActNode(a,
                                        init_N(sol.init_N, pomcp.problem, h, a),
                                        init_V(sol.init_V, pomcp.problem, h, a),
                                        Dict{O,ObsNode{B,A,O}}())
            end
            if length(h.children) <= 1
                if depth > 0
                    return POMDPs.discount(pomcp.problem)^depth * estimate_value(pomcp.solved_estimate, pomcp.problem, s, h, depth)
                else
                    return 0.
                end
            end
        end
    else # run through all the actions
        if isempty(h.children)
            action_space_iter = POMDPs.iterator(POMDPs.actions(pomcp.problem))
            for a in action_space_iter
                h.children[a] = ActNode(a,
                                        init_N(sol.init_N, pomcp.problem, h, a),
                                        init_V(sol.init_V, pomcp.problem, h, a),
                                        Dict{O,ObsNode{B,A,O}}())
            end

            if depth > 0 # no need for a rollout if this is the root node
                return POMDPs.discount(pomcp.problem)^depth * estimate_value(pomcp.solved_estimate, pomcp.problem, s, h, depth)
            else
                return 0.
            end
        end
        total_N = h.N
    end

    # Calculate UCT
    best_criterion_val = -Inf
    local best_node
    for (a,node) in h.children
        if node.N == 0 && total_N <= 1
            criterion_value = node.V
        elseif node.N == 0 && node.V == -Inf
            criterion_value = Inf
        else
            criterion_value = node.V + sol.c*sqrt(log(total_N)/node.N)
        end
        if criterion_value >= best_criterion_val
            best_criterion_val = criterion_value
            best_node = node
        end
    end
    a = best_node.label

    if length(best_node.children) <= sol.k_observation*(best_node.N^sol.alpha_observation)
        state_was_generated = true

        # observation progressive widening
        (sp, o, r) = POMDPs.generate_sor(pomcp.problem, s, a, sol.rng)

        if r == Inf
            warn("POMCP: +Inf reward. This is not recommended and may cause future errors.")
        end

        if haskey(best_node.children, o)
            hao = best_node.children[o]
        else
            if isa(pomcp.node_belief_updater, ParticleReinvigorator)
                hao = ObsNode(o, 0, ParticleCollection{S}(), Dict{A,ActNode{A,O,ObsNode{B,A,O}}}())
            else
                new_belief = update(pomcp.node_belief_updater, h.B, a, o)
                hao = ObsNode(o, 0, new_belief, Dict{A,ActNode{A,O,ObsNode{B,A,O}}}())
            end
            best_node.children[o] = hao
        end

    else
        state_was_generated = false
        # otherwise sample nodes
        os = values(best_node.children)
        wv = WeightVec(Int[node.N for node in os]) # allocation
        hao = sample(sol.rng, os, wv)
        sp = rand(sol.rng, hao.B)
        r = POMDPs.reward(pomcp.problem, s, a, sp)
    end

    if state_was_generated
        if uses_states_from_planner(hao.B)
            push!(hao.B, sp)
        end
        hao.N += 1
    end

    R = r + POMDPs.discount(pomcp.problem)*simulate(pomcp, hao, sp, depth+1)

    best_node.N += 1
    if best_node.V != -Inf
        best_node.V += (R-best_node.V)/best_node.N
    end

    return R
end

make_node{P<:POMCPPlanner, N<:BeliefNode}(::P, b::N) = b

function make_node{S,A,O,B}(pomcp::POMCPPlanner{S,A,O,B}, belief)
    ANodeType = ActNode{A, O, ObsNode{B,A,O}}
    return RootNode{typeof(belief), A, ANodeType}(0, belief, Dict{A,ANodeType}())
end

"""
Add the N's of two nodes - for use in reduce
"""
add_N(a::AbstractActNode, b::AbstractActNode) = a.N + b.N
add_N(a::Int, b::AbstractActNode) = a + b.N


function all_values(abaco, ts, ne, snap)
    ovals = origins_vals(abaco, ne, ts)
    merge(snap.vals, ovals)
end

function deep_count(abaco, ne, domains)
    #@debug "deep_count: $ne, domains: $domains"
    count = 0
    # TODO: filter on tag value
    if length(domains) === 1
        return length(filter(el->el.tag == domains[1], abaco.origins[ne]))
    else
        for node in abaco.origins[ne]
            count += deep_count(abaco, node.ne, domains[2:end])
        end
    end
    count
end

function origins_vals(abaco::Context, ne, ts)
    vals = Dict()
    index = abaco.interval == -1 ? 1 : mvindex(ts, abaco.interval, abaco.ages)

    if haskey(abaco.origins, ne)
        for origin_elem in abaco.origins[ne]
            ropts = span(ts, abaco.interval)
            snap = origin_elem.snap[index]
            # @debug "$ropts --> origin_elem [$(origin_elem.tag).$(origin_elem.ne)]: $snap"

            for (var, val) in all_values(abaco, ts, origin_elem.ne, snap)
                newvar = "$(origin_elem.tag).$var"
                #@debug "[$ne] getting var $newvar"
                if !haskey(vals, newvar)
                    nodes = split(newvar, ".")[1:end-1]
                    value = LValue(deep_count(abaco, ne, nodes), [])
                    vals[newvar] = value
                end
                #@debug "$ne var $newvar: appending $(val.value) (was $(vals[newvar].value))"
                # if interval == -1 considers all values collected at different times
                if (abaco.interval == -1 || ropts === snap.ts)
                    append!(vals[newvar].value, val.value)
                end
            end
        end
    end
    vals
end

function propagate(abaco::Context, snap, ne, formula_name)
    ts = snap.ts
    (etype, target) = abaco.target[ne]
    var = "$etype.$formula_name"
    #@debug "[$ne]: ts:$ts - propagating $var to node [$target]"
    trigger_formulas(abaco, getsnap(abaco, ts, target), target, var)
end

function trigger_formulas(abaco, snap, ne, var::String)
    poll_formulas(abaco, snap, ne, var)

    if haskey(abaco.target, ne)
        propagate(abaco, snap, ne, var)
    end
end

function trigger_formulas(abaco, snap, ne, vars)
    poll_formulas(abaco, snap, ne, vars)

    for var in vars
        # propagate input variables to target
        if haskey(abaco.target, ne)
            propagate(abaco, snap, ne, var)
        end
    end
end

update_inputs(abaco, snap, ne, formula_name, result::Real) = snap_add(snap, ne, formula_name, result)

function update_inputs(abaco, snap, ne, formula_name, result::PValue)
    snap_add(snap, ne, formula_name, result)
    trigger_formulas(abaco, snap, ne, formula_name)
end

#
#     poll_formulas(abaco, snap, ne, vars)
#
# Check the completion status of all formulas that depends on `vars`.
#
function poll_formulas(abaco, snap, ne, vars)
    if snap === nothing
        return
    end
    tag = etype(abaco, ne)
    cfg = snapsetting(abaco, tag)
    formulas = dependents(abaco, tag, vars)

    @debug "[$ne] formulas that depends on [$vars]: $formulas"
    for formula_name in formulas
        fstate = get(snap.outputs, formula_name, FormulaState(false, formula_name))

        # pull up the origins variables
        ovals = origins_vals(abaco, ne, snap.ts)
        allvals = merge(snap.vals, ovals)
        result = poll(cfg, fstate, allvals)

        #@debug "[$ne] allvals: $allvals"
        @debug "[$ne] poll($formula_name) = $result"
        if result !== nothing

            #snap_add(snap, ne, formula_name, result)
            update_inputs(abaco, snap, ne, formula_name, result)

            # propagate to target
            if haskey(abaco.target, ne)
                #@info "propagate $ne --> $formula_name"
                propagate(abaco, snap, ne, formula_name)
            end

            if cfg.oncomplete !== nothing
                if cfg.handle === nothing
                    cfg.oncomplete(snap.ts,
                                   ne,
                                   formula_name,
                                   result,
                                   allvals)
                else
                    cfg.oncomplete(cfg.handle,
                                   snap.ts,
                                   ne,
                                   formula_name,
                                   result,
                                   allvals)
                end
            end
        end
    end
end

function trigger_formulas(abaco, ::Nothing, ne, var, value)
end

#
#     poll(abaco, formula_state::FormulaState, vals::Dict)
#
# Returns the formula value if the formula is computable.
#
# Returns `NaN` when some input variables are missing because
# the formula cannot run to completion.
#
function poll(setting::SnapsSetting, formula_state::FormulaState, vals::Dict)
    formula = setting.formula[formula_state.f]

    poll_simple(setting, formula, formula_state, vals)
end

function poll_simple(setting::SnapsSetting, formula::Formula, formula_state::FormulaState, vals::Dict)
    # Step 1: decide if the formula may be evaluated: all variables are collected
    for v in formula.inputs
        #if !(haskey(vals, v) && isready(vals[v]))
        if !(haskey(vals, v))
            return nothing
        end
    end

    # Step 2: evaluate the formula
    if !formula_state.done
        if setting.emitone
            formula_state.done = true
        end
        result = eval(formula, vals)
        return result
    end
    return nothing
end


using DynamicGrids, DynamicGridsInteract, Test, Colors, ColorSchemes, ImageMagick

# life glider sims

init =  Bool[0 0 0 0 0 0
             0 0 0 0 0 0
             0 0 0 0 0 0
             0 0 0 1 1 1
             0 0 0 0 0 1
             0 0 0 0 1 0]

test3 = [0 0 0 0 0 0
         0 0 0 0 0 0
         0 0 0 0 1 1
         0 0 0 1 0 1
         0 0 0 0 0 1
         0 0 0 0 0 0]

test5 = [0 0 0 0 0 0
         0 0 0 0 0 0
         1 0 0 0 1 1
         1 0 0 0 0 0
         0 0 0 0 0 1
         0 0 0 0 0 0]

processor = ColorProcessor(
    scheme=ColorSchemes.leonardo, zerocolor=nothing, maskcolor=nothing, textconfig=nothing,
)

@testset "InteractOutput" begin
    l0 = ARGB32(get(ColorSchemes.leonardo, 0))
    l1 = ARGB32(get(ColorSchemes.leonardo, 1))

    leonardo2 = [l0 l0 l0 l0 l0 l0
                 l0 l0 l0 l0 l0 l0
                 l1 l0 l0 l0 l1 l1
                 l1 l0 l0 l0 l0 l0
                 l0 l0 l0 l0 l0 l1
                 l0 l0 l0 l0 l0 l0]

    ruleset = Ruleset(Life(); boundary=Wrap())
    output = InteractOutput(init; tspan=1:2, ruleset=ruleset, store=true, processor=processor);
    sim!(output, ruleset)
    sleep(10)
    resume!(output, ruleset; tstop=5)
    sleep(2)

    @test output[1] == init
    @test output[3] == test3
    @test output[5] == test5

    @testset "output image matches colorscheme" begin
        @test output.image_obs[].children.tail[1] == leonardo2
    end

    @testset "output works with store=false" begin
        output = InteractOutput(init; ruleset=ruleset, tspan=1:3, store=false, processor=processor);
        sim!(output, ruleset)
        output.graphicconfig.stoppedframe
        DynamicGrids.stoppedframe(output)
        sleep(10)
        resume!(output, ruleset; tstop=5)
        sleep(2)
        @test output[1] == test5
    end

end

if !Sys.islinux() # No graphic head loaded in CI: TODO add this
    @testset "ElectronOutput" begin
        ruleset = Ruleset(Life(); boundary=Wrap())
        output = ElectronOutput(init; ruleset=ruleset, tspan=1:3, store=true, processor=processor)
        DynamicGrids.setrunning!(output, false)
        sim!(output.interface, ruleset)
        sleep(10)
        DynamicGrids.setrunning!(output, false)
        resume!(output.interface, ruleset; tstop=5)
        sleep(2)
        @test output[3] == test3
        @test output[5] == test5
        close(output.window)
    end
end

@testset "ServerOutput" begin
    ruleset = Ruleset(Life(); boundary=Wrap())
    ServerOutput(init; ruleset=ruleset, port=8080, processor=processor)
    # TODO: test the server somehow
end


export NFSPAgentManager

## definition
"""
    NFSPAgentManager(; agents::Dict{Any, NFSPAgent})

A special MultiAgentManager in which all agents use NFSP policy to play the game.
"""
mutable struct NFSPAgentManager <: AbstractPolicy
    agents::Dict{Any, NFSPAgent}
end

## interactions between the policy and env.
function (π::NFSPAgentManager)(env::AbstractEnv)
    player = current_player(env)
    if player == chance_player(env)
        env |> legal_action_space |> rand
    else
        env |> π.agents[player]
    end
end

RLBase.prob(π::NFSPAgentManager, env::AbstractEnv, args...) = prob(π.agents[current_player(env)], env, args...)

## update NFSPAgentManager
function RLBase.update!(π::NFSPAgentManager, env::AbstractEnv)
    while current_player(env) == chance_player(env)
        env |> legal_action_space |> rand |> env
    end
    update!(π.agents[current_player(env)], env)
end

function (π::NFSPAgentManager)(stage::Union{PreEpisodeStage, PostEpisodeStage}, env::AbstractEnv)
    @sync for (player, agent) in π.agents
        @async agent(stage, env, player)
    end
end


using SFML

# Create the window
window = RenderWindow("Image and Texture Example", 800, 600)
# Set the framerate limit
set_framerate_limit(window, 60)
event = Event()

# Load the texture for the image
texture = Texture(joinpath(dirname(@__FILE__),"..","..","assets","sfmljl_logo.png"))
set_smooth(texture, true)
texture_size = get_size(texture)

# Create the logo sprite and add the texture to it
logo = Sprite()
set_texture(logo, texture)
set_position(logo, Vector2f(400, 300))
set_origin(logo, Vector2f(texture_size.x / 2, texture_size.y / 2))
# Scale it down a bit
scale(logo, Vector2f(0.3, 0.3))

while isopen(window)
    # Check for any events
    while pollevent(window, event)
        if get_type(event) == EventType.CLOSED
            close(window)
        end
    end

    # Rotate the sprite by 2 degrees
    rotate(logo, 2)

    # Clear the screen and redraw the sprite
    clear(window, SFML.white)
    draw(window, logo)
    display(window)
end


# Model of "Asap3" for VPR. Only big difference is that memories occupy the
# whole bottom row, instead of an isthmus of processors.

function asap3_vpr(style = Rectangular(2,1))
    arch = TopLevel{2}("asap3_vpr")

    ####################
    # Normal Processor #
    ####################

    # Get a processor tile and instantiate it.
    processor = build_processor_tile(style)
    for r in 0:29, c in 0:31
        add_child(arch, processor, CartesianIndex(r,c))
    end

	####################
	# Memory Processor #
	####################
	memory_processor = build_processor_tile(style, include_memory = true)
	for r = 30, c = 0:31
        add_child(arch, memory_processor, CartesianIndex(r,c))
	end

	#################
	# 2 Port Memory #
	#################
    memory_2port = build_memory(2)
	for r = 31, c in (0:2:30)
        add_child(arch, memory_2port, CartesianIndex(r,c))
	end

	#################
	# Input Handler #
	#################
	input_handler = build_input_handler(style)
    add_child(arch, input_handler, CartesianIndex(0,-1))

	##################
	# Output Handler #
	##################
	output_handler = build_output_handler(style)
    add_child(arch, output_handler, CartesianIndex(0,32))

	connect_processors(arch, style)
    connect_io(arch, style)
	connect_memories(arch, style)
	return arch
end


function set_wally(nodes, bdcon, wally, cellcenter, cellxmax, cellymax, icell)
	# serch wall
	swith_wall = zeros(4)   # x-, x+, y-, y+を検索する
							# 1:壁
	for i in 1:4
		if Int(bdcon[i][1]) == 2 || Int(bdcon[i][1]) == 7
			swith_wall[i] = 1
		end
	end

	# cal number of ponits
	nop = 0
	if swith_wall[1] == 1
		nop = nop + (cellymax-1)
	end
	if swith_wall[2] == 1
		nop = nop + (cellymax-1)
	end
	if swith_wall[3] == 1
		nop = nop + (cellxmax-1)
	end
	if swith_wall[4] == 1
		nop = nop + (cellxmax-1)
	end

	wallpoint = zeros(nop, 2) # (格子数, 座標)
	ite = 1

	# x-
	if swith_wall[1] == 1
		for j in 1+icell:cellymax-icell
			wallpoint[ite,1] = nodes[icell+1,j,1]
			wallpoint[ite,2] = nodes[icell+1,j,2]
			ite = ite + 1
		end
	end
	# x+
	if swith_wall[2] == 1
		for j in 1+icell:cellymax-icell
			wallpoint[ite,1] = nodes[cellxmax-icell+1,j,1]
			wallpoint[ite,2] = nodes[cellxmax-icell+1,j,2]
			ite = ite + 1
		end
	end
	# y-
	if swith_wall[3] == 1
		for i in 2:cellxmax-1
			wallpoint[ite,1] = nodes[i,icell+1,1]
			wallpoint[ite,2] = nodes[i,icell+1,2]
			ite = ite + 1
		end
	end
	# y+
	if swith_wall[4] == 1
		for i in 2:cellxmax-1
			wallpoint[ite,1] = nodes[i,cellymax-icell+1,1]
			wallpoint[ite,2] = nodes[i,cellymax-icell+1,2]
			ite = ite + 1
		end
	end

	# 距離格納用
	distance_list = zeros(nop,2)
	distance = 0

	# 最大値を2つ抽出
	for j in 1+icell:cellymax-icell
		for i in 1+icell:cellxmax-icell
			for n in 1:nop
				distance = ((cellcenter[i,j,1]-wallpoint[n,1])^2 +
							(cellcenter[i,j,2]-wallpoint[n,2])^2
							)^0.5
				distance_list[n,1] = n
				distance_list[n,2] = distance
			end

			qdistance_list = copy(distance_list)
			qdistance_list = quicksort(qdistance_list,1,nop,2,2)

			# wall_point
			point1 = Int(qdistance_list[1,1])
			point2 = Int(qdistance_list[2,1])

			dis_bottom = ((wallpoint[point1, 1] - wallpoint[point2, 1])^2 + (wallpoint[point1, 2] - wallpoint[point2, 2])^2)^0.5

			# O型格子の関係で同じ点が二つあるため
			if dis_bottom == 0.0
				point1 = Int(qdistance_list[1,1])
				point2 = Int(qdistance_list[3,1])

				dis_bottom = ((wallpoint[point1, 1] - wallpoint[point2, 1])^2 + (wallpoint[point1, 2] - wallpoint[point2, 2])^2)^0.5
			end

			dis1 = distance_list[point1, 2]
			dis2 = distance_list[point2, 2]

			# ヘロンの公式から距離を算出
			s = (dis_bottom + dis1 + dis2)*0.5
			surface = (s*(s-dis_bottom)*(s-dis1)*(s-dis2))^0.5
			wally[i,j] = surface/dis_bottom * 2
		end
	end

	# 仮想セルへの代入
	for i in 1+icell:cellxmax-icell
		for j in 1:icell
			wally[i,j] = wally[i,icell+1]
			wally[i,cellymax-j+1] = wally[i,cellymax-icell]
		end
	end
	for j in 1+icell:cellymax-icell
		for i in 1:icell
			wally[i,j] = wally[icell+1,j]
			wally[cellxmax-i+1,j] = wally[cellxmax-icell,j]
		end
	end

	return wally, swith_wall
end

function cal_yplus(yplus, Qbase, wally, swith_wall, mu, cellxmax, cellymax, vecAx, vecAy, volume, icell)
	tvAxx = 0.0
    tvAxy = 0.0
    Axx = 0.0
    Axy = 0.0
    tvAyx = 0.0
    tvAyy = 0.0
    Ayx = 0.0
    Ayy = 0.0

	if swith_wall[1] == 1
		for j in 1+icell:cellymax-icell
			# 壁に並行な方向の速度を求めるためx-の壁面に対しvecAyを使用
			for i in 1:cellxmax
				tvAyx = 0.5*(vecAy[i,j-1,1]+vecAy[i,j,1])
				tvAyy = 0.5*(vecAy[i,j-1,2]+vecAy[i,j,2])
				Ayx = tvAyx / (tvAyx^2+tvAyy^2)^0.5
				Ayy = tvAyy / (tvAyx^2+tvAyy^2)^0.5
				nu = mu[i,j] / Qbase[i,j,1]
				ubar = abs(Ayx*Qbase[i,j,2] + Ayy*Qbase[i,j,3])
				yplus[i,j], up = cal_wall_val_spalding(ubar, wally[i,j], nu)
			end
		end
	end
	if swith_wall[2] == 1
		for j in 1+:cellymax-icell
			# 壁に並行な方向の速度を求めるためx+の壁面に対しvecAyを使用
			for i in 1:cellxmax
				tvAyx = 0.5*(vecAy[i,j-1,1]+vecAy[i,j,1])
				tvAyy = 0.5*(vecAy[i,j-1,2]+vecAy[i,j,2])
				Ayx = tvAyx / (tvAyx^2+tvAyy^2)^0.5
				Ayy = tvAyy / (tvAyx^2+tvAyy^2)^0.5
				nu = mu[i,j] / Qbase[i,j,1]
				ubar = abs(Ayx*Qbase[i,j,2] + Ayy*Qbase[i,j,3])
				yplus[i,j], up = cal_wall_val_spalding(ubar, wally[i,j], nu)
			end
		end
	end
	if swith_wall[3] == 1
		for j in 1:cellymax
			for i in 1+icell:cellxmax-icell
			# 壁に並行な方向の速度を求めるためy-の壁面に対しvecAxを使用

				tvAxx = 0.5*(vecAx[i-1,j,1]+vecAx[i,j,1])
				tvAxy = 0.5*(vecAx[i-1,j,2]+vecAx[i,j,2])
				Axx = tvAxx / (tvAxx^2+tvAxy^2)^0.5
				Axy = tvAxy / (tvAxx^2+tvAxy^2)^0.5
				nu = mu[i,j] / Qbase[i,j,1]
				ubar = abs(Axx*Qbase[i,j,2] + Axy*Qbase[i,j,3])
				yplus[i,j], up = cal_wall_val_spalding(ubar, wally[i,j], nu)

				if yplus[i,j] > 500
					break
				end
			end
		end
	end
	if swith_wall[4] == 1
		for j in 1:cellymax
			# 壁に並行な方向の速度を求めるためy-の壁面に対しvecAxを使用
			for i in 1+icell:cellxmax-icell
				tvAxx = 0.5*(vecAx[i-1,j,1]+vecAx[i,j,1])
				tvAxy = 0.5*(vecAx[i-1,j,2]+vecAx[i,j,2])
				Axx = tvAxx / (tvAxx^2+tvAxy^2)^0.5
				Axy = tvAxy / (tvAxx^2+tvAxy^2)^0.5
				nu = mu[i,j] / Qbase[i,j,1]
				ubar = abs(Axx*Qbase[i,j,2] + Axy*Qbase[i,j,3])
				yplus[i,j], up = cal_wall_val_spalding(ubar, wally[i,j], nu)
			end
		end
	end

	return yplus
end

function cal_wall_val_spalding(u,y,nu)
	# spalding 定数
	k = 0.4
	B = 5.5

	# ニュートン法
	up = 200.0
	for i in 1:100
		old = up

		up = up - spalding(up,u,y,nu,k,B)/spalding_dash(up,u,y,nu,k,B)

		if abs(up-old)/up < 10^(-4)
			break
		end
	end

	yplus = u*y/nu/up
	return yplus, up
end

function spalding(up,u,y,nu,k,B)
	t = k*up
	F = - 1/up* u*y/nu + up + exp(-k*B)*(exp(t)-1-t-t^2/2-t^3/6)
	return F
end

function spalding_dash(up,u,y,nu,k,B)
	t = k*up
	Fdash = 1/up^2* u*y/nu + 1 + exp(-k*B)*(k*exp(t)-k-k^2*up-k^3*up^2/2)
	return Fdash
end

# test
#=
uu = 6.64037328
yy = 0.06299976860
nunu = 1.406107133e-5
println(cal_wall_val_spalding(uu,yy,nunu))
throw(UndefVarError(:x))
=#

function wallf_Van_Driest(yplus)
	f = 1 - exp(-yplus/26)
	return f
end


const ISO4217 = Dict{Symbol, Tuple{Int, String, Int}}(
    :AED => (2,"UAE Dirham",784),
    :AFN => (2,"Afghani",971),
    :ALL => (2,"Lek",8),
    :AMD => (2,"Armenian Dram",51),
    :ANG => (2,"Netherlands Antillean Guilder",532),
    :AOA => (2,"Kwanza",973),
    :ARS => (2,"Argentine Peso",32),
    :AUD => (2,"Australian Dollar",36),
    :AWG => (2,"Aruban Florin",533),
    :AZN => (2,"Azerbaijanian Manat",944),
    :BAM => (2,"Convertible Mark",977),
    :BBD => (2,"Barbados Dollar",52),
    :BDT => (2,"Taka",50),
    :BGN => (2,"Bulgarian Lev",975),
    :BHD => (3,"Bahraini Dinar",48),
    :BIF => (0,"Burundi Franc",108),
    :BMD => (2,"Bermudian Dollar",60),
    :BND => (2,"Brunei Dollar",96),
    :BOB => (2,"Boliviano",68),
    :BOV => (2,"Mvdol",984),
    :BRL => (2,"Brazilian Real",986),
    :BSD => (2,"Bahamian Dollar",44),
    :BTN => (2,"Ngultrum",64),
    :BWP => (2,"Pula",72),
    :BYR => (0,"Belarussian Ruble",974),
    :BZD => (2,"Belize Dollar",84),
    :CAD => (2,"Canadian Dollar",124),
    :CDF => (2,"Congolese Franc",976),
    :CHE => (2,"WIR Euro",947),
    :CHF => (2,"Swiss Franc",756),
    :CHW => (2,"WIR Franc",948),
    :CLF => (4,"Unidad de Fomento",990),
    :CLP => (0,"Chilean Peso",152),
    :CNY => (2,"Yuan Renminbi",156),
    :COP => (2,"Colombian Peso",170),
    :COU => (2,"Unidad de Valor Real",970),
    :CRC => (2,"Costa Rican Colon",188),
    :CUC => (2,"Peso Convertible",931),
    :CUP => (2,"Cuban Peso",192),
    :CVE => (2,"Cabo Verde Escudo",132),
    :CZK => (2,"Czech Koruna",203),
    :DJF => (0,"Djibouti Franc",262),
    :DKK => (2,"Danish Krone",208),
    :DOP => (2,"Dominican Peso",214),
    :DZD => (2,"Algerian Dinar",12),
    :EGP => (2,"Egyptian Pound",818),
    :ERN => (2,"Nakfa",232),
    :ETB => (2,"Ethiopian Birr",230),
    :EUR => (2,"Euro",978),
    :FJD => (2,"Fiji Dollar",242),
    :FKP => (2,"Falkland Islands Pound",238),
    :GBP => (2,"Pound Sterling",826),
    :GEL => (2,"Lari",981),
    :GHS => (2,"Ghana Cedi",936),
    :GIP => (2,"Gibraltar Pound",292),
    :GMD => (2,"Dalasi",270),
    :GNF => (0,"Guinea Franc",324),
    :GTQ => (2,"Quetzal",320),
    :GYD => (2,"Guyana Dollar",328),
    :HKD => (2,"Hong Kong Dollar",344),
    :HNL => (2,"Lempira",340),
    :HRK => (2,"Kuna",191),
    :HTG => (2,"Gourde",332),
    :HUF => (2,"Forint",348),
    :IDR => (2,"Rupiah",360),
    :ILS => (2,"New Israeli Sheqel",376),
    :INR => (2,"Indian Rupee",356),
    :IQD => (3,"Iraqi Dinar",368),
    :IRR => (2,"Iranian Rial",364),
    :ISK => (0,"Iceland Krona",352),
    :JMD => (2,"Jamaican Dollar",388),
    :JOD => (3,"Jordanian Dinar",400),
    :JPY => (0,"Yen",392),
    :KES => (2,"Kenyan Shilling",404),
    :KGS => (2,"Som",417),
    :KHR => (2,"Riel",116),
    :KMF => (0,"Comoro Franc",174),
    :KPW => (2,"North Korean Won",408),
    :KRW => (0,"Won",410),
    :KWD => (3,"Kuwaiti Dinar",414),
    :KYD => (2,"Cayman Islands Dollar",136),
    :KZT => (2,"Tenge",398),
    :LAK => (2,"Kip",418),
    :LBP => (2,"Lebanese Pound",422),
    :LKR => (2,"Sri Lanka Rupee",144),
    :LRD => (2,"Liberian Dollar",430),
    :LSL => (2,"Loti",426),
    :LYD => (3,"Libyan Dinar",434),
    :MAD => (2,"Moroccan Dirham",504),
    :MDL => (2,"Moldovan Leu",498),
    :MGA => (2,"Malagasy Ariary",969),
    :MKD => (2,"Denar",807),
    :MMK => (2,"Kyat",104),
    :MNT => (2,"Tugrik",496),
    :MOP => (2,"Pataca",446),
    :MRO => (2,"Ouguiya",478),
    :MUR => (2,"Mauritius Rupee",480),
    :MVR => (2,"Rufiyaa",462),
    :MWK => (2,"Kwacha",454),
    :MXN => (2,"Mexican Peso",484),
    :MXV => (2,"Mexican Unidad de Inversion (UDI)",979),
    :MYR => (2,"Malaysian Ringgit",458),
    :MZN => (2,"Mozambique Metical",943),
    :NAD => (2,"Namibia Dollar",516),
    :NGN => (2,"Naira",566),
    :NIO => (2,"Cordoba Oro",558),
    :NOK => (2,"Norwegian Krone",578),
    :NPR => (2,"Nepalese Rupee",524),
    :NZD => (2,"New Zealand Dollar",554),
    :OMR => (3,"Rial Omani",512),
    :PAB => (2,"Balboa",590),
    :PEN => (2,"Nuevo Sol",604),
    :PGK => (2,"Kina",598),
    :PHP => (2,"Philippine Peso",608),
    :PKR => (2,"Pakistan Rupee",586),
    :PLN => (2,"Zloty",985),
    :PYG => (0,"Guarani",600),
    :QAR => (2,"Qatari Rial",634),
    :RON => (2,"Romanian Leu",946),
    :RSD => (2,"Serbian Dinar",941),
    :RUB => (2,"Russian Ruble",643),
    :RWF => (0,"Rwanda Franc",646),
    :SAR => (2,"Saudi Riyal",682),
    :SBD => (2,"Solomon Islands Dollar",90),
    :SCR => (2,"Seychelles Rupee",690),
    :SDG => (2,"Sudanese Pound",938),
    :SEK => (2,"Swedish Krona",752),
    :SGD => (2,"Singapore Dollar",702),
    :SHP => (2,"Saint Helena Pound",654),
    :SLL => (2,"Leone",694),
    :SOS => (2,"Somali Shilling",706),
    :SRD => (2,"Surinam Dollar",968),
    :SSP => (2,"South Sudanese Pound",728),
    :STD => (2,"Dobra",678),
    :SVC => (2,"El Salvador Colon",222),
    :SYP => (2,"Syrian Pound",760),
    :SZL => (2,"Lilangeni",748),
    :THB => (2,"Baht",764),
    :TJS => (2,"Somoni",972),
    :TMT => (2,"Turkmenistan New Manat",934),
    :TND => (3,"Tunisian Dinar",788),
    :TOP => (2,"Pa’anga",776),
    :TRY => (2,"Turkish Lira",949),
    :TTD => (2,"Trinidad and Tobago Dollar",780),
    :TWD => (2,"New Taiwan Dollar",901),
    :TZS => (2,"Tanzanian Shilling",834),
    :UAH => (2,"Hryvnia",980),
    :UGX => (0,"Uganda Shilling",800),
    :USD => (2,"US Dollar",840),
    :USN => (2,"US Dollar (Next day)",997),
    :UYI => (0,"Uruguay Peso en Unidades Indexadas (URUIURUI)",940),
    :UYU => (2,"Peso Uruguayo",858),
    :UZS => (2,"Uzbekistan Sum",860),
    :VEF => (2,"Bolivar",937),
    :VND => (0,"Dong",704),
    :VUV => (0,"Vatu",548),
    :WST => (2,"Tala",882),
    :XAF => (0,"CFA Franc BEAC",950),
    :XAG => (-1,"Silver",961),
    :XAU => (-1,"Gold",959),
    :XBA => (-1,"Bond Markets Unit European Composite Unit (EURCO)",955),
    :XBB => (-1,"Bond Markets Unit European Monetary Unit (E.M.U.-6)",956),
    :XBC => (-1,"Bond Markets Unit European Unit of Account 9 (E.U.A.-9)",957),
    :XBD => (-1,"Bond Markets Unit European Unit of Account 17 (E.U.A.-17)",958),
    :XCD => (2,"East Caribbean Dollar",951),
    :XDR => (-1,"SDR (Special Drawing Right)",960),
    :XOF => (0,"CFA Franc BCEAO",952),
    :XPD => (-1,"Palladium",964),
    :XPF => (0,"CFP Franc",953),
    :XPT => (-1,"Platinum",962),
    :XSU => (-1,"Sucre",994),
    :XTS => (-1,"TESTING",963),
    :XUA => (-1,"ADB Unit of Account",965),
    :XXX => (-1,"NO CURRENCY",999),
    :YER => (2,"Yemeni Rial",886),
    :ZAR => (2,"Rand",710),
    :ZMW => (2,"Zambian Kwacha",967),
    :ZWL => (2,"Zimbabwe Dollar",932))


abstract type AbstractSearchQuery end

struct CASQuery <: AbstractSearchQuery
    ID::String
    cas::Tuple{Int32,Int16,Int16}
end

function CASQuery(str::String)
    cas = cas_parse(str)
    return CASQuery(str,cas)
end

function CASQuery(val::NTuple{3,Integer})
    n1 = val[1]
    n2 = val[2]
    n3 = val[3]
    @assert n1 >=0
    @assert n2 >=0
    @assert n3 >=0
    v1 = convert(Int32,n1)
    v2 = convert(Int16,n2)
    v3 = convert(Int16,n3)
    cas = (v1,v2,v3)
    return CASQuery("",cas)
end

function tuple_to_casstr(n1,n2,n3)

    return string(n1) * "-" * string(n2) * "-" * string(n3)
end

tuple_to_casstr(val::NTuple{3,Integer}) = tuple_to_casstr(val...)
cas(id::CASQuery) = id.cas

struct InChIQuery <: AbstractSearchQuery
    ID::String
end

struct InChIKeyQuery <: AbstractSearchQuery
    ID::String
end

struct PubChemIDQuery <: AbstractSearchQuery
    ID::String
    val::Int
end

function PubChemIDQuery(str::String)
    val = parse(Int,str)
    @assert val >0
    return PubChemIDQuery(str,val)
end

function PubChemIDQuery(val::Integer)
    @assert val >0
    return PubChemIDQuery("",val)
end

id_num(ID::PubChemIDQuery) = ID.val

struct SMILESQuery <: AbstractSearchQuery
    ID::String
end

struct FormulaQuery <: AbstractSearchQuery
    ID::String
end

struct ElementQuery <: AbstractSearchQuery
    ID::String
end

struct AnyQuery <: AbstractSearchQuery
    ID::String
end

value(id::AbstractSearchQuery)::String = strip(id.ID)

function value(id::SMILESQuery)::String
    id = strip(id.ID)
    return replace(id,"SMILES=" =>"")
end

function value(id::InChIQuery)::String
    id_raw = strip(id.ID)
    id_lower = lowercase(id_raw)
    re1 = r"^inchi=1s/"
    re2 =  r"^inchi=1/"
    t1 = occursin(re1,id_lower)
    t2 = occursin(re2,id_lower)
    if t1
        return chop(id_raw,head=9,tail=0)
    elseif t2
        return chop(id_raw,head=8,tail=0)
    else
        throw("incorrect InChI passed")
    end
end

"""
    is_cas(str)::Bool
check if a given string is a cas number and returns true or false accordingly.

"""
function is_cas(str)
    str = strip(str)
    #regex from https://gist.github.com/KhepryQuixote/00946f2f7dd5f89324d8#file-pypubchemxtractor-py-L22
    cas_regex = r"^[1-9][0-9]{1,6}\\-[0-9]{2}\\-[0-9]$"
    return occursin(cas_regex,str)
end

"""
    is_inchi(str)::Bool
check if a given string is an InChI name and returns true or false accordingly.

"""
function is_inchi(str)
    str = strip(str)
    #regex from https://chemistry.stackexchange.com/a/86892
    inchi_regex = r"^InChI\=1S?\/[^\s]+(\s|$)"
    return occursin(inchi_regex,str)
end

"""
    is_inchikey(str)::Bool
check if a given string is an InChI key and returns true or false accordingly.

"""
function is_inchikey(str)
    str = strip(str)
    #14(A-Z)-10(A-Z)-1(A-Z)
    inchikey_regex  =r"[A-Z]{14}-[A-Z]{10}-[A-Z]{1}"
    return occursin(inchikey_regex,str)
end

"""
    is_element(str)::Bool
check if a given string is an element symbol and returns true or false accordingly.

"""
function is_element(str)
    return false
end

"""
    is_pubchemid(str)::Bool
check if a given string is a PubChem ID and returns true or false accordingly.
A PubChem ID is a positive integer.
"""
function is_pubchemid(str)
    str = strip(str)
    res = tryparse(Int,str)
   if res === nothing
    return false
   elseif res <= 0
    return false
   else
    return true
   end
end

"""
    is_smiles(str)::Bool
check if a given string is an element symbol and returns true or false accordingly.

"""
function is_smiles(str)
    str = strip(str)
    str = uppercase(str)
    smiles_regex  =r"^SMILES="
    return occursin(smiles_regex,str)
end


# run the algorithm to test 2 layer boostNet

dataset = "set5" # "set2", "set5"
train_data, train_label, validation_data, validation_label, test_data, test_label =
  share(load_data(dataset))

model = boostNet(train_data, train_label', m=2, T=1000, B=10)

say()
p = predict(model, train_data)
say("train error: ", mean(p .!= train_label'))

p = predict(model, test_data)
say("test error: ", mean(p .!= test_label'))

say("love huifang")


module DungBase

using Dates, UUIDs, StaticArrays
import IntervalSets: width, (..), AbstractInterval, leftendpoint
using StatsBase, StructArrays, Dierckx, AngleBetweenVectors, LinearAlgebra

export VideoFile, WholeVideo, FragmentedVideo, DisjointVideo, AbstractTimeLine, AbstractPeriod, Instantaneous, Prolonged, Temporal, Board, Calibration, POI, Metadata, Run, Experiment, Track, Common, Point, PointCollection

export start, duration, stop, files, filenames, point, pointcollection, homing, searching, searchcenter, turningpoint

stop(x) = start(x) + duration(x)

include("videos.jl")

include("temporals.jl")

include("calibrations.jl")

include("pois.jl")

include("track.jl")

# include(joinpath(@__DIR__, "experimental_setup.jl"))

struct Metadata
    setup::Dict{Symbol, Any}
    comment::String
    date::DateTime
end

# Base.getproperty(x::Metadata, level::Symbol) = get(x.setup, level, missing)
#=getcoord(x::Run, poi::Symbol) = get(x.pois, poi, missing)
_getlevel(x::Metadata, level::Symbol) = get(x, level, missing)
getlevel(x::Run, level::Symbol) = _getlevel(x.metadata, level)=#

struct Run{T, M}
    data::T # pois::Dict{Symbol, T} or Track
    metadata::M
end

#=struct Run
    setup::Dict{Symbol, Any}
    pois::Dict{Symbol, POI}
    comment::String
    date::DateTime
end=#

struct Experiment
    runs::Vector{Run}
    description::String
end

end # module


# XXX create optimization model after anyModel has been initialized
"""
Create constraints of the model's underlying optimization problem.
"""
function createOptModel!(anyM::anyModel)

	# <editor-fold desc="create technology related variables and constraints"

	techIdx_arr = collect(keys(anyM.parts.tech))
    parDef_dic = defineParameter(anyM.options,anyM.report)

    # XXX get dimension of expansion and capacity variables and mapping of capacity constraints
    tsYear_dic = Dict(zip(anyM.supTs.step,collect(0:anyM.options.shortExp:(length(anyM.supTs.step)-1)*anyM.options.shortExp)))
    prepVar_dic = Dict{Int,Dict{Symbol,NamedTuple}}()
    prepareTechs!(techIdx_arr,prepVar_dic,tsYear_dic,anyM)
    if any(getindex.(anyM.report,1) .== 3) print(getElapsed(anyM.options.startTime)); errorTest(anyM.report,anyM.options) end

    # remove technologies without any potential capacity variables
    techIdx_arr = collect(keys(prepVar_dic))
    foreach(x -> delete!(anyM.parts.tech, x),setdiff(collect(keys(anyM.parts.tech)),techIdx_arr))

    # XXX create all technology related elements

    # creates dictionary that assigns combination of superordinate dispatch timestep and dispatch level to dispatch timesteps
    allLvlTsDis_arr = unique(getfield.(values(anyM.cInfo),:tsDis))
    ts_dic = Dict((x[1], x[2]) => anyM.sets[:Ts].nodes[x[1]].lvl == x[2] ? [x[1]] : getDescendants(x[1],anyM.sets[:Ts],false,x[2]) for x in Iterators.product(anyM.supTs.step,allLvlTsDis_arr))

    # creates dictionary that assigns combination of expansion region and dispatch level to dispatch region
    allLvlR_arr = union(getindex.(getfield.(getfield.(values(anyM.parts.tech),:balLvl),:exp),2),map(x -> x.rDis,values(anyM.cInfo)))

    allRExp_arr = union([getfield.(getNodesLvl(anyM.sets[:R],x),:idx) for x in allLvlR_arr]...)
    r_dic = Dict((x[1], x[2]) => (anyM.sets[:R].nodes[x[1]].lvl <= x[2] ? getDescendants(x[1], anyM.sets[:R],false,x[2]) : getAncestors(x[1],anyM.sets[:R],:int,x[2])[end]) |> (z -> typeof(z) <: Array ? z : [z]) for x in Iterators.product(allRExp_arr,allLvlR_arr))

    produceMessage(anyM.options,anyM.report, 3," - Determined dimension of expansion and capacity variables for technologies")

    # constraints for technologies are prepared in threaded loop and stored in an array of dictionaries
	techCnsDic_arr = Array{Dict{Symbol,cnsCont}}(undef,length(techIdx_arr))
	tech_itr = collect(enumerate(techIdx_arr))

	@threads for (idx,t) in tech_itr
		techCnsDic_arr[idx] = createTech!(t,anyM.parts.tech[t],prepVar_dic[t],copy(parDef_dic),ts_dic,r_dic,anyM)
	end

    # loops over array of dictionary with constraint container for each technology to create actual jump constraints
    for (idx,cnsDic) in enumerate(techCnsDic_arr), cnsSym in keys(cnsDic)
        anyM.parts.tech[techIdx_arr[idx]].cns[cnsSym] = createCns(cnsDic[cnsSym],anyM.optModel)
    end
    produceMessage(anyM.options,anyM.report, 1," - Created variables and constraints for all technologies")

	# </editor-fold>

	# <editor-fold desc="create exchange related variables and constraints"

	prepExc_dic = Dict{Symbol,NamedTuple}()
	partExc = anyM.parts.exc
	partLim = anyM.parts.lim

	# obtain dimensions of expansion variables for exchange
	potExc_df = prepareExcExpansion!(partExc,partLim,prepExc_dic,tsYear_dic,anyM)

	# obtain capacity dimensions solely based on expansion variables
	prepareCapacity!(partExc,prepExc_dic,prepExc_dic[:expExc].var,:capaExc,anyM)
	addResidualCapaExc!(partExc,prepExc_dic,potExc_df,anyM)

	if !all(map(x -> isempty(x),values(prepExc_dic[:capaExc])))
		# create expansion and capacity variables
		createExpCap!(partExc,prepExc_dic,anyM)
		# create capacity constraint
		createCapaExcCns!(partExc,anyM)
		produceMessage(anyM.options,anyM.report, 2," - Created all variables and constraints related to expansion and capacity for exchange")
		# create dispatch related variables
		createExcVar!(partExc,ts_dic,anyM)
		produceMessage(anyM.options,anyM.report, 2," - Created all dispatch variables for exchange")
		# create capacity restrictions
		createRestrExc!(ts_dic,partExc,anyM)
		produceMessage(anyM.options,anyM.report, 2," - Created all capacity restrictions for exchange")
		produceMessage(anyM.options,anyM.report, 1," - Created variables and constraints for exchange")
	end

	# </editor-fold>

	createTradeVarCns!(anyM.parts.trd,anyM)
	createEnergyBal!(techIdx_arr,anyM)
	createLimitCns!(techIdx_arr,anyM.parts.lim,anyM)

	produceMessage(anyM.options,anyM.report, 1," - Completed model creation")
end


# Recreates the analysis from the original data files

datapath = "dataset"
projectname = "MEG3"
experimentdirnames = ["FKBP12-mTOR"]

samplenames1 = ["E", "G"]
samplenames23 = ["E", "F"]
samplenames4 = ["C", "F"]
samplenames5 = ["A", "B"]

nreplicates = 5
ncells = 10

mc_iterations = 10000

outputdir = "output"
mkpath(outputdir)
outputdatapath = joinpath(outputdir, "resultspos.jld2")

using Distributed
currentworkers = addprocs(exeflags = "--project")
@everywhere using SMLMAssociationAnalysis_NCB
using Printf
using LocalizationMicroscopy
using Statistics
using FileIO

experimentresults = Vector{Vector{Vector{Result}}}[]
for experimentdirname ∈ experimentdirnames
    println("Starting experiment $experimentdirname.")
    experimentpath = joinpath(datapath, experimentdirname)
    replicateresults = Vector{Vector{Result}}[]
    for i ∈ 1:nreplicates
        sampleresults = Vector{Result}[]
        println("    Starting replicate $i.")
        replicatepath = joinpath(experimentpath, "Replicate $i")
        samplenames = i == 5 ? samplenames5 :
                      i == 4 ? samplenames4 :
                      i == 1 ? samplenames1 :
                               samplenames23
        for samplename ∈ samplenames
            results = Result[]
            println("        Starting sample $samplename.")
            for j ∈ 1:ncells
                jcell = j
                if i == 1 && samplename == "E"
                    jcell += 10
                end

                println("            Starting cell $jcell.")
                cellpath = joinpath(replicatepath, "$samplename $(Printf.@sprintf("%03i", jcell)).bin.txt")
                localizations = loadlocalizations(cellpath, LocalizationMicroscopy.nikonelementstext)
                ch1_name = i == 2 ? "488" : "647"
                ch2_name = i == 2 ? "647" : "488"  # the channels were reversed for replicate 2. Unsure if this caused other issues.

                ch1_startframe = 1
                ch2_startframe = 11001

                ch1_task = remotecall(getmolecules, currentworkers[1], localizations,
                    ch1_name,
                    ch1_startframe,
                    11000,
                    100,
                    10,
                    34.2,
                    500,
                    200,
                )
                ch2_task = remotecall(getmolecules, currentworkers[2], localizations,
                    ch2_name,
                    ch2_startframe,
                    11000,
                    100,
                    10,
                    34.2,
                    500,
                    200,
                )

                if i == 2
                    ch1_task, ch2_task = ch2_task, ch1_task
                end

                ch1_molecules, ch1_localizations = fetch(ch1_task)
                ch2_molecules, ch2_localizations = fetch(ch2_task)

                ch1_neighbors, ch2_neighbors, distances = exclusivenearestneighbors(ch1_molecules, ch2_molecules)

                percentileranks = montecarloaffinity(
                    ch1_molecules,
                    ch2_molecules,
                    ch1_neighbors,
                    ch2_neighbors,
                    distances,
                    800,
                    mc_iterations
                )

                positivecontrol_percentileranks, positivecontrol_distances = simulate100(ch1_molecules, ch2_molecules, ch1_neighbors, ch2_neighbors, 80, 800, mc_iterations)
                # I found that the variance in running the negative control drops smoothly as a function of neighbor counts.
                # This heuristic cuts down on computational time without sacrificing quality of the correction.
                negative_iterations = length(distances) < 250 ? 30 :
                                      length(distances) < 500 ? 20 :
                                      length(distances) < 1000 ? 10 :
                                      length(distances) < 2000 ? 5 :
                                      1
                negativecontrol = [simulate0(ch1_molecules, ch2_molecules, 800, 10000) for i in 1:negative_iterations]
                negativecontrol_percentileranks = [first(i) for i in negativecontrol]
                negativecontrol_distances = [last(i) for i in negativecontrol]

                mediandistance = length(distances) == 0 ? NaN : median(distances)

                println("                Done: $(length(distances)) neighbors from $(length(ch1_molecules)) and $(length(ch2_molecules)) molecules, $(length(ch1_localizations)) and $(length(ch2_localizations)) localizations; median distance $mediandistance")
                ch1_data = ChannelData(ch1_name, ch1_molecules, ch1_neighbors)
                ch2_data = ChannelData(ch2_name, ch2_molecules, ch2_neighbors)
                result = Result(projectname, experimentdirname, 1, cellpath, i,
                                [ch1_data, ch2_data], distances, mediandistance, percentileranks,
                                positivecontrol_distances, negativecontrol_distances,
                                positivecontrol_percentileranks, negativecontrol_percentileranks)

                push!(results, result)
            end
            push!(sampleresults, results)
        end
        push!(replicateresults, sampleresults)
    end
    push!(experimentresults, replicateresults)
end

save(outputdatapath, "experimentresults", experimentresults)

rmprocs(currentworkers)


using Flux

struct Path{T}
    sequence::Array{T}
    log_probability::AbstractFloat
    score::AbstractFloat
end

function expand(
    path::Path{T},
    predict::Function,
    vocabulary::AbstractVector{T},
    length_normalization::AbstractFloat
)::AbstractVector{Path{T}} where T
    # Compute probabilities for next token.
    log_word_probabilities = predict(path.sequence |> gpu)

    # Combine with own log probability.
    # Addition because log(p1 * p2) = log(p1) + log(p2).
    log_probabilities = log_word_probabilities .+ path.log_probability

    # Compute length penalty for normalization.
    α = length_normalization
    length_penalty = ((5 + (length(path.sequence) + 1))^α) / ((5 + 1)^α)

    # Compute scores.
    scores = log_probabilities ./ length_penalty

    # Return expanded paths.
    return [
        Path(
            [path.sequence..., vocabulary[i]],
            log_probabilities[i],
            scores[i],
        )
        for i ∈ 1:length(vocabulary)
    ]
end

function expand(
    path::Path{T},
    predict::Function,
    vocabulary::AbstractVector{T},
    expandable::Function,
    length_normalization::AbstractFloat
)::AbstractVector{Path{T}} where T
    if expandable(path.sequence)
        # Expand next token.
        expand(path, predict, vocabulary, length_normalization)
    else
        # Otherwise, return as is.
        [path]
    end
end

function has_redundant_trigrams(sequence::AbstractVector)::Bool
    if length(sequence) <= 3
        return false
    end
    trigram = sequence[end - 2:end]
    for i ∈ 1:length(sequence) - 3
        if sequence[i:i + 2] == trigram
            @debug "Sequence $sequence has redundant trigram $trigram at position $i."
            return true
        end
    end
    return false
end

has_redundant_trigrams(path::Path)::Bool = has_redundant_trigrams(path.sequence)

has_no_redundant_trigrams(path::Path)::Bool = !has_redundant_trigrams(path.sequence)

function block_redundant_trigrams!(
    paths::AbstractVector{Path{T}}
)::AbstractVector{Path{T}} where T
    return filter!(has_no_redundant_trigrams, paths)
end

# Beam search for generating sequences.
function beam_search(
    width::Integer,
    vocabulary::AbstractVector{T},
    # Function returning a column vector of log probabilities for each word
    # in the vocabulary, given the current sequence.
    predict::Function,
    # Predicate to check whether a sequence is expandable.
    # One might for example check for an end token.
    expandable::Function;
    initial_sequence::AbstractVector{T}=[],
    length_normalization::AbstractFloat=0.0,
)::AbstractVector{T} where T
    # Best current path.
    best_path::Path{T} = Path(initial_sequence, 1.0, -Inf)
    # Paths to be considered as start for a single iteration.
    # Start with one initial path, the empty sequence.
    paths::AbstractVector{Path{T}} = [best_path]

    # Expand iteratively until no expandable path is left.
    i = 1
    while length(paths) > 0 && any(path -> expandable(path.sequence), paths)
        @debug "Predicting next token for iteration $i."

        # Calculate paths (hypotheses and probabilities).
        next_paths::AbstractVector{Path{T}} = vcat(
            [
                expand(path, predict, vocabulary, expandable, length_normalization)
                for path ∈ paths
            ]...
        )

        # Block redundant trigrams to avoid repetition.
        # TODO Make configurable.
        # TODO Paulus et al. (2017, p.4) apply this only on the test set.
        block_redundant_trigrams!(next_paths)

        # Sort paths by descending score.
        sort!(next_paths, by=path -> path.score, rev=true)

        # Select best paths.
        paths = next_paths[1:min(length(next_paths), width)]

        local_best_path::Path{T} = first(paths)
        @info "Found best sequence for iteration $i with probability $(exp(local_best_path.log_probability)) (score $(local_best_path.score))."

        if local_best_path.score > best_path.score
            best_path = local_best_path
        end

        i += 1
    end

    @info "Found best sequence with probability $(exp(best_path.log_probability)) (score $(best_path.score))."
    return best_path.sequence
end

# Beam search for generating sequences.
function beam_search(
    width::Integer,
    vocabulary::AbstractVector{T},
    # Function returning a column vector of log probabilities for each word
    # in the vocabulary, given the current sequence.
    predict::Function,
    # Max steps to expand.
    steps::Integer;
    initial_sequence::AbstractVector{T}=[],
    length_normalization::AbstractFloat=0.0,
)::AbstractVector{T} where T
    beam_search(
        width,
        vocabulary,
        predict,
        sequence -> length(sequence) < length(initial_sequence) + steps,
        initial_sequence=initial_sequence,
        length_normalization=length_normalization,
    )
end

# Beam search for generating sequences.
function beam_search(
    width::Integer,
    vocabulary::AbstractVector{T},
    # Function returning a column vector of log probabilities for each word
    # in the vocabulary, given the current sequence.
    predict::Function,
    # Predicate to check whether a sequence is expandable.
    # One might for example check for an end token.
    expandable::Function,
    # Max steps to expand.
    steps::Integer;
    initial_sequence::AbstractVector{T}=[],
    length_normalization::AbstractFloat=0.0,
)::AbstractVector{T} where T
    beam_search(
        width,
        vocabulary,
        predict,
        sequence -> expandable(sequence) && (length(sequence) < length(initial_sequence) + steps),
        initial_sequence=initial_sequence,
        length_normalization=length_normalization,
    )
end

# Like beam search, but only the single locally best path is considered.
function greedy_search(
    vocabulary::AbstractVector{T},
    # Function returning a column vector of log probabilities for each word
    # in the vocabulary, given the current sequence and input.
    predict::Function,
    # Predicate to check whether a sequence is expandable.
    # One might for example check for an end token.
    expandable::Function;
    initial_sequence::AbstractVector{T}=[],
)::AbstractVector{T} where T
    beam_search(
        1,
        vocabulary,
        predict,
        expandable,
        initial_sequence=initial_sequence,
    )
end

# Like beam search, but only the single locally best path is considered.
function greedy_search(
    vocabulary::AbstractVector{T},
    # Function returning a column vector of log probabilities for each word
    # in the vocabulary, given the current sequence and input.
    predict::Function,
    # Max steps to expand.
    steps::Integer;
    initial_sequence::AbstractVector{T}=[]
)::AbstractVector{T} where T
    beam_search(
        1,
        vocabulary,
        predict,
        steps,
        initial_sequence=initial_sequence,
    )
end

# Like beam search, but only the single locally best path is considered.
function greedy_search(
    vocabulary::AbstractVector{T},
    # Function returning a column vector of log probabilities for each word
    # in the vocabulary, given the current sequence and input.
    predict::Function,
    # Predicate to check whether a sequence is expandable.
    # One might for example check for an end token.
    expandable::Function,
    # Max steps to expand.
    steps::Integer;
    initial_sequence::AbstractVector{T}=[]
)::AbstractVector{T} where T
    beam_search(
        1,
        vocabulary,
        predict,
        expandable,
        steps,
        initial_sequence=initial_sequence,
    )
end


module materials
using ..Types
export dummyMat
#penetrator
steel_4140 = MaterialPenetrator{ST}(Properties(293, 0.283, 24.0e6, 30.0e6, 17117, 166717, 146500, 15079, 2023, 125958, 77201, 96254, 1168, 1.0))
steel_1018 = MaterialPenetrator{ST}(Properties(200, 0.283, 24.0e6, 30.0e6, 17117, 113800, 100000, 15079, 1672, 79524, 54693, 65702, 965, 1.0))
steel_4130 = MaterialPenetrator{ST}(Properties(300, 0.283, 24.0e6, 30.0e6, 17117, 170700, 150000, 15079, 2047, 129154, 78828, 98554, 1182, 1.0))
aluminium_2024 = MaterialPenetrator{AL}(Properties(120, 0.100, 10.0e6, 10.0e6, 18024, 68280, 60000, 16374, 2178, 53288, 40572, 39421, 1258, 1.0))
tungsten_alloy_GTE = MaterialPenetrator{WF}(Properties(285, 0.639, 41.667e6, 52.5e6, 15804, 162165, 142500,13222, 1328, 85.0e3, 49075, 93626, 767, 1.0))
tungsten_alloy_WHA_1597 = MaterialPenetrator{WF}(Properties(264, 0.625, 39.68e6, 50.0e6, 15700, 150216, 132000, 13047, 1292, 86.0e3, 49652,86727, 746, 1.0))
tungsten_alloy_WHA_1598 = MaterialPenetrator{WF}(Properties(280, 0.620, 39.68e6, 50.0e6, 15664, 159320, 140000, 13099, 1263, 119.0e3, 68705, 91983, 439, 1.2))
depleted_uranium = MaterialPenetrator{DU}(Properties(231, 0.671, 14.5e6, 24.0e6, 10500, 131439, 115500, 7612, 1167, 140.0e3, 80829, 75886, 674, 1.0))
tantalum = MaterialPenetrator{TA}(Properties(255, 0.602, 30.0e6, 27.0e6, 12500, 145095, 127500, 11559, 1294, 40.0e3, 23094, 83770, 747, 1.0))
titanium_6AI = MaterialPenetrator{TI}(Properties(285, 0.163, 19.2e6, 16.0e6, 19002, 162165, 142500, 17771, 2629, 120.0e3, 69282, 93626, 1518, 1.0))
Aluminium_teflon = MaterialPenetrator{AT}(Properties(6, 0.082, 2.69e6, 2.69e6, 13127, 3414, 3000, 9378, 538, 1500, 866, 1971, 311, 1.0))

#target
steel = MetallicParts{ST}(Properties(150, 0.283, 24.0e6, 30.0e6, 17117, 85350, 75000, 15079, 1448, 51507, 41911, 49277, 836, 1.0))
face_hardened_steel = MetallicParts{FH}(Properties(530, 0.283, 24.0e6, 30.0e6, 17117, 301570, 265000, 15079, 2721, 210894, 127105, 174112, 1571, 1.0))
aluminium = MetallicParts{AL}(Properties(120, 0.100, 10.0e6, 10.6e6, 18024, 68280, 60000, 16374, 2178, 53288, 40572, 39421, 1258, 1.0))
titanium = MetallicParts{TI}(Properties(285, 0.163, 19.2e6, 16.0e6, 19002, 162165, 142500, 17771, 2629, 120.0e3, 69282, 93626, 1518, 1.0))
magnesium = MetallicParts{MG}(Properties(55, 0.064, 4.2e6, 6.5e6, 15847, 31295, 27500, 13265, 1843, 29.0e3, 22.0e3, 18068, 1064, 1.0))
cast_iron = MetallicParts{CI}(Properties(185, 0.257, 15.4e6, 23.0e6, 15435, 105265, 92500, 12675, 1687, 61190, 43645, 60775, 974, 1.0))
copper = MetallicParts{CU}(Properties(85, 0.320, 19.0e6, 17.0e6, 15394, 48365, 42500, 12617, 1025, 45965, 28565, 27924, 592, 1.0))
lead = NonMetallicParts{PB}(Properties(8, 0.394, 2.0e6, 2.4e6, 9144, 4552, 4000, 3689, 283, 2189.5, 870.0, 2628, 164, 1.0))
tuballoy = MetallicParts{TU}(Properties(240, 0.674, 20.7e6, 45.0e6, 12914, 136560, 120000, 9074, 1187, 140.0e3, 80829, 78843, 685, 1.0))
doron = NonMetallicParts{DO}(Properties(100, 0.073, 2.3e6, 2.3e6, 12996, 56900, 50000, 9191, 2327, 34.0e3, 19630, 32851, 1344, 1.0))
phenolic = NonMetallicParts{PH}(Properties(120, 0.069, 1.48e6, 4.0e6, 12472, 21622, 19000, 7583, 1476, 19.0e3, 12.0e3, 12483, 852, 1.0))
graphite_epoxy =MetallicParts{GE}(Properties(70, 0.057, 1.6e6, 1.6e6, 12635, 39830, 35000, 8675, 2204, 2000, 17500, 22996, 1272, 1.0))
unbonded_nylon = NonMetallicParts{UN}(Properties(20, 0.027, 0.9e6, 1.07e6, 13180, 11380, 10000, 9454, 1711, 9178.5, 29145, 6570, 988, 1.0))
bonded_nylon = NonMetallicParts{BN}(Properties(38, 0.034, 0.1e6, 87.0e3, 8528, 21622, 19000, 2808, 2102, 19430, 29145, 12483, 1214, 1.0))
lexan = NonMetallicParts{LX}(Properties(20, 0.046, 0.5e6, 0.3e6, 10341, 11380, 10000, 5398, 1311, 10.0e3, 9.0e3, 6570, 757, 1.0))
cast_plexiglas = NonMetallicParts{CP}(Properties(18, 0.042, 0.76e6, 0.4e6, 11438, 10242, 9000, 6965, 1302, 10.0e3, 11.0e3, 5913, 752, 1.0))
stretched_plexiglas = NonMetallicParts{SP}(Properties(23, 0.044, 0.76e6, 0.4e6, 11326, 13087, 11500, 6805, 1438, 10.0e3, 11.0e3, 7.556, 830, 1.0))
bullet_resistant_plexiglas = MetallicParts{BG}(Properties(1000, 0.089, 6.45e6, 10.0e6, 19320, 11380, 10000, 13939, 943, 8000, 70.0e3, 6570, 544, 1.0))
pine = NonMetallicParts{PN}(Properties(1, 0.013, 0.28e6, 1.3e6, 11881, 569, 500, 7599, 552, 310.0, 900.0, 329, 318, 1.0))
oak = NonMetallicParts{OK}(Properties(2, 0.025, 1.04e6, 1.8e6, 13955, 1138, 1000, 10561, 562, 800.0, 1780.0, 657, 325, 1.0))

rubber = MaterialUnkown{Unkown}()
#flesh = MaterialUnkown{Unkown}()
flesh = MaterialHuman{Body}()

dummyMat = MaterialUnkown{Unkown}()

end