text:
export BlockSparseTensor,
       blockview,
       addblock!

#
# BlockSparseTensor (Tensor using BlockSparse storage)
#

const BlockSparseTensor{ElT,N,StoreT,IndsT} = Tensor{ElT,N,StoreT,IndsT} where {StoreT<:BlockSparse}

# Special version for BlockSparseTensor
# Generic version doesn't work since BlockSparse us parametrized by
# the Tensor order
function similar_type(::Type{<:Tensor{ElT,NT,<:BlockSparse{ElT,VecT},<:Any}},
                      ::Type{IndsR}) where {NT,ElT,VecT,IndsR}
  NR = length(IndsR)
  return Tensor{ElT,NR,BlockSparse{ElT,VecT,NR},IndsR}
end

function similar_type(::Type{<:Tensor{ElT,NT,<:BlockSparse{ElT,VecT},<:Any}},
                      ::Type{IndsR}) where {NT,ElT,VecT,IndsR<:NTuple{NR}} where {NR}
  return Tensor{ElT,NR,BlockSparse{ElT,VecT,NR},IndsR}
end

new_block_pos(T::BlockSparseTensor{ElT,N},
              block::Block{N}) where {ElT,N} = new_block_pos(blockoffsets(T),block)

function BlockSparseTensor(::Type{ElT},
                           ::UndefInitializer,
                           boffs::BlockOffsets{N},
                           inds) where {ElT<:Number,N}
  nnz_tot = nnz(boffs,inds)
  storage = BlockSparse(ElT,undef,boffs,nnz_tot)
  return tensor(storage,inds)
end

function BlockSparseTensor(::UndefInitializer,
                           blockoffsets::BlockOffsets{N},
                           inds) where {N}
  return BlockSparseTensor(Float64,undef,blockoffsets,inds)
end

function BlockSparseTensor(::Type{ElT},
                           blockoffsets::BlockOffsets{N},
                           inds) where {ElT<:Number,N}
  nnz_tot = nnz(blockoffsets,inds)
  storage = BlockSparse(ElT,blockoffsets,nnz_tot)
  return tensor(storage,inds)
end

function BlockSparseTensor(blockoffsets::BlockOffsets,
                           inds)
  return BlockSparseTensor(Float64,blockoffsets,inds)
end

"""
BlockSparseTensor(::UndefInitializer,
                  blocks::Vector{Block{N}},
                  inds)

Construct a block sparse tensor with uninitialized memory
from indices and locations of non-zero blocks.
"""
BlockSparseTensor(::UndefInitializer,
                  blocks::Blocks,
                  inds) = BlockSparseTensor(Float64,undef,blocks,inds)

function BlockSparseTensor(::Type{ElT},
                           ::UndefInitializer,
                           blocks::Blocks,
                           inds) where {ElT}
  boffs,nnz = blockoffsets(blocks,inds)
  storage = BlockSparse(ElT,undef,boffs,nnz)
  return tensor(storage,inds)
end

#function BlockSparseTensor(::UndefInitializer,
#                           blocks::Blocks{N},
#                           inds::Vararg{DimT,N}) where {DimT,N}
#  return BlockSparseTensor(undef,blocks,inds)
#end

"""
BlockSparseTensor(inds)

Construct a block sparse tensor with no blocks.
"""
function BlockSparseTensor(inds)
  return BlockSparseTensor(BlockOffsets{length(inds)}(),inds)
end

function BlockSparseTensor(::Type{ElT},
                           inds) where {ElT<:Number,N}
  return BlockSparseTensor(ElT,BlockOffsets{length(inds)}(),inds)
end

"""
BlockSparseTensor(inds)

Construct a block sparse tensor with no blocks.
"""
function BlockSparseTensor(inds::Vararg{DimT,N}) where {DimT,N}
  return BlockSparseTensor(BlockOffsets{N}(),inds)
end

"""
BlockSparseTensor(blocks::Vector{Block{N}},
                  inds)

Construct a block sparse tensor with the specified blocks.
Defaults to setting structurally non-zero blocks to zero.
"""
BlockSparseTensor(blocks::Blocks,
                  inds) = BlockSparseTensor(Float64,blocks,inds)

function BlockSparseTensor(::Type{ElT},
                           blocks::Blocks,
                           inds) where {ElT}
  boffs,nnz = blockoffsets(blocks,inds)
  storage = BlockSparse(ElT,boffs,nnz)
  return tensor(storage,inds)
end

#function BlockSparseTensor(blocks::Vector{Block{N}},
#                           inds) where {N}
#  blockoffsets,nnz = blockoffsets(blocks,inds)
#  storage = BlockSparse(blockoffsets,nnz)
#  return tensor(storage,inds)
#end

"""
BlockSparseTensor(blocks::Vector{Block{N}},
                  inds::BlockDims...)

Construct a block sparse tensor with the specified blocks.
Defaults to setting structurally non-zero blocks to zero.
"""
function BlockSparseTensor(blocks::Blocks{N},
                           inds::Vararg{BlockDim,N}) where {DimT,N}
  return BlockSparseTensor(blocks,inds)
end

function Base.similar(::BlockSparseTensor{ElT,N},
                      blockoffsets::BlockOffsets{N},
                      inds) where {ElT,N}
  return BlockSparseTensor(ElT,undef,blockoffsets,inds)
end

function Base.similar(::Type{<:BlockSparseTensor{ElT,N}},
                      blockoffsets::BlockOffsets{N},
                      inds) where {ElT,N}
  return BlockSparseTensor(ElT,undef,blockoffsets,inds)
end

function Base.zeros(::BlockSparseTensor{ElT,N},
                    blockoffsets::BlockOffsets{N},
                    inds) where {ElT,N}
  return BlockSparseTensor(ElT,blockoffsets,inds)
end

function Base.zeros(::Type{<:BlockSparseTensor{ElT,N}},
                    blockoffsets::BlockOffsets{N},
                    inds) where {ElT,N}
  return BlockSparseTensor(ElT,blockoffsets,inds)
end

# Basic functionality for AbstractArray interface
Base.IndexStyle(::Type{<:BlockSparseTensor}) = IndexCartesian()

# Get the CartesianIndices for the range of indices
# of the specified
function blockindices(T::BlockSparseTensor{ElT,N},
                      block) where {ElT,N}
  return CartesianIndex(blockstart(T,block)):CartesianIndex(blockend(T,block))
end

"""
indexoffset(T::BlockSparseTensor,i::Int...) -> offset,block,blockoffset

Get the offset in the data of the specified
CartesianIndex. If it falls in a block that doesn't
exist, return nothing for the offset.
Also returns the block the index is found in and the offset
within the block.
"""
function indexoffset(T::BlockSparseTensor{ElT,N},
                     i::Vararg{Int,N}) where {ElT,N}
  index_within_block,block = blockindex(T,i...)
  block_dims = blockdims(T,block)
  offset_within_block = LinearIndices(block_dims)[CartesianIndex(index_within_block)]
  offset_of_block = offset(T,block)
  offset_of_i = isnothing(offset_of_block) ? nothing : offset_of_block+offset_within_block
  return offset_of_i,block,offset_within_block
end

# TODO: Add a checkbounds
# TODO: write this nicer in terms of blockview?
#       Could write:
#       block,index_within_block = blockindex(T,i...)
#       return blockview(T,block)[index_within_block]
Base.@propagate_inbounds function Base.getindex(T::BlockSparseTensor{ElT,N},
                                                i::Vararg{Int,N}) where {ElT,N}
  offset,_ = indexoffset(T,i...)
  isnothing(offset) && return zero(ElT)
  return store(T)[offset]
end

Base.@propagate_inbounds function Base.getindex(T::BlockSparseTensor{ElT,0}) where {ElT}
  return store(T)[]
end

# These may not be valid if the Tensor has no blocks
#Base.@propagate_inbounds Base.getindex(T::BlockSparseTensor{<:Number,1},ind::Int) = store(T)[ind]

#Base.@propagate_inbounds Base.getindex(T::BlockSparseTensor{<:Number,0}) = store(T)[1]

# Add the specified block to the BlockSparseTensor
# Insert it such that the blocks remain ordered.
# Defaults to adding zeros.
function addblock!(T::BlockSparseTensor{ElT,N},
                   newblock::Block{N}) where {ElT,N}
  newdim = blockdim(T,newblock)
  newpos = new_block_pos(T,newblock)
  newoffset = 0
  if newpos!=1
    newoffset = offset(T,newpos-1)+blockdim(T,newpos-1)
  end
  # Insert new block into blockoffsets list
  insert!(blockoffsets(T),newpos,BlockOffset{N}(newblock,newoffset))
  # Insert new block into data
  splice!(data(store(T)),newoffset+1:newoffset,zeros(ElT,newdim))
  # Shift the offsets of the block after the inserted one
  for i in newpos+1:nnzblocks(T)
    block_i,offset_i = blockoffsets(T)[i]
    blockoffsets(T)[i] = BlockOffset{N}(block_i,offset_i+newdim)
  end
  return newoffset
end

# TODO: Add a checkbounds
Base.@propagate_inbounds function Base.setindex!(T::BlockSparseTensor{ElT,N},
                                                 val,
                                                 i::Vararg{Int,N}) where {ElT,N}
  offset,block,offset_within_block = indexoffset(T,i...)
  if isnothing(offset)
    offset_of_block = addblock!(T,block)
    offset = offset_of_block+offset_within_block
  end
  store(T)[offset] = val
  return T
end

"""
blockview(T::BlockSparseTensor,pos::Int)

Given a specified position in the block-offset list, return a Dense Tensor
that is a view to the data in that block (to avoid block lookup if the position
is known already).
"""
function blockview(T::BlockSparseTensor,
                   pos::Union{Int,Nothing})
  # TODO: don't allow nothing input
  isnothing(pos) && error("Block must be structurally non-zero to get a view")
  blockT,offsetT = blockoffsets(T)[pos]
  return blockview(T,BlockOffset(blockT,offsetT))
end

function blockview(T::BlockSparseTensor,
                   bof::BlockOffset)
  blockT,offsetT = bof
  blockdimsT = blockdims(T,blockT)
  blockdimT = prod(blockdimsT)
  dataTslice = @view data(store(T))[offsetT+1:offsetT+blockdimT]
  return tensor(Dense(dataTslice),blockdimsT)
end

# TODO: this is not working right now
#struct EachBlock{ElT,N,StoreT,IndsT}
#  T::BlockSparseTensor{ElT,N,StoreT,IndsT}
#end
#
#function Base.iterate(iter::EachBlock,state=(blockview(iter.T,1),1))
#  block,ind = state
#  ind > nnzblocks(iter.T) && return nothing
#  return blockview(iter.T,ind),ind+1
#end
#
#function eachnzblock(T::BlockSparseTensor)
#  return EachBlock(T)
#end

# convert to Dense
function dense(T::TensorT) where {TensorT<:BlockSparseTensor}
  R = zeros(dense(TensorT),dense(inds(T)))
  for (block,offset) in blockoffsets(T)
    # TODO: make sure this assignment is efficient
    R[blockindices(T,block)] = blockview(T,block)
  end
  return R
end

#
# Operations
#

# TODO: extend to case with different block structures
function Base.:+(T1::BlockSparseTensor,T2::BlockSparseTensor)
  inds(T1) ≠ inds(T2) && error("Cannot add block sparse tensors with different block structure")
  return tensor(store(T1)+store(T2),inds(T1))
end

function Base.permutedims(T::BlockSparseTensor{<:Number,N},
                          perm::NTuple{N,Int}) where {N}
  blockoffsetsR,indsR = permutedims(blockoffsets(T),inds(T),perm)
  R = similar(T,blockoffsetsR,indsR)
  permutedims!(R,T,perm)
  return R
end

function _permute_combdims(combdims::NTuple{NC,Int},
                           perm::NTuple{NP,Int}) where {NC,NP}
  res = MVector{NC,Int}(undef)
  iperm = invperm(perm)
  for i in 1:NC
    res[i] = iperm[combdims[i]]
  end
  return Tuple(res)
end

#
# These are functions to help with combining and uncombining
#

# Note that combdims is expected to be contiguous and ordered
# smallest to largest
function combine_dims(blocks::Blocks{N},
                      inds,
                      combdims::NTuple{NC,Int}) where {N,NC}
  nblcks = nblocks(inds,combdims)
  blocks_comb = Blocks{N-NC+1}(undef,nnzblocks(blocks))
  for (i,block) in enumerate(blocks)
    blocks_comb[i] = combine_dims(block,inds,combdims)
  end
  return blocks_comb
end

function combine_dims(block::Block,
                      inds,
                      combdims::NTuple{NC,Int}) where {NC}
  nblcks = nblocks(inds,combdims)
  slice = getindices(block,combdims)
  slice_comb = LinearIndices(nblcks)[slice...]
  block_comb = deleteat(block,combdims)
  block_comb = insertafter(block_comb,tuple(slice_comb),minimum(combdims)-1)
  return block_comb
end

# In the dimension dim, permute the blocks
function perm_blocks(blocks::Blocks{N},
                     dim::Int,
                     perm) where {N}
  blocks_perm = Blocks{N}(undef,nnzblocks(blocks))
  iperm = invperm(perm)
  for (i,block) in enumerate(blocks)
    blocks_perm[i] = setindex(block,iperm[block[dim]],dim)
  end
  return blocks_perm
end

# In the dimension dim, permute the block
function perm_block(block::Block,
                    dim::Int,
                    perm) where {N}
  iperm = invperm(perm)
  return setindex(block,iperm[block[dim]],dim)
end

# In the dimension dim, combine the specified blocks
function combine_blocks(blocks::Blocks,
                        dim::Int,
                        blockcomb::Vector{Int})
  blocks_comb = copy(blocks)
  nnz_comb = nnzblocks(blocks)
  for (i,block) in enumerate(blocks)
    dimval = block[dim]
    blocks_comb[i] = setindex(block,blockcomb[dimval],dim)
  end
  unique!(blocks_comb)
  return blocks_comb
end

function permutedims_combine_output(T::BlockSparseTensor{ElT,N},
                                    is,
                                    perm::NTuple{N,Int},
                                    combdims::NTuple{NC,Int},
                                    blockperm::Vector{Int},
                                    blockcomb::Vector{Int}) where {ElT,N,NC}
  # Permute the indices
  indsT = inds(T)
  inds_perm = permute(indsT,perm)

  # Now that the indices are permuted, compute
  # which indices are now combined
  combdims_perm = sort(_permute_combdims(combdims,perm))

  # Permute the nonzero blocks (dimension-wise)
  blocks = nzblocks(T)
  blocks_perm = permutedims(blocks,perm)

  # Combine the nonzero blocks (dimension-wise)
  blocks_perm_comb = combine_dims(blocks_perm,inds_perm,combdims_perm)

  # Permute the blocks (within the newly combined dimension)
  comb_ind_loc = minimum(combdims_perm)
  blocks_perm_comb = perm_blocks(blocks_perm_comb,comb_ind_loc,blockperm)
  blocks_perm_comb = sort(blocks_perm_comb;lt=isblockless)

  # Combine the blocks (within the newly combined and permuted dimension)
  blocks_perm_comb = combine_blocks(blocks_perm_comb,comb_ind_loc,blockcomb)

  return BlockSparseTensor(ElT,blocks_perm_comb,is)
end

function permutedims_combine(T::BlockSparseTensor{ElT,N},
                             is,
                             perm::NTuple{N,Int},
                             combdims::NTuple{NC,Int},
                             blockperm::Vector{Int},
                             blockcomb::Vector{Int}) where {ElT,N,NC}

  R = permutedims_combine_output(T,is,perm,combdims,blockperm,blockcomb)

  # Permute the indices
  inds_perm = permute(inds(T),perm)

  # Now that the indices are permuted, compute
  # which indices are now combined
  combdims_perm = sort(_permute_combdims(combdims,perm))
  comb_ind_loc = minimum(combdims_perm)

  # Determine the new index before combining
  inds_to_combine = getindices(inds_perm,combdims_perm)
  ind_comb = ⊗(inds_to_combine...)
  ind_comb = permuteblocks(ind_comb,blockperm)

  for bof in blockoffsets(T)
    Tb = blockview(T,bof)
    b = nzblock(bof)
    b_perm = permute(b,perm)
    b_perm_comb = combine_dims(b_perm,inds_perm,combdims_perm)
    b_perm_comb = perm_block(b_perm_comb,comb_ind_loc,blockperm)
    b_in_combined_dim = b_perm_comb[comb_ind_loc]
    new_b_in_combined_dim = blockcomb[b_in_combined_dim]
    offset = 0
    pos_in_new_combined_block = 1
    while b_in_combined_dim-pos_in_new_combined_block > 0 &&
            blockcomb[b_in_combined_dim-pos_in_new_combined_block] == new_b_in_combined_dim
      offset += blockdim(ind_comb,b_in_combined_dim-pos_in_new_combined_block)
      pos_in_new_combined_block += 1
    end
    b_new = setindex(b_perm_comb,new_b_in_combined_dim,comb_ind_loc)

    Rb_total = blockview(R,b_new)
    dimsRb_tot = dims(Rb_total)
    subind = ntuple(i->i==comb_ind_loc ? range(1+offset,stop=offset+blockdim(ind_comb,b_in_combined_dim)) : range(1,stop=dimsRb_tot[i]),N-NC+1)
    Rb = @view array(Rb_total)[subind...]
    #Rb = reshape(Rb,permute(dims(Tb),perm))
    #permutedims!(Rb,Tb,perm)
    Tb_perm = permutedims(Tb,perm)
    copyto!(Rb,Tb_perm)
  end

  return R
end

# TODO: optimize by avoiding findfirst
function _number_uncombined(blockval::Int,
                            blockcomb::Vector{Int})
  if blockval == blockcomb[end]
    return length(blockcomb)-findfirst(==(blockval),blockcomb)+1
  end
  return findfirst(==(blockval+1),blockcomb)-findfirst(==(blockval),blockcomb)
end

# TODO: optimize by avoiding findfirst
function _number_uncombined_shift(blockval::Int,
                                  blockcomb::Vector{Int})
  if blockval == 1
    return 0
  end
  ncomb_shift = 0
  for i = 1:blockval-1
    ncomb_shift += findfirst(==(i+1),blockcomb) - findfirst(==(i),blockcomb) - 1
  end
  return ncomb_shift
end

# Uncombine the blocks along the dimension dim
# according to the pattern in blockcomb (for example, blockcomb
# is [1,2,2,3] and dim = 2, so the blocks (1,2),(2,3) get
# split into (1,2),(1,3),(2,4))
function uncombine_blocks(blocks::Blocks{N},
                          dim::Int,
                          blockcomb::Vector{Int}) where {N}
  blocks_uncomb = Blocks{N}()
  ncomb_tot = 0
  for i in 1:length(blocks)
    block = blocks[i]
    blockval = block[dim]
    ncomb = _number_uncombined(blockval,blockcomb)
    ncomb_shift = _number_uncombined_shift(blockval,blockcomb)
    push!(blocks_uncomb,setindex(block,blockval+ncomb_shift,dim))
    for j in 1:ncomb-1
      push!(blocks_uncomb,setindex(block,blockval+ncomb_shift+j,dim))
    end
  end
  return blocks_uncomb
end

function uncombine_block(block::Block{N},
                         dim::Int,
                         blockcomb::Vector{Int}) where {N}
  blocks_uncomb = Blocks{N}()
  ncomb_tot = 0
  blockval = block[dim]
  ncomb = _number_uncombined(blockval,blockcomb)
  ncomb_shift = _number_uncombined_shift(blockval,blockcomb)
  push!(blocks_uncomb,setindex(block,blockval+ncomb_shift,dim))
  for j in 1:ncomb-1
    push!(blocks_uncomb,setindex(block,blockval+ncomb_shift+j,dim))
  end
  return blocks_uncomb
end

function uncombine_output(T::BlockSparseTensor{ElT,N},
                          is,
                          combdim::Int,
                          blockperm::Vector{Int},
                          blockcomb::Vector{Int}) where {ElT<:Number,N}
  ind_uncomb_perm = ⊗(setdiff(is,inds(T))...)
  inds_uncomb_perm = insertat(inds(T),ind_uncomb_perm,combdim)
  # Uncombine the blocks of T
  blocks_uncomb = uncombine_blocks(nzblocks(T),combdim,blockcomb)
  blocks_uncomb_perm = perm_blocks(blocks_uncomb,combdim,invperm(blockperm))
  boffs_uncomb_perm,nnz_uncomb_perm = blockoffsets(blocks_uncomb_perm,inds_uncomb_perm)
  T_uncomb_perm = tensor(BlockSparse(ElT,boffs_uncomb_perm,nnz_uncomb_perm),inds_uncomb_perm)
  R = reshape(T_uncomb_perm,is)
  return R
end

function Base.reshape(blockT::Block{NT},
                      indsT,
                      indsR) where {NT}
  nblocksT = nblocks(indsT)
  nblocksR = nblocks(indsR)
  blockR = Tuple(CartesianIndices(nblocksR)[LinearIndices(nblocksT)[CartesianIndex(blockT)]])
  return blockR
end

function uncombine(T::BlockSparseTensor{<:Number,NT},
                   is,
                   combdim::Int,
                   blockperm::Vector{Int},
                   blockcomb::Vector{Int}) where {NT}
  NR = length(is)
  R = uncombine_output(T,is,combdim,blockperm,blockcomb)
  invblockperm = invperm(blockperm)

  # This is needed for reshaping the block
  # It is already calculated in uncombine_output, use it from there
  ind_uncomb_perm = ⊗(setdiff(is,inds(T))...)
  ind_uncomb = permuteblocks(ind_uncomb_perm,blockperm)
  # Same as inds(T) but with the blocks uncombined
  inds_uncomb = insertat(inds(T),ind_uncomb,combdim)
  inds_uncomb_perm = insertat(inds(T),ind_uncomb_perm,combdim)

  for bof in blockoffsets(T)
    b = nzblock(bof)
    Tb_tot = blockview(T,bof)
    dimsTb_tot = dims(Tb_tot)

    bs_uncomb = uncombine_block(b,combdim,blockcomb)

    offset = 0
    for i in 1:length(bs_uncomb)
      b_uncomb = bs_uncomb[i]
      b_uncomb_perm = perm_block(b_uncomb,combdim,invblockperm)
      b_uncomb_perm_reshape = reshape(b_uncomb_perm,inds_uncomb_perm,is)

      Rb = blockview(R,b_uncomb_perm_reshape)

      b_uncomb_in_combined_dim = b_uncomb_perm[combdim]

      start = offset+1
      stop = offset+blockdim(ind_uncomb_perm,b_uncomb_in_combined_dim)
      subind = ntuple(i->i==combdim ? range(start,stop=stop) : range(1,stop=dimsTb_tot[i]),NT)

      offset = stop

      Tb = @view array(Tb_tot)[subind...]

      copyto!(Rb,Tb)
    end
  end
  return R
end

# TODO: handle case with different element types in R and T
#function permutedims!!(R::BlockSparseTensor{<:Number,N},
#                       T::BlockSparseTensor{<:Number,N},
#                       perm::NTuple{N,Int}) where {N}
#  blockoffsetsTp,indsTp = permute(blockoffsets(T),inds(T),perm)
#  if blockoffsetsTp == blockoffsets(R)
#    R = permutedims!(R,T,perm)
#    return R
#  end
#  R = similar(T,blockoffsetsTp,indsTp)
#  permutedims!(R,T,perm)
#  return R
#end

function Base.copyto!(R::BlockSparseTensor,
                      T::BlockSparseTensor)
  for bof in blockoffsets(T)
    copyto!(blockview(R, nzblock(bof)), blockview(T, bof))
  end
  return R
end

# TODO: handle case where:
# f(zero(ElR),zero(ElT)) != promote_type(ElR,ElT)
function permutedims!!(R::BlockSparseTensor{ElR,N},
                       T::BlockSparseTensor{ElT,N},
                       perm::NTuple{N,Int},
                       f::Function=(r,t)->t) where {ElR,ElT,N}
  # TODO: write a custom function for merging two sorted
  # lists with no repeated elements
  nzblocksRR = unique!(sort(vcat(nzblocks(R),permutedims(nzblocks(T),perm));lt=isblockless))
  RR = BlockSparseTensor(promote_type(ElR,ElT),nzblocksRR,inds(R))
  copyto!(RR,R)
  permutedims!(RR,T,perm,f)
  return RR
end

function Base.permutedims!(R::BlockSparseTensor{<:Number,N},
                           T::BlockSparseTensor{<:Number,N},
                           perm::NTuple{N,Int},
                           f::Function=(r,t)->t) where {N}
  for (blockT,_) in blockoffsets(T)
    # Loop over non-zero blocks of T/R
    Tblock = blockview(T,blockT)
    Rblock = blockview(R,permute(blockT,perm))
    permutedims!(Rblock,Tblock,perm,f)
  end
  return R
end

#function Base.permutedims!(R::BlockSparseTensor{<:Number,N},
#                           T::BlockSparseTensor{<:Number,N},
#                           perm::NTuple{N,Int},
#                           f::Function) where {N}
#  for (blockT,_) in blockoffsets(T)
#    # Loop over non-zero blocks of T/R
#    Tblock = blockview(T,blockT)
#    Rblock = blockview(R,permute(blockT,perm))
#    permutedims!(Rblock,Tblock,perm,f)
#  end
#  return R
#end

#
# Contraction
#

# TODO: complete this function: determine the output blocks from the input blocks
# Also, save the contraction list (which block-offsets contract with which),
# may not be generic with other contraction functions!
function contraction_output(T1::TensorT1,
                            T2::TensorT2,
                            indsR::IndsR) where {TensorT1<:BlockSparseTensor,
                                                 TensorT2<:BlockSparseTensor,
                                                 IndsR}
  TensorR = contraction_output_type(TensorT1,TensorT2,IndsR)
  return similar(TensorR,blockoffsetsR,indsR)
end

"""
find_matching_positions(t1,t2) -> t1_to_t2

In a tuple of length(t1), store the positions in t2
where the element of t1 is found. Otherwise, store 0
to indicate that the element of t1 is not in t2.

For example, for all t1[pos1] == t2[pos2], t1_to_t2[pos1] == pos2,
otherwise t1_to_t2[pos1] == 0.
"""
function find_matching_positions(t1,t2)
  t1_to_t2 = @MVector zeros(Int,length(t1))
  for pos1 in 1:length(t1)
    for pos2 in 1:length(t2)
      if t1[pos1] == t2[pos2]
        t1_to_t2[pos1] = pos2
      end
    end
  end
  return Tuple(t1_to_t2)
end

function contract_labels(labels1,labels2,labelsR)
  labels1_to_labels2 = find_matching_positions(labels1,labels2)
  labels1_to_labelsR = find_matching_positions(labels1,labelsR)
  labels2_to_labelsR = find_matching_positions(labels2,labelsR)
  return labels1_to_labels2,labels1_to_labelsR,labels2_to_labelsR
end

function are_blocks_contracted(block1::Block{N1},
                               block2::Block{N2},
                               labels1_to_labels2::NTuple{N1,Int}) where {N1,N2}
  for i1 in 1:N1
    i2 = labels1_to_labels2[i1]
    if i2 > 0
      # This dimension is contracted
      if block1[i1] != block2[i2]
        return false
      end
    end
  end
  return true
end

function contract_blocks(block1::Block{N1},
                         labels1_to_labelsR,
                         block2::Block{N2},
                         labels2_to_labelsR,
                         ::Val{NR}) where {N1,N2,NR}
  blockR = @MVector zeros(Int,NR)
  for i1 in 1:N1
    iR = labels1_to_labelsR[i1]
    if iR > 0
      blockR[iR] = block1[i1]
    end
  end
  for i2 in 1:N2
    iR = labels2_to_labelsR[i2]
    if iR > 0
      blockR[iR] = block2[i2]
    end
  end
  return Tuple(blockR)
end

function contract_blockoffsets(boffs1::BlockOffsets{N1},inds1,labels1,
                               boffs2::BlockOffsets{N2},inds2,labels2,
                               indsR,labelsR) where {N1,N2}
  NR = length(labelsR)
  ValNR = ValLength(labelsR)
  labels1_to_labels2,labels1_to_labelsR,labels2_to_labelsR = contract_labels(labels1,labels2,labelsR)
  blocksR = Block{NR}[]
  contraction_plan = Tuple{Int,Int,Int}[]
  for (pos1,(block1,offset1)) in enumerate(boffs1)
    for (pos2,(block2,offset2)) in enumerate(boffs2)
      if are_blocks_contracted(block1,block2,labels1_to_labels2)
        blockR = contract_blocks(block1,labels1_to_labelsR,
                                 block2,labels2_to_labelsR,
                                 ValNR)
        push!(contraction_plan,(pos1,pos2,0))
        push!(blocksR,blockR)
      end
    end
  end

  sorted_blocksR = sort(blocksR; lt=isblockless)
  unique!(sorted_blocksR)
  blockoffsetsR = BlockOffsets{NR}(undef,length(sorted_blocksR))
  nnzR = 0
  for (i,blockR) in enumerate(sorted_blocksR)
    blockoffsetsR[i] = BlockOffset(blockR, nnzR)
    nnzR += blockdim(indsR, blockR)
  end

  # Now get the locations of the output blocks
  # in the sorted block-offsets list
  for (i,blockR) in enumerate(blocksR)
    posR = findblock(blockoffsetsR, blockR)
    pos1,pos2,_ = contraction_plan[i]
    contraction_plan[i] = (pos1, pos2, posR)
  end

  return blockoffsetsR,contraction_plan
end

function contraction_output(T1::TensorT1,
                            labelsT1,
                            T2::TensorT2,
                            labelsT2,
                            labelsR) where {TensorT1<:BlockSparseTensor,
                                            TensorT2<:BlockSparseTensor}
  indsR = contract_inds(inds(T1),labelsT1,inds(T2),labelsT2,labelsR)
  TensorR = contraction_output_type(TensorT1,TensorT2,typeof(indsR))
  blockoffsetsR,contraction_plan = contract_blockoffsets(blockoffsets(T1),inds(T1),labelsT1,
                                                         blockoffsets(T2),inds(T2),labelsT2,
                                                         indsR,labelsR)
  R = similar(TensorR,blockoffsetsR,indsR)
  return R,contraction_plan
end

function contract(T1::BlockSparseTensor{<:Any,N1},
                  labelsT1,
                  T2::BlockSparseTensor{<:Any,N2},
                  labelsT2,
                  labelsR = contract_labels(labelsT1,labelsT2)) where {N1,N2}
  R,contraction_plan = contraction_output(T1,labelsT1,T2,labelsT2,labelsR)
  R = contract!(R,labelsR,T1,labelsT1,T2,labelsT2,contraction_plan)
  return R
end

function contract!(R::BlockSparseTensor{<:Number,NR},
                   labelsR,
                   T1::BlockSparseTensor{<:Number,N1},
                   labelsT1,
                   T2::BlockSparseTensor{<:Number,N2},
                   labelsT2,
                   contraction_plan) where {N1,N2,NR}
  already_written_to = fill(false,nnzblocks(R))
  # In R .= α .* (T1 * T2) .+ β .* R
  α = 1
  for (pos1,pos2,posR) in contraction_plan
    blockT1 = blockview(T1,pos1)
    blockT2 = blockview(T2,pos2)
    blockR = blockview(R,posR)
    β = 1
    if !already_written_to[posR]
      already_written_to[posR] = true
      # Overwrite the block of R
      β = 0
    end
    contract!(blockR,labelsR,
              blockT1,labelsT1,
              blockT2,labelsT2,
              α,β)
  end
  return R
end

const IntTuple = NTuple{N,Int} where N
const IntOrIntTuple = Union{Int,IntTuple}

function permute_combine(inds::IndsT,
                         pos::Vararg{IntOrIntTuple,N}) where {IndsT,N}
  IndT = eltype(IndsT)
  # Using SizedVector since setindex! doesn't
  # work for MVector when eltype not isbitstype
  newinds = SizedVector{N,IndT}(undef)
  for i ∈ 1:N
    pos_i = pos[i]
    newind_i = inds[pos_i[1]]
    for p in 2:length(pos_i)
      newind_i = newind_i ⊗ inds[pos_i[p]]
    end
    newinds[i] = newind_i
  end
  IndsR = similar_type(IndsT,Val{N})
  indsR = IndsR(Tuple(newinds))
  return indsR
end

"""
Indices are combined according to the grouping of the input,
for example (1,2),3 will combine the first two indices.
"""
function combine(inds::IndsT,
                 com::Vararg{IntOrIntTuple,N}) where {IndsT,N}
  IndT = eltype(IndsT)
  # Using SizedVector since setindex! doesn't
  # work for MVector when eltype not isbitstype
  newinds = SizedVector{N,IndT}(undef)
  i_orig = 1
  for i ∈ 1:N
    newind_i = inds[i_orig]
    i_orig += 1
    for p in 2:length(com[i])
      newind_i = newind_i ⊗ inds[i_orig]
      i_orig += 1
    end
    newinds[i] = newind_i
  end
  IndsR = similar_type(IndsT,Val{N})
  indsR = IndsR(Tuple(newinds))
  return indsR
end

function permute_combine(boffs::BlockOffsets,
                         inds::IndsT,
                         pos::Vararg{IntOrIntTuple,N}) where {IndsT,N}
  perm = flatten(pos...)
  boffsp,indsp = permutedims(boffs,inds,perm)
  indsR = combine(indsp,pos...)
  boffsR = reshape(boffsp,indsp,indsR)
  return boffsR,indsR
end

function Base.reshape(boffsT::BlockOffsets{NT},
                      indsT,
                      indsR) where {NT}
  NR = length(indsR)
  boffsR = BlockOffsets{NR}(undef,nnzblocks(boffsT))
  nblocksT = nblocks(indsT)
  nblocksR = nblocks(indsR)
  for (i,(blockT,offsetT)) in enumerate(boffsT)
    blockR = Tuple(CartesianIndices(nblocksR)[LinearIndices(nblocksT)[CartesianIndex(blockT)]])
    boffsR[i] = blockR => offsetT
  end
  return boffsR
end

function Base.reshape(boffsT::BlockOffsets{NT},
                      blocksR::Vector{Block{NR}}) where {NR,NT}
  boffsR = BlockOffsets{NR}(undef,nnzblocks(boffsT))
  # TODO: check blocksR is ordered and are properly reshaped
  # versions of the blocks of boffsT
  for (i,(blockT,offsetT)) in enumerate(boffsT)
    blockR = blocksR[i]
    boffsR[i] = BlockOffset(blockR,offsetT)
  end
  return boffsR
end

function Base.reshape(T::BlockSparse,
                      boffsR::BlockOffsets)
  return BlockSparse(data(T),boffsR)
end

function Base.reshape(T::BlockSparseTensor,
                      boffsR::BlockOffsets,
                      indsR)
  storeR = reshape(store(T),boffsR)
  return tensor(storeR,indsR)
end

function Base.reshape(T::BlockSparseTensor,
                      indsR)
  # TODO: add some checks that the block dimensions
  # are consistent (e.g. nnzblocks(T) == nnzblocks(R), etc.)
  boffsR = reshape(blockoffsets(T),inds(T),indsR)
  R = reshape(T,boffsR,indsR)
  return R
end

function permute_combine(T::BlockSparseTensor{ElT,NT,IndsT},
                         pos::Vararg{IntOrIntTuple,NR}) where {ElT,NT,IndsT,NR}
  boffsR,indsR = permute_combine(blockoffsets(T),inds(T),pos...)

  perm = flatten(pos...)

  length(perm)≠NT && error("Index positions must add up to order of Tensor ($NT)")
  isperm(perm) || error("Index positions must be a permutation")

  if !is_trivial_permutation(perm)
    Tp = permutedims(T,perm)
  else
    Tp = copy(T)
  end
  NR==NT && return Tp
  R = reshape(Tp,boffsR,indsR)
  return R
end

#
# Print block sparse tensors
#

#function Base.summary(io::IO,
#                      T::BlockSparseTensor{ElT,N}) where {ElT,N}
#  println(io,Base.dims2string(dims(T))," ",typeof(T))
#  for (dim,ind) in enumerate(inds(T))
#    println(io,"Dim $dim: ",ind)
#  end
#  println(io,"Number of nonzero blocks: ",nnzblocks(T))
#end

#function Base.summary(io::IO,
#                      T::BlockSparseTensor{ElT,N}) where {ElT,N}
#  println(io,typeof(T))
#  println(io,Base.dims2string(dims(T))," ",typeof(T))
#  for (dim,ind) in enumerate(inds(T))
#    println(io,"Dim $dim: ",ind)
#  end
#  println("Number of nonzero blocks: ",nnzblocks(T))
#end

function _range2string(rangestart::NTuple{N,Int},
                       rangeend::NTuple{N,Int}) where {N}
  s = ""
  for n in 1:N
    s = string(s,rangestart[n],":",rangeend[n])
    if n < N
      s = string(s,", ")
    end
  end
  return s
end

function Base.show(io::IO,
                   mime::MIME"text/plain",
                   T::BlockSparseTensor)
  summary(io,T)
  println(io)
  for (block,_) in blockoffsets(T)
    blockdimsT = blockdims(T,block)
    # Print the location of the current block
    println(io,"Block: ",block)
    println(io," [",_range2string(blockstart(T,block),blockend(T,block)),"]")
    #println(io,"Start: ",Tuple(blockstart(T,block)))
    #println(io,"End: ",Tuple(blockend(T,block)))
    # Print the dimension of the current block
    #println(io,"   (",Base.dims2string(blockdimsT),")")
    print_tensor(io,blockview(T,block))
    println(io)
    println(io)
  end
end

Base.show(io::IO, T::BlockSparseTensor) = show(io,MIME("text/plain"),T)


text:
### A Pluto.jl notebook ###
# v0.15.1

using Markdown
using InteractiveUtils

# ╔═╡ 3fbe924e-f792-11ea-1cb2-5d1fb71e77b3
using Pkg

# ╔═╡ 3fbec6ec-f792-11ea-0a62-83a272a57a5b
begin
	using StanSample
	using StatisticalRethinking
	using PrettyTables
end

# ╔═╡ b3057494-f791-11ea-3ec6-1587562245eb
md"## Intro-logpdf-1s.jl"

# ╔═╡ 3fbf493c-f792-11ea-06b2-adf07fe86716
begin
	df = CSV.read(sr_datadir("Howell1.csv"), DataFrame; delim=';')
	df2 = filter(row -> row[:age] >= 18, df)
end;

# ╔═╡ 3fca0eaa-f792-11ea-2a3d-75f1f0db58b2
md"##### This is an alternative way of writing the Stan language model for the heights example."

# ╔═╡ 3fd3ad14-f792-11ea-2edc-6f151b03b780
md"##### This is referred to in chapter 9. The model block resembles the way loglik() functions used for Optim are constructed."

# ╔═╡ 3fd42ba4-f792-11ea-22ef-272a936150d1
heightsmodel = "
// Inferring a Rate
data {
  int<lower=1> N;
  real<lower=0> h[N];
}
parameters {
  real mu;
  real<lower=0,upper=250> sigma;
}
model {
  // Priors for mu and sigma
  target += normal_lpdf(mu | 178, 20);

  // Observed heights, add loglikelihood to target
  target += normal_lpdf(h | mu, sigma);
}
";

# ╔═╡ 3fddb2dc-f792-11ea-2602-f16d4a595972
sm = SampleModel("heights", heightsmodel);

# ╔═╡ 3fde4396-f792-11ea-3fb2-336c91f459d3
heightsdata = Dict("N" => length(df2[:, :height]), "h" =>df2.height);

# ╔═╡ 3fe6c480-f792-11ea-1faa-cdec32a7d865
rc = stan_sample(sm, data=heightsdata);

# ╔═╡ 3fefbbe4-f792-11ea-3602-612bbf22dc17
if success(rc)
 	chns = read_samples(sm, :dataframe)
	chns
end

# ╔═╡ 1288ad5f-1a74-4ed4-9065-420c9bb28139
axiskeys(chns)

# ╔═╡ cdd365e8-92f0-442c-9714-d005d0897a7b
HTML(pretty_table(String, chns[1:3, 1:2, :], backend=:html))

# ╔═╡ 3ff6ee3c-f792-11ea-1aa9-e129fcaf7171
md"## End intro-logpdf-1s.jl"

# ╔═╡ Cell order:
# ╟─b3057494-f791-11ea-3ec6-1587562245eb
# ╠═3fbe924e-f792-11ea-1cb2-5d1fb71e77b3
# ╠═3fbec6ec-f792-11ea-0a62-83a272a57a5b
# ╠═3fbf493c-f792-11ea-06b2-adf07fe86716
# ╟─3fca0eaa-f792-11ea-2a3d-75f1f0db58b2
# ╟─3fd3ad14-f792-11ea-2edc-6f151b03b780
# ╠═3fd42ba4-f792-11ea-22ef-272a936150d1
# ╠═3fddb2dc-f792-11ea-2602-f16d4a595972
# ╠═3fde4396-f792-11ea-3fb2-336c91f459d3
# ╠═3fe6c480-f792-11ea-1faa-cdec32a7d865
# ╠═3fefbbe4-f792-11ea-3602-612bbf22dc17
# ╠═1288ad5f-1a74-4ed4-9065-420c9bb28139
# ╠═cdd365e8-92f0-442c-9714-d005d0897a7b
# ╟─3ff6ee3c-f792-11ea-1aa9-e129fcaf7171

ext: jl


text:
module Generic

import Base: getindex, setindex!
import Base: read, write
import Base.==

using Avro.Schemas
using Avro.Io

export GenericRecord,
       GenericEnumSymbol,
       GenericFixed,
       read,
       write,
       SchemaNotFoundError,
       UnknownSchema

struct SchemaNotFoundError <: Exception
    datum_schema
end

Base.showerror(io, err::SchemaNotFoundError) = print(io, "Schema not found in union: $(err.datum_schema)")

struct UnknownSchema <: Exception end

Base.showerror(io, err::UnknownSchema) = print(io, "Unknown datum type")

# Generic writers

write(encoder::Encoder, ::NullSchema, value::Nothing) = encode_null(encoder, value)
write(encoder::Encoder, ::BooleanSchema, value::Bool) = encode_boolean(encoder, value)
write(encoder::Encoder, ::IntSchema, value::Int32) = encode_int(encoder, value)
write(encoder::Encoder, ::LongSchema, value::Int64) = encode_long(encoder, value)
write(encoder::Encoder, ::FloatSchema, value::Float32) = encode_float(encoder, value)
write(encoder::Encoder, ::DoubleSchema, value::Float64) = encode_double(encoder, value)
write(encoder::Encoder, ::BytesSchema, value::UInt8) = encode_byte(encoder, value)
write(encoder::Encoder, ::BytesSchema, value::Vector{UInt8}) = encode_bytes(encoder, value)
write(encoder::Encoder, ::StringSchema, value::String) = encode_string(encoder, value)

"""
Writes an array of Avro objects if there is a
write(Encoder, typeof(ArraySchema.items), T) method.
"""
function write(encoder::Encoder, schema::ArraySchema, value::Vector{T}) where T
    bytes_written = encode_long(encoder, Int64(length(value)))
    for item in value
        bytes_written += write(encoder, schema.items, item)
    end
    bytes_written += encode_byte(encoder, zero(UInt8))
    bytes_written
end

"""
Writes a map of Avro objects if there is a
write(Encoder, typeof(MapSchema.values), T) method.
"""
function write(encoder::Encoder, schema::MapSchema, value::Dict{String, T}) where T
    bytes_written = encode_long(encoder, Int64(length(value)))
    for (k, v) in value
        bytes_written += encode_string(encoder, k)
        bytes_written += write(encoder, schema.values, v)
    end
    bytes_written += encode_byte(encoder, zero(UInt8))
    bytes_written
end

"""
Writes a value as a Union if the schema of the value is in the union.
"""
function write(encoder::Encoder, schema::Schemas.UnionSchema, value)
    datum_schema = resolve_schema(value)
    index = findfirst(isequal(datum_schema), schema.schemas) % Int64
    if index == 0
        throw(SchemaNotFoundError(datum_schema))
    end
    union_tag = index - one(index)
    encode_long(encoder, union_tag) + write(encoder, schema.schemas[index], value)
end

# Generic readers

read(decoder::Decoder, ::NullSchema) = decode_null(decoder)
read(decoder::Decoder, ::BooleanSchema) = decode_boolean(decoder)
read(decoder::Decoder, ::IntSchema) = decode_int(decoder)
read(decoder::Decoder, ::LongSchema) = decode_long(decoder)
read(decoder::Decoder, ::FloatSchema) = decode_float(decoder)
read(decoder::Decoder, ::DoubleSchema) = decode_double(decoder)
read(decoder::Decoder, ::BytesSchema) = decode_bytes(decoder)
read(decoder::Decoder, ::StringSchema) = decode_string(decoder)

"""
Reads array of Avro objects into generic instances.
"""
function read(decoder::Decoder, schema::ArraySchema)
    result = Any[]
    block_count = decode_long(decoder)
    while block_count != 0
        if block_count < 0
            block_count = -block_count
            #  Consume the block_size which is not used
            decode_long(decoder)
        end
        for i in 1:block_count
            push!(result, read(decoder, schema.items))
        end
        block_count = decode_long(decoder)
    end
    result
end

"""
Read maps of Avro objects into generic instances.
"""
function read(decoder::Decoder, schema::MapSchema)
    result = Dict{String, Any}()
    block_count = decode_long(decoder)
    while block_count != 0
        if block_count < 0
            block_count = -block_count
            #  Consume the block_size which is not used
            decode_long(decoder)
        end
        for i in 1:block_count
            key = decode_string(decoder)
            value = read(decoder, schema.values)
            result[key] = value
        end
        block_count = decode_long(decoder)
    end
    result
end

"""
Read unions of Avro objects into generic instances.
"""
function read(decoder::Decoder, schema::UnionSchema)
    union_tag = decode_long(decoder)
    schema = schema.schemas[union_tag + one(union_tag)]
    read(decoder, schema)
end

# Generic implementations for derived types

"""
Contains data for Avro records.
"""
struct GenericRecord
    schema::RecordSchema
    values::Vector{Any}
end

function ==(a::GenericRecord, b::GenericRecord)
    a.schema == b.schema && a.values == b.values
end

function put(record::GenericRecord, key::Symbol, v)
    field = findfirst(field -> field.name == key, record.schema.fields)
    record.values[field.position + one(field.position)] = v
end

put(record::GenericRecord, v, i::Int) = record[i + one(i)] = v
setindex!(record::GenericRecord, v, i::Int) = record.values[i] = v

function get(record::GenericRecord, key::Symbol)
    field = findfirst(field -> field.name == key, record.schema.fields)
    record.values[field.position + one(field.position)]
end

get(record::GenericRecord, i::Int) = record[i + one(i)]
getindex(record::GenericRecord, i::Int) = record.values[i]

function write(encoder::Encoder, schema::RecordSchema, datum::GenericRecord)
    bytes_written = 0
    for field in schema.fields
        bytes_written += write(encoder, field.schema, get(datum, field.position))
    end
    bytes_written
end

function read(decoder::Decoder, schema::RecordSchema)
    n = length(schema.fields)
    values = Array{Any}(undef, n)
    for i in 1:n
        values[i] = read(decoder, schema.fields[i].schema)
    end
    GenericRecord(schema, values)
end

"""
An enum symbol.
"""
struct GenericEnumSymbol
    schema::EnumSchema
    symbol::String
end

function ==(a::GenericEnumSymbol, b::GenericEnumSymbol)
    a.schema == b.schema && a.symbol == b.symbol
end

function write(encoder::Encoder, schema::EnumSchema, datum::GenericEnumSymbol)
    symbol_index = findfirst(isequal(datum.symbol), schema.symbols) - one(Int32)
    encode_int(encoder, symbol_index % Int32)
end

function read(decoder::Decoder, schema::EnumSchema)
    symbol_index = decode_int(decoder)
    GenericEnumSymbol(schema, schema.symbols[symbol_index + one(symbol_index)])
end

"""
Contains data for Avro fixed objects.
"""
struct GenericFixed
    schema::FixedSchema
    bytes::Vector{UInt8}
end

function ==(a::GenericFixed, b::GenericFixed)
    a.schema == b.schema && a.bytes == b.bytes
end

function write(encoder::Encoder, schema::FixedSchema, datum::GenericFixed)
    @assert schema.size == length(datum.bytes)
    encode_fixed(encoder, datum.bytes)
end

function read(decoder::Decoder, schema::FixedSchema)
    GenericFixed(schema, decode_fixed(decoder, schema.size))
end

# Utilities for generic implementations

resolve_schema(::Nothing) = Schemas.NULL
resolve_schema(::Bool) = Schemas.BOOLEAN
resolve_schema(::Int32) = Schemas.INT
resolve_schema(::Int64) = Schemas.LONG
resolve_schema(::Float32) = Schemas.FLOAT
resolve_schema(::Float64) = Schemas.DOUBLE
resolve_schema(::Vector{UInt8}) = Schemas.BYTES
resolve_schema(::String) = Schemas.STRING
resolve_schema(::Vector) = ArraySchema
resolve_schema(::Dict) = MapSchema
resolve_schema(datum::GenericRecord) = datum.schema
resolve_schema(datum::GenericEnumSymbol) = datum.schema
resolve_schema(datum::GenericFixed) = datum.schema
resolve_schema(::Any) = throw(UnknownSchema())

end


text:
# Autogenerated wrapper script for Hwloc_jll for i686-linux-gnu
export libhwloc

JLLWrappers.@generate_wrapper_header("Hwloc")
JLLWrappers.@declare_library_product(libhwloc, "libhwloc.so.15")
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_library_product(
        libhwloc,
        "lib/libhwloc.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
# This might go to Gridap in the future. We keep it here for the moment.
function change_axes(a::Algebra.ArrayCounter,axes)
  @notimplemented
end

# This might go to Gridap in the future. We keep it here for the moment.
function change_axes(a::Algebra.CounterCOO{T,A}, axes::A) where {T,A}
  b=Algebra.CounterCOO{T}(axes)
  b.nnz = a.nnz
  b
end

# This might go to Gridap in the future. We keep it here for the moment.
function change_axes(a::Algebra.AllocationCOO{T,A}, axes::A) where {T,A}
  counter=change_axes(a.counter,axes)
  Algebra.AllocationCOO(counter,a.I,a.J,a.V)
end

function local_views(a)
  @abstractmethod
end

function local_views(a::AbstractVector,rows)
  @notimplemented
end

function local_views(a::AbstractMatrix,rows,cols)
  @notimplemented
end

function consistent_local_views(a,ids,isconsistent)
  @abstractmethod
end

function local_views(a::AbstractPData)
  a
end

function local_views(a::PRange)
  a.partition
end

function local_views(a::PVector)
  a.values
end

function local_views(a::PVector,rows::PRange)
  PArrays.local_view(a,rows)
end

function local_views(a::PSparseMatrix)
  a.values
end

function local_views(a::PSparseMatrix,rows::PRange,cols::PRange)
  PArrays.local_view(a,rows,cols)
end

function change_ghost(a::PVector,ids_fespace::PRange)
  if a.rows === ids_fespace
    a_fespace = a
  else
    a_fespace = similar(a,eltype(a),ids_fespace)
    a_fespace .= a
  end
  a_fespace
end

function consistent_local_views(a::PVector,ids_fespace::PRange,isconsistent)
  a_fespace = change_ghost(a,ids_fespace)
  if ! isconsistent
    exchange!(a_fespace)
  end
  a_fespace.values
end

function Algebra.allocate_vector(::Type{<:PVector{T,A}},ids::PRange) where {T,A}
  values = map_parts(ids.partition) do ids
    Tv = eltype(A)
    Tv(undef,num_lids(ids))
  end
  PVector(values,ids)
end

struct FullyAssembledRows end
struct SubAssembledRows end

# For the moment we use COO format even though
# it is quite memory consuming.
# We need to implement communication in other formats in
# PartitionedArrays.jl

struct PSparseMatrixBuilderCOO{T,B}
  local_matrix_type::Type{T}
  par_strategy::B
end

function Algebra.nz_counter(
  builder::PSparseMatrixBuilderCOO{A}, axs::Tuple{<:PRange,<:PRange}) where A
  rows, cols = axs
  counters = map_parts(rows.partition,cols.partition) do r,c
    axs = (Base.OneTo(num_lids(r)),Base.OneTo(num_lids(c)))
    Algebra.CounterCOO{A}(axs)
  end
  DistributedCounterCOO(builder.par_strategy,counters,rows,cols)
end

"""
"""
struct DistributedCounterCOO{A,B,C,D} <: GridapType
  par_strategy::A
  counters::B
  rows::C
  cols::D
  function DistributedCounterCOO(
    par_strategy,
    counters::AbstractPData{<:Algebra.CounterCOO},
    rows::PRange,
    cols::PRange)
    A = typeof(par_strategy)
    B = typeof(counters)
    C = typeof(rows)
    D = typeof(cols)
    new{A,B,C,D}(par_strategy,counters,rows,cols)
  end
end

function local_views(a::DistributedCounterCOO)
  a.counters
end

function local_views(a::DistributedCounterCOO,rows,cols)
  @check rows === a.rows
  @check cols === a.cols
  a.counters
end

function Algebra.nz_allocation(a::DistributedCounterCOO)
  allocs = map_parts(nz_allocation,a.counters)
  DistributedAllocationCOO(a.par_strategy,allocs,a.rows,a.cols)
end

struct DistributedAllocationCOO{A,B,C,D} <:GridapType
  par_strategy::A
  allocs::B
  rows::C
  cols::D
  function DistributedAllocationCOO(
    par_strategy,
    allocs::AbstractPData{<:Algebra.AllocationCOO},
    rows::PRange,
    cols::PRange)
    A = typeof(par_strategy)
    B = typeof(allocs)
    C = typeof(rows)
    D = typeof(cols)
    new{A,B,C,D}(par_strategy,allocs,rows,cols)
  end
end

function change_axes(a::DistributedAllocationCOO{A,B,<:PRange,<:PRange},
                     axes::Tuple{<:PRange,<:PRange}) where {A,B}
  local_axes=map_parts(axes[1].partition,axes[2].partition) do rows,cols
    (Base.OneTo(num_lids(rows)), Base.OneTo(num_lids(cols)))
  end
  allocs=map_parts(change_axes,a.allocs,local_axes)
  DistributedAllocationCOO(a.par_strategy,allocs,axes[1],axes[2])
end

function local_views(a::DistributedAllocationCOO)
  a.allocs
end

function local_views(a::DistributedAllocationCOO,rows,cols)
  @check rows === a.rows
  @check cols === a.cols
  a.allocs
end

function first_gdof_from_ids(ids)
  num_oids(ids)>0 ? Int(ids.lid_to_gid[ids.oid_to_lid[1]]) : 1
end

function find_gid_and_part(hid_to_hdof,dofs)
  hid_to_ldof = view(dofs.hid_to_lid,hid_to_hdof)
  hid_to_gid = dofs.lid_to_gid[hid_to_ldof]
  hid_to_part = dofs.lid_to_part[hid_to_ldof]
  hid_to_gid, hid_to_part
end

function Algebra.create_from_nz(a::PSparseMatrix)
  # For FullyAssembledRows the underlying Exchanger should
  # not have ghost layer making assemble! do nothing (TODO check)
  assemble!(a)
  a
end

function Algebra.create_from_nz(a::DistributedAllocationCOO{<:FullyAssembledRows})
  f(x) = nothing
  A, = _fa_create_from_nz_with_callback(f,a)
  A
end

function _fa_create_from_nz_with_callback(callback,a)

  # Recover some data
  I,J,C = map_parts(a.allocs) do alloc
    alloc.I, alloc.J, alloc.V
  end
  parts = get_part_ids(a.allocs)
  rdofs = a.rows # dof ids of the test space
  cdofs = a.cols # dof ids of the trial space
  ngrdofs = length(rdofs)
  ngcdofs = length(cdofs)
  nordofs = map_parts(num_oids,rdofs.partition)
  nocdofs = map_parts(num_oids,cdofs.partition)
  first_grdof = map_parts(first_gdof_from_ids,rdofs.partition)
  first_gcdof = map_parts(first_gdof_from_ids,cdofs.partition)
  cneigs_snd = cdofs.exchanger.parts_snd
  cneigs_rcv = cdofs.exchanger.parts_rcv

  # This one has not ghost rows
  rows = PRange(
    parts,
    ngrdofs,
    nordofs,
    first_grdof)

  callback_output = callback(rows)

  # convert I and J to global dof ids
  to_gids!(I,rdofs)
  to_gids!(J,cdofs)

  # Find the ghost cols
  hcol_to_hcdof = touched_hids(cdofs,J)
  hcol_to_gid, hcol_to_part = map_parts(
    find_gid_and_part,hcol_to_hcdof,cdofs.partition)

  # Create the range for cols
  cols = PRange(
    parts,
    ngcdofs,
    nocdofs,
    first_gcdof,
    hcol_to_gid,
    hcol_to_part,
    cneigs_snd,
    cneigs_rcv)

  # Convert again I,J to local numeration
  to_lids!(I,rows)
  to_lids!(J,cols)

  # Adjust local matrix size to linear system's index sets
  b=change_axes(a,(rows,cols))

  # Compress local portions
  values = map_parts(create_from_nz,b.allocs)

  # Build the matrix exchanger. This can be empty since no ghost rows
  exchanger = empty_exchanger(parts)

  # Finally build the matrix
  A = PSparseMatrix(values,rows,cols,exchanger)

  A, callback_output
end

function Algebra.create_from_nz(a::DistributedAllocationCOO{<:SubAssembledRows})
  f(x) = nothing
  A, = _sa_create_from_nz_with_callback(f,f,a)
  A
end

function _sa_create_from_nz_with_callback(callback,async_callback,a)

  # Recover some data
  I,J,C = map_parts(a.allocs) do alloc
    alloc.I, alloc.J, alloc.V
  end
  parts = get_part_ids(a.allocs)
  rdofs = a.rows # dof ids of the test space
  cdofs = a.cols # dof ids of the trial space
  ngrdofs = length(rdofs)
  ngcdofs = length(cdofs)
  nordofs = map_parts(num_oids,rdofs.partition)
  nocdofs = map_parts(num_oids,cdofs.partition)
  first_grdof = map_parts(first_gdof_from_ids,rdofs.partition)
  first_gcdof = map_parts(first_gdof_from_ids,cdofs.partition)
  rneigs_snd = rdofs.exchanger.parts_snd
  rneigs_rcv = rdofs.exchanger.parts_rcv
  cneigs_snd = cdofs.exchanger.parts_snd
  cneigs_rcv = cdofs.exchanger.parts_rcv

  # convert I and J to global dof ids
  to_gids!(I,rdofs)
  to_gids!(J,cdofs)

  # Find the ghost rows
  hrow_to_hrdof = touched_hids(rdofs,I)
  hrow_to_gid, hrow_to_part = map_parts(
    find_gid_and_part,hrow_to_hrdof,rdofs.partition)

  # Create the range for rows
  rows = PRange(
    parts,
    ngrdofs,
    nordofs,
    first_grdof,
    hrow_to_gid,
    hrow_to_part,
    rneigs_snd,
    rneigs_rcv)

  # Move values to the owner part
  # since we have integrated only over owned cells
  t = async_assemble!(I,J,C,rows)

  # Here we can overlap computations
  # This is a good place to overlap since
  # sending the matrix rows is a lot of data
  callback_output = callback(rows)

  # Wait the transfer to finish
  map_parts(schedule,t)
  map_parts(wait,t)

  # Find the ghost cols
  hcol_to_hcdof = touched_hids(cdofs,J)
  hcol_to_gid, hcol_to_part = map_parts(
    find_gid_and_part,hcol_to_hcdof,cdofs.partition)

  # Create the range for cols
  cols = PRange(
    parts,
    ngcdofs,
    nocdofs,
    first_gcdof,
    hcol_to_gid,
    hcol_to_part,
    cneigs_snd,
    cneigs_rcv)

  # Overlap rhs communications with CSC compression
  t2 = async_callback(callback_output)

  # Convert again I,J to local numeration
  to_lids!(I,rows)
  to_lids!(J,cols)

  # Adjust local matrix size to linear system's index sets
  b=change_axes(a,(rows,cols))

  # Compress the local matrices
  values = map_parts(create_from_nz,b.allocs)

  # Wait the transfer to finish
  if t2 !== nothing
    map_parts(schedule,t2)
    map_parts(wait,t2)
  end

  # Finally build the matrix
  # A matrix exchanger will be created under the hood.
  # Building a exchanger can be costly and not always needed.
  # TODO add a more lazy initzialization of this inside PSparseMatrix
  A = PSparseMatrix(values,rows,cols)

  A, callback_output
end

struct PVectorBuilder{T,B}
  local_vector_type::Type{T}
  par_strategy::B
end

function Algebra.nz_counter(builder::PVectorBuilder,axs::Tuple{<:PRange})
  T = builder.local_vector_type
  rows, = axs
  counters = map_parts(rows.partition) do rows
    axs = (Base.OneTo(num_lids(rows)),)
    nz_counter(ArrayBuilder(T),axs)
  end
  PVectorCounter(builder.par_strategy,counters,rows)
end

struct PVectorCounter{A,B,C}
  par_strategy::A
  counters::B
  rows::C
end

Algebra.LoopStyle(::Type{<:PVectorCounter}) = DoNotLoop()

function local_views(a::PVectorCounter)
  a.counters
end

function local_views(a::PVectorCounter,rows)
  @check rows === a.rows
  a.counters
end

function Arrays.nz_allocation(a::PVectorCounter{<:FullyAssembledRows})
  dofs = a.rows
  values = map_parts(nz_allocation,a.counters)
  PVectorAllocationTrackOnlyValues(a.par_strategy,values,dofs)
end

struct PVectorAllocationTrackOnlyValues{A,B,C}
  par_strategy::A
  values::B
  rows::C
end

function local_views(a::PVectorAllocationTrackOnlyValues)
  a.values
end

function local_views(a::PVectorAllocationTrackOnlyValues,rows)
  @check rows === a.rows
  a.values
end

function Algebra.create_from_nz(a::PVectorAllocationTrackOnlyValues{<:FullyAssembledRows})
  # Create PRange for the rows of the linear system
  parts = get_part_ids(a.values)
  ngdofs = length(a.rows)
  nodofs = map_parts(num_oids,a.rows.partition)
  first_grdof = map_parts(first_gdof_from_ids,a.rows.partition)

  # This one has no ghost rows
  rows = PRange(parts,ngdofs,nodofs,first_grdof)

  _rhs_callback(a,rows)
end

function Algebra.create_from_nz(a::PVectorAllocationTrackOnlyValues{<:SubAssembledRows})
  # This point MUST NEVER be reached. If reached there is an inconsistency
  # in the parallel code in charge of vector assembly
  @assert false
end

function _rhs_callback(c_fespace,rows)
  # The ghost values in b_fespace are aligned with the FESpace
  # but not with the ghost values in the rows of A
  b_fespace = PVector(c_fespace.values,c_fespace.rows)

  # This one is aligned with the rows of A
  b = similar(b_fespace,eltype(b_fespace),(rows,))

  # First transfer owned values
  b .= b_fespace

  # Now transfer ghost
  function transfer_ghost(b,b_fespace,ids,ids_fespace)
    for hid in 1:num_hids(ids)
      lid = ids.hid_to_lid[hid]
      gid = ids.lid_to_gid[lid]
      lid_fespace = ids_fespace.gid_to_lid[gid]
      b[lid] = b_fespace[lid_fespace]
    end
  end
  map_parts(
    transfer_ghost,
    b.values,
    b_fespace.values,
    b.rows.partition,
    b_fespace.rows.partition)

  return b
end

function Algebra.create_from_nz(a::PVector)
  # For FullyAssembledRows the underlying Exchanger should
  # not have ghost layer making assemble! do nothing (TODO check)
  assemble!(a)
  a
end

function Algebra.create_from_nz(
  a::DistributedAllocationCOO{<:FullyAssembledRows},
  c_fespace::PVectorAllocationTrackOnlyValues{<:FullyAssembledRows})

  function callback(rows)
    _rhs_callback(c_fespace,rows)
  end

  A,b = _fa_create_from_nz_with_callback(callback,a)
  A,b
end

struct PVectorAllocationTrackTouchedAndValues{A,B,C}
  allocations::A
  values::B
  rows::C
end

function Algebra.create_from_nz(
  a::DistributedAllocationCOO{<:SubAssembledRows},
  c_fespace::PVectorAllocationTrackOnlyValues{<:SubAssembledRows})

  function callback(rows)
    _rhs_callback(c_fespace,rows)
  end

  function async_callback(b)
    # now we can assemble contributions
    async_assemble!(b)
  end

  A,b = _sa_create_from_nz_with_callback(callback,async_callback,a)
  A,b
end

struct ArrayAllocationTrackTouchedAndValues{A}
  touched::Vector{Bool}
  values::A
end

Gridap.Algebra.LoopStyle(::Type{<:ArrayAllocationTrackTouchedAndValues}) = Gridap.Algebra.Loop()

function local_views(a::PVectorAllocationTrackTouchedAndValues,rows)
  @check rows === a.rows
  a.allocations
end

@inline function Arrays.add_entry!(c::Function,a::ArrayAllocationTrackTouchedAndValues,v,i,j)
  @notimplemented
end
@inline function Arrays.add_entry!(c::Function,a::ArrayAllocationTrackTouchedAndValues,v,i)
  if i>0
    if !(a.touched[i])
      a.touched[i]=true
    end
    if v!=nothing
      a.values[i]=c(v,a.values[i])
    end
  end
  nothing
end
@inline function Arrays.add_entries!(c::Function,a::ArrayAllocationTrackTouchedAndValues,v,i,j)
  @notimplemented
end
@inline function Arrays.add_entries!(c::Function,a::ArrayAllocationTrackTouchedAndValues,v,i)
  if v != nothing
    for (ve,ie) in zip(v,i)
      Arrays.add_entry!(c,a,ve,ie)
    end
  else
    for ie in i
      Arrays.add_entry!(c,a,nothing,ie)
    end
  end
  nothing
end

function Arrays.nz_allocation(a::DistributedCounterCOO{<:SubAssembledRows},
                              b::PVectorCounter{<:SubAssembledRows})
  A = nz_allocation(a)
  dofs = b.rows
  values = map_parts(nz_allocation,b.counters)
  B=PVectorAllocationTrackOnlyValues(b.par_strategy,values,dofs)
  A,B
end

function Arrays.nz_allocation(a::PVectorCounter{<:SubAssembledRows})
  dofs = a.rows
  values = map_parts(nz_allocation,a.counters)
  touched = map_parts(values) do values
     fill!(Vector{Bool}(undef,length(values)),false)
  end
  allocations=map_parts(values,touched) do values,touched
    ArrayAllocationTrackTouchedAndValues(touched,values)
  end
  PVectorAllocationTrackTouchedAndValues(allocations,values,dofs)
end

function local_views(a::PVectorAllocationTrackTouchedAndValues)
  a.allocations
end

function Algebra.create_from_nz(a::PVectorAllocationTrackTouchedAndValues)
   parts = get_part_ids(a.values)
   rdofs = a.rows # dof ids of the test space
   ngrdofs = length(rdofs)
   nordofs = map_parts(num_oids,rdofs.partition)
   first_grdof = map_parts(first_gdof_from_ids,rdofs.partition)
   rneigs_snd = rdofs.exchanger.parts_snd
   rneigs_rcv = rdofs.exchanger.parts_rcv

   # Find the ghost rows
   hrow_to_hrdof=map_parts(local_views(a.allocations),rdofs.partition) do allocation, indices
    lids_touched=findall(allocation.touched)
    nhlids = count((x)->indices.lid_to_ohid[x]<0,lids_touched)
    hlids = Vector{Int32}(undef,nhlids)
    cur=1
    for lid in lids_touched
      hlid=indices.lid_to_ohid[lid]
      if hlid<0
        hlids[cur]=-hlid
        cur=cur+1
      end
    end
    hlids
   end
   hrow_to_gid, hrow_to_part = map_parts(
       find_gid_and_part,hrow_to_hrdof,rdofs.partition)

   # Create the range for rows
   rows = PRange(
           parts,
           ngrdofs,
           nordofs,
           first_grdof,
           hrow_to_gid,
           hrow_to_part,
           rneigs_snd,
           rneigs_rcv)

   b = _rhs_callback(a,rows)
   t2 = async_assemble!(b)

   # Wait the transfer to finish
   if t2 !== nothing
     map_parts(schedule,t2)
     map_parts(wait,t2)
   end
   b
end


text:
using Base.Iterators
using DataStructures
using ArgParse

function readInput(filename::String)
    regex = r"Player (\d+) starting position: (\d+)"
    player1, value1 = -1, -1
    player2, value2 = -1, -1

    open(filename) do file
        for line in readlines(file)
            m = match(regex, line)
            if m[1] == "1"
                player1, value1 = parse.(Int, (m[1], m[2]))
            else
                player2, value2 = parse.(Int, (m[1], m[2]))
            end
        end
    end
    return (player1, value1), (player2, value2)
end

function playGame(p1, p2; winning_score::Int = 1000)
    roll_dice(offset::Int) = sum(((offset:offset+2) .- 1) .% 100 .+ 1)

    p1, p2 = p1 - 1, p2 - 1
    s1, s2, dice_i = 0, 0, 0
    for i in countfrom(1)
        if i % 2 == 1
            p1 = (p1 + roll_dice(dice_i + 1)) % 10
            s1 += p1 + 1

            s1 >= winning_score && return i * s2 *3
        else
            p2 = (p2 + roll_dice(dice_i + 1)) % 10
            s2 += p2 + 1

            s2 >= winning_score && return i * s1 * 3
        end

        dice_i = (dice_i + 3) % 100
    end
end

function dirac_game(p1, p2; winning_score::Int=21)
    s1, p1 = 0, p1 - 1
    s2, p2 = 0, p2 - 1
    wins = [0, 0]

    three_dice = vec(sum.(collect(product(1:3, 1:3, 1:3))))

    prev_step = Dict{Tuple{Int, Int, Int, Int}, Int}((p1, s1, p2, s2) => 1)
    turn = true
    while length(prev_step) != 0
        next_step = DefaultDict{Tuple{Int, Int, Int, Int}, Int}(0)
        for k in keys(prev_step)
            v = prev_step[k]
            p1, s1, p2, s2 = k

            for dice in three_dice
                if turn
                    s1_ = s1 + (p1 + dice) % 10 + 1
                    if s1_ >= winning_score
                        wins[1] += v
                    else
                        p1_ = (p1 + dice) % 10
                        next_step[(p1_, s1_, p2, s2)] += v
                    end
                else
                    s2_ = s2 + (p2 + dice) % 10 + 1
                    if s2_ >= winning_score
                        wins[2] += v
                    else
                        p2_ = (p2 + dice) % 10
                        next_step[(p1, s1, p2_, s2_)] += v
                    end
                end
            end
        end
        turn = !turn
        prev_step = next_step
    end
    return maximum(wins)
end

function parse_cli()
    s = ArgParseSettings()
    @add_arg_table s begin
        "--example"
            help = "Calculate example"
            action = :store_true
    end
    return parse_args(s)
end

args = parse_cli();
if args["example"]
    (_, p1), (_, p2) = readInput("example.txt")

    score = playGame(p1, p2)
    @assert score == 739785

    w = dirac_game(p1, p2)
    @assert w == 444356092776315
end

(_, p1), (_, p2) = readInput("input.txt")
score = playGame(p1, p2)
w = dirac_game(p1, p2)

println(score)
println(w)