text:
@doc raw"""
    L0Smooth <: AbstractImageSmoothAlgorithm
    L0Smooth(; λ=2e-2, κ=2.0, βmax=1e5)

    smooth(img, f::L0Smooth)
    smooth!(out, img, f::L0Smooth)

Smoothing `img` via L0 gradient minimization to approximate prominent structure
in a sparsity-control manner.

# Output

Return `Array{Gray{N0f8}}` for `Gray` input or `Array{RGB{N0f8}}` for `RGB` input.

# Details

Using the strategy of minimizing the L0 norm of image's gradient.

This algorithm works particularly effective for sharpening major edges by
increasing the steepness of amplitude transition while eliminating
low-amplitude structures to some extent. See [1] for more details.

# Options

The function argument is described in detail below.

## `λ::Float64`

The argument `𝜆` is a weight directly controlling the significance of
the L0 gradient norm term, which must be greater than zero.

A larger `𝜆` makes the smoothed image have very few edges.

Default: 2e-2

## `βmax::Float64`

The argument `βmax` is the upper bound of `𝛽` in [1], which must be greater than 1e4.

In this algorithm, two auxiliary variables `ℎ` and `𝑣` are used to approximate
the solution. A large enough `𝛽` ensures that the alternating optimization strategy
based on introducing auxiliary variables is available.

Default: 1e5

## `κ::Float64`

The argument `𝜅` is the iteraiton rate, which must be larger than 1.0.

This algorithm using an alternating optimization strategy to get the solution.
In each iteration, the argument `𝛽` controls the similarity between gradient pair
(𝛥₁𝑆ₚ, 𝛥₂𝑆ₚ) (denoted by `` (\partial_x S_p, \partial_y S_p) `` in [1]) and auxiliary pair (ℎₚ, 𝑣ₚ).
The argument `𝜅` is used to update `𝛽` as `𝛽 ⟵ 𝜅𝛽`.

Default: 2.0

# Examples

You can use the default arguments for `L0Smooth`, and then use `smooth` to apply
the `AbstractImageSmoothAlgorithm`.

```julia
using TestImages
using ImageSmooth

img = testimage("cameraman")

fₛ = L0Smooth() # using default arguements
imgₛ = smooth(img, fₛ)
```

Manually setting the arguements is also available:

```julia
fₛ = L0Smooth(λ=0.0015, κ=1.05, βmax=2e5) # manually set the arguments
imgₛ = smooth(img, fₛ)
```

See also [`smooth!`](@ref) for in-place operation.

# References

[1] Xu, L., Lu, C., Xu, Y., & Jia, J. (2011, December). Image smoothing via L 0 gradient minimization. In Proceedings of the 2011 SIGGRAPH Asia conference (pp. 1-12). [DOI:10.1145/2024156.2024208](https://doi.org/10.1145/2024156.2024208)

"""
struct L0Smooth <: AbstractImageSmoothAlgorithm
    """smoothing weight"""
    λ::Float64
    """iteration rate"""
    κ::Float64
    """the upper bound of automatically adapting parameter β which controls the iteraiton times"""
    βmax::Float64

    function L0Smooth(λ, κ, βmax)
        λ > zero(λ) || throw(ArgumentError("smoothing weight λ should be positive"))
        κ > one(κ) || throw(ArgumentError("iteration rate κ should be > 1.0"))
        βmax > 1e4 || throw(ArgumentError("the iteration upper bound βmax should be > 1e4"))
        new(λ, κ, βmax)
    end
end

L0Smooth(; λ::Float64=2e-2, κ::Float64=2.0, βmax::Float64=1e5) = L0Smooth(λ, κ, βmax)

function (f::L0Smooth)(out::GenericGrayImage,
                       img::GenericGrayImage)
    𝑆 = float64.(channelview(img))
    𝜆 = f.λ # smoothing weight
    𝜅 = f.κ # iteration rate
    𝛽max = f.βmax # upper bound of 𝛽
    𝛽 = 2 * 𝜆 # define 𝛽₀
    ∂₁ = [1 -1]
    ∂₂ = [1, -1]
    N, M = size(𝑆)
    sizeI2D = (N, M)
    sizeI2D_t = (M, N)
    # Return the frequency-space representation of `ℱ∂₁` and `ℱ∂₂`.
    ℱ∂₁ = freqkernel(centered(∂₁), sizeI2D)
    ℱ∂₂ = transpose(freqkernel(centered(transpose(∂₂)), sizeI2D_t))
    ℱ𝐼 = fft(𝑆, (1, 2))
    Denormin = @. abs(ℱ∂₁)^2 + abs(ℱ∂₂)^2

    𝛥₁𝑆 = similar(𝑆)
    𝛥₂𝑆 = similar(𝑆)
    𝛥₁ᵀℎ = similar(𝑆)
    𝛥₂ᵀ𝑣 = similar(𝑆)

    Normin = similar(ℱ𝐼)
    t¹ = trues(N, M)
    ℱ𝑆 = similar(ℱ𝐼)

    while 𝛽 < 𝛽max
        # Computing (ℎ, 𝑣) via solving equation (9) in [1]
        # We get the solution (12) in [1] through following process
        # Use (𝛥₁𝑆, 𝛥₂𝑆) to demonstrate (ℎ, 𝑣) for convenience
        fdiff!(𝛥₁𝑆, 𝑆, dims = 2, boundary=:periodic)
        fdiff!(𝛥₂𝑆, 𝑆, dims = 1, boundary=:periodic)

        # For each pixel 𝑝 in 𝑆
        # (ℎₚ, 𝑣ₚ) = (0, 0), while (𝛥₁𝑆ₚ^2 + 𝛥₂𝑆ₚ^2) < λ / 𝛽
        # (ℎₚ, 𝑣ₚ) = (𝛥₁𝑆ₚ, 𝛥₂𝑆ₚ), otherwise
        @. t¹ = (𝛥₁𝑆^2 + 𝛥₂𝑆^2) < 𝜆 / 𝛽

        𝛥₁𝑆[t¹] .= 0
        𝛥₂𝑆[t¹] .= 0

        # For equation (8), ℎ = 𝛥₁𝑆, 𝑣 = 𝛥₂𝑆
        # According to Convolution Theorem, ℱ(𝑓₁ * 𝑓₂) = ℱ(𝑓₁) ⋅ ℱ(𝑓₂)
        # ℱ is the FFT operator, * is a convolution operator, ⋅ is a matrix dot times operator
        # We can compute ℱ(∂₁)* ⋅ ℱ(ℎ) and ℱ(∂₂)* ⋅ ℱ(𝑣) by computing ℱ(𝛥₁ᵀℎ) and ℱ(𝛥₂ᵀ𝑣)
        # ∂₁ and ∂₂ are the difference operators along horizontal axis and vertical axis, respectivly
        # ℱ()* denotes the complex conjugate
        # 𝛥₁ᵀ() and 𝛥₂ᵀ() indicate the transposition of forward difference along horizontal axis and vertical axis
        fdiff!(𝛥₁ᵀℎ, 𝛥₁𝑆, dims = 2, rev=true, boundary=:periodic)
        fdiff!(𝛥₂ᵀ𝑣, 𝛥₂𝑆, dims = 1, rev=true, boundary=:periodic)
        @. 𝛥₁ᵀℎ = -𝛥₁ᵀℎ
        @. 𝛥₂ᵀ𝑣 = -𝛥₂ᵀ𝑣

        # Computing S via equation (8) in [1]
        @. Normin = complex(𝛥₁ᵀℎ + 𝛥₂ᵀ𝑣)
        fft!(Normin, (1, 2))
        @. ℱ𝑆 = (ℱ𝐼 + 𝛽 * Normin) / (1 + 𝛽 * Denormin)
        ifft!(ℱ𝑆, (1, 2))
        @. 𝑆 = real(ℱ𝑆)

        𝛽 = 𝛽 * 𝜅
    end
    clamp01!(𝑆)
    out .= colorview(Gray, 𝑆)
    return out
end

function (f::L0Smooth)(out::GenericImage{<:AbstractRGB},
                       img::GenericImage{<:AbstractRGB})
    𝑆 = float64.(channelview(img))
    𝜆 = f.λ # smoothing weight
    𝜅 = f.κ # iteration rate
    𝛽max = f.βmax # upper bound of 𝛽
    𝛽 = 2 * 𝜆 # define 𝛽₀
    ∂₁ = [1 -1]
    ∂₂ = [1, -1]
    D, N, M = size(𝑆)
    sizeI2D = (N, M)
    sizeI2D_t = (M, N)
    # Return the frequency-space representation of `ℱ∂₁` and `ℱ∂₂`.
    ℱ∂₁ = freqkernel(centered(∂₁), sizeI2D)
    ℱ∂₂ = transpose(freqkernel(centered(transpose(∂₂)), sizeI2D_t))
    ℱ𝐼 = fft(𝑆, (2, 3))
    Denormin = @. abs(ℱ∂₁)^2 + abs(ℱ∂₂)^2
    Denormin = reshape(Denormin, 1, size(Denormin)...)
    Denormin = repeat(Denormin, inner=(1, 1, 1), outer=(D, 1, 1))

    𝛥₁𝑆 = similar(𝑆)
    𝛥₂𝑆 = similar(𝑆)
    s³ = similar(𝑆)
    s¹ = zeros(1, N, M)
    𝛥₁ᵀℎ = similar(𝑆)
    𝛥₂ᵀ𝑣 = similar(𝑆)

    Normin = similar(ℱ𝐼)
    t¹ = trues(1, N, M)
    t³ = trues(D, N, M)
    ℱ𝑆 = similar(ℱ𝐼)

    while 𝛽 < 𝛽max
        # Computing (ℎ, 𝑣) via solving equation (9) in [1]
        # We get the solution (12) in [1] through following process
        # Use (𝛥₁𝑆, 𝛥₂𝑆) to demonstrate (ℎ, 𝑣) for convenience
        fdiff!(𝛥₁𝑆, 𝑆, dims = 3, boundary=:periodic)
        fdiff!(𝛥₂𝑆, 𝑆, dims = 2, boundary=:periodic)

        # For each pixel 𝑝 in 𝑆, 𝛴ₙ denotes the sum of three different channels
        # (ℎₚ, 𝑣ₚ) = (0, 0), while 𝛴ₙ(𝛥₁𝑆ₚⁿ^2 + 𝛥₂𝑆ₚⁿ^2) < λ / 𝛽
        # (ℎₚ, 𝑣ₚ) = (𝛥₁𝑆ₚ, 𝛥₂𝑆ₚ), otherwise
        @. s³ = 𝛥₁𝑆^2 + 𝛥₂𝑆^2
        s¹ .= sum(s³, dims=1)
        @. t¹ = s¹ < 𝜆 / 𝛽
        t³ .= repeat(t¹, inner=(1, 1, 1), outer=(D, 1, 1))

        𝛥₁𝑆[t³] .= 0
        𝛥₂𝑆[t³] .= 0

        # For equation (8), ℎ = 𝛥₁𝑆, 𝑣 = 𝛥₂𝑆
        # According to Convolution Theorem, ℱ(𝑓₁ * 𝑓₂) = ℱ(𝑓₁) ⋅ ℱ(𝑓₂)
        # ℱ is the FFT operator, * is a convolution operator, ⋅ is a matrix dot times operator
        # We can compute ℱ(∂₁)* ⋅ ℱ(ℎ) and ℱ(∂₂)* ⋅ ℱ(𝑣) by computing ℱ(𝛥₁ᵀℎ) and ℱ(𝛥₂ᵀ𝑣)
        # ∂₁ and ∂₂ are the difference operators along horizontal axis and vertical axis, respectivly
        # ℱ()* denotes the complex conjugate
        # 𝛥₁ᵀ() and 𝛥₂ᵀ() indicate the transposition of forward difference along horizontal axis and vertical axis
        fdiff!(𝛥₁ᵀℎ, 𝛥₁𝑆, dims = 3, rev=true, boundary=:periodic)
        fdiff!(𝛥₂ᵀ𝑣, 𝛥₂𝑆, dims = 2, rev=true, boundary=:periodic)
        @. 𝛥₁ᵀℎ = -𝛥₁ᵀℎ
        @. 𝛥₂ᵀ𝑣 = -𝛥₂ᵀ𝑣

        # Computing S via equation (8) in [1]
        @. Normin = complex(𝛥₁ᵀℎ + 𝛥₂ᵀ𝑣)
        fft!(Normin, (2, 3))
        @. ℱ𝑆 = (ℱ𝐼 + 𝛽 * Normin) / (1 + 𝛽 * Denormin)
        ifft!(ℱ𝑆, (2, 3))
        @. 𝑆 = real(ℱ𝑆)

        𝛽 = 𝛽 * 𝜅
    end
    clamp01!(𝑆)
    out .= colorview(RGB, 𝑆)
    return out
end

function (f::L0Smooth)(out::OffsetArray, img::OffsetArray)
    axes(out) == axes(img) || throw(ArgumentError("out and img should have the same axes."))
    OffsetArray(f(out.parent, img.parent), out.offsets)
end


text:
@afgc function _stackDftReg(resource::ArrayFireLibs,imgRef::AFArray{Complex{T},N1},imgF::AFArray{Complex{T},N},usfac) where {T,N,N1}
    if N1 != (N - 1)
        error("Reference image has the wrong dimensionality")
    end
    nimg = size(imgF,N)
    if usfac==0
        ## Compute error for no pixel shift
        @afgc CCmax = reshape(sum(imgRef.*conj(imgF),1:N1),nimg)
        @afgc rfzero = sum(abs2(imgRef))
        @afgc rgzero = reshape(sum(abs2(imgF),1:N1),nimg)
        @afgc error = 1 - CCmax.*conj(CCmax)./(rgzero*rfzero)
        @afgc diffphase = atan2(imag(CCmax),real(CCmax),false)
        output = error
    elseif usfac==1
        ## Whole-pixel shift - Compute crosscorrelation by an IFFT and locate the peak
        L = length(imgRef)
        if N1 == 2
            @afgc CC = ifft2(imgRef.*conj(imgF))
        elseif N1 == 3
            @afgc CC = ifft3(imgRef.*conj(imgF))
        end

        @afgc CC =  abs(CC)
        CCmax = maximum(CC,1)

        for i in 2:N1
            CCmax = maximum(CCmax,i)
        end

        loc = find(CC .== CCmax) ## Problem if duplicates in a frame

        @afgc rfzero = sum(abs2(imgRef))/L
        @afgc rgzero = reshape(sum(abs2(imgF),1:N1),nimg)/L
        @afgc error = abs(1 - CCmax.*conj(CCmax)./(rgzero*rfzero))
        @afgc diffphase = Array(vec(atan2(imag(CCmax),real(CCmax),false)))

        indi = size(imgRef)
        ind2 = tuple([div(x,2) for x in indi]...)

        locTemp = ind2sub(size(CC),Int.(Array(loc)))

        locI = [tuple([locTemp[j][i] for j in 1:N1]...) for i in 1:nimg]

        shift = zeros(nimg,length(locI[1]))

        for im in eachindex(locI)
            for i in eachindex(locI[im])
                if locI[im][i]>ind2[i]
                    shift[im,i]=locI[im][i]-indi[i]-1
                else shift[im,i]=locI[im][i]-1
                end
            end
        end

        #output = Dict("error"=>error,"shift"=>shift,"diffphase"=>diffphase)
        output= (shift,diffphase,error)
    else

        @afgc CC = imgRef.*conj(imgF)

        @afgc CC = padAFArray(CC,N1)

        #CC[ranges...] = imgRef.*conj(imgF)
        #afgc()
        ##  Compute crosscorrelation and locate the peak
        if N1 == 2
            @afgc CC = ifft2(CC)
        elseif N1 == 3
            @afgc CC = ifft3(CC)
        end

        @afgc CC =  abs(CC)
        CCmax = maximum(CC,1)

        for i in 2:N1
           @afgc CCmax = maximum(CCmax,i)
        end

        loc = find(CC .== CCmax)
        indi = size(CC)[1:N1]
        ind2 = tuple([div(x,2) for x in indi]...)
        locTemp = ind2sub(size(CC),Int.(Array(loc)))

        locI = [tuple([locTemp[j][i] for j in 1:N1]...) for i in 1:nimg]

        shift = zeros(nimg,length(locI[1]))

        for im in eachindex(locI)
            for i in eachindex(locI[im])
                if locI[im][i]>ind2[i]
                    shift[im,i]=locI[im][i]-indi[i]-1
                else shift[im,i]=locI[im][i]-1
                end
            end
        end

        shift = shift/2

        ## If upsampling > 2, then refine estimate with matrix multiply DFT
        if usfac > 2
            ### DFT Computation ###
            # Initial shift estimate in upsampled grid
            shift = round.(Integer,shift*usfac)/usfac
            dftShift = div(ceil(usfac*1.5),2) ## center of output array at dftshift+1
            ## Matrix multiplies DFT around the current shift estimate
            @afgc CC = conj(dftups(imgF.*conj(imgRef),ceil(Integer,usfac*1.5),usfac,dftShift-shift*usfac,N1))/(prod(ind2)*usfac^N1)
            ## Locate maximum and map back to original pixel grid

            @afgc CC =  abs(CC)
            CCmax = maximum(CC,1)

            for i in 2:N1
                @afgc CCmax = maximum(CCmax,i)
            end

            loc = find(CC .== CCmax)
            indi = size(CC)[1:N1]
            locTemp = ind2sub(size(CC),Int.(Array(loc)))

            locI = [tuple([locTemp[j][i] for j in 1:N1]...) for i in 1:nimg]

            rg00 = Array(dftups(imgRef.*conj(imgRef),1,usfac,zeros(1,N1),N1))[1]/(prod(ind2)*usfac^N)
            rf00 = Array(dftups(imgF.*conj(imgF),1,usfac,zeros(nimg,N1),N1))[repeat([1],inner=[N1])...,:]./(prod(ind2)*usfac^N)

            locI = map((x) -> x .- dftShift .- 1,locI)
            for im in eachindex(locI)
                for i in eachindex(locI[im])
                    shift[im,i]=shift[im,i]+locI[im][i]/usfac
                end
            end

        else
            rg00 = Array(sum(imgRef.*conj(imgRef)))/(prod(indi))
            rf00 = Array(sum(imgF.*conj(imgF),1:N1))/(prod(indi))
        end
        error = Array(1 - vec(CCmax).*vec(conj(CCmax)))./vec(rg00*rf00)
        error = sqrt.(abs.(error))
        diffphase = Array(vec(atan2(imag(CCmax),real(CCmax),false)))
        ## If its only one row or column the shift along that dimension has no effect. Set to zero.
        shift[:,[div(x,2) for x in size(imgRef)].==1]=0
        output = (shift,diffphase,error)
    end
output
end

@afgc function padAFArray(CC,N)
    for r in N:-1:1
        sz = collect(size(CC))
        sz[r]=div(size(CC)[r],2)
        padR = zeros(AFArray{Complex{Float64}},sz...)
        @afgc CC = cat(r,padR,CC)
        @afgc CC = cat(r,CC,padR)
    end
    CC
end

function SubpixelRegistration.stackDftReg(resource::ArrayFireLibs,imgser::AbstractArray{T,N};ref::AbstractArray{T,N1}=reshape(selectdim(imgser,N,1),size(imgser)[1:(N-1)]),ufac::Int=10,chunkSize=100) where {T,N,N1}
    if ((N1 != (N - 1)) & (N1 != N))
        error("Reference image has the wrong dimensionality")
    end
    @afgc ref = fft(AFArray(ref))
    imgF = fft(imgser,(1:N1...,))

    if N1 == (N-1)
        chunks = [(i:min(chunkSize+i-1,size(imgF)[N])) for i in 1:chunkSize:size(imgF)[N]]
        shifts = zeros(size(imgF)[N],N1)
        diffphase = zeros(size(imgF)[N])
        error = zeros(size(imgF)[N])
        for ch in chunks
            shifts[ch,:],diffphase[ch],error[ch] = _stackDftReg(ArrayFireLibs(),ref,AFArray(selectdim(imgF,N,ch)),ufac)
        end
        results = (shifts,diffphase,error)
    else
        results = _stackDftReg(ArrayFireLibs(),ref,AFArray(imgF),ufac)
    end
    results = map(1:length(diffphase)) do i
        Dict("shift" => shifts[i,:],"diffphase" => diffphase[i],"error"=>error[i])
    end
    return results
end

"
`dftups{T,N}(inp::AbstractArray{T,N},no,usfac::Int=1,off=zeros(N))`

Upsampled DFT by matrix multiplication, computes an upsampled DFT in just a small region. `no` is the size of the region in pixels, `offset` the position in the full array, `usfac` the upsampling parameter."
@afgc function dftups(inp::AFArray,no,usfac::Int=1,offset=zeros(ndims(inp)),N=ndims(inp))
    sz = [size(inp)...]
    Nf = ndims(inp)
    nim = sz[Nf]
    permV = 1:N
    permFull = 0:3
    if Nf == N
        for i in permV
            pFull = collect(permFull)
            idxs = [pFull[i];deleteat!(pFull,i)]
            @afgc inp = reorder(inp,idxs...)
            leftK = AFArray((0:(no-1)).-offset[:,i]')
            @afgc leftK = reshape(leftK,no)
            rightK = AFArray(ifftshift(0:(sz[i]-1))-floor(sz[i]/2))
            @afgc rightK = rightK'
            kern = leftK*rightK
            @afgc kern = exp.(-1im*2*pi/(sz[i]*usfac)*kern)
            nsz=size(inp)
            d = nsz[2:Nf]
            #kern = reorder(AFArray(repeat(kern,outer=(1,1,nsz[1]))),0,2,1,4)
            @afgc inp = kern * reshape(inp, Val(2))
            @afgc inp = reshape(inp,no,d...)
        end
    else
        for i in permV
            pFull = collect(permFull)
            idxs = [pFull[i];deleteat!(pFull,i)]
            @afgc inp = reorder(inp,idxs...)
            leftK = AFArray((0:(no-1)).-offset[:,i]')
            @afgc leftK = reshape(leftK,no,1,nim)
            rightK = AFArray(repeat(ifftshift(0:(sz[i]-1)),outer=(1,sz[Nf]))-floor(sz[i]/2))
            @afgc rightK = reshape(rightK,1,sz[i],sz[Nf])

            kern = leftK*rightK
            @afgc kern = exp.(-1im*2*pi/(sz[i]*usfac)*kern)

            nsz=size(inp)
            d = nsz[2:Nf]
            remDim = prod(nsz[2:N])

            @afgc inp = kern * reshape(inp, nsz[1],remDim,sz[Nf])
            @afgc inp = reshape(inp,no,d...)
        end
    end
    newIdx = collect(N:-1:1).-1
    extraIdx = setdiff(0:3,newIdx)
    reorder(inp,newIdx...,extraIdx...)
end

"
`subPixShift(imgft::AbstractArray{Complex{Float64}},shift::Array{Float64,1})`

Shift the image `imgft` (in Fourier space) by the amount provided in the vector `shift`."
@afgc function subPixShift(imgft::AFArray{Complex{T},N},shift::Array{Float64,2},diffphase::Array{Float64,1}) where {T,N}
    sz = [size(imgft)[1:(N-1)]...]
    nim = size(imgft)[N]
    Z=0
    resh = [repeat([1],inner=N-1)...,nim]
    for i in eachindex(sz)
        shifti = AFArray(ifftshift((-div(sz[i],2)):(ceil(Integer,sz[i]/2)-1)).*shift[:,i]'/sz[i])
        reshTemp = copy(resh)
        reshTemp[i] = size(shifti)[1]
        Z = Z .- reshape(shifti,reshTemp...)
    end

    Greg = imgft .* exp.(2im*pi*Z)
    @afgc Greg = Greg .* reshape(exp.(AFArray(1im*diffphase)),resh...)
    Greg
end

function subPixShift(imgft::AFArray{Complex{T}},shift::Array{Float64,1},diffphase) where {T}
    sz = [size(imgft)...]
    N=0
    for i in eachindex(sz)
        shifti = ifftshift((-div(sz[i],2)):(ceil(Integer,sz[i]/2)-1))*shift[i]/sz[i]
        resh = (repeat([1],inner=[i-1])...,length(shifti))
        N = N .- reshape(shifti,resh...)
    end

    Greg = imgft .* exp.(2im*pi*N)
    Greg = Greg .* exp(1im*diffphase)
    Greg
end

function SubpixelRegistration.alignFromDict(resource::ArrayFireLibs,img2reg::AbstractArray{T,N},dftRegRes::Array{Dict{String,Any},1}) where {T,N}
    if (length(dftRegRes) != size(img2reg)[N])
        error("Alignment results and image stack dimensionalities don't match.")
    end

    img2regF = fft(img2reg,(1:(N-1)...,))
    szF = size(img2reg)[1:(N-1)]

    imRes = subPixShift(AFArray(img2regF),SubpixelRegistration.shifts_from_res(dftRegRes),SubpixelRegistration.diffphase_from_res(dftRegRes))

    imRes = real(ifft(Array(imRes),1:(N-1)))

    imRes
end

ext: jl


text:
using ApproxFun, Base.Test

#The first test checks the solution of the integral equation
# u(x) + \int_{-1}^{+1} \frac{e^{y} u(y)}{\sqrt{1-y^2}} dy = f
# on the interval [-1,1].

x=Fun(identity)
w=1/sqrt(1-x^2)
d=domain(x)

Σ=DefiniteIntegral(d)
bandinds(Σ)

@test domainspace(Σ) ==
    JacobiWeight{Chebyshev{Interval{Float64}},Interval{Float64}}(-0.5,-0.5,Chebyshev())

L=I+Σ[exp(x)*w]
bandinds(L)
usol=sin(2x)
f=L*usol
u=L\f
@test norm(u-usol) <= 10eps()

#The second test checks the solution of the integro-differential equation
# u'(x) + x u(x) + \int_{-2}^{+2} sin(y-x) u(y) \sqrt{4-y^2} dy = f
# on the interval [-2,2], with u(-2) = 1.

x=Fun(identity,[-2.,2.])
w=sqrt(4-x^2)
d=domain(x)

D=Derivative(d)
B=ldirichlet(d)
Σ=DefiniteIntegral(.5,.5,d)

@test domainspace(Σ) ==
    JacobiWeight{Ultraspherical{Int,Interval{Float64}},Interval{Float64}}(.5,.5,Ultraspherical(1,d))

K=LowRankFun((x,y)->sin(y-x)*w(y),Ultraspherical(1,d),domainspace(Σ))

L=D+x+Σ[K]
usol=cospi(20x)
f=L*usol
u=[B;L]\[1.;f]

@test norm(u-usol) ≤ 100eps()

Σ = DefiniteIntegral()

f1=Fun(t->cos(cos(t)),[-π,π])
f=Fun(t->cos(cos(t)),Laurent([-π,π]))

@test_approx_eq sum(f1) Σ*f

f1=Fun(t->cos(cos(t))/t,Laurent(Circle()))
f2=Fun(t->cos(cos(t))/t,Fourier(Circle()))
@test_approx_eq Σ*f1 Σ*f2

f1=Fun(t->cos(cos(t)),Laurent([-π,π]))
f2=Fun(t->cos(cos(t)),Fourier([-π,π]))
@test_approx_eq Σ*f1 Σ*f2

## test over arcs

d=exp(im*Interval(0.1,0.2))
x=Fun(d)
w=1/(sqrt(abs(first(d)-x))*sqrt(abs(last(d)-x)))

@test_approx_eq linesum(w) DefiniteLineIntegral()*w

## Volterra integral equation

d = Interval(0.0,1.0)
V = Volterra(d)
K = LowRankFun((x,y)->sin(y-x),d^2)
L = I-V[K]
f = Fun(exp,d)
@test domainspace(L) == Legendre(d)
@test rangespace(L) == Legendre(d)
@test bandrange(V) == -1:0
u = L\f
@test norm(L*u-f) ≤ 20eps()


text:
module Commands
export Command, MotionCommand, OperatorCommand, LineOperatorCommand, InsertCommand
export key

abstract type Command end

struct MotionCommand <: Command
    r1 :: Integer
    motion :: Char
end

"""
Command which operates on a selection of text, for example yank `y`, cut `c`, or delete `d`
"""
struct OperatorCommand <: Command
    r1 :: Integer
    # one of 'y', 'd', 'c'
    operator :: Char
    r2 :: Integer
    # char = motion, string = textobject
    action :: Union{Char, String}
end

"""
    Command which operates on lines by repeating a motion
e.g. "5dd"
"""
struct LineOperatorCommand <: Command
    r1 :: Integer
    operator :: Char
end

"""
Command which changes into Insert mode, possibly preceded by a motion, for example `A`.
"""
struct InsertCommand <: Command
    c :: Char
end

function key(cmd :: Command) :: Char
    error("method `key` not implemented for type $(typeof(cmd))")
end

key(cmd :: MotionCommand) = cmd.motion
key(cmd :: OperatorCommand) = cmd.operator
key(cmd :: LineOperatorCommand) = cmd.operator
key(cmd :: InsertCommand) = cmd.c

end


text:
#This file contains a wrapper object and some helper functions
#relating to endowment files

mutable struct NCCSFile
    df::DataFrame
    meta::MetaDictionary
end

const NNCCSFile = Union{NCCSFile, Nothing}

#constructs the objects from the metaDictionary
function NCCSFile(meta::MetaDictionary; nccsPath::String = NCCS_PATH,
  iStreamer::Function = I_CSV_STREAMER)

  inPath::String = "$nccsPath\\$(meta.fileName).csv.gz"

  df::DataFrame = readSingleTypeCSV(inPath)

  #check and make sure we have the right file
  yearMatch::RegexMatch = match(r"\d{4}", meta.fileName)
  nonprofitType::Symbol = Symbol(meta.fileName[(yearMatch.offset+5):(yearMatch.offset+6)])
  fileYear::Int = parse(Int, yearMatch.match)
  fullFile::Bool = occursin("full", meta.fileName)

  if nonprofitType ≠ meta.nonprofitType || fileYear ≠ meta.fileYear || fullFile ≠ meta.fullFile
    @warn ("File type information does not match for $(meta.fileName).\n
      nonprofitType: $(nonprofitType) | meta.nonprofitType: $(meta.nonprofitType)\n
      fileYear: $(fileYear)  | meta.fileYear: $(fileYear)\n
      fullFile: $(fullFile) | meta.fullFile: $(fullFile)")
  end

  data::NCCSFile = NCCSFile(df, meta)
  preProcessNCCSFile!(data)

  return data
end

function changeNCCSColumnType!(data::NCCSFile, col::Symbol, T::Type)::NCCSFile

  if T <: MString #case 1: String or StringUMissing
    return data #do nothing
  end

  rename!(data.df, [col=>:oldCol])
  oldCol::Vector{R where R <: MString} = data.df[:oldCol]
  NValues::Int = length(oldCol)

  if (T <: CategoricalValue{Union{Symbol, Missing}}) #case 3: pool the categorical data
    data.df[col] = Vector{Union{Symbol, Missing}}(undef, NValues)
    data.df[col] .= ((s::MString->ismissing(s) ? missing : Symbol(s)).(oldCol))
    categorical!(data.df, col)
    #data.df[col] = compress(data.df[col]) NOTE: Consider using this in the future, currently crashes

  elseif T <: Union{MFloat64, MInt, Int, Float64, Bool, MBool} #case 2: concrete numerical type
    #need to parse each entry
    concreteCol::Vector{T} = Vector{T}(undef, NValues)
    if T == Missing #special case where the column is entirely missing
      concreteCol = missings(NValues)
    else
      TSub::Type = typeintersect(T, Union{Float64, Bool, Int}) #get the nonmissable type

      #=@inbounds @simd for i ∈ 1:NValues #parse the column text
        concreteCol[i] = ifelse(ismissing(oldCol[i]),missing, parse("$(oldCol[i])"))
      end=#
      @inbounds @simd for i ∈ 1:NValues
        concreteCol[i] = ismissing(oldCol[i]) ? missing : something(tryparse(TSub, oldCol[i]), missing)
      end

    end

    data.df[col] = concreteCol

  else #Write explicit code for each type of column
    println("Col $col from file $(data.meta.fileName) is of type $T which is not recognized.")
    error("Col $col from file $(data.meta.fileName) is of type $T which is not recognized.")

  end

  #delete the old column
  deletecols!(data.df, :oldCol)

  return data
end

#fixes the type information
function preProcessNCCSFile!(data::NCCSFile)::NCCSFile

  meta::MetaDictionary = data.meta
  expectedCols::Vector{Symbol} = collect(keys(meta.dictionary))
  actualCols::Vector{Symbol} = (s::Symbol->Symbol(lowercase(string(s)))).(names(data.df))
  names!(data.df, actualCols) #make the names lower case

  #form the lookup dictionary for the type. Takes a column name and gets the type, with a default of String
  typeDict::Dict = Dict(
    actualCols[i] => (
    haskey(meta.dictionary, actualCols[i]) ? meta.dictionary[actualCols[i]] : MString)
    for i::Int ∈ 1:(length(actualCols)))

  #remove columns we don't expect or want
  diff::Vector{Symbol} = setdiff(actualCols, expectedCols)
  if length(diff) > 0
    try
      deletecols!(data.df, setdiff(actualCols, expectedCols))
    catch err
      error("$err\nmeta: $(meta.fileName)")
    end
  end

  #set the column types. Would have preferred to do this on the read, but got a weird error.
  for col::Symbol ∈ actualCols
    changeNCCSColumnType!(data, col, typeDict[col])
  end

  data = preFilterNCCSFile!(data)

  return data
end

#holds a limited amount of pre-filtering logic to improve later stage data processing performance
#NOTE: drops full files and co files before 1997
function preFilterNCCSFile!(data::NCCSFile; revenueLowerBound::Float64 = REVENUE_LOWER_BOUND)::NCCSFile
  nonprofitType::Symbol = data.meta.nonprofitType

  if data.meta.fullFile
    filter!(r::DataFrameRow->false, data.df) #ignore these for now
  elseif nonprofitType == :pc
    #data.df = data.df[((ismissing).(data.df[:fundbal]) .== false) .& (data.df[:fundbal] .> assetLowerBound),:]
    filter!(r::DataFrameRow->ismissing(r[:totrev]) == false && r[:totrev] > revenueLowerBound, data.df)
  elseif nonprofitType == :co #the below filter doesn't really work due to dq issues
    if data.meta.fileYear ≥ 2000
      filter!(r::DataFrameRow->ismissing(r[:totrev2]) == false && r[:totrev2] > revenueLowerBound, data.df)
    else
      filter!(r::DataFrameRow->ismissing(r[:grossrec]) == false && r[:grossrec] > revenueLowerBound, data.df)
    end
  elseif nonprofitType == :pf
    filter!(r::DataFrameRow->ismissing(r[:p1totrev]) == false && r[:p1totrev] > revenueLowerBound, data.df)
  end

  return data
end

#serializes a NCCSFile object
function serializeNCCSFile(data::NCCSFile; nccsPath::String = NCCS_PATH,
    oStreamer::Function = O_JLS_STREAMER,verbose::Bool=true)::NNCCSFile

    outPath::String = "$nccsPath\\$(data.meta.fileName).jls.gz"

    oStream = oStreamer(outPath)
    serialize(oStream, data)
    close(oStream)

    return nothing
end

#wrapper object to serialize a data file object
#For convenience, can return the object if requested
function serializeNCCSFile(meta::MetaDictionary;
  verbose::Bool = true, returnNCCSFile::Bool=true)::NNCCSFile

  #these are the column names that we expect

  if verbose
    println("Serializing $(meta.fileName)...")
  end

  data::NCCSFile = NCCSFile(meta) #make the object
  serializeNCCSFile(data) #serialize the object

  if verbose
    println("Serialization of $(meta.fileName) complete.")
  end

  return returnNCCSFile ? data : nothing
end

#deserializes a NCCSFile object
function deserializeNCCSFile(fileName::String; nccsPath::String = NCCS_PATH,
    iStreamer::Function = I_JLS_STREAMER)::NCCSFile

    inPath::String = "$nccsPath\\$(fileName).jls.gz"

    iStream = iStreamer(inPath)
    data::NCCSFile = deserialize(iStream)
    close(iStream)

    return data
end

deserializeNCCSFile(meta::MetaDictionary)::NCCSFile = deserializeNCCSFile(meta.fileName)

#loads an already downloaded NCCS file
function loadNCCSFile(meta::MetaDictionary; nccsPath::String = NCCS_PATH,
  reserializeData::Bool = true, returnNCCSFile::Bool = true)::NNCCSFile

  filePath::String = "$nccsPath\\$(meta.fileName).jls.gz"

  #only re-process the file if necessesary
  if reserializeData || !isfile(filePath)
    data::NNCCSFile = serializeNCCSFile(meta, returnNCCSFile=returnNCCSFile)
  elseif returnNCCSFile
    data = deserializeNCCSFile(meta)
  else #do nothing
    data = nothing #an empty dataframe
  end

  return data
end

#loads and/or serializes a vector of corresponding files
function loadNCCSFiles(metadata::Vector{MetaDictionary}; mode::Symbol = :serial,
  reserializeData::Bool = true,  serializeOnly::Bool = true)::Union{Nothing,Vector{NCCSFile}}

  if mode == :parallel #relaunch the dedicated parallel function
    return loadNCCSFilesPar(metadata::Vector{MetaDictionary}, reserializeData=reserializeData,
      serializeOnly=serializeOnly)
  end

  if serializeOnly
    for meta::MetaDictionary ∈ metadata #do this in a loop format for space conservation
      loadNCCSFile(meta,reserializeData=reserializeData, returnNCCSFile=false)
    end
    return nothing
  else
    return (meta::MetaDictionary->loadNCCSFile(meta,returnNCCSFile=true)).(metadata)
  end
end

#loads many NCCS serial files in parallel
function loadNCCSFilesPar(metadata::Vector{MetaDictionary}; reserializeData::Bool = true,
  serializeOnly::Bool=true, maxSerializationWorkers::Int = MAX_SERIALIZATION_WORKERS)::Union{Nothing, Vector{DataFrame}}

  numFiles::Int = length(metadata)
  pids = workers()
  np::Int = min(length(pids), maxSerializationWorkers)

  #divide up the work
  assignments::Vector{Vector{Int}} = ((i::Int)->Vector{Int}()).(1:np)
  for i::Int ∈ 1:numFiles
    push!(assignments[i%np+1], i) #assigns a file to a pid
  end

  #set up a vector for pushing and receiving assignments
  futures::Vector{Future} = Vector{Future}()
  sizehint!(futures, np)

  for i::Int ∈ 1:np
    if length(assignments[i]) > 0
      push!(futures,
        (@spawnat pids[i] loadNCCSFiles(metadata[assignments[i]], reserializeData=reserializeData,
          serializeOnly=serializeOnly)))
    end
  end

  #get the work and return it if desired
  if serializeOnly
    (fetch).(futures)
    return nothing
  else
    return vcat((fetch).(futures)...)
  end
end

#define this type. We may want to add more later.
mutable struct NCCSData
  df::DataFrame
end

#deals with removing fields
function removeFields!(data::NCCSData, fields::Vector{Symbol})::NCCSData
  deletecols!(data.df, fields)
  #NOTE: If we add processed metadata, deal with that here
  return data
end

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
export H, Z, X, Y, T, S
export SOp
export COp
export CNOT
export SWAP

"""
Constructs gate primitives for use in tensor contractions via Op()
All matrices are of complex float types.

# Gates

- `H` : Hadamard gate
- `Z` : Pauli Z gate
- `X` : Pauli X gate
- `Y` : Pauli Y gate
- `T` : Rotate π/8 gate
- `S` : Rotate π/4 gate

"""
H = (1.0/sqrt(2.0)) * complex([1   1; 1  -1])
Z = complex([1.0   0; 0    -1.0])
X = complex([0   1.0; 1.0     0])
Y = complex([0.0 -im; im    0.0])
T = complex([1.0   0; 0  ℯ^(im*(π/4))])
S = complex([1.0   0; 0      im])

"""
Performs an gate operation on the register by constructing
the tensor product of the operator. Single index only operates on
one qubit, the double index operates on multiple *consecutive* qubits.

In general the operation is
(I ⊗ I ... I ⊗ G_q ⊗ G_q+1 ⊗ ... G_p ⊗ I ... ⊗ I) | Register >
"""
function SOp(G::Array{ComplexF64}, q::Int, reg::Register)

    # Create pre-gate Identity
    op = zeros(ComplexF64, 2^(q-1), 2^(q-1)) + I

    # Operate gate G on qubit q
    op = op ⊗ G

    # Identity operations after qubit q
    ident = zeros(ComplexF64, 2, 2) + I
    for i in q+1:reg.N
        op = op ⊗ ident
    end

    reg.state = op * reg.state
end

function SOp(G::Array{ComplexF64}, q::Int, p::Int, reg::Register)
    @assert q < p
    @assert q != p
    op = zeros(ComplexF64, 2^(q-1), 2^(q-1)) + I

    # Operate 2x2 gate G on qubits q through p
    for i in q:p
        op = op ⊗ G
    end

    # Identity operations after qubit p
    ident = zeros(ComplexF64, 2, 2) + I
    for i in p+1:reg.N
        op = op ⊗ ident
    end

    # Update register
    reg.state = op * reg.state
end

"""
Controlled gate operation
"""
# FIXME: Add case for control below gate
function COp(G::Array{ComplexF64}, c::Int, q::Int, reg::Register)
    @assert c != q

    # Create pre-gate Identity
    if c > 1
        op = zeros(ComplexF64, 2^(c-1), 2^(c-1)) + I
    end

    # Dimension of full controlled gate
    dim = 2^((q-c)+1)
    hf = div(dim,2)

    # Fill in 2x2 gates properly
    ident = zeros(ComplexF64, 2, 2) + I
    con = zeros(ComplexF64, dim, dim)
    for i in 1:2:hf-1
        con[i : 1+i, i : 1+i] = ident
        con[i+hf : 1+i+hf, i+hf : 1+i+hf] = G
    end

    # Complete full gate up to qubit q
    if c > 1
        op = op ⊗ con
    else
        op = con
    end

    # Identity operations after qubit q
    for i in q+1:reg.N
        op = op ⊗ ident
    end

    # Update qubit register
    reg.state = op * reg.state
end

function CNOT(c::Int, q::Int, reg::Register)
    if c > q
        SOp(H, q, c, reg)
        COp(X, q, c, reg)
        SOp(H, q, c, reg)
    else
        COp(X, c, q, reg)
    end
end

function SWAP(q::Int, p::Int, reg::Register)
    CNOT(q, p, reg)
    CNOT(p, q, reg)
    CNOT(q, p, reg)
end


text:
using StatsBase

struct FiveNumberSummary
    minimum::Float64
    lower_quartile::Float64
    median::Float64
    upper_quartile::Float64
    maximum::Float64
end

mutable struct BoxplotGraphics{R<:Number} <: GraphicsArea
    data::Vector{FiveNumberSummary}
    color::Symbol
    width::Int
    left::R
    right::R

    function BoxplotGraphics{R}(
            data::AbstractVector{R},
            width::Int,
            color::Symbol,
            left::R,
            right::R) where R
        width = max(width, 10)
        new{R}(
            [FiveNumberSummary(
                minimum(data),
                percentile(data, 25),
                percentile(data, 50),
                percentile(data, 75),
                maximum(data)
            )], color, width, left, right)
    end
end

nrows(c::BoxplotGraphics) = 3*length(c.data)
ncols(c::BoxplotGraphics) = c.width

function BoxplotGraphics(
        data::AbstractVector{R},
        width::Int;
        color::Symbol = :blue,
        left::R,
        right::R,
        labels::AbstractVector{<:AbstractString} = []) where {R <: Number}
    BoxplotGraphics{R}(data, width, color, left, right)
end

function addseries!(c::BoxplotGraphics, data::AbstractVector{R}) where {R <: Number}
    append!(c.data, [FiveNumberSummary(
        minimum(data),
        percentile(data, 25),
        percentile(data, 50),
        percentile(data, 75),
        maximum(data)
    )])
end

function printrow(io::IO, c::BoxplotGraphics, row::Int)
    0 < row <= nrows(c) || throw(ArgumentError("Argument row out of bounds: $row"))
    series = c.data[Int(ceil(row/3))]

    transform = value -> Int(round((value - c.left)/(c.right - c.left) * c.width))

    seriesRow = Int((row-1) % 3) + 1

    minChar = ['│', '├' , '│'][seriesRow]
    lineChar = [' ', '─' , ' '][seriesRow]
    leftBoxChar = ['┌', '┤' , '└'][seriesRow]
    lineBoxChar = ['─', ' ' , '─'][seriesRow]
    medianChar = ['┬', '│' , '┴'][seriesRow]
    rightBoxChar = ['┐', '├' , '┘'][seriesRow]
    maxChar = ['│', '┤' , '│'][seriesRow]

    line = [' ' for _ in 1:c.width]

    # Draw points first - this is most important, so they'll always be drawn
    # even if there's not enough space
    line[transform(series.minimum)] = minChar
    line[transform(series.lower_quartile)] = leftBoxChar
    line[transform(series.median)] = medianChar
    line[transform(series.upper_quartile)] = rightBoxChar
    line[transform(series.maximum)] = maxChar

    # Fill in gaps with lines
    for i in transform(series.minimum)+1:transform(series.lower_quartile)-1
        line[i] = lineChar
    end
    for i in transform(series.lower_quartile)+1:transform(series.median)-1
        line[i] = lineBoxChar
    end
    for i in transform(series.median)+1:transform(series.upper_quartile)-1
        line[i] = lineBoxChar
    end
    for i in transform(series.upper_quartile)+1:transform(series.maximum)-1
        line[i] = lineChar
    end

    printstyled(io, join(line), color = c.color)
end


text:
abstract type Barrier end

function drop(this::Barrier)
    this.mEnabled = false
end

struct BufferBarrier <: Barrier
    mEnabled::Bool
    mTarget::vk.VkCommandBuffer
    mSrcStage::vk.VkPipelineStageFlags
    mDstStage::vk.VkPipelineStageFlags
    mHandle::vk.VkBufferMemoryBarrier

    function BufferBarrier(cmd::RecordingCommandBuffer, buffer::vk.VkBuffer;
                      srcStage::vk.VkPipelineStageFlags = 0, dstStage::vk.VkPipelineStageFlags = 0,
                      srcAccessMask::vk.VkAccessFlags = 0, dstAccessMask::vk.VkAccessFlags = 0)
        this = new(true, cmd.getBuffer().handle(),
                    srcStage,
                    dstStage,
                    vk.VkBufferMemoryBarrier(
                        vk.VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER, # sType::VkStructureType
                        C_NULL, # pNext::Ptr{Cvoid}
                        srcAccessMask, # srcAccessMask::VkAccessFlags
                        dstAccessMask, # dstAccessMask::VkAccessFlags
                        vk.VK_QUEUE_FAMILY_IGNORED, # srcQueueFamilyIndex::UInt32
                        vk.VK_QUEUE_FAMILY_IGNORED, # dstQueueFamilyIndex::UInt32
                        buffer, # buffer::VkBuffer
                        0, # offset::VkDeviceSize
                        vk.VK_WHOLE_SIZE, # size::VkDeviceSize
                    ))
        return this
    end
end
@class BufferBarrier

function apply(this::BufferBarrier)
    if this.mEnabled
        vk.vkCmdPipelineBarrier(this.mTarget, this.mSrcStage, this.mDstStage,
                                0, # dependencyFlags
                                0, # memoryBarrierCount
                                C_NULL, # pMemoryBarriers
                                1, # bufferMemoryBarrierCount
                                pointer([this.mHandle]), # pBufferMemoryBarriers
                                0, # imageMemoryBarrierCount
                                C_NULL # pImageMemoryBarriers
        )
    end
end

function handle(this::BufferBarrier)
    return this.mHandle
end


text:
using DynamicalSystemsBase
using Test, StaticArrays
using DynamicalSystemsBase: create_jacobian, create_tangent, stateeltype, isinplace

println("\nTesting inference...")

@testset "Inference" begin
    dss = [Systems.towel(), Systems.henon_iip()]
    @testset "IIP = $IIP" for IIP in [false, true]
        ds = IIP ? dss[2] : dss[1]
        @test_nowarn @inferred integrator(ds)

        f = ds.f
        s = get_state(ds)
        p = ds.p
        t = 0
        D = dimension(ds)

        @testset "Create Jac" begin
            # @test_nowarn @inferred create_jacobian(f, Val{IIP}(), s, p, t, Val{D}())
            @test_nowarn @inferred create_tangent(f, ds.jacobian, ds.J, Val{IIP}(), Val{2}())
            @test_nowarn @inferred jacobian(ds)
        end

        # Integrator state inference:
        @testset "Integrator" begin
            @test_nowarn @inferred stateeltype(ds)
            integ = integrator(ds)
            @test_nowarn @inferred stateeltype(integ)
            @test stateeltype(integ) == Float64
        end

        @testset "Tangent" begin
            integ = tangent_integrator(ds, 2)
            @test_nowarn @inferred stateeltype(integ)
            @test stateeltype(integ) == Float64
        end
        @testset "Parallel" begin
            integ = parallel_integrator(ds, [get_state(ds)])
            @test_nowarn @inferred stateeltype(integ)
            @test stateeltype(integ) == Float64
        end
    end
end


text:
using CloudGraphs
using Test

@test isdefined(Main, :CloudGraphs)
@test typeof(CloudGraphs) == Module

#include("QuickPackProtoTest.jl")

include("CloudGraphs.jl")

# Big data tests
# Maybe done: Get BigData tests to run again.
# if !haskey(ENV, "TRAVIS_OS_NAME")
include("BigData.jl")
# else
#   print("[TEST] NOTE: Testing in Travis, skipping the Mongo bigData test for the moment...")
# end

# Return the true exit status from FactCheck
# FactCheck.exitstatus()


text:
"""
# ElectricalEngineering.Tab20bc

Defines colors based on
https://matplotlib.org/users/plotting/colormaps/grayscale_01_04.pdf. As an
alternative `ElectricalEngineering.Gray` can be used instead of `ElectricalEngineering.Tab20bc` in order to create a
gray scale plot instead of a color plot.

# Colors

The following colors are defined:

- `colorBrown1`, `colorBrown2`, `colorBrown3`, `colorBrown4`
- `colorBlue1`, `colorBlue2`, `colorBlue3`, `colorBlue4`
- `colorRed1`, `colorRed2`, `colorRed3`, `colorRed4`
- `colorOrange1`, `colorOrange2`, `colorOrange3`, `colorOrange4`
- `colorGreen1`, `colorGreen2`, `colorGreen3`, `colorGreen4`
- `colorPurple1`, `colorPurple2`, `colorPurple3`, `colorPurple4`
- `colorDarkPurple1`, `colorDarkPurple2`, `colorDarkPurple3`, `colorDarkPurple4`
- `colorMagenta1`, `colorMagenta2`, `colorMagenta3`, `colorMagenta4`
- `colorGray1`, `colorGray2`, `colorGray3`, `colorGray4`

# Example

Copy and paste code:

```julia
using PyPlot,ElectricalEngineering,ElectricalEngineering.Tab20bc
x=collect(0:0.01:1);
figure(figsize=(3.3,2.5))
plot(x,x,color=colorRed1,linestyle=lineStyle1,linewidth=lineWidth1,
label="(1)")
plot(x,x.^2,color=colorRed2,linestyle=lineStyle2,linewidth=lineWidth2,
label="(2)")
plot(x,x.^3,color=colorRed3,linestyle=lineStyle3,linewidth=lineWidth3,
label="(3)")
plot(x,x.^4,color=colorRed4,linestyle=lineStyle4,linewidth=lineWidth4,
label="(4)")
subplots_adjust(left=0.20, right=0.95, top=0.95, bottom=0.20)
legend(loc="upper left",fontsize=legendFontSize)
```

# See also

`ElectricalEngineering.Gray`

"""
module Tab20bc
  using ElectricalEngineering

    export colorBrown1,colorBrown2,colorBrown3,colorBrown4,
        colorBlue1,colorBlue2,colorBlue3,colorBlue4,
        colorRed1,colorRed2,colorRed3,colorRed4,
        colorOrange1,colorOrange2,colorOrange3,colorOrange4,
        colorGreen1,colorGreen2,colorGreen3,colorGreen4,
        colorPurple1,colorPurple2,colorPurple3,colorPurple4,
        colorDarkPurple1,colorDarkPurple2,colorDarkPurple3,colorDarkPurple4,
        colorMagenta1,colorMagenta2,colorMagenta3,colorMagenta4,
        colorGray1,colorGray2,colorGray3,colorGray4

    # Color schemes tab20c and tab20b are taken from
    # https://matplotlib.org/users/plotting/colormaps/grayscale_01_04.pdf
    const colorBrown1="#8C6D31"
    const colorBrown2="#BD9E39"
    const colorBrown3="#E7BA52"
    const colorBrown4="#E7CB94"

    const colorBlue1="#3182BD"
    const colorBlue2="#6BAED6"
    const colorBlue3="#9ECAE1"
    const colorBlue4="#C6DBEF"

    const colorRed1="#843C39"
    const colorRed2="#AD494A"
    const colorRed3="#D6616B"
    const colorRed4="#E7969C"

    const colorOrange1="#E6550D"
    const colorOrange2="#FD8D3C"
    const colorOrange3="#FDAE6B"
    const colorOrange4="#FDD0A2"

    const colorGreen1="#31A354"
    const colorGreen2="#8CA252"
    const colorGreen3="#B5CF6B"
    const colorGreen4="#C7E9C0"

    const colorPurple1="#756BB1"
    const colorPurple2="#9E9AC8"
    const colorPurple3="#BCBDDC"
    const colorPurple4="#DADAEB"

    const colorDarkPurple1="#393B79"
    const colorDarkPurple2="#5254A3"
    const colorDarkPurple3="#6B6ECF"
    const colorDarkPurple4="#9C9EDE"

    const colorMagenta1="#7B4173"
    const colorMagenta2="#A55194"
    const colorMagenta3="#CE6DBD"
    const colorMagenta4="#DE9ED6"

    const colorGray1="#636363"
    const colorGray2="#969696"
    const colorGray3="#BDBDBD"
    const colorGray4="#D9D9D9"
end


text:
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.8.0-DEV
#     language: julia
#     name: julia-1.8
# ---

# %% [markdown]
# https://docs.julialang.org/en/v1/base/arrays/#Base.PermutedDimsArrays.PermutedDimsArray

# %%
using ArraysOfArrays
f(inp) = nestedview(PermutedDimsArray(inp, (2, 3, 1)), 2)
w, h = 5, 5
inp = 0.1(1:3) .+ 10(1:5)' .+ reshape(1:5, 1, 1, :)
out = f(inp)

# %%
out[2]

# %%
using BenchmarkTools
w, h = 1000, 1000
inp = 0.1(1:3) .+ 10(1:5)' .+ reshape(1:5, 1, 1, :)
@show [@view(inp[i, :, :]) for i in 1:3] == collect(eachslice(inp; dims=1)) == f(inp)
@btime [@view($inp[i, :, :]) for i in 1:3]
@btime collect(eachslice($inp; dims=1))
@btime f($inp)
@btime PermutedDimsArray($inp, (2, 3, 1));

# %%

max_stars_repo_path: 0014/collect views.jl

max_stars_repo_name: genkuroki/public

max_issues_repo_path: 0014/collect views.jl

max_issues_repo_name: genkuroki/public

max_forks_repo_path: 0014/collect views.jl

max_forks_repo_name: genkuroki/public


text:
module SpikeStats

using Random
using StatsBase
using Statistics
using DataFrames
using DSP
using Distributions

include("poissonSpikeModels.jl")
include("computeCV.jl")
include("computeFano.jl")
include("spikeCorr.jl")
include("spikeTriggeredAvgFuncs.jl")
include("whiteNoiseGenerator.jl")

export poissonSpikes, computeCV, computeFano, poissonSpikesRefractory, poissonSpikesVarRate,spikeTriggeredAvg, spikeCorr, whiteNoiseGenerator, spikeTriggeredAvgNDimFilter, poissonSpikesStim
end


text:
# # SIMP example: Point Load Cantilever
#-
#md # !!! tip
#md #     This example is also available as a Jupyter notebook:
#md #     [`simp.ipynb`](@__NBVIEWER_ROOT_URL__/examples/simp.ipynb)
#-
# ## Commented Program
#
# What follows is a program spliced with comments.
#md # The full program, without comments, can be found in the next [section](@ref simp-plain-program).

using TopOpt

# ### Define the problem
E = 1.0 # Young’s modulus
v = 0.3 # Poisson’s ratio
f = 1.0; # downward force

nels = (30, 10, 10)
problem = PointLoadCantilever(Val{:Linear}, nels, (1.0, 1.0, 1.0), E, v, f);

# See also the detailed API of `PointLoadCantilever`:
#md # ```@docs
#md # TopOpt.TopOptProblems.PointLoadCantilever
#md # ```

# ### Parameter settings
V = 0.3 # volume fraction
xmin = 1e-6 # minimum density
rmin = 2.0; # density filter radius

# ### Define a finite element solver
penalty = TopOpt.PowerPenalty(3.0)
solver = FEASolver(Direct, problem; xmin=xmin, penalty=penalty)

# ### Define compliance objective
comp = TopOpt.Compliance(problem, solver)
filter = DensityFilter(solver; rmin=rmin)
obj = x -> comp(filter(x))

# ### Define volume constraint
volfrac = TopOpt.Volume(problem, solver)
constr = x -> volfrac(filter(x)) - V

# You can enable the iteration printouts with `Nonconvex.NonconvexCore.show_residuals[] = true`

# ### Define subproblem optimizer
mma_options =
    options = MMAOptions(;
        maxiter=3000, tol=Nonconvex.Tolerance(; x=1e-3, f=1e-3, kkt=0.001)
    )
convcriteria = Nonconvex.KKTCriteria()
x0 = fill(V, length(solver.vars))
optimizer = Optimizer(
    obj, constr, x0, MMA87(); options=mma_options, convcriteria=convcriteria
)

# ### Define SIMP optimizer
simp = SIMP(optimizer, solver, penalty.p);

# ### Solve
result = simp(x0);

@show result.convstate
@show optimizer.workspace.iter
@show result.objval

# ### (Optional) Visualize the result using Makie.jl
# Need to run `using Pkg; Pkg.add(Makie)` first
# ```julia
# using TopOpt.TopOptProblems.Visualization: visualize
# fig = visualize(problem; topology = result.topology,
#     problem; topology = result.topology, default_exagg_scale = 0.07,
#     scale_range = 10.0, vector_linewidth = 3, vector_arrowsize = 0.5,
# )
# Makie.display(fig)
# ```

# or convert it to a Mesh
# Need to run `using Pkg; Pkg.add(GeometryBasics)` first
# ```julia
# import Makie, GeometryBasics
# result_mesh = GeometryBasics.Mesh(problem, result.topology);
# Makie.mesh(result_mesh)
# ```

#md # ## [Plain Program](@id simp-plain-program)
#md #
#md # Below follows a version of the program without any comments.
#md # The file is also available here: [simp.jl](simp.jl)
#md #
#md # ```julia
#md # @__CODE__
#md # ```


text:
# *****************************************************************************
# Written by Ritchie Lee, ritchie.lee@sv.cmu.edu
# *****************************************************************************
# Copyright ã 2015, United States Government, as represented by the
# Administrator of the National Aeronautics and Space Administration. All
# rights reserved.  The Reinforcement Learning Encounter Simulator (RLES)
# platform is licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. You
# may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0. Unless required by applicable
# law or agreed to in writing, software distributed under the License is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the specific language
# governing permissions and limitations under the License.
# _____________________________________________________________________________
# Reinforcement Learning Encounter Simulator (RLES) includes the following
# third party software. The SISLES.jl package is licensed under the MIT Expat
# License: Copyright (c) 2014: Youngjun Kim.
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED
# "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
# NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
# PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
# *****************************************************************************

"""
Monte Carlo tree search for the ACASX problem.
Example usage: config=configure(ACASX_MCTS,"normal","nvn_dasc"); acasx_mcts(;config...)
"""
module ACASX_MCTS

export configure, acasx_mcts

import Compat.ASCIIString
using ExprSearch.MCTS
using Datasets
using Reexport
using JSON, GZip
using RLESUtils, FileUtils, Configure, Observers, Loggers, LogSystems
import RLESTypes.SymbolTable

using GrammarExpts
using ACASXProblem
using DerivTreeVis, MCTSTreeView
import Configure.configure

const CONFIGDIR = joinpath(dirname(@__FILE__), "..", "config")
const RESULTDIR = joinpath(dirname(@__FILE__), "..", "..", "..", "results")

configure(::Type{Val{:ACASX_MCTS}}, configs::AbstractString...) = configure_path(CONFIGDIR, configs...)

"""
Example call:
config=configure(ACASX_MCTS, "nvn_dasc", "normal")
acasx_mcts(; config...)
"""
function acasx_mcts(;outdir::AbstractString=joinpath(RESULTDIR, "ACASX_MCTS"),
                    seed=1,
                    logfileroot::AbstractString="acasx_mcts_log",

                    runtype::Symbol=:nmacs_vs_nonnmacs,
                    data::AbstractString="dasc",
                    manuals::AbstractString="",
                    clusterdataname::AbstractString="",

                    n_iters::Int64=200,
                    searchdepth::Int64=20,
                    explorationconst::Float64=2000.0,
                    q0::Float64=-1000.0,
                    maxsteps::Int64=20,
                    max_neg_reward::Float64=-1000.0,
                    step_reward::Float64=0.0,
                    maxmod::Bool=false, #use the max update mod

                    loginterval::Int64=100,
                    vis::Bool=true,
                    mctstreevis::Bool=false,
                    treevis_interval::Int64=50)
    mkpath(outdir)

    problem = ACASXClustering(runtype, data, manuals, clusterdataname)

    logsys = MCTS.logsystem()
    empty_listeners!(logsys)
    send_to!(STDOUT, logsys, ["verbose1", "result"])
    send_to!(STDOUT, logsys, "current_best"; interval=loginterval)
    logs = TaggedDFLogger()
    send_to!(logs, logsys, ["computeinfo", "parameters", "result"])
    send_to!(logs, logsys,  "current_best"; interval=loginterval)
    send_to!(logs, logsys,  "elapsed_cpu_s"; interval=loginterval)

    #if mctstreevis
        #view, viewstep = viewstep_f(treevis_interval)
        #add_observer(observer, "mcts_tree", viewstep)
    #end

    mcts_params = MCTSESParams(maxsteps, max_neg_reward, step_reward, n_iters, searchdepth,
                             explorationconst, maxmod, q0, seed, logsys;
                             userargs=SymbolTable(:ids=>collect(1:length(problem.Dl))))

    result = exprsearch(mcts_params, problem)

    #manually push! extra info to log
    push!(logs, "parameters", ["seed", seed])
    push!(logs, "parameters", ["runtype", runtype])
    push!(logs, "parameters", ["data", data])
    add_folder!(logs, "expression", [ASCIIString, ASCIIString, ASCIIString],
        ["raw", "pretty", "natural"])
    push!(logs, "expression", [string(result.expr), pretty_string(result.tree, FMT_PRETTY),
        pretty_string(result.tree, FMT_NATURAL, true)])

    #save log
    outfile = joinpath(outdir, "$(logfileroot).txt")
    save_log(outfile, logs)

    if vis
        derivtreevis(get_derivtree(result), joinpath(outdir, "$(logfileroot)_derivtreevis"))
    end

    #save mcts tree
    if mctstreevis
        GZip.open(joinpath(outdir, "mctstreevis.json.gz"), "w") do f
            JSON.print(f, view.steps)
        end
    end

    textfile(joinpath(outdir, "summary.txt"), "mcts", seed=seed, n_iters=n_iters,
           fitness=result.fitness, expr=string(result.expr))

    result
end

end #module


text:
using Test
using StaticArrays

function empty_test(h::HRepresentation{T}) where T
    #v = @inferred doubledescription(h)
    v = doubledescription(h)
    U = T == Float64 ? Float64 : Rational{BigInt}
    v = doubledescription(h)
    @test v isa Polyhedra.Hull{U,Vector{U}}
    @test !haspoints(v)
    @test !hasallrays(v)
end
empty_hyperplane_test(z) = empty_test(intersect(HyperPlane([z, z], 1)))
empty_halfspace_test(z) = empty_test(intersect(HalfSpace([z, z], -1)))
empty_space_test(z) = empty_test(HyperPlane([z, 1], 1) ∩ HyperPlane([z, 1], -1))
empty_range_test(z) = empty_test(HalfSpace([1, z], -1) ∩ HalfSpace([-1, z], -1))

function simplex_test(z)
    h = hrep([HalfSpace([ 1,  1], 1),
              HalfSpace([ 1, -1], z),
              HalfSpace([-1,  z], z)])
    v = doubledescription(h)
    T = z isa AbstractFloat ? Float64 : Rational{BigInt}
    @test v isa Polyhedra.Hull{T,Vector{T}}
    @test collect(points(v)) == [[z, z], [z, 1], [1//2, 1//2]]
    @test !hasallrays(v)
    h = hrep([HalfSpace(zeros(typeof(z), 2), z); h.halfspaces])
    h = hrep([HalfSpace([ 1,  1], 1),
              HalfSpace([ z,  z], z),
              HalfSpace([ 1, -1], z),
              HalfSpace([-1,  z], z)])
    #intersect!(h, )
    v = doubledescription(h)
    @test v isa Polyhedra.Hull{T,Vector{T}}
    @test collect(points(v)) == [[z, z], [z, 1], [1//2, 1//2]]
    @test !hasallrays(v)
end

@testset "Double Description" begin
    @test Polyhedra.polytypefor(Float32) == Float64
    @testset "H-representation -> V-representation" begin
        @testset "Intersection" begin
            @testset "Vector" begin
                @testset "Exact" begin
                    simplex_test(0)
                end
                @testset "Numerical" begin
                    simplex_test(0.0)
                end
            end
            @testset "SVector" begin
                @testset "Exact" begin
                    h = hrep([HalfSpace((@SVector [ 1,  1]), 1),
                              HalfSpace((@SVector [ 1, -1]), 0),
                              HalfSpace((@SVector [-1,  0]), 0)])
                    #v = @inferred doubledescription(h)
                    v = doubledescription(h)
                    @test v isa Polyhedra.Hull{Rational{BigInt},SVector{2,Rational{BigInt}}}
                    @test collect(points(v)) == [(@SVector [0, 0]), (@SVector [0, 1]), (@SVector [1//2, 1//2])]
                    @test !hasallrays(v)
                end
                @testset "Numerical" begin
                    h = hrep([HalfSpace((@SVector [ 1.,  1]), 1),
                              HalfSpace((@SVector [ 1., -1]), 0),
                              HalfSpace((@SVector [-1.,  0]), 0)])
                    #v = @inferred doubledescription(h)
                    v = doubledescription(h)
                    @test v isa Polyhedra.Hull{Float64,SVector{2,Float64}}
                    @test collect(points(v)) == [(@SVector [0.0, 0.0]), (@SVector [0.0, 1.0]), (@SVector [1/2, 1/2])]
                    @test !hasallrays(v)
                end
            end
        end
        @testset "MixedMatHRep" begin
            @testset "Exact" begin
                h = hrep([ 1  1
                           1 -1
                          -1  0],
                         [1, 0, 0])
                #v = @inferred doubledescription(h)
                v = doubledescription(h)
                @test v isa Polyhedra.MixedMatVRep{Rational{BigInt}}
                @test v.V == [0//1 0//1; 0//1 1//1; 1//2 1//2]
                @test v.R == zeros(Rational{BigInt}, 0, 2)
                @test isempty(v.Rlinset)
            end
            @testset "Numerical" begin
                h = hrep([ 1  1
                           1 -1
                          -1  0],
                         [1., 0, 0])
                #v = @inferred doubledescription(h)
                v = doubledescription(h)
                @test v isa Polyhedra.MixedMatVRep{Float64}
                @test v.V == [0 0; 0 1; 1/2 1/2]
                @test v.R == zeros(0, 2)
                @test isempty(v.Rlinset)
            end
        end
        _str(::Float64) = "Numerical"
        _str(::Int) = "Exact"
        @testset "Empty Intersection $(_str(z))" for z in [0, 0.0]
            @testset "0x_1 + 0x_2 = 1" begin
                empty_hyperplane_test(z)
            end
            @testset "0x_1 + 0x_2 ≤ -1" begin
                empty_halfspace_test(z)
            end
            @testset "-1 = 0x_1 + x_2 = 1" begin
                empty_space_test(z)
            end
            @testset "1 ≤ x_1 + 0x_2 ≤ -1" begin
                empty_range_test(z)
            end
        end
        @testset "Simple hyperplane $(_str(o))" for o in [1, 1.0]
            h = intersect(HyperPlane([o, o], o))
            v = doubledescription(h)
            @test nlines(v) == 1
            @test first(lines(v)) == Line([-1, 1])
            @test !hasrays(v) == 1
            @test npoints(v) == 1
            @test first(points(v)) == [1, 0]
        end
        @testset "Quadrilateral $(_str(z))" for z in [0, 0.0]
            h = HalfSpace([2, 1], 4) ∩ HalfSpace([1, 2], 4) ∩
                HalfSpace([-1, z], z) ∩ HalfSpace([z, -1], z)
            v = doubledescription(h)
            @test !haslines(v)
            @test !hasrays(v)
            @test npoints(v) == 4
            @test all(points(v) .≈ [[0, 0], [2, 0], [0, 2], [4//3, 4//3]])
        end
    end
    @testset "V-representation -> H-representation" begin
        @testset "Hull" begin
            @testset "Vector" begin
                @testset "Exact" begin
                    v = conichull([1, 0],
                                  [0, 1])
                    h = doubledescription(v)
                    #h = @inferred doubledescription(v)
                    @test !hashyperplanes(h)
                    @test collect(halfspaces(h)) == [HalfSpace([0, 0], 1), HalfSpace([-1, 0], 0), HalfSpace([0, -1], 0)] # FIXME get rid of (0, 0) 1
                end
                @testset "Numerical" begin
                    v = conichull([1., 0.],
                                  [0., 1.])
                    h = doubledescription(v)
                    #h = @inferred doubledescription(v)
                    @test !hashyperplanes(h)
                    @test collect(halfspaces(h)) == [HalfSpace([0, 0], 1), HalfSpace([-1, 0], 0), HalfSpace([0, -1], 0)] # FIXME get rid of (0, 0) 1
                end
            end
        end
    end
end

max_stars_repo_name: blegat/Polyhedra

max_issues_repo_name: blegat/Polyhedra

max_forks_repo_name: blegat/Polyhedra


text:
# Autogenerated wrapper script for SQLite_jll for x86_64-linux-gnu
export libsqlite, sqlite3

using Zlib_jll
JLLWrappers.@generate_wrapper_header("SQLite")
JLLWrappers.@declare_library_product(libsqlite, "libsqlite3.so.0")
JLLWrappers.@declare_executable_product(sqlite3)
function __init__()
    JLLWrappers.@generate_init_header(Zlib_jll)
    JLLWrappers.@init_library_product(
        libsqlite,
        "lib/libsqlite3.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_executable_product(
        sqlite3,
        "bin/sqlite3",
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
module TestMSQueue

using Reagents
using Test

include("../../../examples/msqueue.jl")

function test_1()
    q = MSQueue{Int}()
    push!(q, 1)
    @test popfirst!(q) == 1
end

function test_123()
    q = MSQueue{Int}()
    push!(q, 1)
    push!(q, 2)
    push!(q, 3)
    @test popfirst!(q) == 1
    @test popfirst!(q) == 2
    @test popfirst!(q) == 3
end

end  # module


text:
# Deltamax Estimator
# Estimate the number of servers using the Deltamax algorithm
# Input: arrival order sorted by departure order
# Output: estimated number of servers
function c_deltamax(out_order::Vector{Int64})
  N = size(out_order, 1)
  currmax = out_order[1]
  deltamax = 0
  for i = 2:N
    nextmax = max(out_order[i], currmax)
    deltamax = max(deltamax, nextmax - currmax)
    currmax = nextmax
  end
  deltamax
end

# Order-based Estimator (Sets)
# Estimate the number of servers using the Order-based algorithm
# Input: arrival order sorted by departure order
# Output: estimated number of servers
function c_order_slow(out_order::Vector{Int64})
    N = size(out_order, 1)
    currmax = 0
    nservers = 0
    for i = 1:N
        currmax = max(out_order[i], currmax)
        custall = Set(1:currmax)
        custdeparted = Set(out_order[1:i])
        custservice = setdiff(custall, custdeparted)
        nservers = max(length(custservice) + 1, nservers)
    end
    nservers
end

# Order-based Estimator (Scalar)
# Estimate the number of servers using the Order-based algorithm
# Input: arrival order sorted by departure order
# Output: estimated number of servers
function c_order(out_order::Vector{Int64})
    N = size(out_order, 1)
    currmax = 0
    noccupied = 0
    nservers = 0
    for i = 1:N
        nextmax = max(out_order[i], currmax)
        if out_order[i] < currmax
            noccupied = noccupied - 1
        else
            noccupied = noccupied + nextmax - currmax - 1
        end
        currmax = nextmax
        nservers = max(nservers, noccupied + 1)
    end
    nservers
end

# WORK IN PROGRESS: Summary Estimator - fast version
# WARNING: PRODUCES INCORRECT OUTPUT
# function c_var_unf_fast(A::Vector{Float64}, D::Vector{Float64}, cmax::Int64)
#   # Initilaize data structures and functions
#   N = size(D, 1)
#   B̂ = zeros(N)
#   temp = zeros(N)
#   VS = zeros(cmax, 2) # var, unf
#   VS[1,:] = [-1,-1]
#   max_undef = 1
#   Dsort = sort(D)
#   ## Grid Search
#   for c = 2:cmax
#     B̂[1:c] = A[1:c]
#     for i = (c + 1):N
#       B̂[i] = max(A[i], Dsort[i - c])
#     end
#     Ŝ = D - B̂
#     if any(Ŝ .< 0)
#       VS[c, 1:2] = -1
#       max_undef = c
#     else
#       VS[c, 1] = var(Ŝ) * (N-1) # var
#     end
#   end
#   (r1, r2) = (99, 99) # initialize result to 99 for debugging
#   # var method
#   try
#     r1 = findin(VS[:, 1], minimum(filter(x -> x >= 0, VS[:, 1])))[1]
#   catch
#     r1 = 1
#   end
#   # uninformed method
#   r2 = max_undef + 1
#   (r1, r2, VS)
# end

# Summary Estimator
# Estimate the number of servers using the Variance and Uninformed methods
# Input: arrival times, departure times, and the max number of servers
# Output: estimated number of servers for each method
function c_var_unf(A::Vector{Float64}, D::Vector{Float64}, cmax::Int64)
  # Initilaize data structures and functions
  N = size(D, 1)
  B̂ = zeros(N)
  temp = zeros(N)
  VS = zeros(cmax, 2) # var, unf
  max_undef = 0
  ## Grid Search
  for c = 1:cmax
    B̂[1:c] = A[1:c]
    for i = (c + 1):N
      Dₖ = sort(D[1:(i - 1)])
      B̂[i] = max(A[i], Dₖ[i - c])
    end
    Ŝ = D - B̂
    if any(Ŝ .< 0)
      VS[c, 1:2] = -1
      max_undef = c
    else
      VS[c, 1] = var(Ŝ) * (N-1) # var
    end
  end
  (r1, r2) = (99, 99) # initialize result to 99 for debugging
  # var method
  try
    r1 = findin(VS[:, 1], minimum(filter(x -> x >= 0, VS[:, 1])))[1]
  catch
    r1 = 1
  end
  # uninformed method
  r2 = max_undef + 1
  (r1, r2, VS)
end


text:
################################################################################
# Coordinates
################################################################################
function minimal_coordinates(joint::JointConstraint, pbody::Node, cbody::Node)
	Δx = minimal_coordinates(joint.translational, current_configuration(pbody.state)..., current_configuration(cbody.state)...)
    Δθ = minimal_coordinates(joint.rotational,    current_configuration(pbody.state)..., current_configuration(cbody.state)...)
	return [Δx; Δθ]
end

function set_minimal_coordinates!(joint::JointConstraint, pnode::Node, cnode::Node,
        timestep;
        Δx::AbstractVector=szeros(input_dimension(joint.translational)),
        Δθ::AbstractVector=szeros(input_dimension(joint.rotational)))

    set_minimal_coordinates!(joint.rotational,    pnode, cnode, timestep; Δθ=Δθ)
    set_minimal_coordinates!(joint.translational, pnode, cnode, timestep; Δx=Δx)

    return nothing
end

function set_minimal_coordinates!(mechanism, joint::JointConstraint{T,N,Nc}, xθ; iter=true) where {T,N,Nc}
    # bodies
    pbody = get_body(mechanism, joint.parent_id)
    cbody = get_body(mechanism, joint.child_id)

    # unpack
    Δx = xθ[SUnitRange(joint.minimal_index[1][1], joint.minimal_index[1][2])]
    Δθ = xθ[SUnitRange(joint.minimal_index[2][1], joint.minimal_index[2][2])]

    # set
    set_minimal_coordinates!(joint, pbody, cbody, mechanism.timestep, Δx=Δx, Δθ=Δθ)

    # recursive update down the kinematic chain
    if iter
        current = get_minimal_coordinates(mechanism)
        for id in recursivedirectchildren!(mechanism.system, joint.id)
            node = get_node(mechanism, id)
            if node isa JointConstraint
                set_minimal_coordinates!(mechanism, node, current[id])
            end
        end
    end
end

################################################################################
# Velocities
################################################################################
function minimal_velocities(joint::JointConstraint, pnode::Node, cnode::Node, timestep)
    Δv = minimal_velocities(joint.translational, pnode, cnode, timestep)
    Δω = minimal_velocities(joint.rotational,    pnode, cnode, timestep)
    return [Δv; Δω]
end

function set_minimal_velocities!(joint::JointConstraint, pnode::Node, cnode::Node, timestep;
        Δv=szeros(input_dimension(joint.translational)),
        Δω=szeros(input_dimension(joint.rotational)))

    # get
	vb, ωb = get_child_velocity(joint,
        initial_configuration_velocity(pnode.state)...,
	 	current_configuration(cnode.state)...,
        timestep,
        Δv=Δv, Δω=Δω)

    # set
	set_maximal_velocities!(cnode; v=vb, ω=ωb)
	set_previous_configuration!(cnode, timestep)

    return nothing
end

function set_minimal_velocities!(mechanism, joint::JointConstraint{T,N,Nc}, vω) where {T,N,Nc}
    # bodies
    pbody = get_body(mechanism, joint.parent_id)
    cbody = get_body(mechanism, joint.child_id)

    # unpack
    Δv = vω[SUnitRange(joint.minimal_index[1][1], joint.minimal_index[1][2])]
    Δω = vω[SUnitRange(joint.minimal_index[2][1], joint.minimal_index[2][2])]

    # set
    set_minimal_velocities!(joint, pbody, cbody, mechanism.timestep, Δv=Δv, Δω=Δω)
    return
end

function get_child_velocity(joint::JointConstraint,
    xa::AbstractVector, va::AbstractVector, qa::Quaternion, ωa::AbstractVector,
    xb::AbstractVector, qb::Quaternion,
    timestep;
    Δv=szeros(input_dimension(joint.translational)),
    Δω=szeros(input_dimension(joint.rotational)))

    rot = joint.rotational
    tra = joint.translational
    pa = tra.vertices[1]
    pb = tra.vertices[2]
    orientation_offset = rot.orientation_offset
    Arot = zerodimstaticadjoint(nullspace_mask(rot))
    Atra = zerodimstaticadjoint(nullspace_mask(tra))

    Δx = minimal_coordinates(joint.translational, xa, qa, xb, qb)
    Δq = inv(orientation_offset) * inv(qa) * qb

    # 1 step backward in time
    xa1 = next_position(xa, -va, timestep)
    qa1 = next_orientation(qa, -ωa, timestep)

    # Finite difference
    Δx1 = Δx .- Δv * timestep
    Δq1 = Δq * inv(axis_angle_to_quaternion(Arot * Δω * timestep))

    qb1 = qa1 * orientation_offset * Δq1
    xb1 = xa1 + vector_rotate(pa + Atra * Δx1, qa1) - vector_rotate(pb, qb1)

    # Finite difference
    vb = (xb - xb1) / timestep
    ωb = angular_velocity(qb1, qb, timestep)

    return vb, ωb
end

################################################################################
# Coordinates and Velocities
################################################################################
function minimal_coordinates_velocities(joint::JointConstraint, pnode::Node, cnode::Node, timestep)
    Δxθ = minimal_coordinates(joint, pnode, cnode)
    Δvω = minimal_velocities(joint, pnode, cnode, timestep)
    return [Δxθ; Δvω]
end

function set_minimal_coordinates_velocities!(mechanism::Mechanism, joint::JointConstraint;
        xmin::AbstractVector=szeros(2*input_dimension(joint)))

    pnode = get_body(mechanism, joint.parent_id)
    cnode = get_body(mechanism, joint.child_id)

    xa = pnode.state.x2
    va = pnode.state.v15
    qa = pnode.state.q2
    ωa = pnode.state.ω15

    za = [xa; va; vector(qa); ωa]

    set_minimal_coordinates_velocities!(joint, pnode, cnode, mechanism.timestep; xmin=xmin, zp=za)
end

function set_minimal_coordinates_velocities!(joint::JointConstraint,
        pbody::Node, cbody::Node,
        timestep;
        zp::AbstractVector=[pbody.state.x2; pbody.state.v15; vector(pbody.state.q2); pbody.state.ω15],
        xmin::AbstractVector=szeros(2*input_dimension(joint)))

    nu = input_dimension(joint)

    Δx = xmin[SUnitRange(joint.minimal_index[1]...)]
    Δθ = xmin[SUnitRange(joint.minimal_index[2]...)]

    Δv = xmin[nu .+ SUnitRange(joint.minimal_index[1]...)]
    Δω = xmin[nu .+ SUnitRange(joint.minimal_index[2]...)]

    rot = joint.rotational
    tra = joint.translational
    pa = tra.vertices[1]
    pb = tra.vertices[2]
    orientation_offset = rot.orientation_offset
    Arot = zerodimstaticadjoint(nullspace_mask(rot))
    Atra = zerodimstaticadjoint(nullspace_mask(tra))

    # parent state
    xa = SVector{3}(zp[1:3])
    va = SVector{3}(zp[3 .+ (1:3)])
    qa = Quaternion(zp[6 .+ (1:4)]...)
    ωa = SVector{3}(zp[10 .+ (1:3)])

    # positions
    Δq = axis_angle_to_quaternion(Arot * Δθ)
	qb = qa * orientation_offset * Δq
    xb = xa + vector_rotate(pa + Atra * Δx, qa) - vector_rotate(pb, qb)

    # previous configuration
    xa1 = next_position(xa, -va, timestep)
    qa1 = next_orientation(qa, -ωa, timestep)

    # finite-difference configuration
    Δx1 = Δx .- Δv * timestep
    Δq1 = Δq * inv(axis_angle_to_quaternion(Arot * Δω * timestep))

    qb1 = qa1 * orientation_offset * Δq1
    xb1 = xa1 + vector_rotate(pa + Atra * Δx1, qa1) - vector_rotate(pb, qb1)

    # finite-difference velocity
    vb = (xb - xb1) / timestep
    ωb = angular_velocity(qb1, qb, timestep)

    # set child state
    cbody.state.x2 = xb
    cbody.state.v15 = vb
    cbody.state.q2 = qb
    cbody.state.ω15 = ωb

    return [xb; vb; vector(qb); ωb]
end

function minimal_coordinates_velocities_jacobian_parent(joint::JointConstraint{T},
    xa::AbstractVector, va::AbstractVector, qa::Quaternion, ωa::AbstractVector,
    timestep;
    Δx=szeros(input_dimension(joint.translational)),
    Δθ=szeros(input_dimension(joint.rotational)),
    Δv=szeros(input_dimension(joint.translational)),
    Δω=szeros(input_dimension(joint.rotational))) where T

    rot = joint.rotational
    tra = joint.translational
    pa = tra.vertices[1]
    pb = tra.vertices[2]
    orientation_offset = rot.orientation_offset
    Arot = zerodimstaticadjoint(nullspace_mask(rot))
    Atra = zerodimstaticadjoint(nullspace_mask(tra))

    # positions
    Δq = axis_angle_to_quaternion(Arot * Δθ)
	qb = qa * orientation_offset * Δq
    xb = xa + vector_rotate(pa + Atra * Δx, qa) - vector_rotate(pb, qb)

    # step backward in time
    xa1 = next_position(xa, -va, timestep)
    qa1 = next_orientation(qa, -ωa, timestep)

    # step backward in time
    Δx1 = Δx .- Δv * timestep
    Δq1 = Δq * inv(axis_angle_to_quaternion(Arot * Δω * timestep))

    qb1 = qa1 * orientation_offset * Δq1
    xb1 = xa1 + vector_rotate(pa + Atra * Δx1, qa1) - vector_rotate(pb, qb1)

    # Jacobians

    ∂xb∂xa = 1.0 * I(3)

    ∂xb∂va = szeros(T, 3, 3)

    ∂xb∂qa = ∂vector_rotate∂q(pa + Atra * Δx, qa)
    ∂xb∂qa += -∂vector_rotate∂q(pb, qb) * Rmat(orientation_offset * axis_angle_to_quaternion(Arot * Δθ))

    ∂xb∂ωa = szeros(T, 3, 3)

    ∂qb∂xa = szeros(T, 4, 3)

    ∂qb∂va = szeros(T, 4, 3)

    ∂qb∂qa = Rmat(orientation_offset * axis_angle_to_quaternion(Arot * Δθ))

    ∂qb∂ωa = szeros(T, 4, 3)

    ∂vb∂xa = szeros(T, 3, 3)

    ∂vb∂va = 1.0 * I(3)

    ∂vb∂qa = 1.0 / timestep * ∂vector_rotate∂q(pa + Atra * Δx, qa)
    ∂vb∂qa -= 1.0 / timestep * ∂vector_rotate∂q(pb, qb) * Rmat(orientation_offset * Δq)
    ∂vb∂qa += -1.0 / timestep * ∂vector_rotate∂q(pa + Atra * Δx1, qa1) * rotational_integrator_jacobian_orientation(qa, -ωa, timestep, attjac=false)
    ∂vb∂qa += 1.0 / timestep * ∂vector_rotate∂q(pb, qb1) * Rmat(orientation_offset * Δq1) * rotational_integrator_jacobian_orientation(qa, -ωa, timestep, attjac=false)

    ∂vb∂ωa = 1.0 / timestep * ∂vector_rotate∂q(pa + Atra * Δx1, qa1) * rotational_integrator_jacobian_velocity(qa, -ωa, timestep)
    ∂vb∂ωa += -1.0 / timestep * ∂vector_rotate∂q(pb, qb1) * Rmat(orientation_offset * Δq1) * rotational_integrator_jacobian_velocity(qa, -ωa, timestep)

    ∂ωb∂xa = szeros(T, 3, 3)

    ∂ωb∂va = szeros(T, 3, 3)

    ∂ωb∂qa = ∂angular_velocity∂q1(qb1, qb, timestep) * Rmat(orientation_offset * Δq1) * rotational_integrator_jacobian_orientation(qa, -ωa, timestep, attjac=false)
    ∂ωb∂qa += ∂angular_velocity∂q2(qb1, qb, timestep) * Rmat(orientation_offset * Δq)

    ∂ωb∂ωa = -1.0 * ∂angular_velocity∂q1(qb1, qb, timestep) * Rmat(orientation_offset * Δq1) * rotational_integrator_jacobian_velocity(qa, -ωa, timestep)

    [
        ∂xb∂xa ∂xb∂va ∂xb∂qa ∂xb∂ωa;
        ∂vb∂xa ∂vb∂va ∂vb∂qa ∂vb∂ωa;
        ∂qb∂xa ∂qb∂va ∂qb∂qa ∂qb∂ωa;
        ∂ωb∂xa ∂ωb∂va ∂ωb∂qa ∂ωb∂ωa;
    ]
end

function set_minimal_coordinates_velocities_jacobian_parent(joint::JointConstraint,
    pnode::Node, cnode::Node, timestep)

    xmin = minimal_coordinates_velocities(joint, pnode, cnode, timestep)
    qb = current_orientation(cnode.state)
    xa, va, qa, ωa = initial_configuration_velocity(pnode.state)
    za = [xa; va; vector(qa); ωa]

    nu = input_dimension(joint)
    Δx = xmin[SUnitRange(joint.minimal_index[1]...)]
    Δθ = xmin[SUnitRange(joint.minimal_index[2]...)]
    Δv = xmin[nu .+ SUnitRange(joint.minimal_index[1]...)]
    Δω = xmin[nu .+ SUnitRange(joint.minimal_index[2]...)]
    J = minimal_coordinates_velocities_jacobian_parent(joint,
                xa, va, qa, ωa,
                timestep;
                Δx=Δx,
                Δθ=Δθ,
                Δv=Δv,
                Δω=Δω)
    J = cat(Diagonal(sones(6)), LVᵀmat(qb)', Diagonal(sones(3)), dims=(1,2)) * J
    J = J * cat(Diagonal(sones(6)), LVᵀmat(qa), Diagonal(sones(3)), dims=(1,2))
    return J
end

function minimal_coordinates_velocities_jacobian_minimal(joint::JointConstraint{T},
    xa::AbstractVector, va::AbstractVector, qa::Quaternion, ωa::AbstractVector,
    timestep;
    Δx=szeros(input_dimension(joint.translational)),
    Δθ=szeros(input_dimension(joint.rotational)),
    Δv=szeros(input_dimension(joint.translational)),
    Δω=szeros(input_dimension(joint.rotational))) where T

    rot = joint.rotational
    tra = joint.translational
    pa = tra.vertices[1]
    pb = tra.vertices[2]
    orientation_offset = rot.orientation_offset
    Arot = zerodimstaticadjoint(nullspace_mask(rot))
    Atra = zerodimstaticadjoint(nullspace_mask(tra))

    # positions
    Δq = axis_angle_to_quaternion(Arot * Δθ)
	qb = qa * orientation_offset * Δq
    xb = xa + vector_rotate(pa + Atra * Δx, qa) - vector_rotate(pb, qb)

    # step backward in time
    xa1 = next_position(xa, -va, timestep)
    qa1 = next_orientation(qa, -ωa, timestep)

    # step backward in time
    Δq = axis_angle_to_quaternion(Arot * Δθ)
    Δx1 = Δx .- Δv * timestep
    Δq1 = Δq * inv(axis_angle_to_quaternion(Arot * Δω * timestep))

    qb1 = qa1 * orientation_offset * Δq1
    xb1 = xa1 + vector_rotate(pa + Atra * Δx1, qa1) - vector_rotate(pb, qb1)

    # Jacobians
    ∂xb∂Δx = rotation_matrix(qa) * Atra
    ∂xb∂Δθ = -∂vector_rotate∂q(pb, qb) * Lmat(qa * orientation_offset) * daxis_angle_to_quaterniondx(Arot * Δθ) * Arot
    ∂xb∂Δv = szeros(T, 3, input_dimension(joint.translational))
    ∂xb∂Δω = szeros(T, 3, input_dimension(joint.rotational))

    ∂qb∂Δx = szeros(T, 4, input_dimension(joint.translational))
    ∂qb∂Δθ = Lmat(qa * orientation_offset) * daxis_angle_to_quaterniondx(Arot * Δθ) * Arot
    ∂qb∂Δv = szeros(T, 4, input_dimension(joint.translational))
    ∂qb∂Δω = szeros(T, 4, input_dimension(joint.rotational))

    ∂vb∂Δx = 1 / timestep * rotation_matrix(qa) * Atra
    ∂vb∂Δx += -1 / timestep * rotation_matrix(qa1) * Atra

    ∂vb∂Δθ = -1.0 / timestep * ∂vector_rotate∂q(pb, qb) * Lmat(qa * orientation_offset) * daxis_angle_to_quaterniondx(Arot * Δθ) * Arot
    ∂vb∂Δθ += 1.0 / timestep * ∂vector_rotate∂q(pb, qb1) * Rmat(inv(axis_angle_to_quaternion(Arot * Δω * timestep))) * Lmat(qa1 * orientation_offset) * daxis_angle_to_quaterniondx(Arot * Δθ) * Arot

    ∂vb∂Δv = rotation_matrix(qa1) * Atra

    ∂vb∂Δω = ∂vector_rotate∂q(pb, qb1) * Lmat(qa1 * orientation_offset * Δq) * Tmat() * daxis_angle_to_quaterniondx(Arot * Δω * timestep) * Arot

    ∂ωb∂Δx = szeros(T, 3, input_dimension(joint.translational))

    ∂ωb∂Δθ = ∂angular_velocity∂q1(qb1, qb, timestep) * Rmat(inv(axis_angle_to_quaternion(Arot * Δω * timestep))) * Lmat(qa1 * orientation_offset) * daxis_angle_to_quaterniondx(Arot * Δθ) * Arot
    ∂ωb∂Δθ += ∂angular_velocity∂q2(qb1, qb, timestep) * Lmat(qa * orientation_offset) * daxis_angle_to_quaterniondx(Arot * Δθ) * Arot

    ∂ωb∂Δv = szeros(3, input_dimension(joint.translational))

    ∂ωb∂Δω = ∂angular_velocity∂q1(qb1, qb, timestep) * Lmat(qa1 * orientation_offset * Δq) * Tmat() * daxis_angle_to_quaterniondx(Arot * Δω * timestep) * Arot * timestep

    [
        ∂xb∂Δx ∂xb∂Δθ ∂xb∂Δv ∂xb∂Δω;
        ∂vb∂Δx ∂vb∂Δθ ∂vb∂Δv ∂vb∂Δω;
        ∂qb∂Δx ∂qb∂Δθ ∂qb∂Δv ∂qb∂Δω;
        ∂ωb∂Δx ∂ωb∂Δθ ∂ωb∂Δv ∂ωb∂Δω;
    ]
end

function set_minimal_coordinates_velocities_jacobian_minimal(joint::JointConstraint,
    pnode::Node, cnode::Node, timestep)

    xmin = minimal_coordinates_velocities(joint, pnode, cnode, timestep)
    qb = current_orientation(cnode.state)
    xa, va, qa, ωa = initial_configuration_velocity(pnode.state)
    za = [xa; va; vector(qa); ωa]

    nu = input_dimension(joint)
    Δx = xmin[SUnitRange(joint.minimal_index[1]...)]
    Δθ = xmin[SUnitRange(joint.minimal_index[2]...)]
    Δv = xmin[nu .+ SUnitRange(joint.minimal_index[1]...)]
    Δω = xmin[nu .+ SUnitRange(joint.minimal_index[2]...)]
    J = minimal_coordinates_velocities_jacobian_minimal(joint,
                xa, va, qa, ωa,
                timestep;
                Δx=Δx,
                Δθ=Δθ,
                Δv=Δv,
                Δω=Δω)
    J = cat(Diagonal(sones(6)), LVᵀmat(qb)', Diagonal(sones(3)), dims=(1,2)) * J
    return J
end


text:
function exploration(ram::Net, raw_locations)
    sz = k = ram.glimpse_net.retina.patch_size
    locations = map(l->denormalize(28,l), raw_locations)
    locations = map(l->reshape(l, 2, :, 1), locations)
    locations = cat(locations..., dims=3)
    from_x = Int.(round.(locations[1,:,:] .- div(k,2) .+ 1))
    from_y = Int.(round.(locations[2,:,:] .- div(k,2) .+ 1))
    to_x, to_y = from_x .+ sz .- 1, from_y .+ sz .- 1
    B, T = size(from_x)
    from_x = reshape(from_x, B, T, 1)
    to_x = reshape(to_x, B, T, 1)
    from_y = reshape(from_y, B, T, 1)
    to_y = reshape(to_y, B, T, 1)
    points = cat(from_x, to_x, from_y, to_y, dims=3);
    grid = zeros(Int,28,28,B);
    points = min.(28, points);
    points = max.(1, points);
    for i = 1:B, t = 1:T
        fx,tx,fy,ty = points[i,t,:]
        grid[fy:ty,fx:tx,i] .+= 1
    end
    return grid
end

max_stars_repo_name: ilkerkesen/RAM

max_issues_repo_name: ilkerkesen/RAM

max_forks_repo_name: ilkerkesen/RAM


text:
# This is an example on training the SCAMPy implementation of the EDMF
# scheme with data generated using PyCLES.
#
# The example seeks to find the optimal values of the entrainment and
# detrainment parameters of the EDMF scheme to replicate the LES profiles
# of the BOMEX experiment.
#
# This example is fully parallelized and can be run in the cluster with
# the included script.

# Import modules to all processes
@everywhere using Pkg
@everywhere Pkg.activate("../..")
@everywhere using Distributions
@everywhere using StatsBase
@everywhere using LinearAlgebra
@everywhere using BlockDiagonals
# Import EKP modules
@everywhere using EnsembleKalmanProcesses.EnsembleKalmanProcessModule
@everywhere using EnsembleKalmanProcesses.Observations
@everywhere using EnsembleKalmanProcesses.ParameterDistributionStorage
@everywhere include(joinpath(@__DIR__, "../../src/helper_funcs.jl"))
using JLD2
using NPZ

#########
#########  Define the parameters and their priors
#########

# Define the parameters that we want to learn
param_names = ["entrainment_factor", "detrainment_factor"]
n_param = length(param_names)

# Prior information: Define transform to unconstrained gaussian space
constraints = [ [bounded(0.01, 0.3)],
                [bounded(0.01, 1.0)]]
# All vars are standard Gaussians in unconstrained space
prior_dist = [Parameterized(Normal(0.0, 1.0))
                for x in range(1, n_param, length=n_param) ]
priors = ParameterDistribution(prior_dist, constraints, param_names)

#########
#########  Define simulation parameters and data directories
#########

# Define observation window (s)
t_starts = [4.0] * 3600  # 4hrs
t_ends = [6.0] * 3600  # 6hrs
# Define variables considered in the loss function
y_names = Array{String, 1}[]
push!(y_names, ["thetal_mean", "ql_mean", "qt_mean", "total_flux_h", "total_flux_qt"])

# Define preconditioning and regularization of inverse problem
perform_PCA = true # Performs PCA on data

# Define directories to fetch data from and store results to
les_names = ["Bomex"]
les_suffixes = ["may18"]
les_root = "/groups/esm/ilopezgo"
scm_names = ["Bomex"]  # same as `les_names` in perfect model setting
scm_data_root = pwd()  # path to folder with `Output.<scm_name>.00000` files
scampy_dir = joinpath(pwd(), "SCAMPy")  # path to SCAMPy
save_full_EDMF_data = false  # if true, save each ensemble output file
outdir_root = pwd()

@assert isdir(les_root) & isdir(scm_data_root) & isdir(scampy_dir)

#########
#########  Retrieve true LES samples from PyCLES data and transform
#########

# Init arrays
yt = zeros(0)
yt_var_list = Array{Float64, 2}[]
yt_big = zeros(0)
yt_var_list_big = Array{Float64, 2}[]
P_pca_list = []
pool_var_list = []
@assert (  # Each entry in these lists correspond to one simulation case
    length(les_names) == length(les_suffixes) == length(scm_names)
    == length(y_names) == length(t_starts) == length(t_ends)
)
for (les_name, les_suffix, scm_name, y_name, tstart, tend) in zip(
        les_names, les_suffixes, scm_names, y_names, t_starts, t_ends
    )
    # Get SCM vertical levels for interpolation
    z_scm = get_profile(joinpath(scm_data_root, "Output.$scm_name.00000"), ["z_half"])
    # Get (interpolated and pool-normalized) observations, get pool variance vector
    les_dir = joinpath(les_root, "Output.$les_name.$les_suffix")
    yt_, yt_var_, pool_var = obs_LES(y_name, les_dir, tstart, tend, z_scm = z_scm)
    push!(pool_var_list, pool_var)
    if perform_PCA
        yt_pca, yt_var_pca, P_pca = obs_PCA(yt_, yt_var_)
        append!(yt, yt_pca)
        push!(yt_var_list, yt_var_pca)
        push!(P_pca_list, P_pca)
    else
        append!(yt, yt_)
        push!(yt_var_list, yt_var_)
        global P_pca_list = nothing
    end
    # Save full dimensionality (normalized) output for error computation
    append!(yt_big, yt_)
    push!(yt_var_list_big, yt_var_)
end
d = length(yt) # Length of data array

# Construct global observational covariance matrix, TSVD
Γy = Matrix(BlockDiagonal(yt_var_list))

yt_var_big = cov_from_cov_list(yt_var_list_big)

n_samples = 1
samples = zeros(n_samples, length(yt))
samples[1,:] = yt

#########
#########  Calibrate: Ensemble Kalman Inversion
#########

algo = Inversion() # Sampler(vcat(get_mean(priors)...), get_cov(priors))
N_ens = 20 # number of ensemble members
N_iter = 10 # number of EKP iterations.
Δt = 1.0 # Artificial time stepper of the EKI.
println("NUMBER OF ENSEMBLE MEMBERS: $N_ens")
println("NUMBER OF ITERATIONS: $N_iter")

initial_params = construct_initial_ensemble(priors, N_ens, rng_seed=rand(1:1000))
ekobj = EnsembleKalmanProcess(initial_params, yt, Γy, algo )

# Define caller function
@everywhere g_(x::Array{Float64,1}) = run_SCAMPy(
    x, $param_names, $y_names, $scampy_dir,
    $scm_data_root, $scm_names, $t_starts, $t_ends,
    P_pca_list = $P_pca_list, norm_var_list = $pool_var_list,
)

# Create output dir
algo_type = typeof(algo) == Sampler{Float64} ? "eks" : "eki"
outdir_path = joinpath(outdir_root, "results_$(algo_type)_dt$(Δt)_p$(n_param)_e$(N_ens)_i$(N_iter)_d$d")
println("Name of outdir path for this EKP is: $outdir_path")
mkpath(outdir_path)

# EKP iterations
g_ens = zeros(N_ens, d)
norm_err_list = []
g_big_list = []
for i in 1:N_iter
    # Note that the parameters are transformed when used as input to SCAMPy
    params_cons_i = deepcopy(transform_unconstrained_to_constrained(priors,
        get_u_final(ekobj)) )
    params = [row[:] for row in eachrow(params_cons_i')]
    @everywhere params = $params
    array_of_tuples = pmap(g_, params) # Outer dim is params iterator
    (sim_dirs_arr, g_ens_arr, g_ens_arr_pca) = ntuple(l->getindex.(array_of_tuples,l),3) # Outer dim is G̃, G
    println(string("\n\nEKP evaluation $i finished. Updating ensemble ...\n"))
    for j in 1:N_ens
        g_ens[j, :] = g_ens_arr_pca[j]
    end
    # Get normalized error for full dimensionality output
    push!(norm_err_list, compute_errors(g_ens_arr, yt_big))
    # Store full dimensionality output
    push!(g_big_list, g_ens_arr)
    # Get normalized error
    if typeof(algo) != Sampler{Float64}
        update_ensemble!(ekobj, Array(g_ens') , Δt_new=Δt)
    else
        update_ensemble!(ekobj, Array(g_ens') )
    end
    println("\nEnsemble updated. Saving results to file...\n")
    # Convert to arrays
    phi_params = Array{Array{Float64,2},1}(transform_unconstrained_to_constrained(priors, get_u(ekobj)))
    phi_params_arr = zeros(i+1, n_param, N_ens)
    for (k,elem) in enumerate(phi_params)
        phi_params_arr[k,:,:] = elem
    end

    # Save EKP information to JLD2 file
    save(joinpath(outdir_path, "ekp.jld2"),
        "ekp_u", transform_unconstrained_to_constrained(priors, get_u(ekobj)),
        "ekp_g", get_g(ekobj),
        "truth_mean", ekobj.obs_mean,
        "truth_cov", ekobj.obs_noise_cov,
        "ekp_err", ekobj.err,
        "truth_mean_big", yt_big,
        "truth_cov_big", yt_var_big,
        "g_big", g_big_list,
        "norm_err", norm_err_list,
        "P_pca", P_pca_list,
        "pool_var", pool_var_list,
        "phi_params", phi_params_arr,
    )
    # Convert to arrays
    phi_params = Array{Array{Float64,2},1}(transform_unconstrained_to_constrained(priors, get_u(ekobj)))
    phi_params_arr = zeros(i+1, n_param, N_ens)
    g_big_arr = zeros(i, N_ens, length(yt_big))
    for (k,elem) in enumerate(phi_params)
        phi_params_arr[k,:,:] = elem
        if k < i + 1
            g_big_arr[k,:,:] = hcat(g_big_list[k]...)'
        end
    end
    norm_err_arr = hcat(norm_err_list...)' # N_iter, N_ens
    # Or you can also save information to numpy files with NPZ
    npzwrite(joinpath(outdir_path,"y_mean.npy"), ekobj.obs_mean)
    npzwrite(joinpath(outdir_path,"Gamma_y.npy"), ekobj.obs_noise_cov)
    npzwrite(joinpath(outdir_path,"ekp_err.npy"), ekobj.err)
    npzwrite(joinpath(outdir_path,"phi_params.npy"), phi_params_arr)
    npzwrite(joinpath(outdir_path,"y_mean_big.npy"), yt_big)
    npzwrite(joinpath(outdir_path,"Gamma_y_big.npy"), yt_var_big)
    npzwrite(joinpath(outdir_path,"norm_err.npy"), norm_err_arr)
    npzwrite(joinpath(outdir_path,"g_big.npy"), g_big_arr)
    for (l, P_pca) in enumerate(P_pca_list)
        npzwrite(joinpath(outdir_path,"P_pca_$(scm_names[l]).npy"), P_pca)
        npzwrite(joinpath(outdir_path,"pool_var_$(scm_names[l]).npy"), pool_var_list[l])
    end

    if save_full_EDMF_data
        # Save full EDMF data from every ensemble
        eki_iter_path = joinpath(outdir_path, "EKI_iter_$i")
        mkpath(eki_iter_path)
        # get a simulation directory `.../Output.SimName.UUID`, and corresponding parameter name
        for (ens_i, sim_dirs) in enumerate(sim_dirs_arr)  # each ensemble returns a list of simulation directories
            ens_i_path = joinpath(eki_iter_path, "ens_$ens_i")
            mkpath(ens_i_path)
            for (scm_name, sim_dir) in zip(scm_names, sim_dirs)
                # Copy simulation data to output directory
                dirname = splitpath(sim_dir)[end]
                @assert dirname[1:7] == "Output."  # sanity check
                # Stats file
                tmp_data_path = joinpath(sim_dir, "stats/Stats.$scm_name.nc")
                save_data_path = joinpath(ens_i_path, "Stats.$scm_name.$ens_i.nc")
                run(`cp $tmp_data_path $save_data_path`)
                # namefile and paramfile
                tmp_namefile_path = joinpath(sim_dir, "$scm_name.in")
                save_namefile_path = joinpath(ens_i_path, "$scm_name.in")
                run(`cp $tmp_namefile_path $save_namefile_path`)
                tmp_paramfile_path = joinpath(sim_dir, "paramlist_$scm_name.in")
                save_paramfile_path = joinpath(ens_i_path, "paramlist_$scm_name.in")
                run(`cp $tmp_paramfile_path $save_paramfile_path`)
            end
        end
    end
end
# EKP results: Has the ensemble collapsed toward the truth?
println("\nEKP ensemble mean at last stage (original space):")
println( mean( transform_unconstrained_to_constrained(priors, get_u_final(ekobj)), dims=2) ) # Parameters are stored as columns


text:
function __init__iris()
    DEPNAME = "Iris"
    LINK = "https://archive.ics.uci.edu/ml/machine-learning-databases/iris/"
    DOCS = "https://archive.ics.uci.edu/ml/datasets/Iris"
    DATA = "iris.data"

    register(DataDep(
        DEPNAME,
        """
        Dataset: The Iris dataset
        Website: $DOCS
        """,
        LINK .* [DATA],
        "6f608b71a7317216319b4d27b4d9bc84e6abd734eda7872b71a458569e2656c0"  # if checksum omitted, will be generated by DataDeps
    ))
end

"""
    Iris(; as_df = true, dir = nothing)

Fisher's classic iris dataset.

Measurements from 3 different species of iris: setosa, versicolor and
virginica. There are 50 examples of each species.

There are 4 measurements for each example: sepal length, sepal width, petal
length and petal width.  The measurements are in centimeters.

The module retrieves the data from the [UCI Machine Learning Repository](https://archive.ics.uci.edu/ml/datasets/iris).

NOTE: no pre-defined train-test split for this dataset.

# Arguments

$ARGUMENTS_SUPERVISED_TABLE

# Fields

$FIELDS_SUPERVISED_TABLE

# Methods

$METHODS_SUPERVISED_TABLE

# Examples

```julia-repl
julia> dataset = Iris()
Iris:
  metadata => Dict{String, Any} with 4 entries
  features => 150×4 DataFrame
  targets => 150×1 DataFrame
  dataframe => 150×5 DataFrame

julia> dataset[1:2]
(2×4 DataFrame
 Row │ sepallength  sepalwidth  petallength  petalwidth
     │ Float64      Float64     Float64      Float64
─────┼──────────────────────────────────────────────────
   1 │         5.1         3.5          1.4         0.2
   2 │         4.9         3.0          1.4         0.2, 2×1 DataFrame
 Row │ class
     │ String15
─────┼─────────────
   1 │ Iris-setosa
   2 │ Iris-setosa)

julia> X, y = Iris(as_df=false)[]
([5.1 4.9 … 6.2 5.9; 3.5 3.0 … 3.4 3.0; 1.4 1.4 … 5.4 5.1; 0.2 0.2 … 2.3 1.8], InlineStrings.String15["Iris-setosa" "Iris-setosa" … "Iris-virginica" "Iris-virginica"])
```
"""
struct Iris <: SupervisedDataset
    metadata::Dict{String, Any}
    features
    targets
    dataframe
end

function Iris(; dir = nothing, as_df = true)
    path = datafile("Iris", "iris.data", dir)
    df = read_csv(path, header=0)
    rename!(df, ["sepallength", "sepalwidth", "petallength", "petalwidth", "class"])

    features = df[!, Not(:class)]
    targets = df[!, [:class]]

    metadata = Dict{String, Any}()
    metadata["path"] = path
    metadata["n_observations"] = size(df, 1)
    metadata["feature_names"] = names(features)
    metadata["target_names"] = names(targets)

    if !as_df
        features = df_to_matrix(features)
        targets = df_to_matrix(targets)
        df = nothing
    end

    return Iris(metadata, features, targets, df)
end

# deprecated in v0.6
function Base.getproperty(::Type{Iris}, s::Symbol)
    if s == :features
        @warn "Iris.features() is deprecated, use `Iris().features` instead."
        return () -> Iris(as_df=false).features
    elseif s == :labels
        @warn "Iris.labels() is deprecated, use `Iris().targets` instead."
        return () -> Iris(as_df=false).targets |> vec
    else
        return getfield(Iris, s)
    end
end


text:
########## Parameters for scripts ##########

######## parameters for loading data ########
parms_io = (
    intraday = true,  # true for intraday data
    # date range of data
    date_start = Dates.DateTime("0000-01-01"), date_end = Dates.DateTime("9999-01-01"),
    # daily time range
    wa = Dates.Hour(9) + Dates.Minute(30), wb = Dates.Hour(15) + Dates.Minute(59),
    adjusted = false,  # adjust intraday data

    # index of the column in the dataframe to work on
    column_index = 1,
)

max_stars_repo_name: yanncalec/FracFin

max_issues_repo_name: yanncalec/FracFin

max_forks_repo_name: yanncalec/FracFin


text:
###############################################################################
# REGRESSORS

KNNRegressor = @load KNNRegressor pkg=NearestNeighborModels verbosity=0
LinearRegressor = @load LinearRegressor pkg=MLJLinearModels verbosity=0
XGBoostRegressor = @load XGBoostRegressor pkg=XGBoost verbosity=0

###############################################################################
# CLASSIFIERS

LogisticClassifier = @load LogisticClassifier pkg=MLJLinearModels verbosity=0
KNNClassifier = @load KNNClassifier pkg=NearestNeighborModels verbosity=0
XGBoostClassifier = @load XGBoostClassifier pkg=XGBoost verbosity=0

###############################################################################
# INTERACTIONS LINEAR MODEL

mutable struct InteractionTransformer <: Unsupervised
    column_pattern
end

function MLJ.fit(model::InteractionTransformer, verbosity::Int, X)
    matching_columns = filter(
        x -> occursin(model.column_pattern, string(x)),
        Tables.columnnames(X)
        )
    n = length(matching_columns)
    interaction_pairs = Vector{Pair{Symbol, Symbol}}(undef, n*(n-1)÷2)
    index = 1
    for i in 1:n-1
        for j in i+1:n
            interaction_pairs[index] = matching_columns[i] => matching_columns[j]
            index += 1
        end
    end
    fitresult = (interaction_pairs=interaction_pairs, ninter=length(interaction_pairs))
    return fitresult, nothing, nothing
end

function MLJ.transform(model::InteractionTransformer, fitresult, X)
    n = nrows(X)
    interactions = Matrix{Float64}(undef, n, fitresult.ninter)
    index = 1
    for (col₁, col₂) in fitresult.interaction_pairs
        interactions[:, index] = Tables.getcolumn(X, col₁) .* Tables.getcolumn(X, col₂)
        index += 1
    end
    names = Tuple(Symbol(string(col₁, "_", col₂)) for (col₁, col₂) in fitresult.interaction_pairs)
    interactions_ndt = NamedTuple{names}([interactions[:, i] for i in 1:fitresult.ninter])
    return merge(Tables.columntable(X), interactions_ndt)
end

struct InteractionLMClassifier <: MLJ.ProbabilisticComposite
    interaction_transformer::InteractionTransformer
    linear_model
end
struct InteractionLMRegressor <: MLJ.DeterministicComposite
    interaction_transformer::InteractionTransformer
    linear_model
end

"""
Unfortunately it seems the definition of a pipeline in the module poses problems.
I thus resort to the definition of this trivial learning-network while a better API for pipelines
if not available (Ongoing work as of now: https://github.com/JuliaAI/MLJBase.jl/pull/664).
"""
InteractionLM = Union{InteractionLMClassifier, InteractionLMRegressor}

InteractionLMClassifier(;column_pattern="^rs[0-9]+", kwargs...) =
    InteractionLMClassifier(InteractionTransformer(Regex(column_pattern)), LogisticClassifier(kwargs...))

InteractionLMRegressor(;column_pattern="^rs[0-9]+", kwargs...) =
    InteractionLMRegressor(InteractionTransformer(Regex(column_pattern)), LinearRegressor(kwargs...))

function MLJ.fit(model::InteractionLM, verbosity::Int, X, y)
    Xs = source(X)
    ys = source(y)

    inter_mach = machine(model.interaction_transformer, Xs)
    Xt = MLJ.transform(inter_mach, Xs)

    pred_mach = machine(model.linear_model, Xt, ys)
    ŷ = MLJ.predict(pred_mach, Xt)

    mach = machine(supertype(supertype(typeof(model)))(), Xs, ys; predict=ŷ)

	return!(mach, model, verbosity)
end

MLJ.target_scitype(model::InteractionLMRegressor) = AbstractVector{<:MLJ.Continuous}
MLJ.target_scitype(model::InteractionLMClassifier) = AbstractVector{<:Finite}