text:
# Unit tests for judiVector
# Philipp Witte (pwitte.slim@gmail.com)
# May 2018
#
# Mathias Louboutin, mlouboutin3@gatech.edu
# Updated July 2020

datapath = joinpath(dirname(pathof(JUDI)))*"/../data/"

# number of sources/receivers
nrec = 120
ns = 251
ftol = 1e-6

################################################# test constructors ####################################################

@testset "judiVector Unit Tests with $(nsrc) sources" for nsrc=[1, 2]
    @timeit TIMEROUTPUT "judiVector (nsrc=$(nsrc))" begin
        # set up judiVector from array
        dsize = (nsrc,)
        rec_geometry = example_rec_geometry(nsrc=nsrc, nrec=nrec)
        data = randn(Float32, ns, nrec)
        d_obs = judiVector(rec_geometry, data)
        @test typeof(d_obs) == judiVector{Float32, Array{Float32, 2}}
        @test isequal(process_input_data(d_obs, rec_geometry), d_obs)
        @test isequal(process_input_data(d_obs), d_obs.data)

        @test isequal(d_obs.nsrc, nsrc)
        @test isequal(typeof(d_obs.data), Array{Array{Float32, 2}, 1})
        @test isequal(typeof(d_obs.geometry), GeometryIC{Float32})
        @test iszero(norm(d_obs.data[1] - d_obs.data[end]))
        @test isequal(size(d_obs), dsize)

        # set up judiVector from cell array
        data = Array{Array{Float32, 2}, 1}(undef, nsrc)
        for j=1:nsrc
            data[j] = randn(Float32, ns, nrec)
        end
        d_obs =  judiVector(rec_geometry, data)

        @test isequal(d_obs.nsrc, nsrc)
        @test isequal(typeof(d_obs.data), Array{Array{Float32, 2}, 1})
        @test isequal(typeof(d_obs.geometry), GeometryIC{Float32})
        @test iszero(norm(d_obs.data - d_obs.data))
        @test isequal(size(d_obs), dsize)
        @test isapprox(convert_to_array(d_obs), vcat([vec(d) for d in data]...); rtol=ftol)

        # contructor for in-core data container
        block = segy_read(datapath*"unit_test_shot_records_$(nsrc).segy"; warn_user=false)
        d_block = judiVector(block; segy_depth_key="RecGroupElevation")
        dsize = (nsrc,)

        @test isequal(d_block.nsrc, nsrc)
        @test isequal(typeof(d_block.data), Array{Array{Float32, 2}, 1})
        @test isequal(typeof(d_block.geometry), GeometryIC{Float32})
        @test isequal(size(d_block), dsize)

        # contructor for in-core data container and given geometry
        rec_geometry = Geometry(block; key="receiver", segy_depth_key="RecGroupElevation")
        d_block = judiVector(rec_geometry, block)

        @test isequal(d_block.nsrc, nsrc)
        @test isequal(typeof(d_block.data), Array{Array{Float32, 2}, 1})
        @test isequal(typeof(d_block.geometry), GeometryIC{Float32})
        @test isequal(rec_geometry, d_block.geometry)
        @test isequal(size(d_block), dsize)

        # contructor for out-of-core data container from single container
        container = segy_scan(datapath, "unit_test_shot_records_$(nsrc)", ["GroupX", "GroupY", "RecGroupElevation", "SourceSurfaceElevation", "dt"])
        d_cont = judiVector(container; segy_depth_key="RecGroupElevation")

        @test typeof(d_cont) == judiVector{Float32, SeisCon}
        @test isequal(d_cont.nsrc, nsrc)
        @test isequal(typeof(d_cont.data), Array{SegyIO.SeisCon, 1})
        @test isequal(typeof(d_cont.geometry), GeometryOOC{Float32})
        @test isequal(size(d_cont), dsize)

        # contructor for single out-of-core data container and given geometry
        d_cont = judiVector(rec_geometry, container)

        @test isequal(d_cont.nsrc, nsrc)
        @test isequal(typeof(d_cont.data), Array{SegyIO.SeisCon, 1})
        @test isequal(typeof(d_cont.geometry), GeometryIC{Float32})
        @test isequal(rec_geometry, d_cont.geometry)
        @test isequal(size(d_cont), dsize)

        # contructor for out-of-core data container from cell array of containers
        container_cell = Array{SegyIO.SeisCon}(undef, nsrc)
        for j=1:nsrc
            container_cell[j] = split(container, j)
        end
        d_cont =  judiVector(container_cell; segy_depth_key="RecGroupElevation")

        @test isequal(d_cont.nsrc, nsrc)
        @test isequal(typeof(d_cont.data), Array{SegyIO.SeisCon, 1})
        @test isequal(typeof(d_cont.geometry), GeometryOOC{Float32})
        @test isequal(size(d_cont), dsize)

        # contructor for out-of-core data container from cell array of containers and given geometry
        d_cont =  judiVector(rec_geometry, container_cell)

        @test isequal(d_cont.nsrc, nsrc)
        @test isequal(typeof(d_cont.data), Array{SegyIO.SeisCon, 1})
        @test isequal(typeof(d_cont.geometry), GeometryIC{Float32})
        @test isequal(rec_geometry, d_cont.geometry)
        @test isequal(size(d_cont), dsize)

    #################################################### test operations ###################################################

        # conj, transpose, adjoint
        @test isequal(size(d_obs), size(conj(d_obs)))
        @test isequal(size(d_block), size(conj(d_block)))
        @test isequal(size(d_cont), size(conj(d_cont)))

        @test isequal(reverse(size(d_obs)), size(transpose(d_obs)))
        @test isequal(reverse(size(d_block)), size(transpose(d_block)))
        @test isequal(reverse(size(d_cont)), size(transpose(d_cont)))

        @test isequal(reverse(size(d_obs)), size(adjoint(d_obs)))
        @test isequal(reverse(size(d_block)), size(adjoint(d_block)))
        @test isequal(reverse(size(d_cont)), size(adjoint(d_cont)))

        # +, -, *, /
        @test iszero(norm(2*d_obs - (d_obs + d_obs)))
        @test iszero(norm(d_obs - (d_obs + d_obs)/2))

        @test iszero(norm(2*d_block - (d_block + d_block)))
        @test iszero(norm(d_block - (d_block + d_block)/2))

        # lmul!, rmul!, ldiv!, rdiv!
        data_ones = Array{Array{Float32, 2}, 1}(undef, nsrc)
        for j=1:nsrc
            data_ones[j] = ones(Float32, ns, nrec)
        end
        d1 =  judiVector(rec_geometry, data_ones)
        lmul!(2f0, d1)
        @test all([all(d1.data[i] .== 2f0) for i = 1:d1.nsrc])
        rmul!(d1, 3f0)
        @test all([all(d1.data[i] .== 6f0) for i = 1:d1.nsrc])
        ldiv!(2f0,d1)
        @test all([all(d1.data[i] .== 3f0) for i = 1:d1.nsrc])
        rdiv!(d1, 3f0)
        @test all([all(d1.data[i] .== 1f0) for i = 1:d1.nsrc])

        # vcat
        d_vcat = [d_block; d_block]
        @test isequal(length(d_vcat), 2*length(d_block))
        @test isequal(d_vcat.nsrc, 2*d_block.nsrc)
        @test isequal(d_vcat.geometry.xloc[1], d_block.geometry.xloc[1])

        # dot, norm, abs
        @test isapprox(norm(d_block), sqrt(dot(d_block, d_block)))
        @test isapprox(norm(d_cont), sqrt(dot(d_cont, d_cont)))
        @test isapprox(abs.(d_block.data[1]), abs(d_block).data[1]) # need to add iterate for JUDI vector

        # vector space axioms
        u = judiVector(rec_geometry, randn(Float32, ns, nrec))
        v = judiVector(rec_geometry, randn(Float32, ns, nrec))
        w = judiVector(rec_geometry, randn(Float32, ns, nrec))
        a = .5f0 + rand(Float32)
        b = .5f0 + rand(Float32)

        @test isapprox(u + (v + w), (u + v) + w; rtol=ftol)
        @test isapprox(u + v, v + u; rtol=ftol)
        @test isapprox(-u, -1f0 * u; rtol=ftol)
        @test iszero(norm(u + u*(-1)))
        @test isapprox(a .* (b .* u), (a * b) .* u; rtol=ftol)
        @test isapprox(u, u .* 1; rtol=ftol)
        @test isapprox(a .* (u + v), a .* u + a .* v; rtol=1f-5)
        @test isapprox((a + b) .* v, a .* v + b .* v; rtol=1f-5)

        # subsamling
        d_block_sub = d_block[1]
        @test isequal(d_block_sub.nsrc, 1)
        @test isequal(typeof(d_block_sub.geometry), GeometryIC{Float32})
        @test isequal(typeof(d_block_sub.data), Array{Array{Float32, 2}, 1})

        inds = nsrc > 1 ? (1:nsrc) : 1
        d_block_sub = d_block[inds]
        @test isequal(d_block_sub.nsrc, nsrc)
        @test isequal(typeof(d_block_sub.geometry), GeometryIC{Float32})
        @test isequal(typeof(d_block_sub.data), Array{Array{Float32, 2}, 1})

        d_cont =  judiVector(container_cell; segy_depth_key="RecGroupElevation")
        d_cont_sub = d_cont[1]
        @test isequal(d_cont_sub.nsrc, 1)
        @test isequal(typeof(d_cont_sub.geometry), GeometryOOC{Float32})
        @test isequal(typeof(d_cont_sub.data), Array{SegyIO.SeisCon, 1})

        d_cont_sub = d_cont[inds]
        @test isequal(d_cont_sub.nsrc, nsrc)
        @test isequal(typeof(d_cont_sub.geometry), GeometryOOC{Float32})
        @test isequal(typeof(d_cont_sub.data), Array{SegyIO.SeisCon, 1})

        # Conversion to SegyIO.Block
        src_geometry = Geometry(block; key="source", segy_depth_key="SourceSurfaceElevation")
        wavelet = randn(Float32, src_geometry.nt[1])
        q = judiVector(src_geometry, wavelet)
        block_out =  judiVector_to_SeisBlock(d_block, q; source_depth_key="SourceSurfaceElevation", receiver_depth_key="RecGroupElevation")

        @test isapprox(block.data, block_out.data)
        @test isapprox(get_header(block, "SourceX"), get_header(block_out, "SourceX"); rtol=1f-6)
        @test isapprox(get_header(block, "GroupX"), get_header(block_out, "GroupX"); rtol=1f-6)
        @test isapprox(get_header(block, "RecGroupElevation"), get_header(block_out, "RecGroupElevation"); rtol=1f-6)
        @test isequal(get_header(block, "ns"), get_header(block_out, "ns"))
        @test isequal(get_header(block, "dt"), get_header(block_out, "dt"))

        block_obs =  judiVector_to_SeisBlock(d_obs, q; source_depth_key="SourceSurfaceElevation", receiver_depth_key="RecGroupElevation")
        d_obs1 = judiVector(block_obs)
        @test isapprox(d_obs1.data, d_obs.data)

        block_q = src_to_SeisBlock(q)
        q_1 = judiVector(block_q)
        for i = 1:nsrc
            vec(q_1.data[i]) == vec(q.data[i])
        end

        # scale
        a = .5f0 + rand(Float32)
        d_scale = deepcopy(d_block)

        # Test norms
        d_ones = judiVector(rec_geometry, 2f0 .* ones(Float32, ns, nrec))
        @test isapprox(norm(d_ones, 2), sqrt(rec_geometry.dt[1]*nsrc*ns*nrec*4))
        @test isapprox(norm(d_ones, 1), rec_geometry.dt[1]*nsrc*ns*nrec*2)
        @test isapprox(norm(d_ones, Inf), 2)

        # Indexing and utilities
        @test isfinite(d_obs)
        @test ndims(judiVector{Float32, Array{Float32, 2}}) == 1
        @test isapprox(sum(d_obs), sum(sum([vec(d) for d in d_obs])))
        d0 = copy(d_obs)
        fill!(d0, 0f0)
        @test iszero(norm(d0))

        @test firstindex(d_obs) == 1
        @test lastindex(d_obs) == nsrc
        @test axes(d_obs) == (Base.OneTo(nsrc),)
        @test ndims(d_obs) == 1

        d0[1] = d_obs.data[1]
        @test isapprox(d0.data[1], d_obs.data[1])

        # broadcast multiplication
        u = judiVector(rec_geometry, randn(Float32, ns, nrec))
        v = judiVector(rec_geometry, randn(Float32, ns, nrec))
        u_scale = deepcopy(u)
        v_scale = deepcopy(v)

        u_scale .*= 2f0
        @test isapprox(u_scale, 2f0 * u; rtol=ftol)
        v_scale .+= 2f0
        @test isapprox(v_scale, 2f0 + v; rtol=ftol)
        u_scale ./= 2f0
        @test isapprox(u_scale, u; rtol=ftol)
        u_scale .= 2f0 .* u_scale .+ v_scale
        @test isapprox(u_scale, 2f0 * u + 2f0 + v; rtol=ftol)
        u_scale .= u .+ v
        @test isapprox(u_scale, u + v)
        u_scale .= u .- v
        @test isapprox(u_scale, u - v)
        u_scale .= u .* v
        @test isapprox(u_scale.data[1], u.data[1].*v.data[1])
        u_scale .= u ./ v
        @test isapprox(u_scale.data[1], u.data[1]./v.data[1])

        # broadcast identity
        u = judiVector(rec_geometry, randn(Float32, ns, nrec))
        v = judiVector(rec_geometry, randn(Float32, ns, nrec))
        u_id = deepcopy(u)
        v_id = deepcopy(v)
        broadcast!(identity, u_id, v_id)    # copy v_id into u_id

        @test isapprox(v, v_id)
        @test isapprox(v, u_id)

        # in-place overwrite
        u = judiVector(rec_geometry, randn(Float32, ns, nrec))
        v = judiVector(rec_geometry, randn(Float32, ns, nrec))
        u_cp = deepcopy(u)
        v_cp = deepcopy(v)
        copy!(v_cp, u_cp)

        @test isapprox(u, u_cp)
        @test isapprox(u, v_cp)

        # similar
        d_zero = similar(d_block, Float32)

        @test isequal(d_zero.geometry, d_block.geometry)
        @test isequal(size(d_zero), size(d_block))

        # retrieve out-of-core data
        d_get = get_data(d_cont)
        @test isapprox(d_block, d_get)

        # Test copies/similar
        w1 = deepcopy(d_obs)
        @test isapprox(w1, d_obs)
        w1 = similar(d_obs)
        @test w1.nsrc == d_obs.nsrc
        w1 .= d_obs
        @test w1.nsrc == d_obs.nsrc
        @test isapprox(w1.data, d_obs.data)

        # Test transducer
        q = judiVector(Geometry(0f0, 0f0, 0f0; dt=2, t=1000), randn(Float32, 251))
        tr = transducer(q, (10, 10), 30, pi/2)
        @test length(tr.geometry.xloc[1]) == 22
        @test tr.geometry.xloc[1][1:11] == range(-30., 30., length=11)
        @test tr.geometry.xloc[1][12:end] == range(-30., 30., length=11)
        @test all(tr.geometry.zloc[1][12:end] .== -10f0)
        @test all(tr.geometry.zloc[1][1:11] .== 0f0)

        q = judiVector(Geometry(0f0, 0f0, 0f0; dt=2, t=1000), randn(Float32, 251))
        tr = transducer(q, (10, 10), 30, pi)
        @test length(tr.geometry.xloc[1]) == 22
        @test isapprox(tr.geometry.zloc[1][1:11], range(30., -30., length=11); atol=1f-14, rtol=1f-14)
        @test isapprox(tr.geometry.zloc[1][12:end], range(30., -30., length=11); atol=1f-14, rtol=1f-14)
        @test isapprox(tr.geometry.xloc[1][12:end], -10f0*ones(11); atol=1f-14, rtol=1f-14)
        @test isapprox(tr.geometry.xloc[1][1:11], zeros(11); atol=1f-14, rtol=1f-14)

        # Test integral & derivative
        refarray = Array{Array{Float32, 2}, 1}(undef, nsrc)
        for j=1:nsrc
            refarray[j] = randn(Float32, ns, nrec)
        end
        d_orig = judiVector(rec_geometry, refarray)

        dt = rec_geometry.dt[1]

        d_cumsum = cumsum(d_orig)
        for i = 1:d_orig.nsrc
            @test isapprox(dt * cumsum(refarray[i],dims=1), d_cumsum.data[i])
        end

        d_diff = diff(d_orig)
        for i = 1:d_orig.nsrc
            @test isapprox(1/dt * refarray[i][1,:], d_diff.data[i][1,:])
            @test isapprox(d_diff.data[i][2:end,:], 1/dt * diff(refarray[i],dims=1))
        end

        @test isapprox(cumsum(d_orig,dims=1),cumsum(d_orig))
        @test isapprox(diff(d_orig,dims=1),diff(d_orig))

        d_cumsum_rec = cumsum(d_orig,dims=2)
        for i = 1:d_orig.nsrc
            @test isapprox(cumsum(refarray[i],dims=2), d_cumsum_rec.data[i])
        end

        d_diff_rec = diff(d_orig,dims=2)
        for i = 1:d_orig.nsrc
            @test isapprox(refarray[i][:,1], d_diff_rec.data[i][:,1])
            @test isapprox(d_diff_rec.data[i][:,2:end], diff(refarray[i],dims=2))
        end
    end
end


text:
const den_signs = Ref([" ", "░", "▒", "▓", "█"])

"""
Unlike the `BrailleCanvas`, the density canvas
does not simply mark a "pixel" as set.
Instead it increments a counter per character
that keeps track of the frequency of pixels
drawn in that character. Together with a variable
that keeps track of the maximum frequency,
the canvas can thus draw the density of datapoints.
"""
mutable struct DensityCanvas <: Canvas
    grid::Array{UInt,2}
    colors::Array{ColorType,2}
    pixel_width::Int
    pixel_height::Int
    origin_x::Float64
    origin_y::Float64
    width::Float64
    height::Float64
    xscale::Union{Symbol,Function}
    yscale::Union{Symbol,Function}
    max_density::Float64
end

@inline pixel_width(c::DensityCanvas) = c.pixel_width
@inline pixel_height(c::DensityCanvas) = c.pixel_height
@inline origin_x(c::DensityCanvas) = c.origin_x
@inline origin_y(c::DensityCanvas) = c.origin_y
@inline width(c::DensityCanvas) = c.width
@inline height(c::DensityCanvas) = c.height
@inline nrows(c::DensityCanvas) = size(c.grid, 2)
@inline ncols(c::DensityCanvas) = size(c.grid, 1)

@inline x_pixel_per_char(::Type{DensityCanvas}) = 1
@inline y_pixel_per_char(::Type{DensityCanvas}) = 2

function DensityCanvas(char_width::Int, char_height::Int;
                       origin_x::Number = 0.,
                       origin_y::Number = 0.,
                       width::Number = 1.,
                       height::Number = 1.,
                       xscale::Union{Symbol,Function} = :identity,
                       yscale::Union{Symbol,Function} = :identity)
    width > 0 || throw(ArgumentError("width has to be positive"))
    height > 0 || throw(ArgumentError("height has to be positive"))
    char_width = max(char_width, 5)
    char_height = max(char_height, 5)
    pixel_width = char_width * x_pixel_per_char(DensityCanvas)
    pixel_height = char_height * y_pixel_per_char(DensityCanvas)
    grid = fill(0, char_width, char_height)
    colors = fill(nothing, char_width, char_height)
    DensityCanvas(grid, colors,
                  pixel_width, pixel_height,
                  Float64(origin_x), Float64(origin_y),
                  Float64(width), Float64(height),
                  xscale, yscale, 1)
end

function pixel_to_char_point(c::DensityCanvas, pixel_x::Number, pixel_y::Number)
    pixel_x = pixel_x < c.pixel_width ? pixel_x : pixel_x - 1
    pixel_y = pixel_y < c.pixel_height ? pixel_y : pixel_y - 1
    cw, ch = size(c.grid)
    char_x = floor(Int, pixel_x / c.pixel_width * cw) + 1
    char_y = floor(Int, pixel_y / c.pixel_height * ch) + 1
    char_x, char_y
end

function pixel!(c::DensityCanvas, pixel_x::Int, pixel_y::Int, color::UserColorType)
    0 <= pixel_x <= c.pixel_width || return c
    0 <= pixel_y <= c.pixel_height || return c
    char_x, char_y = pixel_to_char_point(c, pixel_x, pixel_y)
    c.grid[char_x,char_y] += 1
    c.max_density = max(c.max_density, c.grid[char_x,char_y])
    set_color!(c.colors, char_x, char_y, crayon_256_color(color))
    c
end

function printrow(io::IO, c::DensityCanvas, row::Int)
    0 < row <= nrows(c) || throw(ArgumentError("Argument row out of bounds: $row"))
    signs = den_signs[]
    y = row
    den_sign_count = length(signs)
    val_scale = (den_sign_count - 1) / c.max_density
    for x in 1:ncols(c)
        den_index = round(Int, c.grid[x,y] * val_scale, RoundNearestTiesUp) + 1
        print_color(c.colors[x,y], io, signs[den_index])
    end
    nothing
end

ext: jl


text:
#!/usr/bin/env julia
# Copyright 2017, Andrey Prokopenko, ORNL

include("meshnd.jl")
include("meshsparse.jl")
include("ca_utils.jl")

## Input data
# Note: the matrix is assumed to be symmetric. If that is not true, one would
# need to edit the underlap() function to properly compute the first underlap
# set (other sets are based on this one)
global m
if true
    ## Matrix parameters
    matrix_source = "generator"
    # matrix_source = 'file'

    if matrix_source == "generator"
        # Domain and subdomains parameters
        #  nxyz - number of elements in array sets the problem dimension.
        #  mxyz - number of subdomains in each direction
        #         Note that this array can be larger than the nxyz array, in
        #         which case we reduce its size.
        nxyz = [4 4]
        mxyz = [1 1]

    elseif matrix_source == "file"
        # Number of subdomains
        m = 16
        # matrix_file = 'G3_circuit/G3_circuit.mtx'
        # matrix_file = 'Boeing/msc04515.mtx'
        matrix_file = "A.mm"
    end
end

global ilu_setup
if true
    ## Solver parameters
    # Outer solver parameters
    # NOTE: RAS preconditioners cannot be used with CG (due to nonsymmetry)

    # solver = 'CG'
    # solver = 'GMRES-left'
    solver = "GMRES-right"

    max_its = 1000
    tol     = 1e-12
    restart = 20

    # Subdomain solver parameters
    sub_solver = "direct"

    # sub_solver = 'ilu'
    if sub_solver == "ilu"
        # ilu_setup.type    = "crout"
        ilu_setup.type    = "nofill"
        ilu_setup.milu    = "row"
        ilu_setup.droptol = 0.1
    end
end

## Check input parameters consistency
if matrix_source == "generator"
    # Check that the dimension of subdomain array is at least that of nxyz
    # array (if smaller, last dimension is not used)
    @assert(size(mxyz,2) >= size(nxyz,2))
    mxyz = mxyz[1:size(nxyz,2)]

    # Check that there is at least one point per dimension per subdomain
    # FIXME Julia complains about the assert
    # @assert(size(find(mxyz > nxyz),2) == 0, "Elements of mxyz must be <= than corresponding elements of nxyz")

    m = prod(mxyz)
end

## Construct matrix and rhs
global A
global coords
if matrix_source == "generator"
    coords = zeros(prod(nxyz), 3)

    if size(nxyz,2) == 1
        nx = nxyz[1]
        A = meshsparse(meshnd(nx))
        coords[:,1] = [1:nx]

    elseif size(nxyz,2) == 2
        nx = nxyz[1]
        ny = nxyz[2]

        # Construct the matrix
        # NOTE: we swap x and y dimensions as meshnd uses YX ordering
        A = meshsparse(meshnd(ny, nx), 5);  # 5-point stencil
        # A = meshsparse(meshnd(ny, nx), 9);  # 9-point stencil

        # Construct coordinates
        for j = 1:ny
            for i = 1:nx
                coords[(j-1)*nx + i,:] = [i j 0]
            end
        end

    else
        nx = nxyz[1]
        ny = nxyz[2]
        nz = nxyz[3]

        # Construct the matrix
        # NOTE: we swap x and y dimensions as meshnd uses YXZ ordering
        A = meshsparse(meshnd(ny, nx, nz),  7); #  7-point stencil
        # A = meshsparse(meshnd(ny, nx, nz), 27); # 27-point stencil

        # Construct coordinates
        for k = 1:nz
            for j = 1:ny
                for i = 1:nx
                    coords[(k-1)*ny*nx + (j-1)*nx + i,:] = [i j k]
                end
            end
        end
    end

elseif matrix_source == "file"
    # Read in the matrix
    A = mmread(matrix_file)

    # No coordinates are available
end
# Make sure the matrix is square
@assert(size(A,1) == size(A,2))

n = size(A,1)

# Set seed for reproducibility
# rng(1426)

b = A*rand(n,1)

## Construct partitioning
# row_gids0 and col_gids0 are both cell arrays, and follow Petra notion, i.e.
#   row_gids{I} contains global DOFs of a row    map of subdomain I
#   col_gids{I} contains global DOFs of a column map of subdomain I
# gid2proc is the global ID -> subdomain ID mapping
global row_gids0
global col_gids0
global gid2proc
if matrix_source == "generator"
    row_gids0 = partition_manual(A, nxyz, mxyz)
elseif matrix_source == "file"
    row_gids0 = partition_metis(A, m)
end
col_gids0 = Array{Any}(1, m)
gid2proc  = zeros(Int64, 1, n)
for I = 1:m
    col_gids0[I] = construct_cols(A, row_gids0[I])
    gid2proc[row_gids0[I]] = I
end

# ================ NOT CONVERTED STARTING THIS POINT ================
# % This is a debugging option
# % 0 : plot the residual history for selected preconditioners [default]
# % 1 : plot the connectivity structure of constructed preconditioner
# % 2 : plot the dropped connectivity structure
# global do_stencil_plot
#
# do_stencil_plot = 1
#
# %% Select preconditioners
# % Select a list of precondioners to run
# % Each run can be commented/uncommented independetly
# % WARNING: do not use split_type = 'full' with Schwarz overlap != 0
# precs = {}
# %% None
# % PROROTYPE: (A,name,{schwarz_level,schwarz_method},sub_solver,underlap_level)
# % precs = [precs {get_prec(A, 'none',                         {0},        sub_solver, 0)}]
# % precs = [precs {get_prec(A, 'subdomain Jacobi',             {0},        sub_solver, 0)}]
# % precs = [precs {get_prec(A, 'underlap + diagonal',          {0},        sub_solver, 1)}]
# % precs = [precs {get_prec(A, 'underlap + block diagonal',    {0},        sub_solver, 1)}]
# % precs = [precs {get_prec(A, 'underlap + block GS',          {0},        sub_solver, 1)}]
# % precs = [precs {get_prec(A, 'underlap + diagonal',          {0},        sub_solver, 2)}]
# % precs = [precs {get_prec(A, 'underlap + block diagonal',    {0},        sub_solver, 2)}]
# precs = [precs {get_prec(A, 'underlap + block GS',          {0},        sub_solver, 2)}]
# % precs = [precs {get_prec(A, 'l1 Gauss-Seidel',      {0},        sub_solver, {})}]
#
# %% Solve
# nprecs = size(precs,2)
# labels = cell(nprecs,1)
#
# % Reuse same colors for < 10 plots
# cc = linspecer(max(9,nprecs))
#
# for i = 1:nprecs
    # % Setup preconditioner
    # % Setting up a preconditioner is done by calling the preconditioner
    # % function with an empty array. This constructs all the necessary data for
    # % solve, and returns the preconditioner label which is going to be used in
    # % a legend
    # tic
    # label = precs{i}([])
    # setup_time = toc
#
    # if ~do_stencil_plot
        # disp(['Running: ', label])
        # fprintf('  setup time = %f\n', setup_time)
#
        # tic
        # if strcmp(solver, 'CG')
            # [~,flag,~,~,res_hist] = pcg(A, b, tol, max_its, precs{i})
        # else
            # rst    = min(restart,                size(A,1))
            # mx_its = min(fix(max_its/restart)+1, size(A,1))
            # if strcmp(solver, 'GMRES-left')
                # [~,flag,~,~,res_hist] = gmres(A, b, rst, tol, mx_its, precs{i})
            # elseif strcmp(solver, 'GMRES-right')
                # params = [tol, max_its, 1]
                # x0 = zeros(n, 1)
                # [~,flag,~,res_hist] = gmres_r(A, precs{i}, b, params, x0)
            # end
        # end
        # solve_time = toc
#
        # % Check result for convergence
        # if flag
            # fprintf('%s has not converged: flag = %d\n', solver, flag)
            # continue
        # else
            # fprintf('  solve time = %f\n', solve_time)
        # end
#
        # labels{i} = label
#
        # % Plot residual history
        # semilogy(1:size(res_hist), res_hist/norm(b), '-o', ...
            # 'Color',     cc(i,:), ...
            # 'LineWidth', 3)
        # hold on
    # end
# end
#
# if ~do_stencil_plot
    # % Construct title
    # if strcmp(matrix_source, 'generator')
        # if     size(nxyz,2) == 1, titlestr{1} = 'Laplace1D: '
        # elseif size(nxyz,2) == 2, titlestr{1} = 'Laplace2D: '
        # else                      titlestr{1} = 'Laplace3D: '
        # end
        # titlestr{1} = [titlestr{1}, 'size=[', int2str(nxyz), ']', ...
            # ', subdomains=[', int2str(mxyz), ']']
    # elseif strcmp(matrix_source, 'file')
        # titlestr{1} = [matrix_file, ': size=', int2str(size(A,1)), ...
            # ', subdomains=', int2str(m)]
    # end
    # titlestr{2} = solver
    # if strcmp(solver, 'GMRES-left')
        # titlestr{2} = [titlestr{2}, ': restart=', int2str(restart)]
    # elseif strcmp(solver, 'GMRES-right')
        # titlestr{2} = [titlestr{2}, ': no restart']
    # end
#
    # title(titlestr)
    # xlabel('Iteration number')
    # ylabel('Relative residual')
    # legend(labels)
    # % set(gca,'ygrid','on')
# end
# hold off
#
# end
#
# %% Helper functions
# % underlap() constructs level s underlap
# % It can operator in two modes:
# %   - 'full'  where level 1 underlap contains subdomain nodes within distance 1 from *any* neighbor
# %   - 'split' where 'full' underlap set is split into subsets (the article mode)
# function [sets] = underlap(A, i, s, stype)
# global m
# global gid2proc
# global row_gids0
# global col_gids0
#
# assert(s > 0)
# assert(strcmp(stype, 'full') || strcmp(stype, 'split'))
#
# % Get subdomain row and column maps
# row_gids = row_gids0{i}
# col_gids = col_gids0{i}
#
# % Construct local matrix (equivalent to the distributed variant in Trilinos)
# % It is a rectangular matrix (short and fat)
# lA = A(row_gids, col_gids)
# n  = size(lA,1)
#
# % excl set is used to exclude points which were already identified to be a
# % part of some s-set. We need it as, for instance, underlap level 2 set is
# % connected to both underlap level 1 and underlap level 3. At that point,
# % excl would contain underlap level 1, so we would get underlap level 3.
# excl = []
#
# if strcmp(stype, 'full')
    # [~,cols] = find(lA)
    # cols = setdiff(cols, 1:n)
    # % Find rows with nonzero elements in ghost columns
    # [rows,~] = find(lA(:,cols))
    # rows = unique(rows')
#
    # sets = cell(s,1)
#
    # sets{1} = rows
#
    # excl = sets{1}
    # for l = 2:s
        # if size(excl, 1) == 0
            # sets{l} = []
            # continue
        # end
#
        # % Find rows connected to the last s-set
        # [rows,~] = find(lA(:, sets{l-1}))
        # rows = unique(rows')
#
        # % Construct a new s-set
        # sets{l} = setdiff(rows, excl)
#
        # % Update excl set to include current s-set
        # excl = cat(2, excl, sets{l})
    # end
#
# else % stype = 'split'
    # % Find subdomains within distance s of the current one
    # overlap_gids = partition_schwarz(A, row_gids, s)
    # procs        = gid2proc(overlap_gids)
    # neighs       = setdiff(unique(procs), i)
#
    # %% Construct matrix D
    # % Matrix D is such that D(i,j) is the distance from node i to subdomain j
    # % 1. We consider only neighbor subdomains, the distance to non-neighbors is ignored
    # % 2. D(i,j) = 0 means infinity
    # D = zeros(n,m)
    # for k = 1:size(neighs,2)
        # neigh = neighs(k)
#
        # rows = row_gids0{neigh}
        # for l = 1:s
            # % Find "partition_schwarz(A, row_gids0{k}, l) \ partition_schwarz(A, row_gids0{k}, l-1)"
            # rows1 = partition_schwarz(A, rows, 1)
            # sub_k_over_l_gids = setdiff(rows1, rows)
            # rows = rows1
#
            # % Find local indices of the subdomain that correspond to the
            # % overlapped part of the "neighboring subdomain"
            # sub_i_under_gids  = setdiff(row_gids0{i}, setdiff(row_gids0{i}, sub_k_over_l_gids))
            # sub_i_under_lids  = []
            # for t = 1:size(sub_i_under_gids,2)
                # sub_i_under_lids = [sub_i_under_lids find(row_gids0{i} == sub_i_under_gids(t))]
            # end
#
            # % Fill in correspoding rows of D with distance l
            # D(sub_i_under_lids, k) = l
        # end
    # end
#
    # %% Construct sets
    # % sets{., 1} : contains a level subset (for varios underlap levels)
    # % sets{., 2} : contains distances from each node in the level subset to
    # %              each of the neighboring subdomains
    # sets = cell(0,2)
#
    # % S is a matrix with each row being a unique row of D
    # S = unique(D, 'rows')
    # pos = 1
    # for t = 1:size(S,1)
        # distances = S(t,:)
        # if size(find(distances),2) == 0
            # % Row with all zeros. Nodes corresponding to such rows are futher
            # % away from any neighbor than distance s, therefore we ignore them
            # continue
        # end
#
        # % Level subset = all rows of D with the same distances as the row of S
        # sets{pos,1} = find(ismember(D, distances, 'rows'))'
#
        # % We replace 0 with inf so that we have a partial sorting later
        # distances(distances == 0) = inf
        # sets{pos,2} = distances
#
        # pos = pos+1
    # end
# end
# end
#
# % Standard block Jacobi method
# function [M, row_gids] = prec_block_jacobi(A, sub_index, schwarz_level)
# global row_gids0
#
# if size(A,1) == 0
    # M = 'Block Jacobi'
    # return
# end
#
# row_gids = partition_schwarz(A, row_gids0{sub_index}, schwarz_level)
# M = A(row_gids, row_gids)
# end
#
# % Standard unpreconditioned method (do nothing)
# function [M, row_gids] = prec_none(A, sub_index, schwarz_level)
# global row_gids0
#
# if size(A,1) == 0
    # M = 'None'
    # return
# end
# assert(schwarz_level == 0)
#
# row_gids = partition_schwarz(A, row_gids0{sub_index}, schwarz_level)
# M = eye(size(row_gids,2))
# end
#
# function [M, row_gids] = prec_l1_gs(A, sub_index, schwarz_level)
# global row_gids0
#
# if size(A,1) == 0
    # M = 'l1 Gauss-Seidel'
    # return
# end
# assert(schwarz_level == 0)
#
# row_gids = partition_schwarz(A, row_gids0{sub_index}, schwarz_level)
# M = tril(A(row_gids, row_gids))
#
# end
#
# % SC14 paper method: interior + diagonalized underlap
# function [M, row_gids] = prec_underlap_diag(A, sub_index, schwarz_level, s)
# global m
# global row_gids0
#
# assert(schwarz_level <= 0, ...
    # 'schwarz_level > 0 was most likely incorrect for underlap and was removed')
#
# if size(A,1) == 0
    # M = ['Underlap (', int2str(s), ') + diagonal']
    # return
# end
# assert(s > 0)
#
# % Construct overlapped matrix
# row_gids = row_gids0{sub_index}
# M = A(row_gids, row_gids)
#
# % Construct underlap sets
# uset = underlap(A, sub_index, s, 'full')
#
# % Fix interior part
# n = size(row_gids0{sub_index},2)
# inner = 1:n
# for i = 1:s
    # inner = setdiff(inner, uset{i})
# end
# outer = setdiff(1:n, inner)
#
# M(inner, outer) = 0
# M(outer, inner) = 0
# M(outer, outer) = diag(diag(M(outer,outer)))
#
# end
#
# % New method 1: inner + block diagonal underlap
# function [M, row_gids] = prec_underlap_blockdiag(A, sub_index, schwarz_level, s)
# global m
# global row_gids0
#
# assert(schwarz_level <= 0, ...
    # 'schwarz_level > 0 was most likely incorrect for underlap and was removed')
#
# if size(A,1) == 0
    # M = ['Underlap (', int2str(s), ') + block diagonal']
    # return
# end
# assert(s > 0)
#
# % Construct overlapped matrix
# row_gids = row_gids0{sub_index}
# M = A(row_gids, row_gids)
#
# % Construct underlap sets
# uset       = underlap(A, sub_index, s, 'full')
# uset_split = underlap(A, sub_index, s, 'split')
#
# % Fix interior part
# n = size(row_gids0{sub_index},2)
# inner = 1:n
# for i = 1:s
    # inner = setdiff(inner, uset{i})
# end
#
# for i = 1:s
    # % Zero out bi-directional connections between inner and underlap
    # M(inner,   uset{i}) = 0
    # M(uset{i}, inner) = 0
    # % Zero out bi-directional connections between different underlap levels
    # for j = i+1:s
        # M(uset{i}, uset{j}) = 0
        # M(uset{j}, uset{i}) = 0
    # end
#
    # % Zero out bi-directional connections between subsets within the same underlap level
    # nsets = size(uset_split, 1)
    # for j = 1:nsets
        # set_j = uset_split{j,1}
        # for k = 1:nsets
            # if k == j
                # continue
            # end
            # set_k = uset_split{k,1}
#
            # M(set_j,set_k) = 0
            # M(set_k,set_j) = 0
        # end
    # end
# end
#
# end
#
# % New method 1: inner + block GS underlap
# function [M, row_gids] = prec_underlap_blockgs(A, sub_index, schwarz_level, s)
# global m
# global row_gids0
#
# assert(schwarz_level <= 0, ...
    # 'schwarz_level > 0 was most likely incorrect for underlap and was removed')
#
#
# if size(A,1) == 0
    # M = ['Underlap (', int2str(s), ') + block GS']
    # return
# end
# assert(s > 0)
#
# % Construct overlapped matrix
# row_gids = row_gids0{sub_index}
# M = A(row_gids, row_gids)
#
# % Construct underlap sets
# uset       = underlap(A, sub_index, s, 'full')
# uset_split = underlap(A, sub_index, s, 'split')
#
# % Fix interior part
# n = size(row_gids0{sub_index},2)
# inner = 1:n
# for i = 1:s
    # inner = setdiff(inner, uset{i})
# end
# % M(inner,inner) = tril(M(inner,inner))
#
# for i = 1:s
    # % Zero out uni-directional connections from underlap to inner
    # M(uset{i}, inner) = 0
#
    # % Zero out uni-directional connections between all underlap subsets, and between underlap sets and inner
    # nsets = size(uset_split, 1)
    # for j = 1:nsets
        # set_j = uset_split{j,1}
        # ind_j = uset_split{j,2}
        # for k = 1:nsets
            # if k == j
                # continue
            # end
            # set_k = uset_split{k,1}
            # ind_k = uset_split{k,2}
#
            # % This is a very tricky part
            # % We zero out connections based on a partial sorting algorithms based on array operator <
            # % If we can compare two distance arrays, and one is smaller than
            # % the other, then we zero out connections. An example:
            # % (1,1) is     smaller than (1,2)   => we do     zero out
            # % (1,2) is NOT smaller than (1,1)   => we do not zero out
            # % (1,2) is NOT smaller than (2,1)   => we do not zero out
            # % (1,2) is     smaller than (2,2)   => we do     zero out
            # % (1,2) is     smaller than (2,inf) => we do     zero out
            # % etc
            # [cmp1, cmp2] = compare(ind_j, ind_k)
            # if cmp1 || ...          % ind_j <= ind_k
               # (~cmp1 && ~cmp2)     % ind_j and ind_k are incomparable (in the partial sorting sense)
                # M(set_j,set_k) = 0
            # end
        # end
    # end
# end
#
# end
#
# function [x] = gen_prec(A, schwarz_level, schwarz_method, solver, P, b)
# global ilu_setup
# global m
# global coords
# global do_stencil_plot
# global row_gids0
#
# assert(strcmp(schwarz_method, 'RAS') || strcmp(schwarz_method, 'AS') || strcmp(schwarz_method, 'AAS'))
# assert(strcmp(solver, 'direct') || strcmp(solver, 'ilu'))
#
# persistent M
# persistent L
# persistent U
# persistent row_gids
#
# if size(b,1) == 0
    # % Setup stage
    # row_gids = cell(1,m)
#
    # colors  = ['k', 'b', 'g', 'r', 'm', 'c']
    # ncolors = size(colors,2)
#
    # for I = 1:m
        # tic
        # [M{I}, row_gids{I}] = P(A, I, schwarz_level)
        # if do_stencil_plot
            # hold on
            # lcoords = coords(row_gids{I},:)
            # if do_stencil_plot == 1
                # % Show remaining connections
                # X = M{I}
            # else
                # % Show dropped connections
                # X = A(row_gids{I},row_gids{I}) - M{I}
            # end
            # color = colors(rem(I,ncolors)+1)
            # plot_stencil(X, lcoords, color)
            # % To look at the individual subdomains, uncomment the following lines
            # % Then, step through subdomains in command line with dbcont
            # if 0
                # keyboard
                # if I ~= m
                    # clf
                # end
            # end
        # end
#
        # if strcmp(solver, 'direct')
            # [L{I}, U{I}] = lu(M{I})
        # else
            # % Current approach is not very good, as it does
            # % ILU on the overlapped matrix
            # % What we would like to do is to solve "tails"
            # % exactly and then perform ILU only on the proper
            # % inner part (underlap s)
            # [L{I}, U{I}] = ilu(M{I}, ilu_setup)
        # end
    # end
#
    # % Construct label by attaching Schwarz to the label of P
    # x = P([], -1, -1)
    # if schwarz_level
        # x = [x, ', schwarz=', schwarz_method, ...
            # '(', int2str(schwarz_level), ')']
    # end
    # x = [x, ', solver=', solver]
#
# else
    # % Solve stage
    # n = size(b, 1)
    # x         = zeros(n, 1)
    # indicator = zeros(n, 1)
#
    # for I = 1:m
        # lb = b(row_gids{I})
#
        # lx = U{I} \ (L{I} \ lb)
#
        # if strcmp(schwarz_method, 'RAS')
            # % RAS (Restrictive Additive Schwarz)
            # % Assign only owned DOFs
            # x(row_gids0{I}) = lx(1:size(row_gids0{I},2))
#
        # elseif strcmp(schwarz_method, 'AS') || strcmp(schwarz_method, 'AAS')
            # % AS (Additive Schwarz)
            # % Contributions from multiple subdomains are summed
            # x(row_gids{I}) = x(row_gids{I}) + lx
            # indicator(row_gids{I}) = indicator(row_gids{I})+1
        # end
    # end
    # if strcmp(schwarz_method, 'AAS')
        # % AAS (Averaged Additive Schwarz)
        # % Contributions from multiple subdomains to the same dof are averaged
        # % by the number of contributing subdomain
        # x = x ./ indicator
    # end
# end
# end
#
# function [prec] = get_prec(A, name, schwarz_cell, sub_solver, underlap_level)
# if size(schwarz_cell,2) == 1
    # if schwarz_cell{1} == 0
        # schwarz_cell = [schwarz_cell{1} {'AS'}]
    # else
        # error('Please specify RAS/AS for overlap level > 0')
    # end
# end
#
# if     strcmp(name, 'none')
    # prec = @(b)(gen_prec(A, ...
        # schwarz_cell{1}, schwarz_cell{2}, sub_solver, ...
        # @(X,i,l)(prec_none(X,i,l)), b))
#
# elseif strcmp(name, 'subdomain Jacobi')
    # prec = @(b)(gen_prec(A, ...
        # schwarz_cell{1}, schwarz_cell{2}, sub_solver, ...
        # @(X,i,l)(prec_block_jacobi(X,i,l)), b))
#
# elseif strcmp(name, 'l1 Gauss-Seidel')
    # prec = @(b)(gen_prec(A, ...
        # schwarz_cell{1}, schwarz_cell{2}, sub_solver, ...
        # @(X,i,l)(prec_l1_gs(X,i,l)), b))
#
# elseif strcmp(name, 'underlap + diagonal')
    # prec = @(b)(gen_prec(A, ...
        # schwarz_cell{1}, schwarz_cell{2}, sub_solver, ...
        # @(X,i,l)(prec_underlap_diag(X, i, l, underlap_level)), b))
#
# elseif strcmp(name, 'underlap + block diagonal')
    # prec = @(b)(gen_prec(A, ...
        # schwarz_cell{1}, schwarz_cell{2}, sub_solver, ...
        # @(X,i,l)(prec_underlap_blockdiag(X, i, l, ...
        # underlap_level)), b))
#
# elseif strcmp(name, 'underlap + block GS')
    # prec = @(b)(gen_prec(A, ...
        # schwarz_cell{1}, schwarz_cell{2}, sub_solver, ...
        # @(X,i,l)(prec_underlap_blockgs(X, i, l, ...
        # underlap_level)), b))
#
# else
    # error(['Unknown preconditioner: ', name])
# end
# end

max_stars_repo_name: aprokop/ca-preconditioners

max_issues_repo_name: aprokop/ca-preconditioners

max_forks_repo_name: aprokop/ca-preconditioners


text:
using SparseArrays

function numerical_fluxes(flux, mesh, w, scheme, dt=0.0)
    map(i_face -> numerical_flux(flux, mesh, w, scheme, i_face, dt), collect(inner_faces(mesh)))
end

"""Sparse matrix to sum the flux on both side of a face and get the change in a cell.

In 1D: sparse matrix.
For 2D cartesian mesh, it should be a higher order mesh, but since there is no higher dimensional sparse array, the field needs to be reshaped to a vector...
"""
function inner_faces_to_cells_matrix(mesh::AbstractCartesianMesh)
    D = spzeros(FiniteVolumes.nb_cells(mesh), FiniteVolumes.nb_inner_faces(mesh))
    li = LinearIndices(cell_centers(mesh))
    for (i, face) in enumerate(FiniteVolumes.inner_faces(mesh))
        i_cell_1, i_cell_2 = cells_next_to_inner_face(mesh, face)
        D[li[i_cell_1], i] -= face_area(mesh, face) / cell_volume(mesh, i_cell_1)
        D[li[i_cell_2], i] += face_area(mesh, face) / cell_volume(mesh, i_cell_1)
    end
    D
end

function update!(w, Φ, mesh, dt=0.0)
    for (i, face) in enumerate(inner_faces(mesh))
        i_cell_1, i_cell_2 = cells_next_to_inner_face(mesh, face)
        w[i_cell_1] -= dt*Φ[i] * face_area(mesh, face) / cell_volume(mesh, i_cell_1)
        w[i_cell_2] += dt*Φ[i] * face_area(mesh, face) / cell_volume(mesh, i_cell_2)
    end
end


text:
module PkgMPC

# Write your package code here.
using Reexport
@reexport using Plots
@reexport using CSV
@reexport using JuMP
@reexport using Ipopt
@reexport using DataFrames
@reexport using DifferentialEquations
@reexport using PlotlyJS

include("Extra.jl")
include("main.jl")
include("NMPC.jl")
include("Plant.jl")

export F, Model_Par, Simulation_Data, Simulate_Plant
export Simulate_Plant, Plot_Plant, Save_Plant, Export_Plant

end


text:
function view(::HadamardGate, i::Vector{Int64})
    ["—[H ]—"], i
end

function view(::XGate, i::Vector{Int64})
    ["—[X ]—"], i
end

function view(::YGate, i::Vector{Int64})
    ["—[Y ]—"], i
end

function view(::ZGate, i::Vector{Int64})
    ["—[Z ]—"], i
end

function view(::RxGate, i::Vector{Int64})
    ["—[Rx]—"], i
end

function view(::RyGate, i::Vector{Int64})
    ["—[Ry]—"], i
end

function view(::RzGate, i::Vector{Int64})
    ["—[Rz]—"], i
end

function view(::RotationGate, i::Vector{Int64})
    ["—[Rθ]—"], i
end

function view(::PhaseShiftGate, i::Vector{Int64})
    ["—[Pϕ]—"], i
end

function view(::SGate, i::Vector{Int64})
    ["—[S ]—"], i
end

function view(::TGate, i::Vector{Int64})
    ["—[T ]—"], i
end

function view(::SdagGate, i::Vector{Int64})
    ["—[S†]—"], i
end

function view(::TdagGate, i::Vector{Int64})
    ["—[T†]—"], i
end

function view(::SwapGate, i::Vector{Int64})
    ["——x———", "——x———"], i
end

view(g) = view(g, [1])

function view(g::ControlledGate, i::Vector{Int64})
    gate = fill("——•———", length(i))
    gate[1:num_wires(g.U)] = view(g.U)[1]
    return gate, i
end

function view(::AbstractGate, i::Vector{Int64})
    fill("——□———", length(i)), i
end

function interleave(a::Vector{T}, b::Vector{T}) where {T}
    length(a)-1 == length(b) || error("Vector b must be one element shorter than Vector a")
    c = Vector{T}(undef, length(a) + length(b))
    c[1] = a[1]
    i = 1
    for (x, y) in zip(b, a[2:end])
        c[i += 1] = x
        c[i += 1] = y
    end
    return c
end

function updatecol(g::CircuitGate{M,G}, nw::Vector{String}, ng::Vector{String}) where {M,G}
    iwire = [x for x in g.iwire]
    if length(g.iwire) > 1
        index = min(iwire...):max(iwire...)-1
        ng[index] .= "  |   "
    end
    gate, index = view(g.gate, iwire)
    nw[index] .= gate
    return nw, ng
end

function wire_enum(N::Int)
    wires = lpad.(1:N, 5, " ")
    wires = rpad.(wires, 6, " ")
    gaps = fill("      ", N-1)
    return wires, gaps
end

function momentdiagram(io::IO, m::Moment, N::Int)
    if N == 0
        return
    end
    w, g = wire_enum(N)
    i = Int[]
    w = w .* "...|"
    g = g .* "   |"
    nw = fill("——————", N)
    ng = fill("      ", N-1)
    for gate in m
        r = min(gate.iwire...):max(gate.iwire...)
        if length(intersect(i, r)) == 0
            i = union(i, r)
            nw, ng = updatecol(gate, nw, ng)
        else
            i = r
            w = w .* nw
            g = g .* ng
            nw = fill("——————", N)
            ng = fill("      ", N-1)
            nw, ng = updatecol(gate, nw, ng)
        end
    end
    w = w .* nw .* "|..."
    g = g .* ng .* "|   "
    circuit = interleave(w, g)
    println(io, "")
    for c in reverse(circuit)
        println(io, c)
    end
    println(io, "")
end

function Base.show(io::IO, m::Vector{Moment})
    if length(m) == 0
        return
    end
    N = maximum(req_wires.(m))
    for moment in m
        momentdiagram(io, moment, N)
    end
end

function Base.show(io::IO, m::Moment)
    for gate in m
        println(io, gate)
    end
end

function Base.show(io::IO, cgs::Vector{<:CircuitGate}, N::Union{Nothing,Int}=nothing)
    if isempty(cgs)
        println(io, "[]")
        return
    end
    if isnothing(N)
        N = maximum(req_wires.(cgs))
    end
    w, g = wire_enum(N)
    i = Int[]
    nw = fill("——————", N)
    ng = fill("      ", N-1)
    for cg in cgs
        r = min(cg.iwire...):max(cg.iwire...)
        if length(intersect(i, r)) == 0
            i = union(i, r)
            nw, ng = updatecol(cg, nw, ng)
        else
            i = r
            w = w .* nw
            g = g .* ng
            nw = fill("——————", N)
            ng = fill("      ", N-1)
            nw, ng = updatecol(cg, nw, ng)
        end
    end
    w = w .* nw
    g = g .* ng
    circuit = interleave(w, g)
    println(io, "")
    for c in reverse(circuit)
        println(io, c)
    end
end

function _show(io::IO, c::Vector{Moment}, N::Int)
    cg = getproperty.(c, :gates)
    cg = collect(Base.Iterators.flatten(cg))
    show(io::IO, cg, N)
end

"""
    Base.show(io::IO, c::Circuit) = Base.show(io, c.cgc)

extends base `show` method to draw visual representation of a `Circuit{N}` object.
"""
Base.show(io::IO, c::Circuit{N}) where {N} = _show(io, c.moments, N)


text:
using DigitSetSudoku
using Base.Test

const spec =
"""
400000805
030000000
000700000
020000060
000080400
000010000
000603070
500200000
104000000
"""

puzzle = SudokuPuzzle(spec)

@test (@sprintf "%s" puzzle) ==
"""
4 . . | . . . | 8 . 5
. 3 . | . . . | . . .
. . . | 7 . . | . . .
------+-------+------
. 2 . | . . . | . 6 .
. . . | . 8 . | 4 . .
. . . | . 1 . | . . .
------+-------+------
. . . | 6 . 3 | . 7 .
5 . . | 2 . . | . . .
1 . 4 | . . . | . . .
"""

solution = solve(puzzle)

@test (@sprintf "%s" solution) ==
"""
4 1 7 | 3 6 9 | 8 2 5
6 3 2 | 1 5 8 | 9 4 7
9 5 8 | 7 2 4 | 3 1 6
------+-------+------
8 2 5 | 4 3 7 | 1 6 9
7 9 1 | 5 8 6 | 4 3 2
3 4 6 | 9 1 2 | 7 5 8
------+-------+------
2 8 9 | 6 4 3 | 5 7 1
5 7 3 | 2 9 1 | 6 8 4
1 6 4 | 8 7 5 | 2 9 3
"""


text:
# Cartesian product. See https://en.wikipedia.org/wiki/Cartesian_product

# next_tuple is a helper function which extracts all combinations of arrays by
# using the length of each array Yn and counting up from 0 at each index. For
# example, if we had an array X of lengths [2, 3, 3] next_tuple will adjust indices
# as follows: [0, 0, 0] -> [1, 0, 0] -> [0, 1, 0] -> [1, 1, 0] -> [0, 2, 0] ... [1, 2, 2]

function next_tuple!(indices::Array{Int,1}, sizes::Array{Int,1})::Bool
    for i in length(indices):-1:1
        if indices[i] < sizes[i]
            indices[i] += 1
            return true
        else
            indices[i] = 1
        end
    end
    false
end

# N fold Cartesian product
function cartprod(aa::Array{Array{T, 1}, 1})::Array{Array{T, 1}, 1} where {T}
    result::Array{Array{T, 1}, 1} = []
    indices = ones(Int, length(aa))
    sizes = map(a -> length(a), aa)
    hasMore = length(indices) > 0
    while hasMore
        elt = [aa[i][indices[i]] for i in 1:length(indices)]
        push!(result, elt)
        hasMore = next_tuple!(indices, sizes)
    end

  return result
end

function flatmap(f, a)
    result = []
    for elt in a
        push!(result, f(elt)...)
    end
    result
end

isVar(::Var) = true
isVar(::BoolExpr) = false
isAnd(::And) = true
isAnd(::BoolExpr) = false

children(e::Var)::BoolExprVec = [e]
children(e::Not)::BoolExprVec = [e.expr]
children(e::And)::BoolExprVec = e.exprs
children(e::Or)::BoolExprVec = e.exprs

function children(v::BoolExprVec)
    result = []
    for e in v
        for o in children(e)
            push!(result, o)
        end
    end
    result
end

negate(exprs::BoolExprVec) = map(e -> Not(e), exprs)


text:
# Autogenerated wrapper script for Pangomm_jll for x86_64-linux-gnu-cxx11
export libpangomm

using Cairomm_jll
using Glibmm_jll
using Pango_jll
JLLWrappers.@generate_wrapper_header("Pangomm")
JLLWrappers.@declare_library_product(libpangomm, "libpangomm-2.48.so.1")
function __init__()
    JLLWrappers.@generate_init_header(Cairomm_jll, Glibmm_jll, Pango_jll)
    JLLWrappers.@init_library_product(
        libpangomm,
        "lib/libpangomm-2.48.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
# # Linking structural and evolutionary information
#
#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__cookbook/notebooks/02_Linking_structural_and_evolutionary_information.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__cookbook/notebooks/02_Linking_structural_and_evolutionary_information.ipynb)
#
#
# ## Problem description
#
# It is a very common task to map sequence to structure residue number. For
# example, to link structural information coming from *PDB* and evolutionary
# information calculated from multiple sequence alignments.
#
# The naive way of mapping sequence and structure is to perform global pairwise
# alignment between the sequence and the *PDB* sequence (using the residues in
# *ATOM*). The problem with this approach is that the sequences can have
# missing regions and standard pairwise alignment algorithms often yield
# incorrect assignations around those regions
# [(Velankar et.al. 2013)](https://doi.org/10.1093/nar/gks1258). This is
# particularly important when aligning *PDB* sequences, that can have missing
# residues, and sequences coming from multiple sequence alignments, that can be
# incomplete or have unaligned regions (e.g. insert states).
#
# The [SIFTS](https://www.ebi.ac.uk/pdbe/docs/sifts/index.html)
# (Structure Integration with Function, Taxonomy and Sequences) database solves
# this problem and provides residue level mapping between PDB and other
# databases (e.g. *UniProt* and *Pfam*).
#
# The `SIFTS` module of MIToS has functions to access this residue level mapping
# between *PDB* and other databases. Also, MIToS keeps track of the residue
# number of each residue in a multiple sequence alignment (MSA) using its annotations.
# Both things together, allow the correct mapping of sequence and structure
# without performing error-prone pairwise alignments.
#
# Particular solutions depend on problem details, here we show some common ways
# to use MIToS and *SIFTS* to map evolutionary information calculated in an MSA
# (e.g. entropy) with structural information (e.g. B-factors).
#
# ## PDB and Pfam alignment mapping
#
# This is the easiest problem to solve with the
# [MIToS `Pfam` module](@ref Module-Pfam) because *SIFTS* already has a residue
# level mapping between *PDB* and *Pfam*.
#
# For this example, we are going to map the columns in the multiple sequence
# alignment of the *PF09645* Pfam family and the residues in the chain *A* from
# the *2VQC* *PDB* file. The needed files are available in the MIToS test suite:

using MIToS
pdb_file   = abspath(pathof(MIToS), "..", "..", "test", "data", "2VQC.pdb")
pfam_file  = abspath(pathof(MIToS), "..", "..", "test", "data", "PF09645_full.stockholm")
sifts_file = abspath(pathof(MIToS), "..", "..", "test", "data", "2vqc.xml.gz")
#md nothing # hide

# You can also use `downloadpdb` from `MIToS.PDB`, `downloadpfam` from
# `MIToS.Pfam` and `downloadsifts` from `MIToS.SIFTS` to get the corresponding
# files from those databases.
#
# It is important to read the Pfam MSA file using `generatemapping=true` and
# `useidcoordinates=true` because that allows keeping track of the residue
# number using the MSA annotations.

using MIToS.Pfam
msa = read(pfam_file, Stockholm, generatemapping=true, useidcoordinates=true)

# First, we need to know what is the sequence in the MSA that correspond to the
# PDB we want to link. Luckily, Pfam Stockholm files store the mapping between
# sequences and PDB chains. You can access that mapping using the `getseq2pdb`
# function from `MIToS.Pfam`

seq2pdbs = getseq2pdb(msa)

# The returned dictionary gives you all the PDB chains associated with a
# determined sequence in the MSA. But, in this case, we want to go in the other
# direction to find all the sequences associated with a determined *PDB* chain.
# We are going to use a list comprehension because it is possible for a single
# chain to be associated with more than one sequence in the *Pfam* MSA (e.g.
# domain repeats).

pdb_code  = "2VQC"
pdb_chain = "A"
seq_ids = [ seq for (seq, pdbs) in seq2pdbs if (pdb_code, pdb_chain) in pdbs ]

# In this example, we are going to use the only sequence we found for the *A*
# of *2VQC*.

seq_id = seq_ids[1]

# Finally, we can use the `msacolumn2pdbresidue` function from the Pfam module
# to get a dictionary from the MSA column index to the *PDB* residue number:

pfam_id = "PF09645"
msacol2pdbres = msacolumn2pdbresidue(msa, seq_id, pdb_code, pdb_chain, pfam_id, sifts_file)

# This dictionary has the mapping between MSA column and PDB residue that allows
# the mapping between evolutionary and structural information. For example, to
# measure the correlation between entropy (related to residue variation in an
# MSA column) and the mean B factor of the residue:

using MIToS.Information
Hx = mapcolfreq!(entropy,
				 msa,
				 Counts(ContingencyTable(Int, Val{1}, UngappedAlphabet())))

# To get quick access to each PDB residue based on its residue number, we can
# read the PDB file into a dictionary using the `read` and `residuesdict`
# functions from the MIToS `PDB` module:

using MIToS.PDB
res_dict = residuesdict(read(pdb_file, PDBFile, occupancyfilter=true), "1", "A") # model 1 chain A

# Then, we can iterate the mapping dictionary to link the MSA and PDB based
# values:

using Statistics

x = Float64[]
y = Float64[]

for (col_index, res_number) in msacol2pdbres
	if res_number != "" # i.e. MSA column has an associated PDB residue
		push!(x, Hx[col_index])
		push!(y, mean(parse(Float64, atom.B) for atom in res_dict[res_number].atoms))
	end
end

cor(x, y)

# ##  Unknown sequence coordinates
#
# While *Pfam* alignments have the start and end of the aligned region indicated
# in the sequence name, other multiple sequence alignments don't give any hint
# about that. In those cases, we should use pairwise alignments. However,
# instead of aligning the sequence coming from the MSA and the *PDB* sequence,
# we can align the MSA sequence to the *UniProt* sequence to reduce the
# possibility of mapping errors. Once we have the mapping of the MSA sequence
# to the *UniProt* sequence, we can use *SIFTS* to map the *PDB* sequence to
# the MSA sequence using the *UniProt* numeration.
#
# For this example, we are going to use the following files included in MIToS
# documentation:

using MIToS
pdb_file   = abspath(pathof(MIToS), "..", "..", "docs", "data", "1dur.pdb")
msa_file   = abspath(pathof(MIToS), "..", "..", "docs", "data", "blast_alignment.fa")
sifts_file = abspath(pathof(MIToS), "..", "..", "docs", "data", "1dur.xml.gz")
uniprot_file = abspath(pathof(MIToS), "..", "..", "docs", "data", "P00193.fasta")
#md nothing # hide

# First, we are going to read the MSA file. In this case, we can not use
# `useidcoordinates=true` because the sequence names don't have the sequence
# coordinates in the Pfam format. However, we are going to use
# `generatemapping=true` to get the default mapping for each sequence in the
# alignment (from `1` to the length of the aligned region):

using MIToS.MSA
msa = read(msa_file, FASTA, generatemapping=true)

# After that, we get the first sequence of the MSA, the one we know that
# corresponds to the PDB of interest. We need the sequence as a `String`
# without gaps (unaligned), so we use the `MIToS.MSA` `stringsequence` function
# together with `replace`:

msa_seq = replace(stringsequence(msa, 1), '-' => "")

# Also, we are going to read the *UniProt* sequence. You can easily download the
# sequence from UniProt by doing:
# ```julia
# using MIToS.Utils
# download_file("https://www.uniprot.org/uniprot/P00193.fasta", "P00193.fasta")
# ```
# To read the FASTA file we are going to use the `FastaIO` package:

using FastaIO
uniprot_sequences = readfasta(uniprot_file)

# And get the unique sequence:

uniprot_seq = uniprot_sequences[1][2]

# We can perform a pairwise sequence alignment between both sequences by using
# the [`BioAlignments` package](https://github.com/BioJulia/BioAlignments.jl)
# from the *BioJulia* suite. In this case, we use a semi-global alignment
# (no start/end gap penalty) because we know that the MSA sequence is a region
# of the *UniProt* sequence.

using BioAlignments
costmodel = AffineGapScoreModel(BLOSUM62, gap_open=-10, gap_extend=-1)
aln = pairalign(SemiGlobalAlignment(), msa_seq, uniprot_seq, costmodel)

# Then, we only need to iterate the alignment to designate the positions and
# store the equivalences in a dictionary:

function seq2refnumber(aln)
    seq_pos = 0
	ref_pos = 0
	last_seq_pos = 0
	seq2ref = Dict{Int,Int}()
    for (seq_res, ref_res) in alignment(aln)
        if seq_res != '-'
            seq_pos += 1
		end
        if ref_res != '-'
            ref_pos += 1
		end
		if seq_pos != last_seq_pos
			seq2ref[seq_pos] = ref_pos
			last_seq_pos = seq_pos
    	end
	end
    seq2ref
end

seqnum2uniprotnum = seq2refnumber(aln)

# Then, we can use `getsequencemapping` to go from MSA column number to
# *UniProt* residue, and `siftsmapping` to go from *UniProt* to *PDB*:

seqmap = getsequencemapping(msa, 1)

# -

colnum2uniprotnum = Dict{Int,Int}()
for (colnum, seqnum) in enumerate(seqmap)
	if seqnum != 0 # getsequencemapping returns 0 where there is a gap
		colnum2uniprotnum[colnum] = seqnum2uniprotnum[seqnum]
	end
end
colnum2uniprotnum

# -

using MIToS.SIFTS

uniprotnum2pdbnum = siftsmapping(sifts_file,
    dbUniProt,
    "P00193",
    dbPDB,
    "1dur", # SIFTS stores PDB identifiers in lowercase
    chain="A",
    missings=false) # residues without coordinates aren't used in the mapping

# To finally get the dictionary from MSA column index to PDB residue number

colnum2pdbnum = Dict{Int,String}()
for (colnum, uniprotnum) in colnum2uniprotnum
	pdbresnum = get(uniprotnum2pdbnum, string(uniprotnum), "")
	if pdbresnum != ""
		colnum2pdbnum[colnum] = pdbresnum
	end
end

colnum2pdbnum


text:
# Note ReshapeLayer is NOT a UtilLayer because it computes backward (via some upper layers)
@defstruct ReshapeLayer Layer (
  name :: AbstractString = "reshape",
  (tops :: Vector{Symbol} = [], length(tops) > 0),
  (bottoms :: Vector{Symbol} = [], length(bottoms) == length(tops)),
  (shape :: NTuple = (), eltype(shape) == Int && all(broadcast(>, collect(shape),  0))),
)
@characterize_layer(ReshapeLayer,
  can_do_bp => true, # back-propagate via upper layers
)

type ReshapeLayerState <: LayerState
  layer      :: ReshapeLayer
  blobs      :: Vector{Blob}
  blobs_diff :: Vector{Blob}
end

function setup(backend::Backend, layer::ReshapeLayer, inputs::Vector{Blob}, diffs::Vector{Blob})
  blobs = map(inputs) do blob
    reshape_blob(backend, blob, layer.shape..., get_num(blob))
  end
  blobs_diff = map(diffs) do blob
    if isa(blob, NullBlob)
      NullBlob()
    else
      reshape_blob(backend, blob, layer.shape..., get_num(blob))
    end
  end

  return ReshapeLayerState(layer, blobs, blobs_diff)
end

function forward(backend::Backend, state::ReshapeLayerState, inputs::Vector{Blob})
end
function backward(backend::Backend, state::ReshapeLayerState, inputs::Vector{Blob}, diffs::Vector{Blob})
end
function shutdown(backend::Backend, state::ReshapeLayerState)
end


text:
# prior_baseline.jl
# David Merrell
#
# A simple script for using the prior knowledge
# as though it were a predictor.
#

module PriorPredictor

using JSON
using DataFrames
using CSV
using ArgParse

export julia_main

function load_args(args::Vector{String})

    s = ArgParseSettings()
    @add_arg_table! s begin
        "prior_graph"
            help="path to a reference graph file"
            required = true
            arg_type = String
        "output_file"
            help="Name of output file containing edge probabilities"
            required = true
            arg_type = String
    end

    args = parse_args(args, s)
    return args
end

function make_output(argd::Dict)
    infile = argd["prior_graph"]
    outfile = argd["output_file"]

    ref_adj_df = CSV.read(infile, header=false)
    ref_ps = Vector{Vector{Float64}}()
    adj_mat = convert(Matrix, ref_adj_df)

    for i=1:size(adj_mat,1)
        push!(ref_ps, adj_mat[:,i])
    end

    results = Dict()
    results["edge_conf_key"] = "parent_sets"
    results["parent_sets"] = ref_ps

    #res_json = JSON.json(results)
    f = open(outfile, "w")
    #write(f, res_json)
    JSON.print(f, results)
    close(f)
end

function julia_main()
    argd = load_args(ARGS)
    make_output(argd)
    return 0
end

end

using .PriorPredictor

julia_main()

max_stars_repo_name: gitter-lab/ssps

max_issues_repo_name: gitter-lab/ssps

max_forks_repo_name: gitter-lab/ssps


text:
# SPDX-FileCopyrightText: 2021 Johannes Hampp <johannes.hampp@zeu.jlug.de>
#
# SPDX-License-Identifier: MIT

# Activate julia environment
using Pkg
Pkg.activate("./workflow/envs")
using GlobalEnergyGIS

era_year = parse(Int64, snakemake.wildcards["era_year"])

println("Downloading ERA5 data.")
download_and_convert_era5(era_year)

max_stars_repo_name: euronion/synde

max_issues_repo_name: euronion/synde

max_forks_repo_name: euronion/synde


text:
########
# Dual #
########

struct Dual{T,V<:Real,N} <: Real
    value::V
    partials::Partials{N,V}
end

################
# Constructors #
################

@inline (::Type{Dual{T}})(value::V, partials::Partials{N,V}) where {T,N,V} = Dual{T,V,N}(value, partials)

@inline function (::Type{Dual{T}})(value::A, partials::Partials{N,B}) where {T,N,A,B}
    C = promote_type(A, B)
    return Dual{T}(convert(C, value), convert(Partials{N,C}, partials))
end

@inline (::Type{Dual{T}})(value::Real, partials::Tuple) where {T} = Dual{T}(value, Partials(partials))
@inline (::Type{Dual{T}})(value::Real, partials::Tuple{}) where {T} = Dual{T}(value, Partials{0,typeof(value)}(partials))
@inline (::Type{Dual{T}})(value::Real, partials::Real...) where {T} = Dual{T}(value, partials)
@inline (::Type{Dual{T}})(value::V, ::Chunk{N}, p::Val{i}) where {T,V<:Real,N,i} = Dual{T}(value, single_seed(Partials{N,V}, p))

@inline Dual(args...) = Dual{Void}(args...)

####################
# TagMismatchError #
####################

struct TagMismatchError{X,Y} <: Exception
    x::Dual{X}
    y::Dual{Y}
end

function TagMismatchError(x, y, z)
    if isa(x, Dual) && isa(y, Dual) && tagtype(x) !== tagtype(y)
        return TagMismatchError(x, y)
    elseif isa(x, Dual) && isa(z, Dual) && tagtype(x) !== tagtype(z)
        return TagMismatchError(x, z)
    elseif isa(y, Dual) && isa(z, Dual) && tagtype(y) !== tagtype(z)
        return TagMismatchError(y, z)
    else
        error("the provided arguments have matching tags, or are not Duals")
    end
end

function Base.showerror(io::IO, e::TagMismatchError{X,Y}) where {X,Y}
    print(io, "potential perturbation confusion detected when computing binary operation ",
              "on $(e.x) and $(e.y) (tag mismatch: $X != $Y). ForwardDiff cannot safely ",
              "perform differentiation in this context; see the following issue for ",
              "details: https://github.com/JuliaDiff/jl/issues/83")
end

##############################
# Utility/Accessor Functions #
##############################

@inline value(x::Real) = x
@inline value(d::Dual) = d.value

@inline partials(x::Real) = Partials{0,typeof(x)}(tuple())
@inline partials(d::Dual) = d.partials
@inline partials(x::Real, i...) = zero(x)
@inline partials(d::Dual, i) = d.partials[i]
@inline partials(d::Dual, i, j) = partials(d, i).partials[j]
@inline partials(d::Dual, i, j, k...) = partials(partials(d, i, j), k...)

@inline npartials(::Dual{T,V,N}) where {T,V,N} = N
@inline npartials(::Type{Dual{T,V,N}}) where {T,V,N} = N

@inline order(::Type{V}) where {V} = 0
@inline order(::Type{Dual{T,V,N}}) where {T,V,N} = 1 + order(V)

@inline valtype(::V) where {V} = V
@inline valtype(::Type{V}) where {V} = V
@inline valtype(::Dual{T,V,N}) where {T,V,N} = V
@inline valtype(::Type{Dual{T,V,N}}) where {T,V,N} = V

@inline tagtype(::V) where {V} = Void
@inline tagtype(::Type{V}) where {V} = Void
@inline tagtype(::Dual{T,V,N}) where {T,V,N} = T
@inline tagtype(::Type{Dual{T,V,N}}) where {T,V,N} = T

####################################
# N-ary Operation Definition Tools #
####################################

macro define_binary_dual_op(f, xy_body, x_body, y_body)
    defs = quote
        @inline $(f)(x::Dual, y::Dual) = throw(TagMismatchError(x, y))
        @inline $(f)(x::Dual{T}, y::Dual{T}) where {T} = $xy_body
        @inline $(f)(x::Dual{T,Dual{S,X,N},M}, y::Dual{T,Dual{S,Y,N},M}) where {T,S,X<:Real,Y<:Real,N,M} = $xy_body
        @inline $(f)(x::Dual{T,Dual{S,X,N},M}, y::Dual{S,Y,N})           where {T,S,X<:Real,Y<:Real,N,M} = $x_body
        @inline $(f)(x::Dual{S,X,N},           y::Dual{T,Dual{S,Y,N},M}) where {T,S,X<:Real,Y<:Real,N,M} = $y_body
    end
    for R in REAL_TYPES
        expr = quote
            @inline $(f)(x::Dual{T}, y::$R) where {T} = $x_body
            @inline $(f)(x::$R, y::Dual{T}) where {T} = $y_body
        end
        append!(defs.args, expr.args)
    end
    return esc(defs)
end

macro define_ternary_dual_op(f, xyz_body, xy_body, xz_body, yz_body, x_body, y_body, z_body)
    defs = quote
        @inline $(f)(x::Dual, y::Dual, z::Dual) = throw(TagMismatchError(x, y, z))
        @inline $(f)(x::Dual{T}, y::Dual{T}, z::Dual{T}) where {T} = $xyz_body
        @inline $(f)(x::Dual{T,Dual{S,X,N},M}, y::Dual{T,Dual{S,Y,N},M}, z::Dual{T,Dual{S,Z,N},M}) where {T,S,X<:Real,Y<:Real,Z<:Real,N,M} = $xyz_body
        @inline $(f)(x::Dual{T,Dual{S,X,N},M}, y::Dual{T,Dual{S,Y,N},M}, z::Dual{S,Z,N})           where {T,S,X<:Real,Y<:Real,Z<:Real,N,M} = $xy_body
        @inline $(f)(x::Dual{T,Dual{S,X,N},M}, y::Dual{S,Y,N},           z::Dual{T,Dual{S,Z,N},M}) where {T,S,X<:Real,Y<:Real,Z<:Real,N,M} = $xz_body
        @inline $(f)(x::Dual{S,X,N},           y::Dual{T,Dual{S,Y,N},M}, z::Dual{T,Dual{S,Z,N},M}) where {T,S,X<:Real,Y<:Real,Z<:Real,N,M} = $yz_body
        @inline $(f)(x::Dual{T,Dual{S,X,N},M}, y::Dual{S,Y,N},           z::Dual{S,Z,N})           where {T,S,X<:Real,Y<:Real,Z<:Real,N,M} = $x_body
        @inline $(f)(x::Dual{S,X,N},           y::Dual{T,Dual{S,Y,N},M}, z::Dual{S,Z,N})           where {T,S,X<:Real,Y<:Real,Z<:Real,N,M} = $y_body
        @inline $(f)(x::Dual{S,X,N},           y::Dual{S,Y,N},           z::Dual{T,Dual{S,Z,N},M}) where {T,S,X<:Real,Y<:Real,Z<:Real,N,M} = $z_body
    end
    for R in REAL_TYPES
        expr = quote
            @inline $(f)(x::Dual, y::Dual, z::$R) = throw(TagMismatchError(x, y, z))
            @inline $(f)(x::Dual, y::$R, z::Dual) = throw(TagMismatchError(x, y, z))
            @inline $(f)(x::$R, y::Dual, z::Dual) = throw(TagMismatchError(x, y, z))

            @inline $(f)(x::Dual{T}, y::Dual{T}, z::$R) where {T} = $xy_body
            @inline $(f)(x::Dual{T}, y::$R, z::Dual{T}) where {T} = $xz_body
            @inline $(f)(x::$R, y::Dual{T}, z::Dual{T}) where {T} = $yz_body

            @inline $(f)(x::Dual{T,Dual{S,X,N},M}, y::Dual{T,Dual{S,Y,N},M}, z::$R)                    where {T,S,X<:Real,Y<:Real,N,M} = $xy_body
            @inline $(f)(x::Dual{T,Dual{S,X,N},M}, y::$R,                    z::Dual{T,Dual{S,Z,N},M}) where {T,S,X<:Real,Z<:Real,N,M} = $xz_body
            @inline $(f)(x::$R,                    y::Dual{T,Dual{S,Y,N},M}, z::Dual{T,Dual{S,Z,N},M}) where {T,S,Y<:Real,Z<:Real,N,M} = $yz_body

            @inline $(f)(x::Dual{T,Dual{S,X,N},M}, y::Dual{S,Y,N},           z::$R)                    where {T,S,X<:Real,Y<:Real,N,M} = $x_body
            @inline $(f)(x::Dual{T,Dual{S,X,N},M}, y::$R,                    z::Dual{S,Z,N})           where {T,S,X<:Real,Z<:Real,N,M} = $x_body
            @inline $(f)(x::$R,                    y::Dual{T,Dual{S,Y,N},M}, z::Dual{S,Z,N})           where {T,S,Y<:Real,Z<:Real,N,M} = $y_body
            @inline $(f)(x::Dual{S,X,N},           y::Dual{T,Dual{S,Y,N},M}, z::$R)                    where {T,S,X<:Real,Y<:Real,N,M} = $y_body
            @inline $(f)(x::Dual{S,X,N},           y::$R,                    z::Dual{T,Dual{S,Z,N},M}) where {T,S,X<:Real,Z<:Real,N,M} = $z_body
            @inline $(f)(x::$R,                    y::Dual{S,Y,N},           z::Dual{T,Dual{S,Z,N},M}) where {T,S,Y<:Real,Z<:Real,N,M} = $z_body
        end
        append!(defs.args, expr.args)
        for Q in REAL_TYPES
            Q === R && continue
            expr = quote
                @inline $(f)(x::Dual{T}, y::$R, z::$Q) where {T} = $x_body
                @inline $(f)(x::$R, y::Dual{T}, z::$Q) where {T} = $y_body
                @inline $(f)(x::$R, y::$Q, z::Dual{T}) where {T} = $z_body
            end
            append!(defs.args, expr.args)
        end
        expr = quote
            @inline $(f)(x::Dual{T}, y::$R, z::$R) where {T} = $x_body
            @inline $(f)(x::$R, y::Dual{T}, z::$R) where {T} = $y_body
            @inline $(f)(x::$R, y::$R, z::Dual{T}) where {T} = $z_body
        end
        append!(defs.args, expr.args)
    end
    return esc(defs)
end

function unary_dual_definition(M, f)
    work = CommonSubexpressions.cse(quote
        val = $(f)(x)
        deriv = $(DiffBase.diffrule(f, :x))
    end)
    return quote
        @inline function ($M).$(f)(d::Dual{T}) where T
            x = value(d)
            $work
            return Dual{T}(val, deriv * partials(d))
        end
    end
end

function binary_dual_definition(M, f)
    dvx, dvy = DiffBase.diffrule(f, :vx, :vy)
    xy_work = CommonSubexpressions.cse(quote
        val = $(f)(vx, vy)
        dvx = $dvx
        dvy = $dvy
    end)
    x_work = CommonSubexpressions.cse(quote
        val = $(f)(vx, vy)
        dvx = $dvx
    end)
    y_work = CommonSubexpressions.cse(quote
        val = $(f)(vx, vy)
        dvy = $dvy
    end)
    return quote
        @define_binary_dual_op(
            ($M).$(f),
            begin
                vx, vy = value(x), value(y)
                $xy_work
                return Dual{T}(val, _mul_partials(partials(x), partials(y), dvx, dvy))
            end,
            begin
                vx, vy = value(x), value(y)
                $x_work
                return Dual{T}(val, dvx * partials(x))
            end,
            begin
                vx, vy = value(x), value(y)
                $y_work
                return Dual{T}(val, dvy * partials(y))
            end
        )
    end
end

#####################
# Generic Functions #
#####################

Base.copy(d::Dual) = d

Base.eps(d::Dual) = eps(value(d))
Base.eps(::Type{D}) where {D<:Dual} = eps(valtype(D))

Base.rtoldefault(::Type{D}) where {D<:Dual} = Base.rtoldefault(valtype(D))

Base.floor(::Type{R}, d::Dual) where {R<:Real} = floor(R, value(d))
Base.floor(d::Dual) = floor(value(d))

Base.ceil(::Type{R}, d::Dual) where {R<:Real} = ceil(R, value(d))
Base.ceil(d::Dual) = ceil(value(d))

Base.trunc(::Type{R}, d::Dual) where {R<:Real} = trunc(R, value(d))
Base.trunc(d::Dual) = trunc(value(d))

Base.round(::Type{R}, d::Dual) where {R<:Real} = round(R, value(d))
Base.round(d::Dual) = round(value(d))

Base.hash(d::Dual) = hash(value(d))
Base.hash(d::Dual, hsh::UInt) = hash(value(d), hsh)

function Base.read(io::IO, ::Type{Dual{T,V,N}}) where {T,V,N}
    value = read(io, V)
    partials = read(io, Partials{N,V})
    return Dual{T,V,N}(value, partials)
end

function Base.write(io::IO, d::Dual)
    write(io, value(d))
    write(io, partials(d))
end

@inline Base.zero(d::Dual) = zero(typeof(d))
@inline Base.zero(::Type{Dual{T,V,N}}) where {T,V,N} = Dual{T}(zero(V), zero(Partials{N,V}))

@inline Base.one(d::Dual) = one(typeof(d))
@inline Base.one(::Type{Dual{T,V,N}}) where {T,V,N} = Dual{T}(one(V), zero(Partials{N,V}))

@inline Base.rand(d::Dual) = rand(typeof(d))
@inline Base.rand(::Type{Dual{T,V,N}}) where {T,V,N} = Dual{T}(rand(V), zero(Partials{N,V}))
@inline Base.rand(rng::AbstractRNG, d::Dual) = rand(rng, typeof(d))
@inline Base.rand(rng::AbstractRNG, ::Type{Dual{T,V,N}}) where {T,V,N} = Dual{T}(rand(rng, V), zero(Partials{N,V}))

# Predicates #
#------------#

isconstant(d::Dual) = iszero(partials(d))

for pred in RealInterface.UNARY_PREDICATES
    @eval Base.$(pred)(d::Dual) = $(pred)(value(d))
end

for pred in RealInterface.BINARY_PREDICATES
    @eval begin
        @define_binary_dual_op(
            Base.$(pred),
            $(pred)(value(x), value(y)),
            $(pred)(value(x), y),
            $(pred)(x, value(y))
        )
    end
end

########################
# Promotion/Conversion #
########################

function Base.promote_rule(::Type{Dual{T,A,N}},
                           ::Type{Dual{T,B,N}}) where {T,A<:Real,B<:Real,N}
    return Dual{T,promote_type(A, B),N}
end

for R in (:BigFloat, :Bool, :Irrational, :Real)
    @eval begin
        Base.promote_rule(::Type{R}, ::Type{Dual{T,V,N}}) where {R<:$R,T,V<:Real,N} = Dual{T,promote_type(R, V),N}
        Base.promote_rule(::Type{Dual{T,V,N}}, ::Type{R}) where {T,V<:Real,N,R<:$R} = Dual{T,promote_type(V, R),N}
    end
end

Base.convert(::Type{Dual{T,V,N}}, d::Dual{T}) where {T,V<:Real,N} = Dual{T}(convert(V, value(d)), convert(Partials{N,V}, partials(d)))
Base.convert(::Type{Dual{T,V,N}}, x::Real) where {T,V<:Real,N} = Dual{T}(V(x), zero(Partials{N,V}))
Base.convert(::Type{D}, d::D) where {D<:Dual} = d

Base.promote_array_type(F, ::Type{D}, ::Type{A}) where {D<:Dual,A<:AbstractFloat} = promote_type(D, A)
Base.promote_array_type(F, ::Type{<:Dual}, ::Type{<:AbstractFloat}, ::Type{P}) where {P} = P
Base.promote_array_type(F, ::Type{A}, ::Type{D}) where {D<:Dual,A<:AbstractFloat} = promote_type(D, A)
Base.promote_array_type(F, ::Type{<:AbstractFloat}, ::Type{<:Dual}, ::Type{P}) where {P} = P

Base.float(d::Dual{T,V,N}) where {T,V,N} = Dual{T,promote_type(V, Float16),N}(d)
Base.AbstractFloat(d::Dual{T,V,N}) where {T,V,N} = Dual{T,promote_type(V, Float16),N}(d)

##############
# Arithmetic #
##############

# Addition/Subtraction #
#----------------------#

@define_binary_dual_op(
    Base.:+,
    Dual{T}(value(x) + value(y), partials(x) + partials(y)),
    Dual{T}(value(x) + y, partials(x)),
    Dual{T}(x + value(y), partials(y))
)

@define_binary_dual_op(
    Base.:-,
    Dual{T}(value(x) - value(y), partials(x) - partials(y)),
    Dual{T}(value(x) - y, partials(x)),
    Dual{T}(x - value(y), -partials(y))
)

@inline Base.:-(d::Dual{T}) where {T} = Dual{T}(-value(d), -partials(d))

# Multiplication #
#----------------#

@define_binary_dual_op(
    Base.:*,
    begin
        vx, vy = value(x), value(y)
        Dual{T}(vx * vy, _mul_partials(partials(x), partials(y), vy, vx))
    end,
    Dual{T}(value(x) * y, partials(x) * y),
    Dual{T}(x * value(y), x * partials(y))
)

@inline Base.:*(d::Dual, x::Bool) = x ? d : (signbit(value(d))==0 ? zero(d) : -zero(d))
@inline Base.:*(x::Bool, d::Dual) = d * x

# Division #
#----------#

@define_binary_dual_op(
    Base.:/,
    begin
        vx, vy = value(x), value(y)
        Dual{T}(vx / vy, _div_partials(partials(x), partials(y), vx, vy))
    end,
    Dual{T}(value(x) / y, partials(x) / y),
    begin
        v = value(y)
        divv = x / v
        Dual{T}(divv, -(divv / v) * partials(y))
    end
)

# Exponentiation #
#----------------#

for f in (:(Base.:^), :(NaNMath.pow))
    @eval begin
        @define_binary_dual_op(
            $f,
            begin
                vx, vy = value(x), value(y)
                expv = ($f)(vx, vy)
                powval = vy * ($f)(vx, vy - 1)
                logval = isconstant(y) ? one(expv) : expv * log(vx)
                new_partials = _mul_partials(partials(x), partials(y), powval, logval)
                return Dual{T}(expv, new_partials)
            end,
            begin
                v = value(x)
                expv = ($f)(v, y)
                deriv = y * ($f)(v, y - 1)
                return Dual{T}(expv, deriv * partials(x))
            end,
            begin
                v = value(y)
                expv = ($f)(x, v)
                deriv = expv*log(x)
                return Dual{T}(expv, deriv * partials(y))
            end
        )
    end
end

###################################
# General Mathematical Operations #
###################################

@inline Base.conj(d::Dual) = d

@inline Base.transpose(d::Dual) = d

@inline Base.ctranspose(d::Dual) = d

@inline Base.abs(d::Dual) = signbit(value(d)) ? -d : d

for f in RealInterface.UNARY_MATH
    DiffBase.hasdiffrule(f, 1) && eval(unary_dual_definition(Base, f))
end

for f in RealInterface.UNARY_SPECIAL_MATH
    DiffBase.hasdiffrule(f, 1) && eval(unary_dual_definition(SpecialFunctions, f))
end

for f in RealInterface.BINARY_SPECIAL_MATH
    DiffBase.hasdiffrule(f, 2) && eval(binary_dual_definition(SpecialFunctions, f))
end

function to_nanmath!(x)
    if isa(x, Expr)
        if x.head == :call
            f = x.args[1]
            if in(f, RealInterface.UNARY_NAN_MATH) || in(f, RealInterface.BINARY_NAN_MATH)
                x.args[1] = :(NaNMath.$f)
            end
            foreach(to_nanmath!, x.args[2:end])
        else
            foreach(to_nanmath!, x.args)
        end
    end
    return x
end

for f in RealInterface.UNARY_NAN_MATH
    DiffBase.hasdiffrule(f, 1) && eval(to_nanmath!(unary_dual_definition(NaNMath, f)))
end

#################
# Special Cases #
#################

# hypot #
#-------#

@inline function calc_hypot(x, y, ::Type{T}) where T
    vx = value(x)
    vy = value(y)
    h = hypot(vx, vy)
    return Dual{T}(h, (vx/h) * partials(x) + (vy/h) * partials(y))
end

@define_binary_dual_op(
    Base.hypot,
    calc_hypot(x, y, T),
    calc_hypot(x, y, T),
    calc_hypot(x, y, T)
)

@inline function calc_hypot(x, y, z, ::Type{T}) where T
    vx = value(x)
    vy = value(y)
    vz = value(z)
    h = hypot(vx, vy, vz)
    p = (vx / h) * partials(x) + (vy / h) * partials(y) + (vz / h) * partials(z)
    return Dual{T}(h, p)
end

@define_ternary_dual_op(
    Base.hypot,
    calc_hypot(x, y, z, T),
    calc_hypot(x, y, z, T),
    calc_hypot(x, y, z, T),
    calc_hypot(x, y, z, T),
    calc_hypot(x, y, z, T),
    calc_hypot(x, y, z, T),
    calc_hypot(x, y, z, T),
)

# atan2 #
#-------#

@inline function calc_atan2(y, x, ::Type{T}) where T
    z = y / x
    v = value(z)
    atan2v = atan2(value(y), value(x))
    deriv = inv(one(v) + v*v)
    return Dual{T}(atan2v, deriv * partials(z))
end

@define_binary_dual_op(
    Base.atan2,
    calc_atan2(x, y, T),
    calc_atan2(x, y, T),
    calc_atan2(x, y, T)
)

# fma #
#-----#

@generated function calc_fma_xyz(x::Dual{T,<:Real,N},
                                 y::Dual{T,<:Real,N},
                                 z::Dual{T,<:Real,N}) where {T,N}
    ex = Expr(:tuple, [:(fma(value(x), partials(y)[$i], fma(value(y), partials(x)[$i], partials(z)[$i]))) for i in 1:N]...)
    return quote
        $(Expr(:meta, :inline))
        v = fma(value(x), value(y), value(z))
        return Dual{T}(v, $ex)
    end
end

@inline function calc_fma_xy(x::Dual{T}, y::Dual{T}, z::Real) where T
    vx, vy = value(x), value(y)
    result = fma(vx, vy, z)
    return Dual{T}(result, _mul_partials(partials(x), partials(y), vy, vx))
end

@generated function calc_fma_xz(x::Dual{T,<:Real,N},
                                y::Real,
                                z::Dual{T,<:Real,N}) where {T,N}
    ex = Expr(:tuple, [:(fma(partials(x)[$i], y,  partials(z)[$i])) for i in 1:N]...)
    return quote
        $(Expr(:meta, :inline))
        v = fma(value(x), y, value(z))
        Dual(v, $ex)
    end
end

@define_ternary_dual_op(
    Base.fma,
    calc_fma_xyz(x, y, z),                         # xyz_body
    calc_fma_xy(x, y, z),                          # xy_body
    calc_fma_xz(x, y, z),                          # xz_body
    Base.fma(y, x, z),                             # yz_body
    Dual{T}(fma(value(x), y, z), partials(x) * y), # x_body
    Base.fma(y, x, z),                             # y_body
    Dual{T}(fma(x, y, value(z)), partials(z))      # z_body
)

# sincos #
#--------#

@inline sincos(x) = (sin(x), cos(x))

@inline function sincos(d::Dual{T}) where T
    sd, cd = sincos(value(d))
    return (Dual{T}(sd, cd * partials(d)), Dual{T}(cd, -sd * partials(d)))
end

###################
# Pretty Printing #
###################

function Base.show(io::IO, d::Dual{T,V,N}) where {T,V,N}
    print(io, "Dual{$(repr(T))}(", value(d))
    for i in 1:N
        print(io, ",", partials(d, i))
    end
    print(io, ")")
end


text:
# This file was generated, do not modify it. # hide
describe(data, :mean, :median, :nmissing)


text:
"""
x = pivot(A, b; ...)

Solves non-negative least-squares problem by block principal pivoting method
(Algorithm 1) described in Kim & Park (2011).

Optional arguments:
    tol: tolerance for nonnegativity constraints
    max_iter: maximum number of iterations

References:
	J. Kim and H. Park, Fast nonnegative matrix factorization: An
	active-set-like method and comparisons, SIAM J. Sci. Comput., 33 (2011),
	pp. 3261–3281.
"""
function pivot(A,
               b::AbstractVector{T};
               tol::Float64=1e-8,
               max_iter=30*size(A,2)) where T

    # dimensions, initialize solution
    p,q = size(A)

    x = zeros(T, q) # primal variables
    y = -A'*b    # dual variables

    # parameters for swapping
    α = 3
    β = q+1

    # Store indices for the passive set, P
    #    we want Y[P] == 0, X[P] >= 0
    #    we want X[~P]== 0, Y[~P] >= 0
    P = BitArray(false for _ in 1:q)

    y[(!).(P)] =  A[:,(!).(P)]' * (A[:,P]*x[P] - b)

    # identify indices of infeasible variables
    V = @__dot__ (P & (x < -tol)) | (!P & (y < -tol))
    nV = sum(V)

    # while infeasible (number of infeasible variables > 0)
    while nV > 0

    	if nV < β
    		# infeasible variables decreased
    		β = nV  # store number of infeasible variables
    		α = 3   # reset α
    	else
    		# infeasible variables stayed the same or increased
    		if α >= 1
    			α = α-1 # tolerate increases for α cycles
    		else
    			# backup rule
    			i = findlast(V)
    			V = zeros(Bool,q)
    			V[i] = true
    		end
    	end

    	# update passive set
        #     P & ~V removes infeasible variables from P
        #     V & ~P  moves infeasible variables in ~P to P
		@__dot__ P = (P & !V) | (V & !P)

		# update primal/dual variables
		if !all(!, P)
			x[P] =  A[:,P] \ b
		end
		y[(!).(P)] =  A[:,(!).(P)]' * ((A[:,P]*x[P]) - b)

        # check infeasibility
        @__dot__ V = (P & (x < -tol)) | (!P & (y < -tol))
        nV = sum(V)
    end

    x[(!).(P)] .= zero(eltype(x))
    return x
end

## if multiple right hand sides are provided, solve each problem separately.
function pivot(A,
               B::AbstractMatrix{T};
               use_parallel = true,
               kwargs...) where {T}

    n = size(A,2)
    k = size(B,2)

    # compute result for each column
    if use_parallel && nprocs()>1
        X = @distributed (hcat) for i = 1:k
            pivot(A, B[:,i]; kwargs...)
        end
    else
        X = Array{T}(undef,n,k)
        for i = 1:k
            X[:,i] = pivot(A, B[:,i]; kwargs...)
        end
    end

    return X
end


text:
# ******************************************************************************
# PACKAGES
# ******************************************************************************
using LinearAlgebra
using Random
using Distributions

# ******************************************************************************
# OPTIMIZATION FUNCTION DEFNITIONS
# ******************************************************************************
# A function for constructing a basis vector
basis(i,n) = [k==i ? 1.0 : 0.0 for k in 1:n]

# A function to perform a single step of the Hooke-Jeeves optimization method
function hooke_jeeves(f, x, y, α, n, γ = 0.5)
    improved = false
    x_best, y_best = x, y
    for i in 1:n
        for sign in (-1, 1)
            x′ = x + sign*α*basis(i,n)
            y′ = f(x′)
            if y′ < y_best
                x_best, y_best, improved = x′, y′, true
            end
        end
    end
    x, y = x_best, y_best
    if !improved
        α *= γ
    end
    return x, y, α
end

# A function to randomly sample a positive spanning set according to mesh
# adaptive direct search
function random_positive_spanning_set(α, n)
    δ = round(Int, 1/sqrt(α))
    lower_triangular_matrix = Matrix(Diagonal(δ*rand([1,-1], n)))
    for i in 1:(n-1)
        for j in 1:(i-1)
            lower_triangular_matrix[i,j] = rand(-δ + 1:(δ-1))
        end
    end
    D = lower_triangular_matrix[randperm(n),:][:, randperm(n)]
    D = hcat(D, -sum(D,dims=2))
    return [D[:,i] for i in 1:(n+1)]
end

# A function to perform a single step of mesh adaptive direct search
function mesh_adaptive_direct_search(f, x, y, α, n)
    improved = false
    for (i,d) in enumerate(random_positive_spanning_set(α, n))
        x′ = x + α*d; y′ = f(x′)
        if y′ < y
            x, y, improved = x′, y′, true
            x′ = x + 3*α*d; y′ = f(x′)
            if y′ < y
                x, y = x′, y′
            end
            break
        end
    end
    α = improved ? min(2*α, 1) : α/2
    return x, y, α
end

# A function to perform a single step of the cross entropy method
function cross_entropy(f, P, m=100, m_elite=10)
    samples = rand(P,m)
    order = sortperm([f(samples[:,i]) for i in 1:m])
    P = fit(typeof(P), samples[:, order[1:m_elite]])
    return P
end

# A penalty method mixture between a count and a quadratic penalty function
function penalty(c, x, ρ1, ρ2, δ)
    evals = c(x)
    max_arr = [max(evals[i] + δ, 0) for i = 1:length(evals)]
    greater_than_zero = [evals[i] > 0 for i = 1:length(evals)]
    max_arr = max_arr.^2
    return ρ1*sum(max_arr) + ρ2*sum(greater_than_zero)
end

# HJ parameters: ρ1=1.01, ρ2=0, γ=1.001, δ=0.5; 0.1378383301
# MADS parameters: ρ1=1.01, ρ2=0, γ=1.001, δ=0.5; 0.1395356843
function optimize_simple1(f, g, c, x, num_evals; ρ1=1.01, ρ2=0, γ=1.001, δ=0.5)
    α, y, n, ϵ = 1, f(x) + penalty(c, x, ρ1, ρ2, δ), length(x), 1e-9
    while α > ϵ && count(f, g, c) < (num_evals-8)
        #x, y, α = mesh_adaptive_direct_search(x -> f(x) + penalty(c, x, ρ1, ρ2, δ), x, y, α, n)
        x, y, α = hooke_jeeves(x -> f(x) + penalty(c, x, ρ1, ρ2, δ), x, y, α, n, 0.5)
        ρ1 *= γ; ρ2 *= γ
    end
    return x
end

# HJ parameters: ρ1=100, ρ2=1, γ=1.1, δ=0.75; 1.4993756041
# MADS parameters: ρ1=100, ρ2=1, γ=1.01, δ=1.0; 1.5604945284
function optimize_simple2(f, g, c, x, num_evals; ρ1=100, ρ2=1, γ=1.1, δ=0.75)
    α, y, n, ϵ = 1, f(x) + penalty(c, x, ρ1, ρ2, δ), length(x), 1e-9
    while α > ϵ && count(f, g, c) < (num_evals-8)
        #x, y, α = mesh_adaptive_direct_search(x -> f(x) + penalty(c, x, ρ1, ρ2, δ), x, y, α, n)
        x, y, α = hooke_jeeves(x -> f(x) + penalty(c, x, ρ1, ρ2, δ), x, y, α, n, 0.5)
        ρ1 *= γ; ρ2 *= γ
    end
    return x
end

# HJ parameters: ρ1=1, ρ2=1, γ=1.01, δ=0.5; 0.0631003421
# MADS parameters: ρ1=1, ρ2=1, γ=1.01, δ=0.25; 0.0425930332
function optimize_simple3(f, g, c, x, num_evals; ρ1=1, ρ2=1, γ=1.01, δ=0.25)
    α, y, n, ϵ = 1, f(x) + penalty(c, x, ρ1, ρ2, δ), length(x), 1e-9
    while α > ϵ && count(f, g, c) < (num_evals-13)
        x, y, α = mesh_adaptive_direct_search(x -> f(x) + penalty(c, x, ρ1, ρ2, δ), x, y, α, n)
        #x, y, α = hooke_jeeves(x -> f(x) + penalty(c, x, ρ1, ρ2, δ), x, y, α, n, 0.5)
        ρ1 *= γ; ρ2 *= γ
    end
    return x
end

# ρ1=50, ρ2=1, γ=1.01, δ=0.15, m = 100; num_elite = 20, μ=x; Σ=10; -0.5935202993
function optimize_secret1(f, g, c, x, num_evals; ρ1=50, ρ2=1, γ=1.01, δ=0.15)
    μ = x; Σ = 10; P = MvNormal(μ, Σ)
    m = 100; num_elite = 20
    while count(f, g, c) < (num_evals - 2*m)
        P = cross_entropy(x -> f(x) + penalty(c, x, ρ1, ρ2, δ), P, m, num_elite)
        ρ1 *= γ; ρ2 *= γ
    end
    return P.μ
end

# HJ parameters:  ρ1=1, ρ2=1, γ=1.001, δ=0.15; 0.657090133
# MADS parameters:  ρ1=1, ρ2=1, γ=1.001, δ=0.15; 0.3410080274
function optimize_secret2(f, g, c, x, num_evals; ρ1=1, ρ2=1, γ=1.001, δ=0.15)
    α, y, n, ϵ = 0.1, f(x) + penalty(c, x, ρ1, ρ2, δ), length(x), 1e-9
    while α > ϵ && count(f, g, c) < (num_evals - 308)
        x, y, α = mesh_adaptive_direct_search(x -> f(x) + penalty(c, x, ρ1, ρ2, δ), x, y, α, n)
        #x, y, α = hooke_jeeves(x -> f(x) + penalty(c, x, ρ1, ρ2, δ), x, y, α, n, 0.5)
        ρ1 *= γ; ρ2 *= γ
    end
    return x
end

function test_call()
    print("HELLO WORLD")
end

# *******************************************************************************
# MAIN OPTIMIZATION FUNCTION
# *******************************************************************************
function optimize(f, g, c, x0, n, prob)
    if prob == "simple1"
        x_best = optimize_simple1(f, g, c, x0, n)
    elseif prob == "simple2"
        x_best = optimize_simple2(f, g, c, x0, n)
    elseif prob == "simple3"
        x_best = optimize_simple3(f, g, c, x0, n)
    elseif prob == "secret1"
        x_best = optimize_secret1(f, g, c, x0, n)
    elseif prob == "secret2"
        x_best = optimize_secret2(f, g, c, x0, n)
    else
        print("Invalid Problem")
    end
end


text:
using Images, ImageView, ImageMorphology

# load an image and first convert to grayscale and then binarize
img = load("sample-images/bird-3183441_640.jpg");
img_binary = Gray.(img) .> 0.5;
img_binary = Gray.(img_binary)

# resizing images by half
img = restrict(img)
img_binary = restrict(RGB.(img_binary))

# stacking 2 images together
img_width = size(img, 2)

combined_image = fill(RGB4{Float16}(0.,0.,0.), size(img) .* (1, 2))
combined_image[:, 1:img_width] = img
combined_image[:, img_width+1:img_width*2] = img_binary
imshow(combined_image)


text:
struct KrylovTrustRegion{T <: Real} <: SecondOrderOptimizer
    initial_radius::T
    max_radius::T
    eta::T
    rho_lower::T
    rho_upper::T
    cg_tol::T
end

KrylovTrustRegion(; initial_radius::Real = 1.0,
                    max_radius::Real = 100.0,
                    eta::Real = 0.1,
                    rho_lower::Real = 0.25,
                    rho_upper::Real = 0.75,
                    cg_tol::Real = 0.01) =
                    KrylovTrustRegion(initial_radius, max_radius, eta,
                                  rho_lower, rho_upper, cg_tol)

update_h!(d, state, method::KrylovTrustRegion) = nothing

# TODO: support x::Array{T,N} et al.?
mutable struct KrylovTrustRegionState{T} <: AbstractOptimizerState
    x::Vector{T}
    x_previous::Vector{T}
    f_x_previous::T
    s::Vector{T}
    interior::Bool
    accept_step::Bool
    radius::T
    m_diff::T
    f_diff::T
    rho::T
    r::Vector{T}  # residual vector
    d::Vector{T}  # direction to consider
    cg_iters::Int
end

function initial_state(method::KrylovTrustRegion, options, d, initial_x::Array{T}) where T
    n = length(initial_x)
    # Maintain current gradient in gr
    @assert(method.max_radius > 0)
    @assert(0 < method.initial_radius < method.max_radius)
    @assert(0 <= method.eta < method.rho_lower)
    @assert(method.rho_lower < method.rho_upper)
    @assert(method.rho_lower >= 0)

    value_gradient!!(d, initial_x)

    KrylovTrustRegionState(copy(initial_x),    # Maintain current state in state.x
                           copy(initial_x), # x_previous
                           zero(T),            # f_x_previous
                           similar(initial_x), # Maintain current search direction in state.s
                           true,               # interior
                           true,               # accept step
                           method.initial_radius,
                           zero(T),            # model change
                           zero(T),            # observed f change
                           zero(T),            # state.rho
                           Vector{T}(undef, n),       # residual vector
                           Vector{T}(undef, n),       # direction to consider
                           0)                  # cg_iters
end

function trace!(tr, d, state, iteration, method::KrylovTrustRegion, options, curr_time=time())
    dt = Dict()
    dt["time"] = curr_time
    if options.extended_trace
        dt["radius"] = copy(state.radius)
        dt["interior"] = state.interior
        dt["accept_step"] = state.accept_step
        dt["norm(s)"] = norm(state.s)
        dt["rho"] = state.rho
        dt["m_diff"] = state.m_diff
        dt["f_diff"] = state.f_diff
        dt["cg_iters"] = state.cg_iters
    end
    g_norm = norm(gradient(d), Inf)
    update!(tr,
            iteration,
            value(d),
            g_norm,
            dt,
            options.store_trace,
            options.show_trace,
            options.show_every,
            options.callback)
end

function cg_steihaug!(objective::TwiceDifferentiableHV,
                      state::KrylovTrustRegionState{T},
                      method::KrylovTrustRegion) where T
    n = length(state.x)
    x, g, d, r, z, Hd = state.x, gradient(objective), state.d, state.r, state.s, hv_product(objective)

    fill!(z, 0.0)  # the search direction is initialized to the 0 vector,
    r .= g  # so at first the whole gradient is the residual.
    d .= -r # the first direction is the direction of steepest descent.
    rho0 = 1e100  # just a big number

    state.cg_iters = 0
    for i in 1:n
        state.cg_iters += 1
        hv_product!(objective, x, d)
        dHd = dot(d, Hd)
        if -1e-15 < dHd < 1e-15
            break
        end

        alpha = dot(r, r) / dHd

        if dHd < 0. || norm(z .+ alpha .* d) >= state.radius
            a_ = dot(d, d)
            b_ = 2 * dot(z, d)
            c_ = dot(z, z) - state.radius^2
            tau = (-b_ + sqrt(b_ * b_ - 4 * a_ * c_)) / (2 * a_)
            z .+= tau .* d
            break
        end

        z .+= alpha .* d
        rho_prev = dot(r, r)
        if i == 1
            rho0 = rho_prev
        end
        r .+= alpha * Hd
        rho_next = dot(r, r)
        r_sqnorm_ratio = rho_next / rho_prev
        d[:] = -r + r_sqnorm_ratio * d

        if (rho_next / rho0) < method.cg_tol^2
            break
        end
    end

    hv_product!(objective, x, z)
    return dot(g, z) + 0.5 * dot(z, Hd)
end

function update_state!(objective::TwiceDifferentiableHV,
                          state::KrylovTrustRegionState,
                          method::KrylovTrustRegion)
    state.m_diff = cg_steihaug!(objective, state, method)
    @assert state.m_diff <= 0

    state.f_diff = value(objective, state.x .+ state.s) - value(objective)
    state.rho = state.f_diff / state.m_diff
    state.interior = norm(state.s) < 0.9 * state.radius

    if state.rho < method.rho_lower
        state.radius *= 0.25
    elseif (state.rho > method.rho_upper) && (!state.interior)
        state.radius = min(2 * state.radius, method.max_radius)
    end

    state.accept_step = state.rho > method.eta
    if state.accept_step
        state.x .+= state.s
    end

    return false
end

function update_g!(objective, state::KrylovTrustRegionState, method::KrylovTrustRegion)
    if state.accept_step
        # Update the function value and gradient
        state.f_x_previous = value(objective)
        value_gradient!(objective, state.x)
    end
end

function assess_convergence(state::KrylovTrustRegionState, d, options::Options)
    if !state.accept_step
        return state.radius < options.x_abstol, false, false, false, false
    end

    x_converged, f_converged, f_increased, g_converged = false, false, false, false

    if norm(state.s, Inf) < options.x_abstol
        x_converged = true
    end

    # Absolute Tolerance
    # if abs(f_x - f_x_previous) < f_tol
    # Relative Tolerance
    if abs(state.f_diff) < max(options.f_reltol * (abs(value(d)) + options.f_reltol), eps(abs(value(d))+abs(state.f_x_previous)))
        f_converged = true
    end

    if norm(gradient(d), Inf) < options.g_abstol
        g_converged = true
    end

    converged = x_converged || f_converged || g_converged

    return x_converged, f_converged, g_converged, converged, false
end


text:
"""parse a single line of space separated integers"""
space_sep_ints(s::AbstractString; base=10) = parse.(Int, split(s); base)

function file_extension(path::String)
    if contains(path, ".")
        return path[findlast(isequal('.'),path):end]
    else
        return ""
    end
end