"""

    conv4(w, x; kwargs...)

Execute convolutions or cross-correlations using filters specified
with `w` over tensor `x`.

Currently KnetArray{Float32/64,4/5} and Array{Float32/64,4} are
supported as `w` and `x`.  If `w` has dimensions `(W1,W2,...,I,O)` and
`x` has dimensions `(X1,X2,...,I,N)`, the result `y` will have
dimensions `(Y1,Y2,...,O,N)` where

    Yi=1+floor((Xi+2*padding[i]-Wi)/stride[i])

Here `I` is the number of input channels, `O` is the number of output
channels, `N` is the number of instances, and `Wi,Xi,Yi` are spatial
dimensions.  `padding` and `stride` are keyword arguments that can be
specified as a single number (in which case they apply to all
dimensions), or an array/tuple with entries for each spatial
dimension.

# Keywords

* `padding=0`: the number of extra zeros implicitly concatenated at the start and at the end of each dimension.
* `stride=1`: the number of elements to slide to reach the next filtering window.
* `upscale=1`: upscale factor for each dimension.
* `mode=0`: 0 for convolution and 1 for cross-correlation.
* `alpha=1`: can be used to scale the result.
* `algo=0`: specifies which convolution algorithm shoud be used to compute the results. See the CUDNN User Guide for details.
* `workSpace=C_NULL`: data pointer to GPU memory to a workspace needed to able to execute the specified algorithm.
* `workSpaceSizeInBytes=0`: the size in bytes of the provided workSpace. Default=0.
* `handle`: handle to a previously created cuDNN context. Defaults to a Knet allocated handle.

"""
function conv4{T}(w::KnetArray{T},x::KnetArray{T};
                  handle=cudnnhandle(), algo=0, workSpace=C_NULL, workSpaceSizeInBytes=0, alpha=1,
                  o...) # padding=0, stride=1, upscale=1, mode=0
    y = similar(x, cdims(w,x;o...))
    beta=0 # nonzero beta does not make sense when we create y
    @cuda(cudnn, cudnnConvolutionForward,
          (Cptr,Ptr{T},Cptr,Ptr{T},Cptr,Ptr{T},Cptr,   UInt32,Cptr,     Csize_t,             Ptr{T},Cptr,Ptr{T}),
          handle,Ref(T(alpha)),TD(x),x,FD(w),w,CD(w,x;o...),algo,workSpace,workSpaceSizeInBytes,Ref(T(beta)),TD(y),y)
    return y
end

function conv4x{T}(w::KnetArray{T},x::KnetArray{T},dy::KnetArray{T};
                   handle=cudnnhandle(), alpha=1, algo=0, workSpace=C_NULL, workSpaceSizeInBytes=0,
                   o...) # padding=0, stride=1, upscale=1, mode=0
    beta = 0
    dx = similar(x)
    if cudnnVersion >= 4000
        @cuda(cudnn,cudnnConvolutionBackwardData,
              (Cptr,Ptr{T},Cptr,Ptr{T},Cptr,Ptr{T},Cptr,     UInt32,Cptr,     Csize_t,             Ptr{T},Cptr,Ptr{T}),
              handle,Ref(T(alpha)),FD(w),w,TD(dy),dy,CD(w,x;o...),algo,workSpace,workSpaceSizeInBytes,Ref(T(beta)),TD(dx),dx)
    elseif cudnnVersion >= 3000
        @cuda(cudnn,cudnnConvolutionBackwardData_v3,
              (Cptr,Ptr{T},Cptr,Ptr{T},Cptr,Ptr{T},Cptr,     UInt32,Cptr,     Csize_t,             Ptr{T},Cptr,Ptr{T}),
              handle,Ref(T(alpha)),FD(w),w,TD(dy),dy,CD(w,x;o...),algo,workSpace,workSpaceSizeInBytes,Ref(T(beta)),TD(dx),dx)
    else
        @cuda(cudnn,cudnnConvolutionBackwardData,
              (Cptr,Ptr{T},Cptr,Ptr{T},Cptr,Ptr{T},Cptr,       Ptr{T},Cptr,Ptr{T}),
              handle,Ref(T(alpha)),FD(w),w,TD(dy),dy,CD(w,x;o...),Ref(T(beta)),TD(dx),dx)
    end
    return dx
end

function conv4w{T}(w::KnetArray{T},x::KnetArray{T},dy::KnetArray{T};
                   handle=cudnnhandle(), alpha=1, algo=0, workSpace=C_NULL, workSpaceSizeInBytes=0,
                   o...) # padding=0, stride=1, upscale=1, mode=0
    beta = 0
    dw = similar(w)
    if cudnnVersion >= 4000
        @cuda(cudnn,cudnnConvolutionBackwardFilter,
              (Cptr,Ptr{T},Cptr,Ptr{T},Cptr,Ptr{T},Cptr,     UInt32,Cptr,     Csize_t,             Ptr{T},Cptr,Ptr{T}),
              handle,Ref(T(alpha)),TD(x),x,TD(dy),dy,CD(w,x;o...),algo,workSpace,workSpaceSizeInBytes,Ref(T(beta)),FD(dw),dw)
    elseif cudnnVersion >= 3000
        @cuda(cudnn,cudnnConvolutionBackwardFilter_v3,
              (Cptr,Ptr{T},Cptr,Ptr{T},Cptr,Ptr{T},Cptr,     UInt32,Cptr,     Csize_t,             Ptr{T},Cptr,Ptr{T}),
              handle,Ref(T(alpha)),TD(x),x,TD(dy),dy,CD(w,x;o...),algo,workSpace,workSpaceSizeInBytes,Ref(T(beta)),FD(dw),dw)
    else
        @cuda(cudnn,cudnnConvolutionBackwardFilter,
              (Cptr,Ptr{T},Cptr,Ptr{T},Cptr,Ptr{T},Cptr,       Ptr{T},Cptr,Ptr{T}),
              handle,Ref(T(alpha)),TD(x),x,TD(dy),dy,CD(w,x;o...),Ref(T(beta)),FD(dw),dw)
    end
    return dw
end

@primitive conv4(w,x; o...),dy  conv4w(w,x,dy;o...)  conv4x(w,x,dy;o...)
@zerograd  conv4x(w,x,dy;o...)
@zerograd  conv4w(w,x,dy;o...)

"""

    pool(x; kwargs...)

Compute pooling of input values (i.e., the maximum or average of
several adjacent values) to produce an output with smaller height
and/or width.

Currently 4 or 5 dimensional KnetArrays with `Float32` or `Float64`
entries are supported.  If `x` has dimensions `(X1,X2,...,I,N)`, the
result `y` will have dimensions `(Y1,Y2,...,I,N)` where

    Yi=1+floor((Xi+2*padding[i]-window[i])/stride[i])

Here `I` is the number of input channels, `N` is the number of
instances, and `Xi,Yi` are spatial dimensions.  `window`, `padding`
and `stride` are keyword arguments that can be specified as a single
number (in which case they apply to all dimensions), or an array/tuple
with entries for each spatial dimension.

# Keywords:

* `window=2`: the pooling window size for each dimension.
* `padding=0`: the number of extra zeros implicitly concatenated at the start and at the end of each dimension.
* `stride=window`: the number of elements to slide to reach the next pooling window.
* `mode=0`: 0 for max, 1 for average including padded values, 2 for average excluding padded values.
* `maxpoolingNanOpt=0`: Nan numbers are not propagated if 0, they are propagated if 1.
* `alpha=1`: can be used to scale the result.
* `handle`: Handle to a previously created cuDNN context. Defaults to a Knet allocated handle.

"""
function pool{T}(x::KnetArray{T}; handle=cudnnhandle(), alpha=1,
                 o...) # window=2, padding=0, stride=window, mode=0, maxpoolingNanOpt=0
    y = similar(x, pdims(x; o...))
    beta = 0
    @cuda(cudnn, cudnnPoolingForward,
          (Cptr, Cptr,      Ptr{T},    Cptr,Ptr{T},Ptr{T},   Cptr,Ptr{T}),
          handle,PD(x;o...),Ref(T(alpha)),TD(x),x,    Ref(T(beta)),TD(y),y)
    return y
end

function poolx{T}(x::KnetArray{T},y::KnetArray{T},dy::KnetArray{T}; handle=cudnnhandle(), alpha=1, mode=0,
                  o...) # window=2, padding=0, stride=window, maxpoolingNanOpt=0
    if alpha!=1 && mode==0; error("Gradient of pool(alpha!=1,mode=0) broken in CUDNN"); end
    dx = similar(x)
    beta = 0
    @cuda(cudnn,cudnnPoolingBackward,
          (Cptr,Cptr,Ptr{T},Cptr,Ptr{T},Cptr,Ptr{T},Cptr,Ptr{T},Ptr{T},Cptr,Ptr{T}),
          handle,PD(x;mode=mode,o...),Ref(T(alpha)),TD(y),y,TD(dy),dy,TD(x),x,Ref(T(beta)),TD(dx),dx)
    return dx
end

@primitive pool(x;o...),dy,y  poolx(x,y,dy;o...)
@zerograd  poolx(x,y,dy;o...)

"""

Unpooling; `reverse` of pooling.

    x == pool(unpool(x;o...); o...)

"""
function unpool(x; window=2, alpha=1, o...) # padding=0, stride=window, mode=0, maxpoolingNanOpt=0
    w = prod(psize(window,x))
    y = similar(x,updims(x; window=window, o...))
    poolx(y,x,x.*w; o..., window=window, mode=1, alpha=1/alpha)
end

function unpoolx(dy; window=2, alpha=1, o...) # padding=0, stride=window, mode=0, maxpoolingNanOpt=0
    w = prod(psize(window,dy))
    pool(dy; o..., window=window, mode=1, alpha=1/alpha) * w
end

# @primitive unpool(x;o...),dy,y -pool(-dy;o...)
@primitive  unpool(x;o...),dy,y  unpoolx(dy;o...)

"""

Deconvolution; `reverse` of convolution.

"""
function deconv4(w,x; o...)
    y = similar(x,dcdims(w,x;o...))
    return conv4x(w,y,x;o...)
end

function deconv4w(w,x,dy; o...)
    return conv4w(w,dy,x;o...)
end

function deconv4x(w,x,dy; o...)
    return conv4(w,dy;o...)
end

@primitive deconv4(w,x; o...),dy,y  deconv4w(w,x,dy; o...)  deconv4x(w,x,dy; o...)
@zerograd deconv4w(w,x,dy; o...)
@zerograd deconv4x(w,x,dy; o...)

# cudnn descriptors

type TD; ptr
    function TD(a::KnetArray)
        d = Cptr[0]
        @cuda(cudnn,cudnnCreateTensorDescriptor,(Ptr{Cptr},),d)
        n = ndims(a)
        sz = [Cint(size(a,n-i+1)) for i=1:n]
        st = [Cint(stride(a,n-i+1)) for i=1:n]
        @cuda(cudnn,cudnnSetTensorNdDescriptor,
              (Cptr,UInt32,Cint,Ptr{Cint},Ptr{Cint}),
              d[1], DT(a), n, sz, st)
        td = new(d[1])
        finalizer(td, x->@cuda(cudnn,cudnnDestroyTensorDescriptor,(Cptr,),x.ptr))
        return td
    end
end

type FD; ptr
    function FD(a::KnetArray)
        d = Cptr[0]
        @cuda(cudnn,cudnnCreateFilterDescriptor,(Ptr{Cptr},),d)
        n = ndims(a)
        sz = [Cint(size(a,n-i+1)) for i=1:n]
        if cudnnVersion >= 5000
            @cuda(cudnn,cudnnSetFilterNdDescriptor,
                  (Cptr,UInt32,UInt32,Cint,Ptr{Cint}),
                  d[1], DT(a), 0,     n,   sz)
        elseif cudnnVersion >= 4000
            @cuda(cudnn,cudnnSetFilterNdDescriptor_v4,
                  (Cptr,UInt32,UInt32,Cint,Ptr{Cint}),
                  d[1], DT(a), 0,     n,   sz)
        else
            @cuda(cudnn,cudnnSetFilterNdDescriptor,
                  (Cptr,UInt32,Cint,Ptr{Cint}),
                  d[1], DT(a),    n,   sz)
        end
        fd = new(d[1])
        finalizer(fd, x->@cuda(cudnn,cudnnDestroyFilterDescriptor,(Cptr,),x.ptr))
        return fd
    end
end

type CD; ptr
    function CD(w::KnetArray,x::KnetArray; padding=0, stride=1, upscale=1, mode=0)
        d = Cptr[0]
        @cuda(cudnn,cudnnCreateConvolutionDescriptor,(Ptr{Cptr},),d)
        nd = ndims(x)-2
        if cudnnVersion >= 4000
            @cuda(cudnn,cudnnSetConvolutionNdDescriptor,
                  (Cptr,Cint,Ptr{Cint},Ptr{Cint},Ptr{Cint},UInt32,UInt32),
                  d[1],nd,cdsize(padding,nd),cdsize(stride,nd),cdsize(upscale,nd),mode,DT(x))
        elseif cudnnVersion > 3000 # does not work when cudnnVersion==3000
            @cuda(cudnn,cudnnSetConvolutionNdDescriptor_v3,
                  (Cptr,Cint,Ptr{Cint},Ptr{Cint},Ptr{Cint},UInt32,UInt32),
                  d[1],nd,cdsize(padding,nd),cdsize(stride,nd),cdsize(upscale,nd),mode,DT(x))
        else
            @cuda(cudnn,cudnnSetConvolutionNdDescriptor,
                  (Cptr,Cint,Ptr{Cint},Ptr{Cint},Ptr{Cint},UInt32),
                  d[1],nd,cdsize(padding,nd),cdsize(stride,nd),cdsize(upscale,nd),mode)
        end
        cd = new(d[1])
        finalizer(cd, x->@cuda(cudnn,cudnnDestroyConvolutionDescriptor,(Cptr,),x.ptr))
        return cd
    end
end

type PD; ptr
    function PD(x::KnetArray; window=2, padding=0, stride=window, mode=0, maxpoolingNanOpt=0)
        d = Cptr[0]
        @cuda(cudnn,cudnnCreatePoolingDescriptor,(Ptr{Cptr},),d)
        nd = ndims(x)-2
        if cudnnVersion >= 5000
            @cuda(cudnn,cudnnSetPoolingNdDescriptor,
                  (Cptr,UInt32,UInt32,Cint,Ptr{Cint},Ptr{Cint},Ptr{Cint}),
                  d[1],mode,maxpoolingNanOpt,nd,cdsize(window,nd),cdsize(padding,nd),cdsize(stride,nd))
        elseif cudnnVersion >= 4000
            @cuda(cudnn,cudnnSetPoolingNdDescriptor_v4,
                  (Cptr,UInt32,UInt32,Cint,Ptr{Cint},Ptr{Cint},Ptr{Cint}),
                  d[1],mode,maxpoolingNanOpt,nd,cdsize(window,nd),cdsize(padding,nd),cdsize(stride,nd))
        else
            @cuda(cudnn,cudnnSetPoolingNdDescriptor,
                  (Cptr,UInt32,Cint,Ptr{Cint},Ptr{Cint},Ptr{Cint}),
                  d[1],mode,nd,cdsize(window,nd),cdsize(padding,nd),cdsize(stride,nd))
        end
        pd = new(d[1])
        finalizer(pd, x->@cuda(cudnn,cudnnDestroyPoolingDescriptor,(Cptr,),x.ptr))
        return pd
    end
end

import Base: unsafe_convert
unsafe_convert(::Type{Cptr}, td::TD)=td.ptr
unsafe_convert(::Type{Cptr}, fd::FD)=fd.ptr
unsafe_convert(::Type{Cptr}, cd::CD)=cd.ptr
unsafe_convert(::Type{Cptr}, pd::PD)=pd.ptr

# fill and reverse Cint array with padding etc. for cudnn calls
function cdsize(w, nd)
    if isa(w,Number)
        fill(Cint(w),nd)
    elseif length(w)==nd
        [ Cint(w[nd-i+1]) for i=1:nd ]
    else
        throw(DimensionMismatch("$w $nd"))
    end
end

# convert padding etc. size to an Int array of the right dimension
function psize(p, x)
    nd = ndims(x)-2
    if isa(p,Number)
        fill(Int(p),nd)
    elseif length(p)==nd
        collect(Int,p)
    else
        throw(DimensionMismatch("psize: $p $nd"))
    end
end

DT(::KnetArray{Float32})=UInt32(0)
DT(::KnetArray{Float64})=UInt32(1)
DT(::KnetArray{Float16})=UInt32(2)

function cdims(w,x; padding=0, stride=1, o...)
    N = ndims(x)
    ntuple(N) do i
        if i < N-1
            pi = (if isa(padding,Number); padding; else padding[i]; end)
            si = (if isa(stride,Number); stride; else stride[i]; end)
            1 + div(size(x,i) - size(w,i) + 2*pi, si)
        elseif i == N-1
            size(w,N)
        else # i == N
            size(x,N)
        end
    end
end

function pdims(x; window=2, padding=0, stride=window, o...)
    N = ndims(x)
    ntuple(N) do i
        if i < N-1
            wi = (if isa(window,Number); window; else window[i]; end)
            pi = (if isa(padding,Number); padding; else padding[i]; end)
            si = (if isa(stride,Number); stride; else stride[i]; end)
            1 + div(size(x,i) + 2*pi - wi, si)
        else
            size(x,i)
        end
    end
end

function dcdims(w,x; padding=0, stride=1, o...)
    N = ndims(x)
    ntuple(N) do i
        if i < N-1
            pi = (if isa(padding,Number); padding; else padding[i]; end)
            si = (if isa(stride,Number); stride; else stride[i]; end)
            si*(size(x,i)-1) + size(w,i) - 2*pi
        elseif i == N-1
            size(w,N)
        else
            size(x,N)
        end
    end
end

function updims(x; window=2, padding=0, stride=window, o...)
    window = psize(window,x)
    stride = psize(stride,x)
    padding = psize(padding,x)
    N = ndims(x)
    ntuple(N) do i
        if i < N-1
            (size(x,i)-1)*stride[i]+window[i]-2*padding[i]
        else
            size(x,i)
        end
    end
end

# convolution padding size that preserves the input size when filter size is odd and stride=1
padsize(w)=ntuple(i->div(size(w,i)-1,2), ndims(w)-2)

### CPU convolution using im2col from Mocha.jl

# w=Ww,Hw,Cx,Cy
# x=Wx,Hx,Cx,Nx
# y=Wy,Hy,Cy,Nx
# if we apply im2col to a single image:
# w2=(Ww*Hw*Cx),Cy  ;; simple reshape
# x2=(Wy*Hy),(Ww*Hw*Cx)
# y2=(Wy*Hy),Cy     ;; simple reshape after y2=x2*w2

function conv4{T}(w::Array{T,4}, x::Array{T,4};
                  padding=0, stride=1, upscale=1, mode=0, alpha=1,
                  o...) # Ignoring handle, algo, workSpace, workSpaceSizeInBytes
    if upscale != 1; throw(ArgumentError("CPU conv4 only supports upscale=1.")); end
    if mode != 0 && mode != 1; throw(ArgumentError("conv4 only supports mode=0 or 1.")); end
    Wx,Hx,Cx,Nx = size(x)
    Ww,Hw,C1,C2 = size(w)
    if Cx!=C1; throw(DimensionMismatch()); end
    Wy,Hy,Cy,Ny = cdims(w,x;padding=padding,stride=stride)
    # @assert Cy==C2 && Ny==Nx
    y = similar(x, (Wy,Hy,Cy,Ny))
    x2dims = im2col_dims(w,x,y)
    x2 = similar(x, x2dims)
    (p1,p2) = psize(padding,x)
    (s1,s2) = psize(stride,x)
    M,N,K,Y = Wy*Hy,Cy,Ww*Hw*Cx,Wy*Hy*Cy
    alpha,beta,yidx = T(alpha),T(0),1
    @inbounds for n in 1:Nx
        im2col!(w, x, x2, n, p1, p2, s1, s2, mode)
        gemm!('N','N',M,N,K,alpha,pointer(x2),pointer(w),beta,pointer(y,yidx))
        yidx += Y
    end
    return y
end

function conv4w{T}(w::Array{T,4},x::Array{T,4},dy::Array{T,4};
                   padding=0, stride=1, upscale=1, mode=0, alpha=1,
                   o...) # Ignoring handle, algo, workSpace, workSpaceSizeInBytes
    # dw = x'*dy
    Wx,Hx,Cx,Nx = size(x)
    Ww,Hw,C1,C2 = size(w)
    Wy,Hy,Cy,Ny = size(dy)
    # if upscale != 1; throw(ArgumentError("CPU conv4 only supports upscale=1.")); end
    # if mode != 0 && mode != 1; throw(ArgumentError("conv4 only supports mode=0 or 1.")); end
    # @assert Cx==C1 && Cy==C2 && Ny==Nx
    dw = zeros(w)
    x2dims = im2col_dims(w,x,dy)
    x2 = similar(x, x2dims)
    # op(A) is an m-by-k matrix, op(B) is a k-by-n matrix, C is an m-by-n matrix.
    Y,M,N,K = Wy*Hy*Cy,Ww*Hw*Cx,Cy,Wy*Hy
    alpha,beta = T(alpha),T(1)
    (p1,p2) = psize(padding,x)
    (s1,s2) = psize(stride,x)
    dyi = 1
    @inbounds for n in 1:Nx
        im2col!(w, x, x2, n, p1, p2, s1, s2, mode)
        gemm!('T','N',M,N,K,alpha,pointer(x2),pointer(dy,dyi),beta,pointer(dw))
        dyi += Y
    end
    return dw
end

function conv4x{T}(w::Array{T,4},x::Array{T,4},dy::Array{T,4};
                   padding=0, stride=1, upscale=1, mode=0, alpha=1,
                   o...) # Ignoring handle, algo, workSpace, workSpaceSizeInBytes
    # dx = dy*w'
    Wx,Hx,Cx,Nx = size(x)
    Ww,Hw,C1,C2 = size(w)
    Wy,Hy,Cy,Ny = size(dy)
    # if upscale != 1; throw(ArgumentError("CPU conv4 only supports upscale=1.")); end
    # if mode != 0 && mode != 1; throw(ArgumentError("conv4 only supports mode=0 or 1.")); end
    @assert Cx==C1 && Cy==C2 && Ny==Nx
    dx = similar(x)
    x2dims = im2col_dims(w,x,dy)
    x2 = similar(x, x2dims)
    # op(A) is an m-by-k matrix, op(B) is a k-by-n matrix, C is an m-by-n matrix.
    Y,M,N,K = Wy*Hy*Cy,Wy*Hy,Ww*Hw*Cx,Cy
    alpha,beta = T(alpha),T(0)
    (p1,p2) = psize(padding,x)
    (s1,s2) = psize(stride,x)
    dyi = 1
    @inbounds for n in 1:Nx
        gemm!('N','T',M,N,K,alpha,pointer(dy,dyi),pointer(w),beta,pointer(x2))
        col2im!(w,dx,x2,n,p1,p2,s1,s2,mode)
        dyi += Y
    end
    return dx
end

im2col_dims(w,x,y)=(size(y,1)*size(y,2), size(w,1)*size(w,2)*size(w,3))

# Functions from conv.cpp:

for (T,S) in ((Float32,32), (Float64,64)); @eval begin

    function im2col!(w::Array{$T,4}, x::Array{$T,4}, x2::Array{$T,2},
                     n::Int, p1::Int, p2::Int, s1::Int, s2::Int, mode::Int)
        Wx,Hx,Cx,Nx = size(x)
        Ww,Hw,C1,C2 = size(w)
        xn = pointer(x, Wx*Hx*Cx*(n-1)+1)
        ccall(($("im2col$S"),libknet8),Void,
              (Ptr{$T},Ptr{$T},Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint),
              xn,x2,Wx,Hx,Cx,Ww,Hw,p1,p2,s1,s2,mode)
        return x2
    end

    function col2im!(w::Array{$T,4}, x::Array{$T,4}, x2::Array{$T,2},
                     n::Int, p1::Int, p2::Int, s1::Int, s2::Int, mode::Int)
        Wx,Hx,Cx,Nx = size(x)
        Ww,Hw,C1,C2 = size(w)
        xn = pointer(x, Wx*Hx*Cx*(n-1)+1)
        ccall(($("col2im$S"),libknet8),Void,
              (Ptr{$T},Ptr{$T},Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint),
              x2,xn,Wx,Hx,Cx,Ww,Hw,p1,p2,s1,s2,mode)
        return x
    end

    ### CPU pooling from Mocha.jl

    function pool(x::Array{$T,4}; window=2, padding=0, stride=window, mode=0, maxpoolingNanOpt=0, alpha=1, handle=nothing)
        if maxpoolingNanOpt!=0; throw(ArgumentError("CPU pool only supports maxpoolingNanOpt=0")); end
        Wx,Hx,Cx,Nx = size(x);
        Wy,Hy,Cy,Ny = pdims(x;window=window,padding=padding,stride=stride)
        y = similar(x, (Wy,Hy,Cy,Ny))
        (w1,w2) = psize(window, x)
        (p1,p2) = psize(padding, x)
        (s1,s2) = psize(stride, x)
        if mode == 0
            ccall(($("max_pooling_fwd$S"),libknet8),Void,
                  (Ptr{$T},Ptr{$T},Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint),
                  x,y,Wx,Hx,Cx,Nx,Wy,Hy,w1,w2,p1,p2,s1,s2)
        elseif mode == 1 || (mode == 2 && p1==p2==0)
            ccall(($("mean_pooling_fwd$S"),libknet8),Void,
                  (Ptr{$T},Ptr{$T},Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint),
                  x,y,Wx,Hx,Cx,Nx,Wy,Hy,w1,w2,p1,p2,s1,s2)
        else
            throw(ArgumentError("mode $mode not supported by cpu pool"))
        end
        if alpha != 1; scale!(alpha,y); end
        return y
    end

    function poolx(x::Array{$T,4}, y::Array{$T,4}, dy::Array{$T,4};
                   window=2, padding=0, stride=window, mode=0, maxpoolingNanOpt=0, alpha=1, handle=nothing)
        if maxpoolingNanOpt!=0; throw(ArgumentError("CPU pool only supports maxpoolingNanOpt=0")); end
        Wx,Hx,Cx,Nx = size(x);
        Wy,Hy,Cy,Ny = size(y);
        dx = similar(x)
        (w1,w2) = psize(window, x)
        (p1,p2) = psize(padding, x)
        (s1,s2) = psize(stride, x)
        if mode == 0
            if alpha != 1; y = y ./ alpha; end
            ccall(($("max_pooling_bwd$S"),libknet8),Void,
                  (Ptr{$T},Ptr{$T},Ptr{$T},Ptr{$T},Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint),
                  x,y,dy,dx,Wx,Hx,Cx,Nx,Wy,Hy,w1,w2,p1,p2,s1,s2)
        elseif mode == 1 || (mode == 2 && p1==p2==0)
            ccall(($("mean_pooling_bwd$S"),libknet8),Void,
                  (Ptr{$T},Ptr{$T},Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint),
                  dx,dy,Wx,Hx,Cx,Nx,Wy,Hy,w1,w2,p1,p2,s1,s2)
        else
            throw(ArgumentError("mode $mode not supported by cpu pool"))
        end
        if alpha != 1; scale!(alpha,dx); end
        return dx
    end
end;end


"""
Calculate value, gradient, and hessian of the variational ELBO using a PSF
image rather than a mixture of Gaussians.
"""

module DeterministicVIImagePSF

using StaticArrays, DiffBase

import ..DeterministicVI:
    ElboArgs, ElboIntermediateVariables, maximize_f,
    StarPosParams, GalaxyPosParams, CanonicalParams, VariationalParams,
    SourceBrightness, GalaxyComponent, SkyPatch,
    load_source_brightnesses, add_elbo_log_term!,
    accumulate_source_pixel_brightness!, subtract_kl_all_sources!, KL_HELPER_POOL

import ..Model:
    populate_gal_fsm!, getids, ParamSet, linear_world_to_pix, lidx,
    BvnComponent, GalaxyCacheComponent, GalaxySigmaDerivs,
    get_bvn_cov, galaxy_prototypes, linear_world_to_pix,
    Image, eval_psf, CatalogEntry, init_source

import ..SensitiveFloats:
    SensitiveFloat, zero_sensitive_float_array,
    multiply_sfs!, add_scaled_sfs!, clear!

import ..Infer: load_active_pixels!, get_sky_patches

import ..PSF: get_psf_at_point, trim_psf

import WCS

include("deterministic_vi_image_psf/sensitive_float_fft.jl")
include("deterministic_vi_image_psf/kernels.jl")
include("deterministic_vi_image_psf/fsm_matrices.jl")
include("deterministic_vi_image_psf/elbo_image_psf.jl")

export elbo_likelihood_with_fft!, FSMSensitiveFloatMatrices,
       initialize_fsm_sf_matrices!, initialize_fft_elbo_parameters

function infer_source_fft(images::Vector{Image},
                          neighbors::Vector{CatalogEntry},
                          entry::CatalogEntry)
   cat_local = vcat([entry], neighbors)
   vp = Vector{Float64}[init_source(ce) for ce in cat_local]
   patches = get_sky_patches(images, cat_local)
   load_active_pixels!(images, patches)

   ea_fft, fsm_mat = initialize_fft_elbo_parameters(
       images, vp, patches, [1], use_raw_psf=false)
   elbo_fft_opt = get_fft_elbo_function(ea_fft, fsm_mat)
   maximize_f(elbo_fft_opt, ea_fft)

   vp[1]
end

end


using CUDA.CUDNN:
    cudnnActivationForward,
    cudnnActivationForward!,
    cudnnActivationBackward,
    cudnnActivationDescriptor,
        cudnnActivationDescriptor_t,
        cudnnCreateActivationDescriptor,
        cudnnSetActivationDescriptor,
        cudnnGetActivationDescriptor,
        cudnnDestroyActivationDescriptor,
    cudnnActivationMode_t,
        CUDNN_ACTIVATION_SIGMOID,      # 0
        CUDNN_ACTIVATION_RELU,         # 1
        CUDNN_ACTIVATION_TANH,         # 2
        CUDNN_ACTIVATION_CLIPPED_RELU, # 3
        CUDNN_ACTIVATION_ELU,          # 4
        CUDNN_ACTIVATION_IDENTITY,     # 5
    cudnnNanPropagation_t,
        CUDNN_NOT_PROPAGATE_NAN, # 0
        CUDNN_PROPAGATE_NAN      # 1

@testset "cudnn/activation" begin
    @test cudnnActivationDescriptor(C_NULL) isa cudnnActivationDescriptor
    @test Base.unsafe_convert(Ptr, cudnnActivationDescriptor(C_NULL)) isa Ptr
    @test cudnnActivationDescriptor(CUDNN_ACTIVATION_RELU,CUDNN_NOT_PROPAGATE_NAN,0) isa cudnnActivationDescriptor

    (ax,ay) = randn.((10,10))
    (cx,cy) = CuArray.((ax,ay))

    function activationtest(;
        mode=CUDNN_ACTIVATION_SIGMOID,
        nanOpt=CUDNN_NOT_PROPAGATE_NAN,
        coef=1,
        alpha=1,
        beta=0,
    )
        fx = (mode === CUDNN_ACTIVATION_SIGMOID ? 1 ./ (1 .+ exp.(-ax)) :
              mode === CUDNN_ACTIVATION_RELU ? max.(0,ax) :
              mode === CUDNN_ACTIVATION_TANH ? tanh.(ax) :
              mode === CUDNN_ACTIVATION_CLIPPED_RELU ? clamp.(ax,0,coef) :
              mode === CUDNN_ACTIVATION_ELU ? (x->(x >= 0 ? x : coef*(exp(x)-1))).(ax) :
              error("Unknown activation"))
        d = cudnnActivationDescriptor(mode,nanOpt,Cfloat(coef))
        y0 = alpha * fx
        y1 = y0 .+ beta * ay
        @test y0 ≈ cudnnActivationForward(cx; mode, nanOpt, coef, alpha) |> Array
        @test y0 ≈ cudnnActivationForward(cx, d; alpha) |> Array
        @test y1 ≈ cudnnActivationForward!(copy(cy), cx; mode, nanOpt, coef, alpha, beta) |> Array
        @test y1 ≈ cudnnActivationForward!(copy(cy), cx, d; alpha, beta) |> Array
    end

    activationtest(mode=CUDNN_ACTIVATION_SIGMOID)
    activationtest(mode=CUDNN_ACTIVATION_RELU)
    activationtest(mode=CUDNN_ACTIVATION_TANH)
    activationtest(mode=CUDNN_ACTIVATION_CLIPPED_RELU)
    activationtest(mode=CUDNN_ACTIVATION_ELU)
    activationtest(nanOpt=CUDNN_PROPAGATE_NAN)
    activationtest(coef=2,mode=CUDNN_ACTIVATION_CLIPPED_RELU)
    activationtest(coef=2,mode=CUDNN_ACTIVATION_ELU)
    activationtest(alpha=2)
    activationtest(beta=2)
end


export RandomVariable, randomvar

import Base: show

## Random variable implementation

mutable struct RandomVariable <: AbstractVariable
    name                :: Symbol
    collection_type     :: AbstractVariableCollectionType
    input_messages      :: Vector{MessageObservable{AbstractMessage}}
    output_messages     :: Vector{MessageObservable{Message}}
    output_initialised  :: Bool
    output_cache        :: Union{Nothing, EqualityChain}
    marginal            :: MarginalObservable
    pipeline            :: AbstractPipelineStage
    prod_constraint
    prod_strategy
    form_constraint
    form_check_strategy
end

Base.show(io::IO, randomvar::RandomVariable) = print(io, "RandomVariable(", indexed_name(randomvar), ")")

function randomvar(name::Symbol, collection_type::AbstractVariableCollectionType = VariableIndividual(); pipeline = EmptyPipelineStage(), prod_constraint = ProdAnalytical(), prod_strategy = FoldLeftProdStrategy(), form_constraint = UnspecifiedFormConstraint(), form_check_strategy = FormConstraintCheckPickDefault())
    return RandomVariable(name, collection_type, Vector{MessageObservable{AbstractMessage}}(), Vector{MessageObservable{Message}}(), false, nothing, MarginalObservable(), pipeline, prod_constraint, prod_strategy, form_constraint, form_check_strategy)
end

function randomvar(name::Symbol, dims::Tuple; pipeline = EmptyPipelineStage(), prod_constraint = ProdAnalytical(), prod_strategy = FoldLeftProdStrategy(), form_constraint = UnspecifiedFormConstraint(), form_check_strategy = FormConstraintCheckPickDefault())
    return randomvar(name, dims...; pipeline = pipeline, prod_constraint = prod_constraint, prod_strategy = prod_strategy, form_constraint = form_constraint, form_check_strategy = form_check_strategy)
end

function randomvar(name::Symbol, length::Int; pipeline = EmptyPipelineStage(), prod_constraint = ProdAnalytical(), prod_strategy = FoldLeftProdStrategy(), form_constraint = UnspecifiedFormConstraint(), form_check_strategy = FormConstraintCheckPickDefault())
    vars = Vector{RandomVariable}(undef, length)
    @inbounds for i in 1:length
        vars[i] = randomvar(name, VariableVector(i); pipeline = pipeline, prod_constraint = prod_constraint, prod_strategy = prod_strategy, form_constraint = form_constraint, form_check_strategy = form_check_strategy)
    end
    return vars
end

function randomvar(name::Symbol, dims::Vararg{Int}; pipeline = EmptyPipelineStage(), prod_constraint = ProdAnalytical(), prod_strategy = FoldLeftProdStrategy(), form_constraint = UnspecifiedFormConstraint(), form_check_strategy = FormConstraintCheckPickDefault())
    vars = Array{RandomVariable}(undef, dims)
    @inbounds for i in CartesianIndices(axes(vars))
        vars[i] = randomvar(name, VariableArray(i); pipeline = pipeline, prod_constraint = prod_constraint, prod_strategy = prod_strategy, form_constraint = form_constraint, form_check_strategy = form_check_strategy)
    end
    return vars
end

degree(randomvar::RandomVariable)              = length(randomvar.input_messages)
name(randomvar::RandomVariable)                = randomvar.name
collection_type(randomvar::RandomVariable)     = randomvar.collection_type
equality_chain(randomvar::RandomVariable)      = randomvar.equality_chain
prod_constraint(randomvar::RandomVariable)     = randomvar.prod_constraint
prod_strategy(randomvar::RandomVariable)       = randomvar.prod_strategy
form_constraint(randomvar::RandomVariable)     = randomvar.form_constraint
form_check_strategy(randomvar::RandomVariable) = _form_check_strategy(randomvar.form_check_strategy, randomvar)

_form_check_strategy(::FormConstraintCheckPickDefault, randomvar::RandomVariable) = default_form_check_strategy(form_constraint(randomvar))
_form_check_strategy(form_check_strategy, randomvar::RandomVariable)              = form_check_strategy

messages_prod_fn(randomvar::RandomVariable) = messages_prod_fn(prod_strategy(randomvar), prod_constraint(randomvar), UnspecifiedFormConstraint(), default_form_check_strategy(UnspecifiedFormConstraint()))
marginal_prod_fn(randomvar::RandomVariable) = marginal_prod_fn(prod_strategy(randomvar), prod_constraint(randomvar), form_constraint(randomvar), form_check_strategy(randomvar))

getlastindex(randomvar::RandomVariable) = degree(randomvar) + 1

messagein(randomvar::RandomVariable, index::Int)  = @inbounds randomvar.input_messages[index]

function messageout(randomvar::RandomVariable, index::Int)
    if randomvar.output_initialised === false
        initialize_output_messages!(randomvar)
    end
    return @inbounds randomvar.output_messages[index]
end

get_pipeline_stages(randomvar::RandomVariable)        = randomvar.pipeline
add_pipeline_stage!(randomvar::RandomVariable, stage) = randomvar.pipeline = (randomvar.pipeline + stage)

_getmarginal(randomvar::RandomVariable)              = randomvar.marginal
_setmarginal!(randomvar::RandomVariable, observable) = connect!(_getmarginal(randomvar), observable)
_makemarginal(randomvar::RandomVariable)             = collectLatest(AbstractMessage, Marginal, randomvar.input_messages, marginal_prod_fn(randomvar))

function setmessagein!(randomvar::RandomVariable, index::Int, messagein)
    if index === degree(randomvar) + 1
        push!(randomvar.input_messages, messagein)
    else
        error("Inconsistent state in setmessagein! function for random variable $(randomvar). `index` should be equal to `degree(randomvar) + 1 = $(degree(randomvar) + 1)`, $(index) is given instead")
    end
end

function activate!(model, randomvar::RandomVariable)
    # `5` here is empirical observation, maybe we can come up with better heuristic?
    # in case if number of connections is large we use cache equality nodes chain structure
    if degree(randomvar) > 5
        chain_pipeline  = schedule_on(global_reactive_scheduler(getoptions(model)))
        chain_prod_fn   = messages_prod_fn(randomvar)
        randomvar.output_cache = EqualityChain(randomvar.input_messages, chain_pipeline, chain_prod_fn)
    end

    _setmarginal!(randomvar, _makemarginal(randomvar))

    return nothing
end

initialize_output_messages!(randomvar::RandomVariable) = initialize_output_messages!(randomvar.output_cache, randomvar)

# Generic fallback for variables with small number of connected nodes, somewhere <= 5
# We do not create equality chain in this cases, but simply do eager product
function initialize_output_messages!(::Nothing, randomvar::RandomVariable)
    d          = degree(randomvar)
    outputmsgs = randomvar.output_messages
    inputmsgs  = randomvar.input_messages
    prod_fn    = messages_prod_fn(randomvar)

    resize!(outputmsgs, d)

    @inbounds for i in 1:d
        outputmsgs[i] = MessageObservable(Message)
        outputmsg     = collectLatest(AbstractMessage, Message, skipindex(inputmsgs, i), prod_fn)
        connect!(outputmsgs[i], outputmsg)
    end

    randomvar.output_initialised = true

    return nothing
end

# Equality chain initialisation for variables with large number of connected nodes, somewhere > 5
# In this cases it is more efficient to create an equality chain structure, but it does allocate way more memory
function initialize_output_messages!(chain::EqualityChain, randomvar::RandomVariable)
    d          = degree(randomvar)
    outputmsgs = randomvar.output_messages

    resize!(outputmsgs, d)

    @inbounds for i in 1:d
        outputmsgs[i] = MessageObservable(Message)
    end

    initialize!(chain, outputmsgs)
    randomvar.output_initialised = true

    return nothing
end


# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/Materials.jl/blob/master/LICENSE
using Test, Tensors

let parameters = MemoryParameterState(E = 200.0e3,
                                      nu = 0.3,
                                      R0 = 100.0,
                                      Kn = 20.0,
                                      nn = 3.0,
                                      C1 = 10000.0,
                                      D1 = 100.0,
                                      C2 = 50000.0,
                                      D2 = 1000.0,
                                      Q0 = 100.0,
                                      QM = 500.0,
                                      mu = 100.0,
                                      b = 30.0,
                                      eta = 0.5,
                                      m = 0.5,
                                      pt = 0.0,
                                      xi = 0.3),
    mat = Memory(parameters=parameters),

    tostrain(tens::Symm2) = copy(tovoigt(tens; offdiagscale=2.0)),
    tostress(tens::Symm2) = copy(tovoigt(tens)),

    n_cycles = 30,
    points_per_cycle = 40,
    t = range(0.0; stop=Float64(n_cycles), length=n_cycles * points_per_cycle + 1),
    dtime = t[end] / (length(t) - 1),

    # We initialize these manually to automatically get the correct type.
    times = [copy(mat.drivers.time)],
    stresses = [tostress(mat.variables.stress)],
    strains = [tostrain(mat.drivers.strain)],
    plastic_strains = [tostrain(mat.variables.plastic_strain)],
    cumeqs = [copy(mat.variables.cumeq)],
    qs = [copy(mat.variables.q)],
    Rs = [copy(mat.variables.R)],
    zetas = [tostrain(mat.variables.zeta)]

    function snapshot!()
        push!(times, mat.drivers.time)
        push!(stresses, tostress(mat.variables.stress))
        push!(strains, tostrain(mat.drivers.strain))
        push!(plastic_strains, tostrain(mat.variables.plastic_strain))
        push!(cumeqs, copy(mat.variables.cumeq))
        push!(qs, copy(mat.variables.q))
        push!(Rs, copy(mat.variables.R))
        push!(zetas, tostrain(mat.variables.zeta))
    end

    # Amplitude 1
    ea = 0.003
    strains11 = ea * sin.(2*pi*t)
    for dstrain11 in diff(strains11)
        uniaxial_increment!(mat, dstrain11, dtime)
        update_material!(mat)
        snapshot!()
    end
    R1 = copy(Rs[end])

    # Amplitude 2
    ea = 0.005
    strains11 = ea * sin.(2*pi*t)
    for dstrain11 in diff(strains11)
        uniaxial_increment!(mat, dstrain11, dtime)
        update_material!(mat)
        snapshot!()
    end
    R2 = copy(Rs[end])

    # Amplitude 3
    ea = 0.007
    strains11 = ea * sin.(2*pi*t)
    for dstrain11 in diff(strains11)
        uniaxial_increment!(mat, dstrain11, dtime)
        update_material!(mat)
        snapshot!()
    end
    R3 = copy(Rs[end])

    # Amplitude 4 - evanescence
    ea = 0.003
    strains11 = ea * sin.(2*pi*t)
    for _ in 1:3
        for dstrain11 in diff(strains11)
            uniaxial_increment!(mat, dstrain11, dtime)
            update_material!(mat)
            snapshot!()
        end
    end

    R4 = copy(Rs[end])

    @test R2 > R1
    @test R3 > R2
    @test R4 < R3
    @test isapprox(R1, R4; atol=1.0)
end


## Util.jl##

tupletypelength(a)=length(a.parameters)

macro replaceBase(fs...)
    b = Expr(:block)
    for f in fs
        if f == :/
            fa = :div_float
        elseif f == :^
            fa = :pow
        elseif f == :+
            fa = :vadd
        elseif f == :-
            fa = :vsub
        elseif f == :*
            fa = :vmul
        else
            fa = f
        end
        arg_consumed = false
        if fa in (:ceil,:floor,:sqrt,:rsqrt,:rec,
                  :exp,:exp2,:expm1,:log,:log1p,:log2,:log10,
                  :sin,:sinpi,:cos,:cospi,:tan,:tanpi,:asin,:acos,:atan,
                  :sinh,:cosh,:tanh,:asinh,:acosh,:atanh,
                  :trunc,:round,:exponent,:abs,:sincos,:cis)
            e = quote
                (Base.$f)(X::Array{T}) where {T <: Union{Float64,Float32}} = ($fa)(X)
                (Base.$f)(X::Union{Float64,Float32}) = ($fa)([X])[1]
            end
            arg_consumed = true
        end
        if fa in (:copysign,:atan,:pow,:rem,:div_float, :vadd, :vsub, :vmul)
            e = quote
                (Base.$f)(X::Array{T},Y::Array{T}) where {T <: Union{Float32,Float64}} = ($fa)(X,Y)
                (Base.$f)(X::T,Y::T) where {T <: Union{Float32,Float64}} = ($fa)([X],[Y])[1]
            end
            arg_consumed = true
        end
        if !arg_consumed
            error("Function $f not defined by AppleAccelerate.jl")
        end
        push!(b.args,e)
    end
    b
end


#=
  Package: FishEBM
  File: move.jl
  Devin Rose
  Functions for simulating geographical movement
  Created: June 2016
=#

"""
  Description: This function uses known information from the environment
    surrounding each agent as well as known movements to move agents around
    the environment during runtime.

  Returns: Operates directly on agent_db

  Last update: September 2016
"""
function move!(agent_db::Vector{EnviroAgent}, agent_a::AgentAssumptions,
  enviro_a::EnvironmentAssumptions, current_week::Int64)

  #@assert(0.<= agent_a.autonomy[stage] <=1., "Autonomy level for stage $stage must be between 0 and 1")

  lifeStages = Array(Int64, length(agent_db[1].alive)); lifeStages[:] = 0;
  totalHeight = size(enviro_a.habitat)[1]
  totalEnviroSize = size(enviro_a.habitat)[1]* size(enviro_a.habitat)[2]

  stageWeeks = [agent_a.growth[4], agent_a.growth[3], agent_a.growth[2], agent_a.growth[1]]

  #find the age and stage of each current cohort
  for m = 1:length(lifeStages)
    lifeStages[m] = findCurrentStage(current_week, agent_db[1].weekNum[m], agent_a.growth)
  end

  #For each agent
  for n = 1:length(agent_db)
    #simply the location id of the enviro agent
    id = agent_db[n].locationID

    #Check if the enviro agent is empty before preceeding to movement prep
    if isEmpty(agent_db[n]) == false
      #find local movement avalibility
      moveChoices = Array(Float64, 9)
      moveChoices = [
        id-totalHeight-1, id-1, id+totalHeight-1,
        id-totalHeight, id, id+totalHeight,
        id-totalHeight+1, id+1, id+totalHeight+1]

      moveChoices = hcat(moveChoices,[1,4,7,2,5,8,3,6,9])

      #remove all non water and out of enviro choices
      moveChoices = moveChoices[moveChoices[:,1] .> 0, :]
      moveChoices = moveChoices[moveChoices[:,1] .< totalEnviroSize+1, :]
      moveChoices = moveChoices[enviro_a.habitat[moveChoices[:,1]] .> 0, :]

      #for each cohort in the agent database
      for cohort = 1:length(lifeStages)
        stage = lifeStages[cohort]

        if stage > 1
          choices = deepcopy(moveChoices)

          if stage == 4
            #Create a periodicMovement function for readability
            periodicWeek = current_week%52
            moveX = -5*cosd((180*periodicWeek)/27)+1
            moveY = -5*sind((180*periodicWeek)/27)+1
            moveRadius = sqrt(moveX^2 + moveY^2)
            moveArray = Array(Float64, 3, 3)
            fill!(moveArray, 1.)

            if periodicWeek < 14
              moveArray[2,1] = -moveX
              moveArray[3,1] = moveRadius
              moveArray[3,2] = -moveY
            elseif periodicWeek > 13 && periodicWeek < 27
              moveArray[2,3] = moveX
              moveArray[3,2] = -moveY
              moveArray[3,3] = moveRadius
            elseif periodicWeek > 26 && periodicWeek < 40
              moveArray[1,2] = moveY
              moveArray[1,3] = moveRadius
              moveArray[2,3] = moveX
            elseif periodicWeek > 39
              moveArray[1,1] = moveRadius
              moveArray[1,2] = moveY
              moveArray[2,1] = -moveX
            end
          else
            moveArray = agent_a.movement[stage]
          end

          #match the moveChoices with the corresponding movement array
          for moveNum = 1:size(choices)[1]
            choices[moveNum, 2] = round(Int, moveArray[choices[moveNum, 2]])
          end

          #Match natural mortality rate by location, habitat type, and fish age
          choices = hcat(choices, 1-agent_a.naturalmortality[enviro_a.habitat[choices[:,1]], stage])

          #Normalize the choices
          choices[:,2]=choices[:,2]/sum(choices[:,2])
          choices[:,3]=choices[:,3]/sum(choices[:,3])

          moveDistrib = Multinomial(1, choices[:,2]*(1-agent_a.autonomy[stage])
            + choices[:,3]*(agent_a.autonomy[stage]))

          for aliveAges = 1:agent_db[n].alive[cohort]
            newLocation = round(Int, (choices[findfirst(rand(moveDistrib)), 1]))
            if agent_db[n].locationID != newLocation
              newAgentNum = IDToAgentNum(agent_db, newLocation, length(agent_db), 1)
              agent_db[n].alive[cohort] -= 1 # move from
              agent_db[newAgentNum].alive[cohort] += 1 # move to
            end #if doesn't move
          end #for all alive
        end #if stage > 1
      end #for each cohort
    end #if enviro empty
  end #for number agent
end

#Add a periodic movement function


using DataFrames

@testset "Synthesis" begin

# get data for use in tests
data = GRASS.SolarData()
wavt = data.wav[(:c, :mu10)][:,1]
bist = data.bis[(:c, :mu10)][:,1]
dept = data.dep[(:c, :mu10)][:,1]
widt = data.wid[(:c, :mu10)][:,1]

# set trimming parameters
top = NaN
dep = 0.75

@testset "Testing function definitions" begin
    @test isdefined(GRASS, :trim_bisector_chop!)
    @test isdefined(GRASS, :line_profile!)
    @test isdefined(GRASS, :line_from_bis!)
    @test isdefined(GRASS, :synthesize_spectra)
end

@testset "Testing line trimming" begin
    # copy the data
    wavt1 = copy(wavt)
    bist1 = copy(bist)
    dept1 = copy(dept)
    widt1 = copy(widt)

    # do the trimming
    GRASS.trim_bisector_chop!(dep, wavt1, bist1, dept1, widt1, top=top)

    # test that it wasn't changed
    @test wavt1 !== wavt
    @test bist1 !== bist
    @test dept1 !== dept
    @test widt1 !== widt

    # top and bottom
    @test bist1[1] == 1.0 - dep
    @test bist1[end] == 1.0
    @test wavt1[1] == wavt[findfirst(bist .>= 1.0 - dep)]
end

@testset "Testing line synthesis" begin
    # copy the data
    wavt1 = copy(wavt)
    bist1 = copy(bist)
    dept1 = copy(dept)
    widt1 = copy(widt)

    # do the trimming
    GRASS.trim_bisector_chop!(dep, wavt1, bist1, dept1, widt1, top=top)

    # allocate memory for line
    λs = range(5434.0, 5435.0, length=1000)
    lwavgrid = zeros(100)
    rwavgrid = zeros(100)
    allwavs = zeros(200)
    allints = zeros(200)
    prof = ones(length(λs))

    # synthesize line
    mid = 5434.5
    GRASS.line_from_bis!(mid, λs, prof, wavt1, dept1, widt1, lwavgrid, rwavgrid, allwavs, allints)

    # test it
    @test !all(prof .== 1.0)
    @test prof[1] == 1.0
    @test prof[end] == 1.0
    @test minimum(prof) == 1.0 - dep
    @test maximum(prof) == 1.0
end

@testset "Testing bisector in/out" begin
    # read in data
    wavt1 = copy(wavt)
    bist1 = copy(bist)
    dept1 = copy(dept)
    widt1 = copy(widt)

    # do the trimming
    GRASS.trim_bisector_chop!(dep, wavt1, bist1, dept1, widt1, top=top)

    # allocate memory for line
    λs = range(5434.0, 5435.0, length=1000)
    lwavgrid = zeros(100)
    rwavgrid = zeros(100)
    allwavs = zeros(200)
    allints = zeros(200)
    prof = ones(length(λs))

    # synthesize line
    mid = 5434.5
    GRASS.line_from_bis!(mid, λs, prof, wavt1, dept1, widt1, lwavgrid, rwavgrid, allwavs, allints)

    # measure the bisector of synth line
    wavo, biso = GRASS.measure_bisector(λs, prof, interpolate=false)
    wavo .-= mid

    # test that output bisector matches input bisector
    @test all(isapprox.(wavo[2:99], wavt1[2:99], atol=1e-3))
end

@testset "Testing disk-integrated spectrum synthesis" begin
    # set params
    N = 132
    Nt = 5
    res = 7e5

    # simulate spectra
    spec = SpecParams(lines=[5434.5], depths=[dep], resolution=res)
    disk = DiskParams(N=N, Nt=Nt)
    wavs, flux = synthesize_spectra(spec, disk)

    # ensure spectra have correct properties
    @test size(flux,2) == 5
    @test maximum(flux[:,1]) == 1.0
    # @test isapprox(minimum(flux[:,1]), 1.0 - dep, atol=1e-1) # TODO: fix?

    # get velocities
    v_grid, ccf1 = calc_ccf(wavs, flux, spec, normalize=true)
    rvs, sigs = calc_rvs_from_ccf(v_grid, ccf1)

    # test the velocities
    @test calc_rms(rvs) < 1.0
end

end


"""
A named tuple that contains the magic values that are used globally for
whatever purposes.
"""
const _constants = (
    default_stoich_show_size = 50_000,
    default_reaction_bound = 1e3,
    tolerance = 1e-6,
    sampling_keep_iters = 100,
    sampling_size = 1000,
    loopless_strict_inequality_tolerance = 1,
    extracellular_suffixes = ["_e", "[e]", "(e)"],
    exchange_prefixes = ["EX_", "Exch_", "Ex_", "R_EX_", "R_Ex", "R_Exch_"],
    biomass_strings = ["BIOMASS", "biomass", "Biomass"],
    keynames = (
        rxns = ["reactions", "rxns", "RXNS", "REACTIONS", "Reactions", "Rxns"],
        mets = ["metabolites", "mets", "METS", "METABOLITES", "Metabolites", "Mets"],
        genes = ["genes", "GENES", "Genes"],
        lbs = ["lbs", "lb", "lowerbounds", "lower_bounds"],
        ubs = ["ubs", "ub", "upperbounds", "upper_bounds"],
        stochiometry = ["S"],
        balance = ["b"],
        objective = ["c"],
        grrs = ["gene_reaction_rules", "grRules", "rules"],
        ids = ["id", "description"],
    ),
    gene_annotation_checks = (
        "ncbigene",
        "ncbigi",
        "refseq_locus_tag",
        "refseq_name",
        "refseq_synonym",
        "uniprot",
    ),
    reaction_annotation_checks = (
        "bigg.reaction",
        "biocyc",
        "ec-code",
        "kegg.reaction",
        "metanetx.reaction",
        "rhea",
        "sabiork",
        "seed.reaction",
    ),
    metabolite_annotation_checks = (
        "kegg.compound",
        "bigg.metabolite",
        "chebi",
        "inchi_key",
        "sabiork",
        "hmdb",
        "seed.compound",
        "metanetx.chemical",
        "reactome.compound",
        "biocyc",
    ),
    T = 298.15, # Kelvin
    R = 8.31446261815324e-3, # kJ/K/mol
)

const MAX_SENSE = MOI.MAX_SENSE
const MIN_SENSE = MOI.MIN_SENSE


import Base: display, length

struct ChequeredBlocking <: BlockingSchedule end

struct Block{S1,S2,S3,S4}
    Ls::Vector{S1}         # observation operators for both sets of blocks
    vs::Vector{S2}         # copied over observations of the process
    Σs::Vector{S3}         # covariance matrix of the noise for both sets of blocks
    aux_flags::Vector{S4}
    # two sets of knots
    knots::Vector{Int64}
    # two sets of blocks, where each block are indicies of intervals that make up a block
    blocks::Vector{Vector{Int64}}
    # info about the points at which to switch between the systems of ODEs
    change_pts::Vector{ODEChangePt}

    function Block(Ls::Vector{S1}, vs::Vector{S2}, Σs::Vector{S3},
                   aux_flags::Vector{S4}, knots::Vector{Int64},
                   blocks::Vector{Vector{Int64}},
                   change_pts::Vector{ODEChangePt}) where {S1,S2,S3,S4}
        new{S1,S2,S3,S4}(Ls, vs, Σs, aux_flags, knots, blocks, change_pts)
    end
end

length(bl::Block) = length(bl.blocks)
length(::NoBlocking) = 1

"""
    find_end_pts(𝔅::ChequeredBlocking, XX, idx)

Determine the observations for the update of the `idx`-th set of blocks. In
particular, on each block with interval indices [a₁,…,aₙ], observations vᵢ with
i∈{1,…,n-1} are made, whereas the last obesrvation is an exactly observed
process Xₜ at the terminal time t of the aₙ-th interval.
"""
function find_end_pts(bl::Block, XX)
    [( k in bl.knots ? X.yy[end] : bl.vs[k]) for (k,X) in enumerate(XX)]
end

"""
    display(𝔅::NoBlocking)

Nothing particular to display
"""
function display(bl::NoBlocking)
    print("No blocking...\n")
end

"""
    display(𝔅::ChequeredBlocking)

Display the pattern of blocks
"""
function display(bl::Block)
    function print_blocks(knots, idxLast, m)
        M = length(knots)
        get_knot(knots, i) = (M >= i > 0 ? knots[i] : idxLast * (i>0))
        function print_range(from, to)
            for i in from:to
                print("|", get_knot(knots,i-1), "|----",
                      get_knot(knots,i)-get_knot(knots,i-1), "----")
            end
        end

        if m > div(M, 2)
            print_range(1, M+1)
            print("|", idxLast, "|")
        else
            print_range(1,m)
            print("|", get_knot(knots, m) ,"|   ...   ")
            print_range(M+2-m,M+1)
            print("|", idxLast, "|")
        end
        print("  (number of blocks: ", M+1,")")
    end
    print("Chequered Blocking scheme\n",
          "-------------------------\n",
          "Format:\n",
          "block sizes in A: ")
    print_blocks(𝔅.knots[1], length(𝔅.vs), 3)
    print("\nblock sizes in B: ")
    print_blocks(𝔅.knots[2], length(𝔅.vs), 3)
    print("\n")
end

function _print_info(info)
    m = length(info)
    for i in 1:m
        print("b", i, ": ", round(info[i], digits=3), " | ")
    end
    print("\n- - - - - - - - - - - - - -\n")
end

function create_blocks(::ChequeredBlocking, P, params)
    knots, ϵ, change_pt = params
    find_knots(mod, rem) = [k for (i,k) in enumerate(knots) if i % mod == rem]
    knotsA, knotsB = find_knots(2, 1), find_knots(2, 0)

    m, d = size(P[end].L)

    find_L(knots) = [( k in knots ? SMatrix{d,d}(1.0*I) : p.L)
                     for (k,p) in enumerate(P)]
    LsA, LsB = find_L(knotsA), find_L(knotsB)

    vs = [p.v for p in P]

    find_Σ(knots) = [(k in knots ? SMatrix{d,d}(ϵ*I) : p.Σ)
                     for (k,p) in enumerate(P)]
    ΣsA, ΣsB = find_Σ(knotsA), find_Σ(knotsB)

    find_ch_pt(knots) = [(k in knots ? deepcopy(change_pt) : p.change_pt)
                         for (k,p) in enumerate(P)]
    chpA, chpB = find_ch_pt(knotsA), find_ch_pt(knotsB)

    find_aux_flag(knots) = [( k in knots ? Nothing : get_aux_flag(p.Pt))
                            for (k,p) in enumerate(P)]
    auxA, auxB = find_aux_flag(knotsA), find_aux_flag(knotsB)

    """
        knots_to_blocks(knots, idxLast, i)

    Given a list of `knots` fetch the indices of the intervals that together
    make up the `i`-th block. `idxLast` is the index of the last interval.
    """
    function knots_to_blocks(knots, idx_last, i)
        M = length(knots)
        @assert M > 0
        if M >= i > 1
            return (knots[i-1]+1):knots[i]
        elseif i == 1
            return 1:knots[1]
        else
            return (knots[M]+1):idx_last
        end
    end
    blocks = ([collect(knots_to_blocks(knotsA, length(P), i)) for i in 1:length(knotsA)+1],
              [collect(knots_to_blocks(knotsB, length(P), i)) for i in 1:length(knotsB)+1])

    (Block(LsA, vs, ΣsA, auxA, knotsA, blocks[1], chpA),
     Block(LsB, vs, ΣsB, auxB, knotsB, blocks[2], chpB))
end

create_blocks(blocking::NoBlocking, P, params) = blocking


using RobotZoo
using RobotDynamics
using Test
using StaticArrays
using BenchmarkTools
using Rotations
using Random

@testset "Allocations" begin
    include("models_test.jl")
end

@testset "LinearModels" begin
    include("linear_models_test.jl")
    include("linear_models_examples_test.jl")
end


using ImageCore, Colors, FixedPointNumbers, OffsetArrays
using Test, Random

@info "Deprecations are expected"
@testset "permuteddimsview" begin
    a = [1 3; 2 4]
    v = permuteddimsview(a, (1,2))
    @test v == a
    v = permuteddimsview(a, (2,1))
    @test v == a'
    a = rand(3,7,5)
    v = permuteddimsview(a, (2,3,1))
    @test v == permutedims(a, (2,3,1))
end


struct UniformPeriodicReconstructedFV1D{BalanceLaw, T, Fnum, Reconstruction,
                                        EdgeU, Fluxes, Parallel} <: FiniteVolumeSemidiscretisation
    balance_law::BalanceLaw
    meshx::UniformPeriodicMesh1D{T}
    fnum::Fnum
    reconstruction::Reconstruction

    edge_u::EdgeU
    fluxes::Fluxes
    parallel::Parallel

    function UniformPeriodicReconstructedFV1D(
            balance_law::BalanceLaw, meshx::UniformPeriodicMesh1D{T},
            fnum::Fnum, reconstruction::Reconstruction,
            edge_u::EdgeU, fluxes::Fluxes, parallel::Parallel) where {BalanceLaw, T, Fnum, Reconstruction, EdgeU, Fluxes, Parallel}
        @assert typeof(fluxes) <: AbstractArray{variables(balance_law), 1}
        @assert typeof(edge_u) <: AbstractArray{variables(balance_law), 2}
        @assert size(fluxes, 1) == size(edge_u, 2) == numcells(meshx)
        @assert size(edge_u, 1) == 2
        new{BalanceLaw, T, Fnum, Reconstruction, EdgeU, Fluxes, Parallel}(
            balance_law, meshx, fnum, reconstruction, edge_u, fluxes, parallel)
    end
end

function UniformPeriodicReconstructedFV1D(balance_law, meshx, fnum, reconstruction,
                                          parallel=Val{:serial}())
    fluxes = Array{variables(balance_law)}(undef, numcells(meshx))
    edge_u = Array{variables(balance_law)}(undef, 2, numcells(meshx))
    UniformPeriodicReconstructedFV1D(balance_law, meshx, fnum, reconstruction,
                                     edge_u, fluxes, parallel)
end

function Base.show(io::IO, fv::UniformPeriodicReconstructedFV1D)
    print(io, "Finite volume method ",
            "\n  Balance law:    ", fv.balance_law,
            "\n  Mesh:           ", fv.meshx,
            "\n  Numerical flux: ", fv.fnum,
            "\n  Reconstruction: ", fv.reconstruction)
end

function evaluate_coefficients(u, fv::UniformPeriodicReconstructedFV1D,
                               npoints=2*order(fv.reconstruction))
    evaluate_coefficients(u, fv.balance_law, fv.meshx, fv.reconstruction, npoints)
end

@noinline function (semidisc::UniformPeriodicReconstructedFV1D)(du, u, p, t)
    @boundscheck begin
        if size(u) != size(du)
            error("size(u) = $(size(u)) != $(size(du)) = size(du)")
        end
        @assert length(u) == numcells(semidisc.meshx)
        @assert length(u) >= stencil_width(semidisc.reconstruction) - 1
        if eltype(u) != variables(semidisc.balance_law)
            error("eltype(u) == $(eltype(u)) != $(variables(semidisc.balance_law)) == variables(semidisc.balance_law)")
        end
    end

    fill!(du, zero(eltype(du)))
    add_numerical_fluxes!(du, u, semidisc, t)

    nothing
end

function add_numerical_fluxes!(du, u, fv::UniformPeriodicReconstructedFV1D, t)
    @unpack balance_law, meshx, fnum, reconstruction, edge_u, fluxes, parallel = fv

    reconstruct!(edge_u, u, balance_law, meshx, reconstruction, stencil_width_val(reconstruction), parallel)
    compute_numerical_fluxes!(fluxes, edge_u, balance_law, meshx, fnum, parallel)
    compute_du!(du, fluxes, meshx, parallel)
end

function compute_numerical_fluxes!(fluxes, edge_u, balance_law, meshx::UniformPeriodicMesh1D,
                                   fnum, parallel)
    fluxes[1] = fnum(edge_u[2,end], edge_u[1,1], balance_law)
    for edge in 2:length(fluxes)
        fluxes[edge] = fnum(edge_u[2,edge-1], edge_u[1,edge], balance_law)
    end
    nothing
end

function compute_numerical_fluxes!(fluxes, edge_u, balance_law, meshx::UniformPeriodicMesh1D,
                                   fnum, parallel::Val{:threads})
    @inbounds fluxes[1] = fnum(edge_u[2,end], edge_u[1,1], balance_law)
    @inbounds Threads.@threads for edge in 2:length(fluxes)
        fluxes[edge] = fnum(edge_u[2,edge-1], edge_u[1,edge], balance_law)
    end
    nothing
end

function compute_du!(du, fluxes, meshx::UniformPeriodicMesh1D, parallel::Val{:serial})
    @inbounds for cell in 1:numcells(meshx)-1
        du[cell] -= ( fluxes[cell+1] - fluxes[cell] ) / volume(cell, meshx)
    end
    @inbounds du[end] -= ( fluxes[1] - fluxes[end] ) / volume(numcells(meshx), meshx)
    nothing
end

function compute_du!(du, fluxes, meshx::UniformPeriodicMesh1D, parallel::Val{:threads})
    @inbounds Threads.@threads for cell in 1:numcells(meshx)-1
        du[cell] -= ( fluxes[cell+1] - fluxes[cell] ) / volume(cell, meshx)
    end
    @inbounds du[end] -= ( fluxes[1] - fluxes[end] ) / volume(numcells(meshx), meshx)
    nothing
end

################################################################################

function semidiscretise(fv::UniformPeriodicReconstructedFV1D, u₀func, tspan)
    u₀ = compute_coefficients(u₀func, fv.meshx, order(fv.reconstruction))

    ODEProblem(fv, u₀, tspan)
end

"""
    max_dt(t, u, fv::UniformPeriodicReconstructedFV1D, cfl=0.5)

Compute the maximal time step `dt` satisfying the CFL condition `dt <= cfl * dx / 2λ`,
where `dx` is the length of a cell and `λ` the greatest absolute value of the
speed in this cell.
"""
function max_dt(t, u, fv::UniformPeriodicReconstructedFV1D, cfl=0.5)
    @unpack balance_law, meshx = fv
    dt = mapreduce(cell->volume(cell,meshx) / max_abs_speed(u[cell],balance_law),
                    min, cell_indices(meshx), init=typemax(t))
    if dt == Inf
        dt = mapreduce(cell->volume(cell,meshx), min, cell_indices(meshx), init=dt)
    end
    cfl * dt
end


"""
    ncdim_def(ds, name, val, attrib = Dict())
    ncdim_def(ds, dim::NcDim)
    ncdim_def(ds, dims::Vector{NcDim})

# Examples
```julia
## 1. by existing nc file
# f = "/mnt/j/Researches/Obs/Obs_ChinaHW_cluster/ncell_16/clusterId_HItasmax_movTRS_Obs.nc"
dims = nc_dims(f)
ds = nc_open("f2.nc", "c")
ncdim_def(ds, dims);
println(ds)
close(ds)

## 2. by assigned variables
# dates = nc_date(ds)
ds = nc_open("f3.nc", "c")
ncdim_def(ds, "time",
    CFTime.timeencode.(dates, "days since 1970-01-01", eltype(dates)),
    Dict("units" => "days since 1970-01-01",
        "calendar" => "proleptic_gregorian", "long_name" => "time"))
close(ds)
```
"""
function ncdim_def(ds, name, val, attrib = Dict(); verbose = false)
    # x = NcDim(name, val, attrib)
    # val = val |> collect
    if name in keys(ds.dim)
        verbose && @warn "Dimension `$name`: exist!"
        return
    end
    defDim(ds, name, length(val))
    defVar(ds, name, val, (name,); attrib = attrib)
end

# NcDim
function ncdim_def(ds, dim::NcDim)
    ncdim_def(ds, dim.name, dim.vals, dim.atts)
end

function ncdim_def(ds, dims::Vector{NcDim})
    for i = 1:length(dims)
        dim = dims[i]
        ncdim_def(ds, dim.name, dim.vals, dim.atts)
    end
end


export get_agent_view, AbstractGridWorld
export get_actions
abstract type AbstractGridWorld end

function get_agent_view end
function get_agent end

Base.convert(::Type{GridWorldBase}, w::AbstractGridWorld) = w.world
get_object(w::AbstractGridWorld) = get_object(convert(GridWorldBase, w))
get_object(w::AbstractGridWorld, x::Type{<:AbstractObject}) = filter(o -> o isa x, get_object(w))
get_object(w::AbstractGridWorld, x::Type{Agent}) = w.agent
get_pos(w::AbstractGridWorld, ::Type{Agent}) = w.agent_pos

get_agent(w::AbstractGridWorld) = get_object(w, Agent)
get_agent_pos(w::AbstractGridWorld) = get_pos(w, Agent)
get_agent_dir(w::AbstractGridWorld) = w |> get_agent |> get_dir

function get_agent_view(w::AbstractGridWorld, agent_view_size=(7,7))
    wb = convert(GridWorldBase, w)
    v = BitArray{3}(undef, size(wb, 1), agent_view_size...)
    fill!(v, false)
    get_agent_view!(v, w)
end

function (w::AbstractGridWorld)(action::Union{TurnRight, TurnLeft})
    a = get_agent(w)
    set_dir!(a, action(get_dir(a)))
    w
end

get_actions(w::AbstractGridWorld) = (MOVE_FORWARD, TURN_LEFT, TURN_RIGHT)

get_agent_view_inds(w::AbstractGridWorld, s=(7,7)) = get_agent_view_inds(get_agent_pos(w).I, s, get_agent_dir(w))

get_agent_view!(v::BitArray{3}, w::AbstractGridWorld) = get_agent_view!(v, convert(GridWorldBase, w), get_agent_pos(w), get_agent_dir(w))


module TimeIt

if VERSION >= v"0.7.0"
    using Printf
end

export @timeit

macro timen(ex, n)
    quote
        local t0 = time_ns()
        for i = 1:$(esc(n))
            local val = $(esc(ex))
        end
        local t1 = time_ns()
        (t1 - t0) / 1.e9
    end
end

macro timeit(ex)
    quote
        local val = $(esc(ex))  # Warm up
        t = zeros(3)

        # Determine number of loops so that total time > 0.1s.
        n = 1
        for i = 0:9
            n = 10^i
            t[1] = @timen $(esc(ex)) n
            if t[1] > 0.1
                break
            end
        end

        # Two more production runs.
        for i = 2:3
            t[i] = @timen $(esc(ex)) n
        end
        best = minimum(t) / n

        # Format to nano-, micro- or milliseconds.
        factor = 1.
        if best < 1e-6
            factor = 1e9
            pre = "n"
        elseif best < 1e-3
            factor = 1e6
            pre = "\u00b5"
        elseif best < 1.
            factor = 1e3
            pre = "m"
        else
            factor = 1.
            pre = ""
        end
        @printf "%d loops, best of 3: %4.2f %ss per loop\n" n best*factor pre
        best
    end
end

end  # module


using YaoBase, YaoArrayRegister, Random
export Measure, AllLocs, ComputationalBasis, chmeasureoperator

"""
    Measure{N, K, OT, RNG, IT<:Integer} <: PrimitiveBlock{N}
    Measure(n::Int; rng=Random.GLOBAL_RNG, operator=ComputationalBasis(), locs=1:n, collapseto=nothing, remove=false)

Measure operator.
"""
mutable struct Measure{N,K,OT,RNG,IT<:Integer} <: PrimitiveBlock{N}
    rng::RNG
    operator::OT
    locations::Union{NTuple{K,Int},AllLocs}
    collapseto::Union{BitStr64{N},Nothing}
    remove::Bool
    results::Vector{IT}
    function Measure{N,K,OT,RNG,IT}(rng::RNG, operator, locations, collapseto, remove) where {RNG, N, K, OT, IT}
        locations isa AllLocs || @assert_locs_safe N locations
        if collapseto !== nothing && remove == true
            error("invalid keyword combination, expect collapseto or remove, got (collapseto=$collapseto, remove=true)")
        end
        new{N,K,OT,RNG,IT}(rng, operator, locations, collapseto, remove)
    end
end

"""
    chmeasureoperator(m::Measure, op::AbstractBlock)

change the measuring `operator`. It will also discard existing measuring results.
"""
function chmeasureoperator(m::Measure{N}, op::AbstractBlock) where {N}
    Measure(
        N;
        rng = m.rng,
        operator = op,
        locs = m.locations,
        collapseto = m.collapseto,
        remove = m.remove,
    )
end

function Base.:(==)(m1::Measure, m2::Measure)
    res = m1.rng == m2.rng &&
          m1.operator == m2.operator &&
          m1.locations == m2.locations && m1.collapseto == m2.collapseto && m1.remove == m2.remove
    res = res && isdefined(m1, :results) == isdefined(m2, :results)
    res && (!isdefined(m1, :results) || m1.results == m2.results)
end

@interface nqubits_measured(::Measure{N,K}) where {N,K} = K

"""
    Measure(n::Int; rng=Random.GLOBAL_RNG, operator=ComputationalBasis(), locs=AllLocs(), collapseto=nothing, remove=false)

Create a `Measure` block with number of qubits `n`.

# Example

You can create a `Measure` block on given basis (default is the computational basis).

```jldoctest; setup=:(using YaoBlocks)
julia> Measure(4)
Measure(4)
```

Or you could specify which qubits you are going to measure

```jldoctest; setup=:(using YaoBlocks)
julia> Measure(4; locs=1:3)
Measure(4;locs=(1, 2, 3))
```

by default this will collapse the current register to measure results.

```jldoctest; setup=:(using YaoBlocks, YaoArrayRegister, BitBasis, Random; Random.seed!(2))
julia> r = rand_state(3)
ArrayReg{1, Complex{Float64}, Array...}
    active qubits: 3/3

julia> state(r)
8×1 Array{Complex{Float64},2}:
    0.21633342515406265 - 0.21776267239802458im
   -0.17798384008375148 - 0.5040979387214165im
   -0.19761243345925425 + 0.16281482444784728im
   -0.25200691415025867 + 0.15153595884416518im
     0.3650977378140692 + 0.3419566592091794im
  -0.027207023333497483 - 0.3780181361735894im
 -0.0034728372576413743 + 0.1693915490059622im
   -0.19898587237095824 - 0.07607057769761456im

julia> r |> Measure(3)
Measure(3)

julia> state(r)
8×1 Array{Complex{Float64},2}:
                0.0 + 0.0im
                0.0 + 0.0im
                0.0 + 0.0im
                0.0 + 0.0im
 0.7298587746534583 + 0.6835979586433478im
                0.0 + 0.0im
                0.0 + 0.0im
                0.0 + 0.0im
```

But you can also specify the target bit configuration you want to collapse to with keyword `collapseto`.

```jldoctest; setup=:(using YaoBlocks; using BitBasis)
julia> m = Measure(4; collapseto=bit"0101")
Measure(4;collapseto=0101 ₍₂₎)

julia> m.collapseto
0101 ₍₂₎
```
"""
function Measure(n::Int; rng::RNG=Random.GLOBAL_RNG, operator::OT=ComputationalBasis(), locs=AllLocs(), collapseto=nothing, remove=false, result_dtype=BitStr64{n}) where {OT, RNG}
    if locs isa AllLocs
        Measure{n, n, OT, RNG, result_dtype}(rng, operator, locs, collapseto, remove)
    else
        Measure{n, length(locs), OT, RNG, result_dtype}(rng, operator, tuple(locs...), collapseto, remove)
    end
end

Measure(
    ;
    rng = Random.GLOBAL_RNG,
    locs = AllLocs(),
    operator = ComputationalBasis(),
    collapseto = nothing,
    remove = false,
) where {K} = @λ(
    n -> Measure(
        n;
        rng = rng,
        locs = locs,
        operator = operator,
        collapseto = collapseto,
        remove = remove,
    )
)
mat(x::Measure) = error("use BlockMap to get its matrix.")

function apply!(r::AbstractRegister, m::Measure{N}) where {N}
    _check_size(r, m)
    if m.collapseto !== nothing
        m.results = measure_collapseto!(m.rng, m.operator, r, m.locations; config = m.collapseto)
    elseif m.remove
        m.results = measure_remove!(m.rng, m.operator, r, m.locations)
    else
        m.results = measure!(m.rng, m.operator, r, m.locations)
    end
    return m
end


using IterTools

cd(dirname(@__FILE__))
const lines = open("four.txt", "r") do f
    readlines(f)
end

parse_line(l::String)::Array{String} =
    [split(s,":")[1] for s in split(l)]

join_by_space(l) = join(l, " ")

valid_set = Set(["hcl", "iyr", "eyr", "ecl", "pid", "byr", "hgt"])
is_valid(s) = issubset(valid_set, s)

n = (groupby(!isempty, lines)
    |> l -> map(join_by_space, l)
    |> l -> filter(!isempty, l)
    |> l -> map(parse_line, l)
    |> s -> map(is_valid, s)
    |> sum)

println(n)


module SimpleSparseArrays
using Base.Cartesian
using LinearAlgebra

include("core.jl")
include("testing.jl")

end#module


struct DDEProblem{uType,tType,lType,lType2,isinplace,P,F,H,C} <:
                          AbstractDDEProblem{uType,tType,lType,isinplace}
  f::F
  u0::uType
  h::H
  tspan::tType
  p::P
  constant_lags::lType
  dependent_lags::lType2
  callback::C
  neutral::Bool
  @add_kwonly function DDEProblem{iip}(f::AbstractDDEFunction{iip},
                                 u0,h,tspan,p=nothing;
                                 constant_lags=[],
                                 dependent_lags=[],
                                 neutral = f.mass_matrix == I ?
                                           false : det(f.mass_matrix)!=1,
                                 callback = nothing) where {iip}
    _tspan = promote_tspan(tspan)
    new{typeof(u0),typeof(_tspan),
               typeof(constant_lags),typeof(dependent_lags),
               isinplace(f),typeof(p),
               typeof(f),typeof(h),typeof(callback)
               }(f,u0,h,_tspan,p,
                 constant_lags,
                 dependent_lags,callback,
                 neutral)
  end

  function DDEProblem{iip}(f,u0,h,tspan,p=nothing;kwargs...) where {iip}
    DDEProblem(convert(DDEFunction{iip},f),u0,h,tspan,p;kwargs...)
  end

end

function DDEProblem(f::AbstractDDEFunction,u0,h,tspan,p=nothing;kwargs...)
  DDEProblem{isinplace(f)}(f,u0,h,tspan,p;kwargs...)
end

function DDEProblem(f,u0,h,tspan,p=nothing;kwargs...)
  DDEProblem(convert(DDEFunction,f),u0,h,tspan,p;kwargs...)
end


module LyapunovExponents

export LEProblem, ContinuousLEProblem, DiscreteLEProblem, lyapunov_exponents,
    phase_state, report

# Re-export methods from DifferentialEquations extended here:
export init, solve, solve!, step!
import DiffEqBase: init, step!
using DiffEqBase: solve!, set_u!
# solve and solve! are imported/extended in stages.jl
using LinearAlgebra
using Printf
using Markdown

include("documents.jl")
include("stages.jl")
include("utils.jl")
include("online_stats.jl")
include("types.jl")
include("core.jl")
include("terminators.jl")
include("coevolve.jl")
include("continuous_exponents.jl")
include("discrete_exponents.jl")
include("rode.jl")
include("tangent_utils.jl")
include("clv/covariant_vectors.jl")
include("interface.jl")
include("implementations.jl")
include("reports.jl")
include("examples/examples.jl")
include("test.jl")
include("diffeq_hack.jl")
include("plotting.jl")

end # module


using MAT
import Nettle

images_url = "http://www.rctn.org/bruno/sparsenet/IMAGES.mat"
images_file = "$(@__DIR__)/IMAGES.mat"
images_md5 = "b20cb33935e95b556cb83bef5e9afcf1"

file_ok = false

if isfile(images_file)
    h = Nettle.Hasher("md5")
    Nettle.update!(h, read(images_file))
    if Nettle.hexdigest!(h) == images_md5
        file_ok = true
    end
end

if !file_ok
    println("Downloading $images_url")
    download(images_url, images_file)
end

function nat_images()
    f = matopen(images_file)
    read(f, "IMAGES")
end


include("./Algos.jl")
using .Algos
using CSV
using DataFrames
using Plots

MAP_NAME = "map2"
TRAINING_EPS = 300000
EVALUATING_EPS = 100000
SAMPLING_RATE = 0.001

df = DataFrame(CSV.File("maps/$(MAP_NAME).csv"));
tiles = Matrix(df);
agent = Sarsa(tiles, learningRate = 0.1);
y1 = train(agent, TRAINING_EPS, samplingRate = SAMPLING_RATE);
Algos.visualizeOptimalPolicy(agent.env, agent.qTable, filename = "assets/sarsa_$(MAP_NAME)_optimal.png");

yt1 = Algos.evaluate(agent, EVALUATING_EPS, samplingRate = SAMPLING_RATE);

tiles = Matrix(df);
agent = QLearning(tiles, learningRate = 0.1);
y2 = train(agent, TRAINING_EPS, samplingRate = SAMPLING_RATE);
Algos.visualizeOptimalPolicy(agent.env, agent.qTable, filename = "assets/QL_$(MAP_NAME)_optimal.png");

yt2 = Algos.evaluate(agent, EVALUATING_EPS, samplingRate = SAMPLING_RATE);

fig = plot(collect(1:length(y1)) .* trunc(Int, 1 / SAMPLING_RATE), hcat(y1, y2), xlabel = "Episodes", ylabel = "Average Reward per Episode",
            label = ["SARSA" "Q-Learning"], lw = 2, dpi = 300, xlabelfontsize = 8, ylabelfontsize = 8,
            legend=:right, title = "Training", titlefontsize = 10);

savefig(fig, "assets/evaluation.png");

fig = plot(collect(1:length(yt1)) .* trunc(Int, 1 / SAMPLING_RATE), hcat(yt1, yt2), xlabel = "Episodes", ylabel = "Average Reward per Episode",
            label = ["SARSA" "Q-Learning"], lw = 2, dpi = 300, xlabelfontsize = 10, ylabelfontsize = 10,
            legend=:right, title = "Testing", titlefontsize = 10);

savefig(fig, "assets/evaluation_trained.png");


#test file

x = 3
y = 4

x + y

using DataFrames
df = DataFrame(a = 1:20,b = rand([2,5,10],20), x = rand(20))
using Plots
plot(rand(10))


using ThreeBodyDecay
using Test

Ps = ThreeBodyParities('+','-','+'; P0='+')

@testset "Three Body Parities structure" begin
    #
    @test Ps.P1 == Ps[1] == '+'
    @test Ps.P2 == Ps[2] == '-'
    @test Ps.P3 == Ps[3] == '+'
    @test Ps.P0 == Ps[0] == Ps[4] == '+'
    #
    @test_throws ErrorException Ps[5]
    @test_throws ErrorException ThreeBodyParities(0,1,0; two_js0=1)
    @test_throws ErrorException ThreeBodyParities('+','+','+')
end

@testset "operations and interate" begin
    @test sum(i == '+' for i in Ps) == 3
end


module Verlet

export VerletIntegrator
include("./types.jl")
include("./main.jl")

end


# updates.jl
#
# Defines the functions that perform the diagonal update, and also
# that build the linked list and operator cluster update

#  (-2,i) is an off-diagonal site operator h(sigma^+_i + sigma^-_i)
#  (-1,i) is a diagonal site operator h
#  (0,0) is the identity operator I - NOT USED IN THE PROJECTOR CASE
#  (i,j) is a diagonal bond operator J(sigma^z_i sigma^z_j)
@inline isdiagonal(op::NTuple{2,Int}) = (op[1] != -2)
@inline isidentity(op::NTuple{2,Int}) = (op[1] == 0)
@inline issiteoperator(op::NTuple{2,Int}) = (op[1] < 0)
@inline isbondoperator(op::NTuple{2,Int}) = (op[1] > 0)

function mc_step!(f::Function, qmc_state::BinaryGroundState, H::TFIM)
    diagonal_update!(qmc_state, H)

    cluster_data = linked_list_update(qmc_state, H)

    f(cluster_data, qmc_state, H)

    cluster_update!(cluster_data, qmc_state, H)
end

mc_step!(qmc_state, H) = mc_step!((args...) -> nothing, qmc_state, H)

########################## finite-beta #######################################

function mc_step_beta!(f::Function, qmc_state::BinaryThermalState, H::TFIM, beta::Real; eq::Bool = false)
    num_ops = diagonal_update_beta!(qmc_state, H, beta; eq = eq)

    cluster_data = linked_list_update_beta(qmc_state, H)

    f(cluster_data, qmc_state, H)

    cluster_update_beta!(cluster_data, qmc_state, H)

    return num_ops
end

mc_step_beta!(qmc_state, H, beta; eq = false) = mc_step_beta!((args...) -> nothing, qmc_state, H, beta; eq = eq)

# returns true if operator insertion succeeded
function insert_diagonal_operator!(qmc_state::BinaryQMCState{N}, H::TFIM{N}, spin_prop::BitArray{N}, n::Int) where N
    site1, site2 = op = rand(H.op_sampler)
    @inbounds if issiteoperator(op) || spin_prop[site1] == spin_prop[site2]
        qmc_state.operator_list[n] = op
        return true
    else
        return false
    end
end

#############################################################################

function diagonal_update!(qmc_state::BinaryGroundState{N}, H::TFIM{N}) where N
    spin_prop = copyto!(qmc_state.propagated_config, qmc_state.left_config)  # the propagated spin state

    for (n, op) in enumerate(qmc_state.operator_list)
        if !isdiagonal(op)
            @inbounds spin_prop[op[2]] ⊻= 1  # spinflip
        else
            success = false
            while !success
                success = insert_diagonal_operator!(qmc_state, H, spin_prop, n)
            end
        end
    end

    # DEBUG
    # if spin_prop != qmc_state.right_config  # check the spin propagation for error
    #     error("Basis state propagation error in diagonal update!")
    # end
end

#############################################################################

function linked_list_update(qmc_state::BinaryGroundState{N}, H::TFIM{N}) where N
    Ns = nspins(H)
    spin_left = qmc_state.left_config

    len = 2 * Ns
    @simd for op in qmc_state.operator_list
        len += ifelse(issiteoperator(op), 2, 4)
    end

    # retrieve linked list data structures
    LinkList = qmc_state.linked_list  # needed for cluster update
    LegType = qmc_state.leg_types

    # A diagonal bond operator has non trivial associates for cluster building
    Associates = qmc_state.associates

    First = qmc_state.first

    # The first N elements of the linked list are the spins of the LHS basis state
    @inbounds @simd for i in 1:Ns
        LegType[i] = spin_left[i]
        First[i] = i
        Associates[i] = (0, 0, 0)
    end

    spin_prop = copyto!(qmc_state.propagated_config, spin_left)  # the propagated spin state
    idx = Ns

    # Now, add the 2M operators to the linked list. Each has either 2 or 4 legs
    @inbounds for op in qmc_state.operator_list
        if issiteoperator(op)
            site = op[2]
            # lower or left leg
            idx += 1
            LinkList[idx] = First[site]
            LegType[idx] = spin_prop[site]
            current_link = idx

            if !isdiagonal(op)  # off-diagonal site operator
                spin_prop[site] ⊻= 1  # spinflip
            end

            LinkList[First[site]] = current_link  # completes backwards link
            First[site] = current_link + 1
            Associates[idx] = (0, 0, 0)

            # upper or right leg
            idx += 1
            LegType[idx] = spin_prop[site]
            Associates[idx] = (0, 0, 0)
        else  # diagonal bond operator
            site1, site2 = op

            # lower left
            idx += 1
            LinkList[idx] = First[site1]
            LegType[idx] = spin_prop[site1]
            current_link = idx

            LinkList[First[site1]] = current_link  # completes backwards link
            First[site1] = current_link + 2
            vertex1 = current_link
            Associates[idx] = (vertex1 + 1, vertex1 + 2, vertex1 + 3)

            # lower right
            idx += 1
            LinkList[idx] = First[site2]
            LegType[idx] = spin_prop[site2]
            current_link = idx

            LinkList[First[site2]] = current_link  # completes backwards link
            First[site2] = current_link + 2
            Associates[idx] = (vertex1, vertex1 + 2, vertex1 + 3)

            # upper left
            idx += 1
            LegType[idx] = spin_prop[site1]
            Associates[idx] = (vertex1, vertex1 + 1, vertex1 + 3)

            # upper right
            idx += 1
            LegType[idx] = spin_prop[site2]
            Associates[idx] = (vertex1, vertex1 + 1, vertex1 + 2)
        end
    end

    # The last N elements of the linked list are the final spin state
    @inbounds for i in 1:Ns
        idx += 1
        LinkList[idx] = First[i]
        LegType[idx] = spin_prop[i]
        LinkList[First[i]] = idx
        Associates[idx] = (0, 0, 0)
    end

    # DEBUG
    # if spin_prop != spin_right
    #     @debug "Basis state propagation error: LINKED LIST"
    # end

    return len
end

#############################################################################

function cluster_update!(lsize::Int, qmc_state::BinaryGroundState{N}, H::TFIM{N}) where N
    Ns = nspins(H)
    spin_left, spin_right = qmc_state.left_config, qmc_state.right_config
    operator_list = qmc_state.operator_list

    LinkList = qmc_state.linked_list
    LegType = qmc_state.leg_types
    Associates = qmc_state.associates

    in_cluster = zeros(Int, lsize)
    cstack = Stack{Int}()  # This is the stack of vertices in a cluster
    ccount = 0  # cluster number counter

    @inbounds for i in 1:lsize
        # Add a new leg onto the cluster
        if (in_cluster[i] == 0 && Associates[i] === (0, 0, 0))
            ccount += 1
            push!(cstack, i)
            in_cluster[i] = ccount

            flip = rand(Bool)  # flip a coin for the SW cluster flip
            if flip
                LegType[i] ⊻= 1  # spinflip
            end

            while !isempty(cstack)
                leg = LinkList[pop!(cstack)]

                if in_cluster[leg] == 0
                    in_cluster[leg] = ccount  # add the new leg and flip it
                    if flip
                        LegType[leg] ⊻= 1
                    end

                    # now check all associates and add to cluster
                    assoc = Associates[leg]  # a 3-tuple
                    if assoc !== (0, 0, 0)
                        for a in assoc
                            push!(cstack, a)
                            in_cluster[a] = ccount
                            if flip
                                LegType[a] ⊻= 1
                            end
                        end
                    end
                end

            end
        end
    end

    # map back basis states and operator list
    @inbounds for i in 1:Ns
        spin_left[i] = LegType[i]  # left basis state
        spin_right[i] = LegType[lsize-Ns+i]  # right basis state
    end

    ocount = Ns + 1  # next on is leg Ns + 1
    @inbounds for (n, op) in enumerate(operator_list)
        if isbondoperator(op)
            ocount += 4
        else
            if LegType[ocount] == LegType[ocount+1]  # diagonal
                operator_list[n] = (-1, op[2])
            else  # off-diagonal
                operator_list[n] = (-2, op[2])
            end
            ocount += 2
        end
    end

end

#############################################################################
#############  FINITE BETA FUNCTIONS BELOW ##################################
#############################################################################

function resize_op_list!(qmc_state::BinaryThermalState, new_size::Int)
    operator_list = filter!(!isidentity, qmc_state.operator_list)
    len = length(operator_list)

    if len < new_size
        tail = init_op_list(new_size - len)
        append!(operator_list, tail)
    end

    len = 4*length(operator_list)
    # these are going to be overwritten by linked_list_update_beta
    # which will be called right after the diagonal update that called this function
    resize!(qmc_state.linked_list, len)
    resize!(qmc_state.leg_types, len)
    resize!(qmc_state.associates, len)
end

function diagonal_update_beta!(qmc_state::BinaryThermalState, H::TFIM, beta::Real; eq::Bool = false)

    # define the Metropolis probability as a constant
    # https://pitp.phas.ubc.ca/confs/sherbrooke2012/archives/Melko_SSEQMC.pdf
    # equation 1.42
    P_norm = beta * H.P_normalization

    num_ids = count(isidentity, qmc_state.operator_list)
    P_remove = (num_ids + 1) / P_norm
    P_accept = P_norm / num_ids

    spin_prop = copyto!(qmc_state.propagated_config, qmc_state.left_config)  # the propagated spin state

    @inbounds for (n, op) in enumerate(qmc_state.operator_list)
        if !isdiagonal(op)
            spin_prop[op[2]] ⊻= 1  # spinflip
        elseif !isidentity(op)
            if rand() < P_remove
                qmc_state.operator_list[n] = (0, 0)
                num_ids += 1
                P_remove = (num_ids + 1) / P_norm
                P_accept = P_norm / num_ids
            end
        else
            if rand() < P_accept
                success = insert_diagonal_operator!(qmc_state, H, spin_prop, n)

                if success
                    # save one operation lol
                    P_remove = num_ids / P_norm
                    num_ids -= 1
                    P_accept = P_norm / num_ids
                end
            end
        end
    end

    # DEBUG
    # if spin_prop != qmc_state.right_config  # check the spin propagation for error
    #     error("Basis state propagation error in diagonal update!")
    # end

    total_list_size = length(qmc_state.operator_list)
    num_ops = total_list_size - num_ids

    if eq && 1.2*num_ops > total_list_size
        resize_op_list!(qmc_state, round(Int, 1.5*num_ops, RoundUp))
    end

    return num_ops
end

#############################################################################

function linked_list_update_beta(qmc_state::BinaryThermalState, H::TFIM)
    Ns = nspins(H)
    spin_left, spin_right = qmc_state.left_config, qmc_state.right_config

    len = 0
    @simd for op in qmc_state.operator_list
        if issiteoperator(op)
            len += 2
        elseif isbondoperator(op)
            len += 4
        end
    end

    # initialize linked list data structures
    LinkList = qmc_state.linked_list  # needed for cluster update
    LegType = qmc_state.leg_types

    # A diagonal bond operator has non trivial associates for cluster building
    Associates = qmc_state.associates

    First = fill!(qmc_state.first, 0)  #initialize the First list
    Last = fill!(qmc_state.last, 0)   #initialize the Last list
    idx = 0

    spin_prop = copyto!(qmc_state.propagated_config, spin_left)  # the propagated spin state

    # Now, add the 2M operators to the linked list. Each has either 2 or 4 legs
    @inbounds for op in qmc_state.operator_list
        if issiteoperator(op)
            site = op[2]
            # lower or left leg
            idx += 1
            LinkList[idx] = First[site]
            LegType[idx] = spin_prop[site]
            current_link = idx

            if !isdiagonal(op)  # off-diagonal site operator
                spin_prop[site] ⊻= 1  # spinflip
            end

            if First[site] != 0
                LinkList[First[site]] = current_link  # completes backwards link
            else
                Last[site] = current_link
            end
            First[site] = current_link + 1
            Associates[idx] = (0, 0, 0)

            # upper or right leg
            idx += 1
            LegType[idx] = spin_prop[site]
            Associates[idx] = (0, 0, 0)
        elseif isbondoperator(op)  # diagonal bond operator
            site1, site2 = op

            # lower left
            idx += 1
            LinkList[idx] = First[site1]
            LegType[idx] = spin_prop[site1]
            current_link = idx

            if First[site1] != 0
                LinkList[First[site1]] = current_link  # completes backwards link
            else
                Last[site1] = current_link
            end

            First[site1] = current_link + 2
            vertex1 = current_link
            Associates[idx] = (vertex1 + 1, vertex1 + 2, vertex1 + 3)

            # lower right
            idx += 1
            LinkList[idx] = First[site2]
            LegType[idx] = spin_prop[site2]
            current_link = idx

            if First[site2] != 0
                LinkList[First[site2]] = current_link  # completes backwards link
            else
                Last[site2] = current_link
            end

            First[site2] = current_link + 2
            Associates[idx] = (vertex1, vertex1 + 2, vertex1 + 3)

            # upper left
            idx += 1
            LegType[idx] = spin_prop[site1]
            Associates[idx] = (vertex1, vertex1 + 1, vertex1 + 3)

            # upper right
            idx += 1
            LegType[idx] = spin_prop[site2]
            Associates[idx] = (vertex1, vertex1 + 1, vertex1 + 2)
        end
    end

    #Periodic boundary conditions for finite-beta
    @inbounds for i in 1:Ns
        if First[i] != 0  #This might be encountered at high temperatures
            LinkList[First[i]] = Last[i]
            LinkList[Last[i]] = First[i]
        end
    end

    # DEBUG
    # if spin_prop != spin_right
    #     @debug "Basis state propagation error: LINKED LIST"
    # end

    return len
end

#############################################################################

function cluster_update_beta!(lsize::Int, qmc_state::BinaryThermalState, H::TFIM)
    Ns = nspins(H)
    spin_left, spin_right = qmc_state.left_config, qmc_state.right_config
    operator_list = qmc_state.operator_list

    LinkList = qmc_state.linked_list
    LegType = qmc_state.leg_types
    Associates = qmc_state.associates

    in_cluster = zeros(Int, lsize)
    cstack = Stack{Int}()  # This is the stack of vertices in a cluster
    ccount = 0  # cluster number counter

    @inbounds for i in 1:lsize
        # Add a new leg onto the cluster
        if (in_cluster[i] == 0 && Associates[i] === (0, 0, 0))
            ccount += 1
            push!(cstack, i)
            in_cluster[i] = ccount

            flip = rand(Bool)  # flip a coin for the SW cluster flip
            if flip
                LegType[i] ⊻= 1  # spinflip
            end

            while !isempty(cstack)
                leg = LinkList[pop!(cstack)]

                if in_cluster[leg] == 0
                    in_cluster[leg] = ccount  # add the new leg and flip it
                    if flip
                        LegType[leg] ⊻= 1
                    end

                    # now check all associates and add to cluster
                    assoc = Associates[leg]  # a 3-tuple
                    if assoc !== (0, 0, 0)
                        for a in assoc
                            push!(cstack, a)
                            in_cluster[a] = ccount
                            if flip
                                LegType[a] ⊻= 1
                            end
                        end
                    end
                end

            end
        end
    end

    #println(in_cluster)

    # map back basis states and operator list
    First = qmc_state.first
    Last = qmc_state.last
    @inbounds for i in 1:Ns
        if First[i] != 0
            spin_left[i] = LegType[Last[i]]  # left basis state
            spin_right[i] = LegType[First[i]]  # right basis state
        else
            #randomly flip spins not connected to operators
            spin_left[i] = spin_right[i] = rand(Bool)
        end

    end

    ocount = 1  # first leg
    @inbounds for (n, op) in enumerate(operator_list)
        if isbondoperator(op)
            ocount += 4
        elseif !isidentity(op)
            if LegType[ocount] == LegType[ocount+1]  # diagonal
                operator_list[n] = (-1, op[2])
            else  # off-diagonal
                operator_list[n] = (-2, op[2])
            end
            ocount += 2
        end
    end

end


module TestKVStore
using MXNet
using Base.Test

using ..Main: rand_dims

SHAPE = rand_dims()
KEYS  = [5,7,11]

function init_kv()
  kv = mx.KVStore()
  mx.init!(kv, 3, mx.zeros(SHAPE))

  vals = [mx.zeros(SHAPE) for k in KEYS]
  mx.init!(kv, KEYS, vals)
  return kv
end

function test_kv_basic()
  info("KVStore::basic")

  kv = init_kv()
  @test mx.get_type(kv) == :local
  @test mx.get_rank(kv) == 0
  @test mx.get_num_workers(kv) == 1
end

function test_single_kv_pair()
  info("KVStore::single")

  kv = init_kv()
  mx.push!(kv, 3, mx.ones(SHAPE))
  val = mx.empty(SHAPE)
  mx.pull!(kv, 3, val)
  @test maximum(abs(copy(val) - 1)) == 0
end

function test_aggregator()
  info("KVStore::aggregator")

  kv = init_kv()

  num_devs = 4
  devs = [mx.Context(mx.CPU, i) for i=0:num_devs-1]
  vals = [mx.ones(SHAPE, dev) for dev in devs]

  mx.push!(kv, 3, vals)
  mx.pull!(kv, 3, vals)
  for v in vals
    @test maximum(abs(copy(v)) - num_devs) == 0
  end

  # list
  vals = [mx.NDArray[mx.ones(SHAPE, dev)*2 for dev in devs] for k in KEYS]
  mx.push!(kv, KEYS, vals)
  mx.pull!(kv, KEYS, vals)

  for vv in vals
    for v in vv
      @test maximum(abs(copy(v)) - 2num_devs) == 0
    end
  end
end

test_kv_basic()
test_single_kv_pair()
test_aggregator()

end


for (func, elty) in [(:CLBlastSsymm, Float32), (:CLBlastDsymm, Float64),
                     (:CLBlastCsymm, ComplexF32), (:CLBlastZsymm, ComplexF64)]
    #TODO: (:CLBlastHsymm, Float16)

    @eval function $func(layout::CLBlastLayout,
                         side::CLBlastSide, triangle::CLBlastTriangle,
                         m::Integer, n::Integer,
                         alpha::$elty,
                         a_buffer::cl.CL_mem, a_offset::Integer, a_ld::Integer,
                         b_buffer::cl.CL_mem, b_offset::Integer, b_ld::Integer,
                         beta::$elty,
                         c_buffer::cl.CL_mem, c_offset::Integer, c_ld::Integer,
                         queue::cl.CmdQueue, event::cl.Event)
        err = ccall(
            ($(string(func)), libCLBlast),
            cl.CL_int,
            (Cint, Cint, Cint, Csize_t, Csize_t, $elty, Ptr{Cvoid}, Csize_t, Csize_t,
              Ptr{Cvoid}, Csize_t, Csize_t, $elty, Ptr{Cvoid}, Csize_t, Csize_t, Ptr{Cvoid}, Ptr{Cvoid}),
            Cint(layout), Cint(side), Cint(triangle), m, n, alpha, a_buffer, a_offset, a_ld,
              b_buffer, b_offset, b_ld, beta, c_buffer, c_offset, c_ld, Ref(queue), Ref(event)
        )
        if err != cl.CL_SUCCESS
            println(stderr, "Calling function $(string($func)) failed!")
            throw(CLBlastError(err))
        end
        return err
    end

    @eval function symm!(_side::Char, uplo::Char, α::Number, A::cl.CLArray{$elty,2},
                         B::cl.CLArray{$elty,2}, β::Number, C::cl.CLArray{$elty,2};
                         queue::cl.CmdQueue=cl.queue(C))
        # check and convert arguments
        if _side == 'L'
            side = CLBlastSideLeft
        elseif _side == 'R'
            side = CLBlastSideRight
        else
            throw(ArgumentError("Side marker `side` is $(_side) but only 'L' and 'R' are allowed."))
        end
        if uplo == 'U'
            triangle = CLBlastTriangleUpper
        elseif uplo == 'L'
            triangle = CLBlastTriangleLower
        else
            throw(ArgumentError("Upper/lower marker `uplo` is $(uplo) but only 'U' and 'L' are allowed."))
        end
        m, n = size(C)
        j = size(A,1)
        if j != size(A,2)
            throw(DimensionMismatch("`A` has dimensions $(size(A)) but must be square."))
        end
        if j != (_side == 'L' ? m : n)
            throw(DimensionMismatch("`A` has size $(size(A)), `C` has size ($m,$n)."))
        end
        if size(B,2) != n
            throw(DimensionMismatch("`B` has second dimension $(size(B,2)) but needs to match second dimension of `C`, $n."))
        end
        alpha = convert($elty, α)
        beta  = convert($elty, β)
        layout = CLBlastLayoutColMajor

        # output event
        event::cl.Event = cl.Event(C_NULL)

        $func(layout, side, triangle,
              m, n,
              alpha,
              pointer(A), 0, size(A,1),
              pointer(B), 0, size(B,1),
              beta,
              pointer(C), 0, size(C,1),
              queue, event)

        # wait for kernel
        cl.wait(event)

        C
    end

end


import Base.+, Base.-, Base./, Base.*, Base.^
import SchumakerSpline.evaluate
function evaluate(f::Undefined_Function, point::Union{Real,Date,DateTime,DatePeriod})
    missing
end
function (s::Undefined_Function)(x::Union{Real,Date,DateTime,DatePeriod})
    return missing
end

"""
    derivative(f::UnivariateFunction)
This calculates the derivative of the function and returns it as a `UnivariateFunction`.

### Inputs
* `f` - A `UnivariateFunction`.
### Returns
* A `UnivariateFunction`.
"""
function derivative(f::Undefined_Function)
    return f
end

"""
    indefinite_integral(f::UnivariateFunction)
This calculates the indefinite integral of a `UnivariateFunction`.

### Inputs
* `f` - A `UnivariateFunction`.
### Returns
* A `UnivariateFunction`.
"""
function indefinite_integral(f::Undefined_Function)
    return f
end

function +(f::Undefined_Function, number::Real)
    return f
end
function -(f::Undefined_Function, number::Real)
    return f
end
function *(f::Undefined_Function, number::Real)
    return f
end
function /(f::Undefined_Function, number::Real)
    return f
end
function ^(f::Undefined_Function, number::Integer)
    return f
end

function +(f1::Undefined_Function, f2::Undefined_Function)
    return f1
end
function +(f1::Undefined_Function, f2::PE_Function)
    return f1
end
function +(f1::Undefined_Function, f2::Sum_Of_Functions)
    return f1
end
function +(f1::Undefined_Function, f2::Piecewise_Function)
    return f1
end

function +(f1::PE_Function, f2::Undefined_Function)
    return +(f2,f1)
end
function +(f1::Sum_Of_Functions, f2::Undefined_Function)
    return +(f2,f1)
end
function +(f1::Piecewise_Function, f2::Undefined_Function)
    return +(f2,f1)
end

function -(f1::Undefined_Function, f2::Undefined_Function)
    return f1
end
function -(f1::Undefined_Function, f2::PE_Function)
    return f1
end
function -(f1::Undefined_Function, f2::Sum_Of_Functions)
    return f1
end
function -(f1::Undefined_Function, f2::Piecewise_Function)
    return f1
end

function -(f1::PE_Function, f2::Undefined_Function)
    return -(f2,f1)
end
function -(f1::Sum_Of_Functions, f2::Undefined_Function)
    return -(f2,f1)
end
function -(f1::Piecewise_Function, f2::Undefined_Function)
    return -(f2,f1)
end

function *(f1::Undefined_Function, f2::Undefined_Function)
    return f1
end
function *(f1::Undefined_Function, f2::PE_Function)
    return f1
end
function *(f1::Undefined_Function, f2::Sum_Of_Functions)
    return f1
end
function *(f1::Undefined_Function, f2::Piecewise_Function)
    return f1
end

function *(f1::PE_Function, f2::Undefined_Function)
    return *(f2,f1)
end
function *(f1::Sum_Of_Functions, f2::Undefined_Function)
    return *(f2,f1)
end
function *(f1::Piecewise_Function, f2::Undefined_Function)
    return *(f2,f1)
end

function /(f1::Undefined_Function, f2::Undefined_Function)
    return f1
end
function /(f1::Undefined_Function, f2::PE_Function)
    return f1
end
function /(f1::Undefined_Function, f2::Sum_Of_Functions)
    return f1
end
function /(f1::Undefined_Function, f2::Piecewise_Function)
    return f1
end

function /(f1::PE_Function, f2::Undefined_Function)
    return /(f2,f1)
end
function /(f1::Sum_Of_Functions, f2::Undefined_Function)
    return /(f2,f1)
end
function /(f1::Piecewise_Function, f2::Undefined_Function)
    return /(f2,f1)
end