abstract type AbstractBasis{dim,T<:Number} <: AbstractMatrix{T} end

@pure Base.size(::AbstractBasis{dim}) where {dim} = (dim, dim)
Base.getindex(ℬ::AbstractBasis, i::Integer, j::Integer) = getindex(vecbasis(ℬ, :cov), i, j)
@pure Base.eltype(::Type{AbstractBasis{dim,T}}) where {dim,T} = T
@pure getdim(::AbstractBasis{dim}) where {dim} = dim

function subscriptnumber(i::Integer)
    if i < 0
        c = [Char(0x208B)]
    else
        c = []
    end
    for d in reverse(digits(abs(i)))
        push!(c, Char(0x2080+d))
    end
    return join(c)
end

function superscriptnumber(i::Integer)
    if i < 0
        c = [Char(0x207B)]
    else
        c = []
    end
    for d in reverse(digits(abs(i)))
        if d == 0 push!(c, Char(0x2070)) end
        if d == 1 push!(c, Char(0x00B9)) end
        if d == 2 push!(c, Char(0x00B2)) end
        if d == 3 push!(c, Char(0x00B3)) end
        if d > 3 push!(c, Char(0x2070+d)) end
    end
    return join(c)
end

"""
    Basis(v::AbstractMatrix{T}, ::Val{:cov})
    Basis{dim, T<:Number}()
    Basis(θ::T<:Number, ϕ::T<:Number, ψ::T<:Number)

Basis built from a square matrix `v` where columns correspond either to
- primal vectors ie `eᵢ=v[:,i]` if `var=:cov` as by default
- dual vectors ie `eⁱ=v[:,i]` if `var=:cont`.

Basis without any argument refers to the canonical basis (`CanonicalBasis`) in `Rᵈⁱᵐ` (by default `dim=3` and `T=Sym`)

Basis can also be built from Euler angles (`RotatedBasis`) `θ` in 2D and `(θ, ϕ, ψ)` in 3D

The attributes of this object can be obtained by
- `vecbasis(ℬ, :cov)`: square matrix defining the primal basis `eᵢ=e[:,i]`
- `vecbasis(ℬ, :cont)`: square matrix defining the dual basis `eⁱ=E[:,i]`
- `metric(ℬ, :cov)`: square matrix defining the covariant components of the metric tensor `gᵢⱼ=eᵢ⋅eⱼ=g[i,j]`
- `metric(ℬ, :cont)`: square matrix defining the contravariant components of the metric tensor `gⁱʲ=eⁱ⋅eʲ=gⁱʲ[i,j]`

# Examples
```julia
julia> v = Sym[1 0 0; 0 1 0; 0 1 1] ; ℬ = Basis(v)
Basis{3, Sym}
# basis: 3×3 Tensor{2, 3, Sym, 9}:
 1  0  0
 0  1  0
 0  1  1
# dual basis: 3×3 Tensor{2, 3, Sym, 9}:
 1  0   0
 0  1  -1
 0  0   1
# covariant metric tensor: 3×3 SymmetricTensor{2, 3, Sym, 6}:
 1  0  0
 0  2  1
 0  1  1
# contravariant metric tensor: 3×3 SymmetricTensor{2, 3, Sym, 6}:
 1   0   0
 0   1  -1
 0  -1   2

julia> θ, ϕ, ψ = symbols("θ, ϕ, ψ", real = true) ; ℬʳ = Basis(θ, ϕ, ψ) ; display(vecbasis(ℬʳ, :cov))
3×3 Tensor{2, 3, Sym, 9}:
 -sin(ψ)⋅sin(ϕ) + cos(θ)⋅cos(ψ)⋅cos(ϕ)  -sin(ψ)⋅cos(θ)⋅cos(ϕ) - sin(ϕ)⋅cos(ψ)  sin(θ)⋅cos(ϕ)
  sin(ψ)⋅cos(ϕ) + sin(ϕ)⋅cos(θ)⋅cos(ψ)  -sin(ψ)⋅sin(ϕ)⋅cos(θ) + cos(ψ)⋅cos(ϕ)  sin(θ)⋅sin(ϕ)
                        -sin(θ)⋅cos(ψ)                          sin(θ)⋅sin(ψ)         cos(θ)
```
"""
struct Basis{dim,T} <: AbstractBasis{dim,T}
    eᵢ::Matrix{T} # Primal basis `eᵢ=eᵢ[:,i]`
    eⁱ::Matrix{T} # Dual basis `eⁱ=eⁱ[:,i]`
    gᵢⱼ::Symmetric{T,Matrix{T}} # Metric tensor `gᵢⱼ=eᵢ⋅eⱼ=gᵢⱼ[i,j]`
    gⁱʲ::Symmetric{T,Matrix{T}} # Inverse of the metric tensor `gⁱʲ=eⁱ⋅eʲ=gⁱʲ[i,j]`
    function Basis(
        eᵢ::AbstractMatrix{T},
        eⁱ::AbstractMatrix{T},
        gᵢⱼ::AbstractMatrix{T},
        gⁱʲ::AbstractMatrix{T},
    ) where {T}
        dim = size(eᵢ, 1)
        @assert dim == size(eᵢ, 2) "v should be a square matrix"
        if isidentity(eᵢ)
            return CanonicalBasis{dim,T}()
        elseif isidentity(gᵢⱼ)
            return RotatedBasis(eᵢ)
        elseif isdiag(gᵢⱼ)
            χ = sqrt.(diag(gᵢⱼ))
            return OrthogonalBasis(RotatedBasis(eᵢ .* transpose(inv.(χ))), χ)
        else
            eᵢ = Matrix(eᵢ)
            gᵢⱼ = Symmetric(gᵢⱼ)
            gⁱʲ = Symmetric(gⁱʲ)
            eⁱ = Matrix(eⁱ)
            new{dim,T}(eᵢ, eⁱ, gᵢⱼ, gⁱʲ)
        end
    end
    function Basis(ℬ::AbstractBasis{dim,T}, χᵢ::V) where {dim,T,V}
        Χ = collect(χᵢ)
        invΧ = inv.(Χ)
        if Χ == one.(Χ)
            return ℬ
        else
            if ℬ isa OrthonormalBasis
                return OrthogonalBasis(ℬ, Χ)
            else
                eᵢ = vecbasis(ℬ, :cov) .* transpose(Χ)
                eⁱ = vecbasis(ℬ, :cont) .* transpose(invΧ)
                gᵢⱼ = Symmetric(Χ .* metric(ℬ, :cov) .* transpose(Χ))
                gⁱʲ = Symmetric(invΧ .* metric(ℬ, :cont) .* transpose(invΧ))
                new{dim,T}(eᵢ, eⁱ, gᵢⱼ, gⁱʲ)
            end
        end
    end
    function Basis(eᵢ::AbstractMatrix{T}, ::Val{:cov}) where {T}
        dim = size(eᵢ, 1)
        @assert dim == size(eᵢ, 2) "v should be a square matrix"
        if isidentity(eᵢ)
            return CanonicalBasis{dim,T}()
        else
            eᵢ = Matrix(eᵢ)
            gᵢⱼ = simplifyif(Symmetric(transpose(eᵢ) * eᵢ))
            if isidentity(gᵢⱼ)
                return RotatedBasis(eᵢ)
            elseif isdiagonal(gᵢⱼ)
                χ = sqrt.(diag(gᵢⱼ))
                return OrthogonalBasis(RotatedBasis(eᵢ .* transpose(inv.(χ))), χ)
            else
                if T == Sym
                    gⁱʲ = simplifyif(Symmetric(inv(Matrix(gᵢⱼ))))
                else
                    gⁱʲ = simplifyif(inv(gᵢⱼ))
                end
                eⁱ = simplifyif(eᵢ * transpose(gⁱʲ))
                new{dim,T}(eᵢ, eⁱ, gᵢⱼ, gⁱʲ)
            end
        end
    end
    function Basis(eⁱ::AbstractMatrix{T}, ::Val{:cont}) where {T}
        dim = size(eⁱ, 1)
        @assert dim == size(eⁱ, 2) "v should be a square matrix"
        if isidentity(eⁱ)
            return CanonicalBasis{dim,T}()
        else
            eⁱ = Matrix(eⁱ)
            gⁱʲ = simplifyif(Symmetric(transpose(eⁱ) * eⁱ))
            if isidentity(gⁱʲ)
                return RotatedBasis(eⁱ)
            elseif isdiagonal(gⁱʲ)
                uχ = inv.(sqrt.(diag(gⁱʲ)))
                return OrthogonalBasis(RotatedBasis(eⁱ .* transpose(uχ)), uχ)
            else
                if T == Sym
                    gᵢⱼ = simplifyif(Symmetric(inv(Matrix(gⁱʲ))))
                else
                    gᵢⱼ = simplifyif(inv(gⁱʲ))
                end
                eᵢ = simplifyif(eⁱ * transpose(gᵢⱼ))
                new{dim,T}(eᵢ, eⁱ, gᵢⱼ, gⁱʲ)
            end
        end
    end
    Basis(v::AbstractMatrix{T}, var::Symbol) where {T} = Basis(v, Val(var))
    Basis(v::AbstractMatrix{T}) where {T} = Basis(v, :cov)
    Basis(θ::T1, ϕ::T2, ψ::T3 = 0) where {T1,T2,T3} = RotatedBasis(θ, ϕ, ψ)
    Basis(θ::T) where {T} = RotatedBasis(θ)
    Basis{dim,T}() where {dim,T} = CanonicalBasis{dim,T}()
    Basis() = CanonicalBasis()
end

"""
    CanonicalBasis{dim, T}

Canonical basis of dimension `dim` (default: 3) and type `T` (default: Sym)

The attributes of this object can be obtained by
- `vecbasis(ℬ, :cov)`: square matrix defining the primal basis `eᵢ=e[:,i]=δᵢⱼ`
- `vecbasis(ℬ, :cont)`: square matrix defining the dual basis `eⁱ=E[:,i]=δᵢⱼ`
- `metric(ℬ, :cov)`: square matrix defining the covariant components of the metric tensor `gᵢⱼ=eᵢ⋅eⱼ=g[i,j]=δᵢⱼ`
- `metric(ℬ, :cont)`: square matrix defining the contravariant components of the metric tensor `gⁱʲ=eⁱ⋅eʲ=gⁱʲ[i,j]=δᵢⱼ`

# Examples
```julia
julia> ℬ = CanonicalBasis()
CanonicalBasis{3, Sym}
# basis: 3×3 TensND.LazyIdentity{3, Sym}:
 1  0  0
 0  1  0
 0  0  1
# dual basis: 3×3 TensND.LazyIdentity{3, Sym}:
 1  0  0
 0  1  0
 0  0  1
# covariant metric tensor: 3×3 TensND.LazyIdentity{3, Sym}:
 1  0  0
 0  1  0
 0  0  1
# contravariant metric tensor: 3×3 TensND.LazyIdentity{3, Sym}:
 1  0  0
 0  1  0
 0  0  1

julia> ℬ₂ = CanonicalBasis{2, Float64}()
CanonicalBasis{2, Float64}
# basis: 2×2 TensND.LazyIdentity{2, Float64}:
 1.0  0.0
 0.0  1.0
# dual basis: 2×2 TensND.LazyIdentity{2, Float64}:
 1.0  0.0
 0.0  1.0
# covariant metric tensor: 2×2 TensND.LazyIdentity{2, Float64}:
 1.0  0.0
 0.0  1.0
# contravariant metric tensor: 2×2 TensND.LazyIdentity{2, Float64}:
 1.0  0.0
 0.0  1.0
```
"""
struct CanonicalBasis{dim,T} <: AbstractBasis{dim,T} end
CanonicalBasis() = CanonicalBasis{3,Sym}()

"""
    RotatedBasis(θ::T<:Number, ϕ::T<:Number, ψ::T<:Number)
    RotatedBasis(θ::T<:Number)

Orthonormal basis of dimension `dim` (default: 3) and type `T` (default: Sym) built from Euler angles `θ` in 2D and `(θ, ϕ, ψ)` in 3D

# Examples
```julia
julia> θ, ϕ, ψ = symbols("θ, ϕ, ψ", real = true) ; ℬʳ = RotatedBasis(θ, ϕ, ψ) ; display(vecbasis(ℬʳ, :cov))
3×3 Tensor{2, 3, Sym, 9}:
 -sin(ψ)⋅sin(ϕ) + cos(θ)⋅cos(ψ)⋅cos(ϕ)  -sin(ψ)⋅cos(θ)⋅cos(ϕ) - sin(ϕ)⋅cos(ψ)  sin(θ)⋅cos(ϕ)
  sin(ψ)⋅cos(ϕ) + sin(ϕ)⋅cos(θ)⋅cos(ψ)  -sin(ψ)⋅sin(ϕ)⋅cos(θ) + cos(ψ)⋅cos(ϕ)  sin(θ)⋅sin(ϕ)
                        -sin(θ)⋅cos(ψ)                          sin(θ)⋅sin(ψ)         cos(θ)
```
"""
struct RotatedBasis{dim,T} <: AbstractBasis{dim,T}
    eᵢ::Matrix{T} # Primal basis `eᵢ=e[:,i]`
    eⁱ::Matrix{T} # Dual basis `eⁱ=E[:,i]`
    angles::NamedTuple
    function RotatedBasis(R::AbstractMatrix{T}) where {T<:Number}
        dim = size(R, 1)
        if isidentity(R)
            return CanonicalBasis{dim,T}()
        else
            eᵢ = eⁱ = Matrix(R)
            return new{dim,T}(eᵢ, eⁱ, angles(R))
        end
    end
    function RotatedBasis(θ::T1, ϕ::T2, ψ::T3 = 0) where {T1<:Number,T2<:Number,T3<:Number}
        T = promote_type(T1, T2, T3)
        dim = 3
        R = RotZYZ(ϕ, θ, ψ)
        if isidentity(R)
            return CanonicalBasis{dim,T}()
        else
            eᵢ = eⁱ = Matrix(R)
            return new{dim,T}(eᵢ, eⁱ, angles(R))
        end
    end
    function RotatedBasis(θ::T) where {T<:Number}
        dim = 2
        cθ = cos(θ)
        sθ = sin(θ)
        eᵢ = eⁱ = [cθ -sθ; sθ cθ]
        if isidentity(eᵢ)
            return CanonicalBasis{dim,T}()
        else
            return new{dim,T}(eᵢ, eⁱ, angles(eᵢ))
        end
    end
    function RotatedBasis(θ::Sym)
        dim = 2
        cθ = cos(θ)
        sθ = sin(θ)
        eᵢ = eⁱ = [cθ -sθ; sθ cθ]
        if isidentity(eᵢ)
            return CanonicalBasis{dim,Sym}()
        else
            return new{dim,Sym}(eᵢ, eⁱ, (θ = θ,))
        end
    end
end

const OrthonormalBasis{dim,T} = Union{CanonicalBasis{dim,T},RotatedBasis{dim,T}}

struct OrthogonalBasis{dim,T} <: AbstractBasis{dim,T}
    parent::OrthonormalBasis{dim,T}
    λ::Vector{T}
    eᵢ::Matrix{T}
    eⁱ::Matrix{T}
    gᵢⱼ::Diagonal{T,Vector{T}}
    gⁱʲ::Diagonal{T,Vector{T}}
    function OrthogonalBasis(parent::OrthonormalBasis{dim,T}, λ::Vector) where {dim,T}
        λ = T.(λ)
        if λ == one.(λ)
            return parent
        else
            eᵢ = [λ[j] * parent[i, j] for i ∈ 1:dim, j ∈ 1:dim]
            eⁱ = [parent[i, j] / λ[j] for i ∈ 1:dim, j ∈ 1:dim]
            return new{dim,T}(parent, λ, eᵢ, eⁱ, Diagonal(λ .^ 2), Diagonal(inv.(λ) .^ 2))
        end
    end
end

relevant_OrthonormalBasis(ℬ::OrthogonalBasis) = ℬ.parent
relevant_OrthonormalBasis(ℬ::OrthonormalBasis) = ℬ
relevant_OrthonormalBasis(::Basis{dim,T}) where {dim,T} = CanonicalBasis{dim,T}()

@inline CylindricalBasis(θ) = RotatedBasis(0, θ, 0)

@inline SphericalBasis(θ, ϕ) = RotatedBasis(θ, ϕ, 0)

const AllOrthogonalBasis{dim,T} = Union{OrthonormalBasis{dim,T},OrthogonalBasis{dim,T}}

angles(M::AbstractMatrix{T}, ::Val{2}) where {T} =
    (θ = atan(M[2, 1] - M[1, 2], M[1, 1] + M[2, 2]),)
function angles(M::AbstractMatrix{T}, ::Val{3}) where {T}
    R = RotZYZ(M)
    return (θ = R.theta2, ϕ = R.theta1, ψ = R.theta3)
end

"""
    angles(M::AbstractMatrix{T})

Determines the Euler angles corresponding to the input matrix supposed to be a rotation matrix or at least a similarity

# Examples
```julia
julia> θ, ϕ, ψ = symbols("θ, ϕ, ψ", real = true) ; ℬʳ = RotatedBasis(θ, ϕ, ψ) ; display(vecbasis(ℬʳ, :cov))
3×3 Tensor{2, 3, Sym, 9}:
 -sin(ψ)⋅sin(ϕ) + cos(θ)⋅cos(ψ)⋅cos(ϕ)  -sin(ψ)⋅cos(θ)⋅cos(ϕ) - sin(ϕ)⋅cos(ψ)  sin(θ)⋅cos(ϕ)
  sin(ψ)⋅cos(ϕ) + sin(ϕ)⋅cos(θ)⋅cos(ψ)  -sin(ψ)⋅sin(ϕ)⋅cos(θ) + cos(ψ)⋅cos(ϕ)  sin(θ)⋅sin(ϕ)
                        -sin(θ)⋅cos(ψ)                          sin(θ)⋅sin(ψ)         cos(θ)

julia> angles(ℬʳ)
(θ = θ, ϕ = ϕ, ψ = ψ)
```
"""
angles(M::AbstractMatrix{T}) where {T} = angles(M, Val(size(M, 1)))
angles(ℬ::RotatedBasis) = ℬ.angles

angles(v::AbstractVector{T}, ::Val{2}) where {T} = (θ = atan(v[2], v[1]),)
angles(v::AbstractVector{T}, ::Val{3}) where {T} =
    (θ = atan(√(v[1]^2 + v[2]^2), v[3]), ϕ = atan(v[2], v[1]))
angles(v::AbstractVector{T}) where {T} = angles(v, Val(size(v, 1)))

invvar(::Val{:cov}) = :cont
invvar(::Val{:cont}) = :cov
invvar(var) = invvar(Val(var))

"""
    vecbasis(ℬ::AbstractBasis, var = :cov)

Returns the primal (if `var = :cov`) or dual (if `var = :cont`) basis
"""
vecbasis(ℬ::AbstractBasis, ::Val{:cov}) = ℬ.eᵢ
vecbasis(ℬ::AbstractBasis, ::Val{:cont}) = ℬ.eⁱ
vecbasis(::CanonicalBasis{dim,T}, ::Val{:cov}) where {dim,T} = Id2{dim,T}()
vecbasis(::CanonicalBasis{dim,T}, ::Val{:cont}) where {dim,T} = Id2{dim,T}()

vecbasis(ℬ::AbstractBasis, var) = vecbasis(ℬ, Val(var))
vecbasis(ℬ::AbstractBasis) = vecbasis(ℬ, :cov)
vecbasis(ℬ::AbstractBasis, i::Integer, j::Integer, var = :cov) = vecbasis(ℬ, Val(var))[i, j]
vecbasis(ℬ::AbstractBasis, i::Integer, var = :cov) = vecbasis(ℬ, Val(var))[:, i]

strvecbasis(::AbstractBasis, i::Integer, ::Val{:cov} ; vec = "𝐞") = vec * subscriptnumber(i)
strvecbasis(::AbstractBasis, i::Integer, ::Val{:cont} ; vec = "𝐞") = vec * superscriptnumber(i)
strvecbasis(ℬ::AbstractBasis, i::Integer, var = :cov ; vec = "𝐞") = strvecbasis(ℬ, i, Val(var) ; vec = vec)

const dsubscriptchar = Dict(
    "a" => "ₐ",
    "e" => "ₑ",
    "h" => "ₕ",
    "i" => "ᵢ",
    "j" => "ⱼ",
    "k" => "ₖ",
    "l" => "ₗ",
    "m" => "ₘ",
    "n" => "ₙ",
    "o" => "ₒ",
    "p" => "ₚ",
    "r" => "ᵣ",
    "s" => "ₛ",
    "t" => "ₜ",
    "u" => "ᵤ",
    "v" => "ᵥ",
    "x" => "ₓ",
    "β" => "ᵦ",
    "γ" => "ᵧ",
    "ρ" => "ᵨ",
    "ϕ" => "ᵩ",
    "χ" => "ᵪ",
)

const dsuperscriptchar = Dict(
    "a" => "ᵃ",
    "b" => "ᵇ",
    "c" => "ᶜ",
    "d" => "ᵈ",
    "e" => "ᵉ",
    "f" => "ᶠ",
    "g" => "ᵍ",
    "h" => "ʰ",
    "i" => "ⁱ",
    "j" => "ʲ",
    "k" => "ᵏ",
    "l" => "ˡ",
    "m" => "ᵐ",
    "n" => "ⁿ",
    "o" => "ᴼ",
    "p" => "ᵖ",
    "r" => "ʳ",
    "s" => "ˢ",
    "t" => "ᵗ",
    "u" => "ᶸ",
    "v" => "ᵛ",
    "w" => "ʷ",
    "x" => "ˣ",
    "y" => "ʸ",
    "z" => "ᶻ",
    "β" => "ᵝ",
    "γ" => "ᵞ",
    "ε" => "ᵋ",
    "θ" => "ᶿ",
    "ι" => "ᶥ ",
    "ϕ" => "ᵠ",
    "χ" => "ᵡ",
)

subscriptchar(s::String) = s ∈ keys(dsubscriptchar) ? dsubscriptchar[s] : s ∈ keys(dsuperscriptchar) ? dsuperscriptchar[s] : s
superscriptchar(s::String) = s ∈ keys(dsuperscriptchar) ? dsuperscriptchar[s] : s ∈ keys(dsubscriptchar) ? dsubscriptchar[s] : s

strvecbasis(::AbstractBasis, i::AbstractString, ::Val{:cov} ; vec = "𝐞") = vec * subscriptchar(i)
strvecbasis(::AbstractBasis, i::AbstractString, ::Val{:cont} ; vec = "𝐞") = vec * superscriptchar(i)
strvecbasis(ℬ::AbstractBasis, i::AbstractString, var = :cov ; vec = "𝐞") = strvecbasis(ℬ, i, Val(var) ; vec = vec)

"""
    metric(ℬ::AbstractBasis, var = :cov)

Returns the covariant (if `var = :cov`) or contravariant (if `var = :cont`) metric matrix
"""
metric(ℬ::AbstractBasis, ::Val{:cov}) = ℬ.gᵢⱼ
metric(ℬ::AbstractBasis, ::Val{:cont}) = ℬ.gⁱʲ
metric(::OrthonormalBasis{dim,T}, ::Val{:cov}) where {dim,T} = Id2{dim,T}()
metric(::OrthonormalBasis{dim,T}, ::Val{:cont}) where {dim,T} = Id2{dim,T}()

metric(ℬ::AbstractBasis, var) = metric(ℬ, Val(var))
metric(ℬ::AbstractBasis) = metric(ℬ, :cov)
metric(ℬ::AbstractBasis, i::Integer, j::Integer, var = :cov) = metric(ℬ, Val(var))[i, j]

"""
    normalize(ℬ::AbstractBasis, var = cov)

Builds a basis after normalization of column vectors of input matrix `v` where columns define either
- primal vectors ie `eᵢ=v[:,i]/norm(v[:,i])` if `var = :cov` as by default
- dual vector ie `eⁱ=v[:,i]/norm(v[:,i])` if `var = :cont`.
"""
function LinearAlgebra.normalize(ℬ::AbstractBasis, var = :cov)
    w = copy(vecbasis(ℬ, var))
    for i = 1:size(w, 2)
        w[:, i] /= norm(w[:, i])
    end
    return Basis(w, var)
end

"""
    isorthogonal(ℬ::AbstractBasis)

Checks whether the basis `ℬ` is orthogonal
"""
isorthogonal(ℬ::AbstractBasis) = isdiagonal(metric(ℬ))

isorthogonal(::OrthonormalBasis) = true

isorthogonal(::OrthogonalBasis) = true

"""
    isorthonormal(ℬ::AbstractBasis)

Checks whether the basis `ℬ` is orthonormal
"""
isorthonormal(ℬ::AbstractBasis) = isidentity(metric(ℬ))

isorthonormal(::OrthonormalBasis) = true

#####################
# Display Functions #
#####################
for OP in (:show, :print, :display)
    @eval begin
        function Base.$OP(ℬ::AbstractBasis)
            $OP(typeof(ℬ))
            print("→ basis: ")
            $OP(vecbasis(ℬ, :cov))
            print("→ dual basis: ")
            $OP(vecbasis(ℬ, :cont))
            print("→ covariant metric tensor: ")
            $OP(metric(ℬ, :cov))
            print("→ contravariant metric tensor: ")
            $OP(metric(ℬ, :cont))
        end
    end
end

export Basis, CanonicalBasis, RotatedBasis, CylindricalBasis, SphericalBasis, OrthogonalBasis, AllOrthogonalBasis
export getdim, vecbasis, metric, angles, isorthogonal, isorthonormal


module SkeletonPairsTests

using Test
using Gridap.Arrays
using Gridap.Fields
using Gridap.ReferenceFEs
using Gridap.Geometry

domain = (0,1,0,1)
partition = (3,3)
model = CartesianDiscreteModel(domain,partition)

trian_Γ = SkeletonTriangulation(model)
trian = get_volume_triangulation(trian_Γ)

fun(x) = sin(pi*x[1])*cos(pi*x[2])
q2x = get_cell_map(trian)
funq = compose(fun,q2x)
fun_Γ = restrict(funq,trian_Γ)
@test isa(fun_Γ,SkeletonPair)

end # module


using Gadfly, StatsBase, DSP
import Gadfly.plot, StatsBase.pacf

colorscale = Scale.color_discrete_hue()

function buildlayers(f::Function, colnames, colors)
  [f(coln, colr) for (coln, colr) in zip(colnames, colors)]
end #buildlayers

# Raw data visualization

function plot(ta::TimeArray; subplots=false, title="")
  subplots ? error("Subplotting not yet supported") :
  linecolors = colorscale.f(length(ta.colnames))
  layers = buildlayers(ta.colnames, linecolors) do coln, colr
    layer(x=ta.timestamp, y=ta[coln].values, Geom.line, Theme(default_color=colr))
  end
  plot(layers..., Guide.xlabel(""), Guide.ylabel(""), Guide.title(title), Guide.manual_color_key("", ta.colnames, linecolors))
end #plot

function histogram(ta::TimeArray; subplots=false, title="")
  subplots ? error("Subplotting not yet supported") :
  linecolors = colorscale.f(length(ta.colnames))
  layers = buildlayers(ta.colnames, linecolors) do coln, colr
    layer(x=ta[coln].values, Geom.histogram, Theme(default_color=colr))
  end
  plot(layers..., Guide.xlabel(""), Guide.ylabel("Counts"), Guide.title(title), Guide.manual_color_key("", ta.colnames, linecolors))
end #hist

function density(ta::TimeArray; subplots=false, title="")
  subplots ? error("Subplotting not yet supported") :
  linecolors = colorscale.f(length(ta.colnames))
  layers = buildlayers(ta.colnames, linecolors) do coln, colr
    layer(x=ta[coln].values, Geom.density, Theme(default_color=colr))
  end
  plot(layers..., Guide.xlabel(""), Guide.ylabel("Probability Density"), Guide.title(title), Guide.manual_color_key("", ta.colnames, linecolors))
end # density

# Serial correlation & cross-correlation visualization

function directedscatter(ta::TimeArray; naive=false, subplots=false, nlags=1, title="")
  naive || require_clean(ta)
  subplots ? error("Subplotting not yet supported") :
  linecolors = colorscale.f(length(ta.colnames))
  layers = buildlayers(ta.colnames, linecolors) do coln, colr
    layer(x=lag(ta[coln], nlags).values, y=ta[coln].values[nlags+1:end], Geom.line(preserve_order=true), Theme(default_color=colr))
  end
  plot(layers..., Guide.xlabel("t-"*string(nlags)), Guide.ylabel("t"), Guide.title(title), Guide.manual_color_key("", ta.colnames, linecolors))
end # directedscatter

function laghistogram(ta::TimeArray; naive=false, nlags=1, title="")
  naive || require_clean(ta)
  length(ta.colnames) > 1 ? error("Subplotting not yet supported. Use a single-column TimeArray.") :
  plot(
    x=lag(ta, nlags).values, y=ta.values[nlags+1:end], Geom.histogram2d,
    Guide.xlabel("t-"*string(nlags)), Guide.ylabel("t"), Guide.title(title)
  )
end # laghistogram

function acf(ta::TimeArray; naive=false, lags=1:30)
  naive || require_clean(ta)
  length(ta.colnames) > 1 ? error("Subplotting not yet supported. Use a single-column TimeArray.") :
  conf_int = 1.96 / sqrt(length(ta))
  plot(x=lags, y=autocor(ta.values, lags), xmin=collect(lags)-.1, xmax=lags, yintercept=[-conf_int, +conf_int], Geom.bar, Geom.hline(color=colorant"orange"), Scale.y_continuous(minvalue=-1, maxvalue=1), Guide.xlabel("lags"), Guide.ylabel(""), Guide.title(ta.colnames[1]*" ACF"))
end # acf

function pacf(ta::TimeArray; naive=false, lags=1:30)
  naive || require_clean(ta)
  length(ta.colnames) > 1 ? error("Subplotting not yet supported. Use a single-column TimeArray.") :
  conf_int = 1.96 / sqrt(length(ta))
  plot(x=lags, y=pacf(ta.values, lags), xmin=collect(lags)-.1, xmax=lags, yintercept=[-conf_int, +conf_int], Geom.bar, Geom.hline(color=colorant"orange"), Scale.y_continuous(minvalue=-1, maxvalue=1), Guide.xlabel("lags"), Guide.ylabel(""), Guide.title(ta.colnames[1]*" PACF"))
end # pacf

function ccf(ta::TimeArray; naive=false, lags=-15:15)
  naive || require_clean(ta)
  length(ta.colnames) != 2 ? error("`ccf` requires a 2-column TimeArray") :
  conf_int = 1.96 / sqrt(length(ta))
  plot(x=lags, y=crosscor(ta[ta.colnames[1]].values, ta[ta.colnames[2]].values, lags), xmin=collect(lags)-.1, xmax=lags, yintercept=[-conf_int, +conf_int], Geom.bar, Geom.hline(color=colorant"orange"), Scale.y_continuous(minvalue=-1, maxvalue=1), Guide.xlabel("lags"), Guide.ylabel(""), Guide.title(ta.colnames[1]*", "*ta.colnames[2]*" CCF"))
end # ccf

# Spectral analysis

function cumulative_periodogram(ta::TimeArray; naive=false)
  naive || require_clean(ta)
  length(ta.colnames) > 1 ? error("Subplotting not yet supported. Use a single-column TimeArray.") :
  p = periodogram(ta.values)
  plot(x=freq(p), y=cumsum(power(p)), Geom.line, Guide.xlabel("Frequency"), Guide.ylabel("Cumulative Power"), Guide.title(ta.colnames[1]*" Cumulative Periodogram"))
end # cumulatve_periodogram


## MATLAB interfcae to find_zero
## Main functions are
## * fzero(f, ...) to find _a_ zero of f
## * fzeros(f, ...) to attempt to find all zeros of f

"""
    fzero(f, x0; order=0; kwargs...)

Find zero of a function using an iterative algorithm

* `f`: a scalar function or callable object
* `x0`: an initial guess, finite valued.
* `order`

This is a polyalgorithm redirecting to different algorithms based on the value of `order`. Dispatches to `find_zero(f, x0, OrderN(); kwargs...)`.

"""
function fzero(f, x0::Number; kwargs...)
    x = float(x0)
    isinf(x) && throw(ConvergenceFailed("An initial value must be finite"))
    derivative_free(f, x; kwargs...)
end

"""
    fzero(f, a, b; kwargs...)

Find zero of a function within a bracket, [a,b].

Dispatches to `find_zero(f, (a,b), Bisection())`.

"""
fzero(f, a::Number, b::Number; kwargs...) = find_zero(f, (a,b), Bisection(); kwargs...)
fzero(f, bracket::Vector{T}; kwargs...)  where {T <: Number} = find_zero(f, bracket, Bisection(); kwargs...)
fzero(f, bracket::Tuple{T,S}; kwargs...)  where {T <: Number, S<:Number} = find_zero(f, bracket, Bisection();kwargs...)

"""
    fzero(f, fp, x0; kwargs)

Find zero using Newton's method.

Dispatches to `find_zero((f,fp), x0, Roots.Newton(); kwargs...)`.

"""
fzero(f::Function, fp::Function, x0::Real; kwargs...) = newton(f, fp, float(x0); kwargs...)

# match fzero up with find_zero
@noinline function derivative_free(f, x0; order::Int=0, kwargs...)

    if order == 0
        method = Order0()
    elseif order == 1
        method = Order1()
    elseif order == 2
        method = Order2()
    elseif order == 5
        method = Order5()
    elseif order == 8
        method = Order8()
    elseif order == 16
        method = Order16()
    else
        warn("Invalid order. Valid orders are 0, 1, 2, 5, 8, and 16")
        throw(ArgumentError())
    end

    d = Dict(kwargs)
    for (o, n) in ((:ftol, :atol), (:ftolrel, :rtol),
         (:xtol, :xatol), (:xtolrel, :xrtol))
        if haskey(d, o)
            d[n] = d[o]
        end
    end

    find_zero(f, x0, method; d...)
end

# ##
# """
#     fzero(f, x0, bracket; kwargs...)
#
# Find a zero within a bracket with an initial guess to *possibly* speed things along.
#
# Dispatches to the `A42` method.
#
#"""
@deprecate fzero(f, x0::Real, bracket::Vector; kwargs...)  fzero(f, bracket)

## fzeros
"""

`fzeros(f, a, b; kwargs...)`

Searches for all zeros of `f` within an interval `(a,b)`. Assume neither `a` or `b` is a zero.

Dispatches to `find_zeros(f, a, b; kwargs...)`.
"""
function fzeros(f, a::Number, b::Number; kwargs...)
    find_zeros(f, float(a), float(b); kwargs...)
end
fzeros(f, bracket::Vector{T}; kwargs...) where {T <: Number} = fzeros(f, a, b; kwargs...)
fzeros(f, bracket::Tuple{T,S}; kwargs...) where {T <: Number, S<:Number} = fzeros(f, a, b; kwargs...)


module OPpoly
#
# Module OPpoly contains functions manipulating multivariate polynomials.
# A polynomial is represented as a vector of coefficients.  The indexing
# is arranged (taking a D=3,ord=4 example) as
# c000, c100, c010, c001, c200, c110, c101, c020, c011, c002, c300, c210,
# c201, c120, c111, c102, c030, c021, c012, c003, c400, c310, c301, c220,
# c211, c202, c130, c121, c112, c103, c040, c031, c022, c013, c004
#

export tensorToList, listToTensor, polyrec, polyeval, rateval,
       solveLSPolyCoeffs

# =======================================================================
function tensorToList(t)
#
# Convert between tensor (00,10,01,20,11,02,30,21,12,03,...) and list
# (1,2,3,...) indexing.
#
# Version:        Changes:
# --------        -------------
# 13.04.2021      Original code.
#
# Version:        Verification:
# --------        -------------
# 13.04.2021      Tested on sample problems up to D=4 and ord=3.
#
# Inputs:
# -------
# t               : D-by-N matrix of integers, with D the dimension.
#
# Outputs:
# --------
# L               : Length N vector of integers, containing the indices.

   D = size(t,1)
   N = size(t,2)

   L = ones(Int,N)
   for ip in 1:N
      ord = sum(t[1:D,ip])
      for id in 1:D
         De = D - id + 1
         L[ip] = L[ip] + binomial((ord-1)+De,De)
         ord = ord - t[id,ip]
      end
   end

   return L

end # tensorToList

# =======================================================================
function listToTensor(L, D)
#
# Convert between list (1,2,3,...) and tensor indexing
# (00,10,01,20,11,02,30,21,12,03,...).
#
# Version:        Changes:
# --------        -------------
# 13.04.2021      Original code.
#
# Version:        Verification:
# --------        -------------
# 13.04.2021      Verified that this inverts tensorToList.
#
# Inputs:
# -------
# L               : Index list, vector of length N.
# D               : Dimension, integer.
#
# Outputs:
# --------
# t               : D-by-N matrix of tensor indices.

   N = length(L)

   # The matrix A is consistent for any dimension D.
   A = zeros(Int,D,D)
   for id = 1:D
      A[id,id:D] = ones(Int,1,D-id+1)
   end

   b = zeros(Int,D,N)
   for ip in 1:N

      Le = L[ip]

      for id in 1:D

         De = D - id + 1

         ind = 1
         sav = 0
         flg = true
         while flg
            if (ind >= Le)
               flg = false
               Le = Le - sav
            else
               b[id,ip] = b[id,ip] + 1
               sav = ind
               ind = binomial(De+b[id,ip],De)
            end
         end

      end

   end

   t = zeros(Int,D,N)
   t[:,:] = A\b

   return t

end # listToTensor

# =======================================================================
function polyrec(x, ord; ptype=0)
#
# Recursively evaluate the normalized terms of a polynomial; that is,
# the value of each term when the coefficient is equal to 1.
#
# Version:        Changes:
# --------        -------------
# 22.04.2021      Original code.
#
# Version:        Verification:
# --------        -------------
# 22.04.2021      Hand-checked the first few terms of a D=2 case.
#
# Inputs:
# -------
# x               : Column vector of length D.
# ord             : The order of individual x components, up to which
#                   the polynomial is to be computed.
# ptype           : Specifies the type of polynomial.
#                   0: Nominal   x[n+1] = x x[n]
#                   1: Chebyshev T[n+1] = 2xT[n] - T[n-1]
#                   2: Legendre  P[n+1] = ((2n+1)xP[n] - nP[n-1])/(n+1)
#                   3: Hermite   H[n+1] = (x H[n] - sq(n)H[n-1])/sq(n+1)
#
# Outputs:
# --------
# psi             : A D-by-ord+1 matrix containing the terms' values
#                   at the given x vector.

   D = size(x,1)
   psi = zeros(D,ord+1)
   psi[:,1] = ones(D)

   if (ord > 0)
      if (ptype == 1)
         # Chebyshev
         for id in 1:D
            psi[id,2] = x[id]
            for n in 1:ord-1
               psi[id,n+2] = 2*x[id]*psi[id,n+1] - psi[id,n]
            end
         end
      elseif (ptype == 2)
         # Legendre
         for id in 1:D
            psi[id,2] = x[id]
            for n in 1:ord-1
               psi[id,n+2] = ((2*n+1)*x[id]*psi[id,n+1] - n*psi[id,n])/(n+1)
            end
         end
      elseif (ptype == 3)
         # Hermite
         for id in 1:D
            psi[id,2] = x[id]
            sqn = 1.0
            for n in 1:ord-1
               sqn1 = sqrt(n+1)
               psi[id,n+2] = (x[id]*psi[id,n+1] - sqn*psi[id,n])/sqn1
               sqn = sqn1  # Avoid recomputing sqrt's.
            end
         end
      else
         # Default to nominal.
         for id in 1:D
            psi[id,2] = x[id]
            for n = 1:ord-1
               psi[id,n+2] = x[id]*psi[id,n+1]
            end
         end
      end
   end # if ord > 0.

   return psi

end # polyrec

# =======================================================================
function polyeval(p, t, x; ptype=0)
#
# Evaluate a polynomial defined by coefficients p of the tensor product
# t at the coordinates x.  This is currently designed to be easy to use,
# not optimal in terms of speed and storage.
#
# Version:        Changes:
# --------        -------------
# 22.04.2021      Original code.
#
# Version:        Verification:
# --------        -------------
# 22.04.2021
#
# Inputs:
# -------
# p               : Vector of polynomial coefficients.  See the module
#                   header for the convention used for indexing.
# t               : Tensor specifying the x component orders in each
#                   polynomial term.  D-by-N matrix.
# x               : Column vector or matrix; each column contains an x
#                   vector of length D.
# ptype           : Specifies the type of polynomial.  See polyrec.
#
# Outputs:
# --------
# y               : A vector containing the outputs, one scalar output
#                   for each x vector.

   D = size(x,1)   # Dimension of x.
   M = size(x,2)   # Number of points.
   N = size(p,1)   # Number of coefficients.

   ord = maximum(sum(t,dims=1))

   y = zeros(M)

   for ip in 1:M

      # Compute the individual polynomial terms in a recursive way.
      if (D == 1)
         # This is required since the ":" notation is not permitted
         # if x is a scalar.
         psi = polyrec(x[ip],ord,ptype=ptype)
      else
         psi = polyrec(x[:,ip],ord,ptype=ptype)
      end

      for ic in 1:N
         val = 1.0
         for id = 1:D
            val = val*psi[id,t[id,ic]+1]
         end
         y[ip] = y[ip] + p[ic]*val
      end

   end

   return y

end # polyeval: version including explicit input of the sequence.

function polyeval(p, x; ptype=0)
#
# Evaluate a polynomial defined by coefficients p at the coordinates x.
# This is currently designed to be easy to use, not optimal in terms of
# speed and storage.
#
# This version generates the tensor t automatically, based on the
# assumed sequence of polynomials indicated in the module header.
#
# Version:        Changes:
# --------        -------------
# 22.04.2021      Original code.
#
# Version:        Verification:
# --------        -------------
# 22.04.2021
#

   D = size(x,1)   # Dimension of x.
   N = size(p,1)   # Number of coefficients.

   L = range(1,stop=N)
   t = listToTensor(L,D)

   y = polyeval(p,t,x,ptype=ptype)

   return y

end # polyeval: version with an assumed polynomial sequence.

# =======================================================================
function rateval(p, tp, q, tq, x; ptype=0)
#
# Evaluate a rational function defined by numerator coefficients p and
# denominator coefficients q at the coordinates x.  The order associated
# with each coefficient is given in the matrices tp,tq.
#
#
# Version:        Changes:
# --------        -------------
# 22.04.2021      Original code.
#
# Version:        Verification:
# --------        -------------
# 22.04.2021      Verified some example calculations.
#
# Inputs:
# -------
# p,q             : Coefficients of numerator and denominator polynomials.
# tp,tq           : Tensors specifying the x component orders in each
#                   polynomial term.  D-by-Np,Nq matrices.
# x               : Column vector or matrix; each column contains an x
#                   vector.
#
# Outputs:
# --------
# y               : A vector containing the outputs, one scalar output
#                   for each x vector.

   yp = polyeval(p,tp,x,ptype=ptype)
   yq = polyeval(q,tq,x,ptype=ptype)

   return (yp ./ yq)

end # rateval: explicit sequence.

function rateval(p, q, x; ptype=0)
#
# Evaluate a rational function defined by numerator coefficients p and
# denominator coefficients q at the coordinates x.
#
# Version:        Changes:
# --------        -------------
# 22.04.2021      Original code.
#
# Version:        Verification:
# --------        -------------
# 22.04.2021
#

   yp = polyeval(p,x,ptype=ptype)
   yq = polyeval(q,x,ptype=ptype)

   return (yp ./ yq)

end # rateval: implicit sequence.

#=

# =======================================================================
function polyderiv(p, t, k, x; cof=0)
#
# Evaluate the derivatives of a polynomial at specified points.  At the
# moment this supports only nominal polynomials (ptype = 0).  That is,
# if a column of t contains, say, (1, 0, 3, 2) then the polynomial term
# is (x1 x3^3 x4^2).
#
# Version:        Changes:
# --------        -------------
# 18.05.2021      Original code.
#
# Version:        Verification:
# --------        -------------
# 18.05.2021
#
# Inputs:
# -------
# p               : Vector of polynomial coefficients, length N.  See
#                   the module header for the convention used for
#                   indexing.
# t               : Tensor specifying the x component orders in each
#                   polynomial term.  D-by-N matrix.
# k               : An integer specifying which derivative to take:
#                   1 = x1, 2 = x2, ...
# x               : Column vector or matrix of dimension D-by-M; each
#                   column contains an x vector of length D.
#
# Outputs:
# --------
# pd, td          : Coefficients and tensor describing the derivative
#                   polynomial.

   ptype = 0       # Limited functionality in order to get something
                   # working fast.

   D = size(x,1)   # Dimension of x.
   M = size(x,2)   # Number of points.
   N = size(p,1)   # Number of coefficients.

   ord = maximum(sum(t,dims=1))

   dy = zeros(M,L)

   for ip in 1:M

      # Compute the individual polynomial terms in a recursive way.
      if (D == 1)
         # This is required since the ":" notation is not permitted
         # if x is a scalar.
         psi = polyrec(x[ip],ord,ptype=ptype)
      else
         psi = polyrec(x[:,ip],ord,ptype=ptype)
      end

      for ic in 1:N
         val = 1.0
         for id in 1:D
            val = val*psi[id,t[id,ic]+1]
         end
         y[ip] = y[ip] + p[ic]*val
      end

   end

   return y

end # polyderiv

=#

# =======================================================================
#function polyint()
#
# Evaluate
#
# Version:        Changes:
# --------        -------------
# 13.04.2021      Original code.
#
# Version:        Verification:
# --------        -------------
# 13.04.2021
#
# Inputs:
# -------
#
#
# Outputs:
# --------
#

#end # polyint

# ===================================================================
function solveLSPolyCoeffs(x, y, tp, tq; ptype=0, wgt=1.0)
#
# Given a list of samples, solves a (weighted) least-squares
# polynomial or rational-function fit.
#
# [Note, this should be augmented to accommodate complex coefficients.]
#
# Version:        Changes:
# --------        -------------
# 22.04.2021      Original code.
#
# Version:        Verification:
# --------        -------------
# 22.04.2021      Verified the fit to a known quadratic rational
#                 polynomial function in two dimensions.
#
# Inputs:
# -------
# x               : D-by-Ns vector of coordinates.
# y               : Length Ns vector of sample values.
# tp,tq           : Tensors specifying the x component orders in each
#                   polynomial term.  D-by-Np,Nq matrices.  Note that
#                   tq should always include one all-zeros term as its
#                   first entry.
# ptype           : Specifies the type of polynomial.  See polyrec.
# wgt             : A vector of weights, the square roots of the weight
#                   matrix diagonal.
#
# Outputs:
# --------
# p,q             : Vector of best-fit polynomial coefficients; p
#                   for the numerator polynomials, q for the
#                   denominator polynomials.

   D  = size(x,1)
   Ns = size(x,2)
   Np = size(tp,2)
   Nq = size(tq,2)
   Nc = Np + Nq

   # The maximum polynomial order in the individual variables.
   op = maximum(tp)
   oq = maximum(tq)
   omax = max(op,oq)

   if (omax == 0)
      # We are dealing with fitting a least-squares constant, not a
      # polynomial.
      A = ones(Ns,1) .* wgt
      AT = transpose(A)
      p = (AT*A)\(AT*(wgt .* y))
      q = 1.0
   else

      # The matrix in the equation A [p;q] = y.  The "-1" is for the
      # forced constant = 1 term in the denominator.
      A = ones(Ns,Nc-1)
      for isamp in 1:Ns

         # Compute the individual polynomial terms in a recursive way.
         psi = polyrec(x[:,isamp],omax,ptype=ptype)

         for ip in 1:Np
            for id in 1:D
               A[isamp,ip] = A[isamp,ip]*psi[id,tp[id,ip]+1]
            end
         end
         for iq in 1:Nq-1  # Nq-1: Neglect the constant 1.0 in the denominator.
            ind = Np + iq
            for id in 1:D
               A[isamp,ind] = A[isamp,ind]*psi[id,tq[id,iq+1]+1]
            end
            A[isamp,ind] = -A[isamp,ind]*y[isamp]
         end
      end

      A = A .* wgt

      AT = A'
      c = (AT*A)\(AT*(wgt .* y))
      p = c[1:Np]
      if (oq > 0)
         q = [1.0;c[Np+1:Nc-1]]
      else
         q = 1.0
      end

   end

   return p, q

end  # solveLSPolyCoeffs

end # module


using FastCombinations
import FastCombinations: next_states

using Base.Test

@test next_states(1:6, 1, Val{0}) == ()
@test next_states(1:6, 1, Val{1}) == (1,)
@test next_states(1:6, 1, Val{6}) == (1, 2, 3, 4, 5, 6)

@test_throws ErrorException Combinations{0}(1:1)

@test first(Combinations{1}(1:6)) == (1,)
@test first(Combinations{2}(1:6)) == (1,2)
@test first(Combinations{3}(1:6)) == (1,2,3)
@test first(Combinations{4}(1:6)) == (1,2,3,4)
@test first(Combinations{5}(1:6)) == (1,2,3,4,5)
@test first(Combinations{6}(1:6)) == (1,2,3,4,5,6)
@test collect(Combinations{7}(1:6)) == NTuple{7, Int}[]

@test first(MultiCombinations{1}(1:6)) == (1,)
@test first(MultiCombinations{2}(1:6)) == (1,1)
@test first(MultiCombinations{3}(1:6)) == (1,1,1)
@test first(MultiCombinations{4}(1:6)) == (1,1,1,1)
@test first(MultiCombinations{5}(1:6)) == (1,1,1,1,1)
@test first(MultiCombinations{6}(1:6)) == (1,1,1,1,1,1)

# TODO test more complex iterators

@test length(Combinations{1}(1:1)) == 1 == length(collect(Combinations{1}(1:1)))
@test length(Combinations{1}(1:2)) == 2 == length(collect(Combinations{1}(1:2)))
@test length(Combinations{1}(1:3)) == 3 == length(collect(Combinations{1}(1:3)))

@test length(Combinations{2}(1:2)) == 1 == length(collect(Combinations{2}(1:2)))
@test length(Combinations{2}(1:3)) == 3 == length(collect(Combinations{2}(1:3)))
@test length(Combinations{2}(1:4)) == 6 == length(collect(Combinations{2}(1:4)))

@test length(Combinations{3}(1:3)) == 1 == length(collect(Combinations{3}(1:3)))
@test length(Combinations{3}(1:4)) == 4 == length(collect(Combinations{3}(1:4)))
@test length(Combinations{3}(1:5)) == 10 == length(collect(Combinations{3}(1:5)))

@test length(MultiCombinations{1}(1:1)) == 1  == length(collect(MultiCombinations{1}(1:1)))
@test length(MultiCombinations{1}(1:2)) == 2  == length(collect(MultiCombinations{1}(1:2)))
@test length(MultiCombinations{1}(1:3)) == 3  == length(collect(MultiCombinations{1}(1:3)))

@test length(MultiCombinations{2}(1:1)) == 1  == length(collect(MultiCombinations{2}(1:1)))
@test length(MultiCombinations{2}(1:2)) == 3  == length(collect(MultiCombinations{2}(1:2)))
@test length(MultiCombinations{2}(1:3)) == 6  == length(collect(MultiCombinations{2}(1:3)))

@test length(MultiCombinations{3}(1:1)) == 1  == length(collect(MultiCombinations{3}(1:1)))
@test length(MultiCombinations{3}(1:2)) == 4  == length(collect(MultiCombinations{3}(1:2)))
@test length(MultiCombinations{3}(1:3)) == 10 == length(collect(MultiCombinations{3}(1:3)))


struct Ship
    x::Int64
    y::Int64
    waypoint::NamedTuple{(:rx, :ry), Tuple{Int64, Int64}}
end

function parse_instructions(lines)
    instructions = NamedTuple{(:action, :value), Tuple{Char, Int64}}[]
    for line in lines
        a = line[1]
        v = parse(Int64, line[2:end])
        push!(instructions, (action=a, value=v))
    end
    instructions
end

function rotate_waypoint(ship, value)
    deg = value < 0 ? value + 360 : value
    if deg == 0
        ship
    elseif deg == 90
        Ship(ship.x, ship.y, (rx=-ship.waypoint.ry, ry=ship.waypoint.rx))
    elseif deg == 180
        Ship(ship.x, ship.y, (rx=-ship.waypoint.rx, ry=-ship.waypoint.ry))
    elseif deg == 270
        Ship(ship.x, ship.y, (rx=ship.waypoint.ry, ry=-ship.waypoint.rx))
    end
end

function move(ship, instruction)
    if instruction.action == 'N'
        Ship(ship.x, ship.y, (rx=ship.waypoint.rx, ry=ship.waypoint.ry + instruction.value))
    elseif instruction.action == 'S'
        Ship(ship.x, ship.y, (rx=ship.waypoint.rx, ry=ship.waypoint.ry - instruction.value))
    elseif instruction.action == 'E'
        Ship(ship.x, ship.y, (rx=ship.waypoint.rx + instruction.value, ry=ship.waypoint.ry))
    elseif instruction.action == 'W'
        Ship(ship.x, ship.y, (rx=ship.waypoint.rx - instruction.value, ry=ship.waypoint.ry))
    elseif instruction.action == 'L'
        rotate_waypoint(ship, instruction.value)
    elseif instruction.action == 'R'
        rotate_waypoint(ship, -instruction.value)
    elseif instruction.action == 'F'
        x = ship.x + instruction.value * ship.waypoint.rx
        y = ship.y + instruction.value * ship.waypoint.ry
        Ship(x, y, (rx=ship.waypoint.rx, ry=ship.waypoint.ry))
    end
end

function navigate(instructions)
    initial = Ship(0, 0, (rx=10, ry=1))
    ship = foldl(move, instructions; init=initial)
    abs(ship.x - initial.x) + abs(ship.y - initial.y)
end

function main()
    lines = readlines("../resources/input_12.txt")
    instructions = parse_instructions(lines)
    result = navigate(instructions)
    println(result)
end

main()


using ReachabilityAnalysis, Plots

@taylorize function laubloomis!(dx, x, params, t)
    dx[1] = 1.4*x[3] - 0.9*x[1]
    dx[2] = 2.5*x[5] - 1.5*x[2]
    dx[3] = 0.6*x[7] - 0.8*(x[2]*x[3])
    dx[4] = 2 - 1.3*(x[3]*x[4])
    dx[5] = 0.7*x[1] - (x[4]*x[5])
    dx[6] = 0.3*x[1] - 3.1*x[6]
    dx[7] = 1.8*x[6] - 1.6*(x[2]*x[7])
    return dx
end

function laubloomis(; W=0.01)
    # initial states
    X0c = [1.2, 1.05, 1.5, 2.4, 1.0, 0.1, 0.45]
    X0 = Hyperrectangle(X0c, fill(W, 7))

    # initial-value problem
    prob = @ivp(x' = laubloomis!(x), dim: 7, x(0) ∈ X0)

    return prob
end

prob = laubloomis(W=0.01)
alg = TMJets(abstol=1e-11, orderT=7, orderQ=1, adaptive=true);

sol_1 = solve(prob, T=20.0, alg=alg);

sol_1z = overapproximate(sol_1, Zonotope);

# canonical direction along x₄
const e4 = [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0];

ρ(e4, sol_1z)

ρ(e4, sol_1z) < 4.5

@show ρ(e4, sol_1z[end]) + ρ(-e4, sol_1z[end])

prob = laubloomis(W=0.05)
alg = TMJets(abstol=1e-12, orderT=7, orderQ=1, adaptive=false);

sol_2 = solve(prob, T=20.0, alg=alg)
sol_2z = overapproximate(sol_2, Zonotope);

ρ(e4, sol_2z)

ρ(e4, sol_2z) < 4.5

# width of final box
@show ρ(e4, sol_2z[end]) + ρ(-e4, sol_2z[end])

prob = laubloomis(W=0.1)
alg = TMJets(abstol=1e-12, orderT=7, orderQ=1, adaptive=false);

sol_3 = solve(prob, T=20.0, alg=alg)
sol_3z = overapproximate(sol_3, Zonotope);

ρ(e4, sol_3z)

ρ(e4, sol_3z) < 5.0

# width of final box
@show ρ(e4, sol_3z[end]) + ρ(-e4, sol_3z[end])

fig = plot()
plot!(fig, sol_3z, vars=(0, 4), linecolor="green", color=:green, alpha=0.8)
plot!(fig, sol_2z, vars=(0, 4), linecolor="blue",  color=:blue, alpha=0.8)
plot!(fig, sol_1z, vars=(0, 4), linecolor="yellow", color=:yellow, alpha=0.8,
      tickfont=font(10, "Times"), guidefontsize=15,
      xlab=L"t",
      ylab=L"x_4",
      xtick=[0., 5., 10., 15., 20.], ytick=[1.5, 2., 2.5, 3., 3.5, 4., 4.5, 5.],
      xlims=(0., 20.), ylims=(1.5, 5.02),
      bottom_margin=6mm, left_margin=2mm, right_margin=4mm, top_margin=3mm,
      size=(600, 600))

plot!(fig, x->x, x->4.5, 0., 20., line=2, color="red", linestyle=:dash, legend=nothing)
plot!(fig, x->x, x->5., 0., 20., line=2, color="red", linestyle=:dash, legend=nothing)


#==========================================================================================#

# TYPE

mutable struct IV <: GMM

    method::String
    sample::Microdata
    β::Vector{Float64}
    V::Matrix{Float64}
    W::Matrix{Float64}

    IV() = new()
end

#==========================================================================================#

# CONSTRUCTOR

function IV(MD::Microdata, method::String)
    obj        = IV()
    obj.sample = MD
    obj.method = method
    return obj
end

#==========================================================================================#

# INTERFACE

function fit(::Type{IV}, MD::Microdata; novar::Bool = false, method::String = "TSLS")

    if method == "OLS"

        FSD                   = Microdata(MD)
        FSD.mapping[:control] = vcat(MD.mapping[:treatment], MD.mapping[:control])

        pop!(FSD.mapping, :treatment)
        pop!(FSD.mapping, :instrument)

        obj = OLS(FSD)

        _fit!(obj, getweights(obj))

    elseif method == "Reduced form"

        FSD                   = Microdata(MD)
        FSD.mapping[:control] = vcat(MD.mapping[:instrument], MD.mapping[:control])

        pop!(FSD.mapping, :treatment)
        pop!(FSD.mapping, :instrument)

        obj = OLS(FSD)

        _fit!(obj, getweights(obj))

    elseif method == "First stage"

        FSD                    = Microdata(MD)
        FSD.mapping[:control]  = vcat(MD.mapping[:instrument], MD.mapping[:control])
        FSD.mapping[:response] = MD.mapping[:treatment]

        pop!(FSD.mapping, :treatment)
        pop!(FSD.mapping, :instrument)

        obj = OLS(FSD)

        _fit!(obj, getweights(obj))

    elseif length(MD.mapping[:treatment]) == length(MD.mapping[:instrument])

        obj   = IV(MD, "Method of moments")
        k     = length(MD.mapping[:instrument]) + length(MD.mapping[:control])
        obj.W = Matrix{Float64}(I, k, k)

        _fit!(obj, getweights(obj))

    elseif (method == "TSLS") | (method == "2SLS")

        obj   = IV(MD, "Two-step GMM")
        obj.W = crossprod(getmatrix(obj, :instrument, :control), getweights(obj))

        _fit!(obj, getweights(obj))

    elseif (method == "Two-step GMM") | (method == "Optimal GMM")

        obj   = IV(MD, method) ; _fit!(obj, getweights(obj))
        obj.W = wmatrix(obj, getcorr(obj), getweights(obj))

        _fit!(obj, obj.W, getweights(obj))

    else
        throw("unknown method")
    end

    novar || _vcov!(obj, getcorr(obj), getweights(obj))

    return obj
end

#==========================================================================================#

# ESTIMATION

function _fit!(obj::IV, ::UnitWeights)

    y = getvector(obj, :response)
    x = getmatrix(obj, :treatment, :control)
    z = getmatrix(obj, :instrument, :control)

    if obj.method == "Method of moments"
        obj.β = (z' * x) \ (z' * y)
    else
        zγ    = z * (z \ x)
        obj.β = (zγ' * x) \ (zγ' * y)
    end
end

function _fit!(obj::IV, W::Matrix{Float64}, ::UnitWeights)

    y = getvector(obj, :response)
    x = getmatrix(obj, :treatment, :control)
    z = getmatrix(obj, :instrument, :control)

    zγ    = z * (W \ (z' * x))
    obj.β = (zγ' * x) \ (zγ' * y)
end

function _fit!(obj::IV, w::AbstractWeights)

    y = getvector(obj, :response)
    x = getmatrix(obj, :treatment, :control)
    z = getmatrix(obj, :instrument, :control)
    v = Diagonal(w) * z

    if obj.method == "Method of moments"
        obj.β = (v' * x) \ (v' * y)
    else
        γ     = (v' * z) \ (v' * x)
        vγ    = v * γ
        obj.β = (vγ' * x) \ (vγ' * y)
    end
end

function _fit!(obj::IV, W::Matrix{Float64}, w::AbstractWeights)

    y = getvector(obj, :response)
    x = getmatrix(obj, :treatment, :control)
    z = getmatrix(obj, :instrument, :control)
    v = Diagonal(w) * z

    vγ    = v * (W \ (v' * x))
    obj.β = (vγ' * x) \ (vγ' * y)
end

#==========================================================================================#

# SCORE (MOMENT CONDITIONS)

score(obj::IV) = Diagonal(residuals(obj)) * getmatrix(obj, :instrument, :control)

# EXPECTED JACOBIAN OF SCORE × NUMBER OF OBSERVATIONS

function jacobian(obj::IV, ::UnitWeights)
    x = getmatrix(obj, :treatment, :control)
    z = getmatrix(obj, :instrument, :control)
    return - z' * x
end

function jacobian(obj::IV, w::AbstractWeights)
    x = getmatrix(obj, :treatment, :control)
    z = getmatrix(obj, :instrument, :control)
    v = Diagonal(w) * z
    return - v' * x
end

# VARIANCE MATRIX

function _vcov!(obj::IV, corr::Homoscedastic, ::UnitWeights)

    x  = getmatrix(obj, :treatment, :control)
    z  = getmatrix(obj, :instrument, :control)
    r  = residuals(obj)
    σ² = sum(abs2, r) / dof_residual(obj)
    γ  = z \ x
    V  = x' * z * γ

    obj.V = lmul!(σ², inv(V))
end

function _vcov!(obj::IV, corr::Homoscedastic, w::AbstractWeights)

    x  = getmatrix(obj, :treatment, :control)
    z  = getmatrix(obj, :instrument, :control)
    v  = Diagonal(w) * z
    r  = residuals(obj)
    σ² = sum(abs2.(r), w) / dof_residual(obj)
    γ  = (v' * z) \ (v' * x)
    V  = x' * v * γ

    obj.V = lmul!(σ², inv(V))
end

#==========================================================================================#

# LINEAR PREDICTOR

function predict(obj::IV, MD::Microdata)
    if getnames(obj, :treatment, :control) != getnames(MD, :treatment, :control)
        throw("missing variables")
    end
    getmatrix(MD, :treatment, :control) * obj.β
end

# FITTED VALUES

fitted(obj::IV, MD::Microdata) = predict(obj, MD)

# DERIVATIVE OF FITTED VALUES

jacobexp(obj::IV) = getmatrix(obj, :treatment, :control)

#==========================================================================================#

# UTILITIES

coefnames(obj::IV) = getnames(obj, :treatment, :control)
mtitle(obj::IV)    =  "Linear GMM"

# COEFFICIENT OF DETERMINATION

adjr2(obj::IV) = 1.0 - (1.0 - r2(obj)) * (nobs(obj) - 1) / dof_residual(obj)
r2(obj::IV)    = _r2(obj, getweights(obj))

function _r2(obj::IV, ::UnitWeights)
    y   = response(obj)
    ŷ   = fitted(obj)
    rss = sum(abs2, y .- ŷ)
    tss = sum(abs2, y .- mean(y))
    return 1.0 - rss / tss
end

function _r2(obj::IV, w::AbstractWeights)
    y   = response(obj)
    μ   = mean(y, w)
    ŷ   = fitted(obj)
    rss = sum(abs2.(y .- ŷ), w)
    tss = sum(abs2.(y .- μ), w)
    return 1.0 - rss / tss
end


using FactorPostProcessing, Test, Random
using BeastUtils.Logs, LinearAlgebra

Random.seed!(666)

const TOL = 1e-10
const POS_COLS = [9, 5, 13]

function absmax(x::AbstractArray{T}, y::AbstractArray{T}) where T <: Real
    return maximum(abs.(x - y))
end

function svd_process(L_data::Matrix{Float64}, F_data::Matrix{Float64}, row::Int, k::Int, p::Int, n::Int)
    L_vec = @view L_data[row, :]
    L = Matrix(reshape(L_vec, p, k))
    Lsvd = svd(L)
    F_vec = @view F_data[row, :]
    F = reshape(F_vec, k, n)
    F_rotated = Lsvd.Vt * F
    S = Lsvd.S
    Vt = Lsvd.Vt
    U = Lsvd.U
    F_rotated = Vt * F
    return (sv = Lsvd.S, V = Lsvd.U, F_rotated)
end

states = 1000
K = 3
P = 20
N = 100

L_data = randn(states, K * P)
F_data = randn(states, K * N)

L_labels = vec(["L$j$i" for i = 1:K, j=1:P])
F_labels = vec(["factors.taxon_$i.$j" for j=1:K, i = 1:N])

cols = [["posterior", "prior", "likelihood"];
        ["factorPrecision$i" for i = 1:P]; L_labels; F_labels]
data = [randn(states, 3 + P) L_data F_data]

log_path = joinpath(@__DIR__, "testFactorPostProcessing.log")
svd_path = joinpath(@__DIR__, "testFactorPostProcessingSVD.log")

@assert !isfile(log_path)
@assert !isfile(svd_path)

Logs.make_log(log_path, data, cols)
svd_logs(log_path, svd_path, rotate_factors=true)

svd_cols, svd_data = get_log(svd_path, burnin=0.0)

rm(log_path)
rm(svd_path)

L_cols = findall(startswith("L"), svd_cols)
V_cols = findall(startswith("V"), svd_cols)
sv_cols = findall(startswith("sv"), svd_cols)
F_cols = findall(startswith("factors"), svd_cols)

for state = 1:states
    L_actual = reshape(svd_data[state, L_cols], P, K)
    V_actual = reshape(svd_data[state, V_cols], P, K)
    sv_actual = svd_data[state, sv_cols]
    F_actual = reshape(svd_data[state, F_cols], K, N)
    S_expected, V_expected, F_expected = svd_process(L_data, F_data, state, K, P, N)
    L_expected = V_expected * Diagonal(S_expected)
    @test absmax(sv_actual, S_expected) < TOL

    signs = ones(K)
    for k = 1:K
        if L_expected[POS_COLS[k], k] < 0.0
            signs[k] = -1.0
        end
    end

    D = Diagonal(signs)
    L_expected = L_expected * D
    V_expected = V_expected * D
    F_expected = D * F_expected

    @test absmax(V_actual, V_expected) < TOL
    @test absmax(L_actual, L_expected) < TOL
    @test absmax(F_actual, F_expected) < TOL
end


using DiffEqSensitivity, OrdinaryDiffEq, DiffEqCallbacks, Flux
using Random, Test
using Zygote

function test_hybridNODE(sensealg)
    Random.seed!(12345)
    datalength = 100
    tspan = (0.0,100.0)
    t = range(tspan[1],tspan[2],length=datalength)
    target = 3.0*(1:datalength)./datalength  # some dummy data to fit to
    cbinput = rand(1, datalength) #some external ODE contribution
    pmodel = Chain(
        Dense(2, 10, init=zeros),
        Dense(10, 2, init=zeros))
    p, re = Flux.destructure(pmodel)
    dudt(u,p,t) = re(p)(u)

    # callback changes the first component of the solution every time
    # t is an integer
    function affect!(integrator, cbinput)
        event_index = round(Int,integrator.t)
        integrator.u[1] += 0.2*cbinput[event_index]
    end
    callback = PresetTimeCallback(collect(1:datalength),(int)->affect!(int, cbinput))

    # ODE with Callback
    prob = ODEProblem(dudt,[0.0, 1.0],tspan,p)

    function predict_n_ode(p)
        arr = Array(solve(prob, Tsit5(),
            p=p, sensealg=sensealg, saveat=2.0, callback=callback))[1,2:2:end]
        return arr[1:datalength]
    end

    function loss_n_ode()
        pred = predict_n_ode(p)
        loss = sum(abs2,target .- pred)./datalength
    end

    cb = function () #callback function to observe training
        pred = predict_n_ode(p)
        display(loss_n_ode())
    end
    @show sensealg
    Flux.train!(loss_n_ode, Flux.params(p), Iterators.repeated((), 20), ADAM(0.005), cb = cb)
    @test loss_n_ode() < 0.5
    println("  ")
end

function test_hybridNODE2(sensealg)
    Random.seed!(12345)
    u0 = Float32[2.; 0.; 0.; 0.]
    tspan = (0f0,1f0)

    ## Get goal data
    function trueaffect!(integrator)
        integrator.u[3:4] = -3*integrator.u[1:2]
    end
    function trueODEfunc(dx,x,p,t)
        @views dx[1:2] .= x[1:2] + x[3:4]
        dx[1] += x[2]
        dx[2] += x[1]
        dx[3:4] .= 0f0
    end
    cb_ = PeriodicCallback(trueaffect!,0.1f0,save_positions=(true,true),initial_affect=true)
    prob = ODEProblem(trueODEfunc,u0,tspan)
    sol = solve(prob,Tsit5(),callback=cb_,save_everystep=false,save_start=true)
    ode_data = Array(sol)[1:2,1:end]'

    ## Make model
    dudt2 = Chain(Dense(4,50,tanh),
                    Dense(50,2))
    p,re = Flux.destructure(dudt2) # use this p as the initial condition!
    function affect!(integrator)
        integrator.u[3:4] = -3*integrator.u[1:2]
    end
    function ODEfunc(dx,x,p,t)
        dx[1:2] .= re(p)(x)
        dx[3:4] .= 0f0
    end
    z0 = u0
    prob = ODEProblem(ODEfunc,z0,tspan)
    cb = PeriodicCallback(affect!,0.1f0,save_positions=(true,true),initial_affect=true)

    ## Initialize learning functions
    function predict_n_ode()
        _prob = remake(prob,p=p)
        Array(solve(_prob,Tsit5(),u0=z0,p=p,callback=cb,save_everystep=false,save_start=true,sensealg=sensealg))[1:2,:]
    end
    function loss_n_ode()
        pred = predict_n_ode()[1:2,1:end]'
        loss = sum(abs2,ode_data .- pred)
        loss
    end
    loss_n_ode() # n_ode.p stores the initial parameters of the neural ODE
    cba = function ()  #callback function to observe training
        pred = predict_n_ode()[1:2,1:end]'
        display(sum(abs2,ode_data .- pred))
        return false
    end
    cba()

    ## Learn
    ps = Flux.params(p)
    data = Iterators.repeated((), 25)

    @show sensealg

    Flux.train!(loss_n_ode, ps, data, ADAM(0.0025), cb = cba)

    @test loss_n_ode() < 0.5

    println("  ")
end

mutable struct Affect{T}
      callback_data::T
end
compute_index(t) = round(Int,t)+1
function (cb::Affect)(integrator)
    indx = compute_index(integrator.t)
    integrator.u .= integrator.u .+ @view(cb.callback_data[:, indx, 1]) * (integrator.t - integrator.tprev)
end
function test_hybridNODE3(sensealg)
    u0 = Float32[2.; 0.]
    datasize = 101
    tspan = (0.0f0,10.0f0)

    function trueODEfunc(du,u,p,t)
        du .= -u
    end
    t = range(tspan[1],tspan[2],length=datasize)
    prob = ODEProblem(trueODEfunc,u0,tspan)
    ode_data = Array(solve(prob,Tsit5(),saveat=t))

    true_data = reshape(ode_data,(2,length(t),1))
    true_data = convert.(Float32,true_data)
    callback_data = true_data * 1f-3
    train_dataloader = Flux.Data.DataLoader((true_data = true_data,callback_data = callback_data),batchsize=1)
    dudt2 = Chain(Dense(2,50,tanh),
                 Dense(50,2))
    p,re = Flux.destructure(dudt2)
    function dudt(du,u,p,t)
        du .= re(p)(u)
    end
    z0 = Float32[2.; 0.]
    prob = ODEProblem(dudt,z0,tspan)

    function callback_(callback_data)
        affect! = Affect(callback_data)
        condition(u,t,integrator) = integrator.t > 0
        DiscreteCallback(condition,affect!,save_positions=(false,false))
    end

    function predict_n_ode(true_data_0,callback_data, sense)
        _prob = remake(prob,p=p,u0=true_data_0)
        solve(_prob,Tsit5(),callback=callback_(callback_data),saveat=t,sensealg=sense)
    end

    function loss_n_ode(true_data,callback_data)
        sol = predict_n_ode((vec(true_data[:,1,:])),callback_data,sensealg)
        pred = Array(sol)
        loss = Flux.mse((true_data[:,:,1]),pred)
        loss
    end

    ps = Flux.params(p)
    opt = ADAM(0.1)
    epochs = 10
    function cb1(true_data,callback_data)
        display(loss_n_ode(true_data,callback_data))
        return false
    end

    function train!(loss, ps, data, opt, cb)
        ps = Params(ps)
        for (true_data,callback_data) in data
            gs = gradient(ps) do
                loss(true_data,callback_data)
            end
            Flux.update!(opt, ps, gs)
            cb(true_data,callback_data)
        end
        return nothing
    end

    @Flux.epochs epochs train!(loss_n_ode, Params(ps),train_dataloader, opt, cb1)
    loss =  loss_n_ode(true_data[:,:,1],callback_data)
    @info loss
    @test loss < 0.5
end

@testset "PresetTimeCallback" begin
    test_hybridNODE(ForwardDiffSensitivity())
    test_hybridNODE(BacksolveAdjoint())
    test_hybridNODE(InterpolatingAdjoint())
    test_hybridNODE(QuadratureAdjoint())
end

@testset "PeriodicCallback" begin
    test_hybridNODE2(ReverseDiffAdjoint())
    test_hybridNODE2(BacksolveAdjoint())
    test_hybridNODE2(InterpolatingAdjoint())
    test_hybridNODE2(QuadratureAdjoint())
end

@testset "tprevCallback" begin
    test_hybridNODE3(ReverseDiffAdjoint())
    test_hybridNODE3(BacksolveAdjoint())
    test_hybridNODE3(InterpolatingAdjoint())
    test_hybridNODE3(QuadratureAdjoint())
end


### A Pluto.jl notebook ###
# v0.14.2

using Markdown
using InteractiveUtils

# ╔═╡ 7a0128c0-a13f-11eb-1098-ef9c00274a0d
using Plots

# ╔═╡ 7c5e40c4-f84e-4527-a87e-c6ac395f5d9e
function ansatz(x,α)
	return ifelse(abs(x) > 1,1 - abs(x)^α,0)
end;

# ╔═╡ 007bb513-b173-4723-9d8b-2e858ae54dc6
function probability_density(x,α)
	return (ansatz(x,α)^2)
end;

# ╔═╡ dcc6bdb1-daa8-44f1-a6ec-285ac1916179
function E_local(x,α)
	return ifelse(abs(x)<1,(0.5*α*(α - 1)*abs(x)^(α - 2))/(1 - abs(x)^α),0)
end;

# ╔═╡ 99eff10f-5b1e-42d9-9009-56e5aeb821c2
function metropolis(N,α)
	L = 2
	x = rand(-1:0.001:1)
	E = 0
	E2 = 0
	Eln_avg = 0
	ln_avg = 0
	rejection_ratio = 0
	#start the loop
	for i in 1:N
		x_t = x + 0.01*rand(-L:0.000001:L)
		if probability_density(x_t,α) >= probability_density(x,α)
			x = x_t
		else
			dummy = rand()
			if dummy < probability_density(x_t,α)/probability_density(x,α)
				x = x_t
			else
				rejection_ratio += 1/N
			end
		end

		E += E_local(x,α)/N
		E2 += E_local(x,α)^2/N
		Eln_avg += (E_local(x,α)*((-1*log(abs(x))*abs(x)^α)/(1 - abs(x)^α)))/(N)
		ln_avg += ((-1*log(abs(x))*abs(x)^α)/(1 - abs(x)^α))/N
	end
	return E,E2,Eln_avg,ln_avg,rejection_ratio
end;

# ╔═╡ cbfdf51f-721e-4172-b7d7-1911bf666163
function iterate_alpha()
	α = 1.1
	alpha_iter = 50
	N = 500
	random_walkers = 200
	γ = 0.01
	energy = Vector{Float64}(undef,alpha_iter)
	alpha = Vector{Float64}(undef,alpha_iter)
	variance = Vector{Float64}(undef,alpha_iter)
	for i in 1:alpha_iter
		E = 0
		E2 = 0
		dE_dalpha = 0
		Eln = 0
		ln = 0
		rejection_ratio = 0
		for j in 1:random_walkers
			E_met, E2_met, Eln_met, ln_met, _ = metropolis(N,α)
			E += E_met/random_walkers
			E2 += E2_met/random_walkers
			Eln += Eln_met/random_walkers
			ln += ln_met/random_walkers
			# rejection_ratio += rejections_met/random_walkers
		end

		# Define the next α
		dE_dalpha = 2*(Eln-E*ln)
		α = α + 0.05#γ*dE_dalpha

		# Update the arrays for plotting
		energy[i] = E
		alpha[i] = α
		variance[i] = E2 - E^2
	end
	return energy,alpha,variance
end;

# ╔═╡ 20151696-5be9-4b68-afb7-8a7217c16402
energy,alpha,variance = iterate_alpha();

# ╔═╡ e5f66da9-ef79-4f82-9e1b-6ad7b27ec75d
energy

# ╔═╡ 64f872bd-cb8b-42de-8fa3-3270430e9a7f
alpha

# ╔═╡ 5ca1aa2b-099e-479c-88ad-7c9764fa6f85
begin
	plot(energy,
		label = false,
		title = "Energy expenctation Plot",
		linecolor = :blue,
		ylabel = "<E>",
		xlabel = "timestep"
	)

	hline!([π^2/8],
		label = false
	)
end

# ╔═╡ b21c39ff-68f7-4725-9bdc-4c180b5b368b
plot(variance,
	title = "Variance in energy",
	label = false,
	linecolor = :blue,
	ylabel = "Var(E)",
	xlabel = "timestep"
)

# ╔═╡ bb6cfb6d-325e-4913-9f77-3efcfcfc4368
plot(alpha,
	title = "Variational Parameter",
	label = false,
	linecolor = :blue,
	ylabel = "α",
	xlabel = "timestep"
)

# ╔═╡ 11dd75f3-e142-4712-914d-d8d594acada3
begin
	scatter(alpha,
		energy,
		title = "Energy vs α",
		xlabel = "α",
		ylabel = "<E>",
		label = false,
		markercolor = :blue
	)

	hline!([π^2/8],
		label = false
	)
end

# ╔═╡ 43a28afe-622c-4cfc-9679-9554994be6bf
variance

# ╔═╡ Cell order:
# ╠═7a0128c0-a13f-11eb-1098-ef9c00274a0d
# ╠═7c5e40c4-f84e-4527-a87e-c6ac395f5d9e
# ╠═007bb513-b173-4723-9d8b-2e858ae54dc6
# ╠═dcc6bdb1-daa8-44f1-a6ec-285ac1916179
# ╠═99eff10f-5b1e-42d9-9009-56e5aeb821c2
# ╠═cbfdf51f-721e-4172-b7d7-1911bf666163
# ╠═20151696-5be9-4b68-afb7-8a7217c16402
# ╠═e5f66da9-ef79-4f82-9e1b-6ad7b27ec75d
# ╠═64f872bd-cb8b-42de-8fa3-3270430e9a7f
# ╠═5ca1aa2b-099e-479c-88ad-7c9764fa6f85
# ╠═b21c39ff-68f7-4725-9bdc-4c180b5b368b
# ╠═bb6cfb6d-325e-4913-9f77-3efcfcfc4368
# ╠═11dd75f3-e142-4712-914d-d8d594acada3
# ╠═43a28afe-622c-4cfc-9679-9554994be6bf


export dmrg

function dmrg(H::MPO,
              psi0::MPS,
              sweeps::Sweeps;
              kwargs...)::Tuple{Float64,MPS}

  which_factorization::String = get(kwargs,:which_factorization,"automatic")

  psi = copy(psi0)
  N = length(psi)

  PH = ProjMPO(H)
  position!(PH,psi0,1)
  energy = 0.0

  for sw=1:nsweep(sweeps)
    sw_time = @elapsed begin

    @debug begin
      reset!(timer)
      println("\n\n~~~ This is sweep $sw ~~~\n\n")
    end

    for (b,ha) in sweepnext(N)
      position!(PH,psi,b)

      phi = psi[b]*psi[b+1]

      energy,phi = davidson(PH,phi;kwargs...)

      dir = ha==1 ? "fromleft" : "fromright"
      replaceBond!(psi,b,phi;
                   maxdim=maxdim(sweeps,sw),
                   mindim=mindim(sweeps,sw),
                   cutoff=cutoff(sweeps,sw),
                   dir=dir,
                   which_factorization=which_factorization)
    end
    end
    @printf "After sweep %d energy=%.12f maxDim=%d time=%.3f\n" sw energy maxDim(psi) sw_time
    @debug printTimes(timer)
  end
  return (energy,psi)
end

@doc """
Use the density matrix renormalization group (DMRG) algorithm to
optimize a matrix product state (MPS) to be the eigenvector
of the Hermitian matrix product operator (MPO) H with minimal
eigenvalue.

Inputs:
* `H::MPO` - a Hermitian MPO
* `psi0::MPS` - MPS used to initialize the optimization
* `sweeps::Sweeps` - `Sweeps` object used to control the algorithm

Returns:
* `energy::Float64` - eigenvalue of the optimized MPS
* `psi::MPS` - optimized MPS

""" dmrg


#= addition.jl
Philip Barrett, pobarrett@gmail.com
30may2016, Chicago

Tests the addition functions for the polygon class
=#

## Binary set sum
Z = [ 1.0 -1; 1 1; -1 1; -1 -1 ]
a = Polygon( pts=Z )
h = Polygon( pts=[ 0.0 1 ; 1 0 ; -1 0 ; 0 -1 ] )
NN = 40
exactdirs = vcat([ [ cos(2*(i-1)*pi/NN) sin(2*(i-1)*pi/NN) ] for i in 1:NN ]...)
    # The union of the face directions.  Will give an exact result.
exact = setSum( a, h, exactdirs )
exactinner = setSum( a, h, exactdirs, false )

@test deeDoop(exact.pts) ≈ [ -2 1; -2.0 -1; -1 -2; 1 -2; 2.0 -1; 2 1; 1 2; -1 2 ]
@test deeDoop(exact.pts) ≈ [ exactinner.pts[ end, :]; exactinner.pts[ 1:(end-1), :] ]
    # Arises because we have a vertical left edge.  Annoying.

## Array set sum
exactarray = setSum( [a, h], exactdirs )
@test exact.pts == exactarray.pts


using Unitful
using UnitfulEquivalences
using PhysicalConstants.CODATA2018
using Interpolations

import Unitful:
    nm, μm, mm, cm, m, km,
    mg, g, kg,
    ps, ns, μs, ms, s, minute, hr, d, yr, Hz, kHz, MHz, GHz,
    eV,
    μJ, mJ, J,
	μW, mW, W,
	A, N, mol, mmol, V, L, M

	import PhysicalConstants.CODATA2018: N_A

"""
	struct Fov

Represent a field of view

# Fields
- `d::Unitful.Length`  : diameter of Fov
- `z::Unitful.Length`  : thickness
- `a::Unitful.Area`    : area (computed)
- `v::Unitful.Volume`  : volume (computed)

"""
struct Fov
    d::Unitful.Length
    z::Unitful.Length
	a::Unitful.Area
    v::Unitful.Volume

	function Fov(d,z)
		a = π * (d/2.)^2
		v = a * z
		new(d,z,a,v)
	end
end

"""
	struct Fluorophore

Represent a fluorescent molecule

# Fields
- `expeak::Units (nm)`           : peak excitation
- `empeak::Units (nm)`           : peak emission
- `ϵ::Units(``cm^{-1} M^{-1}``)` : molar extinction coefficient
- `Q::Float64`                   : Quantum efficiency
- `σ::Units(``cm^2``)`           : attenuation cross section
"""
struct Fluorophore
    expeak::typeof(1.0nm)
    empeak::typeof(1.0nm)
    ϵ::typeof(1.0/(cm*M))
    Q::Float64
	σ::typeof(1.0cm^2)
	function Fluorophore(ex, en, ϵ, Q)
		σ = log(10) * uconvert(cm^2/mol, ϵ) / N_A
		new(ex, en, ϵ, Q, σ)
	end
end

"""
	struct Monolayer

Represents a monolayer

# Fields
- `p::Unitful.Length`  : pitch between molecules
- `σ::Units (``molecules/μm^2``)`  : number of molecules per area
"""
struct Monolayer
	p::Unitful.Length
    σ::typeof(1.0*μm^-2)

	function Monolayer(p)
		σ = uconvert(μm^-2, 1.0/p^2)
		new(p, σ)

	end
end

abstract type Laser end

"""
	struct CLaser <: Laser

Simple representation of a continous laser

# Fields
- `λ::typeof(1.0nm)`  : Laser wavelength
- `P::typeof(1.0mW)`  : Power

"""
struct CLaser <: Laser
	λ::Unitful.Length
	P::Unitful.Power
end

"""
	struct PulsedLaser

Simple representation of a pulsed laser

# Fields
- `λ::Unitful.Length`  : Laser wavelength
- `Pk::Unitful.Power`  : peak Power
- `P::Unitful.Power`  : average power
- `f::Unitful.Frequency`  : laser frequency
- `w::Unitful.Time` : pulse width

"""
struct PulsedLaser  <: Laser
	λ::Unitful.Length
	Pk::Unitful.Power
	f::Unitful.Frequency
	P::Unitful.Power
	w::Unitful.Time

	function PulsedLaser(λ::Unitful.Length,
						 Pk::Unitful.Power,
						 f::Unitful.Frequency,
		                 w::Unitful.Time)

		P = uconvert(mW, Pk * w * f)
		new(λ, Pk, f, P,  w)
	end

	function PulsedLaser(λ::Unitful.Length,
						 Pk::Unitful.Power,
						 f::Unitful.Frequency,
						 P::Unitful.Power,
						 w::Unitful.Time)

			new(λ, Pk, f, P,  w)
	end
end

# """
# GaussianLaser

# Representation of a Gaussian laser defined by a laser and a waist

# # Fields
# - `laser::Laser`       : A Laser
# - `w0::Unitful.Length` : Waist of laser

# """
# struct GaussianLaser
# 	laser::Laser
# 	w0::Unitful.Length
# end

"""
	struct TLens

Represents a thin lens

# Fields
- `M::Real`         : Magnification
- `f::Unitful.Length`   : focal distance of the lens
- `d::Unitful.Length`   : diameter of the lens

"""
struct TLens
    M::Real
	f::Unitful.Length
	d::Unitful.Length

end

"""
	propagate(tl::TLens, y1::Real, z1::Real)

Propagate ray from (z1, y1) to (z2, y2) through thin lens tl

# Fields
- `tl::TLens`         : A thin lens
- `z1::Unitful.Length`   : point z1
- `y1::Unitful.Length`   : point y1

"""
function propagate(tl::TLens, y1::Unitful.Length, z1::Unitful.Length)
	y2 = tl.M * y1
	iz2 = (1.0/tl.f) - (1.0/z1)
	y2, 1.0/iz2
end

"""
	struct Objective

Simple representation of a microscope objective

# Fields
- `name::String`       : identifies the objective
- `NA::Float64`        : Numerical aperture
- `M::Float64`         : Magnification
- `f::typeof(1.0mm)`   : focal length
- `d::typeof(1.0mm)`   : entrance pupil diameter

"""
struct Objective
	name::String
    NA::Real
    M::Real
	f::Unitful.Length
	d::Unitful.Length

	function Objective(name::String, NA::Float64,  M::Float64)
		new(name, NA, M, -1.0mm, -1.0mm)
	end

	function Objective(name::String, f::Unitful.Length, d::Unitful.Length, M::Float64)
		ff = f / d   # https://www.eckop.com/resources/optics/numerical-aperture-and-f-number/
		NA = 1.0/(2.0*ff)
		new(name, NA, M, f,d)
	end
end

#CCD is defined in terms of a function which returns the CCD response
#(e.g, function ccd returns a response function, which gives efficiency
# as a function of wavelength)
"""
	ccd(lmin::Float64=350.0, lmax::Float64=1000.0)

Return the efficiency of a generic CCD as a function of wavelength.

# Fields

- `lmin::Float64=350.0` : Minimum wavelength for which efficiency is defined
- `lmax::Float64=350.0` : Maximum wavelength for which efficiency is defined

"""
function ccd(lmin::Float64=350.0, lmax::Float64=1000.0)
	function eff(l::Float64)
		if l < lmin || l > lmax
			return 0.
		else
			wl = 350.:50.:1000.
			ϵ = [0.3, 0.4,0.65,0.78,0.82,0.82,0.8,0.72,0.62,0.5,0.37,
			  0.24,0.12,0.07]
			e = CubicSplineInterpolation(wl, ϵ)
			return e(l)
		end
	end
	return eff
end

#FUNCTIONS

"""
	photon_energy(λ::Unitful.Length)

Given wavelength of photon return its energy.
# Fields

- `λ::Unitful.Length`  : Photon wavelength

"""
function photon_energy(λ::Unitful.Length)
	uconvert(eV, λ, Spectral())
end

"""
	delivered_energy(laser::Laser, t::Unitful.Time)

Delivered energy of a laser in a given time.
# Fields

- `laser::Laser`     : Laser
- `t::Unitful.Time`  : Time in which target is illuminated

"""
function delivered_energy(laser::Laser, t::Unitful.Time)
	laser.P * t
end

"""
	n_photons(laser::Laser)

Rate of photons (number of photons per unit time) produced by a laser
# Fields

- `laser::Laser`     : Laser
- `t::Unitful.Time`  : Time in which target is illuminated

"""
function n_photons(laser::Laser)
	uconvert(Hz, laser.P / photon_energy(laser.λ))
end

"""
	n_photons(λ::Unitful.Length, p::Unitful.Power)

Rate of photons (number of photons per unit time) corresponding to a wavelength
λ and a power P

# Fields

- `λ::Unitful.Length` : photon wavelength
- `p::Unitful.Power`  : Power

"""
function n_photons(λ::Unitful.Length, p::Unitful.Power)
	uconvert(Hz, p / photon_energy(λ))
end

"""
	n_photons_int(laser::Laser, t::Unitful.Time)

Integrated number of photons in a given time emitted by a laser

# Fields

- `laser::Laser`    : Laser
- `t::Unitful.Time` : time of measurement

"""
function n_photons_int(laser::Laser, t::Unitful.Time)
	uconvert(eV,delivered_energy(laser, t)) / photon_energy(laser.λ)
end

"""
	photon_density(λ::Unitful.Length, p::Unitful.Power, a::Unitful.Area)

number of photons per unit time per unit area

# Fields

- `λ::Unitful.Length` : photon wavelength
- `p::Unitful.Power`  : Power
- `a::Unitful.Area`   : Area

"""
function photon_density(λ::Unitful.Length, p::Unitful.Power, a::Unitful.Area)
	return n_photons(λ, p)/ a
end

"""
	photon_density(l::Laser, fov::Fov)

number of photons per unit time per unit area, in a Fov illuminated by a laser

# Fields

- `laser::Laser` : Laser
- `fov::Fov`     : Field of view

"""
function photon_density(laser::Laser, fov::Fov)
	return n_photons(laser) / fov.a
end

"""
	diffraction_limit(l::Laser, obj:: Objective)

Return the diameter of diffractive spot for a laser l
focused with an objective obj

# Fields

- `l::Laser`         : A laser
- `obj:: Objective`  : An objective
"""
function diffraction_limit(l::Laser, obj:: Objective)
    return 1.83 * l.λ/(2 * obj.NA)
end

"""
	geometrical_acceptance(d::Float64, D::Float64)

Compute the fraction of photons that make it through an iris
of diameter D located at a distance d from the emission point.

# Fields

- `d::Float64`   : distance between emission point and iris
- `D::Float64`   : Diameter of iris

"""
function geometrical_acceptance(d::Float64, D::Float64)
	return 0.5(1. - d/sqrt(d^2 + (D/2.)^2))
end

"""
	transmission(objective::Objective)

Compute the transmission of an objective (depends only of NA).

# Fields

- `objective::Objective` : Objective

"""
function transmission(objective::Objective)
	A = objective.NA
	A <=1 ? (1 - sqrt(1 - A^2)) /2 : 0.5
end

"""
	transmission(A::Float64)

Compute the transmission as a function of NA.

# Fields

- `A::Float64` : Numerical acceptance (NA)

"""
function transmission(A::Float64)
	A <=1 ? (1 - sqrt(1 - A^2)) /2 : 0.5
end

# fluorophores

"""
	fluorescence(f::Fluorophore, I::Quantity)

Number of  photons emitted per unit time when fluorosphore F
is illuminated with a laser of photon density I

# Fields

- `f::Fluorophore`    : Define fluorophore
- `I::(``nγ/cm^2``)`  : Photon density
"""
fluorescence(f::Fluorophore, I::typeof(1.0Hz/cm^2)) =  f.σ * f.Q * I


context("Distance Fields") do

context("SignedDistanceField") do
    s1 = SignedDistanceField(HyperRectangle(Vec(1,2),Vec(3,4)), [1 2;3 4])
    @fact typeof(s1) --> GeometryTypes.SignedDistanceField{2,Int,Int}
    s1 = SignedDistanceField(HyperRectangle(Vec(1,2),Vec(3,4)), [1. 2.;3. 4.])
    @fact typeof(s1) --> GeometryTypes.SignedDistanceField{2,Int,Float64}

    # functional
    s2 = SignedDistanceField(HyperRectangle(Vec(0,0,0.),Vec(1,1,1.))) do v
        sqrt(sum(v.*v)) - 1 # sphere
    end
    @fact size(s2) --> (11, 11, 11)
    # functional
    s2 = SignedDistanceField(HyperRectangle(Vec(-1,-1.),Vec(2,3.))) do v
        sqrt(sum(v.*v)) - 1 # circle
    end
    @fact size(s2) --> (21, 31)
    @fact size(s2, 1) --> 21
    @fact size(s2, 2) --> 31

    @fact HyperRectangle(s2) --> HyperRectangle(Vec(-1,-1.),Vec(2,3.))
end

context("getindex") do
    sdf = SignedDistanceField(HyperRectangle(Vec(-1,-1.),Vec(2,2.))) do v
        sqrt(sum(v.*v)) - 1 # circle
    end

    @fact sdf[1,1] --> sdf[21*21] # by circle symmetry
    @fact_throws BoundsError sdf[21*21+1]
end

end


using AeroMDAO
using Test

@testset "2D Panel Method - Doublet-Source Kulfan CST" begin
    # Define Kulfan CST coefficients
    alpha_u = [0.2, 0.3, 0.2, 0.15, 0.2]
    alpha_l = [-0.2, -0.1, -0.1, -0.001]
    dzs     = (0., 0.)

    # Define airfoil
    airfoil = (Foil ∘ kulfan_CST)(alpha_u, alpha_l, dzs, 0.0, 60);

    # Define uniform flow
    uniform = Uniform2D(1., 5.)

    # Evaluate case
    cl, cls, cms, cps, panels = solve_case(airfoil, uniform; num_panels = 80)

    @test cl       ≈  0.84188988 atol = 1e-6
    @test sum(cls) ≈  0.84073703 atol = 1e-6
    @test sum(cms) ≈ -0.26104277 atol = 1e-6
end

@testset "Geometry - Airfoil Processing" begin
    # Import and read airfoil coordinates
    foilpath = joinpath((dirname ∘ dirname ∘ pathof)(AeroMDAO), "test/CRM.dat")
    coords   = read_foil(foilpath)

    # Cosine spacing
    cos_foil = cosine_foil(coords, 51)

    # Split airfoil
    up, low  = split_foil(cos_foil)

    # Convert coordinates to Kulfan CST variables
    num_dv   = 4
    alpha_u, alpha_l = coords_to_CST(up, num_dv), coords_to_CST(low, num_dv)

    # Generate same airfoil using Kulfan CST parametrisation
    cst_foil = (Foil ∘ kulfan_CST)(alpha_u, alpha_l, (0., 0.), 0.0)

    uniform  = Uniform2D(1., 5.)
    cl, cls, cms, cps, panels = solve_case(cst_foil, uniform; num_panels = 80)

    @test cl       ≈  0.85736965 atol = 1e-6
    @test sum(cls) ≈  0.85976886 atol = 1e-6
    @test sum(cms) ≈ -0.29766116 atol = 1e-6
end

@testset "Geometry - Two-Section Trapezoidal Wing" begin
    # Define wing
    wing_right = HalfWing(chords    = [1.0, 0.6, 0.2],
                          twists    = [2.0, 0.0, -0.2],
                          spans     = [5.0, 0.5],
                          dihedrals = [5., 5.],
                          sweep_LEs = [5., 5.]);

    # Get wing info
    b, S, c, AR = info(wing_right)
    λ           = taper_ratio(wing_right)
    wing_mac    = mean_aerodynamic_center(wing_right)

    @test b        ≈ 5.50000000                    atol = 1e-6
    @test S        ≈ 4.19939047                    atol = 1e-6
    @test c        ≈ 0.79841269                    atol = 1e-6
    @test AR       ≈ 7.20342634                    atol = 1e-6
    @test λ        ≈ 0.20000000                    atol = 1e-6
    @test wing_mac ≈ [0.42092866, 1.33432539, 0.0] atol = 1e-6
end

@testset "Vortex Lattice Method - NACA 0012 Rectangular Wing" begin
    # Define wing
    wing = Wing(foils     = Foil.(naca4((0,0,1,2)) for i ∈ 1:2),
                chords    = [0.18, 0.16],
                twists    = [0., 0.],
                spans     = [0.5,],
                dihedrals = [5.],
                sweep_LEs = [1.14])

    # Define reference values
    ρ   = 1.225
    ref = [0.25 * mean_aerodynamic_chord(wing), 0., 0.]
    Ω   = [0.0, 0.0, 0.0]

    # Define freestream condition
    uniform = Freestream(10.0, 2.0, 2.0, Ω)

    # Evaluate stability case
    nf_coeffs, ff_coeffs, dv_coeffs = solve_stability_case(wing, uniform; rho_ref = ρ, r_ref = ref, span_num = 20, chord_num = 5)

    # Test values
    nf_tests = [0.001189, -0.000228, 0.152203, -0.000242, -0.003486, -8.1e-5, 0.0, 0.0, 0.0]
    ff_tests = [0.00123,  -0.000271, 0.152198]
    dv_tests = [ 0.068444 -0.000046 -0.000711  0.023607  0.000337;
                 0.010867 -0.007536  0.129968  0.021929 -0.012086;
                 4.402229 -0.012973 -0.070654  6.833903  0.001999;
                 0.031877 -0.013083  0.460035  0.091216 -0.039146;
                -0.112285 -0.004631  0.105695 -0.852395 -0.007696;
                -0.002218 -0.002115  0.008263 -0.003817  0.001079]

    # Nearfield coefficients test
    [ @test nf_c ≈ nf_t atol = 1e-6 for (nf_c, nf_t) in zip(nf_coeffs, nf_tests) ]
    # Farfield coefficients test
    [ @test ff_c ≈ ff_t atol = 1e-6 for (ff_c, ff_t) in zip(ff_coeffs, ff_tests) ]
    # Stability derivatives' coefficients test
    [ @test dv_c ≈ dv_t atol = 1e-6 for (dv_c, dv_t) in zip(dv_coeffs, dv_tests) ]
end

@testset "Vortex Lattice Method - Vanilla Aircraft" begin
    ## Wing
    wing = Wing(foils = Foil.(fill(naca4((0,0,1,2)), 2)),
                chords    = [1.0, 0.6],
                twists    = [0.0, 0.0],
                spans     = [5.0],
                dihedrals = [11.39],
                sweep_LEs = [0.]);

    # Horizontal tail
    htail_foils = Foil.(fill(naca4((0,0,1,2)), 2))
    htail = Wing(foils     = htail_foils,
                 chords    = [0.7, 0.42],
                 twists    = [0.0, 0.0],
                 spans     = [1.25],
                 dihedrals = [0.],
                 sweep_LEs = [6.39])

    # Vertical tail
    vtail_foils = Foil.(fill(naca4((0,0,0,9)), 2))
    vtail = HalfWing(foils     = vtail_foils,
                     chords    = [0.7, 0.42],
                     twists    = [0.0, 0.0],
                     spans     = [1.0],
                     dihedrals = [0.],
                     sweep_LEs = [7.97])

    ## Assembly
    wing_panels  = panel_wing(wing, 16, 10;
                              spacing = "cosine")
    htail_panels = panel_wing(htail, 6, 6;
                              position = [4., 0, 0],
                              angle    = deg2rad(-2.),
                              axis     = [0., 1., 0.],
                              spacing  = "cosine")
    vtail_panels = panel_wing(vtail, 5, 6;
                              position  = [4., 0, 0],
                              angle     = π/2,
                              axis      = [1., 0., 0.],
                              spacing   = "cosine")

    aircraft = Dict("Wing"            => wing_panels,
                    "Horizontal Tail" => htail_panels,
                    "Vertical Tail"   => vtail_panels)

    ## Reference quantities
    ac_name = "My Aircraft"
    S, b, c = projected_area(wing), span(wing), mean_aerodynamic_chord(wing)
    ρ       = 1.225
    ref     = [0.25c, 0., 0.]
    V, α, β = 1.0, 1.0, 1.0
    Ω       = [0.0, 0.0, 0.0]
    fs      = Freestream(V, α, β, Ω)

    ## Stability case
    dv_data = solve_stability_case(aircraft, fs;
                                   rho_ref   = ρ,
                                   r_ref     = ref,
                                   area_ref  = S,
                                   span_ref  = b,
                                   chord_ref = c,
                                   name      = ac_name);

    nfs, ffs, dvs = dv_data[ac_name]

    nf_tests = [0.000258, -0.006642, 0.074301, -0.003435, 0.075511, 0.001563, 0.0, 0.0, 0.0]
    ff_tests = [0.000375, -0.006685, 0.074281]
    dv_tests = [ 0.016795  0.003460  0.003761   0.093303 -0.000674;
                -0.000863 -0.374410  0.403476   0.000630 -0.253848;
                 5.749765  0.046649 -0.01346   15.571205  0.020396;
                 0.022674 -0.196605  0.660392   0.099065 -0.039688;
                -2.70367  -0.132928  0.070111 -37.372278 -0.064439;
                 0.002034  0.087382  0.014991   0.005840  0.091088]

    # Nearfield coefficients test
    [ @test nf_c ≈ nf_t atol = 1e-6 for (nf_c, nf_t) in zip(nfs, nf_tests) ]
    # Farfield coefficients test
    [ @test ff_c ≈ ff_t atol = 1e-6 for (ff_c, ff_t) in zip(ffs, ff_tests) ]
    # Stability derivatives' coefficients test
    [ @test dv_c ≈ dv_t atol = 1e-6 for (dv_c, dv_t) in zip(dvs, dv_tests) ]
end;


"""
Provides astrodynamical models as `ModelingToolkit.ODESystems`.
Check out the `ModelingToolkit` docs to learn how to use these
systems for orbit propagation with `DifferentialEquations`, or
see `GeneralAstrodynamics` for some convenient orbit propagation
wrappers.

# Extended help

## Exports

$(EXPORTS)

## Imports

$(IMPORTS)
"""
module AstrodynamicalModels

# Export every model!
export R2BP, CR3BP, NBP

# Export every `ODEFunction`!
export R2BPFunction, CR3BPFunction, NBPFunction

using Symbolics
using StaticArrays
using LinearAlgebra
using ModelingToolkit

import Memoize: @memoize

using DocStringExtensions
@template (FUNCTIONS, METHODS, MACROS) =
    """
    $(SIGNATURES)

    $(DOCSTRING)
    """

@template (TYPES, CONSTANTS) =
    """
    $(TYPEDEF)

    $(DOCSTRING)
    """

include("R2BP.jl")
include("CR3BP.jl")
include("NBP.jl")

end # module


using StatsBase
using Distributions
using DelimitedFiles
using Combinatorics: with_replacement_combinations

import LinearAlgebra
import LinearAlgebra: diag
import Statistics: mean

mutable struct FactorGraph{}
    order_list::Array{Int64,1}
    variable_count::Int
    n_alphabets::Int
    terms::Dict{Tuple,Array} # TODO, would be nice to have a stronger tuple type
    #TODO In the sanity check make sure that an kth order hyper edge is mapped to  a k rank tensor
    #variable_names::Union{Vector{String}, Nothing}
    #FactorGraph(a,b,c,d,e) = check_model_data(a,b,c,d,e) ? new(a,b,c,d,e) : error("generic init problem")
end

permutations(items, order::Int; asymmetric::Bool = false) = sort(permutations([], items, order, asymmetric))

function permutations(partial_perm::Array{Any,1}, items, order::Int, asymmetric::Bool)
    """
    All possible permutations of a given size.
    If asymmetric is false, then it returns combinations of items of the given order
    If asymmetric is true it returns all possible tuples of the size given by order from items
    """
    if order == 0
        return [tuple(partial_perm...)]
    else
        perms = []
        for item in items
            if !asymmetric && length(partial_perm) > 0
                if partial_perm[end] >= item
                    continue
                end
            end
            perm = permutations(vcat(partial_perm, item), items, order-1, asymmetric)
            append!(perms, perm)
        end
        return perms
    end
end

function raw_sampler_potts(H::FactorGraph, n_samples::Int, centered::Bool)
    """
    Given the FactorGraph, return samples according to its Gibbs distribution

    """

    n = H.variable_count
    q = H.n_alphabets
    n_config = q^n
    configs = [ digits(i,base=q, pad=n) .+ 1 for i = 0:n_config-1]
    weights = [ exp(Energy_Potts(K, H, centered)) for K in configs ]
    #print(configs, weights/sum(weights))
    raw_samples =  wsample(configs, weights, n_samples)
    return raw_samples

   end

function Energy_Potts(state::Array{Int64,  1},H::FactorGraph, cent::Bool)
    """
    Given a state and a FactorGraph, return its energy
    """
    q =  H.n_alphabets
    b =  -1.0/q
    a = 1.0 -(1.0/q)
    if !(cent)
     E = 0.0

     for (e, theta) in H.terms
        edge=Any[]

        [push!(edge, state[j] ) for j in e]

        E += theta[edge...]
     end
    return E
    end

    if cent
        E = 0.0

     for (e, theta) in H.terms
        clrs=Any[]
        r = length(e) #order of interaction
        alphabet_keys =  permutations(Array(1:q), r, asymmetric=true) #No need to generate this everytime

        [push!(clrs, state[j] ) for j in e]
        ct =  Tuple(clrs)

        [E += a^(sum( ct.==c )) * b^(r -  sum(ct.==c)) *theta[c...] for c in alphabet_keys]
     end
    return E
    end

end

function TVD(truth::Dict{}, est::Dict{}, n_samples::Int)
    """
    Total variation distance between two distributions.
    """
    s = 0.0
    for (k ,v) in est
        if haskey(truth, k)
            s+= abs( v -  truth[k])
        else
            s+= v
        end
    end

    for (k,v) in truth
        if !haskey(est, k)
            s+=v
        end
    end

    return s/(2*n_samples)

end

function conditional_energy(u::Int, state::Array{Int64,  1},H::FactorGraph, cent::Bool)
"""
Given a state and a FactorGraph, return its energy
"""
q =  H.n_alphabets
b =  -1.0/q
a = 1.0 -(1.0/q)
if !(cent)
    E = 0.0

    for (e, theta) in H.terms
        if u in e
            edge=Any[]

            [push!(edge, state[j] ) for j in e]

            E += theta[edge...]
        end
    end
    return E
end

if cent
    E = 0.0

    for (e, theta) in H.terms
        if u in e
            clrs=Any[]
            r = length(e) #order of interaction
            alphabet_keys =  permutations(Array(1:q), r, asymmetric=true) #No need to generate this everytime

            [push!(clrs, state[j] ) for j in e]
            ct =  Tuple(clrs)

            [E += a^(sum( ct.==c )) * b^(r -  sum(ct.==c)) *theta[c...] for c in alphabet_keys]
        end
    end
    return E
    end

end

function pth_order_tensor(r)
    a =  Tuple([2 for i in 1:r])
    M =  zeros(a)
    for i in CartesianIndices(M)
        s = prod(2*[Tuple(i)...] .- 3)
        M[i] = s
    end
    return M
end


using Test, FiniteMesh

cd(@__DIR__)
mesh = Mesh("../mesh/Gmsh/t1.msh")

cells, nodes = read_mesh("../mesh/Gmsh/square.msh")
cellid = extract_cell(cells)
edgeNodes, edgeCells = mesh_face_connectivity_2D(cellid)
cellArea = mesh_cell_area_2D(nodes, cellid)

cells, nodes = read_mesh("../mesh/SU2/naca0012.su2")
add_group!(cells, "../mesh/SU2/naca0012.su2")

rg = [0.1, 0.3, 0.5, 0.7, 0.9]
x = hcat(rg, rg, rg)
rg = [0.1 0.3 0.5]
y = vcat(rg, rg, rg, rg, rg)

cells, points = unstructured_grid(cat(x, y, dims = 3), [0.0, 1.0], [0.0, 0.6])

_points = zeros(4, 2)
_points[2, :] = [1.0, 0.0]
_points[3, :] = [1.0, 1.0]
_points[4, :] = [0.0, 1.0]

cells = triangulate(_points)

FiniteMesh.unit_normal(rand(3), rand(3), rand(3))

cellNeighborsReg = regularize_cell_neighbor(mesh.cellid, mesh.cellNeighbors)
FiniteMesh.check_cell_neighbor(mesh.cellid, cellNeighborsReg)

cellFacesReg = regularize_cell_face(mesh.cellid, mesh.cellFaces, mesh.facePoints)
FiniteMesh.check_cell_face(mesh.cellid, cellFacesReg, mesh.facePoints)


# Shorthand type names
typealias Vec Vector{Float64}
typealias Mat Matrix{Float64}
typealias SymMat Symmetric{Float64,Matrix{Float64}}

# Utility functions
AAt(A::Mat) = Symmetric(A*A')
AtA(A::Mat) = Symmetric(A'*A)

# sumsq
sumsq(x::Vec) = norm(x)^2


"""
    xopt, fopt, converged = fmincon(obj, startval)

Minimize the function obj, starting at startval.

fminunc() with no arguments will run an example, execute edit(fminunc,()) to see the code.
fminunc() uses NLopt.jl to do the actual minimization.

"""
function fmincon(obj, startval, R=[], r=[], lb=[], ub=[]; tol = 1e-10, iterlim=0)
    # the objective is an anonymous function
    function objective_function(x::Vector{Float64}, grad::Vector{Float64})
        obj_func_value = obj(x)[1,1]
        return(obj_func_value)
    end
    # impose the linear restrictions
    function constraint_function(x::Vector, grad::Vector, R, r)
        result = R*x .- r
        return result[1,1]
    end
    opt = Opt(:LN_COBYLA, size(startval,1))
    min_objective!(opt, objective_function)
    # impose lower and/or upper bounds
    if lb != [] lower_bounds!(opt, lb) end
    if ub != [] upper_bounds!(opt, ub) end
    # impose linear restrictions, by looping over the rows
    if R != []
        for i = 1:size(R,1)
            equality_constraint!(opt, (theta, g) -> constraint_function(theta, g, R[i:i,:], r[i]), tol)
        end
    end
    xtol_rel!(opt, tol)
    ftol_rel!(opt, tol)
    maxeval!(opt, iterlim)
    (objvalue, xopt, flag) = NLopt.optimize(opt, startval)
    return xopt, objvalue, flag
end

function fmincon()
    println("with no arguments, fmincon() runs a simple example")
    println("type edit(fmincon, ()) to see the example code")
    # return of objective should be real valued, thus the [1] to pull value out of 1-dim array
    obj = x -> x'x
    x = [2.0, 2.0]
    # sum of params should be 1
    R = [1.0 1.0]
    r = 1.0
    results = fmincon(obj, x, R, r)
end