using Colors

Base.mean{C<:Color}(c1::C,c2::C) = convert(C, (convert(XYZ, c1) + convert(XYZ, c2))/2)  # colorimetrically "proper"

# RGB linear vector space averaging
function meanrgb{C<:AbstractRGB}(c1::C, c2::C)
    C((red(c1)+red(c2))/2,
      (green(c1)+green(c2))/2,
      (blue(c1)+blue(c2))/2)
end

n = 10
colors = distinguishable_colors(n, Array(RGB{Float32},0))
m = 50 # number of repeated pixels in each block
grid = Array(eltype(colors), n*m, n*m)
z = zero(RGB{Float32})
for i = 1:n
    for j = i:n   # i == j is a sanity check
        prop = mean(colors[i], colors[j])
        vs = meanrgb(colors[i], colors[j])
        for l = 1:m, k = 1:m
            grid[(i-1)*m+k,(j-1)*m+l] = grid[(j-1)*m+k,(i-1)*m+l] = k > l ? prop : vs
        end
    end
end
# Put a boundary around the diagonal elements
for i = 1:n
    for k = 1:m
        grid[(i-1)*m+k,(i-1)*m+1] = grid[(i-1)*m+k,(i-1)*m+m] = z
        grid[(i-1)*m+1,(i-1)*m+k] = grid[(i-1)*m+m,(i-1)*m+k] = z
    end
end

using Images

gridu8 = map(Clamp(RGB{U8}), grid)
img = Image(gridu8, spatialorder=["x","y"])
imwrite(img, "comparison.png")


module DWT
export
    # WPD
    wpd,
    wpd!,
    iwpd,
    iwpd!,
    # Transform on all signals
    dwtall,
    wptall,
    wpdall,
    idwtall,
    iwptall,
    iwpdall

using
    Wavelets

using
    ..Utils

# ========== Wavelet Packet Decomposition ==========
# 1D Wavelet Packet Decomposition without allocated output array
"""
    wpd(x, wt[, L; standard])

Returns the wavelet packet decomposition (WPD) for L levels for input signal x.

!!! note
    Nonstandard transform is not yet available for 2D-WPD. Setting `standard = true`
    will currently result in an error message.

# Arguments
- `x::AbstractVector{T} where T<:Number` or `x::AbstractArray{T,2} where T<:Number`: Input
  vector/matrix. A vector input undergoes 1D wavelet decomposition whereas a matrix input
  undergoes 2D wavelet decomposition.
- `wt::OrthoFilter`: Wavelet filter.
- `L::Integer`: (Default: `maxtransformlevels(x)`) Number of levels for wavelet
  decomposition.

# Keyword Arguments
- `standard::Bool`: (Default: `true`) Whether to compute the standard or non-standard
  wavelet transform. Only applicable for 2D signals.

# Returns
- `::Array{T,2}` or `::Array{T,3}`: Decomposed signal. For an input vector `x`, output is a
  2D matrix where each column corresponds to a level of decomposition. For an input matrix
  `x`, output is a 3D array where each slice of dimension 3 corresponds to a level of
  decomposition.

# Examples:
```julia
using Wavelets, WaveletsExt

# 1D wavelet decomposition
x = randn(8)
wt = wavelet(WT.haar)
xw = wpd(x, wt)

# 2D wavelet decomposition
x = randn(8,8)
wt = wavelet(WT.haar)
xw = wpd(x, wt)
```

**See also:** [`wpd!`](@ref), [`iwpd`](@ref)
"""
function wpd(x::AbstractVector{T},
             wt::OrthoFilter,
             L::Integer=maxtransformlevels(x)) where T<:Number
    # Sanity check
    @assert isdyadic(x)
    @assert 0 ≤ L ≤ maxtransformlevels(x)

    # Allocate variable for result
    n = length(x)
    xw = Array{T,2}(undef, (n,L+1))
    # Wavelet packet decomposition
    wpd!(xw, x, wt, L)
    return xw
end

# 2D Wavelet Packet Decomposition without allocated output array
function wpd(x::AbstractArray{T,2},
             wt::OrthoFilter,
             L::Integer = maxtransformlevels(x);
             standard::Bool = true) where T<:Number
    # Sanity check
    @assert 0 ≤ L ≤ maxtransformlevels(x)

    # Allocate variable for result
    sz = size(x)
    xw = Array{T,3}(undef, (sz...,L+1))
    # Wavelet packet decomposition
    wpd!(xw, x, wt, L, standard=standard)
    return xw
end

# 1D Wavelet Packet Decomposition with allocated output array
"""
    wpd!(y, x, wt[, L; standard])

Same as `wpd` but without array allocation.

!!! note
    Nonstandard transform is not yet available for 2D-WPD. Setting `standard = true` will
    currently result in an error message.

# Arguments
- `y::AbstractArray{T,2} where T<:Number` or `y::AbstractArray{T,3} where T<:Number`: An
  allocated array to write the outputs of `x` onto.
- `x::AbstractVector{T} where T<:Number` or `x::AbstractArray{T,2} where T<:Number`: Input
  vector/matrix. A vector input undergoes 1D wavelet decomposition whereas a matrix input
  undergoes 2D wavelet decomposition.
- `wt::OrthoFilter`: Wavelet filter.
- `L::Integer`: (Default: `maxtransformlevels(x)`) Number of levels for wavelet
  decomposition.

# Keyword Arguments
- `standard::Bool`: (Default: `true`) Whether to compute the standard or non-standard
  wavelet transform. Only applicable for 2D signals.

# Returns
- `y::AbstractArray{T,2} where T<:Number` or `y::AbstractArray{T,3} where T<:Number`:
  Decomposed signal. For an input vector, output is a 2D matrix where each column
  corresponds to a level of decomposition. For an input matrix, output is a 3D array where
  each slice of dimension 3 corresponds to a level of decomposition.

# Examples:
```
using Wavelets, WaveletsExt

# 1D wavelet decomposition
x = randn(8)
xw = Array{eltype(x)}(undef, (8,3))
wt = wavelet(WT.haar)
wpd!(xw, x, wt)

# 2D wavelet decomposition
x = randn(8,8)
xw = Array{eltype(x)}(undef, (8,8,3))
wt = wavelet(WT.haar)
wpd!(xw, x, wt)
```

**See also:** [`wpd`](@ref), [`iwpd`](@ref)
"""
function wpd!(y::AbstractArray{T,2},
              x::AbstractArray{T,1},
              wt::DiscreteWavelet,
              L::Integer=maxtransformlevels(x)) where T<:Number
    # Sanity check
    n = length(x)
    @assert 0 ≤ L ≤ maxtransformlevels(x)
    @assert size(y) == (n,L+1)

    # Setup
    g, h = WT.makereverseqmfpair(wt, true)
    y[:,1] = x

    # Wavelet Decomposition
    for i in 0:(L-1)
        nₚ = nodelength(n, i)                   # Parent node length
        for j in 0:((1<<i)-1)
            colₚ = i+1                          # Parent column
            rngₚ = (j*nₚ+1):((j+1)*nₚ)           # Parent range
            @inbounds v = @view y[rngₚ, colₚ]    # Parent node
            colᵣ = colₚ+1                       # Child column
            nᵣ = nₚ÷2                           # Child node length
            rng₁ = (2*j*nᵣ+1):((2*j+1)*nᵣ)      # Approx coef range
            rng₂ = ((2*j+1)*nᵣ+1):(2*(j+1)*nᵣ)  # Detail coef range
            @inbounds w₁ = @view y[rng₁, colᵣ]  # Approx coef node
            @inbounds w₂ = @view y[rng₂, colᵣ]  # Detail coef node
            dwt_step!(w₁, w₂, v, h, g)          # Decompose
        end
    end
    return y
end

# 2D Wavelet Packet Decomposition with allocated output array
function wpd!(y::AbstractArray{T,3},
              x::AbstractArray{T,2},
              wt::OrthoFilter,
              L::Integer = maxtransformlevels(x);
              standard::Bool = true) where T<:Number
    # Sanity check
    m, n = size(x)
    @assert 0 ≤ L ≤ maxtransformlevels(x)
    @assert size(y) == (m, n,L+1)

    # Setup
    g, h = WT.makereverseqmfpair(wt, true)   # low & high pass filters
    temp = Array{T,2}(undef, (m,n))

    # First slice of y is level 0, ie. original signal
    @inbounds y[:,:,1] = x
    # ----- Compute L levels of decomposition -----
    for i in 0:(L-1)
        # Parent node width and height
        mₚ = nodelength(m, i)
        nₚ = nodelength(n, i)
        # Iterate over each nodes at current level
        lrange = 0:((1<<i)-1)
        for j in lrange
            for k in lrange
                # Extract parent node
                sliceₚ = i+1
                rng_row = (j*mₚ+1):((j+1)*mₚ)
                rng_col = (k*nₚ+1):((k+1)*nₚ)
                @inbounds v = @view y[rng_row, rng_col, sliceₚ]
                # Extract children nodes of current parent
                sliceᵣ = sliceₚ+1
                mᵣ = mₚ÷2
                nᵣ = nₚ÷2
                @inbounds w₁ = @view y[(2*j*mᵣ+1):((2*j+1)*mᵣ), (2*k*nᵣ+1):((2*k+1)*nᵣ), sliceᵣ]
                @inbounds w₂ = @view y[(2*j*mᵣ+1):((2*j+1)*mᵣ), ((2*k+1)*nᵣ+1):(2*(k+1)*nᵣ), sliceᵣ]
                @inbounds w₃ = @view y[((2*j+1)*mᵣ+1):(2*(j+1)*mᵣ), (2*k*nᵣ+1):((2*k+1)*nᵣ), sliceᵣ]
                @inbounds w₄ = @view y[((2*j+1)*mᵣ+1):(2*(j+1)*mᵣ), ((2*k+1)*nᵣ+1):(2*(k+1)*nᵣ), sliceᵣ]
                # Extract temp subarray (Same size as parent node)
                @inbounds tempₖ = @view temp[rng_row, rng_col]
                # Perform 1 level of wavelet decomposition
                @inbounds dwt_step!(w₁, w₂, w₃, w₄, v, h, g, tempₖ, standard=standard)
            end
        end
    end
    return y
end

# ========== Inverse Wavelet Packet Decomposition ==========
# IWPD by level without allocation
"""
    iwpd(xw, wt[, L; standard])
    iwpd(xw, wt, tree[; standard])

Computes the inverse wavelet packet decomposition (IWPD) for `L` levels or by given `tree.`

!!! note
    Nonstandard transform is not yet available. Setting `standard = true` will
    currently result in an error message.

# Arguments
- `x̂::AbstractVector{T} where T<:Number` or `x̂::AbstractArray{T,2} where T<:Number`:
  Allocated output vector/matrix for reconstructed signal.
- `xw::AbstractArray{T,2} where T<:Number` or `xw::AbstractArray{T,3} where T<:Number`:
  Input array. A 2D input undergoes 1D wavelet reconstruction whereas a 3D input undergoes
  2D wavelet reconstruction.
- `wt::OrthoFilter`: Wavelet filter.
- `L::Integer`: (Default: `maxtransformlevels(x)`) Number of levels of wavelet
  decomposition.
- `tree::BitVector`: Binary tree or Quadtree for inverse transform to be computed
  accordingly.

# Keyword Arguments
- `standard::Bool`: (Default: `true`) Whether to compute the standard or non-standard
  wavelet transform. Only applicable for 2D signals.

# Returns
- `x̂::Array{T,1}` or `x̂::Array{T,2}`: Reconstructed signal.

# Examples:
```
using Wavelets, WaveletsExt

# 1D wavelet decomposition
x = randn(8)
wt = wavelet(WT.haar)
xw = wpd(x, wt)

# 1D wavelet reconstruction
y = iwpd(xw, wt)

# 2D wavelet decomposition
x = randn(8,8)
wt = wavelet(WT.haar)
xw = wpd(x, wt)

# 2D wavelet reconstruction
y = iwpd(xw, wt)
```

**See also:** [`iwpd!`](@ref), [`iwpt`](@ref)
"""
function iwpd(xw::AbstractArray{T},
              wt::OrthoFilter,
              L::Integer = maxtransformlevels(xw,1);
              kwargs...) where T<:Number
    @assert ndims(xw) ≥ 2
    dim = size(xw)[1:(end-1)]
    x̂ = Array{T}(undef, dim)
    iwpd!(x̂, xw, wt, L, kwargs...)
    return x̂
end

# IWPD by tree without allocation
function iwpd(xw::AbstractArray{T},
              wt::OrthoFilter,
              tree::BitVector;
              kwargs...) where T<:Number
    @assert ndims(xw) ≥ 2
    dim = size(xw)[1:(end-1)]
    x̂ = Array{T}(undef, dim)
    iwpd!(x̂, xw, wt, tree, kwargs...)
    return x̂
end

# 1D IWPD by level with allocation
"""
    iwpd(x̂, xw, wt[, L; standard])
    iwpd(x̂, xw, wt, tree[; standard])

Same as `iwpd` but without array allocation.

!!! note
    Nonstandard transform is not yet available. Setting `standard = true` will
    currently result in an error message.

# Arguments
- `x̂::AbstractVector{T} where T<:Number` or `x̂::AbstractArray{T,2} where T<:Number`:
  Allocated output vector/matrix for reconstructed signal.
- `xw::AbstractArray{T,2} where T<:Number` or `xw::AbstractArray{T,3} where T<:Number`:
  Input array. A 2D input undergoes 1D wavelet reconstruction whereas a 3D input undergoes
  2D wavelet reconstruction.
- `wt::OrthoFilter`: Wavelet filter.
- `L::Integer`: (Default: `maxtransformlevels(x)`) Number of levels for wavelet
  decomposition.
- `tree::BitVector`: Binary tree or Quadtree to transform to be computed accordingly.

# Keyword Arguments
- `standard::Bool`: (Default: `true`) Whether to compute the standard or non-standard
  wavelet transform. Only applicable for 2D signals.

# Returns
- `x̂::Array{T,1}` or `x̂::Array{T,2}`: Reconstructed signal.

# Examples:
```
using Wavelets, WaveletsExt

# 1D wavelet decomposition
x = randn(8)
wt = wavelet(WT.haar)
xw = wpd(x, wt)
y = similar(x)

# 1D wavelet reconstruction
iwpd!(y, xw, wt)

# 2D wavelet decomposition
x = randn(8,8)
wt = wavelet(WT.haar)
xw = wpd(x, wt)
y = similar(x)

# 2D wavelet reconstruction
iwpd!(y, xw, wt)
```

**See also:** [`iwpd`](@ref), [`iwpt`](@ref)
"""
function iwpd!(x̂::AbstractVector{T},
               xw::AbstractArray{T,2},
               wt::OrthoFilter,
               L::Integer = maxtransformlevels(x̂)) where T<:Number
    iwpd!(x̂, xw, wt, maketree(length(x̂), L, :full))
    return x̂
end

# 2D IWPD by level with allocation
function iwpd!(x̂::AbstractArray{T,2},
               xw::AbstractArray{T,3},
               wt::OrthoFilter,
               L::Integer = maxtransformlevels(x̂);
               standard::Bool = true) where T<:Number
    iwpd!(x̂, xw, wt, maketree(size(x̂)..., L, :full), standard=standard)
    return x̂
end

# 1D IWPD by tree with allocation
function iwpd!(x̂::AbstractVector{T},
               xw::AbstractArray{T,2},
               wt::OrthoFilter,
               tree::BitVector) where T<:Number
    # Sanity check
    @assert size(x̂,1) == size(xw,1)
    @assert isvalidtree(x̂, tree)
    # Get basis coefficients then compute iwpt
    w = getbasiscoef(xw, tree)
    iwpt!(x̂, w, wt, tree)
    return x̂
end

# 2D IWPD by tree with allocation
function iwpd!(x̂::AbstractArray{T,2},
               xw::AbstractArray{T,3},
               wt::OrthoFilter,
               tree::BitVector;
               standard::Bool = true) where T<:Number
    # Sanity check
    @assert size(x̂,1) == size(xw,1)
    @assert size(x̂,2) == size(xw,2)
    @assert isvalidtree(x̂, tree)

    # Setup
    m, n, _ = size(xw)
    g, h = WT.makereverseqmfpair(wt, true)
    xwₜ = copy(xw)
    temp = Array{T,2}(undef, (m,n))

    # Inverse transform
    for i in reverse(eachindex(tree))
        # Reconstruct node i if it has children
        if tree[i]
            # Extract parent node
            d = getdepth(i,:quad)
            rows = getrowrange(m, i)
            cols = getcolrange(n, i)
            @inbounds v = d == 0 ? x̂ : @view xwₜ[rows, cols, d+1]
            # Extract children nodes of current parent
            dᵣ = d+1
            rows₁ = getrowrange(m, getchildindex(i,:topleft))
            cols₁ = getcolrange(n, getchildindex(i,:topleft))
            rows₂ = getrowrange(m, getchildindex(i,:topright))
            cols₂ = getcolrange(n, getchildindex(i,:topright))
            rows₃ = getrowrange(m, getchildindex(i,:bottomleft))
            cols₃ = getcolrange(n, getchildindex(i,:bottomleft))
            rows₄ = getrowrange(m, getchildindex(i,:bottomright))
            cols₄ = getcolrange(n, getchildindex(i,:bottomright))
            @inbounds w₁ = @view xwₜ[rows₁, cols₁, dᵣ+1]
            @inbounds w₂ = @view xwₜ[rows₂, cols₂, dᵣ+1]
            @inbounds w₃ = @view xwₜ[rows₃, cols₃, dᵣ+1]
            @inbounds w₄ = @view xwₜ[rows₄, cols₄, dᵣ+1]
            # Extract temp subarray (Same size as parent node)
            @inbounds tempᵢ = @view temp[rows, cols]
            # Perform 1 level of wavelet reconstruction
            @inbounds idwt_step!(v, w₁, w₂, w₃, w₄, h, g, tempᵢ, standard=standard)
        else
            continue
        end
    end
    return x̂
end

# ========== Wavelet Packet Transform ==========
# 2D Wavelet Packet Transform without allocated output array
"""
    wpt(x, wt[, L; standard])
    wpt(x, wt, tree[; standard])

Returns the wavelet packet transform (WPT) for `L` levels or by given quadratic `tree`.

!!! note
    Nonstandard transform is not yet available. Setting `standard = true` will currently
    result in an error message.

# Arguments
- `x::AbstractVector{T} where T<:Number` or `x::AbstractArray{T,2} where T<:Number`: Input
  vector/matrix. A vector input undergoes 1D wavelet decomposition whereas a matrix input
  undergoes 2D wavelet decomposition.
- `wt::OrthoFilter`: Wavelet filter.
- `L::Integer`: (Default: `maxtransformlevels(x)`) Number of levels for wavelet
  decomposition.
- `tree::BitVector`: Quadtree to transform to be computed accordingly.

# Keyword Arguments
- `standard::Bool`: (Default: `true`) Whether to compute the standard or non-standard
  wavelet transform. Only applicable for 2D signals.

# Returns
- `::Array{T,1}` or `::Array{T,2}`: Transformed signal.

# Examples:
```
using Wavelets, WaveletsExt

# 1D wavelet decomposition
x = randn(8)
wt = wavelet(WT.haar)
xw = wpt(x, wt)

# 2D wavelet decomposition
x = randn(8,8)
wt = wavelet(WT.haar)
xw = wpt(x, wt)
```

**See also:** [`wpt!`](@ref), [`maketree`](@ref)
"""
function Wavelets.Transforms.wpt(x::AbstractArray{T,2},
                                 wt::OrthoFilter,
                                 L::Integer = maxtransformlevels(x);
                                 standard::Bool = true) where T<:Number
    return wpt(x, wt, maketree(size(x)..., L, :full), standard=standard)
end

function Wavelets.Transforms.wpt(x::AbstractArray{T,2},
                                 wt::OrthoFilter,
                                 tree::BitVector;
                                 standard::Bool = true) where T<:Number
    y = Array{T}(undef, size(x))
    return wpt!(y, x, wt, tree, standard=standard)
end

# 2D Wavelet Packet Transform with allocated output array
"""
    wpt!(y, x, wt[, L; standard])
    wpt!(y, x, wt, tree[; standard])

Same as `wpt` but without array allocation.

!!! note
    Nonstandard transform is not yet available. Setting `standard = true` will currently
    result in an error message.

# Arguments
- `y::AbstractVector{T} where T<:Number` or `y::AbstractArray{T,2} where T<:Number`:
  Allocated output vector/matrix.
- `x::AbstractVector{T} where T<:Number` or `x::AbstractArray{T,2} where T<:Number`: Input
  vector/matrix. A vector input undergoes 1D wavelet packet transform whereas a matrix input
  undergoes 2D wavelet packet transform.
- `wt::OrthoFilter`: Wavelet filter.
- `L::Integer`: (Default: `maxtransformlevels(x)`) Number of levels for wavelet
  decomposition.
- `tree::BitVector`: Quadtree to transform to be computed accordingly.

# Keyword Arguments
- `standard::Bool`: (Default: `true`) Whether to compute the standard or non-standard
  wavelet transform. Only applicable for 2D signals.

# Returns
- `::Array{T,1}` or `::Array{T,2}`: Transformed signal.

# Examples:
```
using Wavelets, WaveletsExt

# 1D wavelet decomposition
x = randn(8)
xw = similar(x)
wt = wavelet(WT.haar)
wpt!(xw, x, wt)

# 2D wavelet decomposition
x = randn(8,8)
xw = similar(x)
wt = wavelet(WT.haar)
wpt!(xw, x, wt)
```

**See also:** [`wpt`](@ref), [`maketree`](@ref)
"""
function Wavelets.Transforms.wpt!(y::AbstractArray{T,2},
                                  x::AbstractArray{T,2},
                                  wt::OrthoFilter,
                                  L::Integer = maxtransformlevels(x),
                                  standard::Bool = true) where T<:Number
    return wpt!(y, x, wt, maketree(size(x)..., L, :full), standard=standard)
end

function Wavelets.Transforms.wpt!(y::AbstractArray{T,2},
                                  x::AbstractArray{T,2},
                                  wt::OrthoFilter,
                                  tree::BitVector;
                                  standard::Bool = true) where T<:Number
    # Sanity check
    @assert isvalidtree(x, tree)
    @assert size(y) == size(x)

    # ----- Allocation and setup to match Wavelets.jl's function requirements -----
    m, n = size(x)
    g, h = WT.makereverseqmfpair(wt, true)      # low & high pass filters
    yₜ = copy(x)                                 # Placeholder of y
    temp = Array{T,2}(undef, (m,n))             # Temp array
    # ----- Compute transforms based on tree -----
    for i in eachindex(tree)
        # Decompose if node i has children
        if tree[i]
            # Extract parent node
            rows = getrowrange(m, i)
            cols = getcolrange(n, i)
            @inbounds v = @view yₜ[rows, cols]
            # Extract children nodes of current parent
            rows₁ = getrowrange(m, getchildindex(i,:topleft))
            cols₁ = getcolrange(n, getchildindex(i,:topleft))
            rows₂ = getrowrange(m, getchildindex(i,:topright))
            cols₂ = getcolrange(n, getchildindex(i,:topright))
            rows₃ = getrowrange(m, getchildindex(i,:bottomleft))
            cols₃ = getcolrange(n, getchildindex(i,:bottomleft))
            rows₄ = getrowrange(m, getchildindex(i,:bottomright))
            cols₄ = getcolrange(n, getchildindex(i,:bottomright))
            @inbounds w₁ = @view y[rows₁, cols₁]
            @inbounds w₂ = @view y[rows₂, cols₂]
            @inbounds w₃ = @view y[rows₃, cols₃]
            @inbounds w₄ = @view y[rows₄, cols₄]
            # Extract temp subarray (Same size as parent node)
            @inbounds tempᵢ = @view temp[rows, cols]
            # Perform 1 level of wavelet decomposition
            @inbounds dwt_step!(w₁, w₂, w₃, w₄, v, h, g, tempᵢ, standard=standard)
            # If current range not at final iteration, copy to yₜ
            if 4*i<length(tree)
                @inbounds yₜ[rows, cols] = @view y[rows, cols]
            end
        else
            continue
        end
    end
    return y
end

# ========== Inverse Wavelet Packet Transform ==========
# 2D Inverse Wavelet Packet Transform without allocated output array
"""
    iwpt(xw, wt[, L; standard])
    iwpt(xw, wt, tree[; standard])

Returns the inverse wavelet packet transform (IWPT) for `L` levels or by given quadratic
`tree`.

!!! note
    Nonstandard transform is not yet available. Setting `standard = true` will
    currently result in an error message.

# Arguments
- `xw::AbstractVector{T} where T<:Number` or `xw::AbstractArray{T,2} where T<:Number`: Input
  vector/matrix. A vector input undergoes 1D wavelet reconstruction whereas a matrix input
  undergoes 2D wavelet reconstruction.
- `wt::OrthoFilter`: Wavelet filter.
- `L::Integer`: (Default: `maxtransformlevels(x)`) Number of levels for wavelet
  decomposition.
- `tree::BitVector`: Quadtree to transform to be computed accordingly.

# Keyword Arguments
- `standard::Bool`: (Default: `true`) Whether to compute the standard or non-standard
  wavelet transform. Only applicable for 2D signals.

# Returns
- `::Array{T,1}` or `::Array{T,2}`: Reconstructed signal.

# Examples:
```
using Wavelets, WaveletsExt

# 1D wavelet decomposition
x = randn(8)
wt = wavelet(WT.haar)
xw = wpt(x, wt)

# 1D wavelet reconstruction
y = iwpt(xw, wt)

# 2D wavelet decomposition
x = randn(8,8)
wt = wavelet(WT.haar)
xw = wpt(x, wt)

# 2D wavelet reconstruction
y = iwpt(xw, wt)
```

**See also:** [`iwpt!`](@ref), [`wpt!`](@ref), [`maketree`](@ref)
"""
function Wavelets.Transforms.iwpt(xw::AbstractArray{T,2},
                                  wt::OrthoFilter,
                                  L::Integer = maxtransformlevels(xw);
                                  standard::Bool = true) where T<:Number
    return iwpt(xw, wt, maketree(size(xw)..., L, :full), standard=standard)
end

function Wavelets.Transforms.iwpt(xw::AbstractArray{T,2},
                                  wt::OrthoFilter,
                                  tree::BitVector;
                                  standard::Bool = true) where T<:Number
    x̂ = similar(xw)
    return iwpt!(x̂, xw, wt, tree, standard=standard)
end

# 2D Inverse Wavelet Packet Transform with allocated output array
"""
    iwpt!(x̂, xw, wt[, L; standard])
    iwpt!(x̂, xw, wt, tree[; standard])

Same as `iwpt` but without array allocation.

!!! note
    Nonstandard transform is not yet available. Setting `standard = true` will currently
    result in an error message.

# Arguments
- `x̂::AbstractVector{T} where T<:Number` or `x̂::AbstractArray{T,2} where T<:Number`:
  Allocated output vector/matrix.
- `xw::AbstractVector{T} where T<:Number` or `xw::AbstractArray{T,2} where T<:Number`: Input
  vector/matrix. A vector input undergoes 1D wavelet packet reconstruction whereas a matrix
  input undergoes 2D wavelet packet reconstruction.
- `wt::OrthoFilter`: Wavelet filter.
- `L::Integer`: (Default: `maxtransformlevels(x)`) Number of levels for wavelet
  decomposition.
- `tree::BitVector`: Quadtree to transform to be computed accordingly.
- `standard::Bool`: (Default: `true`) Whether to compute the standard or non-standard
  wavelet transform. Only applicable for 2D signals.

# Returns
- `::Array{T,1}` or `::Array{T,2}`: Reconstructed signal.

# Examples:
```
using Wavelets, WaveletsExt

# 1D wavelet decomposition
x = randn(8)
xw = similar(x)
wt = wavelet(WT.haar)
wpt!(xw, x, wt)

# 1D wavelet reconstruction
y = similar(x)
iwpt!(y, xw, wt)

# 2D wavelet decomposition
x = randn(8,8)
xw = similar(x)
wt = wavelet(WT.haar)
wpt!(xw, x, wt)

# 2D wavelet reconstruction
y = similar(x)
iwpt!(y, xw, wt)
```

**See also:** [`iwpt`](@ref), [`wpt!`](@ref) [`maketree`](@ref)
"""
function Wavelets.Transforms.iwpt!(x̂::AbstractArray{T,2},
                                   xw::AbstractArray{T,2},
                                   wt::OrthoFilter,
                                   L::Integer = maxtransformlevels(xw);
                                   standard::Bool = true) where T<:Number
    return iwpt!(x̂, xw, wt, maketree(size(xw)..., L, :full), standard=standard)
end

function Wavelets.Transforms.iwpt!(x̂::AbstractArray{T,2},
                                   xw::AbstractArray{T,2},
                                   wt::OrthoFilter,
                                   tree::BitVector;
                                   standard::Bool = true) where T<:Number
    # Sanity check
    @assert isvalidtree(xw, tree)
    @assert size(x̂) == size(xw)

    # ----- Setup -----
    m, n = size(xw)
    g, h = WT.makereverseqmfpair(wt, true)      # low & high pass filters
    xwₜ = copy(xw)                               # Placeholder for xw
    temp = Array{T,2}(undef, (m,n))             # temp array
    # ----- Compute transforms based on tree -----
    for i in reverse(eachindex(tree))
        # Reconstruct to node i if it has children
        if tree[i]
            # Extract parent node
            rows = getrowrange(m, i)
            cols = getcolrange(n, i)
            @inbounds v = @view x̂[rows, cols]
            # Extract children nodes of current parent
            rows₁ = getrowrange(m, getchildindex(i,:topleft))
            cols₁ = getcolrange(n, getchildindex(i,:topleft))
            rows₂ = getrowrange(m, getchildindex(i,:topright))
            cols₂ = getcolrange(n, getchildindex(i,:topright))
            rows₃ = getrowrange(m, getchildindex(i,:bottomleft))
            cols₃ = getcolrange(n, getchildindex(i,:bottomleft))
            rows₄ = getrowrange(m, getchildindex(i,:bottomright))
            cols₄ = getcolrange(n, getchildindex(i,:bottomright))
            @inbounds w₁ = @view xwₜ[rows₁, cols₁]
            @inbounds w₂ = @view xwₜ[rows₂, cols₂]
            @inbounds w₃ = @view xwₜ[rows₃, cols₃]
            @inbounds w₄ = @view xwₜ[rows₄, cols₄]
            # Extract temp subarray (Same size as parent node)
            @inbounds tempᵢ = @view temp[rows, cols]
            # Perform 1 level of wavelet reconstruction
            @inbounds idwt_step!(v, w₁, w₂, w₃, w₄, h, g, tempᵢ, standard=standard)
            # If current range not at final iteration, copy to temp
            if i > 1
                @inbounds xwₜ[rows, cols] = v
            end
        else
            continue
        end
    end
    return x̂
end

include("dwt_one_level.jl")
include("dwt_all.jl")

end # end module


export Boussinesq

"""
    Boussinesq(param;kwargs)

Define an object of type `AbstractModel` in view of solving the initial-value problem for
`abcd`-Boussinesq models (with `b=d` and `c=0`).
See [Bona, Chen, and Saut](https://doi.org/10.1007/s00332-002-0466-4)

# Argument
`param` is of type `NamedTuple` and must contain
- dimensionless parameters `ϵ` (nonlinearity) and `μ` (dispersion);
- numerical parameters to construct the mesh of collocation points as `mesh = Mesh(param)`.

## Optional keyword arguments
- two parameters `a` (default is `-1/3`) and `b` (default is `+1/3`) which determine the model solved. You need `a+2*b=1/3` for validity as a long wave model (without surface tension).
- `ktol`: tolerance of the low-pass Krasny filter (default is `0`, i.e. no filtering);
- `dealias`: dealiasing with Orlicz rule `1-dealias/(dealias+2)` (default is `0`, i.e. no dealiasing);
- `label`: a label for future references (default is `"Boussinesq"`);
- `verbose`: prints information if `true` (default is `true`).

# Return values
Generate necessary ingredients for solving an initial-value problem via `solve!`:
1. a function `Boussinesq.f!` to be called in explicit time-integration solvers;
2. a function `Boussinesq.mapto` which from `(η,v)` of type `InitialData` provides the raw data matrix on which computations are to be executed;
3. a function `Boussinesq.mapfro` which from such data matrix returns the Tuple of real vectors `(η,v)`, where
    - `η` is the surface deformation;
    - `v` is the derivative of the trace of the velocity potential.

"""
mutable struct Boussinesq <: AbstractModel

	label 	:: String
	f!		:: Function
	mapto	:: Function
	mapfro	:: Function

    function Boussinesq(param::NamedTuple;
						a=-1/3,b=1/3,
						dealias=0,ktol=0,
						label="Boussinesq",
						verbose=true)

		m=WhithamBoussinesq(param;Boussinesq=true,
							a=a,b=b,
							dealias=dealias,ktol=ktol,
							verbose=verbose,label=label)

		new(m.label, m.f!, m.mapto, m.mapfro)
    end
end


# Bayesian logistic regression with independent Gaussian priors
using StatsFuns: logistic

type SteinLogisticRegressionGaussianPrior <: SteinLogisticRegressionPosterior
    X::Array{Float64, 2}
    y::Array{Float64, 1} # we assume the y = -1, 1
    priormu::Array{Float64, 1}
    priorvar::Array{Float64, 1}  # we assume components are independent
    c1::Float64
    c2::Float64
    c3::Float64
end

SteinLogisticRegressionGaussianPrior(X::Array{Float64, 2}, y::Array{Float64, 1}) =
SteinLogisticRegressionGaussianPrior(X, y,
                                     repmat([0.0], size(X, 2)),
                                     repmat([1.0], size(X, 2)))

# Constructor with canonical setting of (c1,c2,c3) = (1,1,1)
SteinLogisticRegressionGaussianPrior(X::Array{Float64, 2},
                                     y::Array{Float64, 1},
                                     priormu::Array{Float64, 1},
                                     priorvar::Array{Float64, 1}) =
SteinLogisticRegressionGaussianPrior(X, y, priormu, priorvar,
                                     1.0, 1.0, 1.0)

function logprior(d::SteinLogisticRegressionGaussianPrior,
                  beta::Array{Float64, 1})
    # Independent Gaussian prior for each component
    logpriorvalue =
        -0.5 * sum((beta .- d.priormu') .^ 2 ./ d.priorvar') -
        -0.5 * sum(log(2 * pi .* d.priorvar))
    logpriorvalue
end

function gradlogprior(d::SteinLogisticRegressionGaussianPrior,
                      beta::Array{Float64, 1})
    -(beta .- d.priormu) ./ d.priorvar
end

function numdimensions(d::SteinLogisticRegressionGaussianPrior)
    length(d.priormu)
end

function priordiaghessian(d::SteinLogisticRegressionGaussianPrior,
                          beta::Array{Float64, 1})
    -1.0 ./ d.priorvar
end


### TEST pflap

## using grid

# generate image GRID
n = 20
mgrid = ones(512, 512);
@views mgrid[1:n:end, :] .= 0;
@views mgrid[:, 1:n:end] .= 0;
img = mgrid

# genrate flow
flow = gen_tiled_flow(size(img), 50);

# generate warpped image
imgw = LAP_julia.interpolation.imWarp(img, real(flow), imag(flow))

## using mandrill image

# genrate image MANDRILL
img = testimage("mandril_gray")
imgr = reverse(img, dims=1)
img = imgr

# genrate flow
flow = gen_tiled_flow(size(img), 30);
showflow(flow)

# generate warpped image
imgw = LAP_julia.interpolation.imWarp(img, real(flow), imag(flow));

# show imgw
imgshow(Float32.(imgw))

imgshow(img)

## using chessboard image

tile_size = 5
board_size = 4 # must be even
mini_board = [zeros(tile_size, tile_size) ones(tile_size, tile_size);
              ones(tile_size, tile_size) zeros(tile_size, tile_size)]
chessboard = repeat(mini_board, outer=(convert(Integer,(board_size/2)), convert(Integer,(board_size/2))))
img = chessboard

img = gen_chess()

# genrate flow
flow = gen_tiled_flow(size(img), 20, 1000);

# generate warpped image
imgw = LAP_julia.interpolation.warp_img(img, real(flow), imag(flow))

# show imgw
imgshow(imgw)

# show flow
showflow(flow)

maximum(filter(!isnan, real(flow)))
minimum(filter(!isnan, real(flow)))

### RUN OPTIFLOW:

u_est, source_reg, figs = pflap(img, imgw);

# compare with a single lap:
u_sin_est = single_lap(img, imgw, 32, [65, 65], 3)
showflow(u_sin_est)

LAP_julia.inpaint.inpaint_nans!(u_sin_est)
single_img = LAP_julia.interpolation.imWarp_replicate(imgw, real(u_sin_est), imag(u_sin_est))
imgshow(single_img)
imgshow(imgw)

### compare results

# estimation
maximum(abs.(imag(u_est)))
showflow(u_est)
showflow(u_est .* (-1) .- flow)

# rewarp check
dewarp = LAP_julia.interpolation.imWarp(imgw, real(u_est), imag(u_est))
imgshow(Float32.(dewarp))
imgshow(Float32.(img))

imgshow(source_reg)

# truth
maximum(abs.(imag(flow)))
sum(abs.(real(flow)))/length(real(flow))
showflow(flow, mag = 1)

# show source_reg
fig = PyPlot.figure(dpi = 300, figsize = (5, 5))
PyPlot.imshow(source_reg, cmap = :gray);
gcf()

# show imgw
fig = PyPlot.figure(dpi = 300, figsize = (5, 5))
PyPlot.imshow(Float64.(imgw), cmap = :gray);
gcf()

# k = (-1:1) * ones(1,3)
# A =      [1     2     3;
#           4     5     6;
#           7     8     9];
#
# transpose(A[:]) * k[:]
#
# sum(transpose(A) * (-1:1))
#
# # l0
#
# l = transpose((-1:1) * ones(1,3))
#
# transpose(A[:]) * l[:]
#
# sum(A * (-1:1))
1


abstract type AbstractMarkovDecisionProcess end
## --------------------------------------------- ##
## GENERIC PART
## --------------------------------------------- ##

"""
    MarkovDecisionProcess is a MDP implementation of AbstractMarkovDecisionProcess.

A Markov decision process is a sequential decision problem with fully observable and stochastic environment with a transition model and rewards function.
The discount factor (gamma variable) describes the preference for current rewards over future rewards.
"""
struct MarkovDecisionProcess{T} <: AbstractMarkovDecisionProcess
    initial::T
    states::Set{T}
    actions::Set{T}
    terminal_states::Set{T}
    transitions::Dict
    gamma::Float64
    reward::Dict

    function MarkovDecisionProcess{T}(initial::T, actions_list::Set{T}, terminal_states::Set{T}, transitions::Dict, states::Union{Nothing, Set{T}}, gamma::Float64) where T
        !(0 < gamma <= 1) ? nothing : error("MarkovDecisionProcess():\nThe gamma variable of an MDP must be between 0 and 1, the constructor was given ", gamma, "!")
        new_states = typeof(states) <: Set ? states : Set{typeof(initial)}()

        return new(initial, new_states, actions_list, terminal_states, transitions, gamma, Dict())
    end
end

function MarkovDecisionProcess(initial, actions_list::Set, terminal_states::Set, transitions::Dict;
                                states::Union{Nothing, Set}=nothing, gamma::Float64=0.9)
    return MarkovDecisionProcess{typeof(initial)}(initial, actions_list, terminal_states, transitions, states, gamma)
end

"""
    reward(mdp::T, state) where {T <: AbstractMarkovDecisionProcess}

Return a reward based on the given 'state'.
"""
function reward(mdp::T, state) where {T <: AbstractMarkovDecisionProcess}
    return mdp.reward[state]
end

"""
    transition_model(mdp::T, state, action) where {T <: AbstractMarkovDecisionProcess}

Return a list of (P(s'|s, a), s') pairs given the state 's' and action 'a'.
"""
function transition_model(mdp::T, state, action) where {T <: AbstractMarkovDecisionProcess}
    length(mdp.transitions) == 0 ? nothing : error("transition_model(): The transition model for the given 'mdp' could not be found!")

    return mdp.transitions[state][action]
end

"""
    actions(mdp::T, state) where {T <: AbstractMarkovDecisionProcess}

Return a set of actions that are possible in the given state.
"""
function actions(mdp::T, state) where {T <: AbstractMarkovDecisionProcess}
    if state in mdp.terminal_states
        return Set{Nothing}([nothing])
    else
        return mdp.actions;
    end
end

## --------------------------------------------- ##
## SPECIFIC PART 4x3 GAME
## --------------------------------------------- ##

"""
    GridMarkovDecisionProcess is a MDP implementation of the grid from chapter 17.

Obstacles in the environment are representedby a nothing.

Arguments:

"""
struct GridMarkovDecisionProcess <: AbstractMarkovDecisionProcess
    initial::Tuple{Int64, Int64}
    states::Set{Tuple{Int64, Int64}}
    actions::Set{Tuple{Int64, Int64}}
    terminal_states::Set{Tuple{Int64, Int64}}
    grid::Array{Union{Nothing, Float64}, 2}
    gamma::Float64
    reward::Dict
    function GridMarkovDecisionProcess(initial::Tuple{Int64, Int64}, terminal_states::Set{Tuple{Int64, Int64}},
                                       grid::Array{Union{Nothing, Float64}, 2};
                                       states::Union{Nothing, Set{Tuple{Int64, Int64}}}=nothing, gamma::Float64=0.9)
        (0 < gamma <= 1) ? nothing : error("MarkovDecisionProcess():\nThe gamma variable of an MDP must be between 0 and 1, the constructor was given ", gamma, "!")
        new_states = typeof(states) <: Set ? states : Set{Tuple{Int64, Int64}}()
        orientations::Set = Set{Tuple{Int64, Int64}}([(1, 0), (0, 1), (-1, 0), (0, -1)])
        reward::Dict = Dict()
        for i in 1:size(grid, 1)
            for j in 1:size(grid, 2)
                reward[(i, j)] = grid[i, j]
                if !(grid[i, j] === nothing)
                    push!(new_states, (i, j))
                end
            end
        end

        return new(initial, new_states, orientations, terminal_states, grid, gamma, reward);
    end
end

"""
    go_to(gmdp::GridMarkovDecisionProcess, state::Tuple{Int64, Int64}, direction::Tuple{Int64, Int64})

Return the next state given the current state and direction. If the next state is not known
"""
function go_to(gmdp::GridMarkovDecisionProcess, state::Tuple{Int64, Int64}, direction::Tuple{Int64, Int64})
    next_state::Tuple{Int64, Int64} = state .+ direction

    return next_state in gmdp.states ? next_state : state
end

"""
    transition_model

Return an array of (P(s'|s, a), s') pairs given the state 's' and action 'a'.
"""
function transition_model(gmdp::GridMarkovDecisionProcess, state::Tuple{Int64, Int64}, action::Nothing)
    return [(0.0, state)];
end

"""
    transition_model

Return an array of (P(s'|s, a), s') pairs given the state 's' and action 'a'.
"""
function transition_model(gmdp::GridMarkovDecisionProcess, state::Tuple{Int64, Int64}, action::Tuple{Int64, Int64})
    return [
            (0.8, go_to(gmdp, state, action)),
            (0.1, go_to(gmdp, state, turn_heading(action, -1))),
            (0.1, go_to(gmdp, state, turn_heading(action, 1)))
            ];
end

# (0, 1) will move the agent rightward.
# (-1, 0) will move the agent upward.
# (0, -1) will move the agent leftward.
# (1, 0) will move the agent downward.
const arrow_characters = Dict((0, 1) => ">", (-1, 0) => "^", (0, -1) => "<", (1, 0) => "v", nothing => ".")

"""
    show_grid(gmdp::GridMarkovDecisionProcess, mapping::Dict)

Returns a string representation of the current grid using [<, >, v, ^] to indication the heading and "." to indicate no action.
"""
function show_grid(gmdp::GridMarkovDecisionProcess, policy::Dict)
    mapping = Dict(state => arrow_characters[action] for (state, action) in policy)
    io = IOBuffer()
    for i in 1:size(gmdp.grid, 1)
        for j in 1:size(gmdp.grid, 2)
            print(io, get(mapping, (i,j), "o"))
        end
        print(io,"\n")
    end
    return String(take!(io))
end

"""
    turn_heading

Given an input heading, return the heading in inc positions further in a clockwise manner

The headings can be specified with H. The heading h should occur in the collection of headings H.
"""
function turn_heading(h::Tuple{Int64,Int64}, inc::Int64; H=[(0,1),(1,0),(0,-1),(-1,0)])
    # get index of current heading
    i = findfirst(x->x==h, H)
    # return incremented heading
    return H[(i + inc + length(H) - 1) % length(H) + 1]
end

## --------------------------------------------- ##
## Solution methods
## --------------------------------------------- ##

"""
    value_iteration(mdp::T; epsilon::Float64=0.001) where {T <: AbstractMarkovDecisionProcess}

Return the utilities of the MDP's states as a Dict by applying the value iteration algorithm (Fig. 17.4)
on the given Markov decision process 'mdp' and a arbitarily small positive number 'epsilon'.
"""
function value_iteration(mdp::T; ϵ::Float64=0.001) where {T <: AbstractMarkovDecisionProcess}
    # initialise utility to zero for all states
    U_prime = Dict(collect(Pair(state, 0.0) for state in mdp.states))
    while true
        U = copy(U_prime)
        δ = 0.0
        for state in mdp.states
            U_prime[state] = reward(mdp, state) + mdp.gamma * maximum(map(x -> sum(p*U[newstate] for (p, newstate) in x),
                                                                          [transition_model(mdp, state, action) for action in actions(mdp, state)]))
            δ = max(δ, abs(U_prime[state] - U[state]));
        end
        if (δ < ((ϵ * (1 - mdp.gamma))/mdp.gamma))
            return U_prime;
        end
    end
end

function policy_extraction(mdp::T, U::Dict) where T<:AbstractMarkovDecisionProcess
	pi = Dict()
	for state in mdp.states
		A = collect(actions(mdp, state))
		expected_utility = map(x -> sum(p*U[newstate] for (p, newstate) in x),  [transition_model(mdp, state, action) for action in A])
		pi[state] = A[argmax(expected_utility)]
	end
	return pi
end

#function expected_utility(mdp::T, U::Dict, state::Tuple{Int64, Int64}, action::Nothing) where {T <: AbstractMarkovDecisionProcess}
#    return sum((p * U[state_prime] for (p, state_prime) in transition_model(mdp, state, action)));
#end

"""
    policy_evaluation(pi::Dict, U::Dict, mdp::T; k::Int64=20) where {T <: AbstractMarkovDecisionProcess}

Return the updated utilities of the MDP's states by applying the modified policy iteration
algorithm on the given Markov decision process 'mdp', utility function 'U', policy 'pi',
and number of Bellman updates to use 'k'.
"""
function policy_evaluation(pi::Dict, U::Dict, mdp::T; k::Int64=20) where {T <: AbstractMarkovDecisionProcess}
    for i in 1:k
        for state in mdp.states
            U[state] = reward(mdp, state) + mdp.gamma * sum(p * U[newstate] for (p, newstate) in transition_model(mdp, state, pi[state]));
        end
    end
    return U;
end

function policy_evaluation(pi::Dict, U::Dict, gmdp::GridMarkovDecisionProcess; k::Int64=20)
    for i in 1:k
        for state in gmdp.states
            U[state] = reward(gmdp, state) + gmdp.gamma * sum(p * U[newstate] for (p, newstate) in transition_model(gmdp, state, pi[state]));
        end
    end
    return U;
end

"""
    policy_iteration(mdp::T) where {T <: AbstractMarkovDecisionProcess}

Return a policy using the policy iteration algorithm (Fig. 17.7) given the Markov decision process 'mdp'.

it initialises the policy at random
"""
function policy_iteration(mdp::T) where {T <: AbstractMarkovDecisionProcess}
    U = Dict(Pair(state, 0.0) for state in mdp.states)
    pi = Dict(Pair(state, rand(actions(mdp, state))) for state in mdp.states)
    while true
        U = policy_evaluation(pi, U, mdp)
        unchanged = true
        for state in mdp.states
            A = collect(actions(mdp, state))
            expected_utility = map(x -> sum(p*U[newstate] for (p, newstate) in x),  [transition_model(mdp, state, action) for action in A])
            action = A[argmax(expected_utility)]
            if action != pi[state]
                pi[state] = action
                unchanged = false
            end
        end
        if unchanged
            return pi, U;
        end
    end
end


using WriteVTK

function vtk_connectivity_map_highorder(Nq1, Nq2 = 1, Nq3 = 1)
    connectivity = Array{Int, 1}(undef, Nq1 * Nq2 * Nq3)
    L = LinearIndices((1:Nq1, 1:Nq2, 1:Nq3))

    corners = (
        (1, 1, 1),
        (Nq1, 1, 1),
        (Nq1, Nq2, 1),
        (1, Nq2, 1),
        (1, 1, Nq3),
        (Nq1, 1, Nq3),
        (Nq1, Nq2, Nq3),
        (1, Nq2, Nq3),
    )
    edges = (
        (2:(Nq1 - 1), 1:1, 1:1),
        (Nq1:Nq1, 2:(Nq2 - 1), 1:1),
        (2:(Nq1 - 1), Nq2:Nq2, 1:1),
        (1:1, 2:(Nq2 - 1), 1:1, 1:1),
        (2:(Nq1 - 1), 1:1, Nq3:Nq3),
        (Nq1:Nq1, 2:(Nq2 - 1), Nq3:Nq3),
        (2:(Nq1 - 1), Nq2:Nq2, Nq3:Nq3),
        (1:1, 2:(Nq2 - 1), Nq3:Nq3),
        (1:1, 1:1, 2:(Nq3 - 1)),
        (Nq1:Nq1, 1:1, 2:(Nq3 - 1)),
        (1:1, Nq2:Nq2, 2:(Nq3 - 1)),
        (Nq1:Nq1, Nq2:Nq2, 2:(Nq3 - 1)),
    )
    faces = (
        (1:1, 2:(Nq2 - 1), 2:(Nq3 - 1)),
        (Nq1:Nq1, 2:(Nq2 - 1), 2:(Nq3 - 1)),
        (2:(Nq1 - 1), 1:1, 2:(Nq3 - 1)),
        (2:(Nq1 - 1), Nq2:Nq2, 2:(Nq3 - 1)),
        (2:(Nq1 - 1), 2:(Nq2 - 1), 1:1),
        (2:(Nq1 - 1), 2:(Nq2 - 1), Nq3:Nq3),
    )
    if Nq2 == Nq3 == 1
        @assert Nq1 > 1
        corners = (corners[1:2]...,)
        edges = (edges[1],)
        faces = ()
    elseif Nq3 == 1
        @assert Nq1 > 1
        @assert Nq2 > 1
        corners = (corners[1:4]...,)
        edges = (edges[1:4]...,)
        faces = (faces[5],)
    end

    # corners
    iter = 1
    for (i, j, k) in corners
        connectivity[iter] = L[i, j, k]
        iter += 1
    end
    # edges
    for (is, js, ks) in edges
        for k in ks, j in js, i in is
            connectivity[iter] = L[i, j, k]
            iter += 1
        end
    end
    # faces
    for (is, js, ks) in faces
        for k in ks, j in js, i in is
            connectivity[iter] = L[i, j, k]
            iter += 1
        end
    end
    # interior
    for k in 2:(Nq3 - 1), j in 2:(Nq2 - 1), i in 2:(Nq1 - 1)
        connectivity[iter] = L[i, j, k]
        iter += 1
    end
    return connectivity
end

#=
This is the 1D WriteMesh routine
=#
function writemesh_highorder(
    base_name,
    x1;
    x2 = nothing,
    x3 = nothing,
    fields = (),
    realelems = 1:size(x1)[end],
)
    (Nq1, _) = size(x1)

    con = vtk_connectivity_map_highorder(Nq1)

    M = MeshCell{VTKCellTypes.VTKCellType, Array{Int, 1}}
    cells = Array{M, 1}(undef, length(realelems))

    for (i, e) in enumerate(realelems)
        offset = (e - 1) * Nq1
        cells[i] = MeshCell(VTKCellTypes.VTK_LAGRANGE_CURVE, offset .+ con)
    end

    if x2 == nothing
        @assert isnothing(x3)
        vtkfile =
            vtk_grid("$(base_name)", @view(x1[:]), cells; compress = false)
    elseif x3 == nothing
        vtkfile = vtk_grid(
            "$(base_name)",
            @view(x1[:]),
            @view(x2[:]),
            cells;
            compress = false,
        )
    else
        vtkfile = vtk_grid(
            "$(base_name)",
            @view(x1[:]),
            @view(x2[:]),
            @view(x3[:]),
            cells;
            compress = false,
        )
    end
    for (name, v) in fields
        vtk_point_data(vtkfile, v, name)
    end
    outfiles = vtk_save(vtkfile)
end

#=
This is the 2D WriteMesh routine
=#
function writemesh_highorder(
    base_name,
    x1,
    x2;
    x3 = nothing,
    fields = (),
    realelems = 1:size(x1)[end],
)
    @assert size(x1) == size(x2)
    (Nq1, Nq2, _) = size(x1)
    @assert Nq1 == Nq2
    con = vtk_connectivity_map_highorder(Nq1, Nq2)

    M = MeshCell{VTKCellTypes.VTKCellType, Array{Int, 1}}
    cells = Array{M, 1}(undef, length(realelems))
    for (i, e) in enumerate(realelems)
        offset = (e - 1) * Nq1 * Nq2
        cells[i] =
            MeshCell(VTKCellTypes.VTK_LAGRANGE_QUADRILATERAL, offset .+ con[:])
    end

    if x3 == nothing
        vtkfile = vtk_grid(
            "$(base_name)",
            @view(x1[:]),
            @view(x2[:]),
            cells;
            compress = false,
        )
    else
        vtkfile = vtk_grid(
            "$(base_name)",
            @view(x1[:]),
            @view(x2[:]),
            @view(x3[:]),
            cells;
            compress = false,
        )
    end
    for (name, v) in fields
        vtk_point_data(vtkfile, v, name)
    end
    outfiles = vtk_save(vtkfile)
end

#=
This is the 3D WriteMesh routine
=#
function writemesh_highorder(
    base_name,
    x1,
    x2,
    x3;
    fields = (),
    realelems = 1:size(x1)[end],
)
    (Nq1, Nq2, Nq3, _) = size(x1)
    @assert Nq1 == Nq2 == Nq3
    con = vtk_connectivity_map_highorder(Nq1, Nq2, Nq3)
    M = MeshCell{VTKCellTypes.VTKCellType, Array{Int, 1}}
    cells = Array{M, 1}(undef, length(realelems))
    for (i, e) in enumerate(realelems)
        offset = (e - 1) * Nq1 * Nq2 * Nq3
        cells[i] = MeshCell(VTKCellTypes.VTK_LAGRANGE_HEXAHEDRON, offset .+ con)
    end

    vtkfile = vtk_grid(
        "$(base_name)",
        @view(x1[:]),
        @view(x2[:]),
        @view(x3[:]),
        cells;
        compress = false,
    )
    for (name, v) in fields
        vtk_point_data(vtkfile, v, name)
    end
    outfiles = vtk_save(vtkfile)
end

#=
This is the 1D WriteMesh routine
=#
function writemesh_raw(
    base_name,
    x1;
    x2 = nothing,
    x3 = nothing,
    fields = (),
    realelems = 1:size(x1)[end],
)
    (Nq1, _) = size(x1)
    Nsubcells = (Nq1 - 1)

    M = MeshCell{VTKCellTypes.VTKCellType, Array{Int, 1}}
    cells = Array{M, 1}(undef, Nsubcells * length(realelems))
    for e in realelems
        offset = (e - 1) * Nq1
        for i in 1:(Nq1 - 1)
            cells[i + (e - 1) * Nsubcells] =
                MeshCell(VTKCellTypes.VTK_LINE, offset .+ [i, i + 1])
        end
    end

    vtkfile = vtk_grid("$(base_name)", @view(x1[:]), cells; compress = false)
    for (name, v) in fields
        vtk_point_data(vtkfile, v, name)
    end
    outfiles = vtk_save(vtkfile)
end

#=
This is the 2D WriteMesh routine
=#
function writemesh_raw(
    base_name,
    x1,
    x2;
    x3 = nothing,
    fields = (),
    realelems = 1:size(x1)[end],
)
    @assert size(x1) == size(x2)
    (Nq1, Nq2, _) = size(x1)
    Nsubcells = (Nq1 - 1) * (Nq2 - 1)

    M = MeshCell{VTKCellTypes.VTKCellType, Array{Int, 1}}
    cells = Array{M, 1}(undef, Nsubcells * length(realelems))
    ind = LinearIndices((1:Nq1, 1:Nq2))
    for e in realelems
        offset = (e - 1) * Nq1 * Nq2
        for j in 1:(Nq2 - 1)
            for i in 1:(Nq1 - 1)
                cells[i + (j - 1) * (Nq1 - 1) + (e - 1) * Nsubcells] = MeshCell(
                    VTKCellTypes.VTK_PIXEL,
                    offset .+ ind[i:(i + 1), j:(j + 1)][:],
                )
            end
        end
    end

    if x3 == nothing
        vtkfile = vtk_grid(
            "$(base_name)",
            @view(x1[:]),
            @view(x2[:]),
            cells;
            compress = false,
        )
    else
        vtkfile = vtk_grid(
            "$(base_name)",
            @view(x1[:]),
            @view(x2[:]),
            @view(x3[:]),
            cells;
            compress = false,
        )
    end
    for (name, v) in fields
        vtk_point_data(vtkfile, v, name)
    end
    outfiles = vtk_save(vtkfile)
end

#=
This is the 3D WriteMesh routine
=#
function writemesh_raw(
    base_name,
    x1,
    x2,
    x3;
    fields = (),
    realelems = 1:size(x1)[end],
)
    (Nq1, Nq2, Nq3, _) = size(x1)
    (N1, N2, N3) = (Nq1 - 1, Nq2 - 1, Nq3 - 1)
    Nsubcells = N1 * N2 * N3

    M = MeshCell{VTKCellTypes.VTKCellType, Array{Int, 1}}
    cells = Array{M, 1}(undef, Nsubcells * length(realelems))
    ind = LinearIndices((1:Nq1, 1:Nq2, 1:Nq3))
    for e in realelems
        offset = (e - 1) * Nq1 * Nq2 * Nq3
        for k in 1:N3
            for j in 1:N2
                for i in 1:N1
                    cells[i + (j - 1) * N1 + (k - 1) * N1 * N2 + (e - 1) * Nsubcells] =
                        MeshCell(
                            VTKCellTypes.VTK_VOXEL,
                            offset .+ ind[i:(i + 1), j:(j + 1), k:(k + 1)][:],
                        )
                end
            end
        end
    end

    vtkfile = vtk_grid(
        "$(base_name)",
        @view(x1[:]),
        @view(x2[:]),
        @view(x3[:]),
        cells;
        compress = false,
    )
    for (name, v) in fields
        vtk_point_data(vtkfile, v, name)
    end
    outfiles = vtk_save(vtkfile)
end


# Alias sampling
struct FastSampler{Tv,Ti}
    F::Array{Tv,1}    # array of p[s.A[i]] / (1 / n)
    A::Array{Ti,1}    # array of indices of elements < (1 / n)
    V::Array{Ti,1}    # arary of indices of elements > (1 / n)
    n::Ti
    rng::AbstractRNG  # the pseudo random number generator
end

# sample a random number
function sample(s::FastSampler{Tv,Ti}) where {Tv,Ti}
    #i = rand(1:s.n)
    i = ceil(Ti,rand(s.rng)*s.n)
    f = rand(s.rng)
    if f < s.F[i]
        return s.A[i]
    else
        return s.V[i]
    end
end

# sample a random number
function sampleMany(s::FastSampler{Tv,Ti},sampCount::Ti) where {Tv,Ti}
    #i = rand(1:s.n)
    samples = Array{Ti,1}(sampCount)
    for j = 1:sampCount
        i = ceil(Ti,rand(s.rng)*s.n)
        f = rand(s.rng)
        if f < s.F[i]
            samples[j] = s.A[i]
        else
            samples[j] = s.V[i]
        end
    end
    return samples
end

"""
    s = FastSampler(p)
"""
function FastSampler(p::Array{Tv,1}; residual::Bool = false, rng::AbstractRNG=Random.GLOBAL_RNG) where Tv

    @assert(minimum(p) > 0, "The probability vector has a negative entry")

    n = length(p)

    # generate the Sampler
    F = Array{Tv}(undef, n)
    A = Array{Int64}(undef, n)
    V = Array{Int64}(undef, n)

    # generate a normalized p, so that it sums up to n
    newp = n * copy(p) / sum(p)

    # fill up the arays A and V
    pos = 0                      # number of filled or half filled buckets
    posFilled = 0                # number of filled buckets

    for i in 1:n
        if newp[i] < 1 #&& !epsequal(newp[i], 1)
            pos = pos + 1
            A[pos] = i
            F[pos] = newp[i]
        end
    end

    err = 0
    for i in 1:n
        if newp[i] >= 1
            val = newp[i]

            while val >= 1
                if posFilled < pos
                    # fill up a bucket already containing a small element, and subtract that value from the current element
                    posFilled = posFilled + 1
                    V[posFilled] = i
                    val = val - (1 - F[posFilled])
                else
                    # create a bucket with only one large element. note posFilled = pos
                    pos = pos + 1
                    posFilled = posFilled + 1
                    A[pos] = V[pos] = i
                    F[pos] = 1
                    val = val - 1
                end
            end

            if val > 0 && pos == n
            	err = err + val
            end

            # treat the case in which val becomes smaller than 1. Have to do the pos < n check because of precision errors
            if val > 0 && pos < n
                pos = pos + 1
                A[pos] = i
                F[pos] = val
            end
        end
    end

    while posFilled < n
        # if we are in this case, posFilled should be equal to pos - 1
        posFilled = posFilled + 1
        V[posFilled] = A[posFilled]
        err += 1-F[posFilled]
    end

    if residual
        return FastSampler(F, A, V, n, rng), err
    else
        return FastSampler(F, A, V, n, rng)
    end
end

"""
    s = blockSample(p; k = length(p))

Compute numbers sampled with probability proportional to the vector p.
They are returned in order.
If need be, they can be permuted with randperm.
"""
function blockSample(p; k = length(p))
    samp = zeros(Int,k)

    r = sort(rand(k))

    cs = cumsum(p)
    cs = cs / cs[end]

    j = 1
    for i in 1:k
        while r[i] > cs[j]
            j += 1
        end

        samp[i] = j
    end

    return samp
end


export Euler

"""
    Euler()

This uses the explicit Euler method for prediction, also known as the
tangent predictor.
"""
struct Euler <: AbstractPredictor end
struct EulerCache <: AbstractPredictorCache end

cache(::Euler, H, x, ẋ, t) = EulerCache()

function predict!(
    xnext,
    cache::EulerCache,
    H::HomotopyWithCache,
    x,
    t,
    Δt,
    ẋ,
    ::JacobianMonitor,
)
    @inbounds for i = 1:length(x)
        xnext[i] = x[i] + Δt * ẋ[i]
    end
    nothing
end

order(::EulerCache) = 2


module TestChess

using Test
import Moves
import Util
import Magic

@testset "moves" verbose=true begin
    @test Moves.rookMoves(0,1) == parse(UInt64,
    "00000001"*
    "00000001"*
    "00000001"*
    "00000001"*
    "00000001"*
    "00000001"*
    "00000001"*
    "11111110",base=2)

    @test Moves.rookMoves(0,64) == parse(UInt64,
    "01111111"*
    "10000000"*
    "10000000"*
    "10000000"*
    "10000000"*
    "10000000"*
    "10000000"*
    "10000000",base=2)

    @test Moves.bishopMoves(0,1) == parse(UInt64,
    "10000000"*
    "01000000"*
    "00100000"*
    "00010000"*
    "00001000"*
    "00000100"*
    "00000010"*
    "00000000",base=2)

    @test Moves.bishopMoves(0,64) == parse(UInt64,
    "00000000"*
    "01000000"*
    "00100000"*
    "00010000"*
    "00001000"*
    "00000100"*
    "00000010"*
    "00000001",base=2)

    @test Moves.queenMoves(0,32) == parse(UInt64,
    "10001000"*
    "10010000"*
    "10100000"*
    "11000000"*
    "01111111"*
    "11000000"*
    "10100000"*
    "10010000",base=2)
end

@testset "check " verbose = true begin
    board = Util.startPositions()
    @test Moves.isCheck(board,'B') == false

    board = Util.emptyPositions()
    board["BP"] = Util.setBit(board["BP"], 11,1)
    board["WK"] = Util.setBit(board["WK"],4,1)
    @test Moves.isCheck(board,'W') == true

    board = Util.emptyPositions()
    board["WN"] = Util.setBit(board["WN"], 52,1)
    board["BK"] = Util.setBit(board["BK"],37,1)
    @test Moves.isCheck(board,'B') == true

    board = Util.emptyPositions()
    board["BQ"] = Util.setBit(board["BQ"], 14,1)
    board["WK"] = Util.setBit(board["WK"],8,1)
    @test Moves.isCheck(board,'B') == false

    board = Util.emptyPositions()
    board["BQ"] = Util.setBit(board["BQ"], 1,1)
    board["WK"] = Util.setBit(board["WK"],5,1)
    board["BR"] = Util.setBit(board["BR"], 9,1)
    @test Moves.isCheckMate(board,'W',[]) == 1

    #white in checkmate surrounded by white pieces
    board = Util.emptyPositions()
    board["WP"] = Util.setBit(board["WP"],13,1)
    board["WP"] = Util.setBit(board["WP"],12,1)
    board["WB"] = Util.setBit(board["WB"],6,1)
    board["WQ"] = Util.setBit(board["WQ"],4,1)
    board["WK"] = Util.setBit(board["WK"],5,1)
    board["BQ"] = Util.setBit(board["BQ"],32,1)
    @test Moves.isCheckMate(board,'W',[]) == 1

    #white in check, king could take a piece, that piece protected
    board = Util.emptyPositions()
    board["WP"] = Util.setBit(board["WP"],14,1)
    board["BR"] = Util.setBit(board["BR"],13,1)
    board["BB"] = Util.setBit(board["BB"],29,1)
    board["BQ"] = Util.setBit(board["BQ"],4,1)
    board["WK"] = Util.setBit(board["WK"],6,1)
    @test Moves.isCheckMate(board,'W',[]) == 1

    #white in check but piece can intervene
    board = Util.emptyPositions()
    board["WP"] = Util.setBit(board["WP"],14,1)
    board["WR"] = Util.setBit(board["WR"],13,1)
    board["BB"] = Util.setBit(board["BB"],29,1)
    board["BQ"] = Util.setBit(board["BQ"],4,1)
    board["WK"] = Util.setBit(board["WK"],6,1)
    @test Moves.isCheckMate(board,'W',[]) == 0
    @test Moves.isCheck(board,'W') == true

    #stalemate
    board = Util.emptyPositions()
    board["WP"] = Util.setBit(board["WP"],50,1)
    board["WQ"] = Util.setBit(board["WR"],42,1)
    board["BK"] = Util.setBit(board["BB"],58,1)
    @test Moves.isCheckMate(board,'B',[]) == 2
    @test Moves.isCheck(board,'B') == false

    #capture out of check
    board = Util.emptyPositions()
    board["BQ"] = Util.setBit(board["BQ"],60,1)
    board["WQ"] = Util.setBit(board["WQ"],59,1)
    board["BK"] = Util.setBit(board["BK"],61,1)
    @test Moves.isCheck(board,'B') == false

    #king capture out of checkmate
    board = Util.emptyPositions()
    board["BQ"] = Util.setBit(0,63,1)
    board["BR"] = Util.setBit(0,49,1)
    board["WK"] = Util.setBit(0,64,1)
    @test Moves.isCheck(board,'W') == true
    @test Moves.isCheckMate(board,'W',[]) == 0

    #cannot move into check

    board = Util.emptyPositions()
    board["BK"] = Util.setBit(0,60,1)
    board["WQ"] = Util.setBit(0,46,1)
    board["BP"] = Util.setBit(0,53,1)
    Util.prettyPrintBoard(board)
    @test length(Moves.getMoves(board,'B',[])) == 5

    board = Dict("BR" => 0x8100000000000000, "WK" => 0x0000000000000008, "WB" => 0x0000000000000004, "WQ" => 0x0000000004000000, "BQ" => 0x0000000000000020, "BB" => 0x2000100000000000, "BN" => 0x4200000000000000, "BK" => 0x1000000000000000, "WR" => 0x0000000000000001, "WN" => 0x0000000000000002, "BP" => 0x00e7001000000000, "WP" => 0x00000042110c8000)
    @test Moves.isCheckMate(board,'W',[]) == 0

end

@testset "pawn promotion" verbose = true begin

    #pawn promotion forward to queen
    board = Util.emptyPositions()
    board["WP"] = Util.setBit(board["WP"],56,1)
    board = Moves.makeMoveUCI("h7h8Q",board,'W')
    @test board["WQ"] > 0
    @test board["WP"] == 0

    #pawn promotion forward to queen
    board = Util.emptyPositions()
    board["BP"] = Util.setBit(board["BP"],9,1)
    @test "a2a1Q" in Moves.getMoves(board,'B',[])
end

@testset "en passant" verbose = true begin

    #simple en passant of white pawn
    history = []
    board = Util.emptyPositions()
    board["WP"] = Util.setBit(board["WP"],37,1)
    board["BP"] = Util.setBit(board["BP"],38,1)
    push!(history,"e3e5")
    moves = Moves.getMoves(board,'B',history)
    @test "f5e4" in moves

    # #not en passant of black pawn
    # history = []
    # board = emptyPositions()
    # board["WP"] = Util.setBit(board["WP"],37,1)
    # board["BP"] = Util.setBit(board["BP"],38,1)
    # push!(history,"e3e5")
    # moves = getMovesUCI(board,'B',history)
    # @test "" not in moves
end

@testset "castling" verbose = true begin
    #simple en passant of white pawn
    history = []

    #black castling king side
    board = Util.emptyPositions()
    board["BR"] = Util.setBit(board["BR"],64,1)
    board["BK"] = Util.setBit(board["BK"],61,1)
    occ = reduce((x,y) -> x | y, values(board))
    @test "e8g8" in Moves.castleMoves(board,'B',history,occ)

    #white castle queenside
    board = Util.emptyPositions()
    board["WR"] = Util.setBit(board["WR"],1,1)
    board["WK"] = Util.setBit(board["WK"],5,1)
    occ = reduce((x,y) -> x | y, values(board))
    @test "e1c1" in Moves.castleMoves(board,'W',history,occ)

    #rook already moved
    board = Util.emptyPositions()
    board["WR"] = Util.setBit(board["WR"],1,1)
    board["WK"] = Util.setBit(board["WK"],5,1)
    occ = reduce((x,y) -> x | y, values(board))
    history = ["a3a1"]
    @test length(Moves.castleMoves(board,'W',history,occ)) == 0

    #king already moved
    board = Util.emptyPositions()
    board["WR"] = Util.setBit(board["WR"],1,1)
    board["WK"] = Util.setBit(board["WK"],5,1)
    occ = reduce((x,y) -> x | y, values(board))
    history = ["e2e1"]
    @test length(Moves.castleMoves(board,'W',history,occ)) == 0

    #piece in the way
    board = Util.emptyPositions()
    board["WR"] = Util.setBit(board["WR"],1,1)
    board["WK"] = Util.setBit(board["WK"],5,1)
    board["WB"] = Util.setBit(board["WB"],3,1)
    occ = reduce((x,y) -> x | y, values(board))
    history = []
    @test length(Moves.castleMoves(board,'W',history,occ)) == 0

    #castle through a check
    board = Util.emptyPositions()
    board["BR"] = Util.setBit(board["BR"],60,1)
    board["BK"] = Util.setBit(board["BK"],64,1)
    board["WQ"] = Util.setBit(board["WQ"],9,1)
    occ = reduce((x,y) -> x | y, values(board))
    history = []
    @test length(Moves.castleMoves(board,'B',history,occ)) == 0
end

@testset "search" verbose = true begin
end

@testset "eval" verbose = true begin

    board = Util.emptyPositions()
    board["WR"] = Util.setBit(0,16,1)
    board["BR"] = Util.setBit(0,64,1)
    board["BB"] = Util.setBit(0,9,1)
    @test count_ones(Moves.getAttackBoard(board,'W',false) & Util.getBlackPieces(board)) == 2

end

@testset "util" verbose = true begin
    #least sig bit
#    @test 18  == leastSignificantBit(1130822006472704)
#     @test leastSignificantBit(1130822006734848) == 10

end

end


@testset "fantome" begin

    n = 10
    ϵ = 0.1

    rng = MersenneTwister(42)

    p = SquareParticles(rng, n, ϵ)

    q = [[0.5       , 0.5       ],
         [0.20006862, 0.38139954],
         [0.32267135, 0.72687046],
         [0.68044674, 0.78095695],
         [0.77473372, 0.51989917],
         [0.4949439 , 0.23201753],
         [0.74204115, 0.28248482],
         [0.63101615, 0.3650337 ],
         [0.5846235 , 0.64552143],
         [0.45656816, 0.76644607]]

    v = [[ 0.,   1.] ,
         [ 1.,   0.] ,
         [-1.,   0.] ,
         [ 0.,  -1.] ,
         [ 1.,   0.] ,
         [ 0.,   1.] ,
         [ 0.,   1.] ,
         [-1.,   0.] ,
         [-1.,   0.] ,
         [ 0.,  -1.]]

    for i in 1:n
        p.q[i] = q[i]
        p.v[i] = v[i]
    end

    c = PeriodicCollisions( p )

    ref = [0 9 9 1 9 9 9 9 1 1;
           0 0 9 1 9 1 9 1 9 1;
           0 0 0 9 9 9 3 9 9 6;
           0 0 0 0 9 1 1 9 9 9;
           0 0 0 0 0 9 9 2 2 9;
           0 0 0 0 0 0 9 1 9 1;
           0 0 0 0 0 0 0 7 9 9;
           0 0 0 0 0 0 0 0 9 9;
           0 0 0 0 0 0 0 0 0 1;
           0 0 0 0 0 0 0 0 0 0]

    @test ref ≈ c.fantome

end


# Rectangular Full Packed Matrices

type SymmetricRFP{T<:BlasFloat} <: AbstractMatrix{T}
    data::Vector{T}
    transr::Char
    uplo::Char
end

function Ac_mul_A_RFP{T<:BlasFloat}(A::Matrix{T})
    n = size(A, 2)
    C = LAPACK.sfrk!('N', 'U', 'T', 1.0, A, 0.0, Array(T, div(n*(n+1),2)))
    SymmetricRFP(C, 'N', 'U')
end

type TriangularRFP{T<:BlasFloat} <: AbstractMatrix{T}
    data::Vector{T}
    transr::Char
    uplo::Char
end
TriangularRFP(A::Matrix) = TriangularRFP(trttf!('N', 'U', A), 'N', 'U')

full(A::TriangularRFP) = (A.uplo=='U' ? triu! : tril!)(LAPACK.tfttr!(A.transr, A.uplo, A.data))

type CholeskyDenseRFP{T<:BlasFloat} <: Factorization{T}
    data::Vector{T}
    transr::Char
    uplo::Char
end

cholfact!{T<:BlasFloat}(A::SymmetricRFP{T}) = CholeskyDenseRFP(LAPACK.pftrf!(A.transr, A.uplo, copy(A.data)), A.transr, A.uplo)
cholfact{T<:BlasFloat}(A::SymmetricRFP{T}) = cholfact!(copy(A))

copy(A::SymmetricRFP) = SymmetricRFP(copy(A.data), A.transr, A.uplo)

# Least squares
\(A::CholeskyDenseRFP, B::VecOrMat) = LAPACK.pftrs!(A.transr, A.uplo, A.data, copy(B))

inv(A::CholeskyDenseRFP)=LAPACK.pftri!(A.transr, A.uplo, copy(A.data))


export Biss_SecA_ls
function Biss_SecA_ls(h :: LineModel, stop_ls :: LS_Stopping;
                     f_meta = LS_Function_Meta(),
                     h₀ = obj(h, 0.0), g₀ = grad(h, 0.0),
                     φ_dφ :: Function = (x, y) -> phi_dphi(x, y),
                     γ :: Float64 = 0.8, verboseLS :: Bool = false,
                     kwargs...)

  state = stop_ls.current_state

  t = 1.0
  update!(state, x = t)
  φt, dφt = φ_dφ(h, state; kwargs...)
  tqnp = 0.0; dφtqnp = h₀ * (1.0 + 0.01)
  ht = obj(h, t); gt = grad(h, t)
  OK = update_and_start!(stop_ls, ht = ht, gt = gt, h0 = h₀, g0 = g₀)

  (ta, φta, dφta, tb, φtb, dφtb) = find_interval_ls(h, stop_ls,
                                                    verboseLS = verboseLS;
                                                    kwargs...)
  verboseLS && println("ta = $ta tb = $tb")

  t = ta
  tp = tb
  tqnp = tb

  iter = 0

  dφt = dφta; dφa = dφta; dφb = dφtb;

  verboseLS && @printf(" iter     tqnp        t         dφtqnp     dφt \n")
  verboseLS && @printf(" %4d %7.2e  %7.2e  %7.2e  %7.2e\n",
                        iter, tqnp, t, dφa, dφb)

  while !OK             # admissible: satisfies Armijo & Wolfe,
                        # tired: exceeds maximum number of iterations
    s = t - tqnp; y = dφt - dφtqnp;
    Γ = 3.0 * (dφt + dφtqnp) * s - 6.0 * (φt - dφtqnp)
    if y * s + Γ < eps(Float64) .* (s ^ 2)
      yt = y + Γ / s
    else
      yt = y
    end
    dN = -dφt * s / yt

    if ((tp - t) * dN > 0.0) && (dN / (tp - t) < γ)
      tplus = t + dN
      verboseLS && println("N")
    else
      tplus = (t + tp) / 2
      verboseLS && println("B")
    end
    update!(state, x = tplus)
    φplus, dφplus = φ_dφ(h, state; kwargs...)

    if t > tp
      if dφplus < 0.0
        tp = t
        tqnp = t
      else
        tqnp = t
      end
    else
      if dφplus > 0.0
        tp = t
        tqnp = t
      else
        tqnp = t
      end
    end

    t = tplus
    dφtqnp = dφt
    dφt = dφplus

    iter += 1
    OK = update_and_stop!(stop_ls, ht = obj(h, t), gt = grad(h, t))

    verboseLS && @printf(" %4d %7.2e  %7.2e  %7.2e  %7.2e\n",
                          iter, tqnp, t, dφtqnp, dφt)
  end

  return (stop_ls.current_state, stop_ls.meta.optimal)
end


"""
$(TYPEDEF)
$(TYPEDFIELDS)

"""
struct LogQuantNormal <: LogNormalQuantDist
    mu::NestedFloat
    sigma::NestedFloat
    q::NestedFloat
end

function logquantnormal(mu::Float64, sigma::Float64, q::Float64)::Float64
    return exp(quantnormal(mu, sigma, q))
end

lognormalquant(x::LogQuantNormal, mu, sigma, q) = logquantnormal(mu, sigma, q)


mutable struct LL_approx
    models::Vector{AbstractModel}
    λ::Vector{Float64}
    indexes::Matrix{Bool}

    X_train::Matrix{Float64}
    Y_train::Matrix{Float64}

    X_test::Matrix{Float64}
    Y_test::Matrix{Float64}

    ys::Vector{Function}

    approx_errors::Vector{Float64}
    scale_values::Matrix{Float64}

end

function LL_approx( X_train,
                    Y_train;
                    models = AbstractModel[],
                    λ = [1e-5 for i=1:size(Y_train, 2)],
                    indexes = zeros(Bool, 0, 0),
                    X_test  = zeros(Float64, 0, 0),
                    Y_test  = zeros(Float64, 0, 0),
                    ys = Function[],
                    p_train = 0.8,
                    approx_errors = [Inf],
                    scale_values = ones(1,size(X_train, 2)))

    if size(X_train, 1) != size(Y_train, 1)
        @error("Check size of the train set (X_train, Y_train)")
    end
    if length(size(scale_values)) == 1
        scale_values = Array(scale_values')
    end

    X_train = X_train .* scale_values

    if (isempty(X_test) || isempty(Y_test))
        n = round(Int, p_train*size(X_train, 1))
        idx = randperm(size(X_train, 1))

        X_test = X_train[idx[n+1:end], :]
        Y_test = Y_train[idx[n+1:end], :]

        X_train = X_train[idx[1:n], :]
        Y_train = Y_train[idx[1:n], :]
    else
        X_test = X_test .* scale_values
    end

    LL_approx(models, λ, indexes, X_train, Y_train, X_test, Y_test, ys,
                approx_errors, scale_values)
end

function fit!(ll_approx::LL_approx)
    empty!(ll_approx.models)
    empty!(ll_approx.ys)

    for i = 1:size(ll_approx.Y_train, 2)
        Yx = view(ll_approx.Y_train, :, i)
        if sum(ll_approx.indexes[i,:]) == 0 || var(Yx) < 1e-8
            y =  x -> mean(Yx)
            model = KernelInterpolation(zeros(1), zeros(1,1))
        else
            Xtr = view(ll_approx.X_train, :, ll_approx.indexes[i,:])
            model = KernelInterpolation(Yx, Xtr, λ = ll_approx.λ[i], kernel=PeriodicKernel())
            train!(model)
            y = approximate(model)

        end
        # approximate(m) returns a function ŷ_i
        push!(ll_approx.models, model)
        push!(ll_approx.ys, y)

    end
end

function update_approximation_errors!(ll_approx)
    empty!(ll_approx.approx_errors)

    s = zeros(size(ll_approx.Y_train, 2))
    for i = 1:size(ll_approx.X_test, 1)
        y_approx = Ψ_approx(view(ll_approx.X_test, i,:), ll_approx, false)
        s += (y_approx - view(ll_approx.Y_test, i, :)).^2
    end

    ll_approx.approx_errors = sqrt.(s)

end

function optimize_indexes!(model::LL_approx, debug=false)
    # D_ul = size(model.X_test, 1)
    # if D_ul > 10
        # indexes_initial = rand(Bool[0, 1], size(X_train, 2))
        # indexes_new, res = hillClimbing(ff, indexes_initial, iters=100, debug=true, r = 2)
    # else
    # end
    bruteforce!(model, debug=debug)

end

function Ψ_approx(x, ll_approx, scale_data=true)
    y = zeros(length(ll_approx.ys))

    if scale_data
        x = x .* ll_approx.scale_values
    end
    # x = xx[1,:]

    for i = 1:length(ll_approx.ys)
        y[i] = ll_approx.ys[i](x[view(ll_approx.indexes, i,:)])
    end
    y
end

function bruteforce!(ll_approx; debug=false)
    D = size(ll_approx.X_train, 2)
    D_ll = size(ll_approx.Y_train, 2)

    approx_errors_old = Inf*ones(D_ll)
    indexes_best = zeros(D_ll, D)

    for t = 1:2^D
        x = parse.(Int, collect(bitstring(t)), base=10)
        reverse!(x)

        x_new = Bool.(x[1:D])
        indexes = repeat(x_new', D_ll)

        ll_approx.indexes = indexes

        fit!(ll_approx)
        update_approximation_errors!(ll_approx)

        for i = 1:D_ll
            if ll_approx.approx_errors[i] < approx_errors_old[i]
                approx_errors_old[i] = ll_approx.approx_errors[i]
                indexes_best[i,:] = indexes[i,:]
            end
        end

        if debug
            println(">>>>>>>>>>>> iter: ", t)
            display(indexes_best)
            println("err: ", approx_errors_old)
        end

    end

    ll_approx.indexes = indexes_best

    fit!(ll_approx)
    update_approximation_errors!(ll_approx)

    indexes_best
end


using Compat.Printf
using Compat.LinearAlgebra

include("MBHFMM2D.jl")

using PyPlot

## parameters

lambda = 1.0e-2

ifcharge = false
ifdipole = false
ifquad = false
ifoct = true
ifpot = true
ifgrad = true
ifhess = true

## define source and target points

ns = 10000
nt = 2000

# sources on an ellipse, targets inside

src = zeros(Float64,2,ns)
rnorm = zeros(Float64,2,ns)
dsdt = zeros(Float64,ns)
targ = zeros(Float64,2,nt)

charge = zeros(Float64,1)
dipstr = zeros(Float64,1)
dipvec = zeros(Float64,2,1)
quadstr = zeros(Float64,1)
quadvec = zeros(Float64,3,1)
octstr = zeros(Float64,ns) # octopole strength
octvec = zeros(Float64,4,ns) # octopole "direction"

# storage for solution

pottarg1 = zeros(Float64,nt)
gradtarg1 = zeros(Float64,2,nt)
hesstarg1 = zeros(Float64,3,nt)
pottarg2 = zeros(Float64,nt)
gradtarg2 = zeros(Float64,2,nt)
hesstarg2 = zeros(Float64,3,nt)

center = [6.1;-1.3]

h = 2.0*pi/ns
a = 1.123
b = 0.7

for i = 1:ns
    t = (i-1)*h
    src[:,i] = center + [a*cos(t);b*sin(t)]
    dx = -a*sin(t)
    dy = b*cos(t)
    dsdt[i] = sqrt(dx^2+dy^2)
    rnorm[1,i] = dy/dsdt[i]
    rnorm[2,i] = -dx/dsdt[i]
    nu = rnorm[:,i]
    tau = [-nu[2];nu[1]]
    octvec[:,i] = boxfmm2d_formoctvec(nu,nu,tau)
    octstr[i] = h*dsdt[i]*cos(5*t)
end

for i = 1:nt
    xt = randn(); yt = randn();
    elnorm = sqrt((xt/a)^2 + (yt/b)^2)
    sc = rand();
    xt = sc*xt/elnorm; yt = sc*yt/elnorm
    targ[:,i] = center + [xt;yt]
end

## make parameters structure

fmmpars = MBHFMM2DParams(lambda,src,targ,ifcharge, ifdipole,
                         ifquad, ifoct, charge, dipstr,
                         dipvec, quadstr, quadvec, octstr,
                         octvec, iprec = 3, ifalltarg = false)

## direct calculation

@printf "PERFORMING DIRECT CALCULATION ...\n"

@time mbhfmm2d_direct!(fmmpars,targ,ifpot,pottarg1,ifgrad,
                 gradtarg1,ifhess,hesstarg1)

## set up fmm

# maxboxes determines size of storage arrays, set it big
# for now
# maxnodes determines max number of sources/targets
# per box

@printf "FORMING FMM ...\n"

@time fmmstor = mbhfmm2d_form(fmmpars,maxnodes=30,maxboxes=100000)

@printf "EVALUATING FMM ...\n"

@time mbhfmm2d_targ!(fmmpars,fmmstor,targ,ifpot,pottarg2,ifgrad,
                 gradtarg2,ifhess,hesstarg2)

println("relative error, pot ",
        norm(pottarg1-pottarg2)/norm(pottarg1))
println("relative error, grad ",
        norm(gradtarg1-gradtarg2)/norm(gradtarg1))
println("relative error, hess ",
        norm(hesstarg1-hesstarg2)/norm(hesstarg1))


# include((@__DIR__)*"/../preamble.jl")
function make_model()
# generic parameters
γ₁ = 11.0; β₁ = 1.0; λ₁ = 12.48; θ₁ = 1.6; κ = 2.6
γ₂ = 22.0; β₂ = 1.0; λ₂ = 16.25; θ₂ = 1.0; xˢ = 1.6

# generator
T = [
    -(γ₁+γ₂)        γ₂          γ₁         0.0
        β₂       -(γ₁+β₂)      0.0          γ₁
        β₁          0.0     -(γ₂+β₁)        γ₂
        0.0         β₁          β₂       -(β₁+β₂)
]

# first fluid rates, cᵢ
C = [
    λ₁-θ₁
    λ₁-θ₁
    -θ₁
    -θ₁
]

# second fluid rates, rᵢ(x)
# lower boundary
r_lwr = [
    # λ₂-κ
    # -κ
    λ₂-κ
    -κ
]

# in the interval (0,xˢ)
_size1(h) = (1,Int(round(1.6/h)))
r_below_xˢ(h) = [
    fill(λ₂-θ₂,_size1(h)...)
    fill(-θ₂,_size1(h)...)
    fill(λ₂-θ₂,_size1(h)...)
    fill(-θ₂,_size1(h)...)
]

# in the interval [xˢ,48) (truncate the model at 48)
_size2(h) = (1,Int(round(48.0/h-1.6/h)))
r_above_xˢ(h) = [
    fill(λ₂,_size2(h)...)
    fill(0.0,_size2(h)...)
    fill(λ₂,_size2(h)...)
    fill(0.0,_size2(h)...)
]

r_interior(h) = [r_below_xˢ(h) r_above_xˢ(h)]

# rates at the upper boundary, 48, due to truncation
# choose some sensible values: same as rates just below boundary
r_upr = [λ₂; 0.0]

b = 48.0
first_fluid = BoundedFluidQueue(T,C,b)
model = first_fluid

h = 0.4
nodes = 0.0:0.4:b

ffq(mesh,h) = FluidFluidQueue(DiscretisedFluidQueue(first_fluid,mesh),r_interior(h),r_lwr,r_upr)

# simulation can be simplified if we use a coarser discretiseation
r_interior_sim = [r_below_xˢ(0.4)[:,1] r_above_xˢ(0.4)[:,end]]
# the true model is unbounded, so choose the upper bound of
# the model we simulate from to be large enough that it
# shouldn't affect the simulation dynamics
big_b = 1_000.0
nodes_sim = [0.0;xˢ;big_b]

first_fluid_sim = BoundedFluidQueue(T,C,big_b)
mesh_sim = DGMesh(nodes_sim,1)
dq_sim = DiscretisedFluidQueue(first_fluid_sim,mesh_sim)

ffq_sim = FluidFluidQueue(dq_sim,r_interior_sim,r_lwr,r_upr)
return ffq_sim
end

const model = make_model()


using CUSPARSE
using CUDArt
using Base.Test

m = 25
n = 35
k = 10
blockdim = 5

for (func,func2,typ) in ((:test_csrilu0,:test_csrilu02,:CudaSparseMatrixCSR),
                         (:test_cscilu0,:test_cscilu02,:CudaSparseMatrixCSC))
    @eval begin
        function $func(elty)
            A = rand(elty,m,m)
            d_A = $typ(sparse(A))
            info = CUSPARSE.sv_analysis('N','G','U',d_A,'O')
            d_B = CUSPARSE.ilu0('N',d_A,info,'O')
            h_B = to_host(d_B)
            Alu = lufact(full(A),Val{false})
            Ac = sparse(Alu[:L]*Alu[:U])
            h_A = ctranspose(h_B) * h_B
            @test_approx_eq(h_A.rowval,Ac.rowval)
            @test reduce(&,isfinite(h_A.nzval))
        end

        function $func2(elty)
            A = rand(elty,m,m)
            A += transpose(A)
            A += m * eye(elty,m)
            d_A = $typ(sparse(A))
            d_B = CUSPARSE.ilu02(d_A,'O')
            h_A = to_host(d_B)
            Alu = lufact(full(A),Val{false})
            Ac = sparse(Alu[:L]*Alu[:U])
            h_A = ctranspose(h_A) * h_A
            @test_approx_eq(h_A.rowval,Ac.rowval)
            @test reduce(&,isfinite(h_A.nzval))
        end
    end
end

types = [Float32,Float64,Complex64,Complex128]
for elty in types
    tic()
    test_csrilu0(elty)
    test_cscilu0(elty)
    println("ilu0 took ", toq(), " for ", elty)
    tic()
    test_csrilu02(elty)
    test_cscilu02(elty)
    println("ilu02 took ", toq(), " for ", elty)
end


using Ipopt, OSQP, DataFrames

results_template_alg = DataFrame(n=Int[], lb=Real[], ub=Real[])

function alternatingminimization_matrixcompletion(coords, entries, M, n, numEntriesSampled, k=sqrt(n), gamma=200*sqrt(n)/(numEntriesSampled),
                                            innerEpochs::Int=20, X_ws::Array{Float64, 2}=zeros(n,n))

  X_init=zeros(n,n)

  theIndices=zeros(n,n)
  for t=1:Int(numEntriesSampled)
      theIndices[coords[entries[t],1], coords[entries[t],2]]=1.0
  end
  #X_init=M.*theIndices
  u_t,sigma,=svd(X_ws)
  u_t=u_t[:, 1:k]
  v_t=zeros(n,k)
  ofv_current=1e10
  ofv_prev=1e12

 # Begin alternating min loop

  start_time = time()

      for epoch in 1:innerEpochs #we already have a value of alpha, from warm-start.

        # Minimization over v for a given u
        m=Model(with_optimizer(Mosek.Optimizer))
        @variable(m, v[1:n, 1:k])
        @objective(m, Min, 1.0/(2.0*gamma)*Compat.dot(u_t*v', u_t*v')+0.5*sum((Compat.dot(u_t[coords[entries[t],1],:], v[coords[entries[t],2],:])-M[coords[entries[t],1], coords[entries[t],2]])^2 for t in 1:Int(numEntriesSampled)))
        @suppress optimize!(m)
        v_t.=value.(v)

        # Minimization over u for a given v
        m2=Model(with_optimizer(Mosek.Optimizer))
        @variable(m2, u[1:n, 1:k])
        @objective(m2, Min, 1.0/(2.0*gamma)*Compat.dot(u*v_t', u*v_t')+0.5*sum((Compat.dot(u[coords[entries[t],1],:], v_t[coords[entries[t],2],:])-M[coords[entries[t],1], coords[entries[t],2]])^2 for t in 1:Int(numEntriesSampled)))
        @suppress optimize!(m2)
        u_t.=value.(u)

        X_t=u_t*v_t'
        ofv_current=1.0/(2.0*gamma)*sum(X_t[i,j]^2 for i=1:n for j=1:n)+sum((X_t[coords[entries[t],1], coords[entries[t],2]]-M[coords[entries[t],1], coords[entries[t],2]])^2 for t in 1:Int(numEntriesSampled))

        if abs.(ofv_current-ofv_prev)<1e-4
            break
        end
        ofv_prev=ofv_current

      end
      @show ofv_current
  return u_t*v_t', ofv_current
end

function alternatingminimization_matrixcompletion_solvesdprelax(coords, entries, M, n, numEntriesSampled, k=sqrt(n), gamma=200*sqrt(n)/(numEntriesSampled),
                                            innerEpochs::Int=500)

  theIndices=zeros(n,n)
  for t=1:Int(numEntriesSampled)
      theIndices[coords[entries[t],1], coords[entries[t],2]]=1.0
  end
  X_t=(1.0+1.0/gamma)^(-1.0)*M.*theIndices #Note: no benefit in projecting onto space of rank-k matrices, will obtain samve svd.

  ofv_current=1e10
  ofv_prev=1e12
  ofv_dual=-1e10
  Y_t=zeros(n,n)
  Y_int=zeros(n,n)
  Y_int_prev=zeros(n,n)
  X_avg=zeros(n,n)
  t_k=(1.0+sqrt(5.0))/2.0
  t_k_prev=1.0
  X_int=X_t
  omega=2.0
  X_int_prev=X_t

  eps_tol=1e-3

  results_convergence=similar(results_template_alg, 0)

 # Begin alternating min loop

  start_time = time()

      for epoch in 1:innerEpochs

        # Minimization over Y for a given X_t
        U,sigma,=svd(X_t)
        sum_sigma=sum(sigma)
        if epoch<=innerEpochs-10
            sigma=max.(sigma, (1e-2/epoch)*ones(n))
        else
            sigma=max.(sigma, 1e-8*ones(n))
        end
        sigma.=(sum_sigma/sum(sigma))*sigma

        size(U)

        m=Model(Mosek.Optimizer)
        set_optimizer_attribute(m, "MSK_DPAR_INTPNT_CO_TOL_REL_GAP", 1e-8) #increase tolerance when nearer the optimal solution (i.e. have run more epochs).
        set_optimizer_attribute(m, "MSK_DPAR_INTPNT_CO_TOL_PFEAS", 1e-8)
        set_optimizer_attribute(m, "MSK_DPAR_INTPNT_QO_TOL_DFEAS", 1e-8)
        set_optimizer_attribute(m, "MSK_DPAR_INTPNT_CO_TOL_INFEAS", 1e-8)

        if epoch<=innerEpochs-10
            @variable(m, theta[1:n]>=max((1e-2/epoch), 1e-5))
        else #obtain some lift by relaxing requirement that Y remains "far" away from center
            @variable(m, theta[1:n]>=1e-5)
        end
        @constraint(m, sum(theta)<=k)
        @constraint(m, theta.<=1.0)
        @variable(m, w[1:n])
        # @variable(m, Y[1:n, 1:n])
        # @constraint(m, Y.==sum(theta[i]*U[:,i]*U[:,i]' for i=1:n))
        @constraint(m, imposeSOCP[i=1:n], [w[i]+theta[i]; w[i]-theta[i]; 2.0*sigma[i]] in SecondOrderCone())
        @objective(m, Min, 1.0/(2.0*gamma)*sum(w))#+(n/50)/(2.0*epoch)*Compat.dot(Y-Y_t, Y-Y_t))
         @suppress optimize!(m)
        Y_int=U*Diagonal(value.(theta))*U'
        Y_t=Y_int+(t_k_prev-1.0)/(t_k)*(Y_int-Y_int_prev)
        Y_int_prev=Y_int
        invtheta=1.0./max.(value.(theta), 1e-5)

        X_int, ofv_current, omega=getX_t_new(M, theIndices, U*Diagonal(invtheta)*U', gamma, X_t, n, epoch, omega)
        X_t=X_int+(t_k_prev-1.0)/(t_k)*(X_int-X_int_prev)
        X_int_prev=X_int
        t_k_prev=t_k
        t_k=(1.0+sqrt(4.0*t_k^2+1.0))/2.0

        if epoch%50==0
            # @show ofv_current, epoch
            # Compute dual objective value and use this in the termination criteria?
            theDualBound=getDualBound(M, theIndices, Y_t, gamma, n, k)
            ofv_dual=max(ofv_dual, theDualBound)
            @show ofv_current, epoch, "duality gap is:", (ofv_current-ofv_dual), time()-start_time
        end
        # Remark: major bottleneck is actually the time to build the model; takes 0.4s to build and 0.04s to solve at n=100.
        # Update objective value, this should be non-increasing compared to the previous iterate.
        # Remark: each iterate gives a feasible projection matrix which can be randomly rounded as needed.
        if (ofv_current-ofv_dual)/(abs(ofv_dual)+1e-4)<1e-4 || ofv_prev-ofv_current<1e-3
            break
        end
        # @show ofv_current, ofv_prev
        #push!(results_convergence, [n, ofv_dual,ofv_current])

        ofv_prev=ofv_current

      end
    #CSV.write("plotconvergence_seqaltmin.csv", results_convergence, append=true)

      # Finally, perform greedy rounding on Y_t, alla the sdp relaxation.
      u,v,w=LinearAlgebra.svd(Y_t)
      diag_v=zeros(n)
      diag_v[sortperm(v, rev=true)[1:k]].=1.0
      Y_rounded=u*Diagonal(diag_v)*u'

      m3=Model(with_optimizer(Mosek.Optimizer))
      @variable(m3, X2[1:n, 1:n])
      @objective(m3, Min, 1.0/(2.0*gamma)*sum(X2[i,j]^2 for i=1:n for j=1:n)+0.5*sum((X2[coords[entries[t],1], coords[entries[t],2]]-M[coords[entries[t],1], coords[entries[t],2]])^2 for t in 1:Int(numEntriesSampled)))
      @constraint(m3, X2.==Y_rounded*X2)
      @suppress optimize!(m3)
      X_final=value.(X2)

  return X_final
end

function getX_t_new(M, theIndices, Y_pinv_t, gamma, X_t, n, epoch, omega)

    X_iter=X_t
    X_opt=X_t
    X_prev=X_t
    residual=1e2
    residual_prev=1e3
    gs_tol=max(1e-3/epoch, 1e-5)
    iters=1
    omega=2.0 #Omega>2 is not recommended in any textbooks, but seems to work well here. Remark: omega is effectively global
    while residual>gs_tol && iters<1000 && residual<1e3 #Use Gauss-Seidel approach
        X_prev=X_iter
        L=(1.0/gamma)*(Y_pinv_t)+(1.0/epoch)*Diagonal(ones(n))
        X_iter=(omega*L)\(theIndices.*(M-X_iter)+(1.0/epoch)*X_t-(1.0-omega)*L*X_iter)
        residual_prev=residual
        residual=norm(((1.0/gamma)*(Y_pinv_t*X_iter)+(1.0/epoch)*(X_iter-X_t)+theIndices.*(X_iter-M)))
        if residual_prev<residual-1e-6 #Increasing omega allows convergence to occur with a worse conditioned system, but more slowly.
            omega*=1.5
        end
        iters+=1
    end
    if residual<gs_tol
        X_opt=X_iter
    else #Not going to converge, so use IPopt
        @show "Warning: Using Solver"
        @show residual, omega, iters

        m2=Model(OSQP.Optimizer)
        set_optimizer_attribute(m2, "max_iter", 1e4)
        set_optimizer_attribute(m2, "eps_prim_inf", 1e-3)
        set_optimizer_attribute(m2, "eps_abs", 1e-4)
        @variable(m2, X[1:n, 1:n])
        @constraint(m2, ((1.0/gamma)*(Y_pinv_t*X)+(1.0/epoch)*(X-X_t)+theIndices.*(X-M)).==0.0) #precondition with X_t^\dag for numerical stability.
        @time JuMP.optimize!(m2)
        X_opt=value.(X)
    end

    Residual = theIndices.*(X_opt.-M)
    ofv_current=0.5*Compat.dot(Residual, Residual)+0.5*Compat.dot(X_opt, -Residual)-(0.5/epoch)*Compat.dot(X_opt-X_t, X_t)
    return X_opt, ofv_current, omega
end

function getDualBound(M, theIndices, Y_current, gamma, n, k)
    # Search for optimal alpha, then compute bound
    #m3=Model(Ipopt.Optimizer)
    # set_optimizer_attribute(m3, "linear_solver", "mumps")
    # set_optimizer_attribute(m3, "tol", 1e-4)
    u,v=svd(Y_current)
    u0=u*Diagonal(sqrt.(v))
    m3=Model(Mosek.Optimizer)
        #set_optimizer_attribute(m3, "tol", 1e-4)
    set_optimizer_attribute(m3, "MSK_DPAR_INTPNT_CO_TOL_REL_GAP", 1e-4)
    set_optimizer_attribute(m3, "MSK_DPAR_INTPNT_CO_TOL_DFEAS", 1e-4)
    set_optimizer_attribute(m3, "MSK_DPAR_INTPNT_CO_TOL_INFEAS", 1e-4)
    set_optimizer_attribute(m3, "MSK_DPAR_INTPNT_CO_TOL_MU_RED", 1e-4)
    set_optimizer_attribute(m3, "MSK_DPAR_INTPNT_CO_TOL_NEAR_REL", 1e5)
    set_optimizer_attribute(m3, "MSK_DPAR_INTPNT_CO_TOL_PFEAS", 1e-4)
    set_optimizer_attribute(m3, "MSK_DPAR_INTPNT_QO_TOL_DFEAS", 1e-4)
    set_optimizer_attribute(m3, "MSK_DPAR_INTPNT_QO_TOL_INFEAS", 1e-6)
    set_optimizer_attribute(m3, "MSK_DPAR_INTPNT_QO_TOL_MU_RED", 1e-4)
    set_optimizer_attribute(m3, "MSK_DPAR_INTPNT_QO_TOL_PFEAS", 1e-4)
    set_optimizer_attribute(m3, "MSK_DPAR_INTPNT_QO_TOL_REL_GAP", 1e-4)
    set_optimizer_attribute(m3, "MSK_DPAR_INTPNT_TOL_DFEAS", 1e-4)
    # m3=Model(OSQP.Optimizer)
    # set_optimizer_attribute(m3, "max_iter", 1e4)

    @variable(m3, alpha[1:n, 1:n])
    @variable(m3, r[1:n, 1:n])
    @constraint(m3, r.==alpha'*u0)
    @constraint(m3, alpha.*(ones(n,n)-theIndices).==0.0) #if not in set of indices then eliminate variable
    @objective(m3, Max, -0.5*sum((alpha[i,j]+M[i,j])^2*theIndices[i,j] for i=1:n for j=1:n)-(gamma/2.0)*Compat.dot(r,r))#-(gamma/2.0)*Compat.dot(alpha*alpha', Y_current))
    #-0.5*-(gamma/2.0)*Compat.dot(alpha_t*alpha_t', Y_t)+0.5*sum((M.*theIndices).^2)

    set_start_value.(alpha, rand(n,n))
    @time optimize!(m3)
    alpha_t=value.(alpha)
    u_dual,sigma_dual,=svd(alpha_t*alpha_t')
    dualBound=-0.5*sum((alpha_t[i,j]+M[i,j])^2*theIndices[i,j] for i=1:n for j=1:n)-(gamma/2.0)*sum(sigma_dual[i] for i=1:k)+0.5*sum((M.*theIndices).^2)

    return dualBound
end


export PCAICA, PCA #, FA

ica() = MultivariateStats.ICA(Array{Float64,1}(),
                              Array{Float64,2}(undef,0,0))
pca() = MultivariateStats.PCA(Array{Float64,1}(),
                              Array{Float64,2}(undef,0,0),
                              Array{Float64,1}(),
                              Float64(NaN),
                              Float64(NaN))
"""
mutable struct PCA <: PopulationModel
Principal Components Analysis
"""
mutable struct PCA <: PopulationModel
    k::Int
    pca::MultivariateStats.PCA
    function PCA(k=0::Int, pca=pca())
        new(k,pca)
    end
end

function fit!(model::PCA, Z::AbstractArray; use_marchenko=false)
    model.pca = fit(MultivariateStats.PCA, Z)
    if use_marchenko
        λmin,λmax = marchenko_thresh(size(Z)...)
        model.k = model.k == 0 ? sum(model.pca.prinvars.>λmax) : model.k
    else
        model.k = length(model.pca.prinvars)
    end
    return model
end

function weights(model::PCA)
    V = projection(model.pca)[:,1:model.k]
    normweightvectors!(V)
    return V
end

"""
mutable struct PCAICA <: PopulationModel
See: Detecting cell assemblies in large neuronal populations. Lopes dos Santos et al 2013.
"""
mutable struct PCAICA <: PopulationModel
    k::Int
    pca::MultivariateStats.PCA
    ica::MultivariateStats.ICA
    seed::AbstractRNG
    function PCAICA(k=0::Int, pca=pca(), ica=ica(), seed=Random.seed!(13))
        new(k, pca, ica, seed)
    end
end

function fit!(model::PCAICA, Z::AbstractArray)
    model.pca = fit(MultivariateStats.PCA, Z)
    λmin,λmax = marchenko_thresh(size(Z)...)
    model.k = model.k == 0 ? sum(model.pca.prinvars.>λmax) : model.k
    Psign = projection(model.pca)[:,1:model.k] # projection of first k pc's
    Zproj = transpose(Psign) * Z # project spikes onto pc subspace
    model.ica = fit(MultivariateStats.ICA, Zproj, model.k,
                    winit=rand(model.seed, size(Zproj,1), model.k), # for reproducibility
                    maxiter=10000, tol=1.0e-6)
end

function weights(model::PCAICA)
    Psign = projection(model.pca)[:,1:model.k]
    V = Psign * model.ica.W
    normweightvectors!(V)
    return V
end


using Graphs, SparseArrays

struct SparseMatrixDiGraph{T,TW} <: AbstractGraph{T}
	A::SparseMatrixCSC{Bool, T}
	X::SparseMatrixCSC{T,T}
	W::TW
end

struct IndexedEdge{T}
	src::T
	dst::T
	idx::T
end

function SparseMatrixDiGraph(A::AbstractMatrix{Bool}, W)
	size(A,1) != size(A,2) && throw(ArgumentError("Matrix should be square"))
	any(A[i,i] for i=1:size(A,1)) && throw(ArgumentError("Self loops are not allowed"))
	A = sparse(A)
	length(W) != nnz(A) && throw(ArgumentError("Number of properties different from number of edges"))
	X = sparse(SparseMatrixCSC(A.m, A.n, A.colptr, A.rowval, collect(1:nnz(A)))')
	SparseMatrixDiGraph(A, X, W)
end

function SparseMatrixDiGraph(A::AbstractMatrix)
	SparseMatrixDiGraph(SparseMatrixCSC(A.m, A.n, A.colptr, A.rowval, fill(true, length(A.nzval))), A.nzval)
end

Graphs.edges(g::SparseMatrixDiGraph) = (IndexedEdge{Int}(i,g.A.rowval[k],k) for i=1:size(g.A,2) for k=nzrange(g.A,i))

Base.eltype(g::SparseMatrixDiGraph{T}) where T = T

Graphs.edgetype(g::SparseMatrixDiGraph{T}) where T = IndexedEdge{T}

Graphs.has_edge(g::SparseMatrixDiGraph, i, j) = g.A[i,j]

Graphs.has_vertex(g::SparseMatrixDiGraph, i) = i ∈ 1:size(g.A, 2)

Graphs.inneighbors(g::SparseMatrixDiGraph, i) = @view g.A.rowval[nzrange(g.A,i)]

Graphs.outneighbors(g::SparseMatrixDiGraph, i) = @view g.X.rowval[nzrange(g.X,i)]

inedges(g::SparseMatrixDiGraph, i) = (IndexedEdge{Int}(i,g.A.rowval[k],k) for k=nzrange(g.A,i))

outedges(g::SparseMatrixDiGraph, i) = (IndexedEdge{Int}(g.X.rowval[k],i,g.X.nzval[k]) for k=nzrange(g.X,i))

property(g::SparseMatrixDiGraph, e::IndexedEdge) = g.W[e.idx]

Graphs.ne(g::SparseMatrixDiGraph) = nnz(g.A)

Graphs.nv(g::SparseMatrixDiGraph) = size(g.A,2)

Graphs.vertices(g::SparseMatrixDiGraph) = 1:size(g.A,2)

Graphs.is_directed(::Type{SparseMatrixDiGraph{T,TW}}) where {T,TW} = true

Graphs.is_directed(g::SparseMatrixDiGraph) = true

Base.zero(g::SparseMatrixDiGraph) = SparseMatrixDiGraph(zero(g.A), similar(g.W, 0))


# Genetic Algorithms
# ==================
#         objfun: Objective fitness function
#              N: Search space dimensionality
# initPopulation: Search space dimension ranges as a vector, or initial population values as matrix,
#                 or generation function which produce individual population entities.
# populationSize: Size of the population
#  crossoverRate: The fraction of the population at the next generation, not including elite children,
#                 that is created by the crossover function.
#   mutationRate: Probability of chromosome to be mutated
#              ɛ: Positive integer specifies how many individuals in the current generation
#                 are guaranteed to survive to the next generation.
#                 Floating number specifies fraction of population.
#
function ga(objfun::Function, N::Int;
            initPopulation::Individual = ones(N),
            lowerBounds::Union(Nothing, Vector) = nothing,
            upperBounds::Union(Nothing, Vector) = nothing,
            populationSize::Int = 50,
            crossoverRate::Float64 = 0.8,
            mutationRate::Float64 = 0.1,
            ɛ::Real = 0,
            selection::Function = ((x,n)->1:n),
            crossover::Function = ((x,y)->(y,x)),
            mutation::Function = (x->x),
            iterations::Integer = 100*N,
            tol = 0.0,
            tolIter = 10,
            verbose = false,
            debug = false,
            interim = false)

    store = Dict{Symbol,Any}()

    # Setup parameters
    elite = isa(ɛ, Int) ? ɛ : @compat round(Int, ɛ * populationSize)
    fitFunc = inverseFunc(objfun)

    # Initialize population
    individual = getIndividual(initPopulation, N)
    fitness = zeros(populationSize)
    population = Array(typeof(individual), populationSize)
    offspring = similar(population)

    # Generate population
    for i in 1:populationSize
        if isa(initPopulation, Vector)
            population[i] = initPopulation.*rand(N)
        elseif isa(initPopulation, Matrix)
            population[i] = initPopulation[:, i]
        elseif isa(initPopulation, Function)
            population[i] = initPopulation(N) # Creation function
        else
            error("Cannot generate population")
        end
        fitness[i] = fitFunc(population[i])
        debug && println("INIT $(i): $(population[i]) : $(fitness[i])")
    end
    fitidx = sortperm(fitness, rev = true)
    keep(interim, :fitness, copy(fitness), store)

    # Generate and evaluate offspring
    itr = 1
    bestFitness = 0.0
    bestIndividual = 0
    fittol = 0.0
    fittolitr = 1
    while true
        debug && println("BEST: $(fitidx)")

        # Select offspring
        selected = selection(fitness, populationSize)

        # Perform mating
        offidx = randperm(populationSize)
        for i in 1:2:populationSize
            j = (i == populationSize) ? i-1 : i+1
            if rand() < crossoverRate
                debug && println("MATE $(offidx[i])+$(offidx[j])>: $(population[selected[offidx[i]]]) : $(population[selected[offidx[j]]])")
                offspring[i], offspring[j] = crossover(population[selected[offidx[i]]], population[selected[offidx[j]]])
                debug && println("MATE >$(offidx[i])+$(offidx[j]): $(offspring[i]) : $(offspring[j])")
            else
                offspring[i], offspring[j] = population[selected[i]], population[selected[j]]
            end
        end

        # Perform mutation
        for i in 1:populationSize
            if rand() < mutationRate
                debug && println("MUTATED $(i)>: $(offspring[i])")
                mutation(offspring[i])
                debug && println("MUTATED >$(i): $(offspring[i])")
            end
        end

        # Elitism
        if elite > 0
            for i in 1:elite
                subs = rand(1:populationSize)
                debug && println("ELITE $(fitidx[i])=>$(subs): $(population[fitidx[i]]) => $(offspring[subs])")
                offspring[subs] = population[fitidx[i]]
            end
        end

        # New generation
        for i in 1:populationSize
            population[i] = offspring[i]
            fitness[i] = fitFunc(offspring[i])
            debug && println("FIT $(i): $(fitness[i])")
        end
        fitidx = sortperm(fitness, rev = true)
        bestIndividual = fitidx[1]
        curGenFitness = @compat Float64(objfun(population[bestIndividual]))
        fittol = abs(bestFitness - curGenFitness)
        bestFitness = curGenFitness

        keep(interim, :fitness, copy(fitness), store)
        keep(interim, :bestFitness, bestFitness, store)

        # Verbose step
        verbose &&  println("BEST: $(bestFitness): $(population[bestIndividual]), G: $(itr)")

        # Terminate:
        #  if fitness tolerance is met for specified number of steps
        if fittol < tol
            if fittolitr > tolIter
                break
            else
                fittolitr += 1
            end
        else
            fittolitr = 1
        end
        # if number of iterations more then specified
        if itr >= iterations
            break
        end
        itr += 1
    end

    return population[bestIndividual], bestFitness, itr, fittol, store
end