const CuDense{ElT,VecT} = Dense{ElT,VecT} where {VecT<:CuVector}
const CuDenseTensor{ElT,N,StoreT,IndsT} = Tensor{ElT,N,StoreT,IndsT} where {StoreT<:CuDense}

function Dense{T,SA}(x::Dense{T,SB}) where {T<:Number,SA<:CuArray,SB<:Array}
  return Dense{T,SA}(CuArray(x))
end
function Dense{T,SA}(x::Dense{T,SB}) where {T<:Number,SA<:Array,SB<:CuArray}
  return Dense{T,SA}(collect(x.data))
end
Dense{T,S}(size::Integer) where {T,S<:CuArray{<:T}} = Dense{T,S}(CUDA.zeros(T, size))
function Dense{T,S}(x::T, size::Integer) where {T,S<:CuArray{<:T}}
  arr = CuArray{T}(undef, size)
  fill!(arr, x)
  return Dense{T,S}(arr)
end
cpu(x::CuDense{T}) where {T<:Number} = Dense(collect(x.data))
function Base.complex(::Type{Dense{ElT,VT}}) where {ElT,VT<:CuArray}
  return Dense{complex(ElT),CuVector{complex(ElT)}}
end

similartype(::Type{<:CuArray{<:Any,N}}, eltype::Type) where {N} = CuArray{eltype,N}
function NDTensors.similar(::Type{<:CuArray{T}}, dims) where {T}
  return CuArray{T,length(dims)}(undef, dims)
end

CuArray(x::CuDense{ElT}) where {ElT} = CuVector{ElT}(data(x))
function CuArray{ElT,N}(x::CuDenseTensor{ElT,N}) where {ElT,N}
  return CuArray{ElT,N}(reshape(data(store(x)), dims(inds(x))...))
end
CuArray(x::CuDenseTensor{ElT,N}) where {ElT,N} = CuArray{ElT,N}(x)

*(D::Dense{T,AT}, x::S) where {T,AT<:CuArray,S<:Number} = Dense(x .* data(D))

function Base.:(==)(
  ::Type{<:CuDense{ElT1,CVec1}}, ::Type{<:CuDense{ElT2,CVec2}}
) where {ElT1,ElT2,CVec1,CVec2}
  return (ElT1 == ElT2)
end
Base.getindex(D::CuDense{<:Number}) = collect(data(D))[]
Base.getindex(D::CuDenseTensor{<:Number,0}) = store(D)[]
LinearAlgebra.norm(T::CuDenseTensor) = norm(data(store(T)))

# This is for type promotion for Scalar*Dense
function Base.promote_rule(
  ::Type{<:Dense{ElT1,CuVector{ElT1}}}, ::Type{ElT2}
) where {ElT1,ElT2<:Number}
  ElR = promote_type(ElT1, ElT2)
  VecR = CuVector{ElR}
  return Dense{ElR,VecR}
end

function Base.permutedims(T::CuDenseTensor{<:Number,N}, perm::NTuple{N,Int}) where {N}
  Tp = NDTensors.similar(T, ITensors.NDTensors.permute(inds(T), perm))
  #Tp = permute(T,perm; always_copy=true)
  permute!(Tp, T)
  return Tp
end

function Base.permutedims!(
  R::CuDenseTensor{<:Number,N}, T::CuDenseTensor{<:Number,N}, perm::NTuple{N,Int}
) where {N}
  return permutedims!!(R, T, perm)
end

function permutedims!!(
  B::Tensor{ElT,N,StoreT,IndsB},
  A::Tensor{ElT,N,StoreT,IndsA},
  perm::NTuple{N,Int},
  f::Function=(r, t) -> permute!(r, t),
) where {N,ElT,IndsB,IndsA,StoreT<:CuDense{ElT}}
  Ais = inds(A)
  Bis = ITensors.NDTensors.permute(inds(A), perm)
  B = f(B, A)
  return B
end

function Base.similar(::Type{<:CuDenseTensor{ElT}}, inds::Dims) where {ElT}
  storage_arr = CuVector{ElT}(undef, dim(inds))
  return Tensor(Dense(storage_arr), inds)
end

function Base.similar(::Type{<:CuDenseTensor{ElT}}, inds) where {ElT}
  storage_arr = CuVector{ElT}(undef, dim(inds))
  return Tensor(Dense(storage_arr), inds)
end

function outer!(R::CuDenseTensor, T1::CuDenseTensor, T2::CuDenseTensor)
  R_dat = vec(array(T1)) * transpose(vec(array(T2)))
  copyto!(data(store(R)), vec(R_dat))
  inds_outer = unioninds(inds(T1), inds(T2))
  return R
end

function contract!!(
  R::CuDenseTensor{<:Number,NR},
  labelsR::NTuple{NR},
  T1::CuDenseTensor{<:Number,N1},
  labelsT1::NTuple{N1},
  T2::CuDenseTensor{<:Number,N2},
  labelsT2::NTuple{N2},
) where {NR,N1,N2}
  if N1 == 0
    # TODO: replace with an add! function?
    # What about doing `R .= T1[] .* PermutedDimsArray(T2,perm)`?
    perm = getperm(labelsR, labelsT2)
    newT2 = Tensor(Dense(data(store(T1)) .* data(store(T2))), inds(T2))
    permute!(R, newT2)
  elseif N2 == 0
    perm = getperm(labelsR, labelsT1)
    newT1 = Tensor(Dense(data(store(T2)) .* data(store(T1))), inds(T1))
    permute!(R, newT1)
  elseif N1 + N2 == NR
    # TODO: permute T1 and T2 appropriately first (can be more efficient
    # then permuting the result of T1⊗T2)
    # TODO: implement the in-place version directly
    R = outer!!(R, T1, T2)
    inds_outer = unioninds(inds(T1), inds(T2))
    R = Tensor(store(R), inds_outer)
  else
    R = _contract!!(R, labelsR, T1, labelsT1, T2, labelsT2)
  end
  return R
end

function permutedims!!(
  B::CuDenseTensor{ElT,0},
  A::CuDenseTensor{ElT,0},
  perm::NTuple{0,Int},
  f=(r, t) -> permute!(r, t),
) where {ElT<:Number}
  Cs = f(B, A)
  return Tensor(Dense(vec(Cs)), IndexSet{0}())
end

function permutedims!!(
  B::CuDenseTensor{ElT,N},
  A::CuDenseTensor{ElT,0},
  perm::NTuple{N,Int},
  f=(r, t) -> permute!(r, t),
) where {N,ElT<:Number}
  Cis = ITensors.NDTensors.permute(inds(B), perm)
  Cs = f(B, A)
  return Tensor(Dense(vec(Cs)), Cis)
end

function _contract_scalar!(
  R::CuDenseTensor{ElR},
  labelsR,
  T1::Number,
  labelsT1,
  T2::Number,
  labelsT2,
  α=one(ElR),
  β=zero(ElR),
) where {ElR}
  if iszero(β)
    copyto!(data(R), [α * T1 * T2])
  elseif iszero(α)
    copyto!(data(R), β .* data(R))
  else
    copyto!(data(R), [α * T1 * T2] .+ β .* data(R))
  end
  return R
end

function _contract_scalar!(
  R::CuDenseTensor{ElR,NR},
  labelsR,
  T₁::CuDenseTensor,
  labelsT₁,
  T₂::CuDenseTensor,
  labelsT₂,
  α=one(ElR),
  β=zero(ElR),
) where {ElR,NR}
  if nnz(T₁) == nnz(T₂) == 1
    new_R = Tensor(Dense(data(store(T₁)) .* data(store(T₂))), inds(R))
    copyto!(store(R), store(new_R))
  elseif nnz(T₁) == 1
    props = ContractionProperties(labelsT₁, labelsT₂, labelsR)
    compute_contraction_properties!(props, T₁, T₂, R)
    R = _contract!(R, T₁, T₂, props, α, β)
    #perm = getperm(labelsR,labelsT₂)
    #newT2 = Tensor(Dense(data(store(T₁)).*data(store(T₂))), inds(T₂))
    #permute!(R,newT2)
  elseif nnz(T₂) == 1
    props = ContractionProperties(labelsT₁, labelsT₂, labelsR)
    compute_contraction_properties!(props, T₁, T₂, R)
    R = _contract!(R, T₁, T₂, props, α, β)
    #perm = getperm(labelsR,labelsT₁)
    #newT1 = Tensor(Dense(data(store(T₁)).*data(store(T₂))), inds(T₁))
    #permute!(R,newT1)
  else
    error("In _contract_scalar!, one tensor must be a scalar")
  end
  return R
end

function _gemm_contract!(
  CT::DenseTensor{El,NC},
  AT::DenseTensor{El,NA},
  BT::DenseTensor{El,NB},
  props::ContractionProperties,
  α::Number=one(El),
  β::Number=zero(El),
) where {El,NC,NA,NB}
  # TODO: directly use Tensor instead of Array
  C = array(CT)
  A = array(AT)
  B = array(BT)

  tA = 'N'
  if props.permuteA
    pA = NTuple{NA,Int}(props.PA)
    Ap = permutedims(A, pA)
    AM = reshape(Ap, props.dmid, props.dleft)
    tA = 'T'
  else
    #A doesn't have to be permuted
    if Atrans(props)
      AM = reshape(A, props.dmid, props.dleft)
      tA = 'T'
    else
      AM = reshape(A, props.dleft, props.dmid)
    end
  end

  tB = 'N'
  if props.permuteB
    pB = NTuple{NB,Int}(props.PB)
    Bp = permutedims(B, pB)
    BM = reshape(Bp, props.dmid, props.dright)
  else
    if Btrans(props)
      BM = reshape(B, props.dright, props.dmid)
      tB = 'T'
    else
      BM = reshape(B, props.dmid, props.dright)
    end
  end

  #TODO: this logic may be wrong
  if props.permuteC
    #Need to copy here since we will be permuting
    #into C later
    CM = reshape(copy(C), props.dleft, props.dright)
  else
    if Ctrans(props)
      CM = reshape(C, props.dleft, props.dright)
      (AM, BM) = (BM, AM)
      if tA == tB
        tA = tB = (tA == 'T' ? 'N' : 'T')
      end
    else
      CM = reshape(C, props.dleft, props.dright)
    end
  end

  CM = CUBLAS.gemm!(tA, tB, El(α), AM, BM, El(β), CM)

  if props.permuteC
    pC = NTuple{NC,Int}(props.PC)
    Cr = reshape(CM, props.newCrange...)
    @strided C .= permutedims(Cr, pC)
  end
  return C
end

function _contract!(
  CT::CuDenseTensor{El,NC},
  AT::CuDenseTensor{El,NA},
  BT::CuDenseTensor{El,NB},
  props::ContractionProperties,
  α::Number=one(El),
  β::Number=zero(El),
) where {El,NC,NA,NB}
  if ndims(CT) > 12 || ndims(BT) > 12 || ndims(AT) > 12
    return _gemm_contract!(CT, AT, BT, props, α, β)
  end
  Ainds = inds(AT)
  Adims = dims(Ainds)
  Binds = inds(BT)
  Bdims = dims(Binds)
  Cinds = inds(CT)
  Cdims = dims(Cinds)
  Adata = reshape(data(store(AT)), Adims...)
  Bdata = reshape(data(store(BT)), Bdims...)
  Cdata = reshape(data(store(CT)), Cdims...)
  contracted = commoninds(Ainds, Binds)
  A_only = uniqueinds(Ainds, Binds)
  B_only = uniqueinds(Binds, Ainds)
  ind_dict = Vector{Index}()
  for (idx, i) in enumerate(contracted)
    push!(ind_dict, i)
  end
  if length(A_only) > 0
    for (idx, i) in enumerate(A_only)
      push!(ind_dict, i)
    end
  end
  if length(B_only) > 0
    for (idx, i) in enumerate(B_only)
      push!(ind_dict, i)
    end
  end
  ctainds = zeros(Int, length(Ainds))
  ctbinds = zeros(Int, length(Binds))
  ctcinds = zeros(Int, length(Cinds))
  for (ii, ia) in enumerate(Ainds)
    ctainds[ii] = findfirst(x -> x == ia, ind_dict)
  end
  for (ii, ib) in enumerate(Binds)
    ctbinds[ii] = findfirst(x -> x == ib, ind_dict)
  end
  for (ii, ic) in enumerate(Cinds)
    ctcinds[ii] = findfirst(x -> x == ic, ind_dict)
  end
  id_op = CUDA.CUTENSOR.CUTENSOR.CUTENSOR_OP_IDENTITY
  dict_key = ""
  for cc in zip(ctcinds, Cdims)
    dict_key *= string(cc[1]) * "," * string(cc[2]) * ","
  end
  for aa in zip(ctainds, Adims)
    dict_key *= string(aa[1]) * "," * string(aa[2]) * ","
  end
  for bb in zip(ctbinds, Bdims)
    dict_key *= string(bb[1]) * "," * string(bb[2]) * ","
  end
  if haskey(ENV, "CUTENSOR_AUTOTUNE") && tryparse(Int, ENV["CUTENSOR_AUTOTUNE"]) == 1
    if haskey(ContractionPlans, dict_key)
      dict_val = ContractionPlans[dict_key]
      algo = dict_val
      #plan  = dict_val[2]
      Cdata = CUDA.CUTENSOR.contraction!(
        α,
        Adata,
        Vector{Char}(ctainds),
        id_op,
        Bdata,
        Vector{Char}(ctbinds),
        id_op,
        β,
        Cdata,
        Vector{Char}(ctcinds),
        id_op,
        id_op;
        algo=algo,
      )
    else
      # loop through all algos
      # pick the fastest one
      # store that plan!
      best_time = 1e6
      best_plan = nothing
      best_algo = nothing
      max_algos = Ref{Int32}(C_NULL)
      CUDA.CUTENSOR.cutensorContractionMaxAlgos(max_algos)
      # fix once the other options are documented
      #algos = collect(Cint(CUDA.CUTENSOR.CUTENSOR_ALGO_GETT):Cint(max_algos[] - 1))
      algos = collect(Cint(CUDA.CUTENSOR.CUTENSOR_ALGO_GETT):Cint(-1))
      for algo in reverse(algos)
        try
          Cdata, this_time, bytes, gctime, memallocs = @timed CUDA.CUTENSOR.contraction!(
            α,
            Adata,
            Vector{Char}(ctainds),
            id_op,
            Bdata,
            Vector{Char}(ctbinds),
            id_op,
            β,
            Cdata,
            Vector{Char}(ctcinds),
            id_op,
            id_op;
            algo=CUDA.CUTENSOR.cutensorAlgo_t(algo),
          )
          if this_time < best_time
            best_time = this_time
            #best_plan = this_plan
            best_algo = CUDA.CUTENSOR.cutensorAlgo_t(algo)
          end
        catch err
          @warn "Algorithm $algo not supported"
        end
      end
      ContractionPlans[dict_key] = best_algo
    end
  else
    Cdata = CUDA.CUTENSOR.contraction!(
      α,
      Adata,
      Vector{Char}(ctainds),
      id_op,
      Bdata,
      Vector{Char}(ctbinds),
      id_op,
      β,
      Cdata,
      Vector{Char}(ctcinds),
      id_op,
      id_op,
    )
  end
  return parent(Cdata)
end

function Base.:+(B::CuDenseTensor, A::CuDenseTensor)
  opC = CUTENSOR.CUTENSOR_OP_IDENTITY
  opA = CUTENSOR.CUTENSOR_OP_IDENTITY
  opAC = CUTENSOR.CUTENSOR_OP_ADD
  Ais = inds(A)
  Bis = inds(B)
  ind_dict = Vector{Index}()
  for (idx, i) in enumerate(inds(A))
    push!(ind_dict, i)
  end
  Adata = data(store(A))
  Bdata = data(store(B))
  reshapeBdata = reshape(Bdata, dims(Bis)...)
  reshapeAdata = reshape(Adata, dims(Ais)...)
  ctainds = zeros(Int, length(Ais))
  ctbinds = zeros(Int, length(Bis))
  for (ii, ia) in enumerate(Ais)
    ctainds[ii] = findfirst(x -> x == ia, ind_dict)
  end
  for (ii, ib) in enumerate(Bis)
    ctbinds[ii] = findfirst(x -> x == ib, ind_dict)
  end
  ctcinds = copy(ctbinds)
  C = CUDA.zeros(eltype(Bdata), dims(Bis)...)
  CUTENSOR.elementwiseBinary!(
    one(eltype(Adata)),
    reshapeAdata,
    ctainds,
    opA,
    one(eltype(Bdata)),
    reshapeBdata,
    ctbinds,
    opC,
    C,
    ctcinds,
    opAC,
  )
  copyto!(data(store(B)), vec(C))
  return B
end

function Base.:+(B::CuDense, Bis::IndexSet, A::CuDense, Ais::IndexSet)
  opA = CUTENSOR.CUTENSOR_OP_IDENTITY
  opC = CUTENSOR.CUTENSOR_OP_IDENTITY
  opAC = CUTENSOR.CUTENSOR_OP_ADD
  ind_dict = Vector{Index}()
  for (idx, i) in enumerate(Ais)
    push!(ind_dict, i)
  end
  Adata = data(A)
  Bdata = data(B)
  reshapeBdata = reshape(Bdata, dims(Bis)...)
  reshapeAdata = reshape(Adata, dims(Ais)...)
  ctainds = zeros(Int, length(Ais))
  ctbinds = zeros(Int, length(Bis))
  for (ii, ia) in enumerate(Ais)
    ctainds[ii] = findfirst(x -> x == ia, ind_dict)
  end
  for (ii, ib) in enumerate(Bis)
    ctbinds[ii] = findfirst(x -> x == ib, ind_dict)
  end
  ctcinds = copy(ctbinds)
  C = CUDA.zeros(eltype(Bdata), dims(Bis)...)
  Cis = Bis
  C = CUTENSOR.elementwiseBinary!(
    1, reshapeAdata, ctainds, opA, 1, reshapeBdata, ctbinds, opC, C, ctcinds, opAC
  )
  copyto!(data(B), vec(C))
  return C
end

function Base.:-(B::CuDenseTensor, A::CuDenseTensor)
  opC = CUTENSOR.CUTENSOR_OP_IDENTITY
  opA = CUTENSOR.CUTENSOR_OP_IDENTITY
  opAC = CUTENSOR.CUTENSOR_OP_ADD
  Ais = inds(A)
  Bis = inds(B)
  ind_dict = Vector{Index}()
  for (idx, i) in enumerate(inds(A))
    push!(ind_dict, i)
  end
  Adata = data(store(A))
  Bdata = data(store(B))
  reshapeBdata = reshape(Bdata, dims(Bis)...)
  reshapeAdata = reshape(Adata, dims(Ais)...)
  ctainds = zeros(Int, length(Ais))
  ctbinds = zeros(Int, length(Bis))
  for (ii, ia) in enumerate(Ais)
    ctainds[ii] = findfirst(x -> x == ia, ind_dict)
  end
  for (ii, ib) in enumerate(Bis)
    ctbinds[ii] = findfirst(x -> x == ib, ind_dict)
  end
  ctcinds = copy(ctbinds)
  C = CUDA.zeros(eltype(Bdata), dims(Bis))
  CUTENSOR.elementwiseBinary!(
    one(eltype(Adata)),
    reshapeAdata,
    ctainds,
    opA,
    -one(eltype(Bdata)),
    reshapeBdata,
    ctbinds,
    opC,
    C,
    ctcinds,
    opAC,
  )
  copyto!(data(store(B)), vec(C))
  return B
end

function Base.:-(A::CuDense, Ais::IndexSet, B::CuDense, Bis::IndexSet)
  opA = CUTENSOR.CUTENSOR_OP_IDENTITY
  opC = CUTENSOR.CUTENSOR_OP_IDENTITY
  opAC = CUTENSOR.CUTENSOR_OP_ADD
  ind_dict = Vector{Index}()
  for (idx, i) in enumerate(Ais)
    push!(ind_dict, i)
  end
  Adata = data(A)
  Bdata = data(B)
  reshapeBdata = reshape(Bdata, dims(Bis)...)
  reshapeAdata = reshape(Adata, dims(Ais)...)
  ctainds = zeros(Int, length(Ais))
  ctbinds = zeros(Int, length(Bis))
  for (ii, ia) in enumerate(Ais)
    ctainds[ii] = findfirst(x -> x == ia, ind_dict)
  end
  for (ii, ib) in enumerate(Bis)
    ctbinds[ii] = findfirst(x -> x == ib, ind_dict)
  end
  ctcinds = copy(ctbinds)
  C = CUDA.zeros(eltype(Bdata), dims(Bis)...)
  Cis = Bis
  C = CUTENSOR.elementwiseBinary!(
    one(eltype(Adata)),
    reshapeAdata,
    ctainds,
    opA,
    -one(eltype(Bdata)),
    reshapeBdata,
    ctbinds,
    opC,
    C,
    ctcinds,
    opAC,
  )
  copyto!(data(B), vec(C))
  return C
end

function Base.permute!(B::CuDenseTensor, A::CuDenseTensor)
  Ais = inds(A)
  Bis = inds(B)
  ind_dict = Vector{Index}()
  for (idx, i) in enumerate(Ais)
    push!(ind_dict, i)
  end
  Adata = data(store(A))
  Bdata = data(store(B))
  reshapeBdata = reshape(Bdata, dims(Bis)...)
  reshapeAdata = reshape(Adata, dims(Ais)...)
  if ndims(A) < 12 # use CUTENSOR
    ctainds = zeros(Int, length(Ais))
    ctbinds = zeros(Int, length(Bis))
    for (ii, ia) in enumerate(Ais)
      ctainds[ii] = findfirst(x -> x == ia, ind_dict)
    end
    for (ii, ib) in enumerate(Bis)
      ctbinds[ii] = findfirst(x -> x == ib, ind_dict)
    end
    CUDA.CUTENSOR.permutation!(
      one(eltype(Adata)),
      reshapeAdata,
      Vector{Char}(ctainds),
      reshapeBdata,
      Vector{Char}(ctbinds),
    )
  else # use GPUArrays
    perm = Int[]
    for aix in Ais
      b_pos = findfirst(bix -> bix == aix, Bis)
      push!(perm, b_pos)
    end
    @assert isperm(perm)
    permutedims!(reshapeBdata, reshapeAdata, invperm(perm))
  end
  return vec(reshapeBdata)
end

function Base.permute!(B::CuDense, Bis::IndexSet, A::CuDense, Ais::IndexSet)
  ind_dict = Vector{Index}()
  for (idx, i) in enumerate(Ais)
    push!(ind_dict, i)
  end
  Adata = data(A)
  Bdata = data(B)
  reshapeBdata = reshape(Bdata, dims(Bis)...)
  reshapeAdata = reshape(Adata, dims(Ais)...)
  ctainds = zeros(Int, length(Ais))
  ctbinds = zeros(Int, length(Bis))
  for (ii, ia) in enumerate(Ais)
    ctainds[ii] = findfirst(x -> x == ia, ind_dict)
  end
  for (ii, ib) in enumerate(Bis)
    ctbinds[ii] = findfirst(x -> x == ib, ind_dict)
  end

  CUDA.CUTENSOR.permutation!(
    one(eltype(Adata)),
    reshapeAdata,
    Vector{Char}(ctainds),
    reshapeBdata,
    Vector{Char}(ctbinds),
  )
  return vec(reshapeBdata)
end

Base.:/(A::CuDenseTensor, x::Number) = A * inv(x)


function cubeside_const_lift(_, id, mousehover)
    h_id, h_index = value(mousehover)
    if h_id == id && h_index >= 1 && h_index <= 6
        side =  CubeSides(mousehover.value[2]-1)
        side == TOP     && return Vec3f0(0,0,1)
        side == BOTTOM  && return Vec3f0(0,0,-1)
        side == FRONT   && return Vec3f0(0,1,0)
        side == BACK    && return Vec3f0(0,-1,0)
        side == LEFT    && return Vec3f0(1,0,0)
        side == RIGHT   && return Vec3f0(-1,0,0)
    end
    Vec3f0(1)
end
import Base: +
Base.clamp(x::T) where {T} = clamp(x, T(0),T(1))
+(a::RGBA{T}, b::Number) where {T} = RGBA{T}(clamp(comp1(a)+b), clamp(comp2(a)+b), clamp(comp3(a)+b), clamp(alpha(a)+b))

function cubeside_color(id, mousehover, startcolors, colortex)
    index = mousehover[2]
    update!(colortex, startcolors)
    if mousehover[1] == id
        (index >= 1 && index <= 6) && (colortex[index] = [startcolors[index] + 0.5])
    end
    nothing
end
function colored_cube()
    xdir    = Vec3f0(1f0,0f0,0f0)
    ydir    = Vec3f0(0f0,1f0,0f0)
    zdir    = Vec3f0(0f0,0f0,1f0)
    origin  = Vec3f0(-0.5)
    const quads = [
        (Quad(origin + zdir,   xdir, ydir), RGBA(1.0f0,0f0,0f0,1f0)), # Top
        (Quad(origin,          ydir, xdir), RGBA(0.5f0,0f0,0f0,1f0)), # Bottom
        (Quad(origin + ydir,   zdir, xdir), RGBA(0f0,1.0f0,0f0,1f0)), #Front
        (Quad(origin,          xdir, zdir), RGBA(0f0,0.5f0,0f0,1f0)), # Back
        (Quad(origin + xdir,   ydir, zdir), RGBA(0f0,0f0,1.0f0,1f0)), # Right
        (Quad(origin,          zdir, ydir), RGBA(0f0,0f0,0.5f0,1f0)), # Left
    ]
    cube_steering = merge(map(GLNormalMesh, quads))
end

Base.middle(r::SimpleRectangle{T}) where {T} = Point{2, T}(r.x+(r.w/T(2)), r.y+(r.mousehover/T(2)))

function get_rotation(m)
    idx = Vec3(1, 2, 3)
    xs = norm(m[1, idx])
    ys = norm(m[2, idx])
    zs = norm(m[3, idx])
    Mat4f0(
        m[1,1]/xs, m[1,2]/xs, m[1,3]/xs, 0,
        m[2,1]/ys, m[2,2]/ys, m[2,3]/ys, 0,
        m[3,1]/zs, m[3,2]/zs, m[3,3]/zs, 0,
        0     , 0     , 0     , 1,
    )
end

"""
Creates a camera which is steered by a cube for `window`.
"""
function cubecamera(
        window, projectiontype=Signal(GLAbstraction.PERSPECTIVE);
        cube_area   = Signal(SimpleRectangle(0,0,150,150)),
        eyeposition = Vec3f0(2),
        lookatv     = Vec3f0(0),
        trans       = Signal(Vec3f0(0)),
        theta       = Signal(Vec3f0(0))
    )
    T = Float32
    @materialize mouse_buttons_pressed, mouseposition, buttons_pressed = window.inputs
    dd = doubleclick(mouse_buttons_pressed, 0.3)
    mousehover = GLWindow.mouse2id(window)
    id = Signal(3)
    should_reset = filter(false, dd) do _
        x = value(mousehover).id == value(id)
    end
    cube_steering = colored_cube()
    new_eyeposition = preserve(const_lift(cubeside_const_lift,
        should_reset, id,
        Signal(mousehover)
    ))

    left_ctrl = Set([GLFW.KEY_LEFT_CONTROL])
    use_cam = map(buttons_pressed, window.inputs[:mouseinside]) do b, mi
        b == left_ctrl && mi
    end
    theta, trans = default_camera_control(
        window.inputs, Signal(0.02f0), Signal(1f0),
        use_cam
    )
    fov, upvector = map(Signal, (43f0, Vec3f0(0,0,1)))
    lookatvec, eyeposition =  Signal(lookatv), Signal(eyeposition)
    far = map(eyeposition, lookatvec) do a,b
        max(norm(b-a) * 5f0, 100f0)
    end
    near = map(eyeposition, lookatvec) do a,b
        norm(b-a) * 0.007f0
    end
    main_cam = PerspectiveCamera(
        theta ,trans, lookatvec, eyeposition, upvector,
        window.area, fov, near, far,
        projectiontype
    )
    preserve(map(new_eyeposition) do neweye_dir
        eypos, lookv = value(eyeposition), value(lookatvec)
        oldlength = norm(eypos - lookv)
        push!(eyeposition, lookv+neweye_dir*oldlength)
        for i=1:3
            unitvec = unit(Vec3f0, i)
            if dot(neweye_dir, unitvec) == 0f0 # find an orthogonal vector
                push!(upvector, unitvec)
            end
        end
    end)
    window.cameras[:perspective] = main_cam

    viewmatrix = map(main_cam.view) do m
        inv(get_rotation(m))
    end

    robj = visualize(cube_steering, model=viewmatrix)

    start_colors = cube_steering.attributes
    color_tex    = robj.children[][:attributes]
    preserve(const_lift(cubeside_color, id, mousehover, Signal(start_colors), Signal(color_tex)))
    preserve(id)
    push!(id, robj.children[].id)
    robj
end


using GLFW, ModernGL

# whenever the window size changed (by OS or user resize) this callback function executes
function framebuffer_size_callback(window::GLFW.Window, width::Cint, height::Cint)
    # make sure the viewport matches the new window dimensions; note that width and
    # height will be significantly larger than specified on retina displays.
	glViewport(0, 0, width, height)
end

# process all input: query GLFW whether relevant keys are pressed/released this frame and react accordingly
function processInput(window::GLFW.Window)
    GLFW.GetKey(window, GLFW.KEY_ESCAPE) == GLFW.PRESS && GLFW.SetWindowShouldClose(window, true)
end

# misc. config
const SCR_WIDTH = 800
const SCR_HEIGHT = 600

const vertexShaderSource = """
	#version 330 core
	layout (location = 0) in vec3 aPos;
	void main(void)
	{
	    gl_Position = vec4(aPos.x, aPos.y, aPos.z, 1.0);
	}"""

const fragmentShader1Source = """
	#version 330 core
	out vec4 FragColor;
	void main(void)
	{
	    FragColor = vec4(1.0f, 0.5f, 0.2f, 1.0f);
	}"""

const fragmentShader2Source = """
	#version 330 core
	out vec4 FragColor;
	void main(void)
	{
	    FragColor = vec4(1.0f, 1.0f, 0.0f, 1.0f);
	}"""

GLFW.WindowHint(GLFW.CONTEXT_VERSION_MAJOR, 3)
GLFW.WindowHint(GLFW.CONTEXT_VERSION_MINOR, 3)
GLFW.WindowHint(GLFW.OPENGL_PROFILE, GLFW.OPENGL_CORE_PROFILE)
GLFW.WindowHint(GLFW.OPENGL_FORWARD_COMPAT, GL_TRUE)

# create window
window = GLFW.CreateWindow(SCR_WIDTH, SCR_HEIGHT, "LearnOpenGL")
window == C_NULL && error("Failed to create GLFW window.")
GLFW.MakeContextCurrent(window)
GLFW.SetFramebufferSizeCallback(window, framebuffer_size_callback)

# build and compile shader program
vertexShader = glCreateShader(GL_VERTEX_SHADER)
fragmentShaderOrange = glCreateShader(GL_FRAGMENT_SHADER)
fragmentShaderYellow = glCreateShader(GL_FRAGMENT_SHADER)
shaderProgramOrange = glCreateProgram()
shaderProgramYellow = glCreateProgram()

glShaderSource(vertexShader, 1, Ptr{GLchar}[pointer(vertexShaderSource)], C_NULL)
glCompileShader(vertexShader)
glShaderSource(fragmentShaderOrange, 1, Ptr{GLchar}[pointer(fragmentShader1Source)], C_NULL)
glCompileShader(fragmentShaderOrange)
glShaderSource(fragmentShaderYellow, 1, Ptr{GLchar}[pointer(fragmentShader2Source)], C_NULL)
glCompileShader(fragmentShaderYellow)

# link shaders
glAttachShader(shaderProgramOrange, vertexShader)
glAttachShader(shaderProgramOrange, fragmentShaderOrange)
glLinkProgram(shaderProgramOrange)

glAttachShader(shaderProgramYellow, vertexShader)
glAttachShader(shaderProgramYellow, fragmentShaderYellow)
glLinkProgram(shaderProgramYellow)

# set up vertex data (and buffer(s)) and configure vertex attributes
firstTriangle = GLfloat[-0.9,  0.5, 0.0,    # left
                         0.0, -0.5, 0.0,    # right
                         0.45, 0.5, 0.0]    # top
secondTriangle = GLfloat[0.0, -0.5, 0.0,    # left
						 0.9, -0.5, 0.0,    # right
					     0.45, 0.5, 0.0]    # top
vbos = zeros(GLuint, 2)
vaos = zeros(GLuint, 2)
glGenVertexArrays(2, vaos)
glGenBuffers(2, vbos)
# first triangle setup
glBindVertexArray(vaos[1])
glBindBuffer(GL_ARRAY_BUFFER, vbos[1])
glBufferData(GL_ARRAY_BUFFER, sizeof(firstTriangle), firstTriangle, GL_STATIC_DRAW)
glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 3 * sizeof(GLfloat), Ptr{Cvoid}(0))
glEnableVertexAttribArray(0)
# second triangle setup
glBindVertexArray(vaos[2])
glBindBuffer(GL_ARRAY_BUFFER, vbos[2])
glBufferData(GL_ARRAY_BUFFER, sizeof(secondTriangle), secondTriangle, GL_STATIC_DRAW)
glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 3 * sizeof(GLfloat), Ptr{Cvoid}(0))
glEnableVertexAttribArray(0)

# uncomment this call to draw in wireframe polygons.
# glPolygonMode(GL_FRONT_AND_BACK, GL_LINE)

# render loop
while !GLFW.WindowShouldClose(window)
    processInput(window)

    # render
    glClearColor(0.2, 0.3, 0.3, 1.0)
    glClear(GL_COLOR_BUFFER_BIT)

	# draw first triangle
	glUseProgram(shaderProgramOrange)
    glBindVertexArray(vaos[1])
	glDrawArrays(GL_TRIANGLES, 0, 3)

	# draw  second triangle
	glUseProgram(shaderProgramYellow)
    glBindVertexArray(vaos[2])
	glDrawArrays(GL_TRIANGLES, 0, 3)

    # swap buffers and poll IO events
    GLFW.SwapBuffers(window)
    GLFW.PollEvents()
end

# optional: de-allocate all resources once they've outlived their purpose
glDeleteVertexArrays(2, vaos)
glDeleteBuffers(2, vbos)

GLFW.DestroyWindow(window)


function FESpaces.FESpace(model::DistributedDiscreteModel,
                          reffe::Tuple{RaviartThomas,Any,Any};
                          conformity=nothing,kwargs...)

  cell_reffes = map_parts(model.models) do m
    basis,reffe_args,reffe_kwargs = reffe
    cell_reffe = ReferenceFE(m,basis,reffe_args...;reffe_kwargs...)
  end
  _common_fe_space_constructor(model,cell_reffes;conformity,kwargs...)
end

function FESpace(model::DistributedDiscreteModel,
                 reffe::GenericRefFE{RaviartThomas};
                 conformity=nothing, kwargs...)
  cell_reffes = map_parts(model.models) do m
    Fill(reffe,num_cells(m))
  end
  _common_fe_space_constructor(model,cell_reffes;conformity,kwargs...)
end

function _common_fe_space_constructor(model,cell_reffes;conformity,kwargs...)
  sign_flips=_generate_sign_flips(model,cell_reffes)
  spaces = map_parts(model.models,sign_flips,cell_reffes) do m,sign_flip,cell_reffe
     conf = Conformity(testitem(cell_reffe),conformity)
     cell_fe = CellFE(m,cell_reffe,conf,sign_flip)
     FESpace(m, cell_fe; kwargs...)
  end
  gids =  generate_gids(model,spaces)
  vector_type = _find_vector_type(spaces,gids)
  DistributedSingleFieldFESpace(spaces,gids,vector_type)
end

function _generate_sign_flips(model,cell_reffes)
  sign_flips=map_parts(model.models,model.gids.partition,cell_reffes) do m, p, cell_reffe

    D = num_cell_dims(m)
    gtopo = get_grid_topology(m)

    # Extract composition among cells and facets
    cell_wise_facets_ids = get_faces(gtopo, D, D - 1)
    cache_cell_wise_facets_ids = array_cache(cell_wise_facets_ids)

    # Extract cells around facets
    cells_around_facets = get_faces(gtopo, D - 1, D)
    cache_cells_around_facets = array_cache(cells_around_facets)

    ncells = num_cells(m)
    ptrs = Vector{Int32}(undef,ncells+1)
    for cell in 1:ncells
      reffe=cell_reffe[cell]
      ptrs[cell+1] = num_dofs(reffe)
    end
    PArrays.length_to_ptrs!(ptrs)
    ndata = ptrs[end]-1
    data  = Vector{Bool}(undef,ndata)
    data .= false

    for cell in p.oid_to_lid
      sign_flip = view(data,ptrs[cell]:ptrs[cell+1]-1)
      reffe=cell_reffe[cell]
      D = num_dims(reffe)
      face_own_dofs = get_face_own_dofs(reffe)
      facet_lid = get_offsets(get_polytope(reffe))[D] + 1
      cell_facets_ids = getindex!(cache_cell_wise_facets_ids,
                                  cell_wise_facets_ids,
                                  cell)
      for facet_gid in cell_facets_ids
          facet_cells_around = getindex!(cache_cells_around_facets,
                cells_around_facets,
                facet_gid)
          is_slave=false
          if (length(facet_cells_around)==1)
            is_slave == false
          else
            mx=maximum(p.lid_to_gid[facet_cells_around])
            is_slave = (p.lid_to_gid[cell] == mx)
          end
          if is_slave
              for dof in face_own_dofs[facet_lid]
                  sign_flip[dof] = true
              end
          end
          facet_lid = facet_lid + 1
      end
    end
    PArrays.Table(data,ptrs)
  end
  exchange!(sign_flips,model.gids.exchanger)
  sign_flips
end


"""
    StateVecs

Provides a state vector given a `NamedTuple` containing
tuples of the variable name and the number of its components.
"""
module StateVecs

using ..Grids
export StateVec, over_sub_domains, Cut

struct FieldsPerElement{T}
  vals::Vector{T}
end

"""
    StateVec{T, I, NT}

A state vector containing the number of subdomains,
`n_subdomains`, a `NamedTuple` variable mapper, a tuple
of the variable names, and a vector of vectors, containing
the values for all of the variables.
"""
struct StateVec{T, I, NT}
  """
  Number of subdomains.
  """
  n_subdomains::I

  """
  A `NamedTuple` that maps the variable name to its unique index.
  """
  var_mapper::NT

  """
  A tuple of all variable names.
  """
  var_names::Tuple{Vararg{Symbol}}

  """
  A vector of vectors containing all values for all variables.
  """
  fields::Vector{FieldsPerElement{T}}
end

"""
    get_var_mapper(vars)

Get a `NamedTuple` that maps a variable name
and subdomain ID to a unique index. Example:
```
vars = ((:ρ_0, 1), (:w, 3), (:a, 3), (:α_0, 1))
var_mapper = get_var_mapper(vars)
((:ρ_0, :w, :a, :α_0), (ρ_0 = [1], w = [2, 3, 4], a = [5, 6, 7], α_0 = [8]))
```
"""
function get_var_mapper(vars)
  var_names = tuple([v for (v, nsd) in vars]...)
  end_index = cumsum([nsd for (v, nsd) in vars])
  start_index = [1,[1+x for x in end_index][1:end-1]...]
  vals = [collect(a:b) for (a,b) in zip(start_index, end_index)]
  var_mapper = NamedTuple{var_names}(vals)
  return var_names, var_mapper
end

"""
    StateVec(vars::Tuple{Vararg{Tuple{Symbol, I}}}, grid::Grid{T}) where {I, T}

Return a state vector, given a tuple of tuples of variable
names and the number of their subdomains.
"""
function StateVec(vars::Tuple{Vararg{Tuple{Symbol, I}}}, grid::Grid{T}) where {I, T}
  n_subdomains = max([nsd for (v, nsd) in vars]...)
  n_vars = sum([nsd for (v, nsd) in vars])
  var_names, var_mapper = get_var_mapper(vars)
  all_vars = [eltype(grid.dz)(0) for v in 1:n_vars]
  fields = [FieldsPerElement(deepcopy(all_vars)) for k in over_elems(grid)]
  return StateVec(n_subdomains, var_mapper, var_names, fields)
end

function Base.show(io::IO, sv::StateVec)
  println(io, "----------------------- State Vector")
  println(io, "n_subdomains    = ", sv.n_subdomains)
  println(io, "var_mapper      = ", sv.var_mapper)
  headers = [ length(over_sub_domains(sv, name))==1 ? string(name) : string(name)*"_"*string(i_sd)
             for name in sv.var_names for i_sd in over_sub_domains(sv, name)]
  n_vars = length(headers)
  n_elem = length(sv.fields)
  data = reshape([sv[name, k, i_sd] for name in sv.var_names for
    i_sd in over_sub_domains(sv, name) for k in 1:n_elem], n_elem, n_vars)
  Base.print_matrix(io, [reshape(Text.(headers), 1, n_vars); data])
  println(io, "-----------------------")
end

"""
    over_sub_domains(state_vec::StateVec)

Get list of indexes from 1 to the maximum subdomain size.
"""
over_sub_domains(state_vec::StateVec) = 1:state_vec.n_subdomains

"""
    over_sub_domains(state_vec::StateVec, j::Int)

Get list of indexes from 1 to the maximum subdomain size, except the given index.
"""
over_sub_domains(state_vec::StateVec, j::Int) = [i for i in 1:state_vec.n_subdomains if !(i==j)]

"""
    over_sub_domains(state_vec::StateVec, name::Symbol)

Get list of indexes over all subdomains for variable `name`.
"""
over_sub_domains(state_vec::StateVec, name::Symbol) = 1:length(state_vec.var_mapper[name])

Base.getindex(sv::StateVec, name::Symbol, k, i_sd=1) = sv.fields[k].vals[sv.var_mapper[name][i_sd]]
function Base.setindex!(sv::StateVec, val, name::Symbol, k, i_sd = 1)
  sv.fields[k].vals[sv.var_mapper[name][i_sd]] = val
end

abstract type AbstractCut{I} end

"""
    Cut{I} <: AbstractCut{I}

A Cut struct used to slice the state
vector along the grid-element dimension.
This is used to as an API to pass Cuts
into local derivative/interpolation routines.
"""
struct Cut{I} <: AbstractCut{I}
  e::I
end
Base.getindex(sv::StateVec, name::Symbol, Cut::Cut, i_sd=1) = [sv[name, k, i_sd] for k in Cut.e-1:Cut.e+1]

Base.isnan(sv::StateVec) = any([any(isnan.(fpe.vals)) for fpe in sv.fields])
Base.isinf(sv::StateVec) = any([any(isinf.(fpe.vals)) for fpe in sv.fields])

end


module ArrayInterfaceTracker

using ArrayInterfaceCore
using Tracker

ArrayInterfaceCore.ismutable(::Type{<:Tracker.TrackedArray}) = false
ArrayInterfaceCore.ismutable(T::Type{<:Tracker.TrackedReal}) = false
ArrayInterfaceCore.can_setindex(::Type{<:Tracker.TrackedArray}) = false
ArrayInterfaceCore.fast_scalar_indexing(::Type{<:Tracker.TrackedArray}) = false
ArrayInterfaceCore.aos_to_soa(x::AbstractArray{<:Tracker.TrackedReal,N}) where {N} = Tracker.collect(x)

end # module


module KUDNN
using Compat
using CUDArt,CUDNN

const libkudnn = Libdl.find_library(["libkudnn"], [Pkg.dir("KUDNN/src")])
isempty(libkudnn) && error("Cannot find libkudnn")

export kudnnConvolutionForward, kudnnConvolutionBackwardFilter, kudnnConvolutionBackwardData, kudnnGetConvolutionNdForwardOutputDim
export kudnnPoolingForward, kudnnPoolingBackward, kudnnGetPoolingNdForwardOutputDim

# high level C interface (following CUDNN.jl)

using CUDNN: cudnnHandle, ConvolutionDescriptor, CUDNN_CROSS_CORRELATION, CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM, TD, FD, cptr

# TODO: do not construct this every time:
kudnnConvolutionDescriptor(n)=ConvolutionDescriptor(; padding=ntuple(i->0,n), stride=ntuple(i->1,n), upscale=ntuple(i->1,n), mode=CUDNN_CROSS_CORRELATION)

function kudnnGetConvolutionNdForwardOutputDim(src, filter; convDesc=kudnnConvolutionDescriptor(ndims(src)-2))
    nd = ndims(src)
    nd == ndims(filter) || error("Dimension mismatch")
    size(src,nd-1) == size(filter,nd-1) || error("Dimension mismatch")
    od = Array(Int,nd)
    for i=1:nd-2
        od[i] = 1 + div(size(src,i) - size(filter,i) + 2*convDesc.padding[i], convDesc.stride[i])
    end
    od[nd-1] = size(filter,nd)
    od[nd] = size(src,nd)
    tuple(od...)
end

function kudnnConvolutionForward(src::AbstractCudaArray, filter::AbstractCudaArray, dest=nothing;
                                 handle=cudnnHandle, alpha=1.0, beta=0.0,
                                 convDesc=kudnnConvolutionDescriptor(ndims(src)-2),
                                 algorithm=CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM,
                                 workSpace=Int8[], workSpaceSizeInBytes=0)
    @assert ndims(filter) == ndims(src)
    @assert eltype(filter) == eltype(src)
    osize = kudnnGetConvolutionNdForwardOutputDim(src,filter;convDesc=convDesc)
    (dest == nothing) && (dest = CudaArray(eltype(src), osize))
    @assert osize == size(dest)
    @assert eltype(dest) == eltype(src)
    kudnnConvolutionForward(handle,
                            cptr(alpha,src),TD(src),src,
                            FD(filter),filter,
                            convDesc,algorithm,workSpace,workSpaceSizeInBytes,
                            cptr(beta,dest),TD(dest),dest)
    return dest
end

function kudnnConvolutionBackwardFilter(src::AbstractCudaArray, diff::AbstractCudaArray, grad::AbstractCudaArray;
                                        handle=cudnnHandle, alpha=1.0, beta=0.0,
                                        convDesc=kudnnConvolutionDescriptor(ndims(src)-2))
    kudnnConvolutionBackwardFilter(handle,
                                   cptr(alpha,src),TD(src),src,
                                   TD(diff),diff,convDesc,
                                   cptr(beta,grad),FD(grad),grad)
    return grad
end

function kudnnConvolutionBackwardData(filter::AbstractCudaArray, diff::AbstractCudaArray, grad::AbstractCudaArray;
                                      handle=cudnnHandle, alpha=1.0, beta=0.0,
                                      convDesc=kudnnConvolutionDescriptor(ndims(filter)-2))
    kudnnConvolutionBackwardData(handle,cptr(alpha,diff),
                                 FD(filter),filter,
                                 TD(diff),diff,convDesc,
                                 cptr(beta,grad),TD(grad),grad)
    return grad
end

function kudnnGetPoolingNdForwardOutputDim(pd::PoolingDescriptor, input::AbstractCudaArray)
    n = ndims(input)
    ntuple(n) do i
        i >= n-1 ? size(input, i) :
        1 + ceil(Int, (size(input, i) + 2*pd.padding[i] - pd.dims[i]) / pd.stride[i])
    end
end

function kudnnPoolingForward(pd::PoolingDescriptor, src::AbstractCudaArray, dest::AbstractCudaArray;
                             handle=cudnnHandle, alpha=1.0, beta=0.0)
    kudnnPoolingForward(handle, pd,
                        cptr(alpha,src), TD(src), src,
                        cptr(beta,dest), TD(dest), dest)
    return dest
end

function kudnnPoolingBackward(pd::PoolingDescriptor, src::AbstractCudaArray, srcDiff::AbstractCudaArray, dest::AbstractCudaArray, destDiff::AbstractCudaArray;
                              handle=cudnnHandle, alpha=1.0, beta=0.0)
    cudnnSetTensor(destDiff, 0)  # TODO: make sure kudnnPoolingBackward zeros out its output
    kudnnPoolingBackward(handle, pd,
                         cptr(alpha,src), TD(src), src,
                         TD(srcDiff), srcDiff,
                         TD(dest), dest,
                         cptr(beta,destDiff), TD(destDiff), destDiff)
    return destDiff
end

# low level C interface (similar to libcudnn.jl)

using CUDNN: cudnnCheck,cudnnStatus_t,cudnnHandle_t,cudnnTensorDescriptor_t,cudnnFilterDescriptor_t,cudnnConvolutionDescriptor_t,cudnnConvolutionFwdAlgo_t,cudnnPoolingDescriptor_t

function kudnnConvolutionForward(handle,alpha,srcDesc,srcData,filterDesc,filterData,convDesc,algo,workSpace,workSpaceSizeInBytes,beta,destDesc,destData)
    cudnnCheck(ccall((:kudnnConvolutionForward,libkudnn),cudnnStatus_t,(cudnnHandle_t,Ptr{Void},cudnnTensorDescriptor_t,Ptr{Void},cudnnFilterDescriptor_t,Ptr{Void},cudnnConvolutionDescriptor_t,cudnnConvolutionFwdAlgo_t,Ptr{Void},Csize_t,Ptr{Void},cudnnTensorDescriptor_t,Ptr{Void}),handle,alpha,srcDesc,srcData,filterDesc,filterData,convDesc,algo,workSpace,workSpaceSizeInBytes,beta,destDesc,destData))
end

function kudnnConvolutionBackwardFilter(handle,alpha,srcDesc,srcData,diffDesc,diffData,convDesc,beta,gradDesc,gradData)
    cudnnCheck(ccall((:kudnnConvolutionBackwardFilter,libkudnn),cudnnStatus_t,(cudnnHandle_t,Ptr{Void},cudnnTensorDescriptor_t,Ptr{Void},cudnnTensorDescriptor_t,Ptr{Void},cudnnConvolutionDescriptor_t,Ptr{Void},cudnnFilterDescriptor_t,Ptr{Void}),handle,alpha,srcDesc,srcData,diffDesc,diffData,convDesc,beta,gradDesc,gradData))
end

function kudnnConvolutionBackwardData(handle,alpha,filterDesc,filterData,diffDesc,diffData,convDesc,beta,gradDesc,gradData)
    cudnnCheck(ccall((:kudnnConvolutionBackwardData,libkudnn),cudnnStatus_t,(cudnnHandle_t,Ptr{Void},cudnnFilterDescriptor_t,Ptr{Void},cudnnTensorDescriptor_t,Ptr{Void},cudnnConvolutionDescriptor_t,Ptr{Void},cudnnTensorDescriptor_t,Ptr{Void}),handle,alpha,filterDesc,filterData,diffDesc,diffData,convDesc,beta,gradDesc,gradData))
end

function kudnnPoolingForward(handle,poolingDesc,alpha,srcDesc,srcData,beta,destDesc,destData)
    cudnnCheck(ccall((:kudnnPoolingForward,libkudnn),cudnnStatus_t,(cudnnHandle_t,cudnnPoolingDescriptor_t,Ptr{Void},cudnnTensorDescriptor_t,Ptr{Void},Ptr{Void},cudnnTensorDescriptor_t,Ptr{Void}),handle,poolingDesc,alpha,srcDesc,srcData,beta,destDesc,destData))
end

function kudnnPoolingBackward(handle,poolingDesc,alpha,srcDesc,srcData,srcDiffDesc,srcDiffData,destDesc,destData,beta,destDiffDesc,destDiffData)
    cudnnCheck(ccall((:kudnnPoolingBackward,libkudnn),cudnnStatus_t,(cudnnHandle_t,cudnnPoolingDescriptor_t,Ptr{Void},cudnnTensorDescriptor_t,Ptr{Void},cudnnTensorDescriptor_t,Ptr{Void},cudnnTensorDescriptor_t,Ptr{Void},Ptr{Void},cudnnTensorDescriptor_t,Ptr{Void}),handle,poolingDesc,alpha,srcDesc,srcData,srcDiffDesc,srcDiffData,destDesc,destData,beta,destDiffDesc,destDiffData))
end

end # module KUDNN


using Distributions
using Statistics

@enum ScenarioLoopPlacement OUTER INNER

"""
    RandomVariable{T}

A RandomVariable can be "assigned" to different model parameters. Its
value can be used directly or applied by adding to, or multiplying by,
reference values for the indicated parameter. Note that the distribution
must be instantiated, e.g., call as RandomVariable(:foo, Normal()), not
RandomVariable(:foo, Normal).
"""
struct RandomVariable{T}
    name::Symbol
    dist::T

    function RandomVariable(name::Symbol, dist::T) where T
        self = new{T}(name, dist)
        return self
    end
end

Base.eltype(rv::RandomVariable) = eltype(rv.dist)

distribution(rv::RandomVariable) = rv.dist

abstract type PseudoDistribution end

"""     ReshapedDistribution
A pseudo-distribution that returns a reshaped array of values from the
stored distribution and dimensions.

Example:
    rd = ReshapedDistribution([5, 5], Dirichlet(25,1))
"""
struct ReshapedDistribution <: PseudoDistribution
    dims::Vector{Int}
    dist::Distribution
end

function Base.rand(rd::ReshapedDistribution, draws::Int=1)
    return [reshape(rand(rd.dist), rd.dims...) for i in 1:draws]
end

# SampleStore is a faux Distribution that implements base.rand()
# to yield stored values.
mutable struct SampleStore{T} <: PseudoDistribution
    values::Vector{T}       # generally Int or Float64
    idx::Int                # index of next value to return
    dist::Union{Nothing, Distribution, PseudoDistribution}  # original distribution, if any

    function SampleStore(values::Vector{T}, dist::Union{Nothing, Distribution, PseudoDistribution}=nothing) where T
        return new{T}(values, 1, dist)
    end
end

Base.eltype(ss::SampleStore) = eltype(ss.values)

#
# TBD: maybe have a different SampleStore subtype for values drawn from a dist
# versus those loaded from a file, which would be treated as immutable?
#

function Statistics.quantile(ss::SampleStore{T}, q::Float64) where T
    return quantile(sort(ss.values), q)
end

# TBD: This interpolates values between those in the vector. Is this reasonable?
# Probably shouldn't use correlation on values loaded from a file rather than
# from a proper distribution.
function Statistics.quantile(ss::SampleStore{T}, probs::AbstractArray) where T
    return quantile(sort(ss.values), probs)
end

Base.length(ss::SampleStore{T}) where T = length(ss.values)

Base.iterate(ss::SampleStore{T}) where T = iterate(ss.values)
Base.iterate(ss::SampleStore{T}, idx) where T = iterate(ss.values, idx)

struct TransformSpec
    paramname::Symbol
    op::Symbol
    rvname::Symbol
    dims::Vector{Any}

    function TransformSpec(paramname::Symbol, op::Symbol, rvname::Symbol, dims::Vector{T}=[]) where T
        if ! (op in (:(=), :(+=), :(*=)))
            error("Valid operators are =, +=, and *= (got $op)")
        end

        return new(paramname, op, rvname, dims)
    end
end

"""
    Base.rand(s::SampleStore{T}, n::Int=1) where T

Pseudo `rand`` function that just returns the next `n` values from the vector of
stored values. Currently does not support cycling through results multiple times,
but you can `reset` the `SampleStore` to reuse it.
"""
function Base.rand(s::SampleStore{T}, n::Int=1) where T
    idx = s.idx
    s.idx = next_idx = idx + n
    return n > 1 ? s.values[idx:(next_idx - 1)] : s.values[idx]
end

function Base.reset(s::SampleStore{T}) where T
    s.idx = 1
    return nothing
end

abstract type AbstractSimulationData end

"""
    SimulationDef

Holds all the data that defines a simulation.
"""
mutable struct SimulationDef{T}
    rvdict::OrderedDict{Symbol, RandomVariable}
    translist::Vector{TransformSpec}
    savelist::Vector{Tuple{Symbol, Symbol}}
    nt_type::Any                    # a generated NamedTuple type to hold data for a single trial
    data::T                         # data specific to a given sensitivity analysis method
    payload::Any                    # opaque (to Mimi) data the user wants access to in callbacks

    function SimulationDef{T}(rvlist::Vector,
                           translist::Vector{TransformSpec},
                           savelist::Vector{Tuple{Symbol, Symbol}},
                           data::T) where T <: AbstractSimulationData
        self = new()
        self.rvdict = OrderedDict([rv.name => rv for rv in rvlist])
        self.translist = translist
        self.savelist = savelist

        names = (keys(self.rvdict)...,)
        types = [eltype(fld) for fld in values(self.rvdict)]
        self.nt_type = NamedTuple{names, Tuple{types...}}

        self.data = data
        self.payload = nothing

        return self
    end
end

"""
    SimulationInstance{T}

Holds all the data that defines simulation results.
"""
mutable struct SimulationInstance{T}
    trials::Int
    current_trial::Int
    current_data::Any               # holds data for current_trial when current_trial > 0
    sim_def::SimulationDef{T} where T <: AbstractSimulationData
    models::Vector{M} where M <: AbstractModel
    results::Vector{Dict{Tuple, DataFrame}}
    payload::Any

    function SimulationInstance{T}(sim_def::SimulationDef{T}) where T <: AbstractSimulationData
        self = new()
        self.trials = 0
        self.current_trial = 0
        self.current_data = nothing
        self.sim_def = deepcopy(sim_def)
        self.payload = deepcopy(self.sim_def.payload)

        # These are parallel arrays; each model has a corresponding results dict
        self.models = Vector{AbstractModel}(undef, 0)
        self.results = [Dict{Tuple, DataFrame}()]

        return self
    end
end

"""
    SimulationDef{T}() where T <: AbstractSimulationData

Allow creation of an "empty" SimulationDef instance.
"""
function SimulationDef{T}() where T <: AbstractSimulationData
    SimulationDef{T}([], TransformSpec[], Tuple{Symbol, Symbol}[], T())
end

"""
    set_payload!(sim_def::SimulationDef, payload)

Attach a user's `payload` to the `SimulationDef`. A copy of the payload object
will be stored in the `SimulationInstance` at run time so it can be
accessed in scenario and pre-/post-trial callback functions. The value
is not used by Mimi in any way; it can be anything useful to the user.
"""
set_payload!(sim_def::SimulationDef, payload) = (sim_def.payload = payload)

"""
    payload(sim_def::SimulationDef)

Return the `payload` value set by the user via `set_payload!()`.
"""
payload(sim_def::SimulationDef) = sim_def.payload

"""
    payload(sim_inst::SimulationInstance)

Return the copy of the `payload` value stored in the `SimulationInstance` set by the user via `set_payload!()`.
"""
payload(sim_inst::SimulationInstance) = sim_inst.payload

struct MCSData <: AbstractSimulationData end

const MonteCarloSimulationDef = SimulationDef{MCSData}
const MonteCarloSimulationInstance = SimulationInstance{MCSData}

struct SimIterator{NT, T}
    sim_inst::SimulationInstance{T}

    function SimIterator{NT, T}(sim_inst::SimulationInstance{T}) where {NT <: NamedTuple, T <: AbstractSimulationData}
        return new{NT, T}(sim_inst)
    end
end

function getdataframe(sim_inst::SimulationInstance, comp_name::Symbol, datum_name::Symbol; model::Int = 1)
    return sim_inst.results[model][(comp_name, datum_name)]
end

function Base.getindex(sim_inst::SimulationInstance, comp_name::Symbol, datum_name::Symbol; model::Int = 1)
    msg = "getindex method for `SimulationInstance` has been replaced with getdataframe function for consistency of return type, please use getdataframe instead"
    Base.depwarn("$msg, $(stacktrace())", :getindex)
    return sim_inst.results[model][(comp_name, datum_name)]
end


reversedims(A::Array) = permutedims(A, ndims(A):-1:1) # @generated probably the "right" way

for MAType in (Float32MultiArray, Float64MultiArray, Int32MultiArray, Int64MultiArray)
    @eval function Base.convert(::Type{$MAType}, A::Array{<:Real})
        Arev = reversedims(A)
        msg = $MAType()
        msg.layout.dim = [MultiArrayDimension("",d,s) for (d,s) in zip(size(A), reverse(cumprod([size(Arev)...])))]
        msg.data = Arev[:]
        msg
    end

    @eval function Base.convert(::Type{Array{T}}, msg::$MAType) where {T}
        if isempty(msg.data)
            T[]
        else
            reversedims(reshape(msg.data, (Int(d.size) for d in msg.layout.dim[end:-1:1])...))
        end
    end
end

mutable struct PredictionOutputHolder
    stale::Bool
    y
    z
    r
    b
end

function prediction_callback(msg::PredictionOutput, pred_container::PredictionOutputHolder)
    pred_container.stale = false
    pred_container.y = Array{Float32}(msg.y)
    pred_container.z = Array{Float32}(msg.z)
    pred_container.r = Array{Float32}(msg.r)
    pred_container.b = Array{Float32}(msg.b)
end

const pred_container = PredictionOutputHolder(true,nothing,nothing,nothing,nothing)
const publisher = Publisher{PredictionInput}("prediction_input", queue_size=10)
const subscriber = Subscriber{PredictionOutput}("prediction_output", prediction_callback, (pred_container,), queue_size=10)

function predict!(car1, car2, extras, traj_lengths, car1_future, sample_ct,
                  output::PredictionOutputHolder=pred_container, timeout=1000)
    output.stale = true
    publish(publisher, PredictionInput(Float32MultiArray(car1),
                                       Float32MultiArray(car2),
                                       Float32MultiArray(extras),
                                       Int32MultiArray(traj_lengths),
                                       Float32MultiArray(car1_future),
                                       Float32MultiArray(Float32[]),
                                       Float32MultiArray(Float32[]),
                                       Int32(sample_ct)))
    for i in 1:timeout
        rossleep(Duration(0.001))
        output.stale || break
    end
    output
end

function predict!(car1, car2, extras, traj_lengths, car1_future_x, car1_future_y, sample_ct,
                  output::PredictionOutputHolder=pred_container, timeout=1000)
    output.stale = true
    publish(publisher, PredictionInput(Float32MultiArray(car1),
                                       Float32MultiArray(car2),
                                       Float32MultiArray(extras),
                                       Int32MultiArray(traj_lengths),
                                       Float32MultiArray(Float32[]),
                                       Float32MultiArray(car1_future_x),
                                       Float32MultiArray(car1_future_y),
                                       Int32(sample_ct)))
    for i in 1:timeout
        rossleep(Duration(0.001))
        output.stale || break
    end
    output
end

######
# const y_top = -1.83284f0
# const y_bot = -6.09485f0
const y_top = -1.5f0
const y_bot = -6.5f0
const xdd_choices = Float32[0, -3, 4, -6]
const y_targets = [y_top, y_bot]
const tibvp = SteeringBVP(TripleIntegratorDynamics(1, Float32), TimePlusQuadraticControl(SMatrix{1,1}(Float32(1e-3))))

trapz(x::Vector, dt, x0=0) = cumsum(vcat(0, 0.5*dt*(x[1:end-1] + x[2:end]))) + x0
trapz(x::Matrix, dt, x0=0) = cumsum(hcat(zeros(size(x,1)), 0.5*dt*(x[:,1:end-1] + x[:,2:end])), 2) .+ x0
vel(xd, yd) = sqrt.(xd.^2 + yd.^2)
acc(xd, xdd, yd, ydd, v=vel(xd,yd)) = (xd.*xdd + yd.*ydd)./v
curv(xd, xdd, yd, ydd, v=vel(xd,yd)) = (xd.*ydd - xdd.*yd)./v.^3

@inbounds @views function all_future_xy(xdd_choices, y_targets, state0, tibvp,
                                        a0 = [(1,1)], ply=4, t_hold=3, dt=Float32(0.1);
                                        des_length=length(a0)+ply*t_hold)
    Nx = length(xdd_choices)
    Ny = length(y_targets)
    X = zeros(Float32, Nx^ply, 1 + des_length, 3)
    Y = zeros(Float32, Ny^ply, 1 + des_length, 3)
    if a0 isa Vector{Tuple{Int,Int}}
        X[:,1,:] .= state0[1:2:end]'
        Y[:,1,:] .= state0[2:2:end]'
        t0 = 1 + length(a0)
        for (t, (xi, yi)) in enumerate(a0)
            X[:,t+1,3] .= xdd_choices[xi]
            q0 = SVector(Y[1,t,1],
                         Y[1,t,2],
                         Y[1,t,3])
            qf = SVector(y_targets[yi], 0f0, 0f0)
            c, u = tibvp(q0, qf, 100f0)
            if dt > duration(u)
                q = qf
            else
                q = u.x(q0, u.y, u.t, dt)
            end
            Y[:,t+1,1] .= q[1]
            Y[:,t+1,2] .= q[2]
            Y[:,t+1,3] .= q[3]
        end
    elseif a0 isa Tuple
        Ẍ0, Ÿ0 = a0
        X[:,1,:] .= state0[1:2:end]'
        Y[:,1,:] .= state0[2:2:end]'
        X[:,2:1+length(Ẍ0),3] .= Ẍ0'
        Y[:,2:1+length(Ÿ0),3] .= Ÿ0'
        t0 = 1 + length(Ÿ0)
        cumsum!(Y[:,2:t0,2], Y[:,1:t0-1,3]*dt, 2)    # X integrated below (joint with the case above)
        Y[:,2:t0,2] .+= Y[:,1,2]
        cumsum!(Y[:,2:t0,1], (Y[:,1:t0-1,2] .+ Y[:,2:t0,2])*dt./2, 2)
        Y[:,2:t0,1] .+= Y[:,1,1]
    else
        error("Incompatible a0 type.")
    end

    # xdd, xd, x
    for (r, ci) in enumerate(CartesianRange(Tuple(Nx for i in 1:ply)))
        for (p, v) in enumerate(reverse(ci.I))    # reverse so that later indices (actions for later timesteps) iterate fastest
            for t in 1:t_hold
                X[r,t0+t_hold*(p-1)+t,3] = xdd_choices[v]
            end
        end
    end
    for t in (t0+ply*t_hold+1):1+des_length
        X[:,t,3] = X[:,t-1,3]
    end
    cumsum!(X[:,2:end,2], X[:,1:end-1,3]*dt, 2)
    X[:,2:end,2] .+= X[:,1,2]
    cumsum!(X[:,2:end,1], (X[:,1:end-1,2] .+ X[:,2:end,2])*dt./2, 2)
    X[:,2:end,1] .+= X[:,1,1]

    # ydd, yd, y
    for (r, ci) in enumerate(CartesianRange(Tuple(Ny for i in 1:ply)))
        for (p, v) in enumerate(reverse(ci.I))    # reverse so that later indices (actions for later timesteps) iterate fastest
            curr_step = t0 + (p-1)*t_hold
            q0 = SVector(Y[r,curr_step,1],
                         Y[r,curr_step,2],
                         Y[r,curr_step,3])
            qf = SVector(y_targets[v], 0f0, 0f0)
            c, u = tibvp(q0, qf, 100f0)
            for t in 1:t_hold
                if dt*t > duration(u)
                    q = qf
                else
                    q = u.x(q0, u.y, u.t, t*dt)
                end
                Y[r,curr_step+t,1] = q[1]
                Y[r,curr_step+t,2] = q[2]
                Y[r,curr_step+t,3] = q[3]
            end
        end
    end
    for t in (t0+ply*t_hold+1):1+des_length
        Y[:,t,3] = Y[:,t-1,3]
        Y[:,t,2] = Y[:,t-1,2] + Y[:,t-1,3]*dt
        Y[:,t,1] = Y[:,t-1,1] + Y[:,t-1,2]*dt
    end
    X, Y
end

function propagate(xdd_choices, y_targets, car1_curr_state, tibvp, action_sequence)
    pfX, pfY = all_future_xy(xdd_choices, y_targets, car1_curr_state, tibvp, action_sequence, 0)
    [SVector{6,Float32}(pfX[1,i,1], pfY[1,i,1],
                        pfX[1,i,2], pfY[1,i,2],
                        pfX[1,i,3], pfY[1,i,3]) for i in 1:size(pfX, 2)]
end


module TrackError

using ..Util, ..Tagging, ..Conditioning, ..Traits
import ..Conditioning
export applytrackerr, condvar, condvarapply, tagerror, initerror

# Options.
#> condition on logerr, explicitly

# Options
# 1. Always initialize with correct type -- best performance
# 2. Have nothing, small union ⊥ ∪ Float64
# 3. Vector

"""
`applytrackerr(x, ω, initerr)`

Is `ω` in the domain of `x`?
Returns `(x(ω), v)` where `v = ω ∈ dom(x)`
"""
function applytrackerr(x, ω, initerr)
  ω_ = tagerror(ω, initerr) # zt: surprisingly expensive
  fx = x(ω_)
  (fx = fx, err = ω_.tags.err.val)
end

"Equivalent to `condvar(x)(ω)`, but more efficient"
condvarapply(x, ω, initerr) = applytrackerr(x, ω, initerr).err
condvarapply(x, ω, errtype::Type{T}) where T =
  condvarapply(x, ω, initerror(T))

logcondvarapply(x, ω) = condvarapply(x, ω)

function initerror end
initerror(::Type{Bool}) = true
"""
`condvar(x)`

Random variable that `x` is conditioned on.
I.e. if `x` can be expressed as `x | y`, then `y = condvar(x)`
"""
condvar(x, initerr) = ω -> condvarapply(x, ω, initerr)

"Update `err` by conjoining current error with `b`"
conjoinerror!(err::Box, b) = err.val &= b

# Track error for this ω
dotrack(ω) = !haskey(ω.tags, :donttrack) && haskey(ω.tags, :err)

function Conditioning.condf(::trait(Err), ω, x, y)
  dotrack(ω) && conjoinerror!(ω.tags.err, y(ω))
  x(ω)
end

function Conditioning.cond!(::trait(Err), ω, bool)
  dotrack(ω) && conjoinerror!(ω.tags.err, bool)
  bool
end

# conjoinerror!(ω.tags.err, y(ω))
tagerror(::trait(Err), ω, initerr) = tag(ω,  (err = Box(initerr),))
tagerror(traits, ω, initerr) = tag(ω, (err = Box(initerr),))

"Tag `ω` with `err`"
tagerror(ω::Ω, initerr) where Ω = tagerror(traits(Ω), ω, initerr)

function combinetags(::Type{Val{:err}}, a, b)
  # @show a
  # @show b
  # @assert false
  @show a[]
  @show a[] & b[]
end

end


using ShirleyRayTracer
using Images

#==

# launch using 40 threads
 julia -t 40 --project=.. -L RandomScene.jl  -e "main()" > random_scene.ppm

# launch force using 1 thread
 julia -t 1 --project=.. -L RandomScene.jl  -e "main()" > random_scene.ppm

# launch using thread count from environment
 julia --project=.. -L RandomScene.jl  -e "main()" > random_scene.ppm

==#

function add_random_scene!(scene::Scene)

	add!(scene, Sphere(Point3(0,-1000,0), 1000, Lambertian(0.5, 0.5, 0.5)))

	rand_material(p) = if p < 0.8
				Lambertian()
			elseif p < 0.95
				rf = randf(0.5, 1)
				Metal(rf, rf, rf, 0.5rand())
			else
				Dielectric(1.5)
			end

	for a in -11:10, b in -11:10
		center = Point3(a + 0.9rand(), 0.2, b + 0.9rand())
		if ShirleyRayTracer.magnitude(center - Point3(4, 0.2, 0)) > 0.9
			add!(scene, Sphere(center, 0.2, rand_material(rand())))
		end
	end

	add!(scene, Sphere(Point3(0, 1, 0), 1.0, Dielectric(1.5)))
	add!(scene, Sphere(Point3(-4, 1, 0), 1.0, Lambertian(0.4,0.2,0.1)))
	add!(scene, Sphere(Point3(4, 1, 0), 1.0, Metal(0.7,0.6,0.5, 0.0)))
end

function main(;filename="render.jpg", image_width=1200, aspect_ratio=16/9, samples_per_pixel=10, max_depth=50)
	image_height = round(Int, image_width / aspect_ratio)
	world = Scene(Camera(Point3(13.,2.,3.), zero(Point3), Vec3(0,1,0), 20, aspect_ratio, 0.1, 10.0))
	add_random_scene!(world)
	image = render(world, image_width, image_height, samples_per_pixel, max_depth)
	save(filename, image)
end


const DCLR_EX_DAYS = 9
const DIST_DAYS = 12
const DIST_RET = :net
const DIST_RET_BOUND = 0.2

const DIST_MIN_ALLOWED = 1#DIST_DAYS*2
const DIST_DAYS_TO_STALE_PRICE = 22

function getidentifierscrspevent(dist::DataFrame; Fdate::Symbol = :exdate)
  local ccm::DataFrame = loadccm()
  ccm = preprocessccm!(ccm)

  dist.eid = collect(1:size(dist,1)) #assign an id to each event

  dist = join(dist, ccm, on=:permno=>:lpermno)
  filter!(r::DataFrameRow -> r[Fdate] ∈ r.linkeffdate:Day(1):r.linkenddate, dist)
  @assert allunique(dist.eid) #no dups

  #rename!(dist, :lpermno=>:permno) #sanitize permno

  return dist
end

function readcrspevent(;fname::String = CRSP_DISTRIBUTIONS_NAME,
  path::String=EVENT_PATH,
  incsvstream::Function = IN_CSV_STREAM,
  crspdateformat::DateFormat = CRSP_DATE_FORMAT,
  Fdate::Symbol=:exdate,
  distdays::Int = DIST_DAYS)

  dist::DataFrame = incsvstream("$path\\$fname.csv.lz4") |> CSV.read |> DataFrame

  names!(dist, (s::Symbol->Symbol(lowercase(string(s)))).(names(dist)))

  rename!(dist, [:dclrdt=>:dclrdate, :exdt=>:exdate, :paydt=>:paydate])

  for f ∈ [:dclrdate, :exdate, :paydate]
    dist[!,f] = (i::MInt->ismissing(i) ? missing : Date("$i", CRSP_DATE_FORMAT)).(dist[!,f])
  end

  essentialfields::Vector{Symbol} = [:permno, :dclrdate, :exdate, :distcd, :divamt]
  dist = dist[completecases(dist[!, essentialfields]), :]

  #time between exdate and dclrdate is important
  dist.dclrexdays = (d::Day->d.value).(dist.exdate .- dist.dclrdate)

  #ensures integrity
  dist = dist[dist.dclrexdays .≥ 0, :]

  #dist = getidentifierscrspevent(dist, Fdate=Fdate) NOTE: We will need to call this method- just not here

  return dist
end

function processcrspevent!(dist::DataFrame,
    dclrexdays::Int = DCLR_EX_DAYS,
    Fret::Symbol = DIST_RET,
    retbound::Float64=DIST_RET_BOUND,
    minallowed = DIST_MIN_ALLOWED,
    daystostaleprice::Int=DIST_DAYS_TO_STALE_PRICE,
    days::Int = DIST_DAYS)

  local sdists::GroupedDataFrame

  #now compute the dividend yield from the declaration date
  dist.yield = Vector{MFloat64}(undef, size(dist,1))
  dist.winsorizedyield = falses(size(dist,1))
  dist.dclrextradingdays = Vector{MInt}(undef, size(dist,1))

  lowerdays::Int = max(daystostaleprice,days)
  sdists = groupby(dist, :eid)
  @mpar for i ∈ 1:length(sdists)
    sdist::SubDataFrame = sdists[i]
    #some basic integrity checks
    @assert maximum(sdist.divamt) == minimum(sdist.divamt)
    @assert size(sdist,1) == days + lowerdays + 1
    @assert maximum(sdist.dclrexdays) == minimum(sdist.dclrexdays)
    @assert sdist.dclrdate[1] == sdist.date[lowerdays + 1] - Day(sdist.dclrexdays[1])

    #find the row for the declaration date
    dclridx::Int = searchsortedfirst(sdist.date, sdist.dclrdate[1])
    (dclridx ≤ 0 || dclridx>size(sdist,1) || sdist.date[dclridx] ≠ sdist.dclrdate[1]) && continue

    #compute the yield and winsorize if necessary
    yield::MFloat64 = sdist.divamt[1]/sdist.price[dclridx]
    (ismissing(yield) || (yield < 0.0)) && continue
    sdist.dclrextradingdays .=  abs(sdist.day[dclridx])
    sdist.winsorizedyield .= yield > retbound
    sdist.yield .= min(yield,retbound)
  end
  #now determine the dividend yields from the exiration

  filter!(r::DataFrameRow->(r.day ≥ -days), dist)

  dist.primary = Vector{Bool}(undef, size(dist,1))
  dist.match = Vector{Bool}(undef, size(dist,1))
  dist.ordinary = Vector{Bool}(undef, size(dist,1))
  dist.dclrexok = Vector{Bool}(undef, size(dist,1))

  sdists = groupby(dist, :eid)
  @mpar for i ∈ 1:length(sdists)
    sdist = sdists[i]
    r::DataFrameRow = sdist[1, :] #examine fields common to each eid

    notmissing::Int = sum((!ismissing).(sdist[!,Fret]))
    sdist.match .= (notmissing ≥ minallowed) && (!ismissing(r.yield))
    sdist.ordinary .= r.distcd == 1232
    sdist.dclrexok .= r.dclrexdays ≥ dclrexdays
  end

  dist.primary .= dist.ordinary .* dist.match .* dist.nooverlap .* dist.dclrexok

  dist[(end-10_000):end,:] |> CSV.write("output\\distest-$(REPLICATION_TYPE[]).csv")
end


#!/usr/bin/env julia
using CSV, DataFrames, PGFPlotsX
push!(PGFPlotsX.CUSTOM_PREAMBLE, raw"\usepackage{fontspec}")
push!(PGFPlotsX.CUSTOM_PREAMBLE, raw"\setmainfont{Asana-Math}")
csvform = tempname()
run(pipeline(`./yyg_se_0804.sh`, stdout=csvform))
yyg_se_0804 = CSV.File(csvform) |> DataFrame
pds = [
    Dict(:c => :predicted_deaths_mean, :lab => "Medel",),
    Dict(:c => :predicted_deaths_lower, :lab => "2,5 percentil",),
    Dict(:c => :predicted_deaths_upper, :lab => "97,5 percentil",)]
p = @pgf Axis({title = "Döda Sverige 2020-08-04 YYG",
    date_coordinates_in = "x",
    x_tick_label_style = "{rotate=90}",
    y_tick_label_style = {"/pgf/number format/use comma"},
    xlabel = "Prediktionsdatum", ylabel = "Kumulativt antal döda",
    xmajorgrids, ymajorgrids})
for pd in pds
    @pgf push!(p, PlotInc({no_markers}, Table([yyg_se_0804[:date], yyg_se_0804[pd[:c]]])),
        LegendEntry(pd[:lab]))
end
pgfsave("../../images/yyg_se_0804.svg", p)


# This file was generated, do not modify it. # hide
fit!(ŷ, rows=train)
rms(y[test], ŷ(rows=test))


module Time

using JavaCall: JString

using ..OrekitWrapper.TimeWrapper

import Dates

export AbsoluteDate, get_tai

TimeWrapper.get_tai() = get_tai(TimeScalesFactory)

Base.show(io::IO, tai::TAIScale) = print(io, to_string(tai))

function TimeWrapper.AbsoluteDate(year::Int, month::Int, day::Int, hour::Int, minute::Int, second::Float64, scale)
    return AbsoluteDate(Int32(year), Int32(month), Int32(day),
        Int32(hour), Int32(minute), second, convert(TimeScale, scale))
end

TimeWrapper.AbsoluteDate(str::String, scale) = AbsoluteDate(JString(str), convert(TimeScale, scale))

function TimeWrapper.AbsoluteDate(dt::Dates.DateTime, scale)
    y = Dates.year(dt)
    m = Dates.month(dt)
    d = Dates.day(dt)
    h = Dates.hour(dt)
    mn = Dates.minute(dt)
    s = Dates.millisecond(dt) * 1e-3 + Dates.second(dt)
    return AbsoluteDate(y, m, d, h, mn, s, scale)
end

function Dates.DateTime(dt::AbsoluteDate, scale)
    comp = get_components(dt, convert(TimeScale, scale))
    date = Dates.Date(get_date(comp))
    time = Dates.Time(get_time(comp))
    return Dates.DateTime(date, time)
end

function Dates.Date(date::DateComponents)
    y = get_year(date)
    m = get_month(date)
    d = get_day(date)
    return Dates.Date(y, m, d)
end

function Dates.Time(time::TimeComponents)
    h = get_hour(time)
    m = get_minute(time)
    s = get_second(time)
    sec = round(Int, s)
    ms = (s - sec) * 1e3
    milli = round(Int, ms)
    μs = (ms - milli) * 1e3
    micro = round(Int, μs)
    ns = (μs - micro) * 1e3
    nano = round(Int, ns)
    return Dates.Time(h, m, sec, milli, micro, nano)
end

end


# This file was generated, do not modify it. # hide
using UrlDownload
using DataFrames
using PrettyPrinting
using PyPlot
using MLJ
ioff() # hide
MLJ.color_off(); # hide


struct BlockIndices
    i1::UnitRange{Int}
    i2::UnitRange{Int}
    isdiag::Bool
end

# Remove isdiag field from comparison for dict
Base.isequal(b1::BlockIndices, b2::BlockIndices) = b1.i1 == b2.i1 && b1.i2 == b2.i2
Base.hash(b::BlockIndices, h::UInt) = hash((b.i1, b.i2), h)

"""
    SparseBlocks

A type for speeding up expressions like the following:

    A[rows, cols] .= data
    A[rows, cols] .+= data

where `A` is a `SparseMatrixCSC` and `rows` and `cols` are `UnitRange`s. The
`data` is usually a `Matrix` but will work with any `AbstractArray` or
collection that supports broadcasting.

The key assumption is that the sparsity in `A` is known and fixed, specified by
a list of "blocks," or a chunk of consecutive rows and columns, addressable by
two `UnitRange` objects.

Usage consists of two phases: 1) initialization, where the sparsity structure is
specified, and 2) usage, where the precomputed data in this struct is used to
speed up indexing operations.

## Initialization
To specify the sparsity structure, use [`addblock!`](@ref). For example, we could
do the following:

    sb = SparseBlocks(10,10)  # blocks for a 10x10 sparse matrix
    addblock!(sb, 1:2, 1:2)   # nonzeros block in top left corner
    addblock!(sb, 1:3, 8:10)  # nonzeros block in top right corner

After adding all of our blocks, we initialize the sparsity structure via
[`initialize!`](@ref):

    A = spzeros(10,10)
    initialize!(sb, A)

After calling `initialize!`, the matrix `A` will have the correct number of nonzero
entries. The speed of the indexing operations is given by using cached conversions
from the entries of each block to the storage location in the nonzeros vector of `A`,
so it's critical that the sparsity structure of `A` does not change!

## Usage
Once initialized, we can use the cached information in this object to speed up
writing to our sparse matrix `A`. Basic usage is as simple as wrapping our indices
with an indexing call to our `SparseBlocks` object:

    A[sb[1:2, 1:2]] .= randn(2,2)

Note that these indices must exactly match one of the argument pairs passed to
[`addblock`](@ref). For best performance, you can cache the [`SparseBlockIndex`](@ref)
returned from this object:

    inds = sb[1:2, 1:2]  # cache this somewhere for future indexing calls
    A[inds] .= randn(2,2)

All of these methods avoid dynamic memory allocations and are about 5-20x faster
indexing directly into the sparse matrix.
"""
struct SparseBlocks
    n::Int
    m::Int
    inds::Dict{BlockIndices,Matrix{Int}}
    SparseBlocks(n,m) = new(n,m,Dict{BlockIndices,Matrix{Int}}())
end

"""
    addblock!(blocks, rows, cols)

Add a nonzero block for the sparse matrix, spanning the block defined by
   consecutive `rows` and `cols` ranges. All calls to `addblock!` should be
   made before calling [`initialize!`](@ref).
"""
function addblock!(blocks::SparseBlocks, i1::UnitRange, i2::UnitRange, isdiag::Bool=false)
    if i1.start < 1 || i1.stop > blocks.n || i2.start < 1 || i2.stop > blocks.m
        throw(BoundsError(spzeros(blocks.n, blocks.m), (i1, i2)))
    end
    block = BlockIndices(i1, i2, isdiag)
    if isdiag
        blocks.inds[block] = zeros(Int, min(length(i1), length(i2)), 1)
    else
        blocks.inds[block] = zeros(Int, length(i1), length(i2))
    end
    blocks
end

"""
    initialize!(blocks::SparseBlocks, A::SparseMatrixCSC)

Initialize the sparsity structure specified by `blocks` (via previous calls
to [`addblock!`](@ref)). Writes epsilon values to the nonzero blocks in `A`
to initialize the data storage for the sparse matrix. The indices into this
storage is computed and stored in `blocks`.
"""
function initialize!(blocks::SparseBlocks, A::SparseMatrixCSC{T}) where T
    n,m = blocks.n, blocks.m
    if size(A) != (n,m)
        throw(DimensionMismatch("Dimension of sparse matrix doesn't match expected size of ($n,$m)."))
    end
    @assert size(A) == (blocks.n, blocks.m)
    for block in keys(blocks.inds)
        if block.isdiag
            for (i,j) in zip(block.i1, block.i2)
                A[i,j] = eps(T)
            end
        else
            A[block.i1, block.i2] .= eps(T)
        end
    end
    inds = blocks.inds
    for (block,inds) in pairs(blocks.inds)
        rows = block.i1
        cols = block.i2
        n = length(rows)
        if block.isdiag
            for (j,(r,c)) in enumerate(zip(rows, cols))
                nzind = getnzind(A, r, c)
                if !isempty(nzind)
                    inds[j] = nzind[1]
                end
            end
        else
            for (j,col) in enumerate(cols)
                nzind = getnzind(A, rows[1], col)
                if !isempty(nzind)
                    inds[:,j] .= nzind[1] - 1 .+ (1:n)
                end
            end
        end
    end
    blocks
end

# Get the index into the nonzero vector for the given row and column
function getnzind(A,row,col)
    rows = view(A.rowval, nzrange(A,col))
    istart = searchsorted(rows, row) .+ A.colptr[col] .- 1
    return istart
end

"""
    SparseBlockIndex

A custom index for sparse arrays with a fixed sparsity structure, especially for
    those whose nonzeros appear as dense blocks. This is created by indexing
    into a [`SparseBlocks`] object, using indices to one of the cached blocks.

This object can then be used as a index for those same indices, creating a
    [`SparseBlockView`](@ref) object that efficiently maps cartesian indices
    into the nonzeros vector of the sparse matrix.
"""
struct SparseBlockIndex
    block::BlockIndices
    nzinds::Matrix{Int}
end

function Base.getindex(blocks::SparseBlocks, i1::UnitRange, i2::UnitRange)
    block0 = BlockIndices(i1, i2, false)
    block = getkey(blocks.inds, block0, nothing)
    if isnothing(block)
        throw(KeyError((i1, i2)))
    end
    SparseBlockIndex(block, blocks.inds[block])
end

"""
    SparseBlockView

A custom view into a sparse matrix with known sparsity, and whose nonzero entries
appears as dense blocks. Can be used a normal array, where operations on the elements
write and read directly from the nonzeros vector of the sparse matrix.
"""
struct SparseBlockView{Tv,Ti} <: AbstractMatrix{Tv}
    data::SparseMatrixCSC{Tv,Ti}
    block::SparseBlockIndex
end

# Needed to support expressions like A[B] .= data
Broadcast.dotview(A::AbstractSparseMatrix, block::SparseBlockIndex) = SparseBlockView(A, block)

# Define this type when indexing into a SparseMatrixCSC with a SparseBlockIndex
Base.getindex(A::AbstractSparseMatrix, block::SparseBlockIndex) = SparseBlockView(A, block)

# Array interface
Base.size(B::SparseBlockView) = size(B.block.nzinds)
Base.getindex(B::SparseBlockView, index) = getindex(B.data.nzval, B.block.nzinds[index])
Base.getindex(B::SparseBlockView, I::Vararg{Int,2}) = getindex(B.data.nzval, B.block.nzinds[I...])
Base.setindex!(B::SparseBlockView, v, index) = setindex!(B.data.nzval, v, B.block.nzinds[index])
Base.setindex!(B::SparseBlockView, v, I::Vararg{Int,2}) = setindex!(B.data.nzval, v, B.block.nzinds[I...])
Base.IndexStyle(::SparseBlockView) = IndexLinear()

# Broadcasting interface
struct SparseViewStyle <: Broadcast.BroadcastStyle end
Base.BroadcastStyle(::Type{<:SparseBlockView}) = SparseViewStyle()

# Always use the new broadcasting style over others
Base.BroadcastStyle(::SparseViewStyle, ::Broadcast.BroadcastStyle) = SparseViewStyle()

# Handle generic expressions like B .= f.(arg)
function Broadcast.materialize!(B::SparseBlockView, bc::Broadcast.Broadcasted{Broadcast.DefaultMatrixStyle})
    for i in eachindex(bc)
        B[i] = bc[i]
    end
    B
end

# Handle generic expressions like B .= f.(B, g.(args...))
function Broadcast.materialize!(B::SparseBlockView, bc::Broadcast.Broadcasted{SparseViewStyle})
    for i in eachindex(bc)
        B[i] = bc.f(B[i], bc[i])
    end
    B
end

# Specialization of B .= data::AbstractMatrix
function Broadcast.materialize!(B::SparseBlockView, bc::Broadcast.Broadcasted{Broadcast.DefaultMatrixStyle,Nothing,typeof(identity),<:Tuple{<:AbstractMatrix}})
    data = bc.args[1]
    for i in eachindex(B)
        B[i] = data[i]
    end
    B
end

# Specialization of B .= f.(B, data::AbstractMatrix)
#   which includes B .+= data::AbstractMatrix and variants
function Broadcast.materialize!(B::SparseBlockView, bc::Broadcast.Broadcasted{SparseViewStyle,Nothing,<:Any,<:Tuple{<:SparseBlockView,<:AbstractArray}})
    data = bc.args[2]
    for i in eachindex(B)
        B[i] = bc.f(B[i], data[i])
    end
    B
end

# Specialization of B .= D::Diagonal
function Broadcast.materialize!(B::SparseBlockView, bc::Broadcast.Broadcasted{<:LinearAlgebra.StructuredMatrixStyle{<:Diagonal}, Nothing, typeof(identity), <:Tuple{<:Diagonal}})
    D = bc.args[1]
    if B.block.block.isdiag
        for i in eachindex(B)
            B[i] = D.diag[i]
        end
    else
        for i = 1:minimum(size(B))
            B[i,i] = D.diag[i]
        end
    end
end

function Broadcast.materialize!(B::SparseBlockView, bc::Broadcast.Broadcasted{SparseViewStyle,Nothing,<:Any,<:Tuple{<:SparseBlockView,<:Diagonal}})
    D = bc.args[2]
    f = bc.f
    if B.block.block.isdiag
        for i in eachindex(B)
            B[i] = f(B[i], D.diag[i])
        end
    else
        for i = 1:minimum(size(B))
            B[i,i] = f(B[i], D.diag[i])
        end
    end
end

# Specialize B' .= bc
function Broadcast.materialize!(B::Adjoint{T, <:SparseBlockView{T}}, bc::Broadcast.Broadcasted{Broadcast.DefaultMatrixStyle,Nothing,typeof(identity),<:Tuple{<:AbstractMatrix}}) where T
    data = bc.args[1]
    for i in eachindex(bc)
        it = CartesianIndex(i[2], i[1])
        B.parent[it] = data[i]
    end
    B
end

# Specialize B = UpperTriangular(data)
function Broadcast.materialize!(B::SparseBlockView, bc::Broadcast.Broadcasted{<:LinearAlgebra.StructuredMatrixStyle{<:UpperTriangular}})
    for r = 1:size(B,1)
        for c = r:size(B,2)
            i = CartesianIndex(r,c)
            B[r,c] = bc[i]
        end
    end
end

# Specialize B = LowerTriangular(data)
function Broadcast.materialize!(B::SparseBlockView, bc::Broadcast.Broadcasted{<:LinearAlgebra.StructuredMatrixStyle{<:LowerTriangular}})
    for c = 1:size(B,2)
        for r = c:size(B,1)
            i = CartesianIndex(r,c)
            B[r,c] = bc[i]
        end
    end
end


using InteractiveUtils: code_llvm
using Test

include("../benchmark/map.jl")

"""
    llvm_ir(f, args) :: String

Get LLVM IR of `f(args...)` as a string.
"""
llvm_ir(f, args) = sprint(code_llvm, f, Base.typesof(args...))

nmatches(r, s) = count(_ -> true, eachmatch(r, s))

twice(x) = 2x

@testset "setindex!! (@inbounds)" begin
    @test mapbb(twice, [1:3;], Val(false)) == [2, 4, 6]

    ir_noinbounds = llvm_ir(mapbb, (twice, Float64[], Val(false)))
    if Base.JLOptions().check_bounds == 2  # --check-bounds=no
        @test nmatches(r"fmul <4 x double>", ir_noinbounds) >= 4
    else
        @test nmatches(r"fmul <4 x double>", ir_noinbounds) == 0
    end
    ir_inbounds = llvm_ir(mapbb, (twice, Float64[], Val(true)))
    @test nmatches(r"fmul <4 x double>", ir_inbounds) >= 4
end


module Packing
using GeometryBasics

include("rectangle.jl")
include("guillotine.jl")

export RectanglePacker, GuillotinePacker

end # module


printstyled("  nanfill, ungap\n", color=:light_green)

# Check that nanfill does not error
S = randSeisData()
Ev = SeisEvent(hdr=randSeisHdr(), data=convert(EventTraceData, deepcopy(S)))
C = deepcopy(S[1])

for i = 1:S.n
  L = length(S.x[i])
  inds = rand(1:L, div(L,2))
  S.x[i][inds] .= NaN
end
nanfill!(S)
for i = 1:S.n
  x = getindex(getfield(S, :x), i)
  @test isempty(findall(isnan.(x)))
end

for i = 1:Ev.data.n
  L = length(Ev.data.x[i])
  inds = rand(1:L, div(L,2))
  Ev.data.x[i][inds] .= NaN
end
nanfill!(Ev.data)
for i = 1:Ev.data.n
  x = getindex(getfield(getfield(Ev, :data), :x), i)
  @test isempty(findall(isnan.(x)))
end

L = length(C.x)
inds = rand(1:L, div(L,2))
C.x[inds] .= NaN
nanfill!(C)
@test isempty(findall(isnan.(C.x)))

# Test that traces of all NaNs becomes traces of all zeros
C = SeisChannel()
C.x = fill!(zeros(Float32, 128), NaN32)
nanfill!(C)
@test C.x == zeros(Float32, 128)

S = randSeisData()
U = deepcopy(S)
for i = 1:S.n
  x = getindex(getfield(S, :x), i)
  nx = lastindex(x)
  T = eltype(x)
  fill!(x, T(NaN))

  u = getindex(getfield(U, :x), i)
  fill!(u, zero(T))
end
nanfill!(S)
for i = 1:S.n
  x = getindex(getfield(S, :x), i)
  nx = lastindex(x)
  T = eltype(x)
  u = getindex(getfield(U, :x), i)
  @test T == eltype(u)
  @test S.x[i] == x == U.x[i] == u
end

# Test that ungap calls nanfill properly
Ev2 = ungap(Ev.data, tap=true)
ungap!(Ev.data, tap=true)
@test Ev.data == Ev2
ungap!(C, tap=true)
ungap!(S, tap=true)

# Ensure one segment is short enough to invoke bad behavior in ungap
Ev = randSeisEvent()
Ev.data.fs[1] = 100.0
Ev.data.x[1] = rand(1024)
Ev.data.t[1] = vcat(Ev.data.t[1][1:1,:], [5 2*ceil(S.fs[1])*sμ], [8 2*ceil(S.fs[1])*sμ], [1024 0])

redirect_stdout(out) do
  ungap!(Ev.data)
end


"""
`row` and `col` must be in [1,size(mat)], i.e indexing starts at 1 (Julia).

# Implementation
For some unkwnown reason, calling `MatSetValue` fails.
"""
function Base.setindex!(mat::PetscMat, value::Number, row::Integer, col::Integer)
    MatSetValues(mat, PetscInt[row - 1], PetscInt[col - 1], PetscScalar[value], INSERT_VALUES)
end

# This is stupid but I don't know how to do better yet
Base.setindex!(mat::PetscMat, values, row::Integer, cols) = MatSetValues(mat, [row-1], collect(cols) .- 1, values, INSERT_VALUES)
Base.setindex!(mat::PetscMat, values, rows, col::Integer) = MatSetValues(mat, collect(rows) .- 1, [col-1], values, INSERT_VALUES)

Base.ndims(::Type{PetscMat}) = 2

"""
    create_matrix(nrows, ncols, nrows_loc = PETSC_DECIDE, ncols_loc = PETSC_DECIDE; auto_setup = false)

Create a `PetscMat` matrix of global size `(nrows, ncols)`.

Use `auto_setup = true` to immediatly call `set_from_options!` and `set_up!`.
"""
function create_matrix(nrows, ncols, nrows_loc = PETSC_DECIDE, ncols_loc = PETSC_DECIDE; auto_setup = false, comm::MPI.Comm = MPI.COMM_WORLD)
    mat = MatCreate()
    MatSetSizes(mat::PetscMat, nrows_loc, ncols_loc, nrows, ncols)

    if (auto_setup)
        set_from_options!(mat)
        set_up!(mat)
    end
    return mat
end

"""
Wrapper to `MatCreateComposite` using the "alternative construction" from the PETSc documentation.
"""
function create_composite_add(matrices)
    N, M = MatGetSize(matrices[1])
    n, m = MatGetLocalSize(matrices[1])
    mat = create_matrix(N, M, n, m; auto_setup = false, comm = matrices[1].comm)
    MatSetType(mat, "composite")
    for m in matrices
        MatCompositeAddMat(mat, m)
    end
    assemble!(mat)
    return mat
end

set_global_size!(mat::PetscMat, nrows, ncols) = MatSetSizes(mat, PETSC_DECIDE, PETSC_DECIDE, nrows, ncols)
set_local_size!(mat::PetscMat, nrows, ncols) = MatSetSizes(mat, nrows, ncols, PETSC_DECIDE, PETSC_DECIDE)

set_from_options!(mat::PetscMat) = MatSetFromOptions(mat)

set_up!(mat::PetscMat) = MatSetUp(mat)

"""
    get_range(mat::PetscMat)

Wrapper to `MatGetOwnershipRange`

However, the result `(rstart, rend)` is such that `mat[rstart:rend]` are the rows handled by the local processor.
This is different from the default `PETSc.MatGetOwnershipRange` result where the indexing starts at zero and where
`rend-1` is last row handled by the local processor.
"""
function get_range(mat::PetscMat)
    rstart, rend = MatGetOwnershipRange(mat)
    return (rstart + 1, rend)
end

"""
    get_urange(mat::PetscMat)

Provide a `UnitRange` from the method `get_range`.
"""
function get_urange(mat::PetscMat)
    rstart, rend = MatGetOwnershipRange(mat)
    return rstart+1:rend
end

"""
    Wrapper to `MatAssemblyBegin` and `MatAssemblyEnd` successively.
"""
function assemble!(mat::PetscMat, type::MatAssemblyType = MAT_FINAL_ASSEMBLY)
    MatAssemblyBegin(mat, type)
    MatAssemblyEnd(mat, type)
end

"""
    set_values!(mat::PetscMat, I, J, V, mode = ADD_VALUES)

Set values of `mat` in `SparseArrays` fashion : using COO format:
`mat[I[k], J[k]] = V[k]`.
"""
function set_values!(mat::PetscMat, I::Vector{PetscInt}, J::Vector{PetscInt}, V::Vector{PetscScalar}, mode = ADD_VALUES)
    for (i,j,v) in zip(I, J, V)
        MatSetValue(mat, i - 1, j - 1, v, mode)
    end
end

set_values!(mat, I, J, V, mode = ADD_VALUES) = set_values!(mat, PetscInt.(I), PetscInt.(J), PetscScalar.(V), mode)

function preallocate_MPIAIJ(mat::PetscMat; dnz::PetscInt = 0, dnnz::Vector{PetscInt} = C_NULL, onz::PetscInt = 0, onnz::Vector{PetscInt} = C_NULL)
    MatMPIAIJSetPreallocation(mat, dnz, dnnz, onz, onnz)
end

"""
    mat2file(mat::PetscMat, filename::String, format::PetscViewerFormat = PETSC_VIEWER_ASCII_CSV, type::String = "ascii")

Write a PetscMat to a file.
"""
function mat2file(mat::PetscMat, filename::String, format::PetscViewerFormat = PETSC_VIEWER_ASCII_CSV, type::String = "ascii")
    viewer = PetscViewer(mat.comm, filename, format, type)
    MatView(mat, viewer)
    destroy!(viewer)
end

Base.show(::IO, mat::PetscMat) = MatView(mat)

destroy!(mat::PetscMat) = MatDestroy(mat)


module CryptoMiniSat

import Base.convert
import Base.show
import Base.~

global const cmsat_lib = Sys.isunix() ? "libcryptominisat5" : "cryptominisat5"

export SATSolver
export Lit, LBool
export Slice_Lit, Slice_lbool
export cmsat_new, cmsat_free, nvars, add_clause, add_xor_clause, new_vars, solve,
       get_model, get_conflict, print_stats, set_num_threads, set_verbosity

# structs taken from:
# https://github.com/msoos/cryptominisat/blob/master/src/cryptominisat_c.h.in
struct C_Lit
    x::UInt32
end

struct C_lbool
    x::UInt8
end

struct Slice_Lit
    vals::Ptr{C_Lit}
    num_vals::Csize_t
end

struct Slice_lbool
    vals::Ptr{C_lbool}
    num_vals::Csize_t
end

const l_true = C_lbool(0)
const l_false = C_lbool(1)
const l_undef = C_lbool(2)

const Csatsolver = Ptr{Nothing}
const LBool = Union{Bool, Nothing}

function convert(::Type{LBool}, b::C_lbool)::LBool
    if b == l_undef
        return nothing
    else
        return b.x == 0
    end
end

# adapted from:
# https://github.com/msoos/cryptominisat/blob/master/rust/src/lib.rs
#=
function C_Lit(var::T, negated::Bool) where {T <: Unsigned}
    @assert var < (1 << 31)

    return C_Lit(var << 1 | negated)
end
=#

struct Lit
    var::Int
    negated::Bool
end

Lit(var::Int) = Lit(var, false)
~(l::Lit) = Lit(l.var, ~l.negated)

function Base.show(io::IO, l::Lit)
    if l.negated
        print(io, "¬")
    end
    show(io, l.var)
end

function convert(::Type{C_Lit}, l::Lit)::C_Lit
    @assert l.var < (1 << 31)
    @assert l.var >= 0

    return C_Lit(l.var << 1 | l.negated)
end

function convert(::Type{Lit}, l::C_Lit)::Lit
    return Lit(l.x >> 1, l.x & 0x1)
end

mutable struct SATSolver
    hnd::Csatsolver
    function SATSolver(hnd::Csatsolver)
        a = new(hnd)
        finalizer(cmsat_free, a)
    end
end

SATSolver() = SATSolver(cmsat_new())

function cmsat_free(sat_solver::SATSolver)
    ccall((:cmsat_free, cmsat_lib),
        Nothing,
        (Csatsolver,),
        sat_solver.hnd)
end

function cmsat_new()::Csatsolver
    solver = ccall((:cmsat_new, cmsat_lib), Csatsolver, ())
    return solver
end

function nvars(sat_solver::SATSolver)
    return ccall((:cmsat_nvars, cmsat_lib),
                 Cuint,
                 (Csatsolver,),
                 sat_solver.hnd)
end

function add_clause(sat_solver::SATSolver, clause::Vector{Lit})::Bool
    return ccall((:cmsat_add_clause, cmsat_lib),
                 Bool,
                 (Csatsolver, Ref{C_Lit}, Csize_t),
                 sat_solver.hnd, convert(Vector{C_Lit}, clause), length(clause))
end

function add_xor_clause(sat_solver::SATSolver, vars::Vector{T}, rhs::Bool)::Bool where {T <: Unsigned}
    return ccall((:cmsat_add_xor_clause, cmsat_lib),
          Bool,
          (Csatsolver, Ref{Cuint}, Csize_t, Bool),
          sat_solver.hnd, convert(Vector{Cuint}, vars), length(vars), rhs)
end

function new_vars(sat_solver::SATSolver, n_vars::T) where {T <: Integer}
    ccall((:cmsat_new_vars, cmsat_lib),
           Nothing,
           (Csatsolver, Cuint),
           sat_solver.hnd, n_vars)
end

function solve(sat_solver::SATSolver)::Tuple{LBool,Vector{LBool}}
    sat = ccall((:cmsat_solve, cmsat_lib),
                 C_lbool,
                 (Csatsolver,),
                 sat_solver.hnd)
    return sat, get_model(sat_solver)
end

function solve(sat_solver::SATSolver, assumptions::Vector{Lit})::Tuple{LBool,Vector{LBool}}
    sat = ccall((:cmsat_solve_with_assumptions, cmsat_lib),
                 C_lbool,
                 (Csatsolver, Ref{C_Lit}, Csize_t),
                 sat_solver.hnd, convert(Array{C_Lit}, assumptions), length(assumptions))
    return sat, get_model(sat_solver)
end

function get_model(sat_solver::SATSolver)::Vector{LBool}
    c_mdl = ccall((:cmsat_get_model, cmsat_lib),
                  Slice_lbool,
                  (Csatsolver,),
                  sat_solver.hnd)
    return unsafe_wrap(Array{C_lbool}, c_mdl.vals, c_mdl.num_vals)
end

function get_conflict(sat_solver::SATSolver)::Vector{Lit}
    c_conflict = ccall((:cmsat_get_conflict, cmsat_lib),
                       Slice_Lit,
                       (Csatsolver,),
                       sat_solver.hnd)
    clits = unsafe_wrap(Array{C_Lit}, c_conflict.vals, c_conflict.num_vals)
    return clits
end

function print_stats(sat_solver::SATSolver)
    ccall((:cmsat_print_stats, cmsat_lib),
          Nothing,
          (Csatsolver,),
          sat_solver.hnd)
end

function set_num_threads(sat_solver::SATSolver, n_threads::T) where {T <: Integer}
    ccall((:cmsat_set_num_threads, cmsat_lib),
          Nothing,
          (Csatsolver, Cuint),
          sat_solver.hnd, n_threads)
end

function set_verbosity(sat_solver::SATSolver, n::T) where {T <: Integer}
    ccall((:cmsat_set_verbosity, cmsat_lib),
          Nothing,
          (Csatsolver, Cuint),
          sat_solver.hnd, n)
end

end


__precompile__(false)
module UniversalTensorBoard

using TensorFlow

Summaries = TensorFlow.summary

export set_tb_logdir, @tb_log, reset_tb_logs
export tensorboard

mutable struct TensorBoardSession
    logdir::String
    sess::Session
    writer::Summaries.FileWriter
    global_step::Int
end

function TensorBoardSession(logdir; overwrite=false)
    if overwrite
        rm(logdir; force=true, recursive=true)
    end
    mkpath(logdir)

    sess= Session(Graph())
    writer = Summaries.FileWriter(logdir; graph=sess.graph)
    TensorBoardSession(realpath(logdir), sess, writer, 0)
end

TensorBoardSession() = TensorBoardSession("tensorboard_logs", overwrite=true)
# normally the logs don't overwrite, but if you've not given a path, you clearly don't care.

const default_logging_session = Ref(TensorBoardSession())

####################################################################################################
# Interface

"""
    tensorboard()

Opens Tensorboard in the systems default browser.
"""
tensorboard(tb_sess=default_logging_session[]) = visualize(tb_sess.writer)

"""
    set_tb_logdir(logdir, overwrite=false)

Start a new log in the given directory
"""
function set_tb_logdir(logdir, overwrite=false)
    default_logging_session[] = TensorBoardSession(logdir, overwrite=overwrite)
end

"""
    reset_tb_logs()

Reset the current log, deleteing all information
"""
function reset_tb_logs()
    logdir = default_logging_session[].logdir
    default_logging_session[] = TensorBoardSession(logdir, overwrite=true)
end

macro tb_log(name)
    :(_tb_log($(esc(string(name))), $(esc(name))))
end

#################################################################
# Implmentation
function _tb_log(name, value; tb_sess = default_logging_session[])
    placeholder_node, summary_node  = get_nodes!(tb_sess, name, value)
    summaries = run(tb_sess.sess, summary_node,  Dict(placeholder_node=>value))
    write(tb_sess.writer, summaries, tb_sess.global_step)
    tb_sess.global_step+=1
    value
end

name_fix(name)=replace(name, r"[^A-Za-z0-9.]|\s", "_")
placeholder_node_name(name) = name_fix(name *"_input")
summary_node_name(name) = name_fix(name)

has_node(graph, name) = !isnull(TensorFlow.get_node_by_name(graph, name))

function get_nodes!(tb_sess, name, value)

    gg = tb_sess.sess.graph
    if has_node(gg, placeholder_node_name(name))
        (gg[placeholder_node_name(name)], gg[summary_node_name(name)])
    else
        TensorFlow.as_default(gg) do
            create_nodes!(name, value)
        end
    end
end

function create_nodes!(name, value::T) where T<:Number
    placeholder_node = placeholder(T; name=placeholder_node_name(name), shape=[])
    summary_node = Summaries.scalar(name, placeholder_node; name=summary_node_name(name))
    (placeholder_node, summary_node)
end

function create_nodes!(name, value::AbstractArray{T}) where T<:Number
    shape = collect(size(value))
    placeholder_node=placeholder(T; name=placeholder_node_name(name), shape=shape)
    summary_node = Summaries.histogram(name, placeholder_node; name=summary_node_name(name))
    (placeholder_node, summary_node)
end

end # modulee