export
dense,
flatten,
ae,
ae_num,
ae_init,
fc,
fc_num,
fc_init,
fc_to_code,
ae_to_code,
fcx,
bn,
sparse_softmax_cross_entropy_with_logits,
Resnet1D

#----------------------------------------------------
# activation functions

_string2fn = Dict(
    "relu" => relu,
    "tanh" => tanh,
    "sigmoid" => sigmoid,
    "leakyrelu" => leaky_relu,
    "leaky_relu" => leaky_relu,
    "relu6" => relu6,
    "softmax" => softmax,
    "softplus" => softplus,
    "selu" => selu,
    "elu" => elu,
    "concat_elu"=>concat_elu,
    "concat_relu"=>concat_relu,
    "hard_sigmoid"=>hard_sigmoid,
    "swish"=>swish,
    "hard_swish"=>hard_swish,
    "concat_hard_swish"=>concat_hard_swish,
    "sin"=>sin,
    "fourier"=>fourier
)

function get_activation_function(a::Union{Function, String, Nothing})
    if isnothing(a)
        return x->x
    end
    if isa(a, String)
        if haskey(_string2fn, lowercase(a))
            return _string2fn[lowercase(a)]
        else
            error("Activation function $a is not supported, available activation functions:\n$(collect(keys(_string2fn))).")
        end
    else
        return a
    end
end

@doc raw"""
    fcx(x::Union{Array{Float64,2},PyObject}, output_dims::Array{Int64,1},
    θ::Union{Array{Float64,1}, PyObject};
    activation::String = "tanh")

Creates a fully connected neural network with output dimension `o` and inputs $x\in \mathbb{R}^{m\times n}$.

$$x \rightarrow o_1 \rightarrow o_2 \rightarrow \ldots \rightarrow o_k$$

`θ` is the weights and biases of the neural network, e.g., `θ = ae_init(output_dims)`.

`fcx` outputs two tensors:

- the output of the neural network: $u\in \mathbb{R}^{m\times o_k}$.

- the sensitivity of the neural network per sample: $\frac{\partial u}{\partial x}\in \mathbb{R}^{m \times o_k \times n}$

"""
function fcx(x::Union{Array{Float64,2},PyObject}, output_dims::Array{Int64,1},
    θ::Union{Array{Float64,1}, PyObject};
    activation::String = "tanh")
    pth = install("ExtendedNN")
    extended_nn_ = load_op_and_grad(pth, "extended_nn"; multiple=true)
    config = [size(x,2);output_dims]
    x_,config_,θ_ = convert_to_tensor([x,config,θ], [Float64,Int64,Float64])
    x_ = reshape(x_, (-1,))
    u, du = extended_nn_(x_,config_,θ_,activation)
    reshape(u, (size(x,1), config[end])), reshape(du, (size(x,1), config[end], size(x,2)))
end

"""
    ae(x::PyObject, output_dims::Array{Int64}, scope::String = "default";
        activation::Union{Function,String} = "tanh")

Alias: `fc`

Creates a neural network with intermediate numbers of neurons `output_dims`.
"""
function ae(x::Union{Array{Float64},PyObject}, output_dims::Array{Int64}, scope::String = "default";
        activation::Union{Function,String} = "tanh")
    if isa(x, Array)
        x = constant(x)
    end
    flag = false
    if length(size(x))==1
        x = reshape(x, length(x), 1)
        flag = true
    end
    net = x
    variable_scope(scope, reuse=AUTO_REUSE) do
        for i = 1:length(output_dims)-1
            net = dense(net, output_dims[i], activation=activation)
        end
        net = dense(net, output_dims[end])
    end
    if flag && (size(net,2)==1)
        net = squeeze(net)
    end
    return net
end

"""
    ae(x::Union{Array{Float64}, PyObject}, output_dims::Array{Int64}, θ::Union{Array{Float64}, PyObject};
    activation::Union{Function,String, Nothing} = "tanh")

Alias: `fc`

Creates a neural network with intermediate numbers of neurons `output_dims`. The weights are given by `θ`

# Example 1: Explicitly construct weights and biases
```julia
x = constant(rand(10,2))
n = ae_num([2,20,20,20,2])
θ = Variable(randn(n)*0.001)
y = ae(x, [20,20,20,2], θ)
```

# Example 2: Implicitly construct weights and biases
```julia
θ = ae_init([10,20,20,20,2])
x = constant(rand(10,10))
y = ae(x, [20,20,20,2], θ)
```

See also [`ae_num`](@ref), [`ae_init`](@ref).
"""
function ae(x::Union{Array{Float64}, PyObject}, output_dims::Array{Int64}, θ::Union{Array{Float64}, PyObject};
    activation::Union{Function,String, Nothing} = "tanh")
    activation = get_activation_function(activation)
    x = convert_to_tensor(x)
    θ = convert_to_tensor(θ)
    flag = false
    if length(size(x))==1
        x = reshape(x, length(x), 1)
        flag = true
    end
    offset = 0
    net = x

    output_dims = [size(x,2); output_dims]
    for i = 1:length(output_dims)-2
        m = output_dims[i]
        n = output_dims[i+1]
        net = net * reshape(θ[offset+1:offset+m*n], (m, n))  + θ[offset+m*n+1:offset+m*n+n]
        net = activation(net)
        offset += m*n+n
    end
    m = output_dims[end-1]
    n = output_dims[end]
    net = net * reshape(θ[offset+1:offset+m*n], (m, n))  + θ[offset+m*n+1:offset+m*n+n]

    offset += m*n+n
    if offset!=length(θ)
        error("ADCME: the weights and configuration does not match. Required $offset but given $(length(θ)).")
    end
    if flag && (size(net,2)==1)
        net = squeeze(net)
    end
    return net
end

"""
    ae(x::Union{Array{Float64}, PyObject},
        output_dims::Array{Int64},
        θ::Union{Array{Array{Float64}}, Array{PyObject}};
        activation::Union{Function,String} = "tanh")

Alias: `fc`

Constructs a neural network with given weights and biases `θ`

# Example
```julia
x = constant(rand(10,30))
θ = ae_init([30, 20, 20, 5])
y = ae(x, [20, 20, 5], θ) # 10×5
```
"""
function ae(x::Union{Array{Float64}, PyObject},
    output_dims::Array{Int64},
    θ::Union{Array{Array{Float64}}, Array{PyObject}};
    activation::Union{Function,String} = "tanh")
    if isa(θ, Array{Array{Float64}})
        val = []
        for t in θ
            push!(val, θ'[:])
        end
        val = vcat(val...)
    else
        val = []
        for t in θ
            push!(val, reshape(θ, (-1,)))
        end
        vcat(val...)
    end
    ae(x, output_dims, θ, activation=activation)
end

@doc raw"""
    ae_init(output_dims::Array{Int64}; T::Type=Float64, method::String="xavier")
    fc_init(output_dims::Array{Int64})

Return the initial weights and bias values by TensorFlow as a vector. The neural network architecture is

$$o_1 (Input layer) \rightarrow o_2 \rightarrow \cdots \rightarrow o_n (Output layer)$$

Three types of
random initializers are provided

- `xavier` (default). It is useful for `tanh` fully connected neural network.
```math
W^l_i \sim \sqrt{\frac{1}{n_{l-1}}}
```
- `xavier_avg`. A variant of `xavier`
```math
W^l_i \sim \sqrt{\frac{2}{n_l + n_{l-1}}}
```
- `he`. This is the activation aware initialization of weights and helps mitigate the problem
of vanishing/exploding gradients.
```math
W^l_i \sim \sqrt{\frac{2}{n_{l-1}}}
```

# Example
```julia
x = constant(rand(10,30))
θ = ae_init([30, 20, 20, 5])
y = ae(x, [20, 20, 5], θ) # 10×5
```
"""
function ae_init(output_dims::Array{Int64}; T::Type=Float64, method::String="xavier")
    N = ae_num(output_dims)
    W = zeros(T, N)
    offset = 0
    for i = 1:length(output_dims)-1
        m = output_dims[i]
        n = output_dims[i+1]
        if method=="xavier"
            W[offset+1:offset+m*n] = randn(T, m*n) * T(sqrt(1/m))
        elseif method=="xavier_normal"
            W[offset+1:offset+m*n] = randn(T, m*n) * T(sqrt(2/(n+m)))
        elseif method=="xavier_uniform"
            W[offset+1:offset+m*n] = rand(T, m*n) * T(sqrt(6/(n+m)))
        elseif method=="he"
            W[offset+1:offset+m*n] = randn(T, m*n) * T(sqrt(2/(m)))
        else
            error("Method $method not understood")
        end
        offset += m*n+n
    end
    W
end

"""
    ae_num(output_dims::Array{Int64})
    fc_num(output_dims::Array{Int64})

Estimates the number of weights and biases for the neural network. Note the first dimension
should be the feature dimension (this is different from [`ae`](@ref) since in `ae` the feature
dimension can be inferred), and the last dimension should be the output dimension.

# Example
```julia
x = constant(rand(10,30))
θ = ae_init([30, 20, 20, 5])
@assert ae_num([30, 20, 20, 5])==length(θ)
y = ae(x, [20, 20, 5], θ)
```
"""
function ae_num(output_dims::Array{Int64})
    offset = 0
    for i = 1:length(output_dims)-2
        m = output_dims[i]
        n = output_dims[i+1]
        offset += m*n+n
    end
    m = output_dims[end-1]
    n = output_dims[end]
    offset += m*n+n
    return offset
end

function _ae_to_code(d::Dict, scope::String; activation::String)
    i = 0
    nn_code = ""
    while true
        si = i==0 ? "" : "_$i"
        Wkey = "$(scope)backslashfully_connected$(si)backslashweightscolon0"
        bkey = "$(scope)backslashfully_connected$(si)backslashbiasescolon0"
        if haskey(d, Wkey)
            if i!=0
                nn_code *= "    isa(net, Array) ? (net = $activation.(net)) : (net=$activation(net))\n"
                nn_code *= "    #-------------------------------------------------------------------\n"
            end
            nn_code *= "    W$i = aedict$scope[\"$Wkey\"]\n    b$i = aedict$scope[\"$bkey\"];\n"
            nn_code *= "    isa(net, Array) ? (net = net * W$i .+ b$i') : (net = net *W$i + b$i)\n"
            i += 1
        else
            break
        end
    end
    nn_code = """  global nn$scope\n  function nn$scope(net)
$(nn_code)    return net\n  end """
    nn_code
end

"""
    ae_to_code(file::String, scope::String; activation::String = "tanh")

Return the code string from the feed-forward neural network data in `file`. Usually we can immediately evaluate
the code string into Julia session by
```julia
eval(Meta.parse(s))
```
If `activation` is not specified, `tanh` is the default.
"""
function ae_to_code(file::String, scope::String = "default"; activation::String = "tanh")
    d = matread(file)
    s = "let aedict$scope = matread(\"$file\")\n"*_ae_to_code(d, scope; activation = activation)*"\nend\n"
    return s
end

fc_to_code = ae_to_code

# tensorflow layers from contrib
for (op, tfop) in [(:avg_pool2d, :avg_pool2d), (:avg_pool3d, :avg_pool3d),
        (:flatten, :flatten), (:max_pool2d, :max_pool2d), (:max_pool3d, :max_pool3d)]
    @eval begin
        export $op
        $op(args...; kwargs...) = tf.contrib.layers.$tfop(args...; kwargs...)
    end
end

for (op, tfop) in [(:conv1d, :conv1d), (:conv2d, :conv2d), (:conv2d_in_plane, :conv2d_in_plane),
    (:conv2d_transpose, :conv2d_transpose), (:conv3d, :conv3d), (:conv3d_transpose, :conv3d_transpose)]
    @eval begin
        export $op
        function $op(args...;activation = nothing, bias=false,  kwargs...)
            activation = get_activation_function(activation)
            if bias
                biases_initializer = tf.zeros_initializer()
            else
                biases_initializer = nothing
            end
            tf.contrib.layers.$tfop(args...; activation_fn = activation, biases_initializer=biases_initializer, kwargs...)
        end
    end
end

"""
    dense(inputs::Union{PyObject, Array{<:Real}}, units::Int64, args...;
        activation::Union{String, Function} = nothing, kwargs...)

Creates a fully connected layer with the activation function specified by `activation`
"""
function dense(inputs::Union{PyObject, Array{<:Real}}, units::Int64, args...; activation::Union{String, Function, Nothing} = nothing, kwargs...)
    inputs = constant(inputs)
    activation = get_activation_function(activation)
    tf.contrib.layers.fully_connected(inputs, units,  args...; activation_fn=activation, kwargs...)
end

"""
    bn(args...;center = true, scale=true, kwargs...)

`bn` accepts a keyword parameter `is_training`.
# Example
```julia
bn(inputs, name="batch_norm", is_training=true)
```

!!! note
    `bn` should be used with `control_dependency`
    ```julia
    update_ops = get_collection(UPDATE_OPS)
    control_dependencies(update_ops) do
        global train_step = AdamOptimizer().minimize(loss)
    end
    ```
"""
function bn(args...;center = true, scale=true, kwargs...)
    @warn  """
`bn` should be used with `control_dependency`
```julia
update_ops = get_collection(UPDATE_OPS)
control_dependencies(update_ops) do
    global train_step = AdamOptimizer().minimize(loss)
end
```
""" maxlog=1
    kwargs = Dict{Any, Any}(kwargs)
    if :is_training in keys(kwargs)
        kwargs[:training] = kwargs[:is_training]
        delete!(kwargs, :is_training)
    end
    if :scope in keys(kwargs)
        kwargs[:name] = kwargs[:scope]
        delete!(kwargs, :scope)
    end
    tf.layers.batch_normalization(args...;center = center, scale=scale, kwargs...)
end
sparse_softmax_cross_entropy_with_logits(args...;kwargs...) = tf.nn.sparse_softmax_cross_entropy_with_logits(args...;kwargs...)

export group_conv2d
function group_conv2d(inputs::PyObject, filters::Int64, args...;  groups = 1, scope=scope, kwargs...)
    if groups==1
        return conv2d(inputs, filters, args...;scope=scope, kwargs...)
    else
        dims = size(inputs)
        if mod(dims[end], groups)!=0 || mod(filters, groups)!=0
            error("channels and outputs must be the multiples of `groups`")
        end
        n = div(filters, groups)
        in_ = Array{PyObject}(undef, groups)
        out_ = Array{PyObject}(undef, groups)
        for i = 1:groups
py"""
temp = $inputs[:,:,:,$((i-1)*n):$(i*n)]
"""
            in_[i] = py"temp"
            out_[i] = conv2d(in_[i], n, args...;scope=scope*"_group$i", kwargs...)
        end
        out = concat(out_, dims=4)
        return out
    end
end

export separable_conv2d
function separable_conv2d(inputs, num_outputs, args...; activation=nothing, bias=false, kwargs...)
    activation = get_activation_function(activation)
    if bias
        biases_initializer = tf.zeros_initializer()
    else
        biases_initializer = nothing
    end
    tf.contrib.layers.separable_conv2d(inputs, num_outputs, args...;activation_fn = activation,
                        biases_initializer=biases_initializer, kwargs...)
end

export depthwise_conv2d
function depthwise_conv2d(input, num_outputs, args...;
    kernel_size = 3, channel_multiplier = 1, stride = 1, padding="SAME", bias=false, scope="default", reuse=AUTO_REUSE, kwargs...)
    local res, strides
    if isa(kernel_size, Int64)
        kernel_size = (3,3)
    end
    if isa(stride, Int64)
        strides = (1, stride, stride, 1)
    end
    variable_scope(scope, reuse=reuse) do
        filter = get_variable("dconv2d",
                shape=[kernel_size[1], kernel_size[2], size(input,4), channel_multiplier],
                initializer=tf.contrib.layers.xavier_initializer())
        res = tf.nn.depthwise_conv2d(input, filter, strides, padding, args...; kwargs...)
        if bias
            res = tf.contrib.layers.bias_add(res, scope=scope, reuse=AUTO_REUSE)
        end
    end
    return res
end

"""
$(@doc ae)
"""
fc = ae
"""
$(@doc ae_num)
"""
fc_num = ae_num
"""
$(@doc ae_init)
"""
fc_init = ae_init

#------------------------------------------------------------------------------------------
export dropout

"""
    dropout(x::Union{PyObject, Real, Array{<:Real}},
    rate::Union{Real, PyObject}, training::Union{PyObject,Bool} = true; kwargs...)

Randomly drops out entries in `x` with a rate of `rate`.
"""
function dropout(x::Union{PyObject, Real, Array{<:Real}},
    rate::Union{Real, PyObject}, training::Union{PyObject,Bool, Nothing} = nothing ; kwargs...)
    x = constant(x)
    if isnothing(training)
        training = options.training.training
    else
        training = constant(training)
    end
    tf.keras.layers.Dropout(rate, kwargs...)(x, training)
end

export BatchNormalization
mutable struct BatchNormalization
    dims::Int64
    o::PyObject
end

"""
    BatchNormalization(dims::Int64=2; kwargs...)

Creates a batch normalization layer.
# Example
```julia
b = BatchNormalization(2)
x = rand(10,2)
training = placeholder(true)
y = b(x, training)
run(sess, y)
```
"""
function BatchNormalization(dims::Int64=-1; kwargs...)
    local o
    if dims>=1
        o = tf.keras.layers.BatchNormalization(dims-1, kwargs...)
    else
        o = tf.keras.layers.BatchNormalization(dims, kwargs...)
    end
    BatchNormalization(dims, o)
end

function Base.:show(io::IO, b::BatchNormalization)
    print("<BatchNormalization normalization_dim=$(b.dims)>")
end

function (o::BatchNormalization)(x, training=ADCME.options.training.training)
    flag = false
    if get_dtype(x)==Float64
        x = cast(x, Float32)
        flag = true
    end
    out = o.o(x, training)
    if flag
        out = cast(Float64, out)
    end
    return out
end

export Dense, Conv1D, Conv2D, Conv3D, Conv2DTranspose

mutable struct Dense
    hidden_dim::Int64
    activation::Union{String, Function}
    o::PyObject
end

"""
    Dense(units::Int64, activation::Union{String, Function, Nothing} = nothing,
        args...;kwargs...)

Creates a callable dense neural network.
"""
function Dense(units::Int64, activation::Union{String, Function, Nothing} = nothing,
    args...;kwargs...)
    activation = get_activation_function(activation)
    o = tf.keras.layers.Dense(units, activation, args...;kwargs...)
    Dense(units, activation, o)
end

function Base.:show(io::IO, o::Dense)
    print("<Fully connected neural network with $(o.hidden_dim) hidden units and activation function \"$(o.activation)\">")
end

(o::Dense)(x) = o.o(x)

mutable struct Conv1D
    filters
    kernel_size
    strides
    activation
    o::PyObject
end
"""
    Conv1D(filters, kernel_size, strides, activation, args...;kwargs...)

```julia
c = Conv1D(32, 3, 1, "relu")
x = rand(100, 6, 128) # 128-length vectors with 6 timesteps ("channels")
y = c(x) # shape=(100, 4, 32)
```
"""
function Conv1D(filters, kernel_size, strides=1, activation=nothing, args...;kwargs...)
    activation = get_activation_function(activation)
    o = tf.keras.layers.Conv1D(filters, kernel_size, strides,args...; activation = activation, kwargs...)
    Conv1D(filters, kernel_size, strides, activation, o)
end
function Base.:show(io::IO, o::Conv1D)
    print("<Conv1D filters=$(o.filters) kernel_size=$(o.kernel_size) strides=$(o.strides) activation=$(o.activation)>")
end
function (o::Conv1D)(x::Union{PyObject, Array{<:Real,3}})
    x = constant(x)
    @assert length(size(x))==3
    o.o(x)
end

mutable struct Conv2D
    filters
    kernel_size
    strides
    activation
    o::PyObject
end
"""
    Conv2D(filters, kernel_size, strides, activation, args...;kwargs...)

The arrangement is (samples, rows, cols, channels) (data_format='channels_last')
```julia
Conv2D(32, 3, 1, "relu")
```
"""
function Conv2D(filters, kernel_size, strides=1, activation=nothing, args...;kwargs...)
    activation = get_activation_function(activation)
    o = tf.keras.layers.Conv2D(filters, kernel_size, strides,args...; activation = activation, kwargs...)
    Conv2D(filters, kernel_size, strides, activation, o)
end
function Base.:show(io::IO, o::Conv2D)
    print("<Conv2D filters=$(o.filters) kernel_size=$(o.kernel_size) strides=$(o.strides) activation=$(o.activation)>")
end
function (o::Conv2D)(x::Union{PyObject, Array{<:Real,4}})
    x = constant(x)
    @assert length(size(x))==4
    o.o(x)
end

mutable struct Conv2DTranspose
    filters
    kernel_size
    strides
    activation
    o::PyObject
end

function Conv2DTranspose(filters, kernel_size, strides=1, activation=nothing, args...;kwargs...)
    activation = get_activation_function(activation)
    o = tf.keras.layers.Conv2DTranspose(
        filters, kernel_size, strides=strides; activation = activation, kwargs...
    )
    Conv2DTranspose(filters, kernel_size, strides, activation, o)
end

function (o::Conv2DTranspose)(x::Union{PyObject, Array{<:Real,4}})
    x = constant(x)
    @assert length(size(x))==4
    o.o(x)
end

mutable struct Conv3D
    filters
    kernel_size
    strides
    activation
    o::PyObject
end
"""
    Conv3D(filters, kernel_size, strides, activation, args...;kwargs...)

The arrangement is (samples, rows, cols, channels) (data_format='channels_last')
```julia
c = Conv3D(32, 3, 1, "relu")
x = constant(rand(100, 10, 10, 10, 16))
y = c(x)
# shape=(100, 8, 8, 8, 32)
```
"""
function Conv3D(filters, kernel_size, strides=1, activation=nothing, args...;kwargs...)
    activation = get_activation_function(activation)
    o = tf.keras.layers.Conv3D(filters, kernel_size, strides,args...; activation = activation, kwargs...)
    Conv3D(filters, kernel_size, strides, activation, o)
end
function Base.:show(io::IO, o::Conv3D)
    print("<Conv3D filters=$(o.filters) kernel_size=$(o.kernel_size) strides=$(o.strides) activation=$(o.activation)>")
end
function (o::Conv3D)(x::Union{PyObject, Array{<:Real,5}})
    x = constant(x)
    @assert length(size(x))==5
    o.o(x)
end

#------------------------------------------------------------------------------------------
# resnet, adapted from https://github.com/bayesiains/nsf/blob/master/nn/resnet.py
mutable struct ResnetBlock
    use_batch_norm::Bool
    activation::Union{String,Function}
    linear_layers::Array{Dense}
    bn_layers::Array{BatchNormalization}
    dropout_probability::Float64
end

function ResnetBlock(features::Int64; dropout_probability::Float64=0., use_batch_norm::Bool,
     activation::String="relu")
    activation = get_activation_function(activation)
    bn_layers = []
    if use_batch_norm
        bn_layers = [
            BatchNormalization(),
            BatchNormalization()
        ]
    end
    linear_layers = [
        Dense(features, nothing)
        Dense(features, nothing)
    ]
    ResnetBlock(use_batch_norm, activation, linear_layers, bn_layers, dropout_probability)
end

function (res::ResnetBlock)(input)
    x = input
    if res.use_batch_norm
        x = res.bn_layers[1](x)
    end
    x = res.activation(x)
    x = res.linear_layers[1](x)
    if res.use_batch_norm
        x = res.bn_layers[2](x)
    end
    x = res.activation(x)
    x = dropout(x, res.dropout_probability, options.training.training)
    x = res.linear_layers[2](x)
    return x + input
end

mutable struct Resnet1D
    initial_layer::Dense
    blocks::Array{ResnetBlock}
    final_layer::Dense
end

"""
    Resnet1D(out_features::Int64, hidden_features::Int64;
        num_blocks::Int64=2, activation::Union{String, Function, Nothing} = "relu",
        dropout_probability::Float64 = 0.0, use_batch_norm::Bool = false, name::Union{String, Missing} = missing)

Creates a 1D residual network. If `name` is not missing, `Resnet1D` does not create a new entity.
# Example
```julia
resnet = Resnet1D(20)
x = rand(1000,10)
y = resnet(x)
```

# Example: Digit recognition
```
using MLDatasets
using ADCME

# load data
train_x, train_y = MNIST.traindata()
train_x = reshape(Float64.(train_x), :, size(train_x,3))'|>Array
test_x, test_y = MNIST.testdata()
test_x = reshape(Float64.(test_x), :, size(test_x,3))'|>Array

# construct loss function
ADCME.options.training.training = placeholder(true)
x = placeholder(rand(64, 784))
l = placeholder(rand(Int64, 64))
resnet = Resnet1D(10, num_blocks=10)
y = resnet(x)
loss = mean(sparse_softmax_cross_entropy_with_logits(labels=l, logits=y))

# train the neural network
opt = AdamOptimizer().minimize(loss)
sess = Session(); init(sess)
for i = 1:10000
    idx = rand(1:60000, 64)
    _, loss_ = run(sess, [opt, loss], feed_dict=Dict(l=>train_y[idx], x=>train_x[idx,:]))
    @info i, loss_
end

# test
for i = 1:10
    idx = rand(1:10000,100)
    y0 = resnet(test_x[idx,:])
    y0 = run(sess, y0, ADCME.options.training.training=>false)
    pred = [x[2]-1 for x in argmax(y0, dims=2)]
    @info "Accuracy = ", sum(pred .== test_y[idx])/100
end
```
![](https://github.com/ADCMEMarket/ADCMEImages/tree/master/ADCME/assets/resnet.png?raw=true)
"""
function Resnet1D(out_features::Int64, hidden_features::Int64 = 20;
     num_blocks::Int64=2, activation::Union{String, Function, Nothing} = "relu",
     dropout_probability::Float64 = 0.0, use_batch_norm::Bool = false,
     name::Union{String, Missing} = missing)
     if haskey(ADCME.STORAGE, name)
        @info "Reusing $name..."
        return ADCME.STORAGE[name]
     end
     initial_layer = Dense(hidden_features)
     blocks = ResnetBlock[]
     for i = 1:num_blocks
        push!(blocks, ResnetBlock(
            hidden_features,
            dropout_probability = dropout_probability,
            use_batch_norm = use_batch_norm
        ))
     end
     final_layer = Dense(out_features)
     if ismissing(name)
        name = "Resnet1D_"*randstring(10)
     end
     res = Resnet1D(initial_layer, blocks, final_layer)
     ADCME.STORAGE[name] = res
     return res
end

function (res::Resnet1D)(x)
    x = res.initial_layer(x)
    for b in res.blocks
        x = b(x)
    end
    x = res.final_layer(x)
    x
end

function Base.:show(io::IO, res::Resnet1D)
    println("( Input )")
    println("\t↓")
    show(io, res.initial_layer)
    println("\n")
    for i = 1:length(res.blocks)
        show(io, res.blocks[i])
        println("\n")
        println("\t↓")
    end
    show(io, res.final_layer)
    println("\n( Output )")
end


using SpecialFunctions, HypergeometricFunctions

"""
    Mode(modeFunction; compression = :algebraic, kwargs...)
    Mode(modeFunction, gᵢ, gₒ, norm)

Container for a cavity mode.

# Arguments
- `modeFunction`: the wave function of the mode
- `gᵢ`: coupling function for an input cavity
- `gₒ`: coupling function for an output cavity
- `norm`: square integral of the mode up to time t

If only the mode function is provided, then all other quantities are numerically determined.
This is done by mapping ``ℝ → (-1, 1)`` and then numerically solving the resulting integrals.
The argument `compression` picks the map ``ℝ → (-1, 1)`` and can be a [`Compression`](@ref) or
the respective symbol.
Allowed symbols are the keys of the `ChiralWaveguide.Compressions` dictionary.
Additional keyword arguments will be passed to the numerical integrator.
"""
mutable struct Mode
	modeFunction::Function
	gᵢ::Function # Input coupling
	gₒ::Function # Output coupling
	norm::Function # Square integral of the mode function
end

function Mode(mF::Function; compression::Union{Symbol, Compression} = :algebraic, kwargs...)
	gᵢ, gₒ, norm = generateCoupling(mF; compression = compression, kwargs...)
	Mode(mF, gᵢ, gₒ, norm)
end

#-------------------------------------------------------
# Special predefined and numerically accurate modes

"""
    FourierMode(a0, an = [], bn = []; t₀ = 0.0, σ = 1.0)

General Fourier Series for a mode in the time bin `(t₀, t₀ + σ)`, with mode function

``m(t) = a0 / σ + \\frac{2}{σ} \\sum_n a_n \\cos(2π n (t - t₀) / σ) + \\frac{2}{σ} \\sum_n b_n \\sin(2π n (t - t₀) / σ).``

The coefficients `a0`, `an`, `bn` do not have to be normalised to provide the correct coupling rates.
"""
function FourierMode(a0, an = [], bn = []; t₀ = 0.0, σ = 1.0)
  @assert σ > 0.
  function modefunction(t)
    t₀ < t < t₀ + σ || return 0.0

    res  = a0 / √(σ)
    !isempty(an) && (res += sum([√(2/σ) * an[n] * cos(2π * n * (t - t₀) / σ) for n ∈ 1:length(an)]))
    !isempty(bn) && (res += sum([√(2/σ) * bn[n] * sin(2π * n * (t - t₀) / σ) for n ∈ 1:length(bn)]))
    return res
  end

  function norm(t)
    t ≤ t₀     && return 0.0
		t₀ + σ ≤ t && return a0^2 + ( isempty(an) ? 0.0 : sum(abs2, an) ) + ( isempty(bn) ? 0.0 : sum(abs2, bn) )

    res = a0^2 * (t - t₀) / σ
		if !isempty(an)
    	res += sum(
				[an[n]^2 * ((t-t₀)/σ + sin(4π * n * (t-t₀) / σ) / (4π * n)) for n ∈ 1:length(an)]
			)
			res += √(2) * sum(
				[a0 * an[n] * sin(2π * n * (t-t₀) / σ)  / (π * n) for n ∈ 1:length(an)]
			)
			res += sum(
				[an[n] * an[m] * ( sin(2π * (m-n) * (t-t₀) / σ) / (m-n) + sin(2π * (m+n) * (t-t₀) / σ) / (m+n) ) / (2π)
	      for n ∈ 1:length(an), m ∈ 1:length(an) if n != m]
			)
		end

		if !isempty(bn)
    	res += sum(
				[bn[n]^2 * ((t-t₀)/σ - sin(4π * n * (t-t₀) / σ) / (4π * n)) for n ∈ 1:length(bn)]
			)
	    res += √(2) * sum(
				[a0 * bn[n] * (1 - cos(2π * n * (t-t₀) / σ)) / (π * n) for n ∈ 1:length(bn)]
			)
			res += sum(
				[bn[n] * bn[m] * ( sin(2π * (m-n) * (t-t₀) / σ) / (m-n) - sin(2π * (m+n) * (t-t₀) / σ) / (m+n) ) / (2π)
	      for n ∈ 1:length(bn), m ∈ 1:length(bn) if n != m]
			)
		end

		if !isempty(an) && !isempty(bn)
	    res += 2 * sum(
				[an[n] * bn[m] * ( 2m - (m+n) * cos(2π * (m-n) * (t-t₀) / σ) - (m-n) * cos(2π * (m+n) * (t-t₀) / σ) ) / (2π * (m^2 - n^2))
	      for n ∈ 1:length(an), m ∈ 1:length(bn) if n != m]
			)
	    res += 2 * sum(
				[an[n] * bn[n] * sin(2π * n * (t-t₀) / σ)^2 / (2π * n)
	      for n ∈ 1:min(length(an),length(bn))]
			)
		end

    return res
  end

  return Mode(modefunction,
    t -> t₀ < t < t₀ + σ ?  modefunction(t) / √(norm(t₀ + σ) - norm(t)) : 0,
    t -> t₀ < t < t₀ + σ ? -modefunction(t) / √(               norm(t)) : 0,
    norm
  )
end

"""
    HardBoxMode(; t₀ = 0.0, σ = 1.0)

The mode is a box function, i.e., it takes the value ``1 / \\sqrt{σ}`` for ``t₀ < t < t₀ + σ``
and 0 otherwise. This can be regarded as the special case of a [`FourierMode`](@ref), i.e.,
`FourierMode(1, [], [])`.
"""
function HardBoxMode(; t₀ = 0.0, σ = 1.0)
	@assert σ > 0
	function norm(t)
		t < t₀     && return 0.0
		t < t₀ + σ && return t / σ
		return 1.0
	end
  return Mode(
    t -> (t₀ < t) && (t < t₀ + σ) ?  1.0 / √(σ)          : 0.0,
    t -> (t₀ < t) && (t < t₀ + σ) ?  1.0 / √(σ - t + t₀) : 0.0,
    t -> (t₀ < t) && (t < t₀ + σ) ? -1.0 / √(    t - t₀) : 0.0,
		norm
  )
end

"""
    SoftBoxMode(; τ = 0.0, σ = 1.0, n::Int = 10)

The mode is an approximation to the box mode,
``\\mathrm{box}(t) ≈ \\frac{1}{\\sqrt{(2t)^{2n} + 1}}``.
The parameters `τ` and `σ` control the center and width respectively, `n` the exponent.
"""
function SoftBoxMode(; τ = 0.0, σ = 1.0, n::Int = 10)
	@assert σ > 0
  n̄ = 2n
  norm = π * csc(π / n̄) / n̄

	function softBox(t; n::Int = 20)
	  norm = π * csc(π / 2n) / 2n
	  return (1 / √(norm) ) / √( (2t)^(2n) + 1 )
	end

	function softBoxNorm(t)
		return t * _₂F₁(1, 1/n̄, 1.0 + 1/n̄, -(2t)^n̄) / norm + 1/2
	end

  function softBoxCoupling(t)
    if abs(t) < 0.5
      return softBox(t, n = n) / √(1.0 - softBoxNorm(t))
    end
    Θmt = t < 0.0 ? 1.0 : 0.0
    twotⁿ = (2t)^(n̄)

    mainTerm = (-t / (1 - n̄)) * _₂F₁(1, 1, 2 - 1/n̄, 1/(1.0 + twotⁿ) )

    return 1.0 / √(
      mainTerm + norm * (1.0 + twotⁿ) * Θmt
    )
	end

  return Mode(
    t ->  softBox( (t - τ) / σ, n = n)  / √(σ),
    t ->  softBoxCoupling( (t - τ) / σ) / √(σ),
    t -> -softBoxCoupling(-(t - τ) / σ) / √(σ),
		t ->  softBoxNorm(     (t - τ) / σ) / √(σ)
  )
end

"""
    SoftBoxExpMode(; τ = 0.0, σ = 1.0, γ = 10.0)

A box mode with exponential decaying flanks, given by

                       ⌜ 1                     for |t| < 1/2
    softBoxExp(t, γ) = |
                       ⌞ exp(-γ(|t| - 1/2))    for |t| > 1/2

The parameters `τ` and `σ` control the center and width respectively, `γ` the deay rate.
"""
function SoftBoxExpMode(; τ = 0.0, σ = 1.0, γ = 10.0)
	@assert σ > 0
	@assert γ > 0
	function softBoxExp(t; γ = 10.0)
	  t = abs(t)
	  return t > 0.5 ? exp(-γ * (t-1/2)) : 1.0
	end

  function softBoxExpNorm(t)
    if t < -0.5
      return exp(2γ * (t + 1/2)) / 2γ
    elseif t < 0.5
      return 1/(2γ) + t + 1/2
    end

    return 1/γ + 1.0 - exp(-2γ * (t - 1/2)) / 2γ
	end

  function softBoxExpCoupling(t)
    if t < -0.5
      return exp(γ * (t+1/2)) / √(1 + (1/γ) - exp(2γ * (t+1/2))/(2γ))
    elseif t < 0.5
      return 1/√(1/2 + 1/(2γ) - t)
    end

    return √(2γ)
  end

  return Mode(
    t ->  √(γ / (γ * σ + σ)) * softBoxExp( (t - τ) / σ, γ = γ),
    t ->  softBoxExpCoupling( (t - τ) / σ) / √(σ),
    t -> -softBoxExpCoupling(-(t - τ) / σ) / √(σ),
		t ->  √(γ / (γ * σ + σ)) * softBoxExpNorm( (t - τ) / σ)
  )
end

"""
    GaussMode(; τ = 0.0, σ = 1.0)

Numerical stable mode function & couplings of Gaussian wave packet with mean τ and variance σ.
"""
function GaussMode(; τ = 0.0, σ = 1.0)
  return Mode(
    t ->  exp(-(t - τ)^2 / (2 * σ^2) ) / √( σ * √(π)),
    t ->  exp(-(t - τ)^2 / (2 * σ^2) ) / √((σ * √(π)) * erfc( (t - τ) / σ) / 2 + eps(0.0)),
    t -> -exp(-(t - τ)^2 / (2 * σ^2) ) / √((σ * √(π)) * erfc(-(t - τ) / σ) / 2 + eps(0.0)),
		t ->  (1 + erf( (t - τ) / σ)) / 2.0
  )
end

"""
    ExpMode(; t₀ = 0.0, γ = 1.0)

Exponentially raises up to time t₀ or decays after t₀, depending on the sign of the rate γ.
Amplitude decays with rate γ/2, so that the probability density decays with γ.
"""
function ExpMode(; t₀ = 0.0, γ = 1.0)
	@assert γ != 0.0

	if γ < 0.0
		return Mode(
			t -> t > t₀ ?  √(-γ) * exp(γ*(t-t₀)/2) : 0,
			t -> t > t₀ ?  √(-γ) : 0.0,
			t -> t > t₀ ? -√(-γ) * exp(γ*(t-t₀)/2) / √( 1.0 - exp(γ*(t-t₀)) ) : 0.0,
			t -> t > t₀ ? 1.0 - exp(γ*(t-t₀)) : 0.0
		)
	end
	Mode(
		t -> t < t₀ ?  √(γ) * exp(γ*(t-t₀)/2) : 0,
		t -> t < t₀ ?  √(γ) * exp(γ*(t-t₀)/2) / √( 1.0 - exp(γ*(t-t₀)) ) : 0.0,
		t -> t < t₀ ? -√(γ) : 0.0,
		t -> t < t₀ ? exp(γ*(t-t₀)) : 1.0
	)
end


# ========== Transforms On A Group Of Signals ==========
# ----- ACDWT on a set of signals -----
"""
    acdwtall(x, wt[, L])

Computes the autocorrelation discrete wavelet transform (ACDWT) on each slice of signal.

# Arguments
- `x::AbstractArray{T,2} where T<:Number`: Input signals, where each column corresponds to a
  signal.
- `wt::OrthoFilter`: Orthogonal wavelet filter.
- `L::Integer`: (Default: `Wavelets.maxtransformlevels(xᵢ)`) Number of levels of wavelet
  transforms.

# Returns
- `::Array{T,3}`: Slices of transformed signals.

# Examples
```julia
using Wavelets, WaveletsExt

# Generate random signals
x = randn(32, 5)
# Create wavelet
wt = wavelet(WT.db4)

# ACDWT on all signals in x
xw = acdwtall(x, wt)
```

**See also:** [`acdwt`](@ref)
"""
function acdwtall(x::AbstractArray{T,2},
                  wt::OrthoFilter,
                  L::Integer = maxtransformlevels(x,1)) where T<:Number
    # Allocate space for transforms
    n,N = size(x)
    xw = Array{T,3}(undef, (n,L+1,N))
    # Dimension to slice
    dim_xw = 3
    dim_x = 2
    # Compute transforms
    @inbounds begin
        @views for (xwᵢ, xᵢ) in zip(eachslice(xw, dims=dim_xw), eachslice(x, dims=dim_x))
            acdwt!(xwᵢ, xᵢ, wt, L)
        end
    end
    return xw
end

"""
    iacdwtall(xw[, wt])

Computes the inverse autocorrelation discrete wavelet transform (IACDWT) on each slice of
signal.

# Arguments
- `xw::AbstractArray{T,3} where T<:Number`: ACDWT-transformed signal.
- `wt::Union{OrthoFilter, Nothing}`: (Default: `nothing`) Orthogonal wavelet filter.

# Returns
- `::Array{T,2}`: Slices of reconstructed signals.

# Examples
```julia
using Wavelets, WaveletsExt

# Generate random signals
x = randn(32, 5)
# Create wavelet
wt = wavelet(WT.db4)

# ACDWT on all signals in x
xw = acdwtall(x, wt)

# IACDWT on all signals in xw
x̂ = iacdwtall(xw)
```

**See also:** [`iacdwt`](@ref)
"""
function iacdwtall(xw::AbstractArray{T,3}, wt::Union{OrthoFilter,Nothing} = nothing) where
                   T<:Number
    # Allocate space for transforms
    n,_,N = size(xw)
    x = Array{T,2}(undef, (n,N))
    # Dimension to slice
    dim_x = 2
    dim_xw = 3
    # Compute transforms
    @inbounds begin
        @views for (xᵢ, xwᵢ) in zip(eachslice(x, dims=dim_x), eachslice(xw, dims=dim_xw))
            iacdwt!(xᵢ, xwᵢ)
        end
    end
    return x
end

# ----- ACWPT on a set of signals -----
"""
    acwptall(x, wt[, L])

Computes the autocorrelation wavelet packet transform (ACWPT) on each slice of signal.

# Arguments
- `x""AbstractArray{T,2} where T<:Number`: Input signals, where each column corresponds to a
  signal.
- `wt::OrthoFilter`: Orthogonal wavelet filter.
- `L::Integer`: (Default: `Wavelets.maxtransformlevels(xᵢ)`) Number of levels of wavelet
  transforms.

# Returns
- `::Array{T,3}`: Slices of transformed signals.

# Examples
```julia
using Wavelets, WaveletsExt

# Generate random signals
x = randn(32, 5)
# Create wavelet
wt = wavelet(WT.db4)

# ACWPT on all signals in x
xw = acwptall(x, wt)
```

**See also:** [`acwpt`](@ref)
"""
function acwptall(x::AbstractArray{T,2},
                  wt::OrthoFilter,
                  L::Integer = maxtransformlevels(x,1)) where T<:Number
    # Allocate space for transforms
    n,N = size(x)
    xw = Array{T,3}(undef, (n,1<<L,N))
    # Dimension to slice
    dim_xw = 3
    dim_x = 2
    # Compute transforms
    @inbounds begin
        @views for (xwᵢ, xᵢ) in zip(eachslice(xw, dims=dim_xw), eachslice(x, dims=dim_x))
            acwpt!(xwᵢ, xᵢ, wt, L)
        end
    end
    return xw
end

"""
    iacwptall(xw[, wt])

Computes the inverse autocorrelation wavelet packet transform (IACWPT) on each slice of
signal.

# Arguments
- `xw::AbstractArray{T,3} where T<:Number`: ACWPT-transformed signal.
- `wt::Union{OrthoFilter, Nothing}`: (Default: `nothing`) Orthogonal wavelet filter.

# Returns
- `::Array{T,2}`: Slices of reconstructed signals.

# Examples
```julia
using Wavelets, WaveletsExt

# Generate random signals
x = randn(32, 5)
# Create wavelet
wt = wavelet(WT.db4)

# ACWPT on all signals in x
xw = acwptall(x, wt)

# IACWPT on all signals in xw
x̂ = iacwptall(xw)
```

**See also:** [`iacwpt`](@ref)
"""
function iacwptall(xw::AbstractArray{T,3}, wt::Union{OrthoFilter,Nothing} = nothing) where
                   T<:Number
    # Allocate space for transforms
    n,_,N = size(xw)
    x = Array{T,2}(undef, (n,N))
    # Dimension to slice
    dim_xw = 3
    dim_x = 2
    # Compute transforms
    @inbounds begin
        @views for (xᵢ, xwᵢ) in zip(eachslice(x, dims=dim_x), eachslice(xw, dims=dim_xw))
            iacwpt!(xᵢ, xwᵢ)
        end
    end
    return x
end

# ----- ACWPD on a set of signals -----
"""
    acwpdall(x, wt[, L])

Computes the autocorrelation wavelet packet decomposition (ACWPD) on each slice of signal.

# Arguments
- `x""AbstractArray{T,2} where T<:Number`: Input signals, where each column corresponds to a
  signal.
- `wt::OrthoFilter`: Orthogonal wavelet filter.
- `L::Integer`: (Default: `Wavelets.maxtransformlevels(xᵢ)`) Number of levels of wavelet
  transforms.

# Returns
- `::Array{T,3}`: Slices of transformed signals.

# Examples
```julia
using Wavelets, WaveletsExt

# Generate random signals
x = randn(32, 5)
# Create wavelet
wt = wavelet(WT.db4)

# ACWPD on all signals in x
xw = acwpdall(x, wt)
```

**See also:** [`acwpd`](@ref)
"""
function acwpdall(x::AbstractArray{T,2},
                  wt::OrthoFilter,
                  L::Integer = maxtransformlevels(x,1)) where T<:Number
    # Allocate space for transforms
    n,N = size(x)
    xw = Array{T,3}(undef, (n,1<<(L+1)-1,N))
    # Dimension to slice
    dim_xw = 3
    dim_x = 2
    # Compute transforms
    @inbounds begin
        @views for (xwᵢ, xᵢ) in zip(eachslice(xw, dims=dim_xw), eachslice(x, dims=dim_x))
            acwpd!(xwᵢ, xᵢ, wt, L)
        end
    end
    return xw
end

"""
    iacwpdall(xw[, wt, L])
    iacwpdall(xw, L)
    iacwpdall(xw, wt, tree)
    iacwpdall(xw, tree)

Computes the inverse autocorrelation wavelet packet decomposition (IACWPD) on each slice of
signal.

# Arguments
- `xw::AbstractArray{T,3} where T<:Number`: ACWPT-transformed signal.
- `wt::Union{OrthoFilter, Nothing}`: (Default: `nothing`) Orthogonal wavelet filter.
- `L::Integer`: (Default: `Wavelets.maxtransformlevels(xᵢ)`) Number of levels of wavelet
  transforms.
- `tree::BitVector`: Binary tree for inverse transform to be computed accordingly.

# Returns
- `::Array{T,2}`: Slices of reconstructed signals.

# Examples
```julia
using Wavelets, WaveletsExt

# Generate random signals
x = randn(32, 5)
# Create wavelet
wt = wavelet(WT.db4)

# ACWPD on all signals in x
xw = acwpdall(x, wt)

# IACWPD on all signals in xw
x̂ = iacwpdall(xw)
x̂ = iacwpdalll(xw, maketree(x))
x̂ = iacwpdall(xw, 5)
```

**See also:** [`iacwpd`](@ref)
"""
function iacwpdall(xw::AbstractArray{T,3},
                   wt::Union{OrthoFilter,Nothing} = nothing,
                   L::Integer = maxtransformlevels(x,1)) where T<:Number
    return iacwpdall(xw, L)
end

function iacwpdall(xw::AbstractArray{T,3}, L::Integer) where T<:Number
    return iacwpdall(xw, maketree(size(xw,1), L, :full))
end

function iacwpdall(xw::AbstractArray{T,3},
                   wt::Union{OrthoFilter, Nothing},
                   tree::BitVector) where T<:Number
    return iacwpdall(xw, tree)
end

function iacwpdall(xw::AbstractArray{T,3}, tree::BitVector) where T<:Number
    # Allocate space for transforms
    n,_,N = size(xw)
    x = Array{T,2}(undef, (n,N))
    # Dimension to slice
    dim_xw = 3
    dim_x = 2
    # Compute transforms
    @inbounds begin
        @views for (xᵢ, xwᵢ) in zip(eachslice(x, dims=dim_x), eachslice(xw, dims=dim_xw))
            iacwpd!(xᵢ, xwᵢ)
        end
    end
    return x
end


### A Pluto.jl notebook ###
# v0.12.18

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 60941eaa-1aea-11eb-1277-97b991548781
using Revise,PlutoUI,ExtendableGrids

# ╔═╡ df8aefc0-5273-11eb-1a67-3552732deae0
using PyPlot

# ╔═╡ 08fa72d6-5274-11eb-03bd-ef885fdd216d
P=PyPlot

# ╔═╡ e5cec5c8-5273-11eb-0e0b-03960befcc71
X=collect(0:0.5:10)

# ╔═╡ f088d080-5273-11eb-2085-e5d4734512ff
g1=simplexgrid(X)

# ╔═╡ ff091d22-5273-11eb-2645-89630b592fa0
gridplot(g1,Plotter=P,resolution=(500,500))

# ╔═╡ 2dd685fe-5274-11eb-173c-6f2ccaf41aae
gridplot(g1,map(x->sin(x),g1),Plotter=P)

# ╔═╡ 8d1857ae-5274-11eb-265d-09604d9dba94
g2=simplexgrid(X,X)

# ╔═╡ 9526886c-5274-11eb-33a6-d95a16547472
gridplot(g2,Plotter=P)

# ╔═╡ ae81ae16-5274-11eb-36e6-a1dd77e51abb
gridplot(g2,map((x,y)->(sin(x)*exp(-(y-5)^2)),g2),Plotter=P,colormap=:hot)

# ╔═╡ eb380f5a-5274-11eb-2a7a-87c75706a2ae
p=GridPlotContext(Plotter=P,layout=(1,2),fignum=2,resolution=(800,400))

# ╔═╡ 0114cb24-5275-11eb-3013-0f3b170f25de
begin
	gridplot!(p[1,1],g1,map(x->sin(x),g1),title="1D")
	gridplot!(p[1,2],g2,map((x,y)->(sin(x)*exp(-(y-5)^2)),g2),colormap=:hot,title="2D")
	reveal(p)
end

# ╔═╡ 677d67f4-5275-11eb-17d6-1954c7cd0131
g3=simplexgrid(X,X,X)

# ╔═╡ aaab7994-5275-11eb-159b-fdd7f6bee8c6
@bind zplane Slider(0:0.51:10,default=4)

# ╔═╡ 6dbcc196-5275-11eb-0bd5-2941757247b3
gridplot(g3,Plotter=P,zplane=zplane)

# ╔═╡ d6f7eb72-5275-11eb-2fa5-8b7b03a4743b
@bind flevel Slider(0:0.2:20,default=0.5)

# ╔═╡ 77e28b4c-5275-11eb-199f-253cae56c928
gridplot(g3,map((x,y,z)->(sin(x)*exp(0.1*y)*z),g3),Plotter=P,flevel=flevel,zplane=zplane,colormap=:brg)

# ╔═╡ Cell order:
# ╠═60941eaa-1aea-11eb-1277-97b991548781
# ╠═df8aefc0-5273-11eb-1a67-3552732deae0
# ╠═08fa72d6-5274-11eb-03bd-ef885fdd216d
# ╠═e5cec5c8-5273-11eb-0e0b-03960befcc71
# ╠═f088d080-5273-11eb-2085-e5d4734512ff
# ╠═ff091d22-5273-11eb-2645-89630b592fa0
# ╠═2dd685fe-5274-11eb-173c-6f2ccaf41aae
# ╠═8d1857ae-5274-11eb-265d-09604d9dba94
# ╠═9526886c-5274-11eb-33a6-d95a16547472
# ╠═ae81ae16-5274-11eb-36e6-a1dd77e51abb
# ╠═eb380f5a-5274-11eb-2a7a-87c75706a2ae
# ╠═0114cb24-5275-11eb-3013-0f3b170f25de
# ╠═677d67f4-5275-11eb-17d6-1954c7cd0131
# ╠═6dbcc196-5275-11eb-0bd5-2941757247b3
# ╠═aaab7994-5275-11eb-159b-fdd7f6bee8c6
# ╠═77e28b4c-5275-11eb-199f-253cae56c928
# ╠═d6f7eb72-5275-11eb-2fa5-8b7b03a4743b


function resample(x::AbstractArray{T,2}, fact::Integer=10, fun::Function=first) where {T <: Real}
    nrow, ncol = size(x)
    nrow2 = fld(nrow, fact)
    ncol2 = fld(ncol, fact)

    # type = ifelse(fun in [first, last], T, Float32)
    out = zeros(T, nrow2, ncol2)

    @inbounds for i in 1:nrow2
        indx = (i-1)*fact + 1 : i*fact
        @inbounds for j in 1:ncol2
            indy = (j-1)*fact + 1 : j*fact
            # out[i, j] = fun(x[indx, indy])
            out[i, j] = fun(view(x, indx, indy));
        end
    end
    out
end

function resample(x::AbstractArray{T,3}, fact::Integer=10, fun::Function=first) where {T <: Real}
    nrow, ncol, ntime = size(x)
    nrow2 = fld(nrow, fact)
    ncol2 = fld(ncol, fact)

    out = zeros(T, nrow2, ncol2, ntime)
    @inbounds for i in 1:ntime
        # out[:, :, i] = resample(view(x, :, :, i), fact, fun)
        out[:, :, i] = resample(@view(x[:, :, i]), fact, fun)
    end
    out
end

function resample(x::GeoArray, fact::Integer=10, fun::Function=first)
    # reconstruct the a.f
    GeoArray(resample(x.A, fact, fun), x.f, x.crs)
end


function test{T,N}(x::Array{T,N})
  return N, T, size(x)
end

println(test([1,2,3,4]))
println(test([1 2 3 4]))
println(test([1 2; 3 4]))
println(test(zeros(Int, 3, 3, 3)))

function addone(x)
  y = zeros(Int, length(x))
  for i in 1:length(x)
    y[i] = x[i] + 1
  end
  return y
end

println(addone([1,2 ,3]))


distance(p1::Complex, p2::Complex) = hypot(p2 - p1)
distance((p1, p2)::Tuple{Complex, Complex}) = distance(p1, p2)
function distance(p1::AbstractPoint2D, p2::AbstractPoint2D)
    z1 = complex(p1._x, p1._y)
    z2 = complex(p2._x, p2._y)
    return distance(z1, z2)
end

"""
    longedge(edge, tolerance, geom2fcn)

Return true if `edge` has length greater than `tolerance`.
"""
function longedge(edge::DelaunayEdge, tolerance, g2f::Geometry2Function)
    return distance(g2f(edge)) > tolerance
end

"""
    fcn2geom(x, xmin, xmax)

Linearly scale `x`, which exists ∈ (`xmin`, `xmax`) into `RootsAndPoles` domain.
"""
function fcn2geom(x, xmin, xmax)
    return (MAXCOORD - MINCOORD)*(x - xmin)/(xmax - xmin) + MINCOORD
end

# TODO: See pull request #50 on VoronoiDelaunay.jl which handles the space
# mapping automatically.

"""
    fcn2geom(z, rmin, rmax, imin, imax)

Linearly scale complex `z` within real range (`rmin`, `rmax`) and imaginary range (`imin`,
`imax`) into `RootsAndPoles` domain.
"""
function fcn2geom(z, rmin, rmax, imin, imax)
    zr = fcn2geom(real(z), rmin, rmax)
    zi = fcn2geom(imag(z), imin, imax)
    return complex(zr, zi)
end

"""
    geom2fcn(g, xmin, xmax)

Linearly scale `g`, which exists in the `RootsAndPoles` domain into the range of the
original function, bounded by (`xmin`, `xmax`).
"""
function geom2fcn(g, xmin, xmax)
    return (xmax - xmin)*(g - MINCOORD)/(MAXCOORD - MINCOORD) + xmin
end

"""
    geom2fcn(z, rmin, rmax, imin, imax)

Scale complex value `z` from the `RootsAndPoles` geometry domain to the comlex function""
domain with real range (`rmin`, `rmax`) and imaginary range (`imin`, `imax`).
"""
geom2fcn(z, rmin, rmax, imin, imax) = geom2fcn(real(z), imag(z), rmin, rmax, imin, imax)

"""
    geom2fcn(zr, zi, rmin, rmax, imin, imax)

Scale real and complex components `zr`, `zi` from the `RootsAndPoles` geometry domain to the
function domain.
"""
geom2fcn(zr, zi, rmin, rmax, imin, imax) = complex(geom2fcn(zr, rmin, rmax),
                                                   geom2fcn(zi, imin, imax))

"""
    geom2fcn(pt::AbstractPoint2D, rmin, rmax, imin, imax)

Scale `pt` from the geometry domain to the function domain.
"""
geom2fcn(pt::AbstractPoint2D, rmin, rmax, imin, imax) = geom2fcn(getx(pt), gety(pt), rmin,
                                                                 rmax, imin, imax)

"""
    geom2fcn(edge::DelaunayEdge, rmin, rmax, imin, imax)

Scale both points of `edge` from the geometry domain to the function domain.
"""
function geom2fcn(edge::DelaunayEdge{P}, rmin, rmax, imin, imax) where P <: AbstractPoint2D
    return (geom2fcn(geta(edge), rmin, rmax, imin, imax),
        geom2fcn(getb(edge), rmin, rmax, imin, imax))
end

"""
    getplotdata(tess, quadrants, phasediffs, g2f)

Return complex coordinates of every tesselation edge and a vector with the corresponding
index specifiying its phase quadrant.

This function is useful for plotting the final tesselation using results from calling
`grpf` with the `PlotData()` argument.

The edge colors are (where Q means "quadrant"):

| Color index |         Meaning         |  Quadrant phase   |
|:-----------:|:-----------------------:|:-----------------:|
|      1      |           Q 1           | 0 ≤ arg f < π/2   |
|      2      |           Q 2           | π/2 ≤ arg f < π   |
|      3      |           Q 3           | π ≤ arg f < 3π/2  |
|      4      |           Q 4           | 3π/2 ≤ arg f < 2π |
|      5      | phase change (boundary) |         -         |
|      6      | candidate edge          |         -         |

# Example

Although not self contained, this example shows how `getplotdata` can be used:

```julia
roots, poles, quads, phasediffs, tess, g2f = grpf(f, origcoords, PlotData());
z, edgecolors = getplotdata(tess, quadrants, phasediffs, g2f);

using Plots
plot(real(z), imag(z), group=edgecolors)
```
"""
function getplotdata(tess, quadrants, phasediffs, g2f)
    edges = [e for e in delaunayedges(tess)]

    edgecolors = fill(5, length(edges))
    for ii in eachindex(edges)
        if phasediffs[ii] == 2  # candidate edges
            edgecolors[ii] = 6
        elseif phasediffs[ii] == 0
            edgecolors[ii] = quadrants[getindex(geta(edges[ii]))]
        end
    end

    edgecolors = repeat(edgecolors, inner=3)  # x, y is xₐ, xᵦ, NaN, repeat

    x, y = getplotxy(edges)
    z = g2f.(x, y)

    I = sortperm(edgecolors)

    return z[I], edgecolors[I]
end


#####
# UM
#####
export UMOracle

#returns zstar, ustar
function suppFcnUM(xs, lquants, uquants, cut_sense=:Max)
    const d = length(lquants)

    zstar = 0.0
    for i in 1:d
        mn, mx = minmax(lquants[i] * xs[i], uquants[i] * xs[i])
        zstar += ifelse(cut_sense == :Min, mn, mx)
    end

    ustar = Array(Float64, d)
    posxopt = (cut_sense == :Min) ? lquants : uquants
    negxopt = (cut_sense == :Min) ? uquants : lquants
    for i in 1:d
       ustar[i] = xs[i] > 0 ? posxopt[i] : negxopt[i]
    end

    return zstar, ustar
end

function calc_s(data, eps_, alpha)
	N, d = size(data)
	if (1-eps_/d)^N  > alpha / 2d
		return N + 1
	else
		dBin = Binomial(N, 1-eps_/d)
		return quantile(dBin, 1-alpha/2d)
	end
end

######################
type UMOracle <: AbstractOracle
    lquants::Vector{Float64}
    uquants::Vector{Float64}
    cut_tol::Float64  ##defaults to 1e-6

    # Other options
    debug_printcut::Bool
end

suppFcn(xs::Vector, w::UMOracle, cut_sense) =
        suppFcnUM(xs, w.lquants, w.uquants, cut_sense)

#Preferred Interface
function UMOracle(data, lbounds, ubounds, eps_, alpha; cut_tol=1e-6, debug_printcut=false)
	s = calc_s(data, eps_, alpha)
	N = size(data, 1)
	@assert N- s + 1 < s "UM: N not sufficiently big N: $N \t s: $s"
	if s == N + 1
		return UMOracle(lbounds, ubounds, cut_tol, debug_printcut)
	else
		data_sort = sort_data_cols(data)
		return UMOracle(vec(data_sort[N-s+1, :]), vec(data_sort[s, :]), cut_tol, debug_printcut)
	end
end

function setup(w::UMOracle, rm::Model, prefs)
    # Extract preferences we care about
    w.debug_printcut = get(prefs, :debug_printcut, false)
    w.cut_tol        = get(prefs, :cut_tol, w.cut_tol)
    rd = JuMPeR.getRobust(rm)
    @assert (rd.numUncs == length(w.lquants)) "Num Uncertainties $(rd.numUncs) doesn't match columns in lquants $(w.lquants)"
    @assert (rd.numUncs == length(w.uquants)) "Num Uncertainties $(rd.numUncs) doesn't match columns in uquants $(w.uquants)"
    @assert (length(rd.uncertaintyset) == 0) "Auxiliary constraints on uncertainties not yet supported"
end


push!(LOAD_PATH, "./solutions")

using Test

using ArithmeticKoans
@testset "Arithmetic Operators" begin
    @test ArithmeticKoans.sum(1, 2) == 3
    @test ArithmeticKoans.difference(8, 1) == 7
    @test ArithmeticKoans.multiplication(10, 2) == 20
    @test ArithmeticKoans.division(35, 5) == 7.0
    @test ArithmeticKoans.integer_division(35, 5) == 7
    @test ArithmeticKoans.power(2, 2) == 4
    @test ArithmeticKoans.modulo(12, 10) == 2
end

using LogicalKoans
@testset "Logical Operators" begin
    @test LogicalKoans.not(true) == false
    @test LogicalKoans.not(false) == true
    @test LogicalKoans.equals(true, true) == true
    @test LogicalKoans.equals(false, false) == true
    @test LogicalKoans.not_equals(true, false) == true
    @test LogicalKoans.less_than(1, 10) == true
    @test LogicalKoans.greater_than(10, 1) == true
    @test LogicalKoans.less_than_or_equal_to(2, 2) == true
    @test LogicalKoans.less_than_or_equal_to(1, 10) == true
    @test LogicalKoans.less_than_or_equal_to(10, 1) == false
    @test LogicalKoans.greater_than_or_equal_to(2, 2) == true
    @test LogicalKoans.greater_than_or_equal_to(10, 1) == true
    @test LogicalKoans.greater_than_or_equal_to(1, 10) == false
end

using BitwiseKoans
@testset "Bitwise Operators" begin
    @test BitwiseKoans.not(2) == -3
    @test BitwiseKoans.and(3, 5) == 1
    @test BitwiseKoans.or(2, 4) == 6
    @test BitwiseKoans.xor_koan(2, 4) == 6
    @test BitwiseKoans.right_shift(2, 1) == 1
    @test BitwiseKoans.left_shift(2, 1) == 4
end

# using ComplexNumberKoans
# @testset "Complex Numbers" begin
# end

# using StringKoans
# @testset "Strings" begin
# end

# using ConversionKoans
# @testset "Type Conversions" begin
# end

# using ArrayKoans
# @testset "Arrays" begin
# end

# using VectorOperationKoans
# @testset "Vector Operators" begin
# end

# using MultiDimensionalArrayKoans
# @testset "Multi-dimensional Arrays" begin
# end

# using IndexingKoans
# @testset "Indexing" begin
# end

# using IterationKoans
# @testset "Iteration" begin
# end

# using ModuleKoans
# @testset "Modules" begin
# end

# using StructKoans
# @testset "Structs" begin
# end

# using UnionKoans
# @testset "Unions" begin
# end

# using FunctionKoans
# @testset "Functions" begin
# end

using HighOrderFunctionKoans
@testset "High Order Functions" begin
    f = (x, y) -> x + y
    @test HighOrderFunctionKoans.round_list([1.5, 2.3, 4.7]) == [2.0, 2.0, 5.0]
    @test HighOrderFunctionKoans.apply_fn(f, 7) == 12
    @test HighOrderFunctionKoans.add_n([1, 2, 4, 5],  3) == [4, 5, 7, 8]
    @test HighOrderFunctionKoans.remove_multiples([10, 12, 15, 18, 20], 3) == [10, 20]
end

# using MethodKoans
# @testset "Methods" begin
# end

using DestructuringKoans
@testset "Destructuring" begin
    @test DestructuringKoans.add_separately(1+2im, 3-7im) == (4, -5)
    @test DestructuringKoans.add_coordinates((2, 3), (4, -1)) == (6, 2)
end

using MultipleDispatchKoans
@testset "Multiple Dispatch" begin
    dog = MultipleDispatchKoans.Dog("goodboi")
    cat = MultipleDispatchKoans.Cat("Sir Catcelot")
    @test MultipleDispatchKoans.hello() == "hello stranger"
    @test MultipleDispatchKoans.hello("julia") == "hello julia"
    @test MultipleDispatchKoans.pet_talk(dog) == "$(dog.name): bark"
    @test MultipleDispatchKoans.pet_talk(cat) == "$(cat.name): meow"
end

# using ValueTypeKoans
# @testset "Value Types" begin
# end

using PythonInterop
@testset "Python Interop" begin
    @test PythonInterop.py_sinpi(0) == 0
    @test round(PythonInterop.py_sinpi(1/pi), Base.Rounding.RoundToZero, digits=2) == 0.84

    pyfib = PythonInterop.get_pyfib()
    @test PythonInterop.jlfib(1, PythonInterop.jlfib) == 1
    @test PythonInterop.jlfib(5, PythonInterop.jlfib) == 5
    @test PythonInterop.jlfib(5, pyfib) == 5
    @test PythonInterop.jlfib(7, pyfib) == 13
end


using Random
using Plots
using SparseArrays
using GeoStats
using SpectralGaussianSimulation
using LaTeXStrings
#reset past font scaling
Plots.scalefontsizes()
#scale fonts
Plots.scalefontsizes(1.75)
#Extracting stored varibles

lightblue = colorant"rgb(63%,74%,78%)"
orange = colorant"rgb(85%,55%,13%)"
silver = colorant"rgb(69%,67%,66%)"
rust = colorant"rgb(72%,26%,6%)"

include("./CovFuncs.jl")
include("./Utils.jl")

Random.seed!(123)

n = 1000
N = n^2
NTest = 20000
NTrain = N - NTest

l = 0.1
covfunc!(r) = matern32!(r, l)
vario = MaternVariogram(range=l, order=1.5)
x, Y = specSampleGrid2D(n, 1, 1, vario)

α = 0.02
β = 0.02

# random shuffle of indices
shuffleIDs = shuffle(1:N)
# Selecting groups of nodes as test data
innerInds = [i for i in 1:N if (norm(x[:, i] .- 0.3) < 0.05 || norm((x[:,i] .- 0.75) .* [1.0, 0.1]) < 0.01)]
shuffleIDs = vcat(shuffleIDs, innerInds)
shuffleIDs = unique(shuffleIDs[end:-1:1])[end:-1:1]

# Creating data to be used by each run of for loop
#Reordering data
xBase = x[:, shuffleIDs]

n_extrapolation = min(NTest, length(innerInds))
interpolation_indices_base = (NTrain + 1) : (NTrain + NTest - n_extrapolation)
extrapolation_indices_base = (NTrain + NTest - n_extrapolation + 1) : N

x_train = xBase[:, rand(1 : NTrain, round(Int, α * NTrain))]
x_scattered = xBase[:, interpolation_indices_base[rand(1 : length(interpolation_indices_base), round(Int, β * length(interpolation_indices_base)))]]
x_region = xBase[:, extrapolation_indices_base[rand(1 : length(extrapolation_indices_base), round(Int, β * length(extrapolation_indices_base)))]]

outplot = plot(size=(400,400), aspect_ratio=:equal, xticks=false, yticks=false, axis=false, legend=:bottomright)
scatter!(outplot, vec(x_train[1,:]), vec(x_train[2,:]), color=silver, label=L"\mathrm{training}", markerstrokewidth=0.0, markeralpha=0.5)
scatter!(outplot, vec(x_scattered[1,:]), vec(x_scattered[2,:]), color=rust, label=L"\mathrm{scattered}", markerstrokewidth=0.0)
scatter!(outplot, vec(x_region[1,:]), vec(x_region[2,:]), color=orange, label=L"\mathrm{region}", markerstrokewidth=0.0)

savefig(outplot, "./out/plots/FFTPoints.pdf")


@testset "Comparison" begin
    @testset "Graded Lex Order" begin
        Mod.@polyvar x y z
        @test x > y > z
        @test x^2*y > y^3 > z
        @test y^2 >= x
        # Examples from p. 58, 59 of the 4th edition of "Ideals, Varieties, and Algorithms" of Cox, Little and O'Shea
        @test x^1*y^2*z^3 > x^3*y^2
        @test !(x^1*y^2*z^3 < x^3*y^2)
        @test x^1*y^2*z^4 > x^1*y^1*z^5
        @test !(x^1*y^2*z^4 < x^1*y^1*z^5)
        @test x^4*y^7*z > x^4*y^2*z^3
        @test !(x^4*y^7*z < x^4*y^2*z^3)
        @test x*y^5*z^2 < x^4*y*z^3
        @test !(x*y^5*z^2 > x^4*y*z^3)
        @test x^5*y*z > x^4*y*z^2
        @test !(x^5*y*z < x^4*y*z^2)
    end
    @testset "Equality" begin
        @testset "Monomial equality" begin
            Mod.@polyvar x y
            @test 1 == constantmonomial(x)
            @test 2 != constantmonomial(x)
            @test 2 != x
            @test 2 != x*y
            @test x*y != x
            @test 1x*y == x*y
            @test 2x*y != x*y
            @test x == monomial(x)
            @test x*y^0 == x
            @test x != x^0*y
            @test x^2 != x*y
        end
        @testset "Polynomial equality" begin
            Mod.@polyvar x y
            @test isapproxzero(0.0x, ztol=0.0)
            @test polynomial(CustomTerms(x + 1 - x)) isa AbstractPolynomial
            @test MultivariatePolynomials.eqconstant(polynomial(CustomTerms(x + 1 - x)), 1)
            @test MultivariatePolynomials.eqconstant(CustomTerms(x + 1 - x), 1)
            @test 2 * CustomPoly(x + 1 - x) == 2
            @test 2 != CustomTerms(x + 1 - x)
            @test 3x^2 == CustomTerms(x - x + x^2) * 3
            @test CustomPoly(-x + x^2) != x^2
            @test 2*x*y + 3*y^2 == 3*y^2 + 2*y*x
            @test 3*x*y + 2*y^2 != 3*y^2 + 2*y*x
            @test x + y != x * (1 + y)
            @test x*y == 3x + 2x*y - x - x*y - 2x
            @test isapproxzero((1+1e-8)x - x, ztol=1e-7)
            @test !isapproxzero((1+1e-6)x - x, ztol=1e-7)
            @test isapprox((2-1e-3)*x*y + (3+1e-3)*y^2, 3*y^2 + 2*y*x, rtol=1e-2)
            @test !isapprox((2-1e-3)*x*y + (3+1e-1)*y^2, 3*y^2 + 2*y*x, rtol=1e-2)
            @test isapprox(1e-3*x*y + 3*y^2 + x^2, x^2 + 3*y^2, rtol=1e-2, ztol=1e-2)
            @test isapprox(1e-3*x*y + 3*y^2 + x^2, x^2 + 3*y^2, rtol=1e-2, atol=1e-2)
            @test isapprox(3*y^2 + x^2, x^2 + 1e-3*x*y + 3*y^2, rtol=1e-2, ztol=1e-2)
            @test isapprox(3*y^2 + x^2, x^2 + 1e-3*x*y + 3*y^2, rtol=1e-2, atol=1e-2)
            @test !isapprox(3*y^2 + x^2, x^2 + 1e-1*x*y + 3*y^2, rtol=1e-2, ztol=1e-2)
            @test !isapprox(3.0*y^2 + x + x^2, x + 3*y^2, rtol=1e-2, ztol=1e-2)
            @test isapprox(x+1-x, 1+1e-8, rtol=1e-7)
            @test !isapprox(1+1e-8, x+1-x, rtol=1e-9)
        end
        @testset "RationalPoly equality" begin
            Mod.@polyvar x y
            @test (x^2 - x - 6) / (x + 2) != x + 3
            @test x - 3 == (x^2 - x - 6) / (x + 2)
            @test (x^2 - x - 6) / (x - 3) == x + 2
            @test 3 != 4x / 2x
            @test 4x / 2x == 2
            @test 3 != 4x / 2x
            @test 1 - 1/x == (x - 1) / x
            @test 1 - 1/x != (1 - x) / x
            @test x + x/x == 1 + x^2/x
            @test x - x/x == -(1 - x^2/x)
            @test (1+x)/x - 1 == 1/x
            @test isapprox((1+1e-8)x, (x*y)/y, rtol=1e-7)
            @test isapproxzero(((1+1e-8)x - x)/y, ztol=1e-7)
            @test !isapproxzero(((1+1e-8)x - y)/y, ztol=1e-9)
            @test isapprox(((1+1e-8)x*y) / y^2, x / y, rtol=1e-7)
            @test isapprox((2x) / x, 2.001, rtol=1e-2)
            @test !isapprox(2.001,  (2x) / x, rtol=1e-4)
        end
    end
    @testset "Equality between a Polynomial and a type not defining zero #22" begin
        Mod.@polyvar x
        # Polynomial of multiple terms
        p = x + x^2
        @test nothing != p
        @test p != nothing
        @test Dict{Int,Int}() != p
        @test p != Dict{Int,Int}()
        # Polynomial of one term
        p = x + x^2 - x
        @test p != nothing
        @test p != Dict{Int,Int}()
        # Polynomial of no term
        # See https://github.com/blegat/MultivariatePolynomials.jl/issues/22
        p = x - x
        @test p != nothing
        @test p != Dict{Int,Int}()
    end
end


"""
    @fieldequal Supertype

Define a method of `==` for all subtypes of `Supertype`
such that `==` returns true if each pair of the field values
from two instances are equal by `==`.
"""
macro fieldequal(Supertype)
    return esc(quote
        function ==(x::T, y::T) where T <: $Supertype
            f = fieldnames(T)
            getfield.(Ref(x),f) == getfield.(Ref(y),f)
        end
    end)
end

"""
    kron_fastl(A::AbstractMatrix, B::AbstractMatrix)

Compute `kron(I(n), A) * B` without explicitely calculating the Kronecker product
where `I(n)` is an identity matrix with
`n` being the ratio between `size(B, 1)` and `size(A, 2)`.
See also [`kron_fastr`](@ref).

# Reference
Acklam, Peter J. 2003. "MATLAB Array Manipulation Tips and Tricks."
"""
function kron_fastl(A::AbstractMatrix, B::AbstractMatrix)
    p, q = size(A)
    qn, m = size(B)
    n = qn ÷ q
    return reshape(A*reshape(B, q, n*m), p*n, m)
end

"""
    kron_fastr(A::AbstractMatrix, B::AbstractMatrix)

Compute `kron(A, I(n)) * B` without explicitely calculating the Kronecker product
where `I(n)` is an identity matrix with
`n` being the ratio between `size(B, 1)` and `size(A, 2)`.
See also [`kron_fastl`](@ref).

# Reference
Acklam, Peter J. 2003. "MATLAB Array Manipulation Tips and Tricks."
"""
function kron_fastr(A::AbstractMatrix, B::AbstractMatrix)
    p, q = size(A)
    qn, m = size(B)
    n = qn ÷ q
    return reshape(reshape(B', n*m, q)*A', m, p*n)'
end

"""
    getscore(X::AbstractMatrix, resid::AbstractVecOrMat)

Compute the regression scores (the product of the residuals and regressors)
with possibly multiple outcome variables.
"""
function getscore(X::AbstractMatrix, resid::AbstractMatrix)
    T, K = size(X)
    N = size(resid, 2)
    out = similar(X, T, K*N)
    @inbounds for (t, j) in Base.product(1:T, 1:K*N)
        ix = (j-1) ÷ N + 1
        ir = j % N
        ir == 0 && (ir = N)
        out[t,j] = X[t,ix] * resid[t,ir]
    end
    return out
end

getscore(X::AbstractMatrix, resid::AbstractVector) = X .* resid

# Check whether the input data is a column table
function checktable(data)
    Tables.istable(data) ||
        throw(ArgumentError("data of type $(typeof(data)) is not `Tables.jl`-compatible"))
    Tables.columnaccess(data) ||
        throw(ArgumentError("data of type $(typeof(data)) is not a column table"))
end

# Indicate rows with finite and nonmissing data
function _esample!(esample::BitVector, aux::BitVector, v::AbstractVector)
    aux .= isequal.(isfinite.(v), true)
    esample .&= aux
end

"""
    hamilton_filter(y::AbstractVector, h::Integer, p::Integer; subset)
    hamilton_filter(y::AbstractVector, freq::Symbol; subset)

Decompose the time series in `y` into a cyclical component and a trend component
using the method by Hamilton (2018).

Rows in `y` that involve invalid values (e.g., `NaN`, `Inf` and `missing`)
for estimation are skipped.
Returned vectors are of the same length as `y`.

# Arguments
- `y::AbstractVector`: vector storing the time series to be filtered.
- `h::Integer`: horizon of forcasting; for business cycles, should cover 2 years.
- `p::Integer`: number of lags used for estimation, including the contemporary one.
- `freq::Symbol`: use default values of `h` and `p` suggested by Hamilton (2018) based on data frequency; must be `:m`, `:q` or `:y` for monthly, quarterly or annual data.

# Keywords
- `subset::Union{BitVector,Nothing}=nothing`: Boolean indices of rows in `y` to be filtered.

# Returns
- `Vector{Float64}`: the cyclical component.
- `Vector{Float64}`: the trend component.
- `BitVector`: indicators for whether an estimate is obtained for each row of input data `y`.

# Reference
Hamilton, James D. 2018. "Why You Should Never Use the Hodrick-Prescott Filter."
The Review of Economics and Statistics 100 (5): 831-843.
"""
function hamilton_filter(y::AbstractVector, h::Integer, p::Integer;
        subset::Union{BitVector,Nothing}=nothing)
    Tfull = length(y)
    T = Tfull - p - h + 1
    h > 0 || throw(ArgumentError("h must be positive; got $h"))
    p > 0 || throw(ArgumentError("p must be positive; got $p"))
    T > p + h + 1 || throw(ArgumentError("not enough observations for estimation"))

    # Indicate whether a missing value exists in each row
    esample = trues(T)
    # A cache for indicators
    aux = BitVector(undef, T)

    # Construct lag matrix
    X = Matrix{Float64}(undef, T, p+1)
    X[:,1] .= 1.0
    for j = 1:p
        v = view(y, p-j+1:Tfull-h-j+1)
        subset === nothing || (esample .&= view(subset, p-j+1:Tfull-h-j+1))
        _esample!(esample, aux, v)
        copyto!(view(X,esample,j+1), view(v,esample))
    end
    v = view(y, p+h:Tfull)
    subset === nothing || (esample .&= view(subset, p+h:Tfull))
    _esample!(esample, aux, v)
    Y = v[esample]
    T1 = sum(esample)
    if T1 < T
        T1 > size(X, 2) || throw(ArgumentError("not enough valid observations left in sample"))
        X = X[esample, :]
    end
    # OLS
    b = X \ Y
    Xb = X * b
    trend = fill(NaN, Tfull)
    trend[view(p+h:Tfull, esample)] .= Xb
    cycle = fill(NaN, Tfull)
    cycle[view(p+h:Tfull, esample)] .= Y .- Xb
    # Returned esample has the same length as input data
    prepend!(esample, falses(p+h-1))
    return cycle, trend, esample
end

function hamilton_filter(y::AbstractVector, freq::Symbol;
        subset::Union{BitVector,Nothing}=nothing)
    if freq == :m
        N = 12
    elseif freq == :q
        N = 4
    elseif freq == :y
        N = 1
    else
        throw(ArgumentError("invalid input for freq; must be :m, :q or :y"))
    end
    return hamilton_filter(y, 2*N, N, subset=subset)
end

"""
    TransformedVar{T}

Supertype for all types used to specify variable transformation.
"""
abstract type TransformedVar{T} end

_geto(v::TransformedVar) = getfield(v, :o)
size(v::TransformedVar, dim) = size(_geto(v), dim)
length(v::TransformedVar) = length(_geto(v))
vec(v::Union{AbstractVector,TransformedVar}, subset::Union{BitVector,Nothing},
    vartype::Symbol, horz::Int, TF::Type) = v

"""
    Cum{T,S} <: TransformedVar{T}

Cumulatively summed variable for estimating cumulative impulse response.
State dependency is allowed by specifying a second variable.
"""
struct Cum{T,S} <: TransformedVar{T}
    o::T
    s::S
end

Cum(n::Symbol) = Cum(n, nothing)
Cum(i::Integer) = Cum(Int(i), nothing)
Cum(v::AbstractVector) = Cum(v, nothing)

getcolumn(data, c::Cum) =
    Cum(getcolumn(data, _geto(c)), c.s===nothing ? nothing : getcolumn(data, c.s))

function vec(c::Cum{<:AbstractVector}, subset::Union{BitVector,Nothing},
        vartype::Symbol, horz::Int, TF::Type)
    v = c.o
    out = zeros(TF, length(v))
    T = length(v) - horz
    na = convert(TF, NaN)
    ts = vartype === :y ? horz : 0
    if vartype === :y || c.s === nothing
        if subset === nothing
            for h in 0:horz
                for t in 1:T
                    out[ts+t] += coalesce(v[t+h], na)
                end
            end
        else
            for h in 0:horz
                for t in 1:T
                    out[ts+t] += ifelse(subset[t+h], coalesce(v[t+h], na), na)
                end
            end
        end
    else
        if subset === nothing
            for h in 0:horz
                for t in 1:T
                    out[t] += coalesce(c.s[t]*v[t+h], na)
                end
            end
        else
            for h in 0:horz
                for t in 1:T
                    out[t] += ifelse(subset[t+h], coalesce(c.s[t]*v[t+h], na), na)
                end
            end
        end
    end
    if vartype == :y
        out[1:horz] .= na
    else
        out[T+1:end] .= na
    end
    return out
end

_toint(data, ::Nothing) = nothing
_toint(data, c::Cum) = Cum(_toint(data, _geto(c)), _toint(data, c.s))
_toname(data, ::Nothing) = nothing
_toname(data, c::Cum) = Cum(_toname(data, _geto(c)), _toname(data, c.s))

show(io::IO, c::Cum{<:Union{Int,Symbol}}) =
    print(io, typeof(c).name.name, "(", _geto(c), c.s===nothing ? ")" : ", $(c.s))")

"""
    datafile(name::Union{Symbol,String})

Return the file path of the example data file named `name`.csv.gz.
"""
datafile(name::Union{Symbol,String}) = (@__DIR__)*"/../data/$(name).csv.gz"


#!/usr/bin/env julia
doc""" Script to test KullbackLeibler.jl

- FIXME write it!
- XXX Or do it automatically with jldoctest (https://juliadocs.github.io/Documenter.jl/latest/man/doctests/#Doctests-1).

- MIT Licensed, (C) 2018 Lilian Besson (Naereen)
  https://GitHub.com/Naereen/KullbackLeibler.jl
"""

__author__ = "Lilian Besson (Naereen)"
__version__ = "0.1"

include("../julia-src/KullbackLeibler.jl")
# include("KullbackLeibler.jl")
using KullbackLeibler
using Base.Test

# utility for the tests

tolerance = 1e-6

function isclose(x, y, tol=tolerance)
    return abs(x - y) <= tol
end

# write your own tests here

@test isclose(KullbackLeibler.klBern(0.5, 0.5), 0.0)
@test isclose(KullbackLeibler.klBern(0.1, 0.9), 1.757779)
# And this KL is symmetric
@test isclose(KullbackLeibler.klBern(0.9, 0.1), 1.757779)
@test isclose(KullbackLeibler.klBern(0.4, 0.5), 0.020135)
@test isclose(KullbackLeibler.klBern(0.01, 0.99), 4.503217)


# ---
# title: 712. Minimum ASCII Delete Sum for Two Strings
# id: problem712
# author: Indigo
# date: 2021-06-29
# difficulty: Medium
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/minimum-ascii-delete-sum-for-two-strings/description/>
# hidden: true
# ---
#
# Given two strings `s1, s2`, find the lowest ASCII sum of deleted characters to
# make two strings equal.
#
# **Example 1:**
#
#
#
#     Input: s1 = "sea", s2 = "eat"
#     Output: 231
#     Explanation: Deleting "s" from "sea" adds the ASCII value of "s" (115) to the sum.
#     Deleting "t" from "eat" adds 116 to the sum.
#     At the end, both strings are equal, and 115 + 116 = 231 is the minimum sum possible to achieve this.
#
#
# **Example 2:**
#
#
#
#     Input: s1 = "delete", s2 = "leet"
#     Output: 403
#     Explanation: Deleting "dee" from "delete" to turn the string into "let",
#     adds 100[d]+101[e]+101[e] to the sum.  Deleting "e" from "leet" adds 101[e] to the sum.
#     At the end, both strings are equal to "let", and the answer is 100+101+101+101 = 403.
#     If instead we turned both strings into "lee" or "eet", we would get answers of 433 or 417, which are higher.
#
#
# **Note:**
#
# * `0 < s1.length, s2.length <= 1000`.
# * All elements of each string will have an ASCII value in `[97, 122]`.
#
#
## @lc code=start
using LeetCode

function minimum_delete_sum(s1::String, s2::String)
    m, n = length(s1) + 1, length(s2) + 1
    dp = fill(0, m, n)
    dp[2:end, 1] .= cumsum(codeunits(s1))
    dp[1, 2:end] .= cumsum(codeunits(s2))
    for i in 2:m, j in 2:n
        dp[i, j] = (s1[i - 1] == s2[j - 1]) ? (dp[i - 1, j - 1]) :
                   min(dp[i - 1, j] + Int(s1[i - 1]), dp[i, j - 1] + Int(s2[j - 1]))
    end
    return dp[m, n]
end
## @lc code=end


using Turing: VarEstimator, add_sample!, get_var, reset!
using Test

let
    D = 1000
    ve = VarEstimator{Float64}(0, zeros(D), zeros(D))
    ve = reset!(add_sample!(ve, randn(D)))

    # Check that reseting zeros everything.
    @test ve.n === 0
    @test ve.μ == zeros(D)
    @test ve.M == zeros(D)

    # Ensure that asking for the variance doesn't mutate the VarEstimator.
    add_sample!(ve, randn(D))
    add_sample!(ve, randn(D))
    μ, M = deepcopy(ve.μ), deepcopy(ve.M)
    get_var(ve)
    @test ve.μ == μ
    @test ve.M == M
end

# Check that the estimated variance is approximately correct.
let
    D = 10
    ve = VarEstimator{Float64}(0, zeros(D), zeros(D))

    for _ = 1:10000
        s = randn(D)
        add_sample!(ve, s)
    end

    var = get_var(ve)

    @test var ≈ ones(D) atol=0.5
end


using ControlBenchmarks
using ControlSystems
using SparseArrays

# Test the constructor with 1 mass and all values equal to 1
msd = LinearMassSpringDamperChain( 1, 1, 1, 1 )
ben = controlbenchmark( msd )

@test ControlSystems.nstates( ben.sys ) == 2
@test ControlSystems.ninputs( ben.sys ) == 1
@test ControlSystems.noutputs( ben.sys ) == 2

@test issparse( ben.sys.A )
@test issparse( ben.sys.B )
@test issparse( ben.sys.C )
@test issparse( ben.sys.D )

# Test the non-damper constructor with 1 mass and all values equal to 1
ms  = LinearMassSpringChain( 1, 1, 1 )
ben = controlbenchmark( ms )

@test ControlSystems.nstates( ben.sys ) == 2
@test ControlSystems.ninputs( ben.sys ) == 1
@test ControlSystems.noutputs( ben.sys ) == 2

@test issparse( ben.sys.A )
@test issparse( ben.sys.B )
@test issparse( ben.sys.C )
@test issparse( ben.sys.D )

######################################################
# Test the validity checks in the constructor
######################################################

# Give more values than are needed
@test_throws ErrorException LinearMassSpringDamperChain( 1, [1 1 1], 1, 1 )
@test_throws ErrorException LinearMassSpringDamperChain( 1, 1, [1 1 1], 1 )
@test_throws ErrorException LinearMassSpringDamperChain( 1, 1, 1, [1 1] )

# More inputs than masses
@test_throws ErrorException LinearMassSpringDamperChain( 2, 1, 1, 1, inputs = [1 2 3] )

# Input requested on a mass that doesn't exist
@test_throws DomainError LinearMassSpringDamperChain( 2, 1, 1, 1, inputs = [1 3] )

# Give more values than are needed
@test_throws ErrorException LinearMassSpringChain( 1, [1 1 1], 1 )
@test_throws ErrorException LinearMassSpringChain( 1, 1, [1 1] )

# More inputs than masses
@test_throws ErrorException LinearMassSpringChain( 2, 1, 1, inputs = [1 2 3] )

# Input requested on a mass that doesn't exist
@test_throws DomainError LinearMassSpringChain( 2, 1, 1, inputs = [1 3] )

######################################################
# Construct a 5 chain and see if it matches with a computed example from the literature
######################################################

m = 1000
k = 1000
β = 200

ben = controlbenchmark( LinearMassSpringDamperChain( 5, k, β, m ) )

A = ben.sys.A
B = ben.sys.B

# Test values for mass 1
@test A[6, 1] == ( -2 * k / m)
@test A[6, 2] == ( k / m)
@test A[6, 6] == ( -2 * β / m)
@test A[6, 7] == ( β / m)

# Test values for mass 2, 3, 4
for i = 2:4
    @test A[i+5, i-1] == ( k / m)
    @test A[i+5,   i] == ( -2 * k / m)
    @test A[i+5, i+1] == ( k / m)

    @test A[i+5, 5+i-1] == ( β / m)
    @test A[i+5,   5+i] == ( -2 * β / m)
    @test A[i+5, 5+i+1] == ( β / m)
end

# Test values for mass 5
@test A[10,  4] == ( k / m)
@test A[10,  5] == ( -2 * k / m)
@test A[10,  9] == ( β / m)
@test A[10, 10] == ( -2 * β / m)

# There should be only 2 inputs
@test nnz( B ) == 2
@test B[6,  1] == ( -1 / m )
@test B[10, 2] == ( -1 / m )

######################################################
# Do some slight tweaks to the chain
######################################################

m = 1000
k = 1000
β = 200

# No damping and 3 inputs
ben = controlbenchmark( LinearMassSpringDamperChain( 5, k, 0, m; inputs = [1, 3, 5] ) )

A = ben.sys.A
B = ben.sys.B

# No damping was requested

# Test values for mass 1
@test A[6, 6] == 0
@test A[6, 7] == 0

# Test values for mass 2, 3, 4
for i = 2:4
    @test A[i+5, 5+i-1] == 0
    @test A[i+5,   5+i] == 0
    @test A[i+5, 5+i+1] == 0
end

# Test values for mass 5
@test A[10,  9] == 0
@test A[10, 10] == 0

# We should have 3 inputs now
@test nnz( B ) == 3

######################################################
# Test multiple inputs
######################################################
controlbenchmark( LinearMassSpringDamperChain( 5, k, 0, m; inputs = [1, 3, 5] ) )
controlbenchmark( LinearMassSpringDamperChain( 5, k, 0, m; inputs = 1:5 ) )


using Plots, Printf

"""
	shallow_water_2D(; nx, ny, dam_x=true, do_visu=true)

2D shallow water equations solver for an instantaneous dam break.
The Lax-Friedrichs Method was applied to the continuity equation.
Geometry (length of 40 meters, width of 20 meters) and initial conditions
(half of domain has initial water level of 5 meters, other half is dry)
matches BASEMENT version 2.8 Test Case H_1 "Dam break in a closed channel."
Solution of momentum equations requires division by H. Therefore a minimum
H must be set throughout the domain, to avoid numerical instability.

# Arguments
    - `nx`: number of discretised cells for x dimension.
    - `ny`: number of discretised cells for y dimension.
    - `dam_x`: if true, 1D dam break in x-direction, else 1D dam break in y-direction.
    - `do_visu`: if true, each physical time step will be ploted.

# Return values
    - `H`: The solution arrray (Water surface height in m).
"""
@views function shallow_water_2D(;
    # Numerics
    nx = 512,
    ny = 256,
    # Modeling
    dam_x = true,
    # Visualisation
    do_visu = true)
    # Physics
    Lx, Ly = 40.0, 20.0
    g = 9.81
    u_max = 5 #from review of results with very small timesteps
    H_init = 5.0
    time = 0
    # Numerics
    nout = 100
    H_min = 1.0
    # Derived numerics
    dx, dy = Lx / nx, Ly / ny
    nt = 10000#cld(ttot, dt)
    xc, yc = LinRange(dx / 2, Lx - dx / 2, nx), LinRange(dy / 2, Ly - dy / 2, ny)
    # Initial Conditions
    H = zeros(Float64, nx, ny) .+ H_min
    if (dam_x)
        H[1:round(Int64(nx / 2)), :] = zeros(round(Int64(nx / 2)), ny) .+ H_init #1D dam break in x-direction
    else
        H[:, 1:round(Int64(ny / 2))] = zeros(Float64, nx, round(Int64(ny / 2))) .+ H_init #1D dam break in y direction
    end
    #H[1:round(Int64(nx/2)), 1:round(Int64(ny/2))] = zeros(Float64, round(Int64(nx/2)), round(Int64(ny/2))) .+ H_init #2D dam break
    #H[round(Int64(3*nx/8)):round(Int64(5*nx/8)), round(Int64(3*ny/8)):round(Int64(5*ny/8))] = zeros(Float64, round(Int64(nx/4))+1, round(Int64(ny/4))+1) .+ H_init #2D water column
    # Array initialisation
    H_new = copy(H)
    u = zeros(Float64, nx, ny)
    u_new = copy(u)
    v = zeros(Float64, nx, ny)
    v_new = copy(v)
    F = zeros(Float64, nx + 1, ny)
    G = zeros(Float64, nx, ny + 1)
    t_tic = 0.0
    niter = 0
    # Time loop
    for it = 0:nt
        if (it == 11)
            t_tic = Base.time()
            niter = 0
        end
        dt = dx / (max(maximum(abs.(u)), maximum(abs.(v)), u_max)) / 4.1
        time = time + dt
        λx = dx / 2 / dt #x-direction signal speed for Lax-Friedrichs scheme
        λy = dy / 2 / dt #y-direction signal speed for Lax-Friedrichs scheme
        #Continuity
        F[2:end-1, :] .= 1 / 2 .* (H[1:end-1, :] .* u[1:end-1, :] .+ H[2:end, :] .* u[2:end, :]) .- λx / 2 .* (H[2:end, :] .- H[1:end-1, :])
        G[:, 2:end-1] .= 1 / 2 .* (H[:, 1:end-1] .* v[:, 1:end-1] .+ H[:, 2:end] .* v[:, 2:end]) .- λy / 2 .* (H[:, 2:end] .- H[:, 1:end-1])
        H_new[2:end-1, :] .= H[2:end-1, :] .- dt / dx .* (F[3:end-1, :] .- F[2:end-2, :]) .- dt / dy .* (G[2:end-1, 2:end] .- G[2:end-1, 1:end-1])
        #Momentum in x-direction
        F[2:end-1, :] .= 1 / 2 .* ((H[1:end-1, :] .* ((u[1:end-1, :]) .^ 2) .+ 0.5 * g .* ((H[1:end-1, :]) .^ 2)) .+ (H[2:end, :] .* ((u[2:end, :]) .^ 2) .+ 0.5 * g .* ((H[2:end, :]) .^ 2)))
        G[:, 2:end-1] .= 1 / 2 .* ((H[:, 1:end-1] .* u[:, 1:end-1] .* v[:, 1:end-1]) .+ (H[:, 2:end] .* u[:, 2:end] .* v[:, 2:end]))
        u_new[2:end-1, :] .= (1 / 2 .* (H[1:end-2, :] .* u[1:end-2, :] .+ H[3:end, :] .* u[3:end, :]) .- dt / dx .* (F[3:end-1, :] .- F[2:end-2, :]) .- dt / dy .* (G[2:end-1, 2:end] .- G[2:end-1, 1:end-1])) ./ H_new[2:end-1, :]
        #Momentum in y-direction
        F[2:end-1, :] .= 1 / 2 .* ((H[1:end-1, :] .* u[1:end-1, :] .* v[1:end-1, :]) .+ (H[2:end, :] .* u[2:end, :] .* v[2:end, :]))
        G[:, 2:end-1] .= 1 / 2 .* ((H[:, 1:end-1] .* ((v[:, 1:end-1]) .^ 2) .+ 0.5 * g .* ((H[:, 1:end-1]) .^ 2)) .+ (H[:, 2:end] .* ((v[:, 2:end]) .^ 2) .+ 0.5 * g .* ((H[:, 2:end]) .^ 2)))
        v_new[:, 2:end-1] .= (1 / 2 .* (H[:, 1:end-2] .* v[:, 1:end-2] .+ H[:, 3:end] .* v[:, 3:end]) .- dt / dx .* (F[2:end, 2:end-1] .- F[1:end-1, 2:end-1]) .- dt / dy .* (G[:, 3:end-1] .- G[:, 2:end-2])) ./ H_new[:, 2:end-1]
        #No-flux boundary conditions
        H_new[1, :] .= H_new[2, :]
        H_new[end, :] .= H_new[end-1, :]
        u_new[1, :] .= 0
        u_new[end, :] .= 0
        v_new[:, 1] .= 0
        v_new[:, end] .= 0

        H .= H_new
        u .= u_new
        v .= v_new

        niter += 1
        if do_visu && (it % nout == 0)
            p1 = plot(xc, H[:, round(Int64(ny / ny))], xlims = (xc[1], xc[end]), ylims = (-3, 8),
                xlabel = "Lx at y=10 (m)", ylabel = "water surface elevation (m)", label = "h",
                title = "time = $(round(time, sigdigits=3)) s", linewidth = :1.0, framestyle = :box)
            plot!(xc, u[:, round(Int64(ny / ny))], label = "u", linewidth = :1.0)
            plot!(xc, v[:, round(Int64(ny / ny))], label = "v", linewidth = :1.0)
            display(p1)
        end
    end
    t_toc = Base.time() - t_tic
    A_eff = (3 * 2) / 1e9 * nx * sizeof(Float64)  # Effective main memory access per iteration [GB]
    t_it = t_toc / niter                      # Execution time per iteration [s]
    T_eff = A_eff / t_it                       # Effective memory throughput [GB/s]
    @printf("Time = %1.3f sec, T_eff = %1.4f GB/s (niter = %d)\n", t_toc, round(T_eff, sigdigits = 3), niter)
    return H
end

# shallow_water_2D(; nx = 256, ny = 128, do_visu = true)


using .Distributions # Requires.jl replace usual syntax with this

# Define recipe for 2d MvNormal

@recipe function f(dist::MvNormal, alpha=0.0001)
    @assert length(dist) == 2 "MvNormal Plot supports only 2d distribution"

    seriestype --> :contours
    x_mdist = Normal(dist.μ[1], sqrt(dist.Σ[1,1]))
    y_mdist = Normal(dist.μ[2], sqrt(dist.Σ[2,2]))
    xmin, xmax = quantile.(x_mdist, [alpha, 1-alpha])
    ymin, ymax = quantile.(y_mdist, [alpha, 1-alpha])
    (range(xmin, xmax, length=100), range(ymin, ymax, length=100), (x,y)->pdf(dist, [x,y]))
end

function find_limit(mdist_list, prior_p, alpha)
    q_l = [quantile(mdist, alpha) for mdist in mdist_list]
    l = minimum(q_l)
    r = maximum(q_l)

    for i in 1:100
        s = (l+r)/2
        sum_cdf = sum([cdf(mdist, s) * p for (mdist, p) in zip(mdist_list, prior_p)])
        if sum_cdf < alpha
            l = s
        else
            r = s
        end
    end

    left = (l + r) / 2
end

function find_limit2(mdist_list, prior_p, alpha)
    find_limit(mdist_list, prior_p, alpha), find_limit(mdist_list, prior_p, 1-alpha)
end

@recipe function f(mix_dist::MixtureModel{Multivariate, Continuous, <:MvNormal}, alpha=0.0001)
    @assert length(mix_dist.components[1]) == 2 "MixtureModel Plot supports only 2d component distribution"

    seriestype --> :contours

    x_mdist_list = [Normal(comp.μ[1], sqrt(comp.Σ[1,1])) for comp in mix_dist.components]
    y_mdist_list = [Normal(comp.μ[2], sqrt(comp.Σ[2,2])) for comp in mix_dist.components]

    prior_p = mix_dist.prior.p

    xmin, xmax = find_limit2(x_mdist_list, prior_p, alpha)
    ymin, ymax = find_limit2(y_mdist_list, prior_p, alpha)

    (range(xmin, xmax, length=100), range(ymin, ymax, length=100), (x,y)->pdf(mix_dist, [x,y]))
end


@doc doc"""
    LinearAffineMetric <: Metric

The linear affine metric is the metric for symmetric positive definite matrices, that employs
matrix logarithms and exponentials, which yields a linear and affine metric.
"""
struct LinearAffineMetric <: RiemannianMetric end

is_default_metric(::SymmetricPositiveDefinite, ::LinearAffineMetric) = Val(true)

@doc doc"""
    distance(M::SymmetricPositiveDefinite, x, y)
    distance(M::MetricManifold{SymmetricPositiveDefinite,LinearAffineMetric})

Compute the distance on the [`SymmetricPositiveDefinite`](@ref) manifold between `x` and `y`,
as a [`MetricManifold`](@ref) with [`LinearAffineMetric`](@ref). The formula reads

```math
d_{\mathcal P(n)}(x,y)
= \lVert \operatorname{Log}(x^{-\frac{1}{2}}yx^{-\frac{1}{2}})\rVert_{\mathrm{F}}.,
```
where $\operatorname{Log}$ denotes the matrix logarithm and
$\lVert\cdot\rVert_{\mathrm{F}}$ denotes the matrix Frobenius norm.
"""
function distance(M::SymmetricPositiveDefinite{N}, x, y) where {N}
    s = real.(eigvals(x, y))
    return any(s .<= eps()) ? 0 : sqrt(sum(abs.(log.(s)) .^ 2))
end

@doc doc"""
    exp(M::SymmetricPositiveDefinite, x, v)
    exp(M::MetricManifold{SymmetricPositiveDefinite{N},LinearAffineMetric}, x, v)

Compute the exponential map from `x` with tangent vector `v` on the
[`SymmetricPositiveDefinite`](@ref) `M` with its default [`MetricManifold`](@ref) having the
[`LinearAffineMetric`](@ref). The formula reads

```math
\exp_x v = x^{\frac{1}{2}}\operatorname{Exp}(x^{-\frac{1}{2}} v x^{-\frac{1}{2}})x^{\frac{1}{2}},
```

where $\operatorname{Exp}$ denotes to the matrix exponential.
"""
exp(::SymmetricPositiveDefinite, ::Any...)

function exp!(M::SymmetricPositiveDefinite{N}, y, x, v) where {N}
    e = eigen(Symmetric(x))
    U = e.vectors
    S = e.values
    Ssqrt = Diagonal(sqrt.(S))
    SsqrtInv = Diagonal(1 ./ sqrt.(S))
    xSqrt = Symmetric(U * Ssqrt * transpose(U))
    xSqrtInv = Symmetric(U * SsqrtInv * transpose(U))
    T = Symmetric(xSqrtInv * v * xSqrtInv)
    eig1 = eigen(T) # numerical stabilization
    Se = Diagonal(exp.(eig1.values))
    Ue = eig1.vectors
    xue = xSqrt * Ue
    return copyto!(y, xue * Se * transpose(xue))
end

@doc doc"""
    [Ξ,κ] = get_basis(M::SymmetricPositiveDefinite, x, B::DiagonalizingOrthonormalBasis)
    [Ξ,κ] = get_basis(M::MetricManifold{SymmetricPositiveDefinite{N},LinearAffineMetric}, x, B::DiagonalizingOrthonormalBasis)

Return a orthonormal basis `Ξ` as a vector of tangent vectors (of length
[`manifold_dimension`](@ref) of `M`) in the tangent space of `x` on the
[`MetricManifold`](@ref) of [`SymmetricPositiveDefinite`](@ref) manifold `M` with
[`LinearAffineMetric`](@ref) that diagonalizes the curvature tensor $R(u,v)w$
with eigenvalues `κ` and where the direction `B.v` has curvature `0`.
"""
function get_basis(
    M::SymmetricPositiveDefinite{N},
    x,
    B::DiagonalizingOrthonormalBasis,
) where {N}
    xSqrt = sqrt(x)
    eigv = eigen(B.v)
    V = eigv.vectors
    Ξ = [
        (i == j ? 1 / 2 :
             1 / sqrt(2)) * (V[:, i] * transpose(V[:, j]) + V[:, j] * transpose(V[:, i]))
        for i = 1:N
        for j = i:N
    ]
    λ = eigv.values
    κ = [-1 / 4 * (λ[i] - λ[j])^2 for i = 1:N for j = i:N]
    return PrecomputedDiagonalizingOrthonormalBasis(Ξ, κ)
end
function get_basis(
    M::MetricManifold{SymmetricPositiveDefinite{N},LinearAffineMetric},
    x,
    B::DiagonalizingOrthonormalBasis,
) where {N}
    return get_basis(base_manifold(M), x, B)
end

function get_coordinates(
    M::SymmetricPositiveDefinite{N},
    x,
    v,
    B::ArbitraryOrthonormalBasis,
) where {N}
    dim = manifold_dimension(M)
    vout = similar(v, dim)
    @assert size(v) == (N, N)
    @assert dim == div(N * (N + 1), 2)
    k = 1
    for i = 1:N, j = i:N
        scale = ifelse(i == j, 1, sqrt(2))
        @inbounds vout[k] = v[i, j] * scale
        k += 1
    end
    return vout
end
function get_coordinates(
    M::MetricManifold{SymmetricPositiveDefinite{N},LinearAffineMetric},
    x,
    v,
    B::ArbitraryOrthonormalBasis,
) where {N}
    return get_coordinates(base_manifold(M), x, v, B)
end

function get_vector(
    M::SymmetricPositiveDefinite{N},
    x,
    v,
    B::ArbitraryOrthonormalBasis,
) where {N}
    dim = manifold_dimension(M)
    vout = allocate_result(M, get_vector, x)
    @assert size(v) == (div(N * (N + 1), 2),)
    @assert size(vout) == (N, N)
    k = 1
    for i = 1:N, j = i:N
        scale = ifelse(i == j, 1, 1 / sqrt(2))
        @inbounds vout[i, j] = v[k] * scale
        @inbounds vout[j, i] = v[k] * scale
        k += 1
    end
    return vout
end
function get_vector(
    M::MetricManifold{SymmetricPositiveDefinite{N},LinearAffineMetric},
    x,
    v,
    B::ArbitraryOrthonormalBasis,
) where {N}
    return get_vector(base_manifold(M), x, v, B)
end

@doc doc"""
    inner(M::SymmetricPositiveDefinite, x, v, w)
    inner(M::MetricManifold{SymmetricPositiveDefinite,LinearAffineMetric}, x, v, w)

Compute the inner product of `v`, `w` in the tangent space of `x` on
the [`SymmetricPositiveDefinite`](@ref) manifold `M`, as
a [`MetricManifold`](@ref) with [`LinearAffineMetric`](@ref). The formula reads

````math
(v, w)_x = \operatorname{tr}(x^{-1} v x^{-1} w),
````
"""
function inner(M::SymmetricPositiveDefinite, x, v, w)
    F = cholesky(Symmetric(x))
    return tr((F \ Symmetric(v)) * (F \ Symmetric(w)))
end

@doc doc"""
    log(M::SymmetricPositiveDefinite, x, y)
    log(M::MetricManifold{SymmetricPositiveDefinite,LinearAffineMetric}, x, y)

Compute the logarithmic map from `x` to `y` on the [`SymmetricPositiveDefinite`](@ref)
as a [`MetricManifold`](@ref) with [`LinearAffineMetric`](@ref). The formula reads

```math
\log_x y =
x^{\frac{1}{2}}\operatorname{Log}(x^{-\frac{1}{2}}yx^{-\frac{1}{2}})x^{\frac{1}{2}},
```
where $\operatorname{Log}$ denotes to the matrix logarithm.
"""
log(::SymmetricPositiveDefinite, ::Any...)

function log!(M::SymmetricPositiveDefinite{N}, v, x, y) where {N}
    e = eigen(Symmetric(x))
    U = e.vectors
    S = e.values
    Ssqrt = Diagonal(sqrt.(S))
    SsqrtInv = Diagonal(1 ./ sqrt.(S))
    xSqrt = Symmetric(U * Ssqrt * transpose(U))
    xSqrtInv = Symmetric(U * SsqrtInv * transpose(U))
    T = Symmetric(xSqrtInv * y * xSqrtInv)
    e2 = eigen(T)
    Se = Diagonal(log.(max.(e2.values, eps())))
    xue = xSqrt * e2.vectors
    return mul!(v, xue, Se * transpose(xue))
end

@doc doc"""
    vector_transport_to(M::SymmetricPositiveDefinite, x, v, y, ::ParallelTransport)
    vector_transport_to(M::MetricManifold{SymmetricPositiveDefinite,LinearAffineMetric}, x, v, y, ::ParallelTransport)

Compute the parallel transport on the [`SymmetricPositiveDefinite`](@ref) as a
[`MetricManifold`](@ref) with the [`LinearAffineMetric`](@ref).
The formula reads

```math
P_{y\gets x}(v) = x^{\frac{1}{2}}
\operatorname{Exp}\bigl(
\frac{1}{2}x^{-\frac{1}{2}}\log_x(y)x^{-\frac{1}{2}}
\bigr)
x^{-\frac{1}{2}}v x^{-\frac{1}{2}}
\operatorname{Exp}\bigl(
\frac{1}{2}x^{-\frac{1}{2}}\log_x(y)x^{-\frac{1}{2}}
\bigr)
x^{\frac{1}{2}},
```

where $\operatorname{Exp}$ denotes the matrix exponential
and `log` the logarithmic map on [`SymmetricPositiveDefinite`](@ref)
(again with respect to the metric mentioned).
"""
vector_transport_to(::SymmetricPositiveDefinite, ::Any, ::Any, ::Any, ::ParallelTransport)

function vector_transport_to!(
    M::SymmetricPositiveDefinite{N},
    vto,
    x,
    v,
    y,
    ::ParallelTransport,
) where {N}
    distance(M, x, y) < 2 * eps(eltype(x)) && copyto!(vto, v)
    e = eigen(Symmetric(x))
    U = e.vectors
    S = e.values
    Ssqrt = sqrt.(S)
    SsqrtInv = Diagonal(1 ./ Ssqrt)
    Ssqrt = Diagonal(Ssqrt)
    xSqrt = Symmetric(U * Ssqrt * transpose(U))
    xSqrtInv = Symmetric(U * SsqrtInv * transpose(U))
    tv = Symmetric(xSqrtInv * v * xSqrtInv)
    ty = Symmetric(xSqrtInv * y * xSqrtInv)
    e2 = eigen(ty)
    Se = Diagonal(log.(e2.values))
    Ue = e2.vectors
    ty2 = Symmetric(Ue * Se * transpose(Ue))
    e3 = eigen(ty2)
    Sf = Diagonal(exp.(e3.values))
    Uf = e3.vectors
    xue = xSqrt * Uf * Sf * transpose(Uf)
    vtp = Symmetric(xue * tv * transpose(xue))
    return copyto!(vto, vtp)
end


# Range: 1 : 100,000,000
# Time:  36.494s

# Written by Alex Vear in 2019
# Public domain.  No rights reserved.

startval = 1
endval = 100000000

sqrs = [x^2 for x in startval:endval]
tris = [x*(x+1)/2 for x in startval:endval]
println(intersect(sqrs, tris))


# scal!
scal!(a::Array{T}, p_size) where T = LinearAlgebra.BLAS.scal!(length(a), T(12/p_size), a, 1)
scal!(a::CuArray{T}, p_size) where T = CUBLAS.scal!(length(a), T(1/p_size), a)

# Matvec
dispgemm!(tA::Char, α::Float32, A::AbstractArray{Float32, 2}, b::AbstractArray{Float32, 1}, β::Float32, c::Array{Float32, 1}) = LinearAlgebra.BLAS.gemv!(tA, α, A, b, β, c)
dispgemm!(tA::Char, α::Float32, A::AbstractArray{Float32, 2}, b::AbstractArray{Float32, 1}, β::Float32, c::CuArray{Float32, 1}) = CUBLAS.gemv!(tA, α, A, b, β, c)

# MatMat
dispgemm!(tA::Char, tB::Char, α::Float32, A::AbstractArray{Float32, 2}, b::AbstractArray{Float32, 2}, β::Float32, c::Array{Float32, 2}) = LinearAlgebra.BLAS.gemm!(tA, tB, α, A, b, β, c)
dispgemm!(tA::Char, tB::Char, α::Float32, A::AbstractArray{Float32, 2}, b::AbstractArray{Float32, 2}, β::Float32, c::CuArray{Float32, 2}) = CUBLAS.gemm!(tA, tB, α, A, b, β, c)

# Batched mat-mat
Bgemm!(tA, tB, α, A::AbstractArray, B::AbstractArray, β, C::Array) = NNlib.batched_gemm!(tA, tB, α, A, B, β, C)
Bgemm!(tA, tB, α, A::AbstractArray, B::AbstractArray, β, C::CuArray) = CUBLAS.gemm_strided_batched!(tA, tB, α, A, B, β, C)


module ModuleConstruct
using MultipleViewGeometry
using MultipleViewGeometry.ModuleTypes, MultipleViewGeometry.ModuleMathAliases, MultipleViewGeometry.ModuleOperators
using MultipleViewGeometry.ModuleProjection
using StaticArrays, LinearAlgebra
export construct

include("construct_projectionmatrix.jl")
include("construct_fundamentalmatrix.jl")
include("construct_homographymatrix.jl")
include("construct_essentialmatrix.jl")
end


# Helper funtions
get_R(P, G, Δ) = sqrt.(Δ) * G * P * G' * sqrt.(Δ)
get_Q(F, R) = F' * R * F + 1.0
get_K(F, R, Q) = R * F / Q
get_m(m, G, K, e) = G * m + K * e'
get_P(R, K, Q) = R - K * K' * Q
get_S(S, k, e, Q) = S/k + e*e'/Q

function diagnostics(ssm::StateSpaceModel, output)
    μ = output.predicted.μ
    Σ = output.predicted.Σ
    e = output.predicted.e
    u = output.predicted.u
    ME = mean(e; dims = 1)
    MAE = mean(abs.(e); dims = 1)
    MSSE = mean(u.^2; dims = 1)
    LL = mean([logpdf(MvTDist(ssm.ν, μ[t, :], Σ[t, :, :]), y[t, :]) for t in 1:size(ssm.y, 1)])
    return (ME = ME, MAE = MAE, MSSE = MSSE, LogLik = LL)
end