# ---
# title: 1387. Sort Integers by The Power Value
# id: problem1387
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Sort, Graph
# link: <https://leetcode.com/problems/sort-integers-by-the-power-value/description/>
# hidden: true
# ---
#
# The power of an integer `x` is defined as the number of steps needed to
# transform `x` into `1` using the following steps:
#
#   * if `x` is even then `x = x / 2`
#   * if `x` is odd then `x = 3 * x + 1`
#
# For example, the power of x = 3 is 7 because 3 needs 7 steps to become 1 (3
# --> 10 --> 5 --> 16 --> 8 --> 4 --> 2 --> 1).
#
# Given three integers `lo`, `hi` and `k`. The task is to sort all integers in
# the interval `[lo, hi]` by the power value in **ascending order** , if two or
# more integers have **the same** power value sort them by **ascending order**.
#
# Return the `k-th` integer in the range `[lo, hi]` sorted by the power value.
#
# Notice that for any integer `x` `(lo <= x <= hi)` it is **guaranteed** that
# `x` will transform into `1` using these steps and that the power of `x` is
# will **fit** in 32 bit signed integer.
#
#
#
# **Example 1:**
#
#
#
#     Input: lo = 12, hi = 15, k = 2
#     Output: 13
#     Explanation: The power of 12 is 9 (12 --> 6 --> 3 --> 10 --> 5 --> 16 --> 8 --> 4 --> 2 --> 1)
#     The power of 13 is 9
#     The power of 14 is 17
#     The power of 15 is 17
#     The interval sorted by the power value [12,13,14,15]. For k = 2 answer is the second element which is 13.
#     Notice that 12 and 13 have the same power value and we sorted them in ascending order. Same for 14 and 15.
#
#
# **Example 2:**
#
#
#
#     Input: lo = 1, hi = 1, k = 1
#     Output: 1
#
#
# **Example 3:**
#
#
#
#     Input: lo = 7, hi = 11, k = 4
#     Output: 7
#     Explanation: The power array corresponding to the interval [7, 8, 9, 10, 11] is [16, 3, 19, 6, 14].
#     The interval sorted by power is [8, 10, 11, 7, 9].
#     The fourth number in the sorted array is 7.
#
#
# **Example 4:**
#
#
#
#     Input: lo = 10, hi = 20, k = 5
#     Output: 13
#
#
# **Example 5:**
#
#
#
#     Input: lo = 1, hi = 1000, k = 777
#     Output: 570
#
#
#
#
# **Constraints:**
#
#   * `1 <= lo <= hi <= 1000`
#   * `1 <= k <= hi - lo + 1`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


#= for significands in [1, 2), exponents are powers of 2
exponent_max(::Type{Float16})  =   15
exponent_max(::Type{Float32})  =  127
exponent_max(::Type{Float64})  = 1023
exponent_min(::Type{Float16})  = 1 -   15
exponent_min(::Type{Float32})  = 1 -  127
exponent_min(::Type{Float64})  = 1 - 1023
=#
# for significands in [0.5, 1)
exponent_max(::Type{Float16})  =   15+1
exponent_max(::Type{Float32})  =  127+1
exponent_max(::Type{Float64})  = 1023+1
exponent_min(::Type{Float16})  = 2 -   15
exponent_min(::Type{Float32})  = 2 -  127
exponent_min(::Type{Float64})  = 2 - 1023

for (U,F) in ((:UInt64, :Float64), (:UInt32, :Float32), (:UInt16, :Float16))
    @eval begin
        function randf(::Type{$F})
            r = $F(Inf)
            while !isfinite(r)
                r = reinterpret($F, rand($U))
            end
            return r
        end
        function randfloat(::Type{$F}, n::Int=1; emin::Int=exponent_min($F), emax::Int=exponent_max($F), signed::Bool=false)
            emin, emax = minmax(emin, emax)
            emin = max(exponent_min($F), emin)
            emax = min(exponent_max($F), emax)
            n = max(1, n)
            result = rand($F, n)
            # significands in [0.5, 1)
            result *= 0.5
            result += 0.5
            exponent = rand(emin:1:emax, n)
            for i in 1:n
                result[i] = ldexp(result[i], exponent[i])
            end
            if signed
               signs = rand(-1:2:1,n)
               for i in 1:n
                   result[i] = copysign(result[i], signs[i])
               end
            end
            return result
        end
    end
end


include("activate.jl")

import Random
using SparseArrays, LinearAlgebra
using Test
using Plots

const nfeat = 100 * 5
const nobs = 500

# rank of the real data
const r = 30
const Xreal = Matrix{Float64}(undef, nobs, nfeat)

const X_gen_cols = randn(nfeat, r)
const X_gen_rows = randn(r, nobs)
const svals = 100 * rand(r)
for i in 1:nobs
    for j in 1:nfeat
        Xreal[i, j] = sum(X_gen_cols[j, k] * X_gen_rows[k, i] * svals[k] for k in 1:r)
    end
end

nucnorm(Xmat) = sum(abs(σi) for σi in LinearAlgebra.svdvals(Xmat))

@test rank(Xreal) == r

# 0.2 of entries missing
const missing_entries = unique!([(rand(1:nobs), rand(1:nfeat)) for _ in 1:10000])
const present_entries = [(i, j) for i in 1:nobs, j in 1:nfeat if (i, j) ∉ missing_entries]

f(X) = 0.5 * sum((X[i, j] - Xreal[i, j])^2 for (i, j) in present_entries)

function grad!(storage, X)
    storage .= 0
    for (i, j) in present_entries
        storage[i, j] = X[i, j] - Xreal[i, j]
    end
    return nothing
end

const lmo = FrankWolfe.NuclearNormLMO(275_000.0)
const x0 = FrankWolfe.compute_extreme_point(lmo, zero(Xreal))

FrankWolfe.benchmark_oracles(f, grad!, () -> randn(size(Xreal)), lmo; k=100)

# gradient descent
gradient = similar(x0)
xgd = Matrix(x0)
for _ in 1:5000
    @info f(xgd)
    grad!(gradient, xgd)
    xgd .-= 0.01 * gradient
    if norm(gradient) ≤ sqrt(eps())
        break
    end
end

grad!(gradient, x0)
v0 = FrankWolfe.compute_extreme_point(lmo, gradient)
@test dot(v0 - x0, gradient) < 0

const k = 500

xfin, vmin, _, _, traj_data = FrankWolfe.frank_wolfe(
    f,
    grad!,
    lmo,
    x0;
    epsilon=1e7,
    max_iteration=k,
    print_iter=k / 10,
    trajectory=true,
    verbose=true,
    linesearch_tol=1e-7,
    line_search=FrankWolfe.Adaptive(),
    emphasis=FrankWolfe.memory,
    gradient=spzeros(size(x0)...),
)

xfinAFW, vmin, _, _, traj_data = FrankWolfe.away_frank_wolfe(
    f,
    grad!,
    lmo,
    x0;
    epsilon=1e7,
    max_iteration=k,
    print_iter=k / 10,
    trajectory=true,
    verbose=true,
    linesearch_tol=1e-7,
    lazy=true,
    line_search=FrankWolfe.Adaptive(),
    emphasis=FrankWolfe.memory,#,
)

xfinBCG, vmin, _, _, traj_data = FrankWolfe.blended_conditional_gradient(
    f,
    grad!,
    lmo,
    x0;
    epsilon=1e7,
    max_iteration=k,
    print_iter=k / 10,
    trajectory=true,
    verbose=true,
    linesearch_tol=1e-7,
    line_search=FrankWolfe.Adaptive(),
    emphasis=FrankWolfe.memory,#,
)

plot(svdvals(xfin), label="FW", width=3, yaxis=:log)
plot!(svdvals(xfinAFW), label="AFW", width=3, yaxis=:log)
plot!(svdvals(xfinBCG), label="BCG", width=3, yaxis=:log)
plot!(svdvals(xgd), label="Gradient descent", width=3, yaxis=:log)
plot!(svdvals(Xreal), label="Real matrix", linestyle=:dash, width=3, color=:black)
title!("Singular values")


using Test, DelayEmbeddings
using Statistics

println("\nTesting Dataset...")

@testset "Dataset" begin
  data = Dataset(rand(1001,3))
  xs = columns(data)

  @testset "Concatenation/Append" begin
    x, y, z = Dataset(rand(10, 2)), Dataset(rand(10, 2)), rand(10)
    @test Dataset(x) == x
    @test Dataset(x, y) isa Dataset
    @test size(Dataset(x, y)) == (10, 4)
    @test hcat(x, y) isa Dataset
    @test Dataset(x, z) isa Dataset
    @test Dataset(z, x) isa Dataset

    append!(x, y)
    @test length(x) == 20
    w = hcat(x, rand(20))
    @test size(w) == (20, 3)
  end

  @testset "Methods & Indexing" begin
    a = data[:, 1]
    b = data[:, 2]
    c = data[:, 3]

    @test data[1,1] isa Float64
    @test a isa Vector{Float64}
    @test Dataset(a, b, c) == data
    @test size(Dataset(a, b)) == (1001, 2)

    @test data[5] isa SVector{3, Float64}
    @test data[11:20] isa Dataset
    @test data[:, 2:3][:, 1] == data[:, 2]

    @test size(data[1:10,1:2]) == (10,2)
    @test data[1:10,1:2] == Dataset(a[1:10], b[1:10])
    @test data[1:10, SVector(1, 2)] == data[1:10, 1:2]
    e = data[5, SVector(1,2)]
    @test e isa SVector{2, Float64}

    sub = @view data[11:20]
    @test sub isa DelayEmbeddings.SubDataset
    @test sub[2] == data[12]
    @test dimension(sub) == dimension(data)
    d = sub[:, 1]
    @test d isa Vector{Float64}
    e = sub[5, 1:2]
    @test e isa Vector{Float64}
    @test length(e) == 2
    e = sub[5, SVector(1,2)]
    @test e isa SVector{2, Float64}
    f = sub[5:8, 1:2]
    @test f isa Dataset

    # setindex
    data[1] = SVector(0.1,0.1,0.1)
    @test data[1] == SVector(0.1,0.1,0.1)
    @test_throws ErrorException (data[:,1] .= 0)
  end

  @testset "minmax" begin
    mi = minima(data)
    ma = maxima(data)
    mimi, mama = minmaxima(data)
    @test mimi == mi
    @test mama == ma
    for i in 1:3
      @test mi[i] < ma[i]
      a,b = extrema(xs[i])
      @test a == mi[i]
      @test b == ma[i]
    end
  end

  @testset "Conversions" begin
    m = Matrix(data)
    @test Dataset(m) == data

    m = rand(1000, 4)
    @test Matrix(Dataset(m)) == m
  end

  @testset "standardize" begin
    r = standardize(data)
    rs = columns(r)
    for x in rs
      m, s = mean(x), std(x)
      @test abs(m) < 1e-8
      @test abs(s - 1) < 1e-8
    end
  end
end


differ = -(@view(states[1:end-1]), @view(states[2:end]))
emp_var = empirical_var(differ,n, k)#sum(sum([d .^2 for d in differ]))/(2*k)
epsilon = 5*q/√(n*k)
@test abs(emp_var-q) < epsilon


"""
    InfExtendedTime{T} <: TimeType

The type `T` extended with positive and negative infinity.
"""
struct InfExtendedTime{T<:TimeType} <: TimeType
    flag :: InfFlag
    finitevalue :: T

    InfExtendedTime{T}(x::T) where {T<:TimeType} = new{T}(FINITE, x)
    InfExtendedTime{T}(x::Infinite) where {T<:TimeType} = new{T}(x==PosInf ? POSINF : NEGINF)
end

function Base.getproperty(x::InfExtendedTime, s::Symbol)
    if s === :val
        return x.flag != FINITE ? (x.flag == NEGINF ? -∞ : ∞) : x.finitevalue
    elseif s === :instant
        return x.flag != FINITE ? (x.flag == NEGINF ? -∞ : ∞) : x.finitevalue.instant
    else
        return getfield(x, s)
    end
end

InfExtendedTime{T}(x::TimeType) where {T<:TimeType} = InfExtendedTime{T}(convert(T, x))
InfExtendedTime{T}(x::InfExtendedTime) where {T<:TimeType} = InfExtendedTime{T}(convert(T, x.finitevalue))

"""
    InfExtendedTime(T)

Return `InfExtendedTime{T}` for any TimeType `T`.
"""
InfExtendedTime(::Type{T}) where {T<:TimeType} = InfExtendedTime{T}

"""
    InfExtendedTime(x)

Converts `x` to `InfExtendedTime(typeof(x))`.
"""
InfExtendedTime(x::T) where {T<:TimeType} = InfExtendedTime{T}(x)

Utils.isposinf(x::InfExtendedTime) = x.flag == POSINF
Utils.isneginf(x::InfExtendedTime) = x.flag == NEGINF


#-----------------------------------------------------------------------------------------------------
# Measurement
#-----------------------------------------------------------------------------------------------------
export measure
function measure(
    o::Operation,
    v::AbstractVecOrMat
)
    dot(v, o * v)
end
#-----------------------------------------------------------------------------------------------------
function measure(
    o1::Operation,
    o2::Operation,
    v::AbstractVecOrMat
)
    dot(o1 * v, o2 * v)
end
#-----------------------------------------------------------------------------------------------------
export covmat
function covmat(
    ol::Vector{<:Operation},
    v::AbstractVector
)
    n = length(ol)
    am = Vector{Float64}(undef, n)
    cm = Matrix{Float64}(undef, n, n)
    for i=1:n
        am[i] = real(measure(ol[i], v))
    end
    for i=1:n
        for j=i:n
            cm[i,j] = real(measure(ol[i], ol[j], v)) - am[i] * am[j]
        end
    end
    Hermitian(cm)
end
#-----------------------------------------------------------------------------------------------------
function covmat(
    ol::Vector{<:Operation},
    v::AbstractMatrix
)
    n = length(ol)
    am = Vector{Float64}(undef, n)
    cm = Matrix{Float64}(undef, n, n)
    num = size(v, 2)
    for i=1:n
        am[i] = real(measure(ol[i], v))/num
    end
    for i=1:n
        for j=i:n
            cm[i,j] = real(measure(ol[i], ol[j], v))/num - am[i] * am[j]
        end
    end
    Hermitian(cm)
end
#-----------------------------------------------------------------------------------------------------
export covmatc
function covmatc(
    ol::Vector{<:Operation},
    v::AbstractVector
)
    n = length(ol)
    am = Vector{ComplexF64}(undef, n)
    cm = Matrix{ComplexF64}(undef, n, n)
    for i=1:n
        am[i] = measure(ol[i], v)
    end
    for i=1:n
        for j=i:n
            cm[i,j] = measure(ol[i], ol[j], v) - conj(am[i]) * am[j]
        end
    end
    Hermitian(cm)
end
#-----------------------------------------------------------------------------------------------------
function covmatc(
    ol::Vector{<:Operation},
    v::AbstractMatrix
)
    n = length(ol)
    am = Vector{ComplexF64}(undef, n)
    cm = Matrix{ComplexF64}(undef, n, n)
    num = size(v, 2)
    for i=1:n
        am[i] = measure(ol[i], v)/num
    end
    for i=1:n
        for j=i:n
            cm[i,j] = measure(ol[i], ol[j], v)/num - conj(am[i]) * am[j]
        end
    end
    Hermitian(cm)
end
#-----------------------------------------------------------------------------------------------------
# Entropy
#-----------------------------------------------------------------------------------------------------
export entropy
function entropy(
    v::AbstractVector,
    l::Integer,
    i::Integer;
    base=0,
    cutoff=1e-7
)
    base = base==0 ? round(Int64, length(v)^(1/l)) : base
    m = reshape(v, base^i, :)
    s = svdvals(m) .^ 2
    - sum(si * log(si) for si in s if si > cutoff)
end


@testset "histos" begin
    x = [0.2, 6.4, 3.0, 1.6]
    bins = [0.0, 1.0, 2.5, 4.0, 10.0]
    inds = ATools.digitize(x,bins)
    @test all(inds .==[1,4,3,2])

    bins = [0,1,7] # a small and a large bin
    obs = [0.5, 1.5, 1.5, 2.5]; # one observation in the small bin and three in the large
    h  = ATools.hist1d(obs, bins, false) # not normalized
    hn = ATools.hist1d(obs, bins, true) # normalized

    he = ATools.edges(h)
    xn = [h.weights[i] / (he[i+1]- he[i]) for i in 1:length(h.weights)]
    @test all(xn .== hn.weights)

    h= ATools.hist1d(qs, 100, 50.0, 150.0)

    @test all(ATools.edges(h) .== h.edges[1])
    he = ATools.edges(h)
    hc = ATools.centers(h)
    @test hc[1] == (he[1] + he[2]) / 2.0
    @test hc[end] == (he[end-1] + he[end]) / 2.0
    ic = argmax(h.weights)
    @info "position of the maximum" ic
    @info "value of the maximum" hc[ic]
    @test hc[ic] > 95.0 && hc[ic] < 105.0

    h2 = ATools.hist1d(qs, 25, 50.0, 150.0, true)
    h2n = ATools.hist1d(qs, 25, 50.0, 150.0, false)
    ic = argmax(h2.weights)
    r = h2n.weights[ic] / h2.weights[ic]  # area of bins (bin[i+1] - bin[1]) = 4
    #println(r)
    @test r == 4.0

    h2, _ = hist2d(xs, qs, 100, "xs", "qs")
    @test sum(h2.weights) == length(xs)
    @test length(h2.edges[1]) == size(h2.weights)[1] + 1
    max_q = argmax(sum(h2.weights, dims=2))[1]
    @test isapprox(centers(h2)[max_q], mean(qs), atol=2*step(h2.edges[1]))

    profile_nbins = 25
    p1df, _ = ATools.p1df(xs, qs, profile_nbins)
    @test (mean(p1df.y_mean) > 95.0 && mean(p1df.y_mean) < 105.0)
    @test (mean(p1df.x_mean) > 0.3 && mean(p1df.x_mean) < 0.7)
    @test (mean(p1df.y_std) > 8.0 && mean(p1df.y_std) < 12.0)

    p1df_filt, _ = ATools.p1df(xs, qs, profile_nbins, min_proportion=0.1)
    @test nrow(p1df_filt) == profile_nbins

end


#functions
aNb2A = (a,b) -> return(rad2deg(atan(b/a))) #when you know a & b what is A
aNb2B = (a,b) -> return(90-aNb2A(a,b)) #when you know a & b what is B
aNb2c = (a,b) -> return(√(a^2+b^2)) #when you know a & b what is c

aNc2b = (a,c) -> return(√(c^2-a^2)) #when you know a & c what is b
aNc2A = (a,c) -> return(rad2deg(acos(a/c))) #when you know a & c what is A
aNc2B = (a,c) -> return(90-aNc2A(a,c)) #when you know a & c what is B

aNA2b = (a,A) -> return(a*tan(deg2rad(A))) #when you know a & A what is b
aNA2c = (a,A) -> return(a/cos(deg2rad(A))) #when you know a & A what is c
aNA2B = (a,A) -> return(90-A) #when you know a & A what is B

bNA2a = (a,A) -> return(b/tan(deg2rad(A))) #when you know b & A what is a
bNA2c = (a,A) -> return(b/sin(deg2rad(A))) #when you know b & A what is c
bNA2B = (a,A) -> return(90-A) #when you know b & A what is B

bNc2A = (b,c) -> return(rad2deg(asin(b/c))) #when you know b & c what is A
bNc2B = (b,c) -> return(90-bNc2A(b,c)) #when you know b & c what is B
bNc2a = (b,c) -> return(√(c^2-b^2))  #when you know b & c what is a

cNA2a = (c,A) -> return(c*cos(deg2rad(A))) #when you know c & A what is a
cNA2b= (c,A) -> return(c*sin(deg2rad(A))) #when you know c & A what is b
cNA2B = (c,A) -> return(90-A) #when you know c & A what is B
a= 270; b= 175
c =350; A= 30
<<<<<<< Updated upstream
function input(prompt::AbstractString="")
    print(prompt)
    return chomp(readline())
end

response = input("Do you know a & b ? (Y/N)")
println(response)
if(response=="y")
=======

>>>>>>> Stashed changes
    println(aNb2A(a,b))
    println(aNb2B(a,b))
    println(aNb2c(a,b))
    println()
    println(aNc2b(a,c))
    println(aNc2A(a,c))
    println(aNc2B(a,c))
    println()
    println(aNA2b(a,A))
    println(aNA2c(a,A))
    println(aNA2B(a,A))
    println()
    println(bNA2a(a,A))
    println(bNA2c(a,A))
    println(bNA2B(a,A))
    println()
    println(bNc2A(b,c))
    println(bNc2B(b,c))
    println(bNc2a(b,c))
    println()
    println(cNA2a(c,A))
    println(cNA2b(c,A))
    println(cNA2B(c,A))
    println("finished")
<<<<<<< Updated upstream
end
=======
>>>>>>> Stashed changes


using SymPyCall
import SymPyCall: Latexify
using Test

@testset "Latexify" begin
  @syms α
  @test_nowarn Latexify.latexify(- 3/(8*α) + 1/(8*α^2))
  @test  occursin("\\cdot", Latexify.latexify(- 3/(8*α) + 1/(8*α^2), cdot=true))
  @test !occursin("\\cdot", Latexify.latexify(- 3/(8*α) + 1/(8*α^2), cdot=false))
end


using Base.Iterators
using Distributions
using LinearAlgebra
using Random
using RCall

# Define the DGP
function simulate_sample(N0, N1)
    # Set the conditional density
    function dy(xc)
        atoms = [(-xc, 0.2), (xc, 0.2)]
        weights = [0.5, 0.5]
        # weights = [exp(-2xc), 1 - exp(-2xc)]
        MixtureModel(Normal, atoms, weights)
    end

    # Simulate a sample
    x0 = rand(N0)
    y0 = [rand(dy(xc)) for xc in x0]
    X0 = rcopy(R"cbind(1, splines::bs($x0, df = 6, Boundary.knots = c(0, 1)))")

    # Generate the grid points
    y_grid = LinRange(-3.0, 3.0, N1)
    x_grid = [0.25, 0.5, 0.75]
    y1, x1 = Iterators.product(y_grid, x_grid) |> x -> zip(x...) .|> collect
    X1 = rcopy(R"cbind(1, predict(splines::bs($x0, df = 6, Boundary.knots = c(0, 1)), $x1))")
    return dy, y0, x0, X0, y1, x1, X1
end

# # Not run
# Random.seed!(1);
# N0, N1 = 500, 50;
# dy, y0, x0, X0, y1, x1, X1 = simulate_sample(N0, N1);


let part1 = 0, part2 = 1
    # This reads the input in a transposed format — rows here are columns in
    # the file.
    chart = hcat(map(collect, readlines("input.txt"))...)

    nrow, ncol = size(chart)

    for (run, rise) in [(1, 1), (3, 1), (5, 1), (7, 1), (1, 2)]
        trees, row = 0, 1

        for col in 1:rise:ncol
            if chart[row, col] == '#'
                # only count the tree for part1 if the slope is (3, 1)
                part1 += (run == 3 && rise == 1)

                # For part2
                trees += 1
            end

            # # We could use the following if statement which doesn't require
            # # as much mental overhead.
            # row += run
            #
            # # This is the only place where we have to think about the indexing.
            # # If we were using 0-indexing, we'd need `row ≥ nrow`.
            # if row > nrow
            #     row -= nrow
            # end

            # By default, Julia arrays are 1-indexed, so we have to subtract 1
            # before we `%` and add one after. I like this better because its
            # short, fairly clear if you know how `%` works, and doesn't
            # require the overhead of an if-statement.
            row = (row + run - 1) % nrow + 1
        end

        part2 *= trees
    end

    @info "Solution" part1 part2
end


using LightGraphs

struct BudgetedSpanningTreeInstance{T, U}
  graph::AbstractGraph{T}
  rewards::Dict{Edge{T}, Float64}
  weights::Dict{Edge{T}, U}
  budget::U
end

abstract type BudgetedSpanningTreeSolution{T, U}
  # instance::BudgetedSpanningTreeInstance{T, U}
  # tree::Vector{Edge{T}}
end

struct BudgetedSpanningTreeLagrangianSolution{T, U} <: BudgetedSpanningTreeSolution{T, U}
  # Used to store important temporary results from solving the Lagrangian dual.
  instance::BudgetedSpanningTreeInstance{T, U}
  tree::Vector{Edge{T}}
  λ::Float64
  value::Float64
  λmax::Float64 # No dual value higher than this is useful (i.e. they all yield the same solution).
end

struct SimpleBudgetedSpanningTreeSolution{T, U} <: BudgetedSpanningTreeSolution{T, U}
  instance::BudgetedSpanningTreeInstance{T, U}
  tree::Vector{Edge{T}}
end

function _budgeted_spanning_tree_compute_value(i::Union{SpanningTreeInstance{T}, BudgetedSpanningTreeInstance{T}}, solution::Vector{Edge{T}}) where T
  return sum(i.rewards[(e in keys(i.rewards)) ? e : reverse(e)] for e in solution)
end

function _budgeted_spanning_tree_compute_weight(i::BudgetedSpanningTreeInstance{T}, solution::Vector{Edge{T}}) where T
  return sum(i.weights[(e in keys(i.weights)) ? e : reverse(e)] for e in solution)
end

function st_prim_budgeted_lagrangian(i::BudgetedSpanningTreeInstance{T}, λ::Float64) where T
  # Solve the subproblem
  #     l(λ) = \max_{x spanning tree} (rewards + λ weights) x - λ budget.
  sti_rewards = Dict{Edge{T}, Float64}(e => i.rewards[e] + λ * i.weights[e] for e in keys(i.rewards))
  sti = SpanningTreeInstance(i.graph, sti_rewards)
  sti_sol = st_prim(sti)
  sti_value = sum(sti_rewards[(e in keys(i.rewards)) ? e : reverse(e)] for e in sti_sol.tree)
  return sti_value - λ * i.budget, sti_sol.tree
end

function st_prim_budgeted_lagrangian_search(i::BudgetedSpanningTreeInstance{T}, ε::Float64) where T
  # Approximately solve the problem \min_{l ≥ 0} l(λ), where
  #     l(λ) = \max_{x spanning tree} (rewards + λ weights) x - λ budget.
  # This problem is the Lagrangian dual of the budgeted maximum spanning-tree problem:
  #     \max_{x spanning tree} rewards x  s.t.  weights x >= budget.
  # This algorithm provides no guarantee on the optimality of the solution.

  # Initial set of values for λ. The optimum is guaranteed to be contained in this interval.
  weights_norm_inf = maximum(values(i.weights)) # Maximum weight.
  m = nv(i.graph) - 1 # Maximum number of items in a solution. Easy to compute for a spanning tree!
  λmax = Float64(weights_norm_inf * m + 1)

  λlow = 0.0
  λhigh = λmax

  # Perform a golden-ratio search.
  while (λhigh - λlow) > ε
    λmidlow = λhigh - (λhigh - λlow) / MathConstants.φ
    λmidhigh = λlow + (λhigh - λlow) / MathConstants.φ
    vmidlow, _ = st_prim_budgeted_lagrangian(i, λmidlow)
    vmidhigh, _ = st_prim_budgeted_lagrangian(i, λmidhigh)

    if vmidlow < vmidhigh
      λhigh = λmidhigh
      vhigh = vmidhigh
    else
      λlow = λmidlow
      vlow = vmidlow
    end
  end

  vlow, stlow = st_prim_budgeted_lagrangian(i, λlow)
  vhigh, sthigh = st_prim_budgeted_lagrangian(i, λhigh)
  if vlow > vhigh
    return BudgetedSpanningTreeLagrangianSolution(i, stlow, λlow, vlow, λmax)
  else
    return BudgetedSpanningTreeLagrangianSolution(i, sthigh, λhigh, vhigh, λmax)
  end
end

function _solution_symmetric_difference(a::Vector{Edge{T}}, b::Vector{Edge{T}}) where T
  only_in_a = Edge{T}[]
  only_in_b = Edge{T}[]

  for e in a
    if e in b
      continue
    end
    push!(only_in_a, e)
  end

  for e in b
    if e in a
      continue
    end
    push!(only_in_b, e)
  end

  return only_in_a, only_in_b
end

function _solution_symmetric_difference_size(a::Vector{Edge{T}}, b::Vector{Edge{T}}) where T
  only_in_a, only_in_b = _solution_symmetric_difference(a, b)
  return length(only_in_a) + length(only_in_b)
end

function st_prim_budgeted_lagrangian_refinement(i::BudgetedSpanningTreeInstance{T};
                                                ε::Float64=1.0e-3, ζ⁻::Float64=0.2, ζ⁺::Float64=5.0,
                                                stalling⁻::Float64=1.0e-5) where T
  # Check assumptions.
  if ζ⁻ >= 1.0
    error("ζ⁻ must be strictly less than 1.0: the dual multiplier λ is multiplied by ζ⁻ to reach an infeasible solution by less penalising the budget constraint.")
  end
  if ζ⁺ <= 1.0
    error("ζ⁺ must be strictly greater than 1.0: the dual multiplier λ is multiplied by ζ⁺ to reach a feasible solution by penalising more the budget constraint.")
  end

  # Ensure the problem is feasible by only considering the budget constraint.
  feasible_rewards = Dict{Edge{T}, Float64}(e => i.weights[e] for e in keys(i.rewards))
  feasible_instance = SpanningTreeInstance(i.graph, feasible_rewards)
  feasible_solution = st_prim(feasible_instance)
  if _budgeted_spanning_tree_compute_value(feasible_instance, feasible_solution.tree) < i.budget
    # By maximising the left-hand side of the budget constraint, impossible to reach the target budget. No solution!
    return SimpleBudgetedSpanningTreeSolution(i, similar(feasible_solution.tree, 0))
  end

  # Solve the Lagrangian relaxation to optimality.
  lagrangian = st_prim_budgeted_lagrangian_search(i, ε)
  λ0, v0, st0 = lagrangian.λ, lagrangian.value, lagrangian.tree
  λmax = lagrangian.λmax
  b0 = _budgeted_spanning_tree_compute_weight(i, st0) # Budget consumption of this first solution.

  # Find two solutions: one above the budget x⁺ (i.e. respecting the constraint), the other not x⁻.
  x⁺, x⁻ = nothing, nothing

  λi = λ0
  if b0 >= i.budget
    x⁺ = st0
    @assert _budgeted_spanning_tree_compute_weight(i, x⁺) >= i.budget

    stalling = false
    while true
      # Penalise less the constraint: it should no more be satisfied.
      λi *= ζ⁻
      _, sti = st_prim_budgeted_lagrangian(i, λi)
      if _budgeted_spanning_tree_compute_weight(i, sti) < i.budget
        x⁻ = sti
        break
      end

      # Is the process stalling?
      if λi <= stalling⁻
        stalling = true
        break
      end
    end

    # Specific handling of stallings.
    if stalling # First test: don't penalise the constraint at all.
      _, sti = st_prim_budgeted_lagrangian(i, 0.0)
      new_budget = _budgeted_spanning_tree_compute_weight(i, sti)
      if new_budget < i.budget
        x⁻ = sti
        stalling = false
      end
    end

    if stalling # Second test: minimise the left-hand side of the budget constraint, in hope of finding a feasible solution.
      # This process is highly similar to the computation of feasible_solution, but with a reverse objective function.
      infeasible_rewards = Dict{Edge{T}, Float64}(e => - i.weights[e] for e in keys(i.rewards))
      infeasible_solution = st_prim(SpanningTreeInstance(i.graph, infeasible_rewards)).tree

      if _budgeted_spanning_tree_compute_weight(i, infeasible_solution) < i.budget
        x⁻ = infeasible_solution
        stalling = false
      end
    end

    if stalling # Third: decide there is no solution strictly below the budget. No refinement is possible.
      # As x⁺ is feasible, return it.
      return SimpleBudgetedSpanningTreeSolution(i, x⁺)
    end
  else
    x⁻ = st0
    @assert _budgeted_spanning_tree_compute_weight(i, x⁻) < i.budget

    while true
      # Penalise more the constraint: it should become satisfied at some point.
      λi *= ζ⁺
      _, sti = st_prim_budgeted_lagrangian(i, λi)
      if _budgeted_spanning_tree_compute_weight(i, sti) >= i.budget
        x⁺ = sti
        break
      end

      # Is the process stalling? If so, reuse feasible_solution, which is guaranteed to be feasible.
      if λi >= λmax
        x⁺ = feasible_solution.tree
        break
      end
    end
  end

  # Normalise the solutions: the input graph is undirected, the direction of the edges is not important.
  # In case one solution has the edge v -> w and the other one w -> v, make them equal. s
  sort_edge(e::Edge{T}) where T = (src(e) < dst(e)) ? e : reverse(e)
  x⁺ = [sort_edge(e) for e in x⁺]
  x⁻ = [sort_edge(e) for e in x⁻]

  # Iterative refinement. Stop as soon as there is a difference of at most one edge between the two solutions.
  while _solution_symmetric_difference_size(x⁺, x⁻) > 2
    # Enforce the loop invariant.
    @assert x⁺ !== nothing
    @assert x⁻ !== nothing
    @assert _budgeted_spanning_tree_compute_weight(i, x⁺) >= i.budget # Feasible.
    @assert _budgeted_spanning_tree_compute_weight(i, x⁻) < i.budget # Infeasible.

    # Switch elements from one solution to another.
    only_in_x⁺, only_in_x⁻ = _solution_symmetric_difference(x⁺, x⁻)
    e1 = first(only_in_x⁺)
    e2 = first(only_in_x⁻)

    # Create the new solution (don't erase x⁺ nor x⁻: only one of the two will be forgotten, the other will be kept).
    new_x = copy(x⁺)
    filter!(e -> e != e1, new_x)
    push!(new_x, e2)

    # Replace one of the two solutions, depending on whether this solution is feasible (x⁺) or not (x⁻).
    if _budgeted_spanning_tree_compute_weight(i, new_x) >= i.budget
      x⁺ = new_x
    else
      x⁻ = new_x
    end
  end

  # Done!
  return SimpleBudgetedSpanningTreeSolution(i, x⁺)
end


@testset "Tests: Minimize depth for controlled-NOT gate" begin

    params = Dict{String, Any}(
    "num_qubits" => 2,
    "maximum_depth" => 5,
    "elementary_gates" => ["H_1", "H_2", "CNot_1_2", "Identity"],
    "initial_gate" => "Identity",
    "target_gate" => QCO.CNotRevGate(),
    "set_cnot_lower_bound" => 1,
    "set_cnot_upper_bound" => 1,
    "objective" => "minimize_depth",
    "decomposition_type" => "exact"
    )

    model_options = Dict{Symbol, Any}(:model_type => "compact_formulation_1", # Testing incorrect model_type
                                      :all_valid_constraints => 2)            # Testing incorrect all_valid_constraints

    result_qc = QCO.run_QCModel(params, CBC; options = model_options)

    @test result_qc["termination_status"] == MOI.OPTIMAL
    @test result_qc["primal_status"] == MOI.FEASIBLE_POINT
    @test isapprox(result_qc["objective"], 5, atol=tol_0)
    @test isapprox(result_qc["solution"]["z_onoff_var"][1,1], 1, atol=tol_0) || isapprox(result_qc["solution"]["z_onoff_var"][2,1], 1, atol=tol_0)
    @test isapprox(result_qc["solution"]["z_onoff_var"][1,2], 1, atol=tol_0) || isapprox(result_qc["solution"]["z_onoff_var"][2,2], 1, atol=tol_0)
    @test isapprox(result_qc["solution"]["z_onoff_var"][3,3], 1, atol=tol_0)
    @test isapprox(result_qc["solution"]["z_onoff_var"][1,4], 1, atol=tol_0) || isapprox(result_qc["solution"]["z_onoff_var"][2,4], 1, atol=tol_0)
    @test isapprox(result_qc["solution"]["z_onoff_var"][1,5], 1, atol=tol_0) || isapprox(result_qc["solution"]["z_onoff_var"][2,5], 1, atol=tol_0)

end

@testset "Tests: Minimum CNOT swap gate decomposition" begin

    params = Dict{String, Any}(
        "num_qubits" => 2,
        "maximum_depth" => 4,
        "elementary_gates" => ["CNot_1_2", "CNot_2_1", "Identity"],
        "target_gate" => QCO.SwapGate(),
        "set_cnot_lower_bound" => 2,
        "set_cnot_upper_bound" => 3,
        "objective" => "minimize_cnot",
        "decomposition_type" => "exact"
        )

    model_options = Dict{Symbol, Any}(:model_type => "balas_formulation",
                                      :all_valid_constraints => 1,
                                      :unit_magnitude_constraints => true)

    result_qc = QCO.run_QCModel(params, CBC; options = model_options)

    @test result_qc["termination_status"] == MOI.OPTIMAL
    @test result_qc["primal_status"] == MOI.FEASIBLE_POINT
    @test isapprox(result_qc["objective"], 3, atol = tol_0)
    @test isapprox(sum(result_qc["solution"]["z_onoff_var"][1:2,:]), 3, atol=tol_0)

end

@testset "Tests: Minimum depth U3(0,0,π/4) gate" begin

    params = Dict{String, Any}(
    "num_qubits" => 2,
    "maximum_depth" => 2,
    "elementary_gates" => ["U3_1", "U3_2", "Identity"],
    "target_gate" => QCO.kron_single_qubit_gate(2, QCO.U3Gate(0,0,π/4), "q1"),
    "U3_θ_discretization" => [0, π/2],
    "U3_ϕ_discretization" => [0],
    "U3_λ_discretization" => [0, π/4],
    "objective" => "minimize_depth",
    "decomposition_type" => "exact"
    )

    result_qc = QCO.run_QCModel(params, CBC)

    data = QCO.get_data(params)

    @test result_qc["termination_status"] == MOI.OPTIMAL
    @test result_qc["primal_status"] == MOI.FEASIBLE_POINT
    @test isapprox(result_qc["objective"], 1, atol = tol_0)
    @test "Identity" in data["gates_dict"]["3"]["type"]
    @test isapprox(sum(result_qc["solution"]["z_onoff_var"][3,:]), 1, atol = tol_0)
    @test isapprox(sum(result_qc["solution"]["z_onoff_var"][4,:]), 1, atol = tol_0)
    @test data["gates_dict"]["4"]["qubit_loc"] == "qubit_1"
end

@testset "Tests: Minimum depth CU3(0,0,π/4) gate" begin

    params = Dict{String, Any}(
    "num_qubits" => 2,
    "maximum_depth" => 2,
    "elementary_gates" => ["CU3_1_2", "CU3_2_1", "Identity"],
    "target_gate" => QCO.CU3Gate(0, 0, π/4),
    "CU3_θ_discretization" => [0, π/2],
    "CU3_ϕ_discretization" => [0],
    "CU3_λ_discretization" => [0, π/4],
    "objective" => "minimize_depth",
    "decomposition_type" => "exact"
    )

    result_qc = QCO.run_QCModel(params, CBC)

    data = QCO.get_data(params)

    @test result_qc["termination_status"] == MOI.OPTIMAL
    @test result_qc["primal_status"] == MOI.FEASIBLE_POINT
    @test isapprox(result_qc["objective"], 1, atol = tol_0)
    @test "Identity" in data["gates_dict"]["3"]["type"]
    @test isapprox(sum(result_qc["solution"]["z_onoff_var"][3,:]), 1, atol = tol_0)
    @test isapprox(sum(result_qc["solution"]["z_onoff_var"][4,:]), 1, atol = tol_0)
    @test data["gates_dict"]["4"]["qubit_loc"] == "qubit_1_2"
end

@testset "Tests: Minimum depth Rev CU3(0,0,π/4) gate" begin

    params = Dict{String, Any}(
    "num_qubits" => 3,
    "maximum_depth" => 2,
    "elementary_gates" => ["CU3_3_1", "CU3_1_3", "Identity"],
    "target_gate" => QCO.get_full_sized_gate("CU3_3_1", 3, angle = [0, 0, pi/4]),
    "CU3_θ_discretization" => [0, π/2],
    "CU3_ϕ_discretization" => [0],
    "CU3_λ_discretization" => [0, π/4],
    "objective" => "minimize_depth",
    "decomposition_type" => "exact"
    )

    result_qc = QCO.run_QCModel(params, CBC)

    data = QCO.get_data(params)

    @test result_qc["termination_status"] == MOI.OPTIMAL
    @test result_qc["primal_status"] == MOI.FEASIBLE_POINT
    @test isapprox(result_qc["objective"], 1, atol = tol_0)
    @test "Identity" in data["gates_dict"]["3"]["type"]
    @test isapprox(sum(result_qc["solution"]["z_onoff_var"][3,:]), 1, atol = tol_0)
    @test isapprox(sum(result_qc["solution"]["z_onoff_var"][4,:]), 1, atol = tol_0)
    @test data["gates_dict"]["4"]["qubit_loc"] == "qubit_3_1"
end

@testset "Tests: Minimum depth RX, RY, RZ gate decomposition" begin

    params = Dict{String, Any}(
    "num_qubits" => 2,
    "maximum_depth" => 3,
    "elementary_gates" => ["RX_1", "RY_2", "RZ_1", "Identity"],
    "target_gate" => QCO.kron_single_qubit_gate(2, QCO.RXGate(π/4), "q1") * QCO.kron_single_qubit_gate(2, QCO.RYGate(π/4), "q2") * QCO.kron_single_qubit_gate(2, QCO.RZGate(π/4), "q1"),
    "RX_discretization" => [0, π/4],
    "RY_discretization" => [π/4],
    "RZ_discretization" => [π/2, π/4],
    "objective" => "minimize_depth",
    "decomposition_type" => "exact"
    )

    model_options = Dict{Symbol, Any}(:model_type => "balas_formulation",
                                      :commute_gate_constraints => true)

    result_qc = QCO.run_QCModel(params, CBC; options =  model_options)

    @test result_qc["termination_status"] == MOI.OPTIMAL
    @test result_qc["primal_status"] == MOI.FEASIBLE_POINT
    @test isapprox(result_qc["objective"], 3, atol = tol_0)

end

@testset "Tests: Minimum depth CRX, CRY, CRZ gate decomposition" begin

    params = Dict{String, Any}(
    "num_qubits" => 3,
    "maximum_depth" => 3,
    "elementary_gates" => ["CRX_1_2", "CRY_2_3", "CRZ_3_1", "Identity"],
    "target_gate" => QCO.get_full_sized_gate("CRX_1_2", 3, angle = pi/4) * QCO.get_full_sized_gate("CRY_2_3", 3, angle=pi/4) * QCO.get_full_sized_gate("CRZ_3_1", 3, angle=pi/4),
    "CRX_discretization" => [0, π/4],
    "CRY_discretization" => [π/4],
    "CRZ_discretization" => [π/2, π/4],
    "objective" => "minimize_depth",
    "decomposition_type" => "exact"
    )

    model_options = Dict{Symbol, Any}(:model_type => "balas_formulation",
                                      :commute_gate_constraints => true)

    result_qc = QCO.run_QCModel(params, CBC; options =  model_options)

    @test result_qc["termination_status"] == MOI.OPTIMAL
    @test result_qc["primal_status"] == MOI.FEASIBLE_POINT
    @test isapprox(result_qc["objective"], 3, atol = tol_0)

end

@testset "Tests: Minimum depth Rev CRX, CRY, CRZ gate decomposition" begin

    params = Dict{String, Any}(
    "num_qubits" => 3,
    "maximum_depth" => 3,
    "elementary_gates" => ["CRX_3_1", "CRY_3_1", "CRZ_1_3", "Identity"],
    "target_gate" => QCO.get_full_sized_gate("CRX_3_1", 3, angle=pi/4) * QCO.get_full_sized_gate("CRY_3_1", 3, angle=pi/4) * QCO.get_full_sized_gate("CRZ_1_3", 3, angle = pi/2),
    "CRX_discretization" => [0, π/4],
    "CRY_discretization" => [π/4],
    "CRZ_discretization" => [π/2, π/4],
    "objective" => "minimize_depth",
    "decomposition_type" => "exact"
    )

    model_options = Dict{Symbol, Any}(:model_type => "balas_formulation",
                                      :commute_gate_constraints => true)

    result_qc = QCO.run_QCModel(params, CBC; options =  model_options)

    @test result_qc["termination_status"] == MOI.OPTIMAL
    @test result_qc["primal_status"] == MOI.FEASIBLE_POINT
    @test isapprox(result_qc["objective"], 3, atol = tol_0)

end

@testset "Tests: 3-qubit RX gate decomposition" begin

    params = Dict{String, Any}(
    "num_qubits" => 3,
    "maximum_depth" => 2,
    "elementary_gates" => ["U3_1", "U3_2", "U3_3", "Identity"],
    "target_gate" => QCO.kron_single_qubit_gate(3, QCO.RXGate(π/4), "q3"),
    "U3_θ_discretization" => [0, π/4],
    "U3_ϕ_discretization" => [0, -π/2],
    "U3_λ_discretization" => [0, π/2],
    "objective" => "minimize_cnot",
    "decomposition_type" => "exact"
    )

    model_options = Dict{Symbol, Any}(:model_type => "balas_formulation")

    result_qc = QCO.run_QCModel(params, CBC; options =  model_options)

    data = QCO.get_data(params)

    @test result_qc["termination_status"] == MOI.OPTIMAL
    @test result_qc["primal_status"] == MOI.FEASIBLE_POINT
    @test isapprox(result_qc["objective"], 1, atol = tol_0)
    if isapprox(sum(result_qc["solution"]["z_onoff_var"][14,:]), 1, atol=tol_0)
        @test data["gates_dict"]["14"]["qubit_loc"] == "qubit_3"
        @test isapprox(rad2deg(data["gates_dict"]["14"]["angle"]["θ"]),  45, atol=tol_0)
        @test isapprox(rad2deg(data["gates_dict"]["14"]["angle"]["ϕ"]), -90, atol=tol_0)
        @test isapprox(rad2deg(data["gates_dict"]["14"]["angle"]["λ"]),  90, atol=tol_0)
    end

end

@testset "Tests: 3-qubit CRX gate decomposition" begin

    params = Dict{String, Any}(
        "num_qubits" => 3,
        "maximum_depth" => 2,
        "elementary_gates" => ["CU3_1_2", "CU3_2_3", "CU3_1_3", "Identity"],
        "target_gate" => QCO.get_full_sized_gate("CRX_1_3", 3, angle = pi/4),
        "CU3_θ_discretization" => [0, π/4],
        "CU3_ϕ_discretization" => [0, -π/2],
        "CU3_λ_discretization" => [0, π/2],
        "objective" => "minimize_cnot",
        "decomposition_type" => "exact"
    )

    model_options = Dict{Symbol, Any}(:model_type => "balas_formulation")

    result_qc = QCO.run_QCModel(params, CBC; options =  model_options)

    data = QCO.get_data(params)

    @test result_qc["termination_status"] == MOI.OPTIMAL
    @test result_qc["primal_status"] == MOI.FEASIBLE_POINT
    @test isapprox(result_qc["objective"], 1, atol = tol_0)
    if isapprox(sum(result_qc["solution"]["z_onoff_var"][18,:]), 1, atol=tol_0)
        @test data["gates_dict"]["18"]["qubit_loc"] == "qubit_1_3"
        @test isapprox(rad2deg(data["gates_dict"]["18"]["angle"]["θ"]),  45, atol=tol_0)
        @test isapprox(rad2deg(data["gates_dict"]["18"]["angle"]["ϕ"]), -90, atol=tol_0)
        @test isapprox(rad2deg(data["gates_dict"]["18"]["angle"]["λ"]),  90, atol=tol_0)
    end

end

@testset "Tests: feasibility problem" begin
    params = Dict{String, Any}(
        "num_qubits" => 2,
        "maximum_depth" => 3,
        "elementary_gates" => ["H_1", "H_2"],
        "target_gate" => QCO.kron_single_qubit_gate(2, QCO.HGate(), "q1"),
        "objective" => "minimize_depth",
        "decomposition_type" => "exact"
        )

    model_options = Dict{Symbol, Any}(:all_valid_constraints => -1)

    result_qc = QCO.run_QCModel(params, CBC; options =  model_options)

    @test result_qc["termination_status"] == MOI.OPTIMAL
    @test result_qc["primal_status"] == MOI.FEASIBLE_POINT
    data = QCO.get_data(params)

    for i in keys(data["gates_dict"])
        if data["gates_dict"][i]["type"] == "H_1"
            @test isapprox(sum(result_qc["solution"]["z_onoff_var"][parse(Int64, i),:]), 3 , atol = tol_0)
        end
    end

end

@testset "Tests: JuMP set_start_value for z_onoff_var variables" begin
    function input_circuit()
        # [(depth, gate)]
        return [(1, "CNot_2_1"),
                (2, "S_1"),
                (3, "H_2"),
                (4, "S_2")
                ]
    end

    params = Dict{String, Any}(
    "num_qubits" => 2,
    "maximum_depth" => 4,
    "elementary_gates" => ["S_1", "S_2", "H_1", "H_2", "CNot_1_2", "CNot_2_1", "Identity"],
    "target_gate" => QCO.MGate(),
    "input_circuit" => input_circuit(),
    "objective" => "minimize_depth",
    "decomposition_type" => "exact"
    )

    result_qc = QCO.run_QCModel(params, CBC)
    @test result_qc["termination_status"] == MOI.OPTIMAL
    @test result_qc["primal_status"] == MOI.FEASIBLE_POINT
    @test isapprox(sum(result_qc["solution"]["z_onoff_var"][6,:]), 1, atol=tol_0)
    @test isapprox(sum(result_qc["solution"]["z_onoff_var"][5,:]), 0, atol=tol_0)

end

@testset "Tests: Involutory gate constraints" begin

    params = Dict{String, Any}(
    "num_qubits" => 2,
    "maximum_depth" => 4,
    "elementary_gates" => ["S_1", "S_2", "H_1", "H_2", "CNot_1_2", "CNot_2_1", "Identity"],
    "target_gate" => QCO.MGate(),
    "objective" => "minimize_depth",
    "decomposition_type" => "exact"
    )

    gates_dict = QCO.get_data(params)["gates_dict"]
    involutory_gates = QCO.get_involutory_gates(gates_dict)
    @test length(involutory_gates) == 4 #excluding Identity gate

    for i in keys(gates_dict)
        if ("S_1" in gates_dict[i]["type"]) || ("S_2" in gates_dict[i]["type"])
            @test !(parse(Int, i) in involutory_gates)
        end
        if ("H_1" in gates_dict[i]["type"]) || ("H_2" in gates_dict[i]["type"]) || ("CNot_1_2" in gates_dict[i]["type"]) || ("CNot_2_1" in gates_dict[i]["type"])
            @test (parse(Int, i) in involutory_gates)
        end
    end

    result_qc = QCO.run_QCModel(params, CBC)
    @test result_qc["termination_status"] == MOI.OPTIMAL
    @test result_qc["primal_status"] == MOI.FEASIBLE_POINT
    @test isapprox(result_qc["objective"], 4, atol = tol_0)

end

@testset "Tests: TIME_LIMIT for building results dict and log" begin

    params = Dict{String, Any}(
    "num_qubits" => 2,
    "maximum_depth" => 5,
    "elementary_gates" => ["S_1", "S_2", "H_1", "H_2", "CNot_1_2", "CNot_2_1", "Identity"],
    "target_gate" => QCO.MGate(),
    "objective" => "minimize_depth",
    "decomposition_type" => "exact"
    )

    model_options = Dict{Symbol, Any}(:time_limit => 1)

    result_qc = QCO.run_QCModel(params, CBC; options = model_options)
    @test result_qc["termination_status"] == MOI.TIME_LIMIT
    @test result_qc["primal_status"] == MOI.NO_SOLUTION

end

@testset "Tests: constraint_redundant_gate_product_pairs" begin

    function target_gate()
        T1 = QCO.get_full_sized_gate("U3_2", 2, angle = [0,π/2,π])
        T2 = QCO.get_full_sized_gate("U3_1", 2, angle = [π/2,π/2,-π/2])
        return T1*T2
    end

    params = Dict{String, Any}(
               "num_qubits" => 2,
               "maximum_depth" => 2,
               "elementary_gates" => ["U3_1", "U3_2", "Identity"],
               "target_gate" => target_gate(),
               "U3_θ_discretization" => [0, π/2],
               "U3_ϕ_discretization" => [π/2],
               "U3_λ_discretization" => [-π/2, π],
               "objective" => "minimize_depth")

    data = QCO.get_data(params)
    redundant_pairs = QCO.get_redundant_gate_product_pairs(data["gates_dict"])
    @test length(redundant_pairs) == 2
    @test redundant_pairs[1] == (1,4)
    @test redundant_pairs[2] == (5,7)
    @test isapprox(data["gates_dict"]["1"]["matrix"] * data["gates_dict"]["4"]["matrix"], data["gates_dict"]["2"]["matrix"], atol = tol_0)
    @test isapprox(data["gates_dict"]["5"]["matrix"] * data["gates_dict"]["7"]["matrix"], data["gates_dict"]["6"]["matrix"], atol = tol_0)

    result_qc = QCO.run_QCModel(params, CBC)
    @test result_qc["termination_status"] == MOI.OPTIMAL
    @test result_qc["primal_status"] == MOI.FEASIBLE_POINT
    @test isapprox(result_qc["objective"], 2.0, atol = tol_0)

end

@testset "Tests: constraint_idempotent_gates" begin
    params = Dict{String, Any}(
    "num_qubits" => 2,
    "maximum_depth" => 3,
    "elementary_gates" => ["U3_1", "U3_2", "CNot_1_2", "Identity"],
    "target_gate" => QCO.CZGate(),
    "U3_θ_discretization" => [-π/2, 0, π/2],
    "U3_ϕ_discretization" => [0, π/2],
    "U3_λ_discretization" => [0, π/2])

    data = QCO.get_data(params)
    idempotent_pairs = QCO.get_idempotent_gates(data["gates_dict"])
    @test length(idempotent_pairs) == 2

    model_options = Dict{Symbol, Any}(:idempotent_gate_constraints => true)

    result_qc = QCO.run_QCModel(params, CBC; options = model_options)

    @test result_qc["termination_status"] == MOI.OPTIMAL
    @test result_qc["primal_status"] == MOI.FEASIBLE_POINT
    @test isapprox(result_qc["objective"], 3.0, atol = tol_0)

    for i in keys(data["gates_dict"])
        if "Identity" in data["gates_dict"][i]["type"]
            @test isapprox(sum(result_qc["solution"]["z_onoff_var"][parse(Int64, i),:]), 0, atol = tol_0)
        end
    end

end

@testset "Tests: constraint_convex_hull_complex_gates" begin

    function target_gate()
        num_qubits = 4
        CV_1_4 = QCO.get_full_sized_gate("CV_1_4", num_qubits);
        CV_2_4 = QCO.get_full_sized_gate("CV_2_4", num_qubits);
        CV_3_4 = QCO.get_full_sized_gate("CV_3_4", num_qubits);
        CVdagger_3_4 = QCO.get_full_sized_gate("CVdagger_3_4", num_qubits);
        CNot_1_2 = QCO.get_full_sized_gate("CNot_1_2", num_qubits);
        CNot_2_3 = QCO.get_full_sized_gate("CNot_2_3", num_qubits);

        return CV_2_4 * CNot_1_2 * CV_3_4 * CV_1_4 * CNot_2_3 * CVdagger_3_4
    end

    params = Dict{String, Any}(
    "num_qubits" => 4,
    "maximum_depth" => 2,
    "elementary_gates" => ["CV_1_4", "CV_2_4", "CV_3_4", "CVdagger_3_4", "CNot_1_2", "CNot_2_3", "Identity"],
    "target_gate" => target_gate()
    )

    model_options = Dict{Symbol, Any}(:relax_integrality => true,
                                      :convex_hull_gate_constraints => true,
                                      :optimizer_log => false)

    result_qc = QCO.run_QCModel(params, CBC; options = model_options)

    @test result_qc["termination_status"] == MOI.OPTIMAL
    @test result_qc["primal_status"] == MOI.FEASIBLE_POINT
    @test isapprox(result_qc["objective"], 0.8275862068965518, atol = tol_0)
end


# Run tests on block Newton

using NonlinearEigenproblemsTest
using NonlinearEigenproblems
using Test
using LinearAlgebra

@bench @testset "blocknewton" begin
    nep=nep_gallery("dep0",4);
    V = Matrix(1.0I, size(nep,1), 3)
    S0=zeros(3,3);
    S,V=blocknewton(nep,S=S0,X=V,logger=displaylevel,
                    armijo_factor=0.5,
                    maxit=20)

    λv=eigvals(S);
    for λ=λv
        @test minimum(svdvals(compute_Mder(nep,λ)))<sqrt(eps())
    end

end


using BranchAndBound
using JuMP
using Ipopt

const BB = BranchAndBound

# continuous relaxation of MINLP
m = Model(optimizer_with_attributes(Ipopt.Optimizer, "print_level" => 0))

# The following will be supported in the later version of JuMP.
# @variable(m, x >= 0, Int)
# @variable(m, y >= 0, Int)
# undo_relax = JuMP.relax_integrality(m)

@variable(m, x >= 0) # assume integrality
@variable(m, y >= 0) # assume integrality

@objective(m, Min, -x - y)
@constraint(m, -2*x*y + 2*y >= 1)
@constraint(m, -8*x + 10*y <= 13)

# Initialize BNB tree
tree = BB.initialize_tree(m)

# User-defined branch function
function BB.branch!(tree::BB.AbstractTree, node::BB.AbstractNode)
    @info " Node id $(node.id), status $(node.solution_status), bound $(node.bound)"
    if node.bound >= tree.best_incumbent
        @info " Fathomed by bound"
    elseif node.solution_status in [MOI.OPTIMAL, MOI.LOCALLY_SOLVED]
        if !isinteger(node.solution[x])
            @info " Branch at x = $(node.solution[x]), x >= $(ceil(node.solution[x])), x <= $(floor(node.solution[x]))"
            child1 = BB.create_child_node_with_lb(node, x, ceil(node.solution[x]))
            child2 = BB.create_child_node_with_ub(node, x, floor(node.solution[x]))
            BB.push!(tree, child1)
            BB.push!(tree, child2)
        elseif !isinteger(node.solution[y])
            @info " Branch at y = $(node.solution[y]), y >= $(ceil(node.solution[y])), y <= $(floor(node.solution[y]))"
            child1 = BB.create_child_node_with_lb(node, y, ceil(node.solution[y]))
            child2 = BB.create_child_node_with_ub(node, y, floor(node.solution[y]))
            BB.push!(tree, child1)
            BB.push!(tree, child2)
        elseif node.bound < tree.best_incumbent
            @info " New incumbent bound: $(node.bound)"
            tree.best_incumbent = node.bound
        end
    else
        @info " Fathomed by solution status: $(node.solution_status)"
    end
end

function BB.termination(tree::BB.AbstractTree)
    @info "Tree nodes: processed $(length(tree.processed)), left $(length(tree.nodes)), total $(tree.node_counter), best bound $(tree.best_bound), best incumbent $(tree.best_incumbent)"
    if BB.isempty(tree)
        @info "Completed the tree search"
        return true
    end
    if length(tree.processed) >= 20
        @info "Reached node limit"
        return true
    end
    return false
end

BB.run(tree)


push!(LOAD_PATH, string(pwd(),"/src"))
using ModifiedStokesSolver
using PyPlot
using LaTeXStrings

n = 500
zmax = 5
z = linspace(1e-8, zmax, n)
T1S = zeros(n)
T2S = zeros(n)
T3S = zeros(n)
T1L = zeros(n)
T2L = zeros(n)
T3L = zeros(n)

for i=1:n
    T1S[i], T2S[i], T3S[i], T1L[i], T2L[i], T3L[i] = ModifiedStokesSolver.Ti_split(z[i])
end

figure(1)
clf()
plot(z, T1S, label="T1S(z)")
plot(z, T2S, label="T2S(z)")
plot(z, T3S, label="T3S(z)")
plot(z, T1L, label="T1L(z)")
plot(z, T2L, label="T2L(z)")
plot(z, T3L, label="T3L(z)")

xlim(0, zmax)
legend()
xlabel("z")
grid("on")
show()

figure(2)
clf()
f = [1.0, 0.0]
alpha = 1.0
x = linspace(-3.0, 3.0, 10000)
u = map( x -> stokeslet([x, 0.0], f, alpha)[1], x)
plot(x, u, label=L"S(\alpha|x-y|)")
axis([-3,3,0,0.5])
xlabel(L"\alpha |x-y|")
legend()


# auxiliary functions
find_negatives(targets) = findall(vec(targets) .== 0)
find_positives(targets) = findall(vec(targets) .== 1)

@nograd find_negatives, find_positives

# surrogate functions
hinge(x, ϑ::Real = 1) = max(zero(x), 1 + ϑ * x)
quadratic(x, ϑ::Real = 1) = max(zero(x), 1 + ϑ * x)^2

# objectives
function fnr(targets, scores, t, surrogate = quadratic)
    return mean(surrogate.(t .- scores[find_positives(targets)]))
end

function fpr(targets, scores, t, surrogate = quadratic)
    return mean(surrogate.(scores[find_negatives(targets)] .- t))
end

# root finding
function find_root(f::Function, x_0::Real)
    try
        x = Roots.secant_method(f, x_0)
        isnan(x) && error("secant_method failed")
        return x
    catch
        try
            x = Roots.find_zero(f, (x_0, Inf))
            isnan(x) && error("find_zero failed")
            return x
        catch
            xrange = range(x_0; stop = 10*x_0, length = 100)
            xzero, ind  = findmin(abs.(f.(xrange)))
            return xrange[ind]
        end
    end
end


using OrdinaryDiffEq, FFTW, LinearAlgebra, PyPlot

# rhs for the kuramoto - sivashinsky equation
function kse_spectral_rhs!(dy, y, p, t)
    forward_rfft_plan, inverse_rfft_plan, u, sq, A, B = p
    # calc F[u^2] using fft(u^2) = fft(ifft(y)^2)
    # (avoid allocating memory and using the plans)
    mul!(u, inverse_rfft_plan, y)
    u .*= u
    mul!(sq, forward_rfft_plan, u)
    # equation in fourier space:
    @. dy = sq*A + B*y
end

function plot_kse()
    # parameter
    b = 30.0
    dx = 0.4
    dt = 0.01
    tmax = 100.0

    # precompute and preallocate things
    # spatial discretisation
    x = 0:dx:b

    # initial condition
    u0 = @. cos((2*pi/b)*x)

    # fouier transform stuff
    ks = collect(FFTW.rfftfreq(length(u0), 1/dx)) * 2*pi
    forward_rfft_plan = FFTW.plan_rfft(u0)
    y0 = forward_rfft_plan * u0
    inverse_rfft_plan = FFTW.plan_irfft(y0, length(u0))

    # constant coefficients
    A = @. im*ks/2.0
    C = @. im*ks
    B = @. ks^2 - ks^4

    # temporal memory
    u = similar(u0)
    sq = similar(y0)

    # parameters
    p = (forward_rfft_plan, inverse_rfft_plan, u, sq, A, B, C)

    # solve the ode problem with a library solver
    t = 0.0:dt:tmax
    ode_problem = ODEProblem(kse_spectral_rhs!, y0, (0.0, tmax), p)
    @time solution = solve(ode_problem, Tsit5(); saveat=t, maxiters=10^10)
    u = [inverse_rfft_plan * y for y in solution.u]

    # plot the solution
    pcolormesh(x, t, u, cmap="BrBG", shading="nearest")
    xlabel("x")
    ylabel("t")
    colorbar(label="u")
    title("kuramoto - sivashinsky equation")
    tight_layout()
end

plot_kse()


import Base.sign

struct PartialTransposeAtom <: AbstractExpr
    head::Symbol
    id_hash::UInt64
    children::Tuple{AbstractExpr}
    size::Tuple{Int,Int}
    sys::Int
    dims::Vector

    function PartialTransposeAtom(x::AbstractExpr, sys::Int, dims::Vector)
        if x.size[1] ≠ x.size[2]
            throw(ArgumentError("Only square matrices are supported"))
        end
        if !(1 ≤ sys ≤ length(dims))
            throw(
                ArgumentError(
                    "Invalid system, should between 1 and ",
                    length(dims),
                    "; got ",
                    sys,
                ),
            )
        end
        if x.size[1] ≠ prod(dims)
            throw(
                ArgumentError(
                    "Dimension of system doesn't correspond to dimension of subsystems",
                ),
            )
        end
        children = (x,)
        return new(
            :partialtranspose,
            hash(children),
            children,
            x.size,
            sys,
            dims,
        )
    end
end

function sign(x::PartialTransposeAtom)
    return sign(x.children[1])
end

function curvature(x::PartialTransposeAtom)
    return ConstVexity()
end

function monotonicity(x::PartialTransposeAtom)
    return (Nondecreasing(),)
end

function evaluate(x::PartialTransposeAtom)
    return partialtranspose(evaluate(x.children[1]))
end

"""
    partialtranspose(x, sys::Int, dims::Vector)

Returns the partial transpose of `x` over the `sys`th system, where `dims` is a vector of integers encoding the dimensions of each subsystem.
"""
function partialtranspose(x::AbstractMatrix, sys::Int, dims::Vector)
    if size(x, 1) ≠ size(x, 2)
        throw(ArgumentError("Only square matrices are supported"))
    end
    if !(1 ≤ sys ≤ length(dims))
        throw(
            ArgumentError(
                "Invalid system, should between 1 and $(length(dims)); got $sys",
            ),
        )
    end
    if size(x, 1) ≠ prod(dims)
        throw(
            ArgumentError(
                "Dimension of system doesn't correspond to dimension of subsystems",
            ),
        )
    end
    n = length(dims)
    d = prod(dims)
    s = n - sys + 1
    p = collect(1:2n)
    p[s] = n + s
    p[n+s] = s

    rdims = reverse(dims)
    r = reshape(x, (rdims..., rdims...))
    return reshape(permutedims(r, p), (d, d))
end

"""
    permutedims_matrix(dims, p)

Returns a matrix `M` so that for any vector `v` of length `prod(dims)`,

    M*v == vec(permutedims(reshape(v, dims), p))

"""
function permutedims_matrix(dims, p)
    d = prod(dims)
    n = length(dims)
    return sparse(
        reshape(
            PermutedDimsArray(
                reshape(Matrix(I, d, d), (dims..., dims...)),
                (p..., (n+1:2n)...),
            ),
            (d, d),
        ),
    )
end

# borrowing this from transpose.jl:
# Since everything is vectorized, we simply need to multiply x by a permutation
# matrix such that coeff * vectorized(x) - vectorized(x') = 0
function conic_form!(
    x::PartialTransposeAtom,
    unique_conic_forms::UniqueConicForms,
)
    if !has_conic_form(unique_conic_forms, x)
        objective = conic_form!(x.children[1], unique_conic_forms)

        n = length(x.dims)
        d = prod(x.dims)
        s = n - x.sys + 1
        p = collect(1:2n)
        p[s] = n + s
        p[n+s] = s

        rdims = reverse(x.dims)

        partialtranspose_matrix = permutedims_matrix((rdims..., rdims...), p)

        objective = partialtranspose_matrix * objective
        cache_conic_form!(unique_conic_forms, x, objective)
    end
    return get_conic_form(unique_conic_forms, x)
end

function partialtranspose(x::AbstractExpr, sys::Int, dim::Vector)
    return PartialTransposeAtom(x, sys, dim)
end


#==========================================================================
  Regularized Quadratic Discriminant Analysis Solver
==========================================================================#

immutable ModelQDA{T<:BlasReal}
    order::Union{Type{Val{:row}},Type{Val{:col}}}
    W_k::Vector{AbstractMatrix{T}}  # Vector of class whitening matrices
    M::Matrix{T}                    # Matrix of class means (one per row)
    priors::Vector{T}               # Vector of class priors
    gamma::Nullable{T}
    lambda::Nullable{T}
end

function show(io::IO, model::ModelQDA)
    println(io, "Quadratic Discriminant Model")
    println(io, "\nRegularization Parameters:")
    print(io, "    γ = ")
    isnull(model.gamma) ? print(io, "N/A") : showcompact(get(model.gamma))
    print(io, "\n    λ = ")
    isnull(model.lambda) ? print(io, "N/A") : showcompact(get(model.lambda))
    println(io, "\n\nClass Priors:")
    for i in eachindex(model.priors)
        print(io, "    Class ", i, " Probability: ")
        showcompact(io, model.priors[i]*100)
        println(io, "%")
    end
    println(io, "\nClass Means:")
    for j in eachindex(model.priors)
        print(io, "    Class ", j, " Mean: [")
        print(io, join([sprint(showcompact, v) for v in subvector(model.order, model.M, j)], ", "))
        println(io, "]")
    end
    print(io, "\n")
end

function covmatrix{T<:BlasReal}(::Type{Val{:row}}, H::Matrix{T})
    n = size(H,1)
    Σ = H'H
    broadcast!(/, Σ, Σ, n-1)
end

function covmatrix{T<:BlasReal}(::Type{Val{:col}}, H::Matrix{T})
    n = size(H,2)
    Σ = A_mul_Bt(H, H)
    broadcast!(/, Σ, Σ, n-1)
end

for (scheme, dim_obs) in ((:(:row), 1), (:(:col), 2))
    isrowmajor = dim_obs == 1
    dim_param = isrowmajor ? 2 : 1

    H_i      = isrowmajor ? :(H[y .== i, :]) : :(H[:, y .== i])
    _ij, _ji = isrowmajor ? (:i, :j) : (:j, :i)  # Swapped variables for row and column ordering
    _nk, _kn = isrowmajor ? (:n, :k) : (:k, :n)
    _np, _pn = isrowmajor ? (:n, :p) : (:p, :n)
    W_j, H   = isrowmajor ? (:(W_k[j]), :H) : (:H, :(W_k[j]))

    @eval begin
        # λ-regularized QDA - require full covariance matrices
        function classwhiteners!{T<:BlasReal,U}(
                 ::Type{Val{$scheme}},
                H::Matrix{T},
                y::RefVector{U},
                γ::Nullable{T},
                λ::T
            )
            k = convert(Int64, y.k)
            Σ_k = AbstractMatrix{T}[covmatrix(Val{$scheme}, $H_i) for i = 1:k]
            Σ   = covmatrix(Val{$scheme}, H)
            for i in eachindex(Σ_k)
                regularize!(Σ_k[i], λ, Σ)
                Σ_k[i] = whitencov_chol!(Val{$scheme}, Σ_k[i], γ)
            end
            Σ_k
        end

        # No λ-regularization - only need data matrices
        function classwhiteners!{T<:BlasReal,U}(
                 ::Type{Val{$scheme}},
                H::Matrix{T},
                y::RefVector{U},
                γ::Nullable{T}
            )
            k = convert(Int64, y.k)
            if isnull(γ)
                AbstractMatrix{T}[whitendata_qr!(Val{$scheme}, $H_i) for i = 1:k]
            else
                AbstractMatrix{T}[whitendata_svd!(Val{$scheme}, $H_i, get(γ)) for i = 1:k]
            end
        end

        function qda!{T<:BlasReal,U}(
                 ::Type{Val{$scheme}},
                X::Matrix{T},
                M::Matrix{T},
                y::RefVector{U},
                γ::Nullable{T},
                λ::Nullable{T}
            )
            H = centerclasses!(Val{$scheme}, X, M, y)
            if isnull(λ)
                classwhiteners!(Val{$scheme}, H, y, γ)
            else
                classwhiteners!(Val{$scheme}, H, y, γ, get(λ))
            end
        end

        function discriminants_qda{T<:BlasReal}(
                 ::Type{Val{$scheme}},
                W_k::Vector{AbstractMatrix{T}},
                M::Matrix{T},
                priors::Vector{T},
                Z::Matrix{T}
            )
            n = size(Z, $dim_obs)
            p = size(Z, $dim_param)
            k = length(priors)
            D   = Array(T, $_nk, $_kn)  # discriminant function values
            H   = Array(T, $_np, $_pn)  # temporary array to prevent re-allocation k times
            hᵀh = Array(T, n)           # diag(H'H)
            for j = 1:k
                broadcast!(-, H, Z, subvector(Val{$scheme}, M, j))
                dotvectors!(Val{$scheme}, $H * $W_j, hᵀh)
                for i = 1:n
                    D[$_ij, $_ji] = -hᵀh[i]/2 + log(priors[j])
                end
            end
            D
        end
    end
end

doc"`qda(X, y; order, M, lambda, gamma, priors)` Fits a regularized quadratic discriminant model to
the data in `X` based on class identifier `y`."
function qda{T<:BlasReal,U<:Integer}(
        X::Matrix{T},
        y::AbstractVector{U};
        order::Union{Type{Val{:row}},Type{Val{:col}}} = Val{:row},
        M::Matrix{T} = classmeans(order, X, RefVector(y)),
        gamma::Union{T,Nullable{T}}  = Nullable{T}(),
        lambda::Union{T,Nullable{T}} = Nullable{T}(),
        priors::Vector{T} = ones(T,maximum(y))/maximum(y)
    )
    all(priors .> 0) || error("Argument priors must have positive values only")
    isapprox(sum(priors), one(T)) || error("Argument priors must sum to 1")
    yref = isa(y,RefVector) ? y : RefVector(y)
    length(priors) == yref.k || throw(DimensionMismatch("Prior length must match class count"))
    γ = isa(gamma,  Nullable) ? gamma  : Nullable(gamma)
    λ = isa(lambda, Nullable) ? lambda : Nullable(lambda)
    W_k = qda!(order, copy(X), M, yref, γ, λ)
    ModelQDA{T}(order, W_k, M, priors, γ, λ)
end

doc"`discriminants(Model, Z)` Uses `Model` on input `Z` to product the class discriminants."
function discriminants{T<:BlasReal}(mod::ModelQDA{T}, Z::Matrix{T})
    discriminants_qda(mod.order, mod.W_k, mod.M, mod.priors, Z)
end

doc"`classify(Model, Z)` Uses `Model` on input `Z`."
function classify{T<:BlasReal}(mod::ModelQDA{T}, Z::Matrix{T})
    vec(mapslices(indmax, discriminants(mod, Z), mod.order == Val{:row} ? 2 : 1))
end


# unit test the bGrad

include("../src/DMD_Util.jl");
include("../src/DMD_Type.jl");
include("../src/DMD_Func.jl");

# Generate DMD Synthetic Data
#------------------------------------------------------------------------------
# dimensions
m = 100;    # temporal dimension
n = 100;    # spatial dimension
k = 3;      # number of modes
T = Float64;
# generate data
X, t, at, Bt = simDMD(m, n, k, T; seed=123456);
# specify the loss
function lossFunc(r)
	return sum(abs2, r)
end
#
function lossGrad(r)
	conj!(r)
end
#
# Create object
#------------------------------------------------------------------------------
params = DMDParams(k, X, t, lossFunc, lossGrad);
#
# Initialize with true parameter
#------------------------------------------------------------------------------
copyto!(params.a, at);
copyto!(params.B, Bt);
#
# Evaluate bGrad
#------------------------------------------------------------------------------
update_P!(params);
id = rand(1:n);
#
bt = copy(params.br[id]);
Random.randn!(params.br[id]);
#
projb(params, id);
#
err = sum(abs2, bt - params.br[id]);
#
if err < 1e-6
	println("projb: OK");
else
	println("projb: Wrong, err: $err");
end


include("../mpiops.jl")

mpi_init()

function kappa(x, y)
    return 1+x^2+y^2
end

function rhs(x, y)
    -2*x*(1 - x)*(x^2 + y^2 + 1) + 2*x*(-x*y*(1 - y) + y*(1 - x)*(1 - y)) -
             2*y*(1 - y)*(x^2 + y^2 + 1) + 2*y*(-x*y*(1 - x) + x*(1 - x)*(1 - y))
end

function ufunc(x, y)
    x * (1-x) * y * (1-y)
end

mc = MPIConfig(50)
global_to_local, local_to_global = dofmap(mc)
X, Y = get_xy(mc)
f_local = rhs.(X, Y)
using Random; Random.seed!(233)
θ = Variable(fc_init([2,20,20,20,1]))
θ_shared = mpi_bcast(θ)
κ_local = (fc([X'[:] Y'[:]], [20,20,20,1], θ_shared) + 5.0)|>squeeze
κ_local = reshape(κ_local, (mc.n, mc.n))
u_local = poisson_solver(κ_local, f_local, mc)
u = mpi_gather(u_local)[global_to_local]

loss = sum(u^2)
g = gradients(loss, θ)

sess = Session(); init(sess)

function calculate_loss(x)
    run(sess, loss, θ=>x)
end

function calculate_gradients(x)
    run(sess, g, θ=>x)
end

f(x) = (calculate_loss(x), calculate_gradients(x))

test_gradients(f, run(sess, θ), mpi=true)

if mpi_size()>1
    mpi_finalize()
end


export pred_ared_computation

"""
This function computes the actual and predicted reduction for one dimensionnal
TR or ARC algorithm. It also returns the ratio bewtween the two.
Inputs:
  - gₖ, the value of the current derivative of the function
  - fₖ, the value of the current value of the function
  - dersec, the value of the current second derivative of the function
  - d, the value of the descent direction
  - ftestTR, the candiate fonction value (h(x + t * d))
  - gtestTR, the candiate derivative value (h'(x + t * d))
  - (opt) seuil, numerical tolerance for the actual reduction. Default value
    is -1e-10.
"""
function pred_ared_computation(gₖ :: Float64, fₖ :: Float64,  dersec :: Float64,
                               d        :: Float64, ftestTR :: Float64,
                               gtestTR  :: Float64; seuil :: Float64 = -1e-10)

  pred = gₖ * d + 0.5 * dersec * d^2
  if pred > seuil
    ared = (gₖ + gtestTR)*d/2
  else
    ared = ftestTR-fₖ
  end

  ratio=ared/pred
  return (pred,ared,ratio)
end


# Compute the Leonard-Jones potential energy function

function upair(r,data)
  r6 = (r^2)^3
  r12 = (r^2)^6
  return data.eps4*(data.sig12/r12 - data.sig6/r6)

end


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:4] <= 1)

                     @variable(m, -6.019960436071491 <= q <= 11.696817216382863)

                     add_NL_constraint(m, :(gelu(0.9042935565441375 + 0.31236426273908524*gelu(0.5868252904954097 + 0.04712499510197299*$(x[1]) + 0.4005456947892885*$(x[2]) + 0.34544434364772814*$(x[3]) + 0.48957180414623247*$(x[4])) + -0.4128166127064836*gelu(0.5797949667252156 + -0.7245574143823843*$(x[1]) + -0.08207866517990237*$(x[2]) + 0.4631961331947827*$(x[3]) + 0.9179781125579716*$(x[4])) + -0.2476169775927386*gelu(0.8773151656273224 + 0.003673243009740812*$(x[1]) + 0.698013040735662*$(x[2]) + 0.6787695040522581*$(x[3]) + -0.9128403322093082*$(x[4]))) + gelu(0.8600467319105785 + -0.4006275013725942*gelu(0.5868252904954097 + 0.04712499510197299*$(x[1]) + 0.4005456947892885*$(x[2]) + 0.34544434364772814*$(x[3]) + 0.48957180414623247*$(x[4])) + 0.22106887260345687*gelu(0.5797949667252156 + -0.7245574143823843*$(x[1]) + -0.08207866517990237*$(x[2]) + 0.4631961331947827*$(x[3]) + 0.9179781125579716*$(x[4])) + 0.7210213948433983*gelu(0.8773151656273224 + 0.003673243009740812*$(x[1]) + 0.698013040735662*$(x[2]) + 0.6787695040522581*$(x[3]) + -0.9128403322093082*$(x[4]))) + gelu(0.982756228863964 + 0.34792310275142935*gelu(0.5868252904954097 + 0.04712499510197299*$(x[1]) + 0.4005456947892885*$(x[2]) + 0.34544434364772814*$(x[3]) + 0.48957180414623247*$(x[4])) + 0.7924515917531365*gelu(0.5797949667252156 + -0.7245574143823843*$(x[1]) + -0.08207866517990237*$(x[2]) + 0.4631961331947827*$(x[3]) + 0.9179781125579716*$(x[4])) + -0.9399735455074207*gelu(0.8773151656273224 + 0.003673243009740812*$(x[1]) + 0.698013040735662*$(x[2]) + 0.6787695040522581*$(x[3]) + -0.9128403322093082*$(x[4]))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


export make_node, score

import StatsFuns: logπ
import SpecialFunctions: loggamma
import Distributions: Wishart

@node Wishart Stochastic [ out, (ν, aliases = [ df ]), (S, aliases = [ scale ]) ]

@average_energy Wishart (q_out::Any, q_ν::Any, q_S::Any) = begin
    d = dim(q_out)

    m_q_ν   = mean(q_ν)
    m_q_S   = mean(q_S)
    m_q_out = mean(q_out)

    return 0.5 * (
        m_q_ν * (logdet(m_q_S) + d * log(2)) -
        logdet(m_q_out) * (m_q_ν- d - 1.0) +
        tr(cholinv(m_q_S) * m_q_out) + 0.5 * d * (d - 1) * logπ
    ) + mapreduce(i -> loggamma(0.5 * (m_q_ν + 1.0 - i)), +, 1:d)
end


using MTH229

using Test

## test package
@testset "test packages" begin

    ## Roots
    @test fzero(sin, 3, 4)  ≈ pi
    @test fzero(sin, 3.0)  ≈ pi

    ## Zygote
    f(x) = sin(x)
    @test f'(2)  ≈ cos(2)
    @test f''(2)  ≈ -sin(2)

    ## QuadGK
    @test quadgk(sin, 0, pi)[1]  ≈ 2
    @test abs(quadgk(x -> x^2, 0, 1)[2]) <= 1e-14

    ## SymPy
    @test limit(x -> sin(x)/x, 0) == 1
    @test integrate(sin,  0, pi) == 2

end

@testset "test functions" begin

    f(x) = sin(x)
    c = pi/4
    fn = tangent(f, c)
    @test fn(1)  ≈ f(c) + f'(c)*(1 - c)

    fn = secant(f, pi/6, pi/3)
    @test fn(pi/4) <= f(pi/4)

    out = lim(x -> sin(x)/x, 0)
    @test out[end, 2]  ≈ 1

    @test bisection(sin, 3, 4)  ≈ pi
    @test newton(sin, cos, 3.0)  ≈ pi
    @test newton(sin, 3.0)  ≈ pi

    f₁ = x ->  1 + 100x^2 - x^3
    @test isempty(find_zeros(f₁, -100, 100) )
    f₁s = find_zeros(fisheye(f₁), -pi/2, pi/2)
    @test !isempty(f₁s)
    @test abs(f₁(tan(first(f₁s)))) <= 1e-4

    out = sign_chart(x -> (x-1)*(x-2)/(x-3), 0, 4)
    @test all([o[1] for o ∈ out] .≈[1,2,3])

    @test riemann(sin, 0, pi, 10_000)  ≈ 2
    @test fubini((x,y) -> 1, (x->-sqrt(1-x^2), x->sqrt(1-x^2)), (-1,1)) ≈ pi
end

@testset "2d" begin

    x = [[1,2,3], [4,5,6]]
    @test unzip(x)[1] == [1, 4]
    @test unzip(x)[2] == [2, 5]
    @test unzip(x)[3] == [3, 6]

    @test length(unzip(x -> x, 0, 1)[1])  <= 50 # 21
    @test length(unzip(x-> sin(10pi*x), 0, 1)[1]) >= 50 # 233
    @test length(unzip(x -> x, 0, 1, 10)[1])  == 10

    @test uvec([2,2]) == 1/sqrt(2) * [1,1]

end

@testset "plots" begin

    if isinteractive()

        f(x) = 1/x

        trimplot(f, -1, 1)

        plotif(f, f, -1, 1)
        plotif(f, f', -1, 1)

        signchart(sin, 0, 4pi)

        newton_vis(log, 1/2)

        r(t) = [sin(t), cos(t)]
        ts = range(0, stop=pi/2, length=100)
        plot(unzip(r.(ts))...)
        arrow!(r(pi/4), r'(pi/4))

        V(x,y) = [x, x-y]
        plot()
        vectorfieldplot!(V)

    end

end


function Div(Px, Py, bound = "sym", order = 1)
    """
        div - divergence operator

        fd = div(Px,Py, options);
        fd = div(P, options);

          options.bound = 'per' or 'sym'
          options.order = 1 (backward differences)
                        = 2 (centered differences)

          Note that the -div and grad operator are adjoint
          of each other such that
              <grad(f),g>=<f,-div(g)>

          See also: grad.

        Copyright (c) 2007 Gabriel Peyre
    """

    # retrieve number of dimensions
    nbdims = ndims(Px)

    if nbdims >= 3
        if nbdims == 3
            Py = P[:, :, 2]
            Px = P[:, :, 1]
            nbdims = 2
        else
            Pz = P[:, :, :, 3]
            Py = P[:, :, :, 2]
            Px = P[:, :, :, 1]
            nbdims = 3
        end
    end

    if bound == "sym"
        nx = size(Px)[1]
        if order == 1
            fx = Px .- Px[vcat(1, collect(1, nx - 1)), :]
            fx[1, :] = Px[1, :]                        # boundary
            fx[nx, :] = - Px[nx - 1, :]

            if nbdims >= 2
                ny = size(Py)[2]
                fy = Py - Py[:, vcat(1, collect(1, ny - 1))]
                fy[:, 1] = Py[:, 1]                    # boundary
                fy[:, ny] = - Py[:, ny - 1]
            end

            if nbdims >= 3
                nz = size(Pz)[3]
                fz = Pz - Pz[:, :, vcat(1, collect(1 : nz - 1))]
                fz[:, :, 1] = Pz[:, :, 1]                # boundary
                fz[:, :, nz] = - Pz[:, :, nz - 1]
            end
        else
            fx = (Px[vcat(collect(2 : nx), nx), :] - Px[vcat([1], collect(1, nx-1)), :])./2.
            fx[1, :] = + Px[1, :]./2. + Px[1, :]           # boundary
            fx[1, :] = + Px[2, :]./2. - Px[1, :]
            fx[nx, :] = - Px[nx, :] - Px[nx - 1, :]./2.
            fx[nx - 1, :] = + Px[nx, :] - Px[nx - 2, :]./2.

            if nbdims >= 2
                ny = size(Py)[2]
                fy = (Py[:, vcat(collect(2 : ny), ny)] - Py[:, vcat(1, collect(1 : ny-1))])./2.
                fy[:, 1] = + Py[:, 2]./2. + Py[:, 1]       # boundary
                fy[:, 2] = + Py[:, 3]./2. - Py[:,1]
                fy[:, ny] = - Py[:, ny] - Py[:, ny-1]./2.
                fy[:, ny - 1] = + Py[:, ny] - Py[:, ny - 2]./2.
            end

            if nbdims >= 3
                nz = size(Pz)[3]
                fz = (Pz[:, :, vcat(collect(2 : nz), nz)] - Pz[:, :, vcat(1, collect(1 : nz - 1))])./2.
                fz[:, :, 1] = + Pz[:, :, 2]./2. + Pz[:, :, 1] # boundary
                fz[:, :, 2] = + Pz[:, :, 3]./2. - Pz[:, :, 1]
                fz[:, :, ny] = - Pz[:, :, nz] - Pz[:, :, nz - 1]./2.
                fz[:, :, ny - 1] = + Pz[:, :, nz] - Pz[:, :, nz - 2]./2.
            end
        end

    else
        if order == 1
            nx = size(Px)[1]
            fx = Px - Px[vcat(nx, collect(1 : nx-1)), :]

            if nbdims >= 2
                ny = size(Py)[2]
                fy = Py - Py[:, vcat(ny, collect(1 : ny-1))]
            end

            if nbdims>=3
                nz = size(Pz)[3]
                fz = Pz - Pz[:, :, vcat(nz, collect(1 : nz-1))]
            end

        else
            nx = size(Px)[1]
            fx = (Px[vcat(collect(2 : nx), 1), :]) - (Px[vcat(nx, collect(1, nx - 1)), :])

            if nbdims >= 2
                ny = size(Py)[2]
                fy = (Py[:, vcat(collect(2 : ny), 1)]) - (Py[:, vcat(ny, collect(1, ny - 1))])
            end

            if nbdims >= 3
                nz = size(Pz)[3]
                fz = (Pz[:, :, vcat(collect(2, nz), 1)]) - (Pz[:, :, vcat(nz, collect(1 : nz - 1))])
            end
        end
    end

    # gather result
    if nbdims == 3
        fd = fx + fy + fz

    elseif nbdims == 2
        fd = fx + fy

    else
        fd = fx
    end

    return fd
end


# Compute extension velocity field for Fisher-Stefan
# Alex Tam, 09/11/2021

"Compute velocity extension by orthogonal extrapolation"
function extend_velocity(D, dΩ, U, ϕ, par, dx, dy)
    dτ = 0.5*(dx/5 + dy/5) # "Time" step for extrapolation PDE
    vi = interface_velocity(dΩ, U, ϕ, par, dx, dy) # Find speed at interface using Stefan condition
    V = velocity_init(D, dΩ, par, vi) # Obtain "initial" condition for V(x,y,t) using speed at interface
    v = build_vector(V, D) # Create vector using data at interior grid points
    # Solve PDE to extrapolate V outwards from interface
    prob = ODEProblem((du, u, p, t) -> vel_outward_rhs!(du, u, p, t, D, dΩ, vi, par, ϕ, dx, dy), v, (0, par.V_Iterations*dτ))
    sol = solve(prob, Tsit5(), reltol = 1e-3, abstol = 1e-6, saveat = par.V_Iterations*dτ)
    v = sol[:,end] # Update vector of V
    # # Solve PDE to extrapolate V inwards from interface
    prob = ODEProblem((du, u, p, t) -> vel_inward_rhs!(du, u, p, t, D, dΩ, vi, par, ϕ, dx, dy), v, (0, par.V_Iterations*dτ))
    sol = solve(prob, Tsit5(), reltol = 1e-3, abstol = 1e-6, saveat = par.V_Iterations*dτ)
    # Reshape solution to matrix
    V = build_v_matrix(sol[:,end], par, D)
    return V
end

"Apply Stefan condition to find speed at interface points"
function interface_velocity(dΩ, U, ϕ, par, dx, dy)
    vi = Vector{Float64}() # Pre-allocate empty vector of interface speeds
    for ip in dΩ # Loop over interface points
        # Compute du/dx, dϕ/dx, du/dy and dϕ/dy
        if ip.m.xInd != ip.p.xInd # If interface lies in x-direction
            # du/dx
            if (ip.m.Ω == true) && (ip.m.dΩ == false) # Left point lies inside Ω and not close to interface
                if ϕ[ip.m.xInd-1,ip.m.yInd] < 0 # Check adjacent grid point
                    um = U[ip.m.xInd-1,ip.m.yInd] # Set to grid value if inside Ω
                else
                    um = par.u_f # Set to interface value if outisde Ω
                end
                ux = (1/dx)*(ip.m.θxp*um/(ip.m.θxm*(ip.m.θxm+ip.m.θxp)) - (ip.m.θxm+ip.m.θxp)*U[ip.m.xInd,ip.m.yInd]/(ip.m.θxm*ip.m.θxp) + par.u_f*(2*ip.m.θxp+ip.m.θxm)/(ip.m.θxp*(ip.m.θxp+ip.m.θxm))) # Biased stencil
            elseif (ip.m.Ω == true) && (ip.m.dΩ == true) # Left point lies inside Ω and close to interface
                ϕm = ϕ[ip.m.xInd-1,ip.m.yInd]
                ϕmm = ϕ[ip.m.xInd-2,ip.m.yInd]
                if (ϕm < 0) && (ϕmm < 0) # Check two adjacent grid points
                    ux = 1/(2*dx)*(3*U[ip.m.xInd,ip.m.yInd] - 4*U[ip.m.xInd-1,ip.m.yInd] + U[ip.m.xInd-2,ip.m.yInd]) # Second-order
                elseif (ϕm < 0) && (ϕmm >= 0)
                    θ = ϕm/(ϕm - ϕmm)
                    if θ < par.θb # If opposite ghost node is close to interface
                        ux = 0.0
                    else
                        ux = (1/dx)*(par.u_f/(θ*(1+θ)) - (U[ip.m.xInd-1,ip.m.yInd]*(1+θ))/θ + (par.u_f*(2+θ))/(1+θ)) # Second-order difference
                    end
                else
                    ux = 0.0
                end
            elseif (ip.p.Ω == true) && (ip.p.dΩ == false) # Right point lies inside Ω and not close to interface
                if ϕ[ip.p.xInd+1,ip.p.yInd] < 0 # Check adjacent grid point
                    up = U[ip.p.xInd+1,ip.p.yInd] # Set to grid value if inside Ω
                else
                    up = par.u_f # Set to grid value if outside Ω
                end
                ux = (1/dx)*(-par.u_f*(2*ip.p.θxm + ip.p.θxp)/(ip.p.θxm*(ip.p.θxm+ip.p.θxp)) + (ip.p.θxm+ip.p.θxp)*U[ip.p.xInd,ip.p.yInd]/(ip.p.θxm*ip.p.θxp) - ip.p.θxm*up/(ip.p.θxp*(ip.p.θxm+ip.p.θxp))) # Biased stencil
            else # Right point lies inside Ω and close to interface
                ϕp = ϕ[ip.p.xInd+1,ip.p.yInd]
                ϕpp = ϕ[ip.p.xInd+2,ip.p.yInd]
                if (ϕp < 0) && (ϕpp < 0) # Check two adjacent grid points
                    ux = 1/(2*dx)*(-U[ip.p.xInd+2,ip.p.yInd] + 4*U[ip.p.xInd+1,ip.p.yInd] - 3*U[ip.p.xInd,ip.p.yInd]) # Second-order
                elseif (ϕp < 0) && (ϕpp >= 0)
                    θ = ϕp/(ϕp - ϕpp)
                    if θ < par.θb # If opposite ghost node is close to interface
                        ux = 0.0
                    else
                        ux = (1/dx)*(-(par.u_f*(2+θ))/(1+θ) + (U[ip.m.xInd+1,ip.m.yInd]*(1+θ))/θ - par.u_f/(θ*(1+θ))) # Second-order difference
                    end
                else
                    ux = 0.0
                end
            end
            # dϕ/dx
            ϕxl = (ϕ[ip.m.xInd+1, ip.m.yInd]-ϕ[ip.m.xInd-1, ip.m.yInd])/(2*dx)
            ϕxr = (ϕ[ip.p.xInd+1, ip.p.yInd]-ϕ[ip.p.xInd-1, ip.p.yInd])/(2*dx)
            ϕx = ip.m.θxp*ϕxr + (1-ip.m.θxp)*ϕxl
            # dϕ/dy
            ϕp = ip.m.θxp*ϕ[ip.p.xInd, ip.p.yInd+1] + (1-ip.m.θxp)*ϕ[ip.m.xInd, ip.m.yInd+1]
            ϕm = ip.m.θxp*ϕ[ip.p.xInd, ip.p.yInd-1] + (1-ip.m.θxp)*ϕ[ip.m.xInd, ip.m.yInd-1]
            ϕy = (ϕp-ϕm)/(2*dy) # Use ghost point method
            # du/dy
            if ϕp < 0 # Positive ghost point inside Ω
                up =  ip.m.θxp*U[ip.p.xInd, ip.p.yInd+1] + (1-ip.m.θxp)*U[ip.m.xInd, ip.m.yInd+1]
                ϕpp =  ip.m.θxp*ϕ[ip.p.xInd, ip.p.yInd+2] + (1-ip.m.θxp)*ϕ[ip.m.xInd, ip.m.yInd+2] # Second ghost point
                if ϕpp < 0
                    upp =  ip.m.θxp*U[ip.p.xInd, ip.p.yInd+2] + (1-ip.m.θxp)*U[ip.m.xInd, ip.m.yInd+2]
                    uy = (-upp + 4*up - 3*par.u_f)/(2*dy) # Second-order difference
                else
                    θ = ϕp/(ϕp - ϕpp) # Distance between middle ghost point and interface
                    if θ < par.θb
                        uy = 0.0 # Set to zero if ghost point close to interface
                    else
                        uy = ((-(2+θ)*par.u_f)/(1+θ) + (1+θ)*up/θ - par.u_f/(θ*(1+θ)))/dy # Second-order difference
                    end
                end
            elseif ϕm < 0 # Negative ghost point inside Ω
                um =  ip.m.θxp*U[ip.p.xInd, ip.p.yInd-1] + (1-ip.m.θxp)*U[ip.m.xInd, ip.m.yInd-1]
                ϕmm = ip.m.θxp*ϕ[ip.p.xInd, ip.p.yInd-2] + (1-ip.m.θxp)*ϕ[ip.m.xInd, ip.m.yInd-2]
                if ϕmm < 0
                    umm = ip.m.θxp*U[ip.p.xInd, ip.p.yInd-2] + (1-ip.m.θxp)*U[ip.m.xInd, ip.m.yInd-2]
                    uy = (3par.u_f - 4*um + umm)/(2*dy) # Second-order difference
                else
                    θ = ϕm/(ϕm - ϕmm) # Distance between middle ghost point and interface
                    if θ < par.θb
                        uy = 0.0 # Set to zero if ghost point close to interface
                    else
                        uy = ((2+θ)*par.u_f/(1+θ) - (1+θ)*um/θ + par.u_f/(θ*(θ+1)))/dy # Second-order difference
                    end
                end
            else # Both ghost points outside Ω
                uy = 0.0
            end
        else # If interface lies in y-direction
            # du/dx and du/dy
            if (ip.m.Ω == true) && (ip.m.dΩ == false) # Bottom point lies inside Ω and not close to interface
                if ϕ[ip.m.xInd,ip.m.yInd-1] < 0 # Check adjacent grid point
                    ub = U[ip.m.xInd,ip.m.yInd-1] # Set to grid value if inside Ω
                else
                    ub = par.u_f # Set to interface value if outside Ω
                end
                uy = (1/dy)*(ip.m.θyp*ub/(ip.m.θym*(ip.m.θym+ip.m.θyp)) - (ip.m.θym+ip.m.θyp)*U[ip.m.xInd,ip.m.yInd]/(ip.m.θym*ip.m.θyp) + par.u_f*(2*ip.m.θyp+ip.m.θym)/(ip.m.θyp*(ip.m.θyp+ip.m.θym)) ) # Biased stencil
            elseif (ip.m.Ω == true) && (ip.m.dΩ == true) # Bottom point lies inside Ω and close to interface
                ϕm = ϕ[ip.m.xInd,ip.m.yInd-1]
                ϕmm = ϕ[ip.m.xInd,ip.m.yInd-2]
                if (ϕm < 0) && (ϕmm < 0) # Check two adjacent points
                    uy = 1/(2*dy)*(3*U[ip.m.xInd,ip.m.yInd] - 4*U[ip.m.xInd,ip.m.yInd-1] + U[ip.m.xInd,ip.m.yInd-2]) # Second-order
                elseif (ϕm < 0) && (ϕmm >= 0)
                    θ = ϕm/(ϕm - ϕmm)
                    if θ < par.θb # If opposite ghost node is close to interface
                        uy = 0.0
                    else
                        uy = (1/dy)*(par.u_f/(θ*(1+θ)) - (U[ip.m.xInd,ip.m.yInd-1]*(1+θ))/θ + (par.u_f*(2+θ))/(1+θ)) # Second-order difference
                    end
                else
                    uy = 0.0
                end
            elseif (ip.p.Ω == true) && (ip.p.dΩ == false) # Top point lies inside Ω and not close to interface
                if ϕ[ip.p.xInd,ip.p.yInd+1] < 0 # Check adjacent grid point
                    ut = U[ip.p.xInd,ip.p.yInd+1] # Set to grid value if inside Ω
                else
                    ut = par.u_f # Set to interface value if outside Ω
                end
                uy = (1/dy)*(-par.u_f*(2*ip.p.θym + ip.p.θyp)/(ip.p.θym*(ip.p.θym+ip.p.θyp)) + (ip.p.θym+ip.p.θyp)*U[ip.p.xInd, ip.p.yInd]/(ip.p.θym*ip.p.θyp) - ip.p.θym*ut/(ip.p.θyp*(ip.p.θym+ip.p.θyp)) ) # Biased stencil
            else
                ϕp = ϕ[ip.m.xInd,ip.m.yInd+1]
                ϕpp = ϕ[ip.m.xInd,ip.m.yInd+2]
                if (ϕp < 0) && (ϕpp < 0) # Check two adjacent nodes
                    uy = 1/(2*dy)*(-U[ip.p.xInd,ip.p.yInd+2] + 4*U[ip.p.xInd,ip.p.yInd+1] - 3*U[ip.p.xInd,ip.p.yInd]) # Second-order
                elseif (ϕp < 0) && (ϕpp >= 0)
                    θ = ϕp/(ϕp - ϕpp)
                    if θ < par.θb # If opposite ghost node is close to interface
                        uy = 0.0
                    else
                        uy = (1/dy)*(-(par.u_f*(2+θ))/(1+θ) + (U[ip.m.xInd,ip.m.yInd+1]*(1+θ))/θ - par.u_f/(θ*(1+θ))) # Second-order difference
                    end
                else
                    uy = 0.0
                end
            end
            # dϕ/dx
            ϕp = ip.m.θyp*ϕ[ip.p.xInd+1, ip.p.yInd] + (1-ip.m.θyp)*ϕ[ip.m.xInd+1, ip.m.yInd]
            ϕm = ip.m.θyp*ϕ[ip.p.xInd-1, ip.p.yInd] + (1-ip.m.θyp)*ϕ[ip.m.xInd-1, ip.m.yInd]
            ϕx = (ϕp-ϕm)/(2*dx) # Use ghost point method
            # dϕ/dy
            ϕyb = (ϕ[ip.m.xInd, ip.m.yInd+1]-ϕ[ip.m.xInd, ip.m.yInd-1])/(2*dy)
            ϕyt = (ϕ[ip.p.xInd, ip.p.yInd+1]-ϕ[ip.p.xInd, ip.p.yInd-1])/(2*dy)
            ϕy = ip.m.θyp*ϕyt + (1-ip.m.θyp)*ϕyb
            # du/dx
            if ϕp < 0 # Positive ghost point inside Ω
                up = ip.m.θyp*U[ip.p.xInd+1, ip.p.yInd] + (1-ip.m.θyp)*U[ip.m.xInd+1, ip.m.yInd]
                ϕpp = ip.m.θyp*ϕ[ip.p.xInd+2, ip.p.yInd] + (1-ip.m.θyp)*ϕ[ip.m.xInd+2, ip.m.yInd]
                if ϕpp < 0
                    upp = ip.m.θyp*U[ip.m.xInd+2, ip.p.yInd] + (1-ip.m.θyp)*U[ip.m.xInd+2, ip.m.yInd]
                    ux = (-upp + 4*up -3*par.u_f)/(2*dx) # Second-order
                else
                    θ = ϕp/(ϕp - ϕpp) # Distance between middle ghost point and interface
                    if θ < par.θb
                        ux = 0.0 # Set to zero if ghost point close to interface
                    else
                        ux = ((-(2+θ)*par.u_f)/(1+θ) + (1+θ)*up/θ - par.u_f/(θ*(1+θ)))/dx # Second-order
                    end
                end
            elseif ϕm < 0 # Negative ghost point inside Ω
                um = ip.m.θyp*U[ip.m.xInd-1, ip.p.yInd] + (1-ip.m.θyp)*U[ip.m.xInd-1, ip.m.yInd]
                ϕmm = ip.m.θyp*ϕ[ip.p.xInd-2, ip.p.yInd] + (1-ip.m.θyp)*ϕ[ip.m.xInd-2, ip.m.yInd]
                if ϕmm < 0
                    umm = ip.m.θyp*U[ip.p.xInd-2, ip.p.yInd] + (1-ip.m.θyp)*U[ip.m.xInd-2, ip.m.yInd]
                    ux = (3*par.u_f - 4um + umm)/(2*dx) # Second-order
                else
                    θ = ϕm/(ϕm - ϕmm) # Distance between middle ghost point and interface
                    if θ < par.θb
                        ux = 0.0 # Set to zero if ghost point close to interface
                    else
                        ux = ((2+θ)*par.u_f/(1+θ) - (1+θ)*um/θ + par.u_f/(θ*(θ+1)))/dx # Second-order
                    end
                end
            else # Both ghost points outside Ω
                ux = 0.0
            end
        end
        push!(vi, -par.κ*(ux*ϕx + uy*ϕy)) # Append speed at current point to vector
    end
    return vi
end

"Obtain initial condition for V(x,y,t) using speed at interface"
function velocity_init(D, dΩ, par, vi)
    V = zeros(par.Nx, par.Ny) # Pre-allocate matrix of V(x,y,t)
    # Manually overwrite speed close to interface
    for gp in D # Loop over interior grid points
        if (gp.dΩ == true) # If grid point is close to interface
            # Locate closest interface point to current grid point
            if (gp.θxp < gp.θxm) && (gp.θxp < gp.θyp) && (gp.θxp < gp.θym) # Closest point in positive x-direction
                xm = gp.xInd; xp = gp.xInd+1; ym = gp.yInd; yp = gp.yInd
            elseif (gp.θxm < gp.θxp) && (gp.θxm < gp.θyp) && (gp.θxm < gp.θym) # Closest point in negative x-direction
                xm = gp.xInd-1; xp = gp.xInd; ym = gp.yInd; yp = gp.yInd
            elseif (gp.θyp < gp.θxm) && (gp.θyp < gp.θxp) && (gp.θyp < gp.θym) # Closest point in positive y-direction
                xm = gp.xInd; xp = gp.xInd; ym = gp.yInd; yp = gp.yInd+1
            else # Closest point in negative y-direction
                xm = gp.xInd; xp = gp.xInd; ym = gp.yInd-1; yp = gp.yInd
            end
            # Assign V to interface speed at closest point
            for i = 1:length(dΩ)
                if (xm == dΩ[i].m.xInd) && (ym == dΩ[i].m.yInd) && (xp == dΩ[i].p.xInd) && (yp == dΩ[i].p.yInd)
                    V[gp.xInd, gp.yInd] = vi[i]
                end
            end
        end
    end
    return V
end

"Obtain velocity at a specified interface point"
function obtain_interface_velocity(xm, ym, xp, yp, dΩ, vi)
    for i = 1:length(dΩ) # Loop over candidate points
        if (xm == dΩ[i].m.xInd) && (ym == dΩ[i].m.yInd) && (xp == dΩ[i].p.xInd) && (yp == dΩ[i].p.yInd) # If all indices are correct
            return vi[i]
        end
    end
end

"Construct right-hand vector of outward velocity extension PDE, for use in DifferentialEquations.jl"
function vel_outward_rhs!(du, u, p, t, D, dΩ, vi, par, ϕ, dx, dy)
    V = build_v_matrix(u, par, D) # Generate matrix
    for i = 1:length(D) # Loop over interior grid points
        if (D[i].Ω == true) || (D[i].dΩ == true) # Grid points in Ω or close to interface
            du[i] = 0.0 # Set RHS to zero
        else
            # Compute advection coefficients
            ϕx = 1/(2*dx)*(ϕ[D[i].xInd+1, D[i].yInd]-ϕ[D[i].xInd-1, D[i].yInd]) # dϕ/dx
            ϕy = 1/(2*dy)*(ϕ[D[i].xInd, D[i].yInd+1]-ϕ[D[i].xInd, D[i].yInd-1]) # dϕ/dy
            n = sqrt(ϕx^2 + ϕy^2) # Normalisation
            if n == 0 # Prevent division by zero
                a = 0; b = 0
            else
                a = ϕx/n # x-component
                b = ϕy/n # y-component
            end
            # Discretise dV/dx (first-order upwind)
            if a < 0 # Forward difference
                if D[i].θxp == 1.0 # Regular discretisation
                    Vx = (V[D[i].xInd+1, D[i].yInd] - V[D[i].xInd, D[i].yInd])/dx
                else # Irregular discretisation
                    v = obtain_interface_velocity(D[i].xInd, D[i].yInd, D[i].xInd+1, D[i].yInd, dΩ, vi)
                    Vx = (v - V[D[i].xInd, D[i].yInd])/(D[i].θxp*dx)
                end
            else # Backward difference
                if D[i].θxm == 1.0 # Regular discretisation
                    Vx = (V[D[i].xInd, D[i].yInd] - V[D[i].xInd-1, D[i].yInd])/dx
                else # Irregular discretisation
                    v = obtain_interface_velocity(D[i].xInd-1, D[i].yInd, D[i].xInd, D[i].yInd, dΩ, vi)
                    Vx = (V[D[i].xInd, D[i].yInd]-v)/(D[i].θxm*dx)
                end
            end
            # Discretise dV/dy (first-order upwind)
            if b < 0 # Forward difference
                if D[i].θyp == 1.0 # Regular discretisation
                    Vy = (V[D[i].xInd, D[i].yInd+1] - V[D[i].xInd, D[i].yInd])/dy
                else # Irregular discretisation
                    v = obtain_interface_velocity(D[i].xInd, D[i].yInd, D[i].xInd, D[i].yInd+1, dΩ, vi)
                    Vy = (v - V[D[i].xInd, D[i].yInd])/(D[i].θyp*dy)
                end
            else # Backward difference
                if D[i].θym == 1.0
                    Vy = (V[D[i].xInd, D[i].yInd] - V[D[i].xInd, D[i].yInd-1])/dy
                else
                    # Obtain interface velocity
                    v = obtain_interface_velocity(D[i].xInd, D[i].yInd-1, D[i].xInd, D[i].yInd, dΩ, vi)
                    Vy = (V[D[i].xInd, D[i].yInd]-v)/(D[i].θym*dy)
                end
            end
            du[i] = -a*Vx - b*Vy # Update RHS term
        end
    end
end

"Construct right-hand vector of inward velocity extension PDE, for use in DifferentialEquations.jl"
function vel_inward_rhs!(du, u, p, t, D, dΩ, vi, par, ϕ, dx, dy)
    V = build_v_matrix(u, par, D) # Generate matrix
    for i = 1:length(D) # Loop over interior grid points
        if (D[i].Ω == false) || (D[i].dΩ == true) # Grid points outside Ω or close to interface
            du[i] = 0.0 # Set RHS to zero
        else
            # Compute advection coefficients
            ϕx = 1/(2*dx)*(ϕ[D[i].xInd+1, D[i].yInd]-ϕ[D[i].xInd-1, D[i].yInd]) # dϕ/dx
            ϕy = 1/(2*dy)*(ϕ[D[i].xInd, D[i].yInd+1]-ϕ[D[i].xInd, D[i].yInd-1]) # dϕ/dy
            n = sqrt(ϕx^2 + ϕy^2) # Normalisation
            if n == 0 # Prevent division by zero
                a = 0; b = 0
            else
                a = -ϕx/n # x-component
                b = -ϕy/n # y-component
            end
            # Discretise dV/dx (first-order upwind)
            if a < 0 # Forward difference
                if D[i].θxp == 1.0 # Regular discretisation
                    Vx = (V[D[i].xInd+1, D[i].yInd] - V[D[i].xInd, D[i].yInd])/dx
                else # Irregular discretisation
                    v = obtain_interface_velocity(D[i].xInd, D[i].yInd, D[i].xInd+1, D[i].yInd, dΩ, vi)
                    Vx = (v - V[D[i].xInd, D[i].yInd])/(D[i].θxp*dx)
                end
            else # Backward difference
                if D[i].θxm == 1.0 # Regular discretisation
                    Vx = (V[D[i].xInd, D[i].yInd] - V[D[i].xInd-1, D[i].yInd])/dx
                else # Irregular discretisation
                    v = obtain_interface_velocity(D[i].xInd-1, D[i].yInd, D[i].xInd, D[i].yInd, dΩ, vi)
                    Vx = (V[D[i].xInd, D[i].yInd]-v)/(D[i].θxm*dx)
                end
            end
            # Discretise dV/dy (first-order upwind)
            if b < 0 # Forward difference
                if D[i].θyp == 1.0 # Regular discretisation
                    Vy = (V[D[i].xInd, D[i].yInd+1] - V[D[i].xInd, D[i].yInd])/dy
                else # Irregular discretisation
                    v = obtain_interface_velocity(D[i].xInd, D[i].yInd, D[i].xInd, D[i].yInd+1, dΩ, vi)
                    Vy = (v - V[D[i].xInd, D[i].yInd])/(D[i].θyp*dy)
                end
            else # Backward difference
                if D[i].θym == 1.0
                    Vy = (V[D[i].xInd, D[i].yInd] - V[D[i].xInd, D[i].yInd-1])/dy
                else
                    # Obtain interface velocity
                    v = obtain_interface_velocity(D[i].xInd, D[i].yInd-1, D[i].xInd, D[i].yInd, dΩ, vi)
                    Vy = (V[D[i].xInd, D[i].yInd]-v)/(D[i].θym*dy)
                end
            end
            du[i] = -a*Vx - b*Vy # Update RHS term
        end
    end
end


using Base.Iterators: repeated, take, enumerate
using ProtoGrad
using ProtoGrad: Linear, Compose, ReLU, SupervisedObjective, mse
using LinearAlgebra
using Test

@testset "Quadratic objective" begin
    Q = Diagonal([1e1, 1e-1, 1e0, 1e0, 1e0, 1e0, 1e0, 1e0, 1e0, 1e0])
    q = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]

    objective = x -> 0.5*dot(x, Q*x) + dot(x, q)

    x_star = -Q\q
    f_star = objective(x_star)
    L = opnorm(Q)

    w0 = zeros(10)

    @testset "$(name)" for (name, optimizer) in [
        "GradientDescent" => ProtoGrad.GradientDescent(stepsize=1/L),
    ]
        for (k, output) in enumerate(take(optimizer(w0, objective), 200))
            @test objective(output.solution) - f_star <= 2 * L * norm(x_star)^2 / (2 * (k - 1))
        end
    end

    @testset "$(name)" for (name, optimizer) in [
        "Nesterov" => ProtoGrad.Nesterov(stepsize=1/L),
    ]
        for (k, output) in enumerate(take(optimizer(w0, objective), 200))
            @test objective(output.solution) - f_star <= 2 * L * norm(x_star)^2 / k^2
        end
    end

    @testset "$(name)" for (name, optimizer) in [
        "Polyak" => ProtoGrad.Polyak(stepsize=1/L, momentum=0.5),
        "RMSProp" => ProtoGrad.RMSProp(stepsize=1/L, alpha=0.99, epsilon=1e-8),
        "Adam" => ProtoGrad.Adam(stepsize=1/L, beta1=0.9, beta2=0.999, epsilon=1e-8),
    ]
        max_iter = 1000

        iterations = take(optimizer(w0, objective), max_iter)
        solution = ProtoGrad.last(iterations).solution

        @test objective(solution) <= f_star + 0.01 * abs(f_star)
    end
end

@testset "Supervised objective" begin

    @testset "Linear model ($(T))" for T in [Float32, Float64]

        input_size, output_size = 10, 1
        batch_size = 5000

        W_true = randn(T, output_size, input_size)
        b_true = randn(T, output_size)

        x = randn(T, input_size, batch_size)
        y = W_true * x .+ b_true + randn(T, output_size, batch_size) / 10

        data_iter = Iterators.repeated((x, y))
        f = SupervisedObjective(mse, data_iter)
        Lf = 2 * opnorm(x * x') / batch_size
        stepsize = 1 / Lf

        m = Linear(randn(T, output_size, input_size), randn(T, output_size))

        @testset "Basic checks $(name)" for (name, optimizer) in [
            "GradientDescent" => ProtoGrad.GradientDescent(stepsize=stepsize),
            "Polyak" => ProtoGrad.Polyak(stepsize=stepsize, momentum=0.5),
            "Nesterov" => ProtoGrad.Nesterov(stepsize=stepsize),
            "BarzilaiBorwein" => ProtoGrad.BarzilaiBorwein(alpha=stepsize),
            "AdaGrad" => ProtoGrad.AdaGrad(stepsize=stepsize),
            "RMSProp" => ProtoGrad.RMSProp(stepsize=stepsize),
            "AdaDelta" => ProtoGrad.AdaDelta(),
            "Adam" => ProtoGrad.Adam(stepsize=stepsize),
        ]
            m0 = copy(m)
            seq = optimizer(m, f)

            @test Base.IteratorSize(typeof(seq)) == Base.IsInfinite()

            for output in Iterators.take(seq, 5)
                @test length(ProtoGrad.overlap(m, output.solution)) == 0
                @test typeof(output.solution) == typeof(m)
            end

            @test vec(m) == vec(m0)
        end

        @testset "Accuracy $(name)" for (name, optimizer) in [
            "GradientDescent" => ProtoGrad.GradientDescent(stepsize=stepsize),
            "Polyak" => ProtoGrad.Polyak(stepsize=stepsize, momentum=0.5),
            "Nesterov" => ProtoGrad.Nesterov(stepsize=stepsize),
        ]
            seq = optimizer(m, f)

            for (it, output) in enumerate(seq)
                if it >= 1000
                    @test isapprox(output.solution.W, W_true, rtol=5e-2)
                    # @test isapprox(output.solution.b, b_true, rtol=1e-1)
                    break
                end
            end
        end
    end
end