@add_problem affine function affine_negate_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable()
    p = minimize(-x, [x <= 0]; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 0 atol=atol rtol=rtol
        @test evaluate(-x) ≈ 0 atol=atol rtol=rtol
    end
end

@add_problem affine function affine_kron_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = ComplexVariable(3, 3)
    y = [1.0 2.0; 3.0 4.0]
    if test
        @test size(kron(x, y)) == (6, 6)
        @test size(kron(y, x)) == (6, 6)
    end
end

@add_problem affine function affine_multiply_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(1)
    p = minimize(2.0 * x, [x >= 2, x <= 4]; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 4 atol=atol rtol=rtol
        @test (evaluate(2.0x))[1] ≈ 4 atol=atol rtol=rtol
    end

    x = Variable(2)
    A = 1.5 * eye(2)
    p = minimize([2 2] * x, [A * x >= [1.1; 1.1]]; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 2.93333 atol=atol rtol=rtol
        @test (evaluate([2 2] * x))[1] ≈ 2.93333 atol=atol rtol=rtol
        @test vec(evaluate(A * x)) ≈ [1.1; 1.1] atol=atol rtol=rtol
    end

    y = Variable(1)
    x = Variable(3)
    z = [1.0, 2.0, 3.0] * y
    k = -y * [1.0, 2.0, 3.0]
    c = [y <= 3.0, y >= 0.0, x >= ones(3), k <= x, x <= z]
    o = 3 * y
    p = Problem{T}(:minimize, o, c)
    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 3 atol=atol rtol=rtol
    end

    p = Problem{T}(:minimize, o, c...)
    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 3 atol=atol rtol=rtol
    end

    # Check #274
    x = ComplexVariable(2,2)
    p = minimize( real( [1.0im, 0.0]' * x * [1.0im, 0.0] ), [ x == [1.0 0.0; 0.0 1.0] ]; numeric_type = T)

    handle_problem!(p)
    if test
        @test p.optval ≈ 1.0 atol=atol rtol=rtol
    end
end

@add_problem affine function affine_dot_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(2)
    p = minimize(dot([2.0; 2.0], x), x >= [1.1; 1.1]; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 4.4 atol=atol rtol=rtol
        @test (evaluate(dot([2.0; 2.0], x)))[1] ≈ 4.4 atol=atol rtol=rtol
    end
end

@add_problem affine function affine_dot_atom_for_matrix_variables(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(2,2)
    p = minimize(dot(fill(2.0, (2,2)), x), x >= 1.1; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 8.8 atol=atol rtol=rtol
        @test (evaluate(dot(fill(2.0, (2, 2)), x)))[1] ≈ 8.8 atol=atol rtol=rtol
    end
end

@add_problem affine function affine_add_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(1)
    y = Variable(1)
    p = minimize(x + y, [x >= 3, y >= 2]; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 5 atol=atol rtol=rtol
        @test evaluate(x + y) ≈ 5 atol=atol rtol=rtol
    end

    x = Variable(1)
    p = minimize(x, [eye(2) + x >= eye(2)]; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 0 atol=atol rtol=rtol
        @test evaluate(eye(2) + x) ≈ eye(2) atol=atol rtol=rtol
    end

    y = Variable()
    p = minimize(y - 5, y >= -1; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ -6 atol=atol rtol=rtol
        @test evaluate(y - 5) ≈ -6 atol=atol rtol=rtol
    end
end

@add_problem affine function affine_transpose_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(2)
    c = ones(2, 1)
    p = minimize(x' * c, x >= 1; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 2 atol=atol rtol=rtol
        @test (evaluate(x' * c))[1] ≈ 2 atol=atol rtol=rtol
    end

    X = Variable(2, 2)
    c = ones(2, 1)
    p = minimize(c' * X' * c, [X >= ones(2, 2)]; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 4 atol=atol rtol=rtol
        @test (evaluate(c' * X' * c))[1] ≈ 4 atol=atol rtol=rtol
    end

    rows = 2
    cols = 3
    r = rand(rows, cols)
    r_2 = rand(cols, rows)
    x = Variable(rows, cols)
    c = ones(1, cols)
    d = ones(rows, 1)
    p = minimize(c * x' * d + d' * x * c' + (c * x''''' * d)',
                [x' >= r_2, x >= r, x''' >= r_2, x'' >= r]; numeric_type = T)
    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    s = sum(max.(r, r_2')) * 3
    if test
        @test p.optval ≈ s atol=atol rtol=rtol
        @test (evaluate(c * x' * d + d' * x * c' + (c * ((((x')')')')' * d)'))[1] ≈ s atol=atol rtol=rtol
    end
end

@add_problem affine function affine_index_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(2)
    p = minimize(x[1] + x[2], [x >= 1]; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 2 atol=atol rtol=rtol
        @test (evaluate(x[1] + x[2]))[1] ≈ 2 atol=atol rtol=rtol
    end

    x = Variable(3)
    I = [true true false]
    p = minimize(sum(x[I]), [x >= 1]; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 2 atol=atol rtol=rtol
        @test (evaluate(sum(x[I])))[1] ≈ 2 atol=atol rtol=rtol
    end

    rows = 6
    cols = 8
    n = 2
    X = Variable(rows, cols)
    A = randn(rows, cols)
    c = rand(1, n)
    p = minimize(c * X[1:n, 5:5+n-1]' * c', X >= A; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    s = c * A[1:n, 5:5+n-1]' * c'
    if test
        @test p.optval ≈ s[1] atol=atol rtol=rtol
        @test evaluate(c * (X[1:n, 5:(5 + n) - 1])' * c') ≈ s atol=atol rtol=rtol
    end
end

@add_problem affine function affine_sum_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(2,2)
    p = minimize(sum(x), x>=1; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 4 atol=atol rtol=rtol
        @test evaluate(sum(x)) ≈ 4 atol=atol rtol=rtol
    end

    x = Variable(2,2)
    p = minimize(sum(x) - 2*x[1,1], x>=1, x[1,1]<=2; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 1 atol=atol rtol=rtol
        @test (evaluate(sum(x) - 2 * x[1, 1]))[1] ≈ 1 atol=atol rtol=rtol
    end

    x = Variable(10)
    a = rand(10, 1)
    p = maximize(sum(x[2:6]), x <= a; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ sum(a[2:6]) atol=atol rtol=rtol
        @test evaluate(sum(x[2:6])) ≈ sum(a[2:6]) atol=atol rtol=rtol
    end
end

@add_problem affine function affine_diag_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(2,2)
    p = minimize(sum(diag(x,1)), x >= 1; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 1 atol=atol rtol=rtol
        @test evaluate(sum(diag(x, 1))) ≈ 1 atol=atol rtol=rtol
    end

    x = Variable(4, 4)
    p = minimize(sum(diag(x)), x >= 2; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 8 atol=atol rtol=rtol
        @test evaluate(sum(diag(x))) ≈ 8 atol=atol rtol=rtol
    end
end

@add_problem affine function affine_trace_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(2,2)
    p = minimize(tr(x), x >= 1; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 2 atol=atol rtol=rtol
        @test evaluate(tr(x)) ≈ 2 atol=atol rtol=rtol
    end
end

@add_problem affine function affine_dot_multiply_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(3)
    p = maximize(sum(dot(*)(x,[1,2,3])), x<=1; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 6 atol=atol rtol=rtol
        @test evaluate(sum((dot(*))(x, [1, 2, 3]))) ≈ 6 atol=atol rtol=rtol
    end

    x = Variable(3, 3)
    p = maximize(sum(dot(*)(x,eye(3))), x<=1; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 3 atol=atol rtol=rtol
        @test evaluate(sum((dot(*))(x, eye(3)))) ≈ 3 atol=atol rtol=rtol
    end

    x = Variable(5, 5)
    p = minimize(x[1, 1], dot(*)(3,x) >= 3; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 1 atol=atol rtol=rtol
        @test (evaluate(x[1, 1]))[1] ≈ 1 atol=atol rtol=rtol
    end

    x = Variable(3,1)
    p = minimize(sum(dot(*)(ones(3,3), x)), x>=1; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 9 atol=atol rtol=rtol
        @test (evaluate(x[1, 1]))[1] ≈ 1 atol=atol rtol=rtol
    end

    x = Variable(1,3)
    p = minimize(sum(dot(*)(ones(3,3), x)), x>=1; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 9 atol=atol rtol=rtol
        @test (evaluate(x[1, 1]))[1] ≈ 1 atol=atol rtol=rtol
    end

    x = Variable(1, 3, Positive())
    p = maximize(sum(dot(/)(x,[1 2 3])), x<=1; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 11 / 6 atol=atol rtol=rtol
        @test evaluate(sum((dot(/))(x, [1 2 3]))) ≈ 11 / 6 atol=atol rtol=rtol
    end

    # Broadcast fusion works
    x = Variable(5, 5)
    a = 2.0 .* x .* ones(Int, 5)
    if test
        @test a isa Convex.DotMultiplyAtom
    end
end

@add_problem affine function affine_reshape_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    A = rand(2, 3)
    X = Variable(3, 2)
    c = rand()
    p = minimize(sum(reshape(X, 2, 3) + A), X >= c; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ sum(A .+ c) atol=atol rtol=rtol
        @test evaluate(sum(reshape(X, 2, 3) + A)) ≈ sum(A .+ c) atol=atol rtol=rtol
    end

    b = rand(6)
    p = minimize(sum(vec(X) + b), X >= c; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ sum(b .+ c) atol=atol rtol=rtol
        @test evaluate(sum(vec(X) + b)) ≈ sum(b .+ c) atol=atol rtol=rtol
    end

    x = Variable(4, 4)
    c = ones(16, 1)
    reshaped = reshape(x, 16, 1)
    a = collect(1:16)
    p = minimize(c' * reshaped, reshaped >= a; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    # TODO: why is accuracy lower here?
    if test
        @test p.optval ≈ 136 atol=10atol atol=atol rtol=rtol
        @test (evaluate(c' * reshaped))[1] ≈ 136 atol=10atol atol=atol rtol=rtol
    end
end

@add_problem affine function affine_hcat_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(4, 4)
    y = Variable(4, 6)
    p = maximize(sum(x) + sum([y fill(4.0, 4)]), [x y fill(2.0, (4, 2))] <= 2; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 96 atol=atol rtol=rtol
        @test evaluate(sum(x) + sum([y fill(4.0, 4)])) ≈ 96 atol=atol rtol=rtol
        @test evaluate([x y fill(2.0, (4, 2))]) ≈ fill(2.0, (4, 12)) atol=atol rtol=rtol
    end
end

@add_problem affine function affine_vcat_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(4, 4)
    y = Variable(4, 6)

    # TODO: fix dimension mismatch [y 4*eye(4); x -ones(4, 6)]
    p = maximize(sum(x) + sum([y 4*eye(4); x -ones(4, 6)]), [x;y'] <= 2; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    # TODO: why is accuracy lower here?
    if test
        @test p.optval ≈ 104 atol=10atol atol=atol rtol=rtol
        @test evaluate(sum(x) + sum([y 4 * eye(4); x -(ones(4, 6))])) ≈ 104 atol=10atol atol=atol rtol=rtol
        @test evaluate([x; y']) ≈ 2 * ones(10, 4) atol=atol rtol=rtol
    end
end

@add_problem affine function affine_single_hcat_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(4, 4)
    p = maximize(tr(hcat(x)), hcat(x) <= 2; numeric_type = T)

    handle_problem!(p)
    if test
        @test p.optval ≈ 8 atol=atol rtol=rtol
    end
end

@add_problem affine function affine_single_vcat_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(4, 4)
    p = maximize(tr(vcat(x)), vcat(x) <= 2; numeric_type = T)

    handle_problem!(p)
    if test
        @test p.optval ≈ 8 atol=atol rtol=rtol
    end
end

@add_problem affine function affine_Diagonal_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(2, 2)
    if test
        @test_throws ArgumentError Diagonal(x)
    end

    x = Variable(4)
    p = minimize(sum(Diagonal(x)), x == [1, 2, 3, 4]; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 10 atol=atol rtol=rtol
        @test all(abs.(evaluate(Diagonal(x)) - Diagonal([1, 2, 3, 4])) .<= atol)
    end

    x = Variable(3)
    c = [1; 2; 3]
    p = minimize(c' * Diagonal(x) * c, x >= 1, sum(x) == 10; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 21 atol=atol rtol=rtol
    end

    x = Variable(3)
    p = minimize(sum(x), x >= 1, Diagonal(x)[1, 2] == 1; numeric_type = T)

    handle_problem!(p)
    if test
        @test p.status != MOI.OPTIMAL
    end
end

@add_problem affine function affine_conv_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(3)
    h = [1, -1]
    p = minimize(sum(conv(h, x)) + sum(x), x >= 1, x <= 2; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 3 atol=atol rtol=rtol
        @test evaluate(sum(conv(h, x))) ≈ 0 atol=atol rtol=rtol
    end

    x = Variable(3)
    h = [1, -1]
    p = minimize(sum(conv(x, h)) + sum(x), x >= 1, x <= 2; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 3 atol=atol rtol=rtol
        @test evaluate(sum(conv(h, x))) ≈ 0 atol=atol rtol=rtol
    end

end

@add_problem affine function affine_satisfy_problems(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable()
    p = satisfy(x >= 0; numeric_type = T)

    add_constraints!(p, x >= 1)
    add_constraints!(p, [x >= -1, x <= 4])
    handle_problem!(p)
    if test
        @test p.status == MOI.OPTIMAL
    end

    p = satisfy([x >= 0, x >= 1, x <= 2]; numeric_type = T)

    handle_problem!(p)
    if test
        @test p.status == MOI.OPTIMAL
    end

    p = maximize(1, [x >= 1, x <= 2]; numeric_type = T)

    handle_problem!(p)
    if test
        @test p.status == MOI.OPTIMAL
    end

    constr = x >= 0
    constr += x >= 1
    constr += x <= 10
    constr2 = x >= 0
    constr2 += [x >= 2, x <= 3] + constr
    p = satisfy(constr; numeric_type = T)

    handle_problem!(p)
    if test
        @test p.status == MOI.OPTIMAL
    end
end

@add_problem affine function affine_dualvalue(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable()
    p = minimize(x, x >= 0; numeric_type = T)

    handle_problem!(p)
    if test
        @test p.constraints[1].dual ≈ 1 atol=atol rtol=rtol
    end

    x = Variable()
    p = maximize(x, x <= 0; numeric_type = T)

    handle_problem!(p)
    if test
        @test p.constraints[1].dual ≈ 1 atol=atol rtol=rtol
    end

    x = Variable()
    p = minimize(x, x >= 0, x == 2; numeric_type = T)

    handle_problem!(p)
    if test
        @test p.constraints[1].dual ≈ 0 atol=atol rtol=rtol
        @test abs.(p.constraints[2].dual) ≈ 1 atol=atol rtol=rtol
    end

    x = Variable(2)
    A = 1.5 * eye(2)
    p = minimize(dot([2.0; 2.0], x), [A * x >= [1.1; 1.1]]; numeric_type = T)

    handle_problem!(p)
    if test
        dual = [4/3; 4/3]
        @test all(abs.(p.constraints[1].dual - dual) .<= atol)
    end
end

@add_problem affine function affine_Partial_transpose(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    dims = [2,3,4]
    d = prod(dims)
    A = rand(ComplexF64,2,2)
    B = rand(ComplexF64,3,3)
    C = rand(ComplexF64,4,4)
    M = kron(A,B,C)
    Mt1 = kron(transpose(A),B,C)
    Mt2 = kron(A,transpose(B),C)
    Mt3 = kron(A,B,transpose(C))

    Rt1 = ComplexVariable(d,d)
    Rt2 = ComplexVariable(d,d)
    Rt3 = ComplexVariable(d,d)
    S = rand(ComplexF64,d,d)
    handle_problem!(satisfy(partialtranspose(Rt1, 1, dims) == S; numeric_type = T))

    handle_problem!(satisfy(partialtranspose(Rt2, 2, dims) == S; numeric_type = T))

    handle_problem!(satisfy(partialtranspose(Rt3, 3, dims) == S; numeric_type = T))

    if test
        @test partialtranspose(M,1,dims) ≈ Mt1 atol=atol rtol=rtol
        @test partialtranspose(M,2,dims) ≈ Mt2 atol=atol rtol=rtol
        @test partialtranspose(M,3,dims) ≈ Mt3 atol=atol rtol=rtol
        @test partialtranspose(S,1,dims) ≈ evaluate(Rt1) atol=atol rtol=rtol
        @test partialtranspose(S,2,dims) ≈ evaluate(Rt2) atol=atol rtol=rtol
        @test partialtranspose(S,3,dims) ≈ evaluate(Rt3) atol=atol rtol=rtol
    end

    if test
        @test_throws ArgumentError partialtrace(rand(6, 6), 3, [2, 3])
        @test_throws ArgumentError partialtrace(rand(6, 6), 1, [2, 4])
        @test_throws ArgumentError partialtrace(rand(3, 4), 1, [2, 3])
    end
end

@add_problem affine function affine_permuteddims_matrix(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
#this function is used in the partial transpose
    for n in (2, 3, 4, 5)
        dims = ntuple( i -> rand(2:5), n)
        d = prod(dims)
        v = rand(d)
        p = randperm(n)
        out1 = vec(permutedims(reshape(v, dims), p))
        out2 = Convex.permutedims_matrix(dims, p) * v
        if test
            @test out1 ≈ out2 atol=atol rtol=rtol
        end
    end
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    KMM(σ=1.0, B=Inf, ϵ=0.01, λ=0.001)

Kernel Mean Matching (KMM).

## Parameters

* `σ` - Bandwidth of Gaussian kernel (default to `2.0`)
* `B` - Maximum discrepancy allowed (default to `Inf`)
* `ϵ` - Tolerance for unit sum (default to `0.01`)
* `λ` - Regularization parameter (default to `0.001`)

## References

* Huang et al. 2006. Correcting sample selection bias by
  unlabeled data.

### Author

* Júlio Hoffimann (julio.hoffimann@gmail.com)
* Kai Xu (xukai921110@gmail.com)
"""
@with_kw struct KMM{T} <: DensityRatioEstimator
  σ::T=2.0
  B::T=Inf
  ϵ::T=0.01
  λ::T=0.001
end

default_optlib(dre::Type{<:KMM}) = JuMPLib

available_optlib(dre::Type{<:KMM}) = [JuliaLib, JuMPLib]

function _kmm_consts(x_nu, x_de, dre::KMM{T}) where {T<:AbstractFloat}
  @unpack σ, λ = dre

  # Gramian matrices for numerator and denominator
  Kdede = gaussian_gramian(x_de; σ=σ)
  if !iszero(λ)
    Kdede += safe_diagm(Kdede, λ)
  end
  Kdenu = gaussian_gramian(x_de, x_nu; σ=σ)

  # number of denominator and numerator samples
  n_de, n_nu = size(Kdenu)

  Kdede, T(n_de / n_nu) * sum(Kdenu, dims=2)
end

function _densratio(x_nu, x_de, dre::KMM,
                    optlib::Type{<:OptimizationLibrary})
  K, κ = _kmm_consts(x_nu, x_de, dre)
  _kmm_ratios(K, κ, dre, optlib)
end


using Plots

include("helper_functions.jl")

begin
    λ1 = 1
    λ2 = 1
    α = 1
    μ(α) = 1 / (1.5 * (α))
    ν(α) = 1 / (0.7 * (α))
    qs = 0:0.01:1
    parameters = (λ1, λ2, α, μ, ν)
end

ps = numerical_pstar.(qs, parameters...)
a_ps = pstar.(qs, parameters...)
costs = cost.(ps, qs, parameters...)
a_costs = cost.(a_ps, qs, parameters...)

ν_assumption = ν.(α) .> λ1 .* ps .+ λ2 .* qs
μ_assumption = μ.(α) .> λ1 .* (1 .- ps)

a_ν_assumption = ν.(α) .> λ1 .* a_ps .+ λ2 .* qs
a_μ_assumption = μ.(α) .> λ1 .* (1 .- a_ps)

plot(qs, ps, label = "autodiff")
plot!(qs, a_ps, label = "analytic")

# plot!(qs, ν_assumption, colour = :red, label = "pu o ν")
# plot!(qs, μ_assumption, colour = :black, label = "pr o μ")
# plot!(qs, a_ν_assumption, colour = :red, label = "pu o ν")
# plot!(qs, a_μ_assumption, colour = :black, label = "pr o μ")

plot(qs, costs, label = "autodiff_costs")
plot!(qs, a_costs, label = "analytic_costs")
plot!(ylims = (-50, 50))

p = plot(qs, qs, (p, q) -> clamp(cost(p, q, parameters...), -50, 50), st = :surface, zlims =  (-50, 50))

# begin
#     spin_plot = @animate for i in vcat(0:0.5:90, reverse(0:0.5:90))
#         plot(p, camera = (i,i))
#     end
#     gif(spin_plot, "spin_plot.gif", fps = 60)
# end

# gif(spin_plot, "spin_plot.gif", fps = 60)

plot!(shapes_from_intervals(shaded_regions(μ_assumption, qs)), colour=:red, label = "μ violated", alpha=0.3)
plot!(shapes_from_intervals(shaded_regions(ν_assumption, qs)), colour=:blue, label = "ν violated", alpha=0.3)
plot!(shapes_from_intervals(shaded_regions(a_μ_assumption, qs)), colour=:yellow, label = "a_μ violated", alpha=0.3)
plot!(shapes_from_intervals(shaded_regions(a_ν_assumption, qs)), colour=:gray, label = "a_ν violated", alpha=0.3)

plot!(shapes_from_intervals(push!(shaded_regions(μ_assumption, qs), [0, 0.1])), colour=:red, label = "μ violated", alpha=0.3)


using Base.Test
using UnitQuaternions

####################################################################################################
# Constructors
####################################################################################################

print("Testing constructors...")

q = UnitQuaternion([1, 1, 1], 1)
for i = 1:4
    @test_approx_eq q[i] 0.5
end

@test_throws ErrorException UnitQuaternion([1, 1, 1, 1], 1)

@test UnitQuaternion() == UnitQuaternion([0, 0, 0], 1)

for i = 1:100
    a, b, c, d = 5randn(4)
    p = UnitQuaternion(a, b, c, d)
    q = UnitQuaternion([a, b, c, d])
    @test p == UnitQuaternion([a, b, c], d)
    @test q == UnitQuaternion([a, b, c], d)
    @test p == q
    @test_approx_eq sqrt(q.η^2 + norm(q.ϵ)^2) 1
end

@test UnitQuaternion([0, 0, 0]) == UnitQuaternion()
@test UnitQuaternion([π, 0, 0]) == UnitQuaternion(1, 0, 0, 0)
@test UnitQuaternion([2π, 0, 0]) == UnitQuaternion(0, 0, 0, -1)
@test UnitQuaternion([0, π/2, 0]) == UnitQuaternion(0, sqrt(2)/2, 0, sqrt(2)/2)
@test UnitQuaternion([0, 0, π/3]) == UnitQuaternion(0, 0, 0.5, sqrt(3)/2)

println(" done.")

####################################################################################################
# Operators
####################################################################################################

print("Testing operators...")

@test_approx_eq +(UnitQuaternion()) eye(4)
@test_approx_eq +(UnitQuaternion(1, 0, 0, 0)) [0 0 0 1; 0 0 1 0; 0 -1 0 0; -1 0 0 0]
@test_approx_eq +(UnitQuaternion(1, 1, 1, 1)) [ 0.5  0.5 -0.5  0.5;
                                               -0.5  0.5  0.5  0.5;
                                                0.5 -0.5  0.5  0.5;
                                               -0.5 -0.5 -0.5  0.5;]

@test UnitQuaternion(1, 2, 3, 4) + UnitQuaternion() == UnitQuaternion(1, 2, 3, 4)
@test UnitQuaternion() + UnitQuaternion(1, 2, 3, 4) == UnitQuaternion(1, 2, 3, 4)
@test UnitQuaternion(1, 2, 3, 4) ⊕ UnitQuaternion() == UnitQuaternion(1, 2, 3, 4)
@test UnitQuaternion() ⊕ UnitQuaternion(1, 2, 3, 4) == UnitQuaternion(1, 2, 3, 4)

q = UnitQuaternion()
@test q^2 == q
@test q^3 == q
q = UnitQuaternion(sind(90/2), 0, 0, cosd(90/2))
@test q^2 == UnitQuaternion(sind(180/2), 0, 0, cosd(180/2))
@test q^0.5 == UnitQuaternion(sind(45/2), 0, 0, cosd(45/2))

p = UnitQuaternion(sqrt(2)/2, 0, 0, sqrt(2)/2)
q = UnitQuaternion(0, 0, sqrt(2)/2, sqrt(2)/2)
@test p + q == UnitQuaternion(0.5, 0.5, 0.5, 0.5)
@test q ⊕ p == UnitQuaternion(0.5, 0.5, 0.5, 0.5)
@test q + p == UnitQuaternion(0.5, -0.5, 0.5, 0.5)
@test p ⊕ q == UnitQuaternion(0.5, -0.5, 0.5, 0.5)

@test q ⊞ [0, 0, 0] == q
@test p ⊞ (q ⊟ p) == q
@test_approx_eq (q ⊞ [0.1, 0.2, 0.3]) ⊟ q [0.1, 0.2, 0.3]

r = inv(p) + q
@test r == UnitQuaternion(-0.5, -0.5, 0.5, 0.5)
@test_approx_eq q ⊟ p 2atan2(norm(r.ϵ), r.η) * r.ϵ / norm(r.ϵ)

@test p ⊞ [0, 0, π/2] == UnitQuaternion(0.5, 0.5, 0.5, 0.5)

@test_throws BoundsError q[5]
@test_approx_eq q[1] 0
@test_approx_eq q[4] sqrt(2)/2

println(" done.")

####################################################################################################
# Methods
####################################################################################################

print("Testing methods...")

q = UnitQuaternion(0, 0, sqrt(2)/2, sqrt(2)/2)

@test_approx_eq angle(q) π/2

@test_approx_eq axis(q) [0, 0, 1]

qmean = UnitQuaternion(sind(90/2), 0, 0, cosd(90/2))
v = [UnitQuaternion(sind(60/2), 0, 0, cosd(60/2)), UnitQuaternion(sind(120/2), 0, 0, cosd(120/2))]
cov = covariance(v, qmean, [1, 1])
@test_approx_eq rad2deg(sqrt(cov[1])) 30

@test inv(q) == UnitQuaternion(0, 0, -sqrt(2)/2, sqrt(2)/2)

qs = [UnitQuaternion(), UnitQuaternion()]
weights = [rand(), rand()]
@test mean(qs, weights) == UnitQuaternion()
@test mean(qs) == UnitQuaternion()
qs = [UnitQuaternion(1, 0, 0, 0), UnitQuaternion(1, 0, 0, 0)]
weights = [rand(), rand()]
@test mean(qs, weights) == UnitQuaternion(1, 0, 0, 0)
@test mean(qs) == UnitQuaternion(1, 0, 0, 0)
qs = [UnitQuaternion(1.1, 0, 0, 0), UnitQuaternion(0.9, 0, 0, 0)]
weights = [1, 1]
@test mean(qs, weights) == UnitQuaternion(1, 0, 0, 0)
@test mean(qs) == UnitQuaternion(1, 0, 0, 0)
qs = [UnitQuaternion(0, 1.2, 0, 0), UnitQuaternion(0, 0.8, 0, 0)]
weights = [1, 1]
@test mean(qs, weights) == UnitQuaternion(0, 1, 0, 0)
@test mean(qs) == UnitQuaternion(0, 1, 0, 0)
qs = [UnitQuaternion(0, 0, 1.3, 0), UnitQuaternion(0, 0, 0.7, 0)]
weights = [1, 1]
@test mean(qs, weights) == UnitQuaternion(0, 0, 1, 0)
@test mean(qs) == UnitQuaternion(0, 0, 1, 0)
qs = [UnitQuaternion(0, 0, 0, 1.4), UnitQuaternion(0, 0, 0, 0.6)]
weights = [1, 1]
@test mean(qs, weights) == UnitQuaternion(0, 0, 0, 1)
@test mean(qs) == UnitQuaternion(0, 0, 0, 1)

@test_approx_eq log(UnitQuaternion()) [0, 0, 0]
@test_approx_eq log(UnitQuaternion(1, 1, 1, 0)) π/2 * [sqrt(3)/3, sqrt(3)/3, sqrt(3)/3]
@test_approx_eq log(q) [0, 0, π/4]

@test_approx_eq rotatevector(q, [1, 0, 0]) [0, 1, 0]
p = UnitQuaternion(0.5, 0.5, 0.5, 0.5)
@test_approx_eq rotatevector(p, [1, 0, 0]) [0, 1, 0]

for i = 1:10
    r = rand(3)
    p = UnitQuaternion(rand(4))
    q = UnitQuaternion(rand(4))
    @test_approx_eq rotateframe(q, r) rotationmatrix(q) * r
    @test_approx_eq rotateframe(p, r) rotationmatrix(p) * r
    @test_approx_eq rotationmatrix(p + q) rotationmatrix(p) * rotationmatrix(q)
end

p = UnitQuaternion(sind(30/2), 0, 0, cosd(30/2))
q = UnitQuaternion(sind(60/2), 0, 0, cosd(60/2))
@test slerp(p, q, 0.5) == UnitQuaternion(sind(45/2), 0, 0, cosd(45/2))
p = UnitQuaternion(0, sind(10/2), 0, cosd(10/2))
q = UnitQuaternion(0, sind(90/2), 0, cosd(90/2))
@test slerp(p, q, 0.1) == UnitQuaternion(0, sind(18/2), 0, cosd(18/2))
@test slerp(p, q, 0.4) == UnitQuaternion(0, sind(42/2), 0, cosd(42/2))
@test slerp(p, q, 1.1) == UnitQuaternion(0, sind(98/2), 0, cosd(98/2))

q = UnitQuaternion(0, 0, sqrt(2)/2, sqrt(2)/2)
@test_approx_eq vector(UnitQuaternion()) [0, 0, 0, 1]
@test_approx_eq vector(q) [0, 0, sqrt(2)/2, sqrt(2)/2]

println(" done.")


# # [Shor's Algorithm](@id Shor)

# ## References
# * [Neilsen](https://aapt.scitation.org/doi/abs/10.1119/1.1463744?journalCode=ajp)
# * [An Insightful Blog](https://algassert.com/post/1718)

# ## Main Program
# The main program of a Shor's algorithm can be summrized in several lines of code.
# For the theory part, please refer the reference materials above.
# It factorize an integer `L`, and returns one of the factors.
# Here, the input `ver` can be either `Val(:quantum)` or `Val(:classical)`,
# where the classical version is for comparison.

using Yao, BitBasis
using YaoExtensions: KMod, QFTCircuit
using QuAlgorithmZoo: NumberTheory

function shor(L::Int, ver=Val(:quantum); maxtry=100)
    L%2 == 0 && return 2

    ## find short cut solutions like `a^b`
    res = NumberTheory.factor_a_power_b(L)
    res !== nothing && return res[1]

    for i in 1:maxtry
        ## step 1
        x = NumberTheory.rand_primeto(L)

        ## step 2
        r = get_order(ver, x, L; )
        if r%2 == 0 && powermod(x, r÷2, L) != L-1
            ## step 3
            f1, f2 = gcd(powermod(x, r÷2, L)-1, L), gcd(powermod(x, r÷2, L)+1, L)
            if f1!=1
                return f1
            elseif f2!=1
                return f2
            else
                error("Algorithm Fail!")
            end
        end
    end
end

# Except some shortcuts, in each try, the main program can be summarized in several steps
# 1. randomly pick a number that prime to the input numebr `L`, i.e. `gcd(x, L) = 1`.
# The complexity of this algorithm is polynoial.
# 2. get the order `x`, i.e. finding a number `r` that satisfies `mod(x^r, L) = 1`.
# If `r` is even and `x^(r÷2)` is non-trivil, go on, otherwise start another try.
# Here, trivil means equal to `L-1 (mod L)`.
# 3. According to Theorem 5.2 in Neilsen book,
# one of `gcd(x^(r÷2)-1, L)` and `gcd(x^(r÷2)+1, L)` must be a non-trivil (`!=1`) factor of `L`.
# Notice `powermod(x, r÷2, L)` must be `-1` rather than `1`,
# otherwise the order should be `r/2` according to definition.

# The only difference between classical and quantum version is the order finding algorithm.

# ## Order Finding
# We provided a classical order finding algorithm in `NumberTheory`,
# here we focus on the quantum version.
# The algorithm is consisted
# 1. run the circuit to get a bitstring,
# 2. interpret this bitstring in output register as a rational number `s/r`.
# To achieve this, we first interpret it as a floating point number,
# then the continued fraction algorithm can find the best match for us.
#
# When using the quantum version, we have the flexibility to set key word arguments `nshot`,
# `nbit` (size of input data register) and `ncbit` (size of control register, or output register).
# `nbit` can be simply chosen as the minimum register size to store input,
# while `ncbit` can be estimated with the following function
"""estimate the required size of the output register."""
estimate_ncbit(nbit::Int, ϵ::Real) = 2*nbit + 1 + ceil(Int,log2(2+1/2ϵ))

get_order(::Val{:classical}, x::Int, L::Int; kwargs...) = NumberTheory.find_order(x, L)
function get_order(::Val{:quantum}, x::Int, L::Int; nshots::Int=10,
            nbit::Int=bit_length(L-1), ncbit::Int=estimate_ncbit(nbit, 0.25))
    c = order_finding_circuit(x, L; nbit=nbit, ncbit=ncbit)
    reg = join(product_state(nbit, 1), zero_state(ncbit))

    res = measure(copy(reg) |> c; nshots=nshots)
    for r in res
        ## split bit string b into lower bits `k` and higher bits `r`.
        mask = bmask(1:ncbit)
        k,i = r&mask, r>>ncbit
        ## get s/r
        ϕ = bfloat(k)  #
        ϕ == 0 && continue

        ## order_from_float: given a floating point number,
        ## return the closest rational number with bounded number of continued fraction steps.
        order = NumberTheory.order_from_float(ϕ, x, L)
        if order === nothing
            continue
        else
            return order
        end
    end
    return nothing
end

# #### The circuit used for finding order
"""
    order_finding_circuit(x::Int, L::Int; nbit::Int=bit_length(L-1), ncbit::Int=estimate_ncbit(nbit, 0.25)) -> AbstractBlock

Returns the circuit for finding the order of `x` to `L`,
feeding input `|1>⊗|0>` will get the resulting quantum register with the desired "phase" information.
"""
function order_finding_circuit(x::Int, L::Int; nbit::Int, ncbit::Int)
    N = nbit+ncbit
    chain(N, repeat(N, H, 1:ncbit), KMod{N, ncbit}(x, L),
        subroutine(N, QFTCircuit(ncbit)', 1:ncbit))
end

# The circuit for order finding is consisted of three parts
# 1. Hadamard gates,
# 2. `KMod` that computes a classical function `mod(a^k*x, L)`.
# `k` is the integer stored in first `K` (or `ncbit`) qubits and the rest `N-K` qubits stores `a`.
# Notice it is not a basic gate, it should have been compiled to multiple gates, which is not implemented in `Yao` for the moment.
# To learn more about implementing arithmatics on a quantum circuit, please read [this paper](https://arxiv.org/abs/1805.12445).
# 3. Inverse quantum fourier transformation.

# ## Run
# Factorizing `15`, you should see `3` or `5`, please report a bug if it is not...
using Random; Random.seed!(129) #src
shor(15, Val(:quantum))


using ClusterComplex
using Plots
upscale = 1.0
default(size=(400*upscale,400*upscale))
pgfplots()
# gr()

X, H = ClusterComplex.dataset("TwoMoons")
s1 = scatter(X[1,:], X[2,:], 1, leg=false, markersize=1.7);
Plots.pdf(s1, "two-moons-ds")

X, H = ClusterComplex.dataset("Circles")
s2 = scatter(X[1,:], X[2,:], 1, leg=false, markersize=1.7);
Plots.pdf(s2, "circles-ds")

X, H = ClusterComplex.dataset("Sphere")
s3 = plot(X[1,:], X[2,:], X[3,:], seriestype=:scatter3d, legend=:none, markersize=1.7);
Plots.pdf(s3, "sphere-ds")

X, H = ClusterComplex.dataset("OIP300")
S = (rand(1:15000, 1200) |> unique)[1:1000]
s4 = plot(X[1,S], X[2,S], X[3,S], seriestype=:scatter3d, camera=(45, 15), legend=:none, markersize=1);
Plots.pdf(s4, "oip-300-ds")

X, H = ClusterComplex.dataset("OIP15")
S = (rand(1:15000, 1200) |> unique)[1:1000]
s5 = plot(X[1,S], X[2,S], X[3,S], seriestype=:scatter3d, camera=(45, 15), legend=:none, markersize=1);
Plots.pdf(s5, "oip-15-ds")


function f!(resid, x)
    for i in eachindex(x)
       resid[i] = sin(x[i]) + x[i]^3
    end
end
x = ones(5)
@test Sundials.kinsol(f!, x, linear_solver=:Band, jac_upper=0, jac_lower=0) == Sundials.kinsol(f!, x)


julia> sum(diff(Q(1:10^5)) .< 0)
49798


# testing performance of AD
# conclusion is that it scales linearly with dimension;
# however it is a factor 100 slower than the objective.

using BenchmarkTools
using ReverseDiff: GradientTape, GradientConfig, gradient, gradient!, compile
import ForwardDiff

ϕ(r) = exp( - 6 * (abs(r)-1)) - 2 * exp( - 3 * (abs(r)-1))
dϕ(r) = ForwardDiff.derivative(ϕ, r)
ρ(r) = exp( - 4 * (abs(r)-1) )
dρ(r) = ForwardDiff.derivative(ρ, r)
ψ(t) = sqrt(t)
dψ(t) = ForwardDiff.derivative(ψ, t)

function get_rho(x)
   N = length(x)
   rho = zeros(eltype(x), N)
   for n = 2:N
      t = ρ(x[n]-x[n-1])
      rho[n-1] += t
      rho[n] += 1
   end
   for n = 3:N
      t = ρ(x[n] - x[n-2])
      rho[n] += t
      rho[n-2] += t
   end
   return rho
end

function F(x)
   N = length(x)
   E = 0.0
   for j = 1:2, n = (1+j):N
      E += ϕ(x[n]-x[n-j])
   end
   return E + sum(ψ, get_rho(x))
end

function dF(x)
   N = length(x)
   dE = zeros(N)
   F = ψ.(get_rho(x))
   for j = 1:2, n = (1+j):N
      dϕ_ = dϕ(x[n]-x[n-j])
      dρ_ = dρ(x[n]-x[n-j])
      a = dϕ_ + (F[n] - F[n-j]) * dρ_
      dE[n] += a
      dE[n-j] -= a
   end
   return dE
end

x = gresult = compiled_f_tape = nothing
for N in (100, 400, 1600)
   # pre-record a GradientTape for `F` using inputs with Float64 elements
   f_tape = GradientTape(F, rand(N))
   # compile `f_tape` into a more optimized representation
   compiled_f_tape = compile(f_tape)
   # some inputs and work buffers to play around with
   gresult = zeros(N)
   x = collect(1:N) + rand(N)
   println("N = $N")
   print("                 Time for F: ")
   @btime F($x)
   print("  Time for DF (ForwardDiff): ")
   @btime ForwardDiff.gradient($F, $x)
   print("  Time for DF (ReverseDiff): ")
   @btime gradient!($gresult, $compiled_f_tape, $x)
   print("        Time for DF(manual): ")
   @btime dF($x)
end

# RESULTS 03/04/2018, MacBook Pro, Julia v0.6.2,
# - ForwardDiff                   0.7.4
# - ReverseDiff                   0.2.0
#
# N = 8000
#                Time for F:
#   564.137 μs (2 allocations: 62.58 KiB)
#   Time for DF (ForwardDiff):
# N = 100
#                  Time for F:   6.191 μs (1 allocation: 896 bytes)
#   Time for DF (ForwardDiff):   409.391 μs (6043 allocations: 659.25 KiB)
#   Time for DF (ReverseDiff):   451.058 μs (0 allocations: 0 bytes)
#         Time for DF(manual):   17.426 μs (3 allocations: 2.63 KiB)
# N = 400
#                  Time for F:   26.206 μs (1 allocation: 3.25 KiB)
#   Time for DF (ForwardDiff):   6.538 ms (96204 allocations: 10.17 MiB)
#   Time for DF (ReverseDiff):   1.999 ms (0 allocations: 0 bytes)
#         Time for DF(manual):   63.400 μs (3 allocations: 9.75 KiB)
# N = 1600
#                  Time for F:   109.590 μs (1 allocation: 12.63 KiB)
#   Time for DF (ForwardDiff):   155.275 ms (1536804 allocations: 162.25 MiB)
#   Time for DF (ReverseDiff):   15.896 ms (0 allocations: 0 bytes)
#         Time for DF(manual):   260.063 μs (3 allocations: 37.88 KiB)


function gradientdescent(c, yspace, y0=nothing)

    dim = size(yspace,1)
    ymin = yspace[:,1]
    ymax = yspace[:,2]
    N = Int64.(yspace[:,3])

    if y0 === nothing
        M = 100;
        y0 = rand(M, dim)
        for i = 1:dim
            y0[:,i] = ymin[i] .+ y0[:,i]*(ymax[i] - ymin[i])
        end
    end

    # estimate maximum stepsize
    f, df = mpolyder(y0, c, ymin, ymax, N)
    y = sum(y0.*y0, dims=2)
    gamma = sum(df.*df, dims=2)
    alphamax = minimum(sqrt.(y./gamma))
    alphamax = max(alphamax, 10.0)

    beta = 0.8;
    y = 1.0*y0
    M = size(y,1)
    y1 = 1.0*y
    r1 = zeros(M)
    maxstep = zeros(M)
    alpha = alphamax*ones(M)
    a = zeros(dim)
    iter = 0
    while (true)
        iter = iter + 1
        f, df = mpolyder(y, c, ymin, ymax, N)
        gamma = 0.5*sum(df.*df, dims=2)

        if minimum(gamma[:]) <= 1e-12
            fmin, imin = findmin(f)
            sol = y[imin,:]
            return sol, fmin, iter, y, f, df, y0
        end

        for m = 1:M
            for i = 1:dim
                a1 = (y[m,i]-ymax[i])/df[m,i]
                a2 = (y[m,i]-ymin[i])/df[m,i]
                a[i] = (abs(df[m,i])<=1e-12) ? alphamax : max(a1, a2)
            end
            maxstep[m] = min(minimum(a), alphamax)
            alpha[m] = min(alpha[m], maxstep[m])
            y1[m,:] = y[m,:] - alpha[m]*df[m,:]
            r1[m] = f[m] - alpha[m]*gamma[m]
        end

        f1 = mpolyeval(y1, c, ymin, ymax, N)
        for m = 1:M
            fm = f1[m]
            rm = r1[m]

            while (fm > rm)
                alpha[m] = beta*alpha[m]
                rm = f[m] - alpha[m]*gamma[m]
                y1[m,:] = y[m,:] - alpha[m]*df[m,:]
                tm = mpolyeval(reshape(y1[m,:],(1,dim)), c, ymin, ymax, N)
                fm = tm[1]
            end
            y[m,:] = y1[m,:]
            alpha[m] = (1.0/beta)*alpha[m]
        end
    end
end


"""
$(TYPEDEF)

Coordinates as row (vertical), columns (horizontal). Represent the data of each top corner of a pixel.
Topleft for top triangle, bottomleft for bottom triangle
vr, vc: velocity vector of that corner.
Warning: posts/fences: One more corner than number of pixels.
"""
struct Vertex3D
    r::Float64
    c::Float64
    ϕ::Float64

    vr::Float64
    vc::Float64

    Vertex3D(r, c, ϕ, vr, vc) = new(r, c, ϕ, vr, vc)
    Vertex3D() = new(0.0, 0.0, 0.0, 0.0, 0.0)
end

"""
$(SIGNATURES)
"""
function dist(p1::Tuple{Float64,Float64,Float64}, p2::Tuple{Float64,Float64,Float64})
    dr = p2[1] - p1[1]
    dc = p2[2] - p1[2]
    return sqrt(dr^2 + dc^2)
end

"""
$(SIGNATURES)
"""
function dist(p1::Vertex3D, p2::Vertex3D)
    dr = p2.r - p1.r
    dc = p2.c - p1.c
    return sqrt(dr^2 + dc^2)
end

"""
$(SIGNATURES)
"""
function area(v1::Vertex3D, v2::Vertex3D, v3::Vertex3D)
    a = dist(v1, v2)
    b = dist(v2, v3)
    c = dist(v3, v1)
    s = (a + b + c) / 2.0

    # surface_sq = (a + b + c) * (- a + b + c) * (a - b + c) * (a + b - c) / 16.0
    surface_sq = s * (s - a) * (s - b) * (s - c)
    return surface_sq <= 1e-100 ? 0.0 : sqrt(surface_sq)
end

"""
$(SIGNATURES)

Centroid of three points.
"""
centroid(p1::Vertex3D, p2::Vertex3D, p3::Vertex3D) =
    Vertex3D((p1.r + p2.r + p3.r) / 3.0, (p1.c + p2.c + p3.c) / 3.0, 0.0, 0.0, 0.0)

"""
$(SIGNATURES)
"""
function fill_borders!(mat::AbstractMatrix{T}, val::T) where {T}
    mat[begin:end, begin] .= val
    mat[begin:end, end] .= val
    mat[begin, begin:end] .= val
    mat[end, begin:end] .= val
end

"""
$(TYPEDEF)

## Coordinates

Origin is top left, going right and down .`x` goes horizontal and follows columns. x is within 1 and width
`y` goes vertical and follows rows. y is within 1 and height.

# Velocity

Velocity vector along `row` and `col` (velocity along h not used).

Implementation is a struc of arrays. Easier to vectorise.

"""
mutable struct FieldVertex3D
    size::Tuple{Int,Int}

    r::AbstractMatrix{Float64}
    c::AbstractMatrix{Float64}
    vr::AbstractMatrix{Float64}
    vc::AbstractMatrix{Float64}

    # Velocity potential at each corner
    ϕ_top::AbstractMatrix{Float64}
    ϕ_bot::AbstractMatrix{Float64}

    rows_numbers::AbstractMatrix{Float64}
    cols_numbers::AbstractMatrix{Float64}
end

# Information about each corner surrounding a pixel. Posts&fences warning!!!
function FieldVertex3D(height, width)
    # Two matrices comtaining the row and column number of each element of the matrix
    rows_numbers = repeat(Float64.(1:height+1), 1, width + 1)
    cols_numbers = repeat(Float64.(1:width+1)', height + 1, 1)

    mr = rows_numbers
    mc = cols_numbers
    mvr = zeros(Float64, height + 1, width + 1)
    mvc = zeros(Float64, height + 1, width + 1)

    mϕ_top = zeros(Float64, height , width )
    mϕ_bot = zeros(Float64, height , width )

    fv = FieldVertex3D((height, width), mr, mc, mvr, mvc, mϕ_top, mϕ_bot, rows_numbers, cols_numbers)
    reset_border_values!(fv)
    return fv
end

"""
$(SIGNATURES)

Reset the values of the borders corners: all positions at integer values, all velocities at 0.
"""
function reset_border_values!(corners::FieldVertex3D)
    corners.r[1, :] .= corners.rows_numbers[1, :]
    corners.r[end, :] .= corners.rows_numbers[end, :]

    corners.r[:, 1] .= corners.rows_numbers[:, 1]
    corners.r[:, end] .= corners.rows_numbers[:, end]

    corners.c[1, :] .= corners.cols_numbers[1, :]
    corners.c[end, :] .= corners.cols_numbers[end, :]

    corners.c[:, 1] .= corners.cols_numbers[:, 1]
    corners.c[:, end] .= corners.cols_numbers[:, end]

    fill_borders!(corners.vr, 0.0)
    fill_borders!(corners.vc, 0.0)
end

Base.size(fv::FieldVertex3D) = fv.size

"""
$(SIGNATURES)
"""
function Vertex3D(fv::FieldVertex3D, row, col)
    height, width = size(fv)

    if (1 <= row <= height ) && (1 <= col <= width)
        Vertex3D(
            fv.r[row, col],
            fv.c[row, col],
            fv.vr[row, col],
            fv.vc[row, col],
            fv.ϕ_top[min(height, row), col(width, col)],
            fv.ϕ_bot[min(height, row), col(width, col)],
        )
    else
        missing
    end
end

"""
$(TYPEDEF)

The algorithm works by allocating a rectangle from the exit mesh (the face where the light exits the block
of acrylate). Initially, each rectangle representsa single pixel on that face and faces s single pixel on
the caustics picture.

Each rectangle is split into 2 trangles: one upper-left, the other bottom-right like:

```
*------*
|    / |
|   /  |
|  /   |
| /    |
*------*
```
Although the corners of each triangle match corners of the rectangles, it is sometimes easier to think in terms
of triangles rather than corners coming from different rectangles.
"""
struct FaceMesh
    size::Tuple{Int,Int}
    corners::FieldVertex3D

    toptriangles::AbstractMatrix{Tuple{Tuple{Int,Int},Tuple{Int,Int},Tuple{Int,Int}}}
    bottriangles::AbstractMatrix{Tuple{Tuple{Int,Int},Tuple{Int,Int},Tuple{Int,Int}}}

    FaceMesh(s, tl, tt, bt) = new(s, tl, tt, bt)

    function FaceMesh(height::Int, width::Int)
        m_corner = FieldVertex3D(height, width)
        t_triangles = Matrix{Tuple{Tuple{Int,Int},Tuple{Int,Int},Tuple{Int,Int}}}(
            undef,
            height,
            width,
        )
        b_triangles = Matrix{Tuple{Tuple{Int,Int},Tuple{Int,Int},Tuple{Int,Int}}}(
            undef,
            height,
            width,
        )

        for row ∈ 1:height, col ∈ 1:width
            t_triangles[row, col] = ((row, col), (row + 1, col), (row, col + 1))
            b_triangles[row, col] = ((row, col + 1), (row + 1, col), (row + 1, col + 1))
        end

        return FaceMesh((height, width), m_corner, t_triangles, b_triangles)
    end
end

"""
$(SIGNATURES)

Return the dimension of the mesh as the number of rectangles. It does not return the number of corner points.
"""
Base.size(mesh::FaceMesh) = mesh.size

"""
$(SIGNATURES)

Converts a triangle as a triplet of references to mesh vertices to a triplet of 3D coordinates.
"""
function triangle3D(mesh::FaceMesh, row::Int, col::Int, side = Union{:top,:bottom})
    height, width = size(mesh)
    max_distance_squared = (height + 1)^2 + (width + 1)^2

    if 1 <= row <= height && 1 <= col <= width
        if side == :top
            t = mesh.toptriangles[row, col]
        elseif side == :bottom
            t = mesh.bottriangles[row, col]
        end

        t1_row, t1_col = t[1]
        p1 = Vertex3D(mesh.corners, t1_row, t1_col)

        t2_row, t2_col = t[2]
        p2 = Vertex3D(mesh.corners, t2_row, t2_col)

        t3_row, t3_col = t[3]
        p3 = Vertex3D(mesh.corners, t3_row, t3_col)

        return (p1, p2, p3)
    else
        # return (missing, missing, missing)
        @assert false "Coordinates out of bounds at $(row), $(col)"
    end
end

triangle3D(mesh::FaceMesh, ci::CartesianIndex{2}, side = Union{:top,:bottom}) =
    triangle3D(mesh, ci[1], ci[2], side)

"""
$(SIGNATURES)

Centroid of a specific mesh triangle.
"""
function centroid(mesh::FaceMesh, ci::CartesianIndex; side = Union{:top,:bottom})
    if side == :top
        return centroid(top_triangle3D(mesh, ci)...)
    elseif side == :bottom
        return centroid(bot_triangle3D(mesh, ci)...)
    else
        return missing
    end
end

"""
$(SIGNATURES)

Centroid of a specific mesh triangle.
"""
function centroid(mesh::FaceMesh, row::Int, col::Int; side = Union{:top,:bottom})
    if side == :top
        return centroid(top_triangle3D(mesh, row, col)...)
    elseif side == :bottom
        return centroid(bot_triangle3D(mesh, row, col)...)
    else
        return missing
    end
end

"""
$(SIGNATURES)
"""
function area(mesh::FaceMesh, ci::CartesianIndex; side = Union{:top,:bottom})
    if side == :top
        return area(top_triangle3D(mesh, ci)...)
    elseif side == :bottom
        return area(top_triangle3D(mesh, ci)...)
    else
        return missing
    end
end

"""
$(SIGNATURES)
"""
function area(mesh::FaceMesh, row::Int, col::Int; side = Union{:top,:bottom})
    if side == :top
        return area(top_triangle3D(mesh, row, col)...)
    elseif side == :bottom
        return area(top_triangle3D(mesh, row, col)...)
    else
        return missing
    end
end

"""
$(SIGNATURES)
"""
area(t::Tuple{Vertex3D,Vertex3D,Vertex3D}) = area(t...)

"""
$(SIGNATURES)

A Mesh is a collection of triangles. The brightness flowing through a given triangle is just proportional to its
area in the r, r plane. h is ignored.

The function returns a matrix with the quantity of light coming from each 'rectangle'  around a corner. That 'rectangle'
has been shifted and flexed around.
"""
function get_lens_pixels_area(mesh::FaceMesh)
    height, width = size(mesh)

    top_tri_area = get_lens_pixels_area(mesh, :top)
    bot_tri_area =get_lens_pixels_area(mesh, :bottom)

    total_area = top_tri_area + bot_tri_area

    average_energy_per_pixel = average(total_area)
    luminosity_pixels = total_area / average_energy_per_pixel

    @assert abs(sum(luminosity_pixels) - height * width) < 1.0 "Total lens luminosity is incorrect"

    # return luminosity_pixels
    return top_tri_area, bot_tri_area
end

function get_lens_pixels_area(mesh::FaceMesh, side::Union{:top, :bottom})
    height, width = size(mesh)

    # return luminosity_pixels
    return area.([triangle3D(mesh, row, col, side) for row ∈ 1:height, col ∈ 1:width])
end

"""
$(SIGNATURES)

"""
function field_summary(field, fieldname)
    s = round(sum(field), sigdigits = 6)
    max = round(maximum(field), sigdigits = 4)
    min = round(minimum(field), sigdigits = 4)
    avg = round(average(field), sigdigits = 4)
    abs = round(average_absolute(field), sigdigits = 4)

    return """$(fieldname) (Sum/Min/Max/Avg/Avg abs.): $(s) / $(min) / $(max) / $(avg) / $(abs)"""
end

"""
$(SIGNATURES)

"""
field_summary(field) = field_summary(field, "")


module EMD

include("EMDUtil.jl")
using Splines

export findExtrema
export IMF
export IF

#Function to calculate intrinsic mode functions
#N is the maximum number of modes to find
function IMF(y, t, toldev=0.01, tolzero = 0.01, order=4, N=5,window=0)

	if window==0
		window = ones(length(t))
	end

	n = length(y)
	f = zeros(n,N)
	tempy = copy(y)
	eps = 0.00001

	n_modes = 0;

	for i = 1:N
	    avg = zeros(n,1)+1
	    sd = 2*toldev

	    while(mean(abs(avg))>tolzero && sd > toldev)

	    	# Interpolate a spline through the maxima and minima
	        max_ar, min_ar, tmax, tmin = findExtrema(tempy, t)

	        # # Don't use too high an order to interpolate, restrict it if there are not many extrema
	        # p_max = min(order, length(max_ar)-2)
	        # p_min = min(order, length(min_ar)-2)

	        # # Make even order
	        # p_max = Int(2*floor(p_max/2))
	        # p_min = Int(2*floor(p_min/2))

	        # # At least linear
	        # p_max = max(p_max,2)
	        # p_min = max(p_min,2)
	        p_max = (length(max_ar) >= order)? order : 4
	        p_min = (length(min_ar) >= order)? order : 4

	        S1 = Spline(max_ar, tmax, p_max)
	        S2 = Spline(min_ar, tmin, p_min)

	        # Find mean of envelope
	        avg = (S1(t) + S2(t)) / 2
	        avg = avg.*window

	        tempy = tempy-avg

	        sd = mean( (avg.^2)./((y-f[:,i]).^2 + eps) )

	        f[:,i] = f[:,i] + avg

	        #println(sd)
	    end

	    tempy = copy(f[:,i])

	    # Check to see if it's worth continuing or if the remainder is monotone
	    c = mean(abs(tempy))
	    d = diff(tempy)
	    n_modes = n_modes + 1
	    if all(d+c*tolzero .> 0) || all(d-c*tolzero .< 0)
	    	break
	    end
	end

	C = zeros(n,N)
	C[:,1] = y - f[:,1]
	for i = 2:n_modes
		C[:,i] = f[:,i-1]-f[:,i]
	end

	return C
end

#function to calculate instantaneous frequencies
function IF(C, t, x=0, order=4, window=0)

	n,N = size(C)
	n = length(t)

	if x == 0
		x = t+1e-6
	end

	if window == 0
		window = ones(length(x))
	end

	nx = length(x)
	Phi = zeros(nx,N)

	for i =1:N
		H1 = Spline(C[:,i], t, order)

		Phi[:,i] = 1/(2pi)*(H1(x).*H1(x,1,true) - H1(x, 1).*H1(x,0,true))./(H1(x).^2 + H1(x,0,true).^2)
		Phi[:,i] = Phi[:,i].*window
	end

	return Phi

end

end # module


@testset "StabChain unit tests" begin
    a,b = [perm"(1,2,4,3)(5)", perm"(1,2,5,4)"]
    sc = StabilizerChain([a,b])
    @test length(sc) == 1
    pt, S, Δ = sc[1]
    @test pt == 1
    @test S isa Vector{<:Perm}
    @test S == [a,b]
    @test Δ isa Schreier
    @test collect(Δ) == [1,2,4,5,3]

    # first point on the orbit:
    u = Δ[1]
    # the first Schreier generator is trivial
    @test isone(u*a*getinv(Δ, 1^a))
    # but the second is not
    @test u*b*getinv(Δ, 1^b) == perm"(2,5)(3,4)"
    c = perm"(2,5)(3,4)"
    h, new_depth = sift(c, sc)
    @test h == c
    @test new_depth == 2

    # so we extend the base
    @test firstmoved(c) == 2
    push!(sc, firstmoved(c))

    # pushing extends base...
    @test length(sc) == 2
    @test length(sc.base) == 2
    @test sc.base[2] == firstmoved(c)
    @test length(sc.sgs) == 2

    # and allocates generating set and the transversal
    @test sc.sgs[2] == Perm{Int}[] # we've just extended basis here!
    @test length(sc.transversals) == 2 # the sc.transversals[2] is right now a garbage

    # add c to generators at depth new_depth=2
    push!(sc, c, new_depth, recompute=true)
    # pushing the generator adds to the basis ath the given depth
    @test sc.sgs[2] == [c]
    # and (re)computes the orbit at the given depth:
    @test sc.transversals[2] == Schreier([c], sc.base[2])

    # c should sift now to the identity at depth 3
    x, d = sift(c, sc)
    @test isone(x)
    @test d == 3

    u = Δ[2]
    @test u == a
    # Schreier generators are trivial
    @test u*a == Δ[2^a]
    @test u*b == Δ[2^b]
    # so we move to the next point on the orbit
    u = Δ[4] # β = 4
    @test u*a == Δ[4^a]
    @test u*b*getinv(Δ, 4^b) == a^2*b
    d = a^2*b
    @test sift(d, sc) == (d, 2)

    push!(sc, d, 2, recompute=true)

    @test collect(sc.transversals[2]) == [2,5,3,4]

    x, d = sift(a^2*b, sc)
    @test isone(x)
    @test d == 3

    @test order(sc) == 20

    m = match(r"size (\d+)", string(sc))
    @test parse(Int, m.captures[1]) == 20

    m = match(r"Orbit:\s+(\[.*\])", sprint(show, MIME"text/plain"(), sc))
    @test Meta.parse(m.captures[1]) |> eval == collect(sc.transversals[1])
end

@testset "PermGroups" begin
    @test PermGroup([perm"(1)"]) isa PermGroup
	@test PermGroup(perm"(1)") isa PermGroup
    G = PermGroup(perm"(1)")
    @test isdefined(G, :stabchain) == false
    @test order(G) == 1
	@test order(G) isa BigInt
	@test order(Int, G) isa Int
    @test isdefined(G, :stabchain) == true

	G = PermGroup(perm"(1,2,3,4)", perm"(1,2)(4)")
    @test StabilizerChain(G) isa StabilizerChain
    sc = StabilizerChain(G)
    @test length(sc) == 3
    @test order(sc) == factorial(4)
	@test order(sc) isa BigInt
	@test order(Int, sc) isa Int

    @test length(base(G)) == 3
    @test order(G) == factorial(4)

    SN(n) = [Perm(circshift(collect(1:n), -1)), Perm([[2,1];3:n])]

    for N in 6:20
        S = SN(N)
        S = [S; [Perm(randperm(N)) for _ in 1:3]]
        G = PermGroup(S)
        @test order(G) == factorial(N)
    end

    G = PermGroup(perm"(1,2,3,4)", perm"(1,2)")

    m = match(r"order (\d+)", string(G))
    @test m === nothing

    @test perm"(1,3)" in G
    @test perm"(1,5)" ∉ G

    m = match(r"order (\d+)", string(G))
    @test parse(Int, m.captures[1]) == 24

    A = PermGroup(perm"(1,2,3)", perm"(2,3,4)")
    @test order(A) == 12
    m = match(r"order (\d+)", string(A))
    @test parse(Int, m.captures[1]) == 12

    @test perm"(1,2)" ∉ A
    @test perm"(1,3,4)" in A
end

@testset "Iterate over PermGroup" begin
	K1 = PermGroup([perm"(5,6)", perm"(1,2,3,4,5,6)"]) # Symmetric group on 6 symbols
	elements = [g for g in K1]
    @test elements isa Vector{Perm{Int64}}
	uniq_elements = unique(elements)
    @test order(K1) == length(uniq_elements) == 720
    @test uniq_elements == elements

	K2 = PermGroup(Perm{Int16}[perm"(3,4,5)", perm"(1,2,3,4,5)"]) # Alternating group on 5 symbols
	elements = [g for g in K2]
    @test elements isa Vector{Perm{Int16}}
	uniq_elements = unique(elements)
    @test order(K2) == length(uniq_elements) == 60
	@test uniq_elements == elements
end


@doc raw"""
Symmetric quadratic equilibrium in (x,y,z) coordinates with covariant components of
the vector potential given by
```math
A (x,y,z) = \frac{B_0}{4} \big( - y \, (2 + x^2 + y^2) , \, x \, (2 + x^2 + y^2) , \, 0 \big)^T
```
resulting in the magnetic field with covariant components
```math
B(x,y,z) = B_0 \, \begin{pmatrix}
0 \\
0 \\
1 + x^2 + y^2 \\
\end{pmatrix}
```

Parameters: `B₀`
"""
struct SymmetricQuadratic{T <: Number} <: CartesianEquilibrium
    name::String
    B₀::T
    SymmetricQuadratic{T}(B₀::T) where T <: Number = new("SymmetricQuadraticEquilibrium", B₀)
end

SymmetricQuadratic(B₀::T) where T <: Number = SymmetricQuadratic{T}(B₀)

function Base.show(io::IO, equ::SymmetricQuadratic)
    print(io, "Symmetric Quadratic Equilibrium")
end

function R(x::AbstractArray{T,1}, equ::SymmetricQuadratic) where {T <: Number}
    r(x,equ)
end

function r(x::AbstractArray{T,1}, equ::SymmetricQuadratic) where {T <: Number}
    sqrt(r²(x,equ))
end

function r²(x::AbstractArray{T,1}, equ::SymmetricQuadratic) where {T <: Number}
    X(x,equ)^2 + Y(x,equ)^2
end

function θ(x::AbstractArray{T,1}, equ::SymmetricQuadratic) where {T <: Number}
    atan2(Y(x,equ), X(x,equ))
end

function ϕ(x::AbstractArray{T,1}, equ::SymmetricQuadratic) where {T <: Number}
    θ(x,equ)
end

function A₁(x::AbstractArray{T,1}, equ::SymmetricQuadratic) where {T <: Number}
    - equ.B₀ * x[2] * (2 + x[1]^2 + x[2]^2) / 4
end

function A₂(x::AbstractArray{T,1}, equ::SymmetricQuadratic) where {T <: Number}
    + equ.B₀ * x[1] * (2 + x[1]^2 + x[2]^2) / 4
end

function A₃(x::AbstractArray{T,1}, equ::SymmetricQuadratic) where {T <: Number}
    zero(eltype(x))
end

macro symmetric_quadratic_equilibrium()
    generate_equilibrium_code(SymmetricQuadratic(1.); output=false)
end


#=
Train the model using gradient descent.

Training parameters
num_epochs - number of epochs
α - learning rate
μ - momentum
TODO - decay - function of epochs
TODO - early stopping

Store and display parameters
last_train_every - selected epoch intervals for last training batch
full_train_every - selected epoch intervals for full training set
val_every - selected epoch intervals for validation
display - choose whether to display progress
store - choose whether to store progress
=#

################################################################################
#
# Main Training Function
#
################################################################################

function train(net::NeuralNet, num_epochs::Integer, α::Real, μ::Real;
               nesterov::Bool = true, shuffle::Bool=false,
               last_train_every::Integer = 1,
               full_train_every::Integer = num_epochs,
               val_every::Integer = num_epochs)

    # Some counts
    num_obs = size(net.data.X_train)[end]
    num_batches = div(num_obs, net.batch_size)

    # Floating point data type
    F = net.dtype

    # Learning parameters - convert type
    α = F(α)
    μ = F(μ)

    # Nesterov parameters
    a = μ^2
    b = (μ + F(1)) * α

    # Choose momentum gradient descent or Nesterov accelerated gradient descent
    nesterov ? (update! = nag_update!) : (update! = mgd_update!)

    # Choose coffee break function - either regression or classification problem
    coffee_break = choose_coffee_break(net)

    # Permutation of observations - to be shuffled if selected
    perm = collect(1:num_obs)

    # Loop over epochs
    for epoch in 1:num_epochs

        # Shuffle permutation (if selected) and initialise batch unit range
        shuffle && shuffle!(perm)
        batch_range = 1:net.batch_size

        # Loop over mini batches
        for batch in 1:num_batches

            # Update inputs & target with next batch and increment range
            batch_ids = perm[batch_range]
            net.layers[1].A[:] = viewbatch(net.data.X_train, batch_ids)
            net.layers[end].y[:] = viewbatch(net.data.y_train, batch_ids)
            batch_range += net.batch_size

            # Forward propogation through each layer
            fwd_prop!(net)

            # Back propogation through each layer to calculate gradients
            bwd_prop!(net)

            # Update parameters using the gradients
            for ℓ in net.layers[2:end]
                update!(ℓ, μ, α, a, b)
            end

        end

        # Display and store progress updates
        # Last training loss
        last_batch_cost(net, epoch, last_train_every)
        # Coffee break - full training and validation loss
        coffee_break(net, epoch, num_epochs, full_train_every, val_every)

    end

    # Completed training - show final results
    show_completed(net)

    return
end

################################################################################
#
# Momentum Gradient Descent Update
#
################################################################################

function mgd_update!{F<:AbstractFloat}(W::Array{F}, W_vel::Array{F},
                                       ∇W::Array{F}, μ::F, α::F)
    for i in eachindex(W)
        @inbounds W_vel[i] = μ * W_vel[i] - α * ∇W[i]
        @inbounds W[i] += W_vel[i]
    end
end

function mgd_update!{F<:AbstractFloat}(ℓ::NetLayer, μ::F, α::F, a::F, b::F)
    mgd_update!(ℓ.b, ℓ.b_vel, ℓ.∇b, μ, α)
    mgd_update!(ℓ.W, ℓ.W_vel, ℓ.∇W, μ, α)
end

function mgd_update!{F<:AbstractFloat}(ℓ::WordEmbeddingLayer, μ::F, α::F, a::F, b::F)
    mgd_update!(ℓ.W, ℓ.W_vel, ℓ.∇W, μ, α)
end

function mgd_update!{F<:AbstractFloat}(ℓ::PoolLayer, μ::F, α::F, a::F, b::F)
end

################################################################################
#
# Nesterov's Accelerated Gradient Descent Update
#
################################################################################

function nag_update!{F<:AbstractFloat}(W::Array{F}, W_vel::Array{F},
                                       ∇W::Array{F}, μ::F, α::F, a::F, b::F)
    for i in eachindex(W)
        @inbounds W[i] += a * W_vel[i] - b * ∇W[i]
        @inbounds W_vel[i] = μ * W_vel[i] - α * ∇W[i]
    end
end

function nag_update!{F<:AbstractFloat}(ℓ::NetLayer, μ::F, α::F, a::F, b::F)
    nag_update!(ℓ.b, ℓ.b_vel, ℓ.∇b, μ, α, a, b)
    nag_update!(ℓ.W, ℓ.W_vel, ℓ.∇W, μ, α, a, b)
end

function nag_update!{F<:AbstractFloat}(ℓ::WordEmbeddingLayer, μ::F, α::F, a::F, b::F)
    nag_update!(ℓ.W, ℓ.W_vel, ℓ.∇W, μ, α, a, b)
end

function nag_update!{F<:AbstractFloat}(ℓ::PoolLayer, μ::F, α::F, a::F, b::F)
end

################################################################################
#
# Display and store progress
#
################################################################################

########################################
# Last training batch
########################################

# Time
t() = Dates.format(now(), "HH:MM:SS")

# Last batch training error - cost with regularisation
function last_batch_cost(net, epoch, last_train_every)

    if mod(epoch, last_train_every) == 0
        # Cost
        L = net.layers[end]
        J = L.cost(L.A, L.y) / net.batch_size

        # Add regularisation cost
        J += net.reg_cost(net)

        # Add sparsity cost
        J += sparse_cost(net)

        # Display
        print(t())
        @printf(" : Epoch %d, last batch training error (with regⁿ) - %.3f\n",
                epoch, J)

        # Store
        push!(net.last_train_loss, J)
    end

end

########################################
# Coffee Breaks
########################################

# Choos classification or regression
function choose_coffee_break(net)
    L = net.layers[end]
    typeof(L) <: LinOutputLayer ? coffee_break_r : coffee_break_c
end

# Coffee break for classification task - shows accuracy
function coffee_break_c(net, epoch, num_epochs, full_train_every, val_every)

    # Don't run coffee break after last epoch
    epoch == num_epochs && return

    # Both training and validation sets
    if mod(epoch, full_train_every) == 0 && mod(epoch, val_every) == 0
        @printf("\nCoffee break:\n")

        # Cost and display of training set
        J, acc = loss_and_accuracy(net, net.data.X_train, net.data.y_train)
        @printf("Training error (with regⁿ) - %.3f", J)
        @printf("  |  Training accuracy - %.1f\n", acc)
        push!(net.full_train_loss, J)

        # Cost and display of validation set
        J, acc = val_loss_and_accuracy(net, net.data.X_val, net.data.y_val)
        @printf("Validation error (without regⁿ) - %.3f", J)
        @printf("  |  Validation accuracy - %.1f\n\n", acc)
        push!(net.val_loss, J)

    # Training set only
    elseif mod(epoch, full_train_every) == 0
        @printf("\nCoffee break:\n")

        # Cost and display of training set
        J, acc = loss_and_accuracy(net, net.data.X_train, net.data.y_train)
        @printf("Training error (with regⁿ) - %.3f", J)
        @printf("  |  Training accuracy - %.1f\n\n", acc)
        push!(net.full_train_loss, J)

    # Validation set only
    elseif mod(epoch, val_every) == 0
        @printf("\nCoffee break:\n")

        # Cost and display of validation set
        J, acc = val_loss_and_accuracy(net, net.data.X_val, net.data.y_val)
        @printf("Validation error (without regⁿ) - %.3f", J)
        @printf("  |  Validation accuracy - %.1f\n\n", acc)
        push!(net.val_loss, J)

    end
end

# Coffee break for regression task - only loss (no accuracy)
function coffee_break_r(net, epoch, full_train_every, val_every)

    # Don't run coffee break after last epoch
    epoch == num_epochs && return

    # Both training and validation sets
    if mod(epoch, full_train_every) == 0 && mod(epoch, val_every) == 0
        @printf("\nCoffee break:\n")

        # Cost and display of training set
        J, acc = loss(net, net.data.X_train, net.data.y_train)
        @printf("Training error (with regⁿ) - %.3f\n", J)

        # Cost and display of validation set
        J, acc = val_loss_and_accuracy(net, net.data.X_val, net.data.y_val)
        @printf("Validation error (without regⁿ) - %.3f\n\n", J)

    # Training set only
    elseif mod(epoch, full_train_every) == 0
        @printf("\nCoffee break:\n")

        # Cost and display of training set
        J, acc = loss_and_accuracy(net, net.data.X_train, net.data.y_train)
        @printf("Training error (with regⁿ) - %.3f\n\n", J)

    # Validation set only
    elseif mod(epoch, val_every) == 0
        @printf("\nCoffee break:\n")

        # Cost and display of validation set
        J, acc = val_loss_and_accuracy(net, net.data.X_val, net.data.y_val)
        @printf("Validation error (without regⁿ) - %.3f\n\n", J)

    end
end

################################################################################
#
# Completed training - results
#
################################################################################

function show_completed(net)
    if typeof(net.layers[end]) <: LinOutputLayer
        show_completed_regression(net)
    else
        show_completed_classification(net)
    end
end

function show_completed_classification(net)

    @printf "\nCompleted Training:\n"

    if typeof(net.data) <: TData

        # Cost and display of training set
        J, acc = loss_and_accuracy(net, net.data.X_train, net.data.y_train)
        @printf("Training error (with regⁿ) - %.3f", J)
        @printf("  |  Training accuracy - %.1f\n", acc)
        push!(net.full_train_loss, J)

    else
        # Cost and display of training set
        J, acc = loss_and_accuracy(net, net.data.X_train, net.data.y_train)
        @printf("Training error (with regⁿ) - %.3f", J)
        @printf("  |  Training accuracy - %.1f\n", acc)
        push!(net.full_train_loss, J)

        # Cost and display of validation set
        J, acc = val_loss_and_accuracy(net, net.data.X_val, net.data.y_val)
        @printf("Validation error (without regⁿ) - %.3f", J)
        @printf("  |  Validation accuracy - %.1f\n", acc)
        push!(net.val_loss, J)
    end
end

function show_completed_regression(net)

    @printf "\nCompleted Training:\n"

    if typeof(net.data) <: TData

        # Cost and display of training set
        J = loss(net, net.data.X_train, net.data.y_train)
        @printf("Training error (with regⁿ) - %.3f\n", J)
        push!(net.full_train_loss, J)

    else
        # Cost and display of training set
        J = loss(net, net.data.X_train, net.data.y_train)
        @printf("Training error (with regⁿ) - %.3f\n", J)
        push!(net.full_train_loss, J)

        # Cost and display of validation set (without regularisation)
        J = val_loss(net, net.data.X_val, net.data.y_val)
        @printf("Validation error (without regⁿ) - %.3f\n", J)
        push!(net.val_loss, J)
    end
end


using QuantumStatePlots
using QuantumStateBase
using Test

@testset "QuantumStatePlots.jl" begin

ENV["GKSwstype"]="nul"

@testset "plot wigner" begin
    x_range = -5:1.0:5
    p_range = -5:1.0:5
    wf = WignerFunction(x_range, p_range)

    state = VacuumState()
    ws = wf(state)

    file_path = "wigner.png"

    plot_wigner(ws, Heatmap, file_path=file_path)
    @test isfile(file_path)
    isfile(file_path) && rm(file_path)

    plot_wigner(ws, Contour, file_path=file_path)
    @test isfile(file_path)
    isfile(file_path) && rm(file_path)

    plot_wigner(ws, Surface, file_path=file_path)
    @test isfile(file_path)
    isfile(file_path) && rm(file_path)
end

@testset "plot ρ" begin
    file_path = "rho.png"

    plot_ρ(VacuumState(), file_path=file_path)
    @test isfile(file_path)
    isfile(file_path) && rm(file_path)

    plot_ρ(VacuumState(), state_n=5, file_path=file_path)
    @test isfile(file_path)
    isfile(file_path) && rm(file_path)
end

@testset "plot all" begin
    x_range = -5:1.0:5
    p_range = -5:1.0:5
    wf = WignerFunction(x_range, p_range)

    state = VacuumState()
    ws = wf(state)

    file_path = "all.png"

    plot_all(ws, state, file_path=file_path)
    @test isfile(file_path)
    isfile(file_path) && rm(file_path)
end

end


using ITensors
using ITensorGaussianMPS
using LinearAlgebra

# Half filling
N = 20
Nf_up = N ÷ 2
Nf_dn = N ÷ 2
Nf = Nf_up + Nf_dn

@show N, Nf

# Maximum MPS link dimension
_maxlinkdim = 50

@show _maxlinkdim

# DMRG cutoff
_cutoff = 1e-8

# Hopping
t = 1.0

# Electron-electron on-site interaction
U = 1.0

@show t, U

# Make the free fermion Hamiltonian for the up spins
os_up = OpSum()
for n in 1:(N - 1)
  os_up .+= -t, "Cdagup", n, "Cup", n + 1
  os_up .+= -t, "Cdagup", n + 1, "Cup", n
end

# Make the free fermion Hamiltonian for the down spins
os_dn = OpSum()
for n in 1:(N - 1)
  os_dn .+= -t, "Cdagdn", n, "Cdn", n + 1
  os_dn .+= -t, "Cdagdn", n + 1, "Cdn", n
end

# Hopping Hamiltonians for the up and down spins
h_up = hopping_hamiltonian(os_up)
h_dn = hopping_hamiltonian(os_dn)

# Get the Slater determinant
Φ_up = slater_determinant_matrix(h_up, Nf_up)
Φ_dn = slater_determinant_matrix(h_dn, Nf_dn)

# Create an MPS from the slater determinants.
s = siteinds("Electron", N; conserve_qns=true)
println("Making free fermion starting MPS")
@time ψ0 = slater_determinant_to_mps(
  s, Φ_up, Φ_dn; eigval_cutoff=1e-4, cutoff=_cutoff, maxdim=_maxlinkdim
)
@show maxlinkdim(ψ0)

# The total interacting Hamiltonian
os = os_up + os_dn
for n in 1:N
  os .+= U, "Nupdn", n
end
H = MPO(os, s)

println("Free fermion starting state energy")
@show flux(ψ0)
@show inner(ψ0, H, ψ0)

println("\nStart from free fermion state")
sweeps = Sweeps(5)
maxdim!(sweeps, _maxlinkdim)
cutoff!(sweeps, _cutoff)
e, ψ = @time dmrg(H, ψ0, sweeps)
@show e
@show flux(ψ)

using ITensorGaussianMPS: correlation_matrix_to_gmps, correlation_matrix_to_mps, entropy

Λ_up = correlation_matrix(ψ, "Cdagup", "Cup")
Λ_dn = correlation_matrix(ψ, "Cdagdn", "Cdn")
ψ̃0 = correlation_matrix_to_mps(
  s, Λ_up, Λ_dn; eigval_cutoff=1e-2, maxblocksize=4
)
@show inner(ψ̃0, ψ)
@show inner(ψ̃0, H, ψ̃0)


@testset "learning transition small functions" begin
    for x0 ∈ (0.25, 0.5, 0.75,)
        vspace = range(-2.0, stop=2.0, length=5)
        @test Calculus.derivative(_ρ, x0) ≈ ShaleDrillingModel._dρdθρ(x0)
        @test Calculus.derivative(_ρ, x0) ≈ ShaleDrillingModel._dρdσ(x0)

        for u in vspace
            for v in vspace
                @test Calculus.derivative((z) -> _ψ1(u,v,_ρ(z)), x0) ≈ _dψ1dθρ(u,v,_ρ(x0), x0)
                @test Calculus.derivative((z) -> _ψ1(u,v,z), x0) ≈ ShaleDrillingModel._dψ1dρ(u,v,x0)
            end
        end
    end
end

@testset "Frechet derivative of ψ transition matrix" begin

    dpi = similar(tmpv.Πψtmp)
    fdpi = similar(tmpv.Πψtmp)
    h = peturb(σv)
    σ1 = σv - h
    σ2 = σv + h
    hh = σ2 - σ1

    ShaleDrillingModel._βΠψ!(dpi, ψspace, σ1, β)
    ShaleDrillingModel._βΠψ!(fdpi, ψspace, σ2, β)
    fdpi .-= dpi
    fdpi ./= hh
    ShaleDrillingModel._βΠψdθρ!(dpi, ψspace, σv, β)

    @views maxv, idx = findmax(abs.(dpi .- fdpi))
    sub = CartesianIndices(fdpi)[idx]
    @show "worst value is $maxv at $sub for β dΠ/dσ"

    @test 0.0 < maxv < 1e-5
    @test fdpi ≈ dpi

end


"""
# PSH

Module PSH converts seismograms in the ZNE frame into the PSH frame using the free-surface
transform.
"""
module PSH

using SAC

export psh, psh!

"""
    psh_matrix!(α, β, p, M) -> M

Fill in the PSH transfer matrix `M`, given the P-wave velocity `α`, S-wave velocity `β`
and the slowness `p`.  Velocities are in km/s, slowness in s/km.
"""
function psh_matrix!{T}(a, b, p, M::AbstractArray{T,2})
    size(M) == (3,3) || error("Transfer matrix must have size (3,3)")
    isevanescent(a, b, p) &&
        warn("psh_matrix!: Evanescent wave (vp=$a, vs=$b, p=$p)")
    qa = sqrt(1/a^2 - p^2)
    qb = sqrt(1/b^2 - p^2)
    # Kennett, GJI, 1991
    Vpz = -(1 - 2b^2*p^2)/(2a*qa)
    Vpr = p*b^2/a
    Vsz = p*b
    Vsr = (1 - 2b^2*p^2)/(2b*qb)
    Vht = 1/2
    M[:,:] .= [Vpz Vsz  0
               Vpr Vsr  0
                0   0  Vht]
end

"""
    psh_matrix(α, β, p) -> M

Return the PSH transfer matrix M, given the P-wave velocity `α`, S-wave velocity `β`
and the slowness `p`.  Velocities are in km/s, slowness in s/km.
"""
psh_matrix(a, b, p) = psh_matrix!(a, b, p, Array{Complex{Float64}}(3,3))

"""
    isevanescent(α, β, p) -> ::Bool

Return true if the combination of Vp `α`, Vs `β` and slowness `p` lead to an evanescent
wave at the surface.
"""
function isevanescent(a, b, p)
    p = Float64(p)
    a, b = Float64(a), Float64(b)
    1/a^2 - p^2 <= 0.0 || 1/b^2 - p^2 <= 0.0
end

"""
    psh!(Z, R, T, α, β, p) -> P, S, H

Converts three arrays, `SACtr`s or arrays of `SACtr`s in ZRT orientation into
PSH orientation, using the supplied Vp `α`, Vs `β` (km/s) and slowness `p` (s/km).

Traces are updated to be in the order P (like vertical), S (like radial), H (like transverse).
"""
function psh!{F}(Z::Array{F,1}, R::Array{F,1}, T::Array{F,1}, a, b, p)
	length(Z) == length(R) == length(T) || error("psh!: Traces must be same length")
	M = psh_matrix(a, b, p)
    # Kennett 1991
    Z, R, T = M[1,1]*Z + M[1,2]*R + M[1,3]*T,
              M[2,1]*Z + M[2,2]*R + M[2,3]*T,
              M[3,1]*Z + M[3,2]*R + M[3,3]*T
    Z, R, T
end
function psh!(Z::SACtr, R::SACtr, T::SACtr, a, b, p)
    P, S, H = psh(Z.t, R.t, T.t, a, b, p)
    Z.t .= real.(P)
    R.t .= real.(S)
    T.t .= real.(H)
    for t in (Z, R, T) SAC.update_headers!(t) end
    Z.kcmpnm, R.kcmpnm, T.kcmpnm = "P", "S", "H"
    Z, R, T
end
function psh!(Z::Array{SACtr}, R::Array{SACtr}, T::Array{SACtr}, a, b, p)
    for (z, r, t) in zip(Z, R, T) psh!(z, r, t, a, b, p) end
    Z, R, T
end

"""
    psh(Z, R, T, α, β, p) -> P, S, H

Convert three arrays, `SACtr`s or arrays of `SACtr`s in ZRT orientation into
PSH orientation, using the supplied Vp `α`, Vs `β` (km/s) and slowness `p` (s/km),
where `P` is parallel to P-wave motion positive away from the source, `S` is parallel
to SV and `H` is the transverse direction.
"""
psh(Z, R, T, a, b, p) = psh!(deepcopy(Z), deepcopy(R), deepcopy(T), a, b, p)

end # module


using QuadGK
using LambertW
using Base.Threads

"""
Compute range y1, y2 of the Converse Jensen Inequality:
from 0 to y1: ρ1 sum to x1;
from 0 to y2: ρ2 sum to x2.
"""
function converse_jensen_range(x1::Float64, x2::Float64, xmax::Float64=1-1e-10)::Tuple{Float64,Float64}
    # Cap x1, x2 by xmax because LambertW package may have large error near 1/ℯ
    x1 = min(x1, xmax)
    x2 = min(x2, xmax)
    if 1-(1-x1)*(1-log(1-x1)) >= x2
        # case 1: λ2 <= λ1.
        # Solving 1-(1+λ)exp(-λ)=x2 gives λ2=-1-W(-(1-x2)/ℯ) (-1 branch)
        λ2 = -1-lambertw(-(1-x2)/ℯ, -1)
    else
        # case 1: λ2 > λ1
        # Solving 2-(2+λ)exp(-λ)=x1+x2 gives λ2=-2-W(-(2-x1-x2)/ℯ^2) (-1 branch)
        λ2 = -2-lambertw(-(2-x1-x2)/ℯ^2, -1)
    end
    y1 = x1
    y2 = min(1-exp(-λ2), xmax)
    return y1, y2
end

"""Compute the contribution at any y = 1-exp(-λ) to (d/dt) log(d_(x1, x2)(t))."""
function dlog_double_dy(y::Float64, y1::Float64, y2::Float64, x1::Float64, x2::Float64, t::Float64, logd::Float64)::Float64
    ρ1 = (y<=y1) ? 1.            : 0.
    ρ2 = (y<=y2) ? 1-log(1-y)-ρ1 : 0.
    w1 = exp((2*x1+x2)*t+logd)
    w2 = exp((x1+2*x2)*t+logd)
    return (w1*ρ1+w2*ρ2)/((max(w1-1,0.)*ρ1+max(w2-1,0.)*ρ2)*(1-y)+1)
end

"""Compute the value of log(q_x(t)) for all 0 <= t <=1 that is a multiple of dt, for a given x."""
function logq(x::Float64, dx::Float64=1e-3, dt::Float64=1e-3, xmax::Float64=1-1e-10, rtol::Float64=1e-5)::Array{Float64,1}

    x_grid = Array(0.:dx:1-dx)
    nx = length(x_grid)

    t_grid = Array(0.:dt:1-dt)
    nt = length(t_grid)

    logq_grid = zeros(nx,nt)

    @threads for i in 1:nx
        x1, x2 = (x>=x_grid[i]) ? [x, x_grid[i]] : [x_grid[i], x]
        y1, y2 = converse_jensen_range(x1, x2, xmax)
        logd = 0.
        for j in 1:(nt-1)
            t = t_grid[j]
            dlogd, err = quadgk(y -> dlog_double_dy(y, y1, y2, x1, x2, t, logd), 0., xmax, rtol=rtol)
            logd2      = logd-(dlogd-err)*dt
            dlogd, err = quadgk(y -> dlog_double_dy(y, y1, y2, x1, x2, t, logd2), 0., xmax, rtol=rtol)
            logd      -= (dlogd-err)*dt
            logd       = min(logd, -(x1+x2)*(t+dt))   # use trivial bound if numerical bound is worse (for small t)
            logq_grid[i,j+1] = min(logd+(x_grid[i]+dx)*(t+dt), -x*(t+dt))
        end
    end

    return vec(maximum(logq_grid, dims=1))
end

"""Compute match probability bound of an vertex with LP matched level x"""
function level2(x::Float64=1., dx::Float64=1e-3, dt::Float64=1e-3, xmax::Float64=1-1e-10, rtol::Float64=1e-5)::Float64
    log_single = 0.
    t_grid = Array(0.:dt:1-dt)
    nt = length(t_grid)
    logq_grid = logq(x, dx, dt, xmax, rtol)
    for i in 1:nt
        t = t_grid[i]
        y = 1-exp(-x*t+logq_grid[i]-log_single)
        dlog_single = (y==0.) ? x : -log(1-x*y)/y
        log_single2 = log_single - dlog_single*dt
        y = 1-exp(-x*t+logq_grid[i]-log_single2)
        dlog_single = (y==0.) ? x : -log(1-x*y)/y
        log_single -= dlog_single*dt
    end
    return 1-exp(log_single)
end


function cross2d(x,y)
    return x[1]*y[2] - x[2]*y[1]
end

function intersect_line(ls1::NTuple{2,T},ls2::NTuple{2,S}) where {T<:AbstractVector,S<:AbstractVector}
    r = ls1[2] - ls1[1]
    s = ls2[2] - ls2[1]
    qmp = ls2[1]-ls1[1]
    d = cross2d(r,s)
    un = cross2d(qmp,r)
    d == 0.0 && (return un == 0.0)
    return (0 < un/d < 1) && (0 < cross2d(qmp, s)/d < 1)
end

function get_angle(p_pro,p_cur,p_pre)
    r = p_pro - p_cur
    p = p_pre - p_cur
    t = atan(cross2d(p,r),dot(p,r))
    return t > 0 ? t : t + 2pi
end

function intersect_hull(ls, hull)
    v = hull.vertices
    np = length(v)
    if np < 3
        return false
    end
    for i = 1:(np-2)
        if intersect_line(ls,(v[i+1],v[i]))
            return true
        end
    end
    return false
end

# In Julia v0.7 and above, indmax cannot be used with generators.
# In a future version, we should be able to replace this function
# with Base.argmax(f, itr)
function indmax_f(f::Function, itr::AbstractVector)
    imax = firstindex(itr)
    vmax = f(itr[imax])
    for i in (firstindex(itr)+1):lastindex(itr)
        v = f(itr[i])
        if v > vmax
            vmax = v
            imax = i
        end
    end
    imax
end

function concave_hull(tree::KDTree, k::Int)
    npoints = length(tree.data)
    k = clamp(k, 3, npoints-1)
    hull = Hull(eltype(tree.data), k)

    #Start at point with largest x value
    i0 = indmax_f(v -> v[1], tree.data)

    ishull = zeros(Bool,npoints)
    ishull[i0] = true

    #Initiate hull
    p0 = copy(tree.data[i0])
    push!(hull.vertices,p0)

    #Allocate point arrays
    p_cur = copy(p0)
    p_pre = p_cur + [1, 0]
    p_pro = copy(p0)
    d01 = zero(eltype(p0))

    nstep = 1
    npick = 1
    while npick < npoints
        if nstep == 3
            npick = npick - 1
            ishull[i0] = false
        end
        kk = clamp(k, 3 , npoints-npick)
        kind, kdist = knn(tree, p_cur, kk, true, i -> ishull[i])
        angles = collect(get_angle(tree.data[i], p_cur, p_pre) for i in kind)
        w = sortperm(angles,rev=true)
        inter = true
        for i in kind[w]
            #p_pre(0) ---------- d02 --------- p_pro(2)
            #p_pre(0) --d01-- p_cur(1) --d12-- p_pro(2)
            p_pro = tree.data[i]
            d12 = norm(p_pro - p_cur)
            d02 = norm(p_pro - p_pre)
            if ~intersect_hull((p_pro,p_cur),hull) || (d01 + d12) == d02
                p_pre = p_cur
                d01 = d12
                p_cur = tree.data[i]
                ishull[i] = true
                inter = false
                break
            end
        end

        if inter
            if k+1 >= npoints
                @info("Unable to construct concave hull")
                return hull
            end
            return concave_hull(tree, k+1)
        end

        if p_cur != p0
            push!(hull.vertices, copy(p_cur))
            npick = npick + 1
        else
            break
        end
        nstep = nstep + 1
    end

    if any(~in_hull(p,hull) for p in tree.data[.~(ishull)])
        if k+1 >= npoints
            @info("Unable to construct concave hull")
            return hull
        end
        return concave_hull(tree, k+1)
    end

    return Hull(hull.vertices, hull.k, true)
end

function concave_hull(points::Vector, k::Int = 3)
    p = unique(points)
    npoints = length(p)
    npoints < 3 && throw(ArgumentError("Number of unique points should be greater than 2"))
    npoints == 3 && return Hull(points,k, true)
    tree = KDTree(hcat(points...),reorder=false)
    return concave_hull(tree, k)
end


using FFTW:fft
export MelFilters

mutable struct MelFilters
    winlen::Int   # 分帧参数-帧长
    stride::Int   # 分帧参数-帧移
    nbanks::Int   # 梅尔滤波器个数
    nffts ::Int   # 傅立叶变换点数
    alpha::Real   # 预加重系数，默认0.97
    fs::Int       # 采样率，一般16kHz
    maxfreq::Int  # 频域最大有效频率下标
    eps::Real     # 防止下溢系数
    winfunc::Function       # 窗计算函数，如 hanning
    window ::AbstractArray  # 窗函数数值表示
    melbank::AbstractArray  # mel滤波器矩阵

    function MelFilters(;
        winlen::Int = 256,
        stride::Int = winlen>>1,
        nffts ::Int = winlen,
        nbanks::Int = 32,
        alpha::Real = 0.97,
        fs::Int     = 16000,
        eps::Real   = 1e-6,
        winfunc::Function = hanning)
        @assert nffts>=winlen "nffts ≥ winlen, but got nffts=$nffts, winlen=$winlen"
        maxfreq = nffts>>1
        window  = winfunc(winlen)
        melbank = filterbanks(nbanks, nffts, fs)
        new(winlen,stride, nbanks,nffts,alpha,fs, maxfreq,eps, winfunc, window,melbank)
    end
end

function filterbanks(nbanks::Int, nffts::Int, fs::Int)
    MAX   = nffts>>1;                    # 正频率部分的下标最大值
    freq  = (0:(MAX-1))/MAX * fs/2;      # 下标映射到频率
    Fmel  = 2595*log10.(1 .+ freq/700);  # 频率映射到梅尔频率
    dFmel = Fmel[MAX]/(nbanks+1);        # 将Mel带平分
    bank  = zeros(nbanks, MAX);          # 滤波器的频域权重系数
    cFmel = 0.0;                         # 每个Mel频带的中心Mel频率
    for n = 1:nbanks
        cFmel = cFmel + dFmel
        for m = 1:MAX
            if ( Fmel[m] >= cFmel-dFmel ) && ( Fmel[m] <= cFmel+dFmel )
                bank[n,m] = 1.0 - abs( Fmel[m] - cFmel )/dFmel
            end
        end
    end
    return bank
end

function (filter::MelFilters)(wav, func::Union{Function,Nothing}=log)
    W = filter.melbank
    α = filter.alpha
    B = filter.eps
    F = filter.maxfreq
    winlen = filter.winlen
    stride = filter.stride
    window = filter.window
    nffts  = filter.nffts

    d = filterwav(wav, α)                            # 滤波
    X = spectrum(d, window, winlen, stride, nffts)   # 功率谱

    if func !== nothing
        return func.(W * X .+ B)
    else
        return W * X
    end
end

function Base.show(io::IO, f::MelFilters)
    winlen = trunc(f.winlen / f.fs * 1000, digits=3)
    stride = trunc(f.stride / f.fs * 1000, digits=3)
    println(io, "———————————————————————")
    println(io, " winfunc = $(f.winfunc)")
    println(io, "  winlen = $(f.winlen) ($winlen ms)")
    println(io, "  stride = $(f.stride) ($stride ms)")
    println(io, "   nffts = $(f.nffts)")
    println(io, " maxfreq = $(f.maxfreq)")
    println(io, "  nbanks = $(f.nbanks)")
    println(io, "   alpha = $(f.alpha)")
    println(io, "     eps = $(f.eps)")
    println(io, "      fs = $(f.fs) Hz")
    println(io, "———————————————————————")
end


# Author: Lukasz Bratos

# Funkcja zwraca wartosc wyrazenia sin(x) + cos(3x) dla danego x
function f(x :: Float64)
    return sin(x) + cos(3 * x)
end

# Funkcja zwraca wartosc wyrazenia cos(x) - 3sin(3x) dla danego x (co jest pochodna wyrazenia sin(x) + cos(3x))
function derivativeOfF(x :: Float64)
    return cos(x) - 3 * sin(3 * x)
end

# Funkcja wylicza wartosc pochodnej funkcji sin(x) + cos(3x) dla danego x
function derivative(func, x0, h)
    return (func(x0 + h) - func(x0)) / h
end

# Wypisywanie wynikow w formacie pod LaTeXa
for i in 0:54
    println("\$ 2^{-$i} \$ & ", 1.0 + 2.0 ^ (-i), " & ", derivative(f, 1.0, 2.0 ^ (-i)), " & ", abs(derivativeOfF(1.0) - derivative(f, 1.0, 2.0 ^ (-i))), " \\\\\n\\hline")
end

println("f'(1.0) = ", derivativeOfF(1.0))


# (a<1) & (a>0)
# (β<1) & (β>0)
# -a+β < 0

using NLopt
using Revise

function func(x, p)
    # simple linear func
    x*p[1] + p[2]
end

function cost(p, xvec, actual)
    # predicted
    pred = func.(xvec, (p,))

    # cost
    sum((pred .- actual).^2)
end

function constraint1(p)
    p[1] - p[2]
end

function main()
    xvec = Float64.(0.0:0.1:10)
    actualp = [6.0, 5.0]
    actualresponse = func.(xvec, (actualp,))

    opt = Opt(:LN_COBYLA, 2)
    opt.lower_bounds = [-Inf, -Inf]
    opt.upper_bounds = [Inf, Inf]
    opt.xtol_rel = 1e-8

    opt.min_objective = (params, grad) -> cost(params, xvec, actualresponse)

    inequality_constraint!(opt, (p, g) -> constraint1(p))

    (minf,minx,ret) = optimize(opt, [1.234, 3.678])
    numevals = opt.numevals # the number of function evaluations
    println("got $minf at $minx after $numevals iterations (returned $ret)")
end

main()


module MethodOfLines
    using LinearAlgebra
    using SciMLBase
    using DiffEqBase
    using DiffEqOperators
    using ModelingToolkit
    using ModelingToolkit: operation, istree, arguments, variable
    using SymbolicUtils, Symbolics
    using SymbolicUtils: operation, arguments
    using IfElse
    import DomainSets

    include("discretization/fornberg.jl")

    include("grid_types.jl")
    include("MOLFiniteDifference.jl")

    include("discretization/discretize_vars.jl")
    include("MOL_utils.jl")
    include("interiormap.jl")

    include("discretization/differential_discretizer.jl")
    include("bcs/boundary_types.jl")

    include("periodic_map.jl")

    include("discretization/generate_finite_difference_rules.jl")

    include("bcs/generate_bc_eqs.jl")

    include("error_analysis.jl")
    include("discretization/MOL_discretization.jl")

    export MOLFiniteDifference, discretize, symbolic_discretize, ODEFunctionExpr, generate_code, grid_align, edge_align, center_align, get_discrete

end


mutable struct ExactMonotonicRiskRegion <: AbstractRiskRegion
    β::Float64
    surv::Function
end

function ∈(x::VecF64, Ω::ExactMonotonicRiskRegion)
    if Ω.surv(x) < 1 - Ω.β
        return true
    else
        return false
    end
end

mutable struct MonotonicRiskRegion <: AbstractRiskRegion
    β::Float64
    nonrisk_frontier::LinkedList{VecF64}
    risk_frontier::LinkedList{VecF64}
    surv::SurvivorApproximator
    function MonotonicRiskRegion(surv::SurvivorApproximator, β::Float64)
        return new(β,nil(VecF64),nil(VecF64),surv)
    end
end

function add_to_risk_frontier!(Ω::AbstractRiskRegion, x::AbstractVector{T}) where T<:Real
    l = Ω.risk_frontier
    if isa(l, Nil{VecF64})
        return Ω.risk_frontier=list(x[:])
    end

    while !isa(l, Nil{VecF64})
        pf = head(l)
        if vec_isless(x, pf)
            if isa(tail(l),Nil{VecF64})
                return Ω.risk_frontier=list(x[:])
            else
                decapitate!(l)
            end
        end
        l = tail(l)
    end
    return Ω.risk_frontier = cons(x[:], Ω.risk_frontier)
end

function add_to_nonrisk_frontier!(Ω::AbstractRiskRegion, x::AbstractVector{T}) where T<:Real
    l = Ω.nonrisk_frontier
    if isa(l, Nil{VecF64})
        return Ω.nonrisk_frontier=list(x[:])
    end

    while !isa(l, Nil{VecF64})
        pf = head(l)
        if vec_isless(pf, x)
            if isa(tail(l),Nil{VecF64})
                return Ω.nonrisk_frontier=list(x[:])
            else
                decapitate!(l)
            end
        end
        l = tail(l)
    end
    return Ω.nonrisk_frontier = cons(x[:], Ω.nonrisk_frontier)
end

function ∈(x::VecF64, Ω::MonotonicRiskRegion)
    if below_nonrisk_frontier(x, Ω.nonrisk_frontier) return false
    elseif above_risk_frontier(x, Ω.risk_frontier) return true
    else
        prob, err = Ω.surv(x)
        if prob > 1 - Ω.β
            if prob - err > (1 - Ω.β)
                add_to_nonrisk_frontier!(Ω, x)
            end
            return false
        else
            if prob + err < (1-Ω.β)
                add_to_risk_frontier!(Ω, x)
            end
            return true
        end
    end
end

function prob_nonrisk(dist::Distribution{Multivariate, Continuous}, β_list::Vector{Float64}, num_points_surv::Int = 50000, num_points_nonrisk::Int = 10000, α::Float64=0.99)
    sample = rand(dist, num_points_nonrisk)
    surv = SurvivorApproximator(rand(dist, num_points_surv), α)
    Ω = Dict{Float64, MonotonicRiskRegion}()
    probs = Dict{Float64,Float64}()

    for β in β_list
        Ω[β] = MonotonicRiskRegion(surv, β)
        probs[β] = 0.0
    end

    for i in 1:num_points_nonrisk
        point = view(sample, :, i)
        calc_surv_prob = false
        surv_prob = 0.0
        for β in β_list
            if calc_surv_prob
                probs[β] += (surv_prob > 1-β ? 1 : 0)
                continue
            else
                shortcut = false
                # Does sampled point dominate any point in risk region?
                if below_nonrisk_frontier(point, Ω[β].nonrisk_frontier)
                    shortcut=true
                    probs[β] += 1
                elseif above_risk_frontier(point, Ω[β].risk_frontier)
                    shortcut=true
                end

                if !shortcut
                    surv_prob, err = surv(point)
                    if surv_prob > 1 - β
                        probs[β]+=1
                        if surv_prob - err > 1-β
                            add_to_nonrisk_frontier!(Ω[β], point)
                        end
                    else
                        if surv_prob + err < 1-β
                            add_to_risk_frontier!(Ω[β], point)
                        end
                    end
                    calc_surv_prob=true
                end
            end
        end
    end

    for β in β_list
        probs[β]/= num_points_nonrisk
    end
    return probs
end


@diffusion_process JansenRit{T} begin
    :dimensions
    process --> 6
    wiener --> 1
    :parameters
    (A, a, B, b, C, α1, α2, νmax, v, r, μy, σy) --> T
    :additional
    constdiff --> true
end
# θ
#param_names(P::JansenRit) = (:A, :a, :B, :b, :C, :α1, :α2, :νmax, :v, :r, :μy, :σy)
#θ =[3.25, 100.0, 22.0, 50.0, 135.0, 0.8, 0.25, 5.0, 6.0, 0.56, 220.0, 2000.0]
#P = JansenRit(θ...)

sigm(x, P::JansenRit) = P.νmax / (1 + exp(P.r*(P.v - x)))
μy(t, P::JansenRit) = P.μy #constant
C1(P::JansenRit) = P.C
C2(P::JansenRit) = P.α1*P.C
C3(P::JansenRit) = P.α2*P.C
C4(P::JansenRit) = P.α2*P.C

function b(t, x, P::JansenRit)
    @SVector [
        x[4],
        x[5],
        x[6],
        P.A*P.a*(sigm(x[2] - x[3], P)) - 2P.a*x[4] - P.a*P.a*x[1],
        P.A*P.a*(μy(t, P) + C2(P)*sigm(C1(P)*x[1], P)) - 2P.a*x[5] - P.a*P.a*x[2],
        P.B*P.b*(C4(P)*sigm(C3(P)*x[1], P)) - 2P.b*x[6] - P.b*P.b*x[3]]
end

function σ(t, x, P::JansenRit)
    @SMatrix [
        0.0;
        0.0;
        0.0;
        0.0;
        P.σy;
        0.0
    ]
end


export getConvDiffMatrix

"""
function getConvDiffMatrix(M::RegularMesh,sig::Float64,v,iddir::Array{Int,1},idneu::Array{Int,1},
	                      iddirn::Array{Int,1},idneun::Array{Int,1},idint::Array{Int,1})

builds convection diffusion matrix on interior nodes and returns contribution of dirichlet and neuman boundary.
"""
function getConvDiffMatrix(M::RegularMesh,sig::Float64,v,iddir::Array{Int,1},idneu::Array{Int,1}, iddirn::Array{Int,1},idneun::Array{Int,1},idint::Array{Int,1})

	# biuld matrix on padded mesh
	Lap  = getLaplacianMatrix(M)
	Conv = getConvectionMatrix(M,v)
	Apad = sig*Lap + Conv

	# projection for gost point neighbors
	Id          = sparse(1.0I, prod(M.n.+2), prod(M.n.+2))[:,idint]
	Id[iddir,:] = -Id[iddirn,:]  # dirichlect boundary points
	Id[idneu,:] =  Id[idneun,:]  # neuman boundary points

	# matrix for interior nodes
	A    = Apad[idint,:]*Id

	# boundary nodes
	Adir = Apad[idint,iddir]
	Aneu = Apad[idint,idneu]

	return A, Adir, Aneu
end


"""
    SemigroupMethod(epsilon, delta, max_iter)

Semigroup method from Algorithm 1 in [1].

[1] Taghvaei, A., & Mehta, P. G. (2016). Gain function approximation in the feedback particle filter. In 2016 IEEE 55th Conference on Decision and Control (CDC) (pp. 5446–5452). IEEE. https://doi.org/10.1109/CDC.2016.7799105

    SemigroupMethod(epsilon, delta, max_iter, lambda)

Semigroup method with regularization parameter `lambda`.
"""
struct SemigroupMethod <: GainEstimationMethod
    epsilon::Float64
    delta::Float64
    max_iter::Int
    lambda::Float64
    SemigroupMethod(eps, delta, max_iter, lambda=0.) = new(eps, delta, max_iter, lambda)
end

SemigroupMethod(epsilon::Float64, delta::Float64) = SemigroupMethod(epsilon, delta, 100)

function solve!(eq::PoissonEquation, method::SemigroupMethod)
    N = length(eq.positions)
    H = Htilde(eq)
    broadcast!(*, H, H, method.epsilon)

    # compute T operator
    T = zeros(eltype(eq.positions), N, N)
    n = state_dim(eq)
    for i in 1:N
        T[i,i] = one(eltype(eq.positions))
        for j in i+1:N
            sq = zero(eltype(eq.positions))
            for l in 1:n
                sq += (eq.positions[l,i] - eq.positions[l,j])^2
            end
            T[i,j] = exp(-sq/(4*method.epsilon))
            T[j,i] = T[i,j]
        end
    end
    broadcast!(/, T, T, sqrt.(sum(T,dims=1) .* sum(T,dims=2)))
    broadcast!(/, T, T, sum(T,dims=2))

    # add noise to regularize T
    if method.lambda > 0.
        for i in 1:N
            T[i,i] -=  method.lambda * rand(Distributions.Uniform(0.9,1))
        end
    end

    # solve fixed-point equation for potential
    newpotential = copy(eq.potential)::Array{Float64,2}
    fluctuation = 1.
    n = 1
    while fluctuation > method.delta
        LinearAlgebra.mul!(newpotential, eq.potential, T')
        broadcast!(+, newpotential, newpotential, H)
        broadcast!(-, newpotential, newpotential, Statistics.mean(newpotential, dims=2))
        fluctuation = maximum(abs.(newpotential - eq.potential))
        eq.potential .= newpotential
        n += 1
        if n == method.max_iter
            print("!")
            break
        end
    end

    # compute gain from potential
    gainhelper_semigroup!(eq, T)
    broadcast!(/, eq.gain, eq.gain, 2*method.epsilon)
end

function gainhelper_semigroup!(eq::PoissonEquation, T::AbstractMatrix)
    pos     = eq.positions
    pot     = eq.potential
    gain    = eq.gain

    n       = state_dim(eq)
    m       = obs_dim(eq)
    N       = no_of_particles(eq)

    @inbounds for l in 1:n, k in 1:m, i in 1:N
        Tpotpos = zero(eltype(gain))
        Tpot    = zero(eltype(gain))
        Tpos    = zero(eltype(gain))
        for j in 1:N
            Tpot        += T[i,j] * pot[k,j]
            Tpos        += T[i,j] * pos[l,j]
            Tpotpos     += T[i,j] * pot[k,j] * pos[l,j]
        end
        gain[l,i,k] = Tpotpos - Tpot * Tpos
    end
end


using Test, AutoARIMA

@testset "Mathematical Properties" begin

    @test reduce(&, yule_walker(seriesF, 15) .≈ levinson_durbin(seriesF, 15))

    @testset "MA(q) forecast after q steps should be equal to the constant" begin
        μ = 0.1
        ma = MAModel{3}(μ,[0.2,0.3,-0.5],1.0)
        z = [0.2,0.5,0.1]
        f = Float64[]
        for i in 1:10
            zhat = forecast(ma,z)
            push!(z,zhat)
            push!(f,zhat)
        end
        @test all(f[4:end] .≈ μ)
    end

    @testset "BoxCox" begin
        z = rand(1000)
        @test z ≈ inv_boxcox(boxcox(z,0.1),0.1)
        @test z ≈ inv_boxcox(boxcox(z,0.0),0.0)
    end

end


function from_file(filename, sep = '\n')
    s = open(filename) do file
        read(file, String)
    end
    str_sudokus = split(strip(s), sep)
    grids = AbstractArray[]
    for str_sudoku in str_sudokus
        str_sudoku = replace(str_sudoku, "." => "0")
        one_line_grid = parse.(Int, split(str_sudoku, ""))
        grid = reshape(one_line_grid, 9, 9)
        push!(grids, grid)
    end
    return grids
end

function solve_sudoku(grid)
    m = CS.Optimizer(logging = [])

    x = [[MOI.add_constrained_variable(m, MOI.Integer()) for i in 1:9] for j in 1:9]
    for r in 1:9, c in 1:9
        MOI.add_constraint(m, x[r][c][1], MOI.GreaterThan(1.0))
        MOI.add_constraint(m, x[r][c][1], MOI.LessThan(9.0))
    end

    # set variables
    for r in 1:9, c in 1:9
        if grid[r, c] != 0
            sat = [MOI.ScalarAffineTerm(1.0, x[r][c][1])]
            MOI.add_constraint(
                m,
                MOI.ScalarAffineFunction{Float64}(sat, 0.0),
                MOI.EqualTo(convert(Float64, grid[r, c])),
            )
        end
    end
    # sudoku constraints
    for r in 1:9
        MOI.add_constraint(
            m,
            MOI.VectorOfVariables([x[r][c][1] for c in 1:9]),
            CS.AllDifferentSetInternal(9),
        )
    end
    for c in 1:9
        MOI.add_constraint(
            m,
            MOI.VectorOfVariables([x[r][c][1] for r in 1:9]),
            CS.AllDifferentSetInternal(9),
        )
    end
    variables = [MOI.VariableIndex(0) for _ in 1:9]
    for br in 0:2
        for bc in 0:2
            variables_i = 1
            for i in (br * 3 + 1):((br + 1) * 3), j in (bc * 3 + 1):((bc + 1) * 3)
                variables[variables_i] = x[i][j][1]
                variables_i += 1
            end
            MOI.add_constraint(m, variables, CS.AllDifferentSetInternal(9))
        end
    end

    MOI.optimize!(m)
end