module VectorizationBase

import ArrayInterface, LinearAlgebra, Libdl, Hwloc, IfElse
using ArrayInterface:
    StaticInt, Zero, One, StaticBool, True, False,
    contiguous_axis, contiguous_axis_indicator, contiguous_batch_size, stride_rank,
    device, CPUPointer, CPUIndex, eq, ne, lt, le, gt, ge,
    known_length, known_first, known_last, strides, offsets,
    static_first, static_last, static_length
import IfElse: ifelse

asbool(::Type{True}) = true
asbool(::Type{False}) = false
# TODO: see if `@inline` is good enough.
# @inline asvalbool(r) = Val(map(Bool, r))
# @inline asvalint(r) = Val(map(Int, r))
@generated function asvalint(r::T) where {T<:Tuple{Vararg{StaticInt}}}
    t = Expr(:tuple)
    for s ∈ T.parameters
        push!(t.args, s.parameters[1])
    end
    Expr(:call, Expr(:curly, :Val, t))
end
@generated function asvalbool(r::T) where {T<:Tuple{Vararg{StaticBool}}}
    t = Expr(:tuple)
    for b ∈ T.parameters
        push!(t.args, b === True)
    end
    Expr(:call, Expr(:curly, :Val, t))
end
@inline val_stride_rank(A) = asvalint(stride_rank(A))
@inline val_dense_dims(A) = asvalbool(ArrayInterface.dense_dims(A))

# doesn't export `Zero` and `One` by default, as these names could conflict with an AD library
export Vec, Mask, EVLMask, MM, stridedpointer, vload, vstore!, StaticInt, True, False, Bit,
    vbroadcast, mask, vfmadd, vfmsub, vfnmadd, vfnmsub,
    VecUnroll, Unroll, pick_vector_width

using Base: llvmcall, VecElement, HWReal, tail
const LLVMCALL = GlobalRef(Base, :llvmcall)

const FloatingTypes = Union{Float32, Float64} # Float16

const SignedHW = Union{Int8,Int16,Int32,Int64}
const UnsignedHW = Union{UInt8,UInt16,UInt32,UInt64}
const IntegerTypesHW = Union{SignedHW,UnsignedHW}
const IntegerTypes = Union{StaticInt,IntegerTypesHW}

struct Bit; data::Bool; end # Dummy for Ptr
const Boolean = Union{Bit,Bool}
const NativeTypesExceptBit = Union{Bool,HWReal}
const NativeTypes = Union{NativeTypesExceptBit, Bit}

const _Vec{W,T<:Number} = NTuple{W,Core.VecElement{T}}

abstract type AbstractSIMD{W,T <: Union{<:StaticInt,NativeTypes}} <: Real end
abstract type AbstractSIMDVector{W,T} <: AbstractSIMD{W,T} end
struct VecUnroll{N,W,T,V<:Union{NativeTypes,AbstractSIMD{W,T}}} <: AbstractSIMD{W,T}
    data::Tuple{V,Vararg{V,N}}
    @inline (VecUnroll(data::Tuple{V,Vararg{V,N}})::VecUnroll{N,W,T,V}) where {N,W,T,V<:AbstractSIMD{W,T}} = new{N,W,T,V}(data)
    @inline (VecUnroll(data::Tuple{T,Vararg{T,N}})::VecUnroll{N,T,T}) where {N,T<:NativeTypes} = new{N,1,T,T}(data)
    # # following two definitions are for checking that you aren't accidentally creating `VecUnroll{0}`s.
    # @inline (VecUnroll(data::Tuple{V,Vararg{V,N}})::VecUnroll{N,W,T,V}) where {N,W,T,V<:AbstractSIMD{W,T}} = (@assert(N > 0); new{N,W,T,V}(data))
    # @inline (VecUnroll(data::Tuple{T,Vararg{T,N}})::VecUnroll{N,T,T}) where {N,T<:NativeTypes} = (@assert(N > 0); new{N,1,T,T}(data))

    # @inline VecUnroll{N,W,T,V}(data::Tuple{V,Vararg{V,N}}) where {N,W,T,V<:AbstractSIMDVector{W,T}} = new{N,W,T,V}(data)
    # @inline (VecUnroll(data::Tuple{V,Vararg{V,N}})::VecUnroll{N,W,T,Vec{W,T}}) where {N,W,T,V<:AbstractSIMDVector{W,T}} = new{N,W,T,V}(data)
    # @inline (VecUnroll(data::Tuple{V,Vararg{V,N}})::VecUnroll{N,W,T,V}) where {N,W,T,V<:AbstractSIMDVector{W,T}} = new{N,W,T,V}(data)
end
# @inline VecUnroll(data::Tuple) = VecUnroll(promote(data))
# const AbstractSIMD{W,T} = Union{AbstractSIMDVector{W,T},VecUnroll{<:Any,W,T}}
const VecOrScalar = Union{AbstractSIMDVector,NativeTypes}
const NativeTypesV = Union{AbstractSIMD,NativeTypes,StaticInt}
# const NativeTypesV = Union{AbstractSIMD,NativeTypes,StaticInt}
const IntegerTypesV = Union{AbstractSIMD{<:Any,<:IntegerTypes},IntegerTypesHW}
struct Vec{W,T} <: AbstractSIMDVector{W,T}
    data::NTuple{W,Core.VecElement{T}}
    @inline Vec{W,T}(x::NTuple{W,Core.VecElement{T}}) where {W,T<:NativeTypes} = new{W,T}(x)
    @generated function Vec(x::Tuple{Core.VecElement{T},Vararg{Core.VecElement{T},_W}}) where {_W,T<:NativeTypes}
        W = _W + 1
        # @assert W === pick_vector_width(W, T)# || W === 8
        vtyp = Expr(:curly, :Vec, W, T)
        Expr(:block, Expr(:meta,:inline), Expr(:(::), Expr(:call, vtyp, :x), vtyp))
    end
    # @inline function Vec(x::NTuple{W,<:Core.VecElement}) where {W}
    #     T = eltype(x)
    #     @assert W === pick_vector_width(W, T)
    #     # @assert ispow2(W) && (W ≤ max(pick_vector_width(W, T), 8))
    #     new{W,T}(x)
    # end
end

@inline Base.copy(v::AbstractSIMDVector) = v
@inline asvec(x::_Vec) = Vec(x)
@inline asvec(x) = x
@inline data(vu::VecUnroll) = getfield(vu, :data)

@inline unrolleddata(x) = x
@inline unrolleddata(x::VecUnroll) = getfield(x, :data)

@inline _demoteint(::Type{T}) where {T} = T
@inline _demoteint(::Type{Int64}) = Int32
@inline _demoteint(::Type{UInt64}) = UInt32

abstract type AbstractMask{W,U<:UnsignedHW} <: AbstractSIMDVector{W,Bit} end
struct Mask{W,U} <: AbstractMask{W,U}
    u::U
    @inline function Mask{W,U}(u::Unsigned) where {W,U} # ignores U...
        U2 = mask_type(StaticInt{W}())
        new{W,U2}(u % U2)
    end
end
struct EVLMask{W,U} <: AbstractMask{W,U}
    u::U
    evl::UInt32
    @inline function EVLMask{W,U}(u::Unsigned, evl) where {W,U} # ignores U...
        U2 = mask_type(Val{W}())
        new{W,U2}(u % U2, evl % UInt32)
    end
end
@inline Mask{W}(u::U) where {W,U<:Unsigned} = Mask{W,U}(u)
@inline EVLMask{W}(u::U, i) where {W,U<:Unsigned} = EVLMask{W,U}(u, i)
@inline Mask(m::EVLMask{W,U}) where {W,U} = Mask{W,U}(getfield(m,:u))
# Const prop is good enough; added an @inferred test to make sure.
# Removed because confusion can cause more harm than good.

@inline Base.broadcastable(v::AbstractSIMDVector) = Ref(v)

Vec{W,T}(x::Vararg{NativeTypes,W}) where {W,T<:NativeTypes} = Vec(ntuple(w -> Core.VecElement{T}(x[w]), Val{W}()))
Vec{1,T}(x::Union{Float32,Float64}) where {T<:NativeTypes} = T(x)
Vec{1,T}(x::Union{Int8,UInt8,Int16,UInt16,Int32,UInt32,Int64,UInt64,Bool}) where {T<:NativeTypes} = T(x)

@inline Base.length(::AbstractSIMDVector{W}) where W = W
@inline Base.size(::AbstractSIMDVector{W}) where W = (W,)
@inline Base.eltype(::AbstractSIMD{W,T}) where {W,T} = T
@inline Base.conj(v::AbstractSIMDVector) = v # so that things like dot products work.
@inline Base.adjoint(v::AbstractSIMDVector) = v # so that things like dot products work.
@inline Base.transpose(v::AbstractSIMDVector) = v # so that things like dot products work.

# Not using getindex/setindex as names to emphasize that these are generally treated as single objects, not collections.
@generated function extractelement(v::Vec{W,T}, i::I) where {W,I <: IntegerTypesHW,T}
    typ = LLVM_TYPES[T]
    instrs = """
        %res = extractelement <$W x $typ> %0, i$(8sizeof(I)) %1
        ret $typ %res
    """
    call = :($LLVMCALL($instrs, $T, Tuple{_Vec{$W,$T},$I}, data(v), i))
    Expr(:block, Expr(:meta, :inline), call)
end
@generated function insertelement(v::Vec{W,T}, x::T, i::I) where {W,I <: IntegerTypesHW,T}
    typ = LLVM_TYPES[T]
    instrs = """
        %res = insertelement <$W x $typ> %0, $typ %1, i$(8sizeof(I)) %2
        ret <$W x $typ> %res
    """
    call = :(Vec($LLVMCALL($instrs, _Vec{$W,$T}, Tuple{_Vec{$W,$T},$T,$I}, data(v), x, i)))
    Expr(:block, Expr(:meta, :inline), call)
end
@inline (v::AbstractSIMDVector)(i::IntegerTypesHW) = extractelement(v, i - one(i))
@inline (v::AbstractSIMDVector)(i::Integer) = extractelement(v, Int(i) - 1)
Base.@propagate_inbounds (vu::VecUnroll)(i::Integer, j::Integer) = getfield(vu, :data)[j](i)

@inline Base.Tuple(v::Vec{W}) where {W} = ntuple(v, Val{W}())

# Use with care in function signatures; try to avoid the `T` to stay clean on Test.detect_unbound_args

@inline data(v) = v
@inline data(v::Vec) = getfield(v, :data)

function Base.show(io::IO, v::AbstractSIMDVector{W,T}) where {W,T}
    name = typeof(v)
    print(io, "$(name)<")
    for w ∈ 1:W
        print(io, repr(extractelement(v, w-1)))
        w < W && print(io, ", ")
    end
    print(io, ">")
end
Base.bitstring(m::AbstractMask{W}) where {W} = bitstring(data(m))[end-W+1:end]
function Base.show(io::IO, m::AbstractMask{W}) where {W}
    bits = data(m)
    if m isa EVLMask
        print(io, "EVLMask{$W,Bit}<")
    else
        print(io, "Mask{$W,Bit}<")
    end
    for w ∈ 0:W-1
        print(io, bits & 1)
        bits >>= 1
        w < W-1 && print(io, ", ")
    end
    print(io, ">")
end
function Base.show(io::IO, vu::VecUnroll{N,W,T,V}) where {N,W,T,V}
    println(io, "$(N+1) x $V")
    d = data(vu)
    for n in 1:N+1
        show(io, d[n]);
        n > N || println(io)
    end
end

"""
The name `MM` type refers to _MM registers such as `XMM`, `YMM`, and `ZMM`.
`MMX` from the original MMX SIMD instruction set is a [meaningless initialism](https://en.wikipedia.org/wiki/MMX_(instruction_set)#Naming).

The `MM{W,X}` type is used to represent SIMD indexes of width `W` with stride `X`.
"""
struct MM{W,X,I<:Union{HWReal,StaticInt}} <: AbstractSIMDVector{W,I}
    i::I
    @inline MM{W,X}(i::T) where {W,X,T<:Union{HWReal,StaticInt}} = new{W,X::Int,T}(i)
end
@inline MM(i::MM{W,X}) where {W,X} = MM{W,X}(getfield(i, :i))
@inline MM{W}(i::Union{HWReal,StaticInt}) where {W} = MM{W,1}(i)
@inline MM{W}(i::Union{HWReal,StaticInt}, ::StaticInt{X}) where {W,X} = MM{W,X}(i)
@inline data(i::MM) = getfield(i, :i)

@inline extractelement(i::MM{W,X,I}, j) where {W,X,I} = getfield(i, :i) + (X % I) * (j % I)

function Base.getproperty(::AbstractSIMD, ::Symbol)
    throw(ErrorException("""
`Base.getproperty` not defined on AbstractSIMD.
If you wish to work with the data as a tuple, it is recommended to use `Tuple(v)`. Once you have an ordinary tuple, you can access
individual elements normally. Alternatively, you can index using parenthesis, e.g. `v(1)` indexes the first element.
Parenthesis are used instead of `getindex`/square brackets because `AbstractSIMD` objects represent a single number, and
for `x::Number`, `x[1] === x`.

If you wish to perform a reduction on the collection, the naming convention is prepending the base function with a `v`. These functions
are not overloaded, because for `x::Number`, `sum(x) === x`. Functions include `vsum`, `vprod`, `vmaximum`, `vminimum`, `vany`, and `vall`.

If you wish to define a new operation applied to the entire vector, do not define it in terms of operations on the individual eleemnts.
This will often lead to bad code generation -- bad in terms of both performance, and often silently producing incorrect results!
Instead, implement them in terms of existing functions defined on `::AbstractSIMD`. Please feel free to file an issue if you would like
clarification, and especially if you think the function may be useful for others and should be included in `VectorizationBase.jl`.
"""))
end

"""
  pause()

For use in spin-and-wait loops, like spinlocks.
"""
@inline pause() = ccall(:jl_cpu_pause, Cvoid, ())

# notinthreadedregion() = iszero(ccall(:jl_in_threaded_region, Cint, ()))
# function assert_init_has_finished()
#     _init_has_finished[] || throw(ErrorException("bad stuff happened"))
#     return nothing
# end

include("static.jl")
include("cartesianvindex.jl")
include("topology.jl")
include("cpu_info.jl")
# include("cache_inclusivity.jl")
include("early_definitions.jl")
include("promotion.jl")
include("llvm_types.jl")
include("lazymul.jl")
include("strided_pointers/stridedpointers.jl")
# include("strided_pointers/bitpointers.jl")
include("strided_pointers/cartesian_indexing.jl")
include("strided_pointers/grouped_strided_pointers.jl")
include("strided_pointers/cse_stridemultiples.jl")
include("llvm_intrin/binary_ops.jl")
include("vector_width.jl")
include("ranges.jl")
include("llvm_intrin/conversion.jl")
include("llvm_intrin/masks.jl")
include("llvm_intrin/intrin_funcs.jl")
include("llvm_intrin/memory_addr.jl")
include("llvm_intrin/unary_ops.jl")
include("llvm_intrin/vbroadcast.jl")
(VERSION ≥ v"1.6.0-DEV.674") && include("llvm_intrin/verf.jl")
include("llvm_intrin/vector_ops.jl")
include("llvm_intrin/nonbroadcastingops.jl")
include("llvm_intrin/integer_fma.jl")
include("vecunroll/memory.jl")
include("vecunroll/fmap.jl")
include("base_defs.jl")
include("alignment.jl")
include("special/misc.jl")
include("special/double.jl")
include("special/exp.jl")
# include("special/log.jl")

demoteint(::Type{T}, W) where {T} = False()
demoteint(::Type{UInt64}, W::StaticInt) = gt(W, pick_vector_width(UInt64))
demoteint(::Type{Int64}, W::StaticInt) = gt(W, pick_vector_width(Int64))

@generated function simd_vec(::DemoteInt, y::_T, x::Vararg{_T,_W}) where {DemoteInt<:StaticBool,_T,_W}
    W = 1 + _W
    T = DemoteInt === True ? _demoteint(_T) : _T
    trunc = T !== _T
    Wfull = nextpow2(W)
    ty = LLVM_TYPES[T]
    init = W == Wfull ? "undef" : "zeroinitializer"
    instrs = ["%v0 = insertelement <$Wfull x $ty> $init, $ty %0, i32 0"]
    Tup = Expr(:curly, :Tuple, T)
    for w ∈ 1:_W
        push!(instrs, "%v$w = insertelement <$Wfull x $ty> %v$(w-1), $ty %$w, i32 $w")
        push!(Tup.args, T)
    end
    push!(instrs, "ret <$Wfull x $ty> %v$_W")
    llvmc = :($LLVMCALL($(join(instrs,"\n")), _Vec{$Wfull,$T}, $Tup))
    trunc ? push!(llvmc.args, :(y % $T)) : push!(llvmc.args, :y)
    for w ∈ 1:_W
        ref = Expr(:ref, :x, w)
        trunc && (ref = Expr(:call, :%, ref, T))
        push!(llvmc.args, ref)
    end
    quote
        $(Expr(:meta,:inline))
        Vec($llvmc)
    end
end
function vec_quote(demote, W, Wpow2, offset::Int = 0)
    call = Expr(:call, :simd_vec, Expr(:call, demote ? :True : :False)); Wpow2 += offset
    iszero(offset) && push!(call.args, :y)
    foreach(w -> push!(call.args, Expr(:ref, :x, w)), max(1,offset):min(W,Wpow2)-1)
    # foreach(w -> push!(call.args, Expr(:call, :VecElement, Expr(:call, :zero, :T))), W+1:Wpow2)
    call
end
@generated function _vec(::StaticInt{_Wpow2}, ::DemoteInt, y::T, x::Vararg{T,_W}) where {DemoteInt<:StaticBool, _Wpow2, _W, T <: NativeTypes}
    W = _W + 1
    demote = DemoteInt === True
    Wpow2 = demote ? 2_Wpow2 : _Wpow2
    if W ≤ Wpow2
        vec_quote(demote, W, Wpow2)
    else
        tup = Expr(:tuple)
        offset = 0
        while offset < W
            push!(tup.args, vec_quote(demote, W, Wpow2, offset)); offset += Wpow2
        end
        Expr(:call, :VecUnroll, tup)
    end
end
@inline function Vec(y::T, x::Vararg{T,_W}) where {_W, T <: NativeTypes}
    W = StaticInt{_W}() + One()
    _vec(pick_vector_width(W, T), demoteint(T, W), y, x...)
end

@inline reduce_to_onevec(f::F, vu::VecUnroll) where {F} = ArrayInterface.reduce_tup(f, data(vu))

include("precompile.jl")
_precompile_()

function __init__()
    ccall(:jl_generating_output, Cint, ()) == 1 && return
    reset_features!()
    if unwrap(cpu_name()) !== Symbol(Sys.CPU_NAME::String)
        @info "Defining CPU name."
        define_cpu_name()
    end
    safe_topology_load!()
    redefine_attr_count()
    foreach(redefine_cache, 1:4)
    return nothing
end

end # module


"""
    PMMH(n_iters::Int, smc_alg:::SMC, parameters_algs::Tuple{MH})

Particle independant Metropolis–Hastings and
Particle marginal Metropolis–Hastings samplers.

Usage:

```julia
alg = PMMH(100, SMC(20, :v1), MH(1,:v2))
alg = PMMH(100, SMC(20, :v1), MH(1,(:v2, (x) -> Normal(x, 1))))
```
"""
mutable struct PMMH <: InferenceAlgorithm
  n_iters               ::    Int               # number of iterations
  algs                  ::    Tuple             # Proposals for state & parameters
  space                 ::    Set               # sampling space, emtpy means all
  gid                   ::    Int               # group ID
  PMMH(n_iters::Int, algs::Tuple, space::Set, gid::Int) = new(n_iters, algs, space, gid)
  PMMH(n_iters::Int, smc_alg::SMC, parameter_algs...) = begin
      new(n_iters, tuple(parameter_algs..., smc_alg), Set(), 0)
  end
  PMMH(alg::PMMH, new_gid) = new(alg.n_iters, alg.algs, alg.space, new_gid)
end

PIMH(n_iters::Int, smc_alg::SMC) = PMMH(n_iters, tuple(smc_alg), Set(), 0)

function Sampler(alg::PMMH)
  alg_str = "PMMH"
  n_samplers = length(alg.algs)
  samplers = Array{Sampler}(undef, n_samplers)

  space = Set{Symbol}()

  for i in 1:n_samplers
    sub_alg = alg.algs[i]
    if isa(sub_alg, Union{SMC, MH})
      samplers[i] = Sampler(typeof(sub_alg)(sub_alg, i))
    else
      error("[$alg_str] unsupport base sampling algorithm $alg")
    end
    if typeof(sub_alg) == MH && sub_alg.n_iters != 1
      warn("[$alg_str] number of iterations greater than 1 is useless for MH since it is only used for its proposal")
    end
    space = union(space, sub_alg.space)
  end

  # Sanity check for space
  if !isempty(space)
    @assert issubset(Turing._compiler_[:pvars], space) "[$alg_str] symbols specified to samplers ($space) doesn't cover the model parameters ($(Turing._compiler_[:pvars]))"

    if Turing._compiler_[:pvars] != space
      warn("[$alg_str] extra parameters specified by samplers don't exist in model: $(setdiff(space, Turing._compiler_[:pvars]))")
    end
  end

  info = Dict{Symbol, Any}()
  info[:accept_his] = []
  info[:old_likelihood_estimate] = -Inf # Force to accept first proposal
  info[:old_prior_prob] = 0.0
  info[:samplers] = samplers

  Sampler(alg, info)
end

step(model::Function, spl::Sampler{PMMH}, vi::VarInfo, is_first::Bool) = begin
  violating_support = false
  proposal_ratio = 0.0
  new_prior_prob = 0.0
  new_likelihood_estimate = 0.0
  old_θ = copy(vi[spl])

  @debug "Propose new parameters from proposals..."
  for local_spl in spl.info[:samplers][1:end-1]
    @debug "$(typeof(local_spl)) proposing $(local_spl.alg.space)..."
    propose(model, local_spl, vi)
    if local_spl.info[:violating_support] violating_support=true; break end
    new_prior_prob += local_spl.info[:prior_prob]
    proposal_ratio += local_spl.info[:proposal_ratio]
  end

  if !violating_support # do not run SMC if going to refuse anyway
  @debug "Propose new state with SMC..."
  vi = step(model, spl.info[:samplers][end], vi)
  new_likelihood_estimate = spl.info[:samplers][end].info[:logevidence][end]

  @debug "computing accept rate α..."
  is_accept, logα = mh_accept(-(spl.info[:old_likelihood_estimate] + spl.info[:old_prior_prob]), -(new_likelihood_estimate + new_prior_prob); log_proposal_ratio=proposal_ratio)
  end

  @debug "decide whether to accept..."
  if !violating_support && is_accept # accepted
    push!(spl.info[:accept_his], true)
    spl.info[:old_likelihood_estimate] = new_likelihood_estimate
    spl.info[:old_prior_prob] = new_prior_prob
  else                      # rejected
    push!(spl.info[:accept_his], false)
    vi[spl] = old_θ
  end

  vi
end

sample(model::Function, alg::PMMH;
       save_state=false,         # flag for state saving
       resume_from=nothing,      # chain to continue
       reuse_spl_n=0             # flag for spl re-using
      ) = begin

    spl = Sampler(alg)
    alg_str = "PMMH"

    # Number of samples to store
    sample_n = spl.alg.n_iters

    # Init samples
    time_total = zero(Float64)
    samples = Array{Sample}(undef, sample_n)
    weight = 1 / sample_n
    for i = 1:sample_n
        samples[i] = Sample(weight, Dict{Symbol, Any}())
    end

    # Init parameters
    vi = resume_from == nothing ?
              Base.invokelatest(model, VarInfo(), nothing) :
              resume_from.info[:vi]
    n = spl.alg.n_iters

    # PMMH steps
    if PROGRESS spl.info[:progress] = ProgressMeter.Progress(n, 1, "[$alg_str] Sampling...", 0) end
    for i = 1:n
      @debug "$alg_str stepping..."
      time_elapsed = @elapsed vi = step(model, spl, vi, i==1)

      if spl.info[:accept_his][end]     # accepted => store the new predcits
        samples[i].value = Sample(vi).value
      else                              # rejected => store the previous predcits
        samples[i] = samples[i - 1]
      end

      time_total += time_elapsed
      if PROGRESS
        haskey(spl.info, :progress) && ProgressMeter.update!(spl.info[:progress], spl.info[:progress].counter + 1)
      end
    end

    println("[$alg_str] Finished with")
    println("  Running time    = $time_total;")
    accept_rate = sum(spl.info[:accept_his]) / n  # calculate the accept rate
    println("  Accept rate         = $accept_rate;")

    if resume_from != nothing   # concat samples
      pushfirst!(samples, resume_from.value2...)
    end
    c = Chain(0, samples)       # wrap the result by Chain

    if save_state               # save state
      save!(c, spl, model, vi)
    end

    c
end


# This file contains a number of additions to BlockArrays.jl. These are completely
# independent of Stheno.jl, and will (hopefully) move over to BlockArrays.jl at some point.

import Base: +, *, size, getindex, eltype, copy, \, vec, getproperty, zero
import LinearAlgebra: UpperTriangular, LowerTriangular, logdet, Symmetric, transpose,
    adjoint, AdjOrTrans, AdjOrTransAbsMat, cholesky!, logdet, ldiv!, mul!, logabsdet

"""
    blocksizes(X::AbstractBlockVecOrMat)
    blocksizes(X::AbstractArray, d::Int)

Get a vector containing the block sizes over the `d`th dimension of `X`. If `d` not
specified, returns a `Tuple` over blocksizes over each dimension.
"""
blocksizes(X::AbstractBlockArray, d::Int) = blocklengths(axes(X, d))
@nograd blocksizes

cumulsizes(X::AbstractBlockArray, d::Int) = vcat(1, cumsum(blocksizes(X, d)) .+ 1)
@nograd cumulsizes

nblocks(X::BlockArray, d::Int) = size(X.blocks, d)

"""
    are_conformal(A::BlockVecOrMat, B::BlockVecOrMat)

Test whether two block matrices (or vectors) are conformal. This criterion is stricter than
that for general matrices / vectors as we additionally require that each block be conformal
with block of the other matrix with which it will be multiplied. This ensures that the
result is itself straightforwardly representable as `BlockVecOrMat`.
"""
are_conformal(A::AVM, B::AVM) = blocksizes(A, 2) == blocksizes(B, 1)

ZygoteRules.@adjoint function BlockArrays.mortar(_blocks::AbstractArray)
    function mortar_pullback(Δ::NamedTuple{(:blocks, :axes)})
        return (Δ.blocks, )
    end
    function mortar_pullback(Δ::BlockArray)
        return mortar_pullback((blocks = Δ.blocks, axes=nothing))
    end
    return BlockArrays.mortar(_blocks), mortar_pullback
end

ZygoteRules.@adjoint function BlockArrays.Array(X::BlockArray)
    function Array_pullback(Δ::Array)
        ΔX = (blocks=BlockArray(Δ, axes(X)).blocks, axes=nothing)
        return (ΔX,)
    end
    return Array(X), Array_pullback
end


"""
Generates a `Float64` from a `MRG32k3a` instance.
"""
function rand(rng::MRG32k3a, ::Type{Float64})
    rand(rng)
end

"""
Generates a `Float32` from a `MRG32k3a` instance.
"""
function rand(rng::MRG32k3a, ::Type{Float32})
    convert(Float32, rand(rng))
end

"""
Generates a `Float16` from a `MRG32k3a` instance.
"""
function rand(rng::MRG32k3a, ::Type{Float16})
    convert(Float16, rand(rng))
end

"""
Generates a `Int64` in the given range from a `MRG32k3a` instance.
"""
function rand(rng::MRG32k3a, r::UnitRange{Int64})
    r.start + convert(Int64, (r.stop - r.start) * rand(rng, Float64))
end


"""
    readsiggen(file_path::String; T::Type)

Reads a [SigGen](https://github.com/radforddc/icpc_siggen) configuration file (ending in `.config`) in 'file_path'
and returns a dictionary of all parameters. Non-existing parameteres are set to `0`.

## Arguments
* `file_path::String`: File path leading to the [SigGen](https://github.com/radforddc/icpc_siggen) configuration file.

## Keywords
* `T::Type`: Type of the parameters in the output dictionary. Default is `Float64`.
"""
function readsiggen(file_path::String; T::Type=Float64)

    detector_name = split(basename(file_path), ".config")[1]

    config = Dict("name"       => detector_name,
        "verbosity_level"      => 0,
        "xtal_length"          => 0,
        "xtal_radius"          => 0,
        "top_bullet_radius"    => 0, # not implemented in later conversions
        "bottom_bullet_radius" => 0, # not implemented in later conversions
        "pc_length"            => 0,
        "pc_radius"            => 0,
        "bulletize_PC"         => 0,
        "wrap_around_radius"   => 0,
        "ditch_depth"          => 0,
        "ditch_thickness"      => 0,
        "bottom_taper_length"  => 0,
        "hole_length"          => 0,
        "hole_radius"          => 0,
        "outer_taper_length"   => 0,
        "inner_taper_length"   => 0,
        "outer_taper_width"    => 0,
        "inner_taper_width"    => 0,
        "taper_angle"          => 0,
        "taper_length"         => 0,
        "Li_thickness"         => 0,
        "energy"               => 0,
        "xtal_grid"            => 0,
        "impurity_z0"          => 0,
        "impurity_gradient"    => 0,
        "impurity_surface"     => 0,
        "xtal_HV"              => 0,
        "max_iterations"       => 0,
        "write_field"          => 0,
        "write_WP"             => 0,
        "drift_name"           => "",
        "field_name"           => "",
        "wp_name"              => "",
        "xtal_temp"            => 0,
        "preamp_tau"           => 0,
        "time_steps_calc"      => 0,
        "step_time_calc"       => 0,
        "step_time_out"        => 0,
        "charge_cloud_size"    => 0,
        "use_diffusion"        => 0,
        "surface_drift_vel_factor" => 0);

    file = open(file_path) do f
        temp          = readlines(f)
        config_file   = Dict()
        for line in temp
            if line  != "" && string(line[1]) != "#"
                line  = strip(split(line, "#")[1])
                name  = ""
                value = ""
                if length(split(line, "\t")) > 1
                    name  = split(line, "\t")[1]
                    value = strip(split(line, "\t")[2])
                else
                    name  = split(line, " ")[1]
                    value = strip(split(line, " ")[end])
                end

                config_file[name] = value

            end
        end
        return config_file
    end

    config_keys = keys(config)
    file_keys   = keys(file)

    for key in file_keys
        if key in config_keys
            if key == "wp_name" || key == "drift_name" || key == "field_name"
                config[key] = string(file[key])
            elseif key == "taper_angle" && parse(Float64, file["taper_angle"]) != 0
                config["outer_taper_width"] = tan(deg2rad(parse(T, file[key]))) * parse(T, file["outer_taper_length"])
                config["inner_taper_width"] = tan(deg2rad(parse(T, file[key]))) * parse(T, file["inner_taper_length"])
                config[key]                 = parse(T, file[key])
            else
                config[key] = parse(T, file[key])
            end
        end
    end
    return config
end

"""
    siggentodict(config::Dict; units::Dict)

Converts the dictionary containing the parameters from a [SigGen](https://github.com/radforddc/icpc_siggen) configuration file
to a dictionary that can be understood by SolidStateDetectors.jl.

## Arguments
* `config::Dict`: Dictionary containing [SigGen](https://github.com/radforddc/icpc_siggen) parameters (output of `readsiggen()``).

## Keywords
* `units::Dict`: Units used in [SigGen](https://github.com/radforddc/icpc_siggen) configuration file (set to `"mm"`, `"deg"`, `"V"` and `"K"`).
    The dictionary needs the fields `"length"`, `"angle"`, `"potential"` and `"temperature"`.
"""
function siggentodict(config::Dict;
        units::Dict = Dict("length" => "mm","angle" => "deg","potential" => "V","temperature" => "K"))

    #>----------------Settings START-------------------------------------------<

    #>----------------Setting the z and r limits-------------------------------<
    z_limit      = ceil(config["xtal_length"]; sigdigits=1);
    if z_limit - config["xtal_length"] <= 0
        z_limit += 10
    end
    r_limit      = ceil(config["xtal_radius"]; sigdigits=1);
    if r_limit - config["xtal_radius"] <= 0
        r_limit += 10
    end

    #>----------------Setting the Axes & Grid dictionaries---------------------<
    axes         = Dict("r"           => Dict("to"         => r_limit,
                                              "boundaries" => "inf"),
                        "phi"         => Dict("from"       => 0,
                                              "to"         => 0,
                                              "boundaries" => "periodic"),
                        "z"           => Dict("from"       => -10,
                                              "to"         => z_limit,
                                              "boundaries" => Dict("left" => "inf", "right" => "inf"))
                        )

    grid         = Dict("coordinates" => "cylindrical",
                        "axes"        => axes);

    #>----------------Settings END---------------------------------------------<
    #<><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><
    #>----------------Volume START---------------------------------------------<

    #>----------------Main Volume -> Initial Cylinder--------------------------<
    # r and L are given by the input file. Cylindrical symmetry is set to default
    main_vol     = Dict("translate"   => Dict("z"    => config["xtal_length"]/2,
                        "tube"        => Dict(
                        "name"        => "Initial Cylinder",
                        "r"           => Dict("from" => 0.0, "to" => config["xtal_radius"]),
                        "phi"         => Dict("from" => 0.0, "to" => 360.0),
                        "h"           => config["xtal_length"])));

    #>----------------Upper Cone / Taper on top--------------------------------<
    # If not existing, everything will be set to 0
    upper_cone   = Dict("translate"   => Dict("z"      => config["xtal_length"] -
                                            config["outer_taper_length"]/2,
                        "cone"        => Dict(
                        "name"        => "Upper Cone",
                        "r"           => Dict("bottom" => Dict("from" => config["xtal_radius"],
                                                               "to"   => ceil(config["xtal_radius"] + 1.0)),
                                              "top"    => Dict("from" => config["xtal_radius"] -
                                                                            config["outer_taper_width"],
                                                               "to"   => ceil(config["xtal_radius"] + 1.0))),
                        "phi"         => Dict("from"   => 0.0, "to" => 360.0),
                        "h"           => config["outer_taper_length"])));

    #>----------------Lower Cone / Taper on bottom-----------------------------<
    # 45-degree taper at bottom of ORTEC-type crystal
    lower_cone   = Dict("translate"   => Dict("z"      => 0.4*config["taper_length"],
                        "cone"        => Dict(
                        "name"        => "Lower Cone",
                        "r"           => Dict("top"    => Dict("from" => config["xtal_radius"],
                                                               "to"   => ceil(config["xtal_radius"] + 1.0)),
                                              "bottom" => Dict("from" => config["xtal_radius"] -
                                                                            1.2 * config["taper_length"],
                                                               "to"   => ceil(config["xtal_radius"] + 1.0))),
                        "phi"         => Dict("from"   => 0.0, "to" => 360.0),
                        "h"           => 1.2*config["taper_length"])));

    #>----------------Ditch around the bottom contact--------------------------<
    # If not existing, everything will be set to 0
    ditch        = Dict("translate"   => Dict("z"    => config["ditch_depth"]/2,
                        "tube"        => Dict(
                        "name"        => "Ditch",
                        "r"           => Dict("from" => config["wrap_around_radius"] -
                                                            config["ditch_thickness"],
                                              "to"   => config["wrap_around_radius"]),
                        "phi"         => Dict("from" => 0.0, "to" => 360.0),
                        "h"           => config["ditch_depth"])));

    #>----------------Borehole on bottom as contact----------------------------<
    # If not existing, everything will be set to 0
    borehole = if config["bulletize_PC"] == 1
        if config["pc_length"] < config["pc_radius"]
            Dict( "union"   => [
            Dict( "tube"    => Dict(
                  "r"       => config["pc_radius"] - config["pc_length"],
                  "h"       => 1.2*config["pc_length"],
                  "origin"  => Dict("z" => 0.4*config["pc_length"]))),
            Dict( "torus"   => Dict(
                  "r_torus" => config["pc_radius"] - config["pc_length"],
                  "r_tube"  => config["pc_length"]))])
        elseif config["pc_length"] > config["pc_radius"]
            Dict( "union"  => [
            Dict( "tube"   => Dict(
                  "r"      => config["pc_radius"],
                  "h"      => 1.2*config["pc_length"] - config["pc_radius"],
                  "origin" => Dict("z" => (0.8*config["pc_length"] - config["pc_radius"])/2))),
            Dict( "sphere" => Dict(
                  "r"      => config["pc_radius"],
                  "origin" => Dict("z" => config["pc_length"] - config["pc_radius"])))])
        else
            Dict( "sphere" => Dict(
                       "r" => config["pc_radius"]))
        end
    else # not bulletized
            Dict( "tube"   => Dict(
                  "name"   => "Borehole",
                  "r"      => Dict("from" => 0.0, "to" => config["pc_radius"]),
                  "h"      => config["pc_length"],
                  "origin" => Dict("z" => config["pc_length"]/2)));
    end

    #>----------------Hole on top of the detector------------------------------<
    # If not existing, everything will be set to 0
    hole = if config["inner_taper_width"] != 0 && config["inner_taper_length"] != 0
        Dict("translate"   => Dict("z"      => config["xtal_length"] -
                               config["inner_taper_length"]/2,
             "cone"        => Dict(
             "name"        => "Inner Cone",
             "r"           => Dict("top"    => Dict("from" => 0,
                                                   "to"   => config["hole_radius"] +
                                                                config["inner_taper_width"]),
                                  "bottom" => Dict("from" => 0,
                                                   "to"   => config["hole_radius"])),
             "phi"         => Dict("from"   => 0.0, "to" => 360.0),
             "h"           => config["inner_taper_length"])));

    else
        Dict("translate"   => Dict("z"    => config["xtal_length"] -
                                               0.8*config["hole_length"]/2,
             "tube"        => Dict(
             "name"        => "Top Hole",
             "r"           => Dict("from" => 0.0, "to" => config["hole_radius"]),
             "phi"         => Dict("from" => 0.0, "to" => 360.0),
             "h"           => 1.2*config["hole_length"])));
    end

    #>----------------Subtract volumes-----------------------------------------<
    geometry = main_vol;
    if config["outer_taper_length"] != 0 || config["outer_taper_width"] != 0
        geometry = Dict("difference" => [geometry, upper_cone])
    end
    if config["inner_taper_length"] != 0 || config["inner_taper_width"] != 0 || config["hole_radius"] != 0 || config["hole_length"] != 0
        geometry = Dict("difference" => [geometry, hole])
    end
    if config["taper_length"] != 0
        geometry = Dict("difference" => [geometry, lower_cone])
    end
    if config["ditch_thickness"] != 0
        geometry = Dict("difference" => [geometry, ditch])
    end
    if config["bulletize_PC"] == 1 || (config["pc_radius"] != 0 && config["pc_length"] >= 0.1) # below this, it is the thickness of the layer
        geometry = Dict("difference" => [geometry, borehole])
    end

    #>----------------Volume END-----------------------------------------------<
    #<><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><
    #>----------------Contacts START-------------------------------------------<

    #>----------------Area/volume of point contact-----------------------------<
    #>----------------Contact 1------------------------------------------------<
    # Depending on whether the contact is a borehole or not

    geometry_1 = if config["bulletize_PC"] == 1.0
        if config["pc_length"] < config["pc_radius"]
            Dict( "union"      => [
            Dict( "tube"       => Dict(
                  "r"          => config["pc_radius"] - config["pc_length"],
                  "h"          => 0,
                  "origin"     => Dict("z" => config["pc_length"]))),
            Dict( "torus"      => Dict(
                  "r_torus"    => config["pc_radius"] - config["pc_length"],
                  "r_tube"     => Dict("from" => config["pc_length"],
                                       "to"   => config["pc_length"]),
                  "theta"      => Dict("from" => 0.0, "to" => 90.0)))])
        elseif config["pc_length"] > config["pc_radius"]
            Dict( "union"      => [
            Dict( "tube"       => Dict(
                  "r"          => Dict("from" => config["pc_radius"], "to" => config["pc_radius"]),
                  "h"          => config["pc_length"] - config["pc_radius"],
                  "origin"     => Dict("z" => (config["pc_length"] - config["pc_radius"])/2))),
            Dict( "difference" => [
            Dict( "sphere"     => Dict(
                  "r"          => Dict("from" => config["pc_radius"],
                                       "to"   => config["pc_radius"]),
                  "origin"     => Dict("z" => config["pc_length"] - config["pc_radius"]))),
            Dict( "tube"       => Dict(
                  "r"          => 1.2*config["pc_radius"],
                  "h"          => 1.2*config["pc_radius"],
                  "origin"     => Dict("z" => config["pc_length"] - 1.6*config["pc_radius"])))])])
        else
            Dict( "difference" => [
            Dict( "sphere"     => Dict(
                  "r"          => Dict("from" => config["pc_radius"],
                                       "to"   => config["pc_radius"]))),
            Dict( "tube"       => Dict(
                  "r"          => 1.2*config["pc_radius"],
                  "h"          => 1.2*config["pc_radius"],
                  "origin"     => Dict("z" => -0.6 * config["pc_radius"])))])
        end
    elseif config["pc_radius"] != 0 && config["pc_length"] >= 5 # below this, it is the thickness of the layer
            Dict( "union"      => [
            # borehole_wall
            Dict( "tube"       => Dict(
                  "r"          => Dict("from" => config["pc_radius"],
                                       "to"   => config["pc_radius"]),
                  "h"          => config["pc_length"],
                  "origin"     => Dict("z" => config["pc_length"]/2))),
            # borehole_top
            Dict( "tube"       => Dict(
                  "r"          => Dict("from" => 0.0,
                                       "to"   => config["pc_radius"]),
                  "h"          => 0.0,
                  "origin"     => Dict("z" => config["pc_length"]))),
            # borehole_around
            Dict( "tube"       => Dict(
                  "r"          => Dict("from" => config["pc_radius"],
                                       "to"   => config["wrap_around_radius"] - config["ditch_thickness"]),
                  "h"          => 0.0))])
    elseif config["pc_radius"] != 0 && config["pc_length"] < 5
            Dict( "tube"       => Dict(
                  "r"          => config["pc_radius"],
                  "h"          => config["pc_length"],
                  "origin"     => Dict("z" => config["pc_length"]/2)))
    end

    #>----------------Define contact 1-----------------------------------------<

    contact_1 = Dict("material"    => "HPGe",
                     "id"          => 1,
                     "potential"   => 0.0,
                     "geometry"    => geometry_1);

    #>----------------Contact 2------------------------------------------------<
    geometry_2 = Dict[]

    if config["wrap_around_radius"] != 0.0
        bottom = Dict("tube" => Dict(
                      "r"    => Dict("from" => config["wrap_around_radius"],
                                       "to" => config["xtal_radius"] - config["taper_length"]),
                      "phi"  => Dict("from" => 0.0,
                                       "to" => 360.0),
                      "h"    => 0.0))
        push!(geometry_2, bottom)
    end

    wall   = Dict("translate" => Dict("z" => config["xtal_length"]/2 + config["taper_length"]/2 - config["outer_taper_length"]/2,
                  "tube" => Dict(
                  "r"    => Dict("from" => config["xtal_radius"],
                                 "to"   => config["xtal_radius"]),
                  "phi"  => Dict("from" => 0.0,
                                   "to" => 360.0),
                  "h"    => config["xtal_length"] - config["taper_length"] - config["outer_taper_length"])))
    push!(geometry_2, wall)

    top    = Dict("translate" => Dict("z" => config["xtal_length"],
                  "tube" => Dict(
                  "r"    => Dict("from" => config["hole_radius"] + config["inner_taper_width"],
                                   "to" => config["xtal_radius"] - config["outer_taper_width"]),
                  "phi"  => Dict("from" => 0.0,
                                   "to" => 360.0),
                  "h"    => 0.0)))
    push!(geometry_2, top)
    if config["hole_radius"] != 0 && config["hole_length"] != 0 && config["inner_taper_width"] == 0 && config["inner_taper_length"] == 0
        hole   = Dict("translate" => Dict("z" => config["xtal_length"] - config["hole_length"]/2,
                      "tube" => Dict(
                      "r"    => Dict("from" => config["hole_radius"],
                                       "to" => config["hole_radius"]),
                      "phi"  => Dict("from" => 0.0,
                                       "to" => 360.0),
                      "h"    => config["hole_length"])))
        inner_hole_bottom =  Dict("translate" => Dict("z" => config["xtal_length"] - config["hole_length"],
                      "tube" => Dict(
                      "r"    => Dict("from" => 0.0,
                                       "to" => config["hole_radius"]),
                      "phi"  => Dict("from" => 0.0,
                                       "to" => 360.0),
                      "h"    => 0.0)))
        push!(geometry_2, hole)
        push!(geometry_2, inner_hole_bottom)
    end
    if config["hole_radius"] != 0 && config["hole_length"] != 0 && config["inner_taper_width"] != 0 && config["inner_taper_length"] != 0
        inner_taper = Dict("translate" => Dict("z" => config["xtal_length"] - config["inner_taper_length"]/2,
                      "cone" => Dict(
                      "r"    => Dict("top"    => Dict("from" => config["hole_radius"] + config["inner_taper_width"],
                                                      "to"   => config["hole_radius"] + config["inner_taper_width"]),
                                     "bottom" => Dict("from" => config["hole_radius"],
                                                      "to"   => config["hole_radius"])),
                      "phi"  => Dict("from" => 0.0, "to" => 360.0),
                      "h"    => config["inner_taper_length"])))
        inner_hole_bottom = Dict("translate" => Dict("z" => config["xtal_length"] - config["hole_length"],
                      "tube" => Dict(
                      "r"    => Dict("from" => 0.0, "to" => config["hole_radius"]),
                      "phi"  => Dict("from" => 0.0, "to" => 360.0),
                      "h"    => 0.0)))
        push!(geometry_2, inner_taper)
        push!(geometry_2, inner_hole_bottom)
    end
    if config["outer_taper_width"] != 0 && config["outer_taper_length"] != 0
        outer_taper = Dict("translate" => Dict("z" => config["xtal_length"] - config["outer_taper_length"]/2,
                      "cone" => Dict(
                      "r"    => Dict("top"    => Dict("from" => config["xtal_radius"] - config["outer_taper_width"],
                                                      "to"   => config["xtal_radius"] - config["outer_taper_width"]),
                                     "bottom" => Dict("from" => config["xtal_radius"],
                                                      "to"   => config["xtal_radius"])),
                      "phi"  => Dict("from"   => 0.0, "to" => 360.0),
                      "h"    => config["outer_taper_length"])))
        push!(geometry_2, outer_taper)
    end
    if config["taper_length"] != 0
        taper = Dict("translate" => Dict("z" => config["taper_length"]/2,
                      "cone" => Dict(
                      "r"    => Dict("top"    => Dict("from" => config["xtal_radius"],
                                                      "to"   => config["xtal_radius"]),
                                     "bottom" => Dict("from" => config["xtal_radius"] - config["taper_length"],
                                                      "to"   => config["xtal_radius"] - config["taper_length"])),
                      "phi"   => Dict("from"   => 0.0, "to" => 360.0),
                      "h"     => config["taper_length"])))
        push!(geometry_2, taper)
    end

    #>----------------Define contact 2-----------------------------------------<

    contact_2 = Dict("material"    => "HPGe",
                     "id"          => 2,
                     "potential"   => config["xtal_HV"],
                     "geometry"    => Dict("union" => geometry_2));

    #>----------------Contacts END---------------------------------------------<
    #<><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><
    #>----------------Final configuration--------------------------------------<

    impurity_z0            = config["impurity_z0"]*1e10
    impurity_gradient      = config["impurity_gradient"]*1e10
    if units["length"]     == "mm"
        impurity_z0        /= 1e3
        impurity_gradient  /= 1e4
    elseif units["length"] == "m"
        impurity_z0        *= 1e6
        impurity_gradient  *= 1e8
    end

    objects = Dict("semiconductor" => Dict(
                    "material"    => "HPGe",
                    "temperature" => config["xtal_temp"],
                    "impurity_density" => Dict(
                                "name"     => "cylindrical",
                                "r"        => Dict("init"     => 0.0,
                                                   "gradient" => 0.0),
                                "z"        => Dict("init"     => impurity_z0,
                                                   "gradient" => impurity_gradient)),
                    "geometry"    => geometry),
                    "contacts" => [ contact_1, contact_2 ]);

    conf = Dict("name"    => config["name"],
                "units"   => units,
                "grid"    => grid,
                "medium"  => "vacuum",
                "detectors" => [objects]);

    return conf
end


abstract type FindBestAlg end
struct BruteForce <: FindBestAlg end
struct MaybeFaster <: FindBestAlg end

best_acronym(af::AcronymFamily, params::ScoreParams) = best_acronym(BruteForce(), af, params)

best_acronym(::BruteForce, af::AcronymFamily, params::ScoreParams) = maximum_by(last, (a, score(a, params)) for a in acronyms(af))

best_acronym(::MaybeFaster, af::AcronymFamily, params::ScoreParams) = maximum_by(last, _best_faster(af, trace, params) for trace in _family_traces(af))

function _best_faster(af::AcronymFamily, trace::AbstractVector{NTuple{2, Int}}, params::ScoreParams)
    ac_words = Union{String, Nothing}[]
    s = 0.
    group_seen = falses(length(params.word_groups))

    for (i, n) in trace
        words = af.components[n, i]

        if isempty(words)
            push!(ac_words, nothing)
            s -= params.gap_penalty

        else
            best_word = nothing
            best_score = -Inf

            for word in words
                g = params.word_to_group[word]
                ws = params.group_scores[g]

                if g in params.required_groups && !group_seen[g]
                    best_word = word
                    best_score = 0.
                    break
                elseif ws > best_score
                    best_word = word
                    best_score = ws
                end
            end

            push!(ac_words, best_word)
            s += best_score
            group_seen[params.word_to_group[best_word]] = true
        end
    end

    all(group_seen[g] for g in params.required_groups) || (s = -Inf)

    a = Acronym(af.word, ac_words, last.(trace))
    return a, s
end


module ArrayInterfaceCUDA

using Adapt
using ArrayInterface, ArrayInterfaceGPUArrays
using CUDA

function ArrayInterface.lu_instance(A::CuMatrix{T}) where {T}
    CUDA.CUSOLVER.CuQR(similar(A, 0, 0), similar(A, 0))
end

ArrayInterface.device(::Type{<:CUDA.CuArray}) = ArrayInterface.GPU()

end # module


#=
Copyright (c) 2015, Intel Corporation

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

    * Redistributions of source code must retain the above copyright notice,
      this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of Intel Corporation nor the names of its contributors
      may be used to endorse or promote products derived from this software
      without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
=#

# convert binary into HDF5 data
using HDF5

base_dir = "./"

if length(ARGS) > 0
  base_dir = ARGS[1]
end

datasets = Dict("train" => ["$base_dir/train-labels-idx1-ubyte","$base_dir/train-images-idx3-ubyte"],
                "test" => ["$base_dir/t10k-labels-idx1-ubyte","$base_dir/t10k-images-idx3-ubyte"])

for key in keys(datasets)
  label_fn, data_fn = datasets[key]
  label_f = open(label_fn)
  data_f  = open(data_fn)

  label_header = read(label_f, Int32, 2)
  @assert ntoh(label_header[1]) == 2049
  n_label = round(Int, ntoh(label_header[2]))
  data_header = read(data_f, Int32, 4)
  @assert ntoh(data_header[1]) == 2051
  n_data = round(Int, ntoh(data_header[2]))
  @assert n_label == n_data
  h = round(Int, ntoh(data_header[3]))
  w = round(Int, ntoh(data_header[4]))

  println("Exporting $n_data digits of size $h x $w")

  h5open("$base_dir/$key.hdf5", "w") do h5
    dset_data = d_create(h5, "data", datatype(Float32), dataspace(w, h, 1, n_data))
    dset_label = d_create(h5, "label", datatype(Float32), dataspace(1, n_data))

    img = readbytes(data_f, n_data * h*w)
    img = convert(Array{Float32},img) / 256 # scale into [0,1)
    class = readbytes(label_f, n_data)
    class = convert(Array{Float32},class)

    idx = 1:n_data
    println("  $idx...")

    idx = collect(idx)
    rp = randperm(length(idx))
    for j = 1:length(idx)
      r_idx = rp[j]
      dset_data[:,:,1,idx[j]] = img[(r_idx-1)*h*w+1:r_idx*h*w]
      dset_label[1,idx[j]] = class[r_idx]
    end
  end

  close(label_f)
  close(data_f)
end


### reinterpret

@pure samesize(::Type{T}, ::Type{S}) where {T,S} = sizeof(T) == sizeof(S)

## Color->Color
# function reinterpretc(::Type{CV1}, a::Array{CV2,1}) where {CV1<:Colorant,CV2<:Colorant}
#     CV = ccolor(CV1, CV2)
#     l = (length(a)*sizeof(CV2))÷sizeof(CV1)
#     l*sizeof(CV1) == length(a)*sizeof(CV2) || throw(ArgumentError("sizes are incommensurate"))
#     reshape(reinterpret(CV, a), (l,))
# end
function reinterpretc(::Type{CV1}, a::AbstractArray{CV2}) where {CV1<:Colorant,CV2<:Colorant}
    CV = ccolor(CV1, CV2)
    if samesize(CV, CV2)
        return reshape(reinterpret(CV, a), size(a))
    end
    throw(ArgumentError("result shape not specified"))
end

## Color->T
function reinterpretc(::Type{T}, a::AbstractArray{CV}) where {T<:Number,CV<:Colorant}
    if samesize(T, CV)
        return reinterpret(T, a)
    end
    axs = axes(a)
    if sizeof(CV) == sizeof(T)*_len(CV)
        return reinterpret(T, reshape(a, Base.OneTo(1), axs...))
    end
    throw(ArgumentError("result shape not specified"))
end
reinterpretc(::Type{T}, a::AbstractArray{CV,0}) where {T<:Number,CV<:Colorant} =
    reinterpret(T, reshape(a, 1))

_len(::Type{C}) where {C} = _len(C, eltype(C))
_len(::Type{C}, ::Type{Any}) where {C} = error("indeterminate type")
_len(::Type{C}, ::Type{T}) where {C,T} = sizeof(C) ÷ sizeof(T)

## T->Color
# We have to distinguish two forms of call:
#   form 1: reinterpretc(RGB{N0f8}, img)
#   form 2: reinterpretc(RGB, img)
function reinterpretc(::Type{CV}, a::AbstractArray{T}) where {CV<:Colorant,T<:Number}
    @noinline throwdm(::Type{C}, ind1) where C =
        throw(DimensionMismatch("indices $ind1 are not consistent with color type $C"))
    CVT = ccolor_number(CV, T)
    if samesize(CVT, T)
        return reinterpret(CVT, a)
    end
    axs = axes(a)
    if axs[1] == Base.OneTo(sizeof(CVT) ÷ sizeof(eltype(CVT)))
        return reshape(reinterpret(CVT, a), tail(axs))
    end
    throwdm(CV, axs[1])
end

# ccolor_number converts form 2 calls to form 1 calls
ccolor_number(::Type{CV}, ::Type{T}) where {CV<:Colorant,T<:Number} =
    ccolor_number(CV, eltype(CV), T)
ccolor_number(::Type{CV}, ::Type{CVT}, ::Type{T}) where {CV,CVT<:Number,T} = CV # form 1
ccolor_number(::Type{CV}, ::Type{Any}, ::Type{T}) where {CV<:Colorant,T} = CV{T} # form 2

### convert
#
# The main contribution here is "concretizing" the colorant type: allow
#    convert(RGB, a)
# rather than requiring
#    convert(RGB{N0f8}, a)
# Where possible the raw element type of the source is retained.
Base.convert(::Type{Array{C}},   img::Array{C,n}) where {C<:Color1,n} = img
Base.convert(::Type{Array{C,n}}, img::Array{C,n}) where {C<:Color1,n} = img
Base.convert(::Type{Array{C}},   img::Array{C,n}) where {C<:Colorant,n} = img
Base.convert(::Type{Array{C,n}}, img::Array{C,n}) where {C<:Colorant,n} = img

function Base.convert(::Type{Array{Cdest}},
                      img::AbstractArray{Csrc,n}) where {Cdest<:Colorant,n,Csrc<:Colorant}
    convert(Array{Cdest,n}, img)
end

function Base.convert(::Type{Array{Cdest,n}},
                      img::AbstractArray{Csrc,n}) where {Cdest<:Colorant,n,Csrc<:Colorant}
    copyto!(Array{ccolor(Cdest, Csrc)}(undef, size(img)), img)
end

function Base.convert(::Type{Array{Cdest}},
                      img::BitArray{n}) where {Cdest<:Color1,n}
    convert(Array{Cdest,n}, img)
end

function Base.convert(::Type{Array{Cdest}},
                      img::AbstractArray{T,n}) where {Cdest<:Color1,n,T<:Real}
    convert(Array{Cdest,n}, img)
end

function Base.convert(::Type{Array{Cdest,n}},
                      img::BitArray{n}) where {Cdest<:Color1,n}
    copyto!(Array{ccolor(Cdest, Gray{Bool})}(undef, size(img)), img)
end

function Base.convert(::Type{Array{Cdest,n}},
                      img::AbstractArray{T,n}) where {Cdest<:Color1,n,T<:Real}
    copyto!(Array{ccolor(Cdest, Gray{T})}(undef, size(img)), img)
end

function Base.convert(::Type{OffsetArray{Cdest,n,A}},
                      img::AbstractArray{Csrc,n}) where {Cdest<:Colorant,n, A <:AbstractArray,Csrc<:Colorant}
    copyto!(OffsetArray{ccolor(Cdest, Csrc)}(undef, axes(img)),img)
end

function Base.convert(::Type{OffsetArray{C,n,A}},
                      img::AbstractArray{C,n}) where {C<:Colorant,n, A <:AbstractArray}
    img
end

# for docstrings in the operations below
shortname(::Type{T}) where {T<:FixedPoint} = (io = IOBuffer(); FixedPointNumbers.showtype(io, T); String(take!(io)))
shortname(::Type{T}) where {T} = string(T)

# float32, float64, etc. Used for conversions like
#     Array{RGB{N0f8}} -> Array{RGB{Float32}},
# since
#    convert(Array{RGB{Float32}}, A)
# is annoyingly verbose for such a common operation.
for (fn,T) in (#(:float16, Float16),   # Float16 currently has promotion problems
               (:float32, Float32), (:float64, Float64),
               (:n0f8, N0f8), (:n6f10, N6f10),
               (:n4f12, N4f12), (:n2f14, N2f14), (:n0f16, N0f16))
    @eval begin
        ($fn)(::Type{C}) where {C<:Colorant} = base_colorant_type(C){$T}
        ($fn)(::Type{S}) where {S<:Number  } = $T
        ($fn)(c::Colorant) = convert(($fn)(typeof(c)), c)
        ($fn)(n::Number)   = convert(($fn)(typeof(n)), n)
        @deprecate ($fn)(A::AbstractArray{C}) where {C<:Colorant} ($fn).(A)
        fname = $(Expr(:quote, fn))
        Tname = shortname($T)
@doc """
    $fname.(img)

converts the raw storage type of `img` to `$Tname`, without changing the color space.
""" $fn

    end
end

"""
    float(x::Colorant)
    float(T::Type{<:Colorant})

convert the storage type of pixel `x` to a floating point data type while
preserving the `Colorant` information.

If the input is Type `T`, then it is equivalent to [`floattype`](@ref).
"""
float(x::Colorant) = floattype(typeof(x))(x)
float(::Type{T}) where T <: Colorant = floattype(T)