#   Copyright (c) 2020-2022 Velexi Corporation
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

"""
Unit tests for the dual(x, y) function
"""

# --- Imports

# Standard library
import InteractiveUtils.subtypes
import LinearAlgebra
using Test

# GeometricAlgebra.jl
using GeometricAlgebra

# --- File inclusions

# Test utilities
include("test_utils.jl")

# --- Tests

# ------ M::Multivector

@testset "dual(M::Multivector, N::Blade)" begin
    @test_skip 1
end

@testset "dual(M::Multivector, N::Pseudoscalar)" begin
    @test_skip 1
end

@testset "dual(M::Multivector, N::Scalar)" begin
    @test_skip 1
end

@testset "dual(M::Multivector, N::One)" begin
    @test_skip 1
end

# --- B::Blade

@testset "dual(B::Blade, C::Blade)" begin
    # --- Preparations

    test_dim = 20

    C = Blade(rand(test_dim, 7))

    # --- Tests

    # ------ dim(B) == dim(C), grade(B) < grade(C)

    for grade_B in 2:5
        B_vectors = basis(C)[:, 1:grade_B] * rand(grade_B, grade_B)
        B = Blade(B_vectors)

        B_dual_C = dual(B, C)

        # --- Check dim, grade, and norm

        @test dim(B_dual_C) == test_dim
        @test grade(B_dual_C) == grade(C) - grade_B
        @test norm(B_dual_C) == norm(B)

        # --- Check that B and dual(B) are orthogonal complements

        @test LinearAlgebra.norm(transpose(basis(B_dual_C)) * basis(B)) <
            10 * eps(Float64)

        # --- Check sign(dual(B, C))

        # Compute sign of I_C formed from basis(B) and basis(dual(B, C))
        sign_Q = sign(LinearAlgebra.det(
            transpose(hcat(basis(B), basis(B_dual_C))) * basis(C)))

        # Compute expected_sign
        expected_sign = sign(B) * sign_Q

        # Account for sign of I_C^{-1} relative to I_C
        if mod(grade(C), 4) >= 2
            expected_sign = -expected_sign
        end

        # Account for reversals required to eliminate B
        if mod(grade(B), 4) >= 2
            expected_sign = -expected_sign
        end

        @test sign(B_dual_C) == expected_sign

        # Check dual(dual(B, C), C) = (-1)^(grade(C) * (grade(C) - 1) / 2) B
        if mod(grade(C), 4) < 2
            @test dual(B_dual_C, C) ≈ B
        else
            @test dual(B_dual_C, C) ≈ -B
        end
    end

    # ------ dim(B) == dim(C), grade(B) == grade(C)

    coefficients = rand(grade(C), grade(C))
    B_vectors = basis(C) * coefficients
    B = Blade(B_vectors)

    relative_orientation =
        sign(LinearAlgebra.det(transpose(basis(B)) * basis(C)))
    expected_result = mod(grade(B), 4) < 2 ?
        Scalar(relative_orientation * volume(B)) :
        Scalar(-relative_orientation * volume(B))

    @test dual(B, C) ≈ expected_result

    # ------ dim(B) != dim(C)

    B = Blade(rand(test_dim + 1, grade(C) - 1))
    @test_throws DimensionMismatch dual(B, C)

    # ------ B not contained in C

    for grade_B in 2:5
        # grade(B) > grade(C)
        B = Blade(rand(test_dim, grade(C) + 1))
        @test_throws ArgumentError("`B` not contained in `C`") dual(B, C)

        # grade(B) <= grade(C), B not contained in C
        B = Blade(rand(test_dim, grade_B))
        while LinearAlgebra.norm(reject(basis(B), C)) < sqrt(eps(Float64))
            B = Blade(rand(test_dim, grade_B))
        end
        @test_throws ArgumentError("`B` not contained in `C`") dual(B, C)
    end
end

@testset "dual(B::Blade, C::Pseudoscalar)" begin
    # --- Preparations

    test_dim = 10

    test_value = rand()
    test_value = rand() > 0.5 ? test_value : -test_value

    # --- Tests

    # dim(B) == dim(C)
    for test_grade in 2:5
        B = Blade(rand(test_dim, test_grade))
        C = Pseudoscalar(test_dim, test_value)

        B_dual_C = dual(B, C)
        @test B_dual_C == dual(B)

        # Check dual(dual(B, C), C) = (-1)^(grade(C) * (grade(C) - 1) / 2) B
        if mod(grade(C), 4) < 2
            @test dual(B_dual_C, C) ≈ B
        else
            @test dual(B_dual_C, C) ≈ -B
        end
    end

    # dim(B) != dim(C)
    B = Blade(rand(test_dim, 3))
    C = Pseudoscalar(test_dim + 1, test_value)
    @test_throws DimensionMismatch dual(B, C)
end

@testset "dual(B::Blade, C::Scalar)" begin
    test_dim = 9
    test_value = rand()
    test_value = rand() > 0.5 ? test_value : -test_value

    for test_grade in 5:8
        B = Blade(randn(test_dim, test_grade))
        C = Scalar(test_value)
        @test_throws ArgumentError("`B` not contained in `C`") dual(B, C)
    end
end

@testset "dual(B::Blade, C::One)" begin
    test_dim = 10
    for test_grade in 5:8
        B = Blade(randn(test_dim, test_grade))
        C = One()
        @test_throws ArgumentError("`B` not contained in `C`") dual(B, C)
    end
end

@testset "dual(B::Blade, C::Real)" begin
    test_dim = 9
    test_value = get_random_value()

    for test_grade in 5:8
        B = Blade(randn(test_dim, test_grade))
        C = test_value
        @test_throws ArgumentError("`B` not contained in `C`") dual(B, C)
    end
end

@testset "dual(B::Blade, C::Vector)" begin
    # --- Preparations

    test_dim = 20
    grade_C = 1

    C = rand(test_dim, grade_C)[:, grade_C]

    # --- Tests

    # ------ dim(B) == dim(C), C is proportional to B

    B_vectors = C * rand()
    B = Blade(B_vectors)

    relative_orientation =
        sign(LinearAlgebra.det(transpose(basis(B)) * reshape(C, length(C), 1)))
    expected_result = Scalar(relative_orientation * volume(B))

    B_dual_C = dual(B, C)
    @test B_dual_C ≈ expected_result

    # ------ dim(B) != dim(C)

    B = Blade(rand(test_dim + 1, grade_C))
    @test_throws DimensionMismatch dual(B, C)

    # ------ B not contained in C

    # grade(B) > grade(C)
    B = Blade(rand(test_dim, grade_C + 1))
    @test_throws ArgumentError("`B` not contained in `C`") dual(B, C)

    # grade(B) == grade(C), B not contained in C
    B_basis = copy(C)
    B_basis[1] += 1
    B = Blade(B_basis)
    @test_throws ArgumentError("`B` not contained in `C`") dual(B, C)
end

# --- B::Pseudoscalar

@testset "dual(B::Pseudoscalar, C::Blade)" begin
    # --- Preparations

    test_dim = 10

    test_value = rand()
    test_value = rand() > 0.5 ? test_value : -test_value

    # --- Tests

    # dim(B) == dim(C)
    for test_grade in 2:5
        B = Pseudoscalar(test_dim, test_value)
        C = Blade(rand(test_dim, test_grade))
        @test_throws ArgumentError("`B` not contained in `C`") dual(B, C)
    end

    # dim(B) != dim(C)
    B = Pseudoscalar(test_dim + 1, test_value)
    C = Blade(rand(test_dim, 3))
    @test_throws DimensionMismatch dual(B, C)
end

@testset "dual(B::Pseudoscalar, C::Pseudoscalar)" begin
    # --- Preparations

    test_value_2 = get_random_value(1)  # add 2 to keep value away from 0

    # --- Tests

    # dim(B) == dim(C), test_value_1 != 1
    test_value_1 = get_random_value(2)  # add 2 to keep value away from 0 and 1

    for test_dim in 10:13
        B = Pseudoscalar(test_dim, test_value_1)
        C = Pseudoscalar(test_dim, test_value_2)
        B_dual_C = dual(B, C)
        @test B_dual_C isa Scalar
        @test B_dual_C == test_value_1

        # Check dual(dual(B, C), C) = (-1)^(grade(C) * (grade(C) - 1) / 2) B
        if mod(grade(C), 4) < 2
            @test dual(B_dual_C, C) == B
        else
            @test dual(B_dual_C, C) == -B
        end
    end

    # dim(B) == dim(C), test_value_1 == 1
    test_value_1 = 1

    for test_dim in 10:13
        B = Pseudoscalar(test_dim, test_value_1)
        C = Pseudoscalar(test_dim, test_value_2)
        B_dual_C = dual(B, C)
        @test B_dual_C isa One

        # Check dual(dual(B, C), C) = (-1)^(grade(C) * (grade(C) - 1) / 2) B
        if mod(grade(C), 4) < 2
            @test dual(B_dual_C, C) == B
        else
            @test dual(B_dual_C, C) == -B
        end
    end

    # dim(B) != dim(C)
    test_dim = 10
    B = Pseudoscalar(test_dim, test_value_1)
    C = Pseudoscalar(test_dim + 1, test_value_2)
    @test_throws DimensionMismatch dual(B, C)
end

@testset "dual(B::Pseudoscalar, C::Scalar)" begin
    # --- Preparations

    test_value_1 = rand() + 2  # add 2 to keep value away from 0 and 1
    test_value_1 = rand() > 0.5 ? test_value_1 : -test_value_1

    test_value_2 = rand() + 2  # add 2 to keep value away from 0 and 1
    test_value_2 = rand() > 0.5 ? test_value_2 : -test_value_2

    # --- Tests

    for test_dim in 10:13
        B = Pseudoscalar(test_dim, test_value_1)
        C = Scalar(test_value_2)
        @test_throws ArgumentError("`B` not contained in `C`") dual(B, C)
    end
end

@testset "dual(B::Pseudoscalar, C::One)" begin
    test_value = rand() + 2  # add 2 to keep value away from 0 and 1
    test_value = rand() > 0.5 ? test_value : -test_value

    for test_dim in 10:13
        B = Pseudoscalar(test_dim, test_value)
        C = One()
        @test_throws ArgumentError("`B` not contained in `C`") dual(B, C)
    end
end

@testset "dual(B::Pseudoscalar, C::Real)" begin
    # --- Preparations

    test_value_1 = get_random_value(2)  # add 2 to keep value away from 0 and 1
    test_value_2 = get_random_value(2)  # add 2 to keep value away from 0 and 1

    # --- Tests

    for test_dim in 10:13
        B = Pseudoscalar(test_dim, test_value_1)
        C = test_value_2
        @test_throws ArgumentError("`B` not contained in `C`") dual(B, C)
    end
end

@testset "dual(B::Pseudoscalar, C::Vector)" begin
    # --- Preparations

    test_dim = 10
    test_value = get_random_value()

    # --- Tests

    # dim(B) == dim(C) > 1
    B = Pseudoscalar(test_dim, test_value)
    C = rand(test_dim, 1)[:, 1]
    @test_throws ArgumentError("`B` not contained in `C`") dual(B, C)

    # dim(B) == dim(C) == 1
    B = Pseudoscalar(1, test_value)
    C = [rand()]

    B_dual_C = dual(B, C)
    @test B_dual_C isa Scalar
    @test dual(B, C) == Scalar(test_value)

    # dim(B) != dim(C)
    B = Pseudoscalar(test_dim + 1, test_value)
    C = rand(test_dim, 1)[:, 1]
    @test_throws DimensionMismatch dual(B, C)
end

# --- B::Scalar

@testset "dual(B::Scalar, C::Blade)" begin
    # --- Preparations

    test_dim = 15

    test_value = rand()
    test_value = rand() > 0.5 ? test_value : -test_value

    # --- Tests

    for test_grade in 5:8
        B = Scalar(test_value)
        C = Blade(randn(test_dim, test_grade))

        B_dual_C = dual(B, C)
        expected_result = mod(test_grade, 4) < 2 ?
            Blade(C, volume=test_value) :
            Blade(C, volume=-test_value)

        @test B_dual_C isa Blade
        @test B_dual_C == expected_result
    end
end

@testset "dual(B::Scalar, C::Pseudoscalar)" begin
    # --- Preparations

    test_value_1 = rand() + 2  # add 2 to keep value away from 0 and 1
    test_value_1 = rand() > 0.5 ? test_value_1 : -test_value_1

    test_value_2 = rand() + 2  # add 2 to keep value away from 0 and 1
    test_value_2 = rand() > 0.5 ? test_value_2 : -test_value_2

    # --- Tests

    for test_dim in 5:8
        B = Scalar(test_value_1)
        C = Pseudoscalar(test_dim, test_value_2)

        expected_result = mod(test_dim, 4) < 2 ?
            Pseudoscalar(test_dim, test_value_1) :
            Pseudoscalar(test_dim, -test_value_1)

        @test dual(B, C) == expected_result
    end
end

@testset "dual(B::Scalar, C::Scalar)" begin
    test_value_1 = get_random_value(2)  # add 2 to keep value away from 0 and 1
    test_value_2 = get_random_value(2)  # add 2 to keep value away from 0 and 1

    B = Scalar(test_value_1)
    C = Scalar(test_value_2)

    B_dual_C = dual(B, C)
    @test B_dual_C isa Scalar
    @test B_dual_C == test_value_1
end

@testset "dual(B::Scalar, C::One)" begin
    test_value = get_random_value(2)  # add 2 to keep value away from 0 and 1
    B = Scalar(test_value)

    C = One()

    B_dual_C = dual(B, C)
    @test B_dual_C === B
end

@testset "dual(B::Scalar, C::Real)" begin
    test_value_1 = get_random_value(2)  # add 2 to keep value away from 0 and 1
    test_value_2 = get_random_value(2)  # add 2 to keep value away from 0 and 1

    B = Scalar(test_value_1)
    C = test_value_2

    B_dual_C = dual(B, C)
    @test B_dual_C isa Scalar
    @test B_dual_C == test_value_1
end

@testset "dual(B::Scalar, C::Vector)" begin
    # --- Preparations

    test_dim = 15
    test_value = get_random_value()

    # --- Tests

    B = Scalar(test_value)
    C = randn(test_dim, 1)[:, 1]

    B_dual_C = dual(B, C)
    expected_result = Blade(C, volume=test_value)

    @test B_dual_C isa Blade
    @test B_dual_C == expected_result
end

# --- B::One

@testset "dual(B::One, C::Blade)" begin
    # Preparations
    test_dim = 12

    for test_grade in 5:8
        B = One()
        C = Blade(randn(test_dim, test_grade))

        B_dual_C = dual(B, C)
        expected_result = mod(test_grade, 4) < 2 ?
            Blade(C, volume=1, copy_basis=false) :
            Blade(C, volume=-1, copy_basis=false)

        @test B_dual_C isa Blade
        @test B_dual_C == expected_result
    end
end

@testset "dual(B::One, C::Pseudoscalar)" begin
    # --- Preparations

    test_value = rand() + 2  # add 2 to keep value away from 0 and 1
    test_value = rand() > 0.5 ? test_value : -test_value

    # --- Tests

    for test_dim in 15:18
        B = One()
        C = Pseudoscalar(test_dim, test_value)

        expected_result = mod(test_dim, 4) < 2 ?
            Pseudoscalar(test_dim, 1) :
            Pseudoscalar(test_dim, -1)

        @test dual(B, C) == expected_result
    end
end

@testset "dual(B::One, C::Scalar)" begin
    B = One()

    test_value = get_random_value(2)  # add 2 to keep value away from 0 and 1
    C = Scalar(test_value)

    B_dual_C = dual(B, C)
    @test B_dual_C === B
end

@testset "dual(B::One, C::One)" begin
    B = One()
    C = One()
    B_dual_C = dual(B, C)
    @test B_dual_C === B
end

@testset "dual(B::One, C::Vector)" begin
    # Preparations
    test_dim = 12

    B = One()
    C = randn(test_dim, 1)[:, 1]

    B_dual_C = dual(B, C)
    expected_result = Blade(C, volume=1)

    @test B_dual_C isa Blade
    @test B_dual_C == expected_result
end

@testset "dual(B::One, C::Real)" begin
    B = One()
    C = get_random_value(2)  # add 2 to keep value away from 0 and 1

    B_dual_C = dual(B, C)
    @test B_dual_C === B
end

# --- B::Zero or C::Zero

@testset "dual(B::Zero, C)" begin
    # Note: the case C::Zero is excluded from this set of tests

    # --- Preparations

    test_value = get_random_value(2)  # add 2 to keep value away from 0 and 1

    B = Zero()
    expected_error = "The dual of Zero is not well-defined"

    # --- Tests

    # C::Real
    C = test_value
    @test_throws ErrorException(expected_error) dual(B, C)

    # C::Vector
    C = rand(4, 1)[:, 1]
    @test_throws ErrorException(expected_error) dual(B, C)

    # C::One
    C = One()
    @test_throws ErrorException(expected_error) dual(B, C)

    # C::Scalar
    C = Scalar(test_value)
    @test_throws ErrorException(expected_error) dual(B, C)

    # C::Blade
    C = Blade(randn(4, 3))
    @test_throws ErrorException(expected_error) dual(B, C)

    # C::Pseudoscalar
    C = Pseudoscalar(5, test_value)
    @test_throws ErrorException(expected_error) dual(B, C)

    # C::Multivector
    # C = Multivector(5, test_value)
    @test_skip 1
end

@testset "dual(B, C::Zero)" begin
    # --- Preparations

    # Test values
    test_value = get_random_value(2)  # add 2 to keep value away from 0 and 1

    C = Zero()
    expected_error = "The dual of anything relative to Zero is not well-defined"

    # --- Tests

    # B::Real
    B = test_value
    @test_throws ErrorException(expected_error) dual(B, C)

    # B::Vector is not zero
    B = rand(4, 1)[:, 1]
    @test_throws ErrorException(expected_error) dual(B, C)

    # B::Vector is zero
    B = [0; 0; 0]
    @test_throws ErrorException(expected_error) dual(B, C)

    # B::Zero
    B = Zero()
    @test_throws ErrorException(expected_error) dual(B, C)

    # B::One
    B = One()
    @test_throws ErrorException(expected_error) dual(B, C)

    # B::Scalar
    B = Scalar(test_value)
    @test_throws ErrorException(expected_error) dual(B, C)

    # B::Blade
    B = Blade(randn(4, 3))
    @test_throws ErrorException(expected_error) dual(B, C)

    # B::Pseudoscalar
    B = Pseudoscalar(5, test_value)
    @test_throws ErrorException(expected_error) dual(B, C)

    # C::Multivector
    # C = Multivector(5, test_value)
    @test_skip 1
end

# --- B::Real

@testset "dual(B::Real, C::Blade)" begin
    # --- Preparations

    test_dim = 15
    test_value = get_random_value()

    # --- Tests

    for test_grade in 5:8
        B = test_value
        C = Blade(randn(test_dim, test_grade))

        B_dual_C = dual(B, C)
        expected_result = mod(test_grade, 4) < 2 ?
            Blade(C, volume=test_value) :
            Blade(C, volume=-test_value)

        @test B_dual_C isa Blade
        @test B_dual_C == expected_result
    end
end

@testset "dual(B::Real, C::Pseudoscalar)" begin
    # --- Preparations

    test_value_1 = get_random_value(2)  # add 2 to keep value away from 0 and 1
    test_value_2 = get_random_value(2)  # add 2 to keep value away from 0 and 1

    # --- Tests

    for test_dim in 5:8
        B = test_value_1
        C = Pseudoscalar(test_dim, test_value_2)

        expected_result = mod(test_dim, 4) < 2 ?
            Pseudoscalar(test_dim, test_value_1) :
            Pseudoscalar(test_dim, -test_value_1)

        @test dual(B, C) == expected_result
    end
end

@testset "dual(B::Real, C::Scalar)" begin
    test_value_1 = get_random_value(2)  # add 2 to keep value away from 0 and 1
    test_value_2 = get_random_value(2)  # add 2 to keep value away from 0 and 1

    B = test_value_1
    C = Scalar(test_value_2)

    B_dual_C = dual(B, C)
    @test B_dual_C isa Scalar
    @test B_dual_C == test_value_1
end

@testset "dual(B::Real, C::One)" begin
    test_value = get_random_value(2)  # add 2 to keep value away from 0 and 1

    B = test_value
    C = One()

    B_dual_C = dual(B, C)
    @test B_dual_C isa Scalar
    @test B_dual_C == test_value
end

# --- B::Vector

@testset "dual(B::Vector, C::Blade)" begin
    # --- Preparations

    test_dim = 20

    C = Blade(rand(test_dim, 7))

    # --- Tests

    # ------ dim(B) == dim(C), grade(C) > 1, B is contained in C

    grade_B = 1
    B = basis(C)[:, grade_B] * rand()

    B_dual_C = dual(B, C)

    # --- Check dim, grade, and norm

    @test dim(B_dual_C) == test_dim
    @test grade(B_dual_C) == grade(C) - grade_B
    @test norm(B_dual_C) == norm(B)

    # --- Check that B and dual(B) are orthogonal complements

    @test LinearAlgebra.norm(transpose(basis(B_dual_C)) * reshape(B, length(B), 1)) <
        10 * eps(Float64)

    # --- Check sign(dual(B, C))

    # Compute sign of I_C formed from basis(B) and basis(dual(B, C))
    expected_sign = sign(LinearAlgebra.det(
        transpose(hcat(B, basis(B_dual_C))) * basis(C)))

    # Account for sign of I_C^{-1} relative to I_C
    if mod(grade(C), 4) >= 2
        expected_sign = -expected_sign
    end

    @test sign(B_dual_C) == expected_sign

    # Check dual(dual(B, C), C) = (-1)^(grade(C) * (grade(C) - 1) / 2) B
    if mod(grade(C), 4) < 2
        @test dual(B_dual_C, C) ≈ B
    else
        @test dual(B_dual_C, C) ≈ -B
    end

    # ------ dim(B) == dim(C), grade(C) == 1, B is contained in C

    C_grade_1 = Blade(rand(test_dim, grade_B))
    coefficients = rand()
    B = basis(C_grade_1)[:, grade_B] * coefficients

    relative_orientation =
        sign(LinearAlgebra.det(reshape(B, 1, length(B)) * basis(C_grade_1)))
    expected_result = Scalar(relative_orientation * LinearAlgebra.norm(B))

    @test dual(B, C_grade_1) ≈ expected_result

    # ------ dim(B) != dim(C)

    B = rand(test_dim + 1, grade_B)[:, grade_B]
    @test_throws DimensionMismatch dual(B, C)

    # ------ B not contained in C

    # grade(B) <= grade(C), B not contained in C
    B = rand(test_dim, grade_B)[:, grade_B]
    while LinearAlgebra.norm(reject(reshape(B, length(B), 1), C)) < sqrt(eps(Float64))
        B = rand(test_dim, grade_B)[:, grade_B]
    end
    @test_throws ArgumentError("`B` not contained in `C`") dual(B, C)

    # ------ B is a vector of 0s

    B = fill(0, test_dim)
    @test_throws ErrorException("The dual of Zero is not well-defined") dual(B, C)
end

@testset "dual(B::Vector, C::Pseudoscalar)" begin
    # --- Preparations

    test_dim = 10
    test_value = get_random_value()

    # --- Tests

    # dim(B) == dim(C)
    B = rand(test_dim, 1)[:, 1]
    C = Pseudoscalar(test_dim, test_value)

    B_dual_C = dual(B, C)
    @test B_dual_C isa Blade
    @test B_dual_C == dual(Blade(B))

    # Check dual(dual(B, C), C) = (-1)^(grade(C) * (grade(C) - 1) / 2) B
    if mod(grade(C), 4) < 2
        @test dual(B_dual_C, C) ≈ B
    else
        @test dual(B_dual_C, C) ≈ -B
    end

    # dim(B) != dim(C)
    B = rand(test_dim, 1)[:, 1]
    C = Pseudoscalar(test_dim + 1, test_value)
    @test_throws DimensionMismatch dual(B, C)
end

@testset "dual(B::Vector, C::Scalar)" begin
    C = Scalar(get_random_value())

    # B is not zero
    B = randn(9, 1)[:, 1]
    @test_throws ArgumentError("`B` not contained in `C`") dual(B, C)

    # B is zero
    B = [0; 0; 0]
    @test_throws ErrorException("The dual of Zero is not well-defined") dual(B, C)
end

@testset "dual(B::Vector, C::One)" begin
    C = One()

    # B is not zero
    B = randn(10, 1)[:, 1]
    @test_throws ArgumentError("`B` not contained in `C`") dual(B, C)

    # B is zero
    B = [0; 0; 0]
    @test_throws ErrorException("The dual of Zero is not well-defined") dual(B, C)
end


using Logging, Random
using JuMP, Gurobi
using DecisionProgramming

Random.seed!(42)

@info("Creating the influence diagram.")
diagram = InfluenceDiagram()

add_node!(diagram, DecisionNode("DP", [], ["0-3 patents", "3-6 patents", "6-9 patents"]))
add_node!(diagram, ChanceNode("CT", ["DP"], ["low", "medium", "high"]))
add_node!(diagram, DecisionNode("DA", ["DP", "CT"], ["0-5 applications", "5-10 applications", "10-15 applications"]))
add_node!(diagram, ChanceNode("CM", ["CT", "DA"], ["low", "medium", "high"]))

generate_arcs!(diagram)

X_CT = ProbabilityMatrix(diagram, "CT")
X_CT[1, :] = [1/2, 1/3, 1/6]
X_CT[2, :] = [1/3, 1/3, 1/3]
X_CT[3, :] = [1/6, 1/3, 1/2]
add_probabilities!(diagram, "CT", X_CT)

X_CM = ProbabilityMatrix(diagram, "CM")
X_CM[1, 1, :] = [2/3, 1/4, 1/12]
X_CM[1, 2, :] = [1/2, 1/3, 1/6]
X_CM[1, 3, :] = [1/3, 1/3, 1/3]
X_CM[2, 1, :] = [1/2, 1/3, 1/6]
X_CM[2, 2, :] = [1/3, 1/3, 1/3]
X_CM[2, 3, :] = [1/6, 1/3, 1/2]
X_CM[3, 1, :] = [1/3, 1/3, 1/3]
X_CM[3, 2, :] = [1/6, 1/3, 1/2]
X_CM[3, 3, :] = [1/12, 1/4, 2/3]
add_probabilities!(diagram, "CM", X_CM)

generate_diagram!(diagram, default_utility=false)

@info("Creating the decision model.")
model = Model()
z = DecisionVariables(model, diagram)

@info("Creating problem specific constraints and expressions")

function variables(model::Model, dims::AbstractVector{Int}; binary::Bool=false)
    v = Array{VariableRef}(undef, dims...)
    for i in eachindex(v)
        v[i] = @variable(model, binary=binary)
    end
    return v
end

# Number of states in each node
n_DP = num_states(diagram, "DP")
n_CT = num_states(diagram, "CT")
n_DA = num_states(diagram, "DA")
n_CM = num_states(diagram, "CM")

n_T = 5                     # number of technology projects
n_A = 5                     # number of application projects
I_t = rand(n_T)*0.1         # costs of technology projects
O_t = rand(1:3,n_T)         # number of patents for each tech project
I_a = rand(n_T)*2           # costs of application projects
O_a = rand(2:4,n_T)         # number of applications for each appl. project

V_A = rand(n_CM, n_A).+0.5 # Value of an application
V_A[1, :] .+= -0.5          # Low market share: less value
V_A[3, :] .+= 0.5           # High market share: more value

x_T = variables(model, [n_DP, n_T]; binary=true)
x_A = variables(model, [n_DP, n_CT, n_DA, n_A]; binary=true)

M = 20                      # a large constant
ε = 0.5*minimum([O_t O_a])  # a helper variable, allows using ≤ instead of < in constraints (28b) and (29b)

q_P = [0, 3, 6, 9]          # limits of the technology intervals
q_A = [0, 5, 10, 15]        # limits of the application intervals
z_dP = z.z[1]
z_dA = z.z[2]

@constraint(model, [i=1:n_DP],
    sum(x_T[i,t] for t in 1:n_T) <= z_dP[i]*n_T)            #(25)
@constraint(model, [i=1:n_DP, j=1:n_CT, k=1:n_DA],
    sum(x_A[i,j,k,a] for a in 1:n_A) <= z_dP[i]*n_A)        #(26)
@constraint(model, [i=1:n_DP, j=1:n_CT, k=1:n_DA],
    sum(x_A[i,j,k,a] for a in 1:n_A) <= z_dA[i,j,k]*n_A)    #(27)

#(28a)
@constraint(model, [i=1:n_DP],
    q_P[i] - (1 - z_dP[i])*M <= sum(x_T[i,t]*O_t[t] for t in 1:n_T))
#(28b)
@constraint(model, [i=1:n_DP],
    sum(x_T[i,t]*O_t[t] for t in 1:n_T) <= q_P[i+1] + (1 - z_dP[i])*M - ε)
#(29a)
@constraint(model, [i=1:n_DP, j=1:n_CT, k=1:n_DA],
    q_A[k] - (1 - z_dA[i,j,k])*M <= sum(x_A[i,j,k,a]*O_a[a] for a in 1:n_A))
#(29b)
@constraint(model, [i=1:n_DP, j=1:n_CT, k=1:n_DA],
    sum(x_A[i,j,k,a]*O_a[a] for a in 1:n_A) <= q_A[k+1] + (1 - z_dA[i,j,k])*M - ε)

@constraint(model, [i=1:n_DP, j=1:n_CT, k=1:n_DA], x_A[i,j,k,1] <= x_T[i,1])
@constraint(model, [i=1:n_DP, j=1:n_CT, k=1:n_DA], x_A[i,j,k,2] <= x_T[i,1])
@constraint(model, [i=1:n_DP, j=1:n_CT, k=1:n_DA], x_A[i,j,k,2] <= x_T[i,2])

@info("Creating model objective.")
patent_investment_cost = @expression(model, [i=1:n_DP], sum(x_T[i, t] * I_t[t] for t in 1:n_T))
application_investment_cost = @expression(model, [i=1:n_DP, j=1:n_CT, k=1:n_DA], sum(x_A[i, j, k, a] * I_a[a] for a in 1:n_A))
application_value = @expression(model, [i=1:n_DP, j=1:n_CT, k=1:n_DA, l=1:n_CM], sum(x_A[i, j, k, a] * V_A[l, a] for a in 1:n_A))
@objective(model, Max, sum( sum( diagram.P(convert.(State, (i,j,k,l))) * (application_value[i,j,k,l] - application_investment_cost[i,j,k]) for j in 1:n_CT, k in 1:n_DA, l in 1:n_CM ) - patent_investment_cost[i] for i in 1:n_DP ))

@info("Starting the optimization process.")
optimizer = optimizer_with_attributes(
    () -> Gurobi.Optimizer(Gurobi.Env()),
    "IntFeasTol"      => 1e-9,
    "LazyConstraints" => 1,
)
set_optimizer(model, optimizer)
optimize!(model)

@info("Extracting results.")
Z = DecisionStrategy(z)
S_probabilities = StateProbabilities(diagram, Z)

@info("Printing decision strategy:")
print_decision_strategy(diagram, Z, S_probabilities)

@info("Extracting path utilities")
struct PathUtility <: AbstractPathUtility
    data::Array{AffExpr}
end
Base.getindex(U::PathUtility, i::State) = getindex(U.data, i)
Base.getindex(U::PathUtility, I::Vararg{State,N}) where N = getindex(U.data, I...)
(U::PathUtility)(s::Path) = value.(U[s...])

path_utility = [@expression(model,
    sum(x_A[s[index_of(diagram, "DP")], s[index_of(diagram, "CT")], s[index_of(diagram, "DA")], a] * (V_A[s[index_of(diagram, "CM")], a] - I_a[a]) for a in 1:n_A) -
    sum(x_T[s[index_of(diagram, "DP")], t] * I_t[t] for t in 1:n_T)) for s in paths(diagram.S)]
diagram.U = PathUtility(path_utility)

@info("Computing utility distribution.")
U_distribution = UtilityDistribution(diagram, Z)

@info("Printing utility distribution.")
print_utility_distribution(U_distribution)

@info("Printing statistics")
print_statistics(U_distribution)


module SIWPD
export
    siwpd,
    makesiwpdtree

using
    Wavelets

using
    ..Utils,
    ..DWT

"""
    siwpd(x, wt[, L=maxtransformlevels(x), d=L])

Computes the Shift-Invariant Wavelet Packet Decomposition originally developed
by Cohen, Raz & Malah on the vector `x` using the discrete wavelet filter `wt`
for `L` levels with depth `d`.
"""
function siwpd(x::AbstractVector{T}, wt::OrthoFilter,
               L::S = maxtransformlevels(x), d::S = L) where {T<:AbstractFloat, S<:Integer}
    # Sanity check
    n = length(x)
    @assert 0 ≤ L ≤ maxtransformlevels(n)
    @assert 1 ≤ d ≤ L

    # Setup
    g, h = WT.makereverseqmfpair(wt, true)
    # y = Matrix{T}(undef, (n, gettreelength(1<<(L+1))))
    y = zeros(T, (n, gettreelength(1<<(L+1))))
    @inbounds y[:,1] = x

    # Decomposition
    for i in axes(y,2)
        lvl = getdepth(i, :binary)
        len = nodelength(n, lvl)
        dₗ = (0 ≤ lvl ≤ L-d) ? d : L-lvl
        siwpd_subtree!(y, h, g, i, dₗ, 0, L=lvl, len=len)
    end
    return y
end

#=======================================================================================
Recursive calls for decomposition while d>0.
If shift is 0, ie. ss=0, the main root will not be further decomposed even if d>0.
=======================================================================================#
function siwpd_subtree!(y::AbstractMatrix{T₁}, h::Vector{T₂}, g::Vector{T₂},
                        i::S, d::S, ss::S;
                        L::S = getdepth(i, :binary),
                        len::S = nodelength(size(y,1), L)) where
                       {T₁<:AbstractFloat, T₂<:AbstractFloat, S<:Integer}
    # Sanity check
    n, k = size(y)
    @assert 0 ≤ L ≤ maxtransformlevels(n)
    @assert 0 ≤ ss < n
    @assert 0 ≤ d ≤ getdepth(k,:binary)-L
    @assert iseven(len) || len == 1

    # --- Base case ---
    d > 0 || return y

    # --- One level of decomposition ---
    v_start = ss*len+1                                      # Start index for parent node
    v_end = (ss+1)*len                                      # End index for parent node
    @inbounds v = @view y[v_start:v_end, i]                 # Parent node
    w_start = ss*(len÷2)+1                                  # Start index for child node
    w_end = w_start+(len÷2)-1                               # End index for child node
    @inbounds w₁ = @view y[w_start:w_end, getchildindex(i, :left)]    # Left child node
    @inbounds w₂ = @view y[w_start:w_end, getchildindex(i, :right)]   # Right child node
    sidwt_step!(w₁, w₂, v, h, g, false)                          # 1 step of non-shifted decomposition
    w_start = (ss+1<<L)*(len÷2)+1                           # Start index for shifted child node
    w_end = w_start+(len÷2)-1                               # End index for shifted child node
    @inbounds w₁ = @view y[w_start:w_end, getchildindex(i, :left)]    # Left child node
    @inbounds w₂ = @view y[w_start:w_end, getchildindex(i, :right)]   # Right child node
    sidwt_step!(w₁, w₂, v, h, g, true)                           # 1 step of shifted decomposition

    # --- Recursive decomposition while d>0 ---
    if ss > 0       # Decomposition of current shift
        siwpd_subtree!(y, h, g, getchildindex(i, :left), d-1, ss, L=L+1, len=len÷2)
        siwpd_subtree!(y, h, g, getchildindex(i, :right), d-1, ss, L=L+1, len=len÷2)
    end
    # Decomposition of shifted version
    siwpd_subtree!(y, h, g, getchildindex(i, :left), d-1, ss+1<<L, L=L+1, len=len÷2)
    siwpd_subtree!(y, h, g, getchildindex(i, :right), d-1, ss+1<<L, L=L+1, len=len÷2)

    return y
end

# Single step of decomposition
function sidwt_step!(w₁::AbstractVector{T}, w₂::AbstractVector{T},
                v::AbstractVector{T},
                h::Vector{S}, g::Vector{S},
                s::Bool) where {T<:AbstractFloat, S<:AbstractFloat}
    # Sanity check
    @assert length(w₁) == length(w₂) == length(v)÷2
    @assert length(h) == length(g)

    # Setup
    n = length(v)           # Parent length
    n₁ = length(w₁)         # Child length
    filtlen = length(h)     # Filter length

    # One step of discrete transform
    for i in 1:n₁
        k₁ = mod1(2*i-1-s, n)   # Start index for low pass filtering
        k₂ = 2*i-s              # Start index for high pass filtering
        @inbounds w₁[i] = g[end] * v[k₁]
        @inbounds w₂[i] = h[1] * v[k₂]
        for j in 2:filtlen
            k₁ = k₁+1 |> k₁ -> k₁>n ? mod1(k₁,n) : k₁
            k₂ = k₂-1 |> k₂ -> k₂≤0 ? mod1(k₂,n) : k₂
            @inbounds w₁[i] += g[end-j+1] * v[k₁]
            @inbounds w₂[i] += h[j] * v[k₂]
        end
    end
    return w₁, w₂
end

"""
    makesiwpdtree(n, L, d)

Returns the multi-level, multi-depth binary tree corresponding to the Shift-
Invariant Wavelet Packet Decomposition.
"""
function makesiwpdtree(n::Integer, L::Integer, d::Integer)
    @assert 0 ≤ L ≤ maxtransformlevels(n)
    @assert 1 ≤ d ≤ L

    tree = Vector{BitVector}(undef, 2^(L+1)-1)
    for i in eachindex(tree)
        level = floor(Int, log2(i))
        len = nodelength(n, level)
        nshift = n ÷ len        # number of possible shifts for subspace Ω(i,j)
        node = falses(nshift)
        k = ceil(Int, nshift / 1<<d)  # "skips" between shifts
        node[1:k:end] .= true         # mark nodes present in tree
        tree[i] = node
    end
    return tree
end

end # end module


using DiffEqBase, DiffEqOperators, ForwardDiff, LinearAlgebra, Test
const A = rand(300,300)
f(du,u) = mul!(du,A,u)
f(u) = A*u
x = rand(300)
v = rand(300)
du = similar(x)

cache1 = ForwardDiff.Dual{DiffEqOperators.JacVecTag}.(x, v)
cache2 = ForwardDiff.Dual{DiffEqOperators.JacVecTag}.(x, v)
@test DiffEqOperators.num_jacvec!(du, f, x, v) ≈ ForwardDiff.jacobian(f,similar(x),x)*v rtol=1e-6
@test DiffEqOperators.num_jacvec!(du, f, x, v, similar(v), similar(v)) ≈ ForwardDiff.jacobian(f,similar(x),x)*v rtol=1e-6
@test DiffEqOperators.num_jacvec(f, x, v) ≈ ForwardDiff.jacobian(f,similar(x),x)*v rtol=1e-6

@test DiffEqOperators.auto_jacvec!(du, f, x, v) ≈ ForwardDiff.jacobian(f,similar(x),x)*v
@test DiffEqOperators.auto_jacvec!(du, f, x, v, cache1, cache2) ≈ ForwardDiff.jacobian(f,similar(x),x)*v
@test DiffEqOperators.auto_jacvec(f, x, v) ≈ ForwardDiff.jacobian(f,similar(x),x)*v

f(du,u,p,t) = mul!(du,A,u)
f(u,p,t) = A*u
L = JacVecOperator(f,x)
@test L*x ≈ DiffEqOperators.auto_jacvec(f, x, x)
@test L*v ≈ DiffEqOperators.auto_jacvec(f, x, v)
@test mul!(du,L,v) ≈ DiffEqOperators.auto_jacvec(f, x, v)
DiffEqBase.update_coefficients!(L,v,nothing,nothing)
@test mul!(du,L,v) ≈ DiffEqOperators.auto_jacvec(f, v, v)

L = JacVecOperator(f,x,autodiff=false)
DiffEqBase.update_coefficients!(L,x,nothing,nothing)
@test L*x ≈ DiffEqOperators.num_jacvec(f, x, x)
@test L*v ≈ DiffEqOperators.num_jacvec(f, x, v)
@test mul!(du,L,v) ≈ DiffEqOperators.num_jacvec(f, x, v)
DiffEqBase.update_coefficients!(L,v,nothing,nothing)
@test mul!(du,L,v) ≈ DiffEqOperators.num_jacvec(f, v, v)

L2 = JacVecOperator{Float64}(f)
DiffEqBase.update_coefficients!(L2,x,nothing,nothing)
@test L2*x ≈ DiffEqOperators.auto_jacvec(f, x, x)
@test L2*v ≈ DiffEqOperators.auto_jacvec(f, v, v)
@test mul!(du,L2,x) ≈ DiffEqOperators.auto_jacvec(f, x, x)

L2 = JacVecOperator{Float64}(f,autodiff=false)
DiffEqBase.update_coefficients!(L2,x,nothing,nothing)
@test L2*x ≈ DiffEqOperators.num_jacvec(f, x, x)
@test L2*v ≈ DiffEqOperators.num_jacvec(f, v, v) rtol=1e-6

using OrdinaryDiffEq
function lorenz(du,u,p,t)
 du[1] = 10.0(u[2]-u[1])
 du[2] = u[1]*(28.0-u[3]) - u[2]
 du[3] = u[1]*u[2] - (8/3)*u[3]
end
u0 = [1.0;0.0;0.0]
tspan = (0.0,100.0)
ff = ODEFunction(lorenz,jac_prototype=JacVecOperator{Float64}(lorenz,u0))
prob = ODEProblem(ff,u0,tspan)
sol = solve(prob,Rosenbrock23())
sol = solve(prob,Rosenbrock23(linsolve=LinSolveGMRES(tol=1e-10)))

ff = ODEFunction(lorenz,jac_prototype=JacVecOperator{Float64}(lorenz,u0,autodiff=false))
prob = ODEProblem(ff,u0,tspan)
sol = solve(prob,Rosenbrock23())
sol = solve(prob,Rosenbrock23(linsolve=LinSolveGMRES(tol=1e-10)))


"""
Tests for CR3BP dynamics.
"""
module CR3BPTests

    using ModelingToolkit, AstrodynamicalModels, Test

    @testset "CR3BP Model Constructors" begin
        model = CR3BP(; stm=false, structural_simplify=false)
        model = CR3BP(; stm=true, structural_simplify=true)
        @test true
    end

    @testset "CR3BP Model Calculations" begin
        vectorfield = CR3BPFunction()

        r = [0.9253021269565836, 0, 0]
        v = [0.0, 0.05852663414965813, 0.0]
        μ = 0.0009536838895767625

        @test isapprox(
                vectorfield(vcat(r,v), [μ], NaN),
                [0.0, 0.05852663414965813,  0.0, 0.053265045303684255, 0.0, -0.0];
                atol=1e-8
        )
    end

end


const APL = AbstractPolynomialLike

function _expectation(μ::Measure{T, BT}, p::APL, f) where {T, BT}
    i = 1
    s = 0
    basis = MB.basis_covering_monomials(BT, monomials(p))
    for (c, m) in zip(coefficients(p, basis), basis)
        while i <= length(μ.basis) && m != μ.basis[i]
            i += 1
        end
        if i > length(μ.basis)
            error("The base function $m has a non-zero multiplier $c in $p,
                  but its expectation is not known in $μ")
        end
        s += f(μ.values[i], c)
        i += 1
    end
    s
end

"""
    MultivariateMoments.expectation(μ::AbstractMeasureLike, p::AbstractPolynomialLike)
    MultivariateMoments.expectation(p::AbstractPolynomialLike, μ::AbstractMeasureLike)

Computes the expectation ``\\mathbb{E}_{\\mu}[p]``.
"""
function expectation end

expectation(μ::Measure, p::APL) = _expectation(μ, p, (*))
expectation(p::APL, μ::Measure) = _expectation(μ, p, (a, b) -> b * a) # a and b may be noncommutative

"""
    dot(μ::AbstractMeasureLike, p::AbstractPolynomialLike)
    dot(p::AbstractPolynomialLike, μ::AbstractMeasureLike)

See [`expectation`](@ref)
"""
LinearAlgebra.dot(μ::AbstractMeasureLike, p::APL) = expectation(μ, p)
LinearAlgebra.dot(p::APL, μ::AbstractMeasureLike) = expectation(p, μ)


#=
Popis
  Funkcia vypocita hodnotu eccentricity centrality pre kazdy vrchol.

Parametre
  g : datova struktura grafu :: GenericGraph
  dist_key : nazov atributu, v ktorom je ulozene ohodnotenie hrany :: String

Navratova hodnota
  eccentricity : slovnik s mapovanim (kluc=vrchol_grafu => hodnota=eccentricita) :: Dict{Int64, Float64}

Poznamky
=#

function centrality_eccentricity(g::GenericGraph; dist_key::AbstractString="weight")

  epi = AttributeEdgePropertyInspector{Float64}(dist_key)
  eccentricity = [v.index => 0.0 for v=vertices(g)]

  for v in vertices(g)

    temp_dists = dijkstra_shortest_paths(g, AttributeEdgePropertyInspector{Float64}(dist_key), [v]).dists
    temp_dists = filter(x -> x != Inf, temp_dists)
    ecc = 1.0/(maximum(temp_dists))

    if ecc != Inf
      eccentricity[v.index] = ecc
    elseif ecc == Inf
      eccentricity[v.index] = 0.0
    end

  end

  return eccentricity

end


using Test

using CompScienceMeshes
using BEAST

"""
    isdivconforming(space)

Returns for every basis functions in the space the total of the absolute
values of non-compensated skeleton flux, i.e. flux flowing into from one
side of an edge and not flowing out on the other side. Should be all zeros
for divergence conforming functions, at least on a closed surface.
"""
function isdivconforming(space)

    geo = geometry(space)
    mesh = geo

    edges = skeleton(mesh,1)
    D = connectivity(edges, mesh, abs)
    rows = rowvals(D)
    vals = nonzeros(D)

    Flux = zeros(Float64, numcells(mesh),3)
    TotalFlux = zeros(Float64, numcells(edges), numfunctions(space))

    for i in 1 : numfunctions(space)

        fill!(Flux, 0)
        bfs = BEAST.basisfunction(space,i)
        for bf in bfs
            c = bf.cellid
            e = bf.refid
            x = bf.coeff
            Flux[c,e] += x
        end

        for E in 1 : numcells(edges)
            for j in nzrange(D,E)
                F = rows[j]
                e = vals[j]
                @assert 1 <= e <= 3
                TotalFlux[E,i] += Flux[F,e]
            end
        end
    end

    return TotalFlux
end

function touches_predicate(mesh)

    verts = skeleton(mesh,0)
    verts = sort(verts.faces, lt=isless)

    edges = skeleton(mesh,1)
    edges = sort(edges.faces, lt=isless)

    function pred(s)

        num_hits = 0
        for i in 1:length(s)
            if !isempty(searchsorted(verts, SVector(s[i]), lt=isless))
                num_hits += 1
            end
        end

        @assert num_hits < 3
        num_hits == 0 && return false
        num_hits == 1 && return true
        isempty(searchsorted(edges, s, lt=isless)) && return true
        return false
    end

    return pred
end

function interior(mesh::Mesh)
    D = dimension(mesh)
    edges = skeleton(mesh, D-1)
    pred = interior_tpredicate(mesh)
    submesh(pred, edges)
end

#meshfile = Pkg.dir("BEAST","test","sphere2.in")
meshfile = joinpath(dirname(@__FILE__),"assets","sphere316.in")
mesh = readmesh(meshfile)
@test numvertices(mesh) == 160
@test numcells(mesh) == 316

rt = raviartthomas(mesh)
@test numfunctions(rt) == 316 * 3 / 2

fine = barycentric_refinement(mesh)
edges = skeleton(mesh, 1)

bc = buffachristiansen(mesh)
@test numfunctions(bc) == 316 * 3 / 2

lc = isdivconforming(rt)
@test maximum(lc) < eps(Float64) * 1000
println("RT space is div-conforming")

lc = isdivconforming(bc);
@test maximum(lc) < eps(Float64) * 1000
println("BC space is div-conforming")

# Now repeat the exercise with an open mesh
mesh = meshrectangle(1.0, 1.0, 0.2);
fine = barycentric_refinement(mesh);

rt = raviartthomas(mesh)
bc = buffachristiansen(mesh)

@test numfunctions(rt) == 65
@test numfunctions(bc) == 65

int_pred = interior_tpredicate(mesh)
bnd_pred(s) = !int_pred(s)

leaky_edges = findall(sum(abs.(isdivconforming(rt)),dims=1) .!= 0)
@test length(leaky_edges) == 0

bnd = boundary(mesh)
bndtch_pred = touches_predicate(bnd)
edges = interior(mesh)
bndtch_edges = findall(bndtch_pred, edges.faces)
leaky_edges = findall(vec(sum(abs.(isdivconforming(bc)),dims=1)) .!= 0)
@test bndtch_edges == leaky_edges

## Test the charge of BC functions
#meshfile = Pkg.dir("BEAST","test","sphere2.in")
meshfile = joinpath(dirname(@__FILE__),"assets","sphere316.in")
mesh = readmesh(meshfile)
bc = buffachristiansen(mesh)
fine = geometry(bc)
charges = zeros(numcells(fine))

for fn in bc.fns
    abs_charge = 0.0
    net_charge = 0.0
    fill!(charges,0)
    for  _sh in fn
        cellid = _sh.cellid

        #cell = simplex(vertices(fine, fine.faces[cellid]))
        net_charge += _sh.coeff
        charges[cellid] += _sh.coeff
    end
    abs_charge = sum(abs.(charges))
    @test net_charge + 1 ≈ 1
    @test abs_charge ≈ 2
end

# THe BC construction function should throw for non-oriented surfaces
# width, height, h = 1.0, 0.5, 0.05
G1 = meshrectangle(1.0, 1.0, 0.25)
G2 = CompScienceMeshes.rotate(G1, 0.5π * x̂)
G = CompScienceMeshes.weld(G1,G2)
@test_throws AssertionError buffachristiansen(G)
# G3 = CompScienceMeshes.rotate(G1, 1.0π * x̂)


using LinearAlgebra
using DynamicalSystems
using Parameters
using Plots
using MKL
using BasicReservoirComputing
include("utilities.jl")

function opt_rc(N, train_data, valid_data; Δt = 0.01)
    """Find the RC parameters through optimization"""
    # define some parameters
    nspin = 200
    ρA = 0.02
    D = size(train_data)[1]
    multithread = Threads.nthreads() > 1 ? true : false

    # Optimize the RC
    rc_opt = opt(train_data, valid_data, nspin, N, ρA,
                 lb = opt_lower_bounds(), ub = opt_upper_bounds())
    params, cost = rc_opt(maxtime=200, multithread=multithread)
    return params
end

function test_rc(system::DynamicalSystem, params::NamedTuple; Δt = 0.01)
    """Test the rc on the given system with params"""
    D=length(system.u0)
    utrain, utest_arr = make_data(system; train_time = 100, test_time = 25, n_test = 100)

    @unpack N, SR, ρA, α, σ, β, σb, nspin = params

    # Build and Train RC
    reservoir = rc(D, N, ρA, Δt, SR, α, σ, σb, β)
    train_RC(reservoir, utrain, spinup=nspin)

    # Forecast RC
    utrue = utest_arr[1][:, nspin:end]
    nsteps = size(utrue)[2]
    upred = forecast_RC(reservoir, nsteps, uspin = utest_arr[1][:, 1:nspin])
    t = collect(0:nsteps-1)*Δt

    p1 = plot_prediction(t, upred, utrue)

    # Test on many points
    vpt = test_RC(reservoir, utest_arr; ϵ=0.3, spinup=nspin)
    p2 = histogram(vpt, legend=false,
                   overwrite_figure=false,
                   show=false)
    xlabel!(p2, "Valid Prediction Time")
    pfinal = plot(p1, p2, layout = (1, 2), show=true)
    display(pfinal)
    return reservoir
end

# First make some data
system = Systems.lorenz()
Δt = 0.01
train_time=100
valid_time = 15
n_valid = 15
utrain, uvalid_arr = make_data(system; train_time = train_time,
                               test_time = valid_time, n_test = n_valid,
                               Δt = Δt)

# Now find the best parameters
N=200
params = opt_rc(N, utrain, uvalid_arr)

# Test the RC with the optimized params
reservoir = test_rc(system, merge(params, (N=N, ρA=0.02, nspin=200)))

# See how the lyapunov exponents compare
renorm = 10 # speeds up the calculation by not doing QR renormalization every step
nspin = 200

rc_LEs = Global_LEs(reservoir, utrain, nspin,
                    num_exp = length(system.u0),
                    renorm_steps=renorm)
sys_LEs = lyapunovspectrum(system, 10000)
println("rc_LEs = $rc_LEs")
println("system_LEs = $sys_LEs")


using ControlSystems: ssdata
# TODO:  Implement Reducing unstable linear control systems via real Schur transformation.
# By temporarily removing unstable dynamics, reduce, add dynamics back.

"""
    frequency_weighted_reduction(G, Wo, Wi; residual=true)

Find Gr such that ||Wₒ(G-Gr)Wᵢ||∞ is minimized.
For a realtive reduction, set Wo = inv(G) and Wi = I.

If `residual = true`, matched static gain is achieved through "residualization", i.e., setting
```math
0 = A_{21}x_{1} + A_{22}x_{2} + B_{2}u
```
where indices 1/2 correspond to the remaining/truncated states respectively. This choice typically results in a better match in the low-frequency region and a smaller overall error.

Ref: Andras Varga and Brian D.O. Anderson, "Accuracy enhancing methods for the frequency-weighted balancing related model reduction"
https://elib.dlr.de/11746/1/varga_cdc01p2.pdf
"""
function frequency_weighted_reduction(G, Wo, Wi, r; residual=true)
    iscontinuous(G) || error("Discrete systems not supported yet.")
    A,B,C,D = ssdata(G)
    n = G.nx
    if Wo == 1 || Wo == I
        Q1 = grampd(G, :o)
    else
        Wo = ss(Wo)
        if issiso(Wo) && size(G, 1) > 1
            Wo = Wo*I(size(G, 1))
        end
        Ao,Bo,Co,Do = ssdata(Wo)
        As = [
            A                                  zeros(size(A,1), size(Ao,2))
            Bo*C                               Ao
        ]
        Bs = [
            B
            Bo*D
        ]
        Cs = [
            Do*C   Co
        ]
        Ds = Do*D
        WoG = ss(As, Bs, Cs, Ds)
        Q1 = grampd(WoG, :o)[1:n, 1:n]
    end

    if Wi == 1 || Wi == I
        P1 = grampd(G, :c)
    else
        Wi = ss(Wi)
        if issiso(Wi) && size(G, 2) > 1
            Wi = Wi*I(size(G, 1))
        end
        Ai,Bi,Ci,Di = ssdata(ss(Wi))
        As = [
            A                                B*Ci
            zeros(size(Ai,1), size(A,2))     Ai
        ]
        Bs = [
            B*Di
            Bi
        ]
        Cs = [
            C   D*Ci
        ]
        Ds = D*Di
        GWi = ss(As, Bs, Cs, Ds)
        P1 = grampd(GWi, :c)[1:n, 1:n]
    end

    R = Q1
    S = P1
    U,Σ,V = svd!(R*S)
    i1 = 1:r
    i2 = r+1:size(A, 1)
    U1 = U[:,i1]
    V1 = V[:,i1]
    U2 = U[:,i2]
    V2 = V[:,i2]
    Σ1 = Σ[i1]
    Y = Matrix(qr!(R'U1).Q)
    X = Matrix(qr!(S*V1).Q)

    @views if residual
        # The code for the residual case is adapted from
        # DescriptorSystems.jl and written by Andreas Varga
        # https://github.com/andreasvarga/DescriptorSystems.jl/blob/dd144828c3615bea2d5b4977d7fc7f9677dfc9f8/src/order_reduction.jl#L622
        # with license https://github.com/andreasvarga/DescriptorSystems.jl/blob/main/LICENSE.md
        ONE = one(eltype(A))
        L = [Y Matrix(qr!(R'U2).Q)]
        Tr = [X Matrix(qr!(S*V2).Q)]
        amin = L'A*Tr
        bmin = L'B
        cmin = C*Tr
        Ar = amin[i1,i1]
        Er = L[:,i1]'*Tr[:,i1]
        Br = bmin[i1,:]
        Cr = cmin[:,i1]
        Dr = copy(D)
        LUF = lu!(amin[i2,i2])
        ldiv!(LUF,amin[i2,i1])
        ldiv!(LUF,bmin[i2,:])
        # apply state residualization formulas
        mul!(Dr,cmin[:,i2],bmin[i2,:],-ONE, ONE)
        mul!(Br,amin[i1,i2],bmin[i2,:],-ONE, ONE)
        mul!(Cr,cmin[:,i2],amin[i2,i1],-ONE, ONE)
        mul!(Ar,amin[i1,i2],amin[i2,i1],-ONE, ONE)
        # determine a standard reduced system
        SV = svd!(Er)
        di2 = Diagonal(1 ./sqrt.(SV.S))
        return ss(di2*SV.U'*Ar*SV.Vt'*di2, di2*(SV.U'*Br), (Cr*SV.Vt')*di2, Dr)
    else
        L = (Y'X)\Y'
        T = X
        A = L*A*T
        B = L*B
        C = C*T
        D = D
    end
    ss(A,B,C,D, G.timeevol)
end

# "
# Lemma 19.1 See Robust and Optimal Control Ch 19.1
# "
# function controller_reduction_weight(P::ExtendedStateSpace, K)
#     A, B1, B2, C1, C2, D11, D12, D21, D22 = ssdata_e(P)
#     Ak,Bk,Ck,Dk = ssdata(K)
#     R = factorize(I - D22*Dk)
#     R̃ = factorize(I - Dk*D22)
#     Aw = [
#             A+B2*Dk*(R\C2) B2*(R̃\Ck)
#             Bk*(R\C2)      Ak+Bk*D22*(R̃\Ck)
#     ]
#     Bw = [
#         B2/R̃
#         Bk*D22/R̃
#     ]
#     Cw = [R\C2   R\D22*Ck]
#     Dw = D22/R̃
#     ss(Aw, Bw, Cw, Dw)
# end

# """
#     controller_reduction(P, K, r, out=false)

# Minimize    ||(K-Kᵣ) W||∞ if out=false
#             ||W (K-Kᵣ)||∞ if out=true
# See Robust and Optimal Control Ch 19.1
# out indicates if the weight will be applied as output or input weight.
# """
# function controller_reduction(P, K, r, out=false)
#     W = controller_reduction_weight(P, K)
#     if out
#         frequency_weighted_reduction(K, W, 1, r)
#     else
#         frequency_weighted_reduction(K, 1, W, r)
#     end
# end

"""
    hsvd(sys::AbstractStateSpace{Continuous})

Return the Hankel singular values of `sys`, computed as the eigenvalues of `QP`
Where `Q` and `P` are the Gramians of `sys`.
"""
function hsvd(sys::AbstractStateSpace{Continuous})
    P = MatrixEquations.plyapc(sys.A, sys.B)
    Q = MatrixEquations.plyapc(sys.A', sys.C')
    e = svdvals(Q * P)
end

@deprecate minreal2 minreal

# slide 189 https://cscproxy.mpi-magdeburg.mpg.de/mpcsc/benner/talks/lecture-MOR.pdf
# This implementation works, but results in a complex-valued system.
# function model_reduction_irka(sys::AbstractStateSpace{Continuous}, r; tol = 1e-6)
#     A,B,C,D = ssdata(sys)
#     all(iszero, D) || error("Nonzero D not supported.")
#     sysr0, _ = baltrunc(sys, n=r)
#     Ah, Bh, Ch = real.(ssdata(sysr0))
#     e, T = eigen(Ah)
#     eold = 100e
#     k = 0
#     while maximum(abs.(e .- eold) ./ abs.(e)) > tol && k < 1000
#         Bt = T'Bh
#         Ct = Ch*T
#         Vv = map(1:r) do i
#             (-e[i]*I - A)\(B*Bt[i, :])
#         end
#         V = reduce(hcat, Vv)
#         Wv = map(1:r) do i
#             (-e[i]*I - A')\(C'*Ct[:,i])
#         end
#         W = reduce(hcat, Wv)
#         V = orthonormal(V) # not sure about those
#         W = orthonormal(W)
#         Ah = (W'V)\W'A*V
#         Bh = (W'V)\W'B
#         Ch = C*V
#         eold = e
#         k += 1
#         e, T = eigen(Ah)
#     end
#     sysr = ss(Ah, Bh, Ch, 0)

# end

# function orthonormal(X)
#     # U,S,V = svd(X, full=true)
#     # R = U*V'
#     Matrix(qr(X).Q) # |> real # this works sometimes
# end


export reluActivation, reluActivation!

"""
 relu activation A = relu(Y)

 Input:

   Y - array of features

 Optional Input:

   doDerivative - flag for computing derivative, set via varargin
                  Ex: reluActivation(Y,true);

 Output:

  A  - activation
  dA - derivatives
"""
function reluActivation(Y::Array{T},doDerivative::Bool=false) where {T}

    A = max.(Y,zero(T));

    if doDerivative
        dA = sign.(A);
    else
        dA = zeros(T,0)
    end

    return A,dA
end

function reluActivation!(A::Array{T},dA,doDerivative::Bool=false) where {T}
    A .= max.(A,zero(T));
    if doDerivative
        if isempty(dA)
            dA = sign.(A);
        else
            dA .= sign.(A);
        end
    end
    return A,dA
end


using MonteCarloMeasurements, Measurements, Plots, BenchmarkTools, OrdinaryDiffEq, PrettyTables, ChangePrecision, LinearAlgebra
# pgfplots()
default(size=(600,400))
color_palette = [
    RGB(0.1399999, 0.1399999, 0.4),
    RGB(1.0, 0.7075, 0.35),
    RGB(0.414999, 1.0, 1.0),
    RGB(0.6, 0.21, 0.534999),
    RGB(0,0.6,0),
]

function sim((±)::F, tspan, plotfun=plot!, args...; kwargs...) where F
    @changeprecision Float32 begin
        g = 9.79 ± 0.02; # Gravitational constant
        L = 1.00 ± 0.01; # Length of the pendulum
        u₀ = [0.0 ± 0.0, π / 3.0 ± 0.02] # Initial speed and initial angle
        # @show typeof(u₀)
        gL = g/L
        #Define the problem
        function simplependulum(du,u,p,t)
            θ  = u[1]
            dθ = u[2]
            du[1] = dθ
            du[2] = -gL * sin(θ)
        end

        prob = ODEProblem(simplependulum, u₀, tspan)
        sol = solve(prob, Tsit5(), reltol = 1e-6)

        plotfun(sol.t, getindex.(sol.u, 2), args...; kwargs...)
    end
end

## Special function needed to construct the parameters as sigmapoints, as they all have to be constructed in one single call
function sigmasim(tspan, plotfun=plot!, args...; kwargs...) where F
    @changeprecision Float32 begin
        g,L,u02 = StaticParticles(sigmapoints([9.79, 1.0, pi/3.0], diagm([0.02, 0.01, 0.02].^2)))
        u₀ = [0, u02] # Initial speed and initial angle
        # @show typeof(u₀)
        gL = g/L
        #Define the problem
        function simplependulum(du,u,p,t)
            θ  = u[1]
            dθ = u[2]
            du[1] = dθ
            du[2] = -gL * sin(θ)
        end

        prob = ODEProblem(simplependulum, u₀, tspan)
        sol = solve(prob, Tsit5(), reltol = 1e-6)

        plotfun(sol.t, getindex.(sol.u, 2), args...; kwargs...)
    end
end

##
tspan = (0.0f0, 0.5f0)
plot()
sim(Measurements.:±, tspan, label = "Linear", xlims=(tspan[2]-2,tspan[2]), color=color_palette[4])
sim(MonteCarloMeasurements.:±, tspan, errorbarplot!, 0.8413, label = "MCM", xlims=(tspan[2]-0.5,tspan[2]), l=(:dot,), color=color_palette[2], xlabel="Time [s]", ylabel="\$\\theta\$")

# savefig("/home/fredrikb/mcm_paper/figs/0-2.pdf")
##
tspan = (0.0f0, 200)
plot()
sim(Measurements.:±, tspan, label = "Linear", xlims=(tspan[2]-5,tspan[2]), color=color_palette[4])
sim(MonteCarloMeasurements.:±, tspan, label = "Monte Carlo", xlims=(tspan[2]-5,tspan[2]), l=(:dot,), color=color_palette[2], xlabel="Time [s]", ylabel="\$\\theta\$")

##
# We now integrated over 200 seconds and look at the last 5 seconds. This result maybe looks a bit confusing, the linear uncertainty propagation is very sure about the amplitude at certain points but not at others, whereas the Monte-Carlo approach is completely unsure. Furthermore, the linear approach thinks that the amplitude at some points is actually much higher than the starting amplitude, implying that energy somehow has been added to the system! The picture might become a bit more clear by plotting the individual trajectories of the particles

tspan = (0.0f0, 200)
plot()
sim(Measurements.:±, tspan, label = "Linear", xlims=(tspan[2]-5,tspan[2]), l=(5,), color=color_palette[4])
sim(MonteCarloMeasurements.:∓, tspan, mcplot!, xlims=(tspan[2]-5,tspan[2]), l=(color_palette[2],0.3), xlabel="Time [s]", ylabel="\$\\theta\$", label="MCM")
sigmasim(tspan, mcplot!, xlims=(tspan[2]-5,tspan[2]), l=(color_palette[3],0.9,2), xlabel="Time [s]", ylabel="\$\\theta\$", label="MCM \\Sigma")
# savefig("/home/fredrikb/mcm_paper/figs/mcplot.pdf")

# It now becomes clear that each trajectory has a constant amplitude (although individual trajectories amplitudes vary slightly due to the uncertainty in the initial angle), but the phase is all mixed up due to the slightly different frequencies!

# These problems grow with increasing uncertainty and increasing integration time. In fact, the uncertainty reported by Measurements.jl goes to infinity as the integration time does the same.

# Of course, the added accuracy from using MonteCarloMeasurements does not come for free, as it costs some additional computation. We have the following timings for integrating the above system 100 seconds using three different uncertainty representations

##
function naive_mc(tspan)
    for i = 1:100
        sim(certain, tspan, (args...;kwargs...)->nothing)
    end
end

tspan = (0.0f0, 100f0)
certain = (x,y)->x+y*randn()
table = Matrix{Any}(undef,3,6)
t1 = @benchmark sim($certain, $tspan, (args...;kwargs...)->nothing)
t2 = @benchmark sim($Measurements.:±, $tspan, (args...;kwargs...)->nothing) samples=500
t3 = @benchmark sim($MonteCarloMeasurements.:∓, $tspan, (args...;kwargs...)->nothing) samples=500
t4 = @benchmark sigmasim($tspan, (args...;kwargs...)->nothing) samples=500
t5 = @benchmark naive_mc($tspan) samples=500

# table[1,1] = ""
table[1,1] = "Time [ms]"
table[2,1] = "Memory [MiB]"
table[3,1] = "k Allocations"

for (i,t) in enumerate((t1,t2,t3,t4,t5))
    table[1,i+1] = time(t)/1000_000
    table[2,i+1] = memory(t)/1000_000
    table[3,i+1] = allocs(t)/1000
end

# pretty_table(table, ["" "Float32" "Linear" "MCM" "MCM \\Sigma" "Naive MC"], backend=:latex, formatters=ft_printf("%5.1f"))
pretty_table(table, ["" "Float32" "Linear" "MCM" "MCM \\Sigma" "Naive MC"], backend=:text, tf=markdown, formatters=ft_printf("%5.1f"))


module Denoising
export
    # denoising
    RelErrorShrink,
    SureShrink,
    denoiseall,
    relerrorthreshold

using
    Wavelets,
    LinearAlgebra,
    Statistics,
    Plots

using
    ..DWT,
    ..SWT,
    ..ACWT,
    ..Utils

# ========== Threshold Determination Methods ==========
"""
    RelErrorShrink(th, t) <: DNFT

Relative Error Shrink method used in their paper "Efficient Approximation and Denoising of
Graph Signals using the Multiscale Basis Dictionary" for IEEE Transactions on Signal and
Information Processing over Networks, Vol 0, No. 0, 2016.

# Attributes
- `th::Wavelets.Threshold.THType`: (Default: `Wavelets.HardTH()`) Threshold type.
- `t::AbstractFloat`: (Default: 1.0) Threshold size.

# Examples
```julia
using Wavelets, WaveletsExt

RelErrorShrink()                    # Using default th and t values
RelErrorShrink(SoftTH())            # Using default t value
RelErrorShrink(HardTH(), 0.5)       # Using user input th and t values
```

**See also:** [`SureShrink`](@ref), [`VisuShrink`](@ref)
"""
struct RelErrorShrink <: DNFT
    # Attributes
    th::Wavelets.Threshold.THType
    t::AbstractFloat
    # Constructor
    RelErrorShrink(th = HardTH(), t = 1.0) = new(th, t)
end

"""
    SureShrink(th, t) <: DNFT

Stein's Unbiased Risk Estimate (SURE) Shrink

# Attributes
- `th::Wavelets.Threshold.THType`: (Default: `Wavelets.HardTH()`) Threshold type.
- `t::AbstractFloat`: (Default: 1.0) Threshold size.

**See also:** [`RelErrorShrink`](@ref), [`VisuShrink`](@ref)
"""
struct SureShrink <: DNFT
    # Attributes
    th::Wavelets.Threshold.THType
    t::AbstractFloat
    # Constructor
    SureShrink(th, t) = new(th, t)
end

"""
    SureShrink(xw[, redundant, tree, th])

Struct constructor for `SureShrink` based on the signal coefficients `xw`.

# Arguments
- `xw::AbstractArray{T} where T<:Number`: Decomposed signal.
- `redundant::Bool`: (Default: `false`) Whether the transform type of `xw` is a redundant
  transform. Autocorrelation and stationary wavelet transforms are examples of redundant
  transforms.
- `tree::Union{BitVector, Nothing}`: (Default: `nothing`) The basis tree for decomposing
  `xw`. Must be provided if `xw` is decomposed using `wpt`, `swpd`, or `acwpd`.
- `th::Wavelets.Threshold.THType`: (Default: `HardTH()`) Threshold type.

# Returns
`::SureShrink`: SUREShrink object.

# Examples
```julia
SureShrink(xw)                  # `xw` is output of dwt, wpt
SureShrink(xw, true)            # `xw` is output of sdwt, acdwt, swpt, acwpt
SureShrink(xw, true, tree)      # `xw` is output of swpd, acwpd
```

**See also:** [`SureShrink`](@ref), [`surethreshold`](@ref)
"""
function SureShrink(xw::AbstractArray{T},
                    redundant::Bool = false,
                    tree::Union{BitVector, Nothing} = nothing,
                    th::Threshold.THType = HardTH()) where T<:Number
    t = surethreshold(xw, redundant, tree)
    return SureShrink(th, t)
end

"""
    VisuShrink(n, th)

Extension to the `VisuShrink` struct constructor from `Wavelets.jl`.

# Arguments
- `n::Integer`: Signal length.
- `th::Wavelets.Threshold.THType`: Threshold type.

# Returns
`::Wavelets.Threshold.VisuShrink`: VisuShrink object.

# Examples
```julia
using Wavelets, WaveletsExt

VisuShrink(128, SoftTH())
```
"""
function Wavelets.Threshold.VisuShrink(n::Integer, th::Wavelets.Threshold.THType)
    return VisuShrink(th, sqrt(2*log(n)))
end

"""
    surethreshold(coef, redundant[, tree])

Determination of the `t` value used for `SureShrink`. `t` is defined as the threshold value
when the standard deviation of the noisy signal is 1.

# Arguments
- `coef::AbstractArray{T} where T<:Number`: Coefficients of decomposed signal.
- `redundant::Bool`: Whether the transform type of `xw` is a redundant transform.
  Autocorrelation and stationary wavelet transforms are examples of redundant transforms.
- `tree::Union{BitVector, Nothing}`: (Default: `nothing`) The basis tree for decomposing
  `xw`. Must be provided if `xw` is decomposed using `wpt`, `swpd`, or `acwpd`.

# Returns
`::AbstractFloat`: `t` value used for `SureShrink`.

**See also:** [`SureShrink`](@ref)
"""
function surethreshold(coef::AbstractArray{T},
                       redundant::Bool,
                       tree::Union{BitVector,Nothing} = nothing) where T<:Number
    # extract necessary coefficients
    if !redundant                              # dwt or wpt
        y = coef
    elseif redundant && isa(tree, Nothing)     # sdwt, acdwt, swpt, or acwpt
        y = reshape(coef, :)
    else                                       # swpd or acwpd
        leaves = getleaf(tree,:binary)
        y = reshape(coef[:, leaves], :)
    end
    a = sort(abs.(y)).^2
    b = cumsum(a)
    n = length(y)
    c = collect(reverse(0:(n-1)))
    s = b + c .* a
    risk = (n .- (2*(1:n)) + s)/n
    i = argmin(risk)
    return sqrt(a[i])
end

# ========== Noise Estimation ==========
"""
    noisest(x, redundant[, tree])

Extension to the `noisest` function from `Wavelets.jl`. Estimates the standard deviation of
a signal's noise assuming that the noise is distributed normally. This function is generally
used in combination with `VisuShrink` and `SureShrink` in the `denoise`/`denoiseall`
functions.

# Arguments
- `x::AbstractArray{T}`: Decomposed signal.
- `redundant::Bool`: Whether the transform type of `xw` is a redundant transform.
  Autocorrelation and stationary wavelet transforms are examples of redundant transforms.
- `tree::Union{BitVector, Nothing}`: (Default: `nothing`) The basis tree for decomposing
  `xw`. Must be provided if `xw` is decomposed using `wpt`, `swpd`, or `acwpd`.

# Returns
`::AbstractFloat`: Estimated standard deviation of the noise of the signal.

# Examples
```julia
using Wavelets, WaveletsExt

x = randn(128)
wt = wavelet(WT.haar)

# noise estimate for dwt transformation
y = dwt(x, wt)
noise = noisest(y, false)

# noise estimate for wpt transformation
tree = maketree(x, :full)
y = wpt(x, wt, tree)
noise = noisest(y, false, tree)

# noise estimate for sdwt transformation
y = sdwt(x, wt)
noise = noisest(y, true)

# noise estimate for swpd transformation
y = swpd(x, wt)
noise = noisest(y, true, tree)
```

**See also:** [`relerrorthreshold`](@ref), [`VisuShrink`](@ref), [`SureShrink`](@ref)
"""
function Wavelets.Threshold.noisest(x::AbstractArray{T},
                                    redundant::Bool,
                                    tree::Union{BitVector,Nothing} = nothing) where T<:Number
    # Sanity check
    @assert isdyadic(size(x,1))

    # Get detail coefficients
    if !redundant && isa(tree, Nothing)        # regular dwt
        n₀ = size(x,1) ÷ 2
        dr = x[(n₀+1):end]
    elseif !redundant && isa(tree, BitVector)  # regular wpt
        dr = x[finestdetailrange(x, tree)]
    elseif redundant && isa(tree, Nothing)     # redundant dwt
        dr = x[:,end]
    else                                       # redundant wpd
        dr = x[finestdetailrange(x, tree, true)...]
    end
    return Wavelets.Threshold.mad!(dr)/0.6745
end

"""
    relerrorthreshold(coef, [redundant, tree, elbows; makeplot])

Takes in a set of expansion coefficients, 'plot' the threshold vs relative error curve and
select the best threshold value based on the elbow method. If one wants to see the resulting
plot from this computation, simply set `makeplot=true`.

# Arguments
- `coef::AbstractArray{T} where T<:Number`: Decomposed signal.
- `redundant::Bool`: (Default: `false`) Whether the transform type of `xw` is a redundant
  transform. Autocorrelation and stationary wavelet transforms are examples of redundant
  transforms.
- `tree::Union{BitVector, Nothing}`: (Default: `nothing`) The basis tree for decomposing
  `xw`. Must be provided if `xw` is decomposed using `swpd` or `acwpd`.
- `elbows::Integer`: (Default: 2) Number of elbows used to determine the best threshold
  value.

# Keyword Arguments
- `makeplot::Bool`: (Default: `false`) Whether to return the plot that was used to determine
  the best threshold value.

# Returns
- `::AbstractFloat`: Best threshold value.
- `::GR.Plot`: Plot that was used to determine the best threshold value. Only returned if
  `makeplot = true`.

# Examples
```julia
x = randn(128)
wt = wavelet(WT.haar)

# noise estimate for dwt transformation
y = dwt(x, wt)
noise = relerrorthreshold(y, false)

# noise estimate for wpt transformation
tree = maketree(x, :full)
y = wpt(x, wt, tree)
noise = relerrorthreshold(y, false, tree)

# noise estimate for sdwt transformation
y = sdwt(x, wt)
noise = relerrorthreshold(y, true)

# noise estimate for swpd transformation
y = swpd(x, wt)
noise = relerrorthreshold(y, true, tree)
```

**See also:** [`noisest`](@ref), [`RelErrorShrink`](@ref)
"""
function relerrorthreshold(coef::AbstractArray{T},
                           redundant::Bool = false,
                           tree::Union{BitVector,Nothing} = nothing,
                           elbows::Integer = 2;
                           makeplot::Bool = false) where T<:Number
    # Sanity check
    @assert elbows >= 1
    # extract necessary coefficients
    if !redundant                              # dwt or wpt
        c = coef
    elseif redundant && isa(tree, Nothing)     # sdwt
        c = reshape(coef, :)
    else                                       # swpd
        leaves = getleaf(tree,:binary)
        c = reshape(coef[:, leaves], :)
    end
    # the magnitudes of the coefficients
    x = sort(abs.(c), rev = true)
    # compute the relative error curve
    r = orth2relerror(c)
    # shift the data points
    push!(x, 0)
    pushfirst!(r, r[1])
    # reorder the data points
    xmax = maximum(x)
    ymax = maximum(r)
    x = x[end:-1:1]/xmax
    y = r[end:-1:1]/ymax
    # compute elbow method
    ix = Vector{Integer}(undef, elbows)
    A = Vector{Vector{T}}(undef, elbows)
    v = Vector{Vector{T}}(undef, elbows)
    ix[1], A[1], v[1] = findelbow(x,y)          # First elbow point
    for i in 2:elbows                           # Second elbow point and beyond
        @inbounds (ix[i], A[i], v[i]) = findelbow(x[1:ix[i-1]], y[1:ix[i-1]])
    end
    # plot relative error curve
    if makeplot
        p = relerrorplot(x*xmax, y*ymax, ix, A, v)
        return x[ix[end]]*xmax, p
    else
        return x[ix[end]]*xmax
    end
end

"""
    orth2relerror(orth)

Given a vector 'orth' of orthonormal expansion coefficients, return a vector of relative
approximation errors when retaining the 1,2,...,N largest coefficients in magnitude.

# Arguments
- `orth::AbstractVector{T} where T<:Number`: Vector of coefficients.

# Returns
`::Vector{T}`: Relative errors.

**See also:** [`RelErrorShrink`](@ref), [`relerrorthreshold`](@ref), [`findelbow`](@ref)
"""
function orth2relerror(orth::AbstractVector{T}) where T<:Number
    # sort the coefficients
    orth = sort(orth.^2, rev = true)
    # compute the relative errors
    return ((abs.(sum(orth) .- cumsum(orth))).^0.5) / sum(orth).^0.5
end

"""
    findelbow(x, y)

Given the x and y coordinates of a curve, return the elbow.

# Arguments
- `x::AbstractVector{T} where T<:Number`: x-coordinates.
- `y::AbstractVector{T} where T<:Number`: y-coordinates.

# Returns
- `::Integer`: Index of the elbow point.
- `::Vector{T}`: Length of adjacent sides.
- `::Vector{T}`: The y-coordinates going in the direction of (x₁, y₁) to (xₙ, yₙ)

**See also:** [`RelErrorShrink`](@ref), [`relerrorthreshold`](@ref), [`orth2relerror`](@ref)
"""
function findelbow(x::AbstractVector{T}, y::AbstractVector{T}) where T<:Number
    # a unit vector pointing from (x1,y1) to (xN,yN)
    v = [x[end] - x[1], y[end] - y[1]]
    v = v/norm(v,2)
    # subtract (x1,y1) from the coordinates
    xy = [x.-x[1] y.-y[1]]
    # the hypothenuse
    H = reshape((sum(xy.^2, dims = 2)).^0.5, :)
    # the adjacent side
    A = xy * v
    # the opposite side
    O = abs.(H.^2 - A.^2).^0.5
    # return the largest distance
    return (findmax(O)[2], A, v)
end

"""
    relerrorplot(x, y, ix, A, v)

Relative error plot used for threshold determination using the elbow rule.

# Arguments
- `x::Vector{T} where T<:Number`: x-coordinates.
- `y::Vector{T} where T<:Number`: y-coordinates.
- `ix::Vector{<:Integer}`: Indices for elbow points.
- `A::Vector{S} where {S<:Vector{T}, T<:Number}`: Length of adjacent sides.
- `v::Vector{S} where {S<:Vector{T}, T<:Number}`: The y-coordinates going in the direction
  of (x[1], y[1]) to (x[ix], y[ix])

# Returns
`::Plot`: Relative error plot.

**See also:** [`relerrorthreshold`](@ref), [`findelbow`](@ref)
"""
function relerrorplot(x::Vector{T},
                      y::Vector{T},
                      ix::Vector{<:Integer},
                      A::Vector{S},
                      v::Vector{S}) where {T<:Number, S<:Vector{T}}
    @assert length(ix) == length(A) == length(v)
    elbows = length(ix)
    # rescale x and y values
    xmax = maximum(x)
    ymax = maximum(y)
    # relative error line
    p = plot(x, y, lw = 2, color = :blue, legend = false)
    plot!(p, xlims = (0, 1.004*xmax), ylims = (0, 1.004*ymax))
    for i in 1:elbows
        col = i + 1
        # diagonal line
        endpoint = i > 1 ? ix[i-1] : length(x)
        plot!([x[1], 1.004*x[endpoint]], [y[1], 1.004*y[endpoint]], lw = 2,
            color = col)
        # perpendicular line
        dropto = [x[1], y[1]] + A[i][ix[i]]*(v[i].*[xmax, ymax])
        plot!(p, [x[ix[i]], dropto[1]], [y[ix[i]], dropto[2]], lw = 2,
            color = col)
        # highlight point
        scatter!(p, [x[ix[i]]], [y[ix[i]]], color = col)
    end
    # add plot labels
    plot!(p, xlabel = "Threshold", ylabel = "Relative Error")
    return p
end

# TODO: add additional threshold determination methods

# ========== Denoising ==========
"""
    denoise(x, inputtype, wt[; L=maxtransformlevels(size(x,1)),
        tree=maketree(size(x,1), L, :dwt), dnt=VisuShrink(size(x,1)),
        estnoise=noisest, smooth=:regular])

Extension of the `denoise` function from `Wavelets.jl`. Denoise a signal of
input type `inputtype`.

# Arguments:
- `x::AbstractArray{<:Number}`: input signals/coefficients.
- `inputtype::Symbol`: input type of `x`. Current accepted types of inputs are
    - `:sig`: original signals; `x` should be a 2-D array with each column
        representing a signal.
    - `:dwt`: `dwt`-transformed signal coefficients; `x` should be a 1-D array
        with each column representing the coefficients of a signal.
    - `:wpt`: `wpt`-transformed signal coefficients; `x` should be a 1-D array
        with each column representing the coefficients of a signal.
    - `:sdwt`: `sdwt`-transformed signal coefficients; `x` should be a 2-D array
        with each column representing the coefficients of a node.
    - `:swpd`: `swpd`-transformed signal coefficients; `x` should be a 2-D array
        with each column representing the coefficients of a node.
    - `:acwt`: `acwt`-transformed signal coefficients from
        AutocorrelationShell.jl; `x` should be a 2-D array with each column
        representing the coefficients of a node.
    - `:acwpt`: `acwpt`-transformed signal coefficients from
        AutocorrelationShell.jl; `x` should be a 2-D array with each column
        representing the coefficients of a node.
- `wt::Union{DiscreteWavelet, Nothing}`: the discrete wavelet to be used for
    decomposition (for input type `:sig`) and reconstruction. `nothing` can
    be supplied if no reconstruction is necessary.
- `L::Integer`: the number of decomposition levels. Necessary for input types
    `:sig`, `:dwt`, and `:sdwt`. Default value is set to be
    `maxtransformlevels(size(x,1))`.
- `tree::BitVector`: the decomposition tree of the signals. Necessary for input
    types `:wpt` and `:swpd`. Default value is set to be
    `maketree(size(x,1), L, :dwt)`.
- `dnt::DNFT`: denoise type. Default type is set to be `VisuShrink(size(x,1))`.
- `estnoise::Union{Function, Vector{<:Number}}`: noise estimation. Input can be
    provided as a function to estimate noise in signal, or a vector of estimated
    noise. Default is set to be the `noisest` function.
- `smooth::Symbol`: the smoothing method used. `:regular` smoothing thresholds
    all given coefficients, whereas `:undersmooth` smoothing does not threshold
    the lowest frequency subspace node of the wavelet transform. Default is set
    to be `:regular`.

**See also:** [`denoiseall`](@ref), [`noisest`](@ref),
    [`relerrorthreshold`](@ref)
"""
function Wavelets.Threshold.denoise(x::AbstractArray{T},
        inputtype::Symbol,
        wt::Union{DiscreteWavelet,Nothing};
        L::Integer=maxtransformlevels(size(x,1)),
        tree::BitVector=maketree(size(x,1), L, :dwt),
        dnt::S=VisuShrink(size(x,1)),
        estnoise::Union{Function,Number}=noisest,   # can be precomputed
        smooth::Symbol=:regular) where {T<:Number, S<:DNFT}

    @assert smooth ∈ [:undersmooth, :regular]
    @assert inputtype ∈ [:sig, :dwt, :wpt, :sdwt, :swpd, :acwt, :acwpt]

    # wavelet transform if inputtype == :sig
    if inputtype == :sig
        wt === nothing && error("inputtype=:sig not supported with wt=nothing")
        x = dwt(x, wt, L)
        inputtype = :dwt
    end

    if inputtype == :dwt                    # regular dwt
        # noise estimation
        σ = isa(estnoise, Function) ? estnoise(x, false, nothing) : estnoise
        # thresholding
        if smooth == :regular
            x̃ = threshold(x, dnt.th, σ*dnt.t)
        else    # :undersmooth
            n₀ = nodelength(size(x,1), L)
            x̃ = [x[1:n₀]; threshold(x[(n₀+1):end], dnt.th, σ*dnt.t)]
        end
        # reconstruction
        y = wt === nothing ? x̃ : idwt(x̃, wt, L)

    elseif inputtype == :wpt                # regular wpt
        # noise estimation
        σ = isa(estnoise, Function) ? estnoise(x, false, tree) : estnoise
        # thresholding
        if smooth == :regular
            x̃ = threshold(x, dnt.th, σ*dnt.t)
        else    # :undersmooth
            crng = coarsestscalingrange(x, tree, false)
            rng = setdiff(1:size(x,1), crng)
            x̃ = [x[crng]; threshold(x[rng], dnt.th, σ*dnt.t)]
        end
        # reconstruction
        y = wt === nothing ? x̃ : iwpt(x̃, wt, tree)

    elseif inputtype == :sdwt               # stationary dwt
        @assert ndims(x) > 1
        # noise estimation
        σ = isa(estnoise, Function) ? estnoise(x, true, nothing) : estnoise
        # thresholding
        if smooth == :regular
            x̃ = copy(x)
            threshold!(x̃, dnt.th, σ*dnt.t)
        else    # :undersmooth
            temp = x[:,2:end]
            x̃ = [x[:,1] threshold!(temp, dnt.th, σ*dnt.t)]
        end
        # reconstruction
        y = wt === nothing ? x̃ : isdwt(x̃, wt)

    elseif inputtype == :swpd               # stationary wpd
        @assert ndims(x) > 1
        # noise estimation
        σ = isa(estnoise, Function) ? estnoise(x, true, tree) : estnoise
        # thresholding
        if smooth == :regular
            leaves = findall(getleaf(tree,:binary))
            x̃ = copy(x)
            @inbounds x̃[:, leaves] = threshold!(x[:,leaves], dnt.th, σ*dnt.t)
        else    # :undersmooth
            x̃ = copy(x)
            leaves = findall(getleaf(tree,:binary))
            _, coarsestnode = coarsestscalingrange(x, tree, true)
            rng = setdiff(leaves, coarsestnode)
            @inbounds x̃[:,rng] = threshold!(x[:,rng], dnt.th, σ*dnt.t)
        end
        # reconstruction
        y = wt === nothing ? x̃ : iswpd(x̃, wt, tree)
    elseif inputtype == :acwt               # autocorrelation dwt
        @assert ndims(x) > 1
        # noise estimation
        σ = isa(estnoise, Function) ? estnoise(x, true, nothing) : estnoise
        # thresholding
        if smooth == :regular
            x̃ = copy(x)
            threshold!(x̃, dnt.th, σ*dnt.t)
        else    # :undersmooth
            temp = x[:,2:end]
            x̃ = [x[:,1] threshold!(temp, dnt.th, σ*dnt.t)]
        end
        # reconstruction
        y = iacdwt(x̃)

    else                                    # autocorrelation wpt
        @assert ndims(x) > 1
        # noise estimation
        σ = isa(estnoise, Function) ? estnoise(x, true, tree) : estnoise
        # thresholding
        if smooth == :regular
            leaves = findall(getleaf(tree,:binary))
            x̃ = copy(x)
            @inbounds x̃[:, leaves] = threshold!(x[:,leaves], dnt.th, σ*dnt.t)
        else    # :undersmooth
            x̃ = copy(x)
            leaves = findall(getleaf(tree,:binary))
            _, coarsestnode = coarsestscalingrange(x, tree, true)
            rng = setdiff(leaves, coarsestnode)
            @inbounds x̃[:,rng] = threshold!(x[:,rng], dnt.th, σ*dnt.t)
        end
        # reconstruction
        y = iacwpd(x̃, tree)
    end
    return y
end

"""
    denoiseall(x, inputtype, wt[; L=maxtransformlevels(size(x,1)),
        tree=maketree(size(x,1), L, :dwt), dnt=VisuShrink(size(x,1)),
        estnoise=noisest, bestTH=nothing, smooth=:regular])

Denoise multiple signals of input type `inputtype`.

# Arguments:
- `x::AbstractArray{<:Number}`: input signals/coefficients.
- `inputtype::Symbol`: input type of `x`. Current accepted types of inputs are
    - `:sig`: original signals; `x` should be a 2-D array with each column
        representing a signal.
    - `:dwt`: `dwt`-transformed signal coefficients; `x` should be a 2-D array
        with each column representing the coefficients of a signal.
    - `:wpt`: `wpt`-transformed signal coefficients; `x` should be a 2-D array
        with each column representing the coefficients of a signal.
    - `:sdwt`: `sdwt`-transformed signal coefficients; `x` should be a 3-D array
        with each 2-D slice representing the coefficients of a signal.
    - `:swpd`: `swpd`-transformed signal coefficients; `x` should be a 3-D array
        with each 2-D slice representing the coefficients of a signal.
    - `:acwt`: `acwt`-transformed signal coefficients from
        AutocorrelationShell.jl; `x` should be a 3-D array with each 2-D slice
        representing the coefficients of a signal.
    - `:acwpt`: `acwpt`-transformed signal coefficients from
        AutocorrelationShell.jl; `x` should be a 3-D array with each 2-D slice
        representing the coefficients of a signal.
- `wt::Union{DiscreteWavelet, Nothing}`: the discrete wavelet to be used for
    decomposition (for input type `:sig`) and reconstruction. `nothing` can
    be supplied if no reconstruction is necessary.
- `L::Integer`: the number of decomposition levels. Necessary for input types
    `:sig`, `:dwt`, and `:sdwt`. Default value is set to be
    `maxtransformlevels(size(x,1))`.
- `tree::BitVector`: the decomposition tree of the signals. Necessary for input
    types `:wpt` and `:swpd`. Default value is set to be
    `maketree(size(x,1), L, :dwt)`.
- `dnt::DNFT`: denoise type. Default type is set to be `VisuShrink(size(x,1))`.
- `estnoise::Union{Function, Vector{<:Number}}`: noise estimation. Input can be
    provided as a function to estimate noise in signal, or a vector of estimated
    noise. Default is set to be the `noisest` function.
- `bestTH::Union{Function, Nothing}`: method to determine the best threshold
    value for a group of signals. If `nothing` is given, then each signal will
    be denoised by its respective best threshold value determined from the
    parameters `dnt` and `estnoise`; otherwise some function can be passed
    to determine the best threshold value from a vector of threshold values, eg:
    `mean` and `median`. Default is set to be `nothing`.
- `smooth::Symbol`: the smoothing method used. `:regular` smoothing thresholds
    all given coefficients, whereas `:undersmooth` smoothing does not threshold
    the lowest frequency subspace node of the wavelet transform. Default is set
    to be `:regular`.

**See alse:** [`denoise`](@ref), [`noisest`](@ref), [`relerrorthreshold`](@ref)
"""
function denoiseall(x::AbstractArray{T1},
        inputtype::Symbol,
        wt::Union{DiscreteWavelet,Nothing};
        L::Integer=maxtransformlevels(size(x,1)),
        tree::BitVector=maketree(size(x,1), L, :dwt),
        dnt::S=VisuShrink(size(x,1)),
        estnoise::Union{Function,Vector{T2}}=noisest,   # can be precomputed
        bestTH::Union{Function,Nothing}=nothing,
        smooth::Symbol=:regular) where {T1<:Number, T2<:Number, S<:DNFT}

    @assert inputtype ∈ [:sig, :dwt, :wpt, :sdwt, :swpd, :acwt, :acwpt]
    @assert smooth ∈ [:regular, :undersmooth]
    @assert ndims(x) > 1
    n = size(x, 1)              # signal length
    N = size(x)[end]            # number of signals
    # wavelet transform if inputtype == :sig
    if inputtype == :sig
        wt === nothing && error("inputtype=:sig not supported with wt=nothing")
        x = hcat([dwt(x[:,i], wt, L) for i in 1:N]...)
        inputtype = :dwt
    end

    # denoise each signal
    y = Array{T1}(undef, (n,N))
    if bestTH === nothing           # using individual threshold values
        for i in axes(x, ndims(x))
            xᵢ = ndims(x) == 2 ? x[:,i] : x[:,:,i]
            # noise estimation
            σ = isa(estnoise, Function) ? estnoise : estnoise[i]
            # denoising
            @inbounds y[:,i] = denoise(xᵢ, inputtype, wt, L=L, tree=tree,
                dnt=dnt, estnoise=σ, smooth=smooth)
        end
    else                            # using summary threshold value
        # noise estimation
        if isa(estnoise, Function)
            σ = Vector{AbstractFloat}(undef, N)
            for i in axes(x, ndims(x))
                xᵢ = ndims(x) == 2 ? x[:,i] : x[:,:,i]
                if inputtype == :dwt
                    @inbounds σ[i] = estnoise(xᵢ, false, nothing)
                elseif inputtype == :wpt
                    @inbounds σ[i] = estnoise(xᵢ, false, tree)
                elseif inputtype == :sdwt
                    @inbounds σ[i] = estnoise(xᵢ, true, nothing)
                else
                    @inbounds σ[i] = estnoise(xᵢ, true, tree)
                end
            end
            σ = bestTH(σ)
        else
            σ = bestTH(estnoise)
        end
        # denoising
        for i in axes(x, ndims(x))
            xᵢ = ndims(x) == 2 ? x[:,i] : x[:,:,i]
            @inbounds y[:,i] = denoise(xᵢ, inputtype, wt, L=L, tree=tree,
                dnt=dnt, estnoise=σ, smooth=smooth)
        end
    end
    return y
end

end # end module


module bingingstests

using GridapPardiso
using Test
using SparseArrays

# Define linear system

mtype = GridapPardiso.MTYPE_REAL_NON_SYMMETRIC

A = sparse([
  0. -2  3 0
  -2  4 -4 1
  -3  5  1 1
  1 -3 0 2 ])

n = A.n

b = Float64[1, 3, 2, 5]

x = zeros(Float64,n)

# Create the pardiso internal handler

pt = new_pardiso_handle()

# pardisoinit!

iparm = Vector{Int32}(new_iparm())

pardisoinit!(pt,mtype,iparm)

# pardiso! (solving the transpose of the system above)

maxfct = 1
mnum = 1
phase = GridapPardiso.PHASE_ANALYSIS_NUMERICAL_FACTORIZATION_SOLVE_ITERATIVE_REFINEMENT
a = A.nzval
ia = Vector{Int32}(A.colptr)
ja = Vector{Int32}(A.rowval)
perm = zeros(Int32,n)
nrhs = 1
msglvl = 0

err = pardiso!(
  pt,maxfct,mnum,mtype,phase,n,a,ia,ja,perm,nrhs,iparm,msglvl,b,x)

tol = 1.0e-13

@test err == 0

@test maximum(abs.(A'*x-b)) < tol

# pardiso! (solving the system above)

iparm[12] = 2

err = pardiso!(
  pt,maxfct,mnum,mtype,phase,n,a,ia,ja,perm,nrhs,iparm,msglvl,b,x)

@test err == 0

@test maximum(abs.(A*x-b)) < tol

@test pardiso_data_type(mtype,iparm) == Float64

# pardiso_getdiag!

iparm[56] = 1

pt = new_pardiso_handle()

err = pardiso!(
  pt,maxfct,mnum,mtype,phase,n,a,ia,ja,perm,nrhs,iparm,msglvl,b,x)

df = zeros(Float64,n)
da = zeros(Float64,n)

err = pardiso_getdiag!(pt,df,da,mnum,mtype,iparm)

@test err == 0

err = pardiso_getdiag!(pt,df,da,mnum)

@test err == 0

if Int == Int64
    # pardiso_64! (solving the transpose of the system above)
    pt = new_pardiso_handle()

    a = A.nzval
    ia = A.colptr
    ja = A.rowval
    perm = zeros(Int64,n)
    iparm = Vector{Int64}(new_iparm())

    err = pardiso_64!(
      pt,maxfct,mnum,mtype,phase,n,a,ia,ja,perm,nrhs,iparm,msglvl,b,x)

    @test err == 0

    # pardiso_64! (solving the transpose of the system above)

    @test maximum(abs.(A'*x-b)) < tol
end

end


function insertionsort!(x::AbstractArray)
    n = length(x)
    for i ∈ 1:n
        @inbounds key = x[i]
        j = i-1
        @inbounds while j > 0 && x[j] > key
            @inbounds x[j+1] = x[j]
            j-=1
        end
        @inbounds x[j+1] = key
    end
    return nothing
end


### A Pluto.jl notebook ###
# v0.12.7

using Markdown
using InteractiveUtils

# ╔═╡ d17ee77a-2333-11eb-088d-933433605dae
md"""
Définissons un dictionnaire dans lequel on stocke le nombre de moles de ``O_2`` et ``N_2``.
```math
\left ( n_i \right ).
```

Ce faisant, on utilise

* Le langage Julia,
* Dans un environment *notebook* (Pluto).

"""

# ╔═╡ f94ab6c6-2333-11eb-2c7a-916758076373
begin
	n = Dict{Symbol,Float64}()
	n[:O₂] = 1
	n[:N₂] = 3.76
	n
end

# ╔═╡ 335650c8-2334-11eb-2698-59a57f16c669
md"""
Commençons par calculer les fractions molaires.
```math
\forall i = 1, \ldots, N, \quad y_i \equiv \frac{n_i}{\sum_j n_j}.
```

"""

# ╔═╡ 53ac85f4-2334-11eb-222f-67ffde76ecab
begin
	y = Dict{Symbol,Float64}()
	for (i, val) in n
		y[i] = val / sum(values(n))
	end
	y
end

# ╔═╡ 8d6ba298-2334-11eb-3d1b-b7dc5ab65b47
md"""
On rappelle la définition des fractions massiques :
```math
\forall i = 1, \ldots, N, \quad Y_i \equiv \frac{m_i}{\sum_j m_j}.
```

On a besoin de définir les masses molaires des espèces en présense.

"""

# ╔═╡ b21ec9ee-2334-11eb-1eee-371cfe344942
begin
	M = Dict{Symbol,Float64}()
	M[:O₂] = 2 * 16
	M[:N₂] = 2 * 14
	M
end

# ╔═╡ ef206274-2334-11eb-315e-b16f16b7333e
md"""
On commence par calculer la masse
```math
\forall i \in 1, \ldots, N, \quad m_i = M_i n_i.
```

"""

# ╔═╡ d19ecd1e-2334-11eb-2660-3113ff91a503
begin
	m = Dict{Symbol,Float64}()
	for i in intersect(keys.((n, M))...)
		m[i] = M[i] * n[i]
	end
	m
end

# ╔═╡ 1778e6d0-2335-11eb-1b8b-239c96f88297
md"""
Il ne reste plus qu'à évaluer les ``\left ( Y_i \right )``.

"""

# ╔═╡ 12702824-2335-11eb-1e27-8d77009d5f2c
begin
	Y = Dict{Symbol,Float64}()
	for (i, val) in m
		Y[i] = val / sum(values(m))
	end
	Y
end

# ╔═╡ 38b0ed0c-2335-11eb-158f-3112b5f75e9a
y, Y

# ╔═╡ Cell order:
# ╟─d17ee77a-2333-11eb-088d-933433605dae
# ╠═f94ab6c6-2333-11eb-2c7a-916758076373
# ╟─335650c8-2334-11eb-2698-59a57f16c669
# ╠═53ac85f4-2334-11eb-222f-67ffde76ecab
# ╟─8d6ba298-2334-11eb-3d1b-b7dc5ab65b47
# ╠═b21ec9ee-2334-11eb-1eee-371cfe344942
# ╟─ef206274-2334-11eb-315e-b16f16b7333e
# ╠═d19ecd1e-2334-11eb-2660-3113ff91a503
# ╟─1778e6d0-2335-11eb-1b8b-239c96f88297
# ╠═12702824-2335-11eb-1e27-8d77009d5f2c
# ╠═38b0ed0c-2335-11eb-158f-3112b5f75e9a


using LinearAlgebra

"""
    initbox(x::AbstractVector{T}) where {T}

Create a bounding box and initialize with a single point.
"""
function initbox(x::AbstractVector{T}) where {T}
    sdim = length(x)
    box = fill(zero(T), 2*sdim)
    for i in 1:sdim
        box[2*i-1] = box[2*i] = x[i];
    end
    return box
end

function initbox(x::AbstractMatrix{T}) where {T}
	box = initbox(x[1, :])
	return updatebox!(box, x)
end

"""
    updatebox!(box::AbstractVector{T}, x::AbstractVector{T}) where {T}

Update a box with another location, or create a new box.

If the  `box` does not have  the correct dimensions,  it is correctly sized.

`box` = bounding box
    for 1-D `box=[minx,maxx]`, or
    for 2-D `box=[minx,maxx,miny,maxy]`, or
    for 3-D `box=[minx,maxx,miny,maxy,minz,maxz]`
`x` = vector defining a point.  The `box` is expanded to include the
    supplied location `x`.
"""
function updatebox!(box::AbstractVector{T}, x::AbstractVector{T}) where {T}
    sdim = length(x)
    for i in 1:sdim
    	box[2*i-1] = min(box[2*i-1],x[i]);
    	box[2*i]   = max(box[2*i],x[i]);
    end
    return box
end

function updatebox!(box::AbstractVector{T}, x::AbstractMatrix{T}) where {T}
    for i in 1:size(x, 1)
    	updatebox!(box, x[i, :])
    end
    return box
end

"""
    boundingbox(x::AbstractArray{T}) where {T}

Compute the bounding box of the points in `x`.

`x` = holds points, one per row.

Returns `box` = bounding box
    for 1-D `box=[minx,maxx]`, or
    for 2-D `box=[minx,maxx,miny,maxy]`, or
    for 3-D `box=[minx,maxx,miny,maxy,minz,maxz]`
"""
function boundingbox(x::AbstractArray{T}) where {T}
	box = initbox(x[1, :])
	for i in 2:size(x, 1)
		updatebox!(box, x)
	end
    return box
end

"""
    inflatebox!(box::AbstractVector{T}, inflatevalue::T) where {T}

Inflate the box by the value supplied.
"""
function inflatebox!(box::AbstractVector{T}, inflatevalue::T) where {T}
    abox = deepcopy(box)
    sdim = Int(length(box)/2);
    for i=1:sdim
        box[2*i-1] = min(abox[2*i-1],abox[2*i]) - inflatevalue;
        box[2*i]   = max(abox[2*i-1],abox[2*i]) + inflatevalue;
    end
    return box
end

"""
    inbox(box::AbstractVector{T}, x::AbstractVector{T}) where {T}

Is the given location inside the box?

- `box` = vector entries arranged as [minx,maxx,miny,maxy,minz,maxz] (or
  adjusted in an obvious way for lower space dimension).

Note: point on the boundary of the box is counted as being inside.
"""
function inbox(box::AbstractVector{T}, x::AbstractVector{T}) where {T}
    inrange(rangelo,rangehi,r) = ((r>=rangelo) && (r<=rangehi));
    sdim=length(x);
    @assert 2*sdim == length(box)
    if !inrange(box[1], box[2], x[1])
        return false # short-circuit
    end
    for i=2:sdim
        if !inrange(box[2*i-1], box[2*i], x[i])
            return false # short-circuit
        end
    end
    return true
end

function inbox(box::AbstractVector{T}, x::AbstractArray{T}) where {T}
    return inbox(box, vec(x))
end

"""
    boxesoverlap(box1::AbstractVector{T}, box2::AbstractVector{T}) where {T}

Do the given boxes overlap?
"""
function boxesoverlap(box1::AbstractVector{T}, box2::AbstractVector{T}) where {T}
    dim=Int(length(box1)/2);
    @assert 2*dim == length(box2) "Mismatched boxes"
    for i=1:dim
        if box1[2*i-1]>box2[2*i]
            return false;
        end
        if box1[2*i]<box2[2*i-1]
            return false;
        end
    end
    return  true;
end

"""
    intersectboxes(box1::AbstractVector{T}, box2::AbstractVector{T}) where {T}

Compute the intersection of two boxes.

The function returns an empty box (length(b) == 0) if the intersection is
empty; otherwise a box is returned.
"""
function intersectboxes(box1::AbstractVector{T}, box2::AbstractVector{T}) where {T}
    @assert length(box1) == length(box2) "Mismatched boxes"
    b = copy(box1)
    dim=Int(length(box1)/2);
    @assert 2*dim == length(box2) "Wrong box data"
    for i=1:dim
        lb = max(box1[2*i-1], box2[2*i-1])
        ub = min(box1[2*i], box2[2*i])
        if (ub <= lb) # intersection is empty
            return eltype(box1)[] # box of length zero signifies empty intersection
        end
        b[2*i-1] = lb
        b[2*i]   = ub
    end
    return  b;
end


"""
# Gives the number of occurrences (counts) of the vector for each
#  unique value in the vector
"""
function countUniqueValues{T<:Number}(inValues::Array{T,1})
  uValues=unique(inValues)
  uCounts=zeros(Int64,(length(uValues),))
  for j=1:length(inValues)
    for i=1:length(uValues)
      (inValues[j]==uValues[i]) ? uCounts[i]+=1 : 0
    end
  end
  return uCounts,uValues
end


using FractionalDiffEq, Plots

a=1; mu=4

function Brusselator(du, u, p, t)
    du[1] = a-(mu+1)*u[1]+u[1]^2*u[2]
    du[2] = mu*u[1]-u[1]^2*u[2]
    #du
end
alpha=[0.8; 0.8]
t0=0; tfinal=50; y0=[0.2; 0.03]
h=0.01; tspan = (t0, tfinal)
prob = FODESystem(Brusselator, alpha, y0, tspan)
#(t, y) = solve(prob, h, FLMMBDF())
#(t, y) = solve(prob, h, FLMMNewtonGregory())
(t, y) = solve(prob, h, FLMMTrap())

plot(y[1, :], y[2, :])


"""
    Ziegler-Biersack-Littmark (ZBL) Potential

See https://docs.lammps.org/pair_zbl.html
"""
mutable struct ZBL <: Potential
    ε0::Float64      # Electrical permittivity of vacuum.
    e::Float64       # Electron charge.
    zi::Float64      # Atomic number of the atom i.
    zj::Float64      # Atomic number of the atom j.
    a::Float64
    rcutfac::Float64 # Rcut
end

"""
    ZBL(params::Dict)

Creates a ZBL potential.
"""
function ZBL(params::Dict)
    # Creates the ZBL model
    ε0 = params["ε0"]
    e = params["e"]
    zi = params["zi"]
    zj = params["zj"]
    a = 0.46850 / (zi^0.23 + zj^0.23)
    rcutfac = params["rcutfac"]
    return ZBL(ε0, e, zi, zj, a, rcutfac)
end

"""
    potential_energy(p::ZBL, r::Position, args...)

Calculates the potential energy of ZBL with a switching function.
"""
function potential_energy(p::ZBL, r::Position, args...)
    return  zbl(p, norm(r)) + S(p, norm(r))
end

"""
    zbl(p::ZBL, r::Float64)

ZBL function without the switching function S.
"""
function zbl(p::ZBL, r::Float64)
    return 1 / (4.0 * π * p.ε0) * p.zi * p.zj * p.e^2 / r * ϕ(p, r / p.a)
end

"""
    S(r)

Switching function that ramps the energy, force, and curvature smoothly to zero
between an inner and outer cutoff. Here, the inner and outer cutoff are the same
for all pairs of atom types.
"""
function S(p::ZBL, r::Float64)
    return r <= p.rcutfac ? 0.5 * (cos(pi * r / p.rcutfac) + 1.0) : 0
end

# TODO: check this function: "exp" or "e"?
"""
    ϕ(p::ZBL, x::Float64)

Auxiliary function, necessary to calculate ZBL.
"""
function ϕ(p::ZBL, x::Float64)
    return   0.18175 * exp(-3.1998*x)  + 0.50986 * exp(-0.94229*x) +
             0.28022 * exp(-0.40290*x) + 0.2817  * exp(-0.20162*x)
end


using Distributions, JuMP, Ipopt, Random
include("KPNormal.jl")
using KP

function rob(cs, muhat, vs, r)
    const sqrt_vmin = sqrt(minimum(vs))
    function f(Gamma)
        xs, lam = KP.x_l2reg(cs, muhat, vs, Gamma)
#        println(Gamma, "\t", Gamma * sqrt_vmin * v_norm(xs, vs) - r)
        return Gamma * sqrt_vmin * v_norm(xs, vs) - r
    end
    println(f(.01), "\t", f(100))
    Gammastar = fzero(f, .01, 100)
    return KP.x_l2reg(cs, muhat, vs, Gammastar)
end

function rob_direct(cs, muhat, vs, r)
    m = Model(solver=IpoptSolver())
    @variable(m, xs[1:n] >= 0)
    @constraint(m, sum(cs[i]* xs[i] for i = 1:n) <= n)
    @NLobjective(m, Max, sum(muhat[i] * xs[i]/n for i = 1:n) - r/n * sqrt( sum(xs[i]^2 for i = 1:n)))
    solve(m)
end

v_norm(xs, vs) = sqrt(sum(xs.^2 ./ vs))

n =  1000
Random.seed!(8675309)
thetas = ones(n)
vs = ones(n)
cs = ones(n)

#From above
theta_l, v_l, c_l = 0.0, .01,  10.001
theta_m, v_m, c_m = 1.0, .75, 10.001
theta_h, v_h, c_h = 0.3,  2, 10.001

thetas[1:3:n] = theta_l
vs[1:3:n] = v_l
cs[1:3:n] = c_l

thetas[2:3:n] = theta_m
vs[2:3:n] = v_m
cs[2:3:n] = c_m

thetas[3:3:n] = theta_h
vs[3:3:n] = v_h
cs[3:3:n] = c_h

muhat = thetas + randn(n) ./ sqrt(vs)
sqrt_vmin = sqrt(minimum(vs))

tic()
xs, lam = rob(cs, muhat, vs, quantile(Normal(0, 1), .99))
toc()

tic()
xs2 = rob_direct(cs, muhat, vs, quantile(Normal(0, 1), .99))
toc()


# matmul_loop_vectorization.jl
module Matmul
export foo!
include("row_major.jl")

# if dont have LoopVectorization pkg, install it as below
# import Pkg
# Pkg.add("LoopVectorization")
using LoopVectorization

function gemmavx(x, y, z)
    @turbo for m ∈ axes(x, 1), n ∈ axes(y, 2)
        zmn = zero(eltype(z))
        for k ∈ axes(x, 2)
            zmn += x[m, k] * y[k, n]
        end
        z[m, n] = zmn
    end
    return z
end

# z is output, should use . to inplace
# julia array is column-major, numpy aray is row-major
# user should transpose julia or numpy's array to keep same behavior
function foo!(x, y, z)
    x = change_input_to_row_major(x)
    y = change_input_to_row_major(y)
    z .= gemmavx(x, y, z)
    z .= change_output_to_row_major(z)
end

end


using Devectorize
function cost(expanded_target_batch, output_layer_state, batchsize)
@devec CE = -sum(expanded_target_batch .* log(output_layer_state + exp(-30))) ./ batchsize
return CE
end