#################### Random Walk Metropolis ####################

#################### Types and Constructors ####################

mutable struct RWMTune <: SamplerTune
  logf::Union{Function, Missing}
  scale::Union{Float64, Vector{Float64}}
  proposal::SymDistributionType

  RWMTune() = new()

  function RWMTune(x::Vector, scale::Real, logf::Union{Function, Missing};
                   proposal::SymDistributionType=Normal)
    new(logf, Float64(scale), proposal)
  end

  function RWMTune(x::Vector, scale::Vector{T},
                  logf::Union{Function, Missing};
                  proposal::SymDistributionType=Normal) where {T<:Real}
    new(logf, convert(Vector{Float64}, scale), proposal)
  end
end

RWMTune(x::Vector, scale::ElementOrVector{T}; args...) where {T<:Real} =
  RWMTune(x, scale, missing; args...)

const RWMVariate = SamplerVariate{RWMTune}

validate(v::RWMVariate) = validate(v, v.tune.scale)

validate(v::RWMVariate, scale::Float64) = v

function validate(v::RWMVariate, scale::Vector)
  n = length(v)
  length(scale) == n ||
    throw(ArgumentError("length(scale) differs from variate length $n"))
  v
end

#################### Sampler Constructor ####################
"""
    RWM(params::ElementOrVector{Symbol},
                  scale::ElementOrVector{T}; args...) where {T<:Real})

Construct a `Sampler` object for RWM sampling. Parameters are assumed to be
continuous, but may be constrained or unconstrained.

Returns a `Sampler{RWMTune}` type object.

* `params`: stochastic node(s) to be updated with the sampler. Constrained parameters are mapped to unconstrained space according to transformations defined by the Stochastic `unlist()` function.

* `scale`: scaling value or vector of the same length as the combined elements of nodes `params` for the `proposal` distribution. Values are relative to the unconstrained parameter space, where candidate draws are generated.

* `args...`: additional keyword arguments to be passed to the `RWMVariate` constructor.
"""
function RWM(params::ElementOrVector{Symbol},
              scale::ElementOrVector{T}; args...) where {T<:Real}

  samplerfx = function(model::Model, block::Integer)
    block = SamplingBlock(model, block, true)
    v = SamplerVariate(block, scale; args...)
    sample!(v, x -> logpdf!(block, x))
    relist(block, v)
  end
  Sampler(params, samplerfx, RWMTune())
end

"""
    RWM(params::ElementOrVector{Symbol}, moves::Array{Symbol}; args...)

Construct the RWM sampler for Trees. If you set moves to :all it will use all
eligible moves to change the tree. These are currently:
NNI, SPR, Slide, Swing, :EdgeLength

Returns a `Sampler{RWMTune}` type object.
"""
function RWM(params::ElementOrVector{Symbol}, moves::ElementOrVector{Symbol}; args...)
  eligible = [:NNI, :SPR, :Slide, :Swing, :EdgeLength]
  to_use = Symbol[]
  if moves == :all
    to_use = eligible
  else
    for i in moves
      !(i in eligible) &&
        throw("$i is not an eligible tree move. The list of eligible tree moves is $eligible")
      push!(to_use, i)
    end
  end
  samplerfx = function(model::Model, block::Integer)
    block = SamplingBlock(model, block, true)
    # scale is not necessary for RWM on trees so just set it to 1
    v = SamplerVariate(block, 1.0; args...)
    sample!(v, x -> logpdf!(block, x), to_use)
    relist(block, v)
  end
  Sampler(params, samplerfx, RWMTune())
end

#################### Sampling Functions ####################

sample!(v::RWMVariate) = sample!(v, v.tune.logf)

"""
    sample!(v::RWMVariate, logf::Function, moves::Array{Symbol})

Propose a new tree by randomly performing a move from the ones specified in `moves`.

Returns `v` updated with simulated values and associated tuning parameters.
"""
function sample!(v::RWMVariate, logf::Function, moves::Array{Symbol})
  tree = v[1]
  tc = deepcopy(tree)
  move = rand(moves)
  if move == :NNI
    NNI!(tree)
  elseif move == :SPR
    tree = SPR(tree)
  elseif move == :Slide
    slide!(tree)
  elseif move == :Swing
    swing!(tree)
  elseif move == :EdgeLength
    change_edge_length!(tree)
  else
    throw("Tree move not elegible ")
  end
  if rand() < exp(logf(tree) - logf(tc))
    v[1] = tree
  else
    v[1] = tc
  end
  v
end

"""
    sample!(v::RWMVariate, logf::Function)

Draw one sample from a target distribution using the RWM sampler. Parameters
are assumed to be continuous and unconstrained.

Returns `v` updated with simulated values and associated tuning parameters.
"""
function sample!(v::RWMVariate, logf::Function)
  x = v + v.tune.scale .* rand(v.tune.proposal(0.0, 1.0), length(v))
  if rand() < exp(logf(x) - logf(v.value))
    v[:] = x
  end
  v
end


@generated function vrange(::Val{W}, ::Type{T}) where {W,T}
    Expr(:block, Expr(:meta, :inline), Expr(:call, :SVec, Expr(:tuple, [Expr(:call, :(Core.VecElement), T(w)) for w ∈ 0:W-1]...)))
end
@generated function vrangeincr(::Val{W}, i::I, ::Val{O}) where {W,I<:Integer,O}
    bytes = I === Int ? min(8, VectorizationBase.prevpow2(VectorizationBase.REGISTER_SIZE ÷ W)) : sizeof(I)
    # bytes = min(8, VectorizationBase.prevpow2(VectorizationBase.REGISTER_SIZE ÷ W))
    bits = 8bytes
    jtypesym = Symbol(:Int, bits)
    iexpr = bytes == sizeof(I) ? :i : Expr(:call, :%, :i, jtypesym)
    typ = "i$(bits)"
    vtyp = "<$W x $typ>"
    rangevec = join(("$typ $(w+O)" for w ∈ 0:W-1), ", ")
    instrs = String[]
    push!(instrs, "%ie = insertelement $vtyp undef, $typ %0, i32 0")
    push!(instrs, "%v = shufflevector $vtyp %ie, $vtyp undef, <$W x i32> zeroinitializer")
    push!(instrs, "%res = add nsw $vtyp %v, <$rangevec>")
    push!(instrs, "ret $vtyp %res")
    quote
        $(Expr(:meta,:inline))
        llvmcall(
            $(join(instrs,"\n")), Vec{$W,$jtypesym}, Tuple{$jtypesym}, $iexpr
        )
    end
end
@generated function vrangeincr(::Val{W}, i::T, ::Val{O}) where {W,T<:FloatingTypes,O}
    typ = llvmtype(T)
    vtyp = "<$W x $typ>"
    rangevec = join(("$typ $(w+O).0" for w ∈ 0:W-1), ", ")
    instrs = String[]
    push!(instrs, "%ie = insertelement $vtyp undef, $typ %0, i32 0")
    push!(instrs, "%v = shufflevector $vtyp %ie, $vtyp undef, <$W x i32> zeroinitializer")
    push!(instrs, "%res = fadd $vtyp %v, <$rangevec>")
    push!(instrs, "ret $vtyp %res")
    quote
        $(Expr(:meta,:inline))
        llvmcall(
            $(join(instrs,"\n")), Vec{$W,$T}, Tuple{$T}, i
        )
    end
end
@generated function vrangemul(::Val{W}, i::I, ::Val{O}) where {W,I<:Integer,O}
    bytes = I === Int ? min(8, VectorizationBase.prevpow2(VectorizationBase.REGISTER_SIZE ÷ W)) : sizeof(I)
    bits = 8bytes
    jtypesym = Symbol(:Int, bits)
    iexpr = bytes == sizeof(I) ? :i : Expr(:call, :%, :i, jtypesym)
    typ = "i$(bits)"
    vtyp = "<$W x $typ>"
    rangevec = join(("$typ $(w+O)" for w ∈ 0:W-1), ", ")
    instrs = String[]
    push!(instrs, "%ie = insertelement $vtyp undef, $typ %0, i32 0")
    push!(instrs, "%v = shufflevector $vtyp %ie, $vtyp undef, <$W x i32> zeroinitializer")
    push!(instrs, "%res = mul nsw $vtyp %v, <$rangevec>")
    push!(instrs, "ret $vtyp %res")
    quote
        $(Expr(:meta,:inline))
        llvmcall(
            $(join(instrs,"\n")), Vec{$W,$jtypesym}, Tuple{$jtypesym}, $iexpr
        )
    end
end
@generated function vrangemul(::Val{W}, i::T, ::Val{O}) where {W,T<:FloatingTypes,O}
    typ = llvmtype(T)
    vtyp = "<$W x $typ>"
    rangevec = join(("$typ $(w+O).0" for w ∈ 0:W-1), ", ")
    instrs = String[]
    push!(instrs, "%ie = insertelement $vtyp undef, $typ %0, i32 0")
    push!(instrs, "%v = shufflevector $vtyp %ie, $vtyp undef, <$W x i32> zeroinitializer")
    push!(instrs, "%res = fmul fast $vtyp %v, <$rangevec>")
    push!(instrs, "ret $vtyp %res")
    quote
        $(Expr(:meta,:inline))
        llvmcall(
            $(join(instrs,"\n")), Vec{$W,$T}, Tuple{$T}, i
        )
    end
end

@inline svrangeincr(::Val{W}, i, ::Val{O}) where {W,O} = SVec(vrangeincr(Val{W}(), i, Val{O}()))
@inline svrangemul(::Val{W}, i, ::Val{O}) where {W,O} = SVec(vrangemul(Val{W}(), i, Val{O}()))

@inline vrange(i::_MM{W}) where {W} = vrangeincr(Val{W}(), i.i, Val{0}())
@inline svrange(i::_MM{W}) where {W} = SVec(vrangeincr(Val{W}(), i.i, Val{0}()))
@inline Base.:(+)(i::_MM{W}, j::_MM{W}) where {W} = SVec(vadd(vrange(i), vrange(j)))
@inline Base.:(+)(i::_MM{W}, j::Vec{W}) where {W} = vadd(vrange(i), j)
@inline Base.:(+)(i::Vec{W}, j::_MM{W}) where {W} = vadd(i, vrange(j))
@inline Base.:(*)(i::_MM{W}, j::Vec{W}) where {W} = vmul(vrange(i), j)
@inline Base.:(*)(i::Vec{W}, j::_MM{W}) where {W} = vmul(i, vrange(j))
@inline Base.:(+)(i::_MM{W}, j::AbstractStructVec{W}) where {W} = SVec(vadd(vrange(i), extract_data(j)))
@inline Base.:(+)(i::AbstractStructVec{W}, j::_MM{W}) where {W} = SVec(vadd(extract_data(i), vrange(j)))
@inline Base.:(*)(i::_MM{W}, j::AbstractStructVec{W}) where {W} = SVec(vmul(vrange(i), extract_data(j)))
@inline Base.:(*)(i::AbstractStructVec{W}, j::_MM{W}) where {W} = SVec(vmul(extract_data(i), vrange(j)))
@inline vadd(i::_MM{W}, j::_MM{W}) where {W} = SVec(vadd(vrange(i), vrange(j)))
@inline vadd(i::_MM{W}, j::Vec{W}) where {W} = vadd(vrange(i), j)
@inline vadd(i::Vec{W}, j::_MM{W}) where {W} = vadd(i, vrange(j))
@inline vadd(i::_MM{W}, j::AbstractStructVec{W}) where {W} = SVec(vadd(vrange(i), extract_data(j)))
@inline vadd(i::AbstractStructVec{W}, j::_MM{W}) where {W} = SVec(vadd(extract_data(i), vrange(j)))
@inline vmul(i::_MM{W}, j::Vec{W}) where {W} = vmul(vrange(i), j)
@inline vmul(i::_MM{W}, j::AbstractStructVec{W}) where {W} = SVec(vmul(vrange(i), extract_data(j)))
@inline vmul(j::Vec{W}, i::_MM{W}) where {W} = vmul(j, vrange(i))
@inline vmul(j::AbstractStructVec{W}, i::_MM{W}) where {W} = SVec(vmul(extract_data(j), vrange(i)))
@inline Base.:(/)(i::_MM, j::T) where {T<:Number} = SVec(vfdiv(vrange(i,T), j))
@inline Base.:(/)(j::T, i::_MM) where {T<:Number} = SVec(vfdiv(j, vrange(i,T)))
@inline Base.:(/)(i::_MM, j::SVec{W,T}) where {W,T<:Number} = SVec(vfdiv(vrange(i,T), j))
@inline Base.:(/)(j::SVec{W,T}, i::_MM) where {W,T<:Number} = SVec(vfdiv(j, vrange(i,T)))
@inline Base.:(/)(i::_MM, j::_MM) = SVec(vfdiv(vrange(i), vrange(j)))
@inline Base.inv(i::_MM) = inv(svrange(i))

@inline vrange(::Val{W}) where {W} = vrange(Val{W}(), Float64)
@inline svrange(::Val{W}) where {W} = svrange(Val{W}(), Float64)

@inline vrange(i::_MM{W}, ::Type{T}) where {W,T} = vrangeincr(Val{W}(), T(i.i), Val{0}())
@inline vrange(i::_MM{W}, ::Type{T}) where {W,T <: Integer} = vrangeincr(Val{W}(), i.i % T, Val{0}())
@inline svrange(i::_MM, ::Type{T}) where {T} = SVec(vrange(i, T))

@inline Base.:(<<)(i::_MM, j::Integer) = svrange(i) << j
@inline Base.:(>>)(i::_MM, j::Integer) = svrange(i) >> j
@inline Base.:(>>>)(i::_MM, j::Integer) = svrange(i) >>> j

@inline Base.:(*)(i::_MM{W}, j::T) where {W,T} = vmul(svrange(i), j)
@inline Base.:(*)(j::T, i::_MM{W}) where {W,T} = vmul(svrange(i), j)
@inline vmul(i::_MM{W}, j::T) where {W,T} = vmul(svrange(i), j)
@inline vmul(j::T, i::_MM{W}) where {W,T} = vmul(svrange(i), j)
@inline vmul(i::_MM{W}, ::Static{j}) where {W,j} = vmul(svrange(i), j)
@inline vmul(::Static{j}, i::_MM{W}) where {W,j} = vmul(svrange(i), j)
@inline vconvert(::Type{SVec{W,T}}, i::_MM{W}) where {W,T} = svrange(i, T)

@inline Base.:(-)(i::Integer, j::_MM{W}) where {W} = vsub(i, svrange(j))
@inline Base.:(-)(::Static{i}, j::_MM{W}) where {W,i} = vsub(i, svrange(j))
@inline Base.:(-)(i::_MM{W}, j::_MM{W}) where {W} = vsub(svrange(i), svrange(j))
@inline Base.:(-)(i::_MM{W}) where {W} = -svrange(i)
@inline vsub(i::Integer, j::_MM{W}) where {W} = vsub(i, svrange(j))
@inline vsub(::Static{i}, j::_MM{W}) where {W,i} = vsub(i, svrange(j))
@inline vsub(i::_MM{W}, j::_MM{W}) where {W} = vsub(svrange(i), svrange(j))
@inline vsub(i::_MM{W}) where {W} = -svrange(i)

for op ∈ [:(<), :(>), :(≥), :(≤), :(==), :(!=), :(&), :(|), :(⊻), :(%)]
    @eval @inline Base.$op(i::_MM, j::Integer) = $op(svrange(i), j)
    @eval @inline Base.$op(i::Integer, j::_MM) = $op(i, svrange(j))
    @eval @inline Base.$op(i::_MM, ::Static{j}) where {j} = $op(svrange(i), j)
    @eval @inline Base.$op(::Static{i}, j::_MM) where {i} = $op(i, svrange(j))
    @eval @inline Base.$op(i::_MM, j::_MM) = $op(svrange(i), svrange(j))
end
@inline Base.:(*)(i::_MM, j::_MM) = SVec(vmul(vrange(i), vrange(j)))
@inline vmul(i::_MM, j::_MM) = SVec(vmul(vrange(i), vrange(j)))

using VectorizationBase: Static, Zero, One
@inline vadd(::_MM{W,Zero}, v::AbstractSIMDVector{W,T}) where {W,T} = vadd(vrange(Val{W}(), T), v)
@inline vadd(v::AbstractSIMDVector{W,T}, ::_MM{W,Zero}) where {W,T} = vadd(vrange(Val{W}(), T), v)
@inline vadd(::_MM{W,Zero}, ::_MM{W,Zero}) where {W} = vrangemul(Val{W}(), 2, Val{0}())
# @inline vmul(::_MM{W,Zero}, i) where {W} = svrangemul(Val{W}(), i, Val{0}())
# @inline vmul(i, ::_MM{W,Zero}) where {W} = svrangemul(Val{W}(), i, Val{0}())

@inline vmul(::_MM{W,Static{N}}, i) where {W,N} = svrangemul(Val{W}(), i, Val{N}())
@inline vmul(i, ::_MM{W,Static{N}}) where {W,N} = svrangemul(Val{W}(), i, Val{N}())


using ADI
using ProgressLogging
using PSFModels: Gaussian
using PyCall
using SiriusB

fits = pyimport("astropy.io.fits")
pro = pyimport("proplot")

pro.use_style("ggplot")
pro.rc["image.cmap"] = "inferno"
pro.rc["image.origin"] = "lower"
pro.rc["grid"] = false

function plot_results(res_cube, angles, curve; fwhm, epoch, label)
    # produce images
    flat_res = collapse(res_cube, angles)
    snr_map = detectionmap(flat_res, fwhm)
    sig_map = detectionmap(significance, flat_res, fwhm)
    stim_map = stimmap(res_cube, angles)

    arcsec = curve.distance .* SiriusB.pxscale

    ticks = range(20, 179, length=5)
    tick_labs = @. string(round(SiriusB.pxscale * (ticks - 99.5), digits=1))
    # start plotting
    # NB: proplot does tricky things with the Axes class that doesn't play nicely
    # with PyCall.jl because it's treated as a Vector in Julia, so use Python
    # syntax instead
    py"""
    import proplot as pro

    layout = [[1, 2, 5, 5],
              [3, 4, 5, 5]]
    fig, axs = pro.subplots(
        layout,
        width="7.5in",
        wspace=("3.25em", "2.5em", "0"),
        hspace="1.25em",
        share=0
    )
    axs[0].imshow($flat_res, colorbar="r", colorbar_kw=dict(space=0))
    axs[0].text(6, 6, "flat residual", color="w")
    axs[2].imshow($stim_map, vmin=0, vmax=1, colorbar="r", colorbar_kw=dict(space=0))
    axs[2].text(6, 6, "STIM prob.", color="w")
    axs[1].imshow($snr_map, vmin=0, vmax=5, colorbar="r", colorbar_kw=dict(space=0))
    axs[1].text(6, 6, "S/N", color="w")
    axs[3].imshow($sig_map, vmin=0, vmax=5, colorbar="r", colorbar_kw=dict(space=0))
    axs[3].text(6, 6, "significance", color="w")

    axs[0].format(xticks="null")
    axs[1].format(xticks="null", yticks="null")
    axs[3].format(yticks="null")
    axs[0].format(
        yticks=$ticks,
        yticklabels=$tick_labs,
        ylabel="y [arcsec]"
    )
    axs[2].format(
        yticks=$ticks,
        yticklabels=$tick_labs,
        ylabel="y [arcsec]"
    )
    axs[2:4].format(
        xticks=$ticks,
        xticklabels=$tick_labs,
        xlabel="x [arcsec]",
    )

    axs[4].plot($arcsec, $(curve.contrast), label="Gaussian", lw=1)
    axs[4].plot($arcsec, $(curve.contrast_corr), color="C0", ls="--", label="Student-t", lw=1)
    axs[4].fill_between($arcsec, $(curve.contrast), $(curve.contrast_corr), color="C0", alpha=0.2)
    axs[4].legend(ncol=1)
    axs[4].format(
        ytickloc="right",
        yformatter="log",
        xlabel="separation [arcsec]",
        ylabel="5σ contrast",
        yscale="log",
        grid=True,
        ylim=(10**(-3.5), 1)
    )
    fig.savefig($(figuredir("reports", make_filename_friendly("$(epoch)_$label") * ".pdf")))
    """
end

function make_filename_friendly(filename)
    tmp = lowercase(filename)
    tmp = replace(tmp, " " => "_")
    tmp = replace(tmp, "(" => "-")
    tmp = replace(tmp,  ")" => "")
end

_epochs = Dict("2020feb04" => "2020-02-04", "2020nov21" => "2020-11-21", "2020nov28" => "2020-11-28")

function plot_curves(contrast_curves; epoch)
    py"""
    import proplot as pro
    fig, axs = pro.subplots(width="7.5in", height="4in")
    """
    i = 0
    for (label, curve) in contrast_curves
        dist = curve.distance .* SiriusB.auscale
        color = "C$i"
        py"""
        axs.plot($dist, $(curve.contrast), color=$color, label=$label)
        axs.plot($dist, $(curve.contrast_corr), color=$color, ls="--")
        """
        i += 1
    end
    py"""
    axs.text(0.25, 5e-4, $(_epochs[epoch]), fontsize=14)
    axs.legend(ncol=2)
    axs.dualx(lambda x: x * $(SiriusB.parallax), label="separation [arcsec]")
    ax2 = axs.alty(
        yticks=[0.09, 0.31, 0.54, 0.77, 1],
        yticklabels=["8", "6", "4", "2", "0"],
        label="Δ mag"
    )

    axs.format(
        ylim=(10**(-3.5), 1),
        yscale="log",
        yformatter="log",
        xlabel="projected separation [AU]",
        ylabel="5σ contrast",
        grid=True
    )
    fig.savefig($(figuredir(make_filename_friendly("$(epoch)_contrast_curves") * ".pdf")))
    """
end

_ncomps = Dict(
    "2020feb04" => 1,
    "2020nov21" => 3,
    "2020nov28" => 2
)

@progress name="epoch" for epoch in ["2020feb04", "2020nov21", "2020nov28"]
    cube = fits.getdata(datadir("epoch_$epoch", "processed", "$(epoch)_sirius-b_cube_calib_registered_crop.fits"))
    angles = fits.getdata(datadir("epoch_$epoch", "processed", "$(epoch)_sirius-b_pa.fits"))

    raw_median = collapse(cube)
    res = gaussian_fit(raw_median, (100.5, 100.5))
    fwhm = res[3]

    psf_model = Gaussian{eltype(cube)}(res[1:2], fwhm)
    @info "FWHM for epoch $epoch is $fwhm px"
    starphot = Metrics.estimate_starphot(cube, fwhm)
    psfphot = Metrics.estimate_starphot(psf_model, fwhm) * exp(res[4])
    @info "star photometry is ≈ $starphot (PSF photometry was $psfphot)"

    # naturally mask out inner FWHM
    av_cube = AnnulusView(cube .- minimum(cube); inner=fwhm)
    # create annular view for potential reductions
    mav_cube = MultiAnnulusView(cube .- minimum(cube), fwhm; inner=fwhm)

    # orient algorithms with appropriate data structures
    ncomp = _ncomps[epoch]
    targets = [
        (Classic(), av_cube, "median"),
        (PCA(ncomp), av_cube, "PCA($ncomp)"),
        (NMF(ncomp), av_cube, "NMF($ncomp)"),
        (GreeDS(ncomp), av_cube, "GreeDS($ncomp)"),
        (Framewise(Classic(), delta_rot=0.5), mav_cube, "annular median"),
        (Framewise(PCA(ncomp), delta_rot=0.5), mav_cube, "annular PCA($ncomp)"),
        (Framewise(NMF(ncomp), delta_rot=0.5), mav_cube, "annular NMF($ncomp)"),
    ]

    contrast_curves = Dict{String, Any}()

    @progress "algorithm" for (alg, target, label) in targets
        @info label
        filename = make_filename_friendly(datadir("epoch_$epoch", "residuals", "$(epoch)_sirius-b_residual_$label.fits"))
        res_cube = load_or_produce(filename) do
            subtract(alg, target; angles, fwhm)
        end
        filename = make_filename_friendly(datadir("epoch_$epoch", "residuals", "$(epoch)_sirius-b_contrast-curve_$label.csv"))
        cc = load_or_produce(filename) do
            contrast_curve(alg, target, angles, psf_model; fwhm, angles, starphot, nbranch=6)
        end
        push!(contrast_curves, label => cc)
        plot_results(res_cube, angles, cc; fwhm, label, epoch)
    end

    plot_curves(contrast_curves; epoch)
end


"""
    AnalogInChannelCount(dev::DWFDevice)

Return the number of analog in channels.
"""
function AnalogInChannelCount(dev::DWFDevice)
    cChannel = Ref(Cint(0))
    FDwfAnalogInChannelCount(dev.hdwf[], cChannel)
    return cChannel[]
end

"""
    AnalogInFrequencyInfo(dev::DWFDevice)

Return the maximum sampling frequency of the analog in channels.
"""
function AnalogInFrequencyInfo(dev::DWFDevice)
    hzFreq = Ref(Cdouble(0))
    FDwfAnalogInFrequencyInfo(dev.hdwf[], C_NULL, hzFreq)
    return hzFreq[]
end


module SIMDPirates

using MacroTools: postwalk, @capture
using VectorizationBase
using VectorizationBase: llvmtype, AbstractSIMDVector, AbstractStructVec, AbstractStridedPointer

export  Vec, SVec, VE,
    @pirate,
    extract_data,
    vbroadcast,
    vload,
    vloada,
    vstore!,
    vstorea!,
    shufflevector,
    vifelse,
    vfma, vmuladd,
    vsqrt, rsqrt, vinv, vabs2,
    vadd, vsub, vmul, vfdiv, vsum, vprod,
    vfmadd, vfnmadd, vfmsub, vfnmsub,
    gather, scatter!

function vector_args(args)
    vecargs =       [isa(arg, Symbol) ? :($arg::Vec{N,T})             : arg for arg ∈ args]
    abstractargs =  [isa(arg, Symbol) ? :($arg::AbstractSIMDVector{N,T})    : arg for arg ∈ args]
    structargs =    [isa(arg, Symbol) ? :($arg::AbstractStructVec{N,T})     : arg for arg ∈ args]
    vecargs, abstractargs, structargs
end

function parse_func_decl(expr::Expr)
    @assert (expr.head === :(=) || expr.head === :function) "Failed to match function declaration: $expr."
    w = expr.args[1]
    body = expr.args[2]
    WT = @view(w.args[2:end])
    func = first(w.args)
    f = first(func.args)
    args = @view(func.args[2:end])
    f, args, WT, body
end
macro vectordef(rename, expr)
    f, args, R, body = parse_func_decl(expr)
    vecargs, abstractargs, structargs = vector_args(args)
    q = quote
        @inline $rename($(vecargs...)) where {$(R...)} = $body
        @inline $rename($(abstractargs...)) where {$(R...)} = SVec($body)
        @inline $f($(structargs...)) where {$(R...)} = SVec($body)
    end
    esc(q)
end
macro evectordef(rename, expr)
    f, args, R, body = parse_func_decl(expr)
    vecargs, abstractargs, structargs = vector_args(args)
    q = quote
        @inline $rename($(vecargs...)) where {$(R...)} = $body
        @inline $rename($(abstractargs...)) where {$(R...)} = SVec($body)
    end
    esc(q)
end

include("type_definitions.jl")
include("llvm_utils.jl")
include("llvmwrap.jl")
include("conditionals.jl")
include("integer_arithmetic.jl")
include("floating_point_arithmetic.jl")
include("memory.jl")
include("shufflevector.jl")
include("special.jl")
include("pirate.jl")
include("precompile.jl")
_precompile_()

function __init__()
    _precompile_()
end

end # module


#cd("/home/user/github/TOF-Tracer2/")

include("$(pwd())/startupAPiTOF.jl")

# fp = "./ExampleFiles/APiTOFDATA/" # All files in this path will be processed
#fp = "/media/wiebke/Extreme SSD/PSM_vs_PTR3/Data/apiTOFdata/bCP/"

# fp = "/media/wiebke/Extreme SSD/PSM_vs_PTR3/Data/apiTOFdata/bCP_14-07/1stFile/"
fp = "/media/wiebke/Extreme SSD/PSM_vs_PTR3/Data/apiTOFdata/bCP_13-07/"
# fp = "/media/wiebke/Extreme SSD/PSM_vs_PTR3/Data/apiTOFdata/NO3-/bCP/"

filefilterRegexp = r"\.h5$"
#rf = "$(fp)Data_2019_07_14-15_19_40.h5"  # The mass scale from this file defines the mass scale of all
# rf = "$(fp)Data_2019_07_14-11_19_39.h5"		# 14.7. first file
rf = "$(fp)Data_2019_07_13-17_34_29.h5"	# 13.7.
# rf = "$(fp)Data_2019_07_29-13_25_54.h5" # NO3-/bCP/

#masslist = MasslistFunctions.loadMasslist("./ExampleFiles/MASSLISTS/exampleMasslist.csv")
#masslist = MasslistFunctions.loadMasslist("/media/wiebke/Extreme SSD/PSM_vs_PTR3/Data/apiTOFdata/masslists/bCP_BAG_APItof.csv")
masslist = MasslistFunctions.loadMasslist("/media/wiebke/Extreme SSD/PSM_vs_PTR3/Data/apiTOFdata/bCP_13-07/masslist-JAN2021_2.csv")
# masslist = MasslistFunctions.loadMasslist("$(fp)results/Masslist.csv") # no3-CIMS
# cr = [37.028406 487.3424 693.546 739.515 959.697] # pos bCP, 14.-16.7.
cr = [37.028406 252.17256 634.344]	# pos bCP, 13.7
#cr = [37.028406 252.17256]
#cr = [62 125] # neg
# cr = [61.9878 187.979] # NO3--CIMS

# alternatively: use an auto generated masslist
# masslistMasses, masslistElements, masslistElementsMasses, masslistCompositions = MasslistFunctions.createMassList(;C=13:120, O=0:20, N=0:1, allowRadicals=true) #
(masslistMasses, masslistElements, masslistElementsMasses, masslistCompositions) = masslist
s = (masslistMasses.>17) .& ( masslistMasses.<1200)
masslistMasses = masslistMasses[s]
masslistCompositions = masslistCompositions[s,:]

####################### END OF SETTINGS ###############################################################

####################### Processing sequence ###########################################################

correctMassScaleAndExtractSumSpec(
    fp,
    masslistMasses,
    masslistElements,
    masslistElementsMasses,
    masslistCompositions,
    rf,
    cr,
    filefilterRegexp=filefilterRegexp,
    onlyUseAverages = false,
    plotControlMass = false,
    testRangeStart = 393.5, # the mass shift of this region will be shown if plot control mass is set true. Should not be part of calibRegions
    testRangeEnd = 396.5,
    recalibInterval = 900,
    resolution = 4500,
    firstNFiles= 0,
    lastNFiles = 0,
    UMRpeaks = true, # create Unit Mass Resolution peaks
    UMRmasses = length(s),
    UMRmassLow = 0.1,
    UMRmassHigh = 0.7,
    umrBinWidth =100,
    #binWidth = 10
    )
#
baselineAndPeakshape(
    fp,
    peakshapeRegions=6,
    peakshapeRegionStretch=1.0,
    peakshapeQuantileValue = 0.7,
    peakfindingNoiseThresholdValue = 5,
    peakWindowWidth = 100,
    peakfindingSignalLimit=0.5
    )
    #
#
mtrx = deconvolute(
    fp,
    calcTransposed = false
    )
#


# Julia wrapper for header: /usr/include/scip/cons_abspower.h
# Automatically generated using Clang.jl wrap_c

function SCIPincludeConshdlrAbspower(scip)
    ccall((:SCIPincludeConshdlrAbspower, libscip), SCIP_RETCODE, (Ptr{SCIP_},), scip)
end

function SCIPcreateConsAbspower(scip, cons, name, x, z, exponent, xoffset, zcoef, lhs, rhs, initial, separate, enforce, check, propagate, _local, modifiable, dynamic, removable, stickingatnode)
    ccall((:SCIPcreateConsAbspower, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{Ptr{SCIP_CONS}}, Cstring, Ptr{SCIP_VAR}, Ptr{SCIP_VAR}, Cdouble, Cdouble, Cdouble, Cdouble, Cdouble, UInt32, UInt32, UInt32, UInt32, UInt32, UInt32, UInt32, UInt32, UInt32, UInt32), scip, cons, name, x, z, exponent, xoffset, zcoef, lhs, rhs, initial, separate, enforce, check, propagate, _local, modifiable, dynamic, removable, stickingatnode)
end

function SCIPcreateConsBasicAbspower(scip, cons, name, x, z, exponent, xoffset, zcoef, lhs, rhs)
    ccall((:SCIPcreateConsBasicAbspower, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{Ptr{SCIP_CONS}}, Cstring, Ptr{SCIP_VAR}, Ptr{SCIP_VAR}, Cdouble, Cdouble, Cdouble, Cdouble, Cdouble), scip, cons, name, x, z, exponent, xoffset, zcoef, lhs, rhs)
end

function SCIPgetNlRowAbspower(scip, cons, nlrow)
    ccall((:SCIPgetNlRowAbspower, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_CONS}, Ptr{Ptr{SCIP_NLROW}}), scip, cons, nlrow)
end

function SCIPgetNonlinearVarAbspower(scip, cons)
    ccall((:SCIPgetNonlinearVarAbspower, libscip), Ptr{SCIP_VAR}, (Ptr{SCIP_}, Ptr{SCIP_CONS}), scip, cons)
end

function SCIPgetLinearVarAbspower(scip, cons)
    ccall((:SCIPgetLinearVarAbspower, libscip), Ptr{SCIP_VAR}, (Ptr{SCIP_}, Ptr{SCIP_CONS}), scip, cons)
end

function SCIPgetExponentAbspower(scip, cons)
    ccall((:SCIPgetExponentAbspower, libscip), Cdouble, (Ptr{SCIP_}, Ptr{SCIP_CONS}), scip, cons)
end

function SCIPgetOffsetAbspower(scip, cons)
    ccall((:SCIPgetOffsetAbspower, libscip), Cdouble, (Ptr{SCIP_}, Ptr{SCIP_CONS}), scip, cons)
end

function SCIPgetCoefLinearAbspower(scip, cons)
    ccall((:SCIPgetCoefLinearAbspower, libscip), Cdouble, (Ptr{SCIP_}, Ptr{SCIP_CONS}), scip, cons)
end

function SCIPgetLhsAbspower(scip, cons)
    ccall((:SCIPgetLhsAbspower, libscip), Cdouble, (Ptr{SCIP_}, Ptr{SCIP_CONS}), scip, cons)
end

function SCIPgetRhsAbspower(scip, cons)
    ccall((:SCIPgetRhsAbspower, libscip), Cdouble, (Ptr{SCIP_}, Ptr{SCIP_CONS}), scip, cons)
end

function SCIPgetViolationAbspower(scip, cons, sol)
    ccall((:SCIPgetViolationAbspower, libscip), Cdouble, (Ptr{SCIP_}, Ptr{SCIP_CONS}, Ptr{SCIP_SOL}), scip, cons, sol)
end


function Inverse()
    return Inverse(())
end

function value(obj::Inverse, arg0::DerivativeStructure)
    return jcall(obj, "value", DerivativeStructure, (DerivativeStructure,), arg0)
end

function value(obj::Inverse, arg0::jdouble)
    return jcall(obj, "value", jdouble, (jdouble,), arg0)
end


using Parameters, Random
import Distributions: logpdf, loglikelihood

struct Paired{T1,T2} <: ContinuousUnivariateDistribution
    d::T1
    parms::T2
end

Paired(;d, parms) = Paired(d, parms)

loglikelihood(d::Paired, data::Array{<:NamedTuple,1}) = logpdf(d, data)

function logpdf(d::Paired, data::Array{<:NamedTuple,1})
    LL = computeLL(data, d.parms; d=d.d)
    return LL
end

function sample_stimuli(stimuli, N)
    return stimuli[1:N]
end

function simulate(all_stimuli, fixed_parms, n_blocks, n_trials, deadline=5.0, isi=5.0; d)
    # sample unique stimului for the simulation
    stimuli = sample_stimuli(all_stimuli, n_trials)
    # initialize declarative memory
    chunks = [Chunk(word=:_,number=1)] |> empty!
    memory = Declarative(;memory=chunks)
    # create act-r object
    actr = ACTR(;declarative=memory, fixed_parms..., d)
    # preallocate data
    data = Array{NamedTuple,1}()
    # iterate through all blocks
    for block in 1:n_blocks
        # randomize stimulus order
        shuffle!(stimuli)
        temp = simulate_block(actr, stimuli, block, deadline, isi)
        push!(data, temp...)
    end
    return vcat(data...)
end

function simulate_block(actr, stimuli, block, deadline=5.0, isi=5.0)
    # initialize data for the current block
    data = Array{NamedTuple,1}()
    start_time = get_time(actr)
    for stimulus in stimuli
        temp = simulate_trial(actr, stimulus, block, deadline)
        push!(data, temp)
        # 5 second response deadline + additional 5 seconds
        start_time += (deadline + isi)
        set_time!(actr, start_time)
    end
    return data
end

function simulate_trial(actr, stimulus, block, deadline=5.0)
    N=0;L=0.0;recent=Float64[];retrieved=:failed;time_created=0.0
    trial_start = get_time(actr)
    # encode and conflict resolution time
    e_time = 0.085 + 0.05 + 0.05
    add_time!(actr, e_time)
    # retrieve paired associate based on word
    chunk = retrieve(actr; word=stimulus.word)
    # retrieval time
    r_time = compute_RT(actr, chunk)
    # add encoding time and retrieval time to reaction time
    rt = e_time + r_time + .05 + 0.300
    # interrupt retrieval if exceeds deadline
    if rt > deadline
        # set to truncated value
        rt = deadline
        # advance current time to deadline
        set_time!(actr, trial_start + deadline)
        retrieved = :truncated
        # retrieval info if no chunk is retrieved
        chunk = get_chunks(actr; stimulus...)
        if !isempty(chunk)
            # get chunk state at time of retrieval
            @unpack N,L,recent,time_created = chunk[1]
            recent = copy(recent)
        end
    elseif !isempty(chunk)
        # add retreival time to current time
        add_time!(actr, r_time)
        # next conflict resolution
        add_time!(actr, .05)
        retrieved = :retrieved
        # get chunk state at time of retrieval
        @unpack N,L,recent,time_created = chunk[1]
        recent = copy(recent)
        # add chunk if new, or update N and time stamps
        add_chunk!(actr; stimulus...)
    else
        # add retreival failure time to current time
        add_time!(actr, r_time)
        # next conflict resolution
        add_time!(actr, .05)
        # retrieval info if no chunk is retrieved
        chunk = get_chunks(actr; stimulus...)
        if !isempty(chunk)
            # get chunk state at time of retrieval
            @unpack N,L,recent,time_created = chunk[1]
            recent = copy(recent)
        end
    end
    # feedback is always presented after the deadline
    set_time!(actr, trial_start + deadline)
    # encode feedback
    add_time!(actr, 0.05 + 0.05 + 0.085)
    # add chunk: create new chunk if does not exist, otherwise increment N and add time stamp
    add_chunk!(actr; stimulus...)
    data = (stimulus=stimulus,N=N,L=L,block=block,time_created=time_created,
        recent=recent,rt=rt,retrieved=retrieved)
    return data
end

function computeLL(data, fixed_parms; d)
    chunk = Chunk(act=zero(d))
    # initialize declarative memory
    memory = Declarative(;memory=[chunk])
    # create ACTR object with declarative memory and parameters
    actr = ACTR(;declarative=memory, fixed_parms..., d)
    # extract parameters
    @unpack s,lf,τ,ter=actr.parms
    # do not add noise activation values
    actr.parms.noise = false
    σ = s * pi / sqrt(3)
    # initialize log likelihood
    LL = 0.0
    for k in data
        cur_time = k.L + k.time_created
        if k.retrieved == :retrieved
            # reproduce the memory state right before retrieval
            modify!(chunk; N=k.N, recent=k.recent, time_created=k.time_created)
            # compute activation
            compute_activation!(actr, chunk, cur_time)
            # adjust activation values for latency factor
            v = [chunk.act,τ] .- log(lf)
            # log normal race distribution object
            dist = LNR(-v, σ, ter)
            # compute log likelihood of retrieving chunk at time k.rt
            LL += logpdf(dist, 1, k.rt)
        elseif k.retrieved == :truncated
            # special case in which chunk has not been created yet
            if k.N == 0
                v = τ .- log(lf)
                # compute log likelihood of failing to respond within deadline
                LL += logccdf(LogNormal(-v, σ), 5)
            else
                # reproduce the memory state right before retrieval
                modify!(chunk; N=k.N, recent=k.recent, time_created=k.time_created)
                # compute activation
                compute_activation!(actr, chunk, cur_time)
                # adjust activation values for latency factor
                v = [chunk.act,τ] .- log(lf)
                # compute log likelihood of failing to respond within deadline
                LL += logccdf(LogNormal(-v[1], σ), 5) + logccdf(LogNormal(-v[2], σ), 5)
            end
        else
            # retrieval failure for novel pair
            if k.N == 0
                dist = LNR(-[τ-log(lf)], σ, ter)
                LL += logpdf(dist, 1, k.rt)
            else
                # reproduce the memory state immediately prior to retrieval
                modify!(chunk; N=k.N, recent=k.recent, time_created=k.time_created)
                # compute activation
                compute_activation!(actr, chunk, cur_time)
                # mean activation values
                v = [chunk.act,τ] .- log(lf)
                # log normal race distribution object
                dist = LNR(-v, σ, ter)
                # compute log likelihood of failing to retrieve chunk at time k.rt
                LL += logpdf(dist, 2, k.rt)
            end
        end
    end
    return LL
end

function show_learning(all_stimuli, parms, n_blocks=8, n_trials=20; d)
    data = simulate(all_stimuli, parms, n_blocks, n_trials; d)
    df = DataFrame(data)
    groups = groupby(df, [:block,:retrieved])
    return combine(groups, :rt=>mean=>:rt)
end


"""
Voxel corner and edge indexing conventions

        Z
        |

        5------5------6              Extra edges not drawn
       /|            /|              -----------
      8 |           6 |              - face diagonals
     /  9          /  10                - 13: 1 to 3
    8------7------7   |                 - 14: 1 to 8
    |   |         |   |                 - 15: 1 to 6
    |   1------1--|---2  -- Y           - 16: 5 to 7
    12 /          11 /                  - 17: 2 to 7
    | 4           | 2                   - 18: 4 to 7
    |/            |/                 - body diagonal
    4------3------3                     - 19: 1 to 7

  /
 X
"""
# this gets vectorized so we want to ensure it is the
# same type as out vertex
@inline function voxCrnrPos(::Type{PT}) where {PT}
    (PT(0, 0, 0),
    PT(0, 1, 0),
    PT(1, 1, 0),
    PT(1, 0, 0),
    PT(0, 0, 1),
    PT(0, 1, 1),
    PT(1, 1, 1),
    PT(1, 0, 1))
end

const voxCrnrPosInt = voxCrnrPos(SVector{3,UInt8})

# the voxel IDs at either end of the tetrahedra edges, by edge ID
const voxEdgeCrnrs = ((0x01, 0x02),
                      (0x02, 0x03),
                      (0x04, 0x03),
                      (0x01, 0x04),
                      (0x05, 0x06),
                      (0x06, 0x07),
                      (0x08, 0x07),
                      (0x05, 0x08),
                      (0x01, 0x05),
                      (0x02, 0x06),
                      (0x03, 0x07),
                      (0x04, 0x08),
                      (0x01, 0x03),
                      (0x01, 0x08),
                      (0x01, 0x06),
                      (0x05, 0x07),
                      (0x02, 0x07),
                      (0x04, 0x07),
                      (0x01, 0x07))

# direction codes:
# 0 => +x, 1 => +y, 2 => +z,
# 3 => +xy, 4 => +xz, 5 => +yz, 6 => +xyz
const voxEdgeDir = (0x01, 0x00, 0x01, 0x00, 0x01, 0x00, 0x01, 0x00, 0x02, 0x02, 0x02, 0x02, 0x03, 0x04, 0x05, 0x03, 0x04, 0x05, 0x06)

# For a pair of corner IDs, the edge ID joining them
# 0 denotes a pair with no edge
const voxEdgeIx = ((0x00, 0x01, 0x0d, 0x04, 0x09, 0x0f, 0x13, 0x0e),
                   (0x01, 0x00, 0x02, 0x00, 0x00, 0x0a, 0x11, 0x00),
                   (0x0d, 0x02, 0x00, 0x03, 0x00, 0x00, 0x0b, 0x00),
                   (0x04, 0x00, 0x03, 0x00, 0x00, 0x00, 0x12, 0x0c),
                   (0x09, 0x00, 0x00, 0x00, 0x00, 0x05, 0x10, 0x08),
                   (0x0f, 0x0a, 0x00, 0x00, 0x05, 0x00, 0x06, 0x00),
                   (0x13, 0x11, 0x0b, 0x12, 0x10, 0x06, 0x00, 0x07),
                   (0x0e, 0x00, 0x00, 0x0c, 0x08, 0x00, 0x07, 0x00))

# voxel corners that comprise each of the six tetrahedra
const subTets = ((0x01, 0x03, 0x02, 0x07),
                 (0x01, 0x08, 0x04, 0x07),
                 (0x01, 0x04, 0x03, 0x07),
                 (0x01, 0x02, 0x06, 0x07),
                 (0x01, 0x05, 0x08, 0x07),
                 (0x01, 0x06, 0x05, 0x07))
# tetrahedron corners for each edge (indices 1-4)
const tetEdgeCrnrs = ((0x01, 0x02),
                      (0x02, 0x03),
                      (0x01, 0x03),
                      (0x01, 0x04),
                      (0x02, 0x04),
                      (0x03, 0x04))

# triangle cases for a given tetrahedron edge code
const tetTri = ((0x01, 0x03, 0x04, 0x00, 0x00, 0x00),
                (0x01, 0x05, 0x02, 0x00, 0x00, 0x00),
                (0x03, 0x05, 0x02, 0x03, 0x04, 0x05),
                (0x02, 0x06, 0x03, 0x00, 0x00, 0x00),
                (0x01, 0x06, 0x04, 0x01, 0x02, 0x06),
                (0x01, 0x05, 0x06, 0x01, 0x06, 0x03),
                (0x04, 0x05, 0x06, 0x00, 0x00, 0x00),
                (0x04, 0x06, 0x05, 0x00, 0x00, 0x00),
                (0x01, 0x06, 0x05, 0x01, 0x03, 0x06),
                (0x01, 0x04, 0x06, 0x01, 0x06, 0x02),
                (0x02, 0x03, 0x06, 0x00, 0x00, 0x00),
                (0x03, 0x02, 0x05, 0x03, 0x05, 0x04),
                (0x01, 0x02, 0x05, 0x00, 0x00, 0x00),
                (0x01, 0x04, 0x03, 0x00, 0x00, 0x00))


module ImageUtilities

export save_image_grid
export image_grid_figure
export image_grid_tensor
export image_grid
export color_image
export color_image_tensor

using FluxGAN: GAN

using Flux: cpu
using CairoMakie
using Images
using FileIO
using Dates
using EllipsisNotation
using Random

###########################
# image utility functions #
###########################

# save image grid figure from model
#
function save_image_grid(model::GAN, output_dir::String;
                         layout=(5, 5),
                         date=false,
                         file_info=[],
                         img_res=150,
                         kws...)
    if file_info == []
        info = ""
    else
        info = mapreduce(tag -> string(tag) * "_", *, file_info)
    end
    file = info * "grid_$(layout[1])_$(layout[2])" *
           (date ? "_" * string(today()) : "")
    imgs = image_grid_tensor(model, *(layout...); kws...)
    fig = image_grid_figure(imgs, layout; img_res=img_res)
    save(output_dir * "/" * file * ".png", fig)
end

# generate and return color images tensor from model
#
function image_grid_tensor(model::GAN, n::Int;
                           uncenter=true,
                           seed=false)
    @assert model.hparams.img_size ≠ undef "Image size not defined!"
    if seed Random.seed!(69) end
    imgs = cpu(model.G)(randn(Float32, model.hparams.latent_dim, n))
    if uncenter imgs = @. (imgs + 1f0) / 2f0 end
    imgs = reshape(imgs, model.hparams.img_size..., :)
    color_image_tensor(imgs)
end

# return image grid makie figure from image grid tensor
#
function image_grid_figure(imgs::Array, layout::Tuple{Int,Int};
                           img_res=150,
                           hflip=true)
    n = size(imgs)[end]
    rows, cols = layout
    @assert rows * cols == n "Number of images ≠ grid size!"

    imgs = [imgs[.., i] for i in 1:n]

    if hflip imgs = map(img -> reverse(img, dims=2), imgs) end

    fig = Figure(resolution=(cols * img_res, rows * img_res))
    for i = 1:rows, j = 1:cols
        ax, = image(fig[i,j], imgs[i + rows * (j - 1)])
        hidedecorations!(ax)
        hidexdecorations!(ax, ticks=false)
        hideydecorations!(ax, ticks=false)
    end
    fig
end

# return image grid from GAN and layout
#
function image_grid(model::GAN, layout::Tuple{Int,Int}; kws...)
    imgs = image_grid_tensor(model, *(layout...); kws...)
    image_grid(imgs, layout)
end

# convert images tensor into an image grid
#
function image_grid(imgs::Array{Color, 3}, layout::Tuple{Int,Int})
    h, w, n = size(imgs)
    rows, cols = layout
    @assert rows * cols == n "Number of images ≠ grid size!"
    grid = Array{Color}(undef, h * rows, w * cols)
    for i = 1:rows, j = 1:cols
        img = imgs[:, :, i + rows*(j - 1)]
        grid[((i - 1)*h + 1):i*h, ((j - 1)*w + 1):j*w] = img'
    end
    grid
end

# return color RGB image from image array
#
function color_image(img::Array)
    @assert length(size(img)) ∈ [2, 3] "Image tensor dimension ∉ [2, 3]!"
    if length(size(img)) == 3
        img = permutedims(img, (3, 1, 2))
        img = Matrix(colorview(RGB, img))
    else
        img = Matrix(colorview(Gray, img))
    end
    img
end

# convert tensor of image data into color image tensor
#
function color_image_tensor(imgs::Array)
    color_imgs = Array{Color}(undef, size(imgs)[1:2]..., size(imgs)[end])
    for i = 1:size(imgs)[end]
        color_imgs[:, :, i] = color_image(imgs[.., i])
    end
    color_imgs
end

end


abstract type AbstractDeviceArray{T, N} <: AbstractArray{T, N} end

Base.IndexStyle(::AbstractDeviceArray) = IndexLinear()

@inline function Base.iterate(A::AbstractDeviceArray, i=1)
    if (i % UInt) - 1 < length(A)
        (@inbounds A[i], i + 1)
    else
        nothing
    end
end

function Base.sum(A::AbstractDeviceArray{T}) where T
    acc = zero(T)
    for elem in A
        acc += elem
    end
    acc
end

const shmem_counter = Ref{Int}(0)

"""
Creates a local static memory shared inside one block.
Equivalent to `__local` of OpenCL or `__shared__ (<variable>)` of CUDA.
"""
macro LocalMemory(state, T, N)
    id = (shmem_counter[] += 1)
    quote
        lémem = LocalMemory($(esc(state)), $(esc(T)), Val($(esc(N))), Val($id))
        AbstractDeviceArray(lémem, $(esc(N)))
    end
end

export @LocalMemory

"""
Creates a block local array pointer with `T` being the element type
and `N` the length. Both T and N need to be static! C is a counter for
approriately get the correct Local mem id in CUDAnative.
This is an internal method which needs to be overloaded by the GPU Array backends
"""
function LocalMemory(state, ::Type{T}, ::Val{N}, ::Val{C}) where {N, T, C}
    error("Not implemented")
end


"""
    CountEdgesTask

This module includes the composite type CountEdgesTask which includes both the
        generic DaemonTask.Info and an AbstractString as a payload
"""
module CountEdgesTask

using JSON
using PyCall
@pyimport neuroglancer.pipeline as pl

using ...SimpleTasks.Types
using ContactAnalysis
using ContactAnalysis.Precomputed.PrecomputedWrapper

import SimpleTasks.Tasks.DaemonTask
import SimpleTasks.Tasks.BasicTask
import SimpleTasks.Services.Datasource

export CountEdgesTaskDetails, NAME, execute

function str_to_slices(s)
    m=match(r"(\d+)-(\d+)_(\d+)-(\d+)_(\d+)-(\d+)",s)
    bounds = map(x->parse(Int,x), m.captures)
    return (bounds[1]:bounds[2], bounds[3]:bounds[4], bounds[5]:bounds[6])
end

function slices_to_str(s)
    t=map(slice_to_str,s)
    return "$(t[1])_$(t[2])_$(t[3])"
end

function slice_to_str(s)
    return "$(s.start)-$(s.stop)"
end

type CountEdgesPayloadInfo
    segmentation_storage::AbstractString
    slices::AbstractString
    scale_idx::Int64
    output_storage::AbstractString
end

function CountEdgesPayloadInfo(d::Dict{String, Any})
    return CountEdgesPayloadInfo(d["segmentation_storage"],
                                 d["slices"],
                                 d["scale_idx"],
                                 d["output_storage"])
end

type CountEdgesTaskDetails <: DaemonTaskDetails
    basic_info::BasicTask.Info
    payload_info::CountEdgesPayloadInfo
end

"""
Parse task from JSON
"""
function CountEdgesTaskDetails(basic_info::BasicTask.Info, d::Dict{String, Any})
    payload_info = CountEdgesPayloadInfo(d)
    return CountEdgesTaskDetails(basic_info, payload_info)
end

const NAME = "CountEdgesTask"

function DaemonTask.prepare(task::CountEdgesTaskDetails,
                                datasource::DatasourceService)
    return true
end

function DaemonTask.execute(task::CountEdgesTaskDetails,
                                datasource::DatasourceService)

    if length(task.payload_info.segmentation_storage) == ""
        return DaemonTask.Result(true, [])
    end

    segmentation_storage = task.payload_info.segmentation_storage
    str_slices = task.payload_info.slices
    slices = str_to_slices(str_slices)
    scale_idx = task.payload_info.scale_idx
    output_storage = task.payload_info.output_storage

    seg = PrecomputedWrapper(segmentation_storage, scale_idx)
    res = seg.val[:_scale]["resolution"]/1000 # convert nm to um
    weights = []
    for (i,r) in enumerate(res) # calculate surface areas of faces
        push!(weights, prod(res[filter(a->!(i in a), 1:length(res))]))
    end
    d = ContactAnalysis.count_edges(seg[slices...], weights)

    out = pl.Storage(output_storage)
    out[:put_file](file_path="$(str_slices)_edge_counts.json",
                        content="$(JSON.json(d))")
    out[:wait]()

    return DaemonTask.Result(true, [str_slices])
end

function DaemonTask.finalize(task::CountEdgesTaskDetails,
                datasource::DatasourceService, result::DaemonTask.Result)
    return true
end

end # module CountEdgesTask


# note: this example only works on pure julia implementations of
# preprocessing elements and learners. Scikitlearn and other
# PyCall functions are not thread-safe and cannot be used inside
# the threads loop.

using AutoMLPipeline
using AutoMLPipeline.Utils
using Base.Threads

begin
  profbdata = getprofb()
  X = profbdata[:,2:end]
  Y = profbdata[:,1] |> Vector;
  head(x)=first(x,5)
  head(profbdata)
end

#### Column selector
catf = CatFeatureSelector();
numf = NumFeatureSelector()
ohe = OneHotEncoder()
#### Learners
rf = RandomForest();
ada = Adaboost()
dt=PrunedTree()

accuracy(X,Y)=score(:accuracy,X,Y)

acc=[]
learners=[rf,ada,dt]
@threads for i in 1:30
  @threads for lr in learners
    println(lr.name)
    pipe=@pipeline ((catf |> ohe) +(numf )) |> lr
    m=crossvalidate(pipe,X,Y,accuracy,10,true)
    push!(acc,(m...,name=lr.name))
    println(m)
  end
end
acc


using Functors, Test

struct Foo
  x
  y
end
@functor Foo

struct Bar
  x
end
@functor Bar

struct Baz
  x
  y
  z
end
@functor Baz (y,)

struct NoChildren
  x
  y
end

@testset "Nested" begin
  model = Bar(Foo(1, [1, 2, 3]))

  model′ = fmap(float, model)

  @test model.x.y == model′.x.y
  @test model′.x.y isa Vector{Float64}
end

@testset "Exclude" begin
  f(x::AbstractArray) = x
  f(x::Char) = 'z'

  x = ['a', 'b', 'c']
  @test fmap(f, x)  == ['z', 'z', 'z']
  @test fmap(f, x; exclude = x -> x isa AbstractArray) == x

  x = (['a', 'b', 'c'], ['d', 'e', 'f'])
  @test fmap(f, x)  == (['z', 'z', 'z'], ['z', 'z', 'z'])
  @test fmap(f, x; exclude = x -> x isa AbstractArray) == x
end

@testset "Walk" begin
  model = Foo((0, Bar([1, 2, 3])), [4, 5])

  model′ = fmapstructure(identity, model)
  @test model′ == (; x=(0, (; x=[1, 2, 3])), y=[4, 5])
end

@testset "Property list" begin
  model = Baz(1, 2, 3)
  model′ = fmap(x -> 2x, model)

  @test (model′.x, model′.y, model′.z) == (1, 4, 3)
end

@testset "fcollect" begin
  m1 = [1, 2, 3]
  m2 = 1
  m3 = Foo(m1, m2)
  m4 = Bar(m3)
  @test all(fcollect(m4) .=== [m4, m3, m1, m2])
  @test all(fcollect(m4, exclude = x -> x isa Array) .=== [m4, m3, m2])
  @test all(fcollect(m4, exclude = x -> x isa Foo) .=== [m4])

  m1 = [1, 2, 3]
  m2 = Bar(m1)
  m0 = NoChildren(:a, :b)
  m3 = Foo(m2, m0)
  m4 = Bar(m3)
  println(fcollect(m4))
  @test all(fcollect(m4) .=== [m4, m3, m2, m1, m0])
end


using FactCheck
using FunctionalCollections

import FunctionalCollections: SparseBitmappedTrie, SparseNode, SparseLeaf,
    bitpos, index, hasindex, arrayof, update

facts("Sparse Bitmapped Vector Tries") do

    context("bitpos and index") do
        l = SparseLeaf{Int}([6, 11, 16, 21, 26],
                            2^5 | 2^10 | 2^15 | 2^20 | 2^25)
        @fact index(l, 6) --> 1
        @fact index(l, 11) --> 2
        @fact index(l, 16) --> 3
        @fact index(l, 21) --> 4
        @fact index(l, 26) --> 5

        @fact hasindex(l, 5)  --> false
        @fact hasindex(l, 6)  --> true
        @fact hasindex(l, 7)  --> false
        @fact hasindex(l, 11) --> true
        @fact hasindex(l, 12) --> false
        @fact hasindex(l, 21) --> true
    end

    context("SparseLeaf update") do
        l = SparseLeaf{Int}()

        @fact update(l, 1, 1)   --> (leaf) -> hasindex(leaf[1], 1)
        @fact update(l, 10, 10) --> (leaf) -> hasindex(leaf[1], 10) &&
                                             !hasindex(leaf[1], 11)

        l = SparseLeaf{Int}([1, 5], 2^0 | 2^4)
        l, _ = update(l, 2, 2)
        @fact arrayof(l) --> [1, 2, 5]
        @fact length(l) --> 3
        @fact hasindex(l, 1) --> true
        @fact hasindex(l, 2) --> true
        @fact hasindex(l, 3) --> false
        @fact hasindex(l, 4) --> false
        @fact hasindex(l, 5) --> true
        @fact hasindex(l, 6) --> false

        @fact arrayof(update(l, 2, 100)[1]) --> [1, 100, 5]
    end

    context("SparseNode update") do
        n, _ = update(SparseNode(AbstractString), 1, "foo")
        @fact length(arrayof(n)) --> 1

        leaf = arrayof(n)[1].arr[1].arr[1].arr[1].arr[1].arr[1].arr[1]
        @fact hasindex(leaf, 1) --> true
        @fact leaf.arr[1] --> "foo"

        n2, _ = update(n, 33, "bar")
        leaf2 = n2.arr[1].arr[1].arr[1].arr[1].arr[1].arr[1].arr[2]
        @fact hasindex(leaf2, 33) --> true
        @fact arrayof(leaf2)[1] --> "bar"
    end

    context("SparseBitmappedTrie get") do
        n, _ = update(SparseNode(Int), 33, 33)
        @fact get(n, 33, "missing") --> 33
        @fact get(update(SparseNode(Int), 12345, 12345)[1], 12345, "missing") --> 12345
        @fact get(n, 12345, "missing") --> "missing"
    end

    context("SparseBitmappedTrie length") do
        n = SparseNode(Int)
        for i=1:1000
            n, _ = update(n, i, i)
        end
        @fact length(n) --> 1000
    end

    context("SparseBitmappedTrie items") do
        n = SparseNode(Int)
        for i=1:1000
            n, _ = update(n, i, i)
        end
        @fact [i for i=n] --> collect(1:1000)
    end

end