text:
using Printf
using LoopVectorization

Base.@kwdef mutable struct FCFV_Mesh
    type   ::Union{String, Missing}          = missing
    nel    ::Union{Int64,  Missing}          = missing
    nf     ::Union{Int64,  Missing}          = missing
    nv     ::Union{Int64,  Missing}          = missing
    nn_el  ::Union{Int64,  Missing}          = missing
    nf_el  ::Union{Int64,  Missing}          = missing
    xv     ::Union{Vector{Float64}, Missing} = missing # node x coordinate
    yv     ::Union{Vector{Float64}, Missing} = missing # node y coordinate
    xf     ::Union{Vector{Float64}, Missing} = missing # face x coordinate
    yf     ::Union{Vector{Float64}, Missing} = missing # face y coordinate
    bc     ::Union{Vector{Int64},   Missing} = missing # face y coordinate
    xc     ::Union{Vector{Float64}, Missing} = missing # cent x coordinate
    yc     ::Union{Vector{Float64}, Missing} = missing # cent y coordinate
    e2v    ::Union{Matrix{Int64},   Missing} = missing # cell 2 node numbering
    e2f    ::Union{Matrix{Int64},   Missing} = missing # cell 2 face numbering
    vole   ::Union{Vector{Float64}, Missing} = missing # volume of element
    n_x    ::Union{Matrix{Float64}, Missing} = missing # normal 2 face x
    n_y    ::Union{Matrix{Float64}, Missing} = missing # normal 2 face y
    dA     ::Union{Matrix{Float64}, Missing} = missing # face length
end

 # Create sides of mesh
     xmin, xmax = 0, 1
     ymin, ymax = 0, 1
     nx, ny     = 50, 50

 # Generates a 2D rectangular mesh of nx*ny cells * 2 triangles
 ncell      = nx*ny*2
 # Twice resolution mesh
 nx2 = 2*nx + 1
 ny2 = 2*ny + 1
 # 1D axis
 dx,dy = (xmax-xmin)/nx, (xmax-xmin)/ny
 xc  = collect(LinRange(xmin+dx/2.0,xmax-dx/2.0,nx))
 yc  = collect(LinRange(ymin+dy/2.0,ymax-dy/2.0,ny))
 x2  = LinRange(xmin,xmax,nx2)
 y2  = LinRange(ymin,ymax,ny2)
 # 2D mesh: fake ngrid for P-T space ;)
 x2d   = repeat(x2, 1, length(y2))'
 y2d   = repeat(y2, 1, length(x2))
 nodes = zeros(Int64, size(x2d))
 # Number active dofs (midfaces) - Face node numbering
 inum = 0
 for i=1:nx2
     for j=1:ny2
         if (mod(j,2)==1 && mod(i,2)==0) || (mod(j,2)==0 && mod(i,2)==1) ||  (mod(j,2)==0 && mod(i,2)==0)
             global inum += 1
             nodes[i,j] = inum
         end
     end
 end
 nface = inum;
 # Number vertices
 inum = 0
 for i=1:nx2
     for j=1:ny2
         if (mod(j,2)==1 && mod(i,2)==1) || (mod(j,2)==1 && mod(i,2)==1)
            global inum += 1
             nodes[i,j] = inum
         end
     end
 end
 # Face nodes
 xf = zeros(nface)
 yf = zeros(nface)
 tf = zeros(Int64,nface) # node type - 1 boundary
 for i=1:nx2
     for j=1:ny2
         if (mod(j,2)==1 && mod(i,2)==0) || (mod(j,2)==0 && mod(i,2)==1) ||  (mod(j,2)==0 && mod(i,2)==0)
             xf[nodes[i,j]] = x2d[i,j]
             yf[nodes[i,j]] = y2d[i,j]
             if (i==1 || i==nx2 || j==1 || j==ny2)
                 tf[nodes[i,j]] = 1
             end
         end
     end
 end
 # Vertices nodes
 xn = zeros( (nx+1)*(ny+1) )
 yn = zeros( (nx+1)*(ny+1) )
 tn = zeros(Int64, (nx+1)*(ny+1) ) # node type - 1 boundary
 for i=1:nx+1
     for j=1:ny+1
         ii = i + (j-1)*(ny+1)
             xn[ii] = xmin + (i-1)*dx
             yn[ii] = ymin + (j-1)*dy
             if (i==1 || i==nx+1 || j==1 || j==ny+1)
                 tn[ii] = 1
             end
     end
 end
 # Cell 2 face numbering - for matrix connectivity - lower triangles
 face = zeros(Int64, 3, ncell)
 @avx for i=1:nx
     for j=1:ny
         k  = j + (i-1)*ny
         jc = i+1 + (i-1)*1
         ic = j+1 + (j-1)*1
         face[1,k] =  nodes[ic,jc-1] # South
         face[2,k] =  nodes[ic-1,jc] # West
         face[3,k] =  nodes[ic  ,jc] # centerh
         # println(face[:,k])
     end
 end
  # Cell 2 face numbering - for matrix connectivity - upper triangles
 @avx for i=1:nx
    for j=1:ny
        k  = j + (i-1)*ny + nx*ny
        jc = i+1 + (i-1)*1
        ic = j+1 + (j-1)*1
        face[1,k] =  nodes[ic+1,jc] # East
        face[2,k] =  nodes[ic,jc+1] # North
        face[3,k] =  nodes[ic  ,jc] # center
        # println(face[:,k])
    end
end
 # Cell 2 vertices - used for visualisation of quads - lower triangles
 vert = zeros(Int64, 3, ncell)
 @avx for i=1:nx
     for j=1:ny
         k  = j + (i-1)*ny
         jc = i+1 + (i-1)*1
         ic = j+1 + (j-1)*1
         vert[1,k] =  nodes[ic-1,jc+1]
         vert[2,k] =  nodes[ic-1,jc-1]
         vert[3,k] =  nodes[ic+1,jc-1]
        #  vert[4,k] =  nodes[ic+1,jc+1]
     end
 end
  # Cell 2 vertices - used for visualisation of quads - lower triangles
  @avx for i=1:nx
      for j=1:ny
          k  = j + (i-1)*ny + nx*ny
          jc = i+1 + (i-1)*1
          ic = j+1 + (j-1)*1
          vert[1,k] =  nodes[ic-1,jc+1]
          vert[2,k] =  nodes[ic+1,jc+1]
          vert[3,k] =  nodes[ic+1,jc-1]
      end
  end
 # Centroids
 xc = zeros(ncell)
 yc = zeros(ncell)
 w  = 1.0/3.0
 @avx for iel=1:ncell
     tempx = 0
     tempy = 0
     for j=1:3
         tempx += w * xf[face[j,iel]]
         tempy += w * yf[face[j,iel]]
     end
     xc[iel] = tempx
     yc[iel] = tempy
     # println(xc[iel])
 end

 # Fill structure
 mesh        = FCFV_Mesh()
 mesh.type   = "triangle_struct"
 mesh.nel    = ncell
 mesh.nf     = nface
 mesh.nf_el  = 3
 mesh.nf_el  = 3
 mesh.nv     = (nx+1)*(ny+1)
 mesh.xv     = xn
 mesh.yv     = yn
 mesh.xf     = xf
 mesh.yf     = yf
 mesh.e2v    = vert' # cell nodes - not dofs! Dofs are in 'cell'
 mesh.e2f    = face'
 mesh.xc     = xc
 mesh.yc     = yc
 mesh.vole   = dx*dy*ones(ncell)/2
 mesh.bc     = tf

 nodeA = [2 1 3]
 nodeB = [3 2 1]

 # Compute normal to faces
 mesh.n_x = zeros(Float64,mesh.nel,mesh.nf_el)
 mesh.n_y = zeros(Float64,mesh.nel,mesh.nf_el)
 mesh.dA  = zeros(Float64,mesh.nel,mesh.nf_el)

  # Assemble FCFV elements
  for iel=1:mesh.nel

     # println("element: ",  iel)

     for ifac=1:mesh.nf_el

         nodei  = mesh.e2f[iel,ifac]

         # Vertices
         vert1  = mesh.e2v[iel,nodeA[ifac]]
         vert2  = mesh.e2v[iel,nodeB[ifac]]
         bc     = mesh.bc[nodei]
         dxl     = abs(mesh.xv[vert1] - mesh.xv[vert2] );
         dyl     = abs(mesh.yv[vert1] - mesh.yv[vert2] );
         dAi    = sqrt(dxl^2 + dyl^2);

        #  if iel==1 || iel==(nx*ny+1)
        #      println(bc)
        #      @printf("face node, x = %2.2e y = %2.2e\n", mesh.xf[nodei], mesh.yf[nodei])
        #      @printf("vert1    , x = %2.2e y = %2.2e\n", mesh.xv[vert1], mesh.yv[vert1])
        #      @printf("vert2    , x = %2.2e y = %2.2e\n", mesh.xv[vert2], mesh.yv[vert2])
        #  end

         # Face normal
         n_x  =  dy/dAi
         n_y  = -dx/dAi

         # Third vector
         v_x  = mesh.xf[nodei] - mesh.xc[iel]
         v_y  = mesh.yf[nodei] - mesh.yc[iel]

         # Check wether the normal points outwards
         dot                 = n_x*v_x + n_y*v_y
         mesh.n_x[iel,ifac]  = (dot>=0.0)*n_x - (dot<0.0)*n_x
         mesh.n_y[iel,ifac]  = (dot>=0.0)*n_y - (dot<0.0)*n_y
         mesh.dA[iel,ifac]   = dAi
     end
 end


text:
using FiniteDiff
using BenchmarkTools

struct MuJoCoModel{T,CX,CU}
    m::jlModel
    d::jlData
    x::Vector{T}
    fx::Matrix{T}
    fu::Matrix{T}
    nx::Int
    nu::Int
    idx_u::Vector{Int}
    idx_pos::Vector{Int}
    idx_vel::Vector{Int}
    idx_next::Vector{Int}
    cache_x::CX
    cache_u::CU
end

function MuJoCoModel(path)
    m = jlModel(path)
    d = jlData(m)

    nq = length(d.qpos)
    nv = length(d.qvel)
    nx = nq + nv
    nu = length(d.ctrl)

    x = zeros(nx)
    u = zeros(nu)
    fx = zeros(nx, nx)
    fu = zeros(nx, nu)

    idx_u = collect(1:nu)
    idx_pos = collect(1:nq)
    idx_vel = collect(nq .+ (1:nv))
    idx_next = collect(1:nq)

    cache_x = FiniteDiff.JacobianCache(x, x)
    cache_u = FiniteDiff.JacobianCache(u, x)

    MuJoCoModel(m, d, x,
        fx, fu,
        nx, nu,
        idx_u, idx_pos, idx_vel, idx_next,
        cache_x, cache_u)
end

# dynamics methods
function f_mujoco!(y, model::MuJoCoModel, x, u)
    model.d.ctrl .= @views u[model.idx_u]
    model.d.qpos[model.idx_next] .= @views x[model.idx_pos]
    model.d.qvel[model.idx_next] .= @views x[model.idx_vel]
    mj_step(model.m, model.d)
    y[model.idx_pos] .= @views model.d.qpos[model.idx_next]
    y[model.idx_vel] .= @views model.d.qvel[model.idx_next]
    return y
end

f_mujoco(model::MuJoCoModel, x, u) = f_mujoco!(model.x, model, x, u)

function fx_mujoco!(dx, model::MuJoCoModel, x, u)
    FiniteDiff.finite_difference_jacobian!(dx, (a, b) -> f_mujoco!(a, model, b, u), x, model.cache_x)
    return dx
end

fx_mujoco(model::MuJoCoModel, x, u) = fx_mujoco!(model.fx, model, x, u)

function fu_mujoco!(du, model::MuJoCoModel, x, u)
    FiniteDiff.finite_difference_jacobian!(du, (a, b) -> f_mujoco!(a, model, x, b), u, model.cache_u)
    return du
end

fu_mujoco(model::MuJoCoModel, x, u) = fu_mujoco!(model.fu, model, x, u)

# x2 = zeros(acrobot.nx)
# x1 = zeros(acrobot.nx)
# u1 = zeros(acrobot.nu)
# f(acrobot, x1, u1)
# @benchmark f($acrobot, $x1, $u1)
# f!(x2, acrobot, x1, u1)
# @benchmark f!($x2, $acrobot, $x1, $u1)
# fx(acrobot, x1, u1)
# @benchmark fx($acrobot, $x1, $u1)
# fu(acrobot, x1, u1)
# @benchmark fu($acrobot, $x1, $u1)

ext: jl


text:
using InvertibleNetworks, LinearAlgebra, Test

# Test affine or additive coupling layer
affine = true

# Input
nx1 = 16
nx2 = 16
nx_in = 2
nx_hidden = 4
batchsize = 2

# Observed data
nd1 = 14
nd2 = 8
nd_in = 1

# Modeling/imaging operator
A = randn(Float32, nd1*nd2*nd_in, nx1*nx2)

# Link function
Ψ(η) = identity(η)

# Unrolled loop
if affine
    CS = AffineCouplingLayerSLIM(nx1, nx2, nx_in, nx_hidden, batchsize, Ψ; logdet=false, permute=true)
else
    CS = AdditiveCouplingLayerSLIM(nx1, nx2, nx_in, nx_hidden, batchsize, Ψ; logdet=false, permute=true)
end

# Initializations
D = randn(Float32, nd1*nd2*nd_in, batchsize)
X = randn(Float32, nx1, nx2, nx_in, batchsize)

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

# Test invertibility
Y = CS.forward(X, D, A)
X_ = CS.inverse(Y, D, A)
@test isapprox(norm(X - X_)/norm(X), 0f0, atol=1e-6)

Y = CS.inverse(X, D, A)
X_ = CS.forward(Y, D, A)
@test isapprox(norm(X - X_)/norm(X), 0f0, atol=1e-6)

Y = CS.forward(X, D, A)
X_ = CS.backward(0f0.*Y, Y, D, A)[3]
@test isapprox(norm(X - X_)/norm(X), 0f0, atol=1e-6)

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

# Initializations
if affine
    CS = AffineCouplingLayerSLIM(nx1, nx2, nx_in, nx_hidden, batchsize, Ψ; logdet=true, permute=true)
else
    CS = AdditiveCouplingLayerSLIM(nx1, nx2, nx_in, nx_hidden, batchsize, Ψ; logdet=true, permute=true)
end

D = randn(Float32, nd1*nd2*nd_in, batchsize) .* 10f0
D0 = randn(Float32, nd1*nd2*nd_in, batchsize) .* 10f0
dD = D - D0

X = randn(Float32, nx1, nx2, nx_in, batchsize)
X0 = randn(Float32, nx1, nx2, nx_in, batchsize)
dX = X - X0

function loss(CS, X, D, A)
    Y, logdet = CS.forward(X, D, A)
    f = -log_likelihood(Y) - logdet
    ΔY = -∇log_likelihood(Y)
    ΔX, ΔD, X_ = CS.backward(ΔY, Y, D, A)
    @test isapprox(norm(X - X_)/norm(X), 0f0; atol=1f-6)
    return f, ΔX, ΔD, CS.RB.W1.grad
end

# Gradient test for input
f0, gX, gD = loss(CS, X0, D0, A)[1:3]
h = 0.1f0
maxiter = 6
err1 = zeros(Float32, maxiter)
err2 = zeros(Float32, maxiter)

print("\nGradient test slim layer input\n")
for j=1:maxiter
    f = loss(CS, X0 + h*dX, D0 + h*dD, A)[1]
    err1[j] = abs(f - f0)
    err2[j] = abs(f - f0  - h*dot(gX, dX) - h*dot(gD, dD))
    print(err1[j], "; ", err2[j], "\n")
    global h = h/2f0
end

@test isapprox(err1[end] / (err1[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err2[end] / (err2[1]/4^(maxiter-1)), 1f0; atol=1f1)

# Gradient test for weights
if affine
    CS = AffineCouplingLayerSLIM(nx1, nx2, nx_in, nx_hidden, batchsize, Ψ; logdet=true, permute=false)
    CS0 = AffineCouplingLayerSLIM(nx1, nx2, nx_in, nx_hidden, batchsize, Ψ; logdet=true, permute=false)
else
    CS = AdditiveCouplingLayerSLIM(nx1, nx2, nx_in, nx_hidden, batchsize, Ψ; logdet=true, permute=false)
    CS0 = AdditiveCouplingLayerSLIM(nx1, nx2, nx_in, nx_hidden, batchsize, Ψ; logdet=true, permute=false)
end
CS.RB.W1.data *= 10; CS0.RB.W1.data *= 10     # make weights larger
CSini = deepcopy(CS0)
dW = CS.RB.W1.data - CS0.RB.W1.data   # just test for 2 parameters

f0, gX, gD, gW = loss(CS0, X, D, A)
h = 0.1f0
maxiter = 4
err3 = zeros(Float32, maxiter)
err4 = zeros(Float32, maxiter)

print("\nGradient test slim layer weights\n")
for j=1:maxiter
    CS0.RB.W1.data = CSini.RB.W1.data + h*dW
    f = loss(CS0, X, D, A)[1]
    err3[j] = abs(f - f0)
    err4[j] = abs(f - f0 - h*dot(dW, gW))
    print(err3[j], "; ", err4[j], "\n")
    global h = h/2f0
end

@test isapprox(err3[end] / (err3[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err4[end] / (err4[1]/4^(maxiter-1)), 1f0; atol=1f1)


text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder, Pkg

name = "librdkafka"
version = v"1.8.2"

# Collection of sources required to complete build
sources = [
    ArchiveSource("https://github.com/edenhill/librdkafka/archive/v$(version).tar.gz",
                  "6a747d293a7a4613bd2897e28e8791476fbe1ae7361f2530a876e0fd483482a6"),
]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir/librdkafka*
if [[ "${target}" != *-freebsd* ]]; then
    rm -f /opt/${target}/${target}/sys-root/usr/lib/libcrypto.*
    rm -f /opt/${target}/${target}/sys-root/usr/lib/libssl.*
    rm -f /opt/${target}/${target}/sys-root/usr/lib/libsasl2.*
fi
mkdir build
cd build/
cmake -DCMAKE_INSTALL_PREFIX=${prefix} \
    -DCMAKE_TOOLCHAIN_FILE=${CMAKE_TARGET_TOOLCHAIN} \
    -DCMAKE_BUILD_TYPE=Release \
    -DRDKAFKA_BUILD_EXAMPLES=OFF \
    -DRDKAFKA_BUILD_TESTS=OFF \
    ..
make -j${nproc}
make install
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = expand_cxxstring_abis(supported_platforms())

# The products that we will ensure are always built
products = [
    LibraryProduct("librdkafka", :librdkafka)
]

# Dependencies that must be installed before this package can be built
dependencies = [
    Dependency(PackageSpec(name="Lz4_jll", uuid="5ced341a-0733-55b8-9ab6-a4889d929147"))
    Dependency(PackageSpec(name="Zstd_jll", uuid="3161d3a3-bdf6-5164-811a-617609db77b4"))
    Dependency(PackageSpec(name="Zlib_jll", uuid="83775a58-1f1d-513f-b197-d71354ab007a"))
    Dependency(PackageSpec(name="OpenSSL_jll", uuid="458c3c95-2e84-50aa-8efc-19380b2a3a95"))
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies;
               julia_compat="1.6")

max_stars_repo_name: MichelJuillard/Yggdrasil

max_issues_repo_name: MichelJuillard/Yggdrasil

max_forks_repo_name: MichelJuillard/Yggdrasil


text:
using Base.Test

import Celeste.Model: patch_ctrs_pix, patch_radii_pix

function test_tile_image()
    blob, ea, three_bodies = gen_three_body_dataset()
    img = blob[3]
    tile_width = 20
    img.epsilon_mat = rand(size(img.pixels))
    img.iota_vec = rand(size(img.pixels, 1))
    tiles = Model.TiledImage(img; tile_width=tile_width).tiles
    @test size(tiles) == (
        ceil(Int, img.H    / tile_width),
        ceil(Int, img.W / tile_width))
    for tile in tiles
        @test tile.b == img.b
        @test tile.pixels == img.pixels[tile.h_range, tile.w_range]
        @test tile.epsilon_mat[2,3] == img.epsilon_mat[tile.h_range, tile.w_range][2,3]
        @test tile.iota_vec[3] == img.iota_vec[tile.h_range][3]
    end

    tile = tiles[2, 2]
    h_width, w_width = size(tile.pixels)
    for h in 1:h_width, w in 1:w_width
        @test tile.pixels[h, w] == img.pixels[tile.h_range[h], tile.w_range[w]]
    end
end

function test_local_sources()
    # Coarse test that local_sources gets the right objects.

    srand(1)
    blob0 = SampleData.load_stamp_blob(datadir, "164.4311-39.0359_2kpsf")
    for b in 1:5
        blob0[b].H, blob0[b].W = 112, 238
        blob0[b].wcs = SampleData.wcs_id
    end

    three_bodies = [
        sample_ce([4.5, 3.6], false),
        sample_ce([60.1, 82.2], true),
        sample_ce([71.3, 100.4], false),
    ]

    blob = Synthetic.gen_blob(blob0, three_bodies)
    tiled_images = TiledImage[TiledImage(img; tile_width=20) for img in blob]

    tile = ImageTile(1, 1, blob[3], 1000)

    ea = make_elbo_args(tiled_images, three_bodies; patch_radius_pix=20.)
    @test ea.S == 3

    patches = vec(ea.patches[:, 3])
    subset1000 = Model.get_local_sources(tile, patch_ctrs_pix(patches),
                                             patch_radii_pix(patches))
    @test subset1000 == [1,2,3]

    tile_width = 10
    tile = ImageTile(1, 1, blob[3], tile_width)
    make_elbo_args(tiled_images, three_bodies; patch_radius_pix=20.)

    patches = vec(ea.patches[:, 3])
    subset10 = Model.get_local_sources(tile, patch_ctrs_pix(patches),
                                           patch_radii_pix(patches))
    @test subset10 == [1]

    last_tile = ImageTile(11, 24, blob[3], tile_width)
    make_elbo_args(tiled_images, three_bodies; patch_radius_pix=20.)

    patches = vec(ea.patches[:, 3])
    last_subset = Model.get_local_sources(last_tile,
                                              patch_ctrs_pix(patches),
                                              patch_radii_pix(patches))
    @test length(last_subset) == 0

    pop_tile = ImageTile(7, 9, blob[3], tile_width)
    make_elbo_args(tiled_images, three_bodies; patch_radius_pix=20.)

    patches = vec(ea.patches[:, 3])
    pop_subset = Model.get_local_sources(pop_tile, patch_ctrs_pix(patches),
                                             patch_radii_pix(patches))

    @test pop_subset == [2,3]
end

function test_local_sources_2()
    # Check that a larger blob gets the same number of objects
    # as a smaller blob.  (This is useful to check edge cases of
    # the polygon logic.)

    srand(1)
    blob0 = SampleData.load_stamp_blob(datadir, "164.4311-39.0359_2kpsf")
    one_body = [sample_ce([50., 50.], true),]
    for b in 1:5 blob0[b].wcs = SampleData.wcs_id end

    for b in 1:5 blob0[b].H, blob0[b].W = 100, 100 end
    small_blob = Synthetic.gen_blob(blob0, one_body)

    for b in 1:5 blob0[b].H, blob0[b].W = 400, 400 end
    big_blob = Synthetic.gen_blob(blob0, one_body)

    ea_small = make_elbo_args(small_blob, one_body, patch_radius_pix=35.)
    small_source_tiles =
      [ sum([ length(s) > 0 for s in source ]) for source in ea_small.tile_source_map ]

    ea_big = make_elbo_args(big_blob, one_body, patch_radius_pix=35.)
    big_source_tiles =
      [ sum([ length(s) > 0 for s in source ]) for source in ea_big.tile_source_map ]

    @test all(big_source_tiles .== small_source_tiles)
end

function test_local_sources_3()
    # Test local_sources using world coordinates.

    srand(1)
    test_b = 3 # Will test using this band only
    pix_loc = Float64[50., 50.]
    blob0 = SampleData.load_stamp_blob(datadir, "164.4311-39.0359_2kpsf")
    body_loc = WCS.pix_to_world(blob0[test_b].wcs, pix_loc)
    one_body = [sample_ce(body_loc, true),]

    # Get synthetic blobs but with the original world coordinates.
    for b in 1:5 blob0[b].H, blob0[b].W = 100, 100 end
    blob = Synthetic.gen_blob(blob0, one_body)
    for b in 1:5 blob[b].wcs = blob0[b].wcs end

    tile_width = 1
    patch_radius_pix = 5.

    ea = make_elbo_args(blob, one_body, patch_radius_pix=patch_radius_pix)

    # Source should be present
    tile = ImageTile(
        round(Int, pix_loc[1] / tile_width),
        round(Int, pix_loc[2] / tile_width),
        blob[test_b],
        tile_width)

    patches = vec(ea.patches[:,test_b])
    @test Model.get_local_sources(tile, patch_ctrs_pix(patches),
                                  patch_radii_pix(patches)) == [1]

    # Source should not match when you're 1 tile and a half away along the diagonal plus
    # the pixel radius from the center of the tile.
    tile = ImageTile(
        ceil(Int, (pix_loc[1] + 1.5 * tile_width * sqrt(2) +
                patch_radius_pix) / tile_width),
        round(Int, pix_loc[2] / tile_width),
        blob[test_b],
        tile_width)
    @test Model.get_local_sources(tile, patch_ctrs_pix(patches),
                                  patch_radii_pix(patches)) == []

    tile = ImageTile(
        round(Int, (pix_loc[1]) / tile_width),
        ceil(Int, (pix_loc[2]  + 1.5 * tile_width * sqrt(2) +
                           patch_radius_pix) / tile_width),
        blob[test_b],
        tile_width)
    @test Model.get_local_sources(tile, patch_ctrs_pix(patches),
                                      patch_radii_pix(patches)) == []

end

function test_sky_noise_estimates()
    blobs = Array(Vector{Image}, 2)
    blobs[1], ea, three_bodies = gen_three_body_dataset()  # synthetic
    blobs[2] = SampleData.load_stamp_blob(datadir, "164.4311-39.0359_2kpsf")  # real

    for blob in blobs
        for b in 1:5
            sdss_sky_estimate = median(blob[b].epsilon_mat) * median(blob[b].iota_vec)
            crude_estimate = median(blob[b].pixels)
            @test_approx_eq_eps sdss_sky_estimate / crude_estimate 1. .3
        end
    end
end

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

test_tile_image()
test_sky_noise_estimates()
test_local_sources()
test_local_sources_2()
test_local_sources_3()


text:
using DataFrames, LinearAlgebra, MixedModels, Random, SparseArrays, StatsModels, Test
using MixedModels: mulαβ!

@testset "mulαβ!" begin
    for (m, p, n, q, k) in (
                            (10, 0.7, 5, 0.3, 15),
                            (100, 0.01, 100, 0.01, 20),
                            (100, 0.1, 100, 0.2, 100),
                           )
        a = sprand(m, n, p)
        b = sprand(n, k, q)
        as = sparse(a')
        bs = sparse(b')
        c = zeros(m, k)
        ab = a * b
        arbt = Array(b')
        aab = Array(a) * Array(b)
        @test aab ≈ mulαβ!(c, a, bs', true, true)
        @test aab ≈ mulαβ!(c, a, bs')
        @test aab ≈ mulαβ!(c, a, arbt')
        @test aab ≈ mulαβ!(c, a, arbt')
        @test aab ≈ mulαβ!(fill!(c, 0.0), a, arbt', true, true)
        @test maximum(abs, mulαβ!(c, a, arbt', -1.0, true)) ≤ sqrt(eps())
        @test maximum(abs.(ab - aab)) < 100*eps()
        @test a*bs' == ab
        @test as'*b == ab
        @test as'*bs' == ab
        f = Diagonal(rand(n))
        @test Array(a*f) == Array(a)*f
        @test Array(f*b) == f*Array(b)
    end
end

@testset "reweight!" begin
    rng = Random.MersenneTwister(1234321)
    df = (Y = randn(rng, 400), A = repeat(['N','Y'], outer=200),
        G = repeat('A':'T', inner = 2, outer=10), H = repeat('a':'j', inner=40))
    lmm1 = fit!(LinearMixedModel(@formula(Y ~ 1+A+(1+A|G)+(1+A|H)), df,
        wts  = ones(400)))
    @test loglikelihood(lmm1) ≈ -578.9080978272708
    MixedModels.reweight!(lmm1, ones(400))
    @test loglikelihood(fit!(lmm1)) ≈ -578.9080978272708
end

@testset "lmulλ!" begin
    gendata(n::Int, ng::Int) = gendata(MersenneTwister(42), n, ng)

    function gendata(rng::AbstractRNG, n::Int, ng::Int)
        df = DataFrame(Y = randn(rng, n),
                       X = rand(rng, n),
                       G = rand(rng, 1:ng, n),
                       H = rand(rng, 1:ng, n))
        categorical!(df, [:G, :H])
    end

    @testset "Adjoint{T,ReMat{T,1}}, BlockedSparse{T,1,2}" begin
        # this is an indirect test of lmulΛ! for a blocking structure found in
        # an example in MixedModels.jl#123
        df = gendata(10000, 500)
        f = @formula(Y ~ (1 + X|H) + (1|G))
        m500 = fit!(LinearMixedModel(f, df))
        # the real test here isn't in the theta comparison but in that the fit
        # completes successfully
        @test m500.theta ≈ [0.07797345952252123, -0.17640571417349787, 0, 0]
    end
end


text:
# ---
# title: 620. Not Boring Movies
# id: problem620
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories:
# link: <https://leetcode.com/problems/not-boring-movies/description/>
# hidden: true
# ---
#
# X city opened a new cinema, many people would like to go to this cinema. The
# cinema also gives out a poster indicating the movies' ratings and
# descriptions.
#
# Please write a SQL query to output movies with an odd numbered ID and a
# description that is not 'boring'. Order the result by rating.
#
#
#
# For example, table `cinema`:
#
#
#
#     +---------+-----------+--------------+-----------+
#     |   id    | movie     |  description |  rating   |
#     +---------+-----------+--------------+-----------+
#     |   1     | War       |   great 3D   |   8.9     |
#     |   2     | Science   |   fiction    |   8.5     |
#     |   3     | irish     |   boring     |   6.2     |
#     |   4     | Ice song  |   Fantacy    |   8.6     |
#     |   5     | House card|   Interesting|   9.1     |
#     +---------+-----------+--------------+-----------+
#
#
# For the example above, the output should be:
#
#
#
#     +---------+-----------+--------------+-----------+
#     |   id    | movie     |  description |  rating   |
#     +---------+-----------+--------------+-----------+
#     |   5     | House card|   Interesting|   9.1     |
#     |   1     | War       |   great 3D   |   8.9     |
#     +---------+-----------+--------------+-----------+
#
#
#
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
using Revise

haskey(ENV, "HOST") || (ENV["HOST"] = "0.0.0.0")
haskey(ENV, "GENIE_ENV") || (ENV["GENIE_ENV"] = "dev")

### EARLY BIND TO PORT FOR HOSTS WITH TIMEOUT ###

import Sockets

const EARLYBINDING = if haskey(ENV, "EARLYBIND") && lowercase(ENV["EARLYBIND"]) == "true" && haskey(ENV, "PORT")
  printstyled("\nEarly binding to host $(ENV["HOST"]) and port $(ENV["PORT"]) \n", color = :light_blue, bold = true)
  try
    Sockets.listen(parse(Sockets.IPAddr, ENV["HOST"]), parse(Int, ENV["PORT"]))
  catch ex
    @show ex

    printstyled("\nFailed early binding!\n", color = :red, bold = true)
    nothing
  end
else
  nothing
end

### OFF WE GO! ###

using Genie

const ROOT_PATH = pwd()

push!(LOAD_PATH, ROOT_PATH, "src")

Genie.load(context = @__MODULE__)
Genie.run(server = EARLYBINDING)


text:
lines = readlines("input.txt")
numbers = parse.(Int, split(lines[1], ','))
fields_s = reshape(lines[2:end], 6, :)[2:end, :]
fields = parse.(Int, reshape(reduce(vcat, split.(fields_s)), 5,5,:))
marks = falses(size(fields))
stack = Vector()

for n = numbers
    marks[findall(x -> x == n, fields)] .= true

    rowhits = mapslices(prod, marks, dims=[1])
    if any(rowhits) # rows
        for hit = findall(isone, rowhits)
            board = hit[3]
            push!(stack, (fields[:,:,board], marks[:,:,board], n))
            fields[:,:,board] = marks[:,:,board] .= 0
        end
    end

    colhits = mapslices(prod, marks, dims=[2])
    if any(colhits) # cols
        for hit = findall(isone, colhits)
            board = hit[3]
            push!(stack, (fields[:,:,board], marks[:,:,board], n))
            fields[:,:,board] = marks[:,:,board] .= 0
        end
    end
end

(field, mark, n) = last(stack) # last board to finish
println(sum(@. field * !mark) * n)


text:
# see benchmark/fsortperm for choice of algorithm
# the best algorithms depends on the RADIX_SIZE
# most of the time it's between 10 and 11
const VT = Pair{UInt32,UInt32}
const VTV = Vector{VT}

function fsortperm(a::AbstractVector{<:Integer}; counting_sort_cutoff = 2048, RADIX_SIZE = 10, alg = :lsd, rev = false)
    minval, maxval = extrema(a)
    rangelen = maxval - minval + 1

    if rangelen <= counting_sort_cutoff
        return fsortperm_int_range_counting_sort(a, rangelen, minval, rev = rev)
    # elseif alg == :msd
    #     # return Int.(fsortperm_msd_hybrid(a, rangelen, minval, RADIX_SIZE, counting_sort_cutoff = -1))
    #     return Int.(fsortperm_int_range_lsd(a, rangelen, minval, RADIX_SIZE))
    # elseif alg == :lsd
    #     return Int.(fsortperm_int_range_lsd(a, rangelen, minval, RADIX_SIZE))
    else
        return fsortperm_int_range_lsd(a, rangelen, minval, RADIX_SIZE, rev = rev)
    end
end

################################################################################
# MSD radix sort and counting sort hybrid
#
# but I found that turning off counting sort (achieved via counting_sort_cutoff = -1)
# achieved better performance
################################################################################
function fsortperm_msd_hybrid(a, rangelen, minval, RADIX_SIZE, counting_sort_cutoff = -1)
    # println(RADIX_SIZE," ", now())
    # @assert 2^32 > rangelen
    # @assert 2^32 >= length(a)
    vs::VTV = Vector{VT}(length(a))
    @inbounds for i in eachindex(a)
        vs[i] = Pair(UInt32(i),UInt32(a[i] - minval))
    end
    # @time @inbounds vs = [Pair(UInt32(i),UInt32(ai)) for (i, ai) in enumerate(a)]

    iters = Int(ceil(log(2, rangelen)/RADIX_SIZE))
    # @assert iters <= ceil(Integer, sizeof(typeof(vs[1].second))*8/RADIX_SIZE)

    _fsortperm_msd_hybrid(vs, 1, length(vs), rangelen, 0, RADIX_SIZE, iters, counting_sort_cutoff)
end

function _fsortperm_msd_hybrid(vs, lo, hi, rangelen, minval, RADIX_SIZE, iters, ts = similar(vs), counting_sort_cutoff::Int = -1)
    len = hi-lo+1
    if len  == 1
        return [Int(vs[lo].first)]
    elseif rangelen <= counting_sort_cutoff
        # println("iters $iters lo:hi $(lo:hi) rangelen $rangelen minval $minval")
        av = [a1.second for a1 in vs[lo:hi]]
        # println("extrema $(extrema(Int.(av)))")
        av_sortperm = fsortperm_int_range_counting_sort(av, rangelen, minval)
        # issorted([Int(vs[j+lo-1].second) for j in av_sortperm])
        res = [Int(vs[j+lo-1].first) for j in av_sortperm]
        # println("ok2")
        return res
    end

    RADIX_MASK = UInt32(1<<RADIX_SIZE-1)

    # Init
    bin = zeros(UInt32, 2^RADIX_SIZE)

    # Histogram for each element, radix
    @inbounds for i = lo:hi
        v = vs[i].second - minval
        idx = Int((v >> (iters-1)*RADIX_SIZE) & RADIX_MASK) + 1
        @inbounds bin[idx] += 1
    end

    # Sort!

    # check for degenerate case
    v = vs[hi].second - minval
    idx = Int((v >> (iters-1)*RADIX_SIZE) & RADIX_MASK) + 1

    # are all values the same at this radix?
    # if bin[idx,iters] == len;  continue;  end
    if bin[idx] == len
        if iters == 1
            return [Int(vsi.first) for vsi in vs[lo:hi]]
        else
            return _fsortperm_msd_hybrid(vs, lo, hi, rangelen, minval, RADIX_SIZE, iters-1, ts, counting_sort_cutoff = counting_sort_cutoff)
        end
    end

    # cbin = cumsum(bin[:])
    cbin  = copy(bin)
    cumsum!(bin, bin)
    cumsumbin = copy(bin)

    # copy the elements to the temporary array first
    @inbounds for i in lo:hi
        ts[i] = vs[i]
    end

    # now assign to original in the right order
    @inbounds for i in lo:hi
        v = ts[i].second - minval
        idx = Int((v >> (iters-1)*RADIX_SIZE) & RADIX_MASK) + 1
        ci = bin[idx]
        vs[lo - 1 + ci] = ts[i]
        bin[idx] -= 1
    end

    if iters == 1
        return [Int(vsi.first) for vsi in vs[lo:hi]]
    end

    # now that it's finished sorting sort each chunk recursive now
    res = Vector{Int}(length(vs))
    start_now = 1
    new_rangelen = 2^((iters-1)*RADIX_SIZE)
    # if false
    #     i = 1
    #     lo = 1
    #     hi = cumsumbin[1]
    #     rangelen = new_rangelen
    #     minval = 1
    #     iters = 1

    #     i = 2
    #     lo = cumsumbin[1] + 1
    #     hi = cumsumbin[2]
    #     hi |> Int
    #     minval = 1+(i-1)new_rangelen

    #     i = 4
    #     lo = cumsumbin[i-1] + 1
    #     hi = cumsumbin[3]
    #     hi |> Int
    #     minval = 1+(i-1)new_rangelen
    #     iters = 1
    # end

    for i = 1:length(cumsumbin)
        if cbin[i] != 0
            # println("iters $iters ith count $i")
            new_minval = minval+(i-1)new_rangelen
            res[start_now:cumsumbin[i]] = _fsortperm_msd_hybrid(vs, start_now + lo - 1, cumsumbin[i] + lo - 1, new_rangelen, new_minval, RADIX_SIZE, iters-1, ts, counting_sort_cutoff = counting_sort_cutoff)
            start_now = cumsumbin[i] + 1
            # println("Ok")
        end
    end
    res
end

################################################################################
# Counting sort; not recommended if the range is greater than 2048
################################################################################
function fsortperm_int_range_counting_sort(a, rangelen, minval; rev = false)
    cnts = fill(0, rangelen)

    # create a count first
    @inbounds for ai in a
        cnts[ai - minval + 1] +=  1
    end

    # create cumulative sum
    cumsum!(cnts, cnts)

    la = length(a)
    res = Vector{Int}(la)
    if rev
        @inbounds for i in la:-1:1
            ai = a[i] - minval + 1
            c = cnts[ai]
            cnts[ai] -= 1
            res[la - c + 1] = i
        end
    else
        @inbounds for i in la:-1:1
            ai = a[i] - minval + 1
            c = cnts[ai]
            cnts[ai] -= 1
            res[c] = i
        end
    end

    res
end

################################################################################
# LSD radix sort
# based on sortperm_int_range_p1
# see https://discourse.julialang.org/t/ironic-observation-about-sort-and-sortperm-speed-for-small-intergers-vs-r/8715/31?u=xiaodai
################################################################################
function fsortperm_int_range_lsd(a::Vector{T}, rangelen, minval, RADIX_SIZE; rev = false) where T
    # println(RADIX_SIZE," ", now())
    # @assert 2^32 > rangelen
    # @assert 2^32 >= length(a)

    vs = Vector{Pair{UInt32, T}}(undef, length(a))
    if rev
        maxval = maximum(a)
        @inbounds for i in eachindex(a)
            vs[i] = Pair(UInt32(i), maxval - a[i])
        end
    else
        @inbounds for i in eachindex(a)
            vs[i] = Pair(UInt32(i), a[i] - minval)
        end
    end
    # @time @inbounds vs = [Pair(UInt32(i),UInt32(ai)) for (i, ai) in enumerate(a)]
    sorttwo!(vs, rangelen, 0, RADIX_SIZE)
    [Int(vs1.first) for vs1 in vs]
end

function fsortandperm_int_range_lsd(a::Vector{T}, rangelen, minval, RADIX_SIZE; rev = false) where T
    vs = Vector{Pair{UInt32, T}}(length(a))
    if rev
        maxval = maximum(a)
        @inbounds for i in eachindex(a)
            vs[i] = Pair(UInt32(i), maxval - a[i])
        end
    else
        @inbounds for i in eachindex(a)
            vs[i] = Pair(UInt32(i), a[i] - minval)
        end
    end

    sorttwo!(vs, rangelen, 0, RADIX_SIZE)
end

# this does the actual work of sorting
function sorttwo!(vs, rangelen::T, minval::S, RADIX_SIZE) where {T <: Integer, S <: Integer}
    RADIX_MASK = UInt32(1<<RADIX_SIZE-1)
    ts = similar(vs)

    # Init
    lo = 1
    hi = length(vs)
    iters = Int(ceil(log(2, rangelen)/RADIX_SIZE))

    # println("iters: $iters")
    # @assert iters <= ceil(Integer, sizeof(typeof(vs[1].second))*8/RADIX_SIZE)
    bin = zeros(UInt32, 2^RADIX_SIZE, iters)

    # Histogram for each element, radix
    @inbounds for i = lo:hi
        v = vs[i].second - minval
        for j = 1:iters
            idx = Int((v >> (j-1)*RADIX_SIZE) & RADIX_MASK) + 1
            @inbounds bin[idx,j] += 1
        end
    end

    # Sort!
    swaps = 0
    len = hi-lo+1
    @inbounds for j = 1:iters
        # Unroll first data iteration, check for degenerate case
        v = vs[hi].second - minval
        idx = Int((v >> (j-1)*RADIX_SIZE) & RADIX_MASK) + 1

        # are all values the same at this radix?
        if bin[idx,j] == len;  continue;  end

        cbin = cumsum(bin[:,j])
        ci = cbin[idx]
        ts[ci] = vs[hi]
        cbin[idx] -= 1

        # Finish the loop...
        @inbounds for i in hi-1:-1:lo
            v = vs[i].second - minval
            idx = Int((v >> (j-1)*RADIX_SIZE) & RADIX_MASK) + 1
            ci = cbin[idx]
            ts[ci] = vs[i]
            cbin[idx] -= 1
        end
        vs,ts = ts,vs
        swaps += 1
    end

    if isodd(swaps)
        vs,ts = ts,vs
        @inbounds for i = lo:hi
            vs[i] = ts[i]
        end
    end

    vs
end


text:
# SteinGaussianWeightedKernel
#
# These represent a family of kernels of the form
#
# k(x, y) = Q(x,y) exp(-||x-y||^2/(2*beta))
#
# where Q(x,y) is such that k is indeed a valid kernel.
# This kernel is combined with the score function u(x) in order
# generate the mean-zero RKHS.

abstract type SteinGaussianWeightedKernel <: SteinKernel end

### METHODS TO IMPLEMENT

# the function that evaluates Q(x,y)
function Q(kernel::SteinGaussianWeightedKernel, x::Array{Float64,1}, y::Array{Float64,1})
    error("Must implement the method Q")
end

# the function that implements d/dxj Q(x,y)
function dxj_Q(kernel::SteinGaussianWeightedKernel, x::Array{Float64,1}, y::Array{Float64,1}, j::Int64)
    error("Must implement the method dxj_Q")
end

# the function that implements d^2/dxjdyj Q(x,y)
function dxjyj_Q(kernel::SteinGaussianWeightedKernel, x::Array{Float64,1}, y::Array{Float64,1}, j::Int64)
    error("Must implement the method dxjyj_Q")
end

### INDUCED METHODS

function gaussian(kernel::SteinGaussianWeightedKernel, x::Array{Float64,1}, y::Array{Float64,1})
    exp(-1 * norm(x-y,2)^2 / (2 * kernel.beta))
end

# the function that implements d/dyj Q(x,y)
# where we assume Q(x,y) = Q(y,x)
function dyj_Q(kernel::SteinGaussianWeightedKernel, x::Array{Float64,1}, y::Array{Float64,1}, j::Int64)
    dxj_Q(kernel, y, x, j)
end

# defintion of the base kernel
function k(kernel::SteinGaussianWeightedKernel, x::Array{Float64, 1}, y::Array{Float64, 1})
    Q(kernel, x, y) * gaussian(kernel, x, y)
end

# implements \grad_x k(x,y) = gaussian(ker, x, y) * [\grad_x Q(x,y) + (y - x) * Q(x,y) / beta]
function gradxk(kernel::SteinGaussianWeightedKernel, x::Array{Float64, 1}, y::Array{Float64, 1})
    d = length(x)
    base = [dxj_Q(kernel, x, y, j) + (y[j] - x[j]) * Q(kernel, x, y) / kernel.beta for j=1:d]
    base .* gaussian(kernel, x, y)
end

# implements <\grad_y, \grad_x k(x,y)>
function gradxyk(kernel::SteinGaussianWeightedKernel, x::Array{Float64, 1}, y::Array{Float64, 1})
    d = length(x)
    Qterm = Q(kernel, x, y)
    gaussianterm = gaussian(kernel, x, y)

    gradxQ = [dxj_Q(kernel, x, y, j) for j=1:d]
    gradyQ = [dyj_Q(kernel, x, y, j) for j=1:d]

    order0 = (d/kernel.beta - norm(x-y,2)^2 / kernel.beta^2) * Qterm * gaussianterm
    order1 = (gaussianterm / kernel.beta) * dot(gradxQ - gradyQ, x-y)
    order2 = gaussianterm * sum([dxjyj_Q(kernel, x, y, j) for j=1:d])

    order0 + order1 + order2
end


text:
@defcomp ch4emissions begin
    region = Index()

    e_globalCH4emissions = Variable(index=[time], unit="Mtonne/year")
    e0_baselineCH4emissions = Parameter(index=[region], unit="Mtonne/year")
    e_regionalCH4emissions = Variable(index=[time,region], unit="Mtonne/year")
    er_CH4emissionsgrowth = Parameter(index=[time,region], unit="%")

    # read in counterfactual GDP in absence of growth effects (gdp_leveleffects) and actual GDP
    gdp = Parameter(index=[time, region], unit="\$M")
    gdp_leveleffect   = Parameter(index=[time, region], unit="\$M")
    emfeed_emissionfeedback = Parameter(unit="none", default=1.)

    function run_timestep(p, v, d, t)

        # eq.4 in Hope (2006) - regional CH4 emissions as % change from baseline
        for r in d.region
            v.e_regionalCH4emissions[t,r] = p.er_CH4emissionsgrowth[t,r] * p.e0_baselineCH4emissions[r] / 100

            # rescale emissions based on GDP deviation from original scenario pathway
            if p.emfeed_emissionfeedback == 1.
                v.e_regionalCH4emissions[t,r] = v.e_regionalCH4emissions[t,r] * (p.gdp[t,r] / p.gdp_leveleffect[t,r])
            end
        end

        # eq. 5 in Hope (2006) - global CH4 emissions are sum of regional emissions
        v.e_globalCH4emissions[t] = sum(v.e_regionalCH4emissions[t,:])
    end
end


text:
# This file contains default implementations for the `MOI.copy_to` function that can be used by a model.

"""
    automatic_copy_to(dest::MOI.ModelLike, src::MOI.ModelLike;
                      copy_names::Bool=true)

Use [`Utilities.supports_default_copy_to`](@ref) and
[`Utilities.supports_allocate_load`](@ref) to automatically choose between
[`Utilities.default_copy_to`](@ref) or [`Utilities.allocate_load`](@ref) to
apply the copy operation.
"""
function automatic_copy_to(dest::MOI.ModelLike, src::MOI.ModelLike;
                           copy_names::Bool=true)
    if supports_default_copy_to(dest, copy_names)
        default_copy_to(dest, src, copy_names)
    elseif supports_allocate_load(dest, copy_names)
        allocate_load(dest, src, copy_names)
    else
        error("Model $(typeof(dest)) does not support copy",
              copy_names ? " with names" : "", ".")
    end
end

"""
    supports_default_copy_to(model::ModelLike, copy_names::Bool)

Return a `Bool` indicating whether the model `model` supports
[`default_copy_to(model, src, copy_names=copy_names)`](@ref) if all the
attributes set to `src` and constraints added to `src` are supported by `model`.

This function can be used to determine whether a model can be loaded into
`model` incrementally or whether it should be cached and copied at once instead.
This is used by JuMP to determine whether to add a cache or not in two
situations:
1. A first cache can be used to store the model as entered by the user as well
   as the names of variables and constraints. This cache is created if this
   function returns `false` when `copy_names` is `true`.
2. If bridges are used, then a second cache can be used to store the bridged
   model with unnamed variables and constraints. This cache is created if this
   function returns `false` when `copy_names` is `false`.

## Examples

If [`MathOptInterface.set`](@ref), [`MathOptInterface.add_variable`](@ref) and
[`MathOptInterface.add_constraint`](@ref) are implemented for a model of type
`MyModel` and names are supported, then [`MathOptInterface.copy_to`](@ref) can
be implemented as
```julia
MOI.Utilities.supports_default_copy_to(model::MyModel, copy_names::Bool) = true
function MOI.copy_to(dest::MyModel, src::MOI.ModelLike; kws...)
    return MOI.Utilities.automatic_copy_to(dest, src; kws...)
end
```
The [`Utilities.automatic_copy_to`](@ref) function automatically redirects to
[`Utilities.default_copy_to`](@ref).

If names are not supported, simply change the first line by
```julia
MOI.supports_default_copy_to(model::MyModel, copy_names::Bool) = !copy_names
```
The [`Utilities.default_copy_to`](@ref) function automatically throws an helpful
error in case `copy_to` is called with `copy_names` equal to `true`.
"""
supports_default_copy_to(model::MOI.ModelLike, copy_names::Bool) = false

struct IndexMap
    varmap::Dict{MOI.VariableIndex, MOI.VariableIndex}
    conmap::Dict{MOI.ConstraintIndex, MOI.ConstraintIndex}
end
IndexMap() = IndexMap(Dict{MOI.VariableIndex, MOI.VariableIndex}(),
                 Dict{MOI.ConstraintIndex, MOI.ConstraintIndex}())
Base.getindex(idxmap::IndexMap, vi::MOI.VariableIndex) = idxmap.varmap[vi]
function Base.getindex(idxmap::IndexMap, ci::MOI.ConstraintIndex{F, S}) where {F, S}
    idxmap.conmap[ci]::MOI.ConstraintIndex{F, S}
end

Base.setindex!(idxmap::IndexMap, vi1::MOI.VariableIndex, vi2::MOI.VariableIndex) = Base.setindex!(idxmap.varmap, vi1, vi2)
function Base.setindex!(idxmap::IndexMap, ci1::MOI.ConstraintIndex{F, S}, ci2::MOI.ConstraintIndex{F, S}) where {F, S}
    Base.setindex!(idxmap.conmap, ci1, ci2)
end

Base.delete!(idxmap::IndexMap, vi::MOI.VariableIndex) = delete!(idxmap.varmap, vi)
Base.delete!(idxmap::IndexMap, ci::MOI.ConstraintIndex) = delete!(idxmap.conmap, ci)

Base.keys(idxmap::IndexMap) = Iterators.flatten((keys(idxmap.varmap), keys(idxmap.conmap)))

"""
    pass_attributes(dest::MOI.ModelLike, src::MOI.ModelLike, copy_names::Bool, idxmap::IndexMap, passattr!::Function=MOI.set)

Pass the model attributes from the model `src` to the model `dest` using `canpassattr` to check if the attribute can be passed and `passattr!` to pass the attribute. Does not copy `Name` if `copy_names` is `false`.

    pass_attributes(dest::MOI.ModelLike, src::MOI.ModelLike, copy_names::Bool, idxmap::IndexMap, vis_src::Vector{MOI.VariableIndex}, passattr!::Function=MOI.set)

Pass the variable attributes from the model `src` to the model `dest` using `canpassattr` to check if the attribute can be passed and `passattr!` to pass the attribute. Does not copy `VariableName` if `copy_names` is `false`.

    pass_attributes(dest::MOI.ModelLike, src::MOI.ModelLike, copy_names::Bool, idxmap::IndexMap, cis_src::Vector{MOI.ConstraintIndex{F, S}}, passattr!::Function=MOI.set) where {F, S}

Pass the constraint attributes of `F`-in-`S` constraints from the model `src` to the model `dest` using `canpassattr` to check if the attribute can be passed and `passattr!` to pass the attribute. Does not copy `ConstraintName` if `copy_names` is `false`.
"""
function pass_attributes end

function pass_attributes(dest::MOI.ModelLike, src::MOI.ModelLike, copy_names::Bool, idxmap::IndexMap, passattr!::Function=MOI.set)
    # Copy model attributes
    attrs = MOI.get(src, MOI.ListOfModelAttributesSet())
    _pass_attributes(dest, src, copy_names, idxmap, attrs, tuple(), tuple(), tuple(), passattr!)
end
function pass_attributes(dest::MOI.ModelLike, src::MOI.ModelLike, copy_names::Bool, idxmap::IndexMap, vis_src::Vector{VI}, passattr!::Function=MOI.set)
    # Copy variable attributes
    attrs = MOI.get(src, MOI.ListOfVariableAttributesSet())
    vis_dest = map(vi -> idxmap[vi], vis_src)
    _pass_attributes(dest, src, copy_names, idxmap, attrs, (VI,), (vis_src,), (vis_dest,), passattr!)
end
function pass_attributes(dest::MOI.ModelLike, src::MOI.ModelLike, copy_names::Bool, idxmap::IndexMap, cis_src::Vector{CI{F, S}}, passattr!::Function=MOI.set) where {F, S}
    # Copy constraint attributes
    attrs = MOI.get(src, MOI.ListOfConstraintAttributesSet{F, S}())
    cis_dest = map(ci -> idxmap[ci], cis_src)
    _pass_attributes(dest, src, copy_names, idxmap, attrs, (CI{F, S},), (cis_src,), (cis_dest,), passattr!)
end

function _pass_attributes(dest::MOI.ModelLike, src::MOI.ModelLike,
                          copy_names::Bool, idxmap::IndexMap, attrs,
                          supports_args, get_args, set_args,
                          pass_attr!::Function=MOI.set)
    for attr in attrs
        @assert MOI.is_copyable(attr)
        if !copy_names && (attr isa MOI.Name || attr isa MOI.VariableName || attr isa MOI.ConstraintName)
            continue
        end
        if attr isa MOI.VariablePrimalStart || attr isa MOI.ConstraintPrimalStart || attr isa MOI.ConstraintDualStart
            # As starting values are simply *hints* for the optimization, not
            # supporting them gives a warning, not an error
            if !MOI.supports(dest, attr, supports_args...)
                Compat.@warn("$attr is not supported by $(typeof(dest)). This ",
                             "information will be discarded.")
                continue
            end
        end
        value = MOI.get(src, attr, get_args...)
        if value !== nothing
            mapped_value = attribute_value_map(idxmap, value)
            pass_attr!(dest, attr, set_args..., mapped_value)
        end
    end
end

"""
    copyconstraints!(dest::MOI.ModelLike, src::MOI.ModelLike, copy_names::Bool, idxmap::IndexMap, ::Type{F}, ::Type{S}) where {F<:MOI.AbstractFunction, S<:MOI.AbstractSet}

Copy the constraints of type `F`-in-`S` from the model `src` the model `dest` and fill `idxmap` accordingly. Does not copy `ConstraintName` if `copy_names` is `false`.
"""
function copyconstraints!(dest::MOI.ModelLike, src::MOI.ModelLike, copy_names::Bool, idxmap::IndexMap, ::Type{F}, ::Type{S}) where {F<:MOI.AbstractFunction, S<:MOI.AbstractSet}
    # Copy constraints
    cis_src = MOI.get(src, MOI.ListOfConstraintIndices{F, S}())
    for ci_src in cis_src
        f_src = MOI.get(src, MOI.ConstraintFunction(), ci_src)
        f_dest = mapvariables(idxmap, f_src)
        s = MOI.get(src, MOI.ConstraintSet(), ci_src)
        ci_dest = MOI.add_constraint(dest, f_dest, s)
        idxmap.conmap[ci_src] = ci_dest
    end

    pass_attributes(dest, src, copy_names, idxmap, cis_src)
end

attribute_value_map(idxmap, f::MOI.AbstractFunction) = mapvariables(idxmap, f)
attribute_value_map(idxmap, attribute_value) = attribute_value
function default_copy_to(dest::MOI.ModelLike, src::MOI.ModelLike)
    Base.depwarn("default_copy_to(dest, src) is deprecated, use default_copy_to(dest, src, true) instead or default_copy_to(dest, src, false) if you do not want to copy names.", :default_copy_to)
    default_copy_to(dest, src, true)
end

"""
    default_copy_to(dest::MOI.ModelLike, src::MOI.ModelLike, copy_names::Bool)

Implements `MOI.copy_to(dest, src)` by adding the variables and then the
constraints and attributes incrementally. The function
[`supports_default_copy_to`](@ref) can be used to check whether `dest` supports
the copying a model incrementally.
"""
function default_copy_to(dest::MOI.ModelLike, src::MOI.ModelLike, copy_names::Bool)
    MOI.empty!(dest)

    idxmap = IndexMap()

    # Copy variables
    vis_src = MOI.get(src, MOI.ListOfVariableIndices())
    for vi in vis_src
        idxmap.varmap[vi] = MOI.add_variable(dest)
    end

    # Copy variable attributes
    pass_attributes(dest, src, copy_names, idxmap, vis_src)

    # Copy model attributes
    pass_attributes(dest, src, copy_names, idxmap)

    # Copy constraints
    for (F, S) in MOI.get(src, MOI.ListOfConstraints())
        # do the rest in copyconstraints! which is type stable
        copyconstraints!(dest, src, copy_names, idxmap, F, S)
    end

    return idxmap
end

# Allocate-Load Interface: 2-pass copy of a MathOptInterface model
# Some solver wrappers (e.g. SCS, ECOS, SDOI) do not supporting copying an optimization model using `MOI.add_constraints`, `MOI.add_variables` and `MOI.set`
# as they first need to figure out some information about a model before being able to pass the problem data to the solver.
#
# During the first pass (called allocate) : the model collects the relevant information about the problem so that
# on the second pass (called load), the constraints can be loaded directly to the solver (in case of SDOI) or written directly into the matrix of constraints (in case of SCS and ECOS).

# To support `MOI.copy_to` using this 2-pass mechanism, implement the allocate-load interface defined below and do:
# MOI.copy_to(dest::ModelType, src::MOI.ModelLike) = MOIU.allocate_load(dest, src)
# In the implementation of the allocate-load interface, it can be assumed that the different functions will the called in the following order:
# 1) `allocate_variables`
# 2) `allocate` and `allocate_constraint`
# 3) `load_variables`
# 4) `load` and `load_constraint`
# The interface is not meant to be used to create new constraints with `allocate_constraint` followed by `load_constraint` after a solve, it is only meant for being used in this order to implement `MOI.copy_to`.

# TODO deprecate this in MOI v0.7
"""
    needs_allocate_load(model::MOI.ModelLike)::Bool
Return a `Bool` indicating whether `model` does not support `add_variables`/`add_constraint`/`set` but supports `allocate_variables`/`allocate_constraint`/`allocate`/`load_variables`/`load_constraint`/`load`.
That is, the allocate-load interface need to be used to copy an model to `model`.
"""
needs_allocate_load(::MOI.ModelLike) = false

"""
    supports_allocate_load(model::MOI.ModelLike, copy_names::Bool)::Bool

Return a `Bool` indicating whether `model` supports
[`allocate_load(model, src, copy_names=copy_names)`](@ref) if all the
attributes set to `src` and constraints added to `src` are supported by `model`.
"""
supports_allocate_load(::MOI.ModelLike, copy_names::Bool) = false

"""
    allocate_variables(model::MOI.ModelLike, nvars::Integer)

Creates `nvars` variables and returns a vector of `nvars` variable indices.
"""
function allocate_variables end

const ALLOCATE_LOAD_NOT_IMPLEMENTED = ErrorException("The Allocate-Load interface is" *
                                                     " not implemented by the model")

"""
    allocate(model::ModelLike, attr::ModelLikeAttribute, value)
    allocate(model::ModelLike, attr::AbstractVariableAttribute, v::VariableIndex, value)
    allocate(model::ModelLike, attr::AbstractConstraintAttribute, c::ConstraintIndex, value)

Informs `model` that `load` will be called with the same arguments after `load_variables` is called.
"""
function allocate(model::MOI.ModelLike, args...)
    MOI.throw_set_error_fallback(model, args...;
                                 error_if_supported=ALLOCATE_LOAD_NOT_IMPLEMENTED)
end

function allocate(model::MOI.ModelLike, attr::Union{MOI.AbstractVariableAttribute, MOI.AbstractConstraintAttribute}, indices::Vector, values::Vector)
    for (index, value) in zip(indices, values)
        allocate(model, attr, index, value)
    end
end

"""
    allocate_constraint(model::MOI.ModelLike, f::MOI.AbstractFunction, s::MOI.AbstractSet)

Returns the index for the constraint to be used in `load_constraint` that will be called after `load_variables` is called.
"""
function allocate_constraint(model::MOI.ModelLike, func::MOI.AbstractFunction,
                             set::MOI.AbstractSet)
    MOI.throw_add_constraint_error_fallback(model, func, set;
                                            error_if_supported=ALLOCATE_LOAD_NOT_IMPLEMENTED)
end

"""
    load_variables(model::MOI.ModelLike, nvars::Integer)

Prepares the `model` for `loadobjective!` and `load_constraint`.
"""
function load_variables end

"""
    load(model::ModelLike, attr::ModelLikeAttribute, value)
    load(model::ModelLike, attr::AbstractVariableAttribute, v::VariableIndex, value)
    load(model::ModelLike, attr::AbstractConstraintAttribute, c::ConstraintIndex, value)

This has the same effect that `set` with the same arguments except that `allocate` should be called first before `load_variables`.
"""
function load(model::MOI.ModelLike, args...)
    MOI.throw_set_error_fallback(model, args...;
                                 error_if_supported=ALLOCATE_LOAD_NOT_IMPLEMENTED)
end

function load(model::MOI.ModelLike, attr::Union{MOI.AbstractVariableAttribute, MOI.AbstractConstraintAttribute}, indices::Vector, values::Vector)
    for (index, value) in zip(indices, values)
        load(model, attr, index, value)
    end
end

"""
    load_constraint(model::MOI.ModelLike, ci::MOI.ConstraintIndex, f::MOI.AbstractFunction, s::MOI.AbstractSet)

Sets the constraint function and set for the constraint of index `ci`.
"""
function load_constraint(model::MOI.ModelLike, func::MOI.AbstractFunction,
                         set::MOI.AbstractSet)
    MOI.throw_add_constraint_error_fallback(model, func, set;
                                            error_if_supported=ALLOCATE_LOAD_NOT_IMPLEMENTED)
end

function allocate_constraints(dest::MOI.ModelLike, src::MOI.ModelLike, copy_names::Bool, idxmap::IndexMap, ::Type{F}, ::Type{S}) where {F<:MOI.AbstractFunction, S<:MOI.AbstractSet}
    # Allocate constraints
    cis_src = MOI.get(src, MOI.ListOfConstraintIndices{F, S}())

    for ci_src in cis_src
        f_src = MOI.get(src, MOI.ConstraintFunction(), ci_src)
        s = MOI.get(src, MOI.ConstraintSet(), ci_src)
        f_dest = mapvariables(idxmap, f_src)
        ci_dest = allocate_constraint(dest, f_dest, s)
        idxmap.conmap[ci_src] = ci_dest
    end

    return pass_attributes(dest, src, copy_names, idxmap, cis_src, allocate)
end

function load_constraints(dest::MOI.ModelLike, src::MOI.ModelLike, copy_names::Bool, idxmap::IndexMap, ::Type{F}, ::Type{S}) where {F<:MOI.AbstractFunction, S<:MOI.AbstractSet}
    # Load constraints
    cis_src = MOI.get(src, MOI.ListOfConstraintIndices{F, S}())
    for ci_src in cis_src
        ci_dest = idxmap[ci_src]
        f_src = MOI.get(src, MOI.ConstraintFunction(), ci_src)
        f_dest = mapvariables(idxmap, f_src)
        s = MOI.get(src, MOI.ConstraintSet(), ci_src)
        load_constraint(dest, ci_dest, f_dest, s)
    end

    return pass_attributes(dest, src, copy_names, idxmap, cis_src, load)
end

"""
    allocate_load(dest::MOI.ModelLike, src::MOI.ModelLike)

Implements `MOI.copy_to(dest, src)` using the Allocate-Load API. The function
[`supports_allocate_load`](@ref) can be used to check whether `dest` supports
the Allocate-Load API.
"""
function allocate_load(dest::MOI.ModelLike, src::MOI.ModelLike, copy_names::Bool)
    MOI.empty!(dest)

    idxmap = IndexMap()

    # Allocate variables
    nvars = MOI.get(src, MOI.NumberOfVariables())
    vis_src = MOI.get(src, MOI.ListOfVariableIndices())
    vis_dest = allocate_variables(dest, nvars)
    for (var_src, var_dest) in zip(vis_src, vis_dest)
        idxmap.varmap[var_src] = var_dest
    end

    # Allocate variable attributes
    pass_attributes(dest, src, copy_names, idxmap, vis_src, allocate)

    # Allocate model attributes
    pass_attributes(dest, src, copy_names, idxmap, allocate)

    # Allocate constraints
    for (F, S) in MOI.get(src, MOI.ListOfConstraints())
        # do the rest in copyconstraints! which is type stable
        allocate_constraints(dest, src, copy_names, idxmap, F, S)
    end

    # Load variables
    load_variables(dest, nvars)

    # Load variable attributes
    pass_attributes(dest, src, copy_names, idxmap, vis_src, load)

    # Load model attributes
    pass_attributes(dest, src, copy_names, idxmap, load)

    # Copy constraints
    for (F, S) in MOI.get(src, MOI.ListOfConstraints())
        # do the rest in copyconstraints! which is type stable
        load_constraints(dest, src, copy_names, idxmap, F, S)
    end

    return idxmap
end


text:
struct ChainedADMM end

function sample_without_replacement!(collection, n::Int64) where T<:Any
    sample = []
    for i in 1:n
        push!(sample, splice!(collection,rand(eachindex(collection))))
    end
    return sample
end

function initialize_chained_trackers(A::Array{MaybeMissing{Float64}}, type_assignment)
    type_tracker = Utilities.IndexTracker{Symbol}()
    if type_assignment == nothing
        Concepts.disjoint_join(type_tracker, Concepts.fit(DistributionDeduction(), A))
    else
        Concepts.disjoint_join(type_tracker, type_assignment)
    end
    Concepts.disjoint_join(type_tracker, Concepts.fit(ObservedOrMissing(), A))
    result = Concepts.groupby(type_tracker, [:Observed, :Missing])
    result[:Observed][:Total] = Base.findall(!ismissing, A)
    result[:Missing][:Total] = Base.findall(ismissing, A)
    result[:SingleView] = type_tracker
    return result
end

function Concepts.complete(model::ChainedADMM;
                           block_size = nothing,
                           imputation_round::Union{Int64, Nothing} = 10,
                           A::Array{MaybeMissing{Float64}} = nothing,
                           α::Float64 = maximum(A[findall(x -> !ismissing(x),A)]),
                           λ::Float64 = 5e-1,
                           μ::Float64 = 5e-4,
                           ρ::Float64 = 0.3,
                           τ::Float64 = 1.618,
                           maxiter::Int64 = 200,
                           stoptol::Float64 = 1e-5,
                           use_autodiff::Bool = false,
                           gd_iter::Int64 = 50,
                           debug_mode::Bool = false,
                           interactive_plot = false,
                           type_assignment = nothing,
                           warmup = nothing,
                           dynamic_ρ = true,
                           user_input_estimators = nothing,
                           project_rank = nothing,
                           io::IO = Base.stdout,
                           eigen_solver = KrylovMethods(),
                           closed_form_update = false)
    o_tracker, o_type_tracker   = initialize_trackers(A, type_assignment)
    tracker = initialize_chained_trackers(A, type_assignment)
    missing_entries = deepcopy(tracker[:Missing][:Total])
    # @show(missing_entries[1:10])
    observed_entries = deepcopy(tracker[:Observed][:Total])
    # @show(observed_entries[1:10])
    if !isnothing(imputation_round)
        block_size = trunc(Int64, Base.ceil(length(missing_entries) / imputation_round))
        @info @sprintf("imputation round enforced, current block size is %d.\n", block_size)
    end

    # @show(imputation_round)
    if isnothing(warmup)
        warmup::Dict{Symbol, Array{Float64}} = initialize_warmup(tracker, A)
    end
    imputed = A
    round = 0
    local result_completed_matrix = nothing
    local result_type_tracker = nothing
    local result_tracker = nothing
    local result_last_var = nothing
    while length(missing_entries) > 0
        round = round + 1
        @info(round)
        block_samples = sample_without_replacement!(missing_entries, min(block_size, length(missing_entries)))
        result_completed_matrix, result_type_tracker, result_tracker, result_last_var = complete(A = imputed,
                                                                                α = α,
                                                                                λ = λ,
                                                                                μ = μ,
                                                                                ρ = ρ,
                                                                                τ = τ,
                                                                                maxiter = maxiter,
                                                                                stoptol = stoptol,
                                                                                use_autodiff = use_autodiff,
                                                                                gd_iter = gd_iter,
                                                                                start_var = result_last_var,
                                                                                debug_mode = debug_mode,
                                                                                interactive_plot = interactive_plot,
                                                                                type_assignment = type_assignment,

                                                                                dynamic_ρ = dynamic_ρ,
                                                                                user_input_estimators = user_input_estimators,
                                                                                project_rank = project_rank,
                                                                                io = io,
                                                                                eigen_solver = eigen_solver,
                                                                                closed_form_update = closed_form_update)
        predicted_matrix = predict(MatrixCompletionModel(),
                                   completed_matrix = result_completed_matrix,
                                   type_tracker     = result_type_tracker)

        # @show("update imputed")
        for index in block_samples
            imputed[index] = predicted_matrix[index]
        end

        # if haskey(tracker, :Bernoulli) && length(tracker[:Bernoulli][:Observed]) > 0

        # if haskey(tracker, :Bernoulli) && length(tracker[:Bernoulli][:Observed]) > 0
        # warmup[:Bernoulli] =
        # end
    end
    return result_completed_matrix, o_type_tracker, o_tracker, imputed
end


text:
# During handling of an execute_request (when execute_msg is !nothing),
# we redirect STDOUT and STDERR into "stream" messages sent to the IPython
# front-end.

const orig_STDIN = STDIN
const orig_STDOUT = STDOUT
const orig_STDERR = STDERR

# logging in verbose mode goes to original stdio streams.  Use macros
# so that we do not even evaluate the arguments in no-verbose modes
macro vprintln(x...)
    quote
        if verbose::Bool
            println(orig_STDOUT, $(x...))
        end
    end
end
macro verror_show(e, bt)
    quote
        if verbose::Bool
            showerror(orig_STDERR, $e, $bt)
        end
    end
end

function send_stream(s::String, name::String)
    if !isempty(s)
        send_ipython(publish,
                     msg_pub(execute_msg, "stream",
                             ["name" => name, "data" => s]))
    end
end

function watch_stream(rd::IO, name::String)
    try
        while true
            s = readavailable(rd) # blocks until something available
	    send_stream(s, name)
            sleep(0.1) # a little delay to accumulate output
        end
    catch e
        # the IPython manager may send us a SIGINT if the user
        # chooses to interrupt the kernel; don't crash on this
        if isa(e, InterruptException)
            watch_stream(rd, name)
        else
            rethrow()
        end
    end
end

const read_stdin, write_stdin = redirect_stdin()
const read_stdout, write_stdout = redirect_stdout()
const read_stderr, write_stderr = redirect_stderr()

# IJulia issue #42: there doesn't seem to be a good way to make a task
# that blocks until there is a read request from STDIN ... this makes
# it very hard to properly redirect all reads from STDIN to pyin messages.
# In the meantime, however, we can just hack it so that readline works:
import Base.readline
function readline(io::Base.Pipe)
    if io == STDIN
        if !execute_msg.content["allow_stdin"]
            error("IJulia: this front-end does not implement stdin")
        end
        send_ipython(raw_input,
                     msg_reply(execute_msg, "input_request",
                               ["prompt" => "STDIN> "]))
        while true
            msg = recv_ipython(raw_input)
            if msg.header["msg_type"] == "input_reply"
                return msg.content["value"]
            else
                error("IJulia error: unknown stdin reply")
            end
        end
    else
        invoke(readline, (Base.AsyncStream,), io)
    end
end

function watch_stdio()
    @async watch_stream(read_stdout, "stdout")
    @async watch_stream(read_stderr, "stderr")
end


text:
module UnionBenchmarks

include(joinpath(dirname(@__FILE__), "..", "utils", "RandUtils.jl"))

using .RandUtils
using BenchmarkTools

const SUITE = BenchmarkGroup()

###########################
# array Union{T, Nothing} #
###########################

g = addgroup!(SUITE, "array")

const VEC_LENGTH = 10000

_zero(::Type{T}) where {T} = zero(T)
_zero(::Type{Union{T, Nothing}}) where {T} = zero(T)

_abs(x) = abs(x)
_abs(::Nothing) = nothing

_mul(x, y) = x * y
_mul(::Nothing, ::Any) = nothing
_mul(::Any, ::Nothing) = nothing
_mul(::Nothing, ::Nothing) = nothing

_isnothing(x) = false
_isnothing(::Nothing) = true

function perf_sum(X::AbstractArray{T}) where T
    s = _zero(T) + _zero(T)
    @inbounds @simd for x in X
        if x !== nothing
            s += x
        end
    end
    s
end

function perf_sum2(X::AbstractArray{T}) where T
    s = _zero(T) + _zero(T)
    @inbounds @simd for x in X
        if !_isnothing(x)
            s += x
        end
    end
    s
end

function perf_sum3(X::AbstractArray{T}) where T
    s = _zero(T) + _zero(T)
    @inbounds for x in X
        if !_isnothing(x)
            s += x
        end
    end
    s
end

function perf_sum4(X::AbstractArray{T}) where T
    s = _zero(T) + _zero(T)
    @inbounds @simd for x in X
        s += ifelse(x === nothing, _zero(T), x)
    end
    s
end

function perf_countnothing(X::AbstractArray)
    n = 0
    @inbounds for x in X
        n += x === nothing
    end
    n
end

function perf_countequals(X::AbstractArray, Y::AbstractArray)
    n = 0
    @inbounds for i in eachindex(X, Y)
        n += isequal(X[i], Y[i])
    end
    n
end

function perf_simplecopy(X::AbstractArray)
    ret = similar(X)
    @inbounds for i in eachindex(X)
        ret[i] = X[i]
    end
    ret
end

function perf_binaryop(op::Function, X::AbstractArray, Y::AbstractArray)
    ret = similar(X, Union{Nothing, typeof(op(_zero(eltype(X)), _zero(eltype(Y))))})
    @inbounds for i in eachindex(X, Y)
        ret[i] = op(X[i], Y[i])
    end
    ret
end

for T in (Bool, Int8, Int64, Float32, Float64, BigInt, BigFloat, Complex{Float64})
    if T == BigInt
        S = Int128
    elseif T == BigFloat
        S = Float64
    else
        S = T
    end

    # 10% of missing values
    X = Vector{Union{T, Nothing}}(Vector{T}(samerand(S, VEC_LENGTH)))
    Y = Vector{Union{T, Nothing}}(Vector{T}(samerand(S, VEC_LENGTH)))
    X2 = Vector{Union{T, Nothing}}(Vector{T}(samerand(S, VEC_LENGTH)))
    Y2 = Vector{Union{T, Nothing}}(Vector{T}(samerand(S, VEC_LENGTH)))
    X2[samerand(VEC_LENGTH) .> .9] .= nothing
    Y2[samerand(VEC_LENGTH) .> .9] .= nothing

    for (M, A) in ((false, X), (true, X2))
        g["perf_sum", T, M] = @benchmarkable perf_sum($A)
        g["perf_sum2", T, M] = @benchmarkable perf_sum2($A)
        g["perf_sum3", T, M] = @benchmarkable perf_sum3($A)
        g["perf_sum4", T, M] = @benchmarkable perf_sum4($A)

        g["perf_countnothing", T, M] = @benchmarkable perf_countnothing($A)

        g["perf_simplecopy", T, M] =
            @benchmarkable perf_simplecopy($A)
        g["map", identity, T, M] =
            @benchmarkable map(identity, $A)
        g["broadcast", identity, T, M] =
            @benchmarkable broadcast(identity, $A)

        g["map", abs, T, M] =
            @benchmarkable map(_abs, $A)
        g["broadcast", abs, T, M] =
            @benchmarkable broadcast(_abs, $A)
    end

    for (M, A, B) in (((false, false), X, Y),
                      ((true, true), X2, Y2),
                      ((false, true), X, Y2))
        g["perf_countequals", string(T)] =
            @benchmarkable perf_countequals($A, $B)

        g["perf_binaryop", *, T, M] =
            @benchmarkable perf_binaryop(_mul, $A, $B)
        g["map", *, T, M] =
            @benchmarkable map(_mul, $A, $B)
        g["broadcast", *, T, M] =
            @benchmarkable broadcast(_mul, $A, $B)
    end

    for (M, A) in ((false, Vector{T}(X)),
                   (false, X),
                   (true, replace(X2, nothing=>missing)))
        g["skipmissing", collect, eltype(A), M] =
            @benchmarkable collect(skipmissing($A))

        g["skipmissing", sum, eltype(A), M] =
            @benchmarkable sum(skipmissing($A))

        if hasmethod(isless, Tuple{T, T})
            g["sort", eltype(A), M] = @benchmarkable sort($A)
        end
    end
end

end # module


text:
###########################################################################
# Utilities for Gibbs updates
###########################################################################

# using StatsBase
# using Distributions

function logp_partition(PP::Vector{Int},T::Vector{Int},Psi::Vector{Float64},
        alpha::Float64,ia_dist::DiscreteDistribution,is_partition::Bool)
    """
    - `PP`: vector of partition block sizes ordered by arrival time
    - `T`: vector of arrival times
    - `Psi`: vector of beta random variables (can be log(Psi))
    - `alpha`: 'discount parameter' in size-biased reinforcement
    - `ia_dist`: distribution object corresponding to i.i.d. interarrivals
    - `is_partition`: Boolean flag for computing binomial coefficients
    """

    if all(Psi .<= 0)
      # warn("Psi in log space.")
      log_Psi = Psi
      log_Psi_c = log.(1 - exp.(log_Psi))
    elseif all(0 .<= Psi .<= 1)
      log_Psi = log.(Psi)
      log_Psi_c = log.(1 - Psi)
    else
      error("Invalid Ψ.")
    end
    # shift distributions with non-zero mass on zero
    zero_shift = Int(minimum(ia_dist) == 0)

    PP_bar = cumsum(PP)
    ia = T[2:end] .- T[1:(end-1)]

    K = size(Psi,1)
    idx = 1:(K-1)
    N = sum(PP)

    log_p = dot((PP[2:end] .- alpha .- 1),log_Psi[2:end]) + dot((PP_bar[1:(end-1)] .- idx.*alpha .- 1), log_Psi_c[2:end])
    # include arrival times
    log_p += sum(logpdf.(ia_dist, ia .- zero_shift))
    N - T[end] > 0 ? log_p += log(1 - cdf(ia_dist, N-T[end]-zero_shift)) : nothing
    log_p += -sum([lbeta(1 - alpha,T[j] - 1 - (j-1)*alpha) for j in 2:K])
    # include binomial coefficients if for a partition
    if is_partition
      log_p += sum([lbinom(PP_bar[j] - T[j],PP[j] - 1) for j in 2:K])
    end

    return log_p

end

function logp_partition(PP::Vector{Int},T::Vector{Int},
        alpha::Float64,ia_dist::DiscreteDistribution,is_partition::Bool)
    f = (x,y) -> ia_dist
    logp_partition(PP,T,alpha,f,is_partition)
end

function logp_partition(PP::Vector{Int},T::Vector{Int},
        alpha::Float64,ia_dist::Function,is_partition::Bool)
    """
    - `PP`: vector of partition block sizes ordered by arrival time
    - `T`: vector of arrival times
    - `alpha`: 'discount parameter' in size-biased reinforcement
    - `ia_dist`: function that creates an interarrival distribution at fixed parameters
    - `is_partition`: flag for computing binomial coefficients
    """

    # shift distributions with non-zero mass on zero
    zero_shift = Int(minimum(ia_dist(1,1)) == 0)

    PP_bar = cumsum(PP)
    # pop!(PP_bar)
    ia = T[2:end] .- T[1:(end-1)]

    K = length(PP)
    idx = 1:(K-1)
    N = sum(PP)

    log_p = log_CPPF(PP,T,alpha)

    if N - T[end] > 0
      p_gt = 1 - cdf(ia_dist(T[end],K), N-T[end]-zero_shift)
      abs(p_gt)<=eps(one(typeof(p_gt))) || p_gt < 0. ? log_p = -Inf : log_p += log(p_gt)
      # log_p += log(1 - cdf(ia_dist(T[end],K), N-T[end]-zero_shift))
    end
    log_p += sum( [logpdf(ia_dist(T[j-1],j-1),T[j]-T[j-1]-zero_shift) for j in 2:K ])
    # include binomial coefficients if for a partition
    if is_partition
      log_p += sum([lbinom(PP_bar[j] - T[j],PP[j] - 1) for j in 2:K])
    end

    return log_p

end

function logp_pred_partition(PP_train::Vector{Int},PP_test::Vector{Int},
        T::Vector{Int},
        alpha::Float64,ia_dist::DiscreteDistribution,
        is_partition_train::Bool,is_partition_test::Bool)
    f = (x,y) -> ia_dist
    logp_pred_partition(PP_train,PP_test,T,alpha,f,is_partition_train,is_partition_test)
end

function logp_pred_partition(PP_train::Vector{Int},PP_test::Vector{Int},
        T::Vector{Int},
        alpha::Float64,ia_dist::Function,
        is_partition_train::Bool,is_partition_test::Bool)

    zero_shift = Int(minimum(ia_dist(1,1)) == 0)
    K_train = length(PP_train)
    K_test = length(PP_test)

    T_train = T[1:K_train]

    PP_train_bar = cumsum(PP_train)
    N_train = PP_train_bar[end]
    logp_train = log_CPPF(PP_train,T_train,alpha)
    # is_partition_train ? logp_train += sum([lbinom(PP_train_bar[j] - T_train[j],PP_train[j] - 1) for j in 2:length(PP_train)]) : nothing

    PP_test_bar = cumsum(PP_test)
    logp_test = log_CPPF(PP_test,T,alpha)
    # is_partition_test ? logp_test += sum([lbinom(PP_test_bar[j] - T_test[j],PP_test[j] - 1) for j in 2:length(PP_test)]) : nothing
    if is_partition_test
      # predicted part of partition is as usual
      logp_test += sum([lbinom(PP_test_bar[j] - T[j],PP_test[j] - 1) for j in (K_train+1):length(PP_test)])
      # conditioned part of partition is constrained
      for j in 1:K_train
        ( PP_test[j] - N_train > 0 && PP_test_bar[j] - PP_train_bar[j] > 0 )?
          logp_test += lbinom(PP_test_bar[j] - N_train,PP_test[j] - PP_train[j]) : nothing
      end
    end

    N = PP_test_bar[end]

    if N - T[end] > 0
      p_gt = 1 - cdf(ia_dist(T[end],K_test), N-T[end]-zero_shift)
      abs(p_gt)<=eps(one(typeof(p_gt))) || p_gt < 0. ? logp_test = -Inf : logp_test += log(p_gt)
    end
    logp_test += sum( [logpdf(ia_dist(T[j-1],j-1),T[j]-T[j-1]-zero_shift) for j in (K_train+1):K_test])

    return logp_test - logp_train
end

# memoize this?
function lbinom(n::Int,k::Int)
    """
    computes log of binomial coefficient {`n` choose `k`} using `lgamma`
    """
    ret = lgamma(n+1) - lgamma(k+1) - lgamma(n - k + 1)
    return ret
end

function tally_ints(Z::Vector{Int},K::Int)
    """
    counts occurrences in `Z` of integers 1 to `K`
    - `Z`: vector of integers
    - `K`: maximum value to count occurences in `Z`
    """
    ret = zeros(Int,K)
    n = size(Z,1)
    idx_all = 1:n
    idx_j = trues(n)
    for j in 1:K
      for i in idx_all[idx_j]
        if Z[i]==j
          ret[j] += 1
          idx_j[i] = false
        end
      end
    end
    return ret
end

function initialize_alpha(prior_dist::ContinuousUnivariateDistribution)
  return [1 - rand(prior_dist)]
end

function update_ntl_alpha!(alpha::Vector{Float64},PP::Vector{Int},T::Vector{Int},log_prior::Function,w::Float64)
  """
  Slice-sampling update (with sampler parameter `w`) of NTL discount parameter
    `alpha`, conditioned on arrival times `T` and arrival-ordered block/vertex
    counts `PP`.
  """
  alpha_trans = log(1 - alpha[1])
  ss = x -> ntl_alpha_trans_logpdf(x,PP,T,log_prior)
  alpha_trans_new = slice_sampling(ss,w,alpha_trans)
  alpha[1] = 1 - exp(alpha_trans_new)
end

function ntl_alpha_logpdf(alpha::Float64,PP::Vector{Int},T::Vector{Int},log_prior::Function)
    """
    calculate unnormalized log-pdf proportional to `alpha` (discount in NTL)
      (Psi_j's are marginalized)

    log_prior is a function that returns the (possibly unnormalized) prior log-probability
      of `alpha`
    """
    # PP_bar = cumsum(PP)
    logp = log_prior(1 - alpha) + lgamma(PP[1] - alpha) - lgamma(sum(PP) - size(PP,1)*alpha) # prior is specified as a distribution on (0,Inf); alpha ∈ (-Inf,1)
    for j in 2:size(PP,1)
      logp +=  lgamma(PP[j] - alpha) - lbeta(1 - alpha,T[j] - 1 - (j-1)*alpha)
      # logp += lbeta(PP[j] - alpha, PP_bar[j-1] - (j-1)*alpha) - lbeta(1-alpha,T[j] - 1 - (j-1)*alpha)
    end
    return logp
end

function ntl_alpha_trans_logpdf(alpha_trans::Float64,PP::Vector{Int},T::Vector{Int},log_prior::Function)
    """
    calculate unnormalized log-pdf proportional to `alpha_trans`,
      the transformed discount (alpha) in NTL such that alpha = 1 - exp(alpha_trans)

      (Psi_j's are marginalized)

    log_prior is a function that returns the (possibly unnormalized) prior log-probability
      of `alpha`
    """
    return ntl_alpha_logpdf(1 - exp(alpha_trans),PP,T,log_prior) + alpha_trans
end

function seq2part(Z::Vector{Int})
    """
    - `Z`: vector of labels corresponding to clusters/blocks in a partition
    """
    # transform Z into an ordered partition and get arrival times
    K = maximum(Z)
    PP = tally_ints(Z,K)

    # T = zeros(Int,K)
    # for j in 1:K
    #   T[j] = findfirst(Z.==j)
    # end
    return PP
end

function get_arrivals(Z::Vector{Int})
    """
    - `Z`: vector of labels corresponding to clusters/blocks in a partition
    """

    K = maximum(Z)
    T = zeros(Int,K)
    for j in 1:K
      T[j] = findfirst(Z.==j)
    end
    return T
end

function logp_label_sequence(Z::Vector{Int},Psi::Vector{Float64},
        alpha::Float64,ia_dist::DiscreteDistribution)
    """
    - `Z`: vector of labels corresponding to clusters/blocks in a partition
    - `Psi`: vector of beta random variables (can be log(Psi))
    - `alpha`: 'discount parameter' in size-biased reinforcement
    - `ia_dist`: distribution object corresponding to i.i.d. interarrivals
    """

    PP = seq2part(Z)
    T = get_arrivals(Z)
    log_p = logp_partition(PP,T,Psi,alpha,ia_dist,false)
    return log_p

end

function cluster_rm!(x::Vector{Vector{Float64}},k::Int)
  """
  removes column k from vector of vectors x (e.g. corresponding cluster params)
  """
  deleteat!(x,k)
end

function cluster_rm!(x::Vector{Float64},k::Int)
  """
  removes k-th entry from x
  """
  deleteat!(x,k)
end

function cluster_add!(x::Vector{Vector{Float64}},x_new::Vector{Float64},k::Int)
  """
  inserts `x_new` into `x` at entry `k`
  """
  insert!(x,k,x_new)
end

function cluster_add!(x::Vector{Float64},x_new::Float64,k::Int)
  """
  inserts `x_new` into `x` at entry `k`
  """
  insert!(x,k,x_new)
end

function cycle_elements_left!(V::Vector,start_idx::Int,end_idx::Int)
    """
    - `V`: Vector whose elements will be cycled
    - `start_idx`: start of cycle (will be moved to end)
    - `end_idx`: end of cycle
    """

    st = V[start_idx]
    for i in start_idx:(end_idx-1)
      V[i] = V[i+1]
    end
    V[end_idx] = st
    return V
end

function cycle_elements_left!(X::Array,start_idx::Int,end_idx::Int)
    """
    - `X`: Array whose columns will be cycled
    - `start_idx`: start of cycle (will be moved to end)
    - `end_idx`: end of cycle
    """

    for i in 1:size(X,1)
      X[i,:] = cycle_elements_left!(X[i,:],start_idx,end_idx)
    end
    return X
end

function cycle_elements_right!(V::Vector,start_idx::Int,end_idx::Int)
    """
    - `V`: Vector whose elements will be cycled
    - `start_idx`: start of cycle (will be moved to end)
    - `end_idx`: end of cycle
    """

    ed = V[end_idx]
    for i in end_idx:-1:(start_idx+1)
      V[i] = V[i-1]
    end
    V[start_idx] = ed
    return V

end

function cycle_elements_right!(X::Array,start_idx::Int,end_idx::Int)
    """
    - `X`: Array whose columns will be cycled
    - `start_idx`: start of cycle (will be moved to end)
    - `end_idx`: end of cycle
    """

    for i in 1:size(X,1)
      X[i,:] = cycle_elements_right!(X[i,:],start_idx,end_idx)
    end
    return X
end

function log_cppf_counts(PP::Vector{Int},alpha::Float64)
    """
    helper function for `log_CPPF` and `update_label_sequence`
    """
    gt1 = PP .> 1
    ret = sum( lgamma.(PP[gt1] .- alpha) ) - sum(gt1)*lgamma(1 - alpha)
    return ret
end

function log_cppf_arrivals(T::Vector{Int},alpha::Float64)
    """
    helper function for `log_CPPF` and `update_label_sequence`
    """
    K = size(T,1)
    ret = sum( lgamma.(T .- (1:K).*alpha) ) - sum( lgamma.(T[2:end] .- 1 .- (1:(K-1)).*alpha) )
    return ret
end

function log_cppf_arrivals(T::Vector{Int},arrival_offset::Int,alpha::Float64)
  """
  helper function for computing predictive log-probabilities

  `arrival_offset=K` indicates that the first element of `T` corresponds to the
    K-th arrival
  """
  K_end = arrival_offset - 1 + size(T,1)
  return sum( lgamma.(T .- (arrival_offset:K_end).*alpha) ) - sum( lgamma.(T .- 1 .- (arrival_offset-1):(K_end-1).*alpha) )
end

function log_CPPF(PP::Vector{Int},T::Vector{Int},alpha::Float64)
    """
    - `PP`: arrival-ordered vector of partition block sizes
    - `T`: arrival times
    - `alpha`: 'discount' parameter
    """

    n = sum(PP)
    K = size(T,1)

    logp = -lgamma(n - K*alpha) + log_cppf_arrivals(T,alpha) + log_cppf_counts(PP,alpha)
    return logp
end

function log_CPPF(Z::Vector{Int},alpha::Float64)
    """
    - `Z`: vector of labels corresponding to clusters/blocks in a partition
    - `alpha`: 'discount' parameter
    """

    T = get_arrivals(Z)
    PP = seq2part(Z)
    logp = log_CPPF(PP,T,alpha)
    return logp
end

function get_num_blocks(Z::Vector{Int})
    """
    - `Z`: vector of labels corresponding to clusters/blocks in a partition
    """

    n = size(Z,1)
    K = zeros(Int64,n)
    K[1] = 1
    max_z = 1
    for i in 2:n
      if Z[i] > max_z
        K[i] = K[i-1] + 1
        max_z += 1
      end
    end
    return K
end

function update_psi_parameters_sequence!(Psi::Vector{Float64},Z::Vector{Int},alpha::Float64)
    """
    - `Psi`: vector of current values of Psi
    - `Z`: vector of labels corresponding to clusters/blocks in a partition
    - `alpha`: 'discount' parameter
    """

    PP = seq2part(Z)
    update_psi_parameters_partition!(Psi,PP,alpha)
    # Psi[:] = [pu[i] for i in 1:size(Psi,1)]
    return Psi

end

function update_psi_parameters_partition!(Psi::Vector{Float64},PP::Vector{Int},alpha::Float64)
    """
    - `Psi`: vector of current values of Psi
    - `PP`: arrival-ordered vector of partition block sizes
    - `alpha`: 'discount' parameter
    """
    K = size(PP,1)
    PP_bar = cumsum(PP)

    Psi[1] = 1
    for j in 2:K
      Psi[j] = rand(Beta(PP[j]-alpha,PP_bar[j-1]-(j-1)*alpha))
    end
    return Psi
end

function log_sum_exp_weights(logw::Vector{Float64})
  """
  -`logw`: log of weights to be combined for a discrete probability distribution
  """

  maxlogw = maximum(logw)
  shift_logw = logw - maxlogw
  p = exp.(shift_logw)./sum(exp.(shift_logw))
  return p
end

function initialize_arrival_times(PP::Vector{Int},alpha::Float64,ia_dist::DiscreteDistribution)
  f = (x,y)->ia_dist
  initialize_arrival_times(PP,alpha,f)
end

function initialize_arrival_times(PP::Vector{Int},alpha::Float64,ia_dist::Function)
    """
    - `PP`: arrival-ordered vector of partition block sizes
    - `alpha`: 'discount' parameter
    - `ia_dist`: function that returns pdf corresponding to interarrivals
    """
    zero_shift = Int(minimum(ia_dist(1,1)) == 0)
    K = size(PP,1)
    PP_bar = cumsum(PP)
    n = PP_bar[end]

    typeof(ia_dist(1,1))==CRPinterarrival ? crp = true : crp = false
    crp ? crp_dist = ia_dist(1,1) : nothing

    T = zeros(Int,K)
    T[1] = 1

    for j in 2:K
      # determine support of interarrival
      supp = 1:(PP_bar[j-1] - T[j-1] + 1)
      # calculate pmf of conditional distribution
      log_p = zeros(Float64,size(supp,1))
      if crp
        log_p += crp_logpdf(crp_dist.theta,crp_dist.alpha,T[j-1],j-1,supp.-zero_shift)
      else
        log_p += logpdf.(ia_dist(T[j-1],j-1),supp.-zero_shift)
      end
      log_p += lbinom.(PP_bar[j] .- T[j-1] .- supp, PP[j] - 1)
      log_p += lgamma.(T[j-1] .+ supp .- j*alpha) .- lgamma.(T[j-1] .+ supp .- 1 .- (j-1)*alpha)
      # sample an update
      p = log_sum_exp_weights(log_p)
      T[j] = T[j-1] + wsample(supp,p)
      # println("initialized T_",j,"=",T[j])
    end

    return T
end

function sample_interarrival(j::Int,T_jm1::Int,T_jp1::Int,ia_dist::Function,
  zero_shift::Int,PP_bar_jm1::Int,PP_j::Int,alpha::Float64)
  """
  Utility function for arrival time updates
  """
  delta2 = T_jp1 - T_jm1
  # check to see if interarrivals are Geometric or CRP to short-cut some computtions
  typeof(ia_dist(1,1))==Distributions.Geometric{Float64} ? geom = true : geom = false
  typeof(ia_dist(1,1))==CRPinterarrival ? crp = true : crp = false
  crp ? crp_dist = ia_dist(1,1) : nothing
  ia_dist(1,1)==ia_dist(10,10) ? iid = true : iid = false

  # determine support
  supp = 1:min(delta2 - 1, PP_bar_jm1 - T_jm1 + 1)
  if size(supp,1) > 1
    # calculate pmf of conditional distribution
    log_p = zeros(Float64,size(supp,1))
    if crp
      log_p += crp_logpdf(crp_dist.theta,crp_dist.alpha,T_jm1,j-1,supp.-zero_shift)
    elseif !geom && iid # geometric interarrival prior is memoryless => doesn't contribute
      tmp = logpdf.(ia_dist(T_jm1,j-1),supp.-zero_shift)
      log_p += tmp .+ tmp[end:-1:1] # iid interarrivals are symmetric
    end
    log_p += lbinom.(PP_bar_jm1 .+ PP_j .- T_jm1 .- supp, PP_j - 1)
    log_p += lgamma.(T_jm1 .+ supp .- j*alpha) .- lgamma.(T_jm1 .+ supp .- 1 .- (j-1)*alpha)
    if !iid
      for s in supp
        if crp
          log_p[s] += crp_logpdf(crp_dist.theta,crp_dist.alpha,T_jm1+s,j,delta2 - s - zero_shift)
        else
          log_p[s] += logpdf(ia_dist(T_jm1+s,j),delta2 - s - zero_shift)
        end
      end
    end
    # sample an update
    p = log_sum_exp_weights(log_p)
    return wsample(supp,p)
  else
    return 1
  end
end

function sample_final_arrival(T_Km1::Int,K::Int,n::Int,ia_dist::Function,
  zero_shift::Int,PP_bar_Km1::Int,PP_K::Int,alpha::Float64)

  if T_Km1==(n-1)
    TK = n
  else
    typeof(ia_dist(1,1))==Distributions.Geometric{Float64} ? geom = true : geom = false
    typeof(ia_dist(1,1))==CRPinterarrival ? crp = true : crp = false
    crp ? crp_dist = ia_dist(1,1) : nothing

    supp = 1:min(n - T_Km1, PP_bar_Km1 - T_Km1 + 1)
    log_p = zeros(Float64,size(supp,1))

    if crp
      log_p += crp_logpdf(crp_dist.theta,crp_dist.alpha, T_Km1, K-1, supp.-zero_shift)
    elseif !geom
      log_p += logpdf.(ia_dist(T_Km1,K-1), supp.-zero_shift)
    end

    log_p += lbinom.(n .- T_Km1 .- supp, PP_K - 1)
    log_p += lgamma.(T_Km1 .+ supp .- K*alpha) .- lgamma.(T_Km1 .+ supp .- 1 .- (K-1)*alpha)

    if !geom
      for s in supp
        if crp
          p_gt = 1. - crp_cdf(crp_dist.theta,crp_dist.alpha,T_Km1+s,K,n-(T_Km1+s-zero_shift))
        else
          p_gt = 1. - cdf(ia_dist(T_Km1+s,K),n-(T_Km1+s-zero_shift)) # this can be arbitrarily close to zero, need to handle numerical instability
        end
        abs(p_gt)<=eps(one(typeof(p_gt))) || p_gt < 0. ? log_p[s] = -Inf : log_p[s] += log(p_gt)
      end
    end
    p = log_sum_exp_weights(log_p)
    TK = T_Km1 + wsample(supp,p)
  end
  return TK
end

function update_arrival_times!(T::Vector{Int},PP::Vector{Int},alpha::Float64,ia_dist::Function)
    """
    Takes advantage of multiple threads if possible.
    - `PP`: arrival-ordered vector of partition block sizes
    - `T`: current arrival times (to be updated)
    - `alpha`: 'discount' parameter
    - `ia_dist`: function that creates a distribution object corresponding to interarrival distribution
    """
    update_arrival_times_st!(T,PP,alpha,ia_dist)

end

function update_arrival_times!(T::Vector{Int},PP::Vector{Int},alpha::Float64,ia_dist::DiscreteDistribution)
  f = (x,y)->ia_dist
  update_arrival_times!(T,PP,alpha,f)
end

function update_arrival_times_mt!(T::Vector{Int},PP::Vector{Int},alpha::Float64,ia_dist::Function)
    """
    multi-threaded version
    - `PP`: arrival-ordered vector of partition block sizes
    - `T`: current arrival times (to be updated)
    - `alpha`: 'discount' parameter
    - `ia_dist`: function that creates a distribution object corresponding to interarrival distribution
    """
    zero_shift = Int(minimum(ia_dist(1,1)) == 0)
    K = size(T,1)
    PP_bar = cumsum(PP)
    n = PP_bar[end]

    evens = 2:2:(K-1)
    odds = 3:2:(K-1)

    for j in evens
      T[j] = T[j-1] + sample_interarrival(j,T[j-1],T[j+1],ia_dist,zero_shift,PP_bar[j-1],PP[j],alpha)
    end

    for j in odds
      T[j] = T[j-1] + sample_interarrival(j,T[j-1],T[j+1],ia_dist,zero_shift,PP_bar[j-1],PP[j],alpha)
    end

    T[K] = sample_final_arrival(T[K-1],K,n,ia_dist,zero_shift,PP_bar[K-1],PP[K],alpha)
end

function update_arrival_times_st!(T::Vector{Int},PP::Vector{Int},alpha::Float64,ia_dist::Function)
    """
    single-threaded version
    - `PP`: arrival-ordered vector of partition block sizes
    - `T`: current arrival times (to be updated)
    - `alpha`: 'discount' parameter
    - `ia_dist`: function that creates a distribution object corresponding to interarrival distribution
    """
    zero_shift = Int(minimum(ia_dist(1,1)) == 0)
    K = size(T,1)
    PP_bar = cumsum(PP)
    n = PP_bar[end]

    for j in 2:(K-1)
      T[j] = T[j-1] + sample_interarrival(j,T[j-1],T[j+1],ia_dist,zero_shift,PP_bar[j-1],PP[j],alpha)
    end

    T[K] = sample_final_arrival(T[K-1],K,n,ia_dist,zero_shift,PP_bar[K-1],PP[K],alpha)

    return T
end

function swap_elements!(x::Vector,i::Int,j::Int)
  """
  swap elements `i` and `j` of `x` in place
  """
  x[i],x[j] = x[j],x[i]
  return x
end

function update_block_order!(perm::Vector{Int},PP::Vector{Int},T::Vector{Int},alpha::Float64)
    """
    - `perm`: permutation of order of entries in `PP` (to be updated)
    - `PP`: partition of arrival-ordered block sizes
    - `T`: arrival times
    - `alpha`: 'discount' parameter
    Update is through a sequence of proposed adjacent transpositions.

    There are likely better (more efficient) ways to do this.
    """

    K = size(PP,1)
    PP_bar = cumsum(PP)

    for j in 1:(K-1)

      j==1 ? ppbar_jm1 = 0 : ppbar_jm1 = PP_bar[j-1]
      if PP[j]==PP[j+1] # swap is a 50-50 flip
        logp_swap = log(0.5)
        logp_noswap = log(0.5)
      elseif ppbar_jm1 + PP[j+1] >= T[j+1] - 1
        logp_swap = lgamma(ppbar_jm1 + PP[j+1] - T[j] + 1) - lgamma(PP_bar[j+1] - PP[j] - T[j+1] + 2)
        logp_noswap = lgamma(ppbar_jm1 + PP[j] - T[j] + 1) - lgamma(PP_bar[j] - T[j+1] + 2)
      else # impossible to swap blocks given arrival times
        logp_swap = -Inf
        logp_noswap = 0.
      end

      swap = wsample([true;false],log_sum_exp_weights([logp_swap;logp_noswap]))
      if swap
        swap_elements!(PP,j,j+1)
        swap_elements!(perm,j,j+1)
        (j == 1) ? PP_bar[j] =  PP[j] : PP_bar[j] = PP_bar[j-1] + PP[j]
        PP_bar[j+1] = PP_bar[j] + PP[j+1]
      end

    end
    return perm
end