include(joinpath("src", "EpidemicModel.jl"))
CONFIG_FOLDER = "config"

@time begin
    config_file = "cfg_main_agent_density.jl"
    include(joinpath(CONFIG_FOLDER, config_file))  # config dictionary created in config script
    EpidemicModel.run_batch(
        configs=configs, replicates=10, batchname="main_agent_density", compress=true
    )
end

@time begin
    config_file = "cfg_main_behavior_reduction_factor.jl"
    include(joinpath(CONFIG_FOLDER, config_file))  # config dictionary created in config script
    EpidemicModel.run_batch(
        configs=configs, replicates=10, batchname="main_behavior_reduction_factor", compress=true
    )
end

@time begin
    config_file = "cfg_main_policy_reach.jl"
    include(joinpath(CONFIG_FOLDER, config_file))  # config dictionary created in config script
    EpidemicModel.run_batch(
        configs=configs, replicates=10, batchname="main_policy_reach", compress=true
    )
end

@time begin
    config_file = "cfg_main_test_probability.jl"
    include(joinpath(CONFIG_FOLDER, config_file))  # config dictionary created in config script
    EpidemicModel.run_batch(
        configs=configs, replicates=10, batchname="main_test_probability", compress=true
    )
end

@time begin
    config_file = "cfg_main_tick_authority_policy.jl"
    include(joinpath(CONFIG_FOLDER, config_file))  # config dictionary created in config script
    EpidemicModel.run_batch(
        configs=configs, replicates=10, batchname="main_tick_authority_policy", compress=true
    )
end

@time begin
    config_file = "cfg_main_tick_authority_recommendation.jl"
    include(joinpath(CONFIG_FOLDER, config_file))  # config dictionary created in config script
    EpidemicModel.run_batch(
        configs=configs, replicates=10, batchname="main_tick_authority_recommendation", compress=true
    )
end

@time begin
    config_file = "cfg_main_transmission_probability.jl"
    include(joinpath(CONFIG_FOLDER, config_file))  # config dictionary created in config script
    EpidemicModel.run_batch(
        configs=configs, replicates=10, batchname="main_transmission_probability", compress=true
    )
end

@time begin
    config_file = "cfg_small_scale_detail.jl"
    include(joinpath(CONFIG_FOLDER, config_file))  # config dictionary created in config script
    EpidemicModel.run_batch(
        configs=configs, replicates=1, batchname="small_scale_detail", compress=true
    )
end


# """
#     isWaveguide(domain::Domain)
# """
# function isWaveguide(domain::Domain)
#     return domain.domain_type ∈ [:planar_waveguide, :planar_waveguide_background, :halfspace, :pc_waveguide, :pc_waveguide_background]
# end

#
# """
#     isBulkWaveguide(domain::Domain))
# """
# function isBulkWaveguide(domain::Domain)
#     return domain.domain_type ∈ [:bulk_planar_waveguide_x, :bulk_planar_waveguide_y, :bulk_pc_waveguide_x, :bulk_pc_waveguide_y]
# end

#
# """
#     isBackground(domain::Domain)
# """
# function isBackground(domain)
#     return domain.domain_type ∈ [:background, :planar_waveguide_background, :pc_waveguide_background]
# end

# """
#     isDefect(domain::Domain))
# """
# function isDefect(domain::Domain)
#     return domain.domain_type ∈ [:defect, :site_defect, :line_defect, :pc_waveguide]
# end

# # """
# #     isPC(dom::Domain)
# # """
# # function isPC(dom::Domain)
# #     return dom.domain_type ∈ [:pc, :pc_waveguide, :pc_waveguide_background]
# # end
#
#
# """
#     isPlanar(dom::Domain)
# """
# function isPlanar(dom::Domain)
#     return dom.domain_type ∈ [:planar_waveguide, :planar_waveguide_background, :bulk_planar_waveguide_x, :bulk_planar_waveguide_y]
# end

"""
    isHalfSpace(dom::Domain)
"""
function isHalfSpace(dom::Domain)
    return dom.domain_type ∈ [:halfspace_waveguide]
end

"""
    isDirichlet(bc)
"""
isDirichlet(bc::DirichletBC) = true
isDirichlet(bc) = false

"""
    isNeumann(bc)
"""
isNeumann(bc::NeumannBC) = true
isNeumann(bc) = false

"""
    isOpen(bc)
"""
isMatched(bc::MatchedBC) = true
isMatched(bc) = false

"""
    isPeriodic(bc)
"""
isFloquet(bc::FloquetBC) = true
isFloquet(bc) = false

"""
    isPMLout(bl)
"""
isPML(bl::PML) = true
isPML(bl) = false

"""
    isPMLin(bl)
"""
iscPML(bl::cPML) = true
iscPML(bl) = false

"""
    isNone(bl)
"""

isnoBL(bl::noBL) = true
isnoBL(bl) = false


   # assume communication to 6 neighbors by default
using MPI
function  get_neighbors(domain::Domain)
   rowMin = !(domain.m_rowLoc == 0)
   rowMax = !(domain.m_rowLoc == domain.m_tp - 1)
   colMin = !(domain.m_colLoc == 0)
   colMax = !(domain.m_colLoc == domain.m_tp - 1)
   planeMin = !(domain.m_planeLoc == 0)
   planeMax = !(domain.m_planeLoc == domain.m_tp - 1)

   return rowMin, rowMax, colMin, colMax, planeMin, planeMax
end

copyto_zero!(dest, doffs, src, soffs, nelems) = copyto!(dest, doffs+1, src, soffs+1, nelems)

# Assume 128 byte coherence
# Assume Float64 is an "integral power of 2" bytes wide
const CACHE_COHERENCE_PAD_REAL = div(128, sizeof(Float64))

function commRecv(domain::Domain, msgType, xferFields, dx, dy, dz, doRecv, planeOnly)
   comm = domain.comm
   if comm === nothing
       return
   end
   comm = comm::MPI.Comm

   # post recieve buffers for all incoming messages
   maxPlaneComm = xferFields * domain.maxPlaneSize
   maxEdgeComm  = xferFields * domain.maxEdgeSize
   pmsg = 0 # plane comm msg
   emsg = 0 # edge comm msg
   cmsg = 0 # corner comm msg

   baseType = MPI.Datatype(Float64) # TODO support Float32

   # assume communication to 6 neighbors by default
   rowMin, rowMax, colMin, colMax, planeMin, planeMax = get_neighbors(domain)

   fill!(domain.recvRequest, MPI.Request())

   myRank = MPI.Comm_rank(comm)

   # post receives
   function irecv!(fromProc, offset, recvCount)
      idx = offset + 1
      data = view(domain.commDataRecv, idx:(idx+recvCount))
      return MPI.Irecv!(MPI.Buffer(data), fromProc, msgType, domain.comm::MPI.Comm)
   end

   # receive data from neighboring domain faces
   if planeMin && doRecv
      # contiguous memory
      fromProc = myRank - domain.m_tp^2
      recvCount = dx * dy * xferFields
      req = irecv!(fromProc, pmsg*maxPlaneComm, recvCount)
      domain.recvRequest[pmsg+1] = req
      pmsg += 1
   end

   if planeMax
      # contiguous memory
      fromProc = myRank + domain.m_tp^2
      recvCount = dx * dy * xferFields
      req = irecv!(fromProc, pmsg*maxPlaneComm, recvCount)
      domain.recvRequest[pmsg+1] = req
      pmsg += 1
   end

   if rowMin && doRecv
      # semi-contiguous memory
      fromProc = myRank - domain.m_tp
      recvCount = dx * dz * xferFields
      req = irecv!(fromProc, pmsg*maxPlaneComm, recvCount)
      domain.recvRequest[pmsg+1] = req
      pmsg += 1
   end

   if rowMax
      # semi-contiguous memory
      fromProc = myRank + domain.m_tp
      recvCount = dx * dz * xferFields
      req = irecv!(fromProc, pmsg*maxPlaneComm, recvCount)
      domain.recvRequest[pmsg+1] = req
      pmsg += 1
   end

   if colMin && doRecv
      # scattered memory
      fromProc = myRank - 1
      recvCount = dy * dz * xferFields
      req = irecv!(fromProc, pmsg*maxPlaneComm, recvCount)
      domain.recvRequest[pmsg+1] = req
      pmsg += 1
   end

   if colMax
      # scattered memory
      fromProc = myRank + 1
      recvCount = dy * dz * xferFields
      req = irecv!(fromProc, pmsg*maxPlaneComm, recvCount)
      domain.recvRequest[pmsg+1] = req
      pmsg += 1
   end

   if !planeOnly
      # receive data from domains connected only by an edge
      if rowMin && colMin && doRecv
         fromProc = myRank - domain.m_tp - 1
         recvCount = dz * xferFields
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         req = irecv!(fromProc, offset, recvCount)
         domain.recvRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      if rowMin && planeMin && doRecv
         fromProc = myRank - domain.m_tp^2 - domain.m_tp
         recvCount = dx * xferFields
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         req = irecv!(fromProc, offset, recvCount)
         domain.recvRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      if colMin && planeMin && doRecv
         fromProc = myRank - domain.m_tp^2 - 1
         recvCount = dy * xferFields
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         req = irecv!(fromProc, offset, recvCount)
         domain.recvRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      if rowMax && colMax
         fromProc = myRank + domain.m_tp + 1
         recvCount = dz * xferFields
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         req = irecv!(fromProc, offset, recvCount)
         domain.recvRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      if rowMax && planeMax
         fromProc = myRank + domain.m_tp^2 + domain.m_tp
         recvCount = dx * xferFields
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         req = irecv!(fromProc, offset, recvCount)
         domain.recvRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      if colMax && planeMax
         fromProc = myRank + domain.m_tp^2 + 1
         recvCount = dy * xferFields
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         req = irecv!(fromProc, offset, recvCount)
         domain.recvRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      if rowMax && colMin
         fromProc = myRank + domain.m_tp - 1
         recvCount = dz * xferFields
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         req = irecv!(fromProc, offset, recvCount)
         domain.recvRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      if rowMin && planeMax
         fromProc = myRank + domain.m_tp^2 - domain.m_tp
         recvCount = dx * xferFields
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         req = irecv!(fromProc, offset, recvCount)
         domain.recvRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      if colMin && planeMax
         fromProc = myRank + domain.m_tp^2 - 1
         recvCount = dy * xferFields
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         req = irecv!(fromProc, offset, recvCount)
         domain.recvRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      if rowMin && colMax && doRecv
         fromProc = myRank - domain.m_tp + 1
         recvCount = dz * xferFields
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         req = irecv!(fromProc, offset, recvCount)
         domain.recvRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      if rowMax && planeMin && doRecv
         fromProc = myRank - domain.m_tp^2 + domain.m_tp
         recvCount = dx * xferFields
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         req = irecv!(fromProc, offset, recvCount)
         domain.recvRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      if colMax && planeMin && doRecv
         fromProc = myRank - domain.m_tp^2 + 1
         recvCount = dy * xferFields
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         req = irecv!(fromProc, offset, recvCount)
         domain.recvRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      # receive data from domains connected only by a corner
      if rowMin && colMin && planeMin && doRecv
         # corner at domain logical coord (0, 0, 0)
         fromProc = myRank - domain.m_tp^2 - domain.m_tp - 1
         recvCount = xferFields
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
         req = irecv!(fromProc, offset, recvCount)
         domain.recvRequest[pmsg+emsg+cmsg+1] = req
         cmsg += 1
      end

      if rowMin && colMin && planeMax
         # corner at domain logical coord (0, 0, 1)
         fromProc = myRank + domain.m_tp^2 - domain.m_tp - 1
         recvCount = xferFields
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
         req = irecv!(fromProc, offset, recvCount)
         domain.recvRequest[pmsg+emsg+cmsg+1] = req
         cmsg += 1
      end

      if rowMin && colMax && planeMin && doRecv
         # corner at domain logical coord (1, 0, 0)
         fromProc = myRank - domain.m_tp^2 - domain.m_tp + 1
         recvCount = xferFields
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
         req = irecv!(fromProc, offset, recvCount)
         domain.recvRequest[pmsg+emsg+cmsg+1] = req
         cmsg += 1
      end

      if rowMin && colMax && planeMax
         # corner at domain logical coord (1, 0, 1)
         fromProc = myRank + domain.m_tp^2 - domain.m_tp + 1
         recvCount = xferFields
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
         req = irecv!(fromProc, offset, recvCount)
         domain.recvRequest[pmsg+emsg+cmsg+1] = req
         cmsg += 1
      end

      if rowMax && colMin && planeMin && doRecv
         # corner at domain logical coord (0, 1, 0)
         fromProc = myRank - domain.m_tp^2 + domain.m_tp - 1
         recvCount = xferFields
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
         req = irecv!(fromProc, offset, recvCount)
         domain.recvRequest[pmsg+emsg+cmsg+1] = req
         cmsg += 1
      end

      if rowMax && colMin && planeMax
         # corner at domain logical coord (0, 1, 1)
         fromProc = myRank + domain.m_tp^2 + domain.m_tp - 1
         recvCount = xferFields
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
         req = irecv!(fromProc, offset, recvCount)
         domain.recvRequest[pmsg+emsg+cmsg+1] = req
         cmsg += 1
      end

      if rowMax && colMax && planeMin && doRecv
         # corner at domain logical coord (1, 1, 0)
         fromProc = myRank - domain.m_tp^2 + domain.m_tp + 1
         recvCount = xferFields
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
         req = irecv!(fromProc, offset, recvCount)
         domain.recvRequest[pmsg+emsg+cmsg+1] = req
         cmsg += 1
      end

      if rowMax && colMax && planeMax
         # corner at domain logical coord (1, 1, 1)
         fromProc = myRank + domain.m_tp^2 + domain.m_tp + 1
         recvCount = xferFields
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
         req = irecv!(fromProc, offset, recvCount)
         domain.recvRequest[pmsg+emsg+cmsg+1] = req
         cmsg += 1
      end
   end
end

function commSend(domain::Domain, msgType, fields,
                  dx, dy, dz, doSend, planeOnly)

   comm = domain.comm
   if comm === nothing
      return
   end

   xferFields = length(fields)

   # post recieve buffers for all incoming messages
   maxPlaneComm = xferFields * domain.maxPlaneSize
   maxEdgeComm  = xferFields * domain.maxEdgeSize
   pmsg = 0 # plane comm msg
   emsg = 0 # edge comm msg
   cmsg = 0 # corner comm msg

   # MPI_Status status[26] ;

   # assume communication to 6 neighbors by default
   rowMin,rowMax, colMin, colMax, planeMin, planeMax = get_neighbors(domain)

   fill!(domain.sendRequest, MPI.Request())
   myRank = MPI.Comm_rank(comm)

   # post sends
   if planeMin || planeMax
      # ASSUMING ONE DOMAIN PER RANK, CONSTANT BLOCK SIZE HERE
      sendCount = dx * dy

      if planeMin
         # contiguous memory
         srcOffset = 0
         offset = pmsg * maxPlaneComm
         for field in fields
            copyto_zero!(domain.commDataSend, offset, field, srcOffset, sendCount)
            offset += sendCount
         end
         idx = pmsg * maxPlaneComm + 1
         src = view(domain.commDataSend, idx:(idx+(xferFields * sendCount)))

         otherRank = myRank - domain.m_tp^2
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+1] = req
         pmsg += 1
      end

      if planeMax && doSend
         # contiguous memory
         srcOffset = dx*dy*(dz - 1)
         offset = pmsg * maxPlaneComm
         for field in fields
            copyto_zero!(domain.commDataSend, offset, field, srcOffset, sendCount)
            offset += sendCount
         end
         idx = pmsg * maxPlaneComm + 1
         src = view(domain.commDataSend, idx:(idx+(xferFields * sendCount)))

         otherRank = myRank + domain.m_tp^2
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+1] = req
         pmsg += 1
      end
   end

   if rowMin | rowMax
      # ASSUMING ONE DOMAIN PER RANK, CONSTANT BLOCK SIZE HERE
      sendCount = dx * dz

      if rowMin
         # contiguous memory
         srcOffset = 0
         offset = pmsg * maxPlaneComm
         for field in fields
            for i in 0:(dz-1)
               copyto_zero!(domain.commDataSend, offset+i*dx, field, srcOffset+i*dx*dy, dx)
            end
            offset += sendCount
         end
         idx = pmsg * maxPlaneComm + 1
         src = view(domain.commDataSend, idx:(idx+(xferFields * sendCount)))

         otherRank = myRank - domain.m_tp
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+1] = req
         pmsg += 1
      end
      if rowMax && doSend
         # contiguous memory
         srcOffset = dx*(dy - 1)
         offset = pmsg * maxPlaneComm
         for field in fields
            for i in 0:(dz-1)
               copyto_zero!(domain.commDataSend, offset+i*dx , field, srcOffset+i*dx*dy, dx)
            end
            offset += sendCount
         end
         idx = pmsg * maxPlaneComm + 1
         src = view(domain.commDataSend, idx:(idx+(xferFields * sendCount)))

         otherRank = myRank + domain.m_tp
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+1] = req
         pmsg += 1
      end
   end

   if colMin | colMax
      # ASSUMING ONE DOMAIN PER RANK, CONSTANT BLOCK SIZE HERE
      sendCount = dy * dz

      if colMin
         # contiguous memory
         srcOffset = 0
         offset = pmsg * maxPlaneComm
         for field in fields
            for i in 0:(dz-1)
               for j in 0:(dy-1)
                  domain.commDataSend[offset+i*dy+j + 1] = field[srcOffset+i*dx*dy+j*dx + 1]
               end
            end
            offset += sendCount
         end
         idx = pmsg * maxPlaneComm + 1
         src = view(domain.commDataSend, idx:(idx+(xferFields * sendCount)))

         otherRank = myRank - 1
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+1] = req
         pmsg += 1
      end

      if colMax && doSend
         # contiguous memory
         srcOffset = dx - 1
         offset = pmsg * maxPlaneComm
         for field in fields
            for i in 0:(dz-1)
               for j in 0:(dy-1)
                  domain.commDataSend[offset+i*dy+j + 1] = field[srcOffset+j*dx + 1]
               end
            end
            offset += sendCount
         end
         idx = pmsg * maxPlaneComm + 1
         src = view(domain.commDataSend, idx:(idx+(xferFields * sendCount)))
         otherRank = myRank + 1
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+1] = req
         pmsg += 1
      end
   end
   if !planeOnly
      if rowMin && colMin
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         for field in fields
            for i in 0:(dz-1)
               domain.commDataSend[offset+i + 1] = field[i*dx*dy + 1]
            end
            offset += dz
         end
         idx = pmsg * maxPlaneComm + emsg * maxEdgeComm + 1
         src = view(domain.commDataSend, idx:(idx+(xferFields * dz)))
         otherRank = myRank - domain.m_tp - 1

         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      if rowMin && planeMin
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         for field in fields
            for i in 0:(dx-1)
               domain.commDataSend[offset+i + 1] = field[i + 1]
            end
            offset += dx
         end
         idx = pmsg * maxPlaneComm + emsg * maxEdgeComm + 1
         src = view(domain.commDataSend, idx:(idx+(xferFields * dx)))
         otherRank = myRank - domain.m_tp^2 - domain.m_tp
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      if colMin && planeMin
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         for field in fields
            for i in 0:(dy-1)
               domain.commDataSend[offset+i + 1] = field[i*dx + 1]
            end
            offset += dy
         end
         idx = pmsg * maxPlaneComm + emsg * maxEdgeComm + 1
         src = view(domain.commDataSend, idx:(idx+(xferFields * dy)))
         otherRank = myRank - domain.m_tp^2 - 1
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      if rowMax && colMax && doSend
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         srcOffset = dx*dy - 1
         for field in fields
            for i in 0:(dz-1)
               domain.commDataSend[offset+i + 1] = field[srcOffset+i*dx*dy + 1]
            end
            offset += dz
         end
         idx = pmsg * maxPlaneComm + emsg * maxEdgeComm + 1
         src = view(domain.commDataSend, idx:(idx+(xferFields * dz)))
         otherRank = myRank + domain.m_tp + 1
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      if rowMax && planeMax && doSend
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         srcOffset = dx*(dy-1) + dx*dy*(dz-1)
         for field in fields
            for i in 0:(dx-1)
               domain.commDataSend[offset+i + 1] = field[srcOffset+i + 1]
            end
            offset += dx
         end
         idx = pmsg * maxPlaneComm + emsg * maxEdgeComm + 1
         src = view(domain.commDataSend, idx:(idx+(xferFields * dx)))
         otherRank = myRank + domain.m_tp^2 + domain.m_tp
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      if colMax && planeMax && doSend
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         srcOffset = dx*dy*(dz-1) + dx - 1
         for field in fields
            for i in 0:(dy-1)
               domain.commDataSend[offset+i + 1] = field[srcOffset+i*dx + 1]
            end
            offset += dy
         end
         idx = pmsg * maxPlaneComm + emsg * maxEdgeComm + 1
         src = view(domain.commDataSend, idx:(idx+(xferFields * dy)))
         otherRank = myRank + domain.m_tp^2 + 1
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      if rowMax && colMin && doSend
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         srcOffset = dx*(dy-1)
         for field in fields
            for i in 0:(dz-1)
               domain.commDataSend[offset+i + 1] = field[srcOffset+i*dx*dy + 1]
            end
            offset += dz
         end
         idx = pmsg * maxPlaneComm + emsg * maxEdgeComm + 1
         src = view(domain.commDataSend, idx:(idx+(xferFields * dz)))
         otherRank = myRank + domain.m_tp - 1
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      if rowMin && planeMax && doSend
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         srcOffset = dx*dy*(dz-1)
         for field in fields
            for i in 0:(dx-1)
               domain.commDataSend[offset+i + 1] = field[srcOffset+i + 1]
            end
            offset += dx
         end
         idx = pmsg * maxPlaneComm + emsg * maxEdgeComm + 1
         src = view(domain.commDataSend, idx:(idx+(xferFields * dx)))
         otherRank = myRank + domain.m_tp^2 - domain.m_tp
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      if colMin && planeMax && doSend
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         srcOffset = dx*dy*(dz-1)
         for field in fields
            for i in 0:(dy-1)
               domain.commDataSend[offset+i + 1] = field[srcOffset+i*dx + 1]
            end
            offset += dy
         end
         idx = pmsg * maxPlaneComm + emsg * maxEdgeComm + 1
         src = view(domain.commDataSend, idx:(idx+(xferFields * dy)))
         otherRank = myRank + domain.m_tp^2 - 1
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      if rowMin && colMax
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         srcOffset = dx - 1
         for field in fields
            for i in 0:(dz-1)
               domain.commDataSend[offset+i + 1] = field[srcOffset+i*dx*dy + 1]
            end
            offset += dz
         end
         idx = pmsg * maxPlaneComm + emsg * maxEdgeComm + 1
         src = view(domain.commDataSend, idx:(idx+(xferFields * dz)))
         otherRank = myRank - domain.m_tp + 1
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      if rowMax && planeMin
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         srcOffset = dx*(dy - 1)
         for field in fields
            for i in 0:(dx-1)
               domain.commDataSend[offset+i + 1] = field[srcOffset+i + 1]
            end
            offset += dx
         end
         idx = pmsg * maxPlaneComm + emsg * maxEdgeComm + 1
         src = view(domain.commDataSend, idx:(idx+(xferFields * dx)))
         otherRank = myRank - domain.m_tp^2 + domain.m_tp
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      if colMax && planeMin
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
         srcOffset = dx - 1
         for field in fields
            for i in 0:(dy-1)
               domain.commDataSend[offset+i + 1] = field[srcOffset+i*dx + 1]
            end
            offset += dy
         end
         idx = pmsg * maxPlaneComm + emsg * maxEdgeComm + 1
         src = view(domain.commDataSend, idx:(idx+(xferFields * dy)))
         otherRank = myRank - domain.m_tp^2 + 1
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+emsg+1] = req
         emsg += 1
      end

      if rowMin && colMin && planeMin
         # corner at domain logical coord (0, 0, 0) */
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm +  cmsg * CACHE_COHERENCE_PAD_REAL
         for (fi, field) in enumerate(fields)
            # fi is one based already
            domain.commDataSend[offset+fi] = field[1]
         end
         idx = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL + 1
         src = view(domain.commDataSend, idx:(idx+xferFields))
         otherRank = myRank - domain.m_tp^2 - domain.m_tp - 1
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+emsg+cmsg+1] = req
         cmsg += 1
      end

      if rowMin && colMin && planeMax && doSend
         # corner at domain logical coord (0, 0, 1)
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
         for (fi, field) in enumerate(fields)
            # fi is one based already
            domain.commDataSend[offset+fi] = field[dx*dy*(dz-1) + 1]
         end
         idx = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL + 1
         src = view(domain.commDataSend, idx:(idx+xferFields))
         otherRank = myRank + domain.m_tp^2 - domain.m_tp - 1
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+emsg+cmsg+1] = req
         cmsg += 1
      end

      if rowMin && colMax && planeMin
         # corner at domain logical coord (1, 0, 0)
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
         for (fi, field) in enumerate(fields)
            # fi is one based already
            domain.commDataSend[offset+fi] = field[(dx - 1) + 1]
         end
         idx = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL + 1
         src = view(domain.commDataSend, idx:(idx+xferFields))
         otherRank = myRank - domain.m_tp^2 - domain.m_tp + 1
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+emsg+cmsg+1] = req
         cmsg += 1
      end

      if rowMin && colMax && planeMax && doSend
         # corner at domain logical coord (1, 0, 1)
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
         for (fi, field) in enumerate(fields)
            # fi is one based already
            domain.commDataSend[offset+fi] = field[dx*dy*(dz - 1) + (dx - 1) + 1]
         end
         idx = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL + 1
         src = view(domain.commDataSend, idx:(idx+xferFields))
         otherRank = myRank + domain.m_tp^2 - domain.m_tp + 1
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+emsg+cmsg+1] = req
         cmsg += 1
      end

      if rowMax && colMin && planeMin
         # corner at domain logical coord (0, 1, 0)
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
         for (fi, field) in enumerate(fields)
            # fi is one based already
            domain.commDataSend[offset+fi] = field[dx*(dy - 1) + 1]
         end
         idx = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL + 1
         src = view(domain.commDataSend, idx:(idx+xferFields))
         otherRank = myRank - domain.m_tp^2 + domain.m_tp - 1
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+emsg+cmsg+1] = req
         cmsg += 1
      end

      if rowMax && colMin && planeMax && doSend
         # corner at domain logical coord (0, 1, 1)
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
         for (fi, field) in enumerate(fields)
            # fi is one based already
            domain.commDataSend[offset+fi] = field[dx*dy*(dz - 1) + dx*(dy - 1)+ 1]
         end
         idx = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL + 1
         src = view(domain.commDataSend, idx:(idx+xferFields))
         otherRank = myRank + domain.m_tp^2 + domain.m_tp - 1
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+emsg+cmsg+1] = req
         cmsg += 1
      end

      if rowMax && colMax && planeMin
         # corner at domain logical coord (1, 1, 0)
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
         for (fi, field) in enumerate(fields)
            # fi is one based already
            domain.commDataSend[offset+fi] = field[dx*dy - 1 + 1]
         end
         idx = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL + 1
         src = view(domain.commDataSend, idx:(idx+xferFields))
         otherRank = myRank - domain.m_tp^2 + domain.m_tp + 1
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+emsg+cmsg+1] = req
         cmsg += 1
      end

      if rowMax && colMax && planeMax && doSend
         # corner at domain logical coord (1, 1, 1)
         offset = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
         for (fi, field) in enumerate(fields)
            # fi is one based already
            domain.commDataSend[offset+fi] = field[dx*dy*dz - 1 + 1]
         end
         idx = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL + 1
         src = view(domain.commDataSend, idx:(idx+xferFields))
         otherRank = myRank + domain.m_tp^2 + domain.m_tp + 1
         req = MPI.Isend(MPI.Buffer(src), otherRank, msgType, comm)
         domain.sendRequest[pmsg+emsg+cmsg+1] = req
         cmsg += 1
      end
   end

   for i in 1:(pmsg+emsg+cmsg)
      MPI.Wait!(domain.sendRequest[i])
   end
   # MPI.Waitall!(domain.sendRequest)
end

function commSBN(domain::Domain, fields)
   comm = domain.comm
   if comm === nothing
      return
   end

   xferFields = length(fields)

   maxPlaneComm = xferFields * domain.maxPlaneSize
   maxEdgeComm  = xferFields * domain.maxEdgeSize
   pmsg = 0 # plane comm msg
   emsg = 0 # edge comm msg
   cmsg = 0 # corner comm msg

   dx = domain.sizeX + 1
   dy = domain.sizeY + 1
   dz = domain.sizeZ + 1

   rowMin, rowMax, colMin, colMax, planeMin, planeMax = get_neighbors(domain)

   myRank = MPI.Comm_rank(comm)

   if planeMin || planeMax
      opCount = dx * dy ;

      if planeMin
         srcaddr = pmsg * maxPlaneComm;
         status = MPI.Wait!(domain.recvRequest[pmsg + 1]);
         for field in fields
            for i in 1:opCount
               field[i] += domain.commDataRecv[srcaddr+i]
            end
            srcaddr += opCount
         end
         pmsg += 1
      end
      if planeMax
         srcaddr = pmsg * maxPlaneComm;
         status = MPI.Wait!(domain.recvRequest[pmsg + 1]);
         for field in fields
            for i in 1:opCount
               field[dx*dy*(dz-1) + i] += domain.commDataRecv[srcaddr+i]
            end
            srcaddr += opCount
         end
         pmsg += 1
      end
   end

   if rowMin || rowMax
      opCount = dx * dz

      if rowMin
         srcaddr = pmsg * maxPlaneComm;
         status = MPI.Wait!(domain.recvRequest[pmsg + 1]);
         for field in fields
            for i in 0:(dz-1)
               for j in 0:(dx-1)
                  field[i*dx*dy + j + 1] += domain.commDataRecv[srcaddr + i*dx + j + 1]
               end
            end
            srcaddr += opCount
         end
         pmsg += 1
      end

      if rowMax
         srcaddr = pmsg * maxPlaneComm;
         status = MPI.Wait!(domain.recvRequest[pmsg + 1]);
         for field in fields
            for i in 0:(dz-1)
               for j in 0:(dx-1)
                  field[dx*(dy-1) + i*dx*dy + j + 1] += domain.commDataRecv[srcaddr + i*dx + j + 1]
               end
            end
            srcaddr += opCount
         end
         pmsg += 1
      end
   end

   if colMin || colMax
      opCount = dy * dz ;

      if colMin
         srcaddr = pmsg * maxPlaneComm;
         status = MPI.Wait!(domain.recvRequest[pmsg + 1]);
         for field in fields
            for i in 0:(dz-1)
               for j in 0:(dy-1)
                  field[i*dx*dy + j*dx + 1] += domain.commDataRecv[srcaddr + i*dy + j + 1]
               end
            end
            srcaddr += opCount
         end
         pmsg += 1
      end

      if colMax
         srcaddr = pmsg * maxPlaneComm;
         status = MPI.Wait!(domain.recvRequest[pmsg + 1]);
         for field in fields
            for i in 0:(dz-1)
               for j in 0:(dy-1)
                  field[dx - 1 + i*dx*dy + j*dx + 1] += domain.commDataRecv[srcaddr + i*dy + j + 1]
               end
            end
            srcaddr += opCount
         end
         pmsg += 1
      end
   end

   if rowMin && colMin
      srcaddr = pmsg * maxPlaneComm + emsg * maxEdgeComm
      status = MPI.Wait!(domain.recvRequest[pmsg+emsg + 1])
      for field in fields
         for i in 0:(dz-1)
            # (domain.*dest)(i*dx*dy) += srcaddr[i] ;
            field[i*dx*dy + 1] += domain.commDataRecv[srcaddr + i + 1]
         end
         srcaddr += dz
      end
      emsg += 1
   end

   if rowMin && planeMin
      srcaddr = pmsg * maxPlaneComm + emsg * maxEdgeComm
      status = MPI.Wait!(domain.recvRequest[pmsg+emsg + 1])
      for field in fields
         for i in 0:(dx-1)
            # (domain.*dest)(i) += srcaddr[i] ;
            field[i + 1] += domain.commDataRecv[srcaddr + i + 1]
         end
         srcaddr += dx
      end
      emsg += 1
   end

   if colMin && planeMin
      srcaddr = pmsg * maxPlaneComm + emsg * maxEdgeComm
      status = MPI.Wait!(domain.recvRequest[pmsg+emsg + 1])
      for field in fields
         for i in 0:(dy-1)
            # (domain.*dest)(i*dx) += srcaddr[i] ;
            field[i*dx + 1] += domain.commDataRecv[srcaddr + i + 1]
         end
         srcaddr += dy
      end
      emsg += 1
   end

   if rowMax && colMax
      srcaddr = pmsg * maxPlaneComm + emsg * maxEdgeComm
      status = MPI.Wait!(domain.recvRequest[pmsg+emsg + 1])
      for field in fields
         for i in 0:(dz-1)
            # (domain.*dest)(dx*dy - 1 + i*dx*dy) += srcaddr[i] ;
            field[dx*dy - 1 + i*dx*dy + 1] += domain.commDataRecv[srcaddr + i + 1]
         end
         srcaddr += dz
      end
      emsg += 1
   end

   if rowMax && planeMax
      srcaddr = pmsg * maxPlaneComm + emsg * maxEdgeComm
      status = MPI.Wait!(domain.recvRequest[pmsg+emsg + 1])
      for field in fields
         for i in 0:(dx-1)
            # (domain.*dest)(dx*(dy-1) + dx*dy*(dz-1) + i) += srcaddr[i] ;
            field[dx*(dy-1) + dx*dy*(dz-1) + i + 1] += domain.commDataRecv[srcaddr + i + 1]
         end
         srcaddr += dx
      end
      emsg += 1
   end

   if colMax && planeMax
      srcaddr = pmsg * maxPlaneComm + emsg * maxEdgeComm
      status = MPI.Wait!(domain.recvRequest[pmsg+emsg + 1])
      for field in fields
         for i in 0:(dy-1)
            # (domain.*dest)(dx*dy*(dz-1) + dx - 1 + i*dx) += srcaddr[i] ;
            field[dx*dy*(dz-1) + dx - 1 + i*dx + 1] += domain.commDataRecv[srcaddr + i + 1]
         end
         srcaddr += dy
      end
      emsg += 1
   end

   if rowMax && colMin
      srcaddr = pmsg * maxPlaneComm + emsg * maxEdgeComm
      status = MPI.Wait!(domain.recvRequest[pmsg+emsg + 1])
      for field in fields
         for i in 0:(dz-1)
            # (domain.*dest)(dx*(dy-1) + i*dx*dy) += srcaddr[i] ;
            field[dx*(dy-1) + i*dx*dy + 1] += domain.commDataRecv[srcaddr + i + 1]
         end
         srcaddr += dz
      end
      emsg += 1
   end

   if rowMin && planeMax
      srcaddr = pmsg * maxPlaneComm + emsg * maxEdgeComm
      status = MPI.Wait!(domain.recvRequest[pmsg+emsg + 1])
      for field in fields
         for i in 0:(dx-1)
            # (domain.*dest)(dx*dy*(dz-1) + i) += srcaddr[i] ;
            field[dx*dy*(dz-1) + i + 1] += domain.commDataRecv[srcaddr + i + 1]
         end
         srcaddr += dx
      end
      emsg += 1
   end

   if colMin && planeMax
      srcaddr = pmsg * maxPlaneComm + emsg * maxEdgeComm
      status = MPI.Wait!(domain.recvRequest[pmsg+emsg + 1])
      for field in fields
         for i in 0:(dy-1)
            # (domain.*dest)(dx*dy*(dz-1) + i*dx) += srcaddr[i] ;
            field[dx*dy*(dz-1) + i*dx + 1] += domain.commDataRecv[srcaddr + i + 1]
         end
         srcaddr += dy
      end
      emsg += 1
   end

   if rowMin && colMax
      srcaddr = pmsg * maxPlaneComm + emsg * maxEdgeComm
      status = MPI.Wait!(domain.recvRequest[pmsg+emsg + 1])
      for field in fields
         for i in 0:(dz-1)
            # (domain.*dest)(dx - 1 + i*dx*dy) += srcaddr[i] ;
            field[dx - 1 + i*dx*dy + 1] += domain.commDataRecv[srcaddr + i + 1]
         end
         srcaddr += dz
      end
      emsg += 1
   end

   if rowMax && planeMin
      srcaddr = pmsg * maxPlaneComm + emsg * maxEdgeComm
      status = MPI.Wait!(domain.recvRequest[pmsg+emsg + 1])
      for field in fields
         for i in 0:(dx-1)
            # (domain.*dest)(dx*(dy - 1) + i) += srcaddr[i] ;
            field[dx*(dy-1) + i + 1] += domain.commDataRecv[srcaddr + i + 1]
         end
         srcaddr += dx
      end
      emsg += 1
   end

   if colMax && planeMin
      srcaddr = pmsg * maxPlaneComm + emsg * maxEdgeComm
      status = MPI.Wait!(domain.recvRequest[pmsg+emsg + 1])
      for field in fields
         for i in 0:(dy-1)
            # (domain.*dest)(dx - 1 + i*dx) += srcaddr[i] ;
            field[dx - 1 + i*dx + 1] += domain.commDataRecv[srcaddr + i + 1]
         end
         srcaddr += dy
      end
      emsg += 1
   end

   if rowMin && colMin && planeMin
      idx = 0
      comBuf = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
      status = MPI.Wait!(domain.recvRequest[pmsg+emsg+cmsg + 1])
      for (fi, field) in enumerate(fields)
         field[idx+1] += domain.commDataRecv[comBuf + fi] # No 1 here as already taken care of by fi
      end
      cmsg += 1
   end

   if rowMin && colMin && planeMax
      idx = dx*dy*(dz - 1)
      comBuf = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
      status = MPI.Wait!(domain.recvRequest[pmsg+emsg+cmsg + 1])
      for (fi, field) in enumerate(fields)
         field[idx+1] += domain.commDataRecv[comBuf + fi] # No 1 here as already taken care of by fi
      end
      cmsg += 1
   end

   if rowMin && colMax && planeMin
      idx = dx - 1
      comBuf = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
      status = MPI.Wait!(domain.recvRequest[pmsg+emsg+cmsg + 1])
      for (fi, field) in enumerate(fields)
         field[idx+1] += domain.commDataRecv[comBuf + fi] # No 1 here as already taken care of by fi
      end
      cmsg += 1
   end

   if rowMin && colMax && planeMax
      idx = dx*dy*(dz - 1) + (dx - 1)
      comBuf = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
      status = MPI.Wait!(domain.recvRequest[pmsg+emsg+cmsg + 1])
      for (fi, field) in enumerate(fields)
         field[idx+1] += domain.commDataRecv[comBuf + fi] # No 1 here as already taken care of by fi
      end
      cmsg += 1
   end

   if rowMax && colMin && planeMin
      idx = dx*(dy - 1)
      comBuf = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
      status = MPI.Wait!(domain.recvRequest[pmsg+emsg+cmsg + 1])
      for (fi, field) in enumerate(fields)
         field[idx+1] += domain.commDataRecv[comBuf + fi] # No 1 here as already taken care of by fi
      end
      cmsg += 1
   end

   if rowMax && colMin && planeMax
      idx = dx*dy*(dz - 1) + dx*(dy - 1)
      comBuf = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
      status = MPI.Wait!(domain.recvRequest[pmsg+emsg+cmsg + 1])
      for (fi, field) in enumerate(fields)
         field[idx+1] += domain.commDataRecv[comBuf + fi] # No 1 here as already taken care of by fi
      end
      cmsg += 1
   end

   if rowMax && colMax && planeMin
      idx = dx*dy - 1
      comBuf = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
      status = MPI.Wait!(domain.recvRequest[pmsg+emsg+cmsg + 1])
      for (fi, field) in enumerate(fields)
         field[idx+1] += domain.commDataRecv[comBuf + fi] # No 1 here as already taken care of by fi
      end
      cmsg += 1
   end

   if rowMax && colMax && planeMax
      idx = dx*dy*dz - 1
      comBuf = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
      status = MPI.Wait!(domain.recvRequest[pmsg+emsg+cmsg + 1])
      for (fi, field) in enumerate(fields)
         field[idx+1] += domain.commDataRecv[comBuf + fi] # No 1 here as already taken care of by fi
      end
      cmsg += 1
   end
end

function commMonoQ(domain::Domain)
   comm = domain.comm
   if comm === nothing
       return
   end

   xferFields = 3 # delv_xi, delv_eta, delv_zeta
   fields = (domain.delv_xi, domain.delv_eta, domain.delv_zeta)
   maxPlaneComm = xferFields * domain.maxPlaneSize

   pmsg = 0 # plane comm msg
   dx = domain.sizeX
   dy = domain.sizeY
   dz = domain.sizeZ

   # assume communication to 6 neighbors by default
   rowMin, rowMax, colMin, colMax, planeMin, planeMax = get_neighbors(domain)

   numElem = domain.numElem

   # point into ghost data area
   fieldOffsets = [numElem, numElem, numElem]

   myRank = MPI.Comm_rank(comm)

   if planeMin || planeMax
      # ASSUMING ONE DOMAIN PER RANK, CONSTANT BLOCK SIZE HERE */
      opCount = dx * dy

      if planeMin
         # contiguous memory
         MPI.Wait!(domain.recvRequest[pmsg+1])

         offset = pmsg * maxPlaneComm
         for (fi, field) in enumerate(fields)
            fOffset = fieldOffsets[fi]
            copyto_zero!(field, fOffset, domain.commDataRecv, offset, opCount)

            offset += opCount
            fieldOffsets[fi] += opCount
         end
         pmsg += 1
      end
      if planeMax
         # contiguous memory
         MPI.Wait!(domain.recvRequest[pmsg+1])
         offset = pmsg * maxPlaneComm
         for (fi, field) in enumerate(fields)
            fOffset = fieldOffsets[fi]
            copyto_zero!(field, fOffset, domain.commDataRecv, offset, opCount)

            offset += opCount
            fieldOffsets[fi] += opCount
         end
         pmsg += 1
      end
   end

   if rowMin || rowMax
      # ASSUMING ONE DOMAIN PER RANK, CONSTANT BLOCK SIZE HERE
      opCount = dx * dz

      if rowMin
         # contiguous memory
         MPI.Wait!(domain.recvRequest[pmsg+1])
         offset = pmsg * maxPlaneComm
         for (fi, field) in enumerate(fields)
            fOffset = fieldOffsets[fi]
            copyto_zero!(field, fOffset, domain.commDataRecv, offset, opCount)

            offset += opCount
            fieldOffsets[fi] += opCount
         end
         pmsg += 1
      end

      if rowMax
         # contiguous memory
         MPI.Wait!(domain.recvRequest[pmsg+1])
         offset = pmsg * maxPlaneComm
         for (fi, field) in enumerate(fields)
            fOffset = fieldOffsets[fi]
            copyto_zero!(field, fOffset, domain.commDataRecv, offset, opCount)
            offset += opCount
            fieldOffsets[fi] += opCount
         end
         pmsg += 1
      end
   end
   if colMin || colMax
      # ASSUMING ONE DOMAIN PER RANK, CONSTANT BLOCK SIZE HERE
      opCount = dy * dz

      if colMin
         # contiguous memory
         MPI.Wait!(domain.recvRequest[pmsg+1])
         offset = pmsg * maxPlaneComm
         for (fi, field) in enumerate(fields)
            fOffset = fieldOffsets[fi]
            copyto_zero!(field, fOffset, domain.commDataRecv, offset, opCount)

            offset += opCount
            fieldOffsets[fi] += opCount
         end
         pmsg += 1
      end
      if colMax
         # contiguous memory
         MPI.Wait!(domain.recvRequest[pmsg+1])
         offset = pmsg * maxPlaneComm
         for (fi, field) in enumerate(fields)
            fOffset = fieldOffsets[fi]
            copyto_zero!(field, fOffset, domain.commDataRecv, offset, opCount)

            offset += opCount
            fieldOffsets[fi] += opCount
         end
         pmsg += 1
      end
   end
   return nothing
end

function commSyncPosVel(domain::Domain)
   comm = domain.comm
   if comm === nothing
      return
   end
   doRecv = false
   xferFields = 6 ; # x, y, z, xd, yd, zd
   fields = (domain.x, domain.y, domain.z, domain.xd, domain.yd, domain.zd)
   maxPlaneComm = xferFields * domain.maxPlaneSize
   maxEdgeComm  = xferFields * domain.maxEdgeSize
   pmsg = 0 # plane comm msg
   emsg = 0 # edge comm msg
   cmsg = 0 # corner comm msg
   dx = domain.sizeX + 1
   dy = domain.sizeY + 1
   dz = domain.sizeZ + 1

   # assume communication to 6 neighbors by default
   rowMin, rowMax, colMin, colMax, planeMin, planeMax = get_neighbors(domain)

   myRank = MPI.Comm_rank(comm)

   if planeMin | planeMax
      # ASSUMING ONE DOMAIN PER RANK, CONSTANT BLOCK SIZE HERE
      opCount = dx * dy

      if planeMin && doRecv
         # contiguous memory
         MPI.Wait!(domain.recvRequest[pmsg+1])
         offset = pmsg * maxPlaneComm
         for field in fields
            copyto_zero!(field, 0, domain.commDataRecv, offset, opCount)
            offset += opCount
         end
         pmsg += 1
      end

      if planeMax
         # contiguous memory
         MPI.Wait!(domain.recvRequest[pmsg+1])

         destOffset = dx*dy*(dz - 1)
         offset = pmsg * maxPlaneComm
         for field in fields
            copyto_zero!(field, destOffset, domain.commDataRecv, offset, opCount)
            offset += opCount
         end
         pmsg += 1
      end

   end
   if rowMin | rowMax
      # ASSUMING ONE DOMAIN PER RANK, CONSTANT BLOCK SIZE HERE
      opCount = dx * dz

      if rowMin && doRecv
         # contiguous memory
         MPI.Wait!(domain.recvRequest[pmsg+1])
         offset = pmsg * maxPlaneComm
         for field in fields
            for i in 0:(dz-1)
               destOffset = i*dx*dy
               copyto_zero!(field, destOffset, domain.commDataRecv, offset + i*dx, dx)
            end
            offset += opCount
         end
         pmsg += 1
      end
      if rowMax
         # contiguous memory
         MPI.Wait!(domain.recvRequest[pmsg+1])
         offset = pmsg * maxPlaneComm
         for field in fields
            for i in 0:(dz-1)
               destOffset = dx*(dy - 1) + i*dx*dy
               for j in 0:(dx-1)
                  field[destOffset + j + 1] = domain.commDataRecv[offset + i*dx + j + 1]
               end
            end
            offset += opCount
         end
         pmsg += 1
      end
   end
   if colMin | colMax
      # ASSUMING ONE DOMAIN PER RANK, CONSTANT BLOCK SIZE HERE
      opCount = dy * dz

      if colMin && doRecv
         # contiguous memory
         MPI.Wait!(domain.recvRequest[pmsg+1])
         offset = pmsg * maxPlaneComm
         for field in fields
            for i in 0:(dz-1)
               for j in 0:(dy-1)
                  field[i*dx*dy+j*dx + 1] = domain.commDataRecv[offset + i*dy + j + 1]
               end
            end
            offset += opCount
         end
         pmsg += 1
      end
      if colMax
         # contiguous memory
         MPI.Wait!(domain.recvRequest[pmsg+1])
         offset = pmsg * maxPlaneComm
         destOffset = dx-1
         for field in fields
            for i in 0:(dz-1)
               for j in 0:(dy-1)
                  field[destOffset + i*dx*dy + j*dx + 1] = domain.commDataRecv[offset + i*dy + j + 1]
               end
            end
            offset += opCount
         end
         pmsg += 1
      end
   end

   if rowMin && colMin && doRecv
      MPI.Wait!(domain.recvRequest[pmsg+emsg+1])
      offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
      for field in fields
         for i in 0:(dz-1)
            field[i*dx*dy + 1] = domain.commDataRecv[offset + i + 1]
         end
         offset += dz
      end
      emsg += 1
   end

   if rowMin && planeMin && doRecv
      MPI.Wait!(domain.recvRequest[pmsg+emsg+1])
      offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
      for field in fields
         copyto_zero!(field, 0, domain.commDataRecv, offset, dx)
         offset += dx
      end
      emsg += 1
   end

   if colMin && planeMin && doRecv
      MPI.Wait!(domain.recvRequest[pmsg+emsg+1])
      offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
      for field in fields
         for i in 0:(dy-1)
            field[i*dx + 1] = domain.commDataRecv[offset + i + 1]
         end
         offset += dy
      end
      emsg += 1
   end

   if rowMax && colMax
      MPI.Wait!(domain.recvRequest[pmsg+emsg+1])
      offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
      destOffset = dx*dy - 1
      for field in fields
         for i in 0:(dz-1)
            field[destOffset + i*dx*dy + 1] = domain.commDataRecv[offset + i + 1]
         end
         offset += dz
      end
      emsg += 1
   end

   if rowMax && planeMax
      MPI.Wait!(domain.recvRequest[pmsg+emsg+1])
      offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
      destOffset = dx*(dy-1) + dx*dy*(dz-1)
      for field in fields
         copyto_zero!(field, destOffset, domain.commDataRecv, offset, dx)
         offset += dx
      end
      emsg += 1
   end

   if colMax && planeMax
      MPI.Wait!(domain.recvRequest[pmsg+emsg+1])
      offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
      destOffset = dx*dy*(dz-1) + dx - 1
      for field in fields
         for i in 0:(dy-1)
            field[destOffset + i*dx + 1] = domain.commDataRecv[offset + i + 1]
         end
         offset += dy
      end
      emsg += 1
   end

   if rowMax && colMin
      MPI.Wait!(domain.recvRequest[pmsg+emsg+1])
      offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
      destOffset = dx*(dy-1)
      for field in fields
         for i in 0:(dz-1)
            field[destOffset + i*dx*dy + 1] = domain.commDataRecv[offset + i + 1]
         end
         offset += dz
      end
      emsg += 1
   end

   if rowMin && planeMax
      MPI.Wait!(domain.recvRequest[pmsg+emsg+1])
      offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
      destOffset = dx*dy*(dz-1)
      for field in fields
         copyto_zero!(field, destOffset, domain.commDataRecv, offset, dx)
         offset += dx
      end
      emsg += 1
   end

   if colMin && planeMax
      MPI.Wait!(domain.recvRequest[pmsg+emsg+1])
      offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
      destOffset = dx*dy*(dz-1)
      for field in fields
         for i in 0:(dy-1)
            field[destOffset + i*dx + 1] = domain.commDataRecv[offset + i + 1]
         end
         offset += dy
      end
      emsg += 1
   end

   if rowMin && colMax && doRecv
      MPI.Wait!(domain.recvRequest[pmsg+emsg+1])
      offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
      destOffset = dx-1
      for field in fields
         for i in 0:(dz-1)
            field[destOffset + i*dx*dy + 1] = domain.commDataRecv[offset + i + 1]
         end
         offset += dz
      end
      emsg += 1
   end

   if rowMax && planeMin && doRecv
      MPI.Wait!(domain.recvRequest[pmsg+emsg+1])
      offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
      destOffset = dx*(dy - 1)
      for field in fields
         copyto_zero!(field, destOffset, domain.commDataRecv, offset, dx)
         offset += dx
      end
      emsg += 1
   end

   if colMax && planeMin && doRecv
      MPI.Wait!(domain.recvRequest[pmsg+emsg+1])
      offset = pmsg * maxPlaneComm + emsg * maxEdgeComm
      destOffset = dx - 1
      for field in fields
         for i in 0:(dy-1)
            field[destOffset + i*dx + 1] = domain.commDataRecv[offset + i + 1]
         end
         offset += dy
      end
      emsg += 1
   end

   if rowMin && colMin && planeMin && doRecv
      # corner at domain logical coord (0, 0, 0)
      MPI.Wait!(domain.recvRequest[pmsg+emsg+cmsg+1])
      offset = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
      for (fi, field) in enumerate(fields)
         # fi is one based
         field[1] = domain.commDataRecv[offset + fi]
      end
      cmsg += 1
   end

   if rowMin && colMin && planeMax
      # corner at domain logical coord (0, 0, 1)
      MPI.Wait!(domain.recvRequest[pmsg+emsg+cmsg+1])
      offset = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
      idx = dx*dy*(dz - 1)
      for (fi, field) in enumerate(fields)
         # fi is one based
         field[idx + 1] = domain.commDataRecv[offset + fi]
      end
      cmsg += 1
   end

   if rowMin && colMax && planeMin && doRecv
      # corner at domain logical coord (1, 0, 0)
      MPI.Wait!(domain.recvRequest[pmsg+emsg+cmsg+1])
      offset = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
      idx = dx - 1
      for (fi, field) in enumerate(fields)
         # fi is one based
         field[idx + 1] = domain.commDataRecv[offset + fi]
      end
      cmsg += 1
   end

   if rowMin && colMax && planeMax
      # corner at domain logical coord (1, 0, 1)
      MPI.Wait!(domain.recvRequest[pmsg+emsg+cmsg+1])
      offset = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
      idx = dx*dy*(dz - 1) + (dx - 1)
      for (fi, field) in enumerate(fields)
         # fi is one based
         field[idx + 1] = domain.commDataRecv[offset + fi]
      end
      cmsg += 1
   end

   if rowMax && colMin && planeMin && doRecv
      # corner at domain logical coord (0, 1, 0)
      MPI.Wait!(domain.recvRequest[pmsg+emsg+cmsg+1])
      offset = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
      idx = dx*(dy - 1)
      for (fi, field) in enumerate(fields)
         # fi is one based
         field[idx + 1] = domain.commDataRecv[offset + fi]
      end
      cmsg += 1
   end

   if rowMax && colMin && planeMax
      # corner at domain logical coord (0, 1, 1)
      MPI.Wait!(domain.recvRequest[pmsg+emsg+cmsg+1])
      offset = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
      idx = dx*dy*(dz - 1) + dx*(dy - 1)
      for (fi, field) in enumerate(fields)
         # fi is one based
         field[idx + 1] = domain.commDataRecv[offset + fi]
      end
      cmsg += 1
   end

   if rowMax && colMax && planeMin && doRecv
      # corner at domain logical coord (1, 1, 0)
      MPI.Wait!(domain.recvRequest[pmsg+emsg+cmsg+1])
      offset = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
      idx = dx*dy - 1
      for (fi, field) in enumerate(fields)
         # fi is one based
         field[idx + 1] = domain.commDataRecv[offset + fi]
      end
      cmsg += 1
   end

   if rowMax && colMax && planeMax
      # corner at domain logical coord (1, 1, 1)
      MPI.Wait!(domain.recvRequest[pmsg+emsg+cmsg+1])
      offset = pmsg * maxPlaneComm + emsg * maxEdgeComm + cmsg * CACHE_COHERENCE_PAD_REAL
      idx = dx*dy*dz - 1
      for (fi, field) in enumerate(fields)
         # fi is one based
         field[idx + 1] = domain.commDataRecv[offset + fi]
      end
      cmsg += 1
   end

   return nothing
end


    ####################################################################
#######                Main Muffin function                  #######
####################################################################

function muffin(;folder="",dataobj="2gauss",datapsf="",nitermax = 500, rhop = 1,
                rhot = 5, rhov = 2, rhos = 1, μt = 5e-1, μv = 1e-0, mueps = 1e-3,
                bw = 5, ws="",parallel="")

println("")
println("MUFFIN initialisation")

                 ##################################
    ################### data initialisation #################
                 ##################################

    if typeof(dataobj) == ASCIIString
        if dataobj == "m31"
            psf = "data/meerkat_m30_25pix.psf.fits"
            obj = "data/M31.fits"
            tmp = string(Pkg.dir("Muffin"))
            psf = string(tmp,tmp[1],psf)
            obj = string(tmp,tmp[1],obj)
        elseif dataobj == "andro"
            psf = "data/meerkat_m30_25pix.psf.fits"
            obj = "data/andro.fits"
            tmp = string(Pkg.dir("Muffin"))
            psf = string(tmp,tmp[1],psf)
            obj = string(tmp,tmp[1],obj)
        elseif dataobj == "2gauss"
            psf = "data/meerkat_m30_25pix.psf.fits"
            obj = "data/2gauss.fits"
            tmp = string(Pkg.dir("Muffin"))
            psf = string(tmp,tmp[1],psf)
            obj = string(tmp,tmp[1],obj)
        elseif dataobj == "chiara"
            psf = "/home/deguignet/Julia/example_sim_psf.fits"
            obj = "/home/deguignet/Julia/example_sim_dirty.fits"
        elseif isempty(folder)
            tmp = pwd()
            psf = string(tmp,tmp[1],datapsf)
            obj = string(tmp,tmp[1],dataobj)
        elseif typeof(folder) == ASCIIString
                 psf = string(folder,folder[1],datapsf)
                 obj = string(folder,folder[1],dataobj)
        else
                 error("data folder is not correct")
        end

    elseif isempty(dataobj)
        psf = "data/meerkat_m30_25pix.psf.fits"
        obj = "data/M31.fits"
    end

println("psf something :"," ",psf)
println("obj :"," ",obj)

                ##################################
    ################# Structure initialisation #################
                ##################################

if dataobj == "chiara"
    ##################################
    println("loading psf chiara...")
    psfst = loadpsf_dirty(psf)
    println("loading sky...")
    skyst = loadsky_dirty(obj,psfst.mypsf,psfst.nu)
    ##################################
else
    ##################################
    println("loading psf other...")
    psfst = loadpsf(psf,bw)
    println("loading sky...")
    skyst = loadsky(obj,psfst.mypsf,psfst.nu)
    ##################################
end

    ##################################
    spatialwlt  = [WT.db1,WT.db2,WT.db3,WT.db4,WT.db5,WT.db6,WT.db7,WT.db8]
    # spatialwlt  = [WT.db1,WT.db2,WT.db3,WT.db4,WT.db5,WT.db6,WT.db7,WT.db8,WT.haar]

    const nspat = length(spatialwlt)
    const nfreq = size(psfst.mypsf)[3]
    const nspec = 1
    const nxy = size(skyst.mydata)[1]
    niter = 0
    lastiter = 0
    #################################

    #################################
    println("loading param...")
    algost = loadparam(nspat,nfreq,nspec,nxy,niter,lastiter,nitermax)

    println("loading arrays...")
    if ws == "true"
        file = string("/home/deguignet/resultats_400ite_woshar.jld")
        # file = string("/Users/deguignet/test.jld")
        admmst = loadarray_ws(rhop,rhot,rhov,rhos,μt,μv,mueps,nspat,nfreq,nxy,
                            skyst.mydata,psfst.mypsfadj,file)
    algost.niter = load(file,"algost.lastiter")
    toolst = loadtools(nitermax,nfreq,nxy)
    toolst.tol1 = load(file,"toolst.tol1")
    toolst.tol2 = load(file,"toolst.tol2")
    toolst.tol3 = load(file,"toolst.tol3")
    toolst.tol4 = load(file,"toolst.tol4")
    toolst.tol5 = load(file,"toolst.tol5")
    skyst.mydata = load(file,"skyst.mydata")

    else
        admmst = loadarray(rhop,rhot,rhov,rhos,μt,μv,mueps,nspat,nfreq,nxy,
                            skyst.mydata,psfst.mypsfadj,parallel)
        toolst = loadtools(nitermax,nfreq,nxy)
    end

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

                 ##################################
    #####################  Main Admm Loop  #####################
                 ##################################
    println("Starting ADMM")
    #################################
    psfst, skyst, algost, admmst, toolst = muffinadmm(psfst, skyst, algost, admmst, toolst)
    #################################

    return psfst, skyst, algost, admmst, toolst
end

####################################################################
#######                   Main Admm Loop                     #######
####################################################################

function muffinadmm(psfst, skyst, algost, admmst, toolst)

    const rhop = admmst.rhop
    const rhot = admmst.rhot
    const rhov = admmst.rhov
    const rhos = admmst.rhos
    const μt = admmst.μt
    const μv = admmst.μv
    const mueps = admmst.mueps
    const tt = admmst.tt
    const mu = admmst.mu
    const nspat = algost.nspat
    const nfreq = algost.nfreq
    const nspec = algost.nspec
    const nxy = algost.nxy
    const fty = admmst.fty
    const nitermax = algost.nitermax

    spatialwlt  = [WT.db1,WT.db2,WT.db3,WT.db4,WT.db5,WT.db6,WT.db7,WT.db8,WT.haar]

    niter = algost.niter

    loop = true

    println("")
    println("                                =======================                                ")
    println("                               | MUFFIN RECONSTRUCTION |                               ")

    println("========================================================================================")
    println("ADMM"," ","|"," ","||x - xm||"," ","|"," ","||x - xp||"," ","|"," ",
            "||Hx - t||"," ","|"," ","||x - s||"," "," |"," ","||sh - v||"," ","|"," ","     time      |")
    println("========================================================================================")
    tic()
        while loop
            tic()
            niter +=1
            println("    "," ","|"," ","          "," ","|"," ","          "," ","|"," ",
                    "          "," ","|"," ","         "," "," |"," ","          "," ","|"," ")
            # println("ADMM iteration: $niter")

    ######################################
    ######################################
            # tic()
            # @parallel for i in 2:nfreq+1
            #     z = i-1
            #     admmst.wlt[:,:,z],admmst.x[:,:,z],admmst.t[:,:,z,:],admmst.taut[:,:,z,:],admmst.p[:,:,z],admmst.taup[:,:,z] =
            #                                 @fetchfrom(i,parallelmuffin(admmst.wlt[:,:,z], admmst.taut[:,:,z,:], admmst.t[:,:,z,:], rhot, admmst.x[:,:,z],
            #                                 psfst.mypsf[:,:,z], admmst.p[:,:,z], admmst.taup[:,:,z],
            #                                 fty[:,:,z], rhop, admmst.taus[:,:,z], admmst.s[:,:,z], rhos, admmst.mu, spatialwlt,
            #                                 μt, nspat))
########################

            @sync @parallel for z in 1:nfreq
                admmst.wlt[:,:,z],admmst.x[:,:,z],admmst.t[:,:,z,:],admmst.taut[:,:,z,:],admmst.p[:,:,z],admmst.taup[:,:,z] =
                                            parallelmuffin(admmst.wlt[:,:,z], admmst.taut[:,:,z,:], admmst.t[:,:,z,:], rhot, admmst.x[:,:,z],
                                            psfst.mypsf[:,:,z], admmst.p[:,:,z], admmst.taup[:,:,z],
                                            fty[:,:,z], rhop, admmst.taus[:,:,z], admmst.s[:,:,z], rhos, admmst.mu, spatialwlt,
                                            μt, nspat)

            end

            ##############################
            ######### prox spec ##########

            tmp = admmst.tauv + rhov*admmst.v

            admmst.s, admmst.sh = estime_ssh(admmst.s,admmst.sh,tmp,nxy,nspec,admmst.spectralwlt,
                                              admmst.x,admmst.taus,rhov,rhos)

            tmp = admmst.sh - admmst.tauv/rhov
            admmst.v = prox_u(tmp,μv/rhov)

            ########################################
            #### update of Lagrange multipliers ####

            admmst.tauv = admmst.tauv + rhov*(admmst.v-admmst.sh)
            admmst.taus = admmst.taus + rhos*(admmst.s-admmst.x)

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

            #
            # ##############################
            # ########## update x ##########
            # for z in 1:nfreq
            #     admmst.wlt[:,:,z] = myidwt((admmst.wlt)[:,:,z], nspat, (admmst.taut)[:,:,z,:], rhot,
            #                         (admmst.t)[:,:,z,:], spatialwlt)
            # end
            #
            # b = fty + admmst.taup + rhop*admmst.p + admmst.taus + rhos*admmst.s
            #
            # admmst.x = estime_x_par(admmst.x,psfst.mypsf,admmst.wlt + b,mu,nfreq)
            # ##############################
            # ######### prox spat ##########
            # tmp1 = 0.0
            # tmp2 = zeros(Float64,nxy,nxy)
            # for z in 1:nfreq
            #     for b in 1:nspat
            #             hx = dwt(admmst.x[:,:,z],wavelet(spatialwlt[b]))
            #             tmp = hx - admmst.taut[:,:,z,b]/rhot
            #             admmst.t[:,:,z,b] = prox_u(tmp,μt/rhot)
            #             admmst.taut[:,:,z,b] = admmst.taut[:,:,z,b] + rhot*(admmst.t[:,:,z,b]-hx)
            #             tmp1 = vecnorm([tmp2 (hx-(admmst.t)[:,:,z,b])],2)
            #             tmp2 = (hx-(admmst.t)[:,:,z,b])
            #     end
            # end
            # tmp2[:] = 0
            # ##############################
            # ###### prox positivity #######
            #
            # tmp = admmst.x-admmst.taup/rhop
            #
            # admmst.p = max(0,tmp)
            #
            # ##############################
            # ######### prox spec ##########
            #
            # tmp = admmst.tauv + rhov*admmst.v
            #
            # admmst.s, admmst.sh = estime_ssh(admmst.s,admmst.sh,tmp,nxy,nspec,admmst.spectralwlt,
            #                                   admmst.x,admmst.taus,rhov,rhos)
            #
            # tmp = admmst.sh - admmst.tauv/rhov
            # admmst.v = prox_u(tmp,μv/rhov)
            #
            #
            # ########################################
            # #### update of Lagrange multipliers ####
            #
            # admmst.taup = admmst.taup + rhop*(admmst.p-admmst.x)
            # admmst.tauv = admmst.tauv + rhov*(admmst.v-admmst.sh)
            # admmst.taus = admmst.taus + rhos*(admmst.s-admmst.x)

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

            ##############################
            ##### computer residues ######

            push!(toolst.tol1,vecnorm(admmst.x - admmst.xmm, 2)^2)
            push!(toolst.tol2,vecnorm(admmst.x - admmst.p, 2)^2)
            # push!(toolst.tol3,tmp1^2)
            push!(toolst.tol4,vecnorm(admmst.x - admmst.s, 2)^2)
            push!(toolst.tol5,vecnorm(admmst.sh - admmst.v, 2)^2)

            # ##############################
            # ####### stopping rule ########

            if (niter >= nitermax) #|| ((toolst.tol1[niter] < 1E-6) && (toolst.tol2[niter] < 1E-4))
                loop = false
                algost.lastiter = niter
            end

            admmst.xmm[:] = admmst.x

            #
            # println("ADMM"," ","|"," ","||x - xm||"," ","|"," ","||x - xp||"," ","|"," ",
            #         "||Hx - t||"," ","|"," ","||x - s||"," ","|"," ","||sh - v||"," ","|"," ","time")
            # println(niter," ","|"," ",toolst.tol1[niter]," ","|"," ",toolst.tol2[niter]," ","|"," ",
            #         toolst.tol3[niter]," ","|"," ",toolst.tol4[niter]," ","|"," ",toolst.tol5[niter]," ","|"," ",toq())

            @printf("%03d  | ", niter)
            @printf("%02.04e | ", toolst.tol1[niter])
            @printf("%02.04e | ", toolst.tol2[niter])
            # @printf("%02.04e | ", toolst.tol3[niter])
            @printf("%02.04e | ", 0)
            @printf("%02.04e | ", toolst.tol4[niter])
            @printf("%02.04e | ", toolst.tol5[niter])
            @printf("%f seconds  \n", toq())

        end
    println("")
    @printf("time for ADMM : %f seconds \n", toq())
    ####################################################################
    ####################################################################
    ####################################################################

    return psfst, skyst, algost, admmst, toolst

end

####################################################################
#######                  Data functions                      #######
####################################################################

##################################
function cubefilter{T<:FloatingPoint}(imagecube::Array{T,3},psfcube::Array{T,3})
    nximg, nyimg, nfreq = size(imagecube)
    nxpsf, nypsf, nfreq = size(psfcube)
    rescube = Array(Float64,nximg,nyimg,nfreq)
    for k in 1:nfreq
        rescube[:,:,k] = imfilter_fft(imagecube[:,:,k],psfcube[:,:,k],"circular")
    end
    return rescube
end

function cubefilter{T<:FloatingPoint}(imagecube::SharedArray{T,3},psfcube::Array{T,3})
    nximg, nyimg, nfreq = size(imagecube)
    nxpsf, nypsf, nfreq = size(psfcube)
    rescube = Array(Float64,nximg,nyimg,nfreq)
    for k in 1:nfreq
        rescube[:,:,k] = imfilter_fft(imagecube[:,:,k],psfcube[:,:,k],"circular")
    end
    return rescube
end

function cubefilter{T<:FloatingPoint}(imagecube::Array{T,2},psfcube::Array{T,3})
    nximg, nyimg = size(imagecube)
    nxpsf, nypsf, nfreq = size(psfcube)
    rescube = Array(Float64,nximg,nyimg,nfreq)
    for k in 1:nfreq
        rescube[:,:,k] = imfilter_fft(imagecube,psfcube[:,:,k],"circular")
    end
    return rescube
end
##################################

##################################
function cubeaverage{T<:FloatingPoint}(imagecube::Array{T,3},M::Int)
    nxpsf, nypsf, nfreq = size(imagecube)
    if nfreq < M
        error("Not enough channels to average!")
    end
    nfreqavg = itrunc(nfreq/M)
    rescube = Array(Float64, nxpsf, nypsf, nfreqavg)
    for k in 1:nfreqavg
        rescube[:,:,k] = sum(imagecube[:,:,(k-1)*M+1:k*M], 3)/M
    end

    return rescube
end
##################################

##################################
function cubefreq(psf::ASCIIString,imagecube::Array,M::Int)
    nxpsf, nypsf, nfreq = size(imagecube)
    if nfreq < M
        error("Not enough channels to average!")
    end
    nfreqavg = itrunc(nfreq/M)
    file = FITS(psf)
    header = read_header(file[1])
    nustart = header["CRVAL4"]
    nustep = header["CDELT4"]
    nu0 = header["RESTFRQ"]
    nu = zeros(Float64,nfreqavg)
    for i in 1:nfreqavg
        nu[i] = nustart + (M-1)*nustep + (i-1)*M*nustep
    end
    close(file)
    return nu,nu0
end
##################################

##################################
function cropcubexy{T<:FloatingPoint}(imagecube::Array{T,3},M::Int)
    nxpsf, nypsf, nfreq = size(imagecube)
    if ~(nxpsf == nypsf)
        error("Image must be square!")
    end
    if nxpsf < M
        println("no crop")
        return imagecube
    end

    rescube = Array(Float64, M, M, nfreq)
    if (iseven(nxpsf) && iseven(M)) | (isodd(nxpsf) && isodd(M))
        nstart = (nxpsf-M)/2 +1
        nend = (nxpsf+M)/2
        rescube = imagecube[nstart:nend,nstart:nend,:]
    else
        error("size of image and crop must have same parity!")
    end
    return rescube
end
##################################

##################################
function sky2cube{T<:FloatingPoint}(sky::Array{T,2},nu::Array{T,1})
    nx, ny = size(sky)
    nbands = length(nu)
    corl = minimum([nx,ny])/5.0
    rho(x,y) = exp(-norm([x, y])/corl)
    tx = linspace(0,nx-1,nx)
    ty = linspace(0,ny-1,ny)
    field = genghf(tx, ty, rho)
    skycube = Array(Float64, nx, ny, nbands)

    sm, sM = extrema(sky)
    fm, fM = extrema(field)

    alpha = 0.8 + 0.4*(sky-sm)/(sM-sm) + 0.4*(field-fm)/(fM-fm)

    nu0 = (nu[end]+nu[1])/2
    for k in 1:nbands
        skycube[:,:,k] = sky.* (nu[k]/nu0) .^(-alpha)
    end
    return skycube,alpha
end
##################################

##################################
function lecture(directory::ASCIIString)
    file = FITS(directory)
    data = float64(read(file[1]))
    close(file)
    data = squeeze(data,find(([size(data)...].==1)))
    data = data[:,:,:]

    return data
end
##################################

####################################################################
#######                  Admm functions                      #######
####################################################################

#################################
######### x estimation ##########
# function estime_x_par(x::Array{Float64,3},mypsf::Array{Float64,3},
#                         wlt_b::Array{Float64,3},mu::Float64,nfreq::Int64)
#
#     nxy = (size(x))[1]
#     nxypsf = (size(mypsf))[1]
#     fftpsf = zeros(Complex64,nxy,nxy,nfreq)
#     psfcbe = zeros(Complex64,nxy,nxy,nfreq)
#     psfpad = zeros(Float64,nxy,nxy,nfreq)
#
#     for z in 1:nfreq
#         psfpad[1:nxypsf,1:nxypsf,z] = fftshift(mypsf[:,:,z])
#         psfcbe[:,:,z] = 1./(abs(fft(psfpad[:,:,z])).^2+mu)
#         x[:,:,z] = real(ifft(psfcbe[:,:,z].*fft(wlt_b[:,:,z])))
#     end
#
#     return x
# end
#
# function estime_x_par(x::SharedArray{Float64,3},mypsf::Array{Float64,3},
#                         wlt_b::SharedArray{Float64,3},mu::Float64,nfreq::Int64)
#
#     nxy = (size(x))[1]
#     nxypsf = (size(mypsf))[1]
#     fftpsf = zeros(Complex64,nxy,nxy,nfreq)
#     psfcbe = zeros(Complex64,nxy,nxy,nfreq)
#     psfpad = zeros(Float64,nxy,nxy,nfreq)
#
#     for z in 1:nfreq
#         psfpad[1:nxypsf,1:nxypsf,z] = fftshift(mypsf[:,:,z])
#         psfcbe[:,:,z] = 1./(abs(fft(psfpad[:,:,z])).^2+mu)
#         x[:,:,z] = real(ifft(psfcbe[:,:,z].*fft(wlt_b[:,:,z])))
#     end
#
#     return x
# end

#################################
###### proximity operators ######
function prox_u(u::Array,μ::Float64)
    return max(0, 1-μ./abs(u)).*u
end

function prox_u(u::SharedArray,μ::Float64)
    return max(0, 1-μ./abs(u)).*u
end

# function prox_u(u::Array,μ::Float64)
# 	res=zeros(size(u));
# 	x=u;
# 	ind = (abs(x).>(.75*μ^(2/3))) ;
# 	res[ind]= 2/3*x[ind].*(1+cos(2*pi/3-2/3*acos(μ/8*(abs(x[ind])/3).^(-3/2)))) ;
#    return res
# end

#################################
####### s / sh estimation #######
function estime_ssh(s::Array{Float64,3},sh::Array{Float64,3},tmp::Array{Float64,3},
                    nxy::Int64,nspec::Int64,spectralwlt::Array{Float64,3},
                    x::Array{Float64,3},taus::Array{Float64,3},rhov::Float64,rhos::Float64)

    spectralwlt = idct(tmp,3)
    s = (spectralwlt + rhos*x - taus)/(rhov*nspec + rhos)
    sh = dct(s,3)

    return s,sh
end

function estime_ssh(s::Array{Float64,3},sh::Array{Float64,3},tmp::Array{Float64,3},
                    nxy::Int64,nspec::Int64,spectralwlt::Array{Float64,3},
                    x::SharedArray{Float64,3},taus::Array{Float64,3},rhov::Float64,rhos::Float64)

    spectralwlt = idct(tmp,3)
    s = (spectralwlt + rhos*x - taus)/(rhov*nspec + rhos)
    sh = dct(s,3)

    return s,sh
end

function myidwt(wlt::Array{Float64,2},nspat::Int,taut::Array{Float64,4},rhot::Float64,t::Array{Float64,4},spatialwlt)
        wlt = idwt(taut[:,:,1,1] + rhot*t[:,:,1,1],wavelet(spatialwlt[1]))
            for b in 2:nspat
                wlt = wlt + idwt(taut[:,:,1,b] + rhot*t[:,:,1,b],wavelet(spatialwlt[b]))
            end
        return wlt
end
function myidwt(wlt,nspat::Int,taut,rhot::Float64,t,spatialwlt)
        wlt = idwt(taut[:,:,1,1] + rhot*t[:,:,1,1],wavelet(spatialwlt[1]))
            for b in 2:nspat
                wlt = wlt + idwt(taut[:,:,1,b] + rhot*t[:,:,1,b],wavelet(spatialwlt[b]))
            end
        return wlt
end

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

##################################
function cubefreqchiara(nfrequencies::Int)
    nfreq = nfrequencies
    nustart = 9.85e8
    nustep = 2e6

    nu = zeros(Float64,nfreq)
    for i in 1:nfreq
        nu[i] = nustart + (i-1)*nustep
    end
    nu0 = (nu[1] + nu[nfreq])/2

    return nu,nu0
end

function parallelmuffin(wlt::Array{Float64,2},taut::Array{Float64,4},t::Array{Float64,4},rhot::Float64,
                        x::Array{Float64,2},psf::Array{Float64,2},p::Array{Float64,2},taup::Array{Float64,2},
                        fty::Array{Float64,2},rhop::Float64,taus::Array{Float64,2},s::Array{Float64,2},rhos::Float64,
                        mu::Float64,spatialwlt,μt::Float64,nspat::Int)

        wlt = myidwt(wlt, nspat, taut[:,:,1,:], rhot, t[:,:,1,:], spatialwlt)
        b = fty + taup + rhop*p + taus + rhos*s
        wlt_b = wlt + b

        nxy = (size(x))[1]
        nxypsf = (size(psf))[1]
        psfcbe = zeros(Complex64,nxy,nxy)
        psfpad = zeros(Float64,nxy,nxy)
        psfpad[1:nxypsf,1:nxypsf] = psf[:,:]
        psfcbe = 1./(abs(fft(psfpad)).^2+mu)
        x = real(ifft(psfcbe.*fft(wlt_b)))

        # tmp1 = 0.0
        # tmp2 = zeros(Float64,nxy,nxy)
        for b in 1:nspat
                    hx = dwt(x,wavelet(spatialwlt[b]))
                    tmp = hx - taut[:,:,1,b]/rhot
                    t[:,:,1,b] = prox_u(tmp,μt/rhot)
                    taut[:,:,1,b] = taut[:,:,1,b] + rhot*(t[:,:,1,b]-hx)
                    # tmp1 = vecnorm([tmp2 (hx-(t)[:,:,z,b])],2)
                    # tmp2 = (hx-(t)[:,:,z,b])
        end
        # tmp2[:] = 0

        tmp = x-taup/rhop
        p = max(0,tmp)
        taup = taup + rhop*(p-x)

    return wlt,x,t,taut,p,taup

end

function parallelmuffin(wlt::SharedArray{Float64,2},taut::SharedArray{Float64,4},t::SharedArray{Float64,4},rhot::Float64,
                        x::SharedArray{Float64,2},psf::Array{Float64,2},p::SharedArray{Float64,2},taup::SharedArray{Float64,2},
                        fty::Array{Float64,2},rhop::Float64,taus::Array{Float64,2},s::Array{Float64,2},rhos::Float64,
                        mu::Float64,spatialwlt,μt::Float64,nspat::Int)

        wlt = myidwt(wlt, nspat, taut[:,:,1,:], rhot, t[:,:,1,:], spatialwlt)
        b = fty + taup + rhop*p + taus + rhos*s
        wlt_b = wlt + b

        nxy = (size(x))[1]
        nxypsf = (size(psf))[1]
        psfcbe = zeros(Complex64,nxy,nxy)
        psfpad = zeros(Float64,nxy,nxy)
        psfpad[1:nxypsf,1:nxypsf] = psf[:,:]
        psfcbe = 1./(abs(fft(psfpad)).^2+mu)
        x = real(ifft(psfcbe.*fft(wlt_b)))

        # tmp1 = 0.0
        # tmp2 = zeros(Float64,nxy,nxy)
        for b in 1:nspat
                    hx = dwt(x,wavelet(spatialwlt[b]))
                    tmp = hx - taut[:,:,1,b]/rhot
                    t[:,:,1,b] = prox_u(tmp,μt/rhot)
                    taut[:,:,1,b] = taut[:,:,1,b] + rhot*(t[:,:,1,b]-hx)
                    # tmp1 = vecnorm([tmp2 (hx-(t)[:,:,z,b])],2)
                    # tmp2 = (hx-(t)[:,:,z,b])
        end
        # tmp2[:] = 0

        tmp = x-taup/rhop
        p = max(0,tmp)
        taup = taup + rhop*(p-x)

    return wlt,x,t,taut,p,taup

end


# Very junky script testing performance against direct binary IO
#
# In this test, it's about 5x slower to write to ply than to directly dump the
# bytes, this seems to be due to the many calls to write()

using BufferedStreams

import PlyIO: save_ply, write_header, Ply, Element, ArrayProperty, ListProperty

ply = Ply()

nverts = 100000

vertex = Element("vertex",
                 ArrayProperty("x", randn(nverts)),
                 ArrayProperty("y", randn(nverts)),
                 ArrayProperty("z", randn(nverts)),
                 ArrayProperty("r", rand(nverts)),
                 ArrayProperty("g", rand(nverts)),
                 ArrayProperty("b", rand(nverts)))
push!(ply, vertex)

# Some triangular faces
vertex_index = ListProperty("vertex_index", Int32, Int32)
for i=1:nverts
   push!(vertex_index, rand(0:nverts-1,3))
end
push!(ply, Element("face", vertex_index))

#=
# Some edges
vertex_index = ListProperty("vertex_index", Int32, Int32)
for i=1:nverts
   push!(vertex_index, rand(0:nverts-1,2))
end
push!(ply, Element("edge", vertex_index))
=#

function save_ply_ref(ply, stream::IO)
    write_header(ply, stream, false)
    x = ply["vertex"]["x"].data::Vector{Float64}
    y = ply["vertex"]["y"].data::Vector{Float64}
    z = ply["vertex"]["z"].data::Vector{Float64}
    r = ply["vertex"]["r"].data::Vector{Float64}
    g = ply["vertex"]["g"].data::Vector{Float64}
    b = ply["vertex"]["b"].data::Vector{Float64}

    vertex_index = ply["face"]["vertex_index"]
    viinds = vertex_index.start_inds::Vector{Int32}
    vidata = vertex_index.data::Vector{Int32}

    #=
    for i=1:length(x)
        write(stream, x[i])
        write(stream, y[i])
        write(stream, z[i])
        write(stream, r[i])
        write(stream, g[i])
        write(stream, b[i])
    end

    for i=1:length(viinds)-1
        len = viinds[i+1] - viinds[i]
        write(stream, len)
        write(stream, vidata[viinds[i]:viinds[i+1]-1])
    end
    =#

    # Buffered reordering
    #=
    write(stream, [x y z r g b]')
    write(stream, viinds)
    write(stream, vidata)
    =#

    # Benchmark against direct binary IO (invalid ply!), which should be about
    # as fast as you can hope for.
    write(stream, x)
    write(stream, y)
    write(stream, z)
    write(stream, r)
    write(stream, g)
    write(stream, b)

    write(stream, viinds)
    write(stream, vidata)

end

function save_ply_ref(ply, filename::AbstractString)
    open(filename, "w") do fid
        save_ply_ref(ply, fid)
    end
end

function save_ply_buffered(ply, filename; kwargs...)
    open(filename, "w") do fid
        save_ply(ply, BufferedOutputStream(fid, 2^16); kwargs...)
    end
end

save_ply(ply, "test.ply", ascii=false)
save_ply_buffered(ply, "test.ply", ascii=false)

save_ply_ref(ply, "test2.ply")

Profile.clear_malloc_data()
@time save_ply(ply, "test.ply", ascii=false)
@time save_ply_buffered(ply, "test.ply", ascii=false)

@time save_ply_ref(ply, "test2.ply")

#run(`displaz -script test.ply`)


" solution status (from ipopt) "
ApplicationReturnStatus = Dict(
    0 => :Solve_Succeeded,
    1 => :Solved_To_Acceptable_Level,
    2 => :Infeasible_Problem_Detected,
    3 => :Search_Direction_Becomes_Too_Small,
    4 => :Diverging_Iterates,
    5 => :User_Requested_Stop,
    6 => :Feasible_Point_Found,
    -1 => :Maximum_Iterations_Exceeded,
    -2 => :Restoration_Failed,
    -3 => :Error_In_Step_Computation,
    -4 => :Maximum_CpuTime_Exceeded,
    -5 => :Optimize_not_called,
    -6 => :Method_not_defined,
    -10 => :Not_Enough_Degrees_Of_Freedom,
    -11 => :Invalid_Problem_Definition,
    -12 => :Invalid_Option,
    -13 => :Invalid_Number_Detected,
    -100 => :Unrecoverable_Exception,
    -102 => :Insufficient_Memory,
    -199 => :Internal_Error,
)


mutable struct PlateExtendState{U,V,W}
    score::Float64
    weight::Float64
    isempty::Bool
    args::U
    nodes::V
    changes::W
end

function extend_existing_trace!(gen::Plate{T,U}, key::Int, state::PlateExtendState, subtraces, retvals) where {T,U}
    node = haskey(state.nodes, key) ? state.nodes[key] : EmptyAssignment()
    if haskey(state.changes, key)
        kernel_change = state.changes[key]
    else
        kernel_change = nothing
    end
    subtrace::U = state.subtraces[key]
    prev_call = get_call_record(subtrace)
    kernel_args = get_args_for_key(state.args, key)
    (subtrace, kernel_weight) = extend(gen.kernel, kernel_args, kernel_change, subtrace, node)
    state.weight += kernel_weight
    call::CallRecord{T} = get_call_record(subtrace)
    state.score += (call.score - prev_call.score)
    subtraces = assoc(trace, key, subtrace)
    retvals = assoc(retvals, key, call.retval::T)
    state.is_empty = state.is_empty && !has_choices(subtraces)
    (subtraces, retvals)
end

function generate_new_trace!(gen::Plate{T,U}, key::Int, state::PlateExtendState, subtraces, retvals, kernel) where {T,U}
    node = haskey(state.nodes, key) ? state.nodes[key] : EmptyAssignment()
    kernel_args = get_args_for_key(state.args, key)
    (subtrace::U, kernel_weight) = generate(gen.kernel, kernel_args, node)
    state.weight += kernel_weight
    subtraces = push(subtraces, subtrace)
    call::CallRecord{T} = get_call_record(subtrace)
    state.score += call.score
    retvals = push(retvals, call.retval::T)
    (subtraces, retvals)
end

function extend_process_constraints!(nodes, key::Int, node, new_length, to_visit)
    if key <= new_length
        push!(to_visit, key)
        nodes[key] = node
    else
        error("Constraints under key $key not part of trace")
    end
end

function extend_process_constraints!(nodes, key, node, new_length, to_visit)
    error("Constraints under key $key not part of trace")
end

"""
Extend with argument change information
"""
function extend(gen::Plate, args, change::PlateChange, trace::VectorTrace, constraints) where {T}

    (new_length, prev_length) = get_prev_and_new_lengths(args, trace)
    if prev_length < new_length
        error("extend cannot remove random choices")
    end

    # calculate which existing applications to visit
    to_visit = Set{Int}()
    changes = Dict{Int, T}()
    for (i, kernel_change) in zip(change.changed_args, change.sub_changes)
        push!(to_visit, i)
        changes[i] = kernel_change
    end

    _extend(gen, args, trace, constraints, to_visit, changes)
end

"""
Extend without argument change information
"""
function extend(gen::Plate{T,U}, args, change::NoChange, trace::VectorTrace{T,U}, constraints) where {T,U}

    (new_length, prev_length) = get_prev_and_new_lengths(args, trace)

    # visit all existing applications
    to_visit = 1:new_length
    changes = Dict{Int, Any}() # not used

    _extend(gen, args, trace, constraints, to_visit, changes)
end

function _extend(gen::Plate, args, trace::VectorTrace, constraints, to_visit, changes)

    (new_length, prev_length) = get_prev_and_new_lengths(args, trace)
    if prev_length < new_length
        error("extend cannot remove random choices")
    end

    # collect constraints, indexed by key, and error if any are not in the new trace
    nodes = Dict{Int, Any}()
    for (key, node) in get_internal_nodes(constraints)
        extend_process_constraints!(nodes, key, node, new_length, to_visit)
    end

    # collect initial state
    state = PlateExtendState(trace.call.score, 0., trace.is_empty, args, nodes, changes)
    subtraces = trace.subtraces
    retvals = trace.call.retvals

    # visit all existing applications
    for key in to_visit
        (subtraces, retvals) = extend_existing_trace!(gen, key, state, subtraces, retvals)
    end

    # visit each new application and run generate on it
    for key=prev_length+1:new_length
        (subtraces, retvals) = generate_new_trace!(gen, key, state, subtraces, retvals)
    end

    call = CallRecord(state.score, retvals, args)
    trace = VectorTrace(subtraces, call, state.is_empty)
    retchange = nothing
    (trace, state.weight, retchange)
end