text:
export greedyfit, gradient_off!, update_component!

"""
$(TYPEDSIGNATURES)

An adaptive routine to estimate a sparse approximation of an `SparseRadialMapComponent` based on  the pair of ensemble matrices `X` (training set) and `Xvalid` (validation set).
"""
function greedyfit(Nx::Int64, poff::Int64, pdiag::Int64,  X::AbstractMatrix{Float64}, Xvalid::AbstractMatrix{Float64}, maxfamiliesoff::Int64,
                   λ::Float64, δ::Float64, γ::Float64; maxpatience::Int64 = 10^5, verbose::Bool=false)

    train_error = Float64[]
    valid_error = Float64[]

    NxX, Ne = size(X)
    # The widths and centers are computed on the entire set.
    Xsort = deepcopy(sort(hcat(X, Xvalid); dims = 2))
    @assert poff > -1 || pdiag > -1 "The order poff and pdiag of the features must be > 0"
    @assert λ == 0 "Greedy fit is only implemented for λ = 0"
    @assert NxX == Nx "Wrong dimension of the ensemble matrix `X`"

    # Initialize a sparse radial map component C with only a diagonal term of order p
    order = fill(-1, Nx)
    order[end] = pdiag
    C = SparseRadialMapComponent(Nx, order)

    center_std!(C, Xsort; γ = γ)
    C, error = optimize(C, X, nothing, λ, δ)
    x_diag = extractcoeff(C)

    # Compute loss on training and validation sets
    push!(train_error, copy(error))
    push!(valid_error, copy(negative_likelihood(C, Xvalid, λ, δ)))

    best_valid_error = Inf
    patience = 0

    # Update patience
    if valid_error[end] >= best_valid_error
        patience +=1
    else
        best_valid_error = copy(valid_error[end])
        patience = 0
    end

    if verbose == true
        println("Active dimensions: "*string(C.activedim)*" - Training error: "*
        string(train_error[end])*", Validation error: "*string(valid_error[end]))
    end

    if Nx>1 || maxfamiliesoff>0

        # Create a radial map with order p for all the entries
        Cfull = SparseRadialMapComponent(Nx, vcat(poff*ones(Int64, Nx-1), pdiag))

        # Compute centers and widths
        center_std!(Cfull, Xsort; γ = γ)

        # Create weights for the training and validation sets
        ψ_off, ψ_diag, dψ_diag = compute_weights(Cfull, X)
        ψ_offvalid, ψ_diagvalid, dψ_diagvalid = compute_weights(Cfull, Xvalid)

        n_off = size(ψ_off,2)
        n_diag = size(ψ_diag,2)

        Asqrt = zero(ψ_diag)
        # Normalize diagtone basis functions
        μψ = copy(mean(ψ_diag, dims=1)[1,:])
        σψ = copy(std(ψ_diag, dims = 1, corrected = false)[1,:])
        # σψ = copy(norm.(eachcol(ψ_diag)))
        ψ_diagscaled = copy(ψ_diag)
        dψ_diagscaled = copy(dψ_diag)

        ψ_diagscaled .-= μψ'
        ψ_diagscaled ./= σψ'
        dψ_diagscaled ./= σψ'

        lhd = LHD(zeros(n_diag,n_diag), dψ_diagscaled, λ, δ)

        ψ_offscaled = copy(ψ_off)
        μψ_off = copy(mean(ψ_off, dims = 1)[1,:])
        # σψ_off = copy(std(ψ_off, dims = 2, corrected = false)
        # σψ_off = copy(norm.(eachcol(ψ_off)))
        σψ_off = copy(std(ψ_off, dims = 1, corrected = false)[1,:])

        ψ_offscaled .-= μψ_off'
        ψ_offscaled ./= σψ_off'

        # rhs = -ψ_diag x_diag
        rhs = zeros(Ne)
        mul!(rhs, ψ_diag, x_diag[2:n_diag+1])
        rhs .+= x_diag[1]
        rmul!(rhs, -1.0)

        # Create updatable QR factorization
        # For the greedy optimization, we don't use a L2 regularization λ||x||^2,
        # since we are already making a greedy selection of the features
        F = qrfactUnblocked(zeros(0,0))
        # Off-diagonal coefficients will be permuted based on the greedy procedure

        candidates = collect(1:Nx-1)
        # Unordered off-diagonal active dimensions
        offdims = Int64[]

        # Compute the gradient of the different basis
        dJ = zeros((poff+1)*(Nx-1))

        x_off = zeros((poff+1)*(Nx-1))
        x_offsparse = Float64[]
        tmp_off = Float64[]
        tmp_diag = zeros(n_diag)

        budget = min(maxfamiliesoff, Nx-1)
        # Compute the norm of the different candidate features
        sqnormfeatures = map(i-> norm(view(ψ_off, (i-1)*(poff+1)+1:i*(poff+1)))^2, candidates)
        cache = zeros(Ne)
        @inbounds for i=1:budget
            # Compute the gradient of the different basis (use the unscaled basis evaluations)
            mul!(rhs, ψ_diag, view(x_diag,2:n_diag+1))
            rhs .+= x_diag[1]
            rmul!(rhs, -1.0)
            gradient_off!(dJ, cache, ψ_off, x_off, rhs, Ne)

            dJ_dim = map(k-> norm(view(dJ, (k-1)*(poff+1)+1:k*(poff+1)))^2/sqnormfeatures[k], candidates)
            _, max_dim = findmax(dJ_dim)
            new_dim = candidates[max_dim]
            push!(offdims, copy(new_dim))
            append!(tmp_off, zeros(poff+1))
            append!(x_offsparse, zeros(poff+1))

            # Update storage in C
            update_component!(C, poff, new_dim)

            # Compute center and std for this new family of features
            # The centers and widths have already been computed in Cfull
            copy!(C.ξ[new_dim], Cfull.ξ[new_dim])
            copy!(C.σ[new_dim], Cfull.σ[new_dim])

            # Update qr factorization with the new family of features
            if i == 1
                F = qrfactUnblocked(ψ_offscaled[:,(new_dim-1)*(poff+1)+1:new_dim*(poff+1)])
            else
                F = updateqrfactUnblocked!(F, view(ψ_offscaled,:,(new_dim-1)*(poff+1)+1:new_dim*(poff+1)))
            end
            Asqrt .= ψ_diagscaled - fast_mul2(ψ_diagscaled, F.Q, Ne, (poff+1)*size(offdims, 1))
            lhd.A .= (1/Ne)*Asqrt'*Asqrt

            if C.p[Nx] == 0
                @assert size(lhd.A)==(1,1) "Quadratic matrix should be a scalar."
                # This equation is equation (A.9) Couplings for nonlinear ensemble filtering
                # uNx(z) = c + α z so α = 1/√κ*
                fill!(tmp_diag, 1.0)
                tmp_diag[1] = sqrt(1/lhd.A[1,1])
            else
                # Update A and b of the loss function lhd
                # Add L-2 regularization

                fill!(tmp_diag, 1.0)

                @inbounds for i=1:n_diag
                    lhd.A[i,i] += (λ/Ne)
                end
                # Update b coefficient
                lhd.b .= δ*sum(lhd.A, dims=2)[:,1]

                tmp_diag,_,_,_ = projected_newton(tmp_diag, lhd, "TrueHessian")
                tmp_diag .+= δ
            end

            cache .= ψ_diagscaled*tmp_diag
            tmp_off .= -F.R\((F.Q'*cache)[1:(poff+1)*size(offdims, 1)])
            # Rescale coefficients
            for j=1:n_diag

                x_diag[j+1] = tmp_diag[j]/σψ[j]
            end

            for (j, offdimj) in enumerate(offdims)
                tmp_off[(j-1)*(poff+1)+1:j*(poff+1)] ./= σψ_off[(offdimj-1)*(poff+1)+1:offdimj*(poff+1)]
            end

            # Compute and store the constant term (expectation of selected terms)
            x_diag[1] = -dot(μψ, x_diag[2:n_diag+1]) #-dot(μψ_off, tmp_off)

            # Make sure that active dim are in the right order when we affect coefficient.
            # For the split and kfold compute the training and validation losses.
            fill!(x_offsparse, 0.0)
            # Compute the permutation associated to offdims
            perm = sortperm(offdims)

            for (j, offdimj) in enumerate(offdims)
                tmp_offj = tmp_off[(j-1)*(poff+1)+1:j*(poff+1)]
                x_diag[1] -= dot(view(μψ_off, (offdimj-1)*(poff+1)+1:offdimj*(poff+1)), tmp_offj)
                view(x_off, (offdimj-1)*(poff+1)+1:offdimj*(poff+1)) .= tmp_offj
            end

            for (j, offdimj) in enumerate(offdims)
                tmp_offj = tmp_off[(perm[j]-1)*(poff+1)+1:perm[j]*(poff+1)]
                view(x_offsparse, (j-1)*(poff+1)+1:j*(poff+1)) .= tmp_offj
            end

            @assert norm(x_off[x_off .!= 0.0] - x_offsparse[x_offsparse .!= 0.0])<1e-10  "Error in x_off"

            modifycoeff!(C, vcat(x_offsparse, x_diag))
            filter!(x-> x!= new_dim, candidates)

            # Compute loss on training and validation sets
            push!(train_error, copy(negative_likelihood(C, X, λ, δ)))
            push!(valid_error, copy(negative_likelihood(C, Xvalid, λ, δ)))

            if verbose == true
                println("Active dimensions: "*string(C.activedim)*" - Training error: "*
                string(train_error[end])*", Validation error: "*string(valid_error[end]))
            end

            # Update patience
            if valid_error[end] >= best_valid_error
                patience +=1
            else
                best_valid_error = copy(valid_error[end])
                patience = 0
            end

            # Check if patience exceeded maximum patience
            if patience >= maxpatience
                break
            end
        end
    end
    return C, (train_error, valid_error)
end

greedyfit(Nx::Int64, p::Int64, X::AbstractMatrix{Float64}, Xvalid::AbstractMatrix{Float64},
          maxfamiliesoff::Int64, λ::Float64, δ::Float64, γ::Float64;
          maxpatience::Int64 = 10^5, verbose::Bool=false) =
          greedyfit(Nx, p, p, X, Xvalid, maxfamiliesoff, λ, δ, γ; maxpatience = maxpatience, verbose = verbose)

"""
$(TYPEDSIGNATURES)

An adaptive routine to estimate a sparse approximation of an `SparseRadialMapComponent` based on  the ensemble matrix `X`.
"""
function greedyfit(Nx::Int64, poff::Int64, pdiag::Int64, X::AbstractMatrix{Float64},
                   maxfamiliesoff::Int64, λ::Float64, δ::Float64, γ::Float64; maxpatience::Int64 = 10^5, verbose::Bool = false)

    NxX, Ne = size(X)
    Xsort = deepcopy(sort(X; dims = 2))
    train_error = Float64[]
    @assert poff > -1 || pdiag > -1 "The order poff and pdiag of the features must be > 0"
    @assert λ == 0 "Greedy fit is only implemented for λ = 0"
    @assert NxX == Nx "Wrong dimension of the ensemble matrix `X`"

    # Initialize a sparse radial map component C with only a diagonal term of order p
    order = fill(-1, Nx)
    order[end] = pdiag
    C = SparseRadialMapComponent(Nx, order)

    center_std!(C, Xsort; γ = γ)
    C, error = optimize(C, X, nothing, λ, δ)
    x_diag = extractcoeff(C)

    push!(train_error, copy(error))

    best_train_error = Inf
    patience = 0

    # Update patience
    if train_error[end] >= best_train_error
        patience +=1
    else
        best_train_error = copy(train_error[end])
        patience = 0
    end

    if verbose == true
        # Compute loss on training and validation sets
        println("Active dimensions: "*string(C.activedim)*" - Training error: "*
                string(train_error[end]))
    end

    if Nx>1 || maxfamiliesoff>0
        # Create a radial map with order p for all the entries
        Cfull = SparseRadialMapComponent(Nx, vcat(poff*ones(Int64, Nx-1), pdiag))

        # Compute centers and widths
        center_std!(Cfull, Xsort; γ = γ)

        # Create weights
        ψ_off, ψ_diag, dψ_diag = compute_weights(Cfull, X)
        n_off = size(ψ_off,2)
        n_diag = size(ψ_diag,2)

        Asqrt = zero(ψ_diag)
        # Normalize diagtone basis functions
        μψ = copy(mean(ψ_diag, dims=1)[1,:])
        σψ = copy(std(ψ_diag, dims = 1, corrected = false)[1,:])
        # σψ = copy(norm.(eachcol(ψ_diag)))
        ψ_diagscaled = copy(ψ_diag)
        dψ_diagscaled = copy(dψ_diag)

        ψ_diagscaled .-= μψ'
        ψ_diagscaled ./= σψ'
        dψ_diagscaled ./= σψ'

        lhd = LHD(zeros(n_diag,n_diag), dψ_diagscaled, λ, δ)

        ψ_offscaled = copy(ψ_off)
        μψ_off = copy(mean(ψ_off, dims = 1)[1,:])
        # σψ_off = copy(std(ψ_off, dims = 2, corrected = false)
        # σψ_off = copy(norm.(eachcol(ψ_off)))
        σψ_off = copy(std(ψ_off, dims = 1, corrected = false)[1,:])

        ψ_offscaled .-= μψ_off'
        ψ_offscaled ./= σψ_off'

        # rhs = -ψ_diag x_diag
        rhs = zeros(Ne)
        mul!(rhs, ψ_diag, x_diag[2:n_diag+1])
        rhs .+= x_diag[1]
        rmul!(rhs, -1.0)

        # Create updatable QR factorization
        # For the greedy optimization, we don't use a L2 regularization λ||x||^2,
        # since we are already making a greedy selection of the features
        F = qrfactUnblocked(zeros(0,0))
        # Off-diagonal coefficients will be permuted based on the greedy procedure

        candidates = collect(1:Nx-1)
        # Unordered off-diagonal active dimensions
        offdims = Int64[]

        # Compute the gradient of the different basis
        dJ = zeros((poff+1)*(Nx-1))

        x_off = zeros((poff+1)*(Nx-1))
        x_offsparse = Float64[]
        tmp_off = Float64[]
        tmp_diag = zeros(n_diag)

        budget = min(maxfamiliesoff, Nx-1)
        # Compute the norm of the different candidate features
        sqnormfeatures = map(i-> norm(view(ψ_off, (i-1)*(poff+1)+1:i*(poff+1)))^2, candidates)
        cache = zeros(Ne)
        @inbounds for i=1:budget
            # Compute the gradient of the different basis (use the unscaled basis evaluations)
            mul!(rhs, ψ_diag, view(x_diag,2:n_diag+1))
            rhs .+= x_diag[1]
            rmul!(rhs, -1.0)
            gradient_off!(dJ, cache, ψ_off, x_off, rhs, Ne)
            # @show dJ
            dJ_dim = map(k-> norm(view(dJ, (k-1)*(poff+1)+1:k*(poff+1)))^2/sqnormfeatures[k], candidates)
            _, max_dim = findmax(dJ_dim)
            new_dim = candidates[max_dim]
            # @show new_dim
            push!(offdims, copy(new_dim))
            append!(tmp_off, zeros(poff+1))
            append!(x_offsparse, zeros(poff+1))

            # Update storage in C
            update_component!(C, poff, new_dim)

            # Compute center and std for this new family of features
            # The centers and widths have already been computed in Cfull
            copy!(C.ξ[new_dim], Cfull.ξ[new_dim])
            copy!(C.σ[new_dim], Cfull.σ[new_dim])

            # Update qr factorization with the new family of features
            if i == 1
                F = qrfactUnblocked(ψ_offscaled[:,(new_dim-1)*(poff+1)+1:new_dim*(poff+1)])
            else
                F = updateqrfactUnblocked!(F, view(ψ_offscaled,:,(new_dim-1)*(poff+1)+1:new_dim*(poff+1)))
            end

            Asqrt .= ψ_diagscaled - fast_mul2(ψ_diagscaled, F.Q, Ne, (poff+1)*size(offdims, 1))
            lhd.A .= (1/Ne)*Asqrt'*Asqrt
            # @show norm(lhd.A)
            if C.p[Nx] == 0
                @assert size(lhd.A)==(1,1) "Quadratic matrix should be a scalar."
                # This equation is equation (A.9) Couplings for nonlinear ensemble filtering
                # uNx(z) = c + α z so α = 1/√κ*
                fill!(tmp_diag, 1.0)
                tmp_diag[1] = sqrt(1/lhd.A[1,1])
            else
                # Update A and b of the loss function lhd
                # Add L-2 regularization

                fill!(tmp_diag, 1.0)

                @inbounds for i=1:n_diag
                    lhd.A[i,i] += (λ/Ne)
                end
                # Update b coefficient
                lhd.b .= δ*sum(lhd.A, dims=2)[:,1]

                tmp_diag,_,_,_ = projected_newton(tmp_diag, lhd, "TrueHessian")
                tmp_diag .+= δ
            end

            cache .= ψ_diagscaled*tmp_diag
            tmp_off .= -F.R\((F.Q'*cache)[1:(poff+1)*size(offdims, 1)])
            # Rescale coefficients
            for j=1:n_diag
                x_diag[j+1] = tmp_diag[j]/σψ[j]
            end

            for (j, offdimj) in enumerate(offdims)
                tmp_off[(j-1)*(poff+1)+1:j*(poff+1)] ./= σψ_off[(offdimj-1)*(poff+1)+1:offdimj*(poff+1)]
            end

            # Compute and store the constant term (expectation of selected terms)
            x_diag[1] = -dot(μψ, x_diag[2:n_diag+1]) #-dot(μψ_off, tmp_off)

            # Make sure that active dim are in the right order when we affect coefficient.
            # For the split and kfold compute the training and validation losses.
            fill!(x_offsparse, 0.0)
            # Compute the permutation associated to offdims
            perm = sortperm(offdims)

            for (j, offdimj) in enumerate(offdims)
                tmp_offj = tmp_off[(j-1)*(poff+1)+1:j*(poff+1)]
                x_diag[1] -= dot(view(μψ_off, (offdimj-1)*(poff+1)+1:offdimj*(poff+1)), tmp_offj)
                view(x_off, (offdimj-1)*(poff+1)+1:offdimj*(poff+1)) .= tmp_offj
            end

            for (j, offdimj) in enumerate(offdims)
                tmp_offj = tmp_off[(perm[j]-1)*(poff+1)+1:perm[j]*(poff+1)]
                view(x_offsparse, (j-1)*(poff+1)+1:j*(poff+1)) .= tmp_offj
            end

            @assert norm(x_off[x_off .!= 0.0] - x_offsparse[x_offsparse .!= 0.0])<1e-10  "Error in x_off"

            modifycoeff!(C, vcat(x_offsparse, x_diag))
            push!(train_error, copy(negative_likelihood(C, X, λ, δ)))
            filter!(x-> x!= new_dim, candidates)

            if verbose == true
                # Compute loss on training and validation sets
                println("Active dimensions: "*string(C.activedim)*" - Training error: "*
                        string(train_error[end]))
            end

            # Update patience
            if train_error[end] >= best_train_error
                patience +=1
            else
                best_train_error = copy(train_error[end])
                patience = 0
            end

            # Check if patience exceeded maximum patience
            if patience >= maxpatience
                break
            end
        end
    end
    return C, train_error
end

greedyfit(Nx::Int64, p::Int64, X::AbstractMatrix{Float64}, maxfamiliesoff::Int64,
          λ::Float64, δ::Float64, γ::Float64; maxpatience::Int64 = 10^5, verbose::Bool=false) =
          greedyfit(Nx, p, p, X, maxfamiliesoff, λ, δ, γ; maxpatience = maxpatience, verbose = verbose)

function gradient_off!(dJ::AbstractVector{Float64}, cache::AbstractVector{Float64}, ψ_off::AbstractMatrix{Float64}, x_off, rhs, Ne::Int64)
    fill!(dJ, 0.0)
    cache .= ψ_off*x_off - rhs
    dJ .= (1/Ne)*ψ_off'*cache
end

gradient_off!(dJ::AbstractVector{Float64}, ψ_off::AbstractMatrix{Float64}, x_off, rhs, Ne::Int64) =
              gradient_off!(dJ, zeros(Ne), ψ_off, x_off, rhs, Ne)

function update_component!(C::SparseRadialMapComponent, p::Int64, new_dim::Int64)
    @assert new_dim <= C.Nx
    if new_dim == C.Nx
        if p == -1
            C.p[new_dim] = p
            C.ξ[new_dim] = Float64[]
            C.σ[new_dim] = Float64[]
            C.coeff[new_dim] = Float64[]
        elseif p == 0
            C.p[new_dim] = p
            push!(C.activedim, new_dim)
            sort!(C.activedim)
            C.ξ[new_dim] = zeros(p)
            C.σ[new_dim] = zeros(p)
            C.coeff[new_dim] = zeros(p+2)
        else
            C.p[new_dim] = p
            push!(C.activedim, new_dim)
            sort!(C.activedim)
            C.ξ[new_dim] = zeros(p+2)
            C.σ[new_dim] = zeros(p+2)
            C.coeff[new_dim] = zeros(p+3)
        end
    else
        if p == -1
            C.p[new_dim] = p
            C.ξ[new_dim] = Float64[]
            C.σ[new_dim] = Float64[]
            C.coeff[new_dim] = Float64[]
        elseif p == 0
            C.p[new_dim] = p
            push!(C.activedim, new_dim)
            sort!(C.activedim)
            C.ξ[new_dim] = zeros(p)
            C.σ[new_dim] = zeros(p)
            C.coeff[new_dim] = zeros(p+1)
        else
            C.p[new_dim] = p
            push!(C.activedim, new_dim)
            sort!(C.activedim)
            C.ξ[new_dim] = zeros(p)
            C.σ[new_dim] = zeros(p)
            C.coeff[new_dim] = zeros(p+1)
        end
    end
end

    # maxfmaily is the maximal number of
    # if p == 0
    #     maxfamily = ceil(Int64, (sqrt(Ne)-(p+1))/(p+1))
    # elseif p > 0
    #     maxfamily = ceil(Int64, (sqrt(Ne)-(p+3))/(p+1))
    # else
    #     error("Wrong value for p")
    # end


text:
using Deconvolution
using Test, Random, FFTW

Random.seed!(42) # Fixed random seed

##### Wiener deconvolution

x = range(0, stop = 10, length = 15)
s = sinpi.(x) .- 1.5cos.(x)
n = rand(length(s))
blurring_ft  = exp.(-0.001 * (range(-div(length(x), 2), length = length(x)) .^ 2) .^ (5 // 6))
blurred_s_ft = fftshift(blurring_ft) .* fft(s) .+ fft(n)
blurred_s    = real(ifft(blurred_s_ft))
blurring     = ifft(fftshift(blurring_ft))
@testset "No blurring with noise scalar parameter" begin
    @test wiener(s + n, s, 0.5) ≈
        [-1.1111473161507488, -0.17041312679112874, -1.3469825775608049, 1.2447585407259756,
          2.578860561637874, 1.09191362235022, 1.5339504088753226, -0.34582554585559355,
          -1.316232389748769, -0.21271451029702804, -0.6157543266519075, -0.3338395657941805,
          2.1195369797597614, 1.060270051101612, 1.4446941593151053]
end

@testset "Blurring with noise scalar parameter" begin
    @test wiener(blurred_s, s, 0.5, blurring) ≈
        [-1.1086114591808454, -0.17138676636801406, -1.3482585663111306, 1.2427663740234156,
         2.5813314219359276, 1.0940885915582341, 1.532411498782748, -0.34834962964626176,
         -1.311080379275804, -0.21869338550577228, -0.6084828281998079, -0.3366932366271657,
         2.117735346098039, 1.0624233782615877, 1.4429704301669912]
end

@testset "No blurring with noise array" begin
    @test wiener(blurred_s, s, ifftshift(n)) ≈ [-1.4080797956919306, -0.49286583995619976,
        -1.6205631520484833, 0.7973315264441291, 2.041487785980698, 0.7068033074341512,
        1.303796832927927, -0.5691518774600562, -1.7799047917607576, -0.7841432287697161,
        -0.9149839099122377, -0.4555342804458623, 1.761794603633166, 0.5975996848388563, 0.9542148100114622]
end

@testset "Blurring with noise array" begin
    @test wiener(blurred_s, s, ifftshift(n), blurring) ≈ [-1.4250670444336333, -0.48221282618655026,
        -1.636981677852502, 0.8020690789792746, 2.053748123047759, 0.6974830808841243, 1.3177198257314031,
        -0.5714565338646487, -1.7922287374907637, -0.7809316185676288, -0.9185450961060387,
        -0.46118637814202684, 1.7785714975739182, 0.5887392863050276, 0.967947374977039]
end

##### Richardson-Lucy deconvolution

@testset "Richardson-Lucy deconvolution tests" begin
    k=-0.001
    blurring_ft  = exp.(-k * (range(-div(length(x), 2), length = length(x)) .^ 2) .^ (5 // 6))
    blurred_s_ft = fftshift(blurring_ft) .* fft(s)
    blurred_s    = real(ifft(blurred_s_ft))
    blurring     = ifft(fftshift(blurring_ft))

    estimated = lucy(blurred_s, blurring, iterations=1)

    @test sum(abs.(s .- estimated)) < sum(abs.(s .- blurred_s))

    k =-0.0015 # should improve also when blur params are not perfectly estimated
    blurring_ft  = exp.(-k * (range(-div(length(x), 2), length = length(x)) .^ 2) .^ (5 // 6))
    blurring     = ifft(fftshift(blurring_ft))

    estimated = lucy(blurred_s, blurring, iterations=1)

    @test sum(abs.(s .- estimated)) < sum(abs.(s .- blurred_s))
end

ext: jl


text:
# https://adventofcode.com/2020/day/16

# You collect the rules for ticket fields, the numbers on your ticket, and the numbers on other
# nearby tickets for the same train service (via the airport security cameras) together into
# a single document you can reference (your puzzle input).
#
# The rules for ticket fields specify a list of fields that exist somewhere on the ticket and
# the valid ranges of values for each field. For example, a rule like class: 1-3 or 5-7 means
# that one of the fields in every ticket is named class and can be any value in the ranges
# 1-3 or 5-7 (inclusive, such that 3 and 5 are both valid in this field, but 4 is not).
#
# Start by determining which tickets are completely invalid; these are tickets that contain
# values which aren't valid for any field. Ignore your ticket for now.

example1 = "
class: 1-3 or 5-7
row: 6-11 or 33-44
seat: 13-40 or 45-50

your ticket:
7,1,14

nearby tickets:
7,3,47
40,4,50
55,2,20
38,6,12
"

split_lines(str) = split(str, '\n', keepempty=false)

function parse_range(string::AbstractString)::AbstractRange
    from, to = parse.(Int, split(strip(string), '-', keepempty=false))
    return from:to
end

@assert parse_range(" 9-13233 ") == 9:13233

function parse_rule(string::AbstractString)
    m = match(r"^([a-z ]*): ([\d or-]*)$", string)
    if isnothing(m)
        error("parsing error")
    end
    name, rule_string = m.captures
    ranges = map(parse_range, split(rule_string, "or", keepempty=false))
    return name, ranges
end

@assert parse_rule("row: 6-11 or 33-44") == ("row", [6:11, 33:44])
@assert parse_rule("foo: 1-100") == ("foo", [1:100])
@assert parse_rule("bar: 6-11 or 33-44 or 55-66 ") == ("bar", [6:11, 33:44, 55:66])

function parse_ticket(string::AbstractString)
    return parse.(Int, split(string, ',', keepempty=false))
end

function parse_input(input::AbstractString)
    rules, my_ticket, nearby_tickets = split(input, "\n\n", keepempty=false)
    rules = Dict(map(parse_rule, split_lines(rules)))
    my_ticket = parse_ticket(split_lines(my_ticket)[2])
    nearby_tickets = map(parse_ticket, split_lines(nearby_tickets)[2:end])
    return rules, my_ticket, nearby_tickets
end

function collect_rules(rules::AbstractDict)
    ranges = []
    for rule in values(rules)
        for range in rule
            push!(ranges, range)
        end
    end
    return ranges
end

@assert collect_rules(parse_input(example1)[1]) == [1:3, 5:7, 6:11, 33:44, 13:40, 45:50]

function invalid_fields(ticket, valid_ranges)
    values = Integer[]
    for value in ticket
        valid = false
        for range in valid_ranges
            valid |= (value in range)
        end
        !valid && push!(values, value)
    end
    return values
end

@assert invalid_fields([7, 3, 47], [1:3, 5:7, 6:11, 33:44, 13:40, 45:50]) == []
@assert invalid_fields([40, 4, 50], [1:3, 5:7, 6:11, 33:44, 13:40, 45:50]) == [4]
@assert invalid_fields([55, 2, 20], [1:3, 5:7, 6:11, 33:44, 13:40, 45:50]) == [55]
@assert invalid_fields([38, 6, 12], [1:3, 5:7, 6:11, 33:44, 13:40, 45:50]) == [12]

function part1(input)
    rules, my_ticket, nearby_tickets = parse_input(input)
    valid_ranges = collect_rules(rules)

    invalid = Integer[]
    for ticket in nearby_tickets
        append!(invalid, invalid_fields(ticket, valid_ranges))
    end
    return sum(invalid)
end

@assert part1(example1) == 71

example2 = "
class: 0-1 or 4-19
row: 0-5 or 8-19
seat: 0-13 or 16-19

your ticket:
11,12,13

nearby tickets:
3,9,18
15,1,5
5,14,9
"

function isvalid(value, rule)
    for range in rule
        if value in range
            return true
        end
    end
    return false
end

@assert isvalid(9, [0:13, 16:19])
@assert !isvalid(15, [0:13, 16:19])

function trypop!(collection, key)
    try
        pop!(collection, key)
    catch KeyError
    end
    return collection
end

@assert trypop!(Set([:a, :b, :c]), :c) == Set([:a, :b])
@assert trypop!(Set([:a, :b, :c]), :d) == Set([:a, :b, :c])

function guess_fields(rules, nearby_tickets)
    number_of_fields = length(rules)
    valid_ranges = collect_rules(rules)

    valid_tickets = filter(ticket -> isempty(invalid_fields(ticket, valid_ranges)), nearby_tickets)
    @assert !isempty(valid_tickets)

    # per each field, leave only the rules that are consistent with the data
    field_candidates = [Set(keys(rules)) for _ = 1:number_of_fields]
    for ticket in valid_tickets
        for (i, field) in enumerate(ticket)
            for (name, rule) in rules
                if !isvalid(field, rule)
                    pop!(field_candidates[i], name)
                end
            end
        end
    end

    # eliminate the rules, assumming that there is a unique mapping
    field_names = Array{String,1}(undef, number_of_fields)
    for _ = 1:number_of_fields
        pos = findfirst(map(length, field_candidates) .== 1)
        name = first(field_candidates[pos])  # extract the only element
        field_names[pos] = name
        map(field -> trypop!(field, name), field_candidates)
    end
    return field_names
end

@assert guess_fields(parse_input(example2)[[1, 3]]...) == ["row", "class", "seat"]

function part2(input)
    rules, my_ticket, nearby_tickets = parse_input(input)
    ticket_fields = guess_fields(rules, nearby_tickets)

    result = 1
    for (i, name) in enumerate(ticket_fields)
        if startswith(name, "departure")
            result *= my_ticket[i]
        end
    end
    return result
end

test = read("data/day-16.txt", String)
println("Part 1: $(result1 = part1(test))")
println("Part 2: $(result2 = part2(test))")

@assert result1 == 19093
@assert result2 == 5311123569883

max_stars_repo_name: twolodzko/advent-of-code

max_issues_repo_name: twolodzko/advent-of-code

max_forks_repo_name: twolodzko/advent-of-code


text:
using Gadfly, DataFrames, DelimitedFiles

set_default_plot_size(6inch, 3inch)

dlm, varnames = readdlm(joinpath(dirname(@__DIR__), "data","issue120.csv"),
                        ',', header=true)
df = DataFrame(dlm, vec(Symbol.(varnames)))
plot(df, x=:x1, y=:x2)


text:
# Auxiliary functions for TimeModeling module
# Author: Philipp Witte, pwitte@eos.ubc.ca
# Date: September 2016
#
# Mathias Louboutin, mlouboutin3@gatech.edu
# Updated July 2020

export ricker_wavelet, get_computational_nt, calculate_dt, setup_grid, setup_3D_grid
export convertToCell, limit_model_to_receiver_area, extend_gradient, remove_out_of_bounds_receivers
export time_resample, remove_padding, subsample, process_input_data
export generate_distribution, select_frequencies
export devito_model, update_dm, pad_sizes, pad_array
export transducer

"""
    devito_model(model, options;dm=nothing)

Creates a python side model strucutre for devito.

Parameters
* `model`: JUDI Model structure.
* `options`: JUDI Options structure.
* `dm`: Squared slowness perturbation (optional), Array or PhysicalParameter.
"""
function devito_model(model::Model, options::Options; dm=nothing)
    pad = pad_sizes(model, options)
    # Set up Python model structure
    m = pad_array(model[:m].data, pad)
    !isnothing(dm) && (dm = pad_array(reshape(getattr(dm, :data, dm), model.n), pad))
    physpar = Dict((n, pad_array(v.data, pad)) for (n, v) in model.params if n != :m)
    modelPy = pm."Model"(model.o, model.d, model.n, m, fs=options.free_surface,
                         nbl=model.nb, space_order=options.space_order, dt=options.dt_comp, dm=dm;
                         physpar...)
    return modelPy
end

"""
    pad_sizes(model, options; so=nothing)

Computes ABC padding sizes according to the model's numbr of abc points and spatial order

Parameters
* `model`: JUDI or Python side Model.
* `options`: JUDI Options structure.
* `so`: Space order (optional) defaults to options.space_order.
"""
function pad_sizes(model, options; so=nothing)
    isnothing(so) && (so = options.space_order)
    try
        return [(nbl + so, nbr + so) for (nbl, nbr)=model.padsizes]
    catch e
        padsizes = [(model.nb + so, model.nb + so) for i=1:length(model.n)]
        if options.free_surface
            padsizes[end] = (so, model.nb + so)
        end
        return padsizes
    end
end

"""
    pad_array(m, nb; mode=:border)

Pads to the input array with either copying the edge value (:border) or zeros (:zeros)

Parameters
* `m`: Array to be padded.
* `nb`: Size of padding. Array of tuple with one (nb_left, nb_right) tuple per dimension.
* `mode`: Padding mode (optional), defaults to :border.
"""
function pad_array(m::Array{DT}, nb::Array{Tuple{Int64,Int64},1}; mode::Symbol=:border) where {DT}
    n = size(m)
    new_size = Tuple([n[i] + sum(nb[i]) for i=1:length(nb)])
    Ei = []
    for i=length(nb):-1:1
        left, right = nb[i]
        push!(Ei, joExtend(n[i], mode;pad_upper=right, pad_lower=left, RDT=DT, DDT=DT))
    end
    padded = joKron(Ei...) * vec(m)
    return reshape(padded, new_size)
end

"""
    remove_padding(m, nb; true_adjoint=False)

Removes the padding from array `m`. This is the adjoint of [`pad_array`](@ref).

Parameters
* `m`: Array to remvove padding from.
* `nb`: Size of padding. Array of tuple with one (nb_left, nb_right) tuple per dimension.
* `true_adjoint`: Unpadding mode, defaults to False. Will sum the padding to the edge point with `true_adjoint=true`
 and should only be used this way for adjoint testing purpose.
"""
function remove_padding(gradient::Array{DT}, nb::Array{Tuple{Int64,Int64},1}; true_adjoint::Bool=false) where {DT}
    # Pad is applied x then y then z, so sum must be done in reverse order x then y then x
    true_adjoint ? mode = :border : mode = :zeros
    n = size(gradient)
    new_size = Tuple([n[i] - sum(nb[i]) for i=1:length(nb)])
    Ei = []
    for i=length(nb):-1:1
        left, right = nb[i]
        push!(Ei, joExtend(new_size[i], mode;pad_upper=right, pad_lower=left, RDT=DT, DDT=DT))
    end
    gradient = reshape(joKron(Ei...)' * vec(gradient), new_size)
    return collect(Float32, gradient)
end

"""
    limit_model_to_receiver_area(srcGeometry, recGeometry, model, buffer; pert=nothing)

Crops the `model` to the area of the source an receiver with an extra buffer. This reduces the size
of the problem when the model si large and the source and receiver located in a small part of the domain.

In the cartoon below, the full model will be cropped to the center area containg the source (o) receivers (x) and
buffer area (*)

--------------------------------------------
| . . . . . . . . . . . . . . . . . . . . . |
| . . . . . . . . . . . . . . . . . . . . . |  - o Source position
| . . . . * * * * * * * * * * * * . . . . . |  - x receiver positions
| . . . . * x x x x x x x x x x * . . . . . |  - * Extra buffer (grid spacing in that simple case)
| . . . . * x x x x x x x x x x * . . . . . |
| . . . . * x x x x x x x x x x * . . . . . |
| . . . . * x x x x x o x x x x * . . . . . |
| . . . . * x x x x x x x x x x * . . . . . |
| . . . . * x x x x x x x x x x * . . . . . |
| . . . . * x x x x x x x x x x * . . . . . |
| . . . . * * * * * * * * * * * * . . . . . |
| . . . . . . . . . . . . . . . . . . . . . |
| . . . . . . . . . . . . . . . . . . . . . |
--------------------------------------------
Parameters
* `srcGeometry`: Geometry of the source.
* `recGeometry`: Geometry of the receivers.
* `model`: Model to be croped.
* `buffer`: Size of the buffer on each side.
* `pert`: Model perturbation (optional) to be cropped as well.
"""
function limit_model_to_receiver_area(srcGeometry::Geometry, recGeometry::Geometry,
                                      model::Model, buffer::Number; pert=nothing)
    # Restrict full velocity model to area that contains either sources and receivers
    ndim = length(model.n)
    # scan for minimum and maximum x and y source/receiver coordinates
    min_x = min(minimum(recGeometry.xloc[1]), minimum(srcGeometry.xloc[1]))
    max_x = max(maximum(recGeometry.xloc[1]), maximum(srcGeometry.xloc[1]))
    if ndim == 3
        min_y = min(minimum(recGeometry.yloc[1]), minimum(srcGeometry.yloc[1]))
        max_y = max(maximum(recGeometry.yloc[1]), maximum(srcGeometry.yloc[1]))
    end

    # add buffer zone if possible
    min_x = max(model.o[1], min_x - buffer) + 1
    max_x = min(model.o[1] + model.d[1]*(model.n[1]-1), max_x + buffer)
    if ndim == 3
        min_y = max(model.o[2], min_y - buffer)
        max_y = min(model.o[2] + model.d[2]*(model.n[2]-1), max_y + buffer)
    end

    # extract part of the model that contains sources/receivers
    nx_min = Int(min_x ÷ model.d[1]) + 1
    nx_max = Int(max_x ÷ model.d[1]) + 1
    inds = [max(1, nx_min):nx_max, 1:model.n[end]]
    if ndim == 3
        ny_min = Int(min_y ÷ model.d[1]) + 1
        ny_max = Int(max_y ÷ model.d[1]) + 1
        insert!(inds, 2, max(1, ny_min):ny_max)
    end

    # Extract relevant model part from full domain
    n_orig = model.n
    for (p, v) in model.params
        typeof(v) <: AbstractArray && (model.params[p] = v[inds...])
    end

    println("N old $(model.n)")
    model.n = model.m.n
    model.o = model.m.o
    println("N new $(model.n)")
    isnothing(pert) && (return model, nothing)

    pert = reshape(pert, n_orig)[inds...]
    return model, vec(pert)
end

"""
    extend_gradient(model_full, model, array)

This operation does the opposite of [`limit_model_to_receiver_area`](@ref) and put back a cropped array into
the full model.

Parameters
* `model_full`: Full domain model.
* `model`: Cropped model.
* `array`: Array to be extended (padded with zero) to the full model size.
"""
function extend_gradient(model_full::Model, model::Model, gradient::Union{Array, PhysicalParameter})
    # Extend gradient back to full model size
    ndim = length(model.n)
    full_gradient = similar(gradient, model_full)
    nx_start = Int(Float32(Float32(model.o[1] - model_full.o[1]) ÷ model.d[1])) + 1
    nx_end = nx_start + model.n[1] - 1
    if ndim == 2
        full_gradient[nx_start:nx_end,:] = gradient
    else
        ny_start = Int((model.o[2] - model_full.o[2])/model.d[2] + 1)
        ny_end = ny_start + model.n[2] - 1
        full_gradient[nx_start:nx_end,ny_start:ny_end,:] = gradient
    end
    return full_gradient
end

"""
    remove_out_of_bounds_receivers(recGeometry, model)

Removes receivers that are positionned outside the computational domain defined by the model.

Parameters
* `recGeometry`: Geometry of receivers in which out of bounds will be removed.
* `model`: Model defining the computational domain.
"""
function remove_out_of_bounds_receivers(recGeometry::Geometry, model::Model)

    # Only keep receivers within the model
    xmin, xmax = model.o[1], model.o[1] + (model.n[1] - 1)*model.d[1]
    if typeof(recGeometry.xloc[1]) <: Array
        idx_xrec = findall(x -> xmax >= x >= xmin, recGeometry.xloc[1])
        recGeometry.xloc[1] = recGeometry.xloc[1][idx_xrec]
        recGeometry.zloc[1] = recGeometry.zloc[1][idx_xrec]
    end

    # For 3D shot records, scan also y-receivers
    if length(model.n) == 3 && typeof(recGeometry.yloc[1]) <: Array
        ymin, ymax = model.o[2], model.o[2] + (model.n[2] - 1)*model.d[2]
        idx_yrec = findall(x -> ymax >= x >= ymin, recGeometry.yloc[1])
        recGeometry.yloc[1] = recGeometry.yloc[1][idx_yrec]
        recGeometry.zloc[1] = recGeometry.zloc[1][idx_yrec]
    end
    return recGeometry
end

"""
    remove_out_of_bounds_receivers(recGeometry, recData, model)

Removes receivers that are positionned outside the computational domain defined by the model.

Parameters
* `recGeometry`: Geometry of receivers in which out of bounds will be removed.
* `recData`: Shot record for that geometry in which traces will be removed.
* `model`: Model defining the computational domain.
"""
function remove_out_of_bounds_receivers(recGeometry::Geometry, recData::Array, model::Model)

    # Only keep receivers within the model
    xmin, xmax = model.o[1], model.o[1] + (model.n[1] - 1)*model.d[1]
    if typeof(recGeometry.xloc[1]) <: Array
        idx_xrec = findall(x -> xmax >= x >= xmin, recGeometry.xloc[1])
        recGeometry.xloc[1] = recGeometry.xloc[1][idx_xrec]
        recGeometry.zloc[1] = recGeometry.zloc[1][idx_xrec]
        recData = recData[:, idx_xrec]
    end

    # For 3D shot records, scan also y-receivers
    if length(model.n) == 3 && typeof(recGeometry.yloc[1]) <: Array
        ymin, ymax = model.o[2], model.o[2] + (model.n[2] - 1)*model.d[2]
        idx_yrec = findall(x -> ymax >= x >= ymin, recGeometry.yloc[1])
        recGeometry.yloc[1] = recGeometry.yloc[1][idx_yrec]
        recGeometry.zloc[1] = recGeometry.zloc[1][idx_yrec]
        recData = recData[:, idx_yrec]
    end
    return recGeometry, recData
end

remove_out_of_bounds_receivers(G::Geometry, ::Nothing, M::Model) = (remove_out_of_bounds_receivers(G, M), nothing)
remove_out_of_bounds_receivers(::Nothing, ::Nothing, M::Model) = (nothing, nothing)
remove_out_of_bounds_receivers(::Nothing, r::Array, M::Model) = (nothing, r)

"""
    convertToCell(x)

Convert an array `x` to a cell array (`Array{Any,1}`) with `length(x)` entries,\\
where the i-th cell contains the i-th entry of `x`.

Parameters
* `x`: Array to be converted into and array of array
"""
function convertToCell(x::Array{T, 1}) where T
    n = length(x)
    y = Array{Array{T, 1}, 1}(undef, n)
    for j=1:n
        y[j] = [x[j]]
    end
    return y
end

convertToCell(x::Array{Array{T, N}, 1}) where {T, N} = x
convertToCell(x::StepRangeLen) = convertToCell(Float32.(x))
convertToCell(x::Number) = convertToCell([x])

# 1D source time function
"""
    source(tmax, dt, f0)

Create seismic Ricker wavelet of length `tmax` (in milliseconds) with sampling interval `dt` (in milliseonds)\\
and central frequency `f0` (in kHz).

"""
function ricker_wavelet(tmax, dt, f0; t0=nothing)
    R = typeof(dt)
    isnothing(t0) ? t0 = R(0) : tmax = R(tmax - t0)
    nt = Int(tmax ÷ dt) + 1
    t = range(t0, stop=tmax, length=nt)
    r = (pi * f0 * (t .- 1 / f0))
    q = zeros(Float32,nt,1)
    q[:,1] = (1f0 .- 2f0 .* r.^2f0) .* exp.(-r.^2f0)
    return q
end

"""
    calculate_dt(model; dt=nothing)

Compute the computational time step based on the CFL condition and physical parameters
in the model.

Parameters
* `model`: Model structure
* `dt`: User defined time step (optional), will be the value returned if provided.
"""
function calculate_dt(model::Model; dt=nothing)
    if ~isnothing(dt)
        return dt
    end
    m = minimum(model[:m])
    epsilon = maximum(get(model.params, :epsilon, 0))
    modelPy = pm."Model"(origin=model.o, spacing=model.d, shape=model.n,
                         m=m, epsilon=epsilon, nbl=0)
    return convert(Float32, modelPy.critical_dt)
end

"""
    get_computational_nt(srcGeometry, recGeoemtry, model; dt=nothing)

Estimate the number of computational time steps. Required for calculating the dimensions\\
of the matrix-free linear modeling operators. `srcGeometry` and `recGeometry` are source\\
and receiver geometries of type `Geometry` and `model` is the model structure of type \\
`Model`.

"""
function get_computational_nt(srcGeometry, recGeometry, model::Model; dt=nothing)
    # Determine number of computational time steps
    if typeof(srcGeometry) <: GeometryOOC
        nsrc = length(srcGeometry.container)
    else
        nsrc = length(srcGeometry.xloc)
    end
    nt = Array{Integer}(undef, nsrc)
    dtComp = calculate_dt(model; dt=dt)
    for j=1:nsrc
        ntRec = Int(recGeometry.dt[j]*(recGeometry.nt[j]-1) ÷ dtComp) + 1
        ntSrc = Int(srcGeometry.dt[j]*(srcGeometry.nt[j]-1) ÷ dtComp) + 1
        nt[j] = max(ntRec, ntSrc)
    end
    return nt
end

"""
    get_computational_nt(Geoemtry, model; dt=nothing)

Estimate the number of computational time steps. Required for calculating the dimensions\\
of the matrix-free linear modeling operators. `srcGeometry` and `recGeometry` are source\\
and receiver geometries of type `Geometry` and `model` is the model structure of type \\
`Model`.

"""
function get_computational_nt(Geometry, model::Model; dt=nothing)
    # Determine number of computational time steps
    if typeof(Geometry) <: GeometryOOC
        nsrc = length(Geometry.container)
    else
        nsrc = length(Geometry.xloc)
    end
    nt = Array{Integer}(undef, nsrc)
    dtComp = calculate_dt(model; dt=dt)
    for j=1:nsrc
        nt[j] = Int(Geometry.dt[j]*(Geometry.nt[j]-1) ÷ dtComp) + 1
    end
    return nt
end

"""
    setup_grid(geometry, n)

Sets up the coordinate arrays for Devito.

Parameters:
* `geometry`: Geometry containing the coordinates
* `n`: Domain size
"""
function setup_grid(geometry, n)
    # 3D grid
    if length(n)==3
        if length(geometry.xloc[1]) > 1
            source_coords = Array{Float32,2}([vec(geometry.xloc[1]) vec(geometry.yloc[1]) vec(geometry.zloc[1])])
        else
            source_coords = Array{Float32,2}([geometry.xloc[1] geometry.yloc[1] geometry.zloc[1]])
        end
    else
    # 2D grid
        if length(geometry.xloc[1]) > 1
            source_coords = Array{Float32,2}([vec(geometry.xloc[1]) vec(geometry.zloc[1])])
        else
            source_coords = Array{Float32,2}([geometry.xloc[1] geometry.zloc[1]])
        end
    end
    return source_coords
end

"""
    setup_3D_grid(x, y, z)

Converts one dimensional input (x, y, z) into three dimensional coordinates. The number of point created
is `length(x)*lenght(y)` with all the x/y pairs and each pair at depth z[idx[x]]. `x` and `z` must have the same size.

Parameters:
* `x`: X coordinates.
* `y`: Y coordinates.
* `z`: Z coordinates.
"""
function setup_3D_grid(xrec::Array{Array{T, 1}, 1},yrec::Array{Array{T, 1},1},zrec::Array{T, 1}) where T<:Real
    # Take input 1d x and y coordinate vectors and generate 3d grid. Input are cell arrays
    nsrc = length(xrec)
    xloc = Array{Array{T}, 1}(undef, nsrc)
    yloc = Array{Array{T}, 1}(undef, nsrc)
    zloc = Array{Array{T}, 1}(undef, nsrc)
    for i=1:nsrc
        nxrec = length(xrec[i])
        nyrec = length(yrec[i])

        xloc[i] = zeros(T, nxrec*nyrec)
        yloc[i] = zeros(T, nxrec*nyrec)
        zloc[i] = zeros(T, nxrec*nyrec)

        idx = 1

        for k=1:nyrec
            for j=1:nxrec
                xloc[i][idx] = xrec[i][j]
                yloc[i][idx] = yrec[i][k]
                zloc[i][idx] = zrec[i]
                idx += 1
            end
        end
    end
    return xloc, yloc, zloc
end

"""
    setup_3D_grid(x, y, z)

Converts one dimensional input (x, y, z) into three dimensional coordinates. The number of point created
is `length(x)*lenght(y)` with all the x/y pairs and each pair at same depth z.

Parameters:
* `x`: X coordinates.
* `y`: Y coordinates.
* `z`: Z coordinate.
"""
function setup_3D_grid(xrec::Array{T, 1},yrec::Array{T,1}, zrec::T) where T<:Real
# Take input 1d x and y coordinate vectors and generate 3d grid. Input are arrays/ranges
    nxrec = length(xrec)
    nyrec = length(yrec)

    xloc = zeros(T, nxrec*nyrec)
    yloc = zeros(T, nxrec*nyrec)
    zloc = zeros(T, nxrec*nyrec)
    idx = 1
    for k=1:nyrec
        for j=1:nxrec
            xloc[idx] = xrec[j]
            yloc[idx] = yrec[k]
            zloc[idx] = zrec
            idx += 1
        end
    end
    return xloc, yloc, zloc
end

setup_3D_grid(xrec, yrec, zrec) = setup_3D_grid(tof32(xrec), tof32(yrec), tof32(zrec))

"""
    time_resample(data, geometry_in, dt_new)

Resample the input data with sinc interpolation from the current time sampling (geometrty_in) to the
new time sampling `dt_new`.

Parameters
* `data`: Data to be reampled. If data is a matrix, resamples each column.
* `geometry_in`: Geometry on which `data` is defined.
* `dt_new`: New time sampling rate to interpolate onto.
"""
function time_resample(data::Array, geometry_in::Geometry, dt_new)

    if dt_new==geometry_in.dt[1]
        return data, geometry_in
    else
        geometry = deepcopy(geometry_in)
        timeAxis = 0:geometry.dt[1]:geometry.t[1]
        timeInterp = 0:dt_new:geometry.t[1]
        dataInterp = Float32.(SincInterpolation(data, timeAxis, timeInterp))

        geometry.dt[1] = dt_new
        geometry.nt[1] = length(timeInterp)
        geometry.t[1] = (geometry.nt[1] - 1)*geometry.dt[1]
        return dataInterp, geometry
    end
end

"""
    time_resample(data, dt_in, geometry_in)

Resample the input data with sinc interpolation from the current time sampling (dt_in) to the
new time sampling `geometry_out`.

Parameters
* `data`: Data to be reampled. If data is a matrix, resamples each column.
* `geometry_out`: Geometry on which `data` is to be interpolated.
* `dt_in`: Time sampling rate of the `data.`
"""
function time_resample(data::Array, dt_in, geometry_out::Geometry)
    if dt_in==geometry_out.dt[1]
        return data
    else
        timeAxis = 0:dt_in:geometry_out.t[1]
        timeInterp = 0:geometry_out.dt[1]:geometry_out.t[1]
        return  Float32.(SincInterpolation(data, timeAxis, timeInterp))
    end
end

"""
    generate_distribution(x; src_no=1)

Generates a probability distribution for the discrete input judiVector `x`.

Parameters
* `x`: judiVector. Usualy a source with a single trace per source position.
* `src_no`: Index of the source to select out of `x`
"""
function generate_distribution(x; src_no=1)
	# Generate interpolator to sample from probability distribution given
	# from spectrum of the input data

	# sampling information
	nt = x.geometry.nt[src_no]
	dt = x.geometry.dt[src_no]
	t = x.geometry.t[src_no]

	# frequencies
	fs = 1/dt	# sampling rate
	fnyq = fs/2	# nyquist frequency
	df = fnyq/nt	# frequency interval
	f = 0:2*df:fnyq	# frequencies

	# amplitude spectrum of data (serves as probability density function)
	ns = nt ÷ 2 + 1
	amp = abs.(fft(x.data[src_no]))[1:ns]	# get first half of spectrum

	# convert to cumulative probability distribution (integrate)
	pd = zeros(ns)
	pd[1] = dt*amp[1]
	for j=2:ns
		pd[j] = pd[j-1] + amp[j]*df
	end
	pd /= pd[end]	# normalize

	return Spline1D(pd, f)
end

"""
    select_frequencies(q_dist; fmin=0f0, fmax=Inf, nf=1)

Selects `nf` frequencies based on the source distribution `q_dist` computed with [`generate_distribution`](@ref).

Parameters
* `q_dist`: Distribution to sample from.
* `f_min`: Minimum acceptable frequency to sample (defaults to 0).
* `f_max`: Maximum acceptable frequency to sample (defaults to Inf).
* `fd`: Number of frequnecies to sample (defaults to 1).
"""
function select_frequencies(q_dist; fmin=0f0, fmax=Inf, nf=1)
	freq = zeros(Float32, nf)
	for j=1:nf
		while (freq[j] <= fmin) || (freq[j] > fmax)
			freq[j] = q_dist(rand(1)[1])[1]
		end
	end
	return freq
end

"""
    process_input_data(input, geometry, info)

Preprocesses input Array into an Array of Array for modeling

Parameters:
* `input`: Input to preprocess.
* `geometry`: Geometry containing physical parameters.
* `info`: Infor structure.
"""
function process_input_data(input::Array{Float32}, geometry::Geometry, info::Info)
    # Input data is pure Julia array: assume fixed no.
    # of receivers and reshape into data cube nt x nrec x nsrc
    nt = Int(geometry.nt[1])
    nrec = length(geometry.xloc[1])
    nsrc = info.nsrc
    data = reshape(input, nt, nrec, nsrc)
    dataCell = Array{Array{Float32, 2}, 1}(undef, nsrc)
    for j=1:nsrc
        dataCell[j] = data[:,:,j]
    end
    return dataCell
end

"""
    process_input_data(input, model, info)

Preprocesses input Array into an Array of Array for modeling

Parameters:
* `input`: Input to preprocess.
* `model`: Model containing physical parameters.
* `info`: Infor structure.
"""
function process_input_data(input::Array{Float32}, model::Model, info::Info)
    ndims = length(model.n)
    dataCell = Array{Array{Float32, ndims}, 1}(undef, info.nsrc)
    if ndims == 2
        input = reshape(input, model.n[1], model.n[2], info.nsrc)
        for j=1:info.nsrc
            dataCell[j] = input[:,:,j]
        end
    elseif ndims == 3
        input = reshape(input, model.n[1], model.n[2], model.n[3], info.nsrc)
        for j=1:info.nsrc
            dataCell[j] = input[:,:,:,j]
        end
    else
        throw("Number of dimensions not supported.")
    end
    return dataCell
end

process_input_data(input::judiVector, ::Geometry, ::Info) = input.data
process_input_data(input::judiVector, ::Info) = input.data
process_input_data(input::judiWeights, ::Model, ::Info) = input.weights

"""
    reshape(x::Array{Float32, 1}, geometry::Geometry)

Reshapes input vector intu a 3D `nt x nrec x nsrc` Array.
"""
function reshape(x::Array{Float32, 1}, geometry::Geometry)
    nt = geometry.nt[1]
    nrec = length(geometry.xloc[1])
    nsrc = Int(length(x) / nt / nrec)
    return reshape(x, nt, nrec, nsrc)
end

"""
    transducer(q, d, r, theta)

Create the JUDI soure for a circular transducer
Theta=0 points downward:

. . . . - - - . . . . . .

. . . . + + + . . . . . .

. . . . . . . . . . . . .

. . . . . . . . . . . . .

Theta=pi/2 points right:

. . . . - + . . . . . . .

. . . . - + . . . . . . .

. . . . - + . . . . . . .

. . . . . . . . . . . . .

2D only, to extend to 3D

"""
function transducer(q::judiVector, d::Tuple, r::Number, theta)
    length(theta) != length(q.geometry.xloc) && throw("Need one angle per source position")
    size(q.data[1], 2) > 1 && throw("Only point sources can be converted to transducer source")
    # Array of source
    nsrc_loc = 11
    nsrc = length(q.geometry.xloc)
    x_base = collect(range(-r, r, length=nsrc_loc))
    y_base = zeros(nsrc_loc)
    y_base_b = zeros(nsrc_loc) .- d[end]

    # New coords and data
    xloc = Array{Array{Float32, 1}, 1}(undef, nsrc)
    yloc = Array{Array{Float32, 1}, 1}(undef, nsrc)
    zloc = Array{Array{Float32, 1}, 1}(undef, nsrc)
    data = Array{Array{Float32, 2}}(undef, nsrc)
    t = q.geometry.t[1]
    dt = q.geometry.dt[1]

    for i=1:nsrc
        # Build the rotated array of dipole
        R = Float32.([cos(theta[i] - pi/2) sin(theta[i] - pi/2);-sin(theta[i] - pi/2) cos(theta[i] - pi/2)])
        # +1 coords
        r_loc = R * [x_base';y_base']
        # -1 coords
        r_loc_b = R * [x_base';y_base_b']
        xloc[i] = q.geometry.xloc[i] .+ vec(vcat(r_loc[1, :], r_loc_b[1, :]))
        zloc[i] = q.geometry.zloc[i] .+ vec(vcat(r_loc[2, :], r_loc_b[2, :]))
        yloc[i] = zeros(Float32, 2*nsrc_loc)
        data[i] = zeros(Float32, length(q.data[i]), 2*nsrc_loc)
        data[i][:, 1:nsrc_loc] .= q.data[i]/nsrc_loc
        data[i][:, nsrc_loc+1:end] .= -q.data[i]/nsrc_loc
    end
    return judiVector(Geometry(xloc, yloc, zloc; t=t, dt=dt), data)
end

########################################### Misc defaults

pad_array(m::Number, ::Array{Tuple{Int64,Int64},1}) = m
pad_array(::Nothing, ::Array{Tuple{Int64,Int64},1}) = nothing

# Vectorization of single variable (not defined in Julia)
vec(x::Float64) = x;
vec(x::Float32) = x;
vec(x::Int64) = x;
vec(x::Int32) = x;
vec(::Nothing) = nothing

SincInterpolation(Y::Array{T, 2}, S, Up) where T<:Real = sinc.( (Up .- S') ./ (S[2] - S[1]) ) * Y

subsample(::Nothing, i) = nothing
subsample(a::Array{Array{T, N}, 1}, i) where {T, N} = a[i]
subsample(a::Array{SeisCon,1}, i::Int64) = a[i]


text:
# MIT License
#
# Copyright (c) 2020 Joshua E Gentry

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# =================================================================================================
# Test reading/writing a JPG image.
@testset "pix_write_jpeg/pix_read_jpeg" begin

    local pix = pix_with()

    # ---------------------------------------------------------------------------------------------
    # To/From a file.
    local filename = safe_tmp_file()
    @test pix_write_jpeg(filename, pix) == true
    @test pix_read_jpeg(filename) != nothing
    rm(filename)

    # ---------------------------------------------------------------------------------------------
    # To/From an IOBuffer.
    local buffer = IOBuffer()
    @test pix_write_jpeg(buffer, pix) == true
    @test position(buffer) != 0
    seekstart(buffer)
    @test pix_read_jpeg(buffer) != nothing

    # ---------------------------------------------------------------------------------------------
    # To/From an IOStream
    local filename = safe_tmp_file()
    local file     = open(filename, read=true, write=true)
    @test pix_write_jpeg(file, pix) == true
    @test position(file) != 0
    seekstart(file)
    @test pix_read_jpeg(file) != nothing
    close(file)
    rm(filename)

    # ---------------------------------------------------------------------------------------------
    # To/From a byte array.
    local data  = jpeg_with()
    local bytes = pix_write_jpeg(pix)
    @test length(bytes) > 0
    @test pix_read_jpeg(bytes) != nothing
end

# =================================================================================================
# Test writing a JPG image at various qualities.
@testset "pix_write_jpeg | quality = $quality, progressive = $progressive" for
        quality in [ 1, 25, 50, 75, 100 ], progressive in [ true, false ]

    local pix = pix_with()

    # ---------------------------------------------------------------------------------------------
    # To/From a file.
    local filename = safe_tmp_file()
    @test pix_write_jpeg(filename, pix; quality = quality, progressive = progressive) == true
    @test pix_read_jpeg(filename) != nothing
    rm(filename)

    # ---------------------------------------------------------------------------------------------
    # To/From an IOBuffer.
    local buffer = IOBuffer()
    @test pix_write_jpeg(buffer, pix; quality = quality, progressive = progressive) == true
    @test position(buffer) != 0
    seekstart(buffer)
    @test pix_read_jpeg(buffer) != nothing

    # ---------------------------------------------------------------------------------------------
    # To/From a byte array.
    local data  = jpeg_with()
    local bytes = pix_write_jpeg(pix; quality = quality, progressive = progressive)
    @test length(bytes) > 0
    @test pix_read_jpeg(bytes) != nothing
end

# =================================================================================================
# Test reading a JPG image /w parameters
@testset "pix_read_jpeg | cmap = $cmap, reduction = $reduction, luminance = $luminance" for
        cmap in [ true, false ], reduction in [ 1, 2, 4, 8 ], luminance in [ true, false ]

    local pix = pix_with()
    @suppress begin
        # -----------------------------------------------------------------------------------------
        # To/From a file.
        local filename = safe_tmp_file()
        @test pix_write_jpeg(filename, pix) == true
        @test pix_read_jpeg(filename; cmap = cmap, reduction = reduction, luminance = luminance) != nothing
        rm(filename)

        # -----------------------------------------------------------------------------------------
        # To/From an IOBuffer.
        local buffer = IOBuffer()
        @test pix_write_jpeg(buffer, pix) == true
        @test position(buffer) != 0
        seekstart(buffer)
        @test pix_read_jpeg(buffer; cmap = cmap, reduction = reduction, luminance = luminance) != nothing

        # -----------------------------------------------------------------------------------------
        # To/From a byte array.
        local data  = jpeg_with()
        local bytes = pix_write_jpeg(pix)
        @test length(bytes) > 0
        @test pix_read_jpeg(bytes; cmap = cmap, reduction = reduction, luminance = luminance) != nothing
    end
end

# =================================================================================================
# Test reading empty data.
@testset "pix_read_jpeg empty data" begin
    @suppress begin
        # -----------------------------------------------------------------------------------------
        # From a file.
        local filename = safe_tmp_file()
        @test pix_read_jpeg(filename) == nothing
        @test pix_read_jpeg("xyzzy.jpg") == nothing
        rm(filename)

        # -----------------------------------------------------------------------------------------
        # From an IOBuffer.
        local buffer = IOBuffer()
        @test pix_read_jpeg(buffer) == nothing

        # -----------------------------------------------------------------------------------------
        # To/From a byte array.
        local bytes = Vector{UInt8}()
        @test pix_read_jpeg(bytes) == nothing
    end
end

# =================================================================================================
# Test pix_write_jpeg supplying bad qualities.
@testset "pix_write_jpeg /w invalid values" begin
    local pix = pix_with()
    local err = "Quality must be between 1 and 100 inclusive."

    # ---------------------------------------------------------------------------------------------
    # To a file.
    local filename = safe_tmp_file()
    @test (@test_logs (:error, err) pix_write_jpeg(filename, pix; quality = 0)) == false
    @test filesize(filename) == 0
    @test (@test_logs (:error, err) pix_write_jpeg(filename, pix; quality = -1)) == false
    @test filesize(filename) == 0
    @test (@test_logs (:error, err) pix_write_jpeg(filename, pix; quality = 200)) == false
    @test filesize(filename) == 0
    rm(filename)

    # ---------------------------------------------------------------------------------------------
    # To an IOBuffer.
    local buffer = IOBuffer()
    @test (@test_logs (:error, err) pix_write_jpeg(buffer, pix; quality = 0)) == false
    @test length(take!(buffer)) == 0
    @test (@test_logs (:error, err) pix_write_jpeg(buffer, pix; quality = -1)) == false
    @test length(take!(buffer)) == 0
    @test (@test_logs (:error, err) pix_write_jpeg(buffer, pix; quality = 200)) == false
    @test length(take!(buffer)) == 0
    # ---------------------------------------------------------------------------------------------
    # To a byte array.
    @test (@test_logs (:error, err) pix_write_jpeg(pix; quality = 0)) == nothing
    @test (@test_logs (:error, err) pix_write_jpeg(pix; quality = -1)) == nothing
    @test (@test_logs (:error, err) pix_write_jpeg(pix; quality = 200)) == nothing
end

# =================================================================================================
# Test reading/writing a JPG image.
@testset "pix_read_jpeg /w invalid values" begin
    local pix = pix_with()
    local err = "Reduction must be 1, 2, 4, or 8."

    # ---------------------------------------------------------------------------------------------
    # To/From a file.
    local filename = safe_tmp_file()
    @test pix_write_jpeg(filename, pix) == true
    @test (@test_logs (:error, err) pix_read_jpeg(filename; reduction = 0)) == nothing
    @test (@test_logs (:error, err) pix_read_jpeg(filename; reduction = 7)) == nothing
    @test (@test_logs (:error, err) pix_read_jpeg(filename; reduction = 16)) == nothing
    rm(filename)

    # ---------------------------------------------------------------------------------------------
    # To/From an IOBuffer.
    local buffer = IOBuffer()
    @test pix_write_jpeg(buffer, pix) == true
    @test position(buffer) != 0
    seekstart(buffer)
    @test (@test_logs (:error, err) pix_read_jpeg(buffer; reduction = 0)) == nothing
    seekstart(buffer)
    @test (@test_logs (:error, err) pix_read_jpeg(buffer; reduction = 7)) == nothing
    seekstart(buffer)
    @test (@test_logs (:error, err) pix_read_jpeg(buffer; reduction = 16)) == nothing

    # ---------------------------------------------------------------------------------------------
    # To/From a byte array.
    local bytes = pix_write_jpeg(pix)
    @test length(bytes) > 0
    @test (@test_logs (:error, err) pix_read_jpeg(bytes; reduction = 0)) == nothing
    @test (@test_logs (:error, err) pix_read_jpeg(bytes; reduction = 7)) == nothing
    @test (@test_logs (:error, err) pix_read_jpeg(bytes; reduction = 16)) == nothing
end

# =================================================================================================
# Test writing a Pix that has been freed.
@testset "pix_write_jpeg /w freed Pix" begin
    local err = "Pix has been freed."
    local pix = pix_with()
    pix_delete!(pix)

    # ---------------------------------------------------------------------------------------------
    # To a file.
    local filename = safe_tmp_file()
    @test (@test_logs (:error, err) pix_write_jpeg(filename, pix)) == false
    rm(filename)

    # ---------------------------------------------------------------------------------------------
    # To an IOBuffer.
    local buffer = IOBuffer()
    @test (@test_logs (:error, err) pix_write_jpeg(buffer, pix)) == false

    # ---------------------------------------------------------------------------------------------
    # To a byte array.
    @test (@test_logs (:error, err) pix_write_jpeg(pix)) == nothing
end


text:
# Quadratic decision function #
# ========================== #
mutable struct QuadraticDecisionConstraintBridge{T, S} <: MOIB.Constraint.AbstractBridge
    constraint::CI{MOI.ScalarQuadraticFunction{T}, S}
    decision_function::QuadraticDecisionFunction{T}
    set::S
end

function MOIB.Constraint.bridge_constraint(::Type{QuadraticDecisionConstraintBridge{T,S}},
                                           model,
                                           f::QuadraticDecisionFunction{T},
                                           set::S) where {T, S}
    # All decisions have been mapped to the
    # variable part terms at this point.
    # Add the bridged constraint
    constraint = MOI.add_constraint(model,
                                    MOI.ScalarQuadraticFunction(
                                        f.variable_part.affine_terms,
                                        f.variable_part.quadratic_terms,
                                        zero(T)),
                                    MOIU.shift_constant(set, -f.variable_part.constant))
    # Save the constraint index, the decision function, and the set, to allow modifications
    return QuadraticDecisionConstraintBridge{T, S}(constraint, f, set)
end

function MOI.supports_constraint(::Type{<:QuadraticDecisionConstraintBridge{T}},
                                 ::Type{<:QuadraticDecisionFunction{T}},
                                 ::Type{<:MOI.AbstractScalarSet}) where T
    return true
end
function MOIB.added_constrained_variable_types(::Type{<:QuadraticDecisionConstraintBridge})
    return Tuple{DataType}[]
end
function MOIB.added_constraint_types(::Type{QuadraticDecisionConstraintBridge{T, S}}) where {T, S}
    return [(MOI.ScalarQuadraticFunction{T}, S)]
end
function MOIB.Constraint.concrete_bridge_type(::Type{<:QuadraticDecisionConstraintBridge{T}},
                                              ::Type{QuadraticDecisionFunction{T}},
                                              S::Type{<:MOI.AbstractScalarSet}) where T
    return QuadraticDecisionConstraintBridge{T, S}
end

MOI.get(b::QuadraticDecisionConstraintBridge{T, S}, ::MOI.NumberOfConstraints{MOI.ScalarQuadraticFunction{T}, S}) where {T, S} = 1
MOI.get(b::QuadraticDecisionConstraintBridge{T, S}, ::MOI.ListOfConstraintIndices{MOI.ScalarQuadraticFunction{T}, S}) where {T, S} = [b.constraint]

function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintFunction,
                 bridge::QuadraticDecisionConstraintBridge{T,S}) where {T,S}
    f = bridge.decision_function
    # Remove mapped variables to properly unbridge
    from_decision(v) = begin
        result = any(t -> t.variable_index == v,
                     f.decision_part.affine_terms)
        result |= any(t -> t.variable_index_1 == v || t.variable_index_2 == v,
                      f.decision_part.quadratic_terms)
        result |= any(t -> t.variable_index_1 == v,
                      f.cross_terms.affine_terms)
    end
    g = QuadraticDecisionFunction(
        MOIU.filter_variables(v -> !from_decision(v), f.variable_part),
        copy(f.decision_part),
        copy(f.cross_terms))
    return g
end

function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintSet,
                 bridge::QuadraticDecisionConstraintBridge{T,S}) where {T,S}
    return bridge.set
end

function MOI.get(model::MOI.ModelLike, attr::Union{MOI.ConstraintPrimal, MOI.ConstraintDual},
                 bridge::QuadraticDecisionConstraintBridge{T,S}) where {T,S}
    return MOI.get(model, attr, bridge.constraint)
end

function MOI.get(model::MOI.ModelLike, ::DecisionIndex,
                 bridge::QuadraticDecisionConstraintBridge{T}) where T
    return bridge.constraint
end

function MOI.delete(model::MOI.ModelLike, bridge::QuadraticDecisionConstraintBridge)
    MOI.delete(model, bridge.constraint)
    return nothing
end

function MOI.set(model::MOI.ModelLike, ::MOI.ConstraintFunction,
                 bridge::QuadraticDecisionConstraintBridge{T,S},
                 f::QuadraticDecisionFunction{T}) where {T,S,F}
    # Update bridge functions and function constant
    bridge.decision_function = f
    # Change the function of the bridged constraints
    MOI.set(model, MOI.ConstraintFunction(), bridge.constraint,
            MOI.ScalarQuadraticFunction(
                f.variable_part.affine_terms,
                f.variable_part.quadratic_terms,
                zero(T)))
    # Shift constraint set
    MOI.set(model, MOI.ConstraintSet(), bridge.constraint,
            MOIU.shift_constant(bridge.set, -f.variable_part.constant))
    return nothing
end

function MOI.set(model::MOI.ModelLike, ::MOI.ConstraintSet,
                 bridge::QuadraticDecisionConstraintBridge{T,S}, change::S) where {T,S}
    f = bridge.decision_function
    bridge.set = change
    # Shift constraint set
    MOI.set(model, MOI.ConstraintSet(), bridge.constraint,
            MOIU.shift_constant(bridge.set, -f.variable_part.constant))
    return nothing
end

function MOI.modify(model::MOI.ModelLike, bridge::QuadraticDecisionConstraintBridge{T,S}, change::MOI.ScalarConstantChange) where {T,S}
    f = bridge.decision_function
    # Modify constant of variable part
    lq.variable_part.constant = change.new_constant
    # Shift constraint set
    MOI.set(model, MOI.ConstraintSet(), bridge.constraint,
            MOIU.shift_constant(bridge.set, -f.variable_part.constant))
    return nothing
end

function MOI.modify(model::MOI.ModelLike, bridge::QuadraticDecisionConstraintBridge{T,S}, change::MOI.ScalarCoefficientChange) where {T,S}
    f = bridge.decision_function
    # Update variable part
    modify_coefficient!(f.variable_part.affine_terms, change.variable, change.new_coefficient)
    # Modify decision part if variable originated from a decision
    i = something(findfirst(t -> t.variable_index == change.variable,
                            f.decision_part.affine_terms), 0)
    if !iszero(i)
        modify_coefficient!(f.decision_part.affine_terms, change.variable, change.new_coefficient)
    end
    return nothing
end

function MOI.modify(model::MOI.ModelLike, bridge::QuadraticDecisionConstraintBridge{T,S}, change::DecisionCoefficientChange) where {T,S}
    f = bridge.decision_function
    variable_terms = f.variable_part.affine_terms
    decision_terms = f.decision_part.affine_terms
    # Check if decision has been mapped
    i = something(findfirst(t -> t.variable_index == change.decision,
                            variable_terms), 0)
    if !iszero(i)
        # Update mapped variable through ScalarCoefficientChange
        MOI.modify(model, bridge, MOI.ScalarCoefficientChange(change.decision, change.new_coefficient))
    else
        i = something(findfirst(t -> t.variable_index == change.decision,
                                decision_terms), 0)
        # Query the fixed decision value
        unbridged = MOIB.unbridged_variable_function(model, change.decision)::MOI.SingleVariable
        decision_value = MOI.get(model, MOI.VariablePrimal(), unbridged.variable)
        # Update decision part of constraint constant
        coefficient = iszero(i) ? zero(T) : decision_terms[i].coefficient
        f.decision_part.constant +=
            (change.new_coefficient - coefficient) * decision_value
    end
    # Update coefficient in decision part
    modify_coefficient!(decision_terms, change.decision, change.new_coefficient)
    return nothing
end


text:
@sumProductRule(:node_type     => GaussianMeanPrecision,
                :outbound_type => Message{GaussianMeanPrecision},
                :inbound_types => (Nothing, Message{PointMass}, Message{PointMass}),
                :name          => SPGaussianMeanPrecisionOutNPP)

@sumProductRule(:node_type     => GaussianMeanPrecision,
                :outbound_type => Message{GaussianMeanPrecision},
                :inbound_types => (Message{PointMass}, Nothing, Message{PointMass}),
                :name          => SPGaussianMeanPrecisionMPNP)

@sumProductRule(:node_type     => GaussianMeanPrecision,
                :outbound_type => Message{GaussianMeanVariance},
                :inbound_types => (Nothing, Message{Gaussian}, Message{PointMass}),
                :name          => SPGaussianMeanPrecisionOutNGP)

@sumProductRule(:node_type     => GaussianMeanPrecision,
                :outbound_type => Message{GaussianMeanVariance},
                :inbound_types => (Message{Gaussian}, Nothing, Message{PointMass}),
                :name          => SPGaussianMeanPrecisionMGNP)

@naiveVariationalRule(:node_type     => GaussianMeanPrecision,
                      :outbound_type => Message{GaussianMeanPrecision},
                      :inbound_types => (Nothing, ProbabilityDistribution, ProbabilityDistribution),
                      :name          => VBGaussianMeanPrecisionOut)

@naiveVariationalRule(:node_type     => GaussianMeanPrecision,
                      :outbound_type => Message{GaussianMeanPrecision},
                      :inbound_types => (ProbabilityDistribution, Nothing, ProbabilityDistribution),
                      :name          => VBGaussianMeanPrecisionM)

@naiveVariationalRule(:node_type     => GaussianMeanPrecision,
                      :outbound_type => Message{Union{Gamma, Wishart}},
                      :inbound_types => (ProbabilityDistribution, ProbabilityDistribution, Nothing),
                      :name          => VBGaussianMeanPrecisionW)

@structuredVariationalRule(:node_type     => GaussianMeanPrecision,
                           :outbound_type => Message{GaussianMeanVariance},
                           :inbound_types => (Nothing, Message{Gaussian}, ProbabilityDistribution),
                           :name          => SVBGaussianMeanPrecisionOutVGD)

@structuredVariationalRule(:node_type     => GaussianMeanPrecision,
                           :outbound_type => Message{GaussianMeanVariance},
                           :inbound_types => (Message{Gaussian}, Nothing, ProbabilityDistribution),
                           :name          => SVBGaussianMeanPrecisionMGVD)

@structuredVariationalRule(:node_type     => GaussianMeanPrecision,
                           :outbound_type => Message{Union{Gamma, Wishart}},
                           :inbound_types => (ProbabilityDistribution, Nothing),
                           :name          => SVBGaussianMeanPrecisionW)

@marginalRule(:node_type => GaussianMeanPrecision,
              :inbound_types => (Message{Gaussian}, Message{Gaussian}, ProbabilityDistribution),
              :name => MGaussianMeanPrecisionGGD)

@marginalRule(:node_type => GaussianMeanPrecision,
              :inbound_types => (Message{Gaussian}, Message{Gaussian}, Nothing), # Precision is marginalized out
              :name => MGaussianMeanPrecisionGGN)


text:
#  point in time observation that has relevance to markets

import Base: show, getindex

immutable MarketSignal <: AbstractTimeSeries
    timestamp::Vector{Date}
    values::BitArray
    colnames::Vector{ASCIIString}
end

function show(io::IO, ms::MarketSignal)
  # variables
  nrow          = size(ms.values, 1)
  ncol          = size(ms.values, 2)
  spacetime     = strwidth(string(ms.timestamp[1])) + 3
  firstcolwidth = strwidth(ms.colnames[1])
  colwidth      = Int[]
      for m in 1:ncol
          push!(colwidth, max(strwidth(ms.colnames[m]), 5))
      end

  # summary line
  print(io, @sprintf("%dx%d %s %s to %s", nrow, ncol, typeof(ms), string(ms.timestamp[1]), string(ms.timestamp[nrow])))
  println(io, "")
  println(io, "")

  # row label line
   print(io, ^(" ", spacetime), ms.colnames[1], ^(" ", colwidth[1] + 2 -firstcolwidth))

   for p in 2:length(colwidth)
     print(io, ms.colnames[p], ^(" ", colwidth[p] - strwidth(ms.colnames[p]) + 2))
   end
   println(io,"")

  # timestamp and values line
    if nrow > 7
        for i in 1:4
            print(io, ms.timestamp[i], " | ")
        for j in 1:ncol
            print(io, rpad(ms.values[i,j], colwidth[j] + 2, " "))
        end
        println(io, "")
        end
        println(io, '\u22EE')
        for i in nrow-3:nrow
            print(io, ms.timestamp[i], " | ")
        for j in 1:ncol
            print(io, rpad(ms.values[i,j], colwidth[j] + 2, " "))
        end
        println(io, "")
        end
    else
        for i in 1:nrow
            print(io, ms.timestamp[i], " | ")
        for j in 1:ncol
            print(io, rpad(ms.values[i,j], colwidth[j] + 2, " "))
        end
        println(io, "")
        end
    end
end

# convert bool to float and squish a continuous stream to unique, discrete occurances

function discretesignal(ba::TimeArray{Bool,1})
#function discretesignal(ms::MarketSignal)
    iba = float64(ba.values)
    idx = [1]
    for b in 2:length(ba)
        if ba.values[b] !== ba.values[b-1]
            push!(idx, b)
        end
    end
    TimeArray(ba.timestamp[idx], iba[idx], ["signal"])
end


text:
# Test on a student T (conditional distribution of Y), gamma (distribution of γ),
# inv-gamma (distribution of ω) simulation data example
module MvtGammaInvgammaTest

using Distributions, InteractiveUtils, LinearAlgebra
using Random, Roots, SpecialFunctions, Test, WiSER

@info "MvT Gamma InvGamma Test"
@info "generate data"
Random.seed!(1234)
rng = MersenneTwister(1234)
# dimensions
m  = 1000 # number of individuals
ns = rand(rng, 20:20, m) # numbers of observations per individual
p  = 5    # number of fixed effects, including intercept
q  = 3    # number of random effects, including intercept
l  = 5    # number of WS variance covariates, including intercept
obsvec = Vector{WSVarLmmObs{Float64}}(undef, m)
# true parameter values
βtrue = [ 0.1; 6.5; -3.5; 1.0; 5  ]
τtrue = [-1.5; 1.5; -0.5; 0.0; 0.0]
Σγ    = [1.5 0.5 0.3;
         0.5 1.0 0.2;
         0.3 0.2 0.5];
δγω   = [0.0; 0.0; 0.0] # for easier generation of (γ, ω)
σω    = 1.0
Σγω   = [Σγ δγω; δγω' σω]
Lγω   = cholesky(Symmetric(Σγω), check = false).L
Lγ    = Lγω[1:q, 1:q]
lγω   = Lγω[q + 1, 1:q]
lω    = Lγω[q + 1, q + 1]
# parameters for Gamma random deviate ω
# If X ∼ Gamma(α, θ), then E[-ln(X)] = - ψ(α) - ln(θ), Var[-ln(X)] = ψ'(α).
# ωi = log(inv(X)) = - log(X)
# We want Var[ωi] = ψ'(α) = σω and E[ωi] = - ψ(α) - ln(θ) = 0
ωα    = Roots.find_zero(x -> trigamma(x) - σω, 1)
ωα > 1 || error("ωα needs to be >1 for the existence of mean of inverse-gamma")
ωθ    = exp(-digamma(ωα))
# parameters for Gamma random deviate γ
# If X ∼ Gamma(α, θ), then E[X] = αθ, Var[X] = αθ^2.
# We want Var[X] = 1 and don't care about mean (will shift it to 0)
γα    = 4             # shape parameter
γθ    = sqrt(inv(γα)) # scale parameter
# degree of freedom for t
ν     = 6
# generate data
γ     = Vector{Float64}(undef, q)
z     = similar(γ) # hold vector of iid std normal
for i in 1:m
    # first column intercept, remaining entries iid std normal
    X = Matrix{Float64}(undef, ns[i], p)
    X[:, 1] .= 1
    @views randn!(rng, X[:, 2:p])
    # first column intercept, remaining entries iid std normal
    Z = Matrix{Float64}(undef, ns[i], q)
    Z[:, 1] .= 1
    @views randn!(rng, Z[:, 2:q])
    # first column intercept, remaining entries iid std normal
    W = Matrix{Float64}(undef, ns[i], l)
    W[:, 1] .= 1
    @views randn!(rng, W[:, 2:l])
    # generate ω ∼ log-inv-gamma(ωα, ωθ)
    ω = -log(rand(rng, Gamma(ωα, ωθ)))
    # generate random effects: γ = Lγ * z
    # z is iid Gamma with variance 1 and shifted to have mean 0
    Distributions.rand!(rng, Gamma(γα, γθ), z)
    z .-= γα * γθ # shift to have mean 0
    mul!(γ, Lγ, z)
    # generate y from t distribution (ν, μy, σ2ϵ)
    μy  = X * βtrue + Z * γ
    σ2ϵ = W * τtrue .+ dot(γ, lγω) .+ ω
    ysd = exp.(0.5 .* (σ2ϵ))
    # note: variance of T(ν) is ν / (ν - 2)
    y = μy + sqrt(((ν - 2) / ν)) .* ysd .* rand(rng, TDist(ν), ns[i])
    # form a VarLmmObs instance
    obsvec[i] = WSVarLmmObs(y, X, Z, W)
end
# form WSVarLmmModel
vlmm = WSVarLmmModel(obsvec);

@testset "fit! (start from LS fit)" begin
println(); println(); println()
for solver in [
    # KNITRO.KnitroSolver(outlev=3), # outlev 0-6
    Ipopt.Optimizer()
    # Ipopt.IpoptSolver(print_level=0, watchdog_shortened_iter_trigger=3, max_iter=100),
    # Ipopt.IpoptSolver(print_level = 0)
    # Ipopt.IpoptSolver(print_level = 3, hessian_approximation = "limited-memory"),
    # Ipopt.IpoptSolver(print_level = 3, obj_scaling_factor = 1 / m) # instable
    # Ipopt.IpoptSolver(print_level = 5, mu_strategy = "adaptive") # same speed
    # NLopt.NLoptSolver(algorithm = :LD_SLSQP, maxeval = 4000)
    # NLopt.NLoptSolver(algorithm = :LD_MMA, maxeval = 4000),
    # NLopt.NLoptSolver(algorithm = :LD_LBFGS, maxeval = 4000)
    # NLopt.NLoptSolver(algorithm = :LN_BOBYQA, ftol_rel = 1e-12, ftol_abs = 1e-8, maxeval = 10000)
    ]
    println("----------")
    @show solver
    println("----------")
    @info "init_ls!"
    @time init_ls!(vlmm)
    println("β")
    display([βtrue vlmm.β]); println()
    println("τ")
    display([τtrue vlmm.τ]); println()
    println("Lγω")
    display(vlmm.Lγ); println()
    display(Lγ); println()

    @info "unweighted obj/gradient/hessian at init_ls"
    @show nlsv_obj!(vlmm, true, true, true)
    # println("∇β")
    # display(vlmm.∇β); println()
    # println("∇τ")
    # display(vlmm.∇τ); println()
    # println("∇Lγ")
    # display(vlmm.∇Lγ); println()
    # println("Hββ")
    # display(vlmm.Hββ); println()
    # println("Hττ")
    # display(vlmm.Hττ); println()
    # println("HτLγ")
    # display(vlmm.HτLγ); println()
    # println("HLγLγ")
    # display(vlmm.HLγLγ); println()
    # println("∇:")
    # ∇ = [vlmm.∇τ; vech(vlmm.∇Lγ)]
    # display(∇); println()
    # println("FIM:")
    # FIM = [vlmm.Hττ vlmm.HτLγ; vlmm.HτLγ' vlmm.HLγLγ]
    # display(FIM); println()
    # @show eigvals(Symmetric(FIM))
    # @show Symmetric(FIM) \ ∇

    @info "init_mom!"
    @time init_mom!(vlmm, solver)
    println("β")
    display([βtrue vlmm.β]); println()
    println("τ")
    display([τtrue vlmm.τ]); println()
    println("Lγω")
    display(vlmm.Lγ); println()
    display(Lγ); println()

    @info "weighted obj/gradient/hessian at init_ls"
    init_ls!(vlmm)
    vlmm.iswtnls[1] = true
    update_wtmat!(vlmm)
    @show nlsv_obj!(vlmm, true, true, true)
    # println("∇β")
    # display(vlmm.∇β); println()
    # println("∇τ")
    # display(vlmm.∇τ); println()
    # println("∇Lγ")
    # display(vlmm.∇Lγ); println()
    # println("Hββ")
    # display(vlmm.Hββ); println()
    # println("Hττ")
    # display(vlmm.Hττ); println()
    # println("HτLγ")
    # display(vlmm.HτLγ); println()
    # println("HLγLγ")
    # display(vlmm.HLγLγ); println()
    println("∇:")
    ∇ = [vlmm.∇τ; vech(vlmm.∇Lγ)]
    display(∇); println()
    # println("FIM:")
    # FIM = [vlmm.Hττ vlmm.HτLγ; vlmm.HτLγ' vlmm.HLγLγ]
    # display(FIM); println()
    # @show eigvals(Symmetric(FIM))
    # @show Symmetric(FIM) \ ∇

    @info "WNLS fitting"
    @time WiSER.fit!(vlmm, solver, runs=2) # start from init_ls by default
    @info "obj at solution"
    @show nlsv_obj!(vlmm, true, true)
    @info "estimates at solution"
    println("β")
    display([βtrue vlmm.β]); println()
    println("τ")
    display([τtrue vlmm.τ]); println()
    println("Lγω")
    display(vlmm.Lγ); println()
    display(Lγ); println()
    # @info "gradient at solution"
    # @show vlmm.∇β
    # @show vlmm.∇τ
    # @show vlmm.∇Lγ
    @test sqrt(abs2(norm(vlmm.∇τ)) + abs2(norm(vlmm.∇Lγ))) < 1e-5
    # @info "Hessian at solution"
    # @show vlmm.Hββ
    # @show vlmm.HLγLγ
    # @show vlmm.HτLγ
    # @show vlmm.Hττ

    @info "inference at solution"
    show(vlmm)
end
end

end


text:
using Stheno: get_f, get_y, Observation, merge
using BlockArrays

@testset "conditioning" begin

# Test Observation functionality.
let
    rng, N, N′, D = MersenneTwister(123456), 5, 6, 2
    X, X′ = ColsAreObs(randn(rng, D, N)), ColsAreObs(randn(rng, D, N′))
    y, y′ = randn(rng, N), randn(rng, N′)
    f = GP(ConstantMean(1), EQ(), GPC())

    fX, fX′ = f(X), f(X′)
    c1, c2 = fX←y, fX′←y′
    @test Observation(fX, y) == c1
    @test get_f(c1) === fX && get_f(c2) === fX′
    @test get_y(c1) === y && get_y(c2) === y′

    c = merge((c1, c2))
    @test get_y(c) == BlockVector([y, y′])
    @test get_f(c).fs == [fX, fX′]
end

let
    rng, N, N′, D = MersenneTwister(123456), 5, 6, 2
    X, X′ = ColsAreObs(randn(rng, D, N)), ColsAreObs(randn(rng, D, N′))
    y = randn(rng, N)

    # Test mechanics for finite conditioned process with single conditioning.
    f = GP(ConstantMean(1), EQ(), GPC())
    f′ = f | (f(X) ← y)
    @test length(f′) == Inf
    @test length(rand(rng, f′(X))) == N
    @test maximum(abs.(rand(rng, f′(X)) - y)) < 1e-3
    @test maximum(abs.(mean_vec(f′(X)) - y)) < 1e-3
    @test all(abs.(cov(f′(X))) .< 1e-6)
end

# Some tests for self-consistency in the posterior when singly and multiply conditioning.
let
    rng, N, N′, D = MersenneTwister(123456), 1, 2, 2
    X_, X′_ = randn(rng, D, N), randn(rng, D, N′)
    X, X′, Z = ColsAreObs(X_), ColsAreObs(X′_), ColsAreObs(randn(rng, D, N + N′))
    μ, k, XX′ = ConstantMean(1.0), EQ(), ColsAreObs(hcat(X_, X′_))
    f = GP(μ, k, GPC())
    y = rand(rng, f(XX′))

    # Construct full posterior
    f′X = f(X) | (f(XX′) ← y)
    f′X′ = f(X′) | (f(XX′) ← y)
    f′XX′ = f(XX′) | (f(XX′) ← y)
    f′ = f | (f(XX′) ← y)

    # Check that each method is self-consistent.
    Σ′XX′ = cov(f′XX′)
    @test Σ′XX′ ≈ cov(f′(XX′))
    @test Σ′XX′[1:N, 1:N] ≈ cov(f′X)[1:N, 1:N]
    @test Σ′XX′[N+1:N+N′, N+1:N+N′] ≈ cov(f′X′)[1:N′, 1:N′]
    # @test maximum(abs.(Σ′XX′[1:N, N+1:N+N′] - xcov(f′X, f′X′, 1:N, 1:N′))) < 1e-12
    # @test maximum(abs.(Σ′XX′[N+1:N+N′, 1:N] - xcov(f′X′, f′X, 1:N′, 1:N))) < 1e-12

    # Test that conditioning works for BlockGPs.
    fb, Xb = BlockGP([f, f]), BlockData([X, X′])
    Zb = BlockData([Z, Z])
    fb′ = fb | (fb(Xb)←BlockArray(y, [N, N′]))
    @test mean_vec(fb′(Zb)) ≈ mean_vec(f′(Zb))
    @test maximum(abs.(cov(fb′(Zb)) - cov(f′(Zb)))) < 1e-6

    f′b = f | (fb(Xb)←BlockArray(y, [N, N′]))
    @test mean_vec(f′b(X)) ≈ mean_vec(f′X)
    @test maximum(abs.(cov(f′b(Zb)) - cov(f′(Zb)))) < 1e-6

    # Test sugar for multiple-conditioning.
    @test mean_vec(fb′(Zb)) ≈ mean_vec((f | (f(X)←y[1:N], f(X′)←y[N+1:end]))(Zb))
    @test maximum(abs.(cov(fb′(Zb)) - cov((f | (f(X)←y[1:N], f(X′)←y[N+1:end]))(Zb)))) < 1e-6

    yX = rand(rng, f(X))
    f′X, f′X′ = (f(X), f(X′)) | (f(X)←yX)
    f′X1, f′X′1 = f(X) | (f(X)←yX), f(X′) | (f(X)←yX)
    @test mean_vec(f′X) ≈ mean_vec(f′X1)
    @test mean_vec(f′X′) ≈ mean_vec(f′X′1)
    @test maximum(abs.(cov(f′X) - cov(f′X1))) < 1e-6
    @test maximum(abs.(cov(f′X′) - cov(f′X′1))) < 1e-6

    yX, yX′ = y[1:N], y[N+1:end]
    f′X, f′X′ = (f(X), f(X′)) | (f(X)←yX, f(X′)←yX′)
    f′X1, f′X′1 = f(X) | (f(X)←yX, f(X′)←yX′), f(X′) | (f(X)←yX, f(X′)←yX′)
    @test mean_vec(f′X) ≈ mean_vec(f′X1)
    @test mean_vec(f′X′) ≈ mean_vec(f′X′1)
    @test maximum(abs.(cov(f′X) - cov(f′X1))) < 1e-6
    @test maximum(abs.(cov(f′X′) - cov(f′X′1))) < 1e-6
end

# Test Titsias implementation by checking that it (approximately) recovers exact inference
# when M = N and Z = X.
let
    rng, N, N′, D, σ² = MersenneTwister(123456), 2, 3, 5, 1e-1
    X_, X′_ = randn(rng, D, N), randn(rng, D, N′)
    X, X′, Z = ColsAreObs(X_), ColsAreObs(X′_), ColsAreObs(randn(rng, D, N + N′))
    μ, k, XX′ = ConstantMean(1.0), EQ(), ColsAreObs(hcat(X_, X′_))

    # Construct toy problem.
    gpc = GPC()
    f = GP(μ, k, gpc)
    y = f + GP(Noise(σ²), gpc)
    ŷ = rand(rng, y(XX′))

    # Compute exact posterior.
    f′XX′ = f(XX′) | (y(XX′)←ŷ)

    # Compute approximate posterior suff. stats.
    μᵤ, Σᵤᵤ = Stheno.optimal_q(f(XX′), ŷ, f(XX′), sqrt(σ²))

    # Check that exact and approx. posteriors are close in this case.
    @test isapprox(μᵤ, mean_vec(f′XX′); rtol=1e-4)
    @test isapprox(Σᵤᵤ, cov(f′XX′); rtol=1e-4)

    # Compute conditioner and exact posterior compute at test points.
    conditioner = Stheno.Titsias(f(XX′), μᵤ, Σᵤᵤ)
    f′Z = f(Z) | (y(XX′)←ŷ)
    f′Z_approx = f(Z) | conditioner

    # Check that exact and approximate posteriors match up.
    @test isapprox(mean_vec(f′Z), mean_vec(f′Z_approx); rtol=1e-4)
    @test isapprox(cov(f′Z), cov(f′Z_approx); rtol=1e-4)

    # Check that Titsias with BlockGP works the same as Titsias with regular GP.
    ŷX, ŷX′ = ŷ[1:N], ŷ[N+1:end]

    fb, ŷb = BlockGP([f(X), f(X′)]), BlockVector([ŷX, ŷX′])
    μb, Σb = Stheno.optimal_q(fb, ŷb, fb, sqrt(σ²))

    @test μb isa BlockVector
    @test Σb isa LazyPDMat
    @test Stheno.unbox(Σb) isa Symmetric
    @test Stheno.unbox(Stheno.unbox(Σb)) isa AbstractBlockMatrix
    @test μb ≈ μᵤ
    @test Σb ≈ Σᵤᵤ

    # Test that conditioning is indifferent to choice of Blocks.
    conditioner_blocked = Stheno.Titsias(fb, μb, Σb)
    f′Zb = f(BlockData([Z])) | conditioner_blocked

    @test isapprox(mean_vec(f′Z), mean_vec(f′Zb); rtol=1e-4)
    @test isapprox(cov(f′Z), cov(f′Zb); rtol=1e-4)
end

end


text:
using HypothesisTests, Test

@testset "Show" begin
# based on power_divergence.jl tests
d = [[762,484] [327,239] [468,477]]
m = PowerDivergenceTest(d)

@test sprint(show, m, context=:compact => true) ==
    """
    Pearson's Chi-square Test
    -------------------------
    Population details:
        parameter of interest:   Multinomial Probabilities
        value under h_0:         [0.255231, 0.19671, 0.11594, 0.0893561, 0.193574, 0.14919]
        point estimate:          [0.276387, 0.175553, 0.118607, 0.0866884, 0.16975, 0.173014]
        95% confidence interval: Tuple{Float64,Float64}[(0.2579, 0.2955), (0.1571, 0.1947), (0.1001, 0.1378), (0.0682, 0.1058), (0.1513, 0.1889), (0.1545, 0.1922)]

    Test summary:
        outcome with 95% confidence: reject h_0
        one-sided p-value:           <1e-6

    Details:
        Sample size:        2757
        statistic:          30.070149095754687
        degrees of freedom: 2
        residuals:          [2.19886, -2.50467, 0.41137, -0.468583, -2.84324, 3.23867]
        std. residuals:     [4.50205, -4.50205, 0.699452, -0.699452, -5.31595, 5.31595]
    """

d = [ 20, 15, 25 ]
m = PowerDivergenceTest(d)

@test sprint(show, m, context=:compact => true) ==
    """
    Pearson's Chi-square Test
    -------------------------
    Population details:
        parameter of interest:   Multinomial Probabilities
        value under h_0:         [0.333333, 0.333333, 0.333333]
        point estimate:          [0.333333, 0.25, 0.416667]
        95% confidence interval: Tuple{Float64,Float64}[(0.2167, 0.481), (0.1333, 0.3976), (0.3, 0.5643)]

    Test summary:
        outcome with 95% confidence: fail to reject h_0
        one-sided p-value:           0.2865

    Details:
        Sample size:        60
        statistic:          2.5
        degrees of freedom: 2
        residuals:          [0.0, -1.11803, 1.11803]
        std. residuals:     [0.0, -1.36931, 1.36931]
    """

# based on t.jl tests
tst = OneSampleTTest(-5:10)

@test sprint(show, tst) ==
    """
    One sample t-test
    -----------------
    Population details:
        parameter of interest:   Mean
        value under h_0:         0
        point estimate:          2.5
        95% confidence interval: (-0.0369, 5.0369)

    Test summary:
        outcome with 95% confidence: fail to reject h_0
        two-sided p-value:           0.0530

    Details:
        number of observations:   16
        t-statistic:              2.100420126042015
        degrees of freedom:       15
        empirical standard error: 1.1902380714238083
    """

end


text:
"""
    module GMMInference

"""
module GMMInference

using JuMP, NLPModels, Distributions, LinearAlgebra, Optim,
  ForwardDiff, ECOS
import Convex

export GMMModel,
  number_parameters,
  number_observations,
  number_moments,
  get_gi,
  gmm_objective,
  gmm_constraint,
  cue_objective,
  gel_nlp_problem,
  gel_jump_problem,
  gel_optim_args,
  gmm_nlp_problem,
  gmm_jump_problem,
  gel_pλ,
  gel_tests,
  IVLogitShare,
  MixturePanel,
  RCLogit

include("gmmmodel.jl")  # defines GMMModel and estimation methods
include("inference.jl") # test statistics and p-values

# example models
include("models/ivlogitshare.jl")
include("models/rclogit.jl")
include("models/mixturepanel.jl")

end # module