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]...)

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

"""
    SamePad()

Passed as an option to convolutional layers (and friends), this causes
the padding to be chosen such that the input and output sizes agree
(on the first `N` dimensions, the kernel or window) when `stride==1`.

See also [`Conv`](@ref), [`MaxPool`](@ref).
"""
struct SamePad end

calc_padding(lt, pad, k::NTuple{N,T}, dilation, stride) where {T,N}= expand(Val(2*N), pad)
function calc_padding(lt, ::SamePad, k::NTuple{N,T}, dilation, stride) where {N,T}
  #Ref: "A guide to convolution arithmetic for deep learning" https://arxiv.org/abs/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 -> [cld(i, 2), fld(i,2)], vcat, pad_amt))
end

"""
    Conv(filter, in => out, σ = identity;
         stride = 1, pad = 0, dilation = 1, groups = 1, [bias, weight, init])

Standard convolutional layer. `filter` is a tuple of integers
specifying the size of the convolutional kernel;
`in` and `out` specify the number of input and output channels.

Image data should be stored in WHCN order (width, height, channels, batch).
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.
This has `N = 2` spatial dimensions, and needs a kernel size like `(5,5)`,
a 2-tuple of integers.

To take convolutions along `N` feature dimensions, this layer expects as input an array
with `ndims(x) == N+2`, where `size(x, N+1) == in` is the number of input channels,
and `size(x, ndims(x))` is (as always) the number of observations in a batch.
Then:
* `filter` should be a tuple of `N` integers.
* Keywords `stride` and `dilation` should each be either single integer,
  or a tuple with `N` integers.
* Keyword `pad` can be:
  - a single integer for equal padding all around,
  - a tuple of `N` integers, to apply the same padding at begin/end of each spatial dimension,
  - a tuple of `2*N` integers, for asymmetric padding, or
  - the singleton `SamePad()`, to calculate padding such that
    `size(output,d) == size(x,d) / stride` (possibly rounded) for each spatial dimension.
* Keyword `groups` is expected to be an `Int`. It specifies the number of groups
  to divide a convolution into.

Keywords to control initialization of the layer:
* `init` - Function used to generate initial weights. Defaults to `glorot_uniform`.
* `weight` - Initial weights of the layer. Typically an array, and can be used to override
  other configurations. By default, these are generated using [`convfilter`](@ref).
* `bias` - Initial bias is zero by default, this can be disabled entirely by setting it to
  [`Flux.Zeros()`](@ref) or equivalently `false`, or another vector provided as
  `bias = randn(Float32, out)`.

See also [`ConvTranspose`](@ref), [`DepthwiseConv`](@ref), [`CrossCor`](@ref).

# Examples
```jldoctest
julia> xs = rand(Float32, 100, 100, 3, 50); # a batch of images

julia> layer = Conv((5,5), 3 => 7, relu; bias = false)
Conv((5, 5), 3 => 7, relu, bias=false)  # 525 parameters

julia> layer(xs) |> size
(96, 96, 7, 50)

julia> Conv((5,5), 3 => 7; stride = 2)(xs) |> size
(48, 48, 7, 50)

julia> Conv((5,5), 3 => 7; stride = 2, pad = SamePad())(xs) |> size
(50, 50, 7, 50)

julia> Conv((1,1), 3 => 7; pad = (20,10,0,0))(xs) |> size
(130, 100, 7, 50)

julia> Conv((5,5), 3 => 7; stride = 2, dilation = 4)(xs) |> size
(42, 42, 7, 50)
```
"""
struct Conv{N,M,F,A,V}
  σ::F
  weight::A
  bias::V
  stride::NTuple{N,Int}
  pad::NTuple{M,Int}
  dilation::NTuple{N,Int}
  groups::Int
end

"""
    Conv(weight::AbstractArray, [bias, activation; stride, pad, dilation])

Constructs a convolutional layer with the given weight and bias.
Accepts the same keywords (and has the same defaults) as the `Conv((4,4), 3 => 7, relu)`
method.

# Examples
```jldoctest
julia> weight = rand(3, 4, 5);

julia> bias = zeros(5);

julia> c1 = Conv(weight, bias, sigmoid)  # expects 1 spatial dimension
Conv((3,), 4 => 5, σ)  # 65 parameters

julia> c1(randn(100, 4, 64)) |> size
(98, 5, 64)

julia> params(c1) |> length
2
```
"""
function Conv(w::AbstractArray{T,N}, b = true, σ = identity;
              stride = 1, pad = 0, dilation = 1, groups = 1) where {T,N}
  stride = expand(Val(N-2), stride)
  dilation = expand(Val(N-2), dilation)
  pad = calc_padding(Conv, pad, size(w)[1:N-2], dilation, stride)
  bias = create_bias(w, b, size(w, N))
  return Conv(σ, w, bias, stride, pad, dilation, groups)
end

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

  Conv(weight, bias, σ; stride, pad, dilation, groups)
end

"""
    convfilter(filter::Tuple, in => out[; init = glorot_uniform])

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...)

@functor Conv

function (c::Conv)(x::AbstractArray)
  σ, 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, groups = c.groups)
  σ.(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)))
  _print_conv_opt(io, l)
  print(io, ")")
end

function _print_conv_opt(io::IO, l)
  l.σ == identity || print(io, ", ", l.σ)
  all(==(0), l.pad) || print(io, ", pad=", _maybetuple_string(l.pad))
  all(==(1), l.stride) || print(io, ", stride=", _maybetuple_string(l.stride))
  all(==(1), l.dilation) || print(io, ", dilation=", _maybetuple_string(l.dilation))
  l.bias == Zeros() && print(io, ", bias=false")
end

"""
    ConvTranspose(filter, in => out, σ=identity; stride=1, pad=0, dilation=1, [bias, init])

Standard convolutional transpose layer. `filter` is a tuple of integers
specifying the size of the convolutional kernel, while
`in` and `out` specify the number of input and output channels.

Note that `pad=SamePad()` here tries to ensure `size(output,d) == size(x,d) * stride`.

Parameters are controlled by additional keywords, with defaults
`init=glorot_uniform` and `bias=true`.

See also [`Conv`](@ref) for more detailed description of keywords.

# Examples
```jldoctest
julia> xs = rand(Float32, 100, 100, 3, 50);  # a batch of 50 RGB images

julia> lay = ConvTranspose((5,5), 3 => 7, relu)
ConvTranspose((5, 5), 3 => 7, relu)  # 532 parameters

julia> lay(xs) |> size
(104, 104, 7, 50)

julia> ConvTranspose((5,5), 3 => 7, stride=2)(xs) |> size
(203, 203, 7, 50)

julia> ConvTranspose((5,5), 3 => 7, stride=3, pad=SamePad())(xs) |> size
(300, 300, 7, 50)
```
"""
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, activation; stride, pad, dilation])

Constructs a layer with the given weight and bias arrays.
Accepts the same keywords as the `ConvTranspose((4,4), 3 => 7, relu)` method.
"""
function ConvTranspose(w::AbstractArray{T,N}, bias = true, σ = 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(ConvTranspose, pad, size(w)[1:N-2], dilation, stride)
  b = create_bias(w, bias, size(w, N-1))
  return ConvTranspose(σ, w, b, stride, pad, 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 = true) where N

  ConvTranspose(weight, bias, σ; stride, pad, 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)
  σ, 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))
  _print_conv_opt(io, l)
  print(io, ")")
end

function calc_padding(::Type{ConvTranspose}, pad::SamePad, k::NTuple{N,T}, dilation, stride) where {N,T}
  calc_padding(Conv, pad, k .- stride .+ 1, dilation, stride)
end

"""
    DepthwiseConv(filter, in => out, σ=identity; stride=1, pad=0, dilation=1, [bias, init])

Depthwise convolutional layer. `filter` is a tuple of integers
specifying the size of the convolutional kernel, while
`in` and `out` specify the number of input and output channels.

Note that `out` must be an integer multiple of `in`.

Parameters are controlled by additional keywords, with defaults
`init=glorot_uniform` and `bias=true`.

See also [`Conv`](@ref) for more detailed description of keywords.

# Examples
```jldoctest
julia> xs = rand(Float32, 100, 100, 3, 50);  # a batch of 50 RGB images

julia> lay = DepthwiseConv((5,5), 3 => 6, relu; bias=false)
DepthwiseConv((5, 5), 3 => 6, relu, bias=false)  # 150 parameters

julia> lay(xs) |> size
(96, 96, 6, 50)

julia> DepthwiseConv((5,5), 3 => 9, stride=2, pad=2)(xs) |> size
(50, 50, 9, 50)
```
"""
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, [activation; stride, pad, dilation])

Constructs a layer with the given weight and bias arrays.
Accepts the same keywords as the `DepthwiseConv((4,4), 3 => 6, relu)` method.
"""
function DepthwiseConv(w::AbstractArray{T,N}, bias = true, σ = 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(DepthwiseConv, pad, size(w)[1:N-2], dilation, stride)
  b = create_bias(w, bias, prod(size(w)[N-1:end]))
  return DepthwiseConv(σ, w, b, stride, pad, dilation)
end

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 = true) where N
  @assert ch[2] % ch[1] == 0 "Output channels must be integer multiple of input channels"
  return DepthwiseConv(weight, bias, σ; stride, pad, dilation)
end

@functor DepthwiseConv

"""
    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 (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]))
  _print_conv_opt(io, l)
  print(io, ")")
end

"""
    CrossCor(filter, in => out, σ=identity; stride=1, pad=0, dilation=1, [bias, init])

Standard cross convolutional layer. `filter` is a tuple of integers
specifying the size of the convolutional kernel;
`in` and `out` specify the number of input and output channels.

Parameters are controlled by additional keywords, with defaults
`init=glorot_uniform` and `bias=true`.

See also [`Conv`](@ref) for more detailed description of keywords.

# Examples

```jldoctest
julia> xs = rand(Float32, 100, 100, 3, 50);  # a batch of 50 RGB images

julia> lay = CrossCor((5,5), 3 => 6, relu; bias=false)
CrossCor((5, 5), 3 => 6, relu, bias=false)  # 450 parameters

julia> lay(xs) |> size
(96, 96, 6, 50)

julia> CrossCor((5,5), 3 => 7, stride=3, pad=(2,0))(xs) |> size
(34, 32, 7, 50)
```
"""
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, activation; stride, pad, dilation])

Constructs a layer with the given weight and bias arrays.
Accepts the same keywords as the `CrossCor((4,4), 3 => 7, relu)` method.
"""
function CrossCor(w::AbstractArray{T,N}, bias = true, σ = 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(CrossCor, pad, size(w)[1:N-2], dilation, stride)
  b = create_bias(w, bias, size(w, N))
  return CrossCor(σ, w, b, stride, pad, 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 = true) where N

  return CrossCor(weight, bias, σ; stride, pad, 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)
  σ, 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)))
  _print_conv_opt(io, l)
  print(io, ")")
end

"""
    AdaptiveMaxPool(out::NTuple)

Adaptive max pooling layer. Calculates the necessary window size
such that its output has `size(y)[1:N] == out`.

Expects as input an array with `ndims(x) == N+2`, i.e. channel and
batch dimensions, after the `N` feature dimensions, where `N = length(out)`.

See also [`MaxPool`](@ref), [`AdaptiveMeanPool`](@ref).

# Examples
```jldoctest
julia> xs = rand(Float32, 100, 100, 3, 50);  # batch of 50 RGB images

julia> AdaptiveMaxPool((25, 25))(xs) |> size
(25, 25, 3, 50)

julia> MaxPool((4,4))(xs) ≈ AdaptiveMaxPool((25, 25))(xs)
true
```
"""
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::NTuple)

Adaptive mean pooling layer. Calculates the necessary window size
such that its output has `size(y)[1:N] == out`.

Expects as input an array with `ndims(x) == N+2`, i.e. channel and
batch dimensions, after the `N` feature dimensions, where `N = length(out)`.

See also [`MaxPool`](@ref), [`AdaptiveMaxPool`](@ref).

# Examples
```jldoctest
julia> xs = rand(Float32, 100, 100, 3, 50);  # batch of 50 RGB images

julia> AdaptiveMeanPool((25, 25))(xs) |> size
(25, 25, 3, 50)

julia> MeanPool((4,4))(xs) ≈ AdaptiveMeanPool((25, 25))(xs)
true
```
"""
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.

See also [`MaxPool`](@ref), [`GlobalMeanPool`](@ref).

```jldoctest
julia> xs = rand(Float32, 100, 100, 3, 50);

julia> m = Chain(Conv((3,3), 3 => 7), GlobalMaxPool());

julia> m(xs) |> size
(1, 1, 7, 50)

julia> GlobalMaxPool()(rand(3,5,7)) |> size  # preserves 2 dimensions
(1, 5, 7)
```
"""
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.

```jldoctest
julia> xs = rand(Float32, 100, 100, 3, 50);

julia> m = Chain(Conv((3,3), 3 => 7), GlobalMeanPool());

julia> m(xs) |> size
(1, 1, 7, 50)
```
"""
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(window::NTuple; pad=0, stride=window)

Max pooling layer, which replaces all pixels in a block of
size `window` with one.

Expects as input an array with `ndims(x) == N+2`, i.e. channel and
batch dimensions, after the `N` feature dimensions, where `N = length(window)`.

By default the window size is also the stride in each dimension.
The keyword `pad` accepts the same options as for the `Conv` layer,
including `SamePad()`.

See also [`Conv`](@ref), [`MeanPool`](@ref), [`AdaptiveMaxPool`](@ref), [`GlobalMaxPool`](@ref).

# Examples

```jldoctest
julia> xs = rand(Float32, 100, 100, 3, 50);  # batch of 50 RGB images

julia> m = Chain(Conv((5, 5), 3 => 7, pad=SamePad()), MaxPool((5, 5), pad=SamePad()))
Chain(
  Conv((5, 5), 3 => 7, pad=2),          # 532 parameters
  MaxPool((5, 5), pad=2),
)

julia> m[1](xs) |> size
(100, 100, 7, 50)

julia> m(xs) |> size
(20, 20, 7, 50)

julia> lay = MaxPool((5,), pad=2, stride=(3,))  # one-dimensional window
MaxPool((5,), pad=2, stride=3)

julia> lay(rand(Float32, 100, 7, 50)) |> size
(34, 7, 50)
```
"""
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(MaxPool ,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)
  all(==(0), m.pad) || print(io, ", pad=", _maybetuple_string(m.pad))
  m.stride == m.k || print(io, ", stride=", _maybetuple_string(m.stride))
  print(io, ")")
end

_maybetuple_string(pad) = string(pad)
_maybetuple_string(pad::Tuple) = all(==(pad[1]), pad) ? string(pad[1])  : string(pad)

"""
    MeanPool(window::NTuple; pad=0, stride=window)

Mean pooling layer, averaging all pixels in a block of size `window`.

Expects as input an array with `ndims(x) == N+2`, i.e. channel and
batch dimensions, after the `N` feature dimensions, where `N = length(window)`.

By default the window size is also the stride in each dimension.
The keyword `pad` accepts the same options as for the `Conv` layer,
including `SamePad()`.

See also [`Conv`](@ref), [`MaxPool`](@ref), [`AdaptiveMeanPool`](@ref).

# Examples

```jldoctest
julia> xs = rand(Float32, 100, 100, 3, 50);

julia> m = Chain(Conv((5,5), 3 => 7), MeanPool((5,5), pad=SamePad()))
Chain(
  Conv((5, 5), 3 => 7),                 # 532 parameters
  MeanPool((5, 5), pad=2),
)

julia> m[1](xs) |> size
(96, 96, 7, 50)

julia> m(xs) |> size
(20, 20, 7, 50)
```
"""
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(MeanPool, 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)
  all(==(0), m.pad) || print(io, ", pad=", _maybetuple_string(m.pad))
  m.stride == m.k || print(io, ", stride=", _maybetuple_string(m.stride))
  print(io, ")")
end


# performance benchmark

using SparseVectors

macro benchfun2(bfun, tfun, descr, xs, ys, nrtimes)
    esc(quote
        function $(bfun)(x, y, nr::Int)
            println(string("Benchmark on ", $(string(descr)), " ..."))
            xmat = convert(SparseMatrixCSC, x)
            ymat = convert(SparseMatrixCSC, y)

            gc();
            gc_enable(false)

            # warming
            $(tfun)(x, y)
            $(tfun)(xmat, ymat)

            # main profile
            et1 = @elapsed for i = 1:nr
                $(tfun)(x, y)
            end
            et2 = @elapsed for i = 1:nr
                $(tfun)(xmat, ymat)
            end

            gc_enable(true)

            # show result
            @printf("  sparse vector: %9.4fms (gain = %.4fx)\n", et1 * 1e3, et2 / et1)
            @printf("  sparse matrix: %9.4fms\n", et2 * 1e3)
        end
        $(bfun)($xs, $ys, $nrtimes)
    end)
end

# data

n = 1_000_000
x = sprand(n, 0.25)
y = sprand(n, 0.25)

@benchfun2(bench_vadd, +, "sparse vector addition", x, y, 20)
@benchfun2(bench_vmul, .*, "sparse vector multiplication", x, y, 20)

A = sprand(1000, 1000, 0.1)
x = sprand(1000, 0.2)

@benchfun2(bench_Ax, *, "sparse matrix-vector product (A * x)", A, x, 100)
@benchfun2(bench_Atx, At_mul_B, "sparse matrix-vector product (A'x)", A, x, 100)


function _gram(kernel::RK,
               knots::Vector{T}) where {T <: AbstractFloat, RK <: ReproducingKernel_1}
    n_1 = length(knots)
    mat = Matrix{T}(undef, n_1, n_1)
    @inbounds for j = 1:n_1
        for i = 1:j
            mat[i,j] = _rk(kernel, knots[i], knots[j])
            mat[j,i] = mat[i,j]
        end
    end
    return mat
end

function _gram_d1(kernel::RK,
                  knots::Vector{T},
                  d1_knots::Vector{T}) where {T <: AbstractFloat, RK <: ReproducingKernel_2}
    n_1 = length(knots)
    n_2 = length(d1_knots)
    n = n_1 + n_2
    mat = Matrix{T}(undef, n, n)
    @inbounds for j = 1:n_1
        for i = 1:j
            mat[i,j] = _rk(kernel, knots[i], knots[j])
            mat[j,i] = mat[i,j]
        end
    end
    n_1_p1 = n_1 + 1
    @inbounds for j = n_1_p1:n
        for i = 1:n_1
            mat[i,j] = _∂rk_∂ξ(kernel, knots[i], d1_knots[j-n_1])
            mat[j,i] = mat[i,j]
        end
    end
    @inbounds for j = n_1_p1:n
        for i = j:n
            mat[i,j] = _∂²rk_∂η∂ξ(kernel, d1_knots[i-n_1], d1_knots[j-n_1])
            mat[j,i] = mat[i,j]
        end
    end
    return mat
end

function _gram_d2(kernel::RK,
                  knots::Vector{T},
                  d2_knots::Vector{T}) where {T <: AbstractFloat, RK <: ReproducingKernel_3}
    n_1 = length(knots)
    n_2 = length(d2_knots)
    n = n_1 + n_2
    mat = Matrix{T}(undef, n, n)
    @inbounds for i = 1:n_1
        for j = 1:i
            mat[i,j] = _rk(kernel, knots[i], knots[j])
            mat[j,i] = mat[i,j]
        end
    end
    n_1_p1 = n_1 + 1
    @inbounds for i = 1:n_1
        for j = n_1_p1:n
            mat[i,j] = _∂²rk_∂ξ²(kernel, knots[i], d2_knots[j-n_1])
            mat[j,i] = mat[i,j]
        end
    end
    @inbounds for i = n_1_p1:n
        for j = i:n
            mat[i,j] = _∂⁴rk_∂η²∂ξ²(kernel, d2_knots[i-n_1], d2_knots[j-n_1])
            mat[j,i] = mat[i,j]
        end
    end
    return mat
end

function _gram12(kernel::RK,
               knots::Vector{T},
               d1_knots::Vector{T},
               d2_knots::Vector{T}) where {T <: AbstractFloat, RK <: ReproducingKernel_3}
    n_1 = length(knots)
    n_2 = length(d1_knots)
    n_3 = length(d2_knots)
    n = n_1 + n_2 + n_3
    mat = Matrix{T}(undef, n, n)
    @inbounds for j = 1:n_1
        for i = 1:j
            mat[i,j] = _rk(kernel, knots[i], knots[j])
            mat[j,i] = mat[i,j]
        end
    end
    n_1_p1 = n_1 + 1
    ne = n_1 + n_2
    @inbounds for j = n_1_p1:ne
        for i = 1:n_1
            mat[i,j] = _∂rk_∂ξ(kernel, knots[i], d1_knots[j-n_1])
            mat[j,i] = mat[i,j]
        end
    end
    ns = n_1 + n_2 + 1
    @inbounds for j = ns:n
        for i = 1:n_1
            mat[i,j] = _∂²rk_∂ξ²(kernel, knots[i], d2_knots[j-ns+1])
            mat[j,i] = mat[i,j]
        end
    end
    ne = n_1 + n_2
    @inbounds for j = n_1_p1:ne
        for i = j:ne
            mat[i,j] = _∂²rk_∂η∂ξ(kernel, d1_knots[i-n_1], d1_knots[j-n_1])
            mat[j,i] = mat[i,j]
        end
    end
    ns = n_1 + n_2 + 1
    ne = n_1 + n_2
    @inbounds for j = ns:n
        for i = n_1_p1:ne
            mat[i,j] = _∂³rk_∂η∂ξ²(kernel, d1_knots[i-n_1], d2_knots[j-ns+1])
            mat[j,i] = mat[i,j]
        end
    end
    ns = n_1 + n_2 + 1
    @inbounds for j = ns:n
        for i = j:n
            mat[i,j] = _∂⁴rk_∂η²∂ξ²(kernel, d2_knots[i-ns+1], d2_knots[j-ns+1])
            mat[j,i] = mat[i,j]
        end
    end
    return mat
end


# # sparseProbVec.jl
#
# export SparseSBPrior, SparseSBPriorP, SparseSBPriorFull,
#   rpost_sparseStickBreak, logSBMmarginal;
#
# struct SparseSBPrior
#   η::Float64
#   p1::Float64
#   μ::Float64
#   M::Float64
# end
#
# mutable struct SparseSBPriorP
#   η::Float64
#   μ::Float64
#   M::Float64
#   a_p1::Float64
#   b_p1::Float64
#   p1_now::Float64
# end
#
# mutable struct SparseSBPriorFull
#   η::Float64
#   M::Float64
#   a_p1::Float64
#   b_p1::Float64
#   a_μ::Float64
#   b_μ::Float64
#   p1_now::Float64
#   μ_now::Float64
# end
#
#
# ### Sparse Stick Breaking mixture
# function h_slice_mu(z2::Vector{Float64}, n2::Int, M::Float64, μ::Float64,
#   logscale::Bool=false)
#
#   a = -n2*(loggamma(M*μ) + loggamma(M*(1.0 - μ)))
#   b = M*μ*sum( log(z2) .- log(1.0 .- z2))
#   out = a + b
#   if (!logscale)
#     out = exp(out)
#   end
#   out
# end
#
#
# function slice_mu(z::Vector{Float64}, ξ::Vector{Int}, μ_old::Float64, M::Float64,
#   a_μ::Float64, b_μ::Float64)
#
#   # assumes xi=1 for group 1, xi=2 for group 2
#
#   n2 = sum( ξ.==2 )
#   betadist = Beta(a_μ, b_μ)
#   if n2 == 0
#     μ_out = rand( betadist )
#   else
#     z2 = copy( z[ findall(ξ.==2) ] )
#     b_u = h_slice_mu(z2, n2, M, μ_old)
#     u = rand( Uniform(0.0, b_u) )
#     lu = log(u)
#
#     ii = 0
#     keepgoing = true
#     μ_cand = 0.0
#     while keepgoing
#       if ii > 1000
#         throw("too many slices!")
#       end
#
#       μ_cand = rand( betadist )
#       lh = h_slice_mu(z2, n2, M, μ_cand, true)
#
#       keepgoing = lu > lh
#       ii += 1
#     end
#
#     μ_out = copy(μ_cand)
#   end
#
#   μ_out
# end
#
#
#
# function rpost_sparseStickBreak(x::Vector{Int}, p1::Float64, η::Float64,
#   μ::Float64, M::Float64, logout::Bool=false)
#
#   K = length(x)
#   n = K - 1
#   rcrx = cumsum( x[range(K, step=-1, length=n)] )[range(n, step=-1, length=n)]  # ∑_{k+1}^K x_k
#   γ = M * μ
#   δ = M * (1.0 - μ)
#
#   ## mixture component 1 update
#   a1 = 1.0 .+ x[1:n]
#   b1 = η .+ rcrx
#
#   ## mixture component 2 update
#   a2 = γ .+ x[1:n]
#   b2 = δ .+ rcrx
#
#   ## calculate posterior mixture weights
#   lgwt1 = log(p1) .- logbeta.(1.0, η) .+ logbeta.(a1, b1)
#   lgwt2 = log(1.0 - p1) .- logbeta.(γ, δ) .+ logbeta.(a2, b2)
#   ldenom = [ logsumexp( [lgwt1[i], lgwt2[i]] ) for i in 1:n ]
#
#   post_p1 = exp.( lgwt1 .- ldenom )
#   is1 = rand(n) .< post_p1
#   ξ = -1*(is1 .* 1 .- 1) .+ 1
#
#   lz = Vector{Float64}(undef, n)
#   for i in 1:n
#     if is1[i]
#       lx1 = log( rand( Distributions.Gamma( a1[i] ) ) )
#       lx2 = log( rand( Distributions.Gamma( b1[i] ) ) )
#       lxm = max(lx1, lx2)
#       lxsum = lxm + log( exp(lx1 - lxm) + exp(lx2 - lxm) ) # logsumexp
#       lz[i] = lx1 - lxsum
#       # z[i] = rand( Beta(a1[i], b1[i]) )
#     else
#       lx1 = log( rand( Distributions.Gamma( a2[i] ) ) )
#       lx2 = log( rand( Distributions.Gamma( b2[i] ) ) )
#       lxm = max(lx1, lx2)
#       lxsum = lxm + log( exp(lx1 - lxm) + exp(lx2 - lxm) ) # logsumexp
#       lz[i] = lx1 - lxsum
#       # z[i] = rand( Beta(a2[i], b2[i]) )
#     end
#   end
#
#   ## break the Stick
#   lw = Vector{Float64}(undef, K)
#   lwhatsleft = 0.0
#
#   for i in 1:n
#     lw[i] = lz[i] + lwhatsleft
#     lwhatsleft += log( 1.0 - exp(lw[i] - lwhatsleft) ) # logsumexp
#   end
#   lw[K] = copy(lwhatsleft)
#
#   z = exp.(lz)
#   w = exp.(lw)
#
#   if logout
#       (lw, lz, ξ)
#     else
#       (w, z, ξ)
#   end
# end
# function rpost_sparseStickBreak(x::Vector{Int}, p1_old::Float64, η::Float64,
#   μ::Float64, M::Float64, a_p1::Float64, b_p1::Float64, logout::Bool=false)
#
#   ## inference for p1 only
#
#   w, z, ξ = rpost_sparseStickBreak(x, p1_old, η, μ, M, logout)
#   p1_now = rand( Beta(a_p1 + sum(ξ==1), b_p1 + sum(ξ==2) ) )
#
#   (w, z, ξ, p1_now) # w and z may be log-valued depending on logout
# end
# function rpost_sparseStickBreak(x::Vector{Int}, p1_old::Float64, η::Float64,
#   μ_old::Float64, M::Float64, a_p1::Float64, b_p1::Float64, a_μ::Float64, b_μ::Float64, logout::Bool=false)
#
#   ## inference for p1 and μ as well
#
#   w, z, ξ = rpost_sparseStickBreak(x, p1_old, η, μ_old, M, logout)
#   if logout
#       μ_now = slice_mu(exp.(z), ξ, μ_old, M, a_μ, b_μ)
#   else
#       μ_now = slice_mu(z, ξ, μ_old, M, a_μ, b_μ)
#   end
#
#   p1_now = rand( Beta(a_p1 + sum(ξ==1), b_p1 + sum(ξ==2) ) )
#
#   (w, z, ξ, μ_now, p1_now) # w and z may be log-valued depending on logout
# end
#
#
#
# """
#     logSBMmarginal(x, p1, η, μ, M)
#
#   Calculate the log of the SBM prior predictive probability mass function.
#
# """
# function logSBMmarginal(x::Vector{Int}, p1::Float64, η::Float64,
#   μ::Float64, M::Float64)
#
#       K = length(x)
#       n = K - 1
#       rcrx = cumsum( x[range(K, step=-1, length=n)] )[range(n, step=-1, length=n)]  # ∑_{k+1}^K x_k
#       γ = M * μ
#       δ = M * (1.0 - μ)
#
#       ## mixture component 1 update
#       a1 = 1.0 .+ x[1:n]
#       b1 = η .+ rcrx
#
#       ## mixture component 2 update
#       a2 = γ .+ x[1:n]
#       b2 = δ .+ rcrx
#
#       ## calculate posterior mixture weights
#       lgwt1 = log(p1) .- logbeta.(1.0, η) .+ logbeta.(a1, b1)
#       lgwt2 = log(1.0 - p1) .- logbeta.(γ, δ) .+ logbeta.(a2, b2)
#       lsum = [ logsumexp( [lgwt1[i], lgwt2[i]] ) for i in 1:n ]
#
#       sum( lsum )
# end


using LinearAlgebra
using UnicodePlots
include("gemm.jl")

function plot_data(ns, data)
  scaling = 2 * ns.^3

  minv = floatmax()
  maxv = floatmin()
  for (s, v) = pairs(data)
    # if s != :ref
    minv = min(minv, minimum(scaling ./ v))
    maxv = max(maxv, maximum(scaling ./ v))
  # end
  end

  plt = lineplot(ns, scaling ./ data.ref, xlim = (ns[1], ns[end]),
                 ylim = (minv, maxv), title="max GFLOPS", xlabel="n",
                 ylabel="time", name = "ref")

  for (s, v) = pairs(data)
    if s != :ref
      lineplot!(plt, ns, scaling ./ v, name = string(s))
    end
  end
  plt
end

using BenchmarkTools

function run(ns, nsamples = 10)
  ref = zeros(length(ns))
  ji_reg = zeros(length(ns))
  ji = zeros(length(ns))

  for (i, n) = enumerate(ns)
    println()
    @show i,n
    m = k = n
    C = rand(m, n)
    A = rand(m, k)
    B = rand(k, n)
    C2 = copy(C)
    C1 = copy(C)

    t = @elapsed for n = 1:nsamples
      mul!(C2, A, B, 1, 1)
    end
    ref[i] = t / nsamples
    @show ref[i], 2*m^3 / ref[i]

    C .= C1
    t = @elapsed for n = 1:nsamples
      mygemm_NRxMR!(C, m, A, m, B, k, Val(8), Val(4))
    end
    ji_reg[i] = t / nsamples
    @show ji_reg[i], 2m^3/ji_reg[i]
    @show extrema(C - C2)

    C .= C1
    t = @elapsed for n = 1:nsamples
      mygemm!(C, A, B, Val(48), Val(256), Val(256), Val(8), Val(4))
    end
    ji[i] = t / nsamples
    @show ji[i], 2m^3/ji[i]
    @show extrema(C - C2)
  end

  data = (ref = ref,
          ji_reg = ji_reg,
          ji = ji,
         )

  return data
end

ns = 48:48:1000
data = run(ns)
plt = plot_data(ns, data)


using Manopt, ManifoldsBase, Test

@testset "Gradient Plan" begin
    io = IOBuffer()
    M = ManifoldsBase.DefaultManifold(2)
    x = [4.0, 2.0]
    o = GradientDescentOptions(
        x; stopping_criterion=StopAfterIteration(20), stepsize=ConstantStepsize(1.0)
    )
    o.gradient = [1.0, 0.0]
    f(M, y) = distance(M, y, x) .^ 2
    gradf(M, y) = -2 * log(M, y, x)
    p = GradientProblem(M, f, gradf)
    @test get_initial_stepsize(p, o) == 1.0
    @test get_stepsize(p, o, 1) == 1.0
    @test get_last_stepsize(p, o, 1) == 1.0
    # Check Fallbacks of Problen
    @test get_cost(p, o.x) == 0.0
    @test get_gradient(p, o.x) == zero_vector(M, x)
    @test_throws MethodError get_proximal_map(p, 1.0, o.x, 1)
    @test_throws MethodError get_subgradient(p, o.x)
    # Additional Specific Debugs
    a1 = DebugGradient(false, io)
    a1(p, o, 1)
    @test String(take!(io)) == "gradF(x):[1.0, 0.0]"
    a1a = DebugGradient("s:", io)
    a1a(p, o, 1)
    @test String(take!(io)) == "s:[1.0, 0.0]"
    a2 = DebugGradientNorm(false, io)
    a2(p, o, 1)
    @test String(take!(io)) == "|gradF(x)|:1.0"
    a2a = DebugGradientNorm("s:", io)
    a2a(p, o, 1)
    @test String(take!(io)) == "s:1.0"
    a3 = DebugStepsize(false, io)
    a3(p, o, 1)
    @test String(take!(io)) == "s:1.0"
    a3a = DebugStepsize("S:", io)
    a3a(p, o, 1)
    @test String(take!(io)) == "S:1.0"
    # Additional Specific Records
    b1 = RecordGradient(o.gradient)
    b1(p, o, 1)
    @test b1.recorded_values == [o.gradient]
    b2 = RecordGradientNorm()
    b2(p, o, 1)
    @test b2.recorded_values == [1.0]
    b3 = RecordStepsize()
    b3(p, o, 1)
    b3(p, o, 2)
    b3(p, o, 3)
    @test b3.recorded_values == [1.0, 1.0, 1.0]
end


using Test
using CellMLToolkit
using DifferentialEquations, Plots

path = @__DIR__
ml = CellModel(path * "/../models/beeler_reuter_1977.cellml.xml")

# @test length(ml.eqs) == 8
# @test ml.iv.op.name == :time

# eqs, vs = CellMLToolkit.flat_equations(ml)
# @test length(vs) == 8

# @test find_V(ml).op.name == :V

prob = ODEProblem(ml, (0,10000.0))
sol1 = solve(prob, Euler(), dt=0.01, saveat=1.0)
sol2 = solve(prob, TRBDF2(), dtmax=0.5, saveat=1.0)
V1 = map(x -> x[1], sol1.u)
V2 = map(x -> x[1], sol2.u)
err = sum(abs.(V1 .- V2)) / length(V1)
@test err < 0.1

# prob = ODEProblem(ml, (0,10000.0); jac=true)
# sol3 = solve(prob, TRBDF2(), dtmax=0.5, saveat=1.0)
# V3 = map(x -> x[1], sol2.u)
# err = sum(abs.(V1 .- V3)) / length(V1)
# @test err < 0.1

p = list_params(ml)
update_list!(p, "IstimPeriod", 280.0)
prob = ODEProblem(ml, (0,10000.0); jac=false, p=p)
sol4 = solve(prob, TRBDF2(), dtmax=0.5, saveat=1.0)
V4 = map(x -> x[1], sol2.u)
err = sum(abs.(V1[1:250] .- V4[1:250])) / 250
@test err < 0.1

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

function test_plot()
    ml = CellModel(path * "/../models/beeler_reuter_1977.cellml.xml")
    tspan = (0, 5000.0)
    prob = ODEProblem(ml, tspan)
    sol = solve(prob, TRBDF2(), dtmax=0.5)
    x = Array(sol)
    plot(x[1,:])
end


include("utils.jl")

@testset "Euclidean" begin
    E = Euclidean(3)
    Ec = Euclidean(3; field=ℂ)
    EM = Manifolds.MetricManifold(E, Manifolds.EuclideanMetric())
    @test repr(E) == "Euclidean(3; field = ℝ)"
    @test repr(Ec) == "Euclidean(3; field = ℂ)"
    @test repr(Euclidean(2, 3; field = ℍ)) == "Euclidean(2, 3; field = ℍ)"
    @test is_default_metric(EM)
    @test is_default_metric(E, Manifolds.EuclideanMetric())
    @test Manifolds.default_metric_dispatch(E,Manifolds.EuclideanMetric()) === Val{true}()
    p = zeros(3)
    @test det_local_metric(EM, p) == one(eltype(p))
    @test log_local_metric_density(EM, p) == zero(eltype(p))
    @test project!(E, p, p) == p
    @test manifold_dimension(Ec) == 2*manifold_dimension(E)
    X = zeros(3)
    X[1] = 1.0
    Y = similar(X)
    project!(E, Y, p, X)
    @test Y==X

    # real manifold does not allow complex values
    @test_throws DomainError is_manifold_point(Ec,[:a,:b,:b], true)
    @test_throws DomainError is_manifold_point(E,[1.0, 1.0im, 0.0], true)
    @test_throws DomainError is_manifold_point(E,[1], true)
    @test_throws DomainError is_tangent_vector(Ec,[:a,:b,:b], [1.0, 1.0, 0.0], true)
    @test_throws DomainError is_tangent_vector(E,[1.0, 1.0im, 0.0], [1.0, 1.0, 0.0], true) # real manifold does not allow complex values
    @test_throws DomainError is_tangent_vector(E,[1], [1.0, 1.0, 0.0], true)
    @test_throws DomainError is_tangent_vector(E, [0.0, 0.0, 0.0], [1.0], true)
    @test_throws DomainError is_tangent_vector(E, [0.0, 0.0, 0.0], [1.0, 0.0, 1.0im], true)
    @test_throws DomainError is_tangent_vector(Ec, [0.0, 0.0, 0.0], [:a, :b, :c], true)

    @test E^2 === Euclidean(3, 2)
    @test ^(E,2) === Euclidean(3, 2)
    @test E^(2,) === Euclidean(3, 2)
    @test Ec^(4,5) === Euclidean(3, 4, 5; field = ℂ)

    manifolds = [ E, EM, Ec ]
    types = [ Vector{Float64}, ]
    TEST_FLOAT32 && push!(types, Vector{Float32})
    TEST_DOUBLE64 && push!(types, Vector{Double64})
    TEST_STATIC_SIZED && push!(types, MVector{3, Float64})

    types_complex = [ Vector{ComplexF64}, ]
    TEST_FLOAT32 && push!(types_complex, Vector{ComplexF32})
    TEST_DOUBLE64 && push!(types_complex, Vector{ComplexDF64})
    TEST_STATIC_SIZED && push!(types_complex, MVector{3, ComplexF64})

    for M in manifolds
        basis_types = if M == E
            (DefaultOrthonormalBasis(), ProjectedOrthonormalBasis(:svd), DiagonalizingOrthonormalBasis([1.0, 2.0, 3.0]))
        elseif M == Ec
            (DefaultOrthonormalBasis(), DiagonalizingOrthonormalBasis([1.0, 2.0, 3.0]))
        else
            ()
        end
        for T in types
            @testset "$M Type $T" begin
                pts = [convert(T, [1.0, 0.0, 0.0]),
                       convert(T, [0.0, 1.0, 0.0]),
                       convert(T, [0.0, 0.0, 1.0])]
                test_manifold(
                    M,
                    pts,
                    test_reverse_diff = isa(T, Vector),
                    test_project_point = true,
                    test_project_tangent = true,
                    test_musical_isomorphisms = true,
                    test_vector_transport = true,
                    test_mutating_rand = isa(T, Vector),
                    point_distributions = [Manifolds.projected_distribution(M, Distributions.MvNormal(zero(pts[1]), 1.0))],
                    tvector_distributions = [Manifolds.normal_tvector_distribution(M, pts[1], 1.0)],
                    basis_types_vecs = basis_types,
                    basis_types_to_from = basis_types,
                    basis_has_specialized_diagonalizing_get = true,
                    test_vee_hat = isa(M, Euclidean),
                )
            end
        end
    end
    for T in types_complex
        @testset "Complex Euclidean, type $T" begin
            pts = [convert(T, [1.0im, -1.0im, 1.0]),
                   convert(T, [0.0, 1.0, 1.0im]),
                   convert(T, [0.0, 0.0, 1.0])]
            test_manifold(
                Ec,
                pts,
                test_reverse_diff = isa(T, Vector),
                test_project_tangent = true,
                test_musical_isomorphisms = true,
                test_vector_transport = true,
                test_vee_hat = false,
            )
        end
    end

    @testset "hat/vee" begin
        E = Euclidean(3, 2)
        p = collect(reshape(1.0:6.0, (3, 2)))
        X = collect(reshape(7.0:12.0, (3, 2)))
        @test hat(E, p, vec(X)) ≈ X
        Y = allocate(X)
        @test hat!(E, Y, p, vec(X)) === Y
        @test Y ≈ X
        @test vee(E, p, X) ≈ vec(X)
        Y = allocate(vec(X))
        @test vee!(E, Y, p, X) === Y
        @test Y ≈ vec(X)
    end

    @testset "Number systems power" begin
        @test ℝ^2 === Euclidean(2)
        @test ℝ^(2,3) === Euclidean(2, 3)

        @test ℂ^2 === Euclidean(2; field = ℂ)
        @test ℂ^(2,3) === Euclidean(2, 3; field = ℂ)

        @test ℍ^2 === Euclidean(2; field = ℍ)
        @test ℍ^(2,3) === Euclidean(2, 3; field = ℍ)
    end
end


"""
Propagate labels to reach the starting partition. Takes an adjacency matrix
as input and returns a `Modular` object as output.

This approach works better for large graphs (it tends to underperform both the
"random" and "single" approaches for small graphs).

The optimization stops when any label propagation fails to increase the
modularity, measured with `Q`.
"""
function partition_lp(A::Array{Int64, 2})
  nr, nc = size(A)
  min_part = minimum(size(A))
  S = eye(Int64, sum(size(A)))
  M = Modular(A, S)
  old_Q = 0.0
  new_Q = 0.000001
  while new_Q > old_Q
    update_order_top = shuffle([1:nc;])
    update_order_bot = shuffle([1:nr;])
    # Top level
    @inbounds for node_top in update_order_top
       neighbors = [1:nr;][M.A[:,node_top] .> 0] .+ nc
       lab_count = sum(M.S[neighbors,:],1)
       for i = 1:length(lab_count)
          if lab_count[i] == maximum(lab_count)
             M.S[node_top,:] = zeros(Int64, size(M.S)[2])
             M.S[node_top,i] = 1
             break
          end
       end
    end
    # Bottom level
    @inbounds for node_bot in update_order_bot
       neighbors = [1:nc;]'[M.A[node_bot,:] .> 0]
       lab_count = sum(M.S[neighbors,:],1)
       for i = 1:length(lab_count)
          if lab_count[i] == maximum(lab_count)
             M.S[node_bot+nc,:] = zeros(Int64, size(M.S)[2])
             M.S[node_bot+nc,i] = 1
             break
          end
       end
    end
    old_Q = new_Q
    new_Q = Q(M)
  end
  no_empty_modules!(M)
  return M
end


using Main.COSMO
using Test, LinearAlgebra, SparseArrays, Random

function create_diagonal_extractor(Tv, n::Ti) where {Ti}
    # Creates a matrix A that slices out the diagonal entries Xii of
    # a upper-triangular-vectorized square matrix x=vec(X) of dimension n
    A = spzeros(Tv, Ti, n, Ti(n*(n + 1)/2))
    idx = 0
    for i = 1:n
        idx += i
        A[i, idx] = one(Tv)
    end
    return A
end

function create_closest_correlation_model(C::AbstractMatrix{Tv}, settings=nothing;
    psdcone_type=COSMO.PsdConeTriangle) where Tv
    # Closest Correlation Matrix problem
    # min   1/2 ||X-C||^2
    # s.t.    Xii = 1
    #         X ⪴ 0

    @assert psdcone_type in [COSMO.PsdConeTriangle COSMO.PsdConeTriangleLanczos]
    n = size(C, 1)
    @assert size(C, 2) == size(C, 1)

    # Contraint #1: Diagonal equal to one
    A1 = create_diagonal_extractor(Tv, n)
    b1 = -ones(Tv, n)
    constraints = [COSMO.Constraint(A1, b1, COSMO.ZeroSet)]

    # Contraint #2: decision variable is a PSD
    n2 =  Int(n*(n + 1)/2)
    A2 = SparseMatrixCSC(one(Tv)*I, n2, n2)
    b2 = zeros(Tv, n2)
    push!(constraints, COSMO.Constraint(A2, b2, psdcone_type))

    P = SparseMatrixCSC(Tv(1/2)*I, n2, n2)
    q = -COSMO.extract_upper_triangle(C)

    model = COSMO.Model()
    if settings != nothing
        COSMO.assemble!(model, P, q, constraints, settings=settings)
    else
        COSMO.assemble!(model, P, q, constraints)
    end

    return model
end


using OrdinaryDiffEq, Zygote
using DiffEqSensitivity
using Random

function dt!(du, u, p, t)
  x, y = u
  α, β, δ, γ = p
  du[1] = dx = α*x - β*x*y
  du[2] = dy = -δ*y + γ*x*y
end

n_par=3
Random.seed!(2)
u0=rand(2,n_par)
u0[:,1] = [1.0,1.0]
tspan = (0.0, 10.0)
p = [2.2, 1.0, 2.0, 0.4]
prob_ode = ODEProblem(dt!, u0[:,1], tspan)

function test_loss(p1,prob)

	function prob_func(prob, i, repeat)
		@show i
		remake(prob,u0=u0[:,i])
	end

	#define ensemble problem
	ensembleprob = EnsembleProblem(prob,prob_func = prob_func)

	u = Array(solve(ensembleprob, Tsit5(), EnsembleThreads(), trajectories=n_par,
	p=p,
	sensealg = ForwardDiffSensitivity(),
	saveat = 0.1, dt=0.001))[:,end,:]
	loss=sum(u)
	return loss
end

@time gs = Zygote.gradient(p) do p
	test_loss(p,prob_ode)
end


#md # !!! note
#
#     This tutorial is under construction, but the code below is already functional.
#
# Driver that computes the lid-driven cavity benchmark at low Reynolds numbers
# when using a mixed FE Q(k)/Pdisc(k-1).

# Load Gridap library
using Gridap

# Discrete model
n = 100
domain = (0,1,0,1)
partition = (n,n)
model = CartesianDiscreteModel(domain, partition)

# Define Dirichlet boundaries
labels = get_face_labeling(model)
add_tag_from_tags!(labels,"diri1",[6,])
add_tag_from_tags!(labels,"diri0",[1,2,3,4,5,7,8])

# Define reference FE (Q2/P1(disc) pair)
order = 2
reffeᵤ = ReferenceFE(lagrangian,VectorValue{2,Float64},order)
reffeₚ = ReferenceFE(lagrangian,Float64,order-1;space=:P)

# Define test FESpaces
V = TestFESpace(model,reffeᵤ,labels=labels,dirichlet_tags=["diri0","diri1"],conformity=:H1)
Q = TestFESpace(model,reffeₚ,conformity=:L2,constraint=:zeromean)
Y = MultiFieldFESpace([V,Q])

# Define trial FESpaces from Dirichlet values
u0 = VectorValue(0,0)
u1 = VectorValue(1,0)
U = TrialFESpace(V,[u0,u1])
P = TrialFESpace(Q)
X = MultiFieldFESpace([U,P])

# Define triangulation and integration measure
degree = order
Ωₕ = Triangulation(model)
dΩ = Measure(Ωₕ,degree)

# Define bilinear and linear form
f = VectorValue(0.0,0.0)
a((u,p),(v,q)) = ∫( ∇(v)⊙∇(u) - (∇⋅v)*p + q*(∇⋅u) )dΩ
l((v,q)) = ∫( v⋅f )dΩ

# Build affine FE operator
op = AffineFEOperator(a,l,X,Y)

# Solve
uh, ph = solve(op)

# Export results to vtk
writevtk(Ωₕ,"results",order=2,cellfields=["uh"=>uh,"ph"=>ph])


export mapflute

"""
 𝒍 = mapflute(𝒇, 𝒅)

Calculate locations of toneholes, and length of flute
 for frequencies 𝒇 (Hz, descending pitch) and hole diameters 𝒅 (mm, same order)
 in mm measured from embouchure hole
"""
function mapflute(𝒇, 𝒅)
  # TODO: assumes standard temperature ϑ, bore diameter ⌀, wall thickness ℎ
  # hole positions
  𝒍 = []
  # number of holes
  ħ = length(𝒇)-1
  # accumulate closed-hole correction
  𝛥ℓᵥ = 0.0
  for h in 1:ħ
    𝑓 = 𝒇[h]    # opened hole frequency
    𝑓ₜ = 𝒇[h+1] # closed hole frequency
    𝑑ₕ = 𝒅[h]   # hole diameter
    ℓₕ = toneholelength(𝑓, 𝑓ₜ, 𝑑ₕ; 𝛥ℓᵥ=𝛥ℓᵥ) # calculate hole position
    𝒍 = [𝒍; ℓₕ] # append to result list
    # closed hole correction effect on remaining holes
    𝛥ℓᵥ += closedholecorrection(𝑓, 𝑓ₜ, 𝑑ₕ; 𝛥ℓᵥ=𝛥ℓᵥ)
  end
  # append flute tube length to result
  ℓₜ = flutelength(𝒇[end]; 𝛥ℓᵥ=𝛥ℓᵥ)
  return [𝒍; ℓₜ]
end


using Bokeh
if VERSION < v"0.4-"
    using Dates
else
    using Base.Dates
end

start = now()
x = map(d -> start + Dates.Hour(d), 1:(30 * 24))
println(typeof(x))
# x: Array{DateTime,1}
day = 1000 * 3600 * 24
y = 0.5sin(int(x) * pi / day) + cos(int(x) * pi / (day * 7))
# note x needs to be an array of type DateTime or Date
plot(x, y)
showplot("dates_xy.html")

start = Date(1980, 1, 1)
x = map(d -> start + Dates.Day(d), 1:(365 * 30))
println(typeof(x))
# x: Array{Date,1}
y = 0.5sin(int(x) * pi / 365) + cos(int(x) * pi / (365 * 7))

plot(y, x)
showplot("dates_yx.html")


using BenchmarkTools
H = read_hypergraph("/Users/aaronkelley/.julia/dev/CoexistenceHoles.jl/test/test_data/Lagoon of Venice/Hp.dat")

@benchmark begin
    f = betti_hypergraph_ripscomplex(H; max_dim=4)
end

@benchmark begin
    hypergraph_subdivide(H)
end;

@time begin
    G = hypergraph_subdivide(H; expansion = false);
end;

@time begin
    M = 2.0*ones(length(H), length(H));
    [ M[g[1], g[2]] = 1 for g in G ];
    [ M[g[2], g[1]] = 1 for g in G ];
    M[diagind(M)] .= 0;
end;

@time begin
    flt = Rips(M; threshold=1.5,sparse=true)
end;

@time begin
    barcodes = ripserer(flt; dim_max = 4)
end;

@time begin
    betti = [ count(x -> x[2] .== Inf, barcodes[i]) for i in 1:length(barcodes)]
end;

function betti_hypergraph_ripscomplex(H; max_dim = 3)

    if length(H) == 0
        return Int.(zeros(max_dim + 1))
    elseif length(H) == 1
        return Int.(vcat(1, zeros(max_dim)))
    end

    _max_dim = min(max_dim, length(H) - 2);

    G = hypergraph_subdivide(H; expansion = false)
    # build distance matrix to give as input to Ripserer.
    M = 2.0*ones(length(H), length(H))
    [ M[g[1], g[2]] = 1 for g in G ]
    [ M[g[2], g[1]] = 1 for g in G ]
    M[diagind(M)] .= 0;

    #call Ripser
    flt = Rips(M; threshold=1, sparse=true)
    barcodes = ripserer(flt; dim_max = _max_dim)
    # interpret the outout. Barcodes is an array of tuples of dim = max_dim + 1.
    # All tuples where the second entry is Inf are the persistent ones.
    # Betti numbers just count the number of those persistent ones.
    betti = [ count(x -> x[2] .== Inf, barcodes[i]) for i in 1:length(barcodes)]

    padding = (max_dim+1) - length(betti);
    return vcat(betti, Int.(zeros(padding)))

    ## old function using Ripser library: https://github.com/mtsch/Ripser.jl#master
    # if length(H) == 0
    #     return zeros(max_dim + 1)
    # else
    #
    #     G = hypergraph_subdivide(H; expansion = false)
    #     # build distance matrix to give as input to Ripser.
    #     M = 2.0*ones(length(H), length(H))
    #     [ M[g[1], g[2]] = 1 for g in G ]
    #
    #     #call Ripser
    #     barcodes = ripser(M; dim_max = max_dim, threshold = 1.5)
    #     # interpret the outout. Barcodes is an array of tuples of dim = max_dim + 1.
    #     # All tuples where the second entry is Inf are the persistent ones.
    #     # Betti numbers just count the number of those persistent ones.
    #     betti = [ count(x -> x[2] .== Inf, barcodes[i]) for i in 1:length(barcodes)]
    #
    #     return betti
    # end
end


using Distributions

# params (computed from empirical data)
LAMBDA = 0.6882627
ALPHA = 1.3441739239406807
SIGMA = 0.4203008

mutable struct Model
    post_list
    top_list
    tick
    n_users
end

mutable struct Post
    id::Int64
    quality::Float64
    votes::Int64
    views::Int64
    timestamp::Int64
    score::Float64
end

function generate_attention(n)
    pois = Poisson(5)
    return rand(pois, n)
end

# Wie legen wir n fest?
# Aus den Daten holen (wie viele Votes werde pro Minute abgegeben?)

# refactor
function user_interaction!(post_list, attention_list)
    for att in attention_list
        for i in 1:att
            post_list[i].views += 1
            if post_list[i].quality > 0.5
                post_list[i].votes += 1
            end
        end
    end
    return post_list
end

function score_single(post, tick)
    age = (tick - post.timestamp) / 60
    votes = post.votes
    return ((votes - 1)^0.8)/((age + 2)^1.8)
end

function hnscore!(post_list, tick)
    for p in post_list
        p.score = score_single(p, tick)
    end
    return post_list
end

function sort_toplist(post_list)
    toplist = sort(post_list, rev=true, by=(post -> post.score))
    return toplist
end

function bias_lambda(slice, alpha=ALPHA, sd=SIGMA)
    sd * (- alpha * sin(2 * pi  * slice))
end

function draw_arrival_rate(tick, ticks_per_day, arrival_rate_mean=LAMBDA; alpha=ALPHA, arrival_rate_sd=SIGMA)
    rel_slice = mod(tick, ticks_per_day)
    lambda = arrival_rate_mean + bias_lambda(rel_slice, alpha, arrival_rate_sd)
    dist = Poisson(lambda)
    return rand(dist)
end

function tick()

end

l = [draw_arrival_rate(100, 1440) for i in 1:100]

histogram(l)

using Plots

post_list = [Post(rand(Uniform(0, 1)), 1, 0, 6, 0) for i in 1:100]
attention_list = generate_attention(100)
user_interaction!(post_list, attention_list)
hnscore!(post_list, 8)
toplist = sort_toplist(post_list)

post_list

# votes + age -> wie verhalten sie sich in Bezug auf score?


immutable PoissonRmathSampler <: Sampleable{Univariate,Discrete}
    mu::Float64
end

rand(s::PoissonRmathSampler) =
    int(ccall((:rpois, "libRmath-julia"), Float64, (Float64,), s.mu))

function poissonpvec(μ::Float64, n::Int)
    # Poisson probabilities, from 0 to n
    pv = Array(Float64, n+1)
    @inbounds pv[1] = p = exp(-μ)
    for i = 1:n
        @inbounds pv[i+1] = (p *= (μ / i))
    end
    return pv
end

# Naive sampler by counting exp variables
#
# Suitable for small μ
#
immutable PoissonCountSampler <: Sampleable{Univariate,Discrete}
    μ::Float64
end

function rand(s::PoissonCountSampler)
    μ = s.μ
    n = 0
    c = randexp()
    while c < μ
        n += 1
        c += randexp()
    end
    return n
end

# Algorithm from:
#
#   J.H. Ahrens, U. Dieter (1982)
#   "Computer Generation of Poisson Deviates from Modified Normal Distributions"
#   ACM Transactions on Mathematical Software, 8(2):163-179
#
#   For μ sufficiently large, (i.e. >= 10.0)
#
immutable PoissonADSampler <: Sampleable{Univariate,Discrete}
    μ::Float64
    s::Float64
    d::Float64
    L::Int
end

PoissonADSampler(μ::Float64) =
    PoissonADSampler(μ,sqrt(μ),6.0*μ^2,ifloor(μ-1.1484))

function rand(s::PoissonADSampler)
    # Step N
    G = s.μ + s.s*randn()

    if G >= 0.0
        K = ifloor(G)
        # Step I
        if K >= s.L
            return K
        end

        # Step S
        U = rand()
        if s.d*U >= (s.μ-K)^3
            return K
        end

        # Step P
        px,py,fx,fy = procf(s.μ,K,s.s)

        # Step Q
        if fy*(1-U) <= py*exp(px-fx)
            return K
        end
    end

    while true
        # Step E
        E = randexp()
        U = 2.0*rand()-1.0
        T = 1.8+copysign(E,U)
        if T <= -0.6744
            continue
        end

        K = ifloor(s.μ + s.s*T)
        px,py,fx,fy = procf(s.μ,K,s.s)
        c = 0.1069/s.μ

        # Step H
        if c*abs(U) <= py*exp(px+E)-fy*exp(fx+E)
            return K
        end
    end
end

# Procedure F
function procf(μ::Float64, K::Int, s::Float64)
    # can be pre-computed, but does not seem to affect performance
    ω = 0.3989422804014327/s
    b1 = 0.041666666666666664/μ
    b2 = 0.3*b1*b1
    c3 = 0.14285714285714285*b1*b2
    c2 = b2 - 15.0*c3
    c1 = b1 - 6.0*b2 + 45.0*c3
    c0 = 1.0 - b1 + 3.0*b2 - 15.0*c3

    if K < 10
        px = -μ
        py = μ^K/factorial(K)
    else
        δ = 0.08333333333333333/K
        δ -= 4.8*δ^3
        V = (μ-K)/K
        px = K*log1pmx(V) - δ # avoids need for table
        py = 0.3989422804014327/sqrt(K)

    end
    X = (K-μ+0.5)/s
    X2 = X^2
    fx = -0.5*X2 # missing negation in pseudo-algorithm, but appears in fortran code.
    fy = ω*(((c3*X2+c2)*X2+c1)*X2+c0)
    return px,py,fx,fy
end


using Primes

function getmersenneprimes(n)
    t1 = time()
    count = 0
    i = 2
    while(n > count)
        if(isprime(i) && ismersenneprime(2^BigInt(i) - 1))
            println("M$i, cumulative time elapsed: $(time() - t1) seconds")
            count += 1
        end
        i += 1
    end
end

getmersenneprimes(50)


if VERSION >= v"0.7.0"
    using Printf
    using LinearAlgebra
end

function descent( z::Array{T,1}, alpha::Array{T,1}, c::SLiSeFilters.Config, lb::R=R(0.0) )

    project = function( z::Array{T,1} )
        mask = abs.(imag(z)) .< lb
        z[mask] = real(z[mask]) + lb*im*sign.(imag(z[mask]));
        z
    end

    maxiter = 20000;
    i = 1;

    while i < maxiter
        dz,da = Fgradient( z, alpha, c );

        a = 0.5;
        b = 0.5;
        t::Float64 = 1*b;

        resd = residual(z, alpha, c);
        #emm = dot(-([da; dz]), [da; dz]);
        emm = -norm([da; dz])^2;

        while( residual( project(z-t .* dz), alpha - t.*da, c) > real(resd + a * t * emm) )
            t = t*b;
            if 1e-16 > t
                println("done...")
                i=maxiter;
                break;
            end
        end

        alpha -= t*da;
        z = project(z - t*dz);

        if mod(i,500) == 0
            @printf "%.16e\t%d\n" resd i
        end
        i+=1
    end

    return z,alpha
end


#___  ____       ___ _____   _   _       _ _
#|  \/  | |     |_  /  __ \ | | | |     (_| |
#| .  . | |       | | /  \/ | | | |_ __  _| |_ ___
#| |\/| | |       | | |     | | | | '_ \| | __/ _ \
#| |  | | |___/\__/ | \__/\ | |_| | | | | | || (_) |
#_______\_____\____/ \____/  _____|_|___|_____\_____ _____ _____
#| ___ \        (_)         | | \ \ / / / __  |  _  / __  |  _  |
#| |_/ _ __ ___  _  ___  ___| |_ \ V /  `' / /| |/' `' / /| |/' |
#|  __| '__/ _ \| |/ _ \/ __| __|/   \    / / |  /| | / / |  /| |
#| |  | | | (_) | |  __| (__| |_/ /^\ \ ./ /__\ |_/ ./ /__\ |_/ /
#\_|  |_|  \___/| |\___|\___|\__\/   \/ \_____/\___/\_____/\___/
#              _/ |
#             |__/
#
# This code is part of the proposal of the team "MLJC UniTo" - University of Turin
# for "ProjectX 2020" Climate Change for AI.
# The code is licensed under MIT 3.0
# Please read readme or comments for credits and further information.

# Compiler: Julia 1.5.3

# Short description of this file: Numerical Approach to diffusion equation

#Warning: Reset Julia Kernel (pre-compilation required) to avoid errors in
#redefinitions of constants.

import Pkg; Pkg.add("OrdinaryDiffEq"); Pkg.add("LinearAlgebra"); Pkg.add("SparseArrays"); Pkg.add("BenchmarkTools");
using OrdinaryDiffEq, LinearAlgebra, SparseArrays, BenchmarkTools, Plots, DifferentialEquations
const α₂ = 1.0
const α₃ = 1.0
const β₁ = 1.0
const β₂ = 1.0
const β₃ = 1.0
const r₁ = 1.0
const r₂ = 1.0
const D = 100.0
const γ₁ = 0.1
const γ₂ = 0.1
const γ₃ = 0.1
const N = 128
const X = reshape([i for i in 1:N for j in 1:N],N,N)
const Y = reshape([j for i in 1:N for j in 1:N],N,N)
const α₁ = 1.0.*(X.>=4*N/5)

const Mx = Tridiagonal([1.0 for i in 1:N-1],[-2.0 for i in 1:N],[1.0 for i in 1:N-1])
const My = copy(Mx)
Mx[2,1] = 2.0
Mx[end-1,end] = 2.0
My[1,2] = 2.0
My[end,end-1] = 2.0

# Define the discretized PDE as an ODE function
const MyA = zeros(N,N)
const AMx = zeros(N,N)
const DA = zeros(N,N)
function f(du,u,p,t)
   A = @view  u[:,:,1]
   B = @view  u[:,:,2]
   C = @view  u[:,:,3]
  dA = @view du[:,:,1]
  dB = @view du[:,:,2]
  dC = @view du[:,:,3]
  mul!(MyA,My,A)
  mul!(AMx,A,Mx)
  @. DA = D*(MyA + AMx)
  @. dA = DA + α₁ - β₁*A - r₁*A*B + r₂*C
  @. dB = α₂ - β₂*B - r₁*A*B + r₂*C
  @. dC = α₃ - β₃*C + r₁*A*B - r₂*C
end

u0 = zeros(N,N,3)
prob = ODEProblem(f,u0,(0.0,10.0))
#CPU: Intel Core i5 4440 - 4 threads
@btime solve(prob, ROCK4(), reltol = 1e-8, abstol=1e-8, saveat=0.1); # 1.111 s (3452 allocations: 285.59 MiB)
@btime solve(prob, Tsit5(), reltol = 1e-8, abstol=1e-8, saveat=0.1); # 2.886 s (275 allocations: 44.27 MiB)
@btime solve(prob, DP5(), reltol = 1e-8, abstol=1e-8, saveat=0.1);   # 2.591 s (269 allocations: 43.14 MiB)
@btime solve(prob, RK4(), relto1 = 1e-8, abstol = 1e-8, saveat=0.1); # 3.471 s (268 allocations: 43.14 MiB)

@btime solve(prob, ROCK4(), reltol = 1e-7, abstol=1e-7, saveat=0.1); # 604.227 ms (2972 allocations: 195.58 MiB)
@btime solve(prob, Tsit5(), reltol = 1e-7, abstol=1e-7, saveat=0.1); # 2.871 s (275 allocations: 44.27 MiB)
@btime solve(prob, DP5(), reltol = 1e-7, abstol=1e-7, saveat=0.1);   # 2.498 s (269 allocations: 43.14 MiB)
@btime solve(prob, RK4(), relto1 = 1e-7, abstol = 1e-7, saveat=0.1); # 3.547 s (268 allocations: 43.14 MiB)

@btime solve(prob, ROCK4(), reltol = 1e-9, abstol=1e-9, saveat=0.1); # 1.289 s (4324 allocations: 449.13 MiB)
@btime solve(prob, Tsit5(), reltol = 1e-9, abstol=1e-9, saveat=0.1); # 2.282 s (275 allocations: 44.27 MiB)
@btime solve(prob, DP5(), reltol = 1e-9, abstol=1e-9, saveat=0.1);   # 2.039 s (269 allocations: 43.14 MiB)
@btime solve(prob, RK4(), relto1 = 1e-9, abstol = 1e-9, saveat=0.1); # 2.834 s (268 allocations: 43.14 MiB)

@btime solve(prob, ROCK4(), reltol = 1e-6, abstol=1e-6, saveat=0.1); # 351.243 ms (2708 allocations: 146.07 MiB)
@btime solve(prob, Tsit5(), reltol = 1e-6, abstol=1e-6, saveat=0.1); # 2.189 s (275 allocations: 44.27 MiB)
@btime solve(prob, DP5(), reltol = 1e-6, abstol=1e-6, saveat=0.1);   # 2.069 s (269 allocations: 43.14 MiB)
@btime solve(prob, RK4(), relto1 = 1e-6, abstol = 1e-6, saveat=0.1); # 2.887 s (268 allocations: 43.14 MiB)

@btime solve(prob, ROCK4(), reltol = 1e-10, abstol=1e-10, saveat=0.1); # 1.962 s (5900 allocations: 744.69 MiB)
@btime solve(prob, Tsit5(), reltol = 1e-10, abstol=1e-10, saveat=0.1); # 2.462 s (275 allocations: 44.27 MiB)
@btime solve(prob, DP5(), reltol = 1e-10, abstol=1e-10, saveat=0.1);   # 2.133 s (269 allocations: 43.14 MiB)
@btime solve(prob, RK4(), relto1 = 1e-10, abstol = 1e-10, saveat=0.1); # 2.882 s (268 allocations: 43.14 MiB)


@inline function checked_sqrt(x::V) where {W, T, V <: Union{SIMDPirates.Vec{W,T},SIMDPirates.SVec{W,T}}}
    y = SIMDPirates.vsqrt(x)
    SIMDPirates.vifelse(SIMDPirates.visfinite(y), y, SIMDPirates.vbroadcast(V, T(0.005)))
end
@inline function checked_sqrt(x::T) where {T <: Number}
    y = @fastmath sqrt(x)
    isfinite(T) ? y : T(0.005)
end

@generated function process_big_prop_points!(X::ResizableMatrix{T}, BPP::AbstractMatrix{T}) where T
    quote
        resize!(X, size(BPP,1))
        @vectorize $T for i ∈ 1:size(BPP,1)
            y1 = BPP[i,1]
            y2 = BPP[i,2]
            y3 = BPP[i,3]

            R11 = one(T) / checked_sqrt(BPP[i,5])
            # R11 = SIMDPirates.vifelse(SIMDPirates.visfinite(R11)
            x1 = R11 * y1
            L21 = R11 * BPP[i,6]
            L31 = R11 * BPP[i,8]
            R22 = one(T) / checked_sqrt(BPP[i,7] - L21*L21)
            x2 = R22 * (y2 - L21*x1)

            L32 = R22 * (BPP[i,9] - L21 * L31)
            R33 = one(T) / checked_sqrt(BPP[i,10] - L31*L31 - L32*L32)

            X[i,1] = x1
            X[i,2] = x2
            X[i,3] = R33 * (y3 - L31*x1 - L32*x2)
        end
    end
end

@generated function process_big_prop_points!(X::AbstractMatrix{T}, BPP::AbstractMatrix{T}) where T
    quote
        @vectorize $T for i ∈ 1:size(BPP,1)
            y1 = BPP[i,1]
            y2 = BPP[i,2]
            y3 = BPP[i,3]

            R11 = one(T) / checked_sqrt(BPP[i,5])
            # R11 = SIMDPirates.vifelse(SIMDPirates.visfinite(R11)
            x1 = R11 * y1
            L21 = R11 * BPP[i,6]
            L31 = R11 * BPP[i,8]
            R22 = one(T) / checked_sqrt(BPP[i,7] - L21*L21)
            x2 = R22 * (y2 - L21*x1)

            L32 = R22 * (BPP[i,9] - L21 * L31)
            R33 = one(T) / checked_sqrt(BPP[i,10] - L31*L31 - L32*L32)

            X[i,1] = x1
            X[i,2] = x2
            X[i,3] = R33 * (y3 - L31*x1 - L32*x2)
        end
    end
end

function generate_rank1covariances!(rank1covs::AbstractVector{InverseWishart{T}}, X::AbstractMatrix{T}) where T
    resize!(rank1covs, size(X,1))
    @inbounds for i ∈ 1:size(X,1)
        rank1covs[i] = InverseWishart(X[i,1], X[i,2], X[i,3])
    end
end

function process_BPP!(X, rank1covs, mahals, BPP)
    process_big_prop_points!(X, BPP)
    generate_rank1covariances!(rank1covs, X)
    MahalanobisDistances!(mahals, X)
    nothing
end

function process_BPP!(X, rank1covs, mahals, icc::InvCholCovar, BPP, t)
    process_big_prop_points!(X, BPP)
    decorrelate_data!(icc, X, t)
    generate_rank1covariances!(rank1covs, X)
    MahalanobisDistances!(mahals, X)
    nothing
end


using JFVM
using PyPlot
include("CRT.jl")

#Calculating error as a function of timesteps

# define the domain and constants
ϕ  = 0.4 # porosity
u  = 1.0/(3600*24) # m/s, eq. 1 m/day
D  = [1e-9, 1e-9, 1e-9] # m^2/s diffusivity
K_eq = 10.0
k_f = 100.0
k_b = k_f/K_eq
nu_A = 1.0 # stoichiometric coefficient of component A
nu_B = 1.0 # stoichiometric coefficient of component B
nu_C = 1.0 # stoichiometric coefficient of component C
a = 1.0 #order of reaction
b = 1.0 #order of reaction
c = 1.0 #order of reaction
N_comp = 3 # number of components

# Defining time-stepping parameters
tol = 1e-10 # convergence tolerance

# boundary conditions
C_bc = [1.0, 0.5, 0.0] # boundary concentration
C_0 = zeros(N_comp) # initial concentration

NX = 100:100:1000
time_array = zeros(5, length(NX))

for i=1:length(NX)
  Lx = 1.0 # [m]
  Nx = NX[i]  # number of grids
  m = createMesh1D(Nx, Lx)  # mesh

  #Time-stepping parameters
  dt = Lx/(u/ϕ)/Nx
  T = Nx/2*dt

  tic()
  C_SNIA = CRT.CRT_SNIA(m, Nx, ϕ, u, D, K_eq, nu_A, nu_B, nu_C, C_bc, C_0, N_comp, dt, T)
  time_array[1, i] = toc()

  tic()
  C_SIA = CRT.CRT_SIA(m, Nx, ϕ, u, D, K_eq, nu_A, nu_B, nu_C, C_bc, C_0, N_comp, dt, T, tol)
  time_array[2, i] = toc()

  tic()
  C_eps = CRT.CRT_GIA_eps(m, Nx, ϕ, u, D, K_eq, nu_A, nu_B, nu_C, C_bc, C_0, N_comp, dt, T, tol)
  time_array[3, i] = toc()

  tic()
  C_GIA_kin_eq = CRT.CRT_GIA_kin_eq(m, Nx, ϕ, u, D, K_eq, nu_A, nu_B, nu_C, a, b, c, C_bc, C_0, N_comp, dt, T, tol)
  time_array[4, i] = toc()

  tic()
  C_GIA_kin_fb = CRT.CRT_GIA_kin_fb(m, Nx, ϕ, u, D, k_f, k_b, nu_A, nu_B, nu_C, a, b, c, C_bc, C_0, N_comp, dt, T, tol)
  time_array[5, i] = toc()

  println(i)
end

figure()
plot(NX, time_array[1,:])
plot(NX, time_array[2,:])
plot(NX, time_array[3,:])
plot(NX, time_array[4,:])
plot(NX, time_array[5,:])
grid("on")
legend(["SNIA", "SIA", "GIA ("L"$\epsilon$"")", "GIA ("L"$K_{eq}$"")", "GIA ("L"$k_fk_b$"")"])
xlabel(L"$N_x$")
ylabel("Time [s]""")
savefig("time_grid_plot.png")

#CRT_SIA title:
#L"$\phi\frac{\partial C_{\alpha}}{\partial t}+\nabla\cdot(uC_{\alpha})-\nabla\cdot(D_{\alpha}\nabla C_{\alpha})+\nu_{\alpha}\epsilon=0,\quad K_{eq}=$""$K_eq"" (SIA)"

#CRT_SNIA title:
#L"$\phi\frac{\partial C_{\alpha}}{\partial t}+\nabla\cdot(uC_{\alpha})-\nabla\cdot(D_{\alpha}\nabla C_{\alpha})+\nu_{\alpha}\epsilon=0,\quad K_{eq}=$""$K_eq"" (SNIA)"

#CRT_GIA_eps title:
#L"$\phi\frac{\partial C_{\alpha}}{\partial t}+\nabla\cdot(uC_{\alpha})-\nabla\cdot(D_{\alpha}\nabla C_{\alpha})+\nu_{\alpha}\epsilon=0,\quad K_{eq}=$""$K_eq"" (FIM)"

#CRT_GIA_kin_eq title:
#L"$\phi\frac{\partial C_{\alpha}}{\partial t}+\nabla\cdot(uC_{\alpha})-\nabla\cdot(D_{\alpha}\nabla C_{\alpha})+\nu_{\alpha}(K_{eq}C_A^aC_B^b-C_C^c)=0,\quad K_{eq}=$""$K_eq"" (FIM)"

#CRT_GIA_kin_fb title:
#L"$\phi\frac{\partial C_{\alpha}}{\partial t}+\nabla\cdot(uC_{\alpha})-\nabla\cdot(D_{\alpha}\nabla C_{\alpha})+\nu_{\alpha}(k_{f}C_A^aC_B^b-k_{b}C_C^c)=0,\quad K_{eq}=$""$K_eq"" (FIM)"


module PolyOpt

using ClusteredLowRankSolver, BasesAndSamples, AbstractAlgebra

export min_f

function invariant_basis(x,y,z, d)
    # create a vector with a precise type
    v = [(x*y*z)^0]
    for deg=1:d, j=0:div(deg,3), i=0:div(deg-3j,2)
        push!(v, (x+y+z)^(deg-2i-3j) * (x*y+y*z+z*x)^i * (x*y*z)^j)
    end
    return v
end

function min_f(d)
    obj = Objective(0, Dict(), Dict(:M => 1))

    FF = RealField
    R, (x,y,z) = PolynomialRing(FF, ["x", "y", "z"])
    # The polynomial f:
    f =  x^4 + y^4 + z^4 - 4x*y*z + x + y + z

    # An invariant basis up to degree d:
    basis = invariant_basis(x, y, z, 2d)
    # For the sum-of-squares polynomials we have to
    # select elements of the basis based on the degree
    degrees = [total_degree(p) for p in basis]

    # generate samples and a good basis
    cheb_points = [vcat(sample_points_chebyshev(2d+k)...) for k=0:2]
    samples_grid = [[cheb_points[1][i+1], cheb_points[2][j+1], cheb_points[3][k+1]]
        for i=0:2d for j=0:2d+1 for k=0:2d+2]
    basis, samples = approximatefekete(basis, samples_grid)

    psd_dict = Dict()
    symmetry_weights = [[[R(1)]],
                        [[R((x-y)*(y-z)*(z-x))]],
                        [[1/sqrt(FF(2))*(2x-y-z),1/sqrt(FF(2))*(2y*z-x*z-x*y)],
                            [sqrt(FF(3)/FF(2))*(y-z),sqrt(FF(3)/FF(2))*(x*z-x*y)]]]
    for swi=1:length(symmetry_weights)
        rank = length(symmetry_weights[swi])
        # This has in general too many entries, so we remove the ones with too high degree.
        vecs = [kron(symmetry_weights[swi][r], basis) for r=1:rank]
        for r=1:rank
            # the length of the symmetric basis we tensored the symmetric weight with
            l = length(basis)
            # we keep the elements with degree at most d
            keep_idx = [i for i=1:length(vecs[r]) if total_degree(symmetry_weights[swi][r][div(i-1,l)+1]) + degrees[(i-1)%l+1] <= d]
            vecs[r] = vecs[r][keep_idx]
        end
        vecs = [v for v in vecs if length(v) > 0]
        if length(vecs) > 0
            psd_dict[Block((:trivariatesos,swi))] = LowRankMatPol([R(1) for r=1:length(vecs)], vecs)
        end
    end

    constr = Constraint(f, psd_dict, Dict(:M => R(1)), samples)
    pol_problem = LowRankPolProblem(true, obj, [constr])
    sdp = ClusteredLowRankSDP(pol_problem)

    solvesdp(sdp)
end

end


function LinearAlgebra.svd(A::Array{T,2}) where T<:ForwardDiff.Dual
    m = size(A,1)
    n = size(A,2)
    @assert m >= n

    Av = map(ForwardDiff.value,A)
    Uv,sv,Vv = svd(Av)

    Ap = map(ForwardDiff.partials,A)
    UvtApVv = Uv'*Ap*Vv
    sp = diag(UvtApVv)

    # if allunique(sv)
        # If all unique, then gradient well defined
        Sv = Matrix(Diagonal(sv))
        svi = 1. ./ sv
        svi[isinf.(svi)] .= 0.
        Svi = Matrix(Diagonal(svi))

        F = zeros(eltype(Av),m,n)
        for i = 1:m
            for j = 1:n
                if i!=j
                    F[i,j] = 1. / (sv[j]^2 - sv[i]^2)
                end
            end
        end
        F[isinf.(F)] .= 0.

        VvtApUv = Vv'*Ap[reshape(1:m*n,(m,n))']*Uv

        Up = Uv*(F .* (UvtApVv*Sv + Sv*VvtApUv)) + (I - Uv*Uv')*Ap*Vv*Svi
        Vp = Vv*(F .* (Sv*UvtApVv+VvtApUv*Sv)) + (I - Vv*Vv')*Ap[reshape(1:m*n,(m,n))']*Uv*Svi
    # else
    #     # Otherwise need to least-squares
    #     dU = zeros(size(Uv,1),size(Uv,2),size(A,1),size(A,2))
    #     dV = zeros(size(Vv,1),size(Vv,2),size(A,1),size(A,2))
    #     D = Array{Float64, 2}(undef, 2, 2)
    #     u = Array{Float64, 1}(undef, 2)
    #     for i = 1:m
    #         for j = 1:n
    #             for k = 1:m
    #                 for l = 1:n
    #                     D[1,1] = sv[l]
    #                     D[1,2] = sv[k]
    #                     D[2,1] = sv[k]
    #                     D[2,2] = sv[l]
    #                     u[1] = Uv[i,k]*Vv[j,l]
    #                     u[2] =  -Uv[i,l]*Vv[j,k]
    #                     if (sv[l] == sv[k])
    #                         Ω = pinv(D) * u
    #                     else
    #                         Ω = D \ u
    #                     end
    #                     dU[k,l,i,j] = Ω[1]
    #                     dV[k,l,i,j] = Ω[2]
    #                 end
    #             end
    #         end
    #     end
    #     Up = zeros(typeof(A[1].partials),size(Uv))
    #     Vp = zeros(typeof(A[1].partials),size(Vv))
    #     for k = 1:size(Up,1)
    #         for l = 1:size(Up,2)
    #             for i = 1:m
    #                 for j = 1:n
    #                     Up[k,l] += dU[k,l,i,j] * A[i,j].partials
    #                 end
    #             end
    #         end
    #     end
    #     for k = 1:size(Vp,1)
    #         for l = 1:size(Vp,2)
    #             for i = 1:m
    #                 for j = 1:n
    #                     Vp[k,l] += dV[k,l,i,j] * A[i,j].partials
    #                 end
    #             end
    #         end
    #     end
    # end

    U = map(T,Uv,Up)
    s = map(T,sv,sp)
    V = map(T,Vv,Vp)

    return U, s, V
end


#PGPE nonlinearity
#out of place
function nlin(c)
    ψ = c2x(c,tinfo)
    return x2c(abs2.(ψ).*ψ,tinfo)
end

#in place
function nlin!(dc,c,tinfo)
    c2x!(ψ,c,tinfo)
    x2c!(dc,abs2.(ψ).*ψ,tinfo)
end

#dPGPE in reservoir frame
#out of place
function Lgp(c,p,t)
     return P.*(-im*(1-im*γ)*((en - μ).*c .+ g*nlin(c)))
end

#in place
function Lgp!(dc,c,p,t)
    nlin!(dc,c,tinfo)
    dc .= P.*(-im*(1-im*γ)*((en - μ).*c .+ g*dc))
end


using LinearAlgebra
using Statistics
using LinearAlgebra.BLAS
import StatsBase.percentile

const Mat{T} = AbstractArray{T, 2}
const Tup{T} = Tuple{Mat{T}, Mat{T}}

abstract type AbstractGBM{T, I} end

"""
General Boltzmann Machine, including parameter matrix (element type T)
and a mask matrix (element type I) for decimation.
"""
mutable struct GBM{T, I} <: AbstractGBM{T, I}
    W::Matrix{T}
    b::Matrix{T}
    W_mask::Array{I, 2}
    PLrange::Dict
end

"""
    GBM(n_unit::Int)

Construct GBM with number of unit `n_unit`.

 * T - type of GBM parameters [weights and biases;] (default: Float64)
 * I - type of Mask-matrix for decimation (default: Int64)
"""
function GBM(n_unit::Int)
    GBM{Float64, Int64}(
        rand(n_unit, n_unit),
        rand(n_unit, 1),
        ones(Int64, (n_unit, n_unit)) - Matrix(I, n_unit, n_unit),
        Dict("min" => 0.0, "max" => 0.0)
    )
end

## gradient calculation

"""
    calc_grad(gmb::GBM, X::AbstractArray, ctx::Dict)

Calculate gradeint of log-pseudo-likelihood with respect to GBM parameter.

# References
- A. Hyvärinen, Neural Comput. **18**, 2283 (2006)
"""
function calc_grad(gbm::GBM{T}, X::Mat{T}, ctx::Dict)::Tup{T} where T

    alpha = T(ctx[:alpha])
    beta = T(ctx[:beta])

    dev = X' - tanh.(beta * (gbm.W * X' .+ gbm.b))
    dW = @getarray(ctx, :dW_buf, size(gbm.W), similar(gbm.W))

    # same as: dW = beta .* dev * X / size(X, 1)
    gemm!('N', 'N', beta / size(X, 1), dev, X, T(0.0), dW)
    symmetrize!(dW)
    axpy!(-alpha, sign.(gbm.W), dW)

    # gradient for bias params.
    db = beta * mean(dev, dims=2)

    return (dW, db)
end

"""
    update(gbm::GBM, grads::Tuple, ctx::Dict)

Update GBM parameters with SGD method.

Update weight parameters of GBM using provided tuple `dtheta = (dW, db)`
of parameter gradients. Optimal method is Stochastic Gradient Descent [SGD].
"""
function update(gbm::GBM{T}, grads::Tup{T}, ctx::Dict) where T
    dW, db = grads
    lr = T(ctx[:lr])

    # same as: gbm.W += lr * dW
    axpy!(lr, dW, gbm.W)
    axpy!(lr, db, gbm.b)

    symmetrize!(gbm.W)
    gbm.W .*= gbm.W_mask
end

## Pseudo likelihood calculation

"""
    pseudo_likelihood(gbm::GBM, X::AbstractArray, ctx::Dict)

Calculate pseudo-likelihood[PL] with spin-example data `X`.
"""
function pseudo_likelihood(gbm::GBM{T}, X::Mat{T}, ctx::Dict)::T where T
    alpha = T(ctx[:alpha])
    beta = T(ctx[:beta])
    H = tr(X * gbm.W * X') + sum(X * gbm.b)
    tmp = sum(log.(cosh.(beta * (gbm.W * X' .+ gbm.b))))
    PL = (beta * H - tmp) / size(X, 1)
    return PL - alpha * sum(abs.(gbm.W))
end

"""
    fit(gbm::GBM, X::AbstractArray; delta=1e-1 [, opts])

Fit the weight parameters of GBM to data `X`.

# Arguments
- `delta::Float`: Threshold for thinning out parameters.

Options that can be provided in the `opts` dictionary:

* `:alpha` - Decay rate of the L1 regularization term. (default: 0.01)
* `:beta` - Inverse temperature. (default: 1.0)
* `:lr` - Learning rate. (default: 1.0)
* `:epochs` - Maximum permissible number of step in SGD. (default: 20)
* `:n_iter` - Number of learning repetitions. (default: 100)

# References
- E. Aurell & M. Ekeberg, Phys. Rev. Lett. **108**, 090201 (2012).

# Examples
```jldoctest
julia> using InverseIsing

julia> samples = [1 -1 -1;] # Spin configuration.

julia> model = GBM(3) # Set the number of units.

julia> fit(model, samples)

julia> W = infer(model); output = decode(W)

julia> output
OrderedCollections.OrderedDict{Tuple{Int64,Int64},Int64} with 3 entries:
  (1, 2) => -1
  (1, 3) => -1
  (2, 3) => 1
```
    The above example means that the interaction between (1, 2) and (1, 3) is
    antiferromagnetic bond and only (2, 3) is ferromagnetic bond.

"""
function fit(gbm::GBM{T}, X::Mat; delta::T=1e-1, opts...) where {T, I}

    ctx = copy(Dict{Any, Any}(opts))
    option_checker!(ctx)
    epochs = ctx[:epochs]
    max_step = ctx[:n_iter]

    X_train = trans_type(T, X)
    variable_check(X_train)

    for n in 1:max_step

        old_likelihood = pseudo_likelihood(gbm, X_train, ctx)
        update(gbm, calc_grad(gbm, X_train, ctx), ctx)

        if n % epochs == 0
            gbm.W_mask[abs.(gbm.W) .< delta] .= 0
            gbm.W_mask .*= gbm.W_mask'
        end

        if isapprox(pseudo_likelihood(gbm, X_train, ctx), old_likelihood)
            break
        end
    end
    return ctx
end

### sub-utility functions ###
"""
    coef(gbm::GBM)

Get weight marix of a trained GBM.

# Examples
```jldoctest
julia> model = GBM(3)

julia> coef(model)
3×3 Array{Float64,2}:
 0.86902   0.4058     0.568231
 0.700701  0.0892165  0.399372
 0.224095  0.234413   0.380341
```
"""
function coef(gbm::GBM{T}) where T
    W = copy(gbm.W)
    # translate -0.0 to 0.0
    W[W .== -T(0.0)] .= T(0.0)
    return W
end

"""
    weights(gbm::GBM)

Same as coef. See also: [`coef`](@ref)
"""
weights = coef

"""
    infer(gbm::GBM)

Infer GBM parameter matrix using trained weight.

# Examples
```jldoctest
julia> using InverseIsing

julia> model = GBM(3)

julia> infer(model)
3×3 Array{Int64,2}:
 1  1  1
 1  1  1
 1  1  1

julia> fit(model, [1 -1 -1;])

julia> infer(model)
3×3 Array{Int64,2}:
  0  -1  -1
 -1   0   1
 -1   1   0
```
"""
function infer(gbm::GBM{T, I}) where {T, I}
    pred = heaviside(coef(gbm))
    return convert(Array{I}, pred)
end


using ApproxFun
using SpecialExponentialFamilies
using Empirikos
using Random
using DiffResults
using ForwardDiff
using UnPack
using StatsBase
using Expectations
using Splines2
using GLM
using Plots
using LinearAlgebra
using Expectations
using Test
using QuadGK

unif_measure = Uniform(-3.0,3.0)

_df = 3
expfam = ExponentialFamily(; basemeasure=unif_measure, df=_df)
@test length(expfam.Q(1.0)) == _df

expfam_ones = expfam(ones(_df))
@test quadgk(μ->pdf(expfam_ones, μ), -3, 3)[1] ≈ 1.0
n = 10_000

Random.seed!(10)
Zs = rand(expfam_ones, n)

disc_Z = summarize(discretize(Zs))

@test sum(weights(disc_Z)) == n

lindsey_fit = StatsBase.fit(LindseyMethod(ef=expfam), Zs)
lindsey_fit.α

Zs_noise = StandardNormalSample.(Zs .+ 1.0 * randn(n))
Zs_noise2 = NormalSample.(response.(Zs_noise), 1.0)
Zs_discr = summarize(discretize(Zs_noise))

tmp = fit(PenalizedMLE(ef=expfam, optim_options=Optim.Options(show_trace=true, show_every=1, iterations=20)), Zs_noise)
tmp2 = fit(PenalizedMLE(ef=expfam, optim_options=Optim.Options(show_trace=true, show_every=1, iterations=20)), Zs_noise2)
tmp3 = fit(PenalizedMLE(ef=expfam, optim_options=Optim.Options(show_trace=true, show_every=1, iterations=20)), Zs_discr)

tmp2.α_opt
tmp3.α_opt
tmp_lindsey = fit(LindseyMethod(ef=expfam), Zs)
tmp.α_opt == tmp2.α_opt

tmp2.α

sqrt.(diag(tmp3.α_covmat))

efm = EFMarginalDistribution3(expfam_ones, StandardNormalSample())
@test quadgk(x -> pdf(efm, x), -Inf, +Inf)[1] ≈ 1.0
@test cdf(efm, Inf) == 1.0
@test ForwardDiff.derivative(x->cdf(efm,x), 1.0) ≈ pdf(efm, 1.0)
@test pdf(expfam_ones, StandardNormalSample(-2.0)) == pdf(efm, -2.0)
#function marginalize(Z::AbstractNormalSample, prior::Normal)
#end


module JuliaNature

# Write your package code here.

using Statistics

struct Bat_model
    pop_size::Int
    dims::Int
    obj_func
    position
    velocity
    freq_pulse
    pulse_rate
    loudness
    freq_max
    freq_min
    loud_max
    loud_min
    type # 0 - min , 1 - max
end

function initialize(size::Int, dims::Int, obj_func, velocity_range, posit_range, freq_range, loudness_range, type)
    bat_velocity = [rand(velocity_range[1]:.001:velocity_range[2],dims) for i in 1:size]
    bat_position = [rand(posit_range[1]:.001:posit_range[2],dims) for i in 1:size]
    frequencies = [rand(freq_range[1]:.001:freq_range[2],dims) for i in 1:size]
    pulse_rates = [rand(0:.001:1,dims) for i in 1:size]
    loudness = [rand(loudness_range[1]:.001:loudness_range[2],dims) for i in 1:size]
    return Bat_model(size, dims, obj_func, bat_position, bat_velocity,  frequencies, pulse_rates, loudness, [freq_range[1]
    for i in 1:dims], [freq_range[2] for i in 1:dims] , [loudness_range[1] for i in 1:dims], [loudness_range[2] for i in 1:dims], type)
end

function optimize(algo::Bat_model, iter=50, tol= 0.3)
    t = 1
    best_ind = 1
    println(best_ind)
    if algo.type == 0
            best = typemax(typeof(1.0))
        else
            best = typemin(typeof(1.0))
        end
        plus_0 = algo.pulse_rate
        for i in 1: algo.pop_size
            if algo.type == 0
                if algo.obj_func(algo.position[i]...) < best
                    best = algo.obj_func(algo.position[i]...)
                    best_ind = i
                end
            else
                if algo.obj_func(algo.position[i]...) > best
                    best = algo.obj_func(algo.position[i]...)
                    best_ind = i
                end
            end
        end
    println(algo.position[best_ind])
    println(best_ind)
    while t < iter
        alpha = 0.9
        yamma = 0.9
        for i in 1: algo.pop_size
            beta = rand(0:.001:1, algo.dims)
            algo.freq_pulse[i] = algo.freq_min + (algo.freq_max - algo.freq_min) .*  beta
            algo.velocity[i] = algo.velocity[i] + (algo.position[i] - algo.position[best_ind]) .* algo.freq_pulse[i]
            algo.position[i] = algo.position[i] + algo.velocity[i]
            r = rand(0:.01:1, algo.dims)
            if r > algo.pulse_rate[i]
                algo.position[i] = algo.position[best_ind] + 0.01 .* r
            end
            if beta < algo.loudness[i]
                continue
            end
            if algo.type == 0
                if algo.obj_func(algo.position[i]...) < best
                    best = algo.obj_func(algo.position[i]...)
                    best_ind = i
                end
            else
                if algo.obj_func(algo.position[i]...) > best
                    best = algo.obj_func(algo.position[i]...)
                    best_ind = i
                end
            end

            #scaling = 0.01
            #avg = mean(algo.position)
            #e = rand(-1:0.01:1, 1)
            #algo.position[i] = algo.position[i] + scaling .* e .* avg
            #algo.loudness[i] = alpha * algo.loudness[i]
            #algo.pulse_rate[i] = plus_0[i] * (1 - exp(-yamma * t))
        end
        t += 1
        if abs(sum(algo.position[best_ind] - [1, 3])) < tol
            break
        end
    end
    return algo.position[best_ind]
end

end # module


function solve()
    v, t, s, d = [parse(Int, x) for x in split(readline())]
    if t ≤ d / v ≤ s
        "No"
    else
        "Yes"
    end
end

println(solve())


include("./octahedron_test.jl")
include("./tetramihexahedron_test.jl")
include("./octahemioctahedron_test.jl")
include("./cubohemioctahedron_prototype.jl")


#functions for calling the matrix vector products
#i.e. function A(y_in,Nx,Ny,h) computes the product A*u and stores it in y

#Right now only for  the second-order accurate SBP operators for constant coefficients from
#Mattsson and Nordstrom, 2004.

#Written for the 2D domain (x,y) \in (a, b) \times (c, d),
#stacking the grid function in the vertical direction,
#so that, e.g., \partial u / \partial x \approx D \kron I and
#\partial u / \partial y \approx I \kron D
#the faces of the 2D domain are label bottom to top, left to right, i.e.
#side 1 is y = c, side 2 is y = d
#side 3 is x = a, side 4 is x = b

#D2x and D2y compute approximations to
#\partial^2 u / \partial x^2 and
#\partial^2 u / \partial y^2, respectively

#FACEtoVOL(face,u_face,Nx,Ny) maps the face value u_face to a full length vector at the
#nodes corresponding to face

#BxSx \approx the traction on faces 3 and 4
#BySy \approx the traction on faces 1 and 2

#BxSx_tran and BySy_tran are the transposes of BxSx and BySy

# variables order: y_in, Nx, Ny, N, hx, hy, y_out(container)

##
# y_in: input Array
# y_out: output Array
##

function D2x_beta!(y_in::Array{Float64,1}, Nx::Int64, Ny::Int64, N::Int64, hx::Float64, hy::Float64, y_out::Array{Float64,1})
	#N = Nx*Ny
	#hx = Float64(1/(Nx-1))
	#hy = Float64(1/(Ny-1))
	#y = similar(y_in)
	#y = similar(y_in)
	@inbounds  for idx = 1:Ny
		y_out[idx] = (y_in[idx] - 2*y_in[Ny + idx] + y_in[2*Ny + idx]) / hx^2
	end

	@inbounds for idx1 = Ny+1:N-Ny
		y_out[idx1] = (y_in[idx1 - Ny] - 2 * y_in[idx1] + y_in[idx1 + Ny]) / hx^2
	end

	@inbounds for idx2 = N-Ny+1:N
		y_out[idx2] = (y_in[idx2 - 2*Ny] -2 * y_in[idx2 - Ny] + y_in[idx2]) / hx^2
	end
	# return y_out
end

function D2y_beta!(y_in::Array{Float64,1}, Nx::Int64, Ny::Int64, N::Int64, hx::Float64, hy::Float64, y_out::Array{Float64,1})
	#N = Nx*Ny
	#hx = Float64(1/(Nx-1))
	#hy = Float64(1/(Ny-1))
	@inbounds for idx = 1:Ny:N-Ny+1
		y_out[idx] = (y_in[idx] - 2 * y_in[idx + 1] + y_in[idx + 2]) / hy^2
	end

	@inbounds for idx1 = Ny:Ny:N
		y_out[idx1] = (y_in[idx1 - 2] - 2 * y_in[idx1 .- 1] + y_in[idx1]) / hy^2
	end

	@inbounds for j = 1:Nx
		@inbounds for idx = 2+(j-1)*Ny:j*Ny-1
			y_out[idx] = (y_in[idx - 1] - 2 * y_in[idx] + y_in[idx + 1]) / hy^2
		end
	end

	# return y_out

end

function Dx_beta!(y_in::Array{Float64,1}, Nx::Int64, Ny::Int64, N::Int64, hx::Float64, hy::Float64, y_out::Array{Float64,1})

	@inbounds for idx = 1:Ny
		y_out[idx] = (y_in[idx + Ny] - y_in[idx]) / hx
	end

	@inbounds for idx1 = Ny+1:N-Ny
		y_out[idx1] = (y_in[idx1 + Ny]-y_in[idx1 - Ny]) / (2*hx)
	end

	@inbounds for idx2 = N-Ny+1:N
		y_out[idx2] = (y_in[idx2]-y_in[idx2 .- Ny]) ./ hx
	end
	# return y_out
end

function Dy_beta!(y_in::Array{Float64,1}, Nx::Int64, Ny::Int64, N::Int64, hx::Float64, hy::Float64, y_out::Array{Float64,1})

	@inbounds for idx = 1:Ny:N-Ny+1
		y_out[idx] = (y_in[idx + 1] - y_in[idx]) / h
	end

	@inbounds for idx = Ny:Ny:N
		y_out[idx] = (y_in[idx] - y_in[idx - 1]) /h
	end

	@inbounds for j = 1:Nx
		@inbounds for idx = 2+(j-1)*Ny:j*Ny-1
		 	y_out[idx] = (y_in[idx + 1] - y_in[idx - 1]) / (2*h)
		end
	end
	# return y_out
end

function Hxinv_beta!(y_in::Array{Float64,1}, Nx::Int64, Ny::Int64, N::Int64, hx::Float64, hy::Float64, coef::Float64, y_out::Array{Float64,1})
	#N = Nx*Ny
	#y = similar(y_in)
	@inbounds for idx = 1:Ny
		y_out[idx] = coef * (2*y_in[idx]) * (1/hx)
	end

	@inbounds for idx1 = Ny+1:N-Ny
		y_out[idx1] = coef * (1*y_in[idx1]) * (1/hx)
	end

	@inbounds for idx2 = N-Ny+1:N
		y_out[idx2] = coef * (2*y_in[idx2]) * (1/hx)
	end

	# return y_out
end

function Hyinv_beta!(y_in::Array{Float64,1}, Nx::Int64, Ny::Int64, N::Int64, hx::Float64, hy::Float64, coef::Float64, y_out::Array{Float64,1})
	#N = Nx*Ny
	#y = similar(y_in)
	#hx = Float64(1/(Nx-1))
	#hy = Float64(1/(Ny-1))
	@inbounds for idx = 1:Ny:N-Ny+1
		y_out[idx] = coef * (2*y_in[idx]) * (1/hy)
	end

	@inbounds for idx1 = Ny:Ny:N
		y_out[idx1] = coef * (2*y_in[idx1]) * (1/hy)
	end

	@inbounds for i = 1:Nx
		@inbounds for idx2 = 2+(i-1)*Ny:i*Ny-1
			y_out[idx2] = coef * (y_in[idx2]) * (1/hy)
		end
	end
	# return y_out
end

function Hx_beta!(y_in::Array{Float64,1}, Nx::Int64, Ny::Int64, N::Int64, hx::Float64, hy::Float64,coef::Float64 ,y_out::Array{Float64,1})
	@inbounds for idx = 1:Ny
		y_out[idx] = coef * hx*y_in[idx]/2
	end

	@inbounds for idx1 = Ny+1:N-Ny
		y_out[idx1] = coef* hx*y_in[idx1]
	end

	@inbounds for idx2 = N-Ny+1:N
		y_out[idx2] = coef * hx*y_in[idx2]/2
	end
	# return y_out
end

function Hy_beta!(y_in::Array{Float64,1}, Nx::Int64, Ny::Int64, N::Int64, hx::Float64, hy::Float64, y_out::Array{Float64,1})
	#N = Nx*Ny
	#y = similar(y_in)

	@inbounds for idx = 1:Ny:N-Ny+1
		y_out[idx] = hy*y_in[idx]/2
	end

	@inbounds for idx1 = Ny:Ny:N
		y_out[idx1] = hy*y_in[idx1]/2
	end

	@inbounds for i = 1:Nx
		@inbounds for idx2 = 2 + (i-1)*Ny:i*Ny-1
			y_out[idx2] = hy*y_in[idx2]
		end
	end
	# return y_out
end

function FACEtoVOL_beta!(y_in_face::Array{Float64,1},  face::Int64, Nx::Int64, Ny::Int64, N::Int64, y_outs::Array{Array{Float64,1},1})
	if face == 1
		idx = 1:Ny:N-Ny+1
	elseif face==2
		idx = Ny:Ny:N
	elseif face==3
		idx = 1:Ny
	elseif face == 4
		idx = N-Ny+1:N
	else
	end
	y_out = y_outs[face]
	y_out[idx] = y_in_face
	# return y_out
end

function VOLtoFACE_beta!(y_in::Array{Float64,1}, face::Int64,Nx::Int64, Ny::Int64, N::Int64, y_outs::Array{Array{Float64,1},1}) ## Has some issue
	if face == 1
			idx = 1:Ny:N-Ny+1
	elseif face == 2
			idx = Ny:Ny:N
	elseif face == 3
			idx = 1:Ny
	elseif face == 4
			idx = N-Ny+1:N
	else
	end
    y_out = y_outs[face]
	y_out[idx] = y_in[idx]

	# return y_out
end

function Bx_beta!(Nx::Int64,Ny::Int64, N::Int64, y_out::Array{Float64,1})

	@inbounds for idx=1:Ny
		y_out[idx] = -1
	end

	@inbounds for idx = N-Ny+1:N
		y_out[idx] = 1
	end
	# return y_out
end

function By_beta!(Nx::Int64,Ny::Int64,N::Int64,y_By::Array{Float64,1})
	@inbounds for idx = 1:Ny:N-Ny+1
		y_By[idx] = -1
	end

	@inbounds for idx1 = Ny:Ny:N
		y_By[idx1] = 1
	end
	# return y_By
end

function BxSx_beta!(y_in::Array{Float64,1},Nx::Int64,Ny::Int64,N::Int64,hx::Float64,hy::Float64,y_out::Array{Float64,1})
	@inbounds for idx = 1:Ny
		y_out[idx] = (1/hx) * (1.5 * y_in[idx] - 2 * y_in[idx + Ny] + 0.5 * y_in[idx + 2*Ny])
		y_out[N-Ny + idx] = (1/hx) * (0.5 * y_in[N-3*Ny + idx] - 2 * y_in[N-2*Ny + idx] + 1.5 * y_in[N-Ny + idx])
	end
	# return y_out
end

function BySy_beta!(y_in::Array{Float64,1}, Nx::Int64, Ny::Int64, N::Int64, hx::Float64, hy::Float64,y_out::Array{Float64,1})
	@inbounds for idx = 1:Ny:N-Ny+1
		y_out[idx] = (1/hy) * (1.5 * y_in[idx] - 2 * y_in[idx .+ 1] + 0.5 * y_in[idx .+ 2])
	end

	@inbounds for idx = Ny:Ny:N
		y_out[idx] = (1/hy) * (0.5 * y_in[idx - 2] - 2 * y_in[idx - 1] + 1.5 * y_in[idx])
	end

	# return y_out
end

function BxSx_tran_beta!(y_in::Array{Float64,1},Nx::Int64,Ny::Int64,N::Int64,hx::Float64,hy::Float64,y_out::Array{Float64,1}) # be careful with += expression
	#hx = Float64(1/(Nx-1))
   	#hy = Float64(1/(Ny-1))
	@inbounds for idx1 = 1:Ny
		y_out[idx1] = (1.5 * y_in[idx1]) * (1/hx)

		# for idx = Ny+1:2*Ny
		y_out[idx1+Ny] = (-2 * y_in[idx1]) * (1/hx)
		# end

		# for idx  = 2*Ny+1:3*Ny
		y_out[idx1+2Ny] = (0.5 * y_in[idx1]) * (1/hx)
		# end
	end

	@inbounds for idxN = N-Ny+1:N
		y_out[idxN] = (1.5 * y_in[idxN]) * (1/hx)

		# for idx = N-2*Ny+1:N-Ny
		y_out[idxN-Ny] = (-2 * y_in[idxN]) * (1/hx)
		# end

		# for idx = N-3*Ny+1:N-2*Ny
		y_out[idxN-2Ny] = (0.5 * y_in[idxN]) * (1/hx)
		# end
	end
	# return y_out
end

function BySy_tran_beta!(y_in::Array{Float64,1}, Nx::Int64, Ny::Int64, N::Int64, hx::Float64, hy::Float64, y_out::Array{Float64,1}) # Be Careful about double foor loops
	@inbounds for idx1 = 1:Ny:N-Ny+1
		y_out[idx1] = (1.5 * y_in[idx1]) * (1/hy)

		# for idx = Ny+1:2*Ny
		y_out[idx1+1] = (-2 * y_in[idx1]) * (1/hy)
		# end

		# for idx  = 2*Ny+1:3*Ny
		y_out[idx1+2] = (0.5 * y_in[idx1]) * (1/hy)
		# end
	end

	@inbounds for idxN = Ny:Ny:N
		y_out[idxN] = (1.5 * y_in[idxN]) * (1/hy)

		# for idx = N-2*Ny+1:N-Ny
		y_out[idxN-1] = (-2 * y_in[idxN]) * (1/hy)
		# end

		# for idx = N-3*Ny+1:N-2*Ny
		y_out[idxN-2] = (0.5 * y_in[idxN]) * (1/hy)
		# end
	end

	# return y_out
end

# Previous function from deriv_ops.jl

# function FACEtoVOL(y_in_face, face, Nx, Ny)
# 	N = Nx*Ny
# 	y = zeros(N)

# 	if face == 1
# 		idx = 1:Ny:N-Ny+1
# 	elseif face == 2
# 		idx = Ny:Ny:N
# 	elseif face == 3
# 		idx = 1:Ny
# 	elseif face == 4
# 		idx = N-Ny+1:N
# 	else
# 	end

# 	y[idx] = u_face

# 	# return y

# end

# function VOLtoFACE(y_in, face, Nx, Ny)
# 	N = Nx*Ny
#         y = zeros(N)

#         if face == 1
#                 idx = 1:Ny:N-Ny+1
#         elseif face == 2
#                 idx = Ny:Ny:N
#         elseif face == 3
#                 idx = 1:Ny
#         elseif face == 4
#                 idx = N-Ny+1:N
#         else
#         end

# 	y[idx] = y_in[idx]
#         # return y
# end

# function Hxinv(y_in, Nx, Ny, h)
# 	N = Nx*Ny
# 	y = zeros(N)

# 	idx = 1:Ny
# 	y[idx] = (2*y_in[idx]) .* (1/h)

# 	idx = Ny+1:N-Ny
# 	y[idx] = (1*y_in[idx]) .* (1/h)

# 	idx = N-Ny+1:N
# 	y[idx] = (2*y_in[idx]) .* (1/h)

# 	# return y
# end

# function Hyinv(y_in, Nx, Ny, h)
# 	N = Nx*Ny
# 	y = zeros(N)

# 	idx = 1:Ny:N-Ny+1
# 	y[idx] = (2*y_in[idx]) .* (1/h)

# 	idx = Ny:Ny:N
# 	y[idx] = (2*y_in[idx]) .* (1/h)

# 	for i = 1:Nx
# 		idx = 2+(i-1).*Ny:i*Ny-1
# 		y[idx] = y_in[idx] .* (1/h)
# 	end

# 	# return y

# end

# function Hx(y_in, Nx, Ny, h)
# 	N = Nx*Ny
#         y = zeros(N)

#         idx = 1:Ny
# 	y[idx] = h .* (1/2)*y_in[idx]

#         idx = Ny+1:N-Ny
#         y[idx] = h .* 1*y_in[idx]

#         idx = N-Ny+1:N
# 	y[idx] = h .* (1/2)*y_in[idx]

#         # return y

# end

# function Hy(y_in, Nx, Ny, h)
# 	N = Nx*Ny
#         y = zeros(N)

#         idx = 1:Ny:N-Ny+1
# 	y[idx] = h .* (1/2)*y_in[idx]

#         idx = Ny:Ny:N
# 	y[idx] = h .* (1/2)*y_in[idx]

#         for i = 1:Nx
#                 idx = 2+(i-1).*Ny:i*Ny-1
#                 y[idx] = h .* y_in[idx]
#         end

#         # return y

# end

# function BxSx_tran(y_in, Nx, Ny, h)
# 	N = Nx*Ny
# 	y = zeros(N)

# 	idx1 = 1:Ny
# 	y[idx1] += (1.5 .* y_in[idx1]) .* (1/h)
# 	idx = Ny+1:2*Ny
# 	y[idx] += (-2 .* y_in[idx1]) .* (1/h)
# 	idx  = 2*Ny+1:3*Ny
# 	y[idx] += (0.5 .* y_in[idx1]) .* (1/h)

# 	idxN = N-Ny+1:N
# 	y[idxN] += (1.5 .* y_in[idxN]) .* (1/h)
# 	idx = N-2*Ny+1:N-Ny
# 	y[idx] += (-2 .* y_in[idxN]) .* (1/h)
# 	idx = N-3*Ny+1:N-2*Ny
# 	y[idx] += (0.5 .* y_in[idxN]) .* (1/h)

# 	# return y
# end


using Smg2s
using SparseArrays
using MatrixMarket
using LinearAlgebra

n = 10
diag_l = -5
diag_u = -3
nbone = 2

function f(idx::Integer, n = n)
    return 10 * cos(((idx-1) / 2) * 2 * π / n) + 5
end

spec = zeros(ComplexF64, n)
Spectrum!(spec, f, n)

genMat = nonsym(nbone, n, diag_l, diag_u, spec)

MatrixMarket.mmwrite("genMat.mtx", genMat)