module NeuralNetworks

# TODO: Add exports

using LinearAlgebra, Zygote, ForwardDiff

##
## Data types
##

"""
A single layer of a neural network.
"""
struct Layer
    # Each row corresponds to a neuron in the previous layer, including bias
    # term.
    # Each column corresponds to a neuron in this layer, not including the bias
    # term.
    # Biases for a given layer are implicitly added by the next layer, and so do
    # not show up in dimensionality of the same layer. Thus, the dimensionality
    # of this matrix is (num_prev_outputs + 1) * (num_curr_outputs).
    weights::AbstractMatrix{<:Number}
    # Vector-valued function (f: R^d -> R^d) mapping the neurons in this layer
    # to their activations.
    activation_fn::Function
end

"""
A complete neural network.
"""
struct NeuralNetwork
    # Layers of the network, applied in order.
    layers::AbstractVector{<:Layer}
    # Single loss function f: (R^n * R^n) -> R applied to the network's final
    # layer output.
    loss_fn::Function
end

"""
A data structure containing the values computed by a single layer during the
forward-pass.
"""
struct LayerForwardPass
    # Inputs from the previous layer. Includes bias.
    inputs::AbstractVector{<:Number}
    # Sums of each neuron before activation; same dimensionality as layer
    # output.
    sums::AbstractVector{<:Number}
    # Activation of each neuron; same dimensionality as layer output.
    activations::AbstractVector{<:Number}
end

"""
A data structure containing the partial derivatives of the loss with respect to
a given layer during the backward-pass.
"""
struct LayerBackwardPass
    # Partial derivatives w.r.t. pre-activation sums.
    dL_dS::AbstractVector{<:Number}
    # Partial derivatives w.r.t. each weight.
    dL_dW::AbstractMatrix{<:Number}
    # Partial derivatives w.r.t. inputs. This gets used by the previous layer in
    # backprop (these inputs are exactly the outputs of the previous layer).
    dL_dI::AbstractVector{<:Number}
end

function Base.show(io::IO, layer::Layer)
    print("  $(size(layer.weights, 2)) neurons\n")
    print("  Weights " *
          "[$(size(layer.weights, 1))x$(size(layer.weights, 2))]: " *
          "$(layer.weights)\n")
end

function Base.show(io::IO, nn::NeuralNetwork)
    num_layers = length(nn.layers)
    print("Neural network with $(num_layers) layers.\n")
    print("Input dimensionality: $(in_ndims_sans_bias(nn.layers[1]))\n")
    print("Output dimensionality: $(out_ndims(nn.layers[num_layers]))\n")
    for l in 1:num_layers
        print("Layer $(l):\n")
        print(nn.layers[l])
    end
end

function Base.show(io::IO, layer_forward_pass::LayerForwardPass)
    print("Input vector: $(layer_forward_pass.inputs), " *
          "$(size(layer_forward_pass.inputs))\n")
    print("Intermediate sums: $(layer_forward_pass.sums), " *
          "$(size(layer_forward_pass.sums))\n")
    print("Output vector: $(layer_forward_pass.activations), " *
          "$(size(layer_forward_pass.activations))\n")
end

function Base.show(io::IO, layer_backward_pass::LayerBackwardPass)
    print("dL_dS: $(layer_backward_pass.dL_dS), " *
          "$(size(layer_backward_pass.dL_dS))\n")
    print("dL_dW: $(layer_backward_pass.dL_dW), " *
          "$(size(layer_backward_pass.dL_dW))\n")
    print("dL_dI: $(layer_backward_pass.dL_dI), " *
          "$(size(layer_backward_pass.dL_dI))\n")
end

"""
Data structure containing loss stats for each epoch of training.
"""
struct TrainingStats
    losses::AbstractVector
end

##
## Constructors
##

"""
    create_layer(in_ndims_sans_bias, out_ndims, activation_fn)

Creates a single layer of a neural network with the given dimensions and
activation function, and randomly assigned weights.

In this package, a "layer" is composed of (1) weights, which are applied to the
previous layer, and (2) activation functions, which are applied after weights.
"""
function create_layer(in_ndims_sans_bias::Integer,
                      out_ndims::Integer,
                      activation_fn::Function)::Layer
    weights = randn((in_ndims_sans_bias+1, out_ndims)) ./ sqrt(in_ndims_sans_bias)
    return Layer(weights, activation_fn)
end

"""
    validate_layer(layer, num_inputs_sans_bias, num_outputs)

Asserts that the layer parameters have the desired dimensionality, not including
the bias.
"""
function validate_layer(layer::Layer,
                        num_inputs_sans_bias::Integer,
                        num_outputs::Integer)
    @assert in_ndims_sans_bias(layer) == num_inputs_sans_bias
    @assert out_ndims(layer) == num_outputs
end

"""
    compose_layers(input_ndims, output_ndims, loss_fn, layers...)

Creates a neural network from the supplied layers.
"""
function compose_layers(input_ndims::Integer,
                        output_ndims::Integer,
                        loss_fn::Function,
                        layers...)::NeuralNetwork
    num_layers = length(layers)
    nn_layers = Vector{Layer}(undef, num_layers)

    # Verify layers are dimensionally compatible.
    for l in 1:num_layers
        expected_in_ndims_sans_bias = (l == 1 ?
                                       input_ndims :
                                       out_ndims(layers[l-1]))
        expected_out_ndims = (l == num_layers ?
                              output_ndims :
                              in_ndims_sans_bias(layers[l+1]))
        validate_layer(layers[l],
                       expected_in_ndims_sans_bias,
                       expected_out_ndims)
    end

    return NeuralNetwork([layers...], loss_fn)
end

"""
    create_nn(layer_dims, activation_fn, loss_fn,
              network_input_dims, network_output_dims)

Creates a neural network by building layers with the specified dimensionality,
`activation_fn`, and `loss_fn`.

- `layer_dims::AbstractVector`: a vector of integers, where the `n`-th value
corresponds to the number of neurons in the `n`-th layer of the network.
- `activation_fn`: the activation function used for all layers in the network.
- `loss_fn`: a single loss function applied at the end of a network, producing a
  scalar loss value.
- `network_input_dims`: the dimensionality of inputs to the network.
- `network_output_dims`: the dimensionality of outputs of the network.

TODO: This function is convenient, but it's probably better to construct and
compose layers individually.
"""
function create_nn(layer_dims::AbstractVector{<:Integer},
                   activation_fn::Function,
                   loss_fn::Function,
                   network_input_ndims::Integer,
                   network_output_ndims::Integer)::NeuralNetwork
    num_layers = length(layer_dims)
    layers = Vector{Layer}(undef, num_layers)
    for l in 1:num_layers
        in_ndims_sans_bias = (l == 1 ?
                              network_input_ndims :
                              layer_dims[l-1])
        out_ndims = layer_dims[l]
        layers[l] = create_layer(in_ndims_sans_bias, out_ndims,
                                 activation_fn)
    end

    return compose_layers(network_input_ndims, network_output_ndims,
                          loss_fn, layers...)
end

##
## Training and inference
##

"""
    run_forward_pass(layer, inputs)

Computes the activations of `layer` on `inputs`, returning the cached results.
"""
function run_forward_pass(layer::Layer,
                          inputs::AbstractVector{<:Number})::LayerForwardPass
    @assert ndims(inputs) == 1
    @assert length(inputs) == in_ndims_sans_bias(layer)

    # Append the bias term.
    inputs = copy(inputs)
    push!(inputs, 1)

    # Perform matrix multiplication `weights`^T * `input`.
    # For the `weights` matrix, the rows index the inputs and the columns index
    # the outputs. `weights` is transposed before multiplication so that the
    # rows index outputs, and the result is a column vector with an entry for
    # each neuron output in this layer.
    sums = transpose(layer.weights) * inputs

    # Apply each neuron's activation function. The activation function is
    # vector-valued, so the result is a vector of the same dimensionality.
    activations = layer.activation_fn(sums)

    return LayerForwardPass(inputs, sums, activations)
end

"""
    compute_loss_gradient(net_output, label, loss_fn)

Computes the gradient of `loss_fn` parameterized by `label`, with respect to
`net_output`, evaluated at `net_output`.

- `loss_fn`: a function f: (R^n * R^n) -> R

Returns a vector with the size of the output layer.
"""
function compute_loss_gradient(net_output::AbstractVector{<:Number},
                               label::AbstractVector{<:Number},
                               loss_fn::Function)::AbstractVector{<:Number}
    @assert length(net_output) == length(label)

    partial_loss(x) = loss_fn(x, label)
    dL_dO = gradient(partial_loss, net_output)[1]

    return dL_dO
end

"""
    run_backward_pass(layer, forward_pass, dL_dO)

Does backprop to compute the loss gradients with respect to the components of a
single layer, using the values from the `forward_pass`.

Returns a LayerBackwardPass containing the gradients, without mutating the
original layer.
"""
function run_backward_pass(layer::Layer,
                           forward_pass::LayerForwardPass,
                           dL_dO::AbstractVector{<:Number})::LayerBackwardPass
    @assert length(dL_dO) == out_ndims(forward_pass)

    # dL_dS|S = dL_dO|O * dO_dS|S
    # dO_dS is just the derivative (Jacobian) of the activation function.
    # (d_out * d_out)
    dO_dS = ForwardDiff.jacobian(layer.activation_fn, forward_pass.sums)

    # dL_dS = dL_dO * dO_dS
    # For activation functions which are elementwise-only, this is actually
    # equivalent to a vector-vector Hadamard product, because the Jacobian is
    # diagonal.
    # (d_out)
    dL_dS = transpose(dO_dS) * dL_dO

    # dL_dW|W = dL_dS|S * dS_dW|W
    # dS_dW for each weight is the value of the output which is scaled by that
    # weight.
    # (d_in+1 * d_out)
    dL_dW = forward_pass.inputs * transpose(dL_dS)

    # dL_dI|I = sum(dL_dS|S * dS_dI|I)
    # dS_dI is just the applied weight; the matrix multiplication applies the
    # vectorized sum.
    # (d_in+1)
    dL_dI = layer.weights * dL_dS

    # Note that the last element in dL_dW corresponds to bias, and so should be
    # omitted by the previous layer during backprop.
    return LayerBackwardPass(dL_dS, dL_dW, dL_dI)
end

"""
    update_layer_weights!(layer, backward_pass, num_samples, learning_rate)

Updates weights according to backprop gradient calculation by applying values
from a previously computed `backward_pass`.
"""
function update_layer_weights!(layer::Layer,
                               backward_pass::LayerBackwardPass,
                               num_samples::Number,
                               learning_rate::Number)
    @assert size(layer.weights) == size(backward_pass.dL_dW)

    layer.weights .-= (learning_rate .* backward_pass.dL_dW ./ num_samples)
end

"""
    train!(nn, samples, labels, learning_rate, batch_size)
"""
function train!(nn::NeuralNetwork,
                samples::AbstractVector,
                labels::AbstractVector,
                learning_rate::Number,
                batch_size::Integer)::TrainingStats
    @assert length(samples) == length(labels)

    num_layers = length(nn.layers)

    # Divide the data into batch_size chunks.
    sample_batches = Iterators.partition(samples, batch_size) |> collect
    label_batches = Iterators.partition(labels, batch_size) |> collect

    num_passes = length(sample_batches)

    losses = Vector(undef, length(sample_batches))

    for (batch_index, input_batch) in enumerate(sample_batches)
        label_batch = label_batches[batch_index]
        num_samples = length(input_batch)

        # Record the results of forward and backward passes for each layer.
        # These are only used for backprop and are overwritten on every batch.
        # Rows index samples; columns index layers.
        forward_passes = Matrix{LayerForwardPass}(undef, (num_samples, num_layers))
        backward_passes = Matrix{LayerBackwardPass}(undef, (num_samples, num_layers))

        batch_loss = 0.0

        # Compute updates for the batch.
        for (point_index, input_point) in enumerate(input_batch)
            point_label = label_batch[point_index]

            # Run forward-pass.
            for l in 1:num_layers
                forward_passes[point_index, l] = run_forward_pass(nn.layers[l], input_point)
                input_point = forward_passes[point_index, l].activations
            end

            # Compute loss.
            batch_loss += nn.loss_fn(input_point, point_label)

            # Backprop.
            loss_grad = compute_loss_gradient(input_point, point_label, nn.loss_fn)

            for l in reverse(1:num_layers)
                backward_passes[point_index, l] =
                    run_backward_pass(nn.layers[l],
                                      forward_passes[point_index, l],
                                      l == num_layers ?
                                      loss_grad :
                                      backward_passes[point_index, l+1].dL_dI[1:end-1])
            end
        end

        # Update weights. The weight updates have already been computed, so this
        # can be done in any order.
        for i in 1:num_samples, l in 1:num_layers
            update_layer_weights!(nn.layers[l],
                                  backward_passes[i, l],
                                  num_samples,
                                  learning_rate)
        end

        batch_loss /= num_samples

        losses[batch_index] = batch_loss
        if (batch_index % 1000 == 0)
            print("Iteration $batch_index / $(length(sample_batches)) : " *
                  "loss=$(batch_loss)\n")
        end
    end

    return TrainingStats(losses)
end

"""
    predict(nn, inputs)

Computes the end-to-end application of `nn` on `inputs`.
"""
function predict(nn::NeuralNetwork, inputs::AbstractVector{<:Number})
    for layer in nn.layers
        inputs = run_forward_pass(layer, inputs).activations
    end
    return inputs
end

##
## Vectorized implementations
##

# It's not ideal that the implementation is forked, but it's rather challenging
# to have a single implementation that handles both vectorized and
# non-vectorized updates. In particular, the vectorized code is significantly
# more complex, so is not optimized for clarity.

"""
Vectorized forward-pass data structure.
"""
struct LayerForwardPassVectorized
    # (n * d_in+1)
    inputs::AbstractMatrix{<:Number}
    # (n * d_out)
    sums::AbstractMatrix{<:Number}
    # (n * d_out)
    activations::AbstractMatrix{<:Number}
end

"""
Vectorized backward-pass data structure.
"""
struct LayerBackwardPassVectorized
    # Allocate space for the Jacobian of the activation function. The partial
    # derivatives can't be hard-coded since activation functions may vary, and
    # each sample requires a (d * d) matrix, since the activation function can
    # map a vector to a vector.
    # (n * d_out * d_out)
    dO_dS::AbstractArray{<:Number, 3}
    # (n * d_out)
    dL_dS::AbstractMatrix{<:Number}
    # Since only the total gradient update is needed, the n-dimension is collapsed.
    # Otherwise, this would be a 3D (n * d_in+1 + d_out) array.
    # (d_in+1 * d_out)
    dL_dW::AbstractMatrix{<:Number}
    # (n * d_in+1)
    dL_dI::AbstractMatrix{<:Number}
end

"""
    run_forward_pass_vectorized!(layer, inputs, num_samples, batch_size,
                                 forward_pass_out)

Vectorized forward-pass.

Content of extra rows when `num_samples` < `batch_size` is undefined.

- `inputs`: (n * d_in) matrix of samples.
- `num_samples`: n_s, number of input samples.
- `batch_size`: n, batch size, corresponding to the number of rows preallocated
  for various matrices; may be larger than `num_samples` for incomplete batches.
"""
function run_forward_pass_vectorized!(layer::Layer,
                                      inputs::AbstractMatrix{<:Number},
                                      num_samples::Integer,
                                      batch_size::Integer,
                                      forward_pass_out::LayerForwardPassVectorized)
    @assert ndims(inputs) == 2
    @assert size(inputs, 2) == in_ndims_sans_bias(layer)

    # Append the bias column (extra dimension for each sample).
    forward_pass_out.inputs .= [inputs ones(batch_size,1)]

    # sums = biased_inputs * layer.weights
    # (n * d_in+1) * (d_in+1 * d_out)
    mul!(forward_pass_out.sums, forward_pass_out.inputs, layer.weights)

    # 1/ One-liner alternative for multiple activation functions:
    #
    # activations = mapslices(row -> map.(layer.activation_fns, row), sums, dims=2)
    #
    # `mapslices(f, m, dims=2)` broadcasts f to each row of m.
    # `map.(f, v)` performs elementwise function application for equal sized
    # vectors of functions and elements.
    # The net result is that the dth function is applied to every element in the
    # dth column.
    #
    # 2/ Straightforward, but optimized, version:
    #
    # Threads.@threads for j in 1:out_ndims(layer)
    #     map!(layer.activation_fns[j],
    #          view(forward_pass_out.activations, 1:num_samples, j),
    #          view(forward_pass_out.sums, 1:num_samples, j))
    # end

    # Skip activation computation for the last (batch_size - num_samples) rows.
    # Poor cache access pattern.
    Threads.@threads for i in 1:num_samples
        view(forward_pass_out.activations, i, :) .=
            layer.activation_fn(view(forward_pass_out.sums, i, :))
    end
end

"""
    compute_loss_gradient_vectorized!(net_outputs, labels, loss_fn,
                                      num_samples, batch_size,
                                      loss_grad_out)

Vectorized gradient computation.

Content of extra rows when n_s < n is undefined.

- `net_outputs`: (n * d) matrix of output predictions (rows are samples).
- `labels`: (n_s * d) matrix of labels (rows are samples); note that there may
  be a size mismatch between n and n_s if the batch is incomplete.
- `loss_fn`: a function that computes a scalar loss from two matrices.
- `num_samples`: n_s, number of input samples.
- `batch_size`: n, batch size, corresponding to the number of rows preallocated
  for various matrices; may be larger than `num_samples` for incomplete batches.

Returns (n * d) matrix of gradients.

This function may return a different (scaled) result compared to
`compute_loss_gradient`. When the values disagree, the non-vectorized version
should be treated as the reference implementation.
"""
function compute_loss_gradient_vectorized!(net_outputs::AbstractMatrix{<:Number},
                                           labels::AbstractMatrix{<:Number},
                                           loss_fn::Function,
                                           num_samples::Integer,
                                           batch_size::Integer,
                                           loss_grad_out::AbstractMatrix{<:Number})
    @assert size(net_outputs, 2) == size(labels, 2)

    partial_loss(xs) = loss_fn(xs, labels)

    # Compute dL_dO.
    if num_samples == batch_size
        loss_grad_out .= gradient(partial_loss, net_outputs)[1]
    else
        view(loss_grad_out, 1:num_samples, :) .=
            gradient(partial_loss, view(net_outputs, 1:num_samples, :))[1]
    end
end

"""
    run_backward_pass_vectorized!(layer, forward_pass, dL_dO, num_samples, batch_size,
                                  backward_pass_out)

Vectorized backward-pass.

Content of extra rows when `num_samples` < `batch_size` is undefined.

- `dL_dO`: (n * d_out) matrix of gradients.
- `num_samples`: n_s, number of input samples.
- `batch_size`: n, batch size, corresponding to the number of rows preallocated
  for various matrices; may be larger than `num_samples` for incomplete batches.
"""
function run_backward_pass_vectorized!(layer::Layer,
                                       forward_pass::LayerForwardPassVectorized,
                                       dL_dO::AbstractMatrix{<:Number},
                                       num_samples::Integer,
                                       batch_size::Integer,
                                       backward_pass_out::LayerBackwardPassVectorized)
    @assert size(dL_dO, 2) == out_ndims(forward_pass)

    # 1/ One-liner alternative for multiple activation functions:
    #
    # Slice by row, and for each, do element-wise activation function
    # application.
    #
    # dO_dS = mapslices(row -> map(gradient, layer.activation_fns, row),
    #                   forward_pass.sums,
    #                   dims=2) |>
    #     vec -> map(t -> t[1], vec)
    #
    # 2/ Optimized version:
    #
    # Threads.@threads for j in 1:length(layer.activation_fns)
    #     for i in 1:num_samples
    #         backward_pass_out.dL_dS[i,j] = gradient(layer.activation_fns[j],
    #                                                 forward_pass.sums[i,j])[1]
    #     end
    # end

    # Skip the last (batch_size - num_samples) rows.
    # Poor cache access pattern.
    # (n * d_out * d_out)
    Threads.@threads for i in 1:num_samples
        ForwardDiff.jacobian!(view(backward_pass_out.dO_dS, i, :, :),
                              layer.activation_fn,
                              view(forward_pass.sums, i, :))
    end

    # dL_dS = dL_dO * dO_dS
    # (n * d_out)
    # Poor cache access pattern.
    Threads.@threads for i in 1:num_samples
        mul!(view(backward_pass_out.dL_dS, i, :),
             transpose(view(backward_pass_out.dO_dS, i, :, :)),
             view(dL_dO, i, :))
    end

    # 1/ One-liner for 3D solution (keep each sample, (n * d_in+1 * d_out))
    # Slice by row (each row is a (d_in+1 * d_out) matrix!), and for each
    # compute the outer product which computes the weight gradients for the
    # given sample.
    #
    # dL_dW = permutedims(
    #     cat([forward_pass.inputs[i,:] * transpose(dL_dS[i,:]) for i in 1:num_samples]...,
    #         dims=3),
    #     (3,1,2))

    # 2/ Vectorized 3D, (n * d_in * d_out)
    #
    # dL_dW = Array{Float64, 3}(undef, (num_samples, # n
    #                                   size(forward_pass.inputs, 2), # d_in+1
    #                                   size(dL_dS, 2))) # d_out
    # for i in 1:num_samples
    #     # Poor cache locality
    #     dL_dW[i,:,:] .= forward_pass.inputs[i,:] * transpose(dL_dS[i,:])
    # end

    # inputs are (n * d_in+1), dL_dS is (n * d_out).
    # (d_in * d_out)
    fill!(backward_pass_out.dL_dW, 0.0)
    # TODO: Parallelize this with thread-local memory, and combine at end
    # Skip the last (batch_size - num_samples) rows.
    for i in 1:num_samples
        # dL_dW .+= forward_pass.inputs[i,:] * transpose(dL_dS[i,:])
        BLAS.axpy!(1,
                   view(forward_pass.inputs,i,:) *
                   transpose(view(backward_pass_out.dL_dS,i,:)),
                   backward_pass_out.dL_dW)
    end

    # Weights are (d_in+1, d_out)
    # (n * d_in+1)
    # dL_dI = dL_dS * transpose(layer.weights)
    # Uses BLAS to avoid allocation for the transpose.
    BLAS.gemm!('N',
               'T',
               1.0,
               backward_pass_out.dL_dS,
               layer.weights,
               0.0,
               backward_pass_out.dL_dI)
end

"""
    update_layer_weights_vectorized!(layer, backward_pass, learning_rate)

Vectorized layer weight update
"""
function update_layer_weights_vectorized!(layer::Layer,
                                          backward_pass::LayerBackwardPassVectorized,
                                          learning_rate::Number)
    @assert size(layer.weights) == size(backward_pass.dL_dW)

    # One-liner for n-dimension collapse:
    # layer.weights .-= (learning_rate .*
    #                    dropdims(sum(backward_pass.dL_dW, dims=1), dims=1))

    BLAS.axpy!(-learning_rate, backward_pass.dL_dW, layer.weights)
end

"""
    train_vectorized!(nn, samples, labels, learning_rate, batch_size)

Vectorized gradient descent.

TODO: The conditional logic involved in properly handling partial batches causes
a noticeable performance regression.
"""
function train_vectorized!(nn::NeuralNetwork,
                           samples::AbstractVector,
                           labels::AbstractVector,
                           learning_rate::Number,
                           batch_size::Integer)::TrainingStats
    @assert length(samples) == length(labels)

    num_layers = length(nn.layers)

    # Divide the data into batch_size chunks.
    sample_batches = Iterators.partition(samples, batch_size) |>
        batches -> map(batch -> transpose(hcat(batch...)), batches) |>
        collect
    label_batches = Iterators.partition(labels, batch_size) |>
        batches -> map(batch -> transpose(hcat(batch...)), batches) |>
        collect

    num_passes = length(sample_batches)

    # Preallocate memory for each batch iteration.
    forward_passes = Vector{LayerForwardPassVectorized}(undef, num_layers)
    for l in 1:num_layers
        inputs = zeros(batch_size, in_ndims_with_bias(nn.layers[l]))
        sums = zeros(batch_size, out_ndims(nn.layers[l]))
        activations = zeros(batch_size, out_ndims(nn.layers[l]))
        forward_passes[l] = LayerForwardPassVectorized(inputs, sums, activations)
    end

    backward_passes = Vector{LayerBackwardPassVectorized}(undef, num_layers)
    for l in 1:num_layers
        dO_dS = zeros(batch_size, out_ndims(nn.layers[l]), out_ndims(nn.layers[l]))
        dL_dS = zeros(batch_size, out_ndims(nn.layers[l]))
        dL_dW = zeros(in_ndims_with_bias(nn.layers[l]),
                      out_ndims(nn.layers[l]))
        dL_dI = zeros(batch_size, in_ndims_with_bias(nn.layers[l]))
        backward_passes[l] = LayerBackwardPassVectorized(dO_dS, dL_dS, dL_dW, dL_dI)
    end

    loss_grad = zeros(batch_size, out_ndims(nn.layers[end]))

    losses = Vector(undef, num_passes)

    for (index, input_batch) in enumerate(sample_batches)
        num_samples = size(input_batch, 1)

        # Individual subroutines expect matrices with size batch_size, but
        # otherwise correctly handle mismatches between num_samples and
        # batch_size.
        if (num_samples != batch_size)
            print("Incomplete batch; $num_samples / $batch_size\n")
            input_batch = vcat(input_batch,
                               zeros(batch_size - num_samples,
                                     size(input_batch, 2)))
        end

        # Run forward-pass.
        for l in 1:num_layers
            run_forward_pass_vectorized!(nn.layers[l],
                                         l == 1 ?
                                             input_batch :
                                             forward_passes[l-1].activations,
                                         num_samples,
                                         batch_size,
                                         forward_passes[l])
        end

        # Compute loss.
        # Avoid taking a view unless there isn't a full batch.
        loss = (num_samples == batch_size) ?
            nn.loss_fn(forward_passes[num_layers].activations,
                       label_batches[index]) :
            nn.loss_fn(
                view(forward_passes[num_layers].activations,1:num_samples),
                label_batches[index])

        # Backprop.
        compute_loss_gradient_vectorized!(forward_passes[num_layers].activations,
                                          label_batches[index],
                                          nn.loss_fn,
                                          num_samples,
                                          batch_size,
                                          loss_grad)

        for l in reverse(1:num_layers)
            run_backward_pass_vectorized!(nn.layers[l],
                                          forward_passes[l],
                                          l == num_layers ?
                                          loss_grad :
                                          # Chop off the bias column.
                                          view(backward_passes[l+1].dL_dI,
                                               :, 1:out_ndims(nn.layers[l])),
                                          num_samples,
                                          batch_size,
                                          backward_passes[l])
        end

        # Update weights. The weight updates have already been computed, so this
        # can be done in any order.
        # This could be multithreaded, but the network would need to be very
        # deep for the parallelism to outweigh the thread overhead, so it's not
        # recommended.
        for l in 1:num_layers
            update_layer_weights_vectorized!(nn.layers[l],
                                             backward_passes[l],
                                             learning_rate)
        end

        losses[index] = loss
        if (index % 1000 == 0)
            print("Iteration $index / $(num_passes) : loss=$(loss)\n")
        end
    end

    return TrainingStats(losses)
end

##
## Miscellaneous
##

"""
    in_ndims_sans_bias(layer)

Returns the number of inputs to `layer`, not including the bias term (in other
words, only the inputs that come from the previous layer).
"""
function in_ndims_sans_bias(layer::Layer)
    return size(layer.weights, 1) - 1
end

"""
    in_ndims_sans_bias(layer)

Returns the number of inputs to `layer`, not including the bias term (in other
words, only the inputs that come from the previous layer).
"""
function in_ndims_sans_bias(layer::LayerBackwardPass)
    @assert length(dL_dI) == size(layer.dL_dW, 1) - 1
    return length(dL_dI)
end

"""
    in_ndims_with_bias(layer)

Returns the number of inputs to `layer`, including the bias term.
"""
function in_ndims_with_bias(layer::Layer)
    return size(layer.weights, 1)
end

"""
    out_ndims(layer)

Returns the number of outputs from `layer`.
"""
function out_ndims(layer::Layer)
    return size(layer.weights, 2)
end

"""
    out_ndims(layer)

Returns the number of outputs from `layer`.
"""
function out_ndims(layer::LayerForwardPass)
    @assert length(layer.sums) == length(layer.activations)
    return length(layer.activations)
end

"""
    out_ndims(layer)

Returns the number of outputs from `layer`.
"""
function out_ndims(layer::LayerForwardPassVectorized)
    @assert size(layer.sums) == size(layer.activations)

    return size(layer.activations, 2)
end

##
## Buyer beware
##

"""
DEPRECATED

    compute_loss(net_output, label, loss_fn)

Computes the loss between `label` and `net_output` using `loss_fn`. Despite the
loss being a scalar, this transforms it into a singleton array for consistency
with the multi-objective function.

- `loss_fn`: a function f: (R^n * R^n) -> R
"""
function compute_loss(net_output::AbstractVector{<:Number},
                      label::AbstractVector{<:Number},
                      loss_fn::Function)::AbstractVector{<:Number}
    return [loss_fn(label, net_output)]
end

"""
DEPRECATED

    compute_multi_objective_loss_gradient(net_output, label, loss_fn)

Computes the gradient of `loss_fn` parameterized by `label`, with respect to
`net_output`, evaluated at `net_output`.

You probably don't want to do this (this package does not provide a good
solution to multi-objective optimization problems).

- `loss_fn`: a function f: (R^n * R^n) -> R^n
"""
function compute_multi_objective_loss_gradient(net_output::AbstractVector{<:Number},
                                               label::AbstractVector{<:Number},
                                               loss_fn::Function)::AbstractVector{<:Number}
    @assert length(net_output) == length(label)

    # Vector of partially-applied loss functions, where each entry lossmap[i]
    # contains f(x) = loss_fn(label[i], x)
    lossmap = map(y -> (x -> loss_fn(y, x)), label)

    # Take the gradient of each loss function with respect to each output,
    # evaluated at the outputs.
    # dL_dO is just the derivative of the loss function.
    dL_dO = map(gradient, lossmap, net_output) |>
        vec -> map(t -> t[1], vec)

    return dL_dO
end

"""
DEPRECATED

    compute_multi_objective_loss(net_output, label, loss_fn)

Computes the loss between `label` and `net_output` using `loss_fn`.

You probably don't want to do this (this package does not provide a good
solution to multi-objective optimization problems).

- `loss_fn`: a function f: (R^n * R^n) -> R^n
"""
function compute_multi_objective_loss(net_output::AbstractVector{<:Number},
                                      label::AbstractVector{<:Number},
                                      loss_fn::Function)::AbstractVector{<:Number}
    # Vector of partially-applied loss functions, where each entry lossmap[i]
    # contains f(x) = loss_fn(label[i], x)
    lossmap = map(y -> (x -> loss_fn(y, x)), label)

    # Returns a vector loss.
    return map.(lossmap, net_output)
end

end # module


### A Pluto.jl notebook ###
# v0.12.20

using Markdown
using InteractiveUtils

# ╔═╡ a6648952-5a6d-11eb-2c81-99d1514983c0
using Econometrics,
RDatasets,
GLM,
GLMNet,
Statistics,
Plots,
DataFrames,
StatsPlots,
StatsBase

# ╔═╡ 8f1873f2-63b0-11eb-289f-4d880cee7e5a
md"Импортируем нужные библиотеки"

# ╔═╡ b842ca4a-63b1-11eb-2190-0b6c834f85e0
md"### Первый взгляд на датасет"

# ╔═╡ 2f38687c-63b3-11eb-384e-c1847038f4c5
md"Для построения регрессионых моделей воспользуемся датасетом **Auto** из `RDatasets`"

# ╔═╡ 093d42ea-5a72-11eb-19b1-6b0abb122099
auto = dataset("ISLR", "Auto")

# ╔═╡ ad79eb3c-63b5-11eb-3d99-1f2b7aab9a18
md"Выведем некоторую статистическую информацию по колонкам"

# ╔═╡ 22840156-5a72-11eb-0e57-ff3178e5ee6f
describe(auto)

# ╔═╡ c24ee52a-63d4-11eb-1815-0d184c506af2
md"**MPG** - это так называемые литры на 100 км. пути"

# ╔═╡ 61db537c-63d9-11eb-013c-9f74989da4d7
md"Сперва давайте посмотрим на графиках на зависимость между **MPG** и четыремя первыми переменными (после **Cylinders**)
"

# ╔═╡ fa02d9c0-63d9-11eb-1867-2f9780b312fe
scatter(auto.Displacement, auto.MPG, xaxis="Displacement", yaxis="MPG", legend = false)

# ╔═╡ 757e36ae-63d9-11eb-277a-9b2ac8f3b483
scatter(auto.Horsepower, auto.MPG, xaxis="Horsepower", yaxis="MPG", legend = false)

# ╔═╡ d46d4030-63d9-11eb-32d0-1780207fae50
scatter(auto.Weight, auto.MPG, xaxis="Weight", yaxis="MPG", legend = false)

# ╔═╡ e2e70718-63d9-11eb-02e5-0dff07a9c6c7
scatter(auto.Acceleration, auto.MPG, xaxis="Acceleration", yaxis="MPG", legend = false)

# ╔═╡ b26f1aac-63da-11eb-2e66-27cd5597e043
md"
Важно также посмотреть на *outliers* - так называемые выбросы. Подробности и тонкости работы с данными представлены в другом туториале, тут же мы посмотрим только на один важный момент :)

Можно построить **box plot** и после удалить выбросы, которые видны на графике
"

# ╔═╡ d987355c-63da-11eb-1170-910084cc4274
boxplot!(auto.Horsepower, title = "Box Plot - Horsepower", ylabel = "Horsepower", legend = false)

# ╔═╡ f0845d38-63db-11eb-0be2-eb5217bf6170
md"### Модели в `GLM` и `Econometrics`"

# ╔═╡ 17e60466-63c4-11eb-1f97-072f5539cb38
md"Для начала воспользуемся пакетом `GLM`. Начнем с простой парной регрессии: $$\widehat{MPG} = \widehat{\beta_0} + \widehat{\beta_1} \cdot Horsepower$$"

# ╔═╡ 8b294048-63bd-11eb-3df5-959b55069357
linear_model = lm(@formula(MPG ~ Horsepower), auto)

# ╔═╡ e86909f8-63c4-11eb-3f85-9dd2450e8c5f
md"Выведем коэффициент детерминации $$R^2$$"

# ╔═╡ 1fc60af8-63c1-11eb-0814-054e7b9be6f3
GLM.r2(linear_model)

# ╔═╡ f42c722a-63c4-11eb-262c-514bf90aa295
md"Horsepower объясняет примерно 60,6% дисперсии разброса MPG. Вполне неплохо, построим график с данными, регрессией и доверительным интервалом для предсказания с уровнем доверия 95%"

# ╔═╡ 4161b2ca-63c1-11eb-1f47-3f66d74863c3
begin
	pred = DataFrame(Horsepower = 45:0.1:230);
	pr = GLM.predict(linear_model, pred, interval = :prediction, level = 0.95);
	plot(xlabel="Horsepower", ylabel="MPG", legend=:best)
	plot!(auto.Horsepower, auto.MPG, label="data", seriestype=:scatter)
	plot!(pred.Horsepower, pr.prediction, label="model", linewidth=3,
        ribbon = (pr.prediction .- pr.lower, pr.upper .- pr.prediction))
end

# ╔═╡ 35163060-63c7-11eb-2013-d1878e8e2c0b
md"По графику видим, что связь возможно гиперболическая. Рассмотрим модель $$\widehat{MPG} = \widehat{\beta_0} + \widehat{\beta_1} \cdot \dfrac{1}{Horsepower}$$"

# ╔═╡ 65664ad0-63c6-11eb-3e2f-7bf9841d5265
hyperbolic_model = lm(@formula(MPG ~ 1 / Horsepower), auto)

# ╔═╡ 7582b77a-63c6-11eb-3ac2-615c4e294aab
GLM.r2(hyperbolic_model)

# ╔═╡ 5ba889e6-63c7-11eb-14a1-216dee432abe
md"Эта модель, похоже, и правда лучше объясняет **MPG**. Построим аналогичный график"

# ╔═╡ 8bb5ec74-63c6-11eb-1c3e-d358b4ebc5c5
begin
	predd = DataFrame(Horsepower = 45:0.1:232);
	prd = GLM.predict(hyperbolic_model, predd, interval = :prediction, level = 0.95);
	plot(xlabel="Horsepower", ylabel="MPG", legend=:best)
	plot!(auto.Horsepower, auto.MPG, label="data", seriestype=:scatter)
	plot!(predd.Horsepower, prd.prediction, label="model", linewidth=3,
        ribbon = (prd.prediction .- prd.lower, prd.upper .- prd.prediction))
end

# ╔═╡ 01019b9e-63dc-11eb-3148-9b08890b53c9
md"Теперь перейдем к множественной регрессии.

Мы видим, что есть линейная связь между **MPG** и **Horsepower** и **MPG** и **Weight**"

# ╔═╡ 0b27def8-63dc-11eb-003c-454b21b38132
model_GLM = lm(@formula(MPG ~ Weight + Horsepower), auto)

# ╔═╡ 38766bac-63dc-11eb-2dbc-978b5fb3a637
md"
Видим достаточно знакомую и родную табличку! Но фанаты **R** и **Stata** могут справедливо заметить, что как-то маловато всего представлено в таблице. Например, не видно коэффициента детерминации. Не унывать! **Julia** - за минимализм. Да и вывести всё, что Вас интересует, невероятно просто. Достаточно лишь указать интересующий метод и применить его к нашей модели.
"

# ╔═╡ 550be2bc-63dc-11eb-331b-d71a3071ec66
R2 = GLM.r2(model_GLM)

# ╔═╡ 6cd53b1e-63dc-11eb-074a-41c43b6dfaaf
coefs = GLM.coef(model_GLM)

# ╔═╡ 7df28c9e-63dc-11eb-3c8e-5732e36a7e5c
md"
Что ещё можно посмотреть, применив метод к модели?

*dof_residual*: степени свободы остатков (если имеет смысл)

*stderror*: стандартные ошибки коэффициентов

*vcov*: ковариационная матрица оценок коэффициентов

И много всего другого! Интересующиеся могут обратиться к документации `GLM`."

# ╔═╡ ce4f3232-63dc-11eb-374d-412a97f8d2c8
md"
А теперь давайте предсказать значение, задав заранее значения **Horsepower** и **Weight**.
"

# ╔═╡ de284b94-63dc-11eb-2d26-1391517ed17d
new_data = DataFrame(Horsepower = [200, 250], Weight = [4000, 3800])

# ╔═╡ 09221186-63dd-11eb-3c1a-a7fa2ef46e3e
prediction_GLM = DataFrame(MPG_pred = GLM.predict(model_GLM, new_data))

# ╔═╡ 3a4f1000-63e3-11eb-3f33-d1590d15b611
md"
Можем также посмотреть на имеющиеся наблюдения и сравнить с предсказанными значениями
"

# ╔═╡ 88873fa4-63e3-11eb-2d3a-e9d30a385d93
md"Посмотрим на распределение ошибок"

# ╔═╡ 76daf6ba-63e3-11eb-3e19-ff67bdb86a0e
begin
	performance = DataFrame(y_actual = auto.MPG, y_predicted =      GLM.predict(model_GLM));
	error = performance[:y_actual] - performance[:y_predicted];
	histogram(error, bins = 50, title = "Error Analysis", ylabel = "Frequency", 		xlabel = "Error",legend = false)
end

# ╔═╡ 92918540-63e3-11eb-1b57-3d49dfa75b09
md"Ошибки распределены нормально, судя по гистограмме"

# ╔═╡ c080265a-63e3-11eb-1aad-8f85c3d18f3e
md"
Также рассмотрим пакет `Econometrics`, который позволяет строить линейные регрессии, но немного отличается от `GLM`. Конечно, основные различия можно увидеть при более глубоком анализе, нежели тот, что мы приводим в нашем туториале, поэтому заинтересовавшиеся могут обратиться к документации пакета.
"

# ╔═╡ e66b840e-63e3-11eb-12f7-8d4155f58990
model_Econometrics = fit(EconometricModel, @formula(MPG ~ Weight + Horsepower), auto)

# ╔═╡ 6c1f9266-63e4-11eb-08bf-cbabc486dd30
md"
Фанаты **R** и **Stata** могут ликовать! Теперь мы видим табличку, в которой есть и коэффициент детерминации, и LR-тест, и много всего интересного для благородных Донов!
"

# ╔═╡ 0255de96-63e4-11eb-3336-fb9a753ae588
md"Сравним с аналогичной моделью из `GLM`"

# ╔═╡ 13df2904-63e4-11eb-0600-d1c1ebfe9d84
model_GLM

# ╔═╡ 2d6038be-63e4-11eb-2320-d7251a567785
md"Как мы видим, все статистические характеристики коэффициентов идентичны, только в `Econometrics` выводится еще разная полезная информация. В целом, пакет `Econometrics` умеет делать то же самое, что и пакет `GLM`"

# ╔═╡ da6bb162-665c-11eb-20f1-fb018a64646b
md"Может случиться такое, что мы построили модель, а коэффициент **Intercept** оказался незначимым, и хочется от него избавиться. Тогда в формуле надо просто записать ноль как слагаемое"

# ╔═╡ 055e78f2-665d-11eb-188e-8f88fb707c74
model_without_constant = fit(EconometricModel, @formula(MPG ~ 0 + Weight + Horsepower), auto)

# ╔═╡ 2afd1c5a-663b-11eb-26bf-5fbcc7b8b14a
md"### Работа с категориальными признаками"

# ╔═╡ 3e865ba6-663b-11eb-31e0-b57ea1e08bac
md"Если мы присмотримся к нашему датасету, то увидим, что у нас есть такой признак, как **Cylinders**, который мы можем рассмотреть, как категориальный. Убедимся в этом с помощью графика"

# ╔═╡ 6d6f598a-63da-11eb-0de2-fb70c755a0ef
bar!(auto.Cylinders, auto.MPG, xaxis="Cylinders", yaxis="Count", legend = false)

# ╔═╡ 0d648614-663c-11eb-0dc7-4de628f3f80d
md"Для того, чтобы рассматривать эту переменную как отдельные категории, создадим новый столбец с данными уже типа *categorical*"

# ╔═╡ deed8534-6635-11eb-3a64-37f4db496443
auto.Cat_cyl = categorical(auto.Cylinders)

# ╔═╡ 673274bc-663c-11eb-3766-03d3d4699611
md"Теперь построим регрессию с этим категориальным признаком"

# ╔═╡ e7fb9120-6635-11eb-2c9f-bdad924f344e
model_with_dummies = fit(EconometricModel, @formula(MPG ~ 1 + Cat_cyl), auto)

# ╔═╡ ab477a9e-663c-11eb-3bdb-7bbcc20dcaf8
md"Может смутить, что на графике четко видно, что **Cylinders** принимает 5 уникальных значений, а в выходе мы видим только четыре значения. Это связано с тем, что мы работает с дамми-переменными, и, чтобы избежать мультиколлинеарности, одну переменную не включают явно в уравнение регрессии (это верно только для модели с константой)"

# ╔═╡ ed569b2c-66be-11eb-3a87-bbc80441bde7
dummies_without_constant = fit(EconometricModel, @formula(MPG ~ 0 + Cat_cyl), auto)

# ╔═╡ 08585c6c-66bf-11eb-07e7-ab4d21eb9de2
md"Тут уже мы видим коэффициенты при всех категориях"

# ╔═╡ 8fca184a-6460-11eb-3685-9d8f282862fb
md"### Доверительные интервалы для коэффициентов"

# ╔═╡ 9b3e3364-6657-11eb-266d-29b8805c4a4c
md"Если вам нужен стандартный доверительный интервал (то есть доверительный интервал с уровнем доверия 95%), то достаточно просто посмотреть на выход модели, или вызвать следующую функцию"

# ╔═╡ 2df89a88-646f-11eb-353d-f9eec629f6e0
GLM.coeftable(model_GLM)

# ╔═╡ c17a0f14-6657-11eb-261a-837f9ca8bf81
md"Если нужен другой уровень доверия, то можно воспользоваться библиотекой `StatsBase`. Например, возьмем уровень доверия 80%"

# ╔═╡ 74743d1e-646f-11eb-3900-1789409e3f54
StatsBase.confint(model_GLM, level = 0.8)

# ╔═╡ e8f4608c-6657-11eb-2760-bbb016fc985e
StatsBase.confint(model_Econometrics, level = 0.8)

# ╔═╡ 20550cdc-6658-11eb-3359-b19b970c50de
md"Как мы видим, данная функция уживается и с `GLM`, и c `Econometrics`, и результаты не различаются, что не может не радовать"

# ╔═╡ 5b7c64c2-665d-11eb-0ea2-fb284873a9ed
md"Если нужен доверительный интервал с уровнем доверия 85% только для константы, тогда можно записать следующим образом"

# ╔═╡ 6ed3411c-665d-11eb-0e16-91bab8d97e52
StatsBase.confint(model_GLM, level = 0.85)[1, :]

# ╔═╡ 835a22cc-665d-11eb-12d1-19401e383788
md"Или для коэффициента при **Weight** при уровне доверия 92%"

# ╔═╡ 8156834e-665d-11eb-0c71-dba7bf9d54cb
StatsBase.confint(model_Econometrics, level = 0.92)[2, :]

# ╔═╡ d8a330c2-63db-11eb-2ad9-6b6a7126d393
md"### F-тест и LR-тест для сравнения моделей"

# ╔═╡ d3efe6bc-63b1-11eb-2ad2-f1df05dd812d
md"Рассмотрим две модели:

$$\widehat{MPG}_1 = \widehat{\beta_0} + \widehat{\beta_1} \cdot Weight + \widehat{\beta_2} \cdot Cylinders$$

$$\widehat{MPG}_2 = \widehat{\gamma_0} + \widehat{\gamma_1} \cdot Weight + \widehat{\gamma_2} \cdot Cylinders + \widehat{\gamma_3} \cdot Horsepower$$

F-тест позволит нам сравнить данные модели, и решить, является ли коэффициент $$\gamma_3$$ статистически значимым. Для большей ясности запишем гипотезы:

$$H_0: \gamma_3 = 0$$

$$H_1: \gamma_3 \neq 0$$

С помощью библиотеки `GLM` проведем F-тест
"

# ╔═╡ 1b8b1ad2-63b2-11eb-0d2e-75df725c1846
unrestricted_model = lm(@formula(MPG ~ Weight + Cylinders + Horsepower), auto)

# ╔═╡ 4617f266-63b2-11eb-257f-63906ea45f97
restricted_model = lm(@formula(MPG ~ Weight + Cylinders), auto)

# ╔═╡ 5a5e4770-63b2-11eb-3228-5dc5ea457b7b
GLM.ftest(unrestricted_model.model, restricted_model.model)

# ╔═╡ ee475c76-63b5-11eb-01ea-4121230e1b6f
md"Как мы видим, p-value, или p(>F), равен 0.0003, то есть он меньше любого разумного уровня значимости, отсюда $$\gamma_3 \neq 0$$ при любом разумном уровне значимости"

# ╔═╡ 52fca18e-665a-11eb-158f-69a0618e9b37
md"Так же мы можем провести LR-тест для сравнения наших моделей. Все те же условия, что и ранее"

# ╔═╡ 3637d41a-665a-11eb-27b7-5960cdc86bf4
GLM.lrtest(unrestricted_model.model, restricted_model.model)

# ╔═╡ 37b0229c-665b-11eb-3e81-4f52cf9babaf
md" Видим, что p-value, или p(>Chisq), очень маленький, то есть при любом разумном уровне значимости лучше выбирать неограниченную модель"

# ╔═╡ d65fc5c2-6215-11eb-3f82-2de3ce633b93
md"### Регуляризация"

# ╔═╡ b7f68a00-63b7-11eb-2789-1dce3762f0a6
md"Так же в **Julia** можно построить модели линейной регрессии с L1-регуляризацией (то есть **LASSO**) и L2-регуляризацией (то есть **Ridge**). Для этого воспользуемся библиотекой `GLMNet`. Заодно введем квадраты четырех переменных, поэлементно возводя их в квадрат (просто потому что)"

# ╔═╡ 17994b30-6216-11eb-05de-d1e8fef75c42
begin
	auto["Displacement_sq"] = auto["Displacement"].^2;
	auto["Horsepower_sq"] = auto["Horsepower"].^2;
	auto["Weight_sq"] = auto["Weight"].^2;
	auto["Acceleration_sq"] = auto["Acceleration"].^2;
end

# ╔═╡ 9ebe9d56-63b8-11eb-204f-fb4007c3392c
md"Важное уточнение: `GLMNet` работает не с датафреймами, а с массивами, поэтому выделим нужные признаки и передевем их в массивы"

# ╔═╡ b77d7afc-61e1-11eb-0e0f-55c1b299680c
X = hcat(Array(auto[:, 3:7]), Array(auto[:, 10:14]))

# ╔═╡ c2e747e2-61e1-11eb-3ac5-d7bc26e4b98e
y = Array(auto[:, 1])

# ╔═╡ e42b986c-63b8-11eb-32c7-45cbb84738bc
md"Сначала инициализируем модель **LASSO**, потом **Ridge**.

Подбираем оптимальные $$\lambda$$ с помощью кросс-валидации"

# ╔═╡ c6253aa6-61e4-11eb-214c-412220911078
lasso = glmnetcv(X, y, alpha = 1)

# ╔═╡ 17f4ee00-62d8-11eb-03f6-e78344c2183b
ridge = glmnetcv(X, y, alpha = 0)

# ╔═╡ fdd471ee-63b8-11eb-1231-f707db1c55d8
md"Хочется сравнить эти модели, напишем простую функцию для рассчета коэффциента детерминации. Тут нам нужна функция mean() из пакета `Statistics`"

# ╔═╡ d35f5116-61ee-11eb-1d2d-a70a56e4f5c4
function r2(model, X, y)
	RSS = sum((GLMNet.predict(model, X) - y).^2);
	TSS = sum((y .- Statistics.mean(y)).^2);
	(TSS - RSS)/TSS;
end

# ╔═╡ fdb7edfa-6212-11eb-07e4-19667961840b
r2(lasso, X, y)

# ╔═╡ 1fe3177c-62d8-11eb-0193-cb9420f3f4c7
r2(ridge, X, y)

# ╔═╡ 1f5c2190-63b9-11eb-302e-75f90c39c114
md"Тут модель **LASSO** показала себя лучше"

# ╔═╡ Cell order:
# ╟─8f1873f2-63b0-11eb-289f-4d880cee7e5a
# ╠═a6648952-5a6d-11eb-2c81-99d1514983c0
# ╟─b842ca4a-63b1-11eb-2190-0b6c834f85e0
# ╟─2f38687c-63b3-11eb-384e-c1847038f4c5
# ╠═093d42ea-5a72-11eb-19b1-6b0abb122099
# ╟─ad79eb3c-63b5-11eb-3d99-1f2b7aab9a18
# ╠═22840156-5a72-11eb-0e57-ff3178e5ee6f
# ╟─c24ee52a-63d4-11eb-1815-0d184c506af2
# ╟─61db537c-63d9-11eb-013c-9f74989da4d7
# ╠═fa02d9c0-63d9-11eb-1867-2f9780b312fe
# ╠═757e36ae-63d9-11eb-277a-9b2ac8f3b483
# ╠═d46d4030-63d9-11eb-32d0-1780207fae50
# ╠═e2e70718-63d9-11eb-02e5-0dff07a9c6c7
# ╟─b26f1aac-63da-11eb-2e66-27cd5597e043
# ╠═d987355c-63da-11eb-1170-910084cc4274
# ╟─f0845d38-63db-11eb-0be2-eb5217bf6170
# ╟─17e60466-63c4-11eb-1f97-072f5539cb38
# ╠═8b294048-63bd-11eb-3df5-959b55069357
# ╟─e86909f8-63c4-11eb-3f85-9dd2450e8c5f
# ╠═1fc60af8-63c1-11eb-0814-054e7b9be6f3
# ╟─f42c722a-63c4-11eb-262c-514bf90aa295
# ╠═4161b2ca-63c1-11eb-1f47-3f66d74863c3
# ╟─35163060-63c7-11eb-2013-d1878e8e2c0b
# ╠═65664ad0-63c6-11eb-3e2f-7bf9841d5265
# ╠═7582b77a-63c6-11eb-3ac2-615c4e294aab
# ╟─5ba889e6-63c7-11eb-14a1-216dee432abe
# ╠═8bb5ec74-63c6-11eb-1c3e-d358b4ebc5c5
# ╟─01019b9e-63dc-11eb-3148-9b08890b53c9
# ╠═0b27def8-63dc-11eb-003c-454b21b38132
# ╟─38766bac-63dc-11eb-2dbc-978b5fb3a637
# ╠═550be2bc-63dc-11eb-331b-d71a3071ec66
# ╠═6cd53b1e-63dc-11eb-074a-41c43b6dfaaf
# ╟─7df28c9e-63dc-11eb-3c8e-5732e36a7e5c
# ╟─ce4f3232-63dc-11eb-374d-412a97f8d2c8
# ╠═de284b94-63dc-11eb-2d26-1391517ed17d
# ╠═09221186-63dd-11eb-3c1a-a7fa2ef46e3e
# ╟─3a4f1000-63e3-11eb-3f33-d1590d15b611
# ╟─88873fa4-63e3-11eb-2d3a-e9d30a385d93
# ╠═76daf6ba-63e3-11eb-3e19-ff67bdb86a0e
# ╟─92918540-63e3-11eb-1b57-3d49dfa75b09
# ╟─c080265a-63e3-11eb-1aad-8f85c3d18f3e
# ╠═e66b840e-63e3-11eb-12f7-8d4155f58990
# ╟─6c1f9266-63e4-11eb-08bf-cbabc486dd30
# ╟─0255de96-63e4-11eb-3336-fb9a753ae588
# ╠═13df2904-63e4-11eb-0600-d1c1ebfe9d84
# ╟─2d6038be-63e4-11eb-2320-d7251a567785
# ╟─da6bb162-665c-11eb-20f1-fb018a64646b
# ╠═055e78f2-665d-11eb-188e-8f88fb707c74
# ╟─2afd1c5a-663b-11eb-26bf-5fbcc7b8b14a
# ╟─3e865ba6-663b-11eb-31e0-b57ea1e08bac
# ╠═6d6f598a-63da-11eb-0de2-fb70c755a0ef
# ╟─0d648614-663c-11eb-0dc7-4de628f3f80d
# ╠═deed8534-6635-11eb-3a64-37f4db496443
# ╟─673274bc-663c-11eb-3766-03d3d4699611
# ╠═e7fb9120-6635-11eb-2c9f-bdad924f344e
# ╟─ab477a9e-663c-11eb-3bdb-7bbcc20dcaf8
# ╠═ed569b2c-66be-11eb-3a87-bbc80441bde7
# ╟─08585c6c-66bf-11eb-07e7-ab4d21eb9de2
# ╟─8fca184a-6460-11eb-3685-9d8f282862fb
# ╟─9b3e3364-6657-11eb-266d-29b8805c4a4c
# ╠═2df89a88-646f-11eb-353d-f9eec629f6e0
# ╟─c17a0f14-6657-11eb-261a-837f9ca8bf81
# ╠═74743d1e-646f-11eb-3900-1789409e3f54
# ╠═e8f4608c-6657-11eb-2760-bbb016fc985e
# ╟─20550cdc-6658-11eb-3359-b19b970c50de
# ╟─5b7c64c2-665d-11eb-0ea2-fb284873a9ed
# ╠═6ed3411c-665d-11eb-0e16-91bab8d97e52
# ╟─835a22cc-665d-11eb-12d1-19401e383788
# ╠═8156834e-665d-11eb-0c71-dba7bf9d54cb
# ╟─d8a330c2-63db-11eb-2ad9-6b6a7126d393
# ╟─d3efe6bc-63b1-11eb-2ad2-f1df05dd812d
# ╠═1b8b1ad2-63b2-11eb-0d2e-75df725c1846
# ╠═4617f266-63b2-11eb-257f-63906ea45f97
# ╠═5a5e4770-63b2-11eb-3228-5dc5ea457b7b
# ╟─ee475c76-63b5-11eb-01ea-4121230e1b6f
# ╟─52fca18e-665a-11eb-158f-69a0618e9b37
# ╠═3637d41a-665a-11eb-27b7-5960cdc86bf4
# ╟─37b0229c-665b-11eb-3e81-4f52cf9babaf
# ╟─d65fc5c2-6215-11eb-3f82-2de3ce633b93
# ╟─b7f68a00-63b7-11eb-2789-1dce3762f0a6
# ╠═17994b30-6216-11eb-05de-d1e8fef75c42
# ╟─9ebe9d56-63b8-11eb-204f-fb4007c3392c
# ╠═b77d7afc-61e1-11eb-0e0f-55c1b299680c
# ╠═c2e747e2-61e1-11eb-3ac5-d7bc26e4b98e
# ╟─e42b986c-63b8-11eb-32c7-45cbb84738bc
# ╠═c6253aa6-61e4-11eb-214c-412220911078
# ╠═17f4ee00-62d8-11eb-03f6-e78344c2183b
# ╟─fdd471ee-63b8-11eb-1231-f707db1c55d8
# ╠═d35f5116-61ee-11eb-1d2d-a70a56e4f5c4
# ╠═fdb7edfa-6212-11eb-07e4-19667961840b
# ╠═1fe3177c-62d8-11eb-0193-cb9420f3f4c7
# ╟─1f5c2190-63b9-11eb-302e-75f90c39c114


"""
This package contains some codes from IterativeSolvers.jl: Copyright (c) 2013--2016 The Julia Language.
For the copyright notice, please check the LICENSE file.
"""
mutable struct BiCGStabIterable{
    opT, solT, vecT, matT, collectT,
    realT <: Real,
    scalarT <: Number}

    A::opT
    l::Int

    x::solT
    r_shadow::solT
    rs::collectT
    us::collectT

    maxiter::Int
    niter::Int
    tol::realT
    residual::realT

    gamma::vecT
    omega::scalarT
    sigma::scalarT
    M::matT

    function BiCGStabIterable(
        A::opT, l::Int, x::solT,
        r_shadow::solT, rs::collectT, us::collectT,
        maxiter::Int, niter::Int, tol::realT, residual::realT,
        gamma::vecT, omega::scalarT, sigma::scalarT,
        M::matT) where {opT,solT,vecT,matT,collectT,realT<:Real,scalarT<:Number}

        new{opT,solT,vecT,matT,collectT,realT,scalarT}(
            A,l,x,r_shadow,rs,us,maxiter,niter,tol,residual,gamma,omega,sigma,M)
    end
end

function init_bicgstab(
    x, A, b, l::Int=2;
    abstol::Real=zero(real(eltype(b))),
    reltol::Real=sqrt(eps(real(eltype(b)))),
    maxiter::Int=2000)

    T = typeof(x)
    N = size(x)
    niter = 0

    # Large vectors.
    r_shadow = T(undef, N)
    rand!(r_shadow)
    rs = Array{T}(undef, l+1)
    us = Array{T}(undef, l+1)
    for idx in 1:(l+1)
        rs[idx] = T(undef, N)
        us[idx] = T(undef, N)
        fill!(us[idx], zero(eltype(x)))
    end

    residual = rs[1]

    # Compute the initial residual = b - A * x
    mul!(residual, A, x)
    sub!(residual, b, residual)
    niter += 1

    gamma = zeros(eltype(x), l)
    omega = sigma = one(eltype(x))

    nrm = norm(residual)

    # For the least-squares problem
    M = zeros(eltype(x), l+1, l+1)

    # Stopping condition based on absolute and relative tolerance.
    tolerance = max(reltol * nrm, abstol)

    BiCGStabIterable(A, l, x, r_shadow, rs, us,
        maxiter, niter, tolerance, nrm,
        gamma, omega, sigma, M
    )
end

@inline converged(it::BiCGStabIterable) = it.residual ≤ it.tol
@inline start(::BiCGStabIterable) = 0
@inline done(it::BiCGStabIterable, iteration::Int) = it.niter ≥ it.maxiter || converged(it)

function iterate(it::BiCGStabIterable, iteration::Int=start(it))
    if done(it, iteration) return nothing end

    T = eltype(it.x)
    L = 2:(it.l + 1)

    it.sigma = -it.omega * it.sigma

    ## BiCG part
    for j = 1:(it.l)
        rho = dot(it.r_shadow, it.rs[j])
        beta = rho / it.sigma

        # us[:, 1 : j] .= rs[:, 1 : j] - β * us[:, 1 : j]
        for idx = 1:j
            sub!(it.us[idx], it.rs[idx], it.us[idx])
        end

        # us[:, j + 1] = Pl \ (A * us[:, j])
        mul!(it.us[j+1], it.A, it.us[j])

        it.sigma = dot(it.r_shadow, it.us[j+1])
        alpha = rho / it.sigma

        for idx = 1:j
            sub_scmul!(it.rs[idx], alpha, it.us[idx+1])
        end

        # rs[:, j + 1] = Pl \ (A * rs[:, j])
        mul!(it.rs[j+1], it.A , it.rs[j])

        # x = x + α * us[:, 1]
        add_scmul!(it.x, alpha, it.us[1])
    end

    # Bookkeeping
    it.niter += 2 * it.l

    ## MR part
    # M = rs' * rs
    for cidx in CartesianIndices(it.M)
        i1 = cidx[1]
        i2 = cidx[2]
        it.M[i1,i2] = dot(it.rs[i1], it.rs[i2])
    end

    # γ = M[L, L] \ M[L, 1]
    F = lu!(view(it.M, L, L))
    ldiv!(it.gamma, F, view(it.M, L, 1))

    for idx in L
        sub_scmul!(it.us[1], it.gamma[idx-1], it.us[idx])
        add_scmul!(it.x, it.gamma[idx-1], it.rs[idx-1])
        sub_scmul!(it.rs[1], it.gamma[idx-1], it.rs[idx])
    end

    it.omega = it.gamma[it.l]
    it.residual = norm(it.rs[1])

    it.residual, iteration + 1
end

"""
l: number of GMRES steps
"""
function bicgstab!(
    x, A, b, l::Int=2;
    abstol::Real=zero(real(eltype(b))),
    reltol::Real=sqrt(eps(real(eltype(b)))),
    maxiter::Int=2000,
    verbose::Bool=false)

    history = zeros(0)

    iterable = init_bicgstab(
        x, A, b, l;
        abstol=abstol, reltol=reltol,
        maxiter=maxiter)

    for (iteration, item) = enumerate(iterable)
        push!(history, iterable.residual)
        verbose && @printf("%3d\t%1.2e\n", iteration, iterable.residual)
    end

    if !converged(iterable)
        push!(history, -1)
        verbose && println("BiCGStab: not converged...")
    end
    # println()

    return history
end


function mcmc_run!(x_current,P_in,p,q_vec,q_mat,u_mat,ind_row::T,ind_chain::T) where {T<:Integer}
    accept_val = 0.0
    P_current = P_in
    x_new = x_current.+(q_mat[ind_row,:,ind_chain].*q_vec)
    p_new = p(x_new)
    A = min(1,p_new/P_current) #calcualte acceptance ratio
    u = u_mat[ind_row,ind_chain]
        if u<=A
            x_current[1:end] = x_new
            P_current = p_new
            accept_val = 1.0
        end

    return P_current,accept_val
end

function n_mcmc_runs!(x_current,P_in,p,q_vec,q_mat,u_mat,ind_row_start::T,ind_chain::T,n_runs::T) where {T<:Integer}
    n_accepts = 0.0
    P_current = P_in
    for i = 1:n_runs
        ind_row=ind_row_start+i
        P_current,accept_val = mcmc_run!(x_current,P_current,p,q_vec,q_mat,u_mat,ind_row,ind_chain)
        n_accepts += accept_val
    end
    return P_current,n_accepts
end

function metropolis_mcmc(x_current_vec,P_current_vec,p,q_vec::Array{Float64,1},file_name::String
    ;n_chains::Integer=7,n_samples::Integer=100,burn_in::Integer=100,sample_freq::Integer=10)

    n_draws = n_samples*sample_freq+burn_in #number of random number per chain
    n_features = length(q_vec)

    u_mat = rand(n_draws,n_chains)
    q_mat = rand(Normal(),(n_draws,n_features,n_chains))

    samples_container = Array{Array{typeof(x_current_vec[1]),1},1}(undef,n_chains)
    P_samples_container = Array{Array{Float64,1},1}(undef,n_chains)

    Threads.@threads for j = 1:n_chains
        core_id = j

        x_current,P_current = copy(x_current_vec[j]), copy(P_current_vec[j])

        #Run burn-in
        println("Core id $(core_id): Start Burn-in")
        P_current,na = n_mcmc_runs!(x_current,P_current,p,q_vec,q_mat,u_mat,0,j,burn_in)
        println("Core id $(core_id): Done Burn-in")

        #Create container for samples
        samples = Array{typeof(x_current),1}(undef,n_samples)
        #Propability value of samples
        P_samples = zeros(n_samples)
        accept_rate = 0.0

        println("Core id $(core_id): Start Sampling")
        #Run samples
        for i = 1:n_samples
            println("Core id $(core_id): ---Sample: $(i)/$(n_samples)")
            ind_row_start = burn_in+(i-1)*sample_freq
            P_current,na = n_mcmc_runs!(x_current,P_current,p,q_vec,q_mat,u_mat,ind_row_start,j,sample_freq)
            accept_rate += na
            #Save sample
            samples[i] = copy(x_current)
            #Save propability values
            P_samples[i] = P_current
        end
        println("Core id $(core_id): Done Sampling")
        accept_rate /= n_samples*sample_freq
        println("Core id $(core_id): accept_rate = $(accept_rate)")

        samples_container[j] = samples
        P_samples_container[j] = P_samples
    end
    save("./output/"*file_name*".jld","samples_container",samples_container,"P_samples_container",P_samples_container)
    return nothing
end


function ssfm_solve(num::Numerics, coeffs::Coefficients)
  time_steps = Int(floor(num.T / num.dt))
  space_steps = Int(floor(num.S / num.ds))
  time = LinRange(0, num.T, time_steps)
  space = LinRange(-num.S / 2, num.S / 2, space_steps)

  # print("Simulation diagnostics:\n")
  # @printf("Time  : from %6.3i ms, to  %6.3i ms; %6d steps \n", time[1] * 1e3, time[time_steps] * 1e3, time_steps)
  # @printf("Space : from %6.3f mm, to  %6.3f mm; %6d steps \n", space[1] * 1e3, space[space_steps] * 1e3, space_steps)

  # Spatial frequency range computation
  k = 2 * pi * LinRange(-1 / (2 * num.ds), 1 / (2 * num.ds), space_steps)
  k = fftshift(k)
  ψ = zeros(ComplexF64, space_steps, time_steps)
  ψ_spect = zeros(ComplexF64, space_steps, time_steps)

  waveform = coeffs.initial.(space)
  # fig0 = plot(space * 1e3,
  #   abs.(waveform) .^ 2,
  #   title="initial configuration",
  #   xlabel="space [mm]",
  #   show=true,
  #   reuse=false,
  #   label="initial state")

  # fig1 = plot(space * 1e3,
  #   abs.(coeffs.β.(space)),
  #   show=true,
  #   reuse=false,
  #   label="potential")

  ## [SPACE INDEX, TIME INDEX]
  ψ[:, 1] = waveform

  fwd_disp = exp.(num.dt / 2 .* coeffs.α * k .^ 2)
  fwd_curvature = exp.(num.dt / 2 .* coeffs.β.(space))
  # display("Gamma at peak")
  # display(coeffs.γ.(ψ[Int(floor(space_steps/2)),1]))
  #@showprogress "Propagating the field... "
  for n = 1:time_steps-1
    # if 1+2*app.as*(app.N - 1) * abs.(ψ[Int(floor(space_steps/2)), n])^2 > 1
    #   display("Collapse detected")
    #   return
    # end
    ψ_spect[:, n] = fft(ψ[:, n])
    ψ_spect[:, n] = ψ_spect[:, n] .* fwd_disp
    ψ[:, n] = ifft(ψ_spect[:, n])
    ψ[:, n+1] = ψ[:, n] .* exp.(num.dt / 2 .* coeffs.γ.(ψ[:, n])) .* fwd_curvature  ## this is an Euler step
  end
  ψ_spect[:, time_steps] = fft(ψ[:, time_steps])

  return time, space, ψ, ψ_spect
end

function ssfm_propagate(num::Numerics, coeffs::Coefficients)
  time_steps = Int(floor(num.T / num.dt))
  space_steps = Int(floor(num.S / num.ds))
  time = LinRange(0, num.T, time_steps)
  space = LinRange(-num.S / 2, num.S / 2, space_steps)

  # Spatial frequency range computation
  k = 2 * pi * LinRange(-1 / (2 * num.ds), 1 / (2 * num.ds), space_steps)
  k = fftshift(k)
  ψ = zeros(ComplexF64, space_steps)
  ψ_spect = zeros(ComplexF64, space_steps)

  waveform = coeffs.initial.(space)

  ## [SPACE INDEX, TIME INDEX]
  ψ = waveform

  fwd_disp = exp.(num.dt / 2 .* coeffs.α * k .^ 2)
  fwd_curvature = exp.(num.dt / 2 .* coeffs.β.(space))
  max_amplitude = maximum(abs.(ψ).^2)
  for n = 1:time_steps-1

    ψ_spect = fft(ψ)
    ψ_spect = ψ_spect .* fwd_disp
    ψ = ifft(ψ_spect)
    ψ= ψ .* exp.(num.dt / 2 .* coeffs.γ.(ψ)) .* fwd_curvature  ## this is an Euler step
    if max_amplitude < maximum(abs.(ψ).^2)
      max_amplitude = maximum(abs.(ψ).^2)
    end
  end
  ψ_spect = fft(ψ)

  return time, space, ψ, ψ_spect, max_amplitude
end

function ground_state_solve(num::Numerics, coeffs::Coefficients)
  time_steps = Int(floor(num.T / num.dt))
  space_steps = Int(floor(num.S / num.ds))
  time = LinRange(0, num.T, time_steps)
  space = LinRange(-num.S / 2, num.S / 2, space_steps)

  # print("Simulation diagnostics:\n")
  # @printf("Time  : from %6.3i ms, to  %6.3i ms; %6d steps \n", time[1] * 1e3, time[time_steps] * 1e3, time_steps)
  # @printf("Space : from %6.3f mm, to  %6.3f mm; %6d steps \n", space[1] * 1e3, space[space_steps] * 1e3, space_steps)

  # Spatial frequency range computation
  k = 2 * pi * LinRange(-1 / (2 * num.ds), 1 / (2 * num.ds), space_steps)
  k = fftshift(k)
  #fig = Plots.plot(k, show=true)
  ψ = zeros(ComplexF64, space_steps)
  ψ_spect = zeros(ComplexF64, space_steps)

  #Implement gaussian initial state
  wave(s::Float64) = sqrt(1 / (sqrt(2 * pi) * num.S / 3)) * exp.(-(s) .^ 2 / (4 * num.S / 3^2))
  integral, error = quadgk(s -> abs(wave(s))^2, -num.S / 2, +num.S / 2)

  waveform = wave.(space) / sqrt(integral)

  ## [SPACE INDEX, TIME INDEX]
  ψ[:] = waveform
  fwd_disp = exp.(-im * num.dt / 2 .* coeffs.α * k .^ 2)
  fwd_curvature = exp.(-im * num.dt / 2 .* coeffs.β.(space))
  # display("Gamma at peak")
  # display(coeffs.γ.(ψ[Int(floor(space_steps/2)),1]))
  target_Δ = 1e-10
  Δ = 1
  let prev, current
    prev = 0
    while Δ > target_Δ
      # if 1+2*app.as*(app.N - 1) * abs.(ψ[Int(floor(space_steps/2)), n])^2 > 1
      #   display("Collapse detected")
      #   return
      # end
      ψ_spect = fft(ψ)
      ψ_spect = ψ_spect .* fwd_disp
      ψ = ifft(ψ_spect)
      ψ = ψ .* exp.(im * num.dt / 2 .* coeffs.γ.(ψ)) .* fwd_curvature  ## this is an Euler step
      # renormalization
      #isplay(sum(abs.(ψ) .^ 2 * num.ds))
      current = sum(abs.(ψ) .^ 2 * num.ds)
      Δ = abs(prev - current)
      ψ = ψ / sqrt(current)
      prev = current
    end
  end
  ψ_spect = fft(ψ)
  return time, space, ψ, ψ_spect
end

function plot_dynamics(time, space, ψ, ψ_spect)
  pyplot()
  # Natural orientation choice for plot
  fig2 = heatmap(time * 1e3,
    space * 1e3,
    abs.(ψ[:, :]) .^ 2,
    show=true,
    title="|ψ|^2",
    ylabel="space [mm]",
    xlabel="time [ms]",
    colorrange=(0, 1),
    reuse=false,
    size=(800, 600))

  # fig2_bis = heatmap(time * 1e3,
  #   space * 1e3,
  #   angle.(ψ[:, :]),
  #   show=true,
  #   title="|ψ|^2",
  #   ylabel="space [mm]",
  #   xlabel="time [ms]",
  #   colorrange=(0, 1),
  #   reuse=false,
  #   size=(800, 600))

  fig3 = plot(space * 1e3,
    abs.(ψ[:, end]) .^ 2,
    title="evolved |ψ|^2",
    xlabel="space [mm]",
    reuse=false,
    label="t=t_max")
  plot!(space * 1e3,
    abs.(ψ[:, 1]) .^ 2,
    show=true,
    reuse=false,
    label="t=0",
    size=(800, 600))
end

function plot_ground_state(time, space, ψ, ψ_spect)
  pyplot()
  # Natural orientation choice for plot
  fig3 = plot(space * 1e3,
    abs.(ψ) .^ 2,
    title="stationary |ψ|^2",
    xlabel="space [mm]",
    reuse=false,
    show=true,
    label="t=t_max")
end


module TestReparametrizedBautin
include("preamble.jl")

using ForwardDiff
using Setfield: compose
using StaticArrays: SVector, SMatrix

using Bifurcations: resolved_points, reparametrize
using Bifurcations.Codim2: first_lyapunov_coefficient
using Bifurcations.Codim2LimitCycle: FoldLimitCycleProblem
using Bifurcations.Continuations: get_prob_cache, get_u0,
    residual!, residual_jacobian
using Bifurcations.Examples.Bautin
using Bifurcations.Examples.Reparametrization: orig_p

# Putting @info so that Travis won't stop testing.
@testset for opts in [
        [],
        [:extra => SVector(-1.0)],
        [:extra => SVector(-1.0), :shift => 0.1, :seed => 1],
        ]
    @info "TestReparametrizedBautin:"
    @info "opts = $opts"

    prob = reparametrize(
        Bautin.make_prob(),
        ; opts...)
    solver1 = init(prob)
    solve!(solver1)

    codim1_points = resolved_points(solver1)
    @test length(codim1_points) == 1
    @test codim1_points[1].point_type === Codim1.PointTypes.hopf
    @test codim1_points[1].u[end] ≈ 0  atol=1e-6

    @info "Continuation of Hopf bifurcation..."
    hopf_prob = BifurcationProblem(
        codim1_points[1],
        solver1,
        compose(orig_p, (@lens _.β₂)),
        (-2.0, 2.0),
    )
    hopf_solver = init(
        hopf_prob;
        verbose = true,
    )
    @time solve!(hopf_solver)

    hopf_β₁ = [u[end-1] for sweep in hopf_solver.super.sol.sweeps for u in sweep.u]
    hopf_β₂ = [u[end]   for sweep in hopf_solver.super.sol.sweeps for u in sweep.u]
    @test all(@. abs(hopf_β₁) < 1e-6)
    @test maximum(hopf_β₂) > 2
    @test minimum(hopf_β₂) < -2

    l1s = [
        first_lyapunov_coefficient(hopf_solver.cache.super.prob_cache, u)
        for sweep in hopf_solver.super.sol.sweeps for u in sweep.u
    ]
    if !all(@. (l1s > 0) == (hopf_β₂ > 0))
        println("[hopf_β₂ l1s]' =")
        display([hopf_β₂ l1s]')
    end
    @test all(@. (l1s > 0) == (hopf_β₂ > 0))

    @info "Resolving Bautin point..."
    @time codim2_points = resolved_points(hopf_solver)
    @test length(codim2_points) == 1
    @test codim2_points[1].point_type === Codim2.PointTypes.bautin
    β_bautin = codim2_points[1].u[end-1:end]
    @test all(@. abs(β_bautin) < 1e-6)

    @info "Creating FoldLimitCycleProblem..."
    @time flc_prob = FoldLimitCycleProblem(
        codim2_points[1],
        hopf_solver;
        num_mesh = 20,
        degree = 3,
    )
    @show flc_prob.super.num_mesh
    @show flc_prob.super.degree
    @test flc_prob.t_domain == ([-2.0, -2.0], [2.0, 2.0])

    @info "Checking Jacobian..."
    @testset "Manual Jacobian vs ForwardDiff" begin
        prob_cache = get_prob_cache(flc_prob)
        u0 = get_u0(flc_prob)
        rng = MersenneTwister(0)
        for _ in 1:5
            # Trying to avoid state values whose abs is larger than 1;
            # fitting polynomial to random numbers is not nice thing
            # to try here...
            u_lc = 0.01 .* tanh.(randn(rng, eltype(u0), length(u0)))
            d = length(flc_prob.xs0)
            u_lc[d + 1] = abs(u_lc[d + 1])  # period > 0
            H_actual, J_actual = residual_jacobian(u_lc, prob_cache)

            H_desired = similar(@view u_lc[1:end-1])
            J_desired = similar(H_desired, (length(H_desired), length(u_lc)))
            ForwardDiff.jacobian!(
                J_desired,
                (y, x) -> residual!(y, x, prob_cache),
                H_desired,
                u_lc,
            )

            @test J_actual ≈ J_desired
            @test H_actual ≈ H_desired
        end
    end

    @info "Continuation of fold bifurcation of a limit cycle..."
    flc_solver = init(
        flc_prob;
        start_from_nearest_root = true,
        max_branches = 0,  # TODO: stop manually doing this
        verbose = true,
    )
    @time solve!(flc_solver)

    flc_β₁ = [u[end-1] for sweep in flc_solver.super.sol.sweeps for u in sweep.u]
    flc_β₂ = [u[end]   for sweep in flc_solver.super.sol.sweeps for u in sweep.u]
    @show maximum(@. abs(4 * flc_β₁ + flc_β₂^2))
    @test all(@. abs(4 * flc_β₁ + flc_β₂^2) < 5e-2)
    @test maximum(flc_β₂) > 2
    @test minimum(flc_β₂) > -5e-2
end

end  # module


using EchelleCCFs
using Test

@testset "Check CCF accuracy" begin  # TODO repeat for multiple mask shapes
    @testset "Tophat mask w/ 1 line" begin
        Δv = 1000
        λ1 = 5025.0
        d1 = 0.5
        σ_v = 5000.0
        λlo = 5000
        λhi = 5050
        npix = 2000
        @test_nowarn BasicLineList([λ1], [d1])
        ll = BasicLineList([λ1], [d1])
        ll = BasicLineList([λ1], [d1])
        @test_nowarn BasicCCFPlan(line_list=ll, mask_shape=TopHatCCFMask(Δv))
        ccfpl = BasicCCFPlan(line_list=ll, mask_shape=TopHatCCFMask(Δv))
        @test_nowarn calc_ccf_v_grid(ccfpl)
        v_grid = calc_ccf_v_grid(ccfpl)
        @test length(v_grid) == calc_length_ccf_v_grid(ccfpl)
        λs = exp.(range(log(λlo),stop=log(λhi),length=npix))
        flux = ones(size(λs))
        flux .*= 1 .- d1.*exp.(-0.5.*((λs.-λ1)./λ1.*(EchelleCCFs.speed_of_light_mps./σ_v)).^2)
        @test_nowarn ccf_1D(λs, flux, ccfpl)
        ccf = ccf_1D(λs, flux, ccfpl)
        @test sum(ccf) > 0
        @test_nowarn  MeasureRvFromCCFGaussian()
        mrfcg = MeasureRvFromCCFGaussian(frac_of_width_to_fit=0.75)
        @test_nowarn measure_rv_from_ccf(v_grid,ccf, alg=mrfcg )
        vfit = measure_rv_from_ccf(v_grid,ccf, alg=mrfcg )
        @test vfit.rv ≈ 0  atol = 0.1   # TODO: tune tolerance better
        var = flux./10^5
        @test_nowarn ccf_1D(λs, flux, var, ccfpl)
        (ccf, ccf_var) = ccf_1D(λs, flux, var, ccfpl)
        @test sum(ccf) > 0
        @test_nowarn measure_rv_from_ccf(v_grid,ccf, ccf_var, alg=mrfcg )
        vfit = measure_rv_from_ccf(v_grid,ccf,ccf_var, alg=mrfcg )  # TODO: Check Why is this doing so much worse?
        #=
        using Plots
        global plt = plot()
        plot!(plt,v_grid,ccf,label=:none)
        scatter!(plt,v_grid,ccf,color=1,label=:none)
        plot!(plt,v_grid,sqrt.(ccf_var),label=:none)
        =#
        @test vfit.rv ≈ 0  atol = 0.5   # TODO: tune tolerance better
    end

    @testset "Tophat mask w/ 2 lines" begin
        Δv = 1000
        λ1 = 5025.0
        λ2 = 5075.0
        d1 = 0.5
        d2 = 0.25
        σ_v = 5000.0
        λlo = 5000
        λhi = 5100
        npix = 8000
        @test_nowarn BasicLineList([λ1,λ2], [d1,d2])
        ll = BasicLineList([λ1,λ2], [d1,d2])
        ll = BasicLineList([λ1,λ2], [d1,d2])
        @test_nowarn BasicCCFPlan(line_list=ll, mask_shape=TopHatCCFMask(Δv))
        ccfpl = BasicCCFPlan(line_list=ll, mask_shape=TopHatCCFMask(Δv))
        @test_nowarn calc_ccf_v_grid(ccfpl)
        v_grid = calc_ccf_v_grid(ccfpl)
        @test length(v_grid) == calc_length_ccf_v_grid(ccfpl)
        λs = exp.(range(log(λlo),stop=log(λhi),length=npix))
        flux = ones(size(λs))
        flux .*= 1 .- d1.*exp.(-0.5.*((λs.-λ1)./λ1.*(EchelleCCFs.speed_of_light_mps./σ_v)).^2)
        flux .*= 1 .- d2.*exp.(-0.5.*((λs.-λ2)./λ2.*(EchelleCCFs.speed_of_light_mps./σ_v)).^2)
        @test_nowarn ccf_1D(λs, flux, ccfpl)
        var = flux./10^5
        @test_nowarn ccf_1D(λs, flux, var, ccfpl)
        (ccf, ccf_var) = ccf_1D(λs, flux, var, ccfpl)

        @test_nowarn  MeasureRvFromCCFGaussian()
        mrfcg = MeasureRvFromCCFGaussian()
        @test_nowarn measure_rv_from_ccf(v_grid,ccf, alg=mrfcg )
        @test_nowarn measure_rv_from_ccf(v_grid,ccf, ccf_var, alg=mrfcg )
        vfit = measure_rv_from_ccf(v_grid,ccf, alg=mrfcg )
        @test vfit.rv ≈ 0  atol = 0.3   # TODO: tune tolerance better
        vfit = measure_rv_from_ccf(v_grid,ccf,ccf_var, alg=mrfcg )  # TODO: Check Why is this doing so much worse?
        #=
        using Plots
        global plt = plot()
        plot!(plt,v_grid,ccf,label=:none)
        scatter!(plt,v_grid,ccf,color=1,label=:none)
        plot!(plt,v_grid,sqrt.(ccf_var),label=:none)
        =#
        @test vfit.rv ≈ 0  atol = 0.3   # TODO: tune tolerance better

    end
    @testset "Gaussian mask" begin
        Δv = 7e3
        λ1 = 5025.0
        λ2 = 5075.0
        d1 = 0.5
        d2 = 0.25
        σ_v = 5000.0
        λlo = 5000
        λhi = 5100
        npix = 4000
        @test_nowarn BasicLineList([λ1,λ2], [d1,d2])
        ll = BasicLineList([λ1,λ2], [d1,d2])
        @test_nowarn BasicCCFPlan(line_list=ll, mask_shape=GaussianCCFMask(σ_v))
        ccfpl = BasicCCFPlan(line_list=ll, mask_shape=GaussianCCFMask(σ_v))
        @test_nowarn calc_ccf_v_grid(ccfpl)
        v_grid = calc_ccf_v_grid(ccfpl)
        @test length(v_grid) == calc_length_ccf_v_grid(ccfpl)
        λs = exp.(range(log(λlo),stop=log(λhi),length=npix))
        flux = ones(size(λs))
        flux .*= 1 .- d1.*exp.(-0.5.*((λs.-λ1)./λ1.*(EchelleCCFs.speed_of_light_mps./σ_v)).^2)
        flux .*= 1 .- d2.*exp.(-0.5.*((λs.-λ2)./λ2.*(EchelleCCFs.speed_of_light_mps./σ_v)).^2)
        @test_nowarn ccf_1D(λs, flux, ccfpl)
        var = flux./10^5
        @test_nowarn ccf_1D(λs, flux, var, ccfpl)
        (ccf, ccf_var) = ccf_1D(λs, flux, var, ccfpl)

        @test_nowarn measure_rv_from_ccf(v_grid,ccf, alg=MeasureRvFromCCFQuadratic() )
        @test_nowarn measure_rv_from_ccf(v_grid,ccf,ccf_var, alg=MeasureRvFromCCFQuadratic() )
        vfit = measure_rv_from_ccf(v_grid,ccf, alg=MeasureRvFromCCFQuadratic() )    # TODO: Check Why is this doing so much worse?
        @test vfit.rv ≈ 0  atol = 0.3   # TODO: tune tolerance better
        vfit = measure_rv_from_ccf(v_grid,ccf,ccf_var, alg=MeasureRvFromCCFQuadratic() )
        @test vfit.rv ≈ 0  atol = 0.3   # TODO: tune tolerance better

        @test_nowarn measure_rv_from_ccf(v_grid,ccf, alg=MeasureRvFromMinCCF() )
        vfit = measure_rv_from_ccf(v_grid,ccf, alg=MeasureRvFromMinCCF())
        @test vfit.rv ≈ 0  atol = 0.1   # TODO: tune tolerance better
        vfit = measure_rv_from_ccf(v_grid,ccf,ccf_var, alg=MeasureRvFromMinCCF())
        @test vfit.rv ≈ 0  atol = 0.5   # TODO: tune tolerance better

        #=
        @test_nowarn measure_rv_from_ccf(v_grid,ccf, alg=MeasureRvFromCCFCentroid() )
        vfit = measure_rv_from_ccf(v_grid,ccf)
        @test vfit.rv ≈ 0  atol = 500   # TODO: tune tolerance better
        =#
    end
end


abstract type CovarianceDesign{T} end

get_Σ(mat) = mat
nfeatures(mat) = size(mat, 1)

function spectrum(mat)
    eigs = eigvals(mat)
    probs = fill(1/length(eigs), length(eigs))
    DiscreteNonParametric(eigs, probs)
end

spectrum(cov::CovarianceDesign) = spectrum(get_Σ(cov))

nfeatures(cov::CovarianceDesign) = cov.p

function simulate_rotated_design(cov, n; rotated_measure = Normal())
    Σ = get_Σ(cov)
    Σ_chol = cholesky(Σ)
    p = nfeatures(cov)
    Z = rand(rotated_measure, n, p)
    X = randn(n, p) * Σ_chol.UL
    X
end

Base.@kwdef struct AR1Design{P<:Union{Missing,Int}} <: CovarianceDesign{P}
    p::P = missing
    ρ = 0.7
end

function get_Σ(cov::AR1Design{Int})
    p = nfeatures(cov)
    ρ = cov.ρ
    Σ = [ρ^(abs(i - j)) for i = 1:p, j = 1:p]
    Σ
end

abstract type DiagonalCovarianceDesign{T} <: CovarianceDesign{T} end

Base.@kwdef struct IdentityCovarianceDesign{P<:Union{Missing,Int}} <:
                   DiagonalCovarianceDesign{P}
    p::P = missing
end

spectrum(::IdentityCovarianceDesign) = DiscreteNonParametric([1.0],[1.0])

function get_Σ(cov::IdentityCovarianceDesign{Int})
    I(cov.p)
end

Base.@kwdef struct UniformScalingCovarianceDesign{P<:Union{Missing,Int}} <:
                   DiagonalCovarianceDesign{P}
    scaling::Float64 = 1.0
    p::P = missing
end

spectrum(unif::UniformScalingCovarianceDesign) = DiscreteNonParametric([unif.scaling],[1.0])

function get_Σ(cov::UniformScalingCovarianceDesign{Int})
    (cov.scaling * I)(cov.p)
end

Base.@kwdef struct ExponentialOrderStatsCovarianceDesign{P<:Union{Missing,Int}} <:
                   DiagonalCovarianceDesign{P}
    p::P = missing
    rate::Float64
end

function spectrum(cov::ExponentialOrderStatsCovarianceDesign)
    p = cov.p
    rate = cov.rate
    tmp = range(1 / (2p); stop = 1 - 1 / (2p), length = p)
    eigs = 1 / rate .* log.(1 ./ tmp)
    DiscreteNonParametric(eigs, fill(1/p, p))
end

get_Σ(cov::ExponentialOrderStatsCovarianceDesign) = Diagonal(support(spectrum(cov)))

struct BlockCovarianceDesign{T, S <: CovarianceDesign{T}, G} <: CovarianceDesign{T}
    blocks::Vector{S}
    groups::G
end

function BlockCovarianceDesign(blocks::Vector{S}) where S<:CovarianceDesign{Missing}
    BlockCovarianceDesign(blocks, missing)
end

nfeatures(cov::BlockCovarianceDesign) = sum(nfeatures.(cov.blocks))

function get_Σ(blockdesign::BlockCovarianceDesign)
    BlockDiagonal(get_Σ.(blockdesign.blocks))
end

function spectrum(blockdesign::BlockCovarianceDesign)
    @unpack blocks, groups = blockdesign
    spectra = spectrum.(blocks)
    mixing_prop = groups.ps ./ groups.p
    MixtureModel(spectra, mixing_prop)
end

function simulate_rotated_design(cov::BlockCovarianceDesign, n; rotated_measure = Normal())
    hcat(simulate_rotated_design.(cov.blocks, n; rotated_measure=rotated_measure)...)
end

# Set groups
function set_groups(design::CovarianceDesign, p::Integer)
    @set design.p = p
end

function set_groups(design::CovarianceDesign, groups::GroupedFeatures)
    set_groups(design, nfeatures(groups))
end

function set_groups(blockdesign::BlockCovarianceDesign, groups::GroupedFeatures)
    updated_blocks = set_groups.(blockdesign.blocks, groups.ps)
    BlockCovarianceDesign(updated_blocks, groups)
end


export IrisEnv, fitness, new_gen!

import RDatasets

"Iris BERL environment"
mutable struct IrisEnv <: BERLenv
    name::String
    memory::Array # Array of storable data
    X::Array{Float64, 2}
    y::Array{Float64, 2}
end

environments["iris"]=IrisEnv

function IrisEnv(cfg::Dict)
    cfg["env"]="Iris"
    cfg["n_in"]=4
    cfg["n_out"]=3
    iris = RDatasets.dataset("datasets", "iris")
    X = convert(Matrix, iris[:, 1:4])'
    X = X ./ maximum(X; dims=2) # normalize
    r = iris[:, 5].refs # labels
    # 1-hot encoding:
    y = zeros(maximum(r), size(X, 2))
    for i in 1:length(r)
        y[r[i], i] = 1.0
    end
    IrisEnv("iris", [], X, y)
end

"Returns the fitness of one Cambrian Individual in this environment"
function fitness(indiv::Cambrian.Individual, env::IrisEnv)
    accuracy = 0.0
    for i in 1:size(env.X, 2)
        out = process(indiv, env.X[:, i])
        if argmax(out) == argmax(env.y[:, i])
            accuracy += 1
        end
    end
    f = [accuracy / size(env.X, 2)]
    try
        Integer.(round.(f))
    catch
        println(f, accuracy, size(env.X, 2))
    end
    f
end

"Sets up the env for the next generation"
function new_gen!(env::IrisEnv, cfg::Dict=nothing)

end


module Metrics

using Photon: Tensor
using Statistics
import Knet

export Metric, BinaryAccuracy, OneHotBinaryAccuracy

abstract type Metric end

"""
Binary accuracy calculation
"""
struct BinaryAccuracy
    threshold

    BinaryAccuracy(;threshold=0.5) = new(threshold)
end

function (a::BinaryAccuracy)(y_pred::Tensor, y_true::Tensor)
    y_pred = y_pred .> a.threshold
    return mean(y_true .== y_pred)
end

"""
Binary accuracy for a onehot classification.
"""
struct OneHotBinaryAccuracy end

function (a::OneHotBinaryAccuracy)(y_pred::Tensor, y_true::Tensor)
    y_pred, y_true = Knet.mat(y_pred), Knet.mat(y_true)

    y_pred = mapslices(argmax,Knet.value(Array(y_pred)),dims=1)
    y_true = Array(y_true)

    mean(y_pred .== y_true)
end

end


module ShapeFromShading

using AlgebraicMultigrid
using Blink
using Distributions
using DSP
using FFTW
using Images
using Interact
using IterativeSolvers
using LinearAlgebra
using Makie
using Optim
using Parameters
using Preconditioners
using SparseArrays
using Statistics

abstract type AbstractSyntheticShape end
@with_kw struct Cake <: AbstractSyntheticShape
    radius::Real = 50
end
@with_kw struct Cake2 <: AbstractSyntheticShape
    radius::Real = 75
end
@with_kw struct Dem <: AbstractSyntheticShape
    radius::Real = 75
end
@with_kw struct Prism <: AbstractSyntheticShape
    radius::Real = 75
end
@with_kw struct Ripple <: AbstractSyntheticShape
    radius::Real = 10
end
@with_kw struct Ripple2 <: AbstractSyntheticShape
    radius::Real = 7.5
end
@with_kw struct SynthGaussian <: AbstractSyntheticShape
    radius::Real = 75
end
@with_kw struct SuperGaussian <: AbstractSyntheticShape
    radius::Real = 100
end
@with_kw struct SynthSphere <: AbstractSyntheticShape
    radius::Real = 50
end
@with_kw struct SynthVase <: AbstractSyntheticShape
    radius::Real = 78.5
end
@with_kw struct Tent <: AbstractSyntheticShape
    radius::Real = 50
end

@doc raw"""
```
(scheme::AbstractIntegrator)(p::AbstractArray, q::AbstractArray)
```
Integrates the gradient field defined by `p` and `q` using various methods as
described in the following sections.
# Output
Returns a heightmap `Z` which defines the reconstructed surface.
# Details
Using the method defined in each type of integrator, each algorithm is run using
parameters passed to the struct when creating the integrator. Results can vary
wildly depending on the chosen integrator.
Available methods are; Frankot, Path, SplitPath, Horn, Durou, Quadratic,
TotalVariation, NonConvex1, NonConvex2, AnisotropicDiffusion and MumfordShah.
See individual documentation for each method for more detials.
# Parameters
Each algorithm has its own set of parameters which are discussed in each the section
for each method.
# Arguments
## `p` and `q`:
`AbstractArray`'s which define the gradient field to be integrated.
# Example
The following example demonstrates  the use of the `Frankot` integrator.
```julia
using ShapeFromShading, Makie

# Generate synthetic gradients
p, q = synthetic_gradient(SynthSphere(38), img_size = 151)

# Create a Frankot() integrator
frankot = Frankot()

# Calculate the heightmap from the gradients
Z = frankot(p, q)

# Normalize to maximum of 1 (not necessary but makes displaying easier)
Z = Z./maximum(Z)

# Display using Makie (Note: Makie can often take several minutes first time)
r = 0.0:0.1:4
surface(r, r, Z)
```
"""
abstract type AbstractIntegrator end
@with_kw struct AnisotropicDiffusion <: AbstractIntegrator
    z::AbstractArray
    μ::Real = 5.0
    ν::Real = 10.0
    λ::AbstractArray = fill(10.0^-6, size(z))
    mask::AbstractArray = fill(1.0, size(z))
    max_iter::Int = 10
end
@with_kw struct Durou <: AbstractIntegrator
    ϵ::Real = 1.0
    max_iter::Real = 1000
end
struct Frankot <: AbstractIntegrator end
@with_kw struct Horn <: AbstractIntegrator
    ϵ::Real = 1.0
    max_iter::Real = 10000
end
@with_kw struct MumfordShah <: AbstractIntegrator
    z::AbstractArray
    μ::Real = 10.0
    ϵ::Real = 0.1
    λ::AbstractArray = fill(10.0^-6, size(z))
    mask::AbstractArray = fill(1.0, size(z))
    max_iter::Int = 50
end
@with_kw struct NonConvex1 <: AbstractIntegrator
    z::AbstractArray
    β::Real = 0.5
    λ::AbstractArray = fill(10.0^-6, size(z))
    mask::AbstractArray = fill(1.0, size(z))
    max_iter::Int = 100
end
@with_kw struct NonConvex2 <: AbstractIntegrator
    z::AbstractArray
    γ::Real = 1.0
    λ::AbstractArray = fill(10.0^-6, size(z))
    mask::AbstractArray = fill(1.0, size(z))
    max_iter::Int = 100
end
struct Path <: AbstractIntegrator end
struct SplitPath <: AbstractIntegrator end
@with_kw struct TotalVariation <: AbstractIntegrator
    z::AbstractArray
    α::Real = 1.0
    λ::AbstractArray = fill(10.0^-6, size(z))
    mask::AbstractArray = fill(1.0, size(z))
    max_iter::Int = 100
end
@with_kw struct Quadratic <: AbstractIntegrator
    z::AbstractArray
    λ::AbstractArray = fill(10.0^-6, size(z))
    mask::AbstractArray = fill(1.0, size(z))
end

abstract type AbstractShapeAlgorithm end
@with_kw struct Deterministic <: AbstractShapeAlgorithm
     p::AbstractArray
     q::AbstractArray
     integration_factor::Real = 10
     smoothness::Real = 20
     β::Real = 1.0
     integration_scheme::AbstractIntegrator = Horn()
end
@with_kw struct DeterministicCustom <: AbstractShapeAlgorithm
    p::AbstractArray
    q::AbstractArray
    integration_factor::Real = 10
    smoothness::Real = 20
    β::Real = 1.0
    integration_scheme::AbstractIntegrator = Horn()
end
@with_kw struct DiscreteShape <: AbstractShapeAlgorithm
    iterations::Int = 2000
    smoothness::Int = 1000
    albedo::Real = Inf
    illumination_direction::Vector{T} where T <: Real = [Inf, Inf, Inf]
end
@with_kw struct DiscreteShapeBound <: AbstractShapeAlgorithm
    iterations::Int = 2000
    smoothness::Int = 1000
    albedo::Real = Inf
    illumination_direction::Vector{T} where T <: Real = [Inf, Inf, Inf]
end
@with_kw struct MultiResolutionHybrid <: AbstractShapeAlgorithm
    T₀::Real = 500.0
    α::Real = 0.99
    integration_factor::Real = 10
    smoothness::Real = 20
    integration_factor2::Real = 10
    smoothness2::Real = 20
end
@with_kw struct Pentland <: AbstractShapeAlgorithm
    slant::Real = Inf
    tilt::Real = Inf
end
@with_kw struct Photometric <: AbstractShapeAlgorithm
    illumination_direction1::Vector{T} where T <: Real = [Inf, Inf, Inf]
    illumination_direction2::Vector{T} where T <: Real = [Inf, Inf, Inf]
    illumination_direction3::Vector{T} where T <: Real = [Inf, Inf, Inf]
    integration_scheme::AbstractIntegrator = Horn()
end
@with_kw struct Shah <: AbstractShapeAlgorithm
    slant::Real = Inf
    tilt::Real = Inf
    iterations::Int = 200
end
@with_kw struct SimulatedAnnealing <: AbstractShapeAlgorithm
    pin::AbstractArray
    qin::AbstractArray
    T₀::Real = 500.0
    α::Real = 0.99
    integration_factor::Real = 10
    smoothness::Real = 20
    max_iter::Int = 40000
    debug::Bool = false
end

include("benchmark.jl")
include("common.jl")
include("estimatealbedo.jl")
include("discreteshape.jl")
include("discreteshapebound.jl")
include("deterministic.jl")
include("integration.jl")
include("multiresolutionhybrid.jl")
include("pentland.jl")
include("photometric.jl")
include("shah.jl")
include("syntheticsurface.jl")
include("simulatedannealing.jl")

export
    # main functions and datatypes
    AbstractIntegrator,
    AbstractShapeAlgorithm,
    AbstractSyntheticShape,
    AnisotropicDiffusion,
    benchmark_image,
    benchmark_integration,
    benchmark_normals,
    benchmark_shape,
    Cake,
    Cake2,
    convert_gradient,
    create_image,
    Cup,
    Dem,
    Deterministic,
    DeterministicCustom,
    DiscreteShape,
    DiscreteShapeBound,
    Durou,
    estimate_img_properties,
    Frankot,
    gen_mask,
    generate_photometric,
    generate_surface,
    generate_surface_data,
    ground_truth,
    Horn,
    MultiResolutionHybrid,
    MumfordShah,
    NonConvex1,
    NonConvex2,
    Path,
    Pentland,
    Photometric,
    Prism,
    retrieve_surface,
    Ripple,
    Ripple2,
    Shah,
    SimulatedAnnealing,
    SplitPath,
    SuperGaussian,
    synthetic_gradient,
    SynthGaussian,
    SynthSphere,
    SynthVase,
    Tent,
    TotalVariation,
    Quadratic,
    MumfordShah2
end # module


@testset "minimal_sum_test" begin
    # sum_nearest {1.0, 2.0, 3.0} = 6.0
    # sum_nearest {1.0, 2.0, NaN, 3.0} = NaN
    # sum_nearest {1.0, -infinity, 2.0, infinity, 3.0} = NaN
end

@testset "minimal_sum_abs_test" begin
    # sum_abs_nearest {1.0, -2.0, 3.0} = 6.0
    # sum_abs_nearest {1.0, -2.0, NaN, 3.0} = NaN
    # sum_abs_nearest {1.0, -infinity, 2.0, infinity, 3.0} = infinity
end

@testset "minimal_sum_sqr_test" begin
    # sum_sqr_nearest {1.0, 2.0, 3.0} = 14.0
    # sum_sqr_nearest {1.0, 2.0, NaN, 3.0} = NaN
    # sum_sqr_nearest {1.0, -infinity, 2.0, infinity, 3.0} = infinity
end

@testset "minimal_dot_test" begin
    # dot_nearest {1.0, 2.0, 3.0} {1.0, 2.0, 3.0} = 14.0
    # dot_nearest {0x10000000000001p0, 0x1p104} {0x0fffffffffffffp0, -1.0} = -1.0
    # dot_nearest {1.0, 2.0, NaN, 3.0} {1.0, 2.0, 3.0, 4.0} = NaN
    # dot_nearest {1.0, 2.0, 3.0, 4.0} {1.0, 2.0, NaN, 3.0} = NaN
    # dot_nearest {1.0, 2.0, 0.0, 4.0} {1.0, 2.0, infinity, 3.0} = NaN
    # dot_nearest {1.0, 2.0, -infinity, 4.0} {1.0, 2.0, 0.0, 3.0} = NaN
end


using MDCDL
using Images, TestImages
using ImageFiltering
using Plots
using LinearAlgebra

function main()
    ### configurations ###
    # dirNsolt = joinpath(Pkg.dir(),"MDCDL","examples","design","sample.json")
    dirNsolt = joinpath(pathof(MDCDL), "..","..","examples", "design", "sample.json")
    dirNsolt = abspath(dirNsolt)
    σ_noise = 1e-2 # variance of AWGN
    λ = 1e-2 # FISTA parameter
    lv = 3 # tree level of NSOLT
    ######################

    ### setup observed image
    # orgImg = testimage("cameraman")
    # u = Array{Float64}(imresize(orgImg,(256,256)))
    orgImg = testimage("lena")
    u = Array{RGB{Float64}}(orgImg)

    # AWGN
    w = mapc.((v)->σ_noise * randn(), u)

    # Degradation (pixel loss)
    P = map((v)-> v > 0.01 ? 1 : 0, rand(Float64,size(u)))

    # Observation
    x = P .* u + w

    ### setup NSOLT
    nsolt = MDCDL.load(dirNsolt)
    println("NSOLT configurations:")
    println(" - Type = $(typeof(nsolt))")
    println(" - Decimation Factor = $(nsolt.decimationFactor)")
    println(" - Number of Channels = $(nsolt.nChannels)")
    println(" - Polyphase Order = $(nsolt.polyphaseOrder)")
    ### show analysis filters

    # setup Multiscale NSOLT
    analyzer = createMultiscaleAnalyzer(nsolt, size(x); level=lv, shape=:vector)
    synthesizer = createMultiscaleSynthesizer(nsolt, size(x); level=lv, shape=:vector)

    # setup FISTA
    y0 = analyze(analyzer, x)

    ∇f(ty) = begin
        - analyze(analyzer, P.*(x - synthesize(synthesizer, ty)))
    end
    prox(ty, η) = MDCDL.softshrink(ty, λ*η)
    viewFcn(itrs, ty, err) = begin
        println("# Iterations=$itrs, error=$err")
    end

    hy = MDCDL.fista(∇f, prox, y0; η=1.0, iterations=400, verboseFunction=viewFcn, absTol=eps())

    # restored image
    ru = synthesize(synthesizer, hy)

    psnr = (a, ref) -> 10*log10(1/(norm(a-ref)^2/length(a)))
    errx = norm(ru - u)

    println("error: $errx")

    plotOrg = plot(u ; xlabel=string("Original Image"), aspect_ratio=:equal);
    plotRes = plot(ru; xlabel=string("Restored Image\nPSNR=",trunc(psnr(ru,u),digits=3)), aspect_ratio=:equal);
    plot(plotOrg, plotObs, plotRes; aspect_ratio=:equal, layout=(1,3))

end

main()


# Local projection
# Example based on Kilian and Kim 2011 dataset

#-----------Load packages------------------------------------------
using VectorAutoregressions
using DelimitedFiles, LinearAlgebra, Statistics, GrowableArrays

#-----------set base-path------------------------------------------
path = joinpath(dirname(pathof(VectorAutoregressions)), "..")

#-----------Load data----------------------------------------------
y      = readdlm(joinpath(path,"test","lp_data.csv"),',')
irfv   = readdlm(joinpath(path,"test","lp_test_var_irf.csv"),',')
irfl   = readdlm(joinpath(path,"test","lp_test_lp_chol_irf.csv"),',')
stdv   = readdlm(joinpath(path,"test","lp_test_var_stdv.csv"),',')
stdl   = readdlm(joinpath(path,"test","lp_test_lp_std.csv"),',')
COVsig = readdlm(joinpath(path,"test","lp_test_var_covsig.csv"),',')
rfirfl = readdlm(joinpath(path,"test","lp_test_lp_rf_irf.csv"),',')

#-----------Hyperparameter-----------------------------------------
const pbar = 12 # max order of lag to test
const H = 24    # horizon
const intercept = true

#-----------Lag-length selection for local projection and red.form IRFs-------------------
#Ex. 1
p = lp_lagorder(y,pbar,H,"aic") # using aic selection criteria
mIRFs = IRFs_localprojection(y, p, H)
#Ex. 2
p = lp_lagorder(y,pbar,H,"bic") # using bic selection criteria
mIRFs = IRFs_localprojection(y, p, H)
#Ex. 3
p = lp_lagorder(y,pbar,H,"aicc") # using aic correced selection criteria
mIRFs = IRFs_localprojection(y, p, H)
#Ex. 4
p = lp_lagorder(y,pbar,H,"hqc") # using Hannan-Quinn selection criteria
mIRFs = IRFs_localprojection(y, p, H)
#Ex. 5
p = 12                                       # fixing the lag-length for all projections
mIRFs = IRFs_localprojection(y, p, H)

#-----------Structural local projection IRFs-------------------
#Ex. 1 -- using a VAR(pbar) as auxiliary model for cholesky identification
p_var = 12
V = VAR(y,p_var,true)
A0inv = V.Σ |> λ -> Array((cholesky(λ)).L)
mStd,mCov_Σ = irf_ci_asymptotic(V, H, V.inter)

mIRFs,_ = IRFs_localprojection(y, p, H, A0inv, mCov_Σ)
IRF,CI = mIRFs.IRF, mIRFs.CI
CIl,CIh = CI.CIl, CI.CIh


function split_adjacency(A)
    graphs = SimpleGraph{Int64}[]

    subinds = connected_components(Graph(A))
    for subset in subinds
        subA = zeros(Int64, size(A)...)
        subA[subset, subset] = A[subset, subset]
        push!(graphs, Graph(subA))
    end

    roots = [subinds[i][1] for i=1:length(subinds)]

    return graphs, roots
end


export intersection
using Dierckx: Spline1D

"""
get the intersection point from direction specified by theta.

Args:
    line (2darray): an array of points.
    theta (float): direction of the intersection line.
    align (len-2 tuple): align to this point in the free dimension.

Returns:
    tuple: the nearest intersection point.
"""
function intersection(line::Vector{PointT}, theta::Real, align::PointT)
    trans(dot, theta) = Base.reverse(rot(dot..., -theta))
    trans_r(dot, theta) = rot(dot[2], dot[1], theta)

    # rotate to y-axis
    rotated_line = trans.(line, theta)
    rotated_xy = trans(align, theta)

    # get segments
    p_pre = rotated_line[1]
    segment = [p_pre]
    segments = Vector{PointT}[]
    direction_pre = 0
    for i = 2:size(rotated_line, 1)
        p = rotated_line[i]
        direction = p[1] > p_pre[1] ? 1 : -1
        if direction * direction_pre < 0
            push!(segments, segment)
            segment = [p_pre, p]
        else
            push!(segment, p)
        end
        direction_pre = direction
        p_pre = p
    end
    push!(segments, segment)

    # intersect segments
    x = rotated_xy[1]
    ys = []
    for segment in segments
        x_ = getindex.(segment, 1)
        y_ = getindex.(segment, 2)
        try
            f = Spline1D(x_, y_; k=1)
            push!(ys, f(x))
        catch e
        end
    end

    # interpolate x
    if isempty(ys)
        error("Can not find connection point!")
    end
    index = argmax(ys)
    trans_r((x, ys[index]), theta)
end


using YaoLang
using YaoLang.Compiler
using YaoLang.Compiler: device_m, enable_timings, timings

# enable_timings()

ex = :(function qft(n::Int)
    1 => H
    for k in 2:n
        @ctrl k 1 => shift(2π / 2^k)
    end

    if n > 1
        2:n => qft(n - 1)
    end
end)

@time device_m(@__MODULE__, ex);

# timings()


using IntervalWavelets
using Test

@testset "Dyadic Rationals" begin
    @testset "Construct Dyadic Rationals" begin
        x = DyadicRational(1, 2)

        @test IntervalWavelets.numerator(x) == 1
        @test resolution(x) == 2

        @test DyadicRational(4, 1) == DyadicRational(2, 0)
        @test DyadicRational(-4, 1) == DyadicRational(-2, 0)
        @test DyadicRational(5, 1) == DyadicRational(5, 1)

        @test_throws DomainError DyadicRational(1, -2)
    end

    @testset "Arithmetic with Dyadic Rationals" begin
        @test DyadicRational(2, 0) + 1 == DyadicRational(3, 0)
        @test DyadicRational(2, 1) + 1 == DyadicRational(4, 1)

        @test DyadicRational(1, 0) + DyadicRational(1, 1) == DyadicRational(3, 1)
        @test DyadicRational(1, 0) - DyadicRational(1, 1) == DyadicRational(1, 1)
        @test DyadicRational(1, 1) - DyadicRational(-1, 1) == DyadicRational(1, 0)

        @test 2 * DyadicRational(2, 0) == DyadicRational(4, 0)
        @test 3 * DyadicRational(1, 2) == DyadicRational(3, 2)
        @test DyadicRational(-3, 1) * DyadicRational(-1, 2) == DyadicRational(3, 3)
    end

    @testset "All Dyadic Rationals" begin
        @test IntervalWavelets.all_dyadic_rationals(0, 1, 0) == DyadicRational.([0 ; 1], 0)
        @test IntervalWavelets.all_dyadic_rationals(0, 1, 1) == DyadicRational.(0:2, 1)

        left = DyadicRational(-1, 1)
        right = DyadicRational(1, 2)
        @test all_dyadic_rationals(left, right, 2) == DyadicRational.(-2:1, 2)
        @test length(all_dyadic_rationals(right, left, 2)) == 0

        @test_throws DomainError IntervalWavelets.all_dyadic_rationals(left, right, 1)
    end

    @testset "Convert Dyadic Rational" begin
        @test Integer(DyadicRational(3, 0)) == 3
        @test_throws InexactError Integer(DyadicRational(3, 2))

        @test float(DyadicRational(3, 0)) == 3.0
        @test float(DyadicRational(3, 2)) == 0.75
    end

    @testset "Compare Dyadic Rationals" begin
        @test DyadicRational(1, 2) < DyadicRational(2, 2)
        @test DyadicRational(1, 2) > DyadicRational(-1, 2)

        @test DyadicRational(1, 2) >= DyadicRational(1, 2)
        @test DyadicRational(1, 2) <= DyadicRational(1, 2)
    end
end


using Base.Test
using RiemannTheta
reload("RiemannTheta")

g, num_vectors = 5, 10
z = [ rand(ComplexF64, g) for i in 1:num_vectors]
mode, ϵ = 0, 1e-6
accuracy_radius = 5.
tttt = rand(5,5)
Y = tttt * tttt'
Ω = Complex.(rand(g, g), Y)

derivs = Vector{ComplexF64}[]
oscillatory_part(z, Ω, 0, ϵ, derivs, accuracy_radius)
@btime oscillatory_part($z, $Ω, 0, $ϵ, $derivs, $accuracy_radius) # 380-390ms / 430 Mb

derivs = [ rand(ComplexF64, 5) for i in 1:5 ]
@btime oscillatory_part($z, $Ω, 0, $ϵ, $derivs, $accuracy_radius) # 6.9s / 5.6Gb

@testset "oscillatory_part derivs are correct" begin
    z₀ = rand(ComplexF64, g) - Complex(0.5, 0.5)
    δ = [Complex(1e-8,1e-8); zeros(ComplexF64, g-1)]
    derivs = Vector{ComplexF64}[]
    ϵ = 1e-8
    # calculate function at slightly shifted z₀
    z = [[z₀] ; [z₀ + circshift(δ, i) for i in 0:g-1]]
    res = oscillatory_part(z, Ω, 0, ϵ, derivs, accuracy_radius)
    dres = [ res[i] - res[1] for i in 2:g+1 ]
    dres

    # and compare to calculated derivates
    for i in 0:g-1
        derivs = [ circshift(δ, i) ]
        res2 = oscillatory_part([z₀], Ω, 0, ϵ, derivs, accuracy_radius)
        @test isapprox(res2[1], dres[i+1], rtol=1e-4)
    end
end


#=
Contains functions/type prototypes for the mesh.
=#

#= Mesh =#
#Subtypes should contain all variables that define a mesh.
#Should make an effort to follow the naming conventions established
#in existing methods in order to give compatibility between methods
abstract type AbstractMesh end
mutable struct Mesh{T} <: AbstractMesh
    G::Matrix{T}
    D::Matrix{T}
    l::Int
    Δᵐ::T
    Δᵖ::T

    # Override constructor for different default meshes for
    # different poll techniques.
    Mesh(N::Int64) = Mesh{Float64}(N)
    function Mesh{T}(N::Int64) where T
        mesh = new()
        mesh.l = 0
        mesh.Δᵐ = min(1, 4.0^(-mesh.l))
        mesh.Δᵖ = 2.0^(-mesh.l)
        mesh.G = Matrix(I,N,N)
        mesh.D = hcat(Matrix(I,N,N),-Matrix(I,N,N))
        return mesh
    end
end

"""
    MeshUpdate!(mesh::Mesh, improvement_found::Bool)

Implements LTMADS update rule from Audet & Dennis 2006 pg. 203 adapted for progressive
barrier constraints with Audet & Dennis 2009 expression 2.4.
"""
function MeshUpdate!(m::Mesh, ::AbstractPoll, result::IterationOutcome)
    if result == Unsuccessful
        m.l += 1
    elseif result == Dominating && m.l > 0
        m.l -= 1
    elseif result == Improving
        m.l == m.l
    end

    m.Δᵐ = min(1, 4.0^(-m.l))
    m.Δᵖ = 2.0^(-m.l)
end