using NNlib: conv, ∇conv_data, depthwiseconv, output_size

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

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

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

"""
    SamePad

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

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

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

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

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

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

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

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

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

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

# Examples

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@functor Conv

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@functor ConvTranspose

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@functor DepthwiseConv

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

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

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

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

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

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

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

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

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

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

# Examples

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

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

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

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

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

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

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

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

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

@functor CrossCor

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

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

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

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

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

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

"""
    AdaptiveMaxPool(out)

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

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

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

"""
    AdaptiveMeanPool(out)

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

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

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

"""
    GlobalMaxPool()

Global max pooling layer.

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

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

  return maxpool(x, pdims)
end

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

"""
    GlobalMeanPool()

Global mean pooling layer.

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

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

  return meanpool(x, pdims)
end

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


export WitnessSet, witness_set, linear_subspace, system, dim, codim, trace_test

"""
    WitnessSet(F, L, S)

Store solutions `S` of the polynomial system `F(x) = L(x) = 0` into a witness set.
"""
struct WitnessSet{
    S<:AbstractSystem,
    Sub<:AbstractSubspace,
    R<:Union{Vector{ComplexF64},PathResult},
}
    F::S
    L::Sub
    # only non-singular results or solutions
    R::Vector{R}
    projective::Bool
end

function WitnessSet(
    F::AbstractSystem,
    L::LinearSubspace,
    R;
    projective::Bool = is_linear(A) && is_homogeneous(System(F)),
)
    WitnessSet(F, L, R, projective)
end

"""
    system(W::WitnessSet)

Get the system stored in `W`.
"""
system(W::WitnessSet) = W.F

"""
    linear_subspace(W::WitnessSet)

Get the linear subspace stored in `W`.
"""
linear_subspace(W::WitnessSet) = W.L

"""
    solutions(W::WitnessSet)

Get the solutions stored in `W`.
"""
solutions(W::WitnessSet{A,B,PathResult}) where {A,B} = solutions(W.R)
solutions(W::WitnessSet{A,B,Vector{ComplexF64}}) where {A,B} = W.R

"""
    results(W::WitnessSet)

Get the results stored in `W`.
"""
results(W::WitnessSet{<:Any,<:Any,PathResult}) = W.R

"""
    degree(W::WitnessSet)

Returns the degree of the witness set `W`. This equals the number of solutions stored.
"""
ModelKit.degree(W::WitnessSet) = length(W.R)

"""
    dim(W::WitnessSet)

The dimension of the algebraic set encoded by the witness set.
"""
dim(W::WitnessSet) = codim(W.L)

"""
    codim(W::WitnessSet)

The dimension of the algebraic set encoded by the witness set.
"""
codim(W::WitnessSet) = dim(W.L)

function Base.show(io::IO, W::WitnessSet)
    print(io, "Witness set for dimension $(dim(W)) of degree $(degree(W))")
end

### Construct witness sets
"""
    witness_set(F; codim = nvariables(F) - length(F), dim = nothing, options...)

Compute a [`WitnessSet`](@ref) for `F` in the given dimension (resp. codimension)
by sampling a random (affine) linear subspace.
After constructing the system this calls [`solve`](@ref) with the provided `options`.

    witness_set(F, L; options...)

Compute [`WitnessSet`](@ref) for `F` and the (affine) linear subspace `L`.

    witness_set(W::WitnessSet, L; options...)

Compute a new [`WitnessSet`](@ref) with the (affine) linear subspace `L` by moving
the linear subspace stored in `W` to `L`.

### Example
```julia-repl
julia> @var x y;
julia> F = System([x^2 + y^2 - 5], [x, y])
System of length 1
 2 variables: x, y

 -5 + x^2 + y^2

julia> W = witness_set(F)
Witness set for dimension 1 of degree 2
```
"""
witness_set(F::System, args...; compile = COMPILE_DEFAULT[], kwargs...) =
    witness_set(fixed(F; compile = compile), args...; kwargs...)
function witness_set(
    F::AbstractSystem;
    target_parameters = nothing,
    dim = nothing,
    codim = nothing,
    options...,
)
    f = System(F)
    projective = is_homogeneous(f)
    if isnothing(dim) && isnothing(codim)
        dim = corank(F) - projective
    elseif !isnothing(codim) && projective
        codim += 1
    end
    L = rand_subspace(size(F, 2); dim = codim, codim = dim, affine = !projective)
    witness_set(
        F,
        L;
        target_parameters = target_parameters,
        projective = projective,
        options...,
    )
end

function witness_set(F::AbstractSystem, L::LinearSubspace; projective = nothing, options...)
    res = solve(F; target_subspace = L, options...)
    if isnothing(projective)
        WitnessSet(F, L, results(res; only_nonsingular = true))
    else
        WitnessSet(F, L, results(res; only_nonsingular = true); projective = projective)
    end
end

corank(F::AbstractSystem; kwargs...) = nvariables(F) - LA.rank(F)
function LinearAlgebra.rank(F::AbstractSystem; target_parameters = nothing)
    m, n = size(F)
    u = zeros(ComplexF64, m)
    U = zeros(ComplexF64, m, n)
    x = randn(ComplexF64, n)
    evaluate_and_jacobian!(u, U, F, x, target_parameters)
    LA.rank(U)
end

### Move witness sets around
function witness_set(W::WitnessSet, L::LinearSubspace; options...)
    if W.projective && !is_linear(L)
        error(
            "The given space is an affine linear subspace (``b ≠ 0``). " *
            " Expected a linear subspace since the given witness set is projective.",
        )
    end
    res = solve(W.F, W.R; start_subspace = W.L, target_subspace = L, options...)
    WitnessSet(W.F, L, results(res; only_nonsingular = true), is_linear(L) && W.projective)
end

"""
    trace_test(W::WitnessSet; options...)

Performs a trace test [^LRS18] to verify whether the given witness set `W` is complete.
Returns the trace of the witness set which should be theoretically be 0 if `W` is complete.
Due to floating point arithmetic this is not the case, thus is has to be manually checked
that the trace is sufficiently small.
Returns `nothing` if the trace test failed due to path tracking failures.
The `options` are the same as for calls to [`witness_set`](@ref).

```julia-repl
julia> @var x y;
julia> F = System([x^2 + y^2 - 5], [x, y])
System of length 1
 2 variables: x, y

 -5 + x^2 + y^2

julia> W = witness_set(F)
Witness set for dimension 1 of degree 2

julia> trace = trace_test(W)
9.981960497718987e-16
```

[^LRS18]: Leykin, Anton, Jose Israel Rodriguez, and Frank Sottile. "Trace test." Arnold Mathematical Journal 4.1 (2018): 113-125.
APA

"""
function trace_test(W₀::WitnessSet; options...)
    L₀ = linear_subspace(W₀)
    F = system(W₀)
    S₀ = solutions(W₀)
    # if we are in the projective setting, we need to make sure that
    # all solutions are on the same affine chart
    # Therefore make the affine chart now
    if W₀.projective
        F = on_affine_chart(F)
        s₀ = sum(s -> set_solution!(s, F, s), S₀)
    else
        s₀ = sum(S₀)
    end

    v = randn(ComplexF64, codim(L₀))
    L₁ = translate(L₀, v)
    L₋₁ = translate(L₀, -v)

    R₁ = solve(F, S₀; start_subspace = L₀, target_subspace = L₁, options...)
    nsolutions(R₁) == degree(W₀) || return nothing

    R₋₁ = solve(F, S₀; start_subspace = L₀, target_subspace = L₋₁, options...)
    nsolutions(R₋₁) == degree(W₀) || return nothing

    s₁ = sum(solutions(R₁))
    s₋₁ = sum(solutions(R₋₁))

    M = [s₋₁ s₀ s₁; 1 1 1]
    singvals = LA.svdvals(M)
    trace = singvals[3] / singvals[1]

    trace
end


function run(s)
    # Your code here
    entries = split(s, "\n")
    entries = map(x -> parse(Int, x), entries)
    for (idx1, el1) in enumerate(entries)
        for (idx2, el2) in enumerate(entries[idx1+1:end])
            if el1 + el2 > 2020
                continue
            end
            for el3 in entries[idx1+idx2+1:end]
                if el1 + el2 + el3 == 2020
                    return el1 * el2 * el3
                end
            end
        end
    end
    return 0
end

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

function main()
    res, time, memory = @timed run(ARGS[1])
    println("_duration:$(time*1000)")
    println(res)
end

main()


include("models.jl")

function test_bernoulli()
    model_bernoulli = bernoulli_example()
    graph_bernoulli = trackdependencies(model_bernoulli)
    @testdependencies(model_bernoulli, w, p, x)
    cond_bernoulli_p = StaticConditional(model_bernoulli, :p)
    @test_nothrow sample(model_bernoulli, Gibbs(cond_bernoulli_p, MH(:w)), 500)

    # Analytic tests
    w = graph_bernoulli[4].value
    p = graph_bernoulli[6].value
    x = graph_bernoulli[2].value
    p_1 = w[1] * pdf(Bernoulli(0.3), x)
    p_2 = w[2] * pdf(Bernoulli(0.7), x)
    Z = p_1 + p_2

    # 𝓅(p | w, x) ∝ 𝓅(p | w) * 𝓅(x | p)
    analytic_conditional = DiscreteNonParametric([0.3, 0.7], [p_1 / Z, p_2 / Z])
    @info "Bernoulli analytic" analytic_conditional
    θ = AutoGibbs.sampled_values(graph_bernoulli)

    local calculated_conditional
    @test_nothrow calculated_conditional = conditionals(graph_bernoulli, @varname(p))[@varname(p)]
    @info "Bernoulli calculated" calculated_conditional(θ)

    @test issimilar(calculated_conditional(θ), analytic_conditional)
end

function test_gmm()
    model_gmm = gmm_example()
    graph_gmm = trackdependencies(model_gmm)
    @testdependencies(model_gmm, μ, w, z, x[1], x[2], x[3])
    cond_gmm_z = StaticConditional(model_gmm, :z)
    @test_nothrow sample(model_gmm, Gibbs(cond_gmm_z, MH(:w, :μ)), 500)
    @test_nothrow sample(model_gmm, Gibbs(cond_gmm_z, HMC(0.01, 10, :w, :μ)), 500)

    # Analytic tests
    x = graph_gmm[2].value
    μ = graph_gmm[7].value
    w = graph_gmm[10].value
    z = graph_gmm[14].value
    p_1 = w[1] .* pdf.(Normal(μ[1], s2_gmm), x)
    p_2 = w[2] .* pdf.(Normal(μ[2], s2_gmm), x)
    (Z1, Z2, Z3) = p_1 .+ p_2

    # 𝓅(zᵢ | μ, w, x, z₋ᵢ) ∝ 𝓅(zᵢ | w) * 𝓅(xᵢ | zᵢ, μ)
    analytic_conditionals = [@varname(z[1]) => Categorical([p_1[1], p_2[1]] ./ Z1),
                             @varname(z[2]) => Categorical([p_1[2], p_2[2]] ./ Z2),
                             @varname(z[3]) => Categorical([p_1[3], p_2[3]] ./ Z3)]
    @info "GMM analytic" analytic_conditionals
    θ = AutoGibbs.sampled_values(graph_gmm)

    local calculated_conditional
    @test_nothrow calculated_conditional = conditionals(graph_gmm, @varname(z))[@varname(z)]
    @info "GMM calculated" calculated_conditional(θ)

    @test issimilar(calculated_conditional(θ), Product([D for (vn, D) in analytic_conditionals]))
end

function test_gmm_loopy()
    model_gmm_loopy = gmm_loopy_example()
    graph_gmm_loopy = trackdependencies(model_gmm_loopy)
    @testdependencies(model_gmm_loopy, μ[1], μ[2], w, z[1], z[2], z[3], x[1], x[2], x[3])
    cond_gmm_loopy_z = StaticConditional(model_gmm_loopy, :z)
    @test_nothrow sample(model_gmm_loopy, Gibbs(cond_gmm_loopy_z, MH(:w, :μ)), 500)
    @test_nothrow sample(model_gmm_loopy, Gibbs(cond_gmm_loopy_z, HMC(0.01, 10, :w, :μ)), 500)

    # Analytic tests
    x = graph_gmm_loopy[2].value
    μ = [graph_gmm_loopy[19].value, graph_gmm_loopy[28].value]
    w = graph_gmm_loopy[32].value
    z = [graph_gmm_loopy[44].value, graph_gmm_loopy[63].value, graph_gmm_loopy[82].value]
    p_1 = w[1] .* pdf.(Normal(μ[1], s2_gmm), x)
    p_2 = w[2] .* pdf.(Normal(μ[2], s2_gmm), x)
    (Z1, Z2, Z3) = p_1 .+ p_2

    # 𝓅(zᵢ | μ, w, x, z₋ᵢ) ∝ 𝓅(zᵢ | w) * 𝓅(xᵢ | zᵢ, μ)
    analytic_conditionals = [@varname(z[1]) => Categorical([p_1[1], p_2[1]] ./ Z1),
                             @varname(z[2]) => Categorical([p_1[2], p_2[2]] ./ Z2),
                             @varname(z[3]) => Categorical([p_1[3], p_2[3]] ./ Z3)]
    @info "Loopy GMM analytic" analytic_conditionals
    θ = AutoGibbs.sampled_values(graph_gmm_loopy)

    local calculated_conditionals
    @test_nothrow calculated_conditionals = conditionals(graph_gmm_loopy, @varname(z))
    @info "Loopy GMM calculated" Dict(vn => cond(θ) for (vn, cond) in calculated_conditionals)

    for (vn, analytic_conditional) in analytic_conditionals
        @test issimilar(calculated_conditionals[vn](θ), analytic_conditional)
    end
end

function test_gmm_shifted()
    model_gmm_shifted = gmm_shifted_example()
    graph_gmm_shifted = trackdependencies(model_gmm_shifted)
    @testdependencies(model_gmm_shifted, μ[1], μ[2], w, z[1], z[2], z[3], x[1], x[2], x[3])
    cond_gmm_shifted_z = StaticConditional(model_gmm_shifted, :z)
    @test_nothrow sample(model_gmm_shifted, Gibbs(cond_gmm_shifted_z, MH(:w, :μ)), 500)
    @test_nothrow sample(model_gmm_shifted, Gibbs(cond_gmm_shifted_z, HMC(0.01, 10, :w, :μ)), 500)
end

function test_hmm()
    model_hmm = hmm_example()
    graph_hmm = trackdependencies(model_hmm)
    @testdependencies(model_hmm, T[1], T[2], m[1], m[2], s[1], s[2], s[3], x[1], x[2], x[3])
    cond_hmm_s = StaticConditional(model_hmm, :s)
    @test_nothrow sample(model_hmm, Gibbs(cond_hmm_s, MH(:m, :T)), 500)
    @test_nothrow sample(model_hmm, Gibbs(cond_hmm_s, HMC(0.01, 10, :m, :T)), 500)

    # Analytic tests
    x = graph_hmm[2].value
    T = [graph_hmm[24].value, graph_hmm[40].value]
    m = [graph_hmm[30].value, graph_hmm[46].value]
    s1 = graph_hmm[50].value
    s2 = graph_hmm[70].value
    s3 = graph_hmm[91].value
    D_obs_1 = Normal(m[1], s2_hmm)
    D_obs_2 = Normal(m[2], s2_hmm)
    p_s1_1 = pdf(Categorical(2), 1) * pdf(Categorical(T[1]), s2) * pdf(D_obs_1, x[1])
    p_s1_2 = pdf(Categorical(2), 2) * pdf(Categorical(T[2]), s2) * pdf(D_obs_2, x[1])
    p_s2_1 = pdf(Categorical(T[s1]), 1) * pdf(Categorical(T[1]), s3) * pdf(D_obs_1, x[2])
    p_s2_2 = pdf(Categorical(T[s1]), 2) * pdf(Categorical(T[2]), s3) * pdf(D_obs_2, x[2])
    p_s3_1 = pdf(Categorical(T[s2]), 1) * pdf(D_obs_1, x[3])
    p_s3_2 = pdf(Categorical(T[s2]), 2) * pdf(D_obs_2, x[3])
    Z_1 = p_s1_1 + p_s1_2
    Z_2 = p_s2_1 + p_s2_2
    Z_3 = p_s3_1 + p_s3_2

    # 𝓅(s₁ | T, m, s₋₁, x) ∝ 𝓅(s₁) 𝓅(s₂ | s₁, T) 𝓅(x₁ | s₁, m)
    #  𝓅(sᵢ | T, m, s₋ᵢ, x) ∝  𝓅(sᵢ | sᵢ₋₁, T) 𝓅(sᵢ₊₁ | sᵢ, T) 𝓅(xᵢ | sᵢ, m) (for i ≥ 2)
    analytic_conditionals = [@varname(s[1]) => Categorical([p_s1_1, p_s1_2] ./ Z_1),
                             @varname(s[2]) => Categorical([p_s2_1, p_s2_2] ./ Z_2),
                             @varname(s[3]) => Categorical([p_s3_1, p_s3_2] ./ Z_3)]
    θ = AutoGibbs.sampled_values(graph_hmm)
    @info "HMM analytic" analytic_conditionals

    local calculated_conditionals
    @test_nothrow calculated_conditionals = conditionals(graph_hmm, @varname(s))
    @info "HMM calculated" Dict(vn => cond(θ) for (vn, cond) in calculated_conditionals)

    for (vn, analytic_conditional) in analytic_conditionals
        # @show vn => probs(calculated_conditionals[vn]), probs(analytic_conditional)
        @test issimilar(calculated_conditionals[vn](θ), analytic_conditional)
    end
end

function test_imm_stick()
    model_imm_stick = imm_stick_example()
    graph_imm_stick = trackdependencies(model_imm_stick)

    # we leave out the μs, because there might be 1--3 of them
    @testdependencies(model_imm_stick,
                      v, # 10 - 1 sticks
                      μ, # 10 cluster centers
                      # z, # 9 data points
                      y[1], y[2], y[3], y[4], y[5], y[6], y[7], y[8], y[9])
    cond_imm_stick = StaticConditional(model_imm_stick, :z)
    @test_nothrow sample(model_imm_stick, Gibbs(cond_imm_stick, MH(:μ, :v)), 500)
    @test_nothrow sample(model_imm_stick, Gibbs(cond_imm_stick, HMC(0.01, 10, :μ, :v)), 500)

    # Analytic tests
    y = graph_imm_stick[2].value
    α = graph_imm_stick[4].value
    K = graph_imm_stick[6].value
    N = graph_imm_stick[7].value
    v = graph_imm_stick[12].value
    # z = graph_imm_stick[16].value
    z = [v.value for v in values(graph_imm_stick.statements)
         if v isa AutoGibbs.Assumption && DynamicPPL.subsumes(@varname(z), v.vn)]
    # μ = graph_imm_stick[19].value
    μ = (v.value for v in values(graph_imm_stick.statements)
         if v isa AutoGibbs.Assumption && DynamicPPL.subsumes(@varname(μ), v.vn)) |> first

    D_w = Categorical(stickbreak(v))
    analytic_conditionals = map(1:N) do n
        p̃ = [exp(logpdf(D_w, z) + logpdf(Normal(μ[z], 1.0), y[n])) for z in support(D_w)]
        @varname(z[n]) => Categorical(p̃ ./ sum(p̃))
    end
    # analytic_conditional = Product([D for (vn, D) in analytic_conditionals])
    # @info "stick-breaking IMM analytic conditionals" analytic_conditionals

    θ = AutoGibbs.sampled_values(graph_imm_stick)

    local calculated_conditionals
    @test_nothrow calculated_conditionals = conditionals(graph_imm_stick, @varname(z))
    # @info "stick-breaking IMM calculated conditionals" Dict(
        # vn => cond(θ) for (vn, cond) in calculated_conditionals)

    for (vn, analytic_conditional) in analytic_conditionals
        # @show vn => probs(calculated_conditionals[vn]), probs(analytic_conditional)
        @test issimilar(calculated_conditionals[vn](θ), analytic_conditional)
    end
end

function test_changepoint()
    model_changepoint = changepoint([1.1, 0.9, 0.2])
    graph_changepoint = trackdependencies(model_changepoint)
    @testdependencies(model_changepoint, λ₁, λ₂, τ, y[1], y[2], y[3])
    @test_nothrow sample(model_changepoint, Gibbs(AutoConditional(:τ), MH(:λ₁, :λ₂)), 20)
end

# function test_reverseps()
#     model_reverse_deps = reverse_deps([0.1, -0.2])
#     graph_reverse_deps = trackdependencies(model_reverse_deps)
#     @testdependencies(model_reverse_deps, m[1], m[2], x)
# end

##########################################################################
#########################################################################
### TEST TOGGLES

test_bernoulli()
test_gmm()
test_gmm_loopy()
test_gmm_shifted()
test_hmm()
test_imm_stick()
test_changepoint()


@testset "simple.jl" begin

    # Load all required packages and define likelihood functions
    ℓ(x) = logpdf(Normal(x, 0.5), 1.0)
    ℓvec(x) = logpdf(MvNormal(x, 0.25 * I), [1.0])
    @everywhere begin
        ℓ(x) = logpdf(Normal(x, 0.5), 1.0)
        ℓvec(x) = logpdf(MvNormal(x, 0.25 * I), [1.0])
    end

    @testset "Scalar model" begin
        Random.seed!(1)

        # model
        prior = Normal(0, 1)

        # true posterior
        μ = 0.8
        σ² = 0.2

        # regular sampling
        samples = sample(ESSModel(prior, ℓ), ESS(), 2_000; progress=false)
        @test samples isa Vector{Float64}
        @test length(samples) == 2_000
        @test mean(samples) ≈ μ atol = 0.05
        @test var(samples) ≈ σ² atol = 0.05

        # parallel sampling
        for alg in (MCMCThreads(), MCMCDistributed(), MCMCSerial())
            samples = sample(ESSModel(prior, ℓ), ESS(), alg, 2_000, 5; progress=false)
            @test samples isa Vector{Vector{Float64}}
            @test length(samples) == 5
            @test all(length(x) == 2_000 for x in samples)
            @test mean(mean, samples) ≈ μ atol = 0.05
            @test mean(var, samples) ≈ σ² atol = 0.05

            # initial parameter
            init_x = randn(5)
            samples = sample(
                ESSModel(prior, ℓ), ESS(), alg, 10, 5; progress=false, init_params=init_x
            )
            @test map(first, samples) == init_x
        end

        # initial parameter
        init_x = randn()
        samples = sample(ESSModel(prior, ℓ), ESS(), 10; progress=false, init_params=init_x)
        @test first(samples) == init_x
    end

    @testset "Scalar model with nonzero mean" begin
        Random.seed!(1)

        # model
        prior = Normal(0.5, 1)

        # true posterior
        μ = 0.9
        σ² = 0.2

        # regular sampling
        samples = sample(ESSModel(prior, ℓ), ESS(), 2_000; progress=false)
        @test samples isa Vector{Float64}
        @test length(samples) == 2_000
        @test mean(samples) ≈ μ atol = 0.05
        @test var(samples) ≈ σ² atol = 0.05

        # parallel sampling
        for alg in (MCMCThreads(), MCMCDistributed(), MCMCSerial())
            samples = sample(ESSModel(prior, ℓ), ESS(), alg, 2_000, 5; progress=false)
            @test samples isa Vector{Vector{Float64}}
            @test length(samples) == 5
            @test all(length(x) == 2_000 for x in samples)
            @test mean(mean, samples) ≈ μ atol = 0.05
            @test mean(var, samples) ≈ σ² atol = 0.05

            # initial parameter
            init_x = randn(5)
            samples = sample(
                ESSModel(prior, ℓ), ESS(), alg, 10, 5; progress=false, init_params=init_x
            )
            @test map(first, samples) == init_x
        end

        # initial parameter
        init_x = randn()
        samples = sample(ESSModel(prior, ℓ), ESS(), 10; progress=false, init_params=init_x)
        @test first(samples) == init_x
    end

    @testset "Scalar model (vectorized)" begin
        Random.seed!(1)

        # model
        prior = MvNormal([0.0], I)

        # true posterior
        μ = [0.8]
        σ² = [0.2]

        # regular sampling
        samples = sample(ESSModel(prior, ℓvec), ESS(), 2_000; progress=false)
        @test samples isa Vector{Vector{Float64}}
        @test length(samples) == 2_000
        @test all(length(x) == 1 for x in samples)
        @test mean(samples) ≈ μ atol = 0.05
        @test var(samples) ≈ σ² atol = 0.05

        # parallel sampling
        for alg in (MCMCThreads(), MCMCDistributed(), MCMCSerial())
            samples = sample(ESSModel(prior, ℓvec), ESS(), alg, 2_000, 5; progress=false)
            @test samples isa Vector{Vector{Vector{Float64}}}
            @test length(samples) == 5
            @test all(length(x) == 2_000 for x in samples)
            @test mean(mean, samples) ≈ μ atol = 0.05
            @test mean(var, samples) ≈ σ² atol = 0.05

            # initial parameter
            init_x = [randn(1) for _ in 1:5]
            samples = sample(
                ESSModel(prior, ℓvec), ESS(), alg, 10, 5; progress=false, init_params=init_x
            )
            @test map(first, samples) == init_x
        end

        # initial parameter
        init_x = randn(1)
        samples = sample(
            ESSModel(prior, ℓvec), ESS(), 10; progress=false, init_params=init_x
        )
        @test first(samples) == init_x
    end

    @testset "Scalar model with nonzero mean (vectorized)" begin
        Random.seed!(1)

        # model
        prior = MvNormal([0.5], I)

        # true posterior
        μ = [0.9]
        σ² = [0.2]

        # regular sampling
        samples = sample(ESSModel(prior, ℓvec), ESS(), 2_000; progress=false)
        @test samples isa Vector{Vector{Float64}}
        @test length(samples) == 2_000
        @test all(length(x) == 1 for x in samples)
        @test mean(samples) ≈ μ atol = 0.05
        @test var(samples) ≈ σ² atol = 0.05

        # parallel sampling
        for alg in (MCMCThreads(), MCMCDistributed(), MCMCSerial())
            samples = sample(ESSModel(prior, ℓvec), ESS(), alg, 2_000, 5; progress=false)
            @test samples isa Vector{Vector{Vector{Float64}}}
            @test length(samples) == 5
            @test all(length(x) == 2_000 for x in samples)
            @test mean(mean, samples) ≈ μ atol = 0.05
            @test mean(var, samples) ≈ σ² atol = 0.05

            # initial parameter
            init_x = [randn(1) for _ in 1:5]
            samples = sample(
                ESSModel(prior, ℓvec), ESS(), alg, 10, 5; progress=false, init_params=init_x
            )
            @test map(first, samples) == init_x
        end

        # initial parameter
        init_x = randn(1)
        samples = sample(
            ESSModel(prior, ℓvec), ESS(), 10; progress=false, init_params=init_x
        )
        @test first(samples) == init_x
    end
end


# MDP parameters, ValueIteration minimizes the cost => cost is positive, reward is negative
no_states = 10
horizon = 5
actions = [:l, :r]
actionCost = 1.
actionsImpact = Base.ImmutableDict(:l => -1, :r => 1)
reward_states = [1, no_states]
reward = -10.
discount_factor = 1.
noise = .6

# MDPs initialization
mdp = FHExample(no_states, horizon, actions, actionCost, actionsImpact, reward_states, reward, discount_factor, noise)

# initialize the solver
# max_iterations: maximum number of iterations value iteration runs for (default is 100)
# belres: the value of Bellman residual used in the solver
solver = ValueIterationSolver(max_iterations=10, belres=1e-3, include_Q=true);

# Solve Value Iteration
VIPolicy = DiscreteValueIteration.solve(solver, mdp);

# Solve Finite Horizon by Value Iteration
FHPolicy = FiniteHorizonPOMDPs.solve(mdp);

# Compare resulting policies
@test all((FiniteHorizonPOMDPs.action(FHPolicy, s) == action(VIPolicy, s) for s in states(mdp)))

# # Evaluates to false because of different elements Types (ExampleState vs FHExampleState)
fh_states = collect(Iterators.flatten([FiniteHorizonPOMDPs.stage_states(mdp, i) for i=1:mdp.horizon]))
ih_states = states(mdp)

z = zip(fh_states, ih_states)

@test all((fh == ih for (fh, ih) in z))


module Forecast

using CSV, Distributions, DataFrames, GZip, LinearAlgebra, Plots,
      RecipesBase, TimeSeries, Statistics

# types
export CCF, STL
# methods
export acf, ccf, d, loess, sma, stl, hma, hmaSymmetricWeights
# datasets
export co2

# source files

## ccf
include("ccf.jl")
include("acf.jl")

include("d.jl")

## stl
include("loess.jl")
include("sma.jl")
include("stl.jl")

## hma
include("hma.jl")

include("utils.jl")

include("datasets.jl")

## recipes
include("plotrecipes.jl")

"""
Collection of methods for Time Series analysis

Methods implemented:

    acf:        Auto-correlation or auto-covariance of univariate series.
    ccf:        Cross-correlation or cros-covariance of two univariate series.
    d:          Lagged differences of a given order for Vector, Array and TimeSeries.
    hma:        Henderson moving average filters.
    loess:      Locally estimated scatterplot smoothing.
    sma:        Simple moving average.
    stl:        Seasonal and Trend decomposition using loess.
"""
Forecast

end


#=

  Euler #41 in Julia.

  Problem 41:
  """
  We shall say that an n-digit number is pandigital if it makes use of all
  the digits 1 to n exactly once. For example, 2143 is a 4-digit pandigital
  and is also prime.

  What is the largest n-digit pandigital prime that exists?
  """

  This Julia model was created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/Julia_progs/

=#

include("Euler.jl")

# 0.03632750s
function euler41a()
    # Simplification:
    # n=9 is not possible since 1+2+3+4+5+6+7+8+9=45 is divisible by 3
    #  n=8 is not possible since 1+2+3+4+5+6+7+8=36 is divisible by 3
    n = 7
    m = 0
    while m == 0 && n >= 4
        p = Array(1:n)
        for pp in sort(all_permutations(p),rev=true)
            v = parse(Int,join(pp,""))
            if isPrime(v)
                return v
            end
        end
        n += 1
    end

    return nothing
end

run_euler(euler41a)


mutable struct NormAttributes
    affine::Bool
    track_stats::Bool
    active::Union{Bool,Nothing}
end

Flux.hasaffine(attr::NormAttributes) = attr.affine

function get_stats(::Val{true}, ::Val{false}, l, x::AbstractArray{T,N}, reduce_dims) where {T,N}
    # testmode with tracked stats
    stats_shape = ntuple(i -> i == N - 1 ? size(x, N - 1) : 1, N)
    return reshape(l.μ, stats_shape), reshape(l.σ², stats_shape)
end

function get_stats(::Val{false}, active, l, x, reduce_dims)
    # trainmode or testmode without tracked stats
    μ = mean(x; dims=reduce_dims)
    diff = x .- μ
    return μ, mean(abs2, diff; dims=reduce_dims)
end

function get_stats(::Val{true}, active::Val{true}, l, x::AbstractArray{T,N}, reduce_dims) where {T,N}
    # trainmode with tracked stats
    μ, σ² = get_stats(Val(false), active, l, x, reduce_dims)
    Zygote.ignore() do
        mtm = l.momentum
        m = prod(size(x)[reduce_dims])  # needed for computing corrected var
        μnew = vec(N ∈ reduce_dims ? μ : mean(μ; dims=N))
        σ²new = vec(N ∈ reduce_dims ? σ² : mean(σ²; dims=N))
        l.μ .= (1 - mtm) .* l.μ .+ mtm .* μnew
        return l.σ² .= (1 - mtm) .* l.σ² .+ mtm .* (m / (m - one(eltype(l.σ²)))) .* σ²new
    end
    return μ, σ²
end

function norm_forward(l, x::AbstractArray{T,N}, reduce_dims, affine_shape) where {T,N}
    μ, σ² = get_stats(Val(l.attrs.track_stats), Val(_isactive(l)), l, x, reduce_dims)
    if hasaffine(l)
        γ = reshape(l.γ, affine_shape)
        β = reshape(l.β, affine_shape)
        return l.λ.(_norm_forward(μ, σ², x, γ, β, l.ϵ))
    else
        return l.λ.(_norm_forward(μ, σ², x, l.ϵ))
    end
end

_norm_forward(μ, σ², x, ϵ) = (x .- μ) ./ sqrt.(σ² .+ ϵ)

_norm_forward(μ, σ², x, γ, β, ϵ) = γ .* (x .- μ) ./ sqrt.(σ² .+ ϵ) .+ β

Zygote.@adjoint function _norm_forward(μ, σ², x, γ, β, ϵ)
    N = ndims(x)

    σ²ϵ = σ² .+ ϵ
    inv_deno = 1 ./ sqrt.(σ²ϵ)
    res_1 = (x .- μ) .* inv_deno
    res_2 = γ .* res_1
    res = res_2 .+ β

    function norm_backward(Δ)
        reduce_dims_affine = filter(i -> isone(size(β, i)), 1:N)
        reduce_dims_stats = filter(i -> isone(size(σ², i)), 1:N)

        Δx = inv_deno .* Δ
        Δμ = -sum(Δx; dims=reduce_dims_stats)
        Δσ² = sum(-eltype(x)(0.5) .* res_2 .* Δ ./ σ²ϵ; dims=reduce_dims_stats)
        Δγ = sum(res_1 .* Δ; dims=reduce_dims_affine)
        Δβ = sum(Δ; dims=reduce_dims_affine)

        return (Δμ, Δσ², Δx, Δγ, Δβ, nothing)
    end

    return res, norm_backward
end

Zygote.@adjoint function _norm_forward(μ, σ², x, ϵ)
    N = ndims(x)

    σ²ϵ = σ² .+ ϵ
    inv_deno = 1 ./ sqrt.(σ²ϵ)
    res = (x .- μ) .* inv_deno

    function norm_backward(Δ)
        reduce_dims_stats = filter(i -> isone(size(σ², i)), 1:N)

        Δx = inv_deno .* Δ
        Δμ = -sum(Δx; dims=reduce_dims_stats)
        Δσ² = sum(-eltype(x)(0.5) .* res .* Δ ./ σ²ϵ; dims=reduce_dims_stats)

        return (Δμ, Δσ², Δx, nothing)
    end

    return res, norm_backward
end


function circle!(c,r)
    θ = LinRange(0,2π,180)
    plot!(r*cos.(θ) .+ c[1], r*sin.(θ) .+ c[2], label = :none)
end

@userplot Circle
@recipe function f(cc::Circle)
    c = cc.args[1]
    r = cc.args[2]

    θ = LinRange(0,2π,180)
    x = r*cos.(θ) .+ c[1]
    y = r*sin.(θ) .+ c[2]
    @series begin
        x, y
    end
end


@with_kw struct Param

    #Temporal variables for simulation
    t_sim_init          :: Float64   = 0.0
    t_sim_final         :: Float64   = 0.0
    out_every_t         :: Float64   = 1.0
    deltat              :: Float64   = 0.05

    #Max possible sim
    max_runtime         ::String    = "70:00:00"

   #Spacial variables
   rmin            :: Real          =  1.0
   rmax            :: Float64       =  30.0
   rnodes          :: Integer       =  128

   thnodes         :: Integer       =  128

    #Initial Conditions
    A0              ::Float64       = 1.0
    σ               ::Float64       = 2.0
    r0              ::Float64       = 5.0
    ω               ::Float64       = 5.0
    m               ::Integer       = 0.0

    #Cylinder Variables
    ϵ               ::Real      = 0.1
    Ω               ::Real      = 0.3
    Z0              ::Real      = 0.001

    #Data variables
    folder          :: String    = "data/"
    fname           :: String    = "default_file_name"

end


@doc raw"""
    pc = PController(Kc)

Construct a Proportional (P) controller by specifying the controller gain defined under the following transfer function representation:

$$g_c(s)=K_c$$

# Arguments
* `Kc::Float64`: controller gain

# Example
```julia
pc = PController(1.0) # specify P controller gain
gc = TransferFunction(pc) # construct transfer function with this P-controller gain
```
"""
mutable struct PController
    Kc::Float64
end

TransferFunction(pc::PController) = TransferFunction([pc.Kc], [1.0])

@doc raw"""
    pic = PIController(Kc, τI)

Construct a Proportional-Integral (PI) controller by specifying the controller gain and integral time constant defined under the following transfer function representation:

$$g_c(s)=K_c \left[1+\frac{1}{\tau_I s}\right]$$

# Arguments
* `Kc::Float64`: controller gain
* `τI::Float64`: integral time constant

# Example
```julia
pic = PIController(1.0, 3.0) # specify PI controller params
gc = TransferFunction(pic) # construct transfer function with these PI-controller params
```
"""
mutable struct PIController
    Kc::Float64
    τI::Float64
end

TransferFunction(pic::PIController) = pic.Kc * (1 + 1 / (pic.τI * s))

mutable struct PIDController
    Kc::Float64
    τI::Float64
    τD::Float64
    α::Float64
end

@doc raw"""
    pidc = PIDController(Kc, τI, τD, α=0.0)

Construct a Proportional-Integral-Derivative (PID) controller by specifying the controller gain, integral time constant, derivative time constant, and derivative filter defined under the following transfer function representation:

$$g_c(s)=K_c \left[1+\frac{1}{\tau_I s}+\tau_D s \frac{1}{\alpha \tau_D s + 1}\right]$$

# Arguments
* `Kc::Float64`: controller gain
* `τI::Float64`: integral time constant
* `τD::Float64`: derivative time constant
* `α::Float64`: derivative filter

# Example
```julia
pidc = PIDController(1.0, 3.0, 0.1) # specify PID controller params
gc = TransferFunction(pidc) # construct transfer function with these PID-controller params
```
"""
PIDController(Kc::Float64, τI::Float64, τD::Float64; α::Float64=0.0) = PIDController(Kc, τI, τD, α)

TransferFunction(pidc::PIDController) = pidc.Kc * (1 + 1.0 / (pidc.τI * s) + pidc.τD * s / (pidc.α * pidc.τD * s + 1))


struct CombinedConfusionMatrix
    matrices
    classlabels

    function CombinedConfusionMatrix(matrices;
                                                    classlabels=nothing)
        if classlabels == nothing
            _, n = size(matrices[1])
            classlabels = ["class $x" for x in 1:n]
        end

        new(matrices, classlabels)
    end

end

function checkedscore(a,b)
    b==0? 0:a/b
end

function score(c)
    m,n = size(c)
    sum([c[i,i] for i in 1:m])
end

function expandmatrix(matrix::Matrix)

    m, n = size(matrix)

    result = Matrix{Number}(m+2, n+2)

    result .= 0

    for row in 1 : m
        for col in 1: n
            result[row,col] = matrix[row,col]
        end

        result[row, n+1] = sum(matrix[row,:])
        rowscore = checkedscore(matrix[row,row], sum(matrix[row,:]))
        result[row, n+2] = rowscore
    end

    for col in 1: n
        result[m+1, col] = sum(matrix[:,col])
        colscore = checkedscore( matrix[col,col], sum(matrix[:,col]))
        result[m+2, col] = colscore
    end

    result[m+1, n+1] = sum(matrix)

    totalscore =   score(matrix) / sum(matrix)
    result[m+2, n+2] = totalscore

    result

end

function tohtml(confusionmatrix::CombinedConfusionMatrix)

    cms = confusionmatrix.matrices
    labels = confusionmatrix.classlabels

    cm = cms[1]
    m,n = size(cms[1])

    cms2 = expandmatrix.(cms)
    d = Float64.(cat(3, cms2...))
    sems = mapslices(StatsBase.sem, d, 3)
    means = mapslices(StatsBase.mean, d, 3)

    stream = IOBuffer()
    write(stream,"""
""")
    write(stream,"
")
    write(stream,"""	""")
    write(stream,"""Predicted class x&#770	""")
    write(stream,"
")
    write(stream,"
")
    for label in labels
        write(stream,"$(label)	")
    end
    write(stream,"Row total	")
    write(stream,"Accuracy  P(x&#770|x)	")
    write(stream,"
")
    write(stream,"
")
    write(stream,"""""")
    write(stream,"""
Actual

class x
""")
    write(stream,"	")
    write(stream,"$(labels[1])	")
    for j in 1:n
        write(stream,"""$(@sprintf("%.2f",means[1,j])) ± $(@sprintf("%.2f",sems[1,j]))	""")
    end
    write(stream,"""$(@sprintf("%.2f",means[1,n+1])) ± $(@sprintf("%.2f",sems[1,n+1]))	""") # Row total
    write(stream,"""$(@sprintf("%.2f",means[1,n+2])) ± $(@sprintf("%.2f",sems[1,n+2]))	""") # Accuracy
    write(stream,"
")
    for i in 2:length(labels)
        write(stream,"
")
        write(stream,"$(labels[i])	")
        for j in 1:n
            write(stream,"""$(@sprintf("%.2f",means[i,j])) ± $(@sprintf("%.2f",sems[i,j]))	""")
        end
        write(stream,"""$(@sprintf("%.2f",means[i,n+1])) ± $(@sprintf("%.2f",sems[i,n+1]))	""") # Row total
        write(stream,"""$(@sprintf("%.2f",means[i,n+2])) ± $(@sprintf("%.2f",sems[i,n+2]))	""") # Accuracy
        write(stream,"
")
    end
    write(stream,"
")
    write(stream,"""	""")
    write(stream,"Column total	")
    for j in 1:length(labels)
        write(stream,"""$(@sprintf("%.2f",means[m+1,j])) ± $(@sprintf("%.2f",sems[m+1,j]))	""")
    end
    write(stream,"""$(@sprintf("%.2f",means[m+1,n+1])) ± $(@sprintf("%.2f",sems[m+1,n+1]))	""")
    write(stream,"
")
    write(stream,"
")
    write(stream,"""	""")
    write(stream,"Reliability P(x|x&#770)	")
    for j in 1:length(labels)
        write(stream,"""$(@sprintf("%.2f",means[m+2,j])) ± $(@sprintf("%.2f",sems[m+2,j]))	""") # Reliability
    end
    write(stream,"""	""")
    write(stream,"""$(@sprintf("%.2f",means[m+2,n+2])) ± $(@sprintf("%.2f",sems[m+2,n+2]))	""")
    write(stream,"
")
    write(stream,"
")
    String(stream)
end

Base.show(io::IO, ::MIME"text/html", c::CombinedConfusionMatrix) =
    print(io, tohtml(c))


using CPUTime
using Plots, Colors

mutable struct Node
    x::Int64
    y::Int64
    cost::Float64
    parent::Tuple{Any, Any}
end

"""移動可能なノード

・なぜか関数dijkstra内で定義したらエラー
・[]が良くない？
"""
function options(x, y, x_max, y_max, gridmap)
    options = []
    if x+1 <= x_max
        if y+1 <= y_max
            if gridmap[y+1, x+1] == 0
                push!(options, (1, 1, sqrt(2)))
            end
        end
        if y >= 2
            if gridmap[y-1, x+1] == 0
                push!(options, (1, -1, sqrt(2)))
            end
        end
        if gridmap[y, x+1] == 0
            push!(options, (1, 0, 1.0))
        end
    end
    if x >= 2
        if y+1 <= y_max
            if gridmap[y+1, x-1] == 0
                push!(options, (-1, 1, sqrt(2)))
            end
        end
        if y >= 2
            if gridmap[y-1, x-1] == 0
                push!(options, (-1, -1, sqrt(2)))
            end
        end
        if gridmap[y, x-1] == 0
            push!(options, (-1, 0, 1.0))
        end
    else
        if y+1 <= y_max
            if gridmap[y+1, x] == 0
                push!(options, (0, 1, 1.0))
            end
        end
        if y >= 2
            if gridmap[y-1, x] == 0
            push!(options, (0, -1, 1.0))
            end
        end
    end

    return options
end

"""ダイクストラで最短経路探索"""
function dijkstra(gridmap::Matrix{Int64}, start, goal)

    x_max = size(gridmap)[1]
    y_max = size(gridmap)[2]

    """一番costが小さいnodeのidを探す"""
    function find_MinCost_id(d)
        for (i, id) in enumerate(keys(d))
            if i==1
                global minid = id
            elseif d[id].cost < d[minid].cost
                minid = id
            end
        end
        minid
    end

    start_node = Node(start[1], start[2], 0, (nothing, nothing))
    goal_node = Node(goal[1], goal[2], typemax(Float64), (nothing, nothing))

    open_set = Dict()
    closed_set = Dict()

    start_id = (start_node.x, start_node.y)
    open_set[start_id] = start_node

    while true
        temp_id = find_MinCost_id(open_set)
        temp_node = open_set[temp_id]

        if temp_node.x == goal_node.x & temp_node.y == goal_node.y
            println("探索終了")
            goal_node.parent = temp_node.parent
            goal_node.cost = temp_node.cost
            break
        end

        delete!(open_set, temp_id)
        closed_set[temp_id] = temp_node

        os = options(temp_node.x, temp_node.y, x_max, y_max, gridmap)
        for o in os
            Δx, Δy, Δcost = o
            new_node = Node(temp_node.x + Δx, temp_node.y + Δy, temp_node.cost + Δcost, temp_id)
            new_id = (new_node.x, new_node.y)

            if haskey(closed_set, new_id)
                """決定済みのとき"""
                continue
            elseif !haskey(open_set, new_id)
                """未探索のとき"""
                open_set[new_id] = new_node
            else
                """未決定のとき"""
                if open_set[new_id].cost >= new_node.cost
                    open_set[new_id] = new_node
                end
            end
        end
    end

    """決定済み集合から最短経路を組む"""
    function compute_optiomal_path(start_node, goal_node, closed_set)

        rx, ry = [goal_node.x], [goal_node.y]
        parent = goal_node.parent
        cost = goal_node.cost

        while parent != (start_node.x, start_node.y)
            node = closed_set[parent]
            push!(rx, node.x)
            push!(ry, node.y)
            cost += node.cost
            parent = node.parent
        end

        push!(rx, start_node.x)
        push!(ry, start_node.y)

        rx, ry, cost
    end

    rx, ry, total_cost = compute_optiomal_path(start_node, goal_node, closed_set)

    rx, ry, total_cost
end

function draw(rx, ry, start, goal, gridmap)
    plot(
        rx, ry,
        aspect_ratio=1, label="optiomal path", legend=:bottomright,
        xlims=(0, 11), ylims=(0, 11)
    )
    scatter!(start, markershape=:diamond, label="start")
    scatter!(goal, markershape=:star6, label="goal")

    # 障害物位置のデータ作成
    obs_id = findall(gridmap .> 0)  # 障害物のindex検索
    obs_x = Vector{Int64}(undef, length(obs_id))
    obs_y = Vector{Int64}(undef, length(obs_id))
    for i in 1:length(obs_id)
        obs_x[i] = obs_id[i][2]
        obs_y[i] = obs_id[i][1]

    end
    scatter!(obs_x, obs_y, marlershape=:star, label="obstacle")
end

function main_1()
    map_example = [
        0 0 0 0 0 0 1 0 0 0
        0 1 0 0 0 0 0 0 1 0
        0 1 1 1 1 0 0 0 1 0
        0 0 0 0 0 0 0 0 1 0
        0 0 0 1 1 0 0 0 1 0
        0 0 0 1 1 0 0 0 1 0
        0 0 0 1 1 0 0 0 0 0
        0 0 0 0 0 0 1 0 0 0
        0 0 0 0 0 0 1 0 0 0
        0 0 0 0 0 0 1 0 0 0
    ]  # mapの例．0は通過可能な点，1は通貨不可能な点を示す

    start = (1, 1)
    goal = (10, 10)

    rx, ry, total_cost = dijkstra(map_example, start, goal)  # ダイクストラ法実行
    draw(rx, ry, start, goal, map_example)

end

@time main_1()


function gamma_posterior(chain)
    gamma_rank = denserank(chain.gamma)
    gamma_unique = sort(unique(chain.gamma))
    gamma_card = length(unique(gamma_unique))
    gamma_freq = counts(gamma_rank, gamma_card)
    pγ = gamma_freq ./ length(chain.gamma)
    N = length(gamma_unique)
    G = length(gamma_unique[1])
    df = DataFrame(
        [zeros(Bool, N) for k in 2:G],
        [Symbol("group$i") for i in 2:G]
    )
    for i in 1:N, j in 2:G
        df[i, j - 1] = gamma_unique[i][j]
    end
    df[!, :prob] = pγ
    return df
end


### A Pluto.jl notebook ###
# v0.16.0

using Markdown
using InteractiveUtils

# ╔═╡ 142b62e1-364e-45d6-9bca-7c69b794f8ce
md"# Tutorial 1"

# ╔═╡ 27f5936d-4043-4cb6-b795-033f6f2a0674
md"Crash course in Julia basics:"

# ╔═╡ 3f08e241-e53d-4d65-935f-ddf6a6bfbbdd
md"""
Arithmetic, arrays, tuples, strings, dictionaries, functions,
iteration, random numbers, package loading, plotting
"""

# ╔═╡ 3e3bcc53-7840-4af7-af2a-25235e544a31
md"(40 min)"

# ╔═╡ 96774cd2-130c-4455-8af5-148d95ea2900
md"## Julia is a calculator:"

# ╔═╡ 8e8a0cbc-96b9-490c-a6bf-47194d7e8e12
1 + 2^3

# ╔═╡ a65fb83a-f977-47ce-828a-4cc485149963
sqrt(1 + 2^3) # do `sqrt(ans)` in REPL

# ╔═╡ dc16182e-b6b3-47ee-9e54-3f7b3ca87ecb
sin(pi)

# ╔═╡ bc1912b8-c66a-4ab5-ba1f-6363f43ec697
md"""
Query a function's document string using `?sin` at the REPL, or in a
notebook:
"""

# ╔═╡ 048a1195-1107-46e9-95ea-2955abd45275
@doc sin

# ╔═╡ ef72056a-9321-4290-b1b3-dd47e367ba54
md"""
I've forgotten how the arcsin is called. Is it `asin` or `arcsin`? I
can search all doc-strings containing "sine" to locate the method:
"""

# ╔═╡ af710d1b-f8c2-4219-8d7e-8e449afd1c48
apropos("sine")

# ╔═╡ 3e690f1c-4e80-4016-a949-7f3bd292fe31
md"Okay, I see it's called `asin`."

# ╔═╡ b541c5ac-3c15-4e71-8514-99938a2932db
asin(1 + 3*im)

# ╔═╡ c00b48c5-d35b-409a-a0de-1721c1bc2df2
md"## Arrays"

# ╔═╡ 3694faf5-d9a4-4089-bca8-7eec7950fd82
md"One dimensional vectors:"

# ╔═╡ cad4d998-43d7-421b-9873-1b1430e635cc
v = [3, 5, 7]

# ╔═╡ 5c43b0f6-30f4-4ef7-b43e-d0ebe87da176
length(v)

# ╔═╡ 26ce1939-e595-496a-9009-4cfb2012998f
md"A \"row vector\" is a 1 x n array:"

# ╔═╡ d66140c7-958b-4347-abd3-cb77d7a79683
row = [3 5 7]

# ╔═╡ 7aee99b6-3c6e-4a71-879e-dc128f3db7b3
md"Multiple row vectors separated by semicolons or new lines define matrices:"

# ╔═╡ 96c1adb9-1b29-433d-a367-9aa3c6cf34d0
A = [3 5 7
     2 4 6
     1 3 5]

# ╔═╡ ef6a776c-8c63-4d30-bf32-94657e65284e
size(A)

# ╔═╡ 4bb2001b-3b46-40f3-9afd-65a5b5facac9
length(A)

# ╔═╡ b97d12e6-2346-410b-b6c8-54f5ed90a964
md"Accessing elements (Julia indices start at 1 not 0):"

# ╔═╡ 177a30bb-473a-4e5e-9292-fe822525fc77
A[1, 2]

# ╔═╡ 109564cb-ea6f-4264-ae5e-4b83dcbc9675
A[1, 2] == A[2]

# ╔═╡ 26714555-a4ad-49cb-8a29-32f11452654a
A[:, 2] # 2nd column

# ╔═╡ d00b2507-e891-4ec5-a5f1-151fcbe229a2
md"Changing elements:"

# ╔═╡ 4d1dd37a-194e-4ed2-81bc-2d5603785a09
A[1, 1] = 42

# ╔═╡ 38477997-58dd-4248-9d87-1d1fbb0e3997
inv(A) # inverse

# ╔═╡ e61c7ae2-8204-4a20-b950-fddef2a1fb10
isapprox(inv(A)*v, A\v) # but RHS more efficient

# ╔═╡ d38b7040-bb41-44be-951b-64cb2a36c8e3
md"## \"Variables\" in Julia *point* to objects"

# ╔═╡ 9f1e3326-b738-4b60-b0e6-441461cc8770
md"Corollary: all passing of function arguments is pass by reference"

# ╔═╡ 494f9e7c-e53a-4c65-8cb0-de601961bc02
md"Like Python; Unlike R, C or FORTRAN."

# ╔═╡ 259bf65c-6868-47b0-a87d-950c50fb2955
begin
  w = v

  w
end

# ╔═╡ 6ef57fdd-3188-4b19-8448-4bcb089096cd
v[1] = 42

# ╔═╡ ef2473e2-bb35-4bd7-a012-a84240254fb6
v

# ╔═╡ 76dc6ac8-03fe-4c6e-bbdd-9799f7bb2e60
w

# ╔═╡ 5b3a3587-31f7-4a00-97a7-88e3af4f10ee
md"## Tuples"

# ╔═╡ 82f84505-336f-4845-b372-8c9866e51852
md"Similar to vectors but of fixed length and immutable (cannot be changed)"

# ╔═╡ 5b71f2a1-5567-4e59-8f3d-b82f1e7b6f7a
begin
  t1 = (1, 2.0, "cat")
  typeof(t1)
end

# ╔═╡ 249c3a08-3ca6-463e-ab04-579e5608ae53
t1[3]

# ╔═╡ f85e0bbc-9db0-4a9f-86cf-35420d9e6995
md"Tuples also come in a *named* variety:"

# ╔═╡ b8b91bc0-0a63-4695-a29a-8d0445351914
t2 = (i = 1, x = 2.0, animal="cat")

# ╔═╡ e0e4f0da-69c7-4941-be65-28d0fcca50a8
t2.x

# ╔═╡ 166a1850-cddb-4f72-9a2e-c468345d87d1
md"## Strings and relatives"

# ╔═╡ 44ef9d71-f75b-4140-b5fa-cb79ebf595ef
begin
  a_string = "the cat"
  a_character = 't'
  a_symbol = :t
end

# ╔═╡ ff4b5618-4be7-4111-91c5-524da38aa391
a_string[1] == a_character

# ╔═╡ fb16016a-87d7-4e9f-ad90-b4405b216771
md"""
A `Symbol` is string-like but
[interned](https://en.wikipedia.org/wiki/String_interning). Generally
use `String` for ordinary textual data, but use `Symbol` for
language reflection (metaprogramming). For example:
"""

# ╔═╡ 20ba7245-15ed-417e-895b-997a92b731e0
isdefined(Main, :z)

# ╔═╡ 542ec737-2c88-4293-a525-e9b04a4bd246
begin
  z = 1 + 2im
  isdefined(Main, :z)
end

# ╔═╡ 8b1e2b91-947c-4046-80f0-ba7781e173cf
z.im

# ╔═╡ 3dcab89a-20ca-462a-9cbc-9a2c39793333
fieldnames(typeof(z))

# ╔═╡ 65eb310b-82e8-4a1d-b887-2b16710e5502
md"Symbols are generalized by *expressions*:"

# ╔═╡ 35998a8b-6fda-4a2f-9452-5298a8a4a401
begin
  ex = :(z == 3)
  eval(ex)
end

# ╔═╡ 2db5e586-237b-45aa-b017-96f1602f2cad
md"If this is confusing, forget it for now."

# ╔═╡ 8b1c08d9-f8dd-4ec9-8be2-2a4017c45345
md"## Dictionaries"

# ╔═╡ 4942a061-18dd-41d0-a7ad-80d84f5b0070
d = Dict('a' => "ant", 'z' => "zebra")

# ╔═╡ d1fb9147-4a06-44fa-8378-5aa686f0b407
d['a']

# ╔═╡ 4f4bf27c-b2f9-4c15-9f43-49763e8691a1
begin
  d['b'] = "bat"
  d
end

# ╔═╡ 29d71401-66e5-484e-bb0e-3a45761c733a
keys(d)

# ╔═╡ 4b80c58d-646a-405b-9979-7044b2f2df33
md"The expression 'a' => \"ant\" is itself a stand-alone object:"

# ╔═╡ 44914181-9dd8-49b4-b2a1-fbbde543f620
begin
  pair = 'a' => "ant"

  first(pair)
end

# ╔═╡ 63badfd3-c501-4d7c-910b-f8a9a217eff2
md"## Functions"

# ╔═╡ e089aa5e-8d6e-48a4-aa36-d2e8546da46a
md"Three ways to define a generic function:"

# ╔═╡ c4c0e393-0c62-4e8e-88a1-a2bf91434413
begin
  foo(x) = x^2 # METHOD 1 (inline)
  foo(3)
end

# ╔═╡ 115b6c86-a0f2-4546-a1cb-54e34396a855
md"or"

# ╔═╡ 3c8e1237-ff55-4dd0-8037-0de5806e1a54
3 |> foo

# ╔═╡ 115b6c86-a0f2-4546-9961-a09632c33fb0
md"or"

# ╔═╡ 169c078b-2f81-4b5e-b7cc-46f4ef988c68
3 |> x -> x^2 # METHOD 2 (anonymous)

# ╔═╡ 115b6c86-a0f2-4546-90fb-7f2721f6fcc7
md"or"

# ╔═╡ 1d690f10-1e44-4e9a-af65-ccd25ecb9818
begin
  function foo2(x) # METHOD 3 (verbose)
      y = x
      z = y
      w = z
      return w^2
  end

  foo2(3)
end

# ╔═╡ 7a568a60-5212-4c87-8890-979691212d9b
md"## Basic iteration"

# ╔═╡ 77d49aef-9644-41c2-a6fa-eaf14df5d44b
md"Here are three ways to square the integers from 1 to 10."

# ╔═╡ c3d721a1-ae10-4c0b-8025-5084804a9f6c
md"METHOD 1 (explicit loop):"

# ╔═╡ d9016c54-9656-4684-9e90-09023d201062
begin
  squares = [] # or Int[] if performance matters
  for x in 1:10
      push!(squares, x^2)
  end

  squares
end

# ╔═╡ a707446f-b9f3-4405-b7bb-2f33ef765cc0
md"METHOD 2 (comprehension):"

# ╔═╡ f86cc305-d2b6-4be3-9625-01172c4a0081
[x^2 for x in 1:10]

# ╔═╡ a0c15cdb-6294-47bc-af4f-11c7de9e21dd
md"METHOD 3 (map):"

# ╔═╡ 7d2bb0cb-fe63-443b-8dba-241d9b744683
map(x -> x^2, 1:10)

# ╔═╡ 4455b804-b818-4fa0-a6e5-2f0b4dca5c59
md"METHOD 4 (broadcasting with dot syntax):"

# ╔═╡ 87188f88-6519-4ffd-8550-20498aa03ed0
(1:10) .^ 2

# ╔═╡ d1f615d1-644a-4358-9e7b-5b943cf6b560
md"## Random numbers"

# ╔═╡ da033108-0b83-4ee8-bce6-0a1379cc1259
typeof(2)

# ╔═╡ c494e441-4853-4a6d-9608-b5032c116833
rand() # sample a Float64 uniformly from interval [0, 1]

# ╔═╡ 2063cdaf-dbb9-4587-b473-4aa968e6937a
rand(3, 4) # do that 12 times and put in a 3 x 4 array

# ╔═╡ a2a2be72-f9fe-4e2c-8d9e-644a1b3cc6b6
randn(3, 4) # use normal distribution instead

# ╔═╡ 6deb462e-668e-475d-ac0a-23ded81445da
rand(Int8) # random elment of type Int8

# ╔═╡ 112438a5-ef2f-4c5d-8535-1a088a6a3228
rand(['a', 'b', 'c'], 10) # 10 random elements from a vector

# ╔═╡ b87a9e77-1e26-4a1f-a39f-7893c73eef39
md"Some standard libraries are needed to do more, for example:"

# ╔═╡ cc2371a4-26be-4325-bcca-51a27994a151
using Random

# ╔═╡ 9ac1f1af-bfdc-4499-9b33-09b4b6661170
randstring(30)

# ╔═╡ cf7971b2-b57f-4386-b45d-88d068bb0fa2
using Statistics  # part of standard library

# ╔═╡ 01f39450-9ce3-46c1-92c8-6db3a590d963
begin
  y = rand(30)
  @show mean(y) quantile(y, 0.75);
end

# ╔═╡ bd256b23-4c27-4ca0-abf5-96f457eb2bdf
md"""
(Use the macro @show before stuff you want printed prefixed by
*what* it is that is being printed.)
"""

# ╔═╡ 133e6f93-bb0b-460d-8a60-2e0214c059f5
md"""
For sampling from more general distributions we need
Distributions.jl package which is not part of the standard library.
"""

# ╔═╡ e7c2edc1-91cb-4ef1-a38a-beb5c7157b57
md"## Loading packages"

# ╔═╡ ab20d999-4cfa-4104-81f5-507803ea9ed6
md"If not in the REPL:"

# ╔═╡ c5e1e96d-691f-4366-9b21-6d4eb642d87e
begin
  using Pkg                        # built-in package manager
  Pkg.activate("env", shared=true) # create a new pkg env
end

# ╔═╡ f7215fcd-0f87-4b4a-b98c-3df1f31879bf
md"""
Add some packages to your environment (latest compatible versions
added by default):
"""

# ╔═╡ fdd4c037-11f7-4777-92b7-4ceba56e97ad
begin
  Pkg.add("Distributions")
  Pkg.add("Plots")
end

# ╔═╡ 64a59295-2630-40c2-b122-214de244406c
md"To load the code for use:"

# ╔═╡ b480add5-8651-4720-8a4b-c83694978e38
begin
  using Distributions
  using Plots

  N = 1000
  samples = rand(Normal(), N);   # equivalent to Julia's built-in `randn(d)`
  samples = (samples).^2;        # square element-wise
end

# ╔═╡ f618c762-a936-433f-a8b6-f7ff516dedc4
g = fit(Gamma, samples)

# ╔═╡ 5db01f65-3d0f-4ba4-81e1-74ef03c2e79e
@show mean(g) median(g) pdf(g, 1)

# ╔═╡ 2a96dfa7-cf5a-4e7f-a04e-0b76409c14a7
md"## Plotting"

# ╔═╡ c89ff154-d5a1-401f-b979-3458f2f26667
begin
  f(x) = pdf(g, x)

  xs = 0:0.1:4 # floats from 0 to 4 in steps of 0.1
  ys = f.(xs)  # apply f element-wise to xs

  plot(xs, ys, xrange=(0,4), yrange=(0,0.2))
  histogram!(samples , normalize=true, alpha=0.4)
end

# ╔═╡ ebe6467e-e767-481b-97e3-6677afc6ca9f
savefig("my_first_plot.png")

# ╔═╡ a7f061b8-d1ed-4b1f-b10e-562ce21caa04
md"# Exercises"

# ╔═╡ 19dae74b-24bd-428f-8f78-ab549ef1544e
md"## Exercise 1"

# ╔═╡ 3af8467f-4988-44dc-a8a4-a53f5149481e
md"Write a function named `total` that adds the elements of its vector input."

# ╔═╡ b04e2fdc-8f1a-4964-8704-dbab8e09b4f1
md"## Exercise 2"

# ╔═╡ 2b88d9c5-61a4-4e1f-a02f-27b6c05e4d02
md"""
Generate a 1000 random samples from the standard normal
distribution. Create a second such sample, and add the two samples
point-wise.  Compute the (sample) mean and variance of the combined
samples. In the same
plot, show a frequency-normalized histogram of the combined samples
and a plot of the pdf for normal distribution with zero mean and
variance `2`.
"""

# ╔═╡ b899ecf2-1e43-4783-bea4-3d467d48781c
md"""
You can use `std` to compute the standard deviation and `sqrt` to
compute square roots.
"""

# ╔═╡ 4f939d4a-e802-4b7d-97c5-2979af8a507d
md"## Exercise 3"

# ╔═╡ 63f410ea-c37a-4761-b639-63c76c72c513
md"The following shows that named tuples share some behaviour with dictionaries:"

# ╔═╡ cb2c42bf-21b5-4e04-8f5a-e1ee9eb5c15c
begin
  t = (x = 1, y = "cat", z = 4.5)
  keys(t)
end

# ╔═╡ b8996fa3-c046-4fb1-adc5-e671db8bca5d
t[:y]

# ╔═╡ 6584c100-ffd6-4fa1-86f1-0ea38de21abb
md"""
Write a function called `dict` that converts a named tuple to an
actual dictionary. You can create an empty dictionary using `Dict()`.
"""

# ╔═╡ 135dac9b-0bd9-4e1d-a55c-0cb44ab816d7
md"""
---

*This notebook was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).*
"""

# ╔═╡ Cell order:
# ╟─142b62e1-364e-45d6-9bca-7c69b794f8ce
# ╟─27f5936d-4043-4cb6-b795-033f6f2a0674
# ╟─3f08e241-e53d-4d65-935f-ddf6a6bfbbdd
# ╟─3e3bcc53-7840-4af7-af2a-25235e544a31
# ╟─96774cd2-130c-4455-8af5-148d95ea2900
# ╠═8e8a0cbc-96b9-490c-a6bf-47194d7e8e12
# ╠═a65fb83a-f977-47ce-828a-4cc485149963
# ╠═dc16182e-b6b3-47ee-9e54-3f7b3ca87ecb
# ╟─bc1912b8-c66a-4ab5-ba1f-6363f43ec697
# ╠═048a1195-1107-46e9-95ea-2955abd45275
# ╟─ef72056a-9321-4290-b1b3-dd47e367ba54
# ╠═af710d1b-f8c2-4219-8d7e-8e449afd1c48
# ╟─3e690f1c-4e80-4016-a949-7f3bd292fe31
# ╠═b541c5ac-3c15-4e71-8514-99938a2932db
# ╟─c00b48c5-d35b-409a-a0de-1721c1bc2df2
# ╟─3694faf5-d9a4-4089-bca8-7eec7950fd82
# ╠═cad4d998-43d7-421b-9873-1b1430e635cc
# ╠═5c43b0f6-30f4-4ef7-b43e-d0ebe87da176
# ╟─26ce1939-e595-496a-9009-4cfb2012998f
# ╠═d66140c7-958b-4347-abd3-cb77d7a79683
# ╟─7aee99b6-3c6e-4a71-879e-dc128f3db7b3
# ╠═96c1adb9-1b29-433d-a367-9aa3c6cf34d0
# ╠═ef6a776c-8c63-4d30-bf32-94657e65284e
# ╠═4bb2001b-3b46-40f3-9afd-65a5b5facac9
# ╟─b97d12e6-2346-410b-b6c8-54f5ed90a964
# ╠═177a30bb-473a-4e5e-9292-fe822525fc77
# ╠═109564cb-ea6f-4264-ae5e-4b83dcbc9675
# ╠═26714555-a4ad-49cb-8a29-32f11452654a
# ╟─d00b2507-e891-4ec5-a5f1-151fcbe229a2
# ╠═4d1dd37a-194e-4ed2-81bc-2d5603785a09
# ╠═38477997-58dd-4248-9d87-1d1fbb0e3997
# ╠═e61c7ae2-8204-4a20-b950-fddef2a1fb10
# ╟─d38b7040-bb41-44be-951b-64cb2a36c8e3
# ╟─9f1e3326-b738-4b60-b0e6-441461cc8770
# ╟─494f9e7c-e53a-4c65-8cb0-de601961bc02
# ╠═259bf65c-6868-47b0-a87d-950c50fb2955
# ╠═6ef57fdd-3188-4b19-8448-4bcb089096cd
# ╠═ef2473e2-bb35-4bd7-a012-a84240254fb6
# ╠═76dc6ac8-03fe-4c6e-bbdd-9799f7bb2e60
# ╟─5b3a3587-31f7-4a00-97a7-88e3af4f10ee
# ╟─82f84505-336f-4845-b372-8c9866e51852
# ╠═5b71f2a1-5567-4e59-8f3d-b82f1e7b6f7a
# ╠═249c3a08-3ca6-463e-ab04-579e5608ae53
# ╟─f85e0bbc-9db0-4a9f-86cf-35420d9e6995
# ╠═b8b91bc0-0a63-4695-a29a-8d0445351914
# ╠═e0e4f0da-69c7-4941-be65-28d0fcca50a8
# ╟─166a1850-cddb-4f72-9a2e-c468345d87d1
# ╠═44ef9d71-f75b-4140-b5fa-cb79ebf595ef
# ╠═ff4b5618-4be7-4111-91c5-524da38aa391
# ╟─fb16016a-87d7-4e9f-ad90-b4405b216771
# ╠═20ba7245-15ed-417e-895b-997a92b731e0
# ╠═542ec737-2c88-4293-a525-e9b04a4bd246
# ╠═8b1e2b91-947c-4046-80f0-ba7781e173cf
# ╠═3dcab89a-20ca-462a-9cbc-9a2c39793333
# ╟─65eb310b-82e8-4a1d-b887-2b16710e5502
# ╠═35998a8b-6fda-4a2f-9452-5298a8a4a401
# ╟─2db5e586-237b-45aa-b017-96f1602f2cad
# ╟─8b1c08d9-f8dd-4ec9-8be2-2a4017c45345
# ╠═4942a061-18dd-41d0-a7ad-80d84f5b0070
# ╠═d1fb9147-4a06-44fa-8378-5aa686f0b407
# ╠═4f4bf27c-b2f9-4c15-9f43-49763e8691a1
# ╠═29d71401-66e5-484e-bb0e-3a45761c733a
# ╟─4b80c58d-646a-405b-9979-7044b2f2df33
# ╠═44914181-9dd8-49b4-b2a1-fbbde543f620
# ╟─63badfd3-c501-4d7c-910b-f8a9a217eff2
# ╟─e089aa5e-8d6e-48a4-aa36-d2e8546da46a
# ╠═c4c0e393-0c62-4e8e-88a1-a2bf91434413
# ╟─115b6c86-a0f2-4546-a1cb-54e34396a855
# ╠═3c8e1237-ff55-4dd0-8037-0de5806e1a54
# ╟─115b6c86-a0f2-4546-9961-a09632c33fb0
# ╠═169c078b-2f81-4b5e-b7cc-46f4ef988c68
# ╟─115b6c86-a0f2-4546-90fb-7f2721f6fcc7
# ╠═1d690f10-1e44-4e9a-af65-ccd25ecb9818
# ╟─7a568a60-5212-4c87-8890-979691212d9b
# ╟─77d49aef-9644-41c2-a6fa-eaf14df5d44b
# ╟─c3d721a1-ae10-4c0b-8025-5084804a9f6c
# ╠═d9016c54-9656-4684-9e90-09023d201062
# ╟─a707446f-b9f3-4405-b7bb-2f33ef765cc0
# ╠═f86cc305-d2b6-4be3-9625-01172c4a0081
# ╟─a0c15cdb-6294-47bc-af4f-11c7de9e21dd
# ╠═7d2bb0cb-fe63-443b-8dba-241d9b744683
# ╟─4455b804-b818-4fa0-a6e5-2f0b4dca5c59
# ╠═87188f88-6519-4ffd-8550-20498aa03ed0
# ╟─d1f615d1-644a-4358-9e7b-5b943cf6b560
# ╠═da033108-0b83-4ee8-bce6-0a1379cc1259
# ╠═c494e441-4853-4a6d-9608-b5032c116833
# ╠═2063cdaf-dbb9-4587-b473-4aa968e6937a
# ╠═a2a2be72-f9fe-4e2c-8d9e-644a1b3cc6b6
# ╠═6deb462e-668e-475d-ac0a-23ded81445da
# ╠═112438a5-ef2f-4c5d-8535-1a088a6a3228
# ╟─b87a9e77-1e26-4a1f-a39f-7893c73eef39
# ╠═cc2371a4-26be-4325-bcca-51a27994a151
# ╠═9ac1f1af-bfdc-4499-9b33-09b4b6661170
# ╠═cf7971b2-b57f-4386-b45d-88d068bb0fa2
# ╠═01f39450-9ce3-46c1-92c8-6db3a590d963
# ╟─bd256b23-4c27-4ca0-abf5-96f457eb2bdf
# ╟─133e6f93-bb0b-460d-8a60-2e0214c059f5
# ╟─e7c2edc1-91cb-4ef1-a38a-beb5c7157b57
# ╟─ab20d999-4cfa-4104-81f5-507803ea9ed6
# ╠═c5e1e96d-691f-4366-9b21-6d4eb642d87e
# ╟─f7215fcd-0f87-4b4a-b98c-3df1f31879bf
# ╠═fdd4c037-11f7-4777-92b7-4ceba56e97ad
# ╟─64a59295-2630-40c2-b122-214de244406c
# ╠═b480add5-8651-4720-8a4b-c83694978e38
# ╠═f618c762-a936-433f-a8b6-f7ff516dedc4
# ╠═5db01f65-3d0f-4ba4-81e1-74ef03c2e79e
# ╟─2a96dfa7-cf5a-4e7f-a04e-0b76409c14a7
# ╠═c89ff154-d5a1-401f-b979-3458f2f26667
# ╠═ebe6467e-e767-481b-97e3-6677afc6ca9f
# ╟─a7f061b8-d1ed-4b1f-b10e-562ce21caa04
# ╟─19dae74b-24bd-428f-8f78-ab549ef1544e
# ╟─3af8467f-4988-44dc-a8a4-a53f5149481e
# ╟─b04e2fdc-8f1a-4964-8704-dbab8e09b4f1
# ╟─2b88d9c5-61a4-4e1f-a02f-27b6c05e4d02
# ╟─b899ecf2-1e43-4783-bea4-3d467d48781c
# ╟─4f939d4a-e802-4b7d-97c5-2979af8a507d
# ╟─63f410ea-c37a-4761-b639-63c76c72c513
# ╠═cb2c42bf-21b5-4e04-8f5a-e1ee9eb5c15c
# ╠═b8996fa3-c046-4fb1-adc5-e671db8bca5d
# ╟─6584c100-ffd6-4fa1-86f1-0ea38de21abb
# ╟─135dac9b-0bd9-4e1d-a55c-0cb44ab816d7


using ChebyshevNodes
using Test

@test intoab(5,10,fromab(5,10,8))  == 8
@test intoab(5,10,fromab(5,10,5))  == 5
@test fromab(5,10,intoab(5,10,8))  == 8
@test fromab(5,10,intoab(5,10,10)) == 10


@testset "counts" begin
    for _ in 1:100
        a1, a2 = rand(1:10), rand(1:10)
        b1, b2 = a1 + rand(0:10), a2 + rand(0:10)
        c1 = rand(1:20, b1 - a1 + 1)
        c2 = rand(1:20, b2 - a2 + 1)
        cap = rand((a1+a2):(b1+b2))
        D = Dict{Int,Int}()
        for (i1, c1) in zip(a1:b1, c1)
            for (i2, c2) in zip(a2:b2, c2)
                k = i1 + i2
                if k ≤ cap
                    D[k] = get(D, k, 0) + c1 * c2
                end
            end
        end
        C = count_combinations(cap, Counts(a1:b1, c1), Counts(a2:b2, c2))
        for (i, c) in pairs(C)
            @test get(D, i, 0) == c
        end
        i1, i2 = extrema(keys(D))
        @test firstindex(C) ≤ i1
        @test i2 ≤ lastindex(C) ≤ cap
    end
end

"Simple function to test iterator."
function naive_capped_indices(cap, I)
    [Tuple(ι) for ι in CartesianIndices(map(i -> 1:i, I)) if sum(Tuple(ι)) ≤ cap]
end

@testset "capped cartesian indices" begin
    for _ in 1:100
        I = Tuple(rand(1:5, rand(1:5)))
        cap = length(I) + rand(0:5)
        iter = CappedCartesianIndices(cap, I)
        naive = naive_capped_indices(cap, I)
        @test length(naive) == @inferred length(iter)
        @test collect(iter) == naive
        @test @inferred(largest_indices(iter)) ==
            Tuple(maximum(mapreduce(x -> [x...], hcat, naive); dims = 2))
    end
    @test_throws ArgumentError CappedCartesianIndices(1, (3, 3,)) # sum too low
    @test_throws ArgumentError CappedCartesianIndices(5, (0, 3,)) # negative
    T = typeof(CappedCartesianIndices(5, (3, 3,)))
    @test Base.IteratorSize(T) ≡ Base.HasLength()
    @test Base.IteratorEltype(T) ≡ Base.HasEltype()
    @test Base.eltype(T) ≡ Tuple{Int,Int}
end


using SymbolicML

using Test
using Statistics
using StatsBase

include("functions/runtests_summary.jl")


## Plot recipe
## define a hueristic to work around asymptotes
## just sort of succesful
@recipe function f(pq::AbstractRationalFunction{T}, a=nothing, b=nothing) where {T}

    xlims = get(plotattributes, :xlims, (nothing,nothing))
    ylims = get(plotattributes, :ylims, (nothing, nothing))
    rational_function_trim(pq, a, b, xlims, ylims)

end

isapproxreal(x::Real) = true
isapproxreal(x::Complex{T}) where {T} = imag(x) <= sqrt(eps(real(T)))
function toobig(pq)
    x -> begin
        y = pq(x)
        isinf(y) && return true
        isnan(y) && return true
        abs(y) > 1e8 && return true
        return false
    end
end

function rational_function_trim(pq, a, b, xlims, ylims)

    p,q = lowest_terms(//(pq...), method=:numerical)
    dpq = derivative(p//q)
    p′,q′ = lowest_terms(dpq)

    λs = Multroot.multroot(q).values
    λs = isempty(λs) ? λs : real.(filter(isapproxreal, λs))

    cps = Multroot.multroot(p′).values
    cps = isempty(cps) ? cps : real.(filter(isapproxreal, cps))
    cps = isempty(cps) ? cps : filter(!toobig(pq), cps)

    a = a == nothing ? xlims[1] : a
    b = b == nothing ? xlims[2] : b

    if a==nothing && b==nothing
        u= poly_interval(p)
        v= poly_interval(q)
        a,b = min(first(u), first(v)), max(last(u), last(v))

        if !isempty(λs)
            a,b = min(a, real(minimum(λs))), max(b, real(maximum(λs)))
        end
        if !isempty(cps)
            a,b = min(a, real(minimum(cps))), max(b, real(maximum(cps)))
        end
        a *= (a > 0 ? 1/1.5 : 1.25)
        b *= (b < 0 ? 1/1.5 : 1.25)
    end

    n = 601
    xs = range(a,stop=b, length=n)
    ys = pq.(xs)
    Mcps = isempty(cps) ? 5 : 3*maximum(abs, pq.(cps))
    M = max(5, Mcps, 1.25*maximum(abs, pq.((a,b))))

    lo = ylims[1] == nothing ? -M : ylims[1]
    hi = ylims[2] == nothing ?  M : ylims[2]
    ys′ = [lo <= yᵢ <= hi ? yᵢ : NaN for yᵢ ∈ ys]
    xs, ys′

end


# convert image to semilogy scaled image
function logimage(A,x)
    @assert size(A,1)==length(x)
    B = copy(A)
    y = logspace(log10(x[1]+0.1),log10(x[end]),length(x))
    for i=1:length(x)
        ind = searchsortedlast(x, y[i])
        ind < 1 && (ind=1)
        B[i,:] = A[ind,:]
    end
    return B
end

# bins of averages on log scale
function logbinavg(p::AbstractArray, f::AbstractArray, bins::Integer)
    @assert length(p) == length(f)
    n = length(p)
    fls = logspace(log10(1), log10(f[end]), bins)
    x = zeros(bins)
    j = 1
    for i = 1:bins
        sum = 0
        c = 0
        while( j <= n && f[j] <= fls[i] )
            sum += p[j]
            c += 1
            j += 1
        end
        if c != 0
            x[i] = sum/c
        else
            x[i] = -200dB
        end
    end

    return x, fls
end

function plotall(s::MCAudio; pgwin=tukey(nsamples(s),0.0003), pgnfft=nextfastfft(nsamples(s)), pgbins::Integer=600)

    fs = s.fs
    n = nsamples(s)
    x = s[:left]

    wl = 256*2
    wl > n && (wl=n>>2)
    gain = 10
    dbr = 80

    # audio
    sv = x
    st = [0:1/fs:(n-1)/fs]

    # periodogram
    P0 = 1
    P = periodogram(x; onesided=true, nfft=pgnfft, fs=fs, window=pgwin)
    Pp, Pf = logbinavg(power(P), freq(P), pgbins)
    Pp = 10*log10(Pp./P0)

    #specrogram
    SG = spectrogram(x, wl, wl>>1; onesided=true, fs=fs, window=hanning)
    SGp = 10*log10(power(SG))
    SGp[find(x->(x<-dbr),SGp)] = -dbr
    SGp[find(x->(x>-gain),SGp)] = -gain
    SGt = time(SG)
    SGf = freq(SG)

    # PLOT

    f, ax = plt.subplots(3, 1, sharex=false)
    ax[1][:plot](st, sv, "b")
    ax[1][:set_xlim]([0,st[end]])
    ax[1][:set_ylim]([-1.1,1.1])
    #ax[1][:set_xlabel]("seconds")

    ax[2][:imshow](flipud(logimage(SGp,SGf)),aspect="auto",interpolation="bilinear",extent=[SGt[1],SGt[end],SGf[1]+0.1,SGf[end]])
    #ax[2][:pcolormesh](repmat(vec(SGt),1,length(SGf)),repmat(vec(SGf),length(SGt),1),flipud(SGp))
    ax[2][:set_ylim]([60,SGf[end]])
    ax[2][:set_yscale]("log")

    ax[2][:set_ylabel]("Hz")
    #ax[2][:set_xlabel]("seconds")

    ax[3][:semilogx](Pf, Pp, "b")
    ax[3][:set_xlim]([10,Pf[end]])
    ax[3][:set_ylim]([-100,2])
    #ax[3][:set_xlabel]("Hz")
    ax[3][:set_ylabel]("dB")
    ax[3][:grid](b=true, axis="x", which="minor")
    ax[3][:grid](b=true, axis="y")
    #f[:set_size_inches](10,7)
    draw()
    return f, ax

end


"A type for the ensemble average scattering coefficients.
Here they are discretised in terms of the depth x of the halfspace"
mutable struct AverageWave{T<:AbstractFloat}
    hankel_order::Int # largest hankel order
    x::Vector{T} # spatial mesh
    amplitudes::Array{Complex{T}} # a matrix of the scattering amplitudes, size(A_mat) = (length(x), 2hankel_order +1)
    # Enforce that the dimensions are correct
    function AverageWave{T}(hankel_order::Int, x::Vector{T}, amplitudes::Array{Complex{T}}) where T <: AbstractFloat
        if (length(x), 2*hankel_order+1) != size(amplitudes)[1:2]
            error("The amplitudes of AverageWave does not satisfy size(amplitudes)[1:2] == (length(X), 2*hankel_order+1)")
        end
        new(hankel_order,x,amplitudes)
    end
end

AverageWave(M::Int, x::AbstractVector{T}, as::AbstractArray{Complex{T}}) where T<:AbstractFloat = AverageWave{T}(M,collect(x),collect(as))

function AverageWave(x::AbstractVector{T}, A_mat::Array{Complex{T}}) where T<:Number
    AverageWave(Int((size(A_mat,2)-1)/2), collect(x), A_mat)
end

"Approximates the error in AverageWave.amplitudes due to the mesh AverageWave.x."
function average_error(avg_w::AverageWave)
    ddf = circshift(avg_w.amplitudes, (2,0,0)) - 2*circshift(avg_w.amplitudes, (1,0,0)) + avg_w.amplitudes
    h = (avg_w.x[2] - avg_w.x[1])
    ddf = ddf[3:end,:,:]/(h^2)
    max_ddf = maximum(abs.(ddf))

    # trapezoidal error
    return max_ddf*h^3/12
end

"Calculates an AverageWave from one EffectiveWave"
function AverageWave(xs::AbstractVector{T}, wave_eff::EffectiveWave{T}) where T<:Number

    amps = wave_eff.amplitudes
    ho = wave_eff.hankel_order
    θ_eff = wave_eff.θ_eff

    S = size(amps,2)

    average_amps = [
        im^T(m)*exp(-im*m*θ_eff)*amps[m+ho+1,s]*exp(im*wave_eff.k_eff*cos(θ_eff)*x)
    for x in xs, m=-ho:ho, s=1:S]

    return AverageWave(ho,xs,average_amps)
end

"Calculates an AverageWave from a vector of EffectiveWave"
function AverageWave(xs::AbstractVector{T}, wave_effs::Vector{EffectiveWave{T}}) where T<:Number

    ho = wave_effs[1].hankel_order
    avg_wave_effs = [AverageWave(xs, wave) for wave in wave_effs]
    amps = sum(avg_wave_effs[i].amplitudes[:,:,:] for i in eachindex(avg_wave_effs))

    return AverageWave(ho, xs, amps)
end

"Numerically solved the integral equation governing the average wave. Optionally can use wave_eff to approximate the wave away from the boundary."
function AverageWave(ω::T, medium::Medium{T}, specie::Specie{T};
        radius_multiplier::T = 1.005,
        x::AbstractVector{T} = [zero(T)],
        tol::T = T(1e-4),
        wave_effs::Vector{EffectiveWave{T}} = [zero(EffectiveWave{T})],
        max_size::Int = 1000,
        kws...
    ) where T<:Number

    k = real(ω/medium.c)

    if x == [zero(T)]
        if maximum(abs(w.k_eff) for w in wave_effs) == zero(T)
            wave_effs = effective_waves(real(ω/medium.c), medium, [specie];
                radius_multiplier=radius_multiplier, tol=tol, mesh_points=2, kws...)
        end
        # estimate a large coarse non-dimensional mesh based on the lowest attenuating effective wave
        a12 = T(2)*radius_multiplier*specie.r
        x = x_mesh(wave_effs[1]; tol = tol,  a12 = a12, max_size=max_size)
    end

    X = x.*k
    (MM_quad,b_mat) = average_wave_system(ω, X, medium, specie; tol = tol, kws...);

    M = Int( (size(b_mat,2) - 1)/2 )
    J = length(collect(X)) - 1

    len = (J + 1) * (2M + 1)
    MM_mat = reshape(MM_quad, (len, len));
    b = reshape(b_mat, (len));

    As = MM_mat\b
    As_mat = reshape(As, (J+1, 2M+1, 1))

    return AverageWave(M, collect(X)./k, As_mat)
end

"note that this uses the non-dimensional X = k*depth"
function average_wave_system(ω::T, X::AbstractVector{T}, medium::Medium{T}, specie::Specie{T};
        θin::Float64 = 0.0, tol::T = 1e-6,
        radius_multiplier::T = 1.005,
        scheme::Symbol = :trapezoidal,
        hankel_order::Int = maximum_hankel_order(ω, medium, [specie]; tol = tol),
        t_vecs::Vector{Vector{Complex{T}}} = t_vectors(ω, medium, [specie]; hankel_order = hankel_order),
        kws...
    ) where T<:AbstractFloat

    k = real(ω/medium.c)
    a12k = radius_multiplier*T(2)*real(k*specie.r);
    M = hankel_order;

    J = length(X) - 1
    h = X[2] - X[1]

    PQ_quad = intergrand_kernel(X, a12k; M = M, θin = θin, scheme=scheme);

    MM_quad = [
        (specie.num_density/(k^2))*t_vecs[1][m+M+1]*PQ_quad[l,m+M+1,j,n+M+1] - ( (m==n && j==l) ? 1.0+0.0im : 0.0+0.0im)
    for  l=1:(J+1), m=-M:M, j=1:(J+1), n=-M:M];

    b_mat = [
        -t_vecs[1][m+M+1]*exp(im*X[l]*cos(θin))*exp(im*m*(pi/2.0 - θin))
    for l = 1:(J+1), m = -M:M]

    return (MM_quad,b_mat)
end

"Returns x the mesh used to discretise the integral equations."
function x_mesh(wave_eff_long::EffectiveWave{T}, wave_eff_short::EffectiveWave{T} = wave_eff_long;
        tol::T = T(1e-5),  a12::T = T(Inf),
        max_size::Int = 1000,
        min_size::Int = 5,
        max_x::T = (-log(tol))/abs(cos(wave_eff_long.θ_eff)*imag(wave_eff_long.k_eff))
) where T<:AbstractFloat

    #= The default min_X result in:
        abs(exp(im*min_X*cos(θ_effs[end])*k_effs[end]/k)) < tol
    =#
    # estimate a reasonable derivative based on more rapidly varying wave_eff_short.
    df = abs(wave_eff_short.k_eff * cos(wave_eff_short.θ_eff))

    # Based on Simpson's rule
        # dX  = (tol*90 / (df^4))^(1/5)
    # Based on trapezoidal integration
        dx  = (tol * T(30) / (df^2))^(1/3)

    # prioritise a mesh fine enough to give accurate discrete integrals
    if dx > a12 dx = a12 end

    if max_x/dx + 1 > max_size
        if dx == a12
            max_x = (max_size - 1)*dx
            @warn("The mesh max_size = $max_size which was too small for tol = $tol. Will shrink meshed region.")
        else # otherwise priortise shrink max_x and increase dx proportionally
            a = sqrt(max_x/((max_size - 1)*dx))
            dx = min(dx*a,a12)
            max_x = (max_size - 1)*dx
            @warn("The mesh max_size = $max_size which was too small for tol = $tol. Will make a smaller and coarser mesh.")
        end
    elseif max_x/dx + 1 < min_size
        max_x = (min_size - 1)*dx
    end

    # if whole correction length a12k was given, then make dX/a12k = integer
    if a12 != T(Inf) && dx < a12
        n = ceil(a12 / dx)
        dx = a12/n
    end

    return zero(T):dx:max_x
end


include("euler/euler.jl")
using .Numbers: is_palindrome
using BenchmarkTools
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

function compute(n::Int)::Int
    result = 0
    for i ∈ 10 ^ n ÷ 11 * 11:-11:10 ^ (n - 1)
        for j ∈ 10 ^ n - 1:-1:10 ^ (n - 1)
            candidate = i * j
            if candidate ≤ result
                break
            elseif is_palindrome(candidate)
                result = candidate
            end
        end
    end
    return result
end

compute(2)

compute(3)

@benchmark compute(3)


"""
Refine Float64 to BigFloat precision.
"""
function iterate!(A::AbstractMatrix{T}, ilo::Integer, ihi::Integer, Q::AbstractM) where {T<:BigFloat}

  # approximation of submatrix eigenvalues in Float64 precision
  r = ilo:ihi
  ev = seigendiag(smalleig(Float64.(view(A, r, r)))[1], 1, ihi-ilo+1)
  ihi = refineprecision!(A, ilo, ihi, Q, ev)
  ilo, ihi
end

"""
Refine all eigenvalues of a list close to estimation value
"""
function refineprecision!(A::AbstractMatrix{T}, ilo::Integer, ihi::Integer, Q::AbstractM, ev::AbstractVector) where {T<:Union{Real,Complex}}

  realcase = T <: Real
  evr = reverse(ev)
  k = length(ev)
  while k >= 1
    evk = ev[k]
    evk2 = ! realcase || isreal(evk) ? T(evk) : Complex{T}(evk)

    ihi0 = ihi
    ihi = refineprecision!(A, ilo, ihi, Q, evk)
    if ihi0 - ihi >= 2 && k > 1 && imag(evk) == -imag(ev[k-1])
      k -= 1
    end
    k -= 1
  end
  ihi
end

"""
Refine one eigenvalue close to estimation value
"""
function refineprecision!(A::AbstractMatrix{T}, ilo::Integer, ihi::Integer, Q::AbstractM, ev::Number) where {T<:Union{Real,Complex}}

  ih2 = ! (T <: Real) || isreal(ev) ? ihi : ihi - 1
  r = ih2:ihi
  while !converged(A, ilo, ihi)
    transform_Hess!(A, ilo, ihi, Q, [ev], ihi, 0)
    ev = eig2(view(A, r, r))[end]
  end
  ihi = ih2 - 1
  if ihi >= ilo
    A[ihi+1,ihi] = 0
  end
  ihi
end

"""
Detect convergence of last eigenvalue block of hessenberg matrix.
"""
function converged(A::AbstractMatrix{T}, ilo::Integer, ihi::Integer) where {T<:Union{Real,Complex}}
  if ilo >= ihi
    return true
  end
  if deflation_criterion(A[ihi,ihi-1], A[ihi-1,ihi-1], A[ihi,ihi])
    return true
  end
  if ! (T <: Real) || discriminant(A, ihi-1, ihi) > 0
    return false
  end
  if ilo >= ihi - 1
    return true
  end
  deflation_criterion(A[ihi-1,ihi-2], A[ihi-2,ihi-2], A[ihi-1,ihi-1]) # TODO
end