##############################################################################
##
## simulated_annealing()
##
## Arguments:
## * cost: Function from states to the real numbers. Often called an energy
##         function, but this algorithm works for both positive and negative
##         costs.
## * s0: The initial state of the system.
## * neighbor: Function from states to states. Produces what the Metropolis
##             algorithm would call a proposal.
## * temperature: Function specifying the temperature at time i.
## * iterations: How many iterations of the algorithm should be run? This is
##               the only termination condition.
## * keep_best: Do we return the best state visited or the last state visited?
##              (Should default to true.)
## * show_trace: Do we show a trace of the system's evolution?
##
##############################################################################

##############################################################################
##
## Off-the-shelf cooling schedules
##
##############################################################################

# Theoretically, SA will converge if the temperature decreases
# according inversely proportional to the current time.
#
# For this to work, an unknown constant must be used.
# In practice, we use 1 instead of this unknown constant.
log_temperature(t::Real) = 1.0 / log(t)

# If the temperature is held constant and the cost function is -log(p),
# SA reduces to the Metropolis algorithm for sampling from a distribution.
constant_temperature(t::Real) = 1.0

# Default neighbor takes Gaussian jumps
function default_neighbor!(x::Vector, storage::Vector)
    for i in 1:length(x)
        storage[i] = x[i] + randn()
    end
end

# NB: The neighbor function is mutating now
function simulated_annealing(cost::Function,
                             s0::Vector;
                             neighbor!::Function = default_neighbor!,
                             temperature::Function = log_temperature,
                             keep_best::Bool = true,
                             tolerance::Real = 1e-8,
                             iterations::Integer = 100_000,
                             store_trace::Bool = false,
                             show_trace::Bool = false)

    # Maintain a trace of the optimization algo's state
    tr = OptimizationTrace()

    # Set our current state to the specified intial state
    s = copy(s0)
    s_n = copy(s0)
    y = cost(s)

    # Track how many function calls we've made
    f_calls = 1

    # Set the best state we've seen to the intial state
    best_s = copy(s0)
    best_y = y

    # Record the number of iterations we perform
    iteration = 0

    # Update our trace information
    os = OptimizationState(copy(s), y, iteration)
    if store_trace
        push!(tr, os)
    end
    if show_trace
        println(os)
    end

    # We always perform a fixed number of iterations
    while iteration < iterations
        # Update the iteration counter
        iteration += 1

        # Call temperature to find the proper temperature at time i
        t = temperature(iteration)

        # Call neighbor to randomly generate a neighbor of our current state
        neighbor!(s, s_n)

        # Evaluate the cost function on our current state and its neighbor
        y_n = cost(s_n)
        f_calls += 1

        if y_n <= y
            # If the proposed new state is superior, we always move to it
            copy!(s, s_n)
            y = y_n

            # If the new state is the best state we have seen,
            #  keep a record of it
            if y_n < best_y
                best_y = y_n
                copy!(best_s, s_n)
            end
        else
            # If the proposed new state is inferior, we move to it with
            #  probability p
            p = exp(-(y_n - y)/t)
            if rand() <= p
                copy!(s, s_n)
                y = y_n
            end
        end

        # Print out the state of the system
        if store_trace || show_trace
            os = OptimizationState(copy(s), y, iteration)
            if store_trace
                push!(tr, os)
            end
            if show_trace
                println(os)
            end
        end
    end

    # If specified by the user, we return the best state we've seen
    # Otherwise, we return the last state we've seen
    if keep_best
        OptimizationResults("Simulated Annealing",
                            s0,
                            best_s,
                            best_y,
                            iterations,
                            false,
                            tr,
                            f_calls,
                            0)
    else
        OptimizationResults("Simulated Annealing",
                            s0,
                            s,
                            y,
                            iterations,
                            false,
                            tr,
                            f_calls,
                            0)
    end
end


using SymPy
using Test
using LinearAlgebra

@testset "Matrix" begin
    ## matrices
    x, y = @vars x y
    A = [x 1; 1 x]
    B = [x 1; 0 2x]
    v = [x, 2]

    ## These fail for older installations of SymPy
    @test simplify(det(A)) == x^2 - 1

    ## we use inverse for A[:inv]()

    # aliased to use inverse
    @test simplify.(inv(A) * A) ==  [1 0; 0 1]
    @test simplify.(A * inv(A)) == [1 0; 0 1]
    ##XXX    @test simplify.(A[:inv]() - inv(A)) == zeros(2, 2)

    @test SymPy.adjoint(B) == [adjoint(x) 0; 1 adjoint(2x)]
    @test SymPy.adjoint(B) == B'
    @test A.dual() == sympy.zeros(2, 2)

    A1 = Sym[25 15 -5; 15 18 0; -5 0 11]
    r = A1.cholesky()
    @test r*r.transpose() == A1

    #    s = LUsolve(A, v)
    s = A.LUsolve(B)
    @test simplify.(A * s) == B

    # norm
    @test norm(A) == sqrt(2 * abs(x)^2 + 2)
    # test norm for different subtypes of AbstractArray
    ## XXX @test norm(A) == norm(Symmetric(A)) LinearAlgebra.Symmetric no long works
    @test norm(A) == norm(view(A, :, :))

    # abs
    @test  all(convert.(Bool, abs.(A) .≧ 0))
    @test  abs.(A) == abs.(view(A, :, :))

    # is_lower, is_square, is_symmetric much slower than julia only counterparts. May deprecate, but for now they are here
    @test A.is_lower == istril(A)
    @test A.is_square == true
    @test A.is_symmetric() != issymmetric(A)

    @vars x real=true
    A = [x 1; 1 x]
    @test A.is_symmetric() == issymmetric(A)

    @test Set(eigvals(A)) == Set([x-1, x+1])

    # issue with transpose being non-typestable
    @test eltype(transpose(A)) == Sym
    @test eltype(Symmetric(A)) == Sym

    #numerical tests
    M = Sym[1 0 0; 0 1 0; 0 0 x]
    evecs = eigvecs(M)
    @test evecs[:,1] == [1, 0, 0]

    A = Sym[1 2 3; 3 6 2; 2 0 1]
    q, r = A.QRdecomposition()
    @test q * r == A
    @test abs(det(q)) == 1

    # for v0.4, the vector type is not correctly inferred
    #L = Vector{Sym}[Sym[2,3,5], Sym[3,6,2], Sym[8,3,6]]
    # L = collect(sympy.Matrix.(([[2,3,5]], [[3,6,2]], [[8,3,6]]]))
    L = collect(sympy.Matrix.([[2,3,5], [3,6,2], [8,3,6]]))
    out = sympy.GramSchmidt(L, true)  # qualify, as L not SymbolicObject type
    for i = 1:3, j = i:3
        @test out[i].dot(out[j]) == (i == j ? 1 : 0)
    end

    A = Sym[4 3; 6 3]
    L, U, _ = A.LUdecomposition()
    @test L == Sym[1 0; 3//2 1]

    A = Sym[1 0; 0 1] * 2
    B = Sym[1 2; 3 4]
    @test A.diagonal_solve(B) == B/2

    M = Sym[1 2 0; 0 3 0; 2 -4 2]
    P, D = M.diagonalize()
    @test D == [1 0 0; 0 2 0; 0 0 3]
    @test P == [-1 0 -1; 0 0 -1; 2 1  2]
    @test D == inv(P) * M * P

    # test SymPy's expm against Julia's expm
    A = [1 2 0; 0 3 0; 2 -4 2]
    M = Sym.(A)
    ## no exp(M)!
    U = M.exp() - exp(A)
    @test maximum(abs.(N.(U))) <= 1e-12

    M = [x y; 1 0]
    @test integrate.(M, x) == [x^2/2 x*y; x 0]
    @test integrate.(M, Ref((x, 0, 2))) == [2 2y; 2 0]

    M = Sym[1 3 0; -2 -6 0; 3 9 6]
    @test M.nullspace()[1] ==  reshape(Sym[-3, 1, 0], 3, 1)

    M = Sym[1 2 0; 0 3 0; 2 -4 2]
    # M.cofactor  uses 0-based indexing!
    i, j = 1, 2
    @test M.cofactor(i, j) == (-1)^(i+j) * det(M[setdiff(1:3, i+1), setdiff(1:3, j+1)])
    @test M.adjugate() / M.det() == M.inv()

    M = Sym[ 6  5 -2 -3;
            -3 -1  3  3;
             2  1 -2 -3;
            -1  1  5  5]

    P, J = M.jordan_form()
    @test J == [2 1 0 0;
                0 2 0 0;
                0 0 2 1;
                0 0 0 2]
    @test J == inv(P) * M * P

    ρ, ϕ = symbols("rho, phi")
    X = [ρ*cos(ϕ), ρ*sin(ϕ), ρ^2]
    Y = [ρ, ϕ]
    @test X.jacobian(Y) ==   [cos(ϕ) -ρ*sin(ϕ);
                              sin(ϕ)  ρ*cos(ϕ);
                                  2ρ       0]
    X = [ρ*cos(ϕ), ρ*sin(ϕ)]
    @test convert(Matrix{Sym}, X.jacobian(Y)) == [cos(ϕ) -ρ*sin(ϕ);
                                                  sin(ϕ)  ρ*cos(ϕ)]
    @test X.jacobian(Y) == view(X, :, :).jacobian(view(Y, :, :))

    ## Issue   #359
    if VERSION >= v"1.2"
        @vars a
        @test eltype([a 1; 1 a] + I) == Sym
        @test eltype([a 1; 1 a] + 2I) == Sym
    end

    ## Issue #397 adjoint losing type
    A = ones(Sym, 1, 1)
    @test eltype(A * A') == Sym

    ## Issue 413 wth matrix exponential; SymMatrix multiplication
    @vars a
    A,A1 = [1 0; 0 a], [1 1; a 2]
    @test exp(A) == [exp(Sym(1)) 0; 0 exp(a)]
    B,B1 = convert(SymMatrix,A), convert(SymMatrix,A1)
    @test B*B1 == convert(SymMatrix, A*A1)
end


## Julia 1.1.0

println("\n\n---------------- general ----------------")
# common
function ×(x, y)
    x * y
end

# simple
∑(x, y) = x + y

x = 10
y = 2
for F in [∑, ×]
    println("$(F)($(x), $(y)) -> $(F(x, y))")
end

# hypot
function hypot(x,y)
    x = abs(x)
    y = abs(y)

    if x > y
        r = y / x
        return x * sqrt(1 + r * r)
    end

    if y == 0
        return zero(x)
    end

    r = x / y
    return y * sqrt(1 + r * r)
end

println("hypot(3, 4) -> $(hypot(3, 4))")

# anonymous
println("\n\n---------------- anonymous ----------------")
f(r) = pi * r^2
g = f
for F in [f, g, x -> x^2 + 2x + 1, x -> x^3 - x^2 - 1]
    println("$(F)($(x)) -> $(F(x))")
end

x = 2
println("(x -> x^2 + 2x + 1)($x) -> $((x -> x^2 + 2x + 1)(x))")

# operators are functions
println("\n\n---------------- operators are functions ----------------")
add = +
x = (1, 2, 3)
for F in [add, +]
    println("$(F)($(x)) -> $(F(x...))")
end

# optional
println("\n\n---------------- optional ----------------")
function opsum(a::Int64, b::Int64=1, c::Int64=1)
	a + b + c
end

for V in [(1), (1, 2), (1, 2, 3)]
    println("opsum($V) -> $(opsum(V...))")
end

# dot
println("\n\n---------------- dot ----------------")

X = [1, 2, 3]
Y = [4, 5, 6]
simple(x, y) = x + 2y
println("simple.(1, $(X)) -> $(simple.(1, X))")
println("simple.($(X), $(Y)) -> $(simple.(X, Y))")


using Test

using MathOptInterface
const MOI = MathOptInterface
const MOIT = MathOptInterface.Test
const MOIU = MathOptInterface.Utilities
const MOIB = MathOptInterface.Bridges

include("../utilities.jl")

mock = MOIU.MockOptimizer(MOIU.UniversalFallback(MOIU.Model{Float64}()))
config = MOIT.TestConfig()

@testset "GeoMeantoRelEntr" begin
    bridged_mock = MOIB.Constraint.GeoMeantoRelEntr{Float64}(mock)

    MOIT.basic_constraint_tests(bridged_mock, config,
        include = [(F, MOI.GeometricMeanCone) for F in [
            MOI.VectorOfVariables, MOI.VectorAffineFunction{Float64}, MOI.VectorQuadraticFunction{Float64}
        ]])

    @testset "geomean1test" begin
        var_primal = [1, 1, 1, 1, 0]
        relentr_dual = [1, 1, 1, 1, -1, -1, -1] / 3
        mock.optimize! = (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(mock, var_primal,
            (MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64}) => [-inv(3)],
            (MOI.VectorOfVariables, MOI.Nonnegatives) => [[1]],
            (MOI.VectorAffineFunction{Float64}, MOI.RelativeEntropyCone) => [relentr_dual])

        MOIT.geomean1ftest(bridged_mock, config)

        var_names = ["t", "x", "y", "z"]
        MOI.set(bridged_mock, MOI.VariableName(), MOI.get(bridged_mock, MOI.ListOfVariableIndices()), var_names)

        nonneg = MOI.get(mock, MOI.ListOfConstraintIndices{MOI.VectorOfVariables, MOI.Nonnegatives}())
        relentr = MOI.get(mock, MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{Float64}, MOI.RelativeEntropyCone}())
        aux = MOI.get(mock, MOI.ListOfVariableIndices())[end]

        @testset "Test mock model" begin
            MOI.set(mock, MOI.VariableName(), aux, "aux")
            @test length(nonneg) == 1
            MOI.set(mock, MOI.ConstraintName(), nonneg[1], "nonneg")
            @test length(relentr) == 1
            MOI.set(mock, MOI.ConstraintName(), relentr[1], "relentr")
            less = MOI.get(mock, MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64}}())
            @test length(less) == 1
            MOI.set(mock, MOI.ConstraintName(), less[1], "less")

            s = """
            variables: t, x, y, z, aux
            less: x + y + z in MathOptInterface.LessThan(3.0)
            nonneg: [aux] in MathOptInterface.Nonnegatives(1)
            relentr: [0.0, x, y, z, t + aux, t + aux, t + aux] in MathOptInterface.RelativeEntropyCone(7)
            maxobjective: t
            """
            model = MOIU.Model{Float64}()
            MOIU.loadfromstring!(model, s)
            MOIU.test_models_equal(mock, model, vcat(var_names, "aux"), ["less", "nonneg", "relentr"])
        end

        @testset "Test bridged model" begin
            geomean = MOI.get(bridged_mock, MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{Float64}, MOI.GeometricMeanCone}())
            @test length(geomean) == 1
            MOI.set(bridged_mock, MOI.ConstraintName(), geomean[1], "geomean")
            less = MOI.get(bridged_mock, MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64}}())
            @test length(less) == 1
            MOI.set(bridged_mock, MOI.ConstraintName(), less[1], "less")

            s = """
            variables: t, x, y, z
            less: x + y + z in MathOptInterface.LessThan(3.0)
            geomean: [1.0t, x, y, z] in MathOptInterface.GeometricMeanCone(4)
            maxobjective: t
            """
            model = MOIU.Model{Float64}()
            MOIU.loadfromstring!(model, s)
            MOIU.test_models_equal(bridged_mock, model, var_names, ["less", "geomean"])
        end

        ci = first(MOI.get(bridged_mock, MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{Float64}, MOI.GeometricMeanCone}()))

        @testset "$attr" for attr in [MOI.ConstraintPrimalStart(), MOI.ConstraintDualStart()]
            @test MOI.supports(bridged_mock, attr, typeof(ci))
            value = (attr isa MOI.ConstraintPrimalStart) ? ones(4) : vcat(-1, fill(inv(3), 3))
            MOI.set(bridged_mock, attr, ci, value)
            @test MOI.get(bridged_mock, attr, ci) ≈ value
            if attr isa MOI.ConstraintPrimalStart
                @test MOI.get(mock, MOI.VariablePrimalStart(), aux) == 0
                @test MOI.get(mock, attr, nonneg[1]) == [0]
                @test MOI.get(mock, attr, relentr[1]) == [0, 1, 1, 1, 1, 1, 1]
            else
                @test MOI.get(mock, attr, nonneg[1]) == [1]
                @test MOI.get(mock, attr, relentr[1]) ≈ relentr_dual
            end
        end

        test_delete_bridge(bridged_mock, ci, 4, (
            (MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64}, 1),
            (MOI.VectorOfVariables, MOI.Nonnegatives, 0),
            (MOI.VectorAffineFunction{Float64}, MOI.RelativeEntropyCone, 0)))
    end

    @testset "geomean2test" begin
        var_primal = vcat(ones(10), 0)
        relentr_dual = vcat(fill(inv(9), 10), fill(-inv(9), 9))
        mock.optimize! = (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(mock, var_primal,
            (MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64}) => fill(-inv(9), 9),
            (MOI.VectorOfVariables, MOI.Nonnegatives) => [[1]],
            (MOI.VectorAffineFunction{Float64}, MOI.RelativeEntropyCone) => [relentr_dual])

        MOIT.geomean2ftest(bridged_mock, config)

        var_names = vcat("t", ["x$(i)" for i in 1:9])
        MOI.set(bridged_mock, MOI.VariableName(), MOI.get(bridged_mock, MOI.ListOfVariableIndices()), var_names)

        nonneg = MOI.get(mock, MOI.ListOfConstraintIndices{MOI.VectorOfVariables, MOI.Nonnegatives}())
        relentr = MOI.get(mock, MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{Float64}, MOI.RelativeEntropyCone}())
        aux = MOI.get(mock, MOI.ListOfVariableIndices())[end]

        @testset "Test mock model" begin
            MOI.set(mock, MOI.VariableName(), aux, "aux")
            @test length(nonneg) == 1
            MOI.set(mock, MOI.ConstraintName(), nonneg[1], "nonneg")
            @test length(relentr) == 1
            MOI.set(mock, MOI.ConstraintName(), relentr[1], "relentr")
            equalto = MOI.get(bridged_mock, MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64}}())
            @test length(equalto) == 9
            equalto_names = ["equalto$(i)" for i in 1:9]
            MOI.set.(bridged_mock, MOI.ConstraintName(), equalto, equalto_names)

            s = """
            variables: t, x1, x2, x3, x4, x5, x6, x7, x8, x9, aux
            equalto1: 1.0x1 in MathOptInterface.EqualTo(1.0)
            equalto2: 1.0x2 in MathOptInterface.EqualTo(1.0)
            equalto3: 1.0x3 in MathOptInterface.EqualTo(1.0)
            equalto4: 1.0x4 in MathOptInterface.EqualTo(1.0)
            equalto5: 1.0x5 in MathOptInterface.EqualTo(1.0)
            equalto6: 1.0x6 in MathOptInterface.EqualTo(1.0)
            equalto7: 1.0x7 in MathOptInterface.EqualTo(1.0)
            equalto8: 1.0x8 in MathOptInterface.EqualTo(1.0)
            equalto9: 1.0x9 in MathOptInterface.EqualTo(1.0)
            nonneg: [aux] in MathOptInterface.Nonnegatives(1)
            relentr: [0.0, x1, x2, x3, x4, x5, x6, x7, x8, x9, t + aux, t + aux, t + aux, t + aux, t + aux, t + aux, t + aux, t + aux, t + aux] in MathOptInterface.RelativeEntropyCone(19)
            maxobjective: t
            """
            model = MOIU.Model{Float64}()
            MOIU.loadfromstring!(model, s)
            MOIU.test_models_equal(mock, model, vcat(var_names, "aux"), vcat(equalto_names, "nonneg", "relentr"))
        end

        @testset "Test bridged model" begin
            geomean = MOI.get(bridged_mock, MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{Float64}, MOI.GeometricMeanCone}())
            @test length(geomean) == 1
            MOI.set(bridged_mock, MOI.ConstraintName(), geomean[1], "geomean")
            equalto = MOI.get(bridged_mock, MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64}}())
            @test length(equalto) == 9
            equalto_names = ["equalto$(i)" for i in 1:9]
            MOI.set.(bridged_mock, MOI.ConstraintName(), equalto, equalto_names)

            s = """
            variables: t, x1, x2, x3, x4, x5, x6, x7, x8, x9
            equalto1: 1.0x1 in MathOptInterface.EqualTo(1.0)
            equalto2: 1.0x2 in MathOptInterface.EqualTo(1.0)
            equalto3: 1.0x3 in MathOptInterface.EqualTo(1.0)
            equalto4: 1.0x4 in MathOptInterface.EqualTo(1.0)
            equalto5: 1.0x5 in MathOptInterface.EqualTo(1.0)
            equalto6: 1.0x6 in MathOptInterface.EqualTo(1.0)
            equalto7: 1.0x7 in MathOptInterface.EqualTo(1.0)
            equalto8: 1.0x8 in MathOptInterface.EqualTo(1.0)
            equalto9: 1.0x9 in MathOptInterface.EqualTo(1.0)
            geomean: [1.0t, x1, x2, x3, x4, x5, x6, x7, x8, x9] in MathOptInterface.GeometricMeanCone(10)
            maxobjective: t
            """
            model = MOIU.Model{Float64}()
            MOIU.loadfromstring!(model, s)
            MOIU.test_models_equal(bridged_mock, model, var_names, vcat(equalto_names, "geomean"))
        end

        ci = first(MOI.get(bridged_mock, MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{Float64}, MOI.GeometricMeanCone}()))

        @testset "$attr" for attr in [MOI.ConstraintPrimalStart(), MOI.ConstraintDualStart()]
            @test MOI.supports(bridged_mock, attr, typeof(ci))
            value = (attr isa MOI.ConstraintPrimalStart) ? ones(10) : vcat(-1, fill(inv(9), 9))
            MOI.set(bridged_mock, attr, ci, value)
            @test MOI.get(bridged_mock, attr, ci) ≈ value
            if attr isa MOI.ConstraintPrimalStart
                @test MOI.get(mock, MOI.VariablePrimalStart(), aux) == 0
                @test MOI.get(mock, attr, nonneg[1]) == [0]
                @test MOI.get(mock, attr, relentr[1]) == vcat(0, ones(18))
            else
                @test MOI.get(mock, attr, nonneg[1]) == [1]
                @test MOI.get(mock, attr, relentr[1]) ≈ relentr_dual
            end
        end

        test_delete_bridge(bridged_mock, ci, 10, (
            (MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64}, 9),
            (MOI.VectorOfVariables, MOI.Nonnegatives, 0),
            (MOI.VectorAffineFunction{Float64}, MOI.RelativeEntropyCone, 0)))
    end

    @testset "geomean3test" begin
        var_primal = [2, 2, 0]
        relentr_dual = [2, 2, -2]
        mock.optimize! = (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(mock, var_primal,
            (MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64}) => [-2],
            (MOI.VectorOfVariables, MOI.Nonnegatives) => [[2]],
            (MOI.VectorAffineFunction{Float64}, MOI.RelativeEntropyCone) => [relentr_dual])

        MOIT.geomean3ftest(bridged_mock, config)

        var_names = ["t", "x"]
        MOI.set(bridged_mock, MOI.VariableName(), MOI.get(bridged_mock, MOI.ListOfVariableIndices()), var_names)

        nonneg = MOI.get(mock, MOI.ListOfConstraintIndices{MOI.VectorOfVariables, MOI.Nonnegatives}())
        relentr = MOI.get(mock, MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{Float64}, MOI.RelativeEntropyCone}())
        aux = MOI.get(mock, MOI.ListOfVariableIndices())[end]

        @testset "Test mock model" begin
            MOI.set(mock, MOI.VariableName(), aux, "aux")
            @test length(nonneg) == 1
            MOI.set(mock, MOI.ConstraintName(), nonneg[1], "nonneg")
            @test length(relentr) == 1
            MOI.set(mock, MOI.ConstraintName(), relentr[1], "relentr")
            less = MOI.get(mock, MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64}}())
            @test length(less) == 1
            MOI.set(mock, MOI.ConstraintName(), less[1], "less")

            s = """
            variables: t, x, aux
            less: 1.0x in MathOptInterface.LessThan(2.0)
            nonneg: [aux] in MathOptInterface.Nonnegatives(1)
            relentr: [0.0, x, t + aux] in MathOptInterface.RelativeEntropyCone(3)
            maxobjective: 2.0t
            """
            model = MOIU.Model{Float64}()
            MOIU.loadfromstring!(model, s)
            MOIU.test_models_equal(mock, model, vcat(var_names, "aux"), ["less", "nonneg", "relentr"])
        end

        @testset "Test bridged model" begin
            geomean = MOI.get(bridged_mock, MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{Float64}, MOI.GeometricMeanCone}())
            @test length(geomean) == 1
            MOI.set(bridged_mock, MOI.ConstraintName(), geomean[1], "geomean")
            less = MOI.get(bridged_mock, MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64}}())
            @test length(less) == 1
            MOI.set(bridged_mock, MOI.ConstraintName(), less[1], "less")

            s = """
            variables: t, x
            less: 1.0x in MathOptInterface.LessThan(2.0)
            geomean: [1.0t, x] in MathOptInterface.GeometricMeanCone(2)
            maxobjective: 2.0t
            """
            model = MOIU.Model{Float64}()
            MOIU.loadfromstring!(model, s)
            MOIU.test_models_equal(bridged_mock, model, var_names, ["less", "geomean"])
        end

        ci = first(MOI.get(bridged_mock, MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{Float64}, MOI.GeometricMeanCone}()))

        @testset "$attr" for attr in [MOI.ConstraintPrimalStart(), MOI.ConstraintDualStart()]
            @test MOI.supports(bridged_mock, attr, typeof(ci))
            value = (attr isa MOI.ConstraintPrimalStart) ? [2, 2] : [-2, 2]
            MOI.set(bridged_mock, attr, ci, value)
            @test MOI.get(bridged_mock, attr, ci) ≈ value
            if attr isa MOI.ConstraintPrimalStart
                @test MOI.get(mock, MOI.VariablePrimalStart(), aux) == 0
                @test MOI.get(mock, attr, nonneg[1]) == [0]
                @test MOI.get(mock, attr, relentr[1]) == [0, 2, 2]
            else
                @test MOI.get(mock, attr, nonneg[1]) == [2]
                @test MOI.get(mock, attr, relentr[1]) ≈ relentr_dual
            end
        end

        test_delete_bridge(bridged_mock, ci, 2, (
            (MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64}, 1),
            (MOI.VectorOfVariables, MOI.Nonnegatives, 0),
            (MOI.VectorAffineFunction{Float64}, MOI.RelativeEntropyCone, 0)))
    end
end


mutable struct Kahan{T <: AbstractFloat}
	sum::T
	c::T
end

Kahan(::Type{T}) where {T <: AbstractFloat} = Kahan{T}(convert(T, 0), convert(T, 0))

#import NumericExtensions.add!
function add!(k::Kahan{T}, x::T) where {T <: AbstractFloat}
	y = x - k.c
	t = k.sum + y
	k.c = (t - k.sum) - y
	k.sum = t
end

import Base.sum
sum(k::Kahan) = k.sum

mutable struct MeanCounter{T <: AbstractFloat}
	n::Int64
	mean::Kahan{T}
end

MeanCounter(::Type{T}) where {T <: AbstractFloat} = MeanCounter{T}(0, Kahan(T))

function add!(m::MeanCounter{T}, x::T) where {T <: AbstractFloat}
	m.n += 1
	add!(m.mean, (x - sum(m.mean)) / m.n)
	return m.mean
end

mean(m::MeanCounter) = sum(m.mean)

export Kahan, MeanCounter, add!, sum, mean


using RayTracer
using Test

@testset "tuple_creation" begin
    @testset "w_1_is_point" begin
        a = RayTracer.tuple(4.3,-4.2,3.1,1)
        @test a.x ≈ 4.3
        @test a.y ≈ -4.2
        @test a.z ≈ 3.1
        @test a.w == 1
        @test RayTracer.is_point(a) == true
        @test RayTracer.is_vector(a) == false
    end

    @testset "w_0_is_vector" begin
        a = RayTracer.tuple(4.3,-4.2,3.1,0)
        @test a.x ≈ 4.3
        @test a.y ≈ -4.2
        @test a.z ≈ 3.1
        @test a.w == 0
        @test RayTracer.is_point(a) == false
        @test RayTracer.is_vector(a) == true
    end

    @testset "point_creates_tuple" begin
        p = RayTracer.point(4,-4,3)
        @test p == RayTracer.tuple(4,-4,3,1)
    end

    @testset "vector_creates_tuple" begin
        v = RayTracer.vector(4,-4,3)
        @test v == RayTracer.tuple(4,-4,3,0)
    end
end

@testset "tuple_operations" begin
    @testset "adding_tuples" begin
        a1 = RayTracer.tuple(3,-2,5,1)
        a2 = RayTracer.tuple(-2,3,1,0)
        @test a1 + a2 == RayTracer.tuple(1,1,6,1)
    end

    @testset "subtracting_two_points" begin
        p1 = RayTracer.point(3,2,1)
        p2 = RayTracer.point(5,6,7)
        @test p1 - p2 == RayTracer.vector(-2,-4,-6)
    end

    @testset "subtracting_vector_from_point" begin
        p = RayTracer.point(3,2,1)
        v = RayTracer.vector(5,6,7)
        @test p - v == RayTracer.point(-2,-4,-6)
    end

    @testset "subtracting_two_vectors" begin
        v1 = RayTracer.vector(3,2,1)
        v2 = RayTracer.vector(5,6,7)
        @test v1 - v2 == RayTracer.vector(-2,-4,-6)
    end

    @testset "subtracting_vector_from_zero_vector" begin
        zero = RayTracer.vector(0,0,0)
        v = RayTracer.vector(1,-2,3)
        @test zero - v == RayTracer.vector(-1,2,-3)
    end

    @testset "negating_tuple" begin
        a = RayTracer.tuple(1,-2,3,-4)
        @test -a == RayTracer.tuple(-1,2,-3,4)
    end

    @testset "multiplying_tuple_by_scalar" begin
        a = RayTracer.tuple(1,-2,3,-4)
        @test a * 3.5 ≈ RayTracer.tuple(3.5,-7,10.5,-4)
        @test a * 0.5 ≈ RayTracer.tuple(0.5,-1,1.5,-4)
    end

    @testset "dividing_tuple_by_scalar" begin
        a = RayTracer.tuple(1,-2,3,-4)
        @test a / 2 ≈ RayTracer.tuple(0.5,-1,1.5,-4)
    end

    @testset "modulus_unit_vectors" begin
        v1 = RayTracer.vector(1,0,0)
        v2 = RayTracer.vector(0,1,0)
        v3 = RayTracer.vector(0,0,1)
        v4 = RayTracer.vector(1,2,3)
        v5 = RayTracer.vector(-1,-2,-3)
        @test RayTracer.modulus(v1) ≈ 1
        @test RayTracer.modulus(v2) ≈ 1
        @test RayTracer.modulus(v3) ≈ 1
        @test RayTracer.modulus(v4) ≈ √14
        @test RayTracer.modulus(v5) ≈ √14
    end

    @testset "normalization_vectors" begin
        v1 = RayTracer.vector(4,0,0)
        v2 = RayTracer.vector(1,2,3)
        @test RayTracer.normalize(v1) ≈ RayTracer.vector(1,0,0)
        @test RayTracer.normalize(v2) ≈ RayTracer.vector(1/√14,2/√14,3/√14)
    end

    @testset "modulus_normalized_vector" begin
        v = RayTracer.vector(1,2,3)
        norm = RayTracer.normalize(v)
        @test RayTracer.modulus(norm) ≈ 1
    end

    @testset "dot_product" begin
        a = RayTracer.vector(1,2,3)
        b = RayTracer.vector(2,3,4)
        @test RayTracer.dot(a,b) ≈ 20
    end

    @testset "cross_product" begin
        a = RayTracer.vector(1,2,3)
        b = RayTracer.vector(2,3,4)
        @test RayTracer.cross(a,b) ≈ RayTracer.vector(-1,2,-1)
        @test RayTracer.cross(b,a) ≈ RayTracer.vector(1,-2,1)
    end
end

@testset "colors" begin
    @testset "clors_are_tuples" begin
        c = RayTracer.color(-0.5,0.4,1.7)
        @test c.red ≈ -0.5
        @test c.green ≈ 0.4
        @test c.blue ≈ 1.7
    end
    @testset "adding colors" begin
        c1 = RayTracer.color(0.9,0.6,0.75)
        c2 = RayTracer.color(0.7,0.1,0.25)
        @test c1 + c2 ≈ RayTracer.color(1.6,0.7,1.0)
    end
    @testset "subtracting_colors" begin
        c1 = RayTracer.color(0.9,0.6,0.75)
        c2 = RayTracer.color(0.7,0.1,0.25)
        @test c1 - c2 ≈ RayTracer.color(0.2,0.5,0.5)
    end
    @testset "multiplying_color_by_scalar" begin
        c = RayTracer.color(0.2,0.3,0.4)
        @test c * 2 ≈ RayTracer.color(0.4,0.6,0.8)
        @test 2 * c ≈ RayTracer.color(0.4,0.6,0.8)
    end
    @testset "multiplying_colors" begin
        c1 = RayTracer.color(1,0.2,0.4)
        c2 = RayTracer.color(0.9,1,0.1)
        @test c1 * c2 ≈ RayTracer.color(0.9,0.2,0.04)
    end
end

@testset "reflection_at_45_degrees" begin
    v = RayTracer.vector(1,-1,0)
    n = RayTracer.vector(0,1,0)
    r = RayTracer.reflect(v,n)
    @test r ≈ RayTracer.vector(1,1,0)
end

@testset "reflection_slanted_surface" begin
    v = RayTracer.vector(0,-1,0)
    n = RayTracer.vector(√2/2,√2/2,0)
    r = RayTracer.reflect(v,n)
    @test ≈(r,RayTracer.vector(1,0,0),1e-10)
end


module StencilTest1

using HPAT

#HPAT.CaptureAPI.set_debug_level(3)
#HPAT.DomainPass.set_debug_level(3)
#HPAT.DistributedPass.set_debug_level(3)
#ParallelAccelerator.set_debug_level(3)
#ParallelAccelerator.CGen.set_debug_level(3)
#using CompilerTools
#CompilerTools.OptFramework.set_debug_level(3)

@acc hpat function test1(file_name)
    t1 = DataSource(DataTable{:userid=Int64, :val2=Float64}, HDF5, file_name)
    avg = stencil(x->(x[-1]+2*x[0]+x[1])/4.0, t1[:val2])
    return avg
end

using HDF5
using MPI
using Base.Test

function main()

    if MPI.Comm_rank(MPI.COMM_WORLD)==0
        h5write("test1_1.hdf5", "/userid", [1, 2, 3, 1, 2])
        h5write("test1_1.hdf5", "/val2", [1.2, 2.1, 3.1, 3.2, 1.9])
    end
    a = test1("test1_1.hdf5")
    println(a)

    @test_approx_eq a [1.2, 2.125, 2.875, 2.85, 1.9]

    if MPI.Comm_rank(MPI.COMM_WORLD)==0
        rm("test1_1.hdf5")
    end

end

end

StencilTest1.main()


abstract type AbstractBitstringProblem <: AbstractProblem end

import Base: length

function fitness(x::Array{Bool,1}, problem)
    fitness(BitArray(x), problem)
end

function fitness(
    x::BitArray,
    problem::T,
    fitnessvalue::Number,
    flipsposition,
) where {T<:AbstractProblem}
    y = copy(x)
    for item ∈ flipsposition
        flip!(y, item)
    end
    fitness(y, problem)
end
"""
    struct OneMax <: AbstractBitstringProblem
        n::Integer
    end

Maximization of the number of `true` in a `BitArray`. The mathematical definition of this problem is
maximizing the number of one-bits in a bitstring, which can be seen in the following formula:

```math
OneMax(x_1,\\dots,x_n)\\coloneqq\\sum_{i=1}^{n}x_i.
```

## Examples
```jldoctest; setup = :(using Pidoh)
julia> problem = OneMax(4);

julia> fitness(BitArray([true,true,false,true]), problem)
3

julia> optimum(problem)
4

```
"""
struct OneMax <: AbstractBitstringProblem
    n::Integer
end
fitness(x::BitArray, problem::OneMax) = sum(x)
optimum(problem::OneMax) = problem.n
length(problem::OneMax) = problem.n

"""
    struct LeadingOnes <: AbstractBitstringProblem
        n::Integer
    end

Maximization of the length of a prefix of `BitArray`  that consits only of `true`. More formally

```math
LeadingOnes(x_1,\\dots,x_n)\\coloneqq\\sum_{i=1}^{n}\\prod_{j=1}^ix_j.
```

## Examples
```jldoctest; setup = :(using Pidoh)
julia> problem = LeadingOnes(4);

julia> fitness(BitArray([true,true,false,true]), problem)
2

julia> optimum(problem)
4

```
"""
struct LeadingOnes <: AbstractBitstringProblem
    n::Integer
end

function fitness(x::BitArray, problem::LeadingOnes)
    lo = 0
    while lo < problem.n && x[lo+1]
        lo += 1
    end
    lo
end

optimum(problem::LeadingOnes) = problem.n
length(problem::LeadingOnes) = problem.n

"""
    struct CliffTwoParameters <: AbstractBitstringProblem
        n::Integer
    end

## Examples
```jldoctest; setup = :(using Pidoh)
julia> problem = CliffTwoParameters(10, 4, 2);

```
"""
struct CliffTwoParameters <: AbstractBitstringProblem
    n::Integer
    m::Integer
    d::Real
end

function fitness(x::BitArray, problem::CliffTwoParameters)
    onemax = sum(x)
    n = problem.n
    m = problem.m
    d = problem.d
    if onemax ≤ n - m
        return onemax * 1.0
    else
        return onemax*1.0 - d - 1
    end
end

optimum(problem::CliffTwoParameters) = problem.n - problem.d - 1
length(problem::CliffTwoParameters) = problem.n

struct ZeroMax <: AbstractBitstringProblem
    n::Integer
end

fitness(x::BitArray, problem::ZeroMax) = problem.n - sum(x)
optimum(problem::ZeroMax) = 0

struct Jump <: AbstractBitstringProblem
    n::Integer
    jumpsize::Integer
end

function fitness(x::BitArray, problem::Jump)
    onemax = sum(x)
    n = problem.n
    k = problem.jumpsize
    jump = -1
    if ((onemax ≤ n - k) || (onemax == n))
        jump = k + onemax
    else
        jump = n - onemax
    end
    jump
end

optimum(problem::Jump) = problem.n + problem.jumpsize
length(problem::Jump) = problem.n

struct GeneralJump <: AbstractBitstringProblem
    n::Integer
    k::Integer
    δ::Integer
end

function fitness(x::BitArray, problem::GeneralJump)
    onemax = sum(x)

    if onemax ≤ problem.n - problem.k || onemax ≥ problem.n-problem.k+problem.δ
        return onemax
    else
        return -onemax
    end
end

optimum(problem::GeneralJump) = problem.n
length(problem::GeneralJump) = problem.n

# function fitness(x::BitArray, problem::OneMax, fitnessvalue:: Number, flipsposition)
#     if length(flipsposition) == 0
#         return fitnessvalue
#     end
#     fit = fitnessvalue
#     for item ∈ flipsposition
#         @inbounds if x[item]
#             fit -= 1
#         else
#             fit += 1
#         end
#     end
#     fit
# end

# function fitness(x::BitArray, problem::Jump, fitnessvalue:: Number, flipsposition)
#     if length(flipsposition) == 0
#         return fitnessvalue
#     end
#     n=problem.n
#     k = problem.jumpsize
#
#     onemax = fitnessvalue
#
#     if fitnessvalue < k
#         onemax = n - fitnessvalue
#     else
#         onemax = fitnessvalue - k
#     end
#
#     if ((onemax ≤ n - k) || (onemax == n))
#         jump =k + onemax
#     else
#         jump = n - onemax
#     end
#
#     for item ∈ flipsposition
#         @inbounds if x[item]
#             onemax -= 1
#         else
#             onemax += 1
#         end
#     end
#
#     jump = -1
#     if ((onemax ≤ n - k) || (onemax == n))
#         jump =k + onemax
#     else
#         jump = n - onemax
#     end
#     jump
# end


module Trapz
    export trapz, @trapz

    include("kernels.jl")

    function trapz(x, y)
        return integrate(x,y)
    end

    function trapz(x, y, val_i::Val{i}) where i
        return integrate(x,y,val_i)
    end

    """
# Syntax for integration over single axis
        trapz(x,y,axis=Val(End))
        Calculates ∫y[..., i (axis) ,...] dx[i]
        for performance reasons x must have the same float type as y.
        1-D Example:
```julia
    vx=0:0.01:1
    vy=(x->x^2).(vx)
    trapz(vx,vy)
```
            Result ≈ 1/3
# Multidimensional Syntax
        2-D Example:
```julia
    vx=0:0.01:1
    vy=0:0.01:2
    z=[x^2+y^2 for x=vx, y=vy]
    trapz((vx,vy),z) # equivalent to trapz(vx, trapz(vy, z))
```
            Result ≈ 4/3
## Partial integration Syntax
        2-D Example: partial integration
```julia
    vx=0:0.01:1
    vy=0:0.01:2
    z=[x^2+y^2 for x=vx, y=vy]
    trapz((vx,:),z) # equivalent to trapz(vx, z, Val(1))
```
            Result ≈ 201-element Array{Float64,1}:
                     0.33335
                     0.3334500000000001
                     0.33375000000000005
                     0.3342499999999999
                     ⋮
                     4.253749999999999
                     4.293449999999999
                     4.333350000000002

        2-D Example: partial integration
```julia
    vx=0:0.01:1
    vy=0:0.01:2
    z=[x^2+y^2 for x=vx, y=vy]
    trapz((:,vy),z) # equivalent to trapz(vy, z, Val(2))
```
            Result ≈ 101-element Array{Float64,1}:
                     2.666700000000002
                     2.6669
                     2.6674999999999995
                     2.6685000000000008
                     ⋮
                     4.587500000000002
                     4.626900000000002
                     4.6667000000000005

    """
    trapz
end # module


module DSP

using Interpolation
using Conv
using Distributions
using DistributedArrays
using DSP # from julia
using SparseArrays
using LinearAlgebra
using FFTW

function chirp!(x; tgrid=nothing,
	fmin=nothing, # minimum fraction of Nyquist sampling
	fmax=nothing, # maximum fraction of Nyquist
	)

	fs=inv(step(tgrid))
	freqmin=(fmin===nothing) ? 0.0 : fmin*fs*0.5
	freqmax=(fmax===nothing) ? fs*0.5 : fmax*fs*0.5

	k=(freqmax-freqmin)*inv(tgrid[end]-tgrid[1])

	for it in 1:length(tgrid)
		t=(it-1)*step(tgrid)
		x[it] = sin(2. * π *(freqmin*t+k/2.0*t*t))
	end

end
"""
Random source signal models for Param.
Generate a band-limited
random signal of a particular maximum time, `tmax`.
Sinc function in the frequency domain corresponds to
a box car function is applied in the time domain.

This method has some extra allocations.
"""
function freq_rand!(x;
	tgrid=nothing,
	fmin=nothing, # minimum fraction of Nyquist sampling
	fmax=nothing, # maximum fraction of Nyquist
	nevent_fracs::Array{Float64}=[0.0],
	rng=Normal(),
	verbose=false,
	)

	nt=size(x,1)
	(tgrid===nothing) && (tgrid=range(0.0, stop=(nt-1)*1.0, step=1.0))  # create dummy tgrid
	fgrid=DSP.rfftfreq(length(tgrid),inv(step(tgrid))) # corresponding fgrid
	(nt≠length(tgrid)) && error("dimension")

	tmax=abs(tgrid[nt]-tgrid[1])

	xfreq=complex.(zeros(length(fgrid))) # allocation in freq domain
	xx=zeros(length(tgrid)) # allocation in time domain

	fs=inv(step(tgrid))
	freqmin=(fmin===nothing) ? 0.0 : fmin*fs*0.5
	freqmax=(fmax===nothing) ? fs*0.5 : fmax*fs*0.5
	(freqmax ≤ freqmin) && error("fmin and fmax")

	Δf = inv(tmax) # resolution in the frequency domain
	(Δf ≤ step(fgrid)) ? error("desired resolution smaller than grid sampling") :
	(Δf ≥ (freqmax-freqmin)) ? error("need to increase tmax") :
	fvec = [f for f in freqmin:Δf:freqmax]
	ifvec = fill(0, size(fvec))

	verbose && println("number of frequencies added to signal:\t", size(fvec,1))
	verbose && println("interval between random variable:\t", Δf)
	verbose && println("minimum frequency added to signal\t",minimum(fvec))
	verbose && println("maximum frequency added to signal\t",maximum(fvec))
	for iff in eachindex(fvec)
		ifvec[iff] =  Interpolation.indminn(fgrid, fvec[iff])[1]
	end

	for iff in eachindex(fvec)
		# translated sinc function
		fshift=fgrid[ifvec[iff]]  # shift for translated sinc function
		X=rand(rng) * exp(im*rand(Uniform(-pi, pi))) # uniformly distributed phase and random amplitude
		Xe=complex(1.0) # intialization for events
		for tfrac in nevent_fracs
			Xe += exp(im*fshift*2*π*tfrac*tgrid[nt]) # phase translations for each event
		end
		X *= Xe # add events to X
		for ifff in 1:length(fgrid)
			α=tmax*(abs(fgrid[ifff] - fshift)) # argument of sinc function
			xfreq[ifff] += (X * sinc(α)) # go
		end
	end
	xx=irfft(xfreq, nt)
	normalize!(xx)
	copyto!(x, xx)
	return x

end

"""
Tapering is necessary to be able to input random signal into finite-difference code
Filtering tapering are applied only if the length of the time series is greater than 10
"""
function get_tapered_random_tmax_signal(tgrid;
					fmin=nothing,
					fmax=nothing,
					tmaxfrac::Float64=1.0,
					dist=Uniform(-2.0, 2.0),
					sparsep=1.0,
					taperperc=20.
					)
	filt_flag=(length(tgrid) > 5) && (!(fmin === nothing)) && (!(fmax===nothing))

	fs = inv(step(tgrid));
	if(filt_flag)
		designmethod = Butterworth(6);
		filtsource = Bandpass(fmin, fmax; fs=fs);
	end

	itind = argmin(abs.(tgrid.-abs(tmaxfrac)*tgrid[end]))
	if(tmaxfrac>0.0)
		its=1:itind
	elseif(tmaxfrac<0.0)
		its=itind:length(tgrid)
	end
	# 20% taper window
	twin = taper(ones(length(its)),taperperc)
	X = zeros(length(its))
	wavsrc = zeros(length(tgrid))
	if(filt_flag)
		X[:] = rand(dist, length(its)) .* twin
	else
		X[:] = rand(dist, length(its))
	end
	if(sparsep ≠ 1.0)
		Xs=sprandn(length(X), sparsep)
		X[findall(Xs.==0.0)].=0.0
	end
	# band limit
	(filt_flag) && (filt!(X, digitalfilter(filtsource, designmethod), X))

	(length(X) ≠ 1) && normalize!(X)
	wavsrc[its] = X
	return wavsrc
end

"""
* `x` : first dimension is time
* `p` : sparsity
* `rng` : Random Number Generator
* `btperc` : Taper Perc
* `etperc` : Taper Perc
* `bfrac` : zeros at the beginning of each coloumn
* `efrac` : zeros fraction at the end of each coloumn
"""
function tapered_rand!(x;p=1.0,rng=Normal(),btperc=0., etperc=0.,bfrac=0.0, efrac=0.0)
	x[:]=0.0
	nt=size(x,1)

	a=1+round(Int,bfrac*nt)
	b=nt-round(Int,efrac*nt)

	dd=size(x)[2:end]
	for i in CartesianRange(dd)
		xx=view(x,a:b,i)
		rand!(rng, xx)
		taper!(xx, bperc=btperc, eperc=etperc)
		nxx=size(xx,1)
		if(p ≠ 1.0)
			xs=sprandn(nxx, p)
			xx[findn(xs.==0.0)]=0.0
		end
	end
	scale!(x, inv(vecnorm(x)))
	return x
end

"""
randomly shift and add x to itself
using circshift; not memory efficient
"""
function cshift_add_rand!(x::AbstractVector; ntimes=1, tminfrac=0.0, tmaxfrac=1.0, rand_amp_flag=true)
	nt=size(x,1)
	a=1+round(Int,tminfrac*nt)
	b=round(Int,tmaxfrac*nt)

	for it in 1:ntimes
		if(rand_amp_flag)
			amp=randn()
		else
			amp=1.0
		end
		# a random shift
		its=rand(DiscreteUniform(a,b))
		xx = circshift(x,(its,)) # pre-allocate for performance?
		for i in eachindex(x)
			x[i] += xx[i] * amp
		end
	end
end

"""
Construct Toy Green's functions
Decaying peaks, control number of events, and their positions, depending on bfrac and efrac.
* `afrac` : control amplitude of peaks
"""
function toy_green!(x;nevents=1,afrac=[inv(2^(ie-1)) for ie in 1:nevents],bfrac=0.0,efrac=0.0)
	nt=size(x,1)
	nr=size(x,2)
	x[:]=0.0
	a=1+round(Int,bfrac*nt)
	b=nt-round(Int,efrac*nt)
	(b-a<nevents+2) && error("not enough samples")

	itvec=zeros(Int64,nevents+1)
	for ir in 1:nr
		itvec[:]=0
		for ie in 1:nevents
			if(ie==1)
				# first event, direct arrival
				itvec[ie]=rand(DiscreteUniform(a,b-nevents-1))
			else
				# all other events after first event
				itvec[ie]=rand(DiscreteUniform(itvec[1]+1,b-1))
			end
			sgn=1.
			# randomly choose sign for all other events
			if(ie≠1)
				sgn=(rand(Bernoulli())==1) ? 1. : -1.
			end
			x[itvec[ie],ir]=sgn*afrac[ie]
		end
	end
	return x
end

function findfreq(
		  x::Array{Float64, ND},
		  tgrid;
		  attrib::Symbol=:peak,
		  threshold::Float64=-50.
		  ) where {ND}

	cx=rfft(x,[1]);
	fgrid=DSP.rfftfreq(length(tgrid),inv(step(tgrid))) # corresponding fgrid

	ax = (abs.(cx).^2); # power spectrum in dB

	if(maximum(ax) == 0.0)
		warn("x is zero"); return 0.0
	else
		ax /= maximum(ax);
		ax = 10. .* log10.(ax)
	end

	if(attrib == :max)
		iii=findlast(ax .>= threshold)
	elseif(attrib == :min)
		iii=findfirst(ax .>= threshold)
	elseif(attrib == :peak)
		iii=argmax(ax)
	end
	ii=CartesianIndices(size(ax))[iii][1]
	return fgrid[ii]

end

"""
Cosine taper a N-dimensional array along its first dimension.

# Arguments
* `x::Array{Float64,N}` :
* `perc::Float64` : taper percentage
"""
function taper(x::AbstractArray,perc::Float64)
	xout=copy(x);
	taper!(xout,perc)
	return xout
end

function taper!(x, perc=0.0; bperc=perc,eperc=perc)

	nt=size(x,1)
	nttb=min(round(Int,nt*bperc*0.01), nt)
	ntte=min(round(Int,nt*eperc*0.01), nt)
	kb=inv(2.0*round(Int,nt*bperc*0.01)-1)*pi
	ke=inv(2.0*round(Int,nt*eperc*0.01)-1)*pi
	for i in CartesianIndices(size(x))
		if(1≤i[1]≤nttb)
			x[i] *= sin((i[1]-1)*kb)
		end
		if(nt-ntte+1≤i[1]≤nt)
			x[i] *= sin((-i[1]+nt)*ke)
		end
	end
end

include("Filters.jl")

end # module


const yosb1 = [8.44e-1, 9.08e-1, 9.32e-1, 9.36e-1, 9.25e-1, 9.08e-1,
               8.81e-1, 8.50e-1, 8.18e-1, 7.84e-1, 7.51e-1, 7.18e-1,
               6.85e-1, 6.58e-1, 6.28e-1, 6.03e-1, 5.80e-1, 5.58e-1,
               5.38e-1, 5.22e-1, 5.06e-1, 4.90e-1, 4.78e-1, 4.67e-1,
               4.57e-1, 4.48e-1, 4.38e-1, 4.31e-1, 4.24e-1, 4.20e-1,
               4.14e-1, 4.11e-1, 4.06e-1]

const osborne1 = let res_init=zeros(33), jac_init=zeros(33,5), x_init=[.5, 1.5, -1.0, .01, .02]

    function res(x, r)
        for i = 1:33
            ti = 1.0e1 * Float64(1-i)
            r[i] = yosb1[i] - (x[1] + x[2] * exp(x[4] * ti) + x[3] * exp(x[5] * ti))
        end
        return r
    end

    function jac(x, j)
        for i = 1:33
            ti = 1.0e1 * Float64(1-i)
            j[i, 1] = -1.e0
            j[i, 2] = -exp(x[4] * ti)
            j[i, 3] = -exp(x[5] * ti)
            j[i, 4] = ti * x[2] * j[i, 2]
            j[i, 5] = ti * x[3] * j[i, 3]
        end
        return j
    end

    f(;scale=1, verbose=false, print_steps=false) = testone("osborne", res, jac,
                                                            res_init, jac_init, x_init;
                                                            scale=scale, verbose=verbose,
                                                            print_steps=print_steps)
end

const yosb2 = [ 1.366e0, 1.191e0, 1.112e0, 1.013e0, 9.91e-1, 8.85e-1,
                8.31e-1, 8.47e-1, 7.86e-1, 7.25e-1, 7.46e-1, 6.79e-1,
                6.08e-1, 6.55e-1, 6.16e-1, 6.06e-1, 6.02e-1, 6.26e-1,
                6.51e-1, 7.24e-1, 6.49e-1, 6.49e-1, 6.94e-1, 6.44e-1,
                6.24e-1, 6.61e-1, 6.12e-1, 5.58e-1, 5.33e-1, 4.95e-1,
                5.00e-1, 4.23e-1, 3.95e-1, 3.75e-1, 3.72e-1, 3.91e-1,
                3.96e-1, 4.05e-1, 4.28e-1, 4.29e-1, 5.23e-1, 5.62e-1,
                6.07e-1, 6.53e-1, 6.72e-1, 7.08e-1, 6.33e-1, 6.68e-1,
                6.45e-1, 6.32e-1, 5.91e-1, 5.59e-1, 5.97e-1, 6.25e-1,
                7.39e-1, 7.10e-1, 7.29e-1, 7.20e-1, 6.36e-1, 5.81e-1,
                4.28e-1, 2.92e-1, 1.62e-1, 9.8e-2, 5.4e-2]

# uftolg is a machine-dependent constant.  it is just slightly
# larger than the log of the smallest positive machine number.
const uftolg = 1.999e0 * log(sqrt(1.001 * 2.2250738585072014e-308))

const osborne2 = let res_init=zeros(65), jac_init=zeros(65,11),
                     x_init=[1.3, .65, .65, .7, .6, 3., 5., 7., 2., 4.5, 5.5]

    function res(x, r)
        for i = 1:65
            ti = 0.1e0 * Float64(1-i)
            ri = x[1] * exp(x[5] * ti)
            for j = 2:4
                t = 0.e0
                theta = -x[j+4] * (ti + x[j+7]) ^ 2
                if theta > uftolg
                    t = exp(theta)
                end
                ri = ri + x[j] * t
            end
            r[i] = yosb2[i] - ri
        end
        return r
    end

    function jac(x, j)
        for i = 1:65
            ti = Float64(1 - i) * 1.e-1
            j[i, 1] = -exp(x[5] * ti)
            j[i, 5] = x[1] * ti * j[i,1]
            for k = 2:4
                t = x[k + 7] + ti
                theta = -x[k+4]*t*t
                r2 = theta > uftolg ? -exp(theta) : 0.e0
                j[i, k] = r2
                r2 = -t * r2 * x[k]
                j[i, k+4] = r2*t
                j[i, k+7] = 2.e0 * x[k+4] * r2
            end
        end
        return j
    end

    f(;scale=1, verbose=false, print_steps=false) = testone("osborne2", res, jac,
                                                            res_init, jac_init, x_init;
                                                            scale=scale, verbose=verbose,
                                                            print_steps=print_steps)
end


using Altro
using TrajectoryOptimization
using StaticArrays
using BenchmarkTools
using RobotZoo
using RobotDynamics
using StaticArrays, LinearAlgebra
using ForwardDiff
using Test
using FiniteDiff
using Random
const TO = TrajectoryOptimization
const RD = RobotDynamics

# Test Functions
function Πsoc(x)
    v = x[1:end-1]
    s = x[end]
    a = norm(v)
    if a <= -s
        return zero(x)
    elseif a <= s
        return x
    elseif a >= abs(s)
        x̄ = append!(v, a)
        return 0.5*(1 + s/a) * x̄
    end
    throw(ErrorException("Invalid second-order cone"))
end
function in_soc(x)
    v = x[1:end-1]
    s = x[end]
    a = norm(v)
	a <= s
end

function alcost(x,z,μ)
    c = push(x, u_bnd)
    proj = Πsoc(z - μ .* c)
    1/2μ*(proj'proj - z'z)
end

function auglag(f,g, x,z,μ)
    # ceq = h(x)
    L0 = f(x) #+ y'ceq + 0.5*μ*ceq'ceq

    # cones = g(x)
    # p = length(cones)
    for i = 1:length(z)
        L0 += alcost(x[uinds[i]], z[i], μ)
        # cone = cones[i]
		# proj = Πsoc(z[i] - μ*cones[i])
        # Iμ = I * inv(μ)
		# pen = proj'proj - z[i]'z[i]
        # L0 += 1 / 2μ * pen
    end
    return L0
end

function DI_objective(Q,R,Qf,xf,dt)
    g = Float64[]
    h = Float64[]
    c = 0.0
    for k = 1:N-1
        append!(g, -Q*xf)
        append!(g, zeros(m))
        append!(h, Q.diag)
        append!(h, R.diag)
        c += 0.5*xf'Q*xf
	end
	# g .*= dt
	# h .*= dtk
	# c *= dt
    append!(g, -Qf*xf)
    # append!(g, zeros(m))
    append!(h, Qf.diag)
    # append!(h, R.diag)
    c += 0.5*xf'Qf*xf
    H = Diagonal(h)
    return H,g,c
end

function DI_dynamics(x0, N; dt=0.1, gravity::Bool=false)
    n = length(x0)
    m = n ÷ 2
    D = m
    NN = N*(n+m)

    #  Continuous dynamics
    Ac = zeros(n,n)
    Bc = zeros(n,m)
    for i = 1:D
        Ac[i,D+i] = 1
        Bc[D+i,i] = 1
    end

    # Euler integration
    Ad = I + Ac*dt
    Bd = Bc*dt

    iA = 1:n
    jA = iA
    jB = n .+ (1:m)
    jA2 = jA .+ (n+m)

    A = zeros(N*n, NN)
    b = zeros(N*n)
    A[iA,jA] .= I(n)
    b[iA] .= -x0
    iA = iA .+ n
    for k = 1:N-1
        A[iA,jA] .= Ad
        A[iA,jB] .= Bd
        A[iA,jA2] .= -I(n)
        if gravity
            b[iA[D]] = -9.81
        end

        # Advance inds
        iA = iA .+ n
        jA = jA .+ (n+m)
        jB = jB .+ (n+m)
        jA2 = jA2 .+ (n+m)
    end
    return A,b
end

function DI_cones(::Val{D},N) where D
    uinds = SVector{D}(1:D) .+ 2D
    qinds = [uinds .+ (k-1)*3D for k = 1:N-1]
end
DI_cones(D,N) = DI_cones(Val(D),N)  # WARNING: type unstable

## Setup
D,N = 2,11
n,m = 2D,D
dt = 0.1
tf = (N-1)*dt
x0 = @SVector fill(0., n)
xf = [(@SVector fill(1.,D)); (@SVector fill(0, D))]
Q = Diagonal(@SVector fill(1.0, n)) * dt
R = Diagonal(@SVector fill(1e-1, m)) * dt
Qf = (N-1)*Q * 100
qinds = DI_cones(D,N)
u_bnd = 6.0

# Define batch functions
H_obj,g_obj,c_obj = DI_objective(Q,R,Qf,xf,dt)
Adyn,bdyn = DI_dynamics(x0,N, dt=dt)

di_obj(x) = 0.5*x'H_obj*x + g_obj'x + c_obj
di_dyn(x) = Adyn*x + bdyn
di_soc(x) = [push(x[qi],u_bnd) for qi in qinds]

# Define using ALTRO
model = RobotZoo.DoubleIntegrator(D)

obj = LQRObjective(Q, R, Qf, xf, N)
cons = ConstraintList(n,m,N)
cone = TO.SecondOrderCone()
# cone = TO.Inequality()
u_bnd = 6.0
connorm = NormConstraint(n, m, u_bnd, cone, :control)
add_constraint!(cons, connorm, 1:N-1)
prob = Problem(model, obj, x0, tf, xf=xf, constraints=cons, integration=RD.Euler(model))
solver = Altro.ALSolver(prob, use_static=Val(true))
conSet = TO.get_constraints(solver)
alobj = TO.get_objective(solver)

# Initialization
inds = LinearIndices(zeros(3D,N))
xinds = [SVector{n}(z[1:n]) for z in eachcol(inds)]
uinds = [SVector{m}(z[n+1:end]) for z in eachcol(inds)][1:end-1]
NN = 3D*N-D
P = (N-1)*m
x0 = rand(NN)
z = rand.(length.(di_soc(x0)))
z0 = deepcopy(z)
μ = 2.4

X0 = [x0[xi] for xi in xinds]
U0 = [x0[ui] for ui in uinds]
Z0 = SampledTrajectory(X0,U0, dt=dt)
initial_trajectory!(solver, Z0)

## Test objective and soc constraints
z = deepcopy(z0)
Altro.reset_duals!(conSet)
@test cost(alobj, Z0) ≈ di_obj(x0)
@test conSet[1].vals ≈ di_soc(x0)

# copy multipliers to ALconSet
cval = conSet[1]
for i = 1:N-1
	cval.λ[i] .= z[i]
	cval.μ[i] .= μ
end

# test augmented Lagrangian value
LA(x) = auglag(di_obj, di_soc, x, z, μ)
@test di_obj(x0) ≈ cost(alobj.obj, Z0)  # unconstrained cost
@test cost(solver) ≈ LA(x0)
@test !(cost(solver) ≈ di_obj(x0))   # make sure the AL cost isn't the same as the normal cost

λbar = z .- μ .* di_soc(x0)
statuses = [TO.cone_status(TO.SecondOrderCone(), λ) for λ in λbar]
@test :below ∈ statuses

# E = TO.QuadraticObjective(n,m,N)
E0 = Altro.CostExpansion{Float64}(n,m,N)
E = Altro.get_ilqr(solver).Efull
let solver = solver.ilqr
    Altro.cost_expansion!(solver.obj, E, solver.Z)
    @test_throws AssertionError Altro.cost_expansion!(solver.obj, E0, solver.Z)
end
# TO.cost_expansion!(E, alobj, Z0)
grad = vcat([e.grad for e in E]...)[1:end-m]
@test grad ≈ FiniteDiff.finite_difference_gradient(LA, x0)

hess_blocks = vcat([[e.xx, e.uu] for e in E]...)
hess = cat(hess_blocks..., dims=(1,2))[1:end-m, 1:end-m]
@test hess ≈ FiniteDiff.finite_difference_hessian(LA, x0) atol=1e-4

@test all(in_soc.(di_soc(x0)))
@test !all(in_soc.(z .- μ*di_soc(x0)))

# Make it go outside the cone
Random.seed!(1)
z = rand.(length.(di_soc(x0)))
z .*= 100
λbar = z .- μ .* di_soc(x0)
statuses = [TO.cone_status(TO.SecondOrderCone(), λ) for λ in λbar]
for i = 1:20
    if :outside in statuses && :in in statuses
        break
    else
        z .= rand.(length.(di_soc(x0)))
        z .*= 100
        λbar .= z .- μ .* di_soc(x0)
        statuses .= [TO.cone_status(TO.SecondOrderCone(), λ) for λ in λbar]
    end
end
@test :outside ∈ statuses
@test :in ∈ statuses
for i = 1:N-1
	cval.λ[i] .= z[i]
	cval.μ[i] .= μ
    cval.μinv[i] .= inv.(μ)
end
LA(x) = auglag(di_obj, di_soc, x, z, μ)

# E = TO.QuadraticObjective(n,m,N)
# E = Altro.CostExpansion{Float64}(n,m,N)
E = Altro.get_ilqr(solver).Efull
@test LA(x0) ≈ cost(solver, Z0)
Altro.cost_expansion!(alobj, E, Z0)
# TO.cost_expansion!(E, alobj, Z0)
alcon = conSet[1]
grad = vcat([e.grad for e in E]...)[1:end-m]
@test grad ≈ FiniteDiff.finite_difference_gradient(LA, x0) atol=1e-6
[grad FiniteDiff.finite_difference_gradient(LA, x0)]

hess_blocks = vcat([[e.xx, e.uu] for e in E]...)
hess = cat(hess_blocks..., dims=(1,2))[1:end-m, 1:end-m]
@test hess ≈ FiniteDiff.finite_difference_hessian(LA, x0) atol=1e-2

## Solve it
prob = Problem(model, obj, zero(SVector{n}), tf, xf=xf, constraints=cons, integration=RD.Euler(model))
solver = Altro.ALSolver(prob, verbose=0, show_summary=false,
	projected_newton=true, penalty_initial=100.0, penalty_scaling=50,
	cost_tolerance_intermediate=1e-1)
initial_controls!(solver, [rand(D) for k = 1:N])
solve!(solver)
@test all(abs.(norm.(controls(solver))[[1,2,N-2,N-1]] .- u_bnd) .< 1e-5)

# prob = Problem(model, obj, xf, tf, x0=x0, constraints=cons)
# solver = ALTROSolver(prob, show_summary=true, verbose=2, projected_newton=false)
# @test benchmark_solve!(solver).allocs == 0
# benchmark_solve!(solver)


#=
Challenge #261 [Easy] verifying 3x3 magic squares

Description

A 3x3 magic square is a 3x3 grid of the numbers 1-9 such that each row, column,
and major diagonal adds up to 15. Here's an example:

8 1 6
3 5 7
4 9 2

The major diagonals in this example are 8 + 5 + 2 and 6 + 5 + 4. (Magic squares
have appeared here on r/dailyprogrammer before, in #65 [Difficult] in 2012.)
Write a function that, given a grid containing the numbers 1-9, determines whether
it's a magic square. Use whatever format you want for the grid, such as a
2-dimensional array, or a 1-dimensional array of length 9, or a function that
takes 9 arguments. You do not need to parse the grid from the program's input,
but you can if you want to. You don't need to check that each of the 9 numbers
appears in the grid: assume this to be true.

Example inputs/outputs

[8, 1, 6, 3, 5, 7, 4, 9, 2] => true
[2, 7, 6, 9, 5, 1, 4, 3, 8] => true
[3, 5, 7, 8, 1, 6, 4, 9, 2] => false
[8, 1, 6, 7, 5, 3, 4, 9, 2] => false

Optional bonus 1

Verify magic squares of any size, not just 3x3.

Optional bonus 2

Write another function that takes a grid whose bottom row is missing, so it only
has the first 2 rows (6 values). This function should return true if it's possible
to fill in the bottom row to make a magic square. You may assume that the numbers
given are all within the range 1-9 and no number is repeated. Examples:

[8, 1, 6, 3, 5, 7] => true
[3, 5, 7, 8, 1, 6] => false

Hint: it's okay for this function to call your function from the main challenge.
This bonus can also be combined with optional bonus 1. (i.e. verify larger magic
squares that are missing their bottom row.)

( source: https://redd.it/4dccix )
=#

read_matrix( data::DataType=Any, input::IO=STDIN, spacing::Char=' ' ) = readdlm(
input, spacing, data, header=false, skipblanks=true, ignore_invalid_chars=true,
quotes=false, comments=false )

function pseudo_magic_square( grid::Matrix{ Int } )
    local value::Bool = loop::Bool = false
    local missing::Int = tmp::Int = 0
    local last::Array{ Int, 1 } = Int[]

    for i in 1:size( grid, 2 )
        column = sum( grid[ :, i ] )
        missing = 15 - column
        push!( last, missing )

        if 9 >= missing
            for j in 1:size( grid, 1 ), k in 1:size( grid, 2 )
                if grid[ j, k ] == missing
                    loop = true
                    break
                end
            end
        else
            loop = true
        end

        if true == loop
            break
        end
    end

    if false == loop
        for i in 1:2
            tmp += grid[ i,i ]
        end
        tmp += last[ 3 ]
        if 15 == tmp
            value = true
        end
    end

    return value
end

function verify_magic_square( grid::Matrix{ Int } )
    local value::Bool = false
    local line::Integer = column::Integer = 0

    if 15 == trace( grid )  #   sums the value of the principal diagonal
        for i in 1:size( grid, 1 )
            line = sum( grid[ i, : ] )
            column = sum( grid[ :, i ] )
            if line != 15 || column != 15
                break
            elseif size( grid, 1 ) == i   #   if is the last verification and it's valid
                value = true
            end
        end
    end

    return value
end

#=
    I've modified the input format. So, that only be possibles input like:

        8 1 6                       8 1 6
        3 5 7           or          3 5 7
        4 9 2
=#
function main( )
    local input::Matrix{ Int } = read_matrix( Int )

    println( @time ( size( input, 1 ) == size( input, 2 ) ) ?
                     verify_magic_square( input ):
                     pseudo_magic_square( input ) )
end

main( )


export fieldterm, xxzmodel, hopping, z_field

_get_interaction(J::AbstractMatrix, i, j) = J[i,j]
_get_interaction(J::Number, i, j) = J

function fieldterm(N, op)
    res = op
    for i in 2:N
        res = 𝟙 ⊗ res + op ⊗ identity_op(i-1)
    end
    res
end

function hopping(N, i, j)
    entries = 2^(N-1)
    A,B,V = zeros(Int, entries), zeros(Int, entries), ones(Int, entries)
    _hopping_coordinates!(A, B, N, i, j)
    return sparse(A,B,V,2^N,2^N)
end

xxzmodel(J::AbstractMatrix, Δ) = xxzmodel(size(J, 1), J, Δ)

function xxzmodel(N,J,Δ)
    # 2^N diagonal entries + N(N-1)/2 * 2^(N-1) offdiagonal terms
    entries = 2^N + 2^(N-2)*N*(N-1)
    A,B,V = zeros(Int, entries), zeros(Int, entries), zeros(float(promote_type(eltype(J), typeof(Δ))), entries)

    @views _zcorrelator_coo!(A[1:2^N],B[1:2^N],V[1:2^N], N, J)
    V[1:2^N] .*= Δ

    @views _hopping!(A[2^N+1:end], B[2^N+1:end], V[2^N+1:end], N, J)
    V[2^N+1:end] .*= 2

    V ./= 4 # pauli -> spin operators
    sparse(A,B,V,2^N,2^N)
    ## In principle using Hermitian and only computing only
    ## half of the off-diagonal speeds up constructing the Hamiltonian.
    ## In practice this slows down everythin else dramatically, as Hermitian
    ## (sometimes) hides the fact that this is a sparse matrix and that leads to
    ## unfortunate interactions (f.e. Hermitian(sparse) ≈ dense takes AGES)
    #Hermitian(sparse(A,B,V,2^N,2^N), :L)
end

"""
    _zcorrelator_coo!(A, B, V, J)

Generate the COO vectors of ∑ᵢⱼJᵢⱼ σzⁱσzʲ.

# Parameters
 - `A`,`B`,`V` vectors for column (`A`) and row indices (`B`) respectively values (`V`)
 - `N` number of sites in the system
 - `J` interaction matrix of size `N`x`N` or scalar. Assumed symmetric!

`A`,`B` and `V` need to have length 2^N.
"""
function _zcorrelator_coo!(A,B,V, N, J)
    A[:] .= 1:2^N
    B[:] .= 1:2^N
    V[:] .= 0
    for i in 1:N
        for j in i+1:N
            _zcorrelator_values!(V, N, _get_interaction(J, j,i), i, j)
        end
    end
end

"""
    _zcorrelator_values!(V, N, Jij, i, j)

*ADD* to `V` the diagonal of Jᵢⱼ σzⁱσzʲ.

# Parameters
 - `V` array of length 2^`N`. Will be mutated!
 - `N` number of sites in the system
 - `Jij` interaction strength
 - `i`,`j` sites to couple
"""
function _zcorrelator_values!(V, N, Jij, i, j)
    # 𝟙(sec1) ⊗ σz ⊗ 𝟙(sec2) ⊗ σz ⊗ 𝟙(sec3)

    #             | 𝟙(sec2) |     0    |   | 𝟙(sec3) |    0     |
    # = 𝟙(sec1) ⊗ | ------- | -------- | ⊗ | ------- | -------- |
    #             |    0    | -𝟙(sec2) |   |    0    | -𝟙(sec3) |

    #             |           [ 𝟙(sec3) |     0    ] |                                  |
    #             | 𝟙(sec2) ⊗ [ ------- | -------- ] |              0                   |
    #             |           [    0    | -𝟙(sec3) ] |                                  |
    # = 𝟙(sec1) ⊗ | -------------------------------- | -------------------------------- |
    #             |                                  |           [ -𝟙(sec3) |    0    ] |
    #             |              0                   | 𝟙(sec2) ⊗ [ -------- | ------- ] |
    #             |                                  |           [    0     | 𝟙(sec3) ] |
    sec1 = 2^(i-1)
    sec2 = 2^(j-i-1)
    sec3 = 2^(N-j)

    blocksize3 = 2sec3 # σz ⊗ 𝟙(sec3)
    blocksize23 = 2sec2*blocksize3 # size of: σz ⊗ 𝟙(sec2) ⊗ σz ⊗ 𝟙(sec3)

    #sec1 loop
    for offset1 in blocksize23 .* (0:sec1-1)
        # upper block
        for offset2 in offset1 .+ blocksize3 .* (0:sec2-1)
            V[1 .+ offset2 .+ (0:sec3-1)] .+= Jij
            V[1 .+ offset2 .+ (sec3:2sec3-1)] .+= -Jij
        end
        # lower block
        for offset2 in offset1 .+ blocksize3 .* (sec2:2sec2-1)
            V[1 .+ offset2 .+ (0:sec3-1)] .+= -Jij
            V[1 .+ offset2 .+ (sec3:2sec3-1)] .+= Jij
        end
    end
end

"""
    _hopping!(A,B,V, J)

Generate the COO vectors for ∑ᵢⱼ Jᵢⱼ (σ₊ⁱσ₋ʲ + σ₋ⁱσ₊ʲ)

# Parameters
 - `A`,`B`,`V` vectors for column (`A`) and row indices (`B`) respectively values (`V`)
 - `J` interaction matrix of size `N`x`N`. Assumed symmetric!

 If `J` is `N`x`N`, then `A`,`B` and `V` need to have length (N-1)N*2^(N-2)
"""
function _hopping!(A,B,V, N, J)
    fill!(V, 1)
    length = 2^(N-1)
    at = 0
    for i in 1:N
        for j in i+1:N
            range = at .+ (1:length)
            @views _hopping_coordinates!(A[range], B[range], N, i, j)
            V[range] .*= _get_interaction(J, j, i)
            at += length
        end
    end
    ## _hopping values! does the transpose
    ## no measureable performance difference
    #copyto!(A,at+1,B,1,at)
    #copyto!(B,at+1,A,1,at)
    #copyto!(V,at+1,V,1,at)
end

"""
    _hopping_coordinates!(A, B, N, i, j)

Compute the coordinates of the matrix entries of the hopping operator between site i and J
in a spin chain. The values are 1 everywhere, so there is no input/output for these in this
function.

Hopping operator: σ₊ⁱσ₋ʲ + σ₋ⁱσ₊ʲ = 𝟙⊗…⊗𝟙⊗σ₊⊗𝟙⊗…⊗𝟙⊗σ₋⊗𝟙⊗…⊗𝟙 + transpose

# Parameters
 - `A`,`B` coordinate vectors. These should have length 2^(N-2) will be overwritten!
 - `N` number of sites in the system
 - `i`,`j` sites the excitation can hop between. 1 ≤ i,j ≤ N
"""
function _hopping_coordinates!(A,B,N,i,j)
    i > j && return _hopping_coordinates!(A,B,N,j,i)
    # 𝟙(sec1) ⊗ σ+ ⊗ 𝟙(sec2) ⊗ σ- ⊗ 𝟙(sec3) + transpose
    #             |                                  |                                  |
    #             |                0                 |                 0                |
    #             |                                  |                                  |
    # = 𝟙(sec1) ⊗ | -------------------------------- | -------------------------------- |
    #             |           [    0    |  𝟙(sec3) ] |                                  |
    #             | 𝟙(sec2) ⊗ [ ------- | -------- ] |                 0                |
    #             |           [    0    |    0     ] |                                  |
    sec1 = 2^(i-1)
    sec2 = 2^(j-i-1)
    sec3 = 2^(N-j)

    blocksize3 = 2sec3 # σz ⊗ 𝟙(sec3)
    blocksize23 = 2sec2*blocksize3 # size of: σz ⊗ 𝟙(sec2) ⊗ σz ⊗ 𝟙(sec3)

    currentIndex = 0
    #sec1 loop
    for offset1 in blocksize23 .* (0:sec1-1)
        for offset2 in blocksize3 .* (0:sec2-1)
            rowOffset2 = offset1 + offset2 # left columns # offset for the row-COORDINATE
            colOffset2 = offset1 + offset2 + blocksize3*sec2 # bottom rows # offset for the col-COORDINATE

            A[currentIndex .+ (1:sec3)] .= colOffset2 .+ (1:sec3) # upper rows
            B[currentIndex .+ (1:sec3)] .= rowOffset2 .+ (sec3+1:2sec3) # right columns
            currentIndex += sec3

            ## now transpose
            B[currentIndex .+ (1:sec3)] .= colOffset2 .+ (1:sec3) # upper rows
            A[currentIndex .+ (1:sec3)] .= rowOffset2 .+ (sec3+1:2sec3) # right columns
            currentIndex += sec3
        end
    end
end

z_field(N::Int) = z_field(ones(N))

function z_field(field_values::Vector)
    N = length(field_values)
    ret = zeros(2^N)
    for (i,hz) in enumerate(field_values)
        _z_field_values!(ret, N, hz, i)
    end
    return Diagonal(ret)
end

"""
    _z_field!(V, N, hz, i)

*ADD* to `V` the diagonal of hz σzⁱ.

# Parameters
 - `V` array of length 2^`N`. Will be mutated!
 - `N` number of sites in the system
 - `hz` field strength
 - `i` site to act on
"""
function _z_field_values!(V, N, hz, i)
    # σz^(i) = 𝟙(sec1) ⊗ σz ⊗  𝟙(sec2)
    sec1 = 2^(i-1)
    sec2 = 2^(N-i)

    blocksize2 = 2sec2 # σz ⊗ 𝟙(sec2)

    for offset1 in blocksize2 .* (0:sec1-1)
        V[offset1 .+ (1:sec2)] .+= hz
        V[offset1 .+ (sec2+1:2sec2)] .-= hz
    end
end


include("IRK8.jl")

function ThreeBodyEnergy(u)
    norm12 = sqrt(u[1]^2+u[2]^2)
    norm23 = sqrt(u[3]^2+u[4]^2)
    norm31 = sqrt(u[5]^2+u[6]^2)
    return (u[7]^2+u[8]^2+u[9]^2+u[10]^2+u[11]^2+u[12]^2)/2 -
           (1/norm12 + 1/norm23 + 1/norm31)
end

function ThreeBodyODE!(du,u,p,t)
    q12 = [u[1], u[2]]
    q23 = [u[3], u[4]]
    q31 = [u[5], u[6]]
    aux12 = q12*dot(q12, q12)^-1.5
    aux23 = q23*dot(q23, q23)^-1.5
    aux31 = q31*dot(q31, q31)^-1.5
    dv1 = aux12 - aux31
    dv2 = aux23 - aux12
    dv3 = aux31 - aux23
    @. du[1:4] = u[9:12] - u[7:10]
    @. du[5:6] = u[7:8] - u[11:12]
    du[7:12] .= [dv1[1], dv1[2], dv2[1], dv2[2], dv3[1], dv3[2]]
    du[13] = 1
    return nothing
end

function ThreeBodyODEGlobalTR!(du,u,p,t)
    q12 = [u[1], u[2]]
    q23 = [u[3], u[4]]
    q31 = [u[5], u[6]]
    dot12 = dot(q12, q12)
    dot23 = dot(q23, q23)
    dot31 = dot(q31, q31)
    norm12 = sqrt(dot12)
    norm23 = sqrt(dot23)
    norm31 = sqrt(dot31)
    aux12 = q12/(dot12*norm12)
    aux23 = q23/(dot23*norm23)
    aux31 = q31/(dot31*norm31)
    dv1 = aux12 - aux31
    dv2 = aux23 - aux12
    dv3 = aux31 - aux23
    @. du[1:4] = u[9:12] - u[7:10]
    @. du[5:6] = u[7:8] - u[11:12]
    du[7:12] .= [dv1[1], dv1[2], dv2[1], dv2[2], dv3[1], dv3[2]]
    A = 1/dot12+1/dot23+1/dot31
    B = 1/norm12+1/norm23+1/norm31
    s = (A*B)^-0.5
    @. du *= s
    du[13] = s
    return nothing
end

abs2rel(u) = [u[1]-u[3], u[2]-u[4], u[3]-u[5], u[4]-u[6], u[5]-u[1], u[6]-u[2],
                 u[7]-u[9], u[8]-u[10], u[9]-u[11], u[10]-u[12], u[11]-u[7], u[12]-u[8]]

function rel2abs(u)
    x1 = (u[5]-u[1])/3.
    x2 = (u[1]-u[3])/3.
    x3 = (u[3]-u[5])/3.
    y1 = (u[6]-u[2])/3.
    y2 = (u[2]-u[4])/3.
    y3 = (u[4]-u[6])/3.

    vx1 = (u[11]-u[7])/3.
    vx2 = (u[7]-u[9])/3.
    vx3 = (u[9]-u[11])/3.
    vy1 = (u[12]-u[8])/3.
    vy2 = (u[8]-u[10])/3.
    vy3 = (u[10]-u[12])/3.

    return [x1, y1, x2, y2, x3, y3, vx1, vy1, vx2, vy2, vx3, vy3]
end

function visualize(u0, odef, p, T0, Tend, n, m, title="", triangles=false)
    odef = ThreeBodyODEGlobalTR!
    p = nothing
    n = 256
    m = 16;
    sol = IRK8(u0, T0, Tend, n, m, odef, p);
    uu = sol.u;
    pl = plot(title=title, aspect_ratio=1)

    uu_abs = [rel2abs(u) for u in uu]
    for j in 1:3
        xx = [u[2*(j-1)+1] for u in uu_abs]
        yy = [u[2*(j-1)+2] for u in uu_abs]

        plot!(xx,yy,legend=false)
    end

    if triangles
        u0_abs = uu_abs[1]
        uT_abs = uu_abs[end]
        plot!([u0_abs[i] for i in [1, 3, 5, 1]], [u0_abs[i] for i in [2, 4, 6, 2]], color="gray")
        plot!([uT_abs[i] for i in [1, 3, 5, 1]], [uT_abs[i] for i in [2, 4, 6, 2]], color="purple")
    end

    display(pl)
    #return sol.retcode
end


#src Introduction to DG methods
using Test: @test #src

# This tutorial is about how to set up a simple way to approximate the solution of a hyperbolic partial
# differential equation. First, we will implement a basic and naive algorithm. Then, we will use predefined
# features from [Trixi.jl](https://github.com/trixi-framework/Trixi.jl) to show how you can use Trixi.jl on your own.

# We will implement the scalar linear advection equation in 1D with the advection velocity $1$.
# ```math
# u_t + u_x = 0,\; \text{for} \;t\in \mathbb{R}^+, x\in\Omega=[-1,1]
# ```
# We define the domain $\Omega$ by setting the boundaries.
coordinates_min = -1.0 # minimum coordinate
coordinates_max = 1.0  # maximum coordinate

# We assume periodic boundaries and the following initial condition.
initial_condition_sine_wave(x) = 1.0 + 0.5 * sin(pi * x)

# ## The discontinuous Galerkin collocation spectral element method (DGSEM)
# ### i. Discretization of the physical domain
# To improve precision we want to approximate the solution on small parts of the physical domain.
# So, we split the domain $\Omega=[-1, 1]$ into elements $Q_l$ of length $dx$.

n_elements = 16 # number of elements

dx = (coordinates_max - coordinates_min) / n_elements # length of one element

# To make the calculation more efficient and storing less information, we transform each element
# $Q_l$ with center point $x_l$ to a reference element $E=[-1, 1]$
# ```math
# Q_l=\Big[x_l-\frac{dx}{2}, x_l+\frac{dx}{2}\Big] \underset{x(\xi)}{\overset{\xi(x)}{\rightleftarrows}} [-1, 1].
# ```
# So, for every element the transformation from the reference domain to the physical domain is defined by
# ```math
# x(\xi) = x_l + \frac{dx}{2} \xi,\; \xi\in[-1, 1]
# ```
# Therefore,
# ```math
# \begin{align*}
# u &= u(x(\xi), t) \\
# u_x &= u_\xi \frac{d\xi}{dx} \\[3pt]
# \frac{d\xi}{dx} &= (x_\xi)^{-1} = \frac{2}{dx} =: J^{-1}. \\
# \end{align*}
# ```
# Here, $J$ is the Jacobian determinant of the transformation.

# Using this transformation, we can transform our equation for all elements $Q_l$.
# ```math
# \frac{dx}{2} u_t^{Q_l} + u_\xi^{Q_l} = 0 \text{, for }t\in\mathbb{R}^+,\; \xi\in[-1, 1]
# ```

# ### ii. Polynomial approach
# Now, we want to approximate the solution in each element $Q_l$ by a polynomial of degree $N$. Since we transformed
# the equation, we can use the same polynomial approach for the reference coordinate $\xi\in[-1, 1]$ in every
# physical element $Q_l$. This saves a lot of ressources by reducing the amount of calculations needed
# and storing less information.

# For DGSEM we choose [Lagrange basis functions](https://en.wikipedia.org/wiki/Lagrange_polynomial)
# $\{l_j\}_{j=0}^N$ as our polynomial basis of degree $N$ in $[-1, 1]$.
# The solution in element $Q_l$ can be approximated by
# ```math
# u(x(\xi), t)\big|_{Q_l} \approx u^{Q_l}(\xi, t) = \sum_{j=0}^N u_j^{Q_l}(t) l_j(\xi)
# ```
# with $N+1$ coefficients $\{u_j^{Q_l}\}_{j=0}^N$.
# By contruction the Lagrange basis has some useful advantages. This basis is defined by $N+1$ nodes, which
# fulfill a Kronecker property at the exact same nodes. Let $\{\xi_i\}_{i=0}^N$ be these nodes.
# ```math
# l_j(\xi_i) = \delta_{i,j} =
# \begin{cases}
# 1, & \text{if } i=j \\
# 0, & \text{else.}
# \end{cases}
# ```
# Because of this property, the polynomial coefficients are exact the values of $u^{Q_l}$ at the nodes
# ```math
# u^{Q_l}(\xi_i, t) = \sum_{j=0}^N u_j^{Q_l}(t) \underbrace{l_j(\xi_i)}_{=\delta_{ij}} = u_i^{Q_l}(t).
# ```

# Next, we want to select the nodes $\{\xi_i\}_{i=0}^N$, which we use for the construction of the Lagrange
# polynomials. We choose the $N+1$ Gauss-Lobatto nodes, which are used for the
# [Gaussian-Lobatto quadrature](https://mathworld.wolfram.com/LobattoQuadrature.html).
# These always contain the boundary points at $-1$ and $+1$ and are well suited as interpolation nodes.
# The corresponding weights will be referred to as $\{w_j\}_{j=0}^N$.
# In Trixi.jl the basis with Lagrange polynomials on Gauss-Lobatto nodes is already defined.
using Trixi
polydeg = 3 #= polynomial degree = N =#
basis = LobattoLegendreBasis(polydeg)
# The Gauss-Lobatto nodes are
nodes = basis.nodes
# with the corresponding weights
weights = basis.weights

# To illustrate how you can integrate using numerical quadrature with this Legendre-Gauss-Lobatto nodes,
# we give an example for $f(x)=x^3$. Since $f$ is of degree $3$, a polynomial interpolation with $N=3$ is exact.
# Therefore, the integral on $[-1, 1]$ can be calculated by
# ```math
# \begin{align*}
# \int_{-1}^1 f(x) dx &= \int_{-1}^1 \Big( \sum_{j=0}^3 f(\xi_j)l_j(x) \Big) dx
# = \sum_{j=0}^3 f(\xi_j) \int_{-1}^1 l_j(x)dx \\
# &=: \sum_{j=0}^3 f(\xi_j) w_j
# = \sum_{j=0}^3 \xi_j^3 w_j
# \end{align*}
# ```
# Let's use our nodes and weights for $N=3$ and plug in
integral = sum(nodes.^3 .* weights)

# Using this polynomial approach leads to the equation
# ```math
# \frac{dx}{2} \dot{u}^{Q_l}(\xi, t) + u^{Q_l}(\xi, t)' = 0
# ```
# with $\dot{u}=\frac{\partial}{\partial t}u$ and $u'=\frac{\partial}{\partial x}u$.
# To approximate the solution, we need to get the polynomial coefficients $\{u_j^{Q_l}\}_{j=0}^N$
# for every element $Q_l$.

# After defining all nodes, we can implement the spatial coordinate $x$ and its initial value $u0$
# for every node.
x = Matrix{Float64}(undef, length(nodes), n_elements)
for element in 1:n_elements
    x_l = coordinates_min + (element - 1) * dx + dx/2
    for i in 1:length(nodes)
        ξ = nodes[i] # nodes in [-1, 1]
        x[i, element] = x_l + dx/2 * ξ
    end
end

u0 = initial_condition_sine_wave.(x)

# To have a look at the initial sinus curve, we plot it.
using Plots
plot(vec(x), vec(u0), label="initial condition", legend=:topleft)

# ### iii. Variational formulation
# After defining the equation and initial condition, we want to implement an algorithm to
# approximate the solution.

# From now on, we only write $u$ instead of $u^{Q_l}$ for simplicity, but consider that all the following
# calculation only concern one element.
# Multiplying the new equation with the smooth Lagrange polynomials $\{l_i\}_{i=0}^N$ (test functions)
# and integrating over the reference element $E=[-1,1]$, we get the variational formulation of our
# transformed partial differential equation for $i=0,...,N$:
# ```math
# \begin{align*}
# \int_{-1}^1 \Big( \frac{dx}{2} \dot{u}(\xi, t) + u'(\xi, t) \Big) l_i(\xi)d\xi
#   &= \underbrace{\frac{dx}{2} \int_{-1}^1 \dot{u}(\xi, t) l_i(\xi)d\xi}_{\text{Term I}} + \underbrace{\int_{-1}^1 u'(\xi, t) l_i(\xi)d\xi}_{\text{Term II}} = 0
# \end{align*}
# ```

# We deal with the two terms separately. We write $\int_{-1, N}^1 \;\cdot\; d\xi$ for the approximation
# of the integral using numerical quadrature with $N+1$ basis points. We use the Gauss-Lobatto nodes
# again. The numerical scalar product $\langle\cdot, \cdot\rangle_N$ is defined by
# $\langle f, g\rangle_N := \int_{-1, N}^1 f(\xi) g(\xi) d\xi$.

# #### Term I:
# In the following calculation we approximate the integral numerically with quadrature on the Gauss-Lobatto
# nodes $\{\xi_i\}_{i=0}^N$ and then use the Kronecker property of the Lagrange polynomials. This approach
# of using the same nodes for the interpolation and quadrature is called collocation.
# ```math
# \begin{align*}
# \frac{dx}{2} \int_{-1}^1 \dot{u}(\xi, t) l_i(\xi)d\xi
# &\approx \frac{dx}{2} \int_{-1, N}^1 \dot{u}(\xi, t) l_i(\xi)d\xi \\
# &= \frac{dx}{2} \sum_{k=0}^N \underbrace{\dot{u}(\xi_k, t)}_{=\dot{u}_k(t)} \underbrace{l_i(\xi_k)}_{=\delta_{k,i}}w_k \\
# &= \frac{dx}{2} \dot{u}_i(t) w_i
# \end{align*}
# ```
# We define the Legendre-Gauss-Lobatto (LGL) mass matrix $M$ and by the Kronecker property follows:
# ```math
# M_{ij} = \langle l_j, l_i\rangle_N = \delta_{ij} w_j,\; i,j=0,...,N.
# ```
using LinearAlgebra
M = diagm(weights)
# Now, we can write the integral with this new matrix.
# ```math
# \frac{dx}{2} \int_{-1, N}^1 \dot{u}(\xi, t) \underline{l}(\xi)d\xi = \frac{dx}{2} M \underline{\dot{u}}(t),
# ```
# where $\underline{\dot{u}} = (\dot{u}_0, ..., \dot{u}_N)^T$ and $\underline{l}$ respectively.

# **Note:** Since the LGL quadrature with $N+1$ nodes is exact up to functions of degree $2N-1$ and
# $\dot{u}(\xi, t) l_i(\xi)$ is of degree $2N$, in general the following holds
# ```math
# \int_{-1}^1 \dot{u}(\xi, t) l_i(\xi) d\xi \neq \int_{-1, N}^1 \dot{u}(\xi, t) l_i(\xi) d\xi.
# ```
# With an exact integration the mass matrix would be dense. Choosing numerical integrating and quadrature
# with the exact same nodes (collocation) leads to the sparse and diagonal mass matrix $M$. This
# is called mass lumping and has the big advantage of an easy invertation of the matrix.

# #### Term II:
# We use spatial partial intergration for the second term:
# ```math
# \int_{-1}^1 u'(\xi, t) l_i(\xi) d\xi = [u l_i]_{-1}^1 - \int_{-1}^1 u l_i'd\xi
# ```
# The resulting integral can be solved exactly with LGL quadrature since the polynomial is now
# of degree $2N-1$.

# Again, we split the calculation in two steps.

# #### Surface term
# As mentioned before, we approximate the solution with a polynomial in every element. Therefore, in
# general the value of this approximation at the interfaces between two elements is not unique. To solve
# this problem we introduce the idea of the numerical flux $u^*$, which will give an exact value at
# the interfaces. One of many different approaches and definitions for the calculation of the
# numerical flux we will deal with in [4. Numerical flux](@ref numerical_flux).
# ```math
# [u l_i]_{-1}^1 = u^*\big|^1 l_i(+1) - u^*\big|_{-1} l_i(-1)
# ```
# Since the Gauss-Lobatto nodes contain the element boundaries $-1$ and $+1$, we can use the
# Kronecker property of $l_i$ for the calculation of $l_i(-1)$ and $l_i(+1)$.
# ```math
# [u \underline{l}]_{-1}^1 = u^*\big|^1 \left(\begin{array}{c} 0 \\ \vdots \\ 0 \\ 1 \end{array}\right)
# - u^*\big|_{-1} \left(\begin{array}{c} 1 \\ 0 \\ \vdots \\ 0\end{array}\right)
# = B \underline{u}^*(t)
# ```
# with the boundary matrix
# ```math
# B = \begin{pmatrix}
# -1 & 0 & \cdots & 0\\
# 0 & 0 & \cdots & 0\\
# \vdots & \vdots & 0 & 0\\
# 0 & \cdots & 0 & 1
# \end{pmatrix}
# \qquad\text{and}\qquad
# \underline{u}^*(t) = \left(\begin{array}{c} u^*\big|_{-1} \\ 0 \\ \vdots \\ 0 \\ u^*\big|^1\end{array}\right).
# ```
B = diagm([-1; zeros(polydeg - 1); 1])

# #### Volume term
# As mentioned before, the new integral can be solved exact since the function inside is of degree $2N-1$.
# ```math
# - \int_{-1}^1 u l_i'd\xi = - \int_{-1, N}^1 u l_i' d\xi
# = - \sum_{k=0}^N u(\xi_k, t) l_i'(\xi_k) w_k
# = - \sum_{k=0}^N u_k(t) D_{ki} w_k
# ```
# where $D$ is the derivative matrix defined by $D_{ki} = l_i'(\xi_k)$ for $i,k=0,...,N$.
D = basis.derivative_matrix

# To show why this matrix is called the derivative matrix, we go back to our example $f(x)=x^3$.
# We calculate the derivation of $f$ at the Gauss-Lobatto nodes $\{\xi_k\}_{k=0}^N$ with $N=8$.
# ```math
# f'|_{x=\xi_k} = \Big( \sum_{j=0}^8 f(\xi_j) l_j(x) \Big)'|_{x=\xi_k} = \sum_{j=0}^8 f(\xi_j) l_j'(\xi_k)
# = \sum_{j=0}^8 f(\xi_j) D_{kj}
# ```
# for $k=0,...,N$ and therefore, $\underline{f}' = D \underline{f}$.
basis_N8 = LobattoLegendreBasis(8)
plot(vec(x), x -> 3 * x^2, label="f'", lw=2)
scatter!(basis_N8.nodes, basis_N8.derivative_matrix * basis_N8.nodes.^3, label="Df", lw=3)

# Combining the volume term for every $i=0,...,N$ results in
# ```math
# \int_{-1}^1 u \underline{l'} d\xi = - D^T M \underline{u}(t)
# ```

# Putting all parts together we get the following equation for the element $Q_l$
# ```math
# \frac{dx}{2} M \underline{\dot{u}}(t) = - B \underline{u}^*(t) + D^T M \underline{u}(t)
# ```
# or equivalent
# ```math
# \underline{\dot{u}}^{Q_l}(t) = \frac{2}{dx} \Big[ - M^{-1} B \underline{u}^{{Q_l}^*}(t) + M^{-1} D^T M \underline{u}^{Q_l}(t)\Big].
# ```
# This is called the weak form of the DGSEM.

# **Note:** For every element $Q_l$ we get a system of $N+1$ ordinary differential equations to
# calculate $N+1$ coefficients. Since the numerical flux $u^*$ is depending on extern values at
# the interfaces, the equation systems of adjacent elements are weakly linked.

# ### [iv. Numerical flux](@id numerical_flux)
# As mentioned above, we still have to handle the problem of different values at the same point at
# the interfaces. This happens with the ideas of the numerical flux $f^*(u)=u^*$. The role of $f^*$
# might seem minor in this simple example, but is important for more complicated problems.
# There are two values at the same spatial coordinate. Let's say we are looking at the interface between
# the elements $Q_l$ and $Q_{l+1}$, while both elements got $N+1$ nodes as defined before. We call
# the first value of the right element $u_R=u_0^{Q_{l+1}}$ and the last one of the left element
# $u_L=u_N^{Q_l}$. So, for the value of the numerical flux on that interface the following holds
# ```math
# u^* = u^*(u_L, u_R).
# ```
# These values are interpreted as start values of a so-called Riemann problem. There are many
# different (approximate) Riemann solvers available and useful for different problems. We will
# use the local Lax-Friedrichs flux.
surface_flux = flux_lax_friedrichs

# The only missing ingredient is the flux calculation at the boundaries $-1$ and $+1$.
# ```math
# u^{{Q_{first}}^*}\big|_{-1} = u^{{Q_{first}}^*}\big|_{-1}(u^{bound}(-1), u_R)
# \quad\text{and}\quad
# u^{{Q_{last}}^*}\big|^1 = u^{{Q_{last}}^*}\big|^1(u_L, u^{bound}(1))
# ```
# The boundaries are periodic, which means that the last value of the last element $u^{Q_{last}}_N$
# is used as $u_L$ at the first interface and accordingly for the other boundary.

# Now, we implement a function, that calculates $\underline{\dot{u}}^{Q_l}$ for the given matrices,
# $\underline{u}$ and $\underline{u}^*$.
function rhs!(du, u, x, t)
    ## Reset du and flux matrix
    du .= zero(eltype(du))
    flux_numerical = copy(du)

    ## Calculate interface and boundary fluxes, $u^* = (u^*|_{-1}, 0, ..., 0, u^*|^1)^T$
    ## Since we use the flux Lax-Friedrichs from Trixi.jl, we have to pass some extra arguments.
    ## Trixi needs the equation we are dealing with and an additional `1`, that indicates the
    ## first coordinate direction.
    equations = LinearScalarAdvectionEquation1D(1.0)
    for element in 2:n_elements-1
        ## left interface
        flux_numerical[1, element] = surface_flux(u[end, element-1], u[1, element], 1, equations)
        flux_numerical[end, element-1] = flux_numerical[1, element]
        ## right interface
        flux_numerical[end, element] = surface_flux(u[end, element], u[1, element+1], 1, equations)
        flux_numerical[1, element+1] = flux_numerical[end, element]
    end
    ## boundary flux
    flux_numerical[1, 1] = surface_flux(u[end, end], u[1, 1], 1, equations)
    flux_numerical[end, end] = flux_numerical[1, 1]

    ## Calculate surface integrals, $- M^{-1} * B * u^*$
    for element in 1:n_elements
        du[:, element] -= (M \ B) * flux_numerical[:, element]
    end

    ## Calculate volume integral, $+ M^{-1} * D^T * M * u$
    for element in 1:n_elements
        flux = u[:, element]
        du[:, element] += (M \ transpose(D)) * M * flux
    end

    ## Apply Jacobian from mapping to reference element
    for element in 1:n_elements
        du[:, element] *= 2 / dx
    end

    return nothing
end

# Combining all definitions and the function that calculates the right-hand side, we define the ODE and
# solve it until `t=2` with OrdinaryDiffEq's `solve` function and the Runge-Kutta method `RDPK3SpFSAL49()`,
# which is optimized for discontinuous Galerkin methods and hyperbolic PDEs. We set some common
# error tolerances `abstol=1.0e-6, reltol=1.0e-6` and pass `save_everystep=false` to avoid saving intermediate
# solution vectors in memory.
using OrdinaryDiffEq
tspan = (0.0, 2.0)
ode = ODEProblem(rhs!, u0, tspan, x)

sol = solve(ode, RDPK3SpFSAL49(), abstol=1.0e-6, reltol=1.0e-6, save_everystep=false)
@test maximum(abs.(u0 - sol.u[end])) < 5e-5 #src

plot(vec(x), vec(sol.u[end]), label="solution at t=$(tspan[2])", legend=:topleft, lw=3)

# ## Alternative Implementation based on Trixi.jl
# Now, we implement the same example. But this time, we directly use the functionality that Trixi.jl
# provides.

using Trixi, OrdinaryDiffEq, Plots

# First, define the equation with a advection_velocity of `1`.
advection_velocity = 1.0
equations = LinearScalarAdvectionEquation1D(advection_velocity)

# Then, create a DG solver with polynomial degree = 3 and (local) Lax-Friedrichs/Rusanov flux as surface flux.
# The implementation of the basis and the numerical flux is now already done.
solver = DGSEM(polydeg=3, surface_flux=flux_lax_friedrichs)

# We will now create a mesh with 16 elements for the physical domain `[-1, 1]` with periodic boundaries.
# We use Trixi.jl's standard mesh [`TreeMesh`](@ref). Since it's limited to hypercube domains, we
# choose `2^4=16` elements. The mesh type supports AMR, that' why `n_cells_max` has to be set, even
# if we don't need AMR here.
coordinates_min = -1.0 # minimum coordinate
coordinates_max = 1.0  # maximum coordinate
mesh = TreeMesh(coordinates_min, coordinates_max,
                initial_refinement_level=4, # number of elements = 2^4
                n_cells_max=30_000) # set maximum capacity of tree data structure (only needed for AMR)

# A semidiscretization collects data structures and functions for the spatial discretization.
# In Trixi, an initial condition has the following parameter structure and is of the type `SVector`.
initial_condition_sine_wave(x, t, equations) = SVector(1.0 + 0.5 * sin(pi * sum(x - equations.advection_velocity * t)))

semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition_sine_wave, solver)

# Again, combining all definitions and the function that calculates the right-hand side, we define the ODE and
# solve it until `t=2` with OrdinaryDiffEq's `solve` function and the Runge-Kutta method `RDPK3SpFSAL49()`.
tspan = (0.0, 2.0)
ode_trixi  = semidiscretize(semi, tspan)

sol_trixi = solve(ode_trixi, RDPK3SpFSAL49(), abstol=1.0e-6, reltol=1.0e-6, save_everystep=false);

# We add a plot of the new approximated solution to the one calculated before.
plot!(sol_trixi, label="solution at t=$(tspan[2]) with Trixi.jl", legend=:topleft, linestyle=:dash, lw=2)
@test maximum(abs.(vec(u0) - sol_trixi.u[end])) ≈ maximum(abs.(u0 - sol.u[end])) #src

# ## Summary of the code
# To sum up, here is the complete code that we used.

# ### Raw implementation
## basis: Legendre-Gauss-Lobatto
using Trixi, LinearAlgebra, OrdinaryDiffEq, Plots
polydeg = 3 #= polynomial degree =#
basis = LobattoLegendreBasis(polydeg)
nodes = basis.nodes # Gauss-Lobatto nodes in [-1, 1]
D = basis.derivative_matrix
M = diagm(basis.weights) # mass matrix
B = diagm([-1; zeros(polydeg - 1); 1])

## mesh
coordinates_min = -1.0 # minimum coordinate
coordinates_max = 1.0  # maximum coordinate
n_elements      = 16   # number of elements

dx = (coordinates_max - coordinates_min) / n_elements # length of one element

x = Matrix{Float64}(undef, length(nodes), n_elements)
for element in 1:n_elements
    x_l = -1 + (element - 1) * dx + dx/2
    for i in 1:length(nodes) # basis points in [-1, 1]
        ξ = nodes[i]
        x[i, element] = x_l + dx/2 * ξ
    end
end

## initial condition
initial_condition_sine_wave(x) = 1.0 + 0.5 * sin(pi * x)
u0 = initial_condition_sine_wave.(x)

plot(vec(x), vec(u0), label="initial condition", legend=:topleft)

## flux Lax-Friedrichs
surface_flux = flux_lax_friedrichs

## rhs! method
function rhs!(du, u, x, t)
    ## reset du
    du .= zero(eltype(du))
    flux_numerical = copy(du)

    ## calculate interface and boundary fluxes
    equations = LinearScalarAdvectionEquation1D(1.0)
    for element in 2:n_elements-1
        ## left interface
        flux_numerical[1, element] = surface_flux(u[end, element-1], u[1, element], 1, equations)
        flux_numerical[end, element-1] = flux_numerical[1, element]
        ## right interface
        flux_numerical[end, element] = surface_flux(u[end, element], u[1, element+1], 1, equations)
        flux_numerical[1, element+1] = flux_numerical[end, element]
    end
    ## boundary flux
    flux_numerical[1, 1] = surface_flux(u[end, end], u[1, 1], 1, equations)
    flux_numerical[end, end] = flux_numerical[1, 1]

    ## calculate surface integrals
    for element in 1:n_elements
        du[:, element] -= (M \ B) * flux_numerical[:, element]
    end

    ## calculate volume integral
    for element in 1:n_elements
        flux = u[:, element]
        du[:, element] += (M \ transpose(D)) * M * flux
    end

    ## apply Jacobian from mapping to reference element
    for element in 1:n_elements
        du[:, element] *= 2 / dx
    end

    return nothing
end

## create ODE problem
tspan = (0.0, 2.0)
ode = ODEProblem(rhs!, u0, tspan, x)

## solve
sol = solve(ode, RDPK3SpFSAL49(), abstol=1.0e-6, reltol=1.0e-6, save_everystep=false)
@test maximum(abs.(vec(u0) - sol_trixi.u[end])) ≈ maximum(abs.(u0 - sol.u[end])) #src

plot(vec(x), vec(sol.u[end]), label="solution at t=$(tspan[2])", legend=:topleft, lw=3)

# ### Alternative Implementation based on Trixi.jl
using Trixi, OrdinaryDiffEq, Plots

## equation with a advection_velocity of `1`.
advection_velocity = 1.0
equations = LinearScalarAdvectionEquation1D(advection_velocity)

## create DG solver with flux lax friedrichs and LGL basis
solver = DGSEM(polydeg=3, surface_flux=flux_lax_friedrichs)

## distretize domain with `TreeMesh`
coordinates_min = -1.0 # minimum coordinate
coordinates_max = 1.0 # maximum coordinate
mesh = TreeMesh(coordinates_min, coordinates_max,
                initial_refinement_level=4, # number of elements = 2^4
                n_cells_max=30_000)

## create initial condition and semidiscretization
initial_condition_sine_wave(x, t, equations) = SVector(1.0 + 0.5 * sin(pi * sum(x - equations.advection_velocity * t)))

semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition_sine_wave, solver)

## solve
tspan = (0.0, 2.0)
ode_trixi  = semidiscretize(semi, tspan)
sol_trixi  = solve(ode_trixi, RDPK3SpFSAL49(), abstol=1.0e-6, reltol=1.0e-6, save_everystep=false);

plot!(sol_trixi, label="solution at t=$(tspan[2]) with Trixi.jl", legend=:topleft, linestyle=:dash, lw=2)
@test maximum(abs.(vec(u0) - sol_trixi.u[end])) ≈ maximum(abs.(u0 - sol.u[end])) #src


import Optim.optimize
import Base: axpy!, scale!

export frank_wolfe_sketched

# type FrankWolfeParams<:OptParams
# 	maxiters::Int
# 	abstol::AbstractFloat
# 	stepsize::StepSizeRule
# end
# stop(p::FrankWolfeParams, eps) = eps < p.abstol ? true : false
# FrankWolfeParams() = FrankWolfeParams(50, 1e-10, DecreasingStepSize(2,1))

function frank_wolfe_sketched(z::AbstractArray, # starting point
	objective::Function,
	grad_objective::Function,
	constraint::Function, # evaluates constraint function
	alpha::Number, # bound on constraint function
	min_lin_st_constraint::Function, # solves min_{c(x)<=alpha} g \dot x (as a function of alpha and g)
	sketch::AbstractSketch = IdentitySketch(x),
	params::FrankWolfeParams = FrankWolfeParams(),
	ch::ConvergenceHistory = ConvergenceHistory("fw_sketched");
	LB = -Inf,
	UB = Inf,
	verbose = false)

	# initialize
	if verbose @printf("%10s%12s%12s%12s%12s%12s\n", "iteration", "UB", "LB", "gap", "rel gap", "time") end

	t = time()
	t0 = copy(t)
	for k=1:params.maxiters

		# function and gradient evaluation
		objval = objective(z)
		G = grad_objective(z)
		# G = A'*g, where g is the gradient wrt the compressed variable
		# so G.factors[1] is the compression operator A
		A = G.factors[1]
		g = vec(G.factors[2])

		# solve fenchel problem
		# Delta is solution to min_{c(X)<=alpha} grad(X_old) \dot X
		Delta, linearized_obj = min_lin_st_constraint(G, alpha)
		dz = A*Delta

		# check stopping condition
		UB = min(UB, objval)
 	   # <G, Delta - X> = <A'g, Delta - X> = <g, A*Delta - A*X> = <g, dz - z>
		LB = max(LB, objval + linearized_obj - dot(g, z))
		old_t, t = t, time()
		if verbose && k%10==0
			cx = constraint(z)
			@printf("%10d%12.4e%12.4e%12.4e%12.4e%12.4e\n", k, UB, LB, UB - LB, (UB - LB)/min(abs(LB), abs(UB)), t - t0)
		end
		update_ch!(ch, t - old_t; obj = UB, dual_obj = LB)
		stop(params, UB, LB) && break
		t = time()

		# choose step size with stepsize rule
		# eg a = 1/k
		a = step(params.stepsize, objective, z, z-dz;
			     objval=objval,
			     normgradsq=UB-LB)

		# take step
		# z = (1-a)*z + a*dz
		scale!(z, 1-a); axpy!(a, dz, z)

		cgd_update!(sketch, Delta, a)
	end
	return reconstruct(sketch)
end