"""
$(SIGNATURES)

"""
function engineer_caustics(source_image)
    imageBW = Float64.(Gray.(source_image))

    # Set global size parameters
    height, width = size(imageBW)
    global N_Pixel_Height = height
    global N_Pixel_Width = width
    global Meters_Per_Pixel = Caustics_Long_Side / N_Pixel_Height
    println("Image size: $(N_Pixel_Height) x $(N_Pixel_Width)")

    # mesh is the same size as the image with an extra row/column to have coordinates to
    # cover each image corner with a triangle.
    mesh = FaceMesh(N_Pixel_Height, N_Pixel_Width)

    # imageBW is normalised to have 1 unit per pixel on average.
    imageBW /= average(imageBW)

    marginal_change_top = marginal_change_bot = nothing
    max_update_top = marginal_change_top = Inf
    counter = 0
    while (abs(max_update_top+marginal_change_bot) > 1e-6 && counter < 10_000)
        counter += 1

        print(
            """
          ================================================================================================================
          STARTING ITERATION $(counter):
              starting ϕ TOP field = $(field_summary(mesh.corners.ϕ_top))
              starting ϕ BOT field = $(field_summary(mesh.corners.ϕ_bot))

            """,
        )

        ϕ_top = mesh.corners.ϕ_top
        ε_top, max_update_top = solve_velocity_potential!(mesh.corners.ϕ_top, imageBW/2., "it$(counter)_top")
        marginal_change_top = mesh.corners.ϕ_top - ϕ_top

        ϕ_bot = mesh.corners.ϕ_bot
        ε_bot, max_update_bot = solve_velocity_potential!(mesh.corners.ϕ_bot, imageBW/2., "it$(counter)_bot")
        marginal_change_bot = mesh.corners.ϕ_bot - ϕ_bot

        print(
            """

            RESULT AT ITERATION $(counter):
                TOP
                Luminosity error = $(field_summary(ε_top))
                Vertical move max update = $(max_update_top)
                Marginal change = $(field_summary(marginal_change_top))
                end ϕ field = $(field_summary(mesh.corners.ϕ_top))

                BOT
                end ϕ field = $(field_summary(mesh.corners.ϕ_bot))
                Luminosity error = $(field_summary(ε_bot))
                Vertical move max update = $(max_update_bot)
                Marginal change = $(field_summary(marginal_change_bot))
                end ϕ field = $(field_summary(mesh.corners.ϕ_bot))

            ================================================================================================================

          """,
        )
        plot_as_quiver(mesh.corners.ϕ_top, n_steps = 50, scale = height / 10, max_length = height / 20)
        average_absolute(marginal_change_top + marginal_change_bot) < 1e-4 && break
    end

    println("\nSTARTING HORIZONTAL ITERATION ---")

    println("\t Horizontal max update TOP = $(max_update_top) /  BOT = $(max_update_bot)")

    # Move the around a nil average.
    mesh.corners.ϕ_top .-= average(mesh.corners.ϕ_top)
    mesh.corners.ϕ_bot .-= average(mesh.corners.ϕ_bot)

    return mesh, imageBW
end

"""
$(SIGNATURES)

ϕ is the _velocity potential_. The velocity represents the direction towards which the mesh vertices should move.
Zones of high luminosity should be covered with more triangles and should attract more vertices.
"""
function solve_velocity_potential!(mesh, image, prefix)
    # Start with a new clean field
    height, width = size(mesh)

    # Illumination only depends on the position of the corners, not their heights. ϕ is not relevant.
    # Get the area of each individual pixel as stretch/shrunk on the lens. Area = energy.
    # _FENCES_SIZED_
    # Illumination only depends on the position of the corners, not their heights. ϕ is not relevant.
    lens_pixels_area_top, lens_pixels_area_bot = get_lens_pixels_area(mesh)

    # Positive error => the triangle needs to shrink (less light). Enforce nil average error.
    ε_top = zeros(Float64, height, width)
    ε_bot = zeros(Float64, height, width)

    # Each triangle should have the luminosity of half a pixel
    ε_top[1:end-1, 1:end-1] = Float64.(lens_pixels_area_top - image/2.)
    ε_top .-= average(ε_top)

    ε_bot[1:end-1, 1:end-1] = Float64.(lens_pixels_area_bot - image/2.)
    ε_bot .-= average(ε_bot)

    println("""
            solve_velocity_potential! before loop:
                $(field_summary(lens_pixels_area_top, "Pixel area"))
                $(field_summary(ε_top, "luminosity error"))
                $(field_summary(lens_pixels_area_bot, "Pixel area"))
                $(field_summary(ε_bot, "luminosity error"))
                """)

    # Save the loss image as a png.
    save_plot_scalar_field!(ε_top+ε_bot, "error_ε_$(prefix)")
    save_plot_scalar_field!(ε_top, "error_εtop_$(prefix)")
    save_plot_scalar_field!(ε_bot, "error_εbot_$(prefix)")

    # Save the generated image as a png.
    lens_pixels_area = lens_pixels_area_top + lens_pixels_area_bot
    save(
        "./examples/img_noluminosityratio_$(prefix).png",
        Gray.(clamp.(lens_pixels_area, 0.0, 1.0)),
    )

    luminosity_ratio = sum(image) / sum(lens_pixels_area)
    save(
        "./examples/img_$(prefix).png",
        Gray.(clamp.(lens_pixels_area * luminosity_ratio, 0.0, 1.0)),
    )

    luminosity_ratio = sum(image/2.) / sum(lens_pixels_area_top)
    save(
        "./examples/img_$(prefix)_top.png",
        Gray.(clamp.(lens_pixels_area_top * luminosity_ratio, 0.0, 1.0)),
    )

    luminosity_ratio = sum(image/2.) / sum(lens_pixels_area_bot)
    save(
        "./examples/img_$(prefix)_bot.png",
        Gray.(clamp.(lens_pixels_area * lens_pixels_area_bot, 0.0, 1.0)),
    )

    # Start with a clean, flat potential field.
    mesh.corners.ϕ_top = zeros(Float64, height, width)
    mesh.corners.ϕ_bot = zeros(Float64, height, width)

    ϕ_b4_top = copy(mesh.corners.ϕ_top)
    ϕ_b4_bot = copy(mesh.corners.ϕ_bot)

    count = 0
    min_update_top = min_update_bot = new_update_top = new_update_bot = 10_000
    old_update_top = 2 * new_update_top
    old_update_bot = 2 * new_update_bot
    while (
        1e-5 < new_update_top + new_update_bot < 1.5 * min_update_top + 1.5 * min_update_bot &&
        1e-4 < (old_update_top - new_update_top) / old_update_top &&
        1e-4 < (old_update_bot - new_update_bot) / old_update_bot &&
        count < 10_000
    )
        count += 1
        old_update_top = new_update_top
        old_update_bot = new_update_bot
        min_update_top = min(min_update_top, new_update_top)
        min_update_bot = min(min_update_bot, new_update_bot)

        new_update_top, ∇²ϕ_est_top, δ_top = propagate_poisson!(mesh.corners.ϕ_top,ε_top)
        new_update_bot, ∇²ϕ_est_bot, δ_bot = propagate_poisson!(mesh.corners.ϕ_bot,ε_bot)

        count % 1_000 == 1 && println("""
                                        TOP
                                        Iteration $(count) - max_update = $(round(new_update_top, sigdigits=4))
                                            $(field_summary(mesh.corners.ϕ_top, "ϕ"))
                                            $(field_summary(∇²ϕ_est_top, "∇²ϕ_est"))
                                            $(field_summary(δ_top, "δ"))
                                        BOTTOM
                                        Iteration $(count) - max_update = $(round(new_update_bot, sigdigits=4))
                                            $(field_summary(mesh.corners.ϕ_bot, "ϕ"))
                                            $(field_summary(∇²ϕ_est_bot, "∇²ϕ_est"))
                                            $(field_summary(δ_bot, "δ"))
                                            """)

        # Save the loss image as a png.
        if count % 1_000 == 1
            save_plot_scalar_field!(∇²ϕ_est_top, "∇²ϕ_est_$(prefix)")
            save_plot_scalar_field!(δ_top, "delta_$(prefix)")
            save_plot_scalar_field!(
                mesh.corners.ϕ_bot + mesh.corners.ϕ_top -ϕ_b4_top -ϕ_b4_bot ,
                "change_mesh.corners.ϕ_$(prefix)",
            )

            save_plot_scalar_field!(∇²ϕ_est_top, "∇²ϕ_est_$(prefix)_top")
            save_plot_scalar_field!(δ_top, "delta_$(prefix)_top")
            save_plot_scalar_field!(
                mesh.corners.ϕ_top - ϕ_b4_top,
                "change_mesh.corners.ϕ_$(prefix)_top",
            )

            save_plot_scalar_field!(∇²ϕ_est_bot, "∇²ϕ_est_$(prefix)_bot")
            save_plot_scalar_field!(δ_bot, "delta_$(prefix)_bot")
            save_plot_scalar_field!(
                mesh.corners.ϕ_bot - ϕ_b4_bot,
                "change_mesh.corners.ϕ_$(prefix)_bot",
            )

            ϕ_b4_top = copy(mesh.corners.ϕ_top)
            ϕ_b4_bot = copy(mesh.corners.ϕ_bot)
        end
    end

    # Now we need to march the mesh row,col corner locations according to this gradient.
    δ_top = march_mesh!(mesh, :top)
    δ_bot = march_mesh!(mesh, :bottom)
    return ε, new_update
end

"""
$(SIGNATURES)

ϕ is the height map.

`march_mesh` flexes the mesh
"""
function march_mesh(mesh, mesh_ϕ)

    # Calculate the gradient of the velocity potential and reverse its direction.
    vr, vc = ∇(mesh_ϕ)
    mesh.corners.vr .= -vr
    mesh.corners.vc .= -vc

    # Clean up the mesh borders
    # reset_border_values!(mesh.corners)

    # For each point in the mesh we need to figure out its velocity
    # However all the nodes located at a border will never move
    # I.e. velocity (Vx, Vy) = (0, 0) and the square of acrylate will remain of the same size.

    # Get the time, at that velocity, for the area of the triangle to be nil.
    # We are only interested by positive values to only move in the direction of the (opposite) gradient
    height, width = size(mesh)
    list_triangles = vcat(
        [triangle3D(mesh, row, col, :top) for row ∈ 1:height, col ∈ 1:width],
        [triangle3D(mesh, row, col, :bottom) for row ∈ 1:height, col ∈ 1:width],
    )

    list_maximum_t = [t for t ∈ find_maximum_t.(list_triangles) if !isnothing(t) && t > 0.0]

    δ = minimum(list_maximum_t) / 1.5
    new_r = mesh.corners.r - δ * vr
    new_c = mesh.corners.c - δ * vc

    return new_r, new_c, δ
end

function march_mesh!(mesh::FaceMesh, side::Union{:top, :bottom})

    # Calculate the gradient of the velocity potential and reverse its direction.
    if side == :top
        mesh.corners.vr, mesh.corners.vc = ∇(mesh.corners.ϕ_top)
    elseif side == :bottom
        mesh.corners.vr, mesh.corners.vc = ∇(mesh.corners.ϕ_bot)
    end
    mesh.corners.vr .*= -1.0
    mesh.corners.vc .*= -1.0

    # Clean up the mesh borders
    reset_border_values!(mesh.corners)

    # For each point in the mesh we need to figure out its velocity
    # However all the nodes located at a border will never move
    # I.e. velocity (Vx, Vy) = (0, 0) and the square of acrylate will remain of the same size.
    mesh_r = copy(mesh.corners.r)
    mesh_c = copy(mesh.corners.c)

    # Get the time, at that velocity, for the area of the triangle to be nil.
    # We are only interested by positive values to only move in the direction of the (opposite) gradient
    height, width = size(mesh)
    list_triangles = [triangle3D(mesh, row, col, side) for row ∈ 1:height, col ∈ 1:width]
    list_maximum_t = [t for t ∈ find_maximum_t.(list_triangles) if !isnothing(t) && t > 0.0]

    δ = minimum(list_maximum_t) / 2.0
    if side == :top
        mesh.corners.r[1:end-1, 1:end-1] .-= δ * mesh.corners.vr[1:end-1, 1:end-1]
        mesh.corners.c[1:end-1, 1:end-1] .-= δ * mesh.corners.vc[1:end-1, 1:end-1]
    elseif
        mesh.corners.r[2:end, 2:end] .-= δ * mesh.corners.vr[1:end-1, 1:end-1]
        mesh.corners.c[2:end, 2:end] .-= δ * mesh.corners.vc[1:end-1, 1:end-1]
    end

    println(
        """

        March mesh with correction_ratio δ = $(δ)
            $(field_summary(mesh.corners.vr, "∇u"))
            $(field_summary(mesh.corners.vc, "∇v"))
            $(field_summary(mesh_r - mesh.corners.r , "new mesh changes on row"))
            $(field_summary(mesh_c - mesh.corners.c, "new mesh changes on col"))
            $(field_summary(mesh_r - mesh.corners.rows_numbers , "total mesh changes on row"))
            $(field_summary(mesh_c - mesh.corners.cols_numbers, "total mesh changes on col"))
            """,
    )

    return δ
end

"""
$(SIGNATURES)

This function implements successive over relaxation for a matrix and its associated error matrix
There is a hardcoded assumption of Neumann boundary conditions -- that the derivative across the
boundary must be zero in all cases. See:
https://math.stackexchange.com/questions/3790299/how-to-iteratively-solve-poissons-equation-with-no-boundary-conditions

- ϕ is the potential to be solved subject to the Poisson equation and is the same size as the corners (_POSTS_SIZED_)
- target is ∇ϕ and is of the size in pixels (_FENCES_SIZED_)

In order of how to think about the flow of the program for the velocity potential,
    - if δ > 0, the Laplacian is too high at that point (∇²ϕ is too low).
    - To decrease the Laplacian ∇²ϕ, note that ∇² is a divergence (of a gradient) that represents net flow going outside.

    - Decrease the Laplacian
    - => decrease a divergence
    - => decrease flow towards outside

    - The flow is a gradient of the potential ϕ.
    - => decrease flow towards outside
    - => either decrease the values of the potential around that point, or increase the value of the potential at that point.
    - => Increase the value of ϕ at that point means a positive change when δ > 0.
    - δ > 0 => correction of the SAME sign.
"""
function propagate_poisson(ϕ::Matrix{Float64}, ε::Matrix{Float64})

    height, width = size(ϕ)

    ∇²ϕ_est = zeros(Float64, height, width)
    δ = zeros(Float64, height, width)
    new_ϕ = ϕ

    ω = 1.99

    for row = 1:height, col = 1:width
        val_up = row == 1 ? 0.0 : new_ϕ[row-1, col]
        val_down = row == height ? 0.0 : new_ϕ[row+1, col]
        val_left = col == 1 ? 0.0 : new_ϕ[row, col-1]
        val_right = col == width ? 0.0 : new_ϕ[row, col+1]

        ∇²ϕ_est[row, col] = val_up + val_down + val_left + val_right - 4 * new_ϕ[row, col]
        δ[row, col] = ω / 4.0 * (∇²ϕ_est[row, col] - ε[row, col])
        new_ϕ[row, col] += δ[row, col]
    end

    return new_ϕ, maximum(abs.(δ))
end

"""
$(SIGNATURES)
"""
function solve_horizontal_dispersion(mesh::FaceMesh, image; f = Focal_Length)
    # height indexes rows, width indexes columns
    height, width = size(mesh)

    # Coordinates on the caustics = simple rectangular values in corners.
    # Coordinates on the lens face comes from corner field.
    # d_row/d_col are _POSTS-SIZED_
    d_row = mesh.corners.rows_numbers - mesh.corners.r
    d_col = mesh.corners.cols_numbers - mesh.corners.c

    # The focal length is in meters. H is in corners.
    H = f / Meters_Per_Pixel

    # true_H is the real distance beteen the surface of the lens and the projection
    # H is the initial distance from the surface of the carved face to the projection screen.
    # H is constant and does not account for the change in heights due to carving.
    # Note: Higher location means closer to the caustics => negative sign
    # true_H is _POSTS_SIZED_
    true_H = zeros(Float64, height + 1, width + 1)
    true_H = -mesh.corners.ϕ .+ H

    # Normals.
    # N_row/N_col are _POSTS_SIZED_
    N_row = zeros(Float64, size(true_H))
    N_col = zeros(Float64, size(true_H))
    N_row = tan.(atan.(d_row ./ true_H) / (n₁ - n₂))
    N_col = tan.(atan.(d_col ./ true_H) / (n₁ - n₂))

    # We need to find the divergence of the Vector field described by Nx and Ny
    # div_row/div_col are _FENCES_SIZED_
    div_row = zeros(Float64, size(true_H))
    div_col = zeros(Float64, size(true_H))
    div_row[1:end-1, 1:end-1] = N_row[2:end, 1:end-1] .- N_row[1:end-1, 1:end-1]
    div_col[1:end-1, 1:end-1] = N_col[1:end-1, 2:end] .- N_col[1:end-1, 1:end-1]

    # Normalised divergence
    # divergence_direction is _FENCES_SIZED_
    divergence_direction = div_row + div_col

    ϕ_next = mesh.corners.ϕ
    max_update = Inf
    for i ∈ 1:1e6
        ϕ_next, max_update = propagate_poisson(ϕ_next, -divergence_direction)

        i % 250 == 1 &&
            println("Convergence horiz. Δ  at counter = $(i) max update = $(max_update)")

        if abs(max_update) <= 1e-6
            println(
                "Convergence horiz. Δ stopped at counter = $(i) with max_update of $(max_update)",
            )
            break
        end
    end

    return ϕ_next, max_update
end


using LinearAlgebra

"""
    SGD(;lr=0.001,gclip=0)
    update!(w,g,p::SGD)
    update!(w,g;lr=0.001)

Container for parameters of the Stochastic gradient descent (SGD)
optimization algorithm used by [`update!`](@ref).

SGD is an optimization technique to minimize an objective function by
updating its weights in the opposite direction of their gradient. The
learning rate (lr) determines the size of the step.  SGD updates the
weights with the following formula:

    w = w - lr * g

where `w` is a weight array, `g` is the gradient of the loss function
w.r.t `w` and `lr` is the learning rate.

If `norm(g) > gclip > 0`, `g` is scaled so that its norm is equal
to `gclip`.  If `gclip==0` no scaling takes place.

SGD is used by default if no algorithm is specified in the two
argument version of `update!`[@ref].
"""
mutable struct SGD
    lr::AbstractFloat
    gclip::AbstractFloat
end

const SGDLR = 0.001

SGD(; lr=SGDLR, gclip=0) = SGD(lr,gclip)

@deprecate Sgd SGD

"""
    Momentum(;lr=0.001, gclip=0, gamma=0.9)
    update!(w,g,p::Momentum)

Container for parameters of the Momentum optimization algorithm used
by [`update!`](@ref).

The Momentum method tries to accelerate SGD by adding a velocity term
to the update.  This also decreases the oscillation between successive
steps. It updates the weights with the following formulas:

    velocity = gamma * velocity + lr * g
    w = w - velocity

where `w` is a weight array, `g` is the gradient of the objective
function w.r.t `w`, `lr` is the learning rate, `gamma` is the momentum
parameter, `velocity` is an array with the same size and type of `w`
and holds the accelerated gradients.

If `norm(g) > gclip > 0`, `g` is scaled so that its norm is equal
to `gclip`.  If `gclip==0` no scaling takes place.

Reference: [Qian,
N. (1999)](http://doi.org/10.1016/S0893-6080(98)00116-6). On the
momentum term in gradient descent learning algorithms.  Neural
Networks : The Official Journal of the International Neural Network
Society, 12(1), 145–151.

"""
mutable struct Momentum
    lr::AbstractFloat
    gclip::AbstractFloat
    gamma::AbstractFloat
    velocity
end

Momentum(; lr=0.001, gclip=0, gamma=0.9)=Momentum(lr, gclip, gamma, nothing)

"""
    Nesterov(; lr=0.001, gclip=0, gamma=0.9)
    update!(w,g,p::Momentum)

Container for parameters of Nesterov's momentum optimization algorithm used
by [`update!`](@ref).

It is similar to standard [`Momentum`](@ref) but with a slightly different update
rule:

    velocity = gamma * velocity_old - lr * g
    w = w_old - velocity_old + (1+gamma) * velocity

where `w` is a weight array, `g` is the gradient of the objective
function w.r.t `w`, `lr` is the learning rate, `gamma` is the momentum
parameter, `velocity` is an array with the same size and type of `w`
and holds the accelerated gradients.

If `norm(g) > gclip > 0`, `g` is scaled so that its norm is equal
to `gclip`.  If `gclip == 0` no scaling takes place.

Reference Implementation : [Yoshua Bengio, Nicolas Boulanger-Lewandowski and Razvan P
ascanu](https://arxiv.org/pdf/1212.0901.pdf)
"""
mutable struct Nesterov
    lr::AbstractFloat
    gclip::AbstractFloat
    gamma::AbstractFloat
    velocity
end

Nesterov(; lr=0.001, gclip=0, gamma=0.9) = Nesterov(lr, gclip, gamma, nothing)

"""
    Adagrad(;lr=0.1, gclip=0, eps=1e-6)
    update!(w,g,p::Adagrad)

Container for parameters of the Adagrad optimization algorithm used by
[`update!`](@ref).

Adagrad is one of the methods that adapts the learning rate to each of
the weights.  It stores the sum of the squares of the gradients to
scale the learning rate.  The learning rate is adapted for each weight
by the value of current gradient divided by the accumulated
gradients. Hence, the learning rate is greater for the parameters
where the accumulated gradients are small and the learning rate is
small if the accumulated gradients are large. It updates the weights
with the following formulas:

    G = G + g .^ 2
    w = w - g .* lr ./ sqrt(G + eps)

where `w` is the weight, `g` is the gradient of the objective function
w.r.t `w`, `lr` is the learning rate, `G` is an array with the same
size and type of `w` and holds the sum of the squares of the
gradients. `eps` is a small constant to prevent a zero value in the
denominator.

If `norm(g) > gclip > 0`, `g` is scaled so that its norm is equal
to `gclip`.  If `gclip==0` no scaling takes place.

Reference: [Duchi, J., Hazan, E., & Singer,
Y. (2011)](http://jmlr.org/papers/v12/duchi11a.html). Adaptive
Subgradient Methods for Online Learning and Stochastic Optimization.
Journal of Machine Learning Research, 12, 2121–2159.

"""
mutable struct Adagrad
    lr::AbstractFloat
    gclip::AbstractFloat
    eps::AbstractFloat
    G
end

Adagrad(; lr=0.1, gclip=0, eps=1e-6)=Adagrad(lr, gclip, eps, nothing)

"""
    Adadelta(;lr=0.01, gclip=0, rho=0.9, eps=1e-6)
    update!(w,g,p::Adadelta)

Container for parameters of the Adadelta optimization algorithm used by
[`update!`](@ref).

Adadelta is an extension of Adagrad that tries to prevent the decrease
of the learning rates to zero as training progresses. It scales the
learning rate based on the accumulated gradients like Adagrad and
holds the acceleration term like Momentum. It updates the weights with
the following formulas:

    G = (1-rho) * g .^ 2 + rho * G
    update = g .* sqrt(delta + eps) ./ sqrt(G + eps)
    w = w - lr * update
    delta = rho * delta + (1-rho) * update .^ 2

where `w` is the weight, `g` is the gradient of the objective function
w.r.t `w`, `lr` is the learning rate, `G` is an array with the same
size and type of `w` and holds the sum of the squares of the
gradients. `eps` is a small constant to prevent a zero value in the
denominator.  `rho` is the momentum parameter and `delta` is an array
with the same size and type of `w` and holds the sum of the squared
updates.

If `norm(g) > gclip > 0`, `g` is scaled so that its norm is equal
to `gclip`.  If `gclip==0` no scaling takes place.

Reference: [Zeiler,
M. D. (2012)](http://arxiv.org/abs/1212.5701). ADADELTA: An Adaptive
Learning Rate Method.

"""
mutable struct Adadelta
    lr::AbstractFloat
    gclip::AbstractFloat
    rho::AbstractFloat
    eps::AbstractFloat
    G
    delta
end

Adadelta(; lr=0.01, gclip=0, rho=0.9, eps=1e-6)=Adadelta(lr, gclip, rho, eps, nothing, nothing)

"""
    Rmsprop(;lr=0.001, gclip=0, rho=0.9, eps=1e-6)
    update!(w,g,p::Rmsprop)

Container for parameters of the Rmsprop optimization algorithm used by
[`update!`](@ref).

Rmsprop scales the learning rates by dividing the root mean squared of
the gradients. It updates the weights with the following formula:

    G = (1-rho) * g .^ 2 + rho * G
    w = w - lr * g ./ sqrt(G + eps)

where `w` is the weight, `g` is the gradient of the objective function
w.r.t `w`, `lr` is the learning rate, `G` is an array with the same
size and type of `w` and holds the sum of the squares of the
gradients. `eps` is a small constant to prevent a zero value in the
denominator.  `rho` is the momentum parameter and `delta` is an array
with the same size and type of `w` and holds the sum of the squared
updates.

If `norm(g) > gclip > 0`, `g` is scaled so that its norm is equal
to `gclip`.  If `gclip==0` no scaling takes place.

Reference: [Tijmen Tieleman and Geoffrey Hinton
(2012)](https://dirtysalt.github.io/images/nn-class-lec6.pdf). "Lecture
6.5-rmsprop: Divide the gradient by a running average of its recent
magnitude."  COURSERA: Neural Networks for Machine Learning 4.2.

"""
mutable struct Rmsprop
    lr::AbstractFloat
    gclip::AbstractFloat
    rho::AbstractFloat
    eps::AbstractFloat
    G
end

Rmsprop(; lr=0.001, gclip=0, rho=0.9, eps=1e-6)=Rmsprop(lr, gclip, rho, eps, nothing)

"""
    Adam(;lr=0.001, gclip=0, beta1=0.9, beta2=0.999, eps=1e-8)
    update!(w,g,p::Adam)

Container for parameters of the Adam optimization algorithm used by
[`update!`](@ref).

Adam is one of the methods that compute the adaptive learning rate. It
stores accumulated gradients (first moment) and the sum of the squared
of gradients (second).  It scales the first and second moment as a
function of time. Here is the update formulas:

    m = beta1 * m + (1 - beta1) * g
    v = beta2 * v + (1 - beta2) * g .* g
    mhat = m ./ (1 - beta1 ^ t)
    vhat = v ./ (1 - beta2 ^ t)
    w = w - (lr / (sqrt(vhat) + eps)) * mhat

where `w` is the weight, `g` is the gradient of the objective function
w.r.t `w`, `lr` is the learning rate, `m` is an array with the same
size and type of `w` and holds the accumulated gradients. `v` is an
array with the same size and type of `w` and holds the sum of the
squares of the gradients. `eps` is a small constant to prevent a zero
denominator. `beta1` and `beta2` are the parameters to calculate bias
corrected first and second moments. `t` is the update count.

If `norm(g) > gclip > 0`, `g` is scaled so that its norm is equal
to `gclip`.  If `gclip==0` no scaling takes place.

Reference: [Kingma, D. P., & Ba,
J. L. (2015)](https://arxiv.org/abs/1412.6980). Adam: a Method for
Stochastic Optimization. International Conference on Learning
Representations, 1–13.

"""
mutable struct Adam
    lr::AbstractFloat
    gclip::AbstractFloat
    beta1::AbstractFloat
    beta2::AbstractFloat
    eps::AbstractFloat
    t::Int
    fstm
    scndm
end

Adam(; lr=0.001, gclip=0, beta1=0.9, beta2=0.999, eps=1e-8)=Adam(lr, gclip, beta1, beta2, eps, 0, nothing, nothing)

"""
    update!(weights, gradients, params)
    update!(weights, gradients; lr=0.001, gclip=0)

Update the `weights` using their `gradients` and the optimization
algorithm parameters specified by `params`.  The 2-arg version
defaults to the [`SGD`](@ref) algorithm with learning rate `lr` and
gradient clip `gclip`.  `gclip==0` indicates no clipping. The
`weights` and possibly `gradients` and `params` are modified in-place.

`weights` can be an individual numeric array or a collection of arrays
represented by an iterator or dictionary.  In the individual case,
`gradients` should be a similar numeric array of `size(weights)` and
`params` should be a single object.  In the collection case, each
individual weight array should have a corresponding params object.
This way different weight arrays can have their own optimization
state, different learning rates, or even different optimization
algorithms running in parallel.  In the iterator case, `gradients` and
`params` should be iterators of the same length as `weights` with
corresponding elements.  In the dictionary case, `gradients` and
`params` should be dictionaries with the same keys as `weights`.

Individual optimization parameters can be one of the following
types. The keyword arguments for each type's constructor and their
default values are listed as well.

* [`SGD`](@ref)`(;lr=0.001, gclip=0)`
* [`Momentum`](@ref)`(;lr=0.001, gclip=0, gamma=0.9)`
* [`Nesterov`](@ref)`(;lr=0.001, gclip=0, gamma=0.9)`
* [`Rmsprop`](@ref)`(;lr=0.001, gclip=0, rho=0.9, eps=1e-6)`
* [`Adagrad`](@ref)`(;lr=0.1, gclip=0, eps=1e-6)`
* [`Adadelta`](@ref)`(;lr=0.01, gclip=0, rho=0.9, eps=1e-6)`
* [`Adam`](@ref)`(;lr=0.001, gclip=0, beta1=0.9, beta2=0.999, eps=1e-8)`

# Example:

    w = rand(d)                 # an individual weight array
    g = lossgradient(w)         # gradient g has the same shape as w
    update!(w, g)               # update w in-place with SGD()
    update!(w, g; lr=0.1)       # update w in-place with SGD(lr=0.1)
    update!(w, g, SGD(lr=0.1))  # update w in-place with SGD(lr=0.1)

    w = (rand(d1), rand(d2))    # a tuple of weight arrays
    g = lossgradient2(w)        # g will also be a tuple
    p = (Adam(), SGD())         # p has params for each w[i]
    update!(w, g, p)            # update each w[i] in-place with g[i],p[i]

    w = Any[rand(d1), rand(d2)] # any iterator can be used
    g = lossgradient3(w)        # g will be similar to w
    p = Any[Adam(), SGD()]      # p should be an iterator of same length
    update!(w, g, p)            # update each w[i] in-place with g[i],p[i]

    w = Dict(:a => rand(d1), :b => rand(d2)) # dictionaries can be used
    g = lossgradient4(w)
    p = Dict(:a => Adam(), :b => SGD())
    update!(w, g, p)

"""
function update! end

for T in (Array{Float32},Array{Float64},KnetArray{Float32},KnetArray{Float64}); @eval begin

    function update!(w::$T, g::$T, p::SGD)
        gclip!(g, p.gclip)
        axpy!(-p.lr, g, w)
    end

    # Two arg defaults to SGD
    function update!(w::$T, g::$T; lr=SGDLR, gclip=0)
        gclip!(g, gclip)
        axpy!(-lr, g, w)
    end

    function update!(w::$T, g::$T, p::Momentum)
        gclip!(g, p.gclip)
        if p.velocity===nothing; p.velocity=zero(w); end
        lmul!(p.gamma, p.velocity)
        axpy!(p.lr, g, p.velocity)
        axpy!(-1, p.velocity, w)
    end

    # https://arxiv.org/pdf/1212.0901.pdf Eq. (7)
    function update!(w::$T, g::$T, p::Nesterov)
        gclip!(g, p.gclip)
        p.velocity ===nothing && (p.velocity = zero(w))
        lmul!(p.gamma, p.velocity)
        axpy!(-1, p.velocity, w)
        axpy!(-p.lr, g, p.velocity)
        axpy!(1+p.gamma, p.velocity, w)
    end

    function update!(w::$T, g::$T, p::Adam)
        gclip!(g, p.gclip)
        if p.fstm===nothing; p.fstm=zero(w); p.scndm=zero(w); end
        p.t += 1
        lmul!(p.beta1, p.fstm)
        axpy!(1-p.beta1, g, p.fstm)
        lmul!(p.beta2, p.scndm)
        axpy!(1-p.beta2, g .* g, p.scndm)
        fstm_corrected = p.fstm / (1 - p.beta1 ^ p.t)
        scndm_corrected = p.scndm / (1 - p.beta2 ^ p.t)
        axpy!(-p.lr, (fstm_corrected ./ (sqrt.(scndm_corrected) .+ p.eps)), w)
    end

    function update!(w::$T, g::$T, p::Adagrad)
        gclip!(g, p.gclip)
        if p.G===nothing; p.G=zero(w); end
        axpy!(1, g .* g, p.G)
        axpy!(-p.lr, g ./ sqrt.(p.G .+ p.eps), w)
    end

    function update!(w::$T, g::$T, p::Adadelta)
        gclip!(g, p.gclip)
        if p.G===nothing; p.G=zero(w); p.delta=zero(w); end
        lmul!(p.rho, p.G)
        axpy!(1-p.rho, g .* g, p.G)
        dw = g .* sqrt.(p.delta .+ p.eps) ./ sqrt.(p.G .+ p.eps)
        lmul!(p.rho, p.delta)
        axpy!(1-p.rho, dw .* dw , p.delta)
        axpy!(-p.lr, dw, w)
    end

    function update!(w::$T, g::$T, p::Rmsprop)
        gclip!(g, p.gclip)
        if p.G===nothing; p.G=zero(w); end
        lmul!(p.rho, p.G)
        axpy!(1-p.rho, g .* g, p.G)
        axpy!(-p.lr, g ./ sqrt.(p.G .+ p.eps), w)
    end

    # If type of g does not match, something may be wrong
    update!(w::$T, g, p)=error("Gradient type mismatch: w::$(typeof(w)) g::$(typeof(g))")
    update!(w::$T, g; o...)=error("Gradient type mismatch: w::$(typeof(w)) g::$(typeof(g))")

    # AutoGrad may return Nothing for a zero gradient
    update!(w::$T, g::Nothing, p)=w
    update!(w::$T, g::Nothing; o...)=w

end; end

# AutoGrad may return Nothing for a zero gradient
update!(w, g::Nothing, p)=w
update!(w, g::Nothing; o...)=w

# This takes care of arrays, tuples, iterators in general.
function update!(w,g,p)
    if !(length(w)==length(g)==length(p))
        error("weight, gradient, and optimization parameters not the same length.")
    end
    if isbitstype(eltype(w))
        error("Bad args: $((typeof(w),typeof(g),typeof(p)))")
    end
    for (wi,gi,pi) in zip(w,g,p)
        update!(wi,gi,pi)
    end
end

# We still need an extra method for Dict.
function update!(w::AbstractDict,g::AbstractDict,p::AbstractDict)
    # g may have some keys missing!
    # if !(length(w)==length(g)==length(p))
    #     error("weight, gradient, and optimization parameters not the same length.")
    # end
    for k in keys(g)
        update!(w[k],g[k],p[k])
    end
end

# Two arg version defaults to SGD.
function update!(w,g;lr=SGDLR,gclip=0)
    if !(length(w)==length(g))
        error("weight, gradient not the same length.")
    end
    for (wi,gi) in zip(w,g)
        update!(wi,gi;lr=lr,gclip=gclip)
    end
end

# Two arg version defaults to SGD.
function update!(w::AbstractDict,g::AbstractDict;lr=SGDLR,gclip=0)
    # g may have some keys missing!
    # if !(length(w)==length(g))
    #     error("weight, gradient not the same length.")
    # end
    for k in keys(g)
        update!(w[k],g[k];lr=lr,gclip=gclip)
    end
end

function gclip!(g, gclip)
    if gclip == 0
        g
    else
        gnorm = norm(g)
        if gnorm <= gclip
            g
        else
            lmul!(gclip/gnorm, g)
        end
    end
end

"""
    optimizers(model, otype; options...)

Given parameters of a `model`, initialize and return corresponding
optimization parameters for a given optimization type `otype` and
optimization options `options`. This is useful because each numeric
array in model needs its own distinct optimization
parameter. `optimizers` makes the creation of optimization parameters
that parallel model parameters easy when all of them use the same type
and options.

"""
optimizers(::KnetArray{<:Number},otype; o...) = otype(;o...)
optimizers(::AbstractArray{<:Number},otype; o...) = otype(;o...)
optimizers(a::AbstractDict,otype; o...)=Dict([ k=>optimizers(v,otype;o...) for (k,v) in a ])
optimizers(a::Tuple,otype; o...)=map(x->optimizers(x,otype;o...), a)
optimizers(a::AbstractArray,otype; o...)=map(x->optimizers(x,otype;o...), a)
optimizers(a,otype;o...)=nothing


#=
    bp_ops: operators used in BP that apply to all sfuncs
=#

using Scruff.Utils
using Folds

export compute_lambda,
       compute_bel,
       send_pi,
       outgoing_pis,
       outgoing_lambdas

    #    @op_impl Scruff function Compute_lambda{SFunc}((range, lambda_msgs))
    #     return normalized_product(lambda_msgs, length(range))
    # end

@impl begin
    struct SFuncComputeLambda end

    function compute_lambda(sf::SFunc, range::__OptVec, lambda_msgs::Vector{<:Score})::Score

        if isempty(lambda_msgs)
            ps = zeros(Float64, length(range))
        else
            lams = [[get_log_score(l, o) for l in lambda_msgs] for o in range]
            ps = [sum(lams[i]) for i in 1:length(range)]
        end
        # avoid underflow
        m = maximum(ps)
        qs = ps .- m
        return LogScore(range, qs)
    end
end

@impl begin
    struct SFuncComputeBel end

    function compute_bel(sf::SFunc{I,O}, range::__OptVec{O},  pi::Dist{O}, lambda::Score)::Dist{<:O} where {I,O}
        ps = [cpdf(pi, (), x) * get_score(lambda, x) for x in range]
        return Cat(range, normalize(ps))
    end
end

@impl begin
    struct SFuncSendPi end

    function send_pi(sf::SFunc{I,O}, range::__OptVec{O}, bel::Dist{O}, lambda_msg::Score)::Dist{<:O} where {I,O}
        # pi_msg = [get_score(bel, x) / max.(1e-8, get_score(lambda_msg, x)) for x in range]
        f(x,y) = y == -Inf ? -Inf : x - y
        ps = [f(logcpdf(bel, (), x), get_log_score(lambda_msg, x)) for x in range]
        # delay exponentiation until after avoiding underflow
        m = maximum(ps)
        qs = ps .- m
        exped = exp.(qs)
        return Cat(range, normalize(exped))
    end
end

@impl begin
    struct SFuncOutgoingPis end

    function outgoing_pis(sf::SFunc, range::__OptVec, bel::Dist,
            incoming_lambdas::__OptVec{<:Score})::Vector{<:Dist}
        return [send_pi(sf, range, bel, l) for l in incoming_lambdas]
    end
end

@impl begin
    struct SFuncOutgoingLambdas end

    function outgoing_lambdas(sf::SFunc{I,O},
        lambda::Score,
        range::__OptVec,
        parranges::NTuple{N,Vector},
        incoming_pis::Tuple)::Vector{<:Score} where {N,I,O}

        lambdas = Score[]
        for i = 1:length(incoming_pis)
            msg = send_lambda(sf, lambda, range, parranges, incoming_pis, i)
            push!(lambdas, msg)
        end
        return lambdas
    end

end


using Test

@testset "/src/ConvexPolytope.jl" begin

using BellScenario: ConvexPolytope

@testset "linear_interpolation()" begin

    @test isapprox(ConvexPolytope.linear_interpolation([1;0;0],[0;1;0],dw=0.2), map(
        x -> reshape(x,(3,1)),
        [[1.0;0;0],[0.8;0.2;0],[0.6;0.4;0],[0.4;0.6;0],[0.2;0.8;0],[0;1.0;0]]
    ))

    @test all(map(
        x -> ConvexPolytope.facet_contains([0 0 0 -1], x),
        ConvexPolytope.linear_interpolation([1;0;0],[0;1;0],dw=0.2)
    ))
end

@testset "facet_contains()" begin

    @test ConvexPolytope.facet_contains([-1 1 1 1],[1/3;1/3;1/3])
    @test !(ConvexPolytope.facet_contains([-1 1 1 1],[1;1;1]))

    @test ConvexPolytope.facet_contains([-1 1 0 -1 0 1 -1],[2/3;1/6;1/6;1/6;2/3;1/6])
end

@testset "violated_facets()" begin

    facets = [[0 -1 0 0],[0 0 -1 0],[0 0 0 -1],[-1 1 1 1]]
    @test ConvexPolytope.violated_facets(facets, [1;1;1]) == [[-1 1 1 1]]
    @test ConvexPolytope.violated_facets(facets, [1/4;1/4;1/4]) == []

    facets = [
        [0 -1 0 0 0 0 0],[0 0 -1 0 0 0 0],[0 0 0 -1 0 0 0],[0 0 0 0 -1 0 0],
        [0 0 0 0 0 -1 0],[0 0 0 0 0 0 -1],[-1 0 0 1 0 0 1],[-1 0 1 0 0 1 0],
        [-1 1 0 0 1 0 0],[-1 -1 0 1 -1 1 0],[-1 -1 1 0 -1 0 1],[-1 0 -1 1 1 -1 0],
        [-1 0 1 -1 1 0 -1], [-1 1 -1 0 0 -1 1], [-1 1 0 -1 0 1 -1]
    ]
    @test ConvexPolytope.violated_facets(facets, [1;0;0;0;1;0]) == [[-1 1 0 -1 0 1 -1]]
end

@testset "facet_violation()" begin

    @test ConvexPolytope.facet_violation([-1 1 0 -1 0 1 -1], [1;1;0;0;0;1;0]) == 1
    @test ConvexPolytope.facet_violation([-1 1 0 -1 0 1 -1], [1;0;0;0;1;0]) == 1

    @test_throws ArgumentError ConvexPolytope.facet_violation([-1 1 0 -1 0 1 -1],[1;1])

    @test ConvexPolytope.facet_violation([-1 1 1 1], [1;1;1]) == 2
end

@testset "facet_distance()" begin
    @test ConvexPolytope.facet_distance([0 -1 0 0],[1;1;1]) ≈ -1
    @test ConvexPolytope.facet_distance([-1 1 1 1], [1;1;1]) ≈ 2/sqrt(3)
end

@testset "contains()" begin

    vertices = [[0;0;0],[1;0;0],[0;1;0],[0;0;1]]
    facets = [[0 -1 0 0],[0 0 -1 0],[0 0 0 -1],[-1 1 1 1]]

    @test !(ConvexPolytope.facets_contain(facets, [1;1;1]))
    @test ConvexPolytope.facets_contain(facets,[1/4;1/4;1/4])
    @test ConvexPolytope.facets_contain(facets,[1/3;1/3;1/3])
    @test all(map(
        (v) -> ConvexPolytope.facets_contain(facets, v),
        vertices
    ))

end

@testset "collect_vertices()" begin
    facets = [[-1 1 0 -1 0 1 -1],[0 -1 0 0 0 0 0]]
    vertices = [
        [1;0;0;0;0;0;0],[1;1;1;1;0;0;0],[1;0;0;0;1;1;1],[1;1;1;0;0;0;0],
        [1;1;1;0;0;0;1],[1;1;0;1;0;0;0],[1;1;0;1;0;1;0],[1;0;1;1;0;0;0],
        [1;0;1;1;1;0;0]
    ]

    facet_dict = ConvexPolytope.collect_vertices(facets, vertices)

    @test length(facet_dict) == 2
    @test facet_dict[1]["facet"] == [-1 1 0 -1 0 1 -1]
    @test facet_dict[1]["vertices"] == [[1;1;1;0;0;0;0],[1;1;0;1;0;1;0]]
    @test facet_dict[2]["facet"] == [0 -1 0 0 0 0 0]
    @test facet_dict[2]["vertices"] == [[1;0;0;0;0;0;0],[1;0;0;0;1;1;1],[1;0;1;1;0;0;0],[1;0;1;1;1;0;0]]
end

@testset "collect_facets()" begin
    facets = [[-1 1 0 -1 0 1 -1],[0 -1 0 0 0 0 0]]
    vertices = [
        [1;0;0;0;0;0;0],[1;0;0;0;1;1;1],[1;1;1;0;0;0;1],[1;0;0;0;1;1;0]
    ]

    facet_dict = ConvexPolytope.collect_facets(vertices, facets)

    @test length(facet_dict) == 4
    @test facet_dict[1]["vertex"] == [1;0;0;0;0;0;0]
    @test facet_dict[1]["facets"] == [[0 -1 0 0 0 0 0]]
    @test facet_dict[2]["vertex"] == [1;0;0;0;1;1;1]
    @test facet_dict[2]["facets"] == [[0 -1 0 0 0 0 0]]
    @test facet_dict[3]["vertex"] == [1;1;1;0;0;0;1]
    @test facet_dict[3]["facets"] == []
    @test facet_dict[4]["vertex"] == [1;0;0;0;1;1;0]
    @test facet_dict[4]["facets"] == [[-1 1 0 -1 0 1 -1],[0 -1 0 0 0 0 0]]
end

@testset "contains()" begin
    facets = [[0 -1 -1 -1 0 0 0], [0 0 0 0 -1 -1 -1], [0 0 0 0 0 0 0]]

    @test ConvexPolytope.contains(facets, [1;0;0;0;0;0;0])
    @test ConvexPolytope.contains(facets, [1;1;1;1;0;0;0])
    @test ConvexPolytope.contains(facets, [1;0.5;0.5;0.5;0.5;0.5;0.5])
    @test !ConvexPolytope.contains(facets, [1;-1;0;0;0;0;0])
end

end


## Parallel Residual Constrained Alternating Minimization
# Function Definitions

function AltMinQR{ET<:Number}(b::AbstractArray{ET}, r::Int, K::Int)
# REAL
# Performs RCAM on a single matricised monochromatic slice
# Calls: minFrobQR

    # Generate L0
    m,n = size(b)
    L = collect(ET, randn(m,r))

    indi = (1:max(m,n))::UnitRange{Int64}
    local R

    # Begin Alterations
    for k = 1:K

        # Solve for R
        R = minFrobQR(L,b,indi)

        # Solve for L
        L = minFrobQR(R,b',indi)

    end

    return L, R
end
export AltMinQR

function minFrobQR{ET<:Number}(L::AbstractArray{ET}, b::AbstractArray{ET}, indi::UnitRange{Int64})
# Solves for each row of input
# Calls: minFrobRsub

    m = size(b,2)::Int
    r = size(L,2)::Int

    # Initialize R
    R = zeros(ET,m,r)

    for k =1:length(indi)

        if indi[k]<=m

            # Solve for the kth row of R
            bsub = b[:,indi[k]]

            # Find the non zero rows of the kth column of b
            ind = find(x -> x != 0, b[:,indi[k]])

            # Keep only the non zero
            bsub = b[ind,indi[k]]

            # QRfact of ind rows of L
            q_qr,r_qr = qr(L[ind,:],thin=false)

            #Extend r to full size
            extension = zeros(size(q_qr,1)-size(r_qr,1),size(r_qr,2))
            r_full = vcat(r_qr,extension)

            R[indi[k],:] = minFrobRsub(q_qr,r_full,bsub)

        end

    end

    R = conj(R)

    return R
end

function minFrobRsub(q::AbstractArray,r_full::AbstractArray,bsub::AbstractArray)
# Solves a single row

    bsub = q'*bsub

    v = backsub(r_full,bsub)

    return v

end

function backsub{ET<:Number}(A::AbstractArray{ET}, y::AbstractArray{ET})

    m = size(A,2)
    x = zeros(ET,m,1)

    x[m] = y[m]/A[m,m]

    for i = m-1:-1:1
        x[i] = (y[i]-dot(A[i,i+1:m], x[i+1:m]))/A[i,i]
    end

    return x

end


using ApproxFunOrthogonalPolynomials, SpecialFunctions, Test
import ApproxFunBase: ldiv_coefficients

## ODEs

d = -20000..20000.0
x = Fun(identity,d)
u = [Dirichlet(d);Derivative(d)^2+I]\[[1,0],0]
u = [Dirichlet(d);Derivative(d)^2+I]\[[1,0],0]
@time u = [Dirichlet(d);Derivative(d)^2+I]\[[1,0],0]
println("Cos/Sin: should be ~0.016920 seconds (3.19 k allocations: 12.593 MiB)")

d = -1000..5.0
x=Fun(identity,d)
u=[Dirichlet(d);Derivative(d)^2-x]\[[1,0],0]
u=[Dirichlet(d);Derivative(d)^2-x]\[[1,0],0]
@time u=[Dirichlet(d);Derivative(d)^2-x]\[[1,0],0]
println("Airy: 0.014356 seconds (1.08 k allocations: 8.015 MB)")

M=cache([Dirichlet(d);Derivative(d)^2-x];padding=true)
@time ApproxFunBase.resizedata!(M,12500,:)
println("Airy construct op: 0.003200 seconds (131 allocations: 3.723 MiB)")

S=Chebyshev()
x=Fun(identity,S)
D=Derivative(S)
L=D^2+(7+2x+6x^2)
B=Dirichlet(S)
n=20000
rhs=fill(1.0,n+2)
u=ldiv_coefficients([B;L],rhs)
u=ldiv_coefficients([B;L],rhs)
@time u=ldiv_coefficients([B;L],rhs)
println("Poly: should be ~0.020926 seconds (3.00 k allocations: 9.416 MiB)")

S=Chebyshev()
x=Fun(identity,S)
D=Derivative(S)
L=D^2+cos(x)
B=Dirichlet(S)
n=2000
rhs=fill(1.0,n+2)
u=ldiv_coefficients([B;L],rhs;maxlength=Inf)
u=ldiv_coefficients([B;L],rhs;maxlength=Inf)
@time u=ldiv_coefficients([B;L],rhs;maxlength=Inf)
println("Cos: should be ~0.0075")

S=Chebyshev()
x=Fun(identity,S)
D=Derivative(S)
L=D^2+sin(x)
B=Dirichlet(S)
n=2000
rhs=fill(1.0,n+2)
u=ldiv_coefficients([B;L],rhs;maxlength=Inf)
u=ldiv_coefficients([B;L],rhs;maxlength=Inf)
@time u=ldiv_coefficients([B;L],rhs;maxlength=Inf)
println("Sin: should be ~0.008663 seconds (660 allocations: 2.987 MB)")

## Bessel

x=Fun(identity,1..2000)
d=domain(x)
B=Dirichlet()
ν=1000.0
D = Derivative()
L=x^2*D^2 + x*D + (x^2 - ν^2)   # our differential operator
u=[B;L]\[[besselj.(ν,endpoints(d))...],0]
@time u=[B;L]\[[besselj.(ν,endpoints(d))...],0]
println("Bessel: should be ~0.008441 seconds (6.14 k allocations: 4.765 MiB)")

x=Fun()
exp(10000*im*x)
@time exp(10000*im*x)
println("Complex exp: Time should be 0.03")

## Piecewise
x=Fun(identity,Domain(-20..15) \ [-10.,-5.,0.,1.])
sp=space(x)
D=Derivative(sp)
B=[Dirichlet(sp);continuity(sp,0:1)]
u=[B;D^2-x]\[[airyai(-20.),0.],zeros(8),0];
@time u=[B;D^2-x]\[[airyai(-20.),0.],zeros(8),0]
println("Piecewise Airy: should be ~0.008")

## Vector
d=ChebyshevInterval()
D=Derivative(d);
B=ldirichlet();
Bn=lneumann();

f=Fun(x->[exp(x),cos(x)],d)

A=[B 0;
   Bn 0;
   0 B;
   D^2-I 2.0I;
   0 D+I];

u=A\Any[0.,0.,0.,f...]
@time u=A\Any[0.,0.,0.,f...]
println("Systems: should be ~0.0008")

d=Interval(-300.,5.)
x=Fun(identity,d)
A=Derivative(d)^2-x
u=nullspace(A)
@test A[1:10,1:10] ≈ transpose(transpose(A)[1:10,1:10])
@time u=nullspace(A)
println("Nullspace Airy: 0.052730 seconds (75.21 k allocations: 56.736 MB)")


export PatchedHomotopy, PatchedHomotopyCache

"""
    PatchedHomotopy(H::AbstractHomotopy, patch, v::PVector)

Augment the homotopy `H` with the given patch `v`. This results in the system `[H(x,t); v ⋅ x - 1]`
"""
struct PatchedHomotopy{H<:AbstractHomotopy,PS<:AbstractAffinePatchState} <: AbstractHomotopy
    homotopy::H
    patch::PS
end

function PatchedHomotopy(H::AbstractHomotopy, p::AbstractAffinePatch, x)
    PatchedHomotopy(H, state(p, x))
end

function Base.size(H::PatchedHomotopy)
    m, n = size(H.homotopy)
    (m + nequations(H.patch), n)
end

struct PatchedHomotopyCache{HC,T} <: AbstractHomotopyCache
    cache::HC
    A::Matrix{T} # intermediate storage of the jacobian
    b::Vector{T} # intermediate storage for the evaluation
end

function cache(ph::PatchedHomotopy, x, t)
    H = ph.homotopy
    c = cache(H, x, t)
    PatchedHomotopyCache(c, jacobian(H, x, t, c), evaluate(H, x, t, c))
end

function evaluate!(u, H::PatchedHomotopy, x, t, c::PatchedHomotopyCache)
    M, N = size(H.homotopy)

    # H(x, t)
    evaluate!(c.b, H.homotopy, x, t, c.cache)
    @inbounds for i = 1:M
        u[i] = c.b[i]
    end
    evaluate!(u, H.patch, x)
    u
end
function jacobian!(U, H::PatchedHomotopy, x::PVector, t, c::PatchedHomotopyCache)
    M, N = size(H.homotopy)
    # J_H(x, t)
    jacobian!(c.A, H.homotopy, x, t, c.cache)
    @inbounds for j = 1:N, i = 1:M
        U[i, j] = c.A[i, j]
    end
    jacobian!(U, H.patch, x)
    U
end

function dt!(u, H::PatchedHomotopy, x::PVector, t, c::PatchedHomotopyCache)
    M, N = size(H.homotopy)
    # [H(x,t); v ⋅ x - 1]/∂t = [∂H(x,t)/∂t; 0]
    dt!(c.b, H.homotopy, x, t, c.cache)
    @inbounds for i = 1:M
        u[i] = c.b[i]
    end
    for i = 1:nequations(H.patch)
        u[M+i] = zero(eltype(u))
    end
    u
end

function evaluate_and_jacobian!(
    u,
    U,
    H::PatchedHomotopy,
    x::PVector,
    t,
    c::PatchedHomotopyCache,
)
    M, N = size(H.homotopy)
    evaluate_and_jacobian!(c.b, c.A, H.homotopy, x, t, c.cache)

    @inbounds for j = 1:N, i = 1:M
        U[i, j] = c.A[i, j]
    end
    jacobian!(U, H.patch, x)

    @inbounds for i = 1:M
        u[i] = c.b[i]
    end
    evaluate!(u, H.patch, x)

    nothing
end

function jacobian_and_dt!(U, u, H::PatchedHomotopy, x::PVector, t, c::PatchedHomotopyCache)
    M, N = size(H.homotopy)
    A, b = c.A, c.b
    jacobian_and_dt!(A, b, H.homotopy, x, t, c.cache)
    # jacobian
    @inbounds for j = 1:N, i = 1:M
        U[i, j] = A[i, j]
    end
    jacobian!(U, H.patch, x)

    # dt
    @inbounds for i = 1:M
        u[i] = b[i]
    end
    for i = 1:nequations(H.patch)
        u[M+i] = zero(eltype(u))
    end

    nothing
end

function evaluate(H::PatchedHomotopy, x, t, c::PatchedHomotopyCache)
    M, N = size(H)
    evaluate!(similar(c.b, M), H, x, t, c)
end
function jacobian(H::PatchedHomotopy, x, t, c::PatchedHomotopyCache)
    M, N = size(H)
    jacobian!(similar(c.A, M, N), H, x, t, c)
end
function dt(H::PatchedHomotopy, x, t, c::PatchedHomotopyCache)
    M, N = size(H)
    dt!(similar(c.b, M), H, x, t, c)
end

basehomotopy(H::PatchedHomotopy) = basehomotopy(H.homotopy)


function set_linearized_b!(b, q′, d::Domain)
    c = CartesianIndices(d)
    l = LinearIndices(d)
    for I = c
        b[l[I]] = q′[I]
    end
    ol = last(l)
    for I = c
        b[ol+l[I]] = 0
    end
end

function set_bψ!(bψ, ψ, ϕ, q, d::Domain)
    c = CartesianIndices(d)
    l = LinearIndices(d)
    for I = c
        bψ[l[I]] = (
            q[I] + ∂x²(I,ϕ,d) + ∂y²(I,ϕ,d) - ∂z²(I,ϕ,d)
          - 2*(∂x²(I,ψ,d)*∂y²(I,ψ,d) - ∂xy(I,ψ,d)^2)
          + ∂xz(I,ψ,d)*∂xz(I,ϕ,d) + ∂yz(I,ψ,d)*∂yz(I,ϕ,d)
        )
    end
    return bψ
end

function set_bϕ!(bϕ, ψ, ϕ, q, d::Domain)
    c = CartesianIndices(d)
    l = LinearIndices(d)
    for I = c
        bϕ[l[I]] = (
            q[I] + ∂x²(I,ψ,d) + ∂y²(I,ψ,d)
          + 2*(∂x²(I,ψ,d)*∂y²(I,ψ,d) - ∂xy(I,ψ,d)^2)
          + ∂xz(I,ψ,d)*∂xz(I,ϕ,d) + ∂yz(I,ψ,d)*∂yz(I,ϕ,d)
        )
    end
    return bϕ
end

function set_cv_bψ!(bψ, q, d::Domain)
    c = CartesianIndices(d)
    l = LinearIndices(d)
    for I = c
        bψ[l[I]] = q[I] - 1
    end
    return bψ
end
function set_cv_bϕ!(bϕ, ψ, d::Domain)
    c = CartesianIndices(d)
    l = LinearIndices(d)
    for I = c
        bϕ[l[I]] = (
            ∂x²(I,ψ,d) + ∂y²(I,ψ,d)
          + 2*(∂x²(I,ψ,d)*∂y²(I,ψ,d) - (∂xy(I,ψ,d)^2))
        )
    end
    return bϕ
end


# a simple LP example
#
#   maximize x + y
#
#   s.t. 50 x + 24 y <= 2400
#        30 x + 33 y <= 2100
#        x >= 45, y >= 5
#
#   solution: x = 45, y = 6.25, objv = 51.25

using Gurobi, Compat.Test, Compat.GC

@testset "LP 01b" begin

    env = Gurobi.Env()
    setparam!(env, "OutputFlag", 0)
    setparams!(env, Method=2)  # using barrier method

    method = getparam(env, "Method")

    model = Gurobi.Model(env, "lp_01", :maximize)

    # add variables
    add_cvars!(model, [1., 1.], [45., 5.], Inf)
    update_model!(model)

    # add constraints
    add_constrs!(model, Cint[1, 3], Cint[1, 2, 1, 2],
        [50., 24., 30., 33.], '<', [2400., 2100.])
    update_model!(model)

    # perform optimization
    optimize(model)

    # show results
    info = get_optiminfo(model)

    sol = get_solution(model)
    @test sol == [45, 6.25]

    objv = get_objval(model)
    @test objv == 51.25

    GC.gc()  # test finalizers
end


function excess_gibbs_free_energy(model::ActivityModel,p,T,z)
    lnγ = log.(activity_coefficient(model,p,T,z))
    return sum(z[i]*R̄*T*lnγ[i] for i ∈ @comps)
end

function bubble_pressure(model::ActivityModel,T,x)
    sat = sat_pure.(model.puremodel,T)
    p_sat = [tup[1] for tup in sat]
    γ     = activity_coefficient(model,1e-4,T,x)
    p     = sum(x.*γ.*p_sat)
    y     = x.*γ.*p_sat ./ p
    return (p,y)
end


## Kernels
## =======

"Real+ -> [0, 1]"
kf1(x, β = 0.0001) = x / (x + β)
kf1β(β) = d -> kf1(d, β)

"(Log) Squared exponential kernel `α = 1/2l^2`, higher α is lower temperature  "
kse(d, α = 1000.0) = - α * d
kseα(α) = d -> kse(d, α)

"γ exponential"
γe(x, l = 1, γ = 1) = exp(-(x/l)^γ)

"Rational Quadratic"
rquad(x, l = 1, α = 1) = (1 + (x^2)/(2*α*l^2))^(-α)

"Power law kernel"
kpow(d, α = 1.0, k = 2) = -k * log(d) + log(α)

burr(x, c = 1, k = 40) =  log(c) + log(k) +  (c - 1) * log(x) - (k + 1)*log(1 + x^c)

# Temperature Modulation
const GLOBALKERNEL_ = Function[kse]

"Global Kernel"
function globalkernel!(k)
  global GLOBALKERNEL_
  GLOBALKERNEL_[1] = k
end

"Retrieve global kernel"
function globalkernel()
  global GLOBALKERNEL_
  GLOBALKERNEL_[1]
end

"Temporarily set global kernel"
function withkernel(thunk, k)
  globalkernel!(k)
  res = thunk()
  globalkernel!(kse)
  res
end

# (Cassette-based) Temperature Modulation
Cassette.@context AlphaCtx

"""
`f(args)` where temperature controlled with temperature `α`

```julia
x = normal(0.0, 1.0)
atα(10, rand, x ==ₛ 0.3)
```
"""
function atα(α, f, args...)
  ctx = AlphaCtx(metadata = α)
  Cassette.overdub(ctx, f, args...)
end

Cassette.overdub(ctx::AlphaCtx, ::typeof(kse), x, α) = kse(x, ctx.metadata)

"""

```julia
x = normal(0.0, 1.0)
@atα 100 rand(y ==ₛ 0.0)
```
"""
macro atα(a, fexpr)
  @pre fexpr.head == :call
  :(atα($(esc(a)), $(esc.(fexpr.args)...)))
end


using DeformableRegistration: Distance, ImageProcessing, Transformation, Interpolation, regOptions
 using Base.Test
 using MicroLogging
 using Interpolations

include("helpers/checkDerivative.jl")

@testset "Distance: distance derivatives - SSD" begin
    # create reference image
    refImg = createImage(100*rand(50,60))

    # measure distance and check derivative (nonparametric)
    centeredGrid = getCellCenteredGrid(refImg)
     centeredGrid.data[:] = centeredGrid.data[:] + 0.3 *rand(size(centeredGrid.data))
     D,dD,d2D = ssdDistance(refImg,refImg,centeredGrid.data,doDerivative=true,doHessian=true)
     Dfunc(x) = ssdDistance(refImg,refImg,x)[1]
     errlin,errquad = checkDerivative(Dfunc,dD',centeredGrid.data,doPlot=false)
     @test checkErrorDecay(errquad)

    options = regOptions()
     options.matrixFree = true;
     D,dD,d2D = ssdDistance(refImg,refImg,centeredGrid.data,doDerivative=true,doHessian=true,options=options)
     Dfunc(x) = ssdDistance(refImg,refImg,x)[1]
     errlin,errquad = checkDerivative(Dfunc,dD',centeredGrid.data)
     @test checkErrorDecay(errquad)

    # measure distance and check derivative (parametic)
    centeredGrid = getCellCenteredGrid(refImg)
     options = DeformableRegistration.regOptions()
     options.parametricOnly=true
     evaluationPoint = [1,0,0,0,1,0]+0.1*rand(6)
     D,dD,d2D = ssdDistance(refImg,refImg,transformGridAffine(centeredGrid,evaluationPoint).data,doDerivative=true,doHessian=true,options=options)
     Dfunc(p) = ssdDistance(refImg,refImg,transformGridAffine(centeredGrid,p).data,options=options)[1]
     errlin,errquad = checkDerivative(Dfunc,dD',evaluationPoint)
     @test checkErrorDecay(errquad)

    options.matrixFree = true;
     D,dD,d2D = ssdDistance(refImg,refImg,transformGridAffine(centeredGrid,evaluationPoint).data,doDerivative=true,doHessian=true,options=options)
     Dfunc(p) = ssdDistance(refImg,refImg,transformGridAffine(centeredGrid,p).data,options=options)[1]
     errlin,errquad = checkDerivative(Dfunc,dD',evaluationPoint)
     @test checkErrorDecay(errquad)
end

@testset "NGF" begin
    @testset "Distance: distance derivatives - NGF - parametric" begin
        refImg = createImage(100*rand(50,60))

        # measure distance and check derivative (nonparametric)
        centeredGrid = getCellCenteredGrid(refImg)
        centeredGrid.data[:] = centeredGrid.data[:] + 0.3 *rand(size(centeredGrid.data))
        D,dD,d2D = ngfDistance(refImg,refImg,centeredGrid.data,doDerivative=true,doHessian=true)
        Dfunc(x) = ngfDistance(refImg,refImg,x)[1]
        errlin,errquad = checkDerivative(Dfunc,dD',centeredGrid.data,doPlot=false)
        @test checkErrorDecay(errquad)

        evaluationPoint = [1,0,0,0,1,0.0]+0.1*rand(6)
        centeredGrid = getCellCenteredGrid(refImg)
        options = DeformableRegistration.regOptions()
        options.edgeParameterR = 0.1
        options.edgeParameterT = 0.1
        options.matrixFree = false
        options.parametricOnly = true
        D,dD,d2D = ngfDistance(refImg,refImg,transformGridAffine(centeredGrid,evaluationPoint).data,doDerivative=true,doHessian=true,options=options)
        Dfunc(p) = ngfDistance(refImg,refImg,transformGridAffine(centeredGrid,p).data,options=options)[1]
        errlin,errquad = checkDerivative(Dfunc,dD',evaluationPoint,doPlot=false)
        @test checkErrorDecay(errquad)

        evaluationPoint = [1,0,0,0,1,0.0]+0.1*rand(6)
        centeredGrid = getCellCenteredGrid(refImg)
        options = DeformableRegistration.regOptions()
        options.edgeParameterR = 0.1
        options.edgeParameterT = 0.1
        options.matrixFree = true
        options.parametricOnly = true
        D,dD,d2D = ngfDistance(refImg,refImg,transformGridAffine(centeredGrid,evaluationPoint).data,doDerivative=true,doHessian=true,options=options)
        Dfunc(p) = ngfDistance(refImg,refImg,transformGridAffine(centeredGrid,p).data,options=options)[1]
        errlin,errquad = checkDerivative(Dfunc,dD',evaluationPoint,doPlot=false)
        @test checkErrorDecay(errquad)
    end

    @testset "Distance: distance derivatives - NGF - nonparametric" begin
        refImg = createImage(100*rand(50,60))

        ##
        centeredGrid = getCellCenteredGrid(refImg)
        centeredGrid.data[:] = centeredGrid.data[:] + 0.3 *rand(size(centeredGrid.data))
        options = DeformableRegistration.regOptions()
        options.edgeParameterR = 0.1
        options.edgeParameterT = 0.1
        options.matrixFree = false
        options.parametricOnly = false
        D,dD,d2D = ngfDistance(refImg,refImg,centeredGrid.data,doDerivative=true,doHessian=true,options=options)
        Dfunc(grid) = ngfDistance(refImg,refImg,grid,options=options)[1]
        errlin,errquad = checkDerivative(Dfunc,dD',centeredGrid.data,doPlot=false)
        @test checkErrorDecay(errquad)

        centeredGrid = getCellCenteredGrid(refImg)
        centeredGrid.data[:] = centeredGrid.data[:] + 0.3 *rand(size(centeredGrid.data))
        options = DeformableRegistration.regOptions()
        options.edgeParameterR = 0.1
        options.edgeParameterT = 0.1
        options.matrixFree = true
        options.parametricOnly = false
        D,dD,d2D = ngfDistance(refImg,refImg,centeredGrid.data,doDerivative=true,doHessian=true,options=options)
        Dfunc(grid) = ngfDistance(refImg,refImg,grid,options=options)[1]
        errlin,errquad = checkDerivative(Dfunc,dD',centeredGrid.data,doPlot=false)
        @test checkErrorDecay(errquad)

    end

    @testset "Distance: ngf epsilon estimation" begin
        testimage = dirname(Base.source_path()) * "/testdata/luebeck.jpg"
        refImg = loadImage(testimage)
        for r = refImg.data
             r = r + 0.01*rand()
        end
        warn("Skipping NGF ϵ estimation test.")
        @test_skip estimateNGFEpsilon(refImg,cutoffPercent=80)[1] >= 0 # seems to run forever
    end
end

@testset "SSD arbitrary grid" begin
    @testset "fixedGrid_IdenticalImagesZeroDeformation" begin
        n = 10

        imageData = ones(n,n)
        referenceImage = createImage(imageData, voxelsize = [1.0,1.0], shift = [0.0, 0.0])
        templateImage = createImage(imageData, voxelsize = [1.0,1.0], shift = [0.0, 0.0])
        transformedGrid = getCellCenteredGrid(referenceImage)
        # voxelsize::Array{Float64,1},shift::Array{Float64,1}, gridSize::Tuple{Int64,Int64}
        functionValue,dFunctionValue,d2FunctionValue,d_transformedImage =
            ssdDistanceArbitraryGrid(referenceImage, templateImage, transformedGrid)
                             # doDerivative::Bool=false,doHessian::Bool=false,
                             # options::regOptions=regOptions(),centeredGrid::scaledArray=scaledArray(zeros(1),(0,),[0,],[0,])))
        @test functionValue ≈ 0 atol=0.0001

        functionValue,dFunctionValue,d2FunctionValue,d_transformedImage =
            ssdDistanceArbitraryGrid(referenceImage, templateImage, transformedGrid, doDerivative=true)
        @test functionValue ≈ 0 atol=0.0001

        imageData = rand(n,n)
        referenceImage = createImage(imageData, voxelsize = [1.0,1.0], shift = [0.0, 0.0])
        templateImage = createImage(imageData, voxelsize = [1.0,1.0], shift = [0.0, 0.0])
        transformedGrid = getCellCenteredGrid(referenceImage)
        # voxelsize::Array{Float64,1},shift::Array{Float64,1}, gridSize::Tuple{Int64,Int64}
        functionValue,dFunctionValue,d2FunctionValue,d_transformedImage =
            ssdDistanceArbitraryGrid(referenceImage, templateImage, transformedGrid)
        @test functionValue ≈ 0 atol=0.0001

        functionValue,dFunctionValue,d2FunctionValue,d_transformedImage =
            ssdDistanceArbitraryGrid(referenceImage, templateImage, transformedGrid, doDerivative=true)
        @test functionValue ≈ 0 atol=0.0001

        imageData = rand(n,n)
        referenceImage = createImage(imageData, voxelsize = [1.0,1.0], shift = [0.0, 0.0])
        templateImage = createImage(imageData, voxelsize = [1.0,1.0], shift = [0.0, 0.0])
        transformedGrid = getCellCenteredGrid(restrictResolutionToLevel(referenceImage,2))
        @debug "-------- transformedGrid"
        @debug transformedGrid
        # voxelsize::Array{Float64,1},shift::Array{Float64,1}, gridSize::Tuple{Int64,Int64}
        functionValue,dFunctionValue,d2FunctionValue,d_transformedImage =
            ssdDistanceArbitraryGrid(referenceImage, templateImage, transformedGrid)
        @test functionValue ≈ 0 atol=0.0001

    end

    @testset "fixedGrid_IdenticalImagesFormats" begin
        n = 10

        imageData = ones(n,n)
        referenceImage = createImage(imageData, voxelsize = [1.0,1.0], shift = [0.0, 0.0])
        templateImage = createImage(imageData, voxelsize = [1.0,1.0], shift = [0.0, 0.0])

        gridSize = (5,5)
        voxelsize = n/gridSize[1]
        transformedGrid = getCellCenteredGrid([voxelsize,voxelsize],[0.0,0.0],gridSize)
        # voxelsize::Array{Float64,1},shift::Array{Float64,1}, gridSize::Tuple{Int64,Int64}
        functionValue,dFunctionValue,d2FunctionValue,d_transformedImage =
            ssdDistanceArbitraryGrid(referenceImage, templateImage, transformedGrid, doDerivative=true)
                             # doDerivative::Bool=false,doHessian::Bool=false,
                             # options::regOptions=regOptions(),centeredGrid::scaledArray=scaledArray(zeros(1),(0,),[0,],[0,])))
        @test functionValue ≈ 0 atol=0.0001
        @test size(dFunctionValue)... == prod(gridSize)*2
    ##
        imageData = rand(n,n)
        referenceImage = createImage(imageData, voxelsize = [1.0,1.0], shift = [0.0, 0.0])
        templateImage = createImage(imageData, voxelsize = [1.0,1.0], shift = [0.0, 0.0])

    ##
        gridSize = (10,20)
        voxelsize = 1.0
        transformedGrid = getCellCenteredGrid([voxelsize,voxelsize],[0.0,0.0],gridSize)
        transformedGrid.data[1:200] = transformedGrid.data[1:200] + 4.0
        # voxelsize::Array{Float64,1},shift::Array{Float64,1}, gridSize::Tuple{Int64,Int64}
        functionValue,dFunctionValue,d2FunctionValue,d_transformedImage =
            ssdDistanceArbitraryGrid(referenceImage, templateImage, transformedGrid, doDerivative=true)
    ##
        @test norm(dFunctionValue[105:195]) ≈ 0
        @test norm(dFunctionValue[305:395]) ≈ 0
        #@test norm(dFunctionValue[5:95]) > 1
        @test size(dFunctionValue)... == prod(gridSize)*2

    end

    @testset "ssdResampling_derivatives" begin

    include("../test/helpers/checkDerivative.jl")
    #include("../src/distances/ssdDistance.jl")

    smoothrand(x::Array{Float64,1}) = smoothArray(rand(size(x)),3, 1.0)

    refImg = createImage( smoothArray(100.0*rand(99,111),3, 3.0) )
        #refImg.data.data = smoothArray(refImg.data.data,11, 1.0)
        centeredGrid = getCellCenteredGrid(refImg)
        centeredGrid.data[:] = centeredGrid.data[:] + 5.0 *smoothrand(centeredGrid.data)
        options = regOptions()
        options.matrixFree = true;
        D,dD,d2D = Distance.ssdDistanceArbitraryGrid(refImg,refImg,centeredGrid, doDerivative=true,
            doHessian=true,options=options, interpolationScheme=BSpline(Linear()))
        Dfunc(x) = Distance.ssdDistanceArbitraryGrid(refImg,refImg,x, doDerivative=false, interpolationScheme=BSpline(Linear()))[1]
        errlin,errquad = checkDerivative(Dfunc,dD',centeredGrid, doPlot=false)
        @test_skip checkErrorDecay(errquad)

    # The transposed interpolation operator is not implemented correctly which is
    # why in the case where the grid resolution changes, the derivative checks are failing.

    gridSize = (25,30)
        voxelsize = size(refImg.data)./gridSize
        vs = voxelsize
        centeredGrid = getCellCenteredGrid([vs[1]; vs[2]],[0.0,0.0],gridSize)
        centeredGrid.data[:] = centeredGrid.data[:] + 5 *smoothrand(centeredGrid.data)
        D,dD,d2D = ssdDistanceArbitraryGrid(refImg,refImg,centeredGrid,doDerivative=true,doHessian=true,options=options, interpolationScheme=BSpline(Linear()))
        Dfunc(x) = ssdDistanceArbitraryGrid(refImg,refImg,x, doDerivative=false, interpolationScheme=BSpline(Cubic(Line())))[1]
        errlin,errquad = checkDerivative(Dfunc,dD',centeredGrid, doPlot=false)
        @test_skip checkErrorDecay(errquad)

    gridSize = (10,10)
        voxelsize = size(refImg.data)./gridSize
        vs = voxelsize
        centeredGrid = getCellCenteredGrid([vs[1]; vs[2]],[0.0,0.0],gridSize)
        centeredGrid.data[:] = centeredGrid.data[:] + 10 * smoothrand(centeredGrid.data)
        D,dD,d2D = ssdDistanceArbitraryGrid(refImg,refImg,centeredGrid,doDerivative=true,doHessian=true,options=options)
        Dfunc(x) = ssdDistanceArbitraryGrid(refImg,refImg,x)[1]
        errlin,errquad = checkDerivative(Dfunc,dD',centeredGrid, doPlot=false)
        @test_skip checkErrorDecay(errquad)

    gridSize = (25,30)
        vs = size(refImg.data)./gridSize
        centeredGrid = getCellCenteredGrid([vs[1]; vs[2]],[0.0,0.0],gridSize)
        disp = 10 * ones(gridSize)
        disp[4:6,4:6]=0.0
        testGrid = 0.0 * centeredGrid + vcat(disp[:], disp[:])
        testGridInterp = interpolateDeformationField(testGrid, getCellCenteredGrid(refImg), interpolationScheme=BSpline(Linear()))
        testGridInterp2 = interpolateDeformationField(testGridInterp, getCellCenteredGrid([vs[1]; vs[2]],[0.0,0.0],gridSize), interpolationScheme=BSpline(Linear()))
         #
         # using PyPlot
         # figure()
         # subplot(1,3,1)
         # imshow(reshape(testGrid.data[1:Int(end/2)],testGrid.dimensions))
         # clim([0,10])
         # subplot(1,3,2)
         # imshow(reshape(testGridInterp.data[1:Int(end/2)],testGridInterp.dimensions))
         # clim([0,10])
         # subplot(1,3,3)
         # imshow(reshape(testGridInterp2.data[1:Int(end/2)],testGridInterp2.dimensions))
         # clim([0,10])

        @test_skip sum(testGrid.data[:]) ≈ sum(testGridInterp2.data[:]) atol=0.1

    end

end


using Oceananigans.Grids

#####
##### Centered fourth-order advection scheme
#####

struct CenteredFourthOrder <: AbstractAdvectionScheme end

const C4 = CenteredFourthOrder

const centered_second_order = CenteredSecondOrder()

@inline ℑ³xᶠᵃᵃ(i, j, k, grid, u) = @inbounds u[i, j, k] - δxᶠᵃᵃ(i, j, k, grid, δxᶜᵃᵃ, u) / 6
@inline ℑ³xᶜᵃᵃ(i, j, k, grid, c) = @inbounds c[i, j, k] - δxᶜᵃᵃ(i, j, k, grid, δxᶠᵃᵃ, c) / 6

@inline ℑ³yᵃᶠᵃ(i, j, k, grid, v) = @inbounds v[i, j, k] - δyᵃᶠᵃ(i, j, k, grid, δyᵃᶜᵃ, v) / 6
@inline ℑ³yᵃᶜᵃ(i, j, k, grid, c) = @inbounds c[i, j, k] - δyᵃᶜᵃ(i, j, k, grid, δyᵃᶠᵃ, c) / 6

@inline ℑ³zᵃᵃᶠ(i, j, k, grid, w) = @inbounds w[i, j, k] - δzᵃᵃᶠ(i, j, k, grid, δzᵃᵃᶜ, w) / 6
@inline ℑ³zᵃᵃᶜ(i, j, k, grid, c) = @inbounds c[i, j, k] - δzᵃᵃᶜ(i, j, k, grid, δzᵃᵃᶠ, c) / 6

# Momentum

@inline fourth_order_momentum_flux_uu(i, j, k, grid, u)    = ℑxᶜᵃᵃ(i, j, k, grid, Ax_ψᵃᵃᶠ, u) * ℑxᶜᵃᵃ(i, j, k, grid, ℑ³xᶠᵃᵃ, u)
@inline fourth_order_momentum_flux_uv(i, j, k, grid, u, v) = ℑxᶠᵃᵃ(i, j, k, grid, Ay_ψᵃᵃᶠ, v) * ℑyᵃᶠᵃ(i, j, k, grid, ℑ³yᵃᶜᵃ, u)
@inline fourth_order_momentum_flux_uw(i, j, k, grid, u, w) = ℑxᶠᵃᵃ(i, j, k, grid, Az_ψᵃᵃᵃ, w) * ℑzᵃᵃᶠ(i, j, k, grid, ℑ³zᵃᵃᶜ, u)

@inline fourth_order_momentum_flux_vu(i, j, k, grid, u, v) = ℑyᵃᶠᵃ(i, j, k, grid, Ax_ψᵃᵃᶠ, u) * ℑxᶠᵃᵃ(i, j, k, grid, ℑ³xᶜᵃᵃ, v)
@inline fourth_order_momentum_flux_vv(i, j, k, grid, v)    = ℑyᵃᶜᵃ(i, j, k, grid, Ay_ψᵃᵃᶠ, v) * ℑyᵃᶜᵃ(i, j, k, grid, ℑ³yᵃᶠᵃ, v)
@inline fourth_order_momentum_flux_vw(i, j, k, grid, v, w) = ℑyᵃᶠᵃ(i, j, k, grid, Az_ψᵃᵃᵃ, w) * ℑzᵃᵃᶠ(i, j, k, grid, ℑ³zᵃᵃᶜ, v)

@inline fourth_order_momentum_flux_wu(i, j, k, grid, u, w) = ℑzᵃᵃᶠ(i, j, k, grid, Ax_ψᵃᵃᶠ, u) * ℑxᶠᵃᵃ(i, j, k, grid, ℑ³xᶜᵃᵃ, w)
@inline fourth_order_momentum_flux_wv(i, j, k, grid, v, w) = ℑzᵃᵃᶠ(i, j, k, grid, Ay_ψᵃᵃᶠ, v) * ℑyᵃᶠᵃ(i, j, k, grid, ℑ³yᵃᶜᵃ, w)
@inline fourth_order_momentum_flux_ww(i, j, k, grid, w)    = ℑzᵃᵃᶜ(i, j, k, grid, Az_ψᵃᵃᵃ, w) * ℑzᵃᵃᶜ(i, j, k, grid, ℑ³zᵃᵃᶠ, w)

# Periodic directions!

@inline momentum_flux_uu(i, j, k, grid, ::C4, u)    = fourth_order_momentum_flux_uu(i, j, k, grid, u)
@inline momentum_flux_uv(i, j, k, grid, ::C4, u, v) = fourth_order_momentum_flux_uv(i, j, k, grid, u, v)
@inline momentum_flux_uw(i, j, k, grid, ::C4, u, w) = fourth_order_momentum_flux_uw(i, j, k, grid, u, w)

@inline momentum_flux_vu(i, j, k, grid, ::C4, u, v) = fourth_order_momentum_flux_vu(i, j, k, grid, u, v)
@inline momentum_flux_vv(i, j, k, grid, ::C4, v)    = fourth_order_momentum_flux_vv(i, j, k, grid, v)
@inline momentum_flux_vw(i, j, k, grid, ::C4, v, w) = fourth_order_momentum_flux_vw(i, j, k, grid, v, w)

@inline momentum_flux_wu(i, j, k, grid, ::C4, u, w) = fourth_order_momentum_flux_wu(i, j, k, grid, u, w)
@inline momentum_flux_wv(i, j, k, grid, ::C4, v, w) = fourth_order_momentum_flux_wv(i, j, k, grid, v, w)
@inline momentum_flux_ww(i, j, k, grid, ::C4, w)    = fourth_order_momentum_flux_ww(i, j, k, grid, w)

# Bounded directions

@inline function momentum_flux_uu(i, j, k, grid::AbstractGrid{FT, <:Bounded}, ::C4, u) where FT
    if i > 1 && i < grid.Nx
        return fourth_order_momentum_flux_uu(i, j, k, grid, u)
    else
        return momentum_flux_uu(i, j, k, grid, centered_second_order, u)
    end
end

@inline function momentum_flux_uv(i, j, k, grid::AbstractGrid{FT, TX, <:Bounded}, ::C4, u, v) where {FT, TX}
    if j > 1 && j < grid.Ny
        return fourth_order_momentum_flux_uv(i, j, k, grid, u, v)
    else
        return momentum_flux_uv(i, j, k, grid, centered_second_order, u, v)
    end
end

@inline function momentum_flux_uw(i, j, k, grid::AbstractGrid{FT, TX, TY, <:Bounded}, ::C4, u, w) where {FT, TX, TY}
    if k > 1 && k < grid.Nz
        return fourth_order_momentum_flux_uw(i, j, k, grid, u, w)
    else
        return momentum_flux_uw(i, j, k, grid, centered_second_order, u, w)
    end
end

@inline function momentum_flux_vu(i, j, k, grid::AbstractGrid{FT, <:Bounded}, ::C4, u, v) where {FT}
    if i > 1 && i < grid.Nx
        return fourth_order_momentum_flux_vu(i, j, k, grid, u, v)
    else
        return momentum_flux_vu(i, j, k, grid, centered_second_order, u, v)
    end
end

@inline function momentum_flux_vv(i, j, k, grid::AbstractGrid{FT, TX, <:Bounded}, ::C4, v) where {FT, TX}
    if j > 1 && j < grid.Ny
        return fourth_order_momentum_flux_vv(i, j, k, grid, v)
    else
        return momentum_flux_vv(i, j, k, grid, centered_second_order, v)
    end
end

@inline function momentum_flux_vw(i, j, k, grid::AbstractGrid{FT, TX, TY, <:Bounded}, ::C4, v, w) where {FT, TX, TY}
    if k > 1 && k < grid.Nz
        return fourth_order_momentum_flux_vw(i, j, k, grid, v, w)
    else
        return momentum_flux_vw(i, j, k, grid, centered_second_order, v, w)
    end
end

@inline function momentum_flux_wu(i, j, k, grid::AbstractGrid{FT, <:Bounded}, ::C4, u, w) where FT
    if i > 1 && i < grid.Nx
        return fourth_order_momentum_flux_wu(i, j, k, grid, u, w)
    else
        return momentum_flux_wu(i, j, k, grid, centered_second_order, u, w)
    end
end

@inline function momentum_flux_wv(i, j, k, grid::AbstractGrid{FT, TX, <:Bounded}, ::C4, v, w) where {FT, TX}
    if j > 1 && j < grid.Ny
        return fourth_order_momentum_flux_wv(i, j, k, grid, v, w)
    else
        return momentum_flux_wv(i, j, k, grid, centered_second_order, v, w)
    end
end

@inline function momentum_flux_ww(i, j, k, grid::AbstractGrid{FT, TX, TY, <:Bounded}, ::C4, w) where {FT, TX, TY}
    if k > 1 && i < grid.Nz
        return fourth_order_momentum_flux_ww(i, j, k, grid, w)
    else
        return momentum_flux_ww(i, j, k, grid, centered_second_order, w)
    end
end

# Tracers

@inline fourth_order_advective_tracer_flux_x(i, j, k, grid, u, c) = Ax_ψᵃᵃᶠ(i, j, k, grid, u) * ℑxᶠᵃᵃ(i, j, k, grid, ℑ³xᶜᵃᵃ, c)
@inline fourth_order_advective_tracer_flux_y(i, j, k, grid, v, c) = Ay_ψᵃᵃᶠ(i, j, k, grid, v) * ℑyᵃᶠᵃ(i, j, k, grid, ℑ³yᵃᶜᵃ, c)
@inline fourth_order_advective_tracer_flux_z(i, j, k, grid, w, c) = Az_ψᵃᵃᵃ(i, j, k, grid, w) * ℑzᵃᵃᶠ(i, j, k, grid, ℑ³zᵃᵃᶜ, c)

@inline advective_tracer_flux_x(i, j, k, grid, ::C4, u, c) = fourth_order_advective_tracer_flux_x(i, j, k, grid, u, c)
@inline advective_tracer_flux_y(i, j, k, grid, ::C4, v, c) = fourth_order_advective_tracer_flux_y(i, j, k, grid, v, c)
@inline advective_tracer_flux_z(i, j, k, grid, ::C4, w, c) = fourth_order_advective_tracer_flux_z(i, j, k, grid, w, c)

@inline function advective_tracer_flux_x(i, j, k, grid::AbstractGrid{FT, <:Bounded}, ::C4, u, c) where FT
    if i > 1 && i < grid.Nx
        return fourth_order_advective_tracer_flux_x(i, j, k, grid, u, c)
    else
        return advective_tracer_flux_x(i, j, k, grid, centered_second_order, u, c)
    end
end

@inline function advective_tracer_flux_y(i, j, k, grid::AbstractGrid{FT, TX, <:Bounded}, ::C4, v, c) where {FT, TX}
    if j > 1 && i < grid.Ny
        return fourth_order_advective_tracer_flux_y(i, j, k, grid, v, c)
    else
        return advective_tracer_flux_y(i, j, k, grid, centered_second_order, v, c)
    end
end

@inline function advective_tracer_flux_z(i, j, k, grid::AbstractGrid{FT, TX, TY, <:Bounded}, ::C4, w, c) where {FT, TX, TY}
    if k > 1 && i < grid.Nz
        return fourth_order_advective_tracer_flux_z(i, j, k, grid, w, c)
    else
        return advective_tracer_flux_z(i, j, k, grid, centered_second_order, w, c)
    end
end


#
# This function fits penalized variance component model with responses y, design matrix X
# and variance components in V by a multiplicative MM algorithm.
#
using Polynomial
function VClasso(y::Vector{Float64}, X::Matrix{Float64},
                 V::Array{Matrix{Float64}, 1}, # Z::Matrix{Float64}, # add Z for Longitudinal design
                 λfactor::Vector{Int64}, # 0 is not penalized, 1 is penalized to 1.0
                       print::Bool=false, λgrid::Vector{Float64}=[0.0])
  n = length(y)   # no. observations
  n2 = n * n
  p = size(X, 2)  # no. predictors
  m = length(V)   # no. variance components
  ngrid = length(λgrid);
  σ2grid = zeros(m, ngrid);
  Itergrid = zeros(ngrid);
  LogLgrid = zeros(ngrid);
  Objgrid = zeros(ngrid);
  # initialize algorithm
  for k = 1:ngrid
   β = X \ y ;
   σ2 = 1e-2 * ones(m);
   σ = sqrt(σ2);
   Ω = zeros(n, n);
   for j = 1:m
     BLAS.axpy!(n2, σ2[j], V[j], 1, Ω, 1)
   end
   Ωchol = cholfact(Ω);
   Ωinv = inv(Ωchol);
   res = y - X * β;
   v = Ωchol \ res ;
   ##
   loglConst = - 0.5 * n * log(2.0 * pi);
   logl = loglConst - 0.5 * logdet(Ωchol) - 0.5 * dot(res, v)
   Iter = 1 :: Int
   logL = 0.0
   Obj = 0.0

   # MM loop
   maxiter = 10000::Int
   tol = 1e-4
   λ = λgrid[k];
   logllasso = -logl + λ * sumabs(λfactor[1:end].*σ[1:end-1]);
   for i = 1:maxiter
    if print
      println(i, " ",  "Obj ", logllasso, " ", "logL", logl)
    end

    # update variance components
    fill!(Ω, 0.0);
    for l = 1: m - 1
      quartic = BLAS.dot(n2, Ωinv, 1, V[l], 1)
      cubic = λ * λfactor[l]
      constant = -σ[l]^4 * ((v'* V[l])*v)[]
      @inbounds σ[l] = min_root(quartic, cubic, constant, l, V, σ, res)
      @inbounds σ2[l] = σ[l] .^2
      @inbounds BLAS.axpy!(n2, σ2[l], V[l], 1, Ω, 1)
    end
    σ2[m] *= BLAS.nrm2(v) / sqrt(trace(Ωinv))
    σ[m] = sqrt(σ2[m])
    BLAS.axpy!(n2, σ2[m], V[m], 1, Ω, 1)
    Ωchol = cholfact(Ω, :L);
    Ωinv = inv(Ωchol);
    Xnew = Ωchol[:L] \ X;
    ynew = Ωchol[:L] \ y;
    # update fixed effects
    β = Xnew \ ynew;
    res = y - X * β;
    v = Ωchol \ res;
    # check convergence
    logllassoold = logllasso
    logl = loglConst - 0.5 * logdet(Ωchol) - 0.5 * dot(res, v)
    logllasso = -logl + λ * sumabs(λfactor[1:end].*σ[1:end-1]);
    if abs(logllasso - logllassoold) < tol * (abs(logllassoold) + 1.0)
      #println(i," ", logl)
      Obj = logllasso
      logL = logl;
      Iter = i
      break
    end
   end # end of loop over maxiter
   σ2grid[:, k] = σ2;
   Itergrid[k] = Iter;
   LogLgrid[k] = logL;
   Objgrid[k] = Obj;
  end # end of loop over λgrid
    # output
  return σ2grid, Itergrid, LogLgrid, Objgrid
end

# this function returns the \sigma that minimize the objective function: negative loglikelihood plus penalty term
function min_root(quartic, cubic, constant, index, V, σ, res)
  # find the real positive root
  complex_root = roots(Poly([quartic, cubic, 0, 0, constant]))
  # keep real roots
  real_root = real(complex_root[imag(complex_root) .== 0])
  real_pos_root = real_root[real_root .> 0]
  # find the real negative root
  complex_root = roots(Poly([quartic, -cubic, 0, 0, constant]))
  real_root = real(complex_root[imag(complex_root) .== 0])
  real_neg_root = real_root[real_root .< 0]
  root = [real_pos_root, 0.0, real_neg_root]
  obj = copy(root);
  m = length(σ)
  n = size(V[1])[1]
  n2 = n^2
  for i = 1:length(root)
    σ[index] = root[i]
    σ2 = σ .^2
    Ω = zeros(n, n);
    for j = 1:m
      BLAS.axpy!(n2, σ2[j], V[j], 1, Ω, 1)
    end
    Ωchol = cholfact(Ω);
    Ωinv = inv(Ωchol);
    v = Ωchol \ res ;
    ##
    obj[i] =  0.5 * logdet(Ωchol) + 0.5 * dot(res, v) + cubic * sumabs(σ[1:end-1])
  end
  return root[indmin(obj)]

end


include("../benchmark.jl")

using JuMP
using Ipopt

function formulation1(lambda,w,mu,mu1,A,B,T, lb_var)
    m = Model(solver = IpoptSolver())
    @variable(m, c[i in A,j in B] >= lb_var)# for (i,j) in T)
    @variable(m, y[i in A,j in B] >= lb_var)# for (i,j) in T )

    @NLobjective(m, Min, -sum( lambda[i,j] * (log(c[i,j]) - (y[i,j]/w[i])^mu1[j] / mu1[j]) for (i,j) = T) )

    for (i,j) in T
      for (p,q) in T
        if i != p || j != q
          @NLconstraint(m,(log(c[i,j]) - (y[i,j]/w[i])^mu1[j] / mu1[j]) - (log(c[p,q]) - (y[p,q]/w[i])^mu1[j] / mu1[j]) >= 0.0)
        end
      end
    end

    @NLconstraint(m, sum(lambda[i,j] * (y[i,j] - c[i,j]) for (i,j) in T)  >= 0.0);
    return m
end

na = 21
nb = 5

A = 1:na;
B = 1:nb;
T = [(i,j) for i in A for j in B];

lambda = ones(na,nb);
wmin = 1;
#wmax = wmin + na - 1;
wmax = 5
w = [wmin + ((wmax-wmin)/(na-1))*(i-1) for i = 1:na]
mu = [1.0, 2.0, 3.0, 5.0, 8.0]
#mu = [1.0, 1.5, 2.0, 2.5, 3.0]

gamma = 0.0

mu1 = mu + 1.0

#lb = 0.0
lb_var = 0.1

srand(1)
m = formulation1(lambda,w,mu,mu1,A,B,T, lb_var)
iter = one_phase_solve(m);

#=
ORG_STDOUT = STDOUT
file = open("test.txt", "w")
redirect_stdout(file)

srand(1)
m = create_tax_problem(10)
one_phase_solve(m)

STDOUT = ORG_STDOUT

close(file)=#


using JuMP, Gurobi
include("function.jl")

SolverBrendon.SolveMIP(m)


using DataFrames
using CSV
using Random
using Plots
using PDMats
using Printf
include("transforms.jl")
include("statistics.jl")
include("validation.jl")
include("plotting.jl")

function gridsearch(f, a, b, numpts=45)
    h = (b-a)/(numpts-1)
    l, m = 0, 0
    for i = a:h:b
        val = f(i);
        if val>m
            m = val; l = i;
        end
    end
    return l, m
end

#macroeconomics time series data
df = DataFrame(CSV.File("data\\TB3MS.csv"))

z = convert(Matrix, df[:,2:2])
n = length(z)
s = collect(1:n) #data points equally spaced in time
portion_train = 0.3

rng = MersenneTwister(1234)
ind = shuffle(rng, s)
id_train = view(ind, 1:Int(ceil(n*portion_train)))
id_test = view(ind, (Int(ceil(n*portion_train))+1):n)

train_x = s[id_train]; test_x = s[id_test]
train_y = z[id_train]; test_y = z[id_test]

X = ones(length(train_x), 1); X0 = [1]
pλ = x -> 1; pθ = x -> 1
range_theta = [0.1 5]
range_lambda = [-5 5]
results = Array{Float64}(undef, length(test_x))#preallocate space for medians and interval widths
intwidths = Array{Float64}(undef, length(test_x))

if true #test and time pdfn and cdfn functions
pdf = model(X, X0, train_x, [test_x[6]], boxCox, boxCoxPrime, pθ, pλ, train_y, range_theta, range_lambda)
cdf = z0 ->  int1D(pdf, 0, z0, "0")
@time begin
un = cdf(30)
end
cdfn = x -> cdf(x)/un #normalized cdf
pdfn = x -> pdf(x)/un #normalized pdf
function testcdfn(x)
    @time begin
        r = cdfn(x)
    end
    return r
end

function testbisection(f, a, b)
    @time begin
        val = bisection(f, a, b)
    end
    return val
end

function testconfidence(f, med)
    @time begin
        val = confidence(f, med)
    end
    return val
end
plt2(pdfn, 0, 20)
end

if false #simply plot at multiple locations
    pltX = zeros(20, length(test_x))
    pltY = zeros(20, length(test_x))
    for i=1:length(test_x)#single-point predictions
        println(i)
        pdf = model(X, X0, train_x, [test_x[i]], boxCox, boxCoxPrime, pθ, pλ, train_y, range_theta, range_lambda)
        cdf = z0 ->  int1D(pdf, 0, z0, "1")
        un = cdf(30)
        cdfn = x -> cdf(x)/un #normalized cdf
        pdfn = x -> pdf(x)/un #normalized pdf
        x, y = plt(pdfn, 5, 18)
        pltX[:, i] = x; pltY[:, i] = y
    end
    display(plot(pltX, pltY, layout = length(test_x), ylim = (0, )))
end

if false #determine medians and confidence intervals
for i=1:length(test_x)#single-point predictions
    println(i)
    @time begin
    pdf = model(X, X0, train_x, [test_x[i]], boxCox, boxCoxPrime, pθ, pλ, train_y, range_theta, range_lambda)
    cdf = z0 ->  int1D(pdf, 0, z0, "0")
    @time begin
    un = cdf(25)
    end
    cdfn = x -> cdf(x)/un #normalized cdf
    pdfn = x -> pdf(x)/un #normalized pdf
    @time begin
    med = bisection(x -> cdfn(x)-0.5, 1e-3, 25, 1e-2, 9)
    #med, val = gridsearch(pdfn, 0, 30, 20) #MAP estimation
    @printf("ith pt med: %f \n", med)
    end
    @time begin
    intwidths[i] = confidence(pdfn, med) #compute width of credible interval corresponding to 95% confidence level
    @printf("ith pt intwidth: %f \n", intwidths[i])
    end
    results[i] = med
    end
end
display(plot(collect(1:n), z));
display(plot!(train_x, train_y, seriestype = scatter));
display(plot!(test_x, results, seriestype = :scatter));
#display(plot!(test_x, test_y, seriestype = :scatter));
display(plot!(repeat(test_x, 2), vcat(results-intwidths, results+intwidths), seriestype = :scatter, color=:blue))
#display(plot!(test_x, test_y))
savefig("treasury6.png");
end

if true #cross validation on training set
    _, Xs, Ys = cross_validate(X, train_x, boxCox, boxCoxPrime, pθ, pλ, train_y, range_theta, range_lambda, 500, 2, 24)
    display(plot(Xs, Ys, layout = length(train_y), legend=nothing, xtickfont = font(10, "Courier"), lw=0.5))
    display(plot!(train_y', layout = length(z), legend=nothing, seriestype = :vline, xtickfont = font(10, "Courier"), lw=0.5))
    savefig("results//treasury//treasury_delete_one_cross_validate2.jpg")
end


# Array subindexing:
function Fgetrow(x::Array,inrow)
    row=Int(inrow[1])
    if row>0
       tmp=zeros(1,size(x,2))
        tmp[:]=x[row,:]
        return (tmp, zeros(size(x)))
    else
        return (zeros(size(x,2)), zeros(size(x)))
    end
end

function Fgetrow_inplace(value,auxvalue,x,inrow)
    row=Int(inrow[1])
    if row>0
        tmp=zeros(1,size(x,2))
        tmp[:]=x[row,:]
        copy!(value,tmp)
        copy!(auxvalue,zeros(size(x)))
    else
        return
        copy!(value,zeros(size(x,2)))
        copy!(auxvalue,zeros(size(x)))
    end
end

function Dgetrow(derivativeIDX,f_c,faux_c,grad_c,grad_n,x,inrow)
    row=Int(inrow[1]) # this is ugly since currently we force all values stored on the graph to be real arrays.
    if derivativeIDX==1
        if row>0
            faux_c[row,:]=grad_c
            axpy!(1.0,faux_c,grad_n)
        end
    end
end

#TODO: gpu version

export Fgetrow

Derivative[Fgetrow]=Dgetrow
Inplace[Fgetrow]=Fgetrow_inplace

import Base.*
#getindex(A::ADnode,:,i::Int)=ADnode(Fgetcol,[A ADint(i)])
getindex(A::ADnode,i::Int,::Colon)=ADnode(Fgetrow,[A ADint(i)])
getrow(A::ADnode,i)=ADnode(Fgetrow,[A ADint(i)])

export getindex
export getrow


@inbounds function _normalize(point::SVector{n}, min_bound::SVector{n}, max_bound::SVector{n}, scale::Real) where {n}
    return (point .- min_bound) ./ scale
end
@inbounds function _normalize(spl::AbstractNormalSpline{n}, point::SVector{n}) where {n}
    return _normalize(point, _get_min_bound(spl), _get_max_bound(spl), _get_scale(spl))
end

@inbounds function _unnormalize(point::SVector{n}, min_bound::SVector{n}, max_bound::SVector{n}, scale::Real) where {n}
    return min_bound .+ scale .* point
end
@inbounds function _unnormalize(spl::AbstractNormalSpline{n}, point::SVector{n}) where {n}
    return _unnormalize(point, _get_min_bound(spl), _get_max_bound(spl), _get_scale(spl))
end

function _normalization_scaling(nodes::AbstractVecOfSVecs)
    min_bound = reduce((x, y) -> min.(x, y), nodes; init = fill(+Inf, eltype(nodes)))
    max_bound = reduce((x, y) -> max.(x, y), nodes; init = fill(-Inf, eltype(nodes)))
    scale = maximum(max_bound .- min_bound)
    return min_bound, max_bound, scale
end

function _normalization_scaling(nodes::AbstractVecOfSVecs, d_nodes::AbstractVecOfSVecs)
    min_bound = min.(reduce((x, y) -> min.(x, y), nodes; init = fill(+Inf, eltype(nodes))), reduce((x, y) -> min.(x, y), d_nodes; init = fill(+Inf, eltype(d_nodes))))
    max_bound = max.(reduce((x, y) -> max.(x, y), nodes; init = fill(-Inf, eltype(nodes))), reduce((x, y) -> max.(x, y), d_nodes; init = fill(-Inf, eltype(d_nodes))))
    scale = maximum(max_bound .- min_bound)
    return min_bound, max_bound, scale
end

function _estimate_accuracy(spl::AbstractNormalSpline{n,T,RK}) where {n,T,RK <: ReproducingKernel_0}
    vmax = maximum(abs, _get_values(spl))
    rmae = zero(T)
    @inbounds for i in 1:length(_get_nodes(spl))
        point = _unnormalize(spl, _get_nodes(spl)[i])
        σ     = _evaluate(spl, point)
        rmae  = max(rmae, abs(_get_values(spl)[i] - σ))
    end
    if vmax > 0
        rmae /= vmax
    end
    rmae = max(rmae, eps(T))
    res  = -floor(log10(rmae)) - 1
    res  = max(res, 0)
    return trunc(Int, res)
end

function _pairwise_sum_norms(nodes::AbstractVecOfSVecs{n,T}) where {n,T}
    ℓ = zero(T)
    @inbounds for i in 1:length(nodes), j in i:length(nodes)
        ℓ += norm(nodes[i] .- nodes[j])
    end
    return ℓ
end

function _pairwise_sum_norms_weighted(nodes::AbstractVecOfSVecs{n,T}, d_nodes::AbstractVecOfSVecs{n,T}, w_d_nodes::T) where {n,T}
    ℓ = zero(T)
    @inbounds for i in 1:length(nodes), j in 1:length(d_nodes)
        ℓ += norm(nodes[i] .- w_d_nodes .* d_nodes[j])
    end
    return ℓ
end

@inline _ε_factor(::RK_H0, ε::T) where {T} = one(T)
@inline _ε_factor(::RK_H1, ε::T) where {T} = T(3)/2
@inline _ε_factor(::RK_H2, ε::T) where {T} = T(2)

@inline _ε_factor_d(::RK_H0, ε::T) where {T} = one(T)
@inline _ε_factor_d(::RK_H1, ε::T) where {T} = T(2)
@inline _ε_factor_d(::RK_H2, ε::T) where {T} = T(5)/2

function _estimate_ε(k::ReproducingKernel_0, nodes)
    ε  = _estimate_ε(nodes)
    ε *= _ε_factor(k, ε)
    k  = typeof(k)(ε)
end

function _estimate_ε(k::ReproducingKernel_0, nodes, d_nodes)
    ε  = _estimate_ε(nodes, d_nodes)
    ε *= _ε_factor_d(k, ε)
    k  = typeof(k)(ε)
end

function _estimate_ε(nodes::AbstractVecOfSVecs{n,T}) where {n,T}
    n₁ = length(nodes)
    ε  = _pairwise_sum_norms(nodes)
    return ε > 0 ? ε * T(n)^T(inv(n)) / T(n₁)^(T(5) / 3) : one(T)
end

function _estimate_ε(nodes::AbstractVecOfSVecs{n,T}, d_nodes::AbstractVecOfSVecs{n,T}, w_d_nodes::T=T(0.1)) where {n,T}
    n₁ = length(nodes)
    n₂ = length(d_nodes)
    ε  = _pairwise_sum_norms(nodes) + _pairwise_sum_norms_weighted(nodes, d_nodes, w_d_nodes) + w_d_nodes * _pairwise_sum_norms(d_nodes)
    return ε > 0 ? ε * T(n)^T(inv(n)) / T(n₁ + n₂)^(T(5) / 3) : one(T)
end

function _estimate_epsilon(nodes::AbstractVecOfSVecs, kernel::ReproducingKernel_0)
    min_bound, max_bound, scale = _normalization_scaling(nodes)
    nodes = _normalize.(nodes, (min_bound,), (max_bound,), scale)
    ε     = _estimate_ε(nodes)
    ε    *= _ε_factor(kernel, ε)
    return ε
end

function _estimate_epsilon(nodes::AbstractVecOfSVecs, d_nodes::AbstractVecOfSVecs, kernel::ReproducingKernel_1)
    min_bound, max_bound, scale = _normalization_scaling(nodes, d_nodes)
    nodes   = _normalize.(nodes, (min_bound,), (max_bound,), scale)
    d_nodes = _normalize.(d_nodes, (min_bound,), (max_bound,), scale)
    ε       = _estimate_ε(nodes, d_nodes)
    ε      *= _ε_factor_d(kernel, ε)
    return ε
end

function _get_gram(nodes::AbstractVecOfSVecs, kernel::ReproducingKernel_0)
    min_bound, max_bound, scale = _normalization_scaling(nodes)
    nodes = _normalize.(nodes, (min_bound,), (max_bound,), scale)
    if kernel.ε == 0
        kernel = _estimate_ε(kernel, nodes)
    end
    return _gram(nodes, kernel)
end

function _get_gram(nodes::AbstractVecOfSVecs, d_nodes::AbstractVecOfSVecs, d_dirs::AbstractVecOfSVecs, kernel::ReproducingKernel_1)
    min_bound, max_bound, scale = _normalization_scaling(nodes, d_nodes)
    nodes   = _normalize.(nodes, (min_bound,), (max_bound,), scale)
    d_nodes = _normalize.(d_nodes, (min_bound,), (max_bound,), scale)
    d_dirs  = d_dirs ./ norm.(d_dirs)
    if kernel.ε == 0
        kernel = _estimate_ε(kernel, nodes, d_nodes)
    end
    return _gram(nodes, d_nodes, d_dirs, kernel)
end

function _get_cond(nodes::AbstractVecOfSVecs, kernel::ReproducingKernel_0)
    T = promote_type(eltype(kernel), eltype(eltype(nodes)))
    gram = _get_gram(nodes, kernel)
    cond = zero(T)
    try
        evs = svdvals!(gram)
        maxevs = maximum(evs)
        minevs = minimum(evs)
        if minevs > 0
            cond = maxevs / minevs
            cond = T(10)^floor(log10(cond))
        end
    catch
    end
    return cond
end

function _get_cond(nodes::AbstractVecOfSVecs, d_nodes::AbstractVecOfSVecs, d_dirs::AbstractVecOfSVecs, kernel::ReproducingKernel_1)
    _get_cond(nodes, kernel)
end

"""
Get estimation of the Gram matrix condition number
Brás, C.P., Hager, W.W. & Júdice, J.J. An investigation of feasible descent algorithms for estimating the condition number of a matrix. TOP 20, 791–809 (2012).
https://link.springer.com/article/10.1007/s11750-010-0161-9
"""
function _estimate_cond(A::AbstractMatrix{T}, F::Factorization, nit = 3) where {T}
    Anorm = norm(A, 1)
    n = size(A, 1)
    x = fill(inv(T(n)), n)
    z = zeros(T, n)
    z′ = zeros(T, n)
    γ = zero(T)
    @inbounds for _ in 1:nit
        ldiv!(z′, F, x)
        γ = zero(T)
        for i in 1:n
            γ += abs(z′[i])
            z′[i] = sign(z′[i])
        end
        ldiv!(z, F, z′)
        zdotx = z ⋅ x
        zmax, imax = T(-Inf), 1
        for i in 1:n
            zᵢ = z[i] = abs(z[i])
            if zᵢ > zmax
                zmax, imax = zᵢ, i
            end
        end
        (zmax <= zdotx) && break
        x .= 0
        x[imax] = 1
    end
    cond = prevpow(T(10), Anorm * γ)
    return cond
end


using DynamicalSystems
using CausalityTools
using Plots
plotlyjs()

s = DynamicalSystems.Systems.coupledstandardmaps(2)
ts_length = 500
orbit = trajectory(s, ts_length - 1, Ttr = 1000)
x1, x2, y1, y2 = ([orbit[:, i] for i = 1:4]...,)

px = plot(x1, label = "x1", lc = :blue)
plot!(x2, label = "x2", lc = :blue, ls = :dash)
ylabel!("Value")

py = plot(y1, label = "y1", lc = :green)
plot!(y2, label = "y2", lc = :green, ls = :dash)
xlabel!("Time step")
ylabel!("Value")

plot(px, py, layout = (2, 1))

E = StateSpaceReconstruction.embed([x1, x2, y2], [1, 2, 3], [0, -10, -5])
E1, E2, E3 = ([E.points[i, :] for i = 1:3]...,)
scatter(E1, E2, E3, legend = false)

using RecipesBase
@recipe function f(E::StateSpaceReconstruction.AbstractEmbedding; vars = [1, 2, 3],
        mc = :black, ms = 2, lc = :black, lw = 1)
    if size(E.points, 1) < 3
        error("Cannot plot reconstruction with < 3 variables.")
    end
    legend --> false
    xlabel --> "x1"
    ylabel --> "x2"
    zlabel --> "x3"
    markercolor --> mc
    markersize --> ms
    ([E.points[i, :] for i in vars]...,)
end

scatter(E, ms = 0.5)

#########################################################
# Helper functions to plot 3D rectangles.
#########################################################
"""
    rectangle3dpts(x, y, z, ϵx, ϵy, ϵz)

Compute the coordinates of the vertices of a 3D rectangle, given the
coordinates of the origin (`x`, `y`, and `z`) and the corresponding edge lengths
(`ϵx`, `ϵy`, `ϵz`).
"""
function rectangle3dpts(x, y, z, ϵx, ϵy, ϵz)
    v1b = [x,    y,    z]
    v2b = [x+ϵx, y,    z]
    v3b = [x+ϵx, y+ϵy, z]
    v4b = [x,    y+ϵy, z]
    v1t = [x,    y,    z+ϵz]
    v2t = [x+ϵx, y,    z+ϵz]
    v3t = [x+ϵx, y+ϵy, z+ϵz]
    v4t = [x,    y+ϵy, z+ϵz]
    (v1b, v2b, v3b, v4b, v1t, v2t, v3t, v4t)
end

"""
    rectangle3dpts(o, ϵ)

Given the origin `o` (3-element vector) and edge lengths `ϵ` (also a
3-element vector) of a 3D rectangle, return the coordinates of its
vertices.
"""
function rectangle3dpts(o, ϵ)
    x, y, z = (o...,)
    ϵx, ϵy, ϵz = (ϵ...,)
    rectangle3dpts(x, y, z, ϵx, ϵy, ϵz)
end

"""
    connectvertices(o, ϵ)

Given the origin `o` (3-element vector) and edge lengths `ϵ` (also a
3-element vector) of a 3D rectangle, return a vector of line segments
(each a dim-by-n_vertices array) that when plotted together yields
the entire rectangle.
"""
function connectvertices(o, ϵ)
    # Get the vertices
    v1b, v2b, v3b, v4b, v1t, v2t, v3t, v4t = rectangle3dpts(o, ϵ)

    # Connect vertices in top and bottom planes
    connect_top = zeros(3, 5)
    connect_bottom = zeros(3, 5)

    # Connect corners of the top and bottom planes
    corner1 = zeros(3, 2)
    corner2 = zeros(3, 2)
    corner3 = zeros(3, 2)
    corner4 = zeros(3, 2)
    connect_bottom[:, 1:5] = [v1b v2b v3b v4b v1b]
    connect_top[:, 1:5]    = [v1t v2t v3t v4t v1t]
    corner1[:, 1:2] = [v1b v1t]
    corner2[:, 1:2] = [v2b v2t]
    corner3[:, 1:2] = [v3b v3t]
    corner4[:, 1:2] = [v4b v4t]

    return [connect_top, connect_bottom,
            corner1, corner2, corner3, corner4]
end

"""
    splitaxes(x)

Return a vector of the individual components of an array of points
provided as an array where each point is a column.
"""
splitaxes(x) = ([x[k, :] for k = 1:size(x, 1)]...,)

"""
    plot_3D_rect!(p, origin, edgelengths;
             lc = :black, lw = 1.0, ls = :solid)

Append a a 3D rectangle to an existing plot `p`, given the
origin `o` (3-element vector) and edge lengths
`ϵ` (3-element vector) of the rectangle.
"""
function plot_3D_rect!(p, o, ϵ;
             lc = :black, lw = 1.0, ls = :solid, lα = 0.7)

    linesegments = connectvertices(o, ϵ)
    for segment in linesegments
        plot!(p, splitaxes(segment),
            lc = lc, lw = lw, ls = ls, lα = lα)
    end
end

"""
Given the `origin` of a hyperrectangle and the `stepsizes` along each
coordinate axis, divide each axis into `N` stepsizes starting
from `origin[i] + (stepsizes[i]/N)/2`
"""
function fill_hyperrectangle_3D(origin, stepsizes, N = 10)
    xs, ys, zs = ([LinRange(origin[i] + stepsizes[i]/N/2, origin[i] + stepsizes[i] - stepsizes[i]/N/2, N) for i = 1:length(origin)]...,)
    pts = zeros(Float64, 3, length(xs)^3)
    i = 0
    for x in xs
        for y in ys
            for z in zs
                i += 1
                pts[:, i] = [x, y, z]
            end
        end
    end
    pts
end

@recipe function f(N, riv::RectangularInvariantMeasure;
        lw = 0.8, lc = :black, lα = 0.5, ls = :do,
        ms = 0.2, mc = :black, mα = 0.3)

    pts = riv.E.points
    ϵ = riv.ϵ
    axisminima, stepsizes = minima_and_stepsizes(pts, ϵ)

    # All bins are assigned a bin origin.
    v = unique(riv.visited_bins_coordinates, dims = 2)

    # Bins may be visited by several times, so we'll get rid
    # of the repetitions.
    n_visited_boxes = size(v, 2)
    D = size(v, 1)

    # Plot the partition grid over the points of the reconstructed
    # orbit.
    legend --> false

    for i = 1:n_visited_boxes
        origin = riv.visited_bins_coordinates[:, i]
        # Get the measure of the box, and fill the box with a
        # number of points scaled to the measure.
        μ = riv.measure.dist[i] * 50
        fillpts = fill_hyperrectangle_3D(origin, stepsizes, N)
        @series begin
            seriestype := :scatter3d
            markersize --> μ/2
            markeralpha --> μ
            markercolor --> mc
            linecolor --> mc
            fillcolor --> :transparent
            splitaxes(fillpts)
        end

        linesegments = connectvertices(origin, stepsizes)

        for segment in linesegments
            @series begin
                seriestype := :path
                linealpha --> lα
                linewidth --> lw
                linecolor --> lc
                splitaxes(segment)
            end
        end
    end
end

@recipe function f(r::RectangularBinningTransferOperator)
    seriestype  :=  :heatmap
    xlabel --> "state #i"
    ylabel --> "state #j"
    r.transfermatrix
end

@recipe function f(id::InvariantDistribution)
    seriestype  :=  :bar
    xlabel --> "state # (i)"
    ylabel --> "invariant density"
    linecolor --> :black
    fillcolor --> :black
    fillalpha --> 0.5
    legend --> :none
    id.dist
end

@recipe function plot_grid(pts::AbstractArray{T, 2}, riv::RectangularInvariantMeasure; N = 5,
        lw = 0.8, lc = :black, lα = 0.5, ls = :do,
        ms = 0.2, mc = :black, mα = 0.3) where {T}

    axisminima, stepsizes = minima_and_stepsizes(pts, riv.ϵ)

    # All bins are assigned a bin origin.
    v = unique(riv.visited_bins_coordinates, dims = 2)

    # Bins may be visited by several times, so we'll get rid
    # of the repetitions.
    n_visited_boxes = size(v, 2)
    D = size(v, 1)

    # Plot the partition grid over the points of the reconstructed
    # orbit.
    legend --> false

    for i = 1:n_visited_boxes
        origin = riv.visited_bins_coordinates[:, i]
        # Get the measure of the box, and fill the box with a
        # number of points scaled to the measure.
        μ = riv.measure.dist[i]

        println(ceil(Int, μ*100))
        fillpts = fill_hyperrectangle_3D(origin, stepsizes, ceil(Int, μ*100))
        @series begin
            seriestype := :scatter3d
            markersize --> μ
            markeralpha --> μ
            markercolor --> mc
            linecolor --> mc
            fillcolor --> mc
            splitaxes(fillpts)
        end

        linesegments = connectvertices(origin, stepsizes)

        for segment in linesegments
            @series begin
                seriestype := :path
                linealpha --> lα
                linewidth --> lw
                linecolor --> lc
                splitaxes(segment)
            end
        end
    end
end

invm = rectangularinvariantmeasure(E, 4)

gr()
p1 = plot(invm.transfermatrix);
p2 = plot(invm.measure);
p3 = scatter(E, ms = 1.0);
p4 = plot(5, invm, lα = 0.2, ls = :dot,
                    ms = 0.5, mα = 0.5, mc = :blue);

plot(plot(p1, p2, link = :x, layout = (2, 1)),
    plot(p3, p4, layout = (2, 1)), layout = (1, 2))
savefig("docs/src/scripts/test.pdf")