text:
using NNlib: conv, ∇conv_data, depthwiseconv, output_size

# pad dims of x with dims of y until ndims(x) == ndims(y)
_paddims(x::Tuple, y::Tuple) = (x..., y[(end - (length(y) - length(x) - 1)):end]...)

_convtransoutdims(isize, ksize, ssize, dsize, pad) = (isize .- 1).*ssize .+ 1 .+ (ksize .- 1).*dsize .- (pad[1:2:end] .+ pad[2:2:end])

expand(N, i::Tuple) = i
expand(N, i::Integer) = ntuple(_ -> i, N)

"""
    SamePad

Padding for convolutional layers will be calculated so that outputshape == inputshape when stride = 1.

For stride > 1 the output shape depends on the type of convolution layer.
"""
struct SamePad end

calc_padding(pad, k::NTuple{N,T}, dilation, stride) where {T,N}= expand(Val(2*N), pad)
function calc_padding(::SamePad, k::NTuple{N,T}, dilation, stride) where {N,T}
  #Ref: "A guide to convolution arithmetic for deep learning" https://arxiv.org/pdf/1603.07285

  # Effective kernel size, including dilation
  k_eff = @. k + (k - 1) * (dilation - 1)
  # How much total padding needs to be applied?
  pad_amt = @. k_eff - 1
  # In case amount of padding is odd we need to apply different amounts to each side.
  return Tuple(mapfoldl(i -> [ceil(Int, i/2), floor(Int, i/2)], vcat, pad_amt))
end

"""
    Conv(filter, in => out, σ = identity; init = glorot_uniform,
         stride = 1, pad = 0, dilation = 1)

    filter = (2,2)
    in = 1
    out = 16
    Conv((2, 2), 1=>16, relu)

Standard convolutional layer. `filter` should be a tuple like `(2, 2)`.
`in` and `out` specify the number of input and output channels respectively.

Data should be stored in WHCN order (width, height, # channels, batch size).
In other words, a 100×100 RGB image would be a `100×100×3×1` array,
and a batch of 50 would be a `100×100×3×50` array.

Accepts keyword arguments `weight` and `bias` to set the corresponding fields.
Setting `bias` to `Flux.Zeros()` will switch bias off for the layer.

Takes the keyword arguments `pad`, `stride` and `dilation`.
For input dimension N,
`pad` should be a single Integer indicating equal padding value for each spatial dimension,
a tuple of length (N-2) to apply symmetric padding or a tuple of length 2*(N-2)
indicating padding values for each spatial dimension at both the ends.
`stride` and `dilation` should be a single Integer or a tuple with N-2 parameters.
Use `pad=SamePad()` to apply padding so that outputsize == inputsize / stride.

# Examples

Apply a `Conv` layer to a 1-channel input using a 2×2 window filter size, giving us a
16-channel output. Output is activated with ReLU.
```julia
filter = (2,2)
in = 1
out = 16
Conv(filter, in => out, relu)
```
"""
struct Conv{N,M,F,A,V}
  σ::F
  weight::A
  bias::V
  stride::NTuple{N,Int}
  pad::NTuple{M,Int}
  dilation::NTuple{N,Int}
end

"""
    Conv(weight::AbstractArray, bias::AbstractArray)
    Conv(weight::AbstractArray, bias::AbstractArray, activation)

Constructs the convolutional layer with user defined weight and bias arrays.

Setting `bias` to `Flux.Zeros()` would switch `bias` off for the layer.

Takes the keyword arguments `pad`, `stride` and `dilation`.
For input dimension N,
`pad` should be a single Integer indicating equal padding value for each spatial dimension,
a tuple of length (N-2) to apply symmetric padding or a tuple of length 2*(N-2)
indicating padding values for each spatial dimension at both the ends.
`stride` and `dilation` should be a single Integer or a tuple with N-2 parameters.

There is also a keyword-only constuctor available for all convoultional
layers.

```julia
weight = rand(Float32, 3, 3, 5)
bias = zeros(Float32, 5)
Conv(weight = weight,
    bias = bias,
    σ = sigmoid)
```
"""
function Conv(w::AbstractArray{T,N}, b::Union{Zeros, AbstractVector{T}}, σ = identity;
              stride = 1, pad = 0, dilation = 1) where {T,N}
  stride = expand(Val(N-2), stride)
  dilation = expand(Val(N-2), dilation)
  pad = calc_padding(pad, size(w)[1:N-2], dilation, stride)
  return Conv(σ, w, b, stride, pad, dilation)
end

function Conv(;weight::AbstractArray{T,N}, bias::Union{Zeros, AbstractVector{T}},
              activation = identity, stride = 1, pad = 0, dilation = 1) where {T,N}
  Conv(weight, bias, activation, stride = stride, pad = pad, dilation = dilation)
end

"""
    convfilter(filter::Tuple, in=>out)

Constructs a standard convolutional weight matrix with given `filter` and
channels from `in` to `out`.

Accepts the keyword `init` (default: `glorot_uniform`) to control the sampling
distribution.

See also: [`depthwiseconvfilter`](@ref)
"""
convfilter(filter::NTuple{N,Integer}, ch::Pair{<:Integer,<:Integer};
          init = glorot_uniform) where N = init(filter..., ch...)

function Conv(k::NTuple{N,Integer}, ch::Pair{<:Integer,<:Integer}, σ = identity;
            init = glorot_uniform,  stride = 1, pad = 0, dilation = 1,
            weight = convfilter(k, ch, init = init), bias = zeros(ch[2])) where N

  Conv(weight, bias, σ,
      stride = stride, pad = pad, dilation = dilation)
end

@functor Conv

function (c::Conv)(x::AbstractArray)
  # TODO: breaks gpu broadcast :(
  # ndims(x) == ndims(c.weight)-1 && return squeezebatch(c(reshape(x, size(x)..., 1)))
  σ, b = c.σ, reshape(c.bias, ntuple(_->1, length(c.stride))..., :, 1)
  cdims = DenseConvDims(x, c.weight; stride=c.stride, padding=c.pad, dilation=c.dilation)
  σ.(conv(x, c.weight, cdims) .+ b)
end

function Base.show(io::IO, l::Conv)
  print(io, "Conv(", size(l.weight)[1:ndims(l.weight)-2])
  print(io, ", ", size(l.weight, ndims(l.weight)-1), "=>", size(l.weight, ndims(l.weight)))
  l.σ == identity || print(io, ", ", l.σ)
  print(io, ")")
end

(a::Conv{<:Any,<:Any,W})(x::AbstractArray{T}) where {T <: Union{Float32,Float64}, W <: AbstractArray{T}} =
  invoke(a, Tuple{AbstractArray}, x)

(a::Conv{<:Any,<:Any,W})(x::AbstractArray{<:Real}) where {T <: Union{Float32,Float64}, W <: AbstractArray{T}} =
  a(T.(x))

"""
    outdims(l::Conv, isize::Tuple)

Calculate the output dimensions given the input dimensions `isize`.
Batch size and channel size are ignored as per [NNlib.jl](https://github.com/FluxML/NNlib.jl).

```julia
m = Conv((3, 3), 3 => 16)
outdims(m, (10, 10)) == (8, 8)
outdims(m, (10, 10, 1, 3)) == (8, 8)
```
"""
outdims(l::Conv, isize) =
  output_size(DenseConvDims(_paddims(isize, size(l.weight)), size(l.weight); stride = l.stride, padding = l.pad, dilation = l.dilation))

"""
    ConvTranspose(filter, in=>out)
    ConvTranspose(filter, in=>out, activation)
    ConvTranspose(filter, in => out, σ = identity; init = glorot_uniform,
                  stride = 1, pad = 0, dilation = 1)

Standard convolutional transpose layer. `filter` should be a tuple like `(2, 2)`.
`in` and `out` specify the number of input and output channels respectively.

Data should be stored in WHCN order (width, height, # channels, batch size).
In other words, a 100×100 RGB image would be a `100×100×3×1` array,
and a batch of 50 would be a `100×100×3×50` array.

Accepts keyword arguments `weight` and `bias` to set the corresponding fields.
Setting `bias` to `Flux.Zeros()` will switch bias off for the layer.

Takes the keyword arguments `pad`, `stride` and `dilation`.
For input dimension N,
`pad` should be a single Integer indicating equal padding value for each spatial dimension,
a tuple of length (N-2) to apply symmetric padding or a tuple of length 2*(N-2)
indicating padding values for each spatial dimension at both the ends.
`stride` and `dilation` should be a single Integer or a tuple with N-2 parameters.
Use `pad=SamePad()` to apply padding so that outputsize == stride * inputsize - stride + 1.
"""
struct ConvTranspose{N,M,F,A,V}
  σ::F
  weight::A
  bias::V
  stride::NTuple{N,Int}
  pad::NTuple{M,Int}
  dilation::NTuple{N,Int}
end

"""
    ConvTranspose(weight::AbstractArray, bias::AbstractArray)
    ConvTranspose(weight::AbstractArray, bias::AbstractArray, activation)

Constructs the convolutional transpose layer with user defined weight and bias arrays.
forward pass.

Setting `bias` to `Flux.Zeros()` would switch `bias` off for the layer.

Takes the keyword arguments `pad`, `stride` and `dilation`.
For input dimension N,
`pad` should be a single Integer indicating equal padding value for each spatial dimension,
a tuple of length (N-2) to apply symmetric padding or a tuple of length 2*(N-2)
indicating padding values for each spatial dimension at both the ends.
`stride` and `dilation` should be a single Integer or a tuple with N-2 parameters.

For keyword-only constuctor, see also [`Conv`](@ref)
"""
function ConvTranspose(w::AbstractArray{T,N}, b::Union{Zeros, AbstractVector{T}}, σ = identity;
                      stride = 1, pad = 0, dilation = 1) where {T,N}
  stride = expand(Val(N-2), stride)
  dilation = expand(Val(N-2), dilation)
  pad = calc_padding(pad, size(w)[1:N-2], dilation, stride)
  return ConvTranspose(σ, w, b, stride, pad, dilation)
end

function ConvTranspose(;weight::AbstractArray{T,N}, bias::Union{Zeros, AbstractVector{T}},
                        activation = identity, stride = 1, pad = 0, dilation = 1) where {T,N}
  ConvTranspose(weight, bias, activation, stride = stride, pad = pad, dilation = dilation)
end

function ConvTranspose(k::NTuple{N,Integer}, ch::Pair{<:Integer,<:Integer}, σ = identity;
                      init = glorot_uniform, stride = 1, pad = 0, dilation = 1,
                      weight = convfilter(k, reverse(ch), init = init), bias = zeros(ch[2])) where N

  ConvTranspose(weight, bias, σ,
              stride = stride, pad = pad, dilation = dilation)
end

@functor ConvTranspose

function conv_transpose_dims(c::ConvTranspose, x::AbstractArray)
    # Calculate size of "input", from ∇conv_data()'s perspective...
    combined_pad = (c.pad[1:2:end] .+ c.pad[2:2:end])
    I = (size(x)[1:end-2] .- 1).*c.stride .+ 1 .+ (size(c.weight)[1:end-2] .- 1).*c.dilation .- combined_pad
    C_in = size(c.weight)[end-1]
    batch_size = size(x)[end]
    # Create DenseConvDims() that looks like the corresponding conv()
    return DenseConvDims((I..., C_in, batch_size), size(c.weight);
                        stride=c.stride,
                        padding=c.pad,
                        dilation=c.dilation,
    )
end

# TODO: Find proper fix for https://github.com/FluxML/Flux.jl/issues/900
@nograd conv_transpose_dims

function (c::ConvTranspose)(x::AbstractArray)
  # ndims(x) == ndims(c.weight)-1 && return squeezebatch(c(reshape(x, size(x)..., 1)))
  σ, b = c.σ, reshape(c.bias, map(_->1, c.stride)..., :, 1)
  cdims = conv_transpose_dims(c, x)
  σ.(∇conv_data(x, c.weight, cdims) .+ b)
end

function Base.show(io::IO, l::ConvTranspose)
  print(io, "ConvTranspose(", size(l.weight)[1:ndims(l.weight)-2])
  print(io, ", ", size(l.weight, ndims(l.weight)), "=>", size(l.weight, ndims(l.weight)-1))
  l.σ == identity || print(io, ", ", l.σ)
  print(io, ")")
end

(a::ConvTranspose{<:Any,<:Any,W})(x::AbstractArray{T}) where {T <: Union{Float32,Float64}, W <: AbstractArray{T}} =
  invoke(a, Tuple{AbstractArray}, x)

(a::ConvTranspose{<:Any,<:Any,W})(x::AbstractArray{<:Real}) where {T <: Union{Float32,Float64}, W <: AbstractArray{T}} =
  a(T.(x))

outdims(l::ConvTranspose{N}, isize) where N = _convtransoutdims(isize[1:2], size(l.weight)[1:N], l.stride, l.dilation, l.pad)

"""
    DepthwiseConv(filter::Tuple, in=>out)
    DepthwiseConv(filter::Tuple, in=>out, activation)
    DepthwiseConv(filter, in => out, σ = identity; init = glorot_uniform,
                  stride = 1, pad = 0, dilation = 1)

Depthwise convolutional layer. `filter` should be a tuple like `(2, 2)`.
`in` and `out` specify the number of input and output channels respectively.
Note that `out` must be an integer multiple of `in`.

Data should be stored in WHCN order (width, height, # channels, batch size).
In other words, a 100×100 RGB image would be a `100×100×3×1` array,
and a batch of 50 would be a `100×100×3×50` array.

Accepts keyword arguments `weight` and `bias` to set the corresponding fields.
Setting `bias` to `Flux.Zeros()` will switch bias off for the layer.

Takes the keyword arguments `pad`, `stride` and `dilation`.
For input dimension N,
`pad` should be a single Integer indicating equal padding value for each spatial dimension,
a tuple of length (N-2) to apply symmetric padding or a tuple of length 2*(N-2)
indicating padding values for each spatial dimension at both the ends.
`stride` and `dilation` should be a single Integer or a tuple with N-2 parameters.
Use `pad=SamePad()` to apply padding so that outputsize == inputsize / stride.
"""
struct DepthwiseConv{N,M,F,A,V}
  σ::F
  weight::A
  bias::V
  stride::NTuple{N,Int}
  pad::NTuple{M,Int}
  dilation::NTuple{N,Int}
end

"""
    DepthwiseConv(weight::AbstractArray, bias::AbstractArray)
    DepthwiseConv(weight::AbstractArray, bias::AbstractArray, activation)

Constructs the `DepthwiseConv` layer with user defined weight and bias arrays.
forward pass.

Setting `bias` to `Flux.Zeros()` would switch `bias` off for the layer.

Takes the keyword arguments `pad`, `stride` and `dilation`.
For input dimension N,
`pad` should be a single Integer indicating equal padding value for each spatial dimension,
a tuple of length (N-2) to apply symmetric padding or a tuple of length 2*(N-2)
indicating padding values for each spatial dimension at both the ends.
`stride` and `dilation` should be a single Integer or a tuple with N-2 parameters.

For keyword-only constuctor, see also [`Conv`](@ref)
"""
function DepthwiseConv(w::AbstractArray{T,N}, b::Union{Zeros, AbstractVector{T}}, σ = identity;
                      stride = 1, pad = 0, dilation = 1) where {T,N}
  stride = expand(Val(N-2), stride)
  dilation = expand(Val(N-2), dilation)
  pad = calc_padding(pad, size(w)[1:N-2], dilation, stride)
  return DepthwiseConv(σ, w, b, stride, pad, dilation)
end

function DepthwiseConv(;weight::AbstractArray{T,N}, bias::Union{Zeros, AbstractVector{T}},
                      activation = identity, stride = 1, pad = 0, dilation = 1) where {T,N}
  DepthwiseConv(weight, bias, activation, stride = stride, pad = pad, dilation = dilation)
end

"""
    depthwiseconvfilter(filter::Tuple, in=>out)

Constructs a depthwise convolutional weight array defined by `filter` and channels
from `in` to `out`.

Accepts the keyword `init` (default: `glorot_uniform`) to control the sampling
distribution.

See also: [`convfilter`](@ref)
"""
depthwiseconvfilter(filter::NTuple{N,Integer}, ch::Pair{<:Integer,<:Integer};
                    init = glorot_uniform) where N = init(filter..., div(ch[2], ch[1]), ch[1])

function DepthwiseConv(k::NTuple{N,Integer}, ch::Pair{<:Integer,<:Integer}, σ = identity;
                      init = glorot_uniform, stride = 1, pad = 0, dilation = 1,
                      weight = depthwiseconvfilter(k, ch, init = init), bias = zeros(ch[2])) where N
  @assert ch[2] % ch[1] == 0 "Output channels must be integer multiple of input channels"

  return DepthwiseConv(
    weight,
    bias,
    σ;
    stride = stride,
    pad = pad,
    dilation = dilation
  )
end

@functor DepthwiseConv

function (c::DepthwiseConv)(x)
  σ, b = c.σ, reshape(c.bias, map(_->1, c.stride)..., :, 1)
  cdims = DepthwiseConvDims(x, c.weight; stride=c.stride, padding=c.pad, dilation=c.dilation)
  σ.(depthwiseconv(x, c.weight, cdims) .+ b)
end

function Base.show(io::IO, l::DepthwiseConv)
  print(io, "DepthwiseConv(", size(l.weight)[1:end-2])
  print(io, ", ", size(l.weight)[end], "=>", prod(size(l.weight)[end-1:end]))
  l.σ == identity || print(io, ", ", l.σ)
  print(io, ")")
end

(a::DepthwiseConv{<:Any,<:Any,W})(x::AbstractArray{T}) where {T <: Union{Float32,Float64}, W <: AbstractArray{T}} =
  invoke(a, Tuple{AbstractArray}, x)

(a::DepthwiseConv{<:Any,<:Any,W})(x::AbstractArray{<:Real}) where {T <: Union{Float32,Float64}, W <: AbstractArray{T}} =
  a(T.(x))

outdims(l::DepthwiseConv, isize) =
  output_size(DepthwiseConvDims(_paddims(isize, (1, 1, size(l.weight)[end], 1)), size(l.weight); stride = l.stride, padding = l.pad, dilation = l.dilation))

"""
    CrossCor(filter, in=>out)
    CrossCor(filter, in=>out, activation)
    CrossCor(filter, in => out, σ = identity; init = glorot_uniform,
             stride = 1, pad = 0, dilation = 1)

Standard cross convolutional layer. `filter` should be a tuple like `(2, 2)`.
`in` and `out` specify the number of input and output channels respectively.

Data should be stored in WHCN order (width, height, # channels, batch size).
In other words, a 100×100 RGB image would be a `100×100×3×1` array,
and a batch of 50 would be a `100×100×3×50` array.

Accepts keyword arguments `weight` and `bias` to set the corresponding fields.
Setting `bias` to `Flux.Zeros()` will switch bias off for the layer.

Takes the keyword arguments `pad`, `stride` and `dilation`.
For input dimension N,
`pad` should be a single Integer indicating equal padding value for each spatial dimension,
a tuple of length (N-2) to apply symmetric padding or a tuple of length 2*(N-2)
indicating padding values for each spatial dimension at both the ends.
`stride` and `dilation` should be a single Integer or a tuple with N-2 parameters.
Use `pad=SamePad()` to apply padding so that outputsize == inputsize / stride.

# Examples

Apply a `CrossCor` layer to a 1-channel input using a 2×2 window filter size, giving us a
16-channel output. Output is activated with ReLU.
```julia
filter = (2,2)
in = 1
out = 16
CrossCor((2, 2), 1=>16, relu)
```
"""
struct CrossCor{N,M,F,A,V}
  σ::F
  weight::A
  bias::V
  stride::NTuple{N,Int}
  pad::NTuple{M,Int}
  dilation::NTuple{N,Int}
end

"""
    CrossCor(weight::AbstractArray, bias::AbstractArray)
    CrossCor(weight::AbstractArray, bias::AbstractArray, activation)

Constructs the standard cross convolutional layer with user defined weight and bias
arrays.

Setting `bias` to `Flux.Zeros()` would switch `bias` off for the layer.

Takes the keyword arguments `pad`, `stride` and `dilation`.
For input dimension N,
`pad` should be a single Integer indicating equal padding value for each spatial dimension,
a tuple of length (N-2) to apply symmetric padding or a tuple of length 2*(N-2)
indicating padding values for each spatial dimension at both the ends.
`stride` and `dilation` should be a single Integer or a tuple with N-2 parameters.

For keyword-only constuctor, see also [`Conv`](@ref)
"""
function CrossCor(w::AbstractArray{T,N}, b::Union{Zeros, AbstractVector{T}}, σ = identity;
                  stride = 1, pad = 0, dilation = 1) where {T,N}
  stride = expand(Val(N-2), stride)
  dilation = expand(Val(N-2), dilation)
  pad = calc_padding(pad, size(w)[1:N-2], dilation, stride)
  return CrossCor(σ, w, b, stride, pad, dilation)
end

function CrossCor(;weight::AbstractArray{T,N}, bias::Union{Zeros, AbstractVector{T}},
                      activation = identity, stride = 1, pad = 0, dilation = 1) where {T,N}
  CrossCor(weight, bias, activation, stride = stride, pad = pad, dilation = dilation)
end

function CrossCor(k::NTuple{N,Integer}, ch::Pair{<:Integer,<:Integer}, σ = identity;
                  init = glorot_uniform, stride = 1, pad = 0, dilation = 1,
                  weight = convfilter(k, ch, init = init), bias = zeros(ch[2])) where N

  CrossCor(weight, bias, σ,
       stride = stride, pad = pad, dilation = dilation)
end

@functor CrossCor

function crosscor(x, w, ddims::DenseConvDims)
  ddims = DenseConvDims(ddims, F=true)
  return conv(x, w, ddims)
end

function (c::CrossCor)(x::AbstractArray)
  # TODO: breaks gpu broadcast :(
  # ndims(x) == ndims(c.weight)-1 && return squeezebatch(c(reshape(x, size(x)..., 1)))
  σ, b = c.σ, reshape(c.bias, map(_->1, c.stride)..., :, 1)
  cdims = DenseConvDims(x, c.weight; stride=c.stride, padding=c.pad, dilation=c.dilation)
  σ.(crosscor(x, c.weight, cdims) .+ b)
end

function Base.show(io::IO, l::CrossCor)
  print(io, "CrossCor(", size(l.weight)[1:ndims(l.weight)-2])
  print(io, ", ", size(l.weight, ndims(l.weight)-1), "=>", size(l.weight, ndims(l.weight)))
  l.σ == identity || print(io, ", ", l.σ)
  print(io, ")")
end

(a::CrossCor{<:Any,<:Any,W})(x::AbstractArray{T}) where {T <: Union{Float32,Float64}, W <: AbstractArray{T}} =
  invoke(a, Tuple{AbstractArray}, x)

(a::CrossCor{<:Any,<:Any,W})(x::AbstractArray{<:Real}) where {T <: Union{Float32,Float64}, W <: AbstractArray{T}} =
  a(T.(x))

outdims(l::CrossCor, isize) =
  output_size(DenseConvDims(_paddims(isize, size(l.weight)), size(l.weight); stride = l.stride, padding = l.pad, dilation = l.dilation))

"""
    AdaptiveMaxPool(out)

Adaptive max pooling layer. `out` is the desired output size (batch and channel dimension excluded).
"""
struct AdaptiveMaxPool{S, O}
  out::NTuple{O, Int}
  AdaptiveMaxPool(out::NTuple{O, Int}) where O = new{O + 2, O}(out)
end

function (a::AdaptiveMaxPool{S})(x::AbstractArray{T, S}) where {S, T}
  insize = size(x)[1:end-2]
  outsize = a.out
  stride = insize .÷ outsize
  k = insize .- (outsize .- 1) .* stride
  pad = 0
  pdims = PoolDims(x, k; padding=pad, stride=stride)
  return maxpool(x, pdims)
end

function Base.show(io::IO, a::AdaptiveMaxPool)
  print(io, "AdaptiveMaxPool(", a.out, ")")
end

"""
    AdaptiveMeanPool(out)

Adaptive mean pooling layer. `out` is the desired output size (batch and channel dimension excluded).
"""
struct AdaptiveMeanPool{S, O}
  out::NTuple{O, Int}
  AdaptiveMeanPool(out::NTuple{O, Int}) where O = new{O + 2, O}(out)
end

function (a::AdaptiveMeanPool{S})(x::AbstractArray{T, S}) where {S, T}
  insize = size(x)[1:end-2]
  outsize = a.out
  stride = insize .÷ outsize
  k = insize .- (outsize .- 1) .* stride
  pad = 0
  pdims = PoolDims(x, k; padding=pad, stride=stride)
  return meanpool(x, pdims)
end

function Base.show(io::IO, a::AdaptiveMeanPool)
  print(io, "AdaptiveMeanPool(", a.out, ")")
end

"""
    GlobalMaxPool()

Global max pooling layer.

Transforms (w,h,c,b)-shaped input into (1,1,c,b)-shaped output,
by performing max pooling on the complete (w,h)-shaped feature maps.
"""
struct GlobalMaxPool end

function (g::GlobalMaxPool)(x)
  # Input size
  x_size = size(x)
  # Kernel size
  k = x_size[1:end-2]
  # Pooling dimensions
  pdims = PoolDims(x, k)

  return maxpool(x, pdims)
end

function Base.show(io::IO, g::GlobalMaxPool)
  print(io, "GlobalMaxPool()")
end

"""
    GlobalMeanPool()

Global mean pooling layer.

Transforms (w,h,c,b)-shaped input into (1,1,c,b)-shaped output,
by performing mean pooling on the complete (w,h)-shaped feature maps.
"""
struct GlobalMeanPool end

function (g::GlobalMeanPool)(x)
  # Input size
  x_size = size(x)
  # Kernel size
  k = x_size[1:end-2]
  # Pooling dimensions
  pdims = PoolDims(x, k)

  return meanpool(x, pdims)
end

function Base.show(io::IO, g::GlobalMeanPool)
  print(io, "GlobalMeanPool()")
end

"""
    MaxPool(k; pad = 0, stride = k)

Max pooling layer. `k` is the size of the window for each dimension of the input.

Use `pad=SamePad()` to apply padding so that outputsize == inputsize / stride.
=======
"""
struct MaxPool{N,M}
  k::NTuple{N,Int}
  pad::NTuple{M,Int}
  stride::NTuple{N,Int}
end

function MaxPool(k::NTuple{N,Integer}; pad = 0, stride = k) where N
  stride = expand(Val(N), stride)
  pad = calc_padding(pad, k, 1, stride)
  return MaxPool(k, pad, stride)
end

function (m::MaxPool)(x)
    pdims = PoolDims(x, m.k; padding=m.pad, stride=m.stride)
    return maxpool(x, pdims)
end

function Base.show(io::IO, m::MaxPool)
  print(io, "MaxPool(", m.k, ", pad = ", m.pad, ", stride = ", m.stride, ")")
end

outdims(l::MaxPool{N}, isize) where N = output_size(PoolDims(_paddims(isize, (l.k..., 1, 1)), l.k; stride = l.stride, padding = l.pad))

"""
    MeanPool(k; pad = 0, stride = k)

Mean pooling layer. `k` is the size of the window for each dimension of the input.

Use `pad=SamePad()` to apply padding so that outputsize == inputsize / stride.
"""
struct MeanPool{N,M}
    k::NTuple{N,Int}
    pad::NTuple{M,Int}
    stride::NTuple{N,Int}
end

function MeanPool(k::NTuple{N,Integer}; pad = 0, stride = k) where N
  stride = expand(Val(N), stride)
  pad = calc_padding(pad, k, 1, stride)
  return MeanPool(k, pad, stride)
end

function (m::MeanPool)(x)
    pdims = PoolDims(x, m.k; padding=m.pad, stride=m.stride)
    return meanpool(x, pdims)
end

function Base.show(io::IO, m::MeanPool)
  print(io, "MeanPool(", m.k, ", pad = ", m.pad, ", stride = ", m.stride, ")")
end

outdims(l::MeanPool{N}, isize) where N = output_size(PoolDims(_paddims(isize, (l.k..., 1, 1)), l.k; stride = l.stride, padding = l.pad))


text:
using MshReader,
      WriteVTK,
      SparseArrays,
      LinearAlgebra

ext: jl


text:
###################
# Pretty Printing #
###################

function prettytime(t)
    if t < 1e3
        value, units = t, "ns"
    elseif t < 1e6
        value, units = t / 1e3, "μs"
    elseif t < 1e9
        value, units = t / 1e6, "ms"
    else
        value, units = t / 1e9, "s"
    end

    if round(value) >= 100
        str = string(@sprintf("%.0f", value), units)
    elseif round(value * 10) >= 100
        str = string(@sprintf("%.1f", value), units)
    else
        str = string(@sprintf("%.2f", value), units)
    end
    return lpad(str, 6, " ")
end

function prettymemory(b)
    if b < 1000
        value = -1
        str = string(round(Int, b), "B")
    elseif b < 1000^2
        value, units = b / 1024, "KiB"
    elseif b < 1000^3
        value, units = b / 1024^2, "MiB"
    else
        value, units = b / 1024^3, "GiB"
    end

    if round(value) >= 100
        str = string(@sprintf("%.0f", value), units)
    elseif round(value * 10) >= 100
        str = string(@sprintf("%.1f", value), units)
    elseif value >= 0
        str = string(@sprintf("%.2f", value), units)
    else
        str = "-"
    end
    return lpad(str, 7, " ")
end

function prettypercent(nominator, denominator)
    value = nominator / denominator * 100

    if denominator == 0 && nominator == 0
        str = " - %"
    elseif denominator == 0
        str = "inf %"
    elseif round(value) >= 100
        str = string(@sprintf("%.0f", value), "%")
    elseif round(value * 10) >= 100
        str = string(@sprintf("%.1f", value), "%")
    else
        str = string(@sprintf("%.2f", value), "%")
    end
    return str
end

function prettycount(t::Int)
    if t < 1000
        return string(t)
    elseif t < 1000^2
        value, units = t / 1000, "k"
    elseif t < 1000^3
        value, units = t / 1e6, "M"
    else
        value, units = t / 1e9, "B"
    end

    if round(value) >= 100
        str = string(@sprintf("%.0f", value), units)
    elseif round(value * 10) >= 100
        str = string(@sprintf("%.1f", value), units)
    else
        str = string(@sprintf("%.2f", value), units)
    end
    return str
end


text:
#Copyright (c) 2015 Rex Computing and Isaac Yonemoto

#see LICENSE.txt

#this work was supported in part by DARPA Contract D15PC00135

#unit-subtraction.jl

#testing subtraction

#__diff_exact - an exact subtraction with the first entity having a higher exponent

wone = convert(Unum{4,6}, 1)
wtwo = convert(Unum{4,6}, 2)
wthr = convert(Unum{4,6}, 3)

eone = Unums.decode_exp(wone)
etwo = Unums.decode_exp(wtwo)
ethr = Unums.decode_exp(wthr)

@test Unums.diff_exact(wone, -wone, eone, eone) == zero(Unum{4,6})   #one plus one is two
@test Unums.diff_exact(wtwo, -wone, etwo, eone) == wone              #two minus one is one
println("===========")
@test Unums.diff_exact(wthr, -wone, ethr, eone) == wtwo

#test that the sign things seem to work.
x = Unum{3,3}(7.191347884985208e19)
y = Unum{3,3}(1.3733397310822187e27)
@test is_negative(x - y)

x = Unum{4,5}(-1.314174462714471e27)
y = Unum{4,5}(1.4517351819998816e20)
@test Unums.is_inward(y, x)
@test is_negative(x - y)

#long subtraction, discovered as an error in rump's formula.

x = Unum{4,7}(0x00000000000000F9, UInt64[0x7d31ee79ca44b643,0x91bb985960000000], 0x0000, 0x0007, 0x0063)
y = Unum{4,7}(0x00000000000000F9, UInt64[0x7d31ee79ca44b643,0x91bb985960000080], 0x0000, 0x0007, 0x0079)
@test y - x == Unum{4,7}(2)

#errors discovered as a result of linear algebra experiments.

x = Unum{3,5}(0x0000000000000007, 0x1000000000000000, 0x0000, 0x0002, 0x001F)
y = Unum{3,5}(0x0000000000000006, 0xA000000000000000, 0x0001, 0x0002, 0x001F)
res = x - y
@test Unums.lub(res) == Unum{3,5}(4)

x = Unum{3,5}(0x0000000000000007, 0x1000000000000000, 0x0000, 0x0002, 0x001F)
y = Ubound(Unum{3,5}(0x0000000000000006, 0x9FFFFFFF00000000, 0x0001, 0x0002, 0x001F), Unum{3,5}(0x0000000000000006, 0xA000000000000000, 0x0001, 0x0002, 0x001F))
res = x - y
@test res > Unum{3,5}(3.75)

x = Unum{3,5}(0x0000000000000007, 0x197E680000000000, 0x0003, 0x0002, 0x0019)
y = Unum{3,5}(0x000000000000000C, 0x04A71A7800000000, 0x0001, 0x0003, 0x001C)
res = x - y
@test is_negative(res.lower)

x = Ubound(Unum{3,5}(0x0000000000000007, 0x197E680000000000, 0x0003, 0x0002, 0x0019), Unum{3,5}(0x0000000000000007, 0x197E67E000000000, 0x0003, 0x0002, 0x001A))
y = Ubound(Unum{3,5}(0x000000000000000C, 0x04A71A6C00000000, 0x0001, 0x0003, 0x001F), Unum{3,5}(0x000000000000000C, 0x04A71A7800000000, 0x0001, 0x0003, 0x001C))
res = x - y
@test res == Ubound(Unum{3,5}(0x000000000000000C, 0x91664E9F00000000, 0x0003, 0x0003, 0x001F), Unum{3,5}(0x000000000000000C, 0x91664E5C00000000, 0x0003, 0x0003, 0x001F))
#problem is that the sign is not properly handled.

#tests discovered Jan 6 2017:  Various problems stemming form size comparison of strange subnormals.
x = Unum{3,5}(0x000000000000000C, 0x37FFFFFF00000000, 0x0001, 0x0003, 0x001F)
y = Unum{3,5}(0x000000000000000C, 0x3800000000000000, 0x0000, 0x0003, 0x001F)
res = x - y
@test(res == Unum{3,5}(0x0000000000000000, 0x0000000000000000, 0x0003, 0x0004, 0x000C))

x = Ubound(Unum{3,5}(0x0000000000000007, 0xD800000000000000, 0x0001, 0x0002, 0x001F), Unum{3,5}(0x000000000000000C, 0x37FFFFFF00000000, 0x0001, 0x0003, 0x001F))
y = Unum{3,5}(0x000000000000000C, 0x3800000000000000, 0x0000, 0x0003, 0x001F)
res = x - y
@test res ==  Ubound(Unum{3,5}(0x0000000000000006, 0x2FFFFFFE00000000, 0x0003, 0x0002, 0x001E), Unum{3,5}(0x0000000000000000, 0x0000000000000000, 0x0003, 0x0000, 0x001F))

x = Unum{3,5}(0x0000000000000001, 0x0000000000000000, 0x0001, 0x0000, 0x001F)
y = Unum{3,5}(0x0000000000000001, 0x0000000000000000, 0x0000, 0x0000, 0x0000)
res = x - y

#=
#11 September 2015 - identified through continuous testing.  The problem here is
#that the two normal floats share their exponent factor, and clears a huge swath
#of zeros in the fraction.
x1 = Unum{4,6}(0x0033, 0x0007, z16, 0x04dcb922d2c6d000, 0x000000015)
x2 = Unum{4,6}(0x0033, 0x0007, Unums.UNUM_SIGN_MASK, 0x04dcb8d6e79b5000, 0x000000015)
de = decode_exp(0x0007, 0x000000015)
@test Unums.__diff_exact(x1, x2, de, de) == Unum{4,6}(0x0016, 0x0008, z16, 0x2facae0000000000, UInt64(0x7b))
#fixed by implementing a subroutine that is triggered only when carry is zero for
#a no-offset subtraction.

#11 septemeber 2015 - identified through continuous testing.  The problem here is
#that two normal floats with an offset one off can also trigger a huge swath of
#zeros in the fraction.
#diagram:  (1)0 000000XXXXX
#            (1)111111XXXXX
#             0 000001XXXXX
x3 = Unum{4,6}(0x0032,0x0006,Unums.UNUM_SIGN_MASK,0x95b579166f2be000,0x0000000000000065)
x4 = Unum{4,6}(0x0033,0x0006,0x0000,0x120df3b26a793000,0x0000000000000066)
dx3 = decode_exp(x3)
dx4 = decode_exp(x4)
@test Unums.__diff_exact(x4, x3, dx4, dx3) == Unum{4,6}(0x002f,0x0006,0x0000,0x1cccdc9ccb8d0000,0x0000000000000064)
#fixed by implementing a subroutine that also identifies this situation.  to keep
#code DRY, created the __shift_many_zeros subroutine.

#11 september 2015 - identified through continuous testing, in parallel to the
#ubound problem as in the unit test for addition.
x5 = Unum{4,6}(0x0032,0x0007,Unums.UNUM_SIGN_MASK,0x4410e89562546000,UInt64(0xf2))
x6 = Unum{4,6}(0x003f,0x0007,z16,0x3db8d7e07e3733ef,UInt64(0xf1))
dx5 = decode_exp(x5)
dx6 = decode_exp(x6)
@test Unums.__diff_exact(x5,x6,dx5,dx6) == Unum{4,6}(0x003f, 0x0007, Unums.UNUM_UBIT_MASK | Unums.UNUM_SIGN_MASK, 0x4a68f94a46718c11, UInt64(0xf1))
#resolved due to an incorrect offset in a fillbit call.
=#

max_stars_repo_name: REX-Computing/unumjl

max_issues_repo_name: REX-Computing/unumjl

max_forks_repo_name: REX-Computing/unumjl


text:
# Formula displays
@testset "Formula Displays" begin
    data = dataset("Ecdat", "Crime")
    # Continuous Response
    f = @formula(CRMRTE ~ 1 + PrbConv)
    model = fit(EconometricModel, f, data)
    @test Econometrics.clean_fm(model) == "Formula: CRMRTE ~ 1 + PrbConv"
    f = @formula(CRMRTE ~ 0 + PrbConv)
    model = fit(EconometricModel, f, data)
    @test Econometrics.clean_fm(model) == "Formula: CRMRTE ~ 0 + PrbConv"
    f = @formula(CRMRTE ~ -1 + PrbConv)
    model = fit(EconometricModel, f, data)
    @test Econometrics.clean_fm(model) == "Formula: CRMRTE ~ -1 + PrbConv"
    f = @formula(CRMRTE ~ -1 + PrbConv + PrbConv & PrbPris^2 + absorb(County + Year))
    model = fit(EconometricModel, f, data)
    @test Econometrics.clean_fm(model) ==
        "Formula: CRMRTE ~ -1 + PrbConv + absorb(County + Year) + PrbConv & (PrbPris ^ 2)"
    f = @formula(log(CRMRTE) ~ 1)
    model = fit(EconometricModel, f, data)
    @test Econometrics.clean_fm(model) == "Formula: log(CRMRTE) ~ 1"
    # Random Effects Model
    f = @formula(CRMRTE ~ 1 + PrbConv)
    model = fit(RandomEffectsEstimator, f, data, panel = :County, time = :Year)
    @test Econometrics.clean_fm(model) == "Formula: CRMRTE ~ PrbConv"
    # Nominal Response Model
    data =
        CSV.read(joinpath(pkgdir(Econometrics), "data", "insure.csv"), DataFrame) |>
        (data -> select(data, [:insure, :age, :male, :nonwhite, :site])) |> dropmissing |>
        (data -> transform!(data, [:insure, :site] .=> categorical, renamecols = false))
    f = @formula(insure ~ 0 + age^2 + male + nonwhite + site)
    model = fit(
        EconometricModel,
        f,
        data,
        contrasts = Dict(:insure => DummyCoding(base = "Uninsure")),
        )
    @test Econometrics.clean_fm(model) == "Formula: insure ~ 0 + (age ^ 2) + male + nonwhite + site"
    # Ordinal Response Model
    data =
        dataset("Ecdat", "Kakadu") |>
        (data -> select(data, [:RecParks, :Sex, :Age, :Schooling]))
    data.RecParks = convert(Vector{Int}, data.RecParks)
    data.RecParks = levels!(categorical(data.RecParks, ordered = true), collect(1:5))
    model = fit(EconometricModel, @formula(RecParks ~ Age * Age + Sex), data)
    @test Econometrics.clean_fm(model) == "Formula: RecParks ~ Age + Sex + Age & Age"
end


text:
"""
    local_interaction(::AbstractFockAddress, u)
    local_interaction(::AbstractFockAddress, ::AbstractFockAddress, v)

Return the sum of (mode-wise) local interactions ``\\frac{u}{2} \\sum_i n_i(n_i-1)`` of a
single component Fock state, or ``v \\sum_i n_{↑,i} n_{↓,i}`` between two Fock states. For a
multi-component Fock state, return the eigenvalue of

```math
\\frac{1}{2}\\sum_{i, σ, τ} u_{σ,τ} a^†_{σ,i}a^†_{τ,i}a^†_{τ,i}a^†_{σ,i} ,
```

where `u::SMatrix` is a symmetric matrix of interaction constants, `i` is a mode index,
and `σ`, `τ` are component indices.

See also [`BoseFS`](@ref), [`FermiFS`](@ref), [`CompositeFS`](@ref).
"""
local_interaction(b::BoseFS, u) = u * bose_hubbard_interaction(b) / 2
local_interaction(f::FermiFS, _) = 0
local_interaction(f::FermiFS, g::FermiFS, v) = v * count_ones(f.bs & g.bs)
local_interaction(f::FermiFS, b::BoseFS, v) = local_interaction(b, f, v)

function local_interaction(a::BoseFS, b::BoseFS, v)
    occ_a = occupied_modes(a)
    occ_b = occupied_modes(b)

    (n_a, i_a, _), st_a = iterate(occ_a)
    (n_b, i_b, _), st_b = iterate(occ_b)

    acc = 0
    while true
        if i_a > i_b
            # b is behind and needs to do a step
            iter = iterate(occ_b, st_b)
            isnothing(iter) && return acc * v
            (n_b, i_b, _), st_b = iter
        elseif i_a < i_b
            # a is behind and needs to do a step
            iter = iterate(occ_a, st_a)
            isnothing(iter) && return acc * v
            (n_a, i_a, _), st_a = iter
        else
            # a and b are at the same position
            acc += n_a * n_b
            # now both need to do a step
            iter = iterate(occ_a, st_a)
            isnothing(iter) && return acc * v
            (n_a, i_a, _), st_a = iter
            iter = iterate(occ_b, st_b)
            isnothing(iter) && return acc * v
            (n_b, i_b, _), st_b = iter
        end
    end
end
function local_interaction(b::BoseFS, f::FermiFS, v)
    acc = 0
    for (n, i) in occupied_modes(b)
        acc += is_occupied(f, i) * n
    end
    return acc * v
end
function local_interaction(fs::CompositeFS, m)
    _interactions(fs.components, m)
end

"""
    _interaction_col(a, bs::Tuple, us::Tuple)

Sum the local interactions of the Fock state `a` with all states in `bs` using the
interaction constants in `us`. This is used to compute all interactions in the column
below the diagonal of the interaction matrix.
"""
@inline _interaction_col(a, ::Tuple{}, ::Tuple{}) = 0
@inline function _interaction_col(a, (b, bs...), (u, us...))
    return local_interaction(a, b, u) + _interaction_col(a, bs, us)
end

"""
    _interactions(addresses, interaction_matrix)

Compute all pairwise interactions in a tuple of `addresses`. The `interaction_matrix` sets the
intraction strengths.

The code is equivalent to the following.

```julia
acc = 0.0
for (i, a) in enumerate(addresses)
    acc += local_interaction(a, interaction_matrix[i, i])
    for (j, b) in enumerate(addresses[i+1:end])
        acc += local_interaction(a, b, interaction_matrix[i, j])
    end
end
return acc
```

It is implemented recursively to ensure type stability.
"""
@inline _interactions(::Tuple{}, ::SMatrix{0,0}) = 0.0
@inline function _interactions(
    (a, as...)::NTuple{N,AbstractFockAddress}, m::SMatrix{N,N}
) where {N}
    # Split the matrix into the column we need now, and the rest.
    (u, column...) = Tuple(m[:, 1])
    # Type-stable way to subset SMatrix:
    rest = SMatrix{N-1,N-1}(view(m, 2:N, 2:N))

    # Get the self-interaction first.
    self = local_interaction(a, u)
    # Get the interactions for the rest of the row.
    row = _interaction_col(a, as, column)
    # Get the interaction for the rest of the rows.
    return self + row + _interactions(as, rest)
end

###
### HubbardRealSpace
###
"""
    HubbardRealSpace(address; u=ones(C, C), t=ones(C), geometry=PeriodicBoundaries(M,))

Hubbard model in real space. Supports single or multi-component Fock state
addresses (with `C` components) and various (rectangular) lattice geometries
in arbitrary dimensions.

```math
  \\hat{H} = -\\sum_{\\langle i,j\\rangle,σ} t_σ a^†_{iσ} a_{jσ} +
  \\frac{1}{2}\\sum_{i,σ} u_{σ,σ} n_{iσ} (n_{iσ} - 1) +
  \\frac{1}{2}\\sum_{i,σ≠τ}u_{σ,τ} n_{iσ} n_{iτ}
```

## Address types

* [`BoseFS`](@ref): Single-component Bose-Hubbard model.
* [`FermiFS`](@ref): Single-component Fermi-Hubbard model. This address only provides a
  single species of (non-interacting) fermions. You probably want to use [`CompositeFS`](@ref).
* [`CompositeFS`](@ref): For multi-component models.

## Geometries

* [`PeriodicBoundaries`](@ref)
* [`HardwallBoundaries`](@ref)
* [`LadderBoundaries`](@ref)

## Other parameters

* `u`: the on-site interaction parameters. Must be a symmetric matrix. `u[i, j]`
  corresponds to the interaction between the `i`-th and `j`-th component. `u[i, i]`
  corresponds to the interaction of a component with itself. Note that `u[i,i]` must
  be zero for fermionic components.
* `t`: the hopping strengths. Must be a vector of length `C`. The `i`-th element of the
  vector corresponds to the hopping strength of the `i`-th component.

"""
struct HubbardRealSpace{
    C,A,G, # C: components
    # The following need to be type params.
    T<:SVector{C,Float64},
    U<:SMatrix{C,C,Float64},
} <: AbstractHamiltonian{Float64}
    address::A
    u::U # interactions
    t::T # hopping strengths
    geometry::G
end

function HubbardRealSpace(
    address;
    u=ones(num_components(address), num_components(address)),
    t=ones(num_components(address)),
    geometry=PeriodicBoundaries((num_modes(address),))
)
    C = num_components(address)

    # Sanity checks
    if prod(size(geometry)) ≠ num_modes(address)
        error("`geometry` does not have the correct number of sites")
    elseif length(u) ≠ 1 && !issymmetric(u)
        error("`u` must be symmetric")
    elseif length(u) ≠ C * C
        error("`u` must be a $C × $C matrix")
    elseif size(t) ≠ (C,)
        error("`t` must be a vector of length $C")
    elseif address isa BoseFS2C
        error("`BoseFS2C` is not supported for this Hamiltonian, use `CompositeFS`")
    end
    warn_fermi_interaction(address, u)

    u_mat = SMatrix{C,C,Float64}(u)
    t_vec = SVector{C,Float64}(t)
    return HubbardRealSpace{C,typeof(address),typeof(geometry),typeof(t_vec),typeof(u_mat)}(
        address, u_mat, t_vec, geometry,
    )
end

"""
    warn_fermi_interaction(address, u)

Warn if interaction matrix `u` does not make sense for `address`.
"""
function warn_fermi_interaction(address::CompositeFS, u)
    C = num_components(address)
    for c in 1:C
        if address.components[c] isa FermiFS && u ≠ ones(C,C) && u[c,c] ≠ 0
            @warn "component $(c) is fermionic, but was given a self-interaction " *
                "strength of $(u[c,c])" maxlog=1
        end
    end
end
function warn_fermi_interaction(address::FermiFS, u)
    if u ≠ ones(1, 1) && u[1, 1] ≠ 0
        @warn "address is fermionic, but was given a self-interaction " *
            "strength of $(u[1,1])" maxlog=1
    end
end
warn_fermi_interaction(_, _) = nothing

LOStructure(::Type{<:HubbardRealSpace}) = IsHermitian()

function Base.show(io::IO, h::HubbardRealSpace)
    println(io, "HubbardRealSpace(")
    println(io, "  ", starting_address(h), ",")
    println(io, "  u = ", Float64.(h.u), ",")
    println(io, "  t = ", Float64.(h.t), ",")
    println(io, "  geometry = ", h.geometry, ",")
    println(io, ")")
end

starting_address(h::HubbardRealSpace) = h.address
diagonal_element(h::HubbardRealSpace, address) = local_interaction(address, h.u)
diagonal_element(h::HubbardRealSpace{1}, address) = local_interaction(address, h.u[1])

###
### Offdiagonals
###
# This may be an inefficient implementation, but it is not actually used anywhere in the
# main algorithm.
get_offdiagonal(h::HubbardRealSpace, add, i) = offdiagonals(h, add)[i]
num_offdiagonals(h::HubbardRealSpace, add) = length(offdiagonals(h, add))

"""
    HubbardRealSpaceCompOffdiagonals{G,A} <: AbstractOffdiagonals{A,Float64}

Offdiagonals for a single address component. Used with [`HubbardRealSpace`](@ref) model
with a single-component address, or a component of a [`CompositeFS`](@ref).
"""
struct HubbardRealSpaceCompOffdiagonals{G,A} <: AbstractOffdiagonals{A,Float64}
    geometry::G
    address::A
    t::Float64
    length::Int
end

function offdiagonals(h::HubbardRealSpace, comp, add)
    neighbours = num_neighbours(h.geometry)
    return HubbardRealSpaceCompOffdiagonals(
        h.geometry, add, h.t[comp], num_occupied_modes(add) * neighbours
    )
end

Base.size(o::HubbardRealSpaceCompOffdiagonals) = (o.length,)

@inline function Base.getindex(o::HubbardRealSpaceCompOffdiagonals, chosen)
    neighbours = num_neighbours(o.geometry)
    particle, neigh = fldmod1(chosen, neighbours)
    src_index = find_occupied_mode(o.address, particle)
    neigh = neighbour_site(o.geometry, src_index.mode, neigh)

    if neigh == 0
        return o.address, 0.0
    else
        dst_index = find_mode(o.address, neigh)
        new_add, value = excitation(o.address, (dst_index,), (src_index,))
        return new_add, -o.t * value
    end
end

# For simple models with one component.
offdiagonals(h::HubbardRealSpace{1,A}, add::A) where {A} = offdiagonals(h, 1, add)

# Multi-component part
"""
    HubbardRealSpaceOffdiagonals{A,T<:Tuple} <: AbstractOffdiagonals{A,Float64}

Offdiagonals of a [`HubbardRealSpace`](@ref) model with a [`CompositeFS`](@ref) address.
"""
struct HubbardRealSpaceOffdiagonals{A,T<:Tuple} <: AbstractOffdiagonals{A,Float64}
    address::A
    parts::T
    length::Int
end

"""
    get_comp_offdiags(h::HubbardRealSpace, add)

Get offdiagonals of all components of address in a type-stable manner.
"""
@inline function get_comp_offdiags(h::HubbardRealSpace, address)
    return _get_comp_offdiags(address.components, h, Val(1))
end

# All steps of recursive function (should) get inlined, creating a type-stable tuple of
# offdiagonals.
@inline function _get_comp_offdiags((a,as...), h, ::Val{I}) where {I}
    return (offdiagonals(h, I, a), _get_comp_offdiags(as, h, Val(I+1))...)
end
@inline _get_comp_offdiags(::Tuple{}, h, ::Val) = ()

function offdiagonals(h::HubbardRealSpace{C,A}, address::A) where {C,A<:CompositeFS}
    parts = get_comp_offdiags(h, address)
    return HubbardRealSpaceOffdiagonals(address, parts, sum(length, parts))
end

Base.size(o::HubbardRealSpaceOffdiagonals) = (o.length,)

# Becomes type unstable without inline for lots of components. Recursive function is used
# because the type of the result of `o.parts[i]` can not be inferred.
@inline function Base.getindex(o::HubbardRealSpaceOffdiagonals{A}, chosen) where {A}
    return _getindex(o.parts, o.address, chosen, Val(1))
end
@inline function _getindex((p, ps...), address::A, chosen, comp::Val{I}) where {A,I}
    if chosen ≤ length(p)
        new_add, val = p[chosen]
        return BitStringAddresses.update_component(address, new_add, comp), val
    else
        chosen -= length(p)
        return _getindex(ps, address, chosen, Val(I + 1))
    end
end


text:
include("CorrFuncs.jl")

function histrecon_u(dims, C2, S2, philen)
    guess, pix = make_rand_im(philen, dims)
    display(heatmap(guess, color=:grays, aspect_ratio=1))
    maxrng = Int64(round(sqrt(dims[1]^2 + dims[2]^2)+1))
    wpick = rand(1:length(pix[1]))
    bpick = rand(1:length(pix[2]))
    kern = CartesianIndices((-1:1, -1:1))

    original = label_components(guess, trues(3,3))
    SN = SN_comp(dims, maxrng)
    typeof(SN)
    S2_BN = S2_initialise(pix, maxrng)
    C2_BN = C2_initialise(guess, maxrng)
    S2n = naninf(S2_finalise(SN, S2_BN, guess))
    C2n = naninf(S2_finalise(SN, C2_BN, guess))
    #print(SN, "\n\n", S2_BN, "\n\n", C2_BN)
    Err1 = 2.0
    count = 1
    cont_break=0;
    T0 = 1.0
    cont_surf_opt=0
    pflag = 0
    true_count = 0
    while Err1 > 0.0001
        if (cont_surf_opt !== 1)
            wpick, bpick = pick_rand_coord(pix)
            #println("stoachastic mode")
        else
            surf_res = surf_opt(guess)
            shortlist = surf_rand(surf_res, pix)
            wpick, bpick = pick_rand_coord(shortlist)
            wpick, bpick = find_equivalent(shortlist, pix, (wpick, bpick))
            #println("surface optimisation mode")
            #println("recon shortlist: \twlength = ", length(shortlist[1]), " blength = ", length(shortlist[2]))
            #println("wpick = ", wpick, "\tbpick = ", bpick)

        end

        pix  = swap_pix(guess, pix, wpick, bpick)
        modified = label_components(guess, trues(3,3))
        S2_BN_d = update_S2_BN(S2_BN, pix, wpick, bpick, philen, maxrng)
        C2_BN_d = update_C2_BN(C2_BN, pix, original, modified, wpick, bpick, maxrng)
        #println(length(S2_BN), "\n", length(S2_BN_d))
        S2_BN .+= S2_BN_d
        C2_BN .+= C2_BN_d
        #println("\n", S2_BN_d, "\n")
        #println("C2_BN_d = ", C2_BN_d)
        S2n = naninf(S2_finalise(SN, S2_BN, guess))
        C2n = naninf(S2_finalise(SN, C2_BN, guess))
        #println("C2n = ", C2n)

        true_count +=1

        Err2 =  ediff_C2(C2n, C2) + (ediff_S2(S2n, S2))# +
        #SN_comp(size(guess), length(guess), maxrng)
        delErr = Err2 - Err1
        if count == 1
            T0=-abs(delErr)/log(0.5)
            println("T0 = ", T0)
        end
        T = T0*0.99^count
        if (rand(0:1) < min(exp(-delErr/T), 1))
            count += 1
            cont_break=0
            #print("change commited at count ", count, "\t")
            original = modified
            Err1 = Err2
        else
            swap_pix(guess, pix, wpick, bpick)
            S2_BN .-= S2_BN_d
            C2_BN .-= C2_BN_d
            cont_break += 1

        end

         # we bake in the change here - this must be undone
        if (mod(true_count, 1000) == 0) & (count !=pflag)
            println("\ncount is ", count, "\t  True count is ", true_count, "\t prob of rejection =", min(exp(delErr/T), 1), "\t  S2 Err = ", ediff_S2_fft(S2n,  S2)*1250/length(guess), "\t C2 Err = ",  ediff_C2(C2n, C2))
            println("S2_delerr = ", ediff_S2(S2n, S2), "\tC2_delerr = ", ediff_C2(C2n, C2))
            println("C2[1] = ", C2n[1], "\t\tS2[1] = ", S2n[1])
            if cont_surf_opt==1
                println("Surface optimisation mode")
            else
                println("Stoachastic mode")
            end
            #println(Err1)
            pflag = count
        end

        if count==4000
            cont_surf_opt=1
            #println("Changed to surface optimisation mode")
        end

        if true_count >= 10000

            break
        end

    end
    return guess, S2n, C2n, S2_BN, C2_BN, SN
end

begin
    test = loadim("uniform_circles3.png")
    dims = size(test)#(200,100)
    maxrng = Int64(round(sqrt(dims[1]^2 + dims[2]^2)+1))
    philen = sum(test)#Int(round(0.3*dims[1]*dims[2]))

    SN = SN_comp(dims, maxrng) #Is SN_comp supposed to be multiplied by 2?

    idx = CartesianIndices(dims)
    wpix = findall(x->x==true, test)
    bpix = setdiff(idx, wpix)
    pix = (wpix, bpix)
    S2N = S2_initialise(pix, maxrng)
    S2 = naninf(S2_finalise(SN, S2N, test))
    #S2[length(S2)÷2:end] .= 0
    C2N = C2_initialise(test, maxrng)
    C2 = naninf(S2_finalise(SN, C2N, test))
end

guess, S2n, C2n, S2_BN1, C2_BN1, SN1 = histrecon_u(dims, C2, S2, philen);
heatmap(guess, color=:grays, aspect_ratio=1)
save("test38.png", colorview(Gray, guess))
imshow(guess)

#dims = (500,500)
#philen = Int(round(0.3*dims[1]*dims[1]))

#@time SN, BN = C2_init_partial(f, 100)
#@time subtract, add = update_C2_BN(BN, pix, original, modified, wpick, bpick)
#C2n, C2_BN = C2_pdist4(f, 100, SN)
a = [0]
dims = (5,5)
maxrng = Int64(round(sqrt(dims[1]^2 + dims[2]^2)+1))
philen = Int(round(0.3*dims[1]*dims[2]))
SN = SN_comp(dims, maxrng)

@time blas_stat5(vec(CartesianIndices(dims)), maxrng)

f, pix = make_rand_im(philen, dims)
heatmap(f, color=:grays, aspect_ratio=1)

    #pix = (wpix, bpix)
wpick = rand(1:length(pix[1]))
bpick = rand(1:length(pix[2]))
kern = CartesianIndices((-1:1, -1:1))

original = label_components(f, trues(3,3))
pix = swap_pix(f, pix, wpick, bpick)
modified = label_components(f, trues(3,3))
#a = pix

S2_BN = S2_initialise(pix, maxrng)
C2_BN = C2_initialise(f, maxrng)
S2n3 = S2_finalise(SN, S2_BN, f)
S2_finalise(SN, C2_BN, f)
S2_BN_d = update_S2_BN(S2_BN, pix, wpick, bpick, philen, maxrng)
C2_BN_d = update_C2_BN(C2_BN, pix, original, modified, wpick, bpick, maxrng)
push!(a, C2_BN_d[1])
blas_stat4(pix[1], 1.0, 100)

original = label_components(f, trues(3,3))
pix = swap_pix(f, pix, wpick, bpick)
modified = label_components(f, trues(3,3))
q = add - subtract
q = add - subtract
BN .= BN .+ q
#c1 = copy(change)
q = [0.0, 2.236068, 1.4142135, 1.0, 2.828427, 2.236068, 2.0, 0.0, 0.0, 1.0, 1.4142135, 1.0, 1.4142135, 0.0, 1.0, 2.0, 0.0, 2.236068, 1.4142135, 1.0, 0.0, 1.0, 1.4142135, 1.0, 1.4142135, 2.236068, 2.0, 2.236068, 0.0, 1.0, 0.0, 1.0, 2.0, 1.0, 1.4142135, 2.236068, 0.0, 1.0, 2.0, 1.0, 1.4142135, 2.236068, 2.0, 2.236068, 2.828427]
ash(q, rng=0:1.0:(maxrng-1), m =1).counts

@time C2_pdist4(f, 100, SN)
heatmap(original, color=:grays, aspect_ratio=1)
heatmap(modified, color=:grays, aspect_ratio=1)
imshow(original)
imshow(modified)
CartesianIndex(1,1)

a = collect(0:4).^2
b = a'
fit(Histogram, vec(sqrt.(a.+b)), 0:1.0:(maxrng)).weights


text:
using SearchDiscoveryReplication
using Test

# Required packages
using Distributions,Random

# Test sets
@testset "Reservation values" begin

	# Search value
	F = Distributions.DiscreteNonParametric([-1/2,1/2],[1/2,1/2])
	@test round(calcZs(0,F,0.1)[1] ,digits=5)== 0.3

	F = Distributions.Normal() ;
	@test round(calcZs(0,F,0.1)[1] ,digits=5)== 0.90235
	@test round.(calcZs([0,1],F,0.1) ,digits=5) == [0.90235,1.90235]

	# Expansion value
	nd = 1 ; cs = 0.1 ; cd = 0.1
	G = Distributions.Normal()
	@test round(calcZd(G,cs,cd,nd),digits=5) == 1.20125

	G = Distributions.Normal(0.2,0.4)
	@test round(calcZd(G,cs,cd,nd),digits=5) == 0.80748

	nd = 2 ; G = Distributions.Normal()
	@test round(calcZd(G,cs,cd,nd),digits=5) == 1.58767

end

@testset "Weitzman model" begin

	m = Weitzman(
		β = [2.0,-0.4], cs = 0.2 , cd = 0.4 ,
		nd = 1, nA0 = 2 , dChars =  MvNormal([1,1],1/2), outDummy=false,
		seed = 234,dV= Normal(0,0), dE = Normal(),cfun = "linear", smo = 25,
		options=estimOptions(nDraws =3)
		)
	nCons = 5
	nProdTot = 10
	nProdCons = 3
	d = genData(m,nCons,nProdTot,nProdCons )
	@test d.path[1] == [3]
	@test d.purch[1] == 3
	@test d.search[1:3] == [1,0,1]

	cfun = SearchDiscoveryReplication.getCfun(m.cfun)
	Random.seed!(m.seed)
	maxPos = maximum(d.pos)

	@test SearchDiscoveryReplication.negLogLik(m,vcat(m.β,m.cs,m.cd),length(m.β),d,cfun,maxPos) ≈ 14.5350 atol = 1e-4

	wd = calcWelfareDemand(m,d)
	@test wd[1]		 	≈ 11.8747 atol = 1e-4
	@test wd[2]	 		≈ 12.9414 atol = 1e-4
	@test wd[4][1]	  	≈ 0.1333 atol = 1e-4
	@test wd[4][end]	≈ 0.5333 atol = 1e-4
end

@testset "FullInfo model" begin
	m = FullInfo(β = [2.0,-0.4],dChars = MvNormal([1,1],1/2), outDummy = false, seed = 234,
					dV= Normal(0,0), dE = Normal(),smo=25,options=estimOptions(nDraws =3))
	nCons = 5
	nProdTot = 10
	nProdCons = 3
	d = genData(m,nCons,nProdTot,nProdCons )
	@test d.purch[1] == 3
	Random.seed!(m.seed)
	@test SearchDiscoveryReplication.negLogLik(m,m.β,d) ≈ 42.8430 atol = 1e-4

	wd = calcWelfareDemand(m,d)
	@test wd[1]		 	≈ 12.9784 atol = 1e-4
	@test wd[4][1]		≈ 0.0666 atol = 1e-4
	@test wd[4][end]	≈ 0.4 atol = 1e-4
end

@testset "RandomSearch model" begin
	m = RandomSearch(β = [2.0,-0.4],cs = 0.3,dChars = MvNormal([1,1],1/2), outDummy = false,
			seed = 234, dV= Normal(0,0), dE = Normal(),cfun = "linear",
			smo=5,options=estimOptions(nDraws =3))
	nCons = 5
	nProdTot = 10
	nProdCons = 3
	d = genData(m,nCons,nProdTot,nProdCons )
	@test d.purch[1] == 3
	@test isequal(d.path[1],[2,3])
	Random.seed!(m.seed)
	cfun = SearchDiscoveryReplication.getCfun(m.cfun)
	@test SearchDiscoveryReplication.negLogLik(m,vcat(m. β,m.cs),d,cfun) ≈ 11.1138 atol = 1e-4

	wd = calcWelfareDemand(m,d)
	@test wd[1]			≈ 10.2291 atol = 1e-4
	@test wd[2]		 	≈ 12.2291  atol = 1e-4
	@test wd[4][1]		≈ 0.1333 atol = 1e-4
	@test wd[4][end]	≈ 0.2666 atol = 1e-4
end

@testset "SearchDiscovery model" begin
	m = SearchDiscovery(β = [1.2,0.4],cs = 0.3,cd = 0.3,dze = -0.1, dChars = MvNormal([0.4,0.5],[0.8,1.2]),
		outDummy = false, seed = 23, nd = 1 , nA0 = 1, dV= Normal(0,1), dE = Normal(),cfun = "linear", zdfun = "linear",
		options=estimOptions(nDraws =200))

	nCons = 30
	nProdTot = 10
	nProdCons = 5
	d = genData(m,nCons,nProdTot,nProdCons)
	@test isequal(d.path[1],[2,3])
	@test d.purch[1] == 3
	@test d.stop[1] == 1

	wd = calcWelfareDemand(m,d)
	@test wd[1]			≈ 45.5269 atol = 1e-4
	@test wd[2]			≈ 65.0179 atol = 1e-4
	@test wd[4][1]		≈ 0.1866 atol = 1e-4
	@test wd[4][end]	≈ 0.1195 atol = 1e-4
end


text:
# This file was generated, do not modify it. # hide
DTR = @load DecisionTreeRegressor pkg=DecisionTree

boston = dataset("MASS", "Boston")

y, X = unpack(boston, ==(:MedV))

train, test = partition(eachindex(y), 0.5, shuffle=true, rng=551);

scitype(X)

max_issues_repo_name: alan-turing-institute/MLJTutorials

max_forks_repo_name: alan-turing-institute/MLJTutorials


text:
abstract type AbstractOracleCost{Mdl} end

oracle_stripe(args...; kwargs...) = oracle_stripe(NoHint(), args...; kwargs...)
oracle_stripe(::AbstractHint, args...; kwargs...) = @assert false
function oracle_stripe(hint::AbstractHint, mdl, A, Π; kwargs...)
    return oracle_stripe(hint, mdl, A; kwargs...)
end

function bound_stripe(A, K, Π, ocl::AbstractOracleCost)
    return bound_stripe(A, K, Π, oracle_model(ocl))
end

function bound_stripe(A, K, ocl::AbstractOracleCost)
    return bound_stripe(A, K, oracle_model(ocl))
end

function bound_stripe(A, K, Π, mdl)
    return bound_stripe(A, K, mdl)
end

@inline cost_type(mdl) = cost_type(typeof(mdl))
@inline cost_type(::Type{<:AbstractOracleCost{Mdl}}) where {Mdl} = cost_type(Mdl)
@inline objective_identity(::typeof(+), T) = zero(T)
@inline objective_identity(::typeof(max), T) = typemin(T)

@inline bottleneck_value(A, Π, Φ, mdl) = compute_objective(max, A, Π, Φ, mdl)
@inline bottleneck_value(A, Φ, mdl) = compute_objective(max, A, Φ, mdl)
@inline total_value(A, Π, Φ, mdl) = compute_objective(+, A, Π, Φ, mdl)
@inline total_value(A, Φ, mdl) = compute_objective(+, A, Φ, mdl)

@inline function compute_objective(g::G, A, Π, Φ::SplitPartition, mdl) where {G}
    return compute_objective(g, A, Π, Φ, oracle_stripe(StepHint(), mdl, A, Π))
end
@inline function compute_objective(g::G, A, Π, Φ, mdl) where {G}
    Φ_dom = convert(DomainPartition, Φ)
    A_prm = A[:, Φ_dom.prm]
    Φ_spl = SplitPartition(length(Φ), Φ_dom.spl)
    return compute_objective(g, A_prm, Π, Φ_spl, mdl)
end
function compute_objective(g::G, A, Π, Φ::SplitPartition, mdl::AbstractOracleCost) where {G}
    cst = objective_identity(g, cost_type(mdl))
    for k = 1:Φ.K
        j = Φ.spl[k]
        j′ = Φ.spl[k + 1]
        cst = g(cst, mdl(j, j′, k))
    end
    return cst
end
@inline function compute_objective(g::G, A, Φ::SplitPartition, mdl) where {G}
    return compute_objective(g, A, Φ, oracle_stripe(StepHint(), mdl, A))
end
@inline function compute_objective(g::G, A, Φ, mdl) where {G}
    Φ_dom = convert(DomainPartition, Φ)
    A_prm = A[:, Φ_dom.prm]
    Φ_spl = SplitPartition(length(Φ), Φ_dom.spl)
    return compute_objective(g, A_prm, Φ_spl, mdl)
end
function compute_objective(g, A, Φ::SplitPartition, mdl::AbstractOracleCost) where {G}
    cst = objective_identity(g, cost_type(mdl))
    for k = 1:Φ.K
        j = Φ.spl[k]
        j′ = Φ.spl[k + 1]
        cst = g(cst, mdl(j, j′, k))
    end
    return cst
end

function compute_weight(A, weight)
    w = oracle_stripe(StepHint(), weight, A)
    m, n = size(A)
    wgt = undefs(cost_type(w), n)
    for j = 1:n
        wgt[j] = Step(w)(Same(j), Next(j + 1))
        Step(w)(Next(j), Same(j + 1))
    end
    return wgt
end

struct Extended{T}
    i::Bool
    x::T
end

infinity(::Type{Float16}) = Inf16
infinity(::Type{Float32}) = Inf32
infinity(::Type{Float64}) = Inf64
infinity(::Type{Extended{T}}) where {T} = Extended(true, zero(T))
infinity(x) = infinity(typeof(x))
infinity(T::Type) = Extended(true, zero(T))

extend(T::Type) = typeof(infinity(T))
#TODO use convert here
extend(x) = (x isa extend(typeof(x))) ? x : Extended(false, x)

#Base.zero(::Extended{T}) where {T} = extend(zero(T))
Base.zero(::Type{Extended{T}}) where {T} = extend(zero(T))
Base.:+(a::Extended, b::Extended) = Extended(a.i | b.i, a.x + b.x)
Base.:*(a::Extended, b::Extended) = Extended(a.i | b.i, a.x * b.x)
Base.typemax(::Type{Extended{T}}) where {T} = infinity(T)
Base.typemin(::Type{Extended{T}}) where {T} = extend(typemin(T))
Base.:<(a::Extended, b::Extended) = (!a.i && b.i) || ((!a.i && !b.i) && (a.x < b.x))
Base.:(==)(a::Extended, b::Extended) = (a.i && b.i) || ((!a.i && !b.i) && (a.x == b.x))
Base.isapprox(a::Extended, b::Extended; kwargs...) = (a.i && b.i) || ((!a.i && !b.i) && isapprox(a.x, b.x; kwargs...))

struct ConstrainedCost{F, W, Tw}
    f::F
    w::W
    w_max::Tw
    function ConstrainedCost{F, W, Tw}(f::F, w::W, w_max) where {F, W, Tw}
        @assert Tw == cost_type(W)
        return new{F, W, Tw}(f, w, w_max)
    end
end

ConstrainedCost(f::F, w::W, w_max) where {F, W} = ConstrainedCost{F, W, cost_type(W)}(f, w, w_max)

cost_type(::Type{ConstrainedCost{F, W, Tw}}) where {F, W, Tw} = extend(cost_type(F))

struct ConstrainedCostOracle{F, W, Tw} <: AbstractOracleCost{ConstrainedCost{F, W, Tw}}
    f::F
    w::W
    w_max::Tw
    function ConstrainedCostOracle{F, W, Tw}(f::F, w::W, w_max) where {F, W, Tw}
        @assert Tw == cost_type(W)
        return new{F, W, Tw}(f, w, w_max)
    end
end

ConstrainedCostOracle(f::F, w::W, w_max) where {F, W} = ConstrainedCostOracle{F, W, cost_type(W)}(f, w, w_max)

oracle_model(ocl::ConstrainedCostOracle) = ConstrainedCost(oracle_model(ocl.f), oracle_model(ocl.w), ocl.w_max)

function oracle_stripe(hint::AbstractHint, cst::ConstrainedCost{F, W, Tw}, A::SparseMatrixCSC, args...; kwargs...) where {F, W, Tw}
    (m, n) = size(A)
    f = oracle_stripe(hint, cst.f, A, args...; kwargs...)
    w = oracle_stripe(hint, cst.w, A, args...; kwargs...)
    w_max = cst.w_max
    return ConstrainedCostOracle(f, w, w_max)
end

@inline function (ocl::ConstrainedCostOracle)(j::Ti, j′::Ti, k...) where {Ti}
    if ocl.w(j, j′, k...) <= ocl.w_max
        return extend(ocl.f(j, j′, k...))
    else
        return infinity(cost_type(ocl.f))
    end
end

#TODO not sure about these two...
bound_stripe(A::SparseMatrixCSC, K, cst::ConstrainedCost) = bound_stripe(A, K, ocl.f)
bound_stripe(A::SparseMatrixCSC, K, ocl::ConstrainedCostOracle) = bound_stripe(A, K, ocl.f)

struct Infeasible end

struct Feasible end

Base.:<(::Feasible, ::Feasible) = false
Base.:<(::Feasible, ::Infeasible) = true
Base.:<(::Infeasible, ::Infeasible) = false
Base.:<(::Infeasible, ::Feasible) = false
Base.:(==)(::Feasible, ::Feasible) = true
Base.:(==)(::Feasible, ::Infeasible) = false
Base.:(==)(::Infeasible, ::Feasible) = false
Base.:(==)(::Infeasible, ::Infeasible) = true

struct FeasibleCost end

@inline (::FeasibleCost)(j, j′, k...) = Feasible()
@inline cost_type(::Type{FeasibleCost}) = Feasible
oracle_model(::FeasibleCost) = FeasibleCost()
oracle_stripe(::AbstractHint, ::FeasibleCost, ::SparseMatrixCSC; kwargs...) = FeasibleCost()
#bound_stripe(A::SparseMatrixCSC, K, ::FeasibleCost) = (Feasible(), Feasible()) #TODO ?

struct Next{T}
    arg::T
end
@inline destep(arg::Next) = arg.arg
struct Same{T}
    arg::T
end
@inline destep(arg::Same) = arg.arg
struct Prev{T}
    arg::T
end
@inline destep(arg::Prev) = arg.arg
struct Jump{T}
    arg::T
end
@inline destep(arg::Jump) = arg.arg

struct Step{Ocl}
    ocl::Ocl
end

@propagate_inbounds (ocl::Step{Ocl})(args...) where {Ocl} = ocl.ocl(maptuple(destep, args...)...)
#function (ocl::Step{Ocl})(args...) where {Ocl}
#    @info "whoopsies" ocl.ocl args
#    assert false
#end


text:
module TestWorld
  using World
  using Helper
  using FactCheck
  import Dots

  facts("Tests world creation") do
    @fact size(World.create(10, 15).data) == (15, 10) --> true
  end
  facts("Tests world creation with zero energy on the plain") do
    plane = World.create(1, 1)
    @fact World.getEnergy(plane, Helper.Point(1,1)) --> Dots.INDEX_EMPTY
    World.destroy(plane)
  end
  # TODO: add test of setting energy into existing dot
  facts("Tests energy reset") do
    plane = World.create(10, 10)
    @fact World.setEnergy(plane, Helper.Point(5, 3), UInt16(12))  --> true
    @fact World.setEnergy(plane, Helper.Point(5, 3), UInt16(123)) --> true
    @fact World.getEnergy(plane, Helper.Point(5, 3)) --> UInt16(123)
    World.destroy(plane)
  end
  facts("Tests energy set to zero") do
    plane = World.create(10, 10)
    @fact World.setEnergy(plane, Helper.Point(5, 3), UInt16(12)) --> true
    @fact World.setEnergy(plane, Helper.Point(5, 3), UInt16(0)) --> true
    @fact World.getEnergy(plane, Helper.Point(5, 3)) --> UInt16(0)
    World.destroy(plane)
  end
  facts("Tests world correct dimentions") do
    plane = World.create(15, 10)
    World.setEnergy(plane, Helper.Point(15, 10), UInt16(123))
    @fact World.getEnergy(plane, Helper.Point(15, 10)) --> UInt16(123)
    World.destroy(plane)
  end
  facts("Tests world correct energy set and get within and outside plane") do
    plane = World.create(15, 10)
    World.setEnergy(plane, Helper.Point(0,   0 ), UInt16(123))
    World.setEnergy(plane, Helper.Point(1,   1 ), UInt16(123))
    World.setEnergy(plane, Helper.Point(15,  1 ), UInt16(123))
    World.setEnergy(plane, Helper.Point(0,   10), UInt16(123))
    World.setEnergy(plane, Helper.Point(10,  15), UInt16(123))
    World.setEnergy(plane, Helper.Point(-10,-15), UInt16(123))
    World.setEnergy(plane, Helper.Point(10,  10), UInt16(123))

    @fact World.getEnergy(plane, Helper.Point(0,   0 )) --> UInt16(0)
    @fact World.getEnergy(plane, Helper.Point(1,   1 )) --> UInt16(123)
    @fact World.getEnergy(plane, Helper.Point(15,  1 )) --> UInt16(123)
    @fact World.getEnergy(plane, Helper.Point(0,   10)) --> UInt16(0)
    @fact World.getEnergy(plane, Helper.Point(10,  15)) --> UInt16(0)
    @fact World.getEnergy(plane, Helper.Point(-10,-15)) --> UInt16(0)
    @fact World.getEnergy(plane, Helper.Point(10,  10)) --> UInt16(123)

    World.destroy(plane)
  end
  facts("Tests getting free position") do
    plane = World.create(2, 1)
    World.setEnergy(plane, Helper.Point(1, 1), UInt16(1))
    pos = World.getFreePos(plane)
    @fact (pos.x === 2 && pos.y === 1) --> true
    #every point are full
    World.setEnergy(plane, Helper.Point(2, 1), UInt16(1))
    @fact World.getFreePos(plane) --> false

    World.destroy(plane)
  end
  facts("Tests grabbing energy") do
    plane = World.create(3, 3)
    World.setEnergy(plane, Helper.Point(2, 2), UInt16(10))
    @fact World.grabEnergy(plane, Helper.Point(2, 2), UInt16(5)) --> UInt16(5)
    @fact World.getEnergy(plane, Helper.Point(2, 2)) --> UInt16(5)
    @fact World.grabEnergy(plane, Helper.Point(2, 2), UInt16(10)) --> UInt16(5)
    @fact World.getEnergy(plane, Helper.Point(2, 2)) --> UInt16(0)
    @fact World.grabEnergy(plane, Helper.Point(2, 2), UInt16(5)) --> UInt16(0)
    World.destroy(plane)
  end
  facts("Tests near free points") do
    plane = World.create(3, 3)
    #
    # 0 1 1
    # 1   1
    # 1 1 1
    #
    World.setEnergy(plane, Helper.Point(1, 1), UInt16(0))
    World.setEnergy(plane, Helper.Point(2, 1), UInt16(1))
    World.setEnergy(plane, Helper.Point(3, 1), UInt16(1))
    World.setEnergy(plane, Helper.Point(1, 2), UInt16(1))
    World.setEnergy(plane, Helper.Point(3, 2), UInt16(1))
    World.setEnergy(plane, Helper.Point(1, 3), UInt16(1))
    World.setEnergy(plane, Helper.Point(2, 3), UInt16(1))
    World.setEnergy(plane, Helper.Point(3, 3), UInt16(1))
    pos = World.getNearFreePos(plane, Helper.Point(2, 2))
    @fact (pos.x === 1 && pos.y === 1) --> true
    #
    # 1 0 1
    # 1   1
    # 1 1 1
    #
    World.setEnergy(plane, Helper.Point(1, 1), UInt16(1))
    World.setEnergy(plane, Helper.Point(2, 1), UInt16(0))
    World.setEnergy(plane, Helper.Point(3, 1), UInt16(1))
    World.setEnergy(plane, Helper.Point(1, 2), UInt16(1))
    World.setEnergy(plane, Helper.Point(3, 2), UInt16(1))
    World.setEnergy(plane, Helper.Point(1, 3), UInt16(1))
    World.setEnergy(plane, Helper.Point(2, 3), UInt16(1))
    World.setEnergy(plane, Helper.Point(3, 3), UInt16(1))
    pos = World.getNearFreePos(plane, Helper.Point(2, 2))
    @fact (pos.x === 2 && pos.y === 1) --> true
    #
    # 1 1 0
    # 1   1
    # 1 1 1
    #
    World.setEnergy(plane, Helper.Point(1, 1), UInt16(1))
    World.setEnergy(plane, Helper.Point(2, 1), UInt16(1))
    World.setEnergy(plane, Helper.Point(3, 1), UInt16(0))
    World.setEnergy(plane, Helper.Point(1, 2), UInt16(1))
    World.setEnergy(plane, Helper.Point(3, 2), UInt16(1))
    World.setEnergy(plane, Helper.Point(1, 3), UInt16(1))
    World.setEnergy(plane, Helper.Point(2, 3), UInt16(1))
    World.setEnergy(plane, Helper.Point(3, 3), UInt16(1))
    pos = World.getNearFreePos(plane, Helper.Point(2, 2))
    @fact (pos.x === 3 && pos.y === 1) --> true
    #
    # 1 1 1
    # 0   1
    # 1 1 1
    #
    World.setEnergy(plane, Helper.Point(1, 1), UInt16(1))
    World.setEnergy(plane, Helper.Point(2, 1), UInt16(1))
    World.setEnergy(plane, Helper.Point(3, 1), UInt16(1))
    World.setEnergy(plane, Helper.Point(1, 2), UInt16(0))
    World.setEnergy(plane, Helper.Point(3, 2), UInt16(1))
    World.setEnergy(plane, Helper.Point(1, 3), UInt16(1))
    World.setEnergy(plane, Helper.Point(2, 3), UInt16(1))
    World.setEnergy(plane, Helper.Point(3, 3), UInt16(1))
    pos = World.getNearFreePos(plane, Helper.Point(2, 2))
    @fact (pos.x === 1 && pos.y === 2) --> true
    #
    # 1 1 1
    # 1   0
    # 1 1 1
    #
    World.setEnergy(plane, Helper.Point(1, 1), UInt16(1))
    World.setEnergy(plane, Helper.Point(2, 1), UInt16(1))
    World.setEnergy(plane, Helper.Point(3, 1), UInt16(1))
    World.setEnergy(plane, Helper.Point(1, 2), UInt16(1))
    World.setEnergy(plane, Helper.Point(3, 2), UInt16(0))
    World.setEnergy(plane, Helper.Point(1, 3), UInt16(1))
    World.setEnergy(plane, Helper.Point(2, 3), UInt16(1))
    World.setEnergy(plane, Helper.Point(3, 3), UInt16(1))
    pos = World.getNearFreePos(plane, Helper.Point(2, 2))
    @fact (pos.x === 3 && pos.y === 2) --> true
    #
    # 1 1 1
    # 1   1
    # 0 1 1
    #
    World.setEnergy(plane, Helper.Point(1, 1), UInt16(1))
    World.setEnergy(plane, Helper.Point(2, 1), UInt16(1))
    World.setEnergy(plane, Helper.Point(3, 1), UInt16(1))
    World.setEnergy(plane, Helper.Point(1, 2), UInt16(0))
    World.setEnergy(plane, Helper.Point(3, 2), UInt16(1))
    World.setEnergy(plane, Helper.Point(1, 3), UInt16(0))
    World.setEnergy(plane, Helper.Point(2, 3), UInt16(1))
    World.setEnergy(plane, Helper.Point(3, 3), UInt16(1))
    pos = World.getNearFreePos(plane, Helper.Point(2, 2))
    @fact (pos.x === 1 && pos.y === 3) --> true
    #
    # 1 1 1
    # 1   1
    # 1 0 1
    #
    World.setEnergy(plane, Helper.Point(1, 1), UInt16(1))
    World.setEnergy(plane, Helper.Point(2, 1), UInt16(1))
    World.setEnergy(plane, Helper.Point(3, 1), UInt16(1))
    World.setEnergy(plane, Helper.Point(1, 2), UInt16(1))
    World.setEnergy(plane, Helper.Point(3, 2), UInt16(1))
    World.setEnergy(plane, Helper.Point(1, 3), UInt16(1))
    World.setEnergy(plane, Helper.Point(2, 3), UInt16(0))
    World.setEnergy(plane, Helper.Point(3, 3), UInt16(1))
    pos = World.getNearFreePos(plane, Helper.Point(2, 2))
    @fact (pos.x === 2 && pos.y === 3) --> true
    #
    # 1 1 1
    # 1   1
    # 1 1 0
    #
    World.setEnergy(plane, Helper.Point(1, 1), UInt16(1))
    World.setEnergy(plane, Helper.Point(2, 1), UInt16(1))
    World.setEnergy(plane, Helper.Point(3, 1), UInt16(1))
    World.setEnergy(plane, Helper.Point(1, 2), UInt16(0))
    World.setEnergy(plane, Helper.Point(3, 2), UInt16(1))
    World.setEnergy(plane, Helper.Point(1, 3), UInt16(1))
    World.setEnergy(plane, Helper.Point(2, 3), UInt16(1))
    World.setEnergy(plane, Helper.Point(3, 3), UInt16(0))
    pos = World.getNearFreePos(plane, Helper.Point(2, 2))
    @fact (pos.x === 3 && pos.y === 3) --> true

    World.destroy(plane)
  end
end

max_stars_repo_name: tmptrash/jevo

max_issues_repo_name: tmptrash/jevo

max_forks_repo_name: tmptrash/jevo


text:
module tf2_ros

    using RobotOS
    @rosimport geometry_msgs.msg: TransformStamped
    rostypegen(@__MODULE__)
    import .geometry_msgs.msg: TransformStamped

    using PyCall
    import Base.==

    const __tf__ = PyCall.PyNULL()
    function __init__()
        copy!(__tf__, pyimport("tf2_ros"))
    end

    export Buffer, TransformBroadcaster
    export sendTransform, lookupTransform, waitForTransform

    """
        Extending the equality operator (==) functionality to the TransformStamped
    type. Useful for testing.
    """
    function Base.:(==)(tf1::TransformStamped, tf2::TransformStamped)
        transform1 = tf1.transform; transform2 = tf1.transform

        if all(transform1.translation.x == transform2.translation.x,
               transform1.translation.y == transform2.translation.y,
               transform1.translation.z == transform2.translation.z,
               transform1.rotation.w == transform2.rotation.w,
               transform1.rotation.x == transform2.rotation.x,
               transform1.rotation.y == transform2.rotation.y,
               transform1.rotation.z == transform2.rotation.z)
            return true
        else
            return false
        end
    end

# ============================================================================#
#               tf2_ros TransformBroadcaster object and functions
# ============================================================================#
    """
        TransformBroadcaster()
    Create a transform broadcaster object.
    """
    struct TransformBroadcaster
        o::PyObject
        function TransformBroadcaster()
            new(__tf__.TransformBroadcaster())
        end
    end

    """
        sendTransform(tf_broadcaster_obj, transform, time, child_frame, parent_frame)
    Broadcast the transformation from tf frame child to parent on ROS topic "/tf".
    """
    function sendTransform(tb::TransformBroadcaster,
                           transform::Any)
        pycall(tb.o.sendTransform, PyAny, transform)
    end

    """
    Publish set of transforms
    """
    function sendTransform(tb::TransformBroadcaster,
                           transforms::Vector{Any})
        pycall(tb.o.sendTransform, PyAny, transforms)
    end

    # ============================================================================#
    #                       tf2_ros Buffer object and functions
    # ============================================================================#
    """
        Buffer(; cache_time=Duration(1000))
    Create a Buffer object, which holds on to transforms for time cache_time
    """
    struct Buffer
        o::PyObject
        function Buffer(; cache_time=Duration(1000))
            pyCache_time = convert(PyObject, cache_time)
            new(__tf__.Buffer(cache_time=pyCache_time))
        end
    end

    """
        generate_error_message(err)
    Return error message string which includes both exception type and error massage information.
    """
    function generate_error_message(err)
        exception_type = err.T.__name__
        error_message = exception_type * ": $(err.val.args[1])"
    end

    """
        lookupTransform(tf_listener_obj, target, source, time)
    Return tuple of (position, quaternion).
    """
    function lookupTransform(tf_buffer::Buffer,
                             target_frame::AbstractString,
                             source_frame::AbstractString,
                             time::Time,
                             timeout::Duration)
        pytime = convert(PyObject, time)
        pytimeout = convert(PyObject, timeout)
        try
            transStamp = pycall(tf_buffer.o.lookup_transform, PyAny,
                                target_frame, source_frame, pytime, pytimeout)
            return transStamp
        catch err
            if isa(err, PyCall.PyError)
               error_message = generate_error_message(err)
               error(error_message)
            else
               rethrow(err)
            end
        end
    end

end


text:
#===============================================================================

    UniqCombinCenters.jl

    Author: Tom Price
    Date:   July 2019

===============================================================================#

import Clustering.ClusteringResult

mutable struct UniqCombinCenters
    uniq_combin_centers ::Set{Vector{Float_T}}
    car                 ::ClusterAnalysisResult
    slht_mean           ::Float_T
    geno_combins        ::Vector{Matrix{Float_T}}
end

const EMPTY_UCC_DICT = OrderedDict{Set{Vector{Float_T}},UniqCombinCenters}()

max_stars_repo_name: MakerButt/chaipcr

max_issues_repo_name: MakerButt/chaipcr

max_forks_repo_name: MakerButt/chaipcr


text:
let
    cost_model = SumOfMakeSpans()
    for (i, f) in enumerate(pctapf_test_problems())
        solver = NBSSolver(
            assignment_model = TaskGraphsMILPSolver(GreedyAssignment()),
            path_planner = PIBTPlanner{NTuple{3,Float64}}()
            )
        let
            pc_tapf = f(solver;cost_function=cost_model,verbose=false)
            base_search_env = pc_tapf.env
            prob = formulate_assignment_problem(solver.assignment_model,pc_tapf)
            sched, _ = solve_assignment_problem!(solver.assignment_model,prob,pc_tapf)

            pibt_planner = route_planner(solver)
            search_env = construct_search_env(solver,deepcopy(sched),base_search_env)
            pc_mapf = PC_MAPF(search_env)
            set_iteration_limit!(pibt_planner,50)
            set_verbosity!(pibt_planner,0)
            solution, valid_flag = pibt!(pibt_planner,deepcopy(pc_mapf))
            # @show i, f, get_cost(solution)
            # @show convert_to_vertex_lists(solution.route_plan)
            @test valid_flag

            reset_solver!(solver)
            solution, cost = solve!(pibt_planner,pc_mapf)
            @test cost != typemax(cost_type(pc_mapf))
        end
    end
end
let
    solver = NBSSolver(
        assignment_model = TaskGraphsMILPSolver(GreedyAssignment()),
        path_planner = PIBTPlanner{NTuple{3,Float64}}(partial=true)
        )
    set_iteration_limit!(solver.path_planner,5)
    rpctapf = replanning_problem_1(solver)
    replan_model = MergeAndBalance()

    base_env = deepcopy(rpctapf.env)
    env = replan!(solver,replan_model,base_env,rpctapf.requests[1])
    pc_tapf = PC_TAPF(env)
    reset_solver!(solver)
    solution, cost = solve!(solver,pc_tapf)
    fix_precutoff_nodes!(solution)
    reset_solver!(solver)
    solution2, cost = solve!(solver,PC_TAPF(solution))

    for v in solution2.cache.active_set
        node = get_node_from_vtx(solution2.schedule,v)
        @test isa(node,GO)
        @test get_destination_location_id(node) == LocationID(-1)
    end

    # prob = formulate_assignment_problem(solver.assignment_model,pc_tapf)
    # sched, _ = solve_assignment_problem!(solver.assignment_model,prob,pc_tapf)
    #
    # pibt_planner = route_planner(solver)
    # pc_mapf = PC_MAPF(construct_search_env(solver,deepcopy(sched),pc_tapf.env))
    # set_iteration_limit!(pibt_planner,5)
    # set_verbosity!(pibt_planner,0)
    # solution, valid_flag = pibt!(pibt_planner,pc_mapf);
    #
    # remap_object_ids!(rpctapf.requests[2].schedule,solution.schedule)
    # env = replan!(solver,replan_model,solution,rpctapf.requests[2])
    # # update_planning_cache!(solver,env)
    # reset_solver!(solver)
    # solution2, cost = solve!(solver,PC_TAPF(env))

end

# function show_cache_state(env,cache=env.cache)
#     for (name,c) in zip(["closed","active"],[cache.closed_set, cache.active_set])
#        vtxs = map(v->string(get_node_from_vtx(env.schedule,v)), collect(c))
#        println("***",name,"***")
#        for node in vtxs
#            println(node)
#        end
#        println("\n")
#    end
#    println("*** FIXED ***")
#    for v in vertices(env.schedule)
#        path_spec = get_path_spec(env.schedule,v)
#        pre = path_spec.fixed ? " - FIXED" : ""
#        println(string(get_node_from_vtx(env.schedule,v)), pre)
#    end
# end


text:
function compute_ideal_step_size(new_objective, new_point, min_value, solution)
    distance_to_solution = LinearAlgebra.norm(new_point - solution)
    obj_gap = (new_objective - min_value)
    if obj_gap > 0.0
        return distance_to_solution^2 / obj_gap
    end
    return 0.0
end

function compute_step_size(new_objective, min_value, diameter)
    return (new_objective - min_value) / diameter^2
end

function create_subgradient_method_parameters(;
    eps_optimal = 1e-6,
    iteration_limit = 100,
    printing_frequency = 50,
    verbose = false,
)
    return ProximalBundleMethod.SubgradientMethodParams(
        eps_optimal,
        iteration_limit,
        verbose,
        printing_frequency,
    )
end

function create_bundle_method_parameters(;
    eps_optimal = 1e-6,
    iteration_limit = 100,
    contraction_factor = 0.5,
    printing_frequency = 50,
    verbose = false,
    full_memory = false,
)
    return ProximalBundleMethod.BundleMethodParams(
        eps_optimal,
        iteration_limit,
        contraction_factor,
        verbose,
        printing_frequency,
        full_memory,
    )
end


text:
using BinDeps

@BinDeps.setup

# set to true to support intel fortran compiler
useIntelFortran = false

# construct absolute path
depsdir  = splitdir(Base.source_path())[1]
builddir = joinpath(depsdir,"builds")
srcdir   = joinpath(depsdir,"src")

println("=== Building ParSpMatVec ===")
println("depsdir  = $depsdir")
println("builddir = $builddir")
println("srcdir   = $srcdir")
println("useIntel = $useIntelFortran")

if !isdir(builddir)
	println("creating build directory")
	mkdir(builddir)
	if !isdir(builddir)
		error("Could not create build directory")
	end
end

@static if is_unix()

	src1 = joinpath(srcdir,"A_mul_B.f90")
	src2 = joinpath(srcdir,"Ac_mul_B.f90")
	src3 = joinpath(srcdir,"a_add_b.f90")
	outfile = joinpath(builddir,"spmatvec")
	@build_steps begin
		if useIntelFortran
			run(`ifort -O3 -xHost -fPIC -fpp -openmp -integer-size 64 -diag-disable=7841 -shared  $src1 $src2 $src3 -o $outfile`)
		else
			println("fortran version")
			run(`gfortran --version`)
			run(`gfortran -O3 -fPIC -cpp -fopenmp -fdefault-integer-8 -shared  $src1 $src2 $src3 -o $outfile`)
		end
	end
end

@static if is_windows()
	src1 = joinpath(srcdir,"A_mul_B.f90")
	src2 = joinpath(srcdir,"Ac_mul_B.f90")
	outfile = joinpath(builddir,"spmatvec.dll")
	@build_steps begin
		println("fortran version")
		run(`gfortran --version`)
		run(`gfortran -O3 -cpp -fopenmp -fdefault-integer-8 -shared -DBUILD_DLL  $src1 $src2 $src3 -o $outfile`)
	end
end

max_stars_repo_name: WneleiGao/ElasticWave


text:
@testset "intrinsiccoregion" begin
    rng = MersenneTwister(123)

    dims = (in=3, out=2, obs=3)
    rank = 1

    A = randn(dims.out, rank)
    B = A * transpose(A) + Diagonal(rand(dims.out))

    X = [(rand(dims.in), rand(1:(dims.out))) for i in 1:(dims.obs)]

    kernel = SqExponentialKernel()
    icoregionkernel = IntrinsicCoregionMOKernel(; kernel=kernel, B=B)

    @test icoregionkernel.B == B
    @test icoregionkernel.kernel == kernel
    @test icoregionkernel(X[1], X[1]) ≈ B[X[1][2], X[1][2]] * kernel(X[1][1], X[1][1])
    @test icoregionkernel(X[1], X[end]) ≈ B[X[1][2], X[end][2]] * kernel(X[1][1], X[end][1])

    KernelFunctions.TestUtils.test_interface(
        icoregionkernel, Vector{Tuple{Float64,Int}}; dim_out=dims.out
    )

    test_ADs(icoregionkernel; dims=dims)

    @test string(icoregionkernel) ==
          string("Intrinsic Coregion Kernel: ", kernel, " with ", dims.out, " outputs")
end


text:
#=

julia julia_callback.jl
  2.457 ms (100002 allocations: 10.68 MiB)

=#
using BenchmarkTools

function rober(t, u)
    k1 = 0.04
    k2 = 3e7
    k3 = 1e4
    y1, y2, y3 = u
    dy1 = -k1 * y1 + k3 * y2 * y3
    dy2 = k1 * y1 - k2 * y2 * y2 - k3 * y2 * y3
    dy3 = k2 * y2 * y2
    return [dy1, dy2, dy3]
end

rober(0, [1.0, 0.0, 0.0])

function call_function(func)
    u = [1.0, 0.0, 0.0]
    dt = 0.1
    times = dt * (0:1e5)
    for time in times
        u = func(time, u)
    end
end

@btime call_function(rober)


text:
function serve(project, dstdir = Path(mktempdir()); format = HTML())
    foreach(reset!, project.rewriters)
    builder = FileBuilder(format, dstdir)
    Pollen.build(project, builder)

    docupdates = Dict{AbstractPath, XNode}
    updatelock = ReentrantLock()
    running = Ref(true)
    # Loop that processes updated/created documents in batches
    @async begin
        while running[]
            lock(updatelock) do
                docs = copy(docupdates)
                foreach(k -> delete!(docupdates, k), keys(docupdates))
            end
            build(project, builder, docs)

            sleep(1/60)
        end
    end

    handlers = filehandlers(project, builder)
    updatech = Channel()
    watcher = watchfiles(handlers, updatech)
    try
        LiveServer.serve(dir=string(dstdir))
    catch e
        rethrow(e)
    finally
        LiveServer.stop(watcher)
        running[] = false
    end

end

function getwatcher(project, dir, builder)
    handlers = getfilehandlers(project, dir, builder)
    return watchfiles(handlers)
end

function watchfiles(handlers, ch)
    watcher = LiveServer.SimpleWatcher() do file
        p = Path(file)
        handlerfns = get(handlers, p, ())
        for f in handlerfns
            f()
        end
    end
    for p in keys(handlers)
        LiveServer.watch_file!(watcher, string(absolute(p)))
    end
    LiveServer.start(watcher)
    return watcher
end

function filehandlers(project::Project, builder::Builder)
    handlers = Dict{AbstractPath, Any}()
    for rewriter in project.rewriters
        for (p, handlerfn) in filehandlers(rewriter, project, builder)
            phandlers = get!(handlers, p, [])
            push!(phandlers, handlerfn)
        end
    end
    return handlers
end

filehandlers(::Rewriter, project, builder) = Dict()


text:
# Autogenerated wrapper script for czmq_jll for armv7l-linux-gnueabihf
export libczmq

using ZeroMQ_jll
JLLWrappers.@generate_wrapper_header("czmq")
JLLWrappers.@declare_library_product(libczmq, "libczmq.so.4")
function __init__()
    JLLWrappers.@generate_init_header(ZeroMQ_jll)
    JLLWrappers.@init_library_product(
        libczmq,
        "lib/libczmq.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
@derived_dimension ElectricField Unitful.𝐈^-1*Unitful.𝐋*Unitful.𝐌*Unitful.𝐓^-3
@derived_dimension Intensity Unitful.𝐌*Unitful.𝐓^-3
@derived_dimension InvIntensity Unitful.𝐌^-1*Unitful.𝐓^3

atomic_units(I::Intensity) = I/(3.5094452e16*u"W"/(u"cm"^2)) |> NoUnits
atomic_units(iI::InvIntensity) = iI*3.5094452e16*u"W"/(u"cm"^2) |> NoUnits
atomic_units(E::ElectricField) = E/(5.14220651e11*u"V"/u"m") |> NoUnits
atomic_units(q) = austrip(q)


text:
@testset "posterior_lti_sde" begin
    rng = MersenneTwister(123456)
    N = 13
    Npr = 15

    N = 3
    Npr = 2

    kernels = vcat(

        # Base kernels.
        (name="base-Matern12Kernel", val=Matern12Kernel()),
        map([Matern32Kernel, Matern52Kernel]) do k
            (name="base-$k", val=k())
        end,

        # Scaled kernels.
        map([1e-1, 1.0, 10.0, 100.0]) do σ²
            (name="scaled-σ²=$σ²", val=σ² * Matern32Kernel())
        end,

        # Stretched kernels.
        map([1e-2, 0.1, 1.0, 10.0, 100.0]) do λ
            (name="stretched-λ=$λ", val=Matern32Kernel() ∘ ScaleTransform(λ))
        end,

        # Summed kernels.
        (
            name="sum-Matern12Kernel-Matern32Kernel",
            val=1.5 * Matern12Kernel() ∘ ScaleTransform(0.1) +
                0.3 * Matern32Kernel() ∘ ScaleTransform(1.1),
        ),
    )

    # Construct a Gauss-Markov model with either dense storage or static storage.
    storages = (
        (name="dense storage Float64", val=ArrayStorage(Float64)),
        (name="static storage Float64", val=SArrayStorage(Float64)),
    )

    # Either regular spacing or irregular spacing in time.
    ts = (
        (name="irregular spacing", val=collect(RegularSpacing(0.0, 0.3, N))),
        (name="regular spacing", val=RegularSpacing(0.0, 0.3, N)),
    )

    σ²s = (
        (name="homoscedastic noise", val=(0.1, ),),
        (name="heteroscedastic noise", val=(rand(rng, N) .+ 1e-1, )),
    )

    @testset "$(kernel.name), $(storage.name), $(t.name), $(σ².name)" for
        kernel in kernels,
        storage in storages,
        t in ts,
        σ² in σ²s

        println("$(kernel.name), $(storage.name), $(t.name), $(σ².name)")

        # Construct Gauss-Markov model.
        f_naive = GP(kernel.val)
        fx_naive = f_naive(collect(t.val), σ².val...)

        f = to_sde(f_naive, storage.val)
        fx = f(t.val, σ².val...)
        model = build_lgssm(fx)

        validate_dims(model)

        y = rand(rng, fx)

        x_pr = rand(rng, Npr) * (maximum(t.val) - minimum(t.val)) .+ minimum(t.val)

        f_post_naive = posterior(fx_naive, y)
        f_post = posterior(fx, y)

        post_obs_var = 0.3
        fx_post_naive = f_post_naive(x_pr, post_obs_var)
        fx_post = f_post(x_pr, post_obs_var)
        y_post = rand(fx_post)

        @test mean.(marginals(fx_post_naive)) ≈ mean.(marginals(fx_post)) rtol=1e-5
        @test std.(marginals(fx_post_naive)) ≈ std.(marginals(fx_post)) rtol=1e-5
        m_and_v = mean_and_var(fx_post)
        @test m_and_v[1] ≈ mean.(marginals(fx_post))
        @test m_and_v[2] ≈ var.(marginals(fx_post))
        @test mean(fx_post) ≈ m_and_v[1] rtol=1e-5
        @test var(fx_post) ≈ m_and_v[2] rtol=1e-5
        @test logpdf(fx_post_naive, y_post) ≈ logpdf(fx_post, y_post) rtol=1e-5
    end
end


text:
@testset "misc/" begin

    include("insert_some_elements.jl")

end