module Test_dss

using DescriptorSystems
using LinearAlgebra
using Polynomials
using Test

println("Test_dss")

@testset "test_dss" begin
# SCALARS
a_2 = [-5 -3; 2 -9];
CSS_111 = dss(-5, 2, 3, [0])
CSS_111_d = dss([3], 2, 1, 1)
CSS_211 = dss(a_2, [1; 2], [1 0], 0)
CSS_221 = dss(a_2, [1 0; 0 2], [1 0], 0)
CSS_222 = dss(a_2, [1 0; 0 2], eye(2), 0)

a_2 = [-5 -3; 2 -9]; e_2 = [1 1; 0 1];
CS_111 = dss(-5, 1, 2, 3, [0])
CS_111_d = dss([3], I, 2, 1, 1)
CS_211 = dss(a_2, e_2, [1; 2], [1 0], 0)
CS_221 = dss(a_2, e_2, [1 0; 0 2], [1 0], 0)
CS_222 = dss(a_2, e_2,[1 0; 0 2], eye(2), 0)

# CONTINUOUS
a_1 = [-5]
C_111 = dss(a_1, [2], [3], [0])
C_211 = dss(a_2, e_2,  [1; 2], [1 0], [0])
C_212 = dss(a_2, [1; 2], eye(2), [0; 0])
C_221 = dss(a_2, e_2, [1 0; 0 2], [1 0], [0 0])
C_222 = dss(a_2, e_2, [1 0; 0 2], eye(2), zeros(Int,2,2))
C_222_d = dss(a_2, [1 0; 0 2], eye(2), eye(2))
C_022 = dss(4.0*eye(2))

# DISCRETE
da_1 = [-0.5]
da_2 = [0.2 -0.8; -0.8 0.07]; de_2 = [1 1; 0 1];
D_111 = dss(da_1, [2], [3], [0], Ts = 0.005)
D_211 = dss(da_2, de_2,[1; 2], [1 0], [0], Ts = 0.005)
D_221 = dss(da_2, [1 0; 0 2], [1 0], [0 0], Ts = 0.005)
D_222 = dss(da_2, [1 0; 0 2], eye(2), zeros(2,2), Ts = 0.005)
D_222_d = dss(da_2, [1 0; 0 2], eye(2), eye(2), Ts = 0.005)
D_022 = dss(4.0*eye(2), Ts = 0.005)

# TESTS
# Contstruct with scalars
@test CS_111 == C_111
@test CS_111_d == dss([3],[2],[1],[1])
@test CS_211 == C_211
@test CS_221 == C_221
@test CS_222 == C_222

# Addition
@test C_111 + C_111 == dss([-5 0; 0 -5],[2; 2],[3 3],[0])
@test C_222 + C_222 == dss([-5 -3 0 0; 2 -9 0 0; 0 0 -5 -3;0 0 2 -9],
[ 1  1  0  0; 0  1  0  0; 0  0  1  1; 0  0  0  1],
[1 0; 0 2; 1 0; 0 2], [1 0 1 0; 0 1 0 1],[0 0; 0 0])
@test C_222 + 1 == dss(a_2,e_2,[1 0; 0 2],[1 0; 0 1],[1 1; 1 1])
@test D_111 + D_111 == dss([-0.5 0; 0 -0.5],[2; 2],[3 3],[0], Ts = 0.005)
@test C_222 + [1 2; 3 4] == dss([-5 -3; 2 -9],[ 1  1; 0 1],
[1 0; 0 2], [1 0; 0 1],[1 2; 3 4])
@test [1 2; 3 4] + C_222  == dss([-5 -3; 2 -9],[ 1  1; 0 1],
[1 0; 0 2], [1 0; 0 1],[1 2; 3 4])
@test C_222 + 2I == dss([-5 -3; 2 -9],[ 1  1; 0 1],
[1 0; 0 2], [1 0; 0 1],[2 0; 0 2])
@test I + C_222  == dss([-5 -3; 2 -9],[ 1  1; 0 1],
[1 0; 0 2], [1 0; 0 1],[1 0; 0 1])

# Subtraction
@test C_111 - C_211 == dss([-5 0 0; 0 -5 -3; 0 2 -9],[1  0  0; 0  1  1; 0  0  1],[2; 1; 2],[3 -1 -0],[0])
@test 1 - C_222 == dss(a_2,e_2,[1 0; 0 2],[-1 -0; -0 -1],[1 1; 1 1])
@test D_111 - D_211 == dss([-0.5 0 0; 0 0.2 -0.8; 0 -0.8 0.07],  [1.0  0.0  0.0; 0.0  1.0  1.0; 0.0  0.0  1.0],
[2; 1; 2], [3 -1 -0],[0], Ts = 0.005)
@test C_222 - [1 2; 3 4] == dss([-5 -3; 2 -9],[ 1  1; 0 1],
[1 0; 0 2], [1 0; 0 1],[-1 -2; -3 -4])
@test [1 2; 3 4] - C_222  == dss([-5 -3; 2 -9],[ 1  1; 0 1],
[1 0; 0 2], [-1 0; 0 -1],[1 2; 3 4])
@test C_222 - 2I == dss([-5 -3; 2 -9],[ 1  1; 0 1],
[1 0; 0 2], [1 0; 0 1],[-2 0; 0 -2])
@test I - C_222  == dss([-5 -3; 2 -9],[ 1  1; 0 1],
[1 0; 0 2], [-1 0; 0 -1],[1 0; 0 1])
@test C_222 - 2 == dss([-5 -3; 2 -9],[ 1  1; 0 1],
[1 0; 0 2], [1 0; 0 1],[-2 -2; -2 -2])

# Multiplication
@test C_111 * C_221 == dss([-5 2 0; 0 -5 -3; 0 2 -9],[1  0  0; 0  1  1; 0  0  1],
[0 0; 1 0; 0 2],[3 0 0],[0 0])
@test C_212 * C_111 == dss([-5 -3 3; 2 -9 6; 0 0 -5],
[0; 0; 2],[1 0 0; 0 1 0],[0; 0])
@test 4*C_222 == dss(a_2,e_2,[1 0; 0 2],[4 0; 0 4],[0 0; 0 0])
@test C_222*4 == dss(a_2,e_2,[4 0; 0 8],[1 0; 0 1],[0 0; 0 0])
@test D_111 * D_221 == dss([-0.5 2 0; 0 0.2 -0.8; 0 -0.8 0.07],
[0 0; 1 0; 0 2],[3 0 0],[0 0],Ts = 0.005)
@test [4 0;0 1]*C_222 == dss(a_2,e_2,[1 0; 0 2],[4 0; 0 1],[0 0; 0 0])
@test C_222*[4 0;0 1] == dss(a_2,e_2,[4 0; 0 2],[1 0; 0 1],[0 0; 0 0])

# Right Division
@test iszero(1/C_222_d - inv(C_222_d))
@test iszero(C_221/C_222_d - C_221*inv(C_222_d))
@test iszero(1/D_222_d - inv(D_222_d))
@test C_222/(4I) == dss(a_2,e_2,[0.25 0; 0 0.5],[1 0; 0 1],[0 0; 0 0])
@test iszero(C_222/4 -C_222*(1/4))

# Left Division
@test iszero(C_222_d\1 - inv(C_222_d))
@test iszero(C_222_d\C_212 - inv(C_222_d)*C_212)
@test iszero(D_222_d\1 - inv(D_222_d))
@test 4I\C_222 == dss(a_2,e_2,[1 0; 0 2],[0.25 0; 0 0.25],[0 0; 0 0])
@test iszero(4\C_222 -(1/4)*C_222)

# Indexing, order, ndims
@test order(C_222) == 2
@test size(C_222) == (2, 2)
@test size(C_212) == (2, 1)
@test ndims(C_222) == 1
@test length(C_222) == 1
@test C_222[1,1] == dss([-5 -3; 2 -9],[1 1; 0 1],[1; 0],[1 0],[0])
@test C_222[1:1,1] == dss([-5 -3; 2 -9],[1 1; 0 1],[1; 0],[1 0],[0])
@test C_222[1,1:2] == C_221
@test size(C_222[1,[]]) == (1,0)
@test C_222[:,2:end]  ≈ dss([-5 -3; 2 -9],[1 1; 0 1],[0; 2],[1 0; 0 1],[0; 0])

A = [-1.0 -2.0; 0.0 -1.0]
E = [1.0 2.0; 0.0 1.0]
B = [0.0; -2.0]
C = [1.0 1.0]
D = 1.0
sys = dss(A, E, B, C, D)
sysd = dss(A, E, B, C, D, Ts = -1)

@test sys + 1.0 == dss(A, E, B, C, D + 1.0)
@test 2.0 + sys == dss(A, E, B, C, D + 2.0)

@test -sys == dss(A, E, B, -C, -D)

# transpose, dual, adjoints
@test sys == transpose(transpose(sys))
@test iszero(sys -gdual(gdual(sys,rev=true)))
@test iszero(sys - adjoint(adjoint(sys)))
@test iszero(sysd - adjoint(adjoint(sysd)))
@test iszero(sysd - gsvselect(sysd,[2,1]))

# Accessing Ts through .Ts
@test D_111.Ts == 0.005

# property names
@test propertynames(C_111) == (:nx, :ny, :nu, :A, :E, :B, :C, :D, :Ts)
@test propertynames(D_111) == (:nx, :ny, :nu, :A, :E, :B, :C, :D, :Ts)

# Building descriptor systems from rational functions and rational matrices
s = Polynomial([0, 1],'s');
g = (s+0.01)/(1+0.01*s);
sys1 = dss(g);
G  = [s^2 s/(s+1); 0 one(s)/s]     # define the 2-by-2 rational matrix G(s)
sys2 = dss(G);

# from Laurent polynomial
t = LaurentPolynomial([1],-1,:z)
p = LaurentPolynomial([24,10,-15,0,1],-2,:z)
q = LaurentPolynomial([1,0,1],-1,:z)
@test p ≈ 24t^2+10t-15+t^(-2)

dss2rm(dss(p,Ts=1))

# Errors
@test_throws ErrorException C_111 + C_222             # Dimension mismatch
@test_throws ErrorException C_111 - C_222             # Dimension mismatch
@test_throws ErrorException C_111 * C_222             # Dimension mismatch
@test_throws ErrorException D_111 + C_111             # Sampling time mismatch
@test_throws ErrorException D_111 - C_111             # Sampling time mismatch
@test_throws ErrorException D_111 * C_111             # Sampling time mismatch
D_diffTs = dss([1], [2], [3], [4], Ts=0.1)
@test_throws ErrorException D_111 + D_diffTs            # Sampling time mismatch
@test_throws ErrorException D_111 - D_diffTs            # Sampling time mismatch
@test_throws ErrorException D_111 * D_diffTs            # Sampling time mismatch
@test_throws ErrorException 1/C_221                     # Not invertible
@test_throws ErrorException 1/C_212                     # Not invertible
@test_throws DimensionMismatch dss([1 2], [1], [2], [3])      # Not square A
@test_throws ErrorException dss([1], [2 0], [1], [2])      # I/0 dim mismatch
@test_throws ErrorException dss([1], [2], [3 4], [1])      # I/0 dim mismatch
@test_throws ErrorException dss([1], [2], [3], [4], Ts = -0.1)  # Negative samping time
@test_throws ErrorException dss(eye(2), eye(2), eye(2), [0]) # Dimension mismatch

try
    s = sprint(show,rdss(1,1,1));
    s = sprint(show,rss(1,1,1));
    s = sprint(show,rdss(1,0,1));
    s = sprint(show,rdss(1,1,0));
    s = sprint(show,rdss(0,1,1));
    s = sprint(show,rdss(0,0,1));
    @test true
catch
    @test false
end

end

end #module


# ----------------------------------------------------------------------------
# This file was created by manually inspecting 2D Brillouin zones.
# ----------------------------------------------------------------------------

const pathsd_2d = Dict(
    :hp1 => [[:Γ, :M, :K, :Γ]],
    # ------------------------------------------------------------------------------------ #
    :tp1 => [[:Γ, :M, :X, :Γ]],
    # ------------------------------------------------------------------------------------ #
    :op1 => [[:Γ, :Y, :S, :X, :Γ]], # op1 and op2 are really the same paths; they differ
    :op2 => [[:Γ, :X, :S, :Y, :Γ]], # only in the definition of :X and :Y points (swapped)
    # ------------------------------------------------------------------------------------ #
    :oc1 => [[:C₀, :Y, :Γ, :Σ₀], [:Γ, :S]], # the k-point labels picked here are plane group
    :oc2 => [[:Δ₀, :Γ, :Y, :F₀], [:Γ, :S]], # 9's; they may not agree with plane group 5
    # TODO: think about if we can do better for plane group 5
    # ------------------------------------------------------------------------------------ #
    :mp1 => [[:Γ, :B], [:Γ, :Y], [:Γ, :A]], # the two settings have different definitions of
    :mp2 => [[:Γ, :B], [:Γ, :Y], [:Γ, :A]]  # :A [(1/2, +1/2) in mp1 and (1/2, -1/2) in mp2]
)

const pointsd_2d = Dict(
    :hp1 => [:Γ => :([0, 0]), :M => :([1/2, 0]), :K => :([1/3, 1/3])],
    # ------------------------------------------------------------------------------------ #
    :tp1 => [:Γ => :([0, 0]), :X => :([0, 1/2]), :M => :([1/2, 1/2])],
    # ------------------------------------------------------------------------------------ #
    :op1 => [:Γ => :([0, 0]), :X => :([0, 1/2]), :S => :([1/2, 1/2]), :Y => :([1/2, 0])],
    :op2 => [:Γ => :([0, 0]), :X => :([1/2, 0]), :S => :([1/2, 1/2]), :Y => :([0, 1/2])],
    # ------------------------------------------------------------------------------------ #
    :oc1 => [:Γ => :([0, 0]), :Y => :([-1/2, 1/2]), :C₀ => :([-δ, 1-δ]), :S => :([0, 1/2]), :Σ₀ => :([δ,  δ])],
    :oc2 => [:Γ => :([0, 0]), :Y => :([1/2, 1/2]),  :F₀ => :([δ,  1-δ]), :S => :([0, 1/2]), :Δ₀ => :([-δ, δ])],
    # ------------------------------------------------------------------------------------ #
    :mp1 => [:Γ => :([0, 0]), :B => :([1/2, 0]), :Y => :([0, 1/2]), :A => :([1/2, +1/2])],
    :mp2 => [:Γ => :([0, 0]), :B => :([1/2, 0]), :Y => :([0, 1/2]), :A => :([1/2, -1/2])]
)

const paramsd_2d = Dict(
   :oc1 => [:δ => :(0.25 * ( 1 + (a*a) / (b*b) ) )],
   :oc2 => [:δ => :(0.25 * ( 1 + (b*b) / (a*a) ) )]
)


############################################################################################
"""
$(TYPEDEF)

Utility struct to help assign boundaries to parameter.

# Fields
$(TYPEDFIELDS)
"""
struct CovarianceMatrix <: AbstractConstraint
    function CovarianceMatrix()
        return new()
    end
end

############################################################################################
function construct_transform(constraint::CovarianceMatrix, val)
    transform = Bijectors.PDBijector()
    return transform, Bijectors.inverse(transform)
end

############################################################################################
function _check(
    _rng::Random.AbstractRNG,
    constraint::CovarianceMatrix,
    val::Matrix{R},
) where {R<:Real}
    ArgCheck.@argcheck LinearAlgebra.issymmetric(val)
    return true
end

idx_upper = tag(zeros(3,3), false, true)
sum(idx_upper)
############################################################################################
#!TODO: Works with flatten/unflatten - but constraint/unconstraint seems to deduce wrong type for ReverseDiff from Bijector - works fine with ForwardDiff/Zygote
#!NOTE: Bijectors map to lower triangular matrix while most AD libraries evaluate upper triangular matrices.
function construct_flatten(
    output::Type{T},
    flattentype::F,
    unflattentype::UnflattenStrict,
    constraint::C,
    x::Matrix{R},
) where {
    T<:AbstractFloat,
    F<:FlattenTypes,
    R<:Real,
    C<:Union{CovarianceMatrix, Distributions.InverseWishart, Bijectors.PDBijector}
}
    #!NOTE: PDBijector seems to unconstrain to a Lower Diagonal Matrix
    idx_upper = tag(x, false, true)
    len = length(x)
    len_unflat = sum(idx_upper)
    function flatten_CovMatrix(x::AbstractMatrix{R}) where {R<:Real}
        ArgCheck.@argcheck length(x) == len
        return Vector{T}(flatten_Symmetric(x, idx_upper))
    end
    buffer_unflat = zeros(R, size(x))
    function CovMatrix_from_vec(v::Union{<:Real,AbstractVector{<:Real}})
        ArgCheck.@argcheck length(v) == len_unflat
        return Symmetric_from_flatten!(buffer_unflat, v, idx_upper)
    end
    return flatten_CovMatrix, CovMatrix_from_vec
end
function construct_flatten(
    output::Type{T},
    flattentype::F,
    unflattentype::UnflattenFlexible,
    constraint::C,
    x::Matrix{R},
) where {
    T<:AbstractFloat,
    F<:FlattenTypes,
    R<:Real,
    C<:Union{CovarianceMatrix, Distributions.InverseWishart, Bijectors.PDBijector}
}
    #!NOTE: PDBijector seems to unconstrain to a Lower Diagonal Matrix
    idx_upper = tag(x, false, true)
    len = length(x)
    len_unflat = sum(idx_upper)
    function flatten_CovMatrix_AD(x::AbstractMatrix{R}) where {R<:Real}
        ArgCheck.@argcheck length(x) == len
        return Vector{R}(flatten_Symmetric(x, idx_upper))
    end
    dims = size(x)
    function CovMatrix_from_vecAD(v::Union{<:Real,AbstractVector{<:Real}})
        ArgCheck.@argcheck length(v) == len_unflat
        buffer_unflat = zeros(eltype(v), dims)
        return Symmetric_from_flatten!(buffer_unflat, v, idx_upper)
    end
    return flatten_CovMatrix_AD, CovMatrix_from_vecAD
end

############################################################################################
#=
!NOTE: For this constraint, we use a Bijector as Transformer, so we do not need to add a new functors
1. MyTransformer <: AbstractTransformer, MyInverseTransformer <: AbstractTransformer
2. define a function construct_transform(MyConstraint, val) -> MyTransformer, MyInverseTransformer
3. overload unconstrain and log_abs_det_jac on MyTransformer.
4. overload constrain on MyInverseTransformer.
=#

############################################################################################
#Export
export
    CovarianceMatrix,
    construct_flatten,
    construct_transform,
    constrain,
    unconstrain,
    log_abs_det_jac


"""
    AbstractInverseRetractionMethod

Abstract type for methods for inverting a retraction (see [`inverse_retract`](@ref)).
"""
abstract type AbstractInverseRetractionMethod end

"""
    AbstractRetractionMethod

Abstract type for methods for [`retract`](@ref)ing a tangent vector to a manifold.
"""
abstract type AbstractRetractionMethod end

"""
    ApproximateInverseRetraction <: AbstractInverseRetractionMethod

An abstract type for representing approximate inverse retraction methods.

"""
abstract type ApproximateInverseRetraction <: AbstractInverseRetractionMethod end

"""
    ApproximateRetraction <: AbstractInverseRetractionMethod

An abstract type for representing approximate retraction methods.
"""
abstract type ApproximateRetraction <: AbstractRetractionMethod end

@doc raw"""
    EmbeddedRetraction{T<:AbstractRetractionMethod} <: AbstractRetractionMethod

Compute a retraction by using the retraction of type `T` in the embedding and projecting the result.

# Constructor

    EmbeddedRetraction(r::AbstractRetractionMethod)

Generate the retraction with retraction `r` to use in the embedding.
"""
struct EmbeddedRetraction{T<:AbstractRetractionMethod} <: AbstractRetractionMethod
    retraction::T
end

"""
    ExponentialRetraction <: AbstractRetractionMethod

Retraction using the exponential map.
"""
struct ExponentialRetraction <: AbstractRetractionMethod end

@doc raw"""
    ODEExponentialRetraction{T<:AbstractRetractionMethod, B<:AbstractBasis} <: AbstractRetractionMethod

Approximate the exponential map on the manifold by evaluating the ODE descripting the geodesic at 1,
assuming the default connection of the given manifold by solving the ordinary differential
equation

```math
\frac{d^2}{dt^2} p^k + Γ^k_{ij} \frac{d}{dt} p_i \frac{d}{dt} p_j = 0,
```

where ``Γ^k_{ij}`` are the Christoffel symbols of the second kind, and
the Einstein summation convention is assumed.

# Constructor

    ODEExponentialRetraction(
        r::AbstractRetractionMethod,
        b::AbstractBasis=DefaultOrthogonalBasis(),
    )

Generate the retraction with a retraction to use internally (for some approaches)
and a basis for the tangent space(s).
"""
struct ODEExponentialRetraction{T<:AbstractRetractionMethod,B<:AbstractBasis} <:
       AbstractRetractionMethod
    retraction::T
    basis::B
end
function ODEExponentialRetraction(r::T) where {T<:AbstractRetractionMethod}
    return ODEExponentialRetraction(r, DefaultOrthonormalBasis())
end
function ODEExponentialRetraction(::T, b::CachedBasis) where {T<:AbstractRetractionMethod}
    return throw(
        DomainError(
            b,
            "Cached Bases are currently not supported, since the basis has to be implemented in a surrounding of the start point as well.",
        ),
    )
end
function ODEExponentialRetraction(r::ExponentialRetraction, ::AbstractBasis)
    return throw(
        DomainError(
            r,
            "You can not use the exponential map as an inner method to solve the ode for the exponential map.",
        ),
    )
end
function ODEExponentialRetraction(r::ExponentialRetraction, ::CachedBasis)
    return throw(
        DomainError(
            r,
            "Neither the exponential map nor a Cached Basis can be used with this retraction type.",
        ),
    )
end

"""
    PolarRetraction <: AbstractRetractionMethod

Retractions that are based on singular value decompositions of the matrix / matrices
for point and tangent vector on a [`AbstractManifold`](@ref)
"""
struct PolarRetraction <: AbstractRetractionMethod end

"""
    ProjectionRetraction <: AbstractRetractionMethod

Retractions that are based on projection and usually addition in the embedding.
"""
struct ProjectionRetraction <: AbstractRetractionMethod end

"""
    QRRetraction <: AbstractRetractionMethod

Retractions that are based on a QR decomposition of the
matrix / matrices for point and tangent vector on a [`AbstractManifold`](@ref)
"""
struct QRRetraction <: AbstractRetractionMethod end

"""
    SoftmaxRetraction <: AbstractRetractionMethod

Describes a retraction that is based on the softmax function.
"""
struct SoftmaxRetraction <: AbstractRetractionMethod end

@doc raw"""
    PadeRetraction{m} <: AbstractRetractionMethod

A retraction based on the Padé approximation of order $m$
"""
struct PadeRetraction{m} <: AbstractRetractionMethod end

function PadeRetraction(m::Int)
    (m < 1) && error(
        "The Padé based retraction is only available for positive orders, not for order $m.",
    )
    return PadeRetraction{m}()
end
@doc raw"""
    CayleyRetraction <: AbstractRetractionMethod

A retraction based on the Cayley transform, which is realized by using the
[`PadeRetraction`](@ref)`{1}`.
"""
const CayleyRetraction = PadeRetraction{1}

@doc raw"""
   EmbeddedInverseRetraction{T<:AbstractInverseRetractionMethod} <: AbstractInverseRetractionMethod

Compute an inverse retraction by using the inverse retraction of type `T` in the embedding and projecting the result

# Constructor

    EmbeddedInverseRetraction(r::AbstractInverseRetractionMethod)

Generate the inverse retraction with inverse retraction `r` to use in the embedding.
"""
struct EmbeddedInverseRetraction{T<:AbstractInverseRetractionMethod} <:
       AbstractInverseRetractionMethod
    inverse_retraction::T
end

"""
    LogarithmicInverseRetraction <: AbstractInverseRetractionMethod

Inverse retraction using the [`log`](@ref)arithmic map.
"""
struct LogarithmicInverseRetraction <: AbstractInverseRetractionMethod end

@doc raw"""
    PadeInverseRetraction{m} <: AbstractRetractionMethod

An inverse retraction based on the Padé approximation of order $m$ for the retraction.
"""
struct PadeInverseRetraction{m} <: AbstractInverseRetractionMethod end

function PadeInverseRetraction(m::Int)
    (m < 1) && error(
        "The Padé based inverse retraction is only available for positive orders, not for order $m.",
    )
    return PadeInverseRetraction{m}()
end
@doc raw"""
    CayleyInverseRetraction <: AbstractInverseRetractionMethod

A retraction based on the Cayley transform, which is realized by using the
[`PadeRetraction`](@ref)`{1}`.
"""
const CayleyInverseRetraction = PadeInverseRetraction{1}

"""
    PolarInverseRetraction <: AbstractInverseRetractionMethod

Inverse retractions that are based on a singular value decomposition of the
matrix / matrices for point and tangent vector on a [`AbstractManifold`](@ref)
"""
struct PolarInverseRetraction <: AbstractInverseRetractionMethod end

"""
    ProjectionInverseRetraction <: AbstractInverseRetractionMethod

Inverse retractions that are based on a projection (or its inversion).
"""
struct ProjectionInverseRetraction <: AbstractInverseRetractionMethod end

"""
    QRInverseRetraction <: AbstractInverseRetractionMethod

Inverse retractions that are based on a QR decomposition of the
matrix / matrices for point and tangent vector on a [`AbstractManifold`](@ref)
"""
struct QRInverseRetraction <: AbstractInverseRetractionMethod end

"""
    NLSolveInverseRetraction{T<:AbstractRetractionMethod,TV,TK} <:
        ApproximateInverseRetraction

An inverse retraction method for approximating the inverse of a retraction using `NLsolve`.

# Constructor

    NLSolveInverseRetraction(
        method::AbstractRetractionMethod[, X0];
        project_tangent=false,
        project_point=false,
        nlsolve_kwargs...,
    )

Constructs an approximate inverse retraction for the retraction `method` with initial guess
`X0`, defaulting to the zero vector. If `project_tangent` is `true`, then the tangent
vector is projected before the retraction using `project`. If `project_point` is `true`,
then the resulting point is projected after the retraction. `nlsolve_kwargs` are keyword
arguments passed to `NLsolve.nlsolve`.
"""
struct NLSolveInverseRetraction{TR<:AbstractRetractionMethod,TV,TK} <:
       ApproximateInverseRetraction
    retraction::TR
    X0::TV
    project_tangent::Bool
    project_point::Bool
    nlsolve_kwargs::TK
    function NLSolveInverseRetraction(m, X0, project_point, project_tangent, nlsolve_kwargs)
        return new{typeof(m),typeof(X0),typeof(nlsolve_kwargs)}(
            m,
            X0,
            project_point,
            project_tangent,
            nlsolve_kwargs,
        )
    end
end
function NLSolveInverseRetraction(
    m,
    X0 = nothing;
    project_tangent::Bool = false,
    project_point::Bool = false,
    nlsolve_kwargs...,
)
    return NLSolveInverseRetraction(m, X0, project_point, project_tangent, nlsolve_kwargs)
end

"""
    SoftmaxInverseRetraction <: AbstractInverseRetractionMethod

Describes an inverse retraction that is based on the softmax function.
"""
struct SoftmaxInverseRetraction <: AbstractInverseRetractionMethod end

"""
    default_inverse_retraction_method(M::AbstractManifold)

The [`AbstractInverseRetractionMethod`](@ref) that is used when calling
[`inverse_retract`](@ref) without specifying the inverse retraction method.
By default, this is the [`LogarithmicInverseRetraction`](@ref).
"""
default_inverse_retraction_method(::AbstractManifold) = LogarithmicInverseRetraction()

"""
    default_retraction_method(M::AbstractManifold)

The [`AbstractRetractionMethod`](@ref) that is used when calling [`retract`](@ref) without
specifying the retraction method. By default, this is the [`ExponentialRetraction`](@ref).
"""
default_retraction_method(::AbstractManifold) = ExponentialRetraction()

"""
    inverse_retract!(M::AbstractManifold, X, p, q[, method::AbstractInverseRetractionMethod])

Compute the inverse retraction, a cheaper, approximate version of the
[`log`](@ref)arithmic map), of points `p` and `q` on the [`AbstractManifold`](@ref) `M`.
Result is saved to `X`.

Inverse retraction method can be specified by the last argument, defaulting to
[`default_inverse_retraction_method`](@ref)`(M)`. See the documentation of respective manifolds for
available methods.

See also [`retract!`](@ref).
"""
function inverse_retract!(
    M::AbstractManifold,
    X,
    p,
    q,
    m::AbstractInverseRetractionMethod = default_inverse_retraction_method(M),
)
    return _inverse_retract!(M, X, p, q, m)
end
function _inverse_retract!(M::AbstractManifold, X, p, q, ::LogarithmicInverseRetraction)
    return log!(M, X, p, q)
end

#
# dispatch to lower level
function _inverse_retract!(M::AbstractManifold, X, p, q, ::CayleyInverseRetraction)
    return inverse_retract_caley!(M, X, p, q)
end
function _inverse_retract!(M::AbstractManifold, X, p, q, m::EmbeddedInverseRetraction)
    return inverse_retract_embedded!(M, X, p, q, m.inverse_retraction)
end
function _inverse_retract!(
    M::AbstractManifold,
    X,
    p,
    q,
    ::PadeInverseRetraction{n},
) where {n}
    return inverse_retract_pade!(M, X, p, q, n)
end
function _inverse_retract!(M::AbstractManifold, X, p, q, ::PolarInverseRetraction)
    return inverse_retract_polar!(M, X, p, q)
end
function _inverse_retract!(M::AbstractManifold, X, p, q, ::ProjectionInverseRetraction)
    return inverse_retract_project!(M, X, p, q)
end
function _inverse_retract!(M::AbstractManifold, X, p, q, ::QRInverseRetraction)
    return inverse_retract_qr!(M, X, p, q)
end
function _inverse_retract!(M::AbstractManifold, X, p, q, ::SoftmaxInverseRetraction)
    return inverse_retract_softmax!(M, X, p, q)
end
function _inverse_retract!(M::AbstractManifold, X, p, q, m::NLSolveInverseRetraction)
    return inverse_retract_nlsolve!(M, X, p, q, m)
end
"""
    inverse_retract_embedded!(M::AbstractManifold, X, p, q, m::AbstractInverseRetractionMethod)

computes the mutating variant of the [`EmbeddedInverseRetraction`](@ref) using
the [`AbstractInverseRetractionMethod`](@ref) `m` in the embedding (see [`get_embedding`](@ref))
and projecting the result back.
"""
function inverse_retract_embedded!(
    M::AbstractManifold,
    X,
    p,
    q,
    m::AbstractInverseRetractionMethod,
)
    return project!(
        M,
        X,
        p,
        inverse_retract(
            get_embedding(M),
            embed(get_embedding(M), p),
            embed(get_embedding(M), q),
            m,
        ),
    )
end

"""
    inverse_retract_caley!(M::AbstractManifold, X, p, q)

computes the mutating variant of the [`CayleyInverseRetraction`](@ref),
which by default calls the first order [`PadeInverseRetraction`§(@ref).
"""
function inverse_retract_caley!(M::AbstractManifold, X, p, q)
    return inverse_retract_pade!(M, X, p, q, 1)
end

"""
    inverse_retract_pade!(M::AbstractManifold, p, q, n)

computes the mutating variant of the [`PadeInverseRetraction`](@ref)`(n)`,
"""
inverse_retract_pade!(M::AbstractManifold, p, q, n)

function inverse_retract_pade! end

"""
    inverse_retract_qr!(M::AbstractManifold, X, p, q)

computes the mutating variant of the [`QRInverseRetraction`](@ref).
"""
inverse_retract_qr!(M::AbstractManifold, X, p, q)

function inverse_retract_qr! end

"""
    inverse_retract_project!(M::AbstractManifold, X, p, q)

computes the mutating variant of the [`ProjectionInverseRetraction`](@ref).
"""
inverse_retract_project!(M::AbstractManifold, X, p, q)

function inverse_retract_project! end

"""
    inverse_retract_polar!(M::AbstractManifold, X, p, q)

computes the mutating variant of the [`PolarInverseRetraction`](@ref).
"""
inverse_retract_polar!(M::AbstractManifold, X, p, q)

function inverse_retract_polar! end

"""
    inverse_retract_nlsolve!(M::AbstractManifold, X, p, q, m::NLSolveInverseRetraction)

computes the mutating variant of the [`NLSolveInverseRetraction`](@ref) `m`.
"""
inverse_retract_nlsolve!(M::AbstractManifold, X, p, q, m::NLSolveInverseRetraction)

function inverse_retract_nlsolve! end

"""
    inverse_retract_softmax!(M::AbstractManifold, X, p, q)

computes the mutating variant of the [`SoftmaxInverseRetraction`](@ref).
"""
inverse_retract_softmax!(M::AbstractManifold, X, p, q)

function inverse_retract_softmax! end

"""
    inverse_retract(M::AbstractManifold, p, q)
    inverse_retract(M::AbstractManifold, p, q, method::AbstractInverseRetractionMethod

Compute the inverse retraction, a cheaper, approximate version of the
[`log`](@ref)arithmic map), of points `p` and `q` on the [`AbstractManifold`](@ref) `M`.

Inverse retraction method can be specified by the last argument, defaulting to
[`default_inverse_retraction_method`](@ref)`(M)`.
For available inverse retractions on certain manifolds see the documentation on the
corresponding manifold.

See also [`retract`](@ref).
"""
function inverse_retract(
    M::AbstractManifold,
    p,
    q,
    m::AbstractInverseRetractionMethod = default_inverse_retraction_method(M),
)
    return _inverse_retract(M, p, q, m)
end
function _inverse_retract(M::AbstractManifold, p, q, ::CayleyInverseRetraction)
    return inverse_retract_caley(M, p, q)
end
function _inverse_retract(M::AbstractManifold, p, q, m::EmbeddedInverseRetraction)
    return inverse_retract_embedded(M, p, q, m.inverse_retraction)
end
_inverse_retract(M::AbstractManifold, p, q, ::LogarithmicInverseRetraction) = log(M, p, q)
function _inverse_retract(M::AbstractManifold, p, q, m::NLSolveInverseRetraction)
    return inverse_retract_nlsolve(M, p, q, m)
end
function _inverse_retract(M::AbstractManifold, p, q, ::PadeInverseRetraction{n}) where {n}
    return inverse_retract_pade(M, p, q, n)
end
function _inverse_retract(M::AbstractManifold, p, q, ::PolarInverseRetraction)
    return inverse_retract_polar(M, p, q)
end
function _inverse_retract(M::AbstractManifold, p, q, ::ProjectionInverseRetraction)
    return inverse_retract_project(M, p, q)
end
function _inverse_retract(M::AbstractManifold, p, q, ::QRInverseRetraction)
    return inverse_retract_qr(M, p, q)
end
function _inverse_retract(M::AbstractManifold, p, q, ::SoftmaxInverseRetraction)
    return inverse_retract_softmax(M, p, q)
end
"""
    inverse_retract_embedded(M::AbstractManifold, p, q, m::AbstractInverseRetractionMethod)

computes the allocating variant of the [`EmbeddedInverseRetraction`](@ref) using
the [`AbstractInverseRetractionMethod`](@ref) `m` in the embedding (see [`get_embedding`](@ref))
and projecting the result back.
"""
function inverse_retract_embedded(
    M::AbstractManifold,
    p,
    q,
    m::AbstractInverseRetractionMethod,
)
    return project(
        M,
        p,
        inverse_retract(
            get_embedding(M),
            embed(get_embedding(M), p),
            embed(get_embedding(M), q),
            m,
        ),
    )
end

"""
    inverse_retract_caley(M::AbstractManifold, p, q)

computes the allocating variant of the [`CayleyInverseRetraction`](@ref),
which by default allocates and calls [`inverse_retract_caley!`](@ref).
"""
function inverse_retract_caley(M::AbstractManifold, p, q)
    X = allocate_result(M, inverse_retract, p, q)
    return inverse_retract_caley!(M, X, p, q)
end
"""
    inverse_retract_pade(M::AbstractManifold, p, q)

computes the allocating variant of the [`PadeInverseRetraction`](@ref)`(n)`,
which by default allocates and calls [`inverse_retract_pade!`](@ref ManifoldsBase.inverse_retract_pade!).
"""
function inverse_retract_pade(M::AbstractManifold, p, q, n)
    X = allocate_result(M, inverse_retract, p, q)
    return inverse_retract_pade!(M, X, p, q, n)
end

"""
    inverse_retract_polar(M::AbstractManifold, p, q)

computes the allocating variant of the [`PolarInverseRetraction`](@ref),
which by default allocates and calls [`inverse_retract_polar!`](@ref ManifoldsBase.inverse_retract_polar!).
"""
function inverse_retract_polar(M::AbstractManifold, p, q)
    X = allocate_result(M, inverse_retract, p, q)
    return inverse_retract_polar!(M, X, p, q)
end
"""
    inverse_retract_project(M::AbstractManifold, p, q)

computes the allocating variant of the [`ProjectionInverseRetraction`](@ref),
which by default allocates and calls [`inverse_retract_project!`](@ref ManifoldsBase.inverse_retract_project!).
"""
function inverse_retract_project(M::AbstractManifold, p, q)
    X = allocate_result(M, inverse_retract, p, q)
    return inverse_retract_project!(M, X, p, q)
end
"""
    inverse_retract_qr(M::AbstractManifold, p, q)

computes the allocating variant of the [`QRInverseRetraction`](@ref),
which by default allocates and calls [`inverse_retract_qr!`](@ref ManifoldsBase.inverse_retract_qr!).
"""
function inverse_retract_qr(M::AbstractManifold, p, q)
    X = allocate_result(M, inverse_retract, p, q)
    return inverse_retract_qr!(M, X, p, q)
end
"""
    inverse_retract_nlsolve(M::AbstractManifold, p, q, m::NLSolveInverseRetraction)

computes the allocating variant of the [`NLSolveInverseRetraction`](@ref) `m`,
which by default allocates and calls [`inverse_retract_nlsolve!`](@ref).
"""
function inverse_retract_nlsolve(M::AbstractManifold, p, q, m::NLSolveInverseRetraction)
    X = allocate_result(M, inverse_retract, p, q)
    return inverse_retract_nlsolve!(M, X, p, q, m)
end
"""
    inverse_retract_softmax(M::AbstractManifold, p, q)

computes the allocating variant of the [`SoftmaxInverseRetraction`](@ref),
which by default allocates and calls [`inverse_retract_softmax!`](@ref ManifoldsBase.inverse_retract_softmax!).
"""
function inverse_retract_softmax(M::AbstractManifold, p, q)
    X = allocate_result(M, inverse_retract, p, q)
    return inverse_retract_softmax!(M, X, p, q)
end

"""
    retract!(M::AbstractManifold, q, p, X)
    retract!(M::AbstractManifold, q, p, X, t::Real=1)
    retract!(M::AbstractManifold, q, p, X, method::AbstractRetractionMethod)
    retract!(M::AbstractManifold, q, p, X, t::Real=1, method::AbstractRetractionMethod)

Compute a retraction, a cheaper, approximate version of the [`exp`](@ref)onential map,
from `p` into direction `X`, scaled by `t`, on the [`AbstractManifold`](@ref) manifold `M`.
Result is saved to `q`.

Retraction method can be specified by the last argument, defaulting to
[`default_retraction_method`](@ref)`(M)`. See the documentation of respective manifolds for available
methods.

See [`retract`](@ref) for more details.
"""
function retract!(
    M::AbstractManifold,
    q,
    p,
    X,
    m::AbstractRetractionMethod = default_retraction_method(M),
)
    return _retract!(M, q, p, X, m)
end
function retract!(
    M::AbstractManifold,
    q,
    p,
    X,
    t::Real,
    method::AbstractRetractionMethod = default_retraction_method(M),
)
    return retract!(M, q, p, t * X, method)
end
# dispatch to lower level
_retract!(M::AbstractManifold, q, p, X, ::CayleyRetraction) = retract_caley!(M, q, p, X)
function _retract!(M::AbstractManifold, q, p, X, m::EmbeddedRetraction)
    return retract_embedded!(M, q, p, X, m.retraction)
end
_retract!(M::AbstractManifold, q, p, X, ::ExponentialRetraction) = exp!(M, q, p, X)
function _retract!(M::AbstractManifold, q, p, X, m::ODEExponentialRetraction)
    return retract_exp_ode!(M, q, p, X, m.retraction, m.basis)
end
_retract!(M::AbstractManifold, q, p, X, ::PolarRetraction) = retract_polar!(M, q, p, X)
function _retract!(M::AbstractManifold, q, p, X, ::ProjectionRetraction)
    return retract_project!(M, q, p, X)
end
_retract!(M::AbstractManifold, q, p, X, ::QRRetraction) = retract_qr!(M, q, p, X)
_retract!(M::AbstractManifold, q, p, X, ::SoftmaxRetraction) = retract_softmax!(M, q, p, X)
function _retract!(M::AbstractManifold, q, p, X, ::PadeRetraction{n}) where {n}
    return retract_pade!(M, q, p, X, n)
end

"""
    retract_embedded!(M::AbstractManifold, X, p, q, m::AbstractRetractionMethod)

computes the mutating variant of the [`EmbeddedRetraction`](@ref) using
the [`AbstractRetractionMethod`](@ref) `m` in the embedding (see [`get_embedding`](@ref))
and projecting the result back.
"""
function retract_embedded!(M::AbstractManifold, q, p, X, m::AbstractRetractionMethod)
    return project!(
        M,
        q,
        retract(
            get_embedding(M),
            embed(get_embedding(M), p),
            embed(get_embedding(M), p, X),
            m,
        ),
    )
end

"""
    retract_caley!(M::AbstractManifold, X, p, q)

computes the mutating variant of the [`CayleyRetraction`](@ref),
which by default falls back to calling the first order [`PadeRetraction`](@ref).
"""
function retract_caley!(M::AbstractManifold, q, p, X)
    return retract_pade!(M, q, p, X, 1)
end

"""
    retract_exp_ode!(M::AbstractManifold, q, p, X, m::AbstractRetractionMethod, B::AbstractBasis)

computes the mutating variant of the [`ODEExponentialRetraction`](@ref)`(m, B)`.
"""
retract_exp_ode!(
    M::AbstractManifold,
    q,
    p,
    X,
    m::AbstractRetractionMethod,
    B::AbstractBasis,
)

function retract_exp_ode! end

"""
    retract_pade!(M::AbstractManifold, q, p, n)

computes the mutating variant of the [`PadeRetraction`](@ref)`(n)`.
"""
retract_pade!(M::AbstractManifold, q, p, n)

function retract_pade! end

"""
    retract_project!(M::AbstractManifold, q, p, X)

computes the mutating variant of the [`ProjectionRetraction`](@ref).
"""
retract_project!(M::AbstractManifold, q, p, X)

function retract_project! end

"""
    retract_polar!(M::AbstractManifold, q, p, X)

computes the mutating variant of the [`PolarRetraction`](@ref).
"""
retract_polar!(M::AbstractManifold, q, p, X)

function retract_polar! end

"""
    retract_qr!(M::AbstractManifold, q, p, X)

computes the mutating variant of the [`QRRetraction`](@ref).
"""
retract_qr!(M::AbstractManifold, q, p, X)

function retract_qr! end

"""
    retract_softmax!(M::AbstractManifold, q, p, X)

computes the mutating variant of the [`SoftmaxRetraction`](@ref).
"""
retract_softmax!(M::AbstractManifold, q, p, X)

function retract_softmax! end

@doc raw"""
    retract(M::AbstractManifold, p, X, method::AbstractRetractionMethod=default_retraction_method(M))
    retract(M::AbstractManifold, p, X, t::Real=1, method::AbstractRetractionMethod=default_retraction_method(M))

Compute a retraction, a cheaper, approximate version of the [`exp`](@ref)onential map,
from `p` into direction `X`, scaled by `t`, on the [`AbstractManifold`](@ref) `M`.

A retraction ``\operatorname{retr}_p: T_p\mathcal M → \mathcal M`` is a smooth map that fulfils

1. ``\operatorname{retr}_p(0) = p``
2. ``D\operatorname{retr}_p(0): T_p\mathcal M \to T_p\mathcal M`` is the identity map,
i.e. ``D\operatorname{retr}_p(0)[X]=X`` holds for all ``X\in T_p\mathcal M``,

where ``D\operatorname{retr}_p`` denotes the differential of the retraction

The retraction is called of second order if for all ``X`` the curves ``c(t) = R_p(tX)``
have a zero acceleration at ``t=0``, i.e. ``c''(0) = 0``.

Retraction method can be specified by the last argument, defaulting to
[`default_retraction_method`](@ref)`(M)`. For further available retractions see the documentation of respective manifolds.

Locally, the retraction is invertible. For the inverse operation, see [`inverse_retract`](@ref).
"""
function retract(
    M::AbstractManifold,
    p,
    X,
    m::AbstractRetractionMethod = default_retraction_method(M),
)
    return _retract(M, p, X, m)
end
function retract(
    M::AbstractManifold,
    p,
    X,
    t::Real,
    m::AbstractRetractionMethod = default_retraction_method(M),
)
    return retract(M, p, t * X, m)
end
function _retract(M::AbstractManifold, p, X, m::EmbeddedRetraction)
    return retract_embedded(M, p, X, m.retraction)
end
_retract(M::AbstractManifold, p, X, ::ExponentialRetraction) = exp(M, p, X)
function _retract(M::AbstractManifold, p, X, m::ODEExponentialRetraction)
    return retract_exp_ode(M, p, X, m.retraction, m.basis)
end
_retract(M::AbstractManifold, p, X, ::PolarRetraction) = retract_polar(M, p, X)
_retract(M::AbstractManifold, p, X, ::ProjectionRetraction) = retract_project(M, p, X)
_retract(M::AbstractManifold, p, X, ::QRRetraction) = retract_qr(M, p, X)
_retract(M::AbstractManifold, p, X, ::SoftmaxRetraction) = retract_softmax(M, p, X)
_retract(M::AbstractManifold, p, X, ::CayleyRetraction) = retract_caley(M, p, X)
function _retract(M::AbstractManifold, p, X, ::PadeRetraction{n}) where {n}
    return retract_pade(M, p, X, n)
end
"""
    retract_embedded(M::AbstractManifold, p, X, m::AbstractRetractionMethod)

computes the allocating variant of the [`EmbeddedRetraction`](@ref) using
the [`AbstractRetractionMethod`](@ref) `m` in the embedding (see [`get_embedding`](@ref))
and projecting the result back.
"""
function retract_embedded(M::AbstractManifold, p, X, m::AbstractRetractionMethod)
    return project(
        M,
        retract(
            get_embedding(M),
            embed(get_embedding(M), p),
            embed(get_embedding(M), p, X),
            m,
        ),
    )
end
"""
    retract_polar(M::AbstractManifold, p, q)

computes the allocating variant of the [`PolarRetraction`](@ref),
which by default allocates and calls [`retract_polar!`](@ref ManifoldsBase.retract_polar!).
"""
function retract_polar(M::AbstractManifold, p, X)
    q = allocate_result(M, retract, p, X)
    return retract_polar!(M, q, p, X)
end
"""
    retract_project(M::AbstractManifold, p, q)

computes the allocating variant of the [`ProjectionRetraction`](@ref),
which by default allocates and calls [`retract_project!`](@ref ManifoldsBase.retract_project!).
"""
function retract_project(M::AbstractManifold, p, X)
    q = allocate_result(M, retract, p, X)
    return retract_project!(M, q, p, X)
end
"""
    retract_qr(M::AbstractManifold, p, q)

computes the allocating variant of the [`QRRetraction`](@ref),
which by default allocates and calls [`retract_qr!`](@ref ManifoldsBase.retract_qr!).
"""
function retract_qr(M::AbstractManifold, p, X)
    q = allocate_result(M, retract, p, X)
    return retract_qr!(M, q, p, X)
end
"""
    retract_exp_ode(M::AbstractManifold, p, q, m::AbstractRetractionMethod, B::AbstractBasis)

computes the allocating variant of the [`ODEExponentialRetraction`](@ref)`(m,B)`,
which by default allocates and calls [`retract_exp_ode!`](@ref ManifoldsBase.retract_exp_ode!).
"""
function retract_exp_ode(
    M::AbstractManifold,
    p,
    X,
    m::AbstractRetractionMethod,
    B::AbstractBasis,
)
    q = allocate_result(M, retract, p, X)
    return retract_exp_ode!(M, q, p, X, m, B)
end
"""
    retract_softmax(M::AbstractManifold, p, q)

computes the allocating variant of the [`SoftmaxRetraction`](@ref),
which by default allocates and calls [`retract_softmax!`](@ref ManifoldsBase.retract_softmax!).
"""
function retract_softmax(M::AbstractManifold, p, X)
    q = allocate_result(M, retract, p, X)
    return retract_softmax!(M, q, p, X)
end

"""
    retract_caley(M::AbstractManifold, p, q)

computes the allocating variant of the [`CayleyRetraction`](@ref),
which by default allocates and calls [`retract_caley!`](@ref ManifoldsBase.retract_caley!).
"""
function retract_caley(M::AbstractManifold, p, X)
    q = allocate_result(M, retract, p, X)
    return retract_caley!(M, q, p, X)
end
"""
    retract_pade(M::AbstractManifold, p, q)

computes the allocating variant of the [`PadeRetraction`](@ref)`(n)`,
which by default allocates and calls [`retract_pade!`](@ref ManifoldsBase.retract_pade!).
"""
function retract_pade(M::AbstractManifold, p, X, n)
    q = allocate_result(M, retract, p, X)
    return retract_pade!(M, q, p, X, n)
end

Base.show(io::IO, ::CayleyRetraction) = print(io, "CayleyRetraction()")
Base.show(io::IO, ::PadeRetraction{m}) where {m} = print(io, "PadeRetraction($(m))")


"""
$(TYPEDEF)

Holds the three parameters of a Lorentzian peak, defined as A / sqrt((ω-ω0)² + Γ²).

# Fields
$(TYPEDFIELDS)

"""
struct Lorentzian
    ω0::Float64
    Γ::Float64
    A::Float64
    Lorentzian(;ω0::Float64, Γ::Float64) = new(ω0, Γ, 1) # default peak height is 1
    Lorentzian(ω0, Γ, A) = new(ω0, Γ, A)
 end

 """
 $(TYPEDEF)

 Holds a set of `Lorentzian` objects belonging to a variable.

 # Fields
 $(TYPEDFIELDS)

 # Constructor
 ```julia
 JacobianSpectrum(res::Result; index::Int, branch::Int)
 ```

 """
 mutable struct JacobianSpectrum
    peaks::Vector{Lorentzian}
 end

"""
$(TYPEDEF)

Holds the compiled response matrix of a system.

# Fields
$(TYPEDFIELDS)

"""
struct ResponseMatrix
    """The response matrix (compiled)."""
    matrix::Matrix{Function}
    """Any symbolic variables in `matrix` to be substituted at evaluation."""
    symbols::Vector{Num}
    """The frequencies of the harmonic variables underlying `matrix`. These are needed to transform
    the harmonic variables to the non-rotating frame."""
    ωs::Vector{Num}

    ResponseMatrix(matrix, symbols, ωs) = new(matrix, symbols, ωs)
end


include_dpo()

# Model
include_model("quadrotor2D")
model = additive_noise_model(model)

# Horizon
T = 51

# Time step
h = 0.1

# Bounds
ul, uu = control_bounds(model, T, 0.0, Inf)

# Initial and final states
x1_nom = [0.0; 1.0; 0.0; 0.0; 0.0; 0.0]
xT_nom = [1.0; 1.0; 0.0; 0.0; 0.0; 0.0]
xl, xu = state_bounds(model, T, x1 = x1_nom, xT = xT_nom)

# Objective
obj = quadratic_tracking_objective(
        [t < T ? Diagonal(ones(model.n)) : Diagonal(10.0 * ones(model.n)) for t = 1:T],
        [Diagonal(1.0e-1 * ones(model.m)) for t = 1:T-1],
        [xT_nom for t = 1:T], [zeros(model.m) for t = 1:T])

# Problem
prob = trajectory_optimization_problem(model,
           obj,
           T,
           h = h,
           xl = xl,
           xu = xu,
           ul = ul,
           uu = uu)

# Trajectory initialization
x0 = linear_interpolation(x1_nom, xT_nom, T) # linear interpolation on state
u0 = [0.1 * ones(model.m) for t = 1:T-1] # random controls

# Pack trajectories into vector
z0 = pack(x0, u0, prob)

# Solve for nominal trajectory
z̄, info = solve(prob, copy(z0))
x̄, ū = unpack(z̄, prob)

using Plots
plot(hcat(x̄...)')
plot(hcat(ū...)', linetype = :steppost)

# DPO

# Linear model
A, B = jacobians(model, x̄, ū, h)

function fd(model::Quadrotor2D{Midpoint, FixedTime}, x⁺, x, u, w, h, t)
    x⁺ - A[t] * x - B[t] * u - w
end

# DPO
β = 1.0
δ = 1.0e-1

x1 = resample(zeros(model.n), Diagonal(ones(model.n)), β)

# Problems
prob_nom = trajectory_optimization(
			model,
			EmptyObjective(),
			T,
			xl = state_bounds(model, T, zeros(model.n), zeros(model.n))[1],
			xu = state_bounds(model, T, zeros(model.n), zeros(model.n))[2],
			ul = control_bounds(model, T, 0.0, 0.0)[1],
			uu = control_bounds(model, T, 0.0, 0.0)[2])

# Mean problem
prob_mean = trajectory_optimization(
				model,
				EmptyObjective(),
				T,
				dynamics = false)

# Sample problems
prob_sample = [trajectory_optimization(
				model,
				EmptyObjective(),
				T,
				dynamics = false,
				xl = state_bounds(model, T, x1 = x1[i])[1],
				xu = state_bounds(model, T, x1 = x1[i])[2]) for i = 1:2 * model.n]

# Sample objective
Q = [(t < T ? Diagonal(10.0 * ones(model.n))
 		: Diagonal(100.0 * ones(model.n))) for t = 1:T]
R = [Diagonal(ones(model.m)) for t = 1:T-1]

obj_sample = sample_objective(Q, R)
policy = linear_feedback(model.n, model.m)
dist = disturbances([Diagonal(δ * ones(model.d)) for t = 1:T-1])
sample = sample_params(β, T)

prob_dpo = dpo_problem(
	prob_nom, prob_mean, prob_sample,
	obj_sample,
	policy,
	dist,
	sample)

z0 = ones(prob_dpo.num_var)

# Solve
if true
	include_snopt()
	z, info = solve(prob_dpo, copy(z0),
		nlp = :SNOPT7,
		tol = 1.0e-7, c_tol = 1.0e-7,
		time_limit = 60 * 2)
	@save joinpath(@__DIR__, "sol_dpo.jld2") z
else
	println("Loading solution...")
    @load joinpath(@__DIR__, "sol_dpo.jld2") z
end

# DPO policy
θ = get_policy(z, prob_dpo)

# TVLQR policy
K, P = tvlqr(A, B, Q, R)

# Policy difference
policy_diff = [norm(vec(θ[t] - K[t])) / norm(vec(K[t])) for t = 1:T-1]
println("policy difference (inf. norm): $(norm(policy_diff, Inf))")


# Mean Field Aligned (MFA) coordinates transformation.

using EmpiricalModeDecomposition
using EmpiricalModeDecomposition: zerocrossing!
using Statistics: mean

export mfa, emd_average

"""
    mfa(x, y, z, Bx, By, Bz; method="MAVG", nw=15, fmax=-1.0, fsample=-1.0, threshold=0.1,
       verbose=false) -> Bμ, Bϕ, Bν

Obtain the time series of rotation matrix from geocentric solar ecliptic (GSE) to Mean
field-aligned (MFA) reference frame, given the positions of the spacecraft `x`, `y`, `z` and
magnetic fields `Bx`, `By`, `Bz` in GSE.

# Optional Arguments
- `method::String="MAVG"`: approach to obtain background ambient field, from moving average
`"MAVG"` with `nw` points, to empirical mode decomposition `"EMD"` with cutoff frequency
`fmax`, sampling frequency `fsample` and power selection `threshold` ∈ (0,1).
"""
function mfa(x::AbstractVector, y::AbstractVector, z::AbstractVector,
   Bx::AbstractVector, By::AbstractVector, Bz::AbstractVector;
   method="MAVG", nw=15, fsample=-1.0, fmax=fsample/50 , threshold=0.1, verbose=false)

   if !(length(x) == length(y) == length(z) ==
      length(Bx) == length(By) == length(Bz))
      throw(DimensionMismatch("requires same vector length for inputs!"))
   end

   μx, μy, μz = similar(x), similar(x), similar(x)
   ϕx, ϕy, ϕz = similar(x), similar(x), similar(x)
   νx, νy, νz = similar(x), similar(x), similar(x)
   Bμ, Bϕ, Bν = similar(x), similar(x), similar(x)

   if method == "MAVG"
     B̄x = sma(Bx, nw)
     B̄y = sma(By, nw)
     B̄z = sma(Bz, nw)
   elseif method == "EMD"
     B̄x = emd_average(Bx, fmax, fsample, threshold; verbose)
     B̄y = emd_average(By, fmax, fsample, threshold; verbose)
     B̄z = emd_average(Bz, fmax, fsample, threshold; verbose)
   end

   for i in axes(x, 1)
      B̄mag = hypot(B̄x[i], B̄y[i], B̄z[i])
      μx[i] = B̄x[i] / B̄mag
      μy[i] = B̄y[i] / B̄mag
      μz[i] = B̄z[i] / B̄mag
      ϕx[i] = y[i]*μz[i] - z[i]*μy[i]
      ϕy[i] = z[i]*μx[i] - x[i]*μz[i]
      ϕz[i] = x[i]*μy[i] - y[i]*μx[i]
      rBmag = hypot(ϕx[i], ϕy[i], ϕz[i])
      ϕx[i] /= rBmag
      ϕy[i] /= rBmag
      ϕz[i] /= rBmag
      νx[i] = μy[i]*ϕz[i] - μz[i]*ϕy[i]
      νy[i] = μz[i]*ϕx[i] - μx[i]*ϕz[i]
      νz[i] = μx[i]*ϕy[i] - μy[i]*ϕx[i]

      Bμ[i] = μx[i]*Bx[i] + μy[i]*By[i] + μz[i]*Bz[i]
      Bϕ[i] = ϕx[i]*Bx[i] + ϕy[i]*By[i] + ϕz[i]*Bz[i]
      Bν[i] = νx[i]*Bx[i] + νy[i]*By[i] + νz[i]*Bz[i]
   end

   if verbose
      @info "method = $method"
   end

   Bμ, Bϕ, Bν
end

"""
    emd_average(y, fmax, fsample, threshold=0.1; verbose=false)

Extracts low frequency component from a signal using Empirical Mode Decomposition (EMD).
The signal is reconstructed by using Intrinsic Mode Functions (IMF) which have an average
frequency lower than a specified cut-off frequency `fmax` and a percentage of the power of
the signal, attributable to frequencies lower than `fmax` and greater than `thresold`.

Reference:
The use of the empirical mode decomposition for the identification of mean field aligned
reference frames, M Regi, Alfredo Del Corpo, Marcello De Lauretis, 2016
https://doi.org/10.4401/ag-7067
"""
function emd_average(y, fmax, fsample, threshold=0.1; verbose=false)

   len = length(y)
   x = range(1, length=len, step=1/fsample)

   imfs = emd(y, x)  # intrinsic mode functions

   if length(imfs) == 1
      verbose && @warn "No oscillation detected!"
      return imfs[1]
   end

   res = imfs[end]
   imf = imfs[1:end-1]

   ## Zero-crossing
   # maximum number of zero-crossing is the number of zero_crossing of a sin wave of
   # frequency fmax in a sampling time interval T0 = len/fsample
   nmax = floor(Int, length(y)*(2*fmax)/fsample+1)

   # number of zeros in each IMF
   crossings = [Int[] for _ = 1:length(imfs)-1] # initialization

   for i in 1:length(crossings)
      zerocrossing!(imf[i], crossings[i])
   end

   nzeros = [length(crossings[i]) for i in 1:length(crossings)]

   ## Average period and signal power
   tmin = floor(Int, fsample/fmax)    # miminum period [# of samples]
   tminhalf = floor(Int, fsample/(2*fmax)) # miminum half period [# of samples]
   #T = zeros(length(crossings))
   ind1_ = Vector{Bool}(undef, length(crossings))

   for i = 1:length(imf)
      # indices of zero crossings plus first and last entries
      ind_ = [1, crossings[i]..., len]
      # lengths of consecutive zero-crossing [# of samples]
      tp = diff(ind_)
      #T[i] = mean(2*tp) # average period of the mode [# of samples]

      p = zeros(length(tp))
      # Evaluating the IMF power in each segment
      for j = 1:length(tp)
         p[j] = sum( abs.( imf[i][ind_[j]:ind_[j+1]-1] ) ) / fsample
      end

      ps = p[tp .> tminhalf] # select powers linked to periods greater than tmin

      # select IMF if the percentage of the selected power is greater than the thresold
      ind1_[i] = sum(ps) / sum(p) > threshold

      verbose && @info "IMF $i, threshold = $threshold, power = $(sum(ps) / sum(p))"
   end

   ## Reconstructing ambient signal
   ind2_ = nzeros .< nmax
   imf_bk = imf[ind1_ .| ind2_] # number of IMFs to add to residue to obtain ambient signal

   y_bk = res + sum(imf_bk)

   return y_bk
end


# test extending type and that various generic methods are defined
# to fulfill the API of AbstractYieldCurve
@testset "generic methods and type extensions" begin
    struct MyYield <: Yields.AbstractYieldCurve
        rate
    end

    Yields.discount(c::MyYield,to) = exp(-c.rate * to)
    # Base.zero(c::MyYield,to) = Continuous(c.rate)

    my = MyYield(0.05)
    @test zero(my,1) ≈ Continuous(0.05)
    @test Yields.forward(my,1) ≈ Continuous(0.05)
    @test Yields.forward(my,1,2) ≈ Continuous(0.05)
    @test discount(my,1,2) ≈ exp(-0.05*1)
    @test accumulation(my,1,2) ≈ exp(0.05*1)
    @test accumulation(my,1) ≈ exp(0.05*1)
    @test Yields.__ratetype(my) == Yields.Rate{Float64,Continuous}
    @test Yields.CompoundingFrequency(my) == Continuous()

    @test Yields.par(my,1) |> Yields.rate > 0
end


# # Stokes flow

#md # ```@meta
#md # CurrentModule = ImmersedLayers
#md # ```

#=
Here, we'll demonstrate a solution of the Stokes equations,
with Dirichlet (i.e. no-slip) boundary conditions on a surface.
The purpose of this case is to demonstrate the use of the tools with
vector-valued data -- in this case, the fluid velocity field.

After taking the curl of the velocity-pressure form of the Stokes equations, the governing
equations for this problem can be written as

$$\mu \nabla^2 \omega - \nabla \times (\delta(\chi) \mathbf{\sigma}) = \nabla\times \nabla \cdot (\delta(\chi)\mathbf{\Sigma})$$
$$\delta^T(\chi) \mathbf{v} = \overline{\mathbf{v}}_b$$
$$\mathbf{v} = \nabla \phi + \nabla\times \psi$$

where $\psi$ and $\phi$ are the solutions of

$$\nabla^2 \psi = -\omega$$
$$\nabla^2 \phi = \delta(\chi) \mathbf{n}\cdot [\mathbf{v}_b]$$

and $\mathbf{\Sigma} = \mu ([\mathbf{v}_b]\mathbf{n} + \mathbf{n} [\mathbf{v}_b])$ is a surface viscous flux tensor;
and $[\mathbf{v}_b] = \mathbf{v}_b^+ - \mathbf{v}_b^-$ and $\overline{\mathbf{v}}_b = (\mathbf{v}_b^+ + \mathbf{v}_b^-)/2$ are
the jump and average of the surface velocities on either side of a surface.

We can discretize and combine these equations into a saddle-point form for $\psi$ and $\sigma$:

$$\begin{bmatrix}L^2 & C_s^T \\ C_s & 0 \end{bmatrix}\begin{pmatrix} \psi \\ \mathbf{\sigma} \end{pmatrix} = \begin{pmatrix} -C D_s [\mathbf{v}_b]  \\ \overline{v}_b - G_s L^{-1} R \mathbf{n}\cdot [\mathbf{v}_b]  \end{pmatrix}$$

Note that $L^2$ represents the discrete biharmonic operator. (We have absorbed the
viscosity $\mu$ into the scaling of the problem.)

We can easily break this down into algorithmic form by the same block-LU decomposition
of other problems. The difference here is that $\mathbf{\sigma}$ is vector-valued.
=#

using ImmersedLayers
using Plots
using LinearAlgebra
using UnPack

#=
## Set up the extra cache and solve function
=#
#=
The problem type is generated with the usual macro call, but now with `vector` type
=#
@ilmproblem StokesFlow vector

#=
As with other cases, the extra cache holds additional intermediate data, as well as
the Schur complement. We need a few more intermediate variable for this problem.
We'll also construct the filtering matrix for showing the traction field.
=#
struct StokesFlowCache{SMT,CMT,RCT,DVT,VNT,ST,VFT,FT} <: AbstractExtraILMCache
   S :: SMT
   C :: CMT
   Rc :: RCT
   dv :: DVT
   vb :: DVT
   vprime :: DVT
   dvn :: VNT
   sstar :: ST
   vϕ :: VFT
   ϕ :: FT
end

#=
Extend the `prob_cache` function to construct the extra cache
=#
function ImmersedLayers.prob_cache(prob::StokesFlowProblem,base_cache::BasicILMCache)
    S = create_CL2invCT(base_cache)
    C = create_surface_filter(base_cache)

    dv = zeros_surface(base_cache)
    vb = zeros_surface(base_cache)
    vprime = zeros_surface(base_cache)

    dvn = ScalarData(dv)
    sstar = zeros_gridcurl(base_cache)
    vϕ = zeros_grid(base_cache)
    ϕ = Nodes(Primal,sstar)

    Rc = RegularizationMatrix(base_cache,dvn,ϕ)

    StokesFlowCache(S,C,Rc,dv,vb,vprime,dvn,sstar,vϕ,ϕ)
end

#=
And finally, extend the `solve` function to do the actual solving.
This form takes as input the the external and internal velocities on the
surface and returns the velocity field, streamfunction field, and surface
traction
=#
function ImmersedLayers.solve(vsplus,vsminus,prob::StokesFlowProblem,sys::ILMSystem)
    @unpack extra_cache, base_cache = sys
    @unpack nrm = base_cache
    @unpack S, C, Rc, dv, vb, vprime, sstar, dvn, vϕ, ϕ  = extra_cache

    σ = zeros_surface(sys)
    s = zeros_gridcurl(sys)
    v = zeros_grid(sys)

    dv .= vsplus-vsminus
    vb .= 0.5*(vsplus + vsminus)

    ## Compute ψ*
    surface_divergence!(v,dv,sys)
    curl!(sstar,v,sys)
    sstar .*= -1.0

    inverse_laplacian!(sstar,sys)
    inverse_laplacian!(sstar,sys)

    ## Adjustment for jump in normal velocity
    pointwise_dot!(dvn,nrm,dv)
    regularize!(ϕ,dvn,Rc)
    inverse_laplacian!(ϕ,sys)
    grad!(vϕ,ϕ,sys)
    interpolate!(vprime,vϕ,sys)
    vprime .= vb - vprime

    ## Compute surface velocity due to ψ*
    curl!(v,sstar,sys)
    interpolate!(vb,v,sys)

    ## Spurious slip
    vprime .-= vb
    σ .= S\vprime

    ## Correction streamfunction
    surface_curl!(s,σ,sys)
    s .*= -1.0
    inverse_laplacian!(s,sys)
    inverse_laplacian!(s,sys)

    ## Correct
    s .+= sstar

    ## Assemble the velocity
    curl!(v,s,sys)
    v .+= vϕ

    ## Filter the traction twice to clean it up a bit
    σ .= C^2*σ

    return v, s, σ
end

#=
## Set up a grid and body
We'll consider a rectangle here
=#
Δx = 0.01
Lx = 4.0
xlim = (-Lx/2,Lx/2)
ylim = (-Lx/2,Lx/2)
g = PhysicalGrid(xlim,ylim,Δx)
Δs = 1.4*cellsize(g)
body = Rectangle(0.5,0.25,Δs,shifted=true)

#=
Set up the problem and the system
=#
prob = StokesFlowProblem(g,body,scaling=GridScaling)
sys = ImmersedLayers.__init(prob)

#=
## Solve the problem
We will consider the flow generated by two basic motions of the body.
In the first, we will simply translate the body to the right at unit velocity.
We set the external $x$ velocity to 1 and $y$ velocity to zero. The
internal velocity we set to zero.
=#
vsplus = zeros_surface(sys)
vsminus = zeros_surface(sys)
vsplus.u .= 1.0
vsplus.v .= 0.0
nothing #hide

#=
Now solve it (and time it)
=#
solve(vsplus,vsminus,prob,sys) #hide
@time v, s, σ = solve(vsplus,vsminus,prob,sys)
nothing #hide

#=
Let's look at the velocity components
=#
plot(v,sys)

#=
Note that the velocity is zero inside, as desired. Let's look at the streamlines here
=#
plot(s,sys)

#=
We'll plot the surface traction components, too.
=#
plot(σ.u,label="σx")
plot!(σ.v,label="σy")

#=
Now, let's apply a different motion, where we rotate it counter-clockwise
at unit angular velocity
=#
pts = points(sys)
vsplus.u .= -pts.v
vsplus.v .= pts.u
nothing #hide

#=
Solve it again and plot the velocity and streamlines
=#
v, s, σ = solve(vsplus,vsminus,prob,sys)
plot(v,sys)

#-
plot(s,sys)


R=8.31451;p=101325;T=293.15

molemass=[
    16.0430,30.0700,44.0970,58.1230,58.1230,72.1500,72.1500,72.1500,
    86.1770,86.1770,86.1770,86.1770,86.1770,100.204,114.231,128.258,
    142.285,28.0540,42.0810,56.1080,56.0180,56.1080,56.1080,70.1340,
    40.0650,54.0920,54.0920,26.0380,70.1340,84.1610,98.1880,84.1610,
    98.1880,112.215,78.1140,92.1410,106.167,106.167,32.0420,48.1090,
    2.01590,18.0153,34.0820,17.0306,27.0260,28.0100,60.0760,76.1430,
    4.00260,20.1797,39.9480,28.0135,31.9988,44.0100,64.0650,44.0129,
    83.8000,131.290]
mol_heat_high=[
    891.090,1561.41,2220.13,2878.57,2869.38,3537.17,3530.24,3516.01,
    4196.58,4188.95,4191.54,4179.15,4186.93,4855.29,5513.88,6173.46,
    6832.31,1411.65,2058.72,2717.75,2711.00,2707.40,2701.10,3376.57,
    1943.53,2594.45,2541.43,1301.21,3320.88,3970.93,4630.19,3954.47,
    4602.35,5264.98,3302.15,3948.84,4608.32,4597.46,764.590,1239.83,
    285.990,44.2240,562.190,383.160,671.600,282.950,548.190,1104.41]
mol_heat_low=[
    802.650,1428.74,2043.23,2657.45,2648.26,3271.83,3264.89,3250.67,
    3887.01,3879.38,3881.97,3869.59,3877.36,4501.49,5115.87,5731.22,
    6345.85,1323.20,1926.05,2540.86,2534.10,2530.50,2524.20,3155.45,
    1855.08,2461.78,2408.76,1256.98,3099.76,3705.59,4320.63,3689.13,
    4292.78,4911.19,3169.48,3771.95,4387.20,4376.34,676.140,1151.39,
    241.760,0.00000,517.970,316.820,649.500,282.950,548.190,1104.41]
z=[
    0.9981,0.9920,0.9834,0.9682,0.9710,0.9450,0.953,0.959,
    0.9190,0.9260,0.9280,0.9350,0.9340,0.8760,0.817,0.735,
    0.6230,0.9940,0.9850,0.9720,0.9690,0.9690,0.972,0.952,
    0.9840,0.9650,0.9730,0.9930,0.9500,0.9270,0.885,0.924,
    0.8940,0.8380,0.9360,0.8920,0.8370,0.8210,0.892,0.978,
    1.0006,0.9520,0.9900,0.9890,0.9200,0.9996,0.988,0.965,
    1.0005,1.0005,0.9993,0.9997,0.9993,0.9944,0.980]
b=[0.0436,0.0894,0.1288,0.1783,0.1703,0.2345,0.2168,0.2025,
0.2846,0.2720,0.2683,0.2550,0.2569,0.3521,0.4278,0.5148,
0.6140,0.0775,0.1225,0.1673,0.1761,0.1761,0.1673,0.2191,
0.1265,0.1871,0.1643,0.0837,0.2236,0.2702,0.3391,0.2757,
0.3256,0.4025,0.2530,0.3286,0.4037,0.4231,0.3286,0.1483,
-0.0051,0.2191,0.1000,0.1049,0.2828,0.0200,0.1095,0.1871,
0.0000,0.0000,0.0265,0.0173,0.0265,0.0728,0.1414]
N=48;N1=55
hhv_mole(x)=sum(mol_heat_high[1:N].*x[1:N])
lhv_mole(x)=sum(mol_heat_low[1:N].*x[1:N])
mole_mass(x)=sum(molemass.*x)
hhv_mass(x)=hhv_mole(x)/mole_mass(x)
lhv_mass(x)=lhv_mole(x)/mole_mass(x)
vol2mol(x)=x[1:N1]./z/sum(x[1:N1]./z)
zmix(x)=1.0-sum(x[1:N1].*b)
lhv_vol(x)=lhv_mole(x)*p/R/T
hhv_vol(x)=hhv_mole(x)*p/R/T
lhv_real_vol(x)=lhv_vol(x)/zmix(x)
hhv_real_vol(x)=hhv_vol(x)/zmix(x)
rel_ρ(x)=sum(x*mole_mass(x)/28.9625)
abs_ρ(x)=sum(x*mole_mass(x))*p/R/T
real_rel_ρ(x)=rel_ρ(x)*0.99963/zmix(x)
real_abs_ρ(x)=abs_ρ(x)/zmix(x)

function calorific(x,out)
    if out==:mol
        return (lhv_mole(x),hhv_mole(x))
    elseif out==:mass
        return (lhv_mass(x),hhv_mass(x))
    elseif out==:vol
        return (lhv_real_vol(x),hhv_real_vol(x))
    elseif out==:rho
        return (real_rel_ρ(x),real_abs_ρ(x))
    end
end

"""
    ng_calorific(mode=:mol,outoption=:mol;CH4=100.0,C2H6=0,C3H8=0,C4H10=0,C4H10_1=0,C5H12=0,C5H12_1=0.,C5H12_2=0,C6H14=0,
    C6H14_1=0.0,C6H14_2=0.0,C6H14_3=0,C6H14_4=0,C7H16=0,C8H18=0.0,C9H20=0.0,C10H22=0.0,C2H4=0.0,
    C3H6=0.0,C4H8=0.0,C4H8_1=0.0,C4H8_2=0.0,C4H8_3=0.0,C5H10=0.0,C3H4=0.0,C4H6=0.0,C4H6_1=0.0,C2H2=0.0,
    C5H10_1=0.0,C6H12=0.0,C7H12=0.0,C6H12_1=0.0,C7H14=0.0,C8H16=0.0,C6H6=0.0,C7H8=0.0,C8H10=0.0,C8H10_1=0.0,
    CH4O=0.0,C2H6HgS2=0.0,H2=0.0,H2O=0.0,H2S=0.0,NH3=0.0,CHN=0.0,CO2=0.0,COS=0.0,CS2=0.0,He=0.0,Ne=0.0,Ar=0.0,
    N2=0.0,O2=0.0,CO2=0.0,O2S=0.0,N2O=0.0,Kr=0.0,Xe=0.0)
# Augument
- `mode`: `:mol` mol heat;`:mass` mass heat;`:vol heat` volume heat
- `outoption` : `:mol` (mole_low_heat,mole_high_heat) kJ/mo1;`:mass` (mass_low_heat,mass_high_heat) MJ/kg;`:vol`(vol_low_heat,vol_high_heat) u"MJ/m3";
                `:rho` (real rel_density,real abs_density u"kg/m3")
"""
function ng_calorific(mode=:mol,outoption=:mol;CH4=1.0,C2H6=0,C3H8=0,nC4H10=0,iC4H10=0,nC5H12=0,iC5H12=0.,
    C5H12_2=0,nC6H14=0,iC6H14=0.0,C6H14_2=0.0,C6H14_3=0,C6H14_4=0,nC7H16=0,nC8H18=0.0,nC9H20=0.0,nC10H22=0.0,
    C2H4=0.0,C3H6=0.0,C4H8=0.0,C4H8_1=0.0,C4H8_2=0.0,C4H8_3=0.0,C5H10=0.0,C3H4=0.0,C4H6=0.0,C4H6_1=0.0,C2H2=0.0,
    C5H10_1=0.0,C6H12=0.0,C7H12=0.0,C6H12_1=0.0,C7H14=0.0,C8H16=0.0,C6H6=0.0,C7H8=0.0,C8H10=0.0,C8H10_1=0.0,
    CH4O=0.0,C2H6HgS2=0.0,H2=0.0,H2O=0.0,H2S=0.0,NH3=0.0,CHN=0.0,CO=0.0,COS=0.0,CS2=0.0,He=0.0,Ne=0.0,Ar=0.0,
    N2=0.0,O2=0.0,CO2=0.0,O2S=0.0,N2O=0.0,Kr=0.0,Xe=0.0)
    x=[CH4,C2H6,C3H8,nC4H10,iC4H10,nC5H12,iC5H12,C5H12_2,nC6H14,
    iC6H14,C6H14_2,C6H14_3,C6H14_4,nC7H16,nC8H18,nC9H20,nC10H22,C2H4,
    C3H6,C4H8,C4H8_1,C4H8_2,C4H8_3,C5H10,C3H4,C4H6,C4H6_1,C2H2,
    C5H10_1,C6H12,C7H12,C6H12_1,C7H14,C8H16,C6H6,C7H8,C8H10,C8H10_1,
    CH4O,C2H6HgS2,H2,H2O,H2S,NH3,CHN,CO,COS,CS2,He,Ne,Ar,
    N2,O2,CO2,O2S,N2O,Kr,Xe]
    if mode==:vol;x=vol2mol(x);end
    calorific(x,outoption)

end

"""
    ng_calorific_cn(mode=:mol,outoption=:mol;甲烷=100,乙烷=0,丙烷=0,丁烷=0,二甲基丙烷=0,戊烷=0,二甲基丁烷=0,
    二二甲基丙烷=0,已烷=0,二甲基戊烷=0,三甲基戊烷=0,二二甲基丁烷=0,二三甲基丁烷=0,庚烷=0,辛烷=0,
    壬烷=0,癸烷=0,乙烯=0,丙烯=0,一丁烯=0,顺二丁烯=0,反二丁烯=0,二甲基丙烯=0,一戊烯=0,丙二烯=0,一二丁二烯=0,
    一三丁二烯=0,乙炔=0,环戊烷=0,甲基环戊烷=0,乙基环戊烷=0,环己烷=0,甲基环己烷=0,乙基环己烷=0,苯=0,
    甲苯=0,乙苯=0,邻二甲苯=0,甲醇=0,甲硫醇=0,氢气=0,水=0,硫化氢=0,氨=0,氰化氢=0,一氧化碳=0,硫氧碳=0,
    二硫化碳=0,氦气=0,氖气=0,氩气=0,氮气=0,氧气=0,二氧化碳=0,二氧化硫=0,一氧化二氮=0,氪气=0,氙气=0)
# Augument
- `mode`: `:mol` mol heat;`:mass` mass heat;`:vol heat` volume heat
- `outoption` : `:mol` (mole_low_heat,mole_high_heat) kJ/mo1;`:mass` (mass_low_heat,mass_high_heat) MJ/kg;`:vol`(vol_low_heat,vol_high_heat) u"MJ/m3";
                `:rho` (real rel_density,real abs_density u"kg/m3")
"""
function ng_calorific_cn(mode=:mol,outoption=:mol;甲烷=1,乙烷=0,丙烷=0,丁烷=0,二甲基丙烷=0,戊烷=0,二甲基丁烷=0,
    二二甲基丙烷=0,已烷=0,二甲基戊烷=0,三甲基戊烷=0,二二甲基丁烷=0,二三甲基丁烷=0,庚烷=0,辛烷=0,
    壬烷=0,癸烷=0,乙烯=0,丙烯=0,一丁烯=0,顺二丁烯=0,反二丁烯=0,二甲基丙烯=0,一戊烯=0,丙二烯=0,一二丁二烯=0,
    一三丁二烯=0,乙炔=0,环戊烷=0,甲基环戊烷=0,乙基环戊烷=0,环己烷=0,甲基环己烷=0,乙基环己烷=0,苯=0,
    甲苯=0,乙苯=0,邻二甲苯=0,甲醇=0,甲硫醇=0,氢气=0,水=0,硫化氢=0,氨=0,氰化氢=0,一氧化碳=0,硫氧碳=0,
    二硫化碳=0,氦气=0,氖气=0,氩气=0,氮气=0,氧气=0,二氧化碳=0,二氧化硫=0,一氧化二氮=0,氪气=0,氙气=0)

    x=[甲烷,乙烷,丙烷,丁烷,二甲基丙烷,戊烷,二甲基丁烷,
    二二甲基丙烷,已烷,二甲基戊烷,三甲基戊烷,二二甲基丁烷,二三甲基丁烷,庚烷,辛烷,
    壬烷,癸烷,乙烯,丙烯,一丁烯,顺二丁烯,反二丁烯,二甲基丙烯,一戊烯,丙二烯,一二丁二烯,
    一三丁二烯,乙炔,环戊烷,甲基环戊烷,乙基环戊烷,环己烷,甲基环己烷,乙基环己烷,苯,
    甲苯,乙苯,邻二甲苯,甲醇,甲硫醇,氢气,水,硫化氢,氨,氰化氢,一氧化碳,硫氧碳,
    二硫化碳,氦气,氖气,氩气,氮气,氧气,二氧化碳,二氧化硫,一氧化二氮,氪气,氙气]
    if mode==:vol;x=vol2mol(x);end
    calorific(x,outoption)
end

"""
    ng_calorific_en(mode=:mol,outoption=:mol;Methane=100.0,Ethane=0.0,Propane=0.0,nButane=0.0,Alkanes=0.0,nPentane=0.0,
    Methylbutane_2=0.0,Dimethylpropane_2_2=0.0,nHexane=0.0,Methylpentane_2=0.0,Methoylpentane_3=0.0,
    Dimethylbutane_2_2=0.0,Dimethylbutane_2_3=0.0,nHeptane=0.0,nOctane=0.0,nNonane=0.0,nDecane=0.0,
    Ethylene=0.0,Propylene=0.0,    Butene_1=0.0,cis_2_Butene=0.0,trans_2_Butene=0.0,Methylpropene_2=0.0,
    Pentene_1=0.0,    Propadiene=0.0,Butadiene_1_2=0.0,Butadiene_1_3=0.0,Acetylene=0.0,Cyclopentane=0.0,
    Methylcyclopentane=0.0,Ethylcyclopentane=0.0,Cyclohexane=0.0,Methylcyclohexane=0.0,Ethylcyclohexane=0.0,
    Benzene=0.0,Toluene=0.0,Ethylbenzene=0.0,o_Xylene=0.0,Methanol=0.0,Methanethiol=0.0,Hydrogen=0.0,
    Water=0.0,Hydrogen_sulfide=0.0,Ammonia=0.0,Hydrogen_cyanide=0.0,Carbon_monoxide=0.0,Carbonyl_sulfide=0.0,
    Carbon_disulfide=0.0,Helium=0.0,Neon=0.0,Argon=0.0,Nitrogen=0.0,Oxygen=0.0,Carbon_dioxide=0.0,Sulfur_Dioxide=0.0,
    Dinitrogen_monoxide=0.0,Krypton=0.0,Xenon=0.0)
# Augument
- `mode`: `:mol` mol heat;`:mass` mass heat;`:vol heat` volume heat
- `outoption` : `:mol` (mole_low_heat,mole_high_heat) kJ/mo1;`:mass` (mass_low_heat,mass_high_heat) MJ/kg;`:vol`(vol_low_heat,vol_high_heat) u"MJ/m3";
                `:rho` (real rel_density,real abs_density u"kg/m3")
"""
function ng_calorific_en(mode=:mol,outoption=:mol;Methane=1.0,Ethane=0.0,Propane=0.0,nButane=0.0,Alkanes=0.0,nPentane=0.0,
    Methylbutane_2=0.0,Dimethylpropane_2_2=0.0,nHexane=0.0,Methylpentane_2=0.0,Methoylpentane_3=0.0,
    Dimethylbutane_2_2=0.0,Dimethylbutane_2_3=0.0,nHeptane=0.0,nOctane=0.0,nNonane=0.0,nDecane=0.0,
    Ethylene=0.0,Propylene=0.0,    Butene_1=0.0,cis_2_Butene=0.0,trans_2_Butene=0.0,Methylpropene_2=0.0,
    Pentene_1=0.0,    Propadiene=0.0,Butadiene_1_2=0.0,Butadiene_1_3=0.0,Acetylene=0.0,Cyclopentane=0.0,
    Methylcyclopentane=0.0,Ethylcyclopentane=0.0,Cyclohexane=0.0,Methylcyclohexane=0.0,Ethylcyclohexane=0.0,
    Benzene=0.0,Toluene=0.0,Ethylbenzene=0.0,o_Xylene=0.0,Methanol=0.0,Methanethiol=0.0,Hydrogen=0.0,
    Water=0.0,Hydrogen_sulfide=0.0,Ammonia=0.0,Hydrogen_cyanide=0.0,Carbon_monoxide=0.0,Carbonyl_sulfide=0.0,
    Carbon_disulfide=0.0,Helium=0.0,Neon=0.0,Argon=0.0,Nitrogen=0.0,Oxygen=0.0,Carbon_dioxide=0.0,Sulfur_Dioxide=0.0,
    Dinitrogen_monoxide=0.0,Krypton=0.0,Xenon=0.0)

    x=[Methane,Ethane,Propane,nButane,Alkanes,nPentane,Methylbutane_2,Dimethylpropane_2_2,nHexane,Methylpentane_2,Methoylpentane_3,
    Dimethylbutane_2_2,Dimethylbutane_2_3,nHeptane,nOctane,nNonane,nDecane,Ethylene,Propylene,Butene_1,cis_2_Butene,trans_2_Butene,
    Methylpropene_2,Pentene_1,Propadiene,Butadiene_1_2,Butadiene_1_3,Acetylene,Cyclopentane,Methylcyclopentane,Ethylcyclopentane,
    Cyclohexane,Methylcyclohexane,Ethylcyclohexane,Benzene,Toluene,Ethylbenzene,o_Xylene,Methanol,Methanethiol,Hydrogen,Water,
    Hydrogen_sulfide,Ammonia,Hydrogen_cyanide,Carbon_monoxide,Carbonyl_sulfide,Carbon_disulfide,Helium,Neon,Argon,Nitrogen,Oxygen,
    Carbon_dioxide,Sulfur_Dioxide,Dinitrogen_monoxide,Krypton,Xenon]
    if mode==:vol;x=vol2mol(x);end
    calorific(x,outoption)
end

export ng_calorific,ng_calorific_cn,ng_calorific_en


using TravelerProblem;
using Test;

@testset "TravelerProblem.jl" begin

        @testset "LP vs Bruteforce" begin
                #Compiling
                tm = TravelerProblem.Test.trcmatrix(4)[1];
                TravelerProblem.solve_with_linear_programming(tm);
                TravelerProblem.solve_with_very_brute_force(tm);
                println();
                printstyled("Testing LP vs Bruteforce",color = :green, bold = true);
                for i in 5:8
                        cm,rr = TravelerProblem.Test.trcmatrix(i);
                        println();
                        printstyled("Matrix dimension ($i,$i)\n", color = :blue)
                        println("Solving with linear programming:")
                        @time lpr = TravelerProblem.solve_with_linear_programming(cm)
                        println("Solving with very brute force:")
                        @time vbf = TravelerProblem.solve_with_very_brute_force(cm);
                        @test lpr.rute == rr || lpr.rute == TravelerProblem.reverserute(rr);
                        @test vbf.rute == rr || vbf.rute == TravelerProblem.reverserute(rr);
                end
                println();
        end

        @testset "LP with bigger matrices" begin

                println();
                printstyled("Testing LP for bigger matrices!!!",color = :green, bold = true);
                for i in 10:10:150
                        cm,rr = TravelerProblem.Test.trcmatrix(i, Int32);
                        println();
                        printstyled("Matrix dimension ($i,$i)\n", color = :blue)
                        println("Solving with linear programming:")
                        @time lpr = TravelerProblem.solve_with_linear_programming(cm)
                        @test lpr.rute == rr || lpr.rute == TravelerProblem.reverserute(rr);
                end
                println();

        end

end


"""
    rmvn(n[, μ], Σ)

Fast parrallel generation of multivariate normal samples.
"""
function rmvn end
for T in (Float64, Float32, Float16)
    @eval function rmvn(n::Int, μ::Vector{$T}, Σ::Matrix{$T})
        μ' .+ _rmvn(n, Σ)
    end
    @eval rmvn(n::Real, μ::Vector{$T}, Σ::Matrix{$T}) = rmvn(Int(n), μ, Σ)
    @eval function rmvn(n::Real, Σ::Matrix{$T})
        d = size(Σ, 2)
        rmvn(Int(n), zeros($T, d), Σ)
    end
end
function rmvn(n::Real, μ::Vector{<:Real}, Σ::Matrix{<:Real})
    rmvn(Int(n), Vector{Float64}(μ), Matrix{Float64}(Σ))
end


# This file is part of GenericSchur.jl, released under the MIT "Expat" license

struct Balancer{T}
    ilo :: Int
    ihi :: Int
    prow :: Vector{Int}
    pcol :: Vector{Int}
    D :: Vector{T}
    trivial :: Bool
end

# intended to be a translation of xGEBAL from LAPACK
# Copyright:
# Univ. of Tennessee
# Univ. of California Berkeley
# Univ. of Colorado Denver
# NAG Ltd.
# w/ some allowance for scheme from R.James et al. 2014 arxiv 1401.5766
"""
    balance!(A; scale=true, permute=true) => Abal, B::Balancer

Balance a matrix so that various operations are more stable.

If `permute`, then row and column permutations are found such
that `Abal` has the block form `[T₁ X Y; 0 C Z; 0 0 T₂]` where
`T₁` and `T₂` are upper-triangular.
If `scale`, then a diagonal similarity transform using powers of 2
is found such that the 1-norm of `Abal` is near unity. The transformations
are encoded into `B` so that they can be inverted for eigenvectors, etc.
Balancing typically improves the accuracy of eigen-analysis.
"""
function balance!(A::AbstractMatrix{T}; scale=true, permute=true,
                  p::Int=1, algo = :pr) where T
    n = checksquare(A)
    RT = real(T)
    β = RT(2)
    factor = RT(0.95)

    sfmin1 = floatmin(RT) / eps(RT)
    sfmax1 = one(RT) / sfmin1
    sfmin2 = sfmin1 * β
    sfmax2 = one(RT) / sfmin2

    D = ones(T,n)
    ilo = 1
    ihi = n
    trivial = true
    sp = fill(0,n)

    if permute
        # look for row permutations which make the lower portion UT
        ihi = n+1
        js = 0
        while ihi > 1
            ihi -= 1
            exch = false
            ms = 0
            for j=ihi:-1:1
                exch = true
                js = j
                for i=1:ihi
                    (i == j) && continue
                    if A[j,i] != 0
                        exch = false
                    end
                end
                if exch
                    ms = ihi
                    break
                end
            end
            if exch
                sp[ms] = js
                if js != ms
                    trivial = false
                    for i=1:ihi
                        A[i,js], A[i,ms] = A[i,ms], A[i,js]
                    end
                    for i=ilo:n
                        A[js,i], A[ms,i] = A[ms,i], A[js,i]
                    end
                end
            else
                break
            end
        end
        # look for column permutations which make the left portion UT
        ilo = 0
        while ilo < n
            ilo += 1
            exch = false
            ms = 0
            js = 0
            for j=ilo:ihi
                js = j
                exch = true
                for i=ilo:ihi
                    (i == j) && continue
                    if A[i,j] != 0
                        exch = false
                    end
                end
                if exch
                    ms = ilo
                    break
                end
            end
            if exch
                sp[ms] = js
                if ms != js
                    trivial = false
                    for i=1:ihi
                        A[i,js], A[i,ms] = A[i,ms], A[i,js]
                    end
                    for i=ilo:n
                        A[js,i], A[ms,i] = A[ms,i], A[js,i]
                    end
                end
            else
                break
            end
        end
    end

    sf(c,r) = (p == 1) ? (c + r) : c*c + r*r
    # abs(c) + abs(r), but args always nonneg

    if scale
        converged = false
        while !converged
            converged = true
            for i=ilo:ihi
                if algo == :pr
                    # LAPACK mixes norms
                    c = norm(view(A,ilo:ihi,i),2)
                    r = norm(view(A,i,ilo:ihi),2)
#                    c = cxnorm(A,ilo,ihi,i,p)
#                    r = rxnorm(A,ilo,ihi,i,p)
                else
                    c = norm(view(A,ilo:ihi,i),p)
                    r = norm(view(A,i,ilo:ihi),p)
                end
                ca,_ = _findamax(view(A,ilo:ihi,i))
                ra,_ = _findamax(view(A,i,ilo:ihi))
                if c == 0 || r == 0
                    continue
                end
                g = r / β
                s = sf(c,r)
                f = one(RT)
                while c < r/β
                    if c >= g || (max(f,c,ca) >= sfmax2) || (min(r,g,ra) <= sfmin2)
                        break
                    end
                    if isnan(c+f+ca+r+g+ra)
                        error("NaN encountered while balancing")
                    end
                    f *= β
                    c *= β
                    ca *= β
                    r /= β
                    g /= β
                    ra /= β
                end
                g = c / β
                while r <= c/β
                    if (g < r) || (max(r,ra) >= sfmax2) || (min(f,c,g,ca) <= sfmin2)
                        break
                    end
                    f /= β
                    c /= β
                    g /= β
                    ca /= β
                    r *= β
                    ra *= β
                end
                if f != one(RT)
                    trivial = false
                end
                (sf(c,r) >= factor * s) && continue
                # TODO: continue on [uo]flow
                converged = false
                D[i] *= f
                A[i,ilo:n] .*= (one(T) / f)
                A[1:ihi,i] .*= f
            end
        end # while !converged
    end # if scale
    prow = sp[1:ilo-1]
    pcol = sp[ihi+1:n]
    B = Balancer(ilo,ihi,prow,pcol,D,trivial)
    return A, B
end

# this is the appropriate transformation for right eigenvectors
# stored as columns of a matrix
function lmul!(B::Balancer, A::StridedMatrix{T}) where T
    n = checksquare(A)
    if B.trivial
        return A
    end
    ilo, ihi = B.ilo, B.ihi
    if ilo != ihi
        lmul!(Diagonal(B.D), A)
    end
    for j=ilo-1:-1:1
        m = B.prow[j]
        if m == j
            continue
        end
        for i=1:n
            A[j,i], A[m,i] = A[m,i], A[j,i]
        end
    end
    for j=ihi+1:n
        m = B.pcol[j-ihi]
        if m == j
            continue
        end
        for i=1:n
            A[j,i], A[m,i] = A[m,i], A[j,i]
        end
    end
    A
end

# this should be the appropriate transformation for left eigenvectors
# stored as columns of a matrix
function ldiv!(B::Balancer, A::StridedMatrix{T}) where T
    n = checksquare(A)
    if B.trivial
        return A
    end
    ilo, ihi = B.ilo, B.ihi
    if ilo != ihi
        rmul!(A, Diagonal(one(T) ./ B.D))
    end
    for j=ilo-1:-1:1
        m = B.prow[j]
        if m == j
            continue
        end
        for i=1:n
            A[j,i], A[m,i] = A[m,i], A[j,i]
        end
    end
    for j=ihi+1:n
        m = B.pcol[j-ihi]
        if m == j
            continue
        end
        for i=1:n
            A[j,i], A[m,i] = A[m,i], A[j,i]
        end
    end
    A
end

# James et al. claim Parlet & Reinsch advocate skipping diagonal elts
# but this is (no longer?) in LAPACK
function cxnorm(A::AbstractMatrix{T},i0,i1,i,p) where T
    RT = real(T)
        s = zero(RT)
        for j=i0:i1
            if j==i
                continue
            end
            if p == 1
                s += abs(A[j,i])
            else
                s += A[j,i]' * A[j,i]
            end
        end
        if p == 2
            s = sqrt(s)
        end
        return s
end
function rxnorm(A::AbstractMatrix{T},i0,i1,i,p) where T
    RT = real(T)
        s = zero(RT)
        for j=i0:i1
            if j==i
                continue
            end
            if p == 1
                s += abs(A[i,j])
            else
                s += A[i,j]' * A[i,j]
            end
        end
        if p == 2
            s = sqrt(s)
        end
        return s
end


module DomainIntegrals

using QuadGK, LinearAlgebra
using GaussQuadrature, FastGaussQuadrature, HCubature

using StaticArrays
using IntervalSets, DomainSets

import Base: convert, checkbounds, length

import DomainSets:
    EmptySpace,
    FullSpace,
    domaintype,
    codomaintype,
    prectype,
    numtype,
    domain

include("common.jl")
include("gauss.jl")
include("measure.jl")
include("singularity.jl")
include("strategy.jl")
include("rules.jl")
include("integral.jl")

end # module


#=
    IndexLike

There are two components to any index, the keys and the values. These differ
from dictionaries in that the values must always be a subtype of
`AbstractUnitRange`. This facilitates a number of convenient traits and
assumptions concerning and index's behavior. When this behavior is overly
restrictive it makes sense use an index as component of some other object.
For example, a dictionary could potentially utilize the keys of an index
to map to the location of the dictionaries values.

# Keys

The only restriction to a collection of keys is that they all be unique to the
other keys in the same collection. This behavior is similar to what's seen if
you use `Set` or the keys of most `AbstractDict` subtypes. Keys of an index
allow the same customization and optimization that is commonly seen throughout
Julia via variety of dynamic, static, or immutable subtypes of `AbstractVector`.
Therefore, the only unique quality of keys for an index is the dichotomy offered
by the `ContinuousTrait` or `DiscreteTrait`.

* Continuous keys : Continuous keys imply an order the inherit to it's type.
  Therefore, all keys that are continuous are assumed to be some sort of range.
  This has the potential to speed up many operations but limits the ability to
  use keys with semantic meaning.
* Discrete keys: All vectors that are not ranges are assumed to be discrete,
  offering increased flexibility at the cost of performance.

It's worth noting that given some customization may help convert a system that
uses discrete keys into one of continuous keys. For example, using the
`Unitful.jl` package can facilitate the use of semantic units of measurement
as a well defined range.

# Values
=#


"""
    HalfWing(foils      :: Vector{Foil},
             chords     :: Vector{Real},
             twists     :: Vector{Real},
             spans      :: Vector{Real},
             dihedrals  :: Vector{Real},
             sweeps     :: Vector{Real})

Definition for a `HalfWing` consisting of ``N+1`` airfoils and their associated chord lengths ``c``, twist angles ``\\iota``, for ``N`` sections with span lengths ``b``, dihedrals ``\\delta`` and sweep angles ``\\Lambda``, with all angles in degrees.
"""
struct HalfWing{T <: Real} <: Aircraft
    foils     :: Vector{Foil{T}}
    chords    :: Vector{T}
    twists    :: Vector{T}
    spans     :: Vector{T}
    dihedrals :: Vector{T}
    sweeps    :: Vector{T}
    function HalfWing(foils :: Vector{Foil{T}}, chords, twists, spans, dihedrals, sweeps) where T <: Real
        # Error handling
        check_wing(foils, chords, twists, spans, dihedrals, sweeps)
        # Convert angles to radians, with adjusting twists to leading edge, and generate HalfWing
        new{T}(foils, chords, -deg2rad.(twists), spans, deg2rad.(dihedrals), deg2rad.(sweeps))
    end
end

function check_wing(foils, chords, twists, spans, dihedrals, sweeps)
    # Check if number of sections match up with number of edges
    @assert (length ∘ zip)(foils, chords, twists) == (length ∘ zip)(spans, dihedrals, sweeps) + 1 "N+1 foil sections, chords and twists are required for N section(s)."
    # Check if lengths are positive
    @assert any(x -> x >= 0., chords) || any(x -> x >= 0., spans) "Chord and span lengths must be positive."
    # Check if dihedrals and sweeps are within bounds
    @assert any(x -> x >= -90. || x <= 90., dihedrals) || any(x -> x >= -90. || x <= 90., sweeps) "Dihedrals and sweep angles must not exceed ±90ᵒ."
end

# Named arguments version for ease, with default NACA-4 0012 airfoil shape
HalfWing(; chords :: Vector{T}, twists :: Vector{T}, spans :: Vector{T}, dihedrals :: Vector{T}, sweep_LEs :: Vector{T}, foils :: Vector{Foil{T}} = fill(Foil(naca4((0,0,1,2)), "NACA 0012"), length(chords))) where T <: Real = HalfWing(foils, chords, twists, spans, dihedrals, sweep_LEs)

HalfWingSection(; span = 1., dihedral = 0., sweep_LE = 0., taper = 1., root_chord = 1., root_twist = 0., tip_twist = 0., root_foil = naca4((0,0,1,2)), tip_foil = naca4((0,0,1,2))) = HalfWing([ Foil(root_foil, "Root"), Foil(tip_foil, "Tip") ], [root_chord, taper * root_chord], [root_twist, tip_twist], [span], [dihedral], [sweep_LE])

# Getters
foils(wing     :: HalfWing) = wing.foils
chords(wing    :: HalfWing) = wing.chords
twists(wing    :: HalfWing) = wing.twists
spans(wing     :: HalfWing) = wing.spans
dihedrals(wing :: HalfWing) = wing.dihedrals
sweeps(wing    :: HalfWing) = wing.sweeps

"""
    span(half_wing :: HalfWing)

Compute the planform span of a `HalfWing`.
"""
span(wing :: HalfWing) = sum(wing.spans)

taper_ratio(wing :: HalfWing) = last(wing.chords) / first(wing.chords)

"""
    projected_area(half_wing :: HalfWing)

Compute the projected area of a `HalfWing` onto the ``x``-``y`` plane.
"""
function projected_area(wing :: HalfWing)
    mean_chords = fwdsum(wing.chords) / 2 # Mean chord lengths of sections.
    mean_twists = fwdsum(wing.twists) / 2 # Mean twist angles of sections
    sum(@. wing.spans * mean_chords * cos(mean_twists))
end

section_macs(wing :: HalfWing) = @views mean_aerodynamic_chord.(wing.chords[1:end-1], fwddiv(wing.chords))

section_projected_areas(wing :: HalfWing) = wing.spans .* fwdsum(wing.chords) / 2

"""
    mean_aerodynamic_chord(half_wing :: HalfWing)

Compute the mean aerodynamic chord of a `HalfWing`.
"""
function mean_aerodynamic_chord(wing :: HalfWing)
    areas = section_projected_areas(wing)
    macs  = section_macs(wing)
    sum(macs .* areas) / sum(areas)
end

function mean_aerodynamic_center(wing :: HalfWing, factor = 0.25)
    # Compute mean aerodynamic chords and projected areas
    macs 		= section_macs(wing)
    areas 		= section_projected_areas(wing)

    # Computing leading edge coordinates
    wing_LE 	= leading_edge(wing)
    x_LEs 		= getindex.(wing_LE, 1)
    y_LEs		= getindex.(wing_LE, 2)

    # Compute x-y locations of MACs
    x_mac_LEs	= @views @. y_mac.(x_LEs[1:end-1], 2 * x_LEs[2:end], wing.chords[2:end] / wing.chords[1:end-1])
    y_macs		= @views @. y_mac.(y_LEs[1:end-1], wing.spans, wing.chords[2:end] / wing.chords[1:end-1])

    mac_coords 	= @. SVector(x_mac_LEs + factor * macs, y_macs, 0.)

    sum(mac_coords .* areas) / sum(areas)
end

camber_thickness(wing :: HalfWing, num) = camber_thickness.(wing.foils, num)

max_thickness_to_chord_ratio_location(wing :: HalfWing, num) = max_thickness_to_chord_ratio_location.(camber_thickness(wing, num))

function max_tbyc_sweeps(wing :: HalfWing, num)
    xs_max_tbyc = max_thickness_to_chord_ratio_location(wing, num)
    max_tbyc 	= last.(xs_max_tbyc)
    xs_temp 	= first.(xs_max_tbyc)
    xs 			= first.(leading_edge(wing)) .+ wing.chords .* first.(xs_temp)
    ds 			= fwddiff(xs)
    widths 		= @. wing.spans / cos(wing.dihedrals)

    sweeps 		= @. atan(ds, widths)
    xs 	  		= fwdsum(xs_temp) / 2
    tbycs 		= fwdsum(max_tbyc) / 2

    xs, tbycs, sweeps
end

"""
    leading_edge(wing :: HalfWing, flip = false)

Compute the leading edge coordinates of a `HalfWing`, with an option to flip the signs of the ``y``-coordinates.
"""
function leading_edge(wing :: HalfWing, flip = false)
    spans, dihedrals, sweeps = wing.spans, wing.dihedrals, wing.sweeps

    sweeped_spans 	 = [ 0; cumsum(@. spans * tan(sweeps)) ]
    dihedraled_spans = [ 0; cumsum(@. spans * tan(dihedrals)) ]
    cum_spans 		 = [ 0; cumsum(spans) ]

    leading 		 = SVector.(sweeped_spans, ifelse(flip, -cum_spans, cum_spans), dihedraled_spans)

    ifelse(flip, leading[end:-1:1], leading)
end

"""
    wing_bounds(wing :: HalfWing, flip = false)

Compute the leading and trailing edge coordinates of a `HalfWing`, with an option to flip the signs of the ``y``-coordinates.
"""
function wing_bounds(wing :: HalfWing, flip = false)
    chords 	 		= wing.chords
    twisted_chords 	= @. chords * sin(wing.twists)

    leading  		= leading_edge(wing, flip)
    trailing 		= SVector.(chords, (zeros ∘ length)(chords), twisted_chords)

    shifted_trailing = ifelse(flip, trailing[end:-1:1], trailing) .+ leading

    leading, shifted_trailing
end

trailing_edge(wing :: HalfWing, flip = false) = wing_bounds(wing, flip)[2]


@testset "Jordan and Rational canonical form" begin

  Qx, x = PolynomialRing(FlintQQ, "x")
  f = x^6 + rand(Qx, 1:5, -3:3)
  M = companion_matrix(f)
  @test minpoly(Qx, M) == f
  @test charpoly(Qx, M) == f

  ff = factor(f)
  g = first(keys(ff.fac))
  g = divexact(g, leading_coefficient(g))
  J = Hecke.jordan_block(g, 2)
  @test minpoly(Qx, J) == g^2
  @test charpoly(Qx, J) == g^2

  M = identity_matrix(FlintQQ, 6)
  for i = 1:6
    for j = i+1:6
      M[i, j] = rand(FlintQQ, -2:2)
    end
  end
  M1 = transpose(M)
  @test Hecke.is_similar(M, M1)
  MT = Hecke.conjugating_matrix(M, M1)
  @test MT*M1*inv(MT) == M
  J, S = jordan_normal_form(M)
  @test S*M*inv(S) == J
  for i = 1:6
    for j = 1:i-1
      @test iszero(J[i, j])
    end
  end
  for i = 1:5
    @test isone(J[i, i+1]) || iszero(J[i, i+1])
  end
  for i = 1:6
    for j = i+2:6
      @test iszero(J[i, i+2])
    end
  end

  f = x^3+ rand(Qx, 1:2, -2:2)
  g = x^3 + rand(Qx, 1:2, -2:2)
  f1 = f
  f2 = f*g
  C = zero_matrix(FlintQQ, 9, 9)
  Hecke._copy_matrix_into_matrix(C, 1, 1, companion_matrix(f1))
  Hecke._copy_matrix_into_matrix(C, 4, 4, companion_matrix(f2))
  CF, TM = rational_canonical_form(C)
  @test CF == C
  @test TM * C * inv(TM) == CF

end

@testset "Spectrum and eigenspaces" begin
  M = matrix(FlintQQ, 4, 4, [ 0, 0, 1, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1 ])
  l = eigvals(M)
  @test length(keys(l)) == 1
  @test haskey(l, one(FlintQQ))
  @test l[one(FlintQQ)] == 2

  K, a = CyclotomicField(3, "a")
  lK = eigvals(M, K)
  @test length(keys(lK)) == 3
  @test haskey(lK, one(K)) && haskey(lK, a) && haskey(lK, -a - 1)
  @test lK[one(K)] == 2
  @test lK[a] == 1
  @test lK[-a - 1] == 1

  M = matrix(K, 4, 4, [ 0, 0, 1, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1 ])
  lK2 = eigvals(M)
  @test lK == lK2

  Eig = eigenspaces(M, side = :right)
  V = zero_matrix(K, 4, 0)
  for (e, v) in Eig
    @test haskey(lK2, e)
    @test lK2[e] == ncols(v)
    @test M*v == e*v
    V = hcat(V, v)
  end
  @test rref!(V) == 4

  Eig = eigenspaces(M, side = :left)
  V = zero_matrix(K, 0, 4)
  for (e, v) in Eig
    @test haskey(lK2, e)
    @test lK2[e] == nrows(v)
    @test v*M == e*v
    V = vcat(V, v)
  end
  @test rref!(V) == 4

  M = [ matrix(QQ, [ 0 1 0 0 0 0;
                    1 0 0 0 0 0;
                    0 0 1 0 0 0;
                    0 0 0 1 0 0;
                    0 0 0 0 1 0;
                    0 0 0 0 0 1; ]),
       matrix(QQ, [ 1 0 0 0 0 0;
                    0 1 0 0 0 0;
                    0 0 0 1 0 0;
                    0 0 1 0 0 0;
                    0 0 0 0 1 0;
                    0 0 0 0 0 1; ]),
       matrix(QQ, [ 1 0 0 0 0 0;
                    0 1 0 0 0 0;
                    0 0 1 0 0 0;
                    0 0 0 1 0 0;
                    0 0 0 0 0 1;
                    0 0 0 0 1 0; ]) ]

  Eig = common_eigenspaces(M, side = :right)
  V = zero_matrix(FlintQQ, 6, 0)
  for (e, v) in Eig
    @test length(e) == 3
    for i = 1:3
      @test M[i]*v == e[i]*v
    end
    V = hcat(V, v)
  end
  @test rref!(V) == 6

  Eig = common_eigenspaces(M, side = :left)
  V = zero_matrix(K, 0, 6)
  for (e, v) in Eig
    @test length(e) == 3
    for i = 1:3
      @test v*M[i] == e[i]*v
    end
    V = vcat(V, v)
  end
  @test rref!(V) == 6

  N = [ matrix(QQ, [ 0 0 1 0 0 0;
                     1 0 0 0 0 0;
                     0 1 0 0 0 0;
                     0 0 0 1 0 0;
                     0 0 0 0 1 0;
                     0 0 0 0 0 1; ]),
        matrix(QQ, [ 1 0 0 0 0 0;
                     0 1 0 0 0 0;
                     0 0 1 0 0 0;
                     0 0 0 0 0 1;
                     0 0 0 1 0 0;
                     0 0 0 0 1 0; ]) ]

  @test_throws ErrorException Hecke.simultaneous_diagonalization(N)

  T, D = Hecke.simultaneous_diagonalization(N, K)
  for i = 1:length(N)
    @test T*N[i]*inv(T) == D[i]
    @test is_invertible(D[i])
    @test is_diagonal(D[i])
  end
end


"""Function to compute exp(2π i x)"""
cis2pi(x) = cispi(2x)

"""Function to compute sin(2π x)"""
sin2pi(x) = sinpi(2x)

"""Function to compute cos(2π x)"""
cos2pi(x) = cospi(2x)


let
        tot = 0
        totR = 0
        open("data.txt") do file
            for line in eachline(file)
                    tokens = split(line, "x")
                    ints = sort(map(x -> parse(Int, x), tokens))
                    a = 3*ints[1]*ints[2] + 2*ints[2]*ints[3] + 2*ints[1]*ints[3]
                    r = 2*ints[1] + 2*ints[2] + ints[1]*ints[2]*ints[3]
                tot += a
                totR += r
            end
        end
        println(tot)
        println(totR)
end


arg_dict = Dict{Any, Any}(
# specify the physics and SBP operator
"physics" => "Euler",  # specify physics to run
"dimensions" => 2,  # this is a two dimensional mesh

# specify temporal and spatial discretization
"run_type" => 30,  # use LSERK 54
"real_time" => true,  # this is an unsteady problem
"t_max" => 10.0,  # make time
"operator_type" => "SBPOmega",  # specify SBP operator
"order" => 1,  # p = 1 operator
"use_DG" => true,  # use discontinuous galerkin solver
"Flux_name" => "RoeFlux",  # numerical flux function used in face integrals
"CFL" => 0.10,  # CFL number

# specify the problem itself
"IC_name" => "ICUnsteadyVortex",  # initial condtiion
"vortex_x0" => 5.0,  # initial location of vortex
"numBC" => 1,  # number of boundary conditions
"BC1" => [0, 1, 2, 3],  # geometric edges to apply the BC to
"BC1_name" => "unsteadyVortexBC",   # name of boundary condition

# specify mesh
"smb_name" => "../../meshes/vortex_s_1_.smb",
"dmg_name" => ".null",

# misc options
"write_vis" => true,  # write paraview files
"output_freq" => 100,  # do output every this many time steps
"do_postproc" => true,  # calculate error at end of run
"exact_soln_func" => "ICUnsteadyVortex",  # use this function for the exact soltuion (to calculate error)
)


using KernelFunctions
using Zygote, ForwardDiff
using Test, LinearAlgebra
using FiniteDifferences
include("utils_AD.jl")
dims = [10,5]

A = rand(dims...)
B = rand(dims...)
K = [zeros(dims[1],dims[1]),zeros(dims[2],dims[2])]
kernels_noparams = [:SqExponentialKernel,:ExponentialKernel,:GammaExponentialKernel,
 :MaternKernel,:Matern32Kernel,:Matern52Kernel,
 :LinearKernel,:PolynomialKernel,
 :RationalQuadraticKernel,:GammaRationalQuadraticKernel,
 :ExponentiatedKernel]
l = 2.0
ds = [0.0,3.0]
vl = l*ones(dims[1])
testfunction(k,A,B) = det(kernelmatrix(k,A,B))
testfunction(k,A) = det(kernelmatrix(k,A))
ADs = [:Zygote,:ForwardDiff]

## Test kappa functions
@testset "Kappa functions" begin
    for AD in ADs
        @testset "$AD" begin
            for k in kernels_noparams
                for d in ds
                    @eval begin @test kappa_AD(Val(Symbol($AD)),$k(),$d) ≈ kappa_fdm($k(),$d) atol=1e-8 end
                end
            end
            # Linear -> C
            # Polynomial -> C,D
            # Gamma (etc) -> gamma
            #
        end
    end
end

@testset "Transform Operations" begin
    for AD in ADs
        @testset "$AD" begin
            @eval begin
            # Scale Transform
            transform_AD(Val(Symbol($AD)),ScaleTransform(l),A)
            # ARD Transform
            transform_AD(Val(Symbol($AD)),ARDTransform(vl),A)
            # LowRankTransform
            transform_AD(Val(Symbol($AD)),LowRankTransform(rand(2,10)),A)
            # Chain Transform
            # transform_AD(Val(Symbol($AD)),LowRankTransform,A)
            end
        end
    end
end

##TODO Eventually store real results in file
@testset "Zygote Automatic Differentiation test" begin
    @testset "ARD" begin
        for k in kernels
            @testset "$k" begin
                @test all(isapprox.(Zygote.gradient(x->testfunction(k(x),A,B),vl)[1], ForwardDiff.gradient(x->testfunction(k(x),A,B),vl)))
                @test  all(isapprox.(Zygote.gradient(x->testfunction(k(vl),x,B),A)[1],ForwardDiff.gradient(x->testfunction(k(vl),x,B),A)))
                @test all(isapprox.(Zygote.gradient(x->testfunction(k(x),A),vl)[1],ForwardDiff.gradient(x->testfunction(k(x),A),vl)))
                @test all(isapprox.(Zygote.gradient(x->testfunction(k(vl),x),A)[1],ForwardDiff.gradient(x->testfunction(k(vl),x),A)))
            end
        end
    end
    @testset "ISO" begin
        for k in kernels
            @testset "$k" begin
                @test all(isapprox.(Zygote.gradient(x->testfunction(k(x),A,B),l)[1],ForwardDiff.gradient(x->testfunction(k(x[1]),A,B),[l])[1]))
                @test all(isapprox.(Zygote.gradient(x->testfunction(k(l),x,B),A)[1],ForwardDiff.gradient(x->testfunction(k(l),x,B),A)))
                @test all(isapprox.(Zygote.gradient(x->testfunction(k(x),A),l)[1],ForwardDiff.gradient(x->testfunction(k(x[1]),A),[l])))
                @test all(isapprox.(Zygote.gradient(x->testfunction(k(l),x),A)[1],ForwardDiff.gradient(x->testfunction(k(l[1]),x),A)))
            end
        end
    end
end

@testset "ForwardDiff AutomaticDifferentation test" begin
    @testset "ARD" begin
        for k in kernels
            @test_nowarn ForwardDiff.gradient(x->testfunction(k(x),A,B),vl)
            @test_nowarn ForwardDiff.gradient(x->testfunction(k(vl),x,B),A)
            @test_nowarn ForwardDiff.gradient(x->testfunction(k(x),A),vl)
            @test_nowarn ForwardDiff.gradient(x->testfunction(k(vl),x),A)
        end
    end
    @testset "ISO" begin
        for k in kernels
            @test_nowarn ForwardDiff.gradient(x->testfunction(k(x[1]),A,B),[l])
            @test_nowarn ForwardDiff.gradient(x->testfunction(k(l),x,B),A)
            @test_nowarn ForwardDiff.gradient(x->testfunction(k(x[1]),A),[l])
            @test_nowarn ForwardDiff.gradient(x->testfunction(k(l[1]),x),A)
        end
    end
end

@testset "Tracker AutomaticDifferentation test" begin
    @testset "ARD" begin
        for k in kernels
            @test_broken all(Tracker.gradient(x->testfunction(k(x),A,B),vl)[1] .≈ ForwardDiff.gradient(x->testfunction(k(x),A,B),vl))
            @test_broken all(Tracker.gradient(x->testfunction(k(vl),x,B),A)[1] .≈ ForwardDiff.gradient(x->testfunction(k(vl),x,B),A))
            @test_broken all(Tracker.gradient(x->testfunction(k(x),A),vl)[1] .≈  ForwardDiff.gradient(x->testfunction(k(x),A),vl))
            @test_broken all.(Tracker.gradient(x->testfunction(k(vl),x),A) .≈ ForwardDiff.gradient(x->testfunction(k(vl),x),A))
        end
    end
    @testset "ISO" begin
        for k in kernels
            @test_broken Tracker.gradient(x->testfunction(k(x[1]),A,B),[l])
            @test_broken Tracker.gradient(x->testfunction(k(l),x,B),A)
            @test_broken Tracker.gradient(x->testfunction(k(x[1]),A),[l])
            @test_broken Tracker.gradient(x->testfunction(k(l),x),A)

        end
    end
end


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

"""
    mean(ensemble)

Mean of `ensemble`.
"""
function mean(ensemble::Ensemble)
  μs = (; (variable => mean(reals) for (variable, reals) in ensemble.reals)...)
  georef(μs, ensemble.domain)
end

"""
    var(ensemble)

Variance of `ensemble`.
"""
function var(ensemble::Ensemble)
  σs = (; (variable => var(reals) for (variable, reals) in ensemble.reals)...)
  georef(σs, ensemble.domain)
end

"""
    quantile(ensemble, p)

`p`-quantile of `ensemble`.
"""
function quantile(ensemble::Ensemble, p::Number)
  cols = []
  for (variable, reals) in ensemble.reals
    quantiles = map(1:nelements(ensemble.domain)) do location
      slice = getindex.(reals, location)
      quantile(slice, p)
    end
    push!(cols, variable => quantiles)
  end
  georef((; cols...), ensemble.domain)
end

quantile(ensemble::Ensemble, ps::AbstractVector) = [quantile(ensemble, p) for p in ps]