@testset "CorrectedMutualInformation" begin

    Gaoetal2011 = joinpath(pwd(), "data", "Gaoetal2011.fasta")

    function gao11_buslje09(measure)
        filename = string("data_Gaoetal2011_soft_Busljeetal2009_measure_", measure, ".txt")
        joinpath(pwd(), "data", filename)
    end

    ## Column numbers for the output of Buslje et. al. 2009
    SCORE = 9
    ZSCORE = 12

    MIToS_SCORE = 2
    MIToS_ZSCORE = 1

    @testset "APC!" begin

        for MI in ([ 0. 2. 4.
                     2. 0. 6.
                     4. 6. 0. ],
                    PairwiseListMatrix([2., 4., 6.]),
                    PairwiseListMatrix([0., 2., 4., 0., 6., 0.], true),
                    NamedArray(PairwiseListMatrix([2., 4., 6.])),
                    NamedArray(PairwiseListMatrix([0., 2., 4., 0., 6., 0.],true)))

            if isa(MI, Matrix{Float64})
                @test Information._mean_column(MI) == [3., 4., 5.]
                @test Information._mean_total(MI) == 4.
            end

            MIp = APC!(MI)

            @test filter(x -> !isnan(x), vec(Matrix(MIp))) ≈ filter(x -> !isnan(x), vec(
                [   NaN -1.0  0.25
                    -1.0  NaN  1.00
                    0.25 1.00   NaN ]))
        end
    end

    @testset "Simple example I" begin

        aln = Residue[ 'A' 'A'
                       'A' 'R' ]

        @testset "MI" begin

            # Fill the Pij matrix of the example
            Pij = zeros(Float64, 20, 20);
            N = 2 # There are 2 sequences
            Pij[1,1] = (1/N) # A A
            Pij[1,2] = (1/N) # A R
            @test sum(Pij) ≈ 1.0 # Is the matrix correct?
            # Fill the marginals
            Pi = dropdims(Base.sum(Pij, dims=2), dims=2)
            @test Pi[1] == 1.0 # Is always A
            Pj = dropdims(Base.sum(Pij, dims=1), dims=1)
            @test Pj[1] == 0.5 # A
            @test Pj[2] == 0.5 # R
            # Start to sum with 0.0
            total = 0.0
            for i in 1:20, j in 1:20
                if Pij[i,j] != 0.0 && Pi[i] != 0.0 && Pj[j] != 0.0
                    total += (Pij[i,j] * log(Pij[i,j]/(Pi[i]*Pj[j]))) # 0.5 * log(0.5/(0.5 * 1.0)) == 0.0
                end
            end
            # Compare with MIToS result
            mi = mapcolpairfreq!(mutual_information, aln,
                    Probabilities(ContingencyTable(Float64,Val{2},UngappedAlphabet())))
            @test mi[1,2] == total
        end

        @testset "MI: Using pseudocount (0.05)" begin

            Pij = zeros(Float64, 20, 20);
            N = (400 * 0.05) + 2 # 0.05 is the pseudocount and there are 2 sequences
            fill!(Pij, 0.05/N);
            Pij[1,1] =(1.05/N) # A A
            Pij[1,2] =(1.05/N) # A R
            @test sum(Pij) ≈ 1.0
            Pi = dropdims(Base.sum(Pij, dims=2), dims=2)
            Pj = dropdims(Base.sum(Pij, dims=1), dims=1)
            total = 0.0
            for i in 1:20, j in 1:20
                total += (Pij[i,j] * log(Pij[i,j]/(Pi[i]*Pj[j])))
            end
            # Compare with MIToS result
            mi = mapcolpairfreq!(mutual_information, aln,
                    Probabilities(ContingencyTable(Float64,Val{2},UngappedAlphabet())),
                    pseudocounts=AdditiveSmoothing(0.05))
            @test mi[1,2] ≈ total

            @testset "APC!" begin

                APC!(mi)
                @test isnan(mi[1,1])
                @test isapprox(mi[1,2], 0.0, rtol=1e-16)

                zerodiagonal = Float64[ 0.0 total
                                        total 0.0 ]

                # MI.. == MI.j == MIi. == total
                # APC = (total * total) / total == total
                # MI - APC == total - total == 0.0
                APC!(zerodiagonal)
                @test zerodiagonal[1,2] ≈ 0.0
            end
        end

        @testset "Z-score" begin

            # 2 Possibilities:  A A   A A
            #                   A R   R A
            # Is almost the same MI, Z score should be 0.0
            results = buslje09(aln, lambda=0.05, clustering=false, apc=false)

            @test results[MIToS_ZSCORE][1,2] ≈ 0.0

            @test aln == Residue[ 'A' 'A'
                                  'A' 'R' ]
        end
    end

    @testset "Simple example II" begin

        aln = Residue[ 'R' 'A'
                       'A' 'R' ]

        @testset "MI" begin

            Pij = zeros(Float64, 20, 20);
            N = 2 # There are 2 sequences
            Pij[2,1] = (1/N) # R A
            Pij[1,2] = (1/N) # A R
            @test sum(Pij) ≈ 1.0
            Pi = dropdims(Base.sum(Pij, dims=2), dims=2)
            @test Pi[2] == 0.5 # R
            @test Pi[1] == 0.5 # A
            Pj = dropdims(Base.sum(Pij, dims=1), dims=1)
            @test Pj[1] == 0.5 # A
            @test Pj[2] == 0.5 # R
            total = 0.0
            for i in 1:20, j in 1:20
                if Pij[i,j] != 0.0 && Pi[i] != 0.0 && Pj[j] != 0.0
                    total += (Pij[i,j] * log(Pij[i,j]/(Pi[i]*Pj[j])))
                end
            end
            @test total == 0.5 * log(2) + 0.5 * log(2) # 0.5/(0.5*0.5) == 2
            # Compare with MIToS result
            mi = mapcolpairfreq!(mutual_information, aln,
                    Probabilities(ContingencyTable(Float64,Val{2},UngappedAlphabet())))
            @test mi[1,2] == total
        end

        @testset "MI: Using pseudocount (0.05)" begin

            Pij = zeros(Float64, 20, 20);
            N = (400 * 0.05) + 2 # 0.05 is the pseudocount and there are 2 sequences
            fill!(Pij, 0.05/N);
            Pij[2,1] =(1.05/N) # R A
            Pij[1,2] =(1.05/N) # A R
            @test sum(Pij) ≈ 1.0
            Pi = dropdims(Base.sum(Pij, dims=2), dims=2)
            Pj = dropdims(Base.sum(Pij, dims=1), dims=1)
            total = 0.0
            for i in 1:20, j in 1:20
                total += (Pij[i,j] * log(Pij[i,j]/(Pi[i]*Pj[j])))
            end
            # Compare with MIToS result
            mi = mapcolpairfreq!(mutual_information, aln,
                    Probabilities(ContingencyTable(Float64,Val{2},UngappedAlphabet())),
                    pseudocounts=AdditiveSmoothing(0.05))
            @test mi[1,2] ≈ total

            @testset "APC!" begin

                APC!(mi)
                @test isnan(mi[1,1])
                @test mi[1,2] ≈ 0.0

                zerodiagonal = Float64[ 0.0 total
                                        total 0.0 ]

                # MI.. == MI.j == MIi. == total
                # APC = (total * total) / total == total
                # MI - APC == total - total == 0.0
                APC!(zerodiagonal)
                @test zerodiagonal[1,2] ≈ 0.0
            end
        end

        @testset "Z-score" begin

            # 4 Possibilities:  R A   R A   A R   A R
            #                   A R   R A   A R   R A
            mi = mapcolpairfreq!(mutual_information, aln,
                    Probabilities(ContingencyTable(Float64,Val{2},UngappedAlphabet())),
                    pseudocounts=AdditiveSmoothing(0.05))
            other_posib = Residue[ 'A' 'R'
                                   'A' 'R' ]
            other_mi = mapcolpairfreq!(mutual_information, other_posib,
                        Probabilities(ContingencyTable(Float64,Val{2},UngappedAlphabet())),
                        pseudocounts=AdditiveSmoothing(0.05))
            # The mean should be:
            r_mean = 0.5 * ( mi[1,2] + other_mi[1,2] )
            # And the std should be similar to:
            r_std = 0.5 * ( sqrt((mi[1,2]-r_mean)^2) + sqrt((other_mi[1,2]-r_mean)^2) )

            results = buslje09(aln, lambda=0.05, clustering=false, apc=false, samples=100)
            @test isapprox(results[MIToS_ZSCORE][1,2], (mi[1,2]-r_mean)/r_std, atol=0.55)

            results = buslje09(aln, lambda=0.05, clustering=false, apc=false, samples=1000)
            @test isapprox(results[MIToS_ZSCORE][1,2], (mi[1,2]-r_mean)/r_std, atol=0.15)

            results = buslje09(aln, lambda=0.05, clustering=false, apc=false, samples=10000)
            @test isapprox(results[MIToS_ZSCORE][1,2], (mi[1,2]-r_mean)/r_std, atol=0.055)

            @test aln == Residue[ 'R' 'A'
                                  'A' 'R' ]
        end
    end

    @testset "Results from Buslje et. al. 2009" begin

        @testset "Simple" begin
            data = readdlm(joinpath(pwd(), "data",
                           "data_simple_soft_Busljeetal2009_measure_MI.txt"), comments=true)
            results = buslje09(joinpath(pwd(), "data", "simple.fasta"), FASTA,
                               lambda=0.0, clustering=false, apc=false)

            @test isapprox(Float64(data[1, SCORE]), results[MIToS_SCORE][1,2], atol=1e-6)
            @test isapprox(Float64(data[1, ZSCORE]), results[MIToS_ZSCORE][1,2], atol=2.)
        end

        @testset "Gaoetal2011" begin
            data = readdlm(gao11_buslje09("MI"), comments=true)
            results = buslje09(Gaoetal2011, FASTA, lambda=0.0, clustering=false, apc=false)

            @test isapprox([ Float64(x) for x in data[:,SCORE] ],
                           matrix2list(results[MIToS_SCORE]), atol=1e-6)
            @test isapprox([ Float64(x) for x in data[:,ZSCORE] ],
                           matrix2list(results[MIToS_ZSCORE]), atol=2.)

            # println(cor(convert(Vector{Float64},data[:,ZSCORE]),matrix2list(results[MIToS_ZSCORE])))

            @testset "cMI" begin
                result = copy(results[1])
                result[diagind(result)] .= 0.0
                @test cumulative(results[1], -Inf) == sum(result, dims=1)
                @test all((cumulative(results[2],0.5) .-
                           [0.0,0.0,1.38629,1.38629,1.38629,0.0]') .< 0.00001)
                           #        1.38629 = 0.693147 + 0.693147
            end

            result_0_05 = buslje09(Gaoetal2011,FASTA,lambda=0.05,clustering=false,apc=false)
            @test isapprox(result_0_05[MIToS_SCORE][1,2], 0.33051006116310444, atol=1e-14)
        end
    end

    @testset "MI + clustering" begin

        data = readdlm(gao11_buslje09("MI_clustering"), comments=true)
        results = buslje09(Gaoetal2011, FASTA, lambda=0.0, clustering=true, apc=false)

        @test isapprox([ Float64(x) for x in data[:,SCORE] ],
                       matrix2list(results[MIToS_SCORE]), atol=1e-6)
        @test isapprox([ Float64(x) for x in data[:,ZSCORE] ],
                       matrix2list(results[MIToS_ZSCORE]), atol=2.)
    end

    @testset "MIp" begin

        data = readdlm(gao11_buslje09("MI_APC"), comments=true)
        results = buslje09(Gaoetal2011, FASTA, lambda=0.0, clustering=false, apc=true)

        @test isapprox([ Float64(x) for x in data[:,SCORE] ],
                       matrix2list(results[MIToS_SCORE]), atol=1e-6)
        @test isapprox([ Float64(x) for x in data[:,ZSCORE] ],
                       matrix2list(results[MIToS_ZSCORE]), atol=2.0)

        @test isapprox(results[MIToS_SCORE][5,6], 0.018484, atol=0.000001)
    end

    @testset "MIp + clustering" begin

        data = readdlm(gao11_buslje09("MI_APC_clustering"), comments=true)
        results = buslje09(Gaoetal2011, FASTA, lambda=0.0, clustering=true, apc=true)

        @test isapprox([ Float64(x) for x in data[:,SCORE] ],
                       matrix2list(results[MIToS_SCORE]), atol=1e-6)
        @test isapprox([ Float64(x) for x in data[:,ZSCORE] ],
                       matrix2list(results[MIToS_ZSCORE]), atol=2.)

        @test isapprox(results[MIToS_SCORE][5,6], 0.018484, atol=0.000001)
    end

    @testset "BLMI" begin

        @testset "Simple" begin
            file = joinpath(pwd(), "data", "simple.fasta")
            busl = buslje09(file, FASTA)
            blmi = BLMI(file, FASTA)

            @test PairwiseListMatrices.getlist(busl[1]) ≈
                  PairwiseListMatrices.getlist(blmi[1])
            @test PairwiseListMatrices.getlist(busl[2]) ≈
                  PairwiseListMatrices.getlist(blmi[2])
        end

        @testset "Gaoetal2011" begin
            msa  = read(Gaoetal2011, FASTA)
            busl = buslje09(Gaoetal2011, FASTA, lambda=0.0, samples=0)
            blmi = BLMI(msa, lambda=0.0, beta=0.0, samples=5)
            # BLMI should be equal to Buslje09 if beta is zero

            @test PairwiseListMatrices.getlist(busl[2]) ≈
                  PairwiseListMatrices.getlist(blmi[2]) # MIapc
            @test msa == read(Gaoetal2011, FASTA)
        end

        @testset "Gaoetal2011, lambda 0.05" begin

            busl = buslje09(Gaoetal2011, FASTA, lambda=0.5, samples=0)
            blmi = BLMI(Gaoetal2011, FASTA, lambda=0.5, beta=0.0, samples=5)
            # BLMI should be equal to Buslje09 if beta is zero

            @test PairwiseListMatrices.getlist(busl[2]) ≈
                  PairwiseListMatrices.getlist(blmi[2]) # MIapc
        end
    end
end


###############################################################################
#
# Yujie's wrapper for CLIMAParameters
#
###############################################################################
# Define a local struct inherited from AbstractEarthParameterSet
struct EarthParameterSet <: AbstractEarthParameterSet end
const EARTH = EarthParameterSet();

""" Avogadro's number `[molecule mol⁻¹]` """
AVOGADRO(FT=Float64) = FT( avogad() );

""" Isobaric specific heat of dry air `[J kg⁻¹ K⁻¹]` """
CP_D(FT=Float64) = FT( cp_d(EARTH) );

""" Isobaric specific heat of dry air `[J mol⁻¹ K⁻¹]` """
CP_D_MOL(FT=Float64) = CP_D(FT) * M_DRYAIR(FT);

""" Isobaric specific heat of liquid water `[J kg⁻¹ K⁻¹]` """
CP_L(FT=Float64) = FT( cp_l(EARTH) );

""" Isobaric specific heat of water vapor `[J kg⁻¹ K⁻¹]` """
CP_V(FT=Float64) = FT( cp_v(EARTH) );

""" Universal gas constant `[J mol⁻¹ K⁻¹]` """
GAS_R(FT=Float64) = FT( gas_constant() );

""" Gravity of the Earth `[m s⁻²]` """
GRAVITY(FT=Float64) = FT( grav(EARTH) );

""" Planck constant `[m² kg s⁻¹]` """
H_PLANCK(FT=Float64) = FT( h_Planck() );

""" Boltzmann constant `[m² kg s⁻² K⁻¹]` """
K_BOLTZMANN(FT=Float64) = FT( k_Boltzmann() );

""" Stefan-Boltzmann constant `[W m⁻² K⁻⁴]` """
K_STEFAN(FT=Float64) = FT( Stefan() );

""" Von Karman constant `[-]` """
K_VON_KARMAN(FT=Float64) = FT( von_karman_const(EARTH) );

""" Latent heat vaporization at ``T_0`` `[K kg⁻¹]` """
LH_V0(FT=Float64) = FT( LH_v0(EARTH) );

""" Light speed in vacuum `[m s⁻¹]` """
LIGHT_SPEED(FT=Float64) = FT( light_speed() );

""" Molar mass of dry air `[kg mol⁻¹]` """
M_DRYAIR(FT=Float64) = FT( molmass_dryair(EARTH) );

""" Molar mass of water `[kg mol⁻¹]` """
M_H₂O(FT=Float64) = FT( molmass_water(EARTH) );

""" Mean atmospheric pressure at sea level `[Pa]` """
P_ATM(FT=Float64) = FT( MSLP(EARTH) );

""" Water vapor pressure at triple temperature `[Pa]` """
PRESS_TRIPLE(FT=Float64) = FT( press_triple(EARTH) );

""" Gas constant water vapor `[J kg⁻¹ K⁻¹]` """
R_V(FT=Float64) = FT( R_v(EARTH) );

""" Gas constant times 298.15 K `[J mol⁻¹]` """
RT_25(FT=Float64) = GAS_R(FT) * T_25(FT);

""" Freezing temperature of water `[K]` """
T_0(FT=Float64) = FT( T_freeze(EARTH) );

""" Kelvin temperature at 25 Celcius `[K]` """
T_25(FT=Float64) = T_0(FT) + 25;

""" Triple temperature of water `[K]` """
T_TRIPLE(FT=Float64) = FT( T_triple(EARTH) );

""" Molar volume of liqiud water """
V_H₂O(FT=Float64) = M_H₂O(FT) / ρ_H₂O(FT);

""" Mean number of days per year [day] """
YEAR_D(FT=Float64) = FT( 365.2422222 );

""" Density of liquid water `[kg m⁻³]` """
ρ_H₂O(FT=Float64) = FT( ρ_cloud_liq(EARTH) );

""" Density of water times gravity `[MPa m⁻¹]` """
ρg_MPa(FT=Float64) = ρ_H₂O(FT) * GRAVITY(FT) * FT(1e-6);


""" Computing in the category of finite sets and relations, and its skeleton.
"""
module FinRelations
export BoolRig, FinOrdRelOb, FinOrdRelation, force

import Base: +, *
using AutoHashEquals

using ...GAT
using ...Theories: DistributiveBicategoryRelations
import ...Theories: dom, codom, id, compose, ⋅, ∘, dagger, dunit, dcounit,
  otimes, ⊗, munit, braid, oplus, ⊕, mzero, swap,
  mcopy, Δ, mmerge, ∇, delete, ◊, create, □, plus, zero, coplus, cozero,
  pair, copair, proj1, proj2, coproj1, coproj2, meet, join, top, bottom
import ..FinSets: force
using ..Matrices
using ..Matrices: zero_matrix

# Rig of booleans
#################

""" The rig of booleans.

This struct is needed because in base Julia, the product of booleans is another
boolean, but the sum of booleans is coerced to an integer: `true + true == 2`.
"""
@auto_hash_equals struct BoolRig <: Number
  value::Bool
end

+(x::BoolRig, y::BoolRig) = BoolRig(x.value || y.value)
*(x::BoolRig, y::BoolRig) = BoolRig(x.value && y.value)

Base.zero(::Type{BoolRig}) = BoolRig(false)
Base.one(::Type{BoolRig}) = BoolRig(true)
Base.float(x::BoolRig) = float(x.value)
Base.show(io::IO, x::BoolRig) = print(io, x.value)

# Category of finite ordinals and relations
###########################################

""" Finite ordinal (natural number).

An object in the category `FinOrdRel` of finite ordinals and relations, which is
the skeleton of the category `FinRel` of finite sets and relations.
"""
@auto_hash_equals struct FinOrdRelOb
  n::Int
end
Base.eachindex(X::FinOrdRelOb) = 1:X.n

""" Binary relation between sets in the form of finite ordinals.

A morphism in the category `FinOrdRel` of finite ordinals and relations, which
is the skeleton of the category `FinRel` of finite sets and relations. The
relation can be represented implicitly by an arbitrary Julia function mapping
pairs of integers to booleans or explicitly by a matrix (dense or sparse) taking
values in the rig of booleans ([`BoolRig`](@ref)).
"""
abstract type FinOrdRelation end

FinOrdRelation(R, dom::FinOrdRelOb, codom::FinOrdRelOb) =
  FinOrdRelation(R, dom.n, codom.n)

""" Relation in FinOrdRel represented by an arbitrary Julia function.
"""
@auto_hash_equals struct FinOrdRelationLazy <: FinOrdRelation
  rel::Function
  dom::Int
  codom::Int
end
FinOrdRelation(R::Function, dom::Int, codom::Int) =
  FinOrdRelationLazy(R, dom, codom)

(R::FinOrdRelationLazy)(x, y) = R.rel(x, y)

""" Relation in FinOrdRel represented by a boolean matrix.

Boolean matrices are also known as logical matrices or relation matrices.
"""
@auto_hash_equals struct FinOrdRelationMatrix{T <: AbstractMatrix{BoolRig}} <: FinOrdRelation
  rel::T
end

FinOrdRelation(R::AbstractMatrix{BoolRig}) = FinOrdRelationMatrix(R)

function FinOrdRelation(R::AbstractMatrix{BoolRig}, dom::Int, codom::Int)
  @assert size(R,2) == dom && size(R,1) == codom
  FinOrdRelationMatrix(R)
end

(R::FinOrdRelationMatrix)(x, y) = R.rel[y, x].value

function force(::Type{T}, R::FinOrdRelation) where T <: AbstractMatrix{BoolRig}
  m, n = dom(R).n, codom(R).n
  M = zero_matrix(T, n, m)
  for i in 1:m, j in 1:n
    if R(i,j)
      M[j,i] = BoolRig(true)
    end
  end
  FinOrdRelationMatrix(M)
end
force(::Type{T}, R::FinOrdRelationMatrix{T}) where T <: AbstractMatrix{BoolRig} = R
force(R::FinOrdRelation) = force(Matrix{BoolRig}, R)

""" FinOrdRel as a distributive bicategory of relations.
"""
@instance DistributiveBicategoryRelations(FinOrdRelOb, FinOrdRelation) begin
  dom(R::FinOrdRelation) = FinOrdRelOb(R.dom)
  codom(R::FinOrdRelation) = FinOrdRelOb(R.codom)

  id(A::FinOrdRelOb) = FinOrdRelation((x,y) -> x == y, A, A)

  function compose(R::FinOrdRelation, S::FinOrdRelation)
    @assert codom(R) == dom(S)
    FinOrdRelation(dom(R), codom(S)) do x,z
      any(R(x,y) && S(y,z) for y in eachindex(codom(R)))
    end
  end

  # Multiplicative monoidal structure

  otimes(A::FinOrdRelOb, B::FinOrdRelOb) = FinOrdRelOb(A.n * B.n)
  munit(::Type{FinOrdRelOb}) = FinOrdRelOb(1)

  function otimes(R::FinOrdRelation, S::FinOrdRelation)
    # Indexing is consistent with that of Kronecker products.
    dom_proj = CartesianIndices((dom(S).n, dom(R).n))
    cod_proj = CartesianIndices((codom(S).n, codom(R).n))
    FinOrdRelation(dom(R)⊗dom(S), codom(R)⊗codom(S)) do x, y
      R(dom_proj[x][2], cod_proj[y][2]) && S(dom_proj[x][1], cod_proj[y][1])
    end
  end

  function braid(A::FinOrdRelOb, B::FinOrdRelOb)
    dom_proj = CartesianIndices((B.n, A.n))
    cod_proj = CartesianIndices((A.n, B.n))
    FinOrdRelation(A⊗B, B⊗A) do x, y
      dom_proj[x][1] == cod_proj[y][2] && dom_proj[x][2] == cod_proj[y][1]
    end
  end

  function mcopy(A::FinOrdRelOb)
    proj = CartesianIndices((A.n, A.n))
    FinOrdRelation((x,y) -> x == proj[y][1] && x == proj[y][2], A, A⊗A)
  end
  function mmerge(A::FinOrdRelOb)
    proj = CartesianIndices((A.n, A.n))
    FinOrdRelation((x,y) -> proj[x][1] == y && proj[x][2] == y, A⊗A, A)
  end
  delete(A::FinOrdRelOb) = FinOrdRelation((x,y) -> true, A, munit(FinOrdRelOb))
  create(A::FinOrdRelOb) = FinOrdRelation((x,y) -> true, munit(FinOrdRelOb), A)

  dagger(R::FinOrdRelation) = FinOrdRelation((x,y) -> R(y,x), codom(R), dom(R))
  dunit(A::FinOrdRelOb) = compose(create(A), mcopy(A))
  dcounit(A::FinOrdRelOb) = compose(mmerge(A), delete(A))

  # Additive monoidal structure

  oplus(A::FinOrdRelOb, B::FinOrdRelOb) = FinOrdRelOb(A.n + B.n)
  mzero(::Type{FinOrdRelOb}) = FinOrdRelOb(0)

  function oplus(R::FinOrdRelation, S::FinOrdRelation)
    m, n = dom(R).n, codom(R).n
    FinOrdRelation(dom(R)⊕dom(S), codom(R)⊕codom(S)) do x, y
      (x <= m && y <= n && R(x, y)) || (x > m && y > n && S(x-m, y-n))
    end
  end

  function swap(A::FinOrdRelOb, B::FinOrdRelOb)
    FinOrdRelation((x,y) -> x == y-B.n || x-A.n == y, A⊕B, B⊕A)
  end

  plus(A::FinOrdRelOb) = FinOrdRelation((x,y) -> x == y || x-A.n == y, A⊕A, A)
  coplus(A::FinOrdRelOb) = FinOrdRelation((x,y) -> x == y || x == y-A.n, A, A⊕A)
  zero(A::FinOrdRelOb) = FinOrdRelation((x,y) -> false, mzero(FinOrdRelOb), A)
  cozero(A::FinOrdRelOb) = FinOrdRelation((x,y) -> false, A, mzero(FinOrdRelOb))

  pair(R::FinOrdRelation, S::FinOrdRelation) = coplus(dom(R))⋅(R⊕S)
  copair(R::FinOrdRelation, S::FinOrdRelation) = (R⊕S)⋅plus(codom(R))
  proj1(A::FinOrdRelOb, B::FinOrdRelOb) = id(A)⊕cozero(B)
  proj2(A::FinOrdRelOb, B::FinOrdRelOb) = cozero(A)⊕id(B)
  coproj1(A::FinOrdRelOb, B::FinOrdRelOb) = id(A)⊕zero(B)
  coproj2(A::FinOrdRelOb, B::FinOrdRelOb) = zero(A)⊕id(B)

  # Lattice structure

  function meet(R::FinOrdRelation, S::FinOrdRelation)
    @assert dom(R) == dom(S) && codom(R) == codom(S)
    FinOrdRelation((x,y) -> R(x,y) && S(x,y), dom(R), codom(R))
  end
  function join(R::FinOrdRelation, S::FinOrdRelation)
    @assert dom(R) == dom(S) && codom(R) == codom(S)
    FinOrdRelation((x,y) -> R(x,y) || S(x,y), dom(R), codom(R))
  end
  top(A::FinOrdRelOb, B::FinOrdRelOb) = FinOrdRelation((x,y) -> true, A, B)
  bottom(A::FinOrdRelOb, B::FinOrdRelOb) = FinOrdRelation((x,y) -> false, A, B)
end

# For relation matrices, delegate to the category of matrices.
const RelMat = FinOrdRelationMatrix

dom(R::RelMat) = FinOrdRelOb(size(R.rel, 2))
codom(R::RelMat) = FinOrdRelOb(size(R.rel, 1))

compose(R::RelMat, S::RelMat) = RelMat(compose(R.rel, S.rel))
otimes(R::RelMat, S::RelMat) = RelMat(otimes(R.rel, S.rel))
oplus(R::RelMat, S::RelMat) = RelMat(oplus(R.rel, S.rel))
dagger(R::RelMat) = RelMat(transpose(R.rel))

meet(R::RelMat, S::RelMat) = RelMat(R.rel .* S.rel)
join(R::RelMat, S::RelMat) = RelMat(R.rel .+ S.rel)
pair(R::RelMat, S::RelMat) = RelMat(pair(R.rel, S.rel))
copair(R::RelMat, S::RelMat) = RelMat(copair(R.rel, S.rel))

end


using LightGraphs
using MetaGraphs
using LinearAlgebra
using SparseArrays
using SuiteSparse
using Combinatorics

include("./misc.jl")

function read_sparsity_pattern(instance_path::String)
    data = load_matpower(instance_path)
    ## Bus load and shunt information
    i_debut, i_fin = find_numarray(1, data)
    checkfor(data, i_debut-1, "mpc.bus")
    nb_bus = i_fin-i_debut+1
    index_bus = Dict( data[i+i_debut-1,1] => i for i=1:nb_bus)
    ## Bus generator information
    i_debut, i_fin = find_numarray(i_fin+1, data)
    checkfor(data, i_debut-1, "mpc.gen")
    #initialize network graph G
    sp = spzeros(nb_bus,nb_bus)
    ## Link information
    i_debut, i_fin = find_numarray(i_fin+1, data)
    checkfor(data, i_debut-1, "mpc.branch")
    for i=i_debut:i_fin
        if data[i, 11] == 0
        else
            orig = index_bus[data[i,1]]
            dest = index_bus[data[i,2]]
            sp[orig,dest] = 1
        end
    end
    sp_sym = sp + sp'
    diag = zeros(nb_bus)
    for i=1:nb_bus
        sum_col = sum(sp_sym[i,:])
        diag[i] = Int(sum_col + 1)
    end
    return sp_sym + Diagonal(diag)
end

function read_sparsity_pattern(graph::AbstractGraph)
    sparsity_pattern = adjacency_matrix(graph)
    for i in 1:nv(graph)
        sparsity_pattern[i, i] = length(neighbors(graph, i)) + 1
    end
    return sparsity_pattern
end


using SecondQuantization
using Test

using SecondQuantization: flatten_term
using Latexify

@testset "Normal Ordering" begin
    @boson a b c
    @parameter ω
    N = a' * a

    @test isequal(normal_order(a * a'), normal_order(a' * a + 1))
    @test isequal(normal_order(a * a', rewriter = threaded_normal_order_simplifier(100)), normal_order(a' * a + 1))

    to_test = to_test = b / ω + b * c + a' * a * b / ω + a' * a * b * c
    @test isequal(normal_order(a * (ω * c+1) * b * a' / ω), to_test)
    @test isequal(normal_order(a * (ω * c+1) * b * a' / ω, rewriter = threaded_normal_order_simplifier(100)), to_test)

    @test isequal(normal_order((a + a')^2), flatten_term(+, 1 + a^2 + flatten_term(*, 2 * a' * a) + a'^2))
    @test isequal(normal_order((a + a')^2, rewriter = threaded_normal_order_simplifier(100)), flatten_term(+, 1 + a^2 + flatten_term(*, 2 * a' * a) + a'^2))

    @test isequal(normal_order(a / a), 1)

    @test isequal(normal_order(a * N * N * N), a + 7 * a' * a^2 + 6 * a'^2 * a^3 + a'^3 * a^4)
end

@testset "Latexify" begin
    @boson a b
    @parameter ω δ

    @test latexify(Expr(:latexifymerge, "\\hat{a}")) == latexify(a)
    @test latexify(Expr(:latexifymerge, "\\hat{a}^\\dagger")) == latexify(a')
    @test latexify(Expr(:latexifymerge, "\\hat{a}^\\dagger\\hat{b}")) == latexify(a' * b)
    @test latexify(Expr(:latexifymerge, "\\hat{a}^\\dagger\\left(1+\\hat{b}\\right)")) == latexify(a' * (1 + b))
    @test latexify(Expr(:latexifymerge, "\\frac{1}{\\omega}\\hat{a}+\\frac{1}{\\omega}\\hat{a}^\\dagger")) == latexify((a + a') / ω)
    @test latexify(Expr(:latexifymerge, "-1+\\hat{a}^\\dagger\\hat{a}")) == latexify(a' * a - 1)
    @test latexify(Expr(:latexifymerge, "\\sin \\left(\\hat{a}+\\hat{a}^\\dagger\\right)")) == latexify(sin(a + a'))
    @test latexify(Expr(:latexifymerge, "\\cos \\left(\\hat{a}+\\hat{a}^\\dagger\\right)")) == latexify(cos(a + a'))
    @test latexify(Expr(:latexifymerge, "e^{\\hat{a}+\\hat{a}^\\dagger}")) == latexify(exp(a + a'))
    @test latexify(Expr(:latexifymerge, "\\hat{a}-\\hat{a}^\\dagger")) == latexify(normal_order(a - a'))
    @test latexify(Expr(:latexifymerge, "{\\hat{a}}^{2}")) == latexify(a^2)
    @test latexify(Expr(:latexifymerge, "\\left(\\hat{a}+\\hat{a}^\\dagger\\right)^{2}")) == latexify((a + a')^2)
    @test latexify(Expr(:latexifymerge, "\\delta\\omega+\\hat{a}")) == latexify(normal_order(a + δ * ω))
    @test latexify(Expr(:latexifymerge, "\\frac{\\delta}{\\omega}+\\hat{a}")) == latexify(a + δ / ω)
    @test latexify(Expr(:latexifymerge, "{\\delta}^{\\omega}+\\hat{a}")) == latexify(a + δ ^ ω)
    @test latexify(Expr(:latexifymerge, "\\left(\\delta+\\omega\\right)^{\\omega}+\\hat{a}")) == latexify(a + (δ+ω) ^ ω)
end

@testset "Commutator" begin
    @boson a b
    @parameter ω δ

    @test isequal(commutator(a, b), 0)
    @test isequal(commutator(a, a), 0)
    @test isequal(commutator(a, a'), 1)
    @test isequal(commutator(a', a), -1)
    @test isequal(commutator(a, a' * ω), ω)
    @test isequal(commutator(a * ω^2 / (ω + δ), a'), ω^2 / (ω + δ))
    @test isequal(commutator(a * ω^2, a' * δ^2), ω^2 * δ^2)
end


################################################################################
#                                 Euler Angles
################################################################################

export angle2dcm,      angle2quat,      angle2rot
export smallangle2dcm, smallangle2quat, smallangle2rot

################################################################################
#                                 Conversions
################################################################################

# Direction Cosine Matrix
# ==============================================================================

"""
    function angle2dcm(angle_r1::Number, angle_r2::Number, angle_r3::Number, rot_seq::Symbol = :ZYX)

Convert the Euler angles `angle_r1`, `angle_r2`, and `angle_r3` [rad] with the
rotation sequence `rot_seq` to a direction cosine matrix.

The rotation sequence is defined by a `:Symbol`. The possible values are:
`:XYX`, `XYZ`, `:XZX`, `:XZY`, `:YXY`, `:YXZ`, `:YZX`, `:YZY`, `:ZXY`, `:ZXZ`,
`:ZYX`, and `:ZYZ`. If no value is specified, then it defaults to `:ZYX`.

# Remarks

This function assigns `dcm = A3 * A2 * A1` in which `Ai` is the DCM related with
the *i*-th rotation, `i Є [1,2,3]`.

# Example

```julia-repl
dcm = angle2dcm(pi/2, pi/3, pi/4, :ZYX)
3×3 StaticArrays.SArray{Tuple{3,3},Float64,2,9}:
  3.06162e-17  0.5       -0.866025
 -0.707107     0.612372   0.353553
  0.707107     0.612372   0.353553
```

"""
function angle2dcm(angle_r1::Number,
                   angle_r2::Number,
                   angle_r3::Number,
                   rot_seq::Symbol = :ZYX)
    # Compute the sines and cosines.
    c1 = cos(angle_r1)
    s1 = sin(angle_r1)

    c2 = cos(angle_r2)
    s2 = sin(angle_r2)

    c3 = cos(angle_r3)
    s3 = sin(angle_r3)

    if rot_seq == :ZYX
        return DCM(       c2*c1    ,        c2*s1    , -s2  ,
                   s3*s2*c1 - c3*s1, s3*s2*s1 + c3*c1, s3*c2,
                   c3*s2*c1 + s3*s1, c3*s2*s1 - s3*c1, c3*c2)'

    elseif rot_seq == :XYX
        return DCM( c2  ,         s1*s2    ,       -c1*s2    ,
                   s2*s3, -s1*c2*s3 + c1*c3, c1*c2*s3 + s1*c3,
                   s2*c3, -s1*c3*c2 - c1*s3, c1*c3*c2 - s1*s3)'

    elseif rot_seq == :XYZ
        return DCM( c2*c3,  s1*s2*c3 + c1*s3, -c1*s2*c3 + s1*s3,
                   -c2*s3, -s1*s2*s3 + c1*c3,  c1*s2*s3 + s1*c3,
                      s2 ,        -s1*c2    ,       c1*c2)'

    elseif rot_seq == :XZX
        return DCM(   c2 ,         c1*s2    ,         s1*s2    ,
                   -s2*c3,  c1*c3*c2 - s1*s3,  s1*c3*c2 + c1*s3,
                    s2*s3, -c1*c2*s3 - s1*c3, -s1*c2*s3 + c1*c3)'

    elseif rot_seq == :XZY
        return DCM(c3*c2, c1*c3*s2 + s1*s3, s1*c3*s2 - c1*s3,
                    -s2 ,        c1*c2    ,        s1*c2    ,
                   s3*c2, c1*s2*s3 - s1*c3, s1*s2*s3 + c1*c3)'

    elseif rot_seq == :YXY
        return DCM(-s1*c2*s3 + c1*c3, s2*s3 , -c1*c2*s3 - s1*c3,
                          s1*s2    ,   c2  ,         c1*s2    ,
                   s1*c3*c2 + c1*s3, -s2*c3,  c1*c3*c2 - s1*s3)'

    elseif rot_seq == :YXZ
        return DCM( c1*c3 + s2*s1*s3, c2*s3, -s1*c3 + s2*c1*s3,
                   -c1*s3 + s2*s1*c3, c2*c3,  s1*s3 + s2*c1*c3,
                        s1*c2       ,  -s2 ,      c2*c1       )'

    elseif rot_seq == :YZX
        return DCM(        c1*c2    ,    s2 ,        -s1*c2    ,
                   -c3*c1*s2 + s3*s1,  c2*c3,  c3*s1*s2 + s3*c1,
                    s3*c1*s2 + c3*s1, -s3*c2, -s3*s1*s2 + c3*c1)'

    elseif rot_seq == :YZY
        return DCM(c1*c3*c2 - s1*s3, s2*c3, -s1*c3*c2 - c1*s3,
                         -c1*s2    ,   c2 ,         s1*s2    ,
                   c1*c2*s3 + s1*c3, s2*s3, -s1*c2*s3 + c1*c3)'

    elseif rot_seq == :ZXY
        return DCM(c3*c1 - s2*s3*s1, c3*s1 + s2*s3*c1, -s3*c2,
                      -c2*s1       ,     c2*c1       ,    s2 ,
                   s3*c1 + s2*c3*s1, s3*s1 - s2*c3*c1,  c2*c3)'

    elseif rot_seq == :ZXZ
        return DCM(-s1*c2*s3 + c1*c3, c1*c2*s3 + s1*c3, s2*s3,
                   -s1*c3*c2 - c1*s3, c1*c3*c2 - s1*s3, s2*c3,
                           s1*s2    ,       -c1*s2    ,  c2)'

    elseif rot_seq == :ZYZ
        return DCM( c1*c3*c2 - s1*s3,  s1*c3*c2 + c1*s3, -s2*c3,
                   -c1*c2*s3 - s1*c3, -s1*c2*s3 + c1*c3,  s2*s3,
                           c1*s2    ,         s1*s2    ,    c2)'
    else
        throw(ArgumentError("The rotation sequence :$rot_seq is not valid."))
    end
end

"""
    function angle2dcm(eulerang::EulerAngles)

Convert the Euler angles `eulerang` (see `EulerAngles`) to a direction cosine
matrix.

# Returns

The direction cosine matrix.

# Remarks

This function assigns `dcm = A3 * A2 * A1` in which `Ai` is the DCM related with
the *i*-th rotation, `i Є [1,2,3]`.

# Example

```julia-repl
julia> angle2dcm(EulerAngles(pi/2, pi/3, pi/4, :ZYX))
3×3 StaticArrays.SArray{Tuple{3,3},Float64,2,9}:
  3.06162e-17  0.5       -0.866025
 -0.707107     0.612372   0.353553
  0.707107     0.612372   0.353553
```

"""
angle2dcm(eulerang::EulerAngles) = angle2dcm(eulerang.a1,
                                             eulerang.a2,
                                             eulerang.a3,
                                             eulerang.rot_seq)

"""
    function smallangle2dcm(θx::Number, θy::Number, θz::Number)

Create a direction cosine matrix from three small rotations of angles `θx`,
`θy`, and `θz` [rad] about the axes X, Y, and Z, respectively.

# Remarks

No process of ortho-normalization is performed with the computed DCM.

# Example

```julia-repl
julia> smallangle2dcm(+0.01, -0.01, +0.01)
3×3 StaticArrays.SArray{Tuple{3,3},Float64,2,9}:
  1.0    0.01  0.01
 -0.01   1.0   0.01
 -0.01  -0.01  1.0
```

"""
smallangle2dcm(θx::Number, θy::Number, θz::Number) = DCM(  1, +θz, -θy,
                                                         -θz,   1, +θx,
                                                         +θy, -θx,   1)'

# Quaternion
# ==============================================================================

"""
    function angle2quat(angle_r1::Number, angle_r2::Number, angle_r3::Number, rot_seq::AbstractString="ZYX")

Convert the Euler angles `angle_r1`, `angle_r2`, and `angle_r3` [rad] with the
rotation sequence `rot_seq` to a quaternion.

The rotation sequence is defined by a `:Symbol`. The possible values are:
`:XYX`, `XYZ`, `:XZX`, `:XZY`, `:YXY`, `:YXZ`, `:YZX`, `:YZY`, `:ZXY`, `:ZXZ`,
`:ZYX`, and `:ZYZ`. If no value is specified, then it defaults to `:ZYX`.

# Remarks

This function assigns `q = q1 * q2 * q3` in which `qi` is the quaternion related
with the *i*-th rotation, `i Є [1,2,3]`.

# Example

```julia-repl
julia> angle2quat(pi/2, pi/3, pi/4, :ZYX)
Quaternion{Float64}:
  + 0.7010573846499779 - 0.09229595564125714.i + 0.5609855267969309.j + 0.43045933457687935.k
```

"""
function angle2quat(angle_r1::Number,
                    angle_r2::Number,
                    angle_r3::Number,
                    rot_seq::Symbol = :ZYX)

    # Compute the sines and cosines of half angle.
    c1 = cos(angle_r1/2)
    s1 = sin(angle_r1/2)

    c2 = cos(angle_r2/2)
    s2 = sin(angle_r2/2)

    c3 = cos(angle_r3/2)
    s3 = sin(angle_r3/2)

    if rot_seq == :ZYX
        q0 = c1*c2*c3 + s1*s2*s3

        s = (q0 < 0) ? -1 : +1

        return Quaternion(s*q0,
                          s*(c1*c2*s3 - s1*s2*c3),
                          s*(c1*s2*c3 + s1*c2*s3),
                          s*(s1*c2*c3 - c1*s2*s3))
    elseif rot_seq == :XYX
        q0 = c1*c2*c3 - s1*c2*s3

        s = (q0 < 0) ? -1 : +1

        return Quaternion(s*q0,
                          s*(c1*c2*s3 + s1*c2*c3),
                          s*(c1*s2*c3 + s1*s2*s3),
                          s*(s1*s2*c3 - c1*s2*s3))
    elseif rot_seq == :XYZ
        q0 = c1*c2*c3 - s1*s2*s3

        s = (q0 < 0) ? -1 : +1

        return Quaternion(s*q0,
                          s*(s1*c2*c3 + c1*s2*s3),
                          s*(c1*s2*c3 - s1*c2*s3),
                          s*(c1*c2*s3 + s1*s2*c3))
    elseif rot_seq == :XZX
        q0 = c1*c2*c3 - s1*c2*s3

        s = (q0 < 0) ? -1 : +1

        return Quaternion(s*q0,
                          s*(c1*c2*s3 + s1*c2*c3),
                          s*(c1*s2*s3 - s1*s2*c3),
                          s*(c1*s2*c3 + s1*s2*s3))
    elseif rot_seq == :XZY
        q0 = c1*c2*c3 + s1*s2*s3

        s = (q0 < 0) ? -1 : +1

        return Quaternion(s*q0,
                          s*(s1*c2*c3 - c1*s2*s3),
                          s*(c1*c2*s3 - s1*s2*c3),
                          s*(c1*s2*c3 + s1*c2*s3))
    elseif rot_seq == :YXY
        q0 = c1*c2*c3 - s1*c2*s3

        s = (q0 < 0) ? -1 : +1

        return Quaternion(s*q0,
                          s*(c1*s2*c3 + s1*s2*s3),
                          s*(c1*c2*s3 + s1*c2*c3),
                          s*(c1*s2*s3 - s1*s2*c3))
    elseif rot_seq == :YXZ
        q0 = c1*c2*c3 + s1*s2*s3

        s = (q0 < 0) ? -1 : +1

        return Quaternion(s*q0,
                          s*(c1*s2*c3 + s1*c2*s3),
                          s*(s1*c2*c3 - c1*s2*s3),
                          s*(c1*c2*s3 - s1*s2*c3))
    elseif rot_seq == :YZX
        q0 = c1*c2*c3 - s1*s2*s3

        s = (q0 < 0) ? -1 : +1

        return Quaternion(s*q0,
                          s*(c1*c2*s3 + s1*s2*c3),
                          s*(s1*c2*c3 + c1*s2*s3),
                          s*(c1*s2*c3 - s1*c2*s3))
    elseif rot_seq == :YZY
        q0 = c1*c2*c3 - s1*c2*s3

        s = (q0 < 0) ? -1 : +1

        return Quaternion(s*q0,
                          s*(s1*s2*c3 - c1*s2*s3),
                          s*(c1*c2*s3 + s1*c2*c3),
                          s*(c1*s2*c3 + s1*s2*s3))
    elseif rot_seq == :ZXY
        q0 = c1*c2*c3 - s1*s2*s3

        s = (q0 < 0) ? -1 : +1

        return Quaternion(s*q0,
                          s*(c1*s2*c3 - s1*c2*s3),
                          s*(c1*c2*s3 + s1*s2*c3),
                          s*(s1*c2*c3 + c1*s2*s3))
    elseif rot_seq == :ZXZ
        q0 = c1*c2*c3 - s1*c2*s3

        s = (q0 < 0) ? -1 : +1

        return Quaternion(s*q0,
                          s*(c1*s2*c3 + s1*s2*s3),
                          s*(s1*s2*c3 - c1*s2*s3),
                          s*(c1*c2*s3 + s1*c2*c3))
    elseif rot_seq == :ZYZ
        q0 = c1*c2*c3 - s1*c2*s3

        s = (q0 < 0) ? -1 : +1

        return Quaternion(s*q0,
                          s*(c1*s2*s3 - s1*s2*c3),
                          s*(c1*s2*c3 + s1*s2*s3),
                          s*(c1*c2*s3 + s1*c2*c3))
    else
        throw(ArgumentError("The rotation sequence :$rot_seq is not valid."))
    end
end

"""
    function angle2quat(eulerang::EulerAngles)

Convert the Euler angles `eulerang` (see `EulerAngles`) to a quaternion.

# Remarks

This function assigns `q = q1 * q2 * q3` in which `qi` is the quaternion related
with the *i*-th rotation, `i Є [1,2,3]`.

# Example

```julia-repl
julia> angle2quat(pi/2, pi/3, pi/4, :ZYX)
Quaternion{Float64}:
  + 0.7010573846499779 - 0.09229595564125714.i + 0.5609855267969309.j + 0.43045933457687935.k
```

"""
angle2quat(eulerang::EulerAngles) = angle2quat(eulerang.a1,
                                               eulerang.a2,
                                               eulerang.a3,
                                               eulerang.rot_seq)

"""
    function smallangle2quat(θx::Number, θy::Number, θz::Number)

Create a quaternion from three small rotations of angles `θx`, `θy`, and `θz`
[rad] about the axes X, Y, and Z, respectively.

# Remarks

The quaternion is normalized.

# Example

```julia-repl
julia> smallangle2quat(+0.01, -0.01, +0.01)
Quaternion{Float64}:
  + 0.9999625021092433 + 0.004999812510546217.i - 0.004999812510546217.j + 0.004999812510546217.k
```

"""
function smallangle2quat(θx::Number, θy::Number, θz::Number)
    q0     = 1
    q1     = θx/2
    q2     = θy/2
    q3     = θz/2
    norm_q = sqrt(q0^2 + q1^2 + q2^2 + q3^2)

    Quaternion(q0/norm_q, q1/norm_q, q2/norm_q, q3/norm_q)
end

################################################################################
#                                     API
################################################################################

"""
    function angle2rot([T,] angle_r1::Number, angle_r2::Number, angle_r3::Number, rot_seq::Symbol = :ZYX)

Convert the Euler angles `angle_r1`, `angle_r2`, and `angle_r3` [rad] with the
rotation sequence `rot_seq` to a rotation description of type `T`, which can be
`DCM` or `Quaternion`. If the type `T` is not specified, then it defaults to
`DCM`.

The rotation sequence is defined by a `:Symbol`. The possible values are:
`:XYX`, `XYZ`, `:XZX`, `:XZY`, `:YXY`, `:YXZ`, `:YZX`, `:YZY`, `:ZXY`, `:ZXZ`,
`:ZYX`, and `:ZYZ`. If no value is specified, then it defaults to `:ZYX`.

# Example

```julia-repl
julia> dcm = angle2rot(pi/2, pi/3, pi/4, :ZYX)
3×3 StaticArrays.SArray{Tuple{3,3},Float64,2,9}:
  3.06162e-17  0.5       -0.866025
 -0.707107     0.612372   0.353553
  0.707107     0.612372   0.353553

julia> q   = angle2rot(Quaternion,pi/2, pi/3, pi/4, :ZYX)
Quaternion{Float64}:
  + 0.7010573846499779 - 0.09229595564125714.i + 0.5609855267969309.j + 0.43045933457687935.k
```

"""
@inline angle2rot(θx::Number, θy::Number, θz::Number, rot_seq::Symbol) =
    angle2dcm(θx, θy, θz, rot_seq)

@inline angle2rot(::Type{DCM},
                  θx::Number,
                  θy::Number,
                  θz::Number,
                  rot_seq::Symbol) =
    angle2dcm(θx, θy, θz, rot_seq)

@inline angle2rot(::Type{Quaternion},
                  θx::Number,
                  θy::Number,
                  θz::Number,
                  rot_seq::Symbol) =
    angle2quat(θx, θy, θz, rot_seq)

"""
    function angle2rot([T,] angle_r1::Number, angle_r2::Number, angle_r3::Number, rot_seq::Symbol = :ZYX)

Convert the Euler angles `eulerang` (see `EulerAngles`) to a rotation
description of type `T`, which can be `DCM` or `Quaternion`. If the type `T` is
not specified, then it defaults to `DCM`.

# Example

```julia-repl
julia> dcm = angle2rot(EulerAngles(pi/2, pi/3, pi/4, :ZYX))
3×3 StaticArrays.SArray{Tuple{3,3},Float64,2,9}:
  3.06162e-17  0.5       -0.866025
 -0.707107     0.612372   0.353553
  0.707107     0.612372   0.353553

julia> q   = angle2rot(Quaternion,EulerAngles(pi/2, pi/3, pi/4, :ZYX))
Quaternion{Float64}:
  + 0.7010573846499779 - 0.09229595564125714.i + 0.5609855267969309.j +
  0.43045933457687935.k
```

"""
@inline angle2rot(eulerang::EulerAngles) =
    angle2dcm(eulerang.a1, eulerang.a2, eulerang.a3, eulerang.rot_seq)

@inline angle2rot(::Type{DCM}, eulerang::EulerAngles) =
    angle2dcm(eulerang.a1, eulerang.a2, eulerang.a3, eulerang.rot_seq)

@inline angle2rot(::Type{Quaternion}, eulerang::EulerAngles) =
    angle2quat(eulerang.a1, eulerang.a2, eulerang.a3, eulerang.rot_seq)

"""
    function smallangle2rot([T,] θx::Number, θy::Number, θz::Number)

Create a rotation description of type `T` from three small rotations of angles
`θx`, `θy`, and `θz` [rad] about the axes X, Y, and Z, respectively.

The type `T` of the rotation description can be `DCM` or `Quaternion`. If the
type `T` is not specified, then if defaults to `DCM`.

# Returns

The rotation description according to the type `T`.

# Example

```julia-repl
julia> dcm = smallangle2rot(+0.01, -0.01, +0.01);

julia> q   = smallangle2rot(Quaternion,+0.01, -0.01, +0.01);

julia> dcm = smallangle2rot(+0.01, -0.01, +0.01)
3×3 StaticArrays.SArray{Tuple{3,3},Float64,2,9}:
  1.0    0.01  0.01
 -0.01   1.0   0.01
 -0.01  -0.01  1.0

julia> q   = smallangle2rot(Quaternion,+0.01, -0.01, +0.01)
Quaternion{Float64}:
  + 0.9999625021092433 + 0.004999812510546217.i - 0.004999812510546217.j + 0.004999812510546217.k
```

"""
@inline smallangle2rot(θx::Number, θy::Number, θz::Number) =
    smallangle2dcm(θx, θy, θz)

@inline smallangle2rot(::Type{DCM}, θx::Number, θy::Number, θz::Number) =
    smallangle2dcm(θx, θy, θz)

@inline smallangle2rot(::Type{Quaternion}, θx::Number, θy::Number, θz::Number) =
    smallangle2quat(θx, θy, θz)


using JuMP, VariationalInequality

# write your own tests here

m = VIPModel()

A = [25; 25; 75; 25; 25]
B = [0.010; 0.010; 0.001; 0.010; 0.010]
T14 = 100
p = 3

@variable(m, h[i=1:p] >= 0)

@NLexpression(m, F1, A[1]+B[1]*h[1]^2 + A[4]+B[4]*(h[1]+h[2])^2 )
@NLexpression(m, F2, A[2]+B[2]*(h[2]+h[3])^2 + A[3]+B[3]*h[2]^2 + A[4]+B[4]*(h[1]+h[2])^2 )
@NLexpression(m, F3, A[2]+B[2]*(h[2]+h[3])^2 + A[5]+B[5]*(h[3])^2 )

@NLconstraint(m, sum{h[i], i=1:p} == T14)

F = [F1, F2, F3]
correspond(m, F, h)
sol, Fval, gap = solveVIP(m, algorithm=:extra_gradient, max_iter=1000, step_size=0.01)

@assert 0<= gap1 < 1e-6

@show sol
@show gap


#=

# 203: Various coordinate systems
([source code](SOURCE_URL))
=#

module Example203_CoordinateSystems

using VoronoiFVM
using LinearAlgebra
using ExtendableGrids
using GridVisualize

function main(;nref=0,r1=0.0, r2=5.0, dim=2,Plotter=nothing)
    h=0.1*2.0^(-nref)
    R=collect(r1:h:r2)
    Z=collect(0:h:2)
    grid=VoronoiFVM.Grid(R)
    if dim==2
        grid=VoronoiFVM.Grid(R,Z)
    end
    circular_symmetric!(grid)

    function symlapcyl(r)
        if r1≈0.0
            return 0.25*(r2*r2-r*r);
        else
            return (log(r)-log(r2))/(log(r1)-log(r2));
        end
    end

    function flux!(f,u0,edge)
        u=unknowns(edge,u0)
        f[1]=u[1,1]-u[1,2]
    end

    function source!(f,node)
        if r1≈0.0
            f[1]=1.0
        else
            f[1]=0.0
        end
    end

    # Create a physics structure
    physics=VoronoiFVM.Physics(num_species=1,flux=flux!,source=source!)
    sys=VoronoiFVM.System(grid,physics,unknown_storage=:dense)
    ispec=1
    enable_species!(sys,ispec,[1])
    ileft=1
    if dim==2
        ileft=4
    end
    if !(r1≈0.0)
        boundary_dirichlet!(sys,ispec,ileft,1.0)
    end
    boundary_dirichlet!(sys,ispec,2,0.0)

    inival=unknowns(sys)
    solution=unknowns(sys)
    inival.=0
    solution.=0

    # Solve stationary problem
    solve!(solution,inival,sys)

    scalarplot(grid,solution[1,:],Plotter=Plotter)

    exact=symlapcyl.(coordinates(grid)[1,:])
    err=norm(solution[1,:]-exact,Inf)
end

#
# Called by unit test
#
function test()
    main() <1.0e-14
end

end


using ITensors,
      LinearAlgebra, # For tr()
      Random,        # To set a seed
      Test

Random.seed!(12345)

digits(::Type{T},x...) where {T} = T(sum([x[length(x)-k+1]*10^(k-1) for k=1:length(x)]))

@testset "Dense ITensor basic functionality" begin

@testset "ITensor constructors" begin
  i = Index(2,"i")
  j = Index(2,"j")
  k = Index(2,"k")

  @testset "Default" begin
    A = ITensor()
    @test store(A) isa Dense{Nothing}
    @test isNull(A)
  end

  @testset "Undef with index" begin
    A = ITensor(undef, i)
    @test store(A) isa Dense{Float64}
    @test !isNull(A)
  end

  @testset "Default with indices" begin
    A = ITensor(i,j)
    @test store(A) isa Dense{Float64}
    @test !isNull(A)
  end

  @testset "Random" begin
    A = randomITensor(i,j)
    @test store(A) isa Dense{Float64}

    @test ndims(A) == order(A) == 2 == length(inds(A))
    @test size(A) == dims(A) == (2,2)

    @test !isNull(A)

    B = randomITensor(IndexSet(i,j))
    @test store(B) isa Dense{Float64}
    @test ndims(B) == order(B) == 2 == length(inds(B))
    @test size(B) == dims(B) == (2,2)
    @test !isNull(B)
  end

  @testset "From matrix" begin
    M = [1 2; 3 4]
    A = ITensor(M,i,j)
    @test store(A) isa Dense{Float64}

    @test M ≈ matrix(A,i,j)
    @test M' ≈ matrix(A,j,i)
    @test_throws DimensionMismatch vector(A)

    @test size(A,1) == size(M,1) == 2
    @test size(A,3) == size(M,3) == 1
    @test_throws ErrorException size(A,0)
    @test_throws ErrorException size(M,0)

    M = [1 2 3; 4 5 6]
    @test_throws DimensionMismatch ITensor(M,i,j)
  end

  @testset "To Matrix" begin
    TM = randomITensor(i,j)

    M1 = matrix(TM)
    for ni in i, nj in j
      @test M1[ni,nj] ≈ TM[i(ni),j(nj)]
    end

    M2 = matrix(TM,j,i)
    for ni in i, nj in j
      @test M2[nj,ni] ≈ TM[i(ni),j(nj)]
    end

    T3 = randomITensor(i,j,k)
    @test_throws DimensionMismatch matrix(T3,i,j)
  end

  @testset "To Vector" begin
    TV = randomITensor(i)

    V = vector(TV)
    for ni in i
      @test V[ni] ≈ TV[i(ni)]
    end

    T2 = randomITensor(i,j)
    @test_throws DimensionMismatch vector(T2)
  end

  @testset "Complex" begin
    A = ITensor(Complex,i,j)
    @test store(A) isa Dense{ComplexF64}
  end

  @testset "Random complex" begin
    A = randomITensor(Complex,i,j)
    @test store(A) isa Dense{ComplexF64}
  end

  @testset "From complex matrix" begin
    M = [1+2im 2; 3 4]
    A = ITensor(M,i,j)
    @test store(A) isa Dense{ComplexF64}
  end

end

@testset "Convert to complex" begin
  i = Index(2,"i")
  j = Index(2,"j")
  A = randomITensor(i,j)
  B = complex(A)
  for ii ∈ dim(i), jj ∈ dim(j)
    @test complex(A[i(ii),j(jj)]) == B[i(ii),j(jj)]
  end
end

@testset "similar" begin
  i = Index(2,"i")
  j = Index(2,"j")
  A = randomITensor(i,j)
  B = similar(A)
  @test inds(B) == inds(A)
  Ac = similar(A, ComplexF32)
  @test store(Ac) isa Dense{ComplexF32}
end

@testset "fill!" begin
  i = Index(2,"i")
  j = Index(2,"j")
  A = randomITensor(i,j)
  fill!(A, 1.0)
  @test all(data(store(A)) .== 1.0)
end

@testset "copyto!" begin
  i = Index(2,"i")
  j = Index(2,"j")
  M = [1 2; 3 4]
  A = ITensor(M,i,j)
  N = 2*M
  B = ITensor(N,i,j)
  copyto!(A, B)
  @test A == B
  @test data(store(A)) == vec(N)
  A = ITensor(M,i,j)
  B = ITensor(N,j,i)
  copyto!(A, B)
  @test A == B
  @test data(store(A)) == vec(transpose(N))
end

@testset "Unary -" begin
  i = Index(2,"i")
  j = Index(2,"j")
  M = [1 2; 3 4]
  A = ITensor(M,i,j)
  @test -A == ITensor(-M, i, j)
end

@testset "dot" begin
  i = Index(2,"i")
  a = [1.0; 2.0]
  b = [3.0; 4.0]
  A = ITensor(a,i)
  B = ITensor(b,i)
  @test dot(A, B) == 11.0
end

@testset "exponentiate" begin
  s1 = Index(2,"s1")
  s2 = Index(2,"s2")
  i1 = Index(2,"i1")
  i2 = Index(2,"i2")
  Amat = rand(2,2,2,2)
  A = ITensor(Amat,i1,i2,s1,s2)

  Aexp = exp(A,(i1,i2),(s1,s2))
  Amatexp = reshape(exp(reshape(Amat,4,4)),2,2,2,2)
  Aexp_from_mat = ITensor(Amatexp,i1,i2,s1,s2)
  @test Aexp ≈ Aexp_from_mat

  #test that exponentiation works when indices need to be permuted
  Aexp = exp(A,(s1,s2),(i1,i2))
  Amatexp = Matrix(exp(reshape(Amat,4,4))')
  Aexp_from_mat = ITensor(reshape(Amatexp,2,2,2,2),s1,s2,i1,i2)
  @test Aexp ≈ Aexp_from_mat

  #test exponentiation when hermitian=true is used
  Amat = reshape(Amat, 4,4)
  Amat = reshape(Amat+Amat'+randn(4,4)*1e-10,2,2,2,2)
  A = ITensor(Amat,i1,i2,s1,s2)
  Aexp = exp(A,(i1,i2),(s1,s2),ishermitian=true)
  Amatexp = reshape(parent(exp(Hermitian(reshape(Amat,4,4)))),
                    2,2,2,2)
  Aexp_from_mat = ITensor(Amatexp,i1,i2,s1,s2)
  @test Aexp ≈ Aexp_from_mat
end

@testset "add and axpy" begin
  i = Index(2,"i")
  a = [1.0; 2.0]
  b = [3.0; 4.0]
  A = ITensor(a,i)
  B = ITensor(b,i)
  c = [5.0; 8.0]
  @test A + B == ITensor([4.0; 6.0], i)
  @test axpy!(2.0, A, B) == ITensor(c, i)
  a = [1.0; 2.0]
  b = [3.0; 4.0]
  A = ITensor(a,i)
  B = ITensor(b,i)
  c = [8.0; 12.0]
  @test add!(A, 2.0, 2.0, B) == ITensor(c, i)

end

@testset "mul! and rmul!" begin
  i = Index(2,"i")
  a = [1.0; 2.0]
  b = [2.0; 4.0]
  A = ITensor(a,i)
  A2, A3 = copy(A), copy(A)
  B = ITensor(b,i)
  @test mul!(A2, A, 2.0) == B == ITensors.add!(A2, 0, 2, A)
  @test rmul!(A, 2.0) == B == ITensors.scale!(A3, 2)

  i = Index(2,"i")
  j = Index(2,"j")
  M = [1 2; 3 4]
  A = ITensor(M,i,j)
  N = 2*M
  B = ITensor(N,j,i)
  @test data(store(mul!(B, A, 2.0))) == 2.0*vec(transpose(M))
end

@testset "show" begin
  i = Index(2,"i")
  a = [1.0; 2.0]
  A = ITensor(a,i)
  s = split(sprint(show, A), '\n')
  @test s[1] == "ITensor ord=1 " * sprint(show, i) * " "
  @test s[2] == "Dense{Float64}"
  @test s[3] == " 1.0"
  @test s[4] == " 2.0"
end

@testset "Test isapprox for ITensors" begin
  m,n = rand(0:20,2)
  i = Index(m)
  j = Index(n)
  realData = rand(m,n)
  complexData = complex(realData)
  A = ITensor(realData, i,j)
  B = ITensor(complexData, i,j)
  @test A≈B
  @test B≈A
  A = permute(A,j,i)
  @test A≈B
  @test B≈A
end

@testset "ITensor tagging and priming" begin
  s1 = Index(2,"Site,s=1")
  s2 = Index(2,"Site,s=2")
  l = Index(3,"Link")
  A1 = randomITensor(s1,l,l')
  A2 = randomITensor(s2,l',l'')
  @testset "findindex(::ITensor,::String)" begin
    @test s1==findindex(A1,"Site")
    @test s1==findindex(A1,"s=1")
    @test s1==findindex(A1,"s=1,Site")
    @test l==findindex(A1,("Link",0))
    @test l'==findindex(A1,("",1))
    @test l'==findindex(A1,("Link",1))
    @test s2==findindex(A2,"Site")
    @test s2==findindex(A2,"s=2")
    @test s2==findindex(A2,"Site")
    @test s2==findindex(A2,("",0))
    @test s2==findindex(A2,("s=2",0))
    @test s2==findindex(A2,("Site",0))
    @test s2==findindex(A2,("s=2,Site",0))
    @test l'==findindex(A2,("",1))
    @test l'==findindex(A2,("Link",1))
    @test l''==findindex(A2,("",2))
    @test l''==findindex(A2,("Link",2))
  end
  @testset "addtags(::ITensor,::String,::String)" begin
    s1u = addtags(s1,"u")
    lu = addtags(l,"u")

    A1u = addtags(A1,"u")
    @test hasinds(A1u,s1u,lu,lu')

    A1u = addtags(A1,"u","Link")
    @test hasinds(A1u,s1,lu,lu')

    A1u = addtags(A1,"u",("",0))
    @test hasinds(A1u,s1u,lu,l')

    A1u = addtags(A1,"u",("Link",0))
    @test hasinds(A1u,s1,lu,l')

    A1u = addtags(A1,"u",("Link",1))
    @test hasinds(A1u,s1,l,lu')
  end
  @testset "removetags(::ITensor,::String,::String)" begin
    A2r = removetags(A2,"Site")
    @test hasinds(A2r,removetags(s2,"Site"),l',l'')

    A2r = removetags(A2,"Link",("",1))
    @test hasinds(A2r,s2,removetags(l,"Link")',l'')
  end
  @testset "replacetags(::ITensor,::String,::String)" begin
    s2tmp = replacetags(s2,"Site","Temp")
    ltmp = replacetags(l,"Link","Temp")

    A2r = replacetags(A2,"Site","Temp")
    @test hasinds(A2r,s2tmp,l',l'')

    A2r = replacetags(A2,"Link","Temp")
    @test hasinds(A2r,s2,ltmp',ltmp'')

    A2r = replacetags(A2,"Link","Temp",("",1))
    @test hasinds(A2r,s2,ltmp',l'')

    A2r = replacetags(A2,("Link",2),("Temp",3))
    @test hasinds(A2r,s2,l',ltmp''')

    A2r = replacetags(A2,("",1),("",5))
    @test hasinds(A2r,s2,prime(l,5),l'')

    #In-place version
    cA2 = copy(A2)
    replacetags!(cA2,("",1),("",5))
    @test hasinds(cA2,s2,prime(l,5),l'')
  end
  @testset "prime(::ITensor,::String)" begin
    A2p = prime(A2)
    @test A2p==A2'
    @test hasinds(A2p,s2',l'',l''')

    A2p = prime(A2,2)
    A2p = A2''
    @test hasinds(A2p,s2'',l''',l'''')

    A2p = prime(A2,"s=2")
    @test hasinds(A2p,s2',l',l'')
  end

  @testset "mapprime" begin
    @test hasinds(mapprime(A2,1,7),s2,l^7,l'')
    @test hasinds(mapprime(A2,0,1),s2',l',l'')
  end
  @testset "setprime" begin
    @test hasinds(setprime(A2,2,s2),s2'',l',l'')
    @test hasinds(setprime(A2,0,l''),s2,l',l)
  end
  @testset "swapprime" begin
    @test hasinds(swapprime(A2,1,3),l''',s2,l'')
  end
end

@testset "Converting Real and Complex Storage" begin

  @testset "Add Real and Complex" begin
    i = Index(2,"i")
    j = Index(2,"j")
    TC = randomITensor(ComplexF64,i,j)
    TR = randomITensor(Float64,i,j)

    S1 = TC+TR
    S2 = TR+TC
    @test typeof(S1.store) == Dense{ComplexF64}
    @test typeof(S2.store) == Dense{ComplexF64}
    for ii=1:dim(i),jj=1:dim(j)
      @test S1[i(ii),j(jj)] ≈ TC[i(ii),j(jj)]+TR[i(ii),j(jj)]
      @test S2[i(ii),j(jj)] ≈ TC[i(ii),j(jj)]+TR[i(ii),j(jj)]
    end
  end

end

@testset "ITensor, Dense{$SType} storage" for SType ∈ (Float64,ComplexF64)
  mi,mj,mk,ml,mα = 2,3,4,5,6,7
  i = Index(mi,"i")
  j = Index(mj,"j")
  k = Index(mk,"k")
  l = Index(ml,"l")
  α = Index(mα,"alpha")
  @testset "Set and get values with IndexVals" begin
    A = ITensor(SType,i,j,k)
    for ii ∈ 1:dim(i), jj ∈ 1:dim(j), kk ∈ 1:dim(k)
      A[k(kk),j(jj),i(ii)] = digits(SType,ii,jj,kk)
    end
    for ii ∈ 1:dim(i), jj ∈ 1:dim(j), kk ∈ 1:dim(k)
      @test A[j(jj),k(kk),i(ii)]==digits(SType,ii,jj,kk)
    end
    @test_throws MethodError A[1]
  end
  @testset "Test permute(ITensor,Index...)" begin
    A = randomITensor(SType,i,k,j)
    permA = permute(A,k,j,i)
    @test k==inds(permA)[1]
    @test j==inds(permA)[2]
    @test i==inds(permA)[3]
    for ii ∈ 1:dim(i), jj ∈ 1:dim(j), kk ∈ 1:dim(k)
      @test A[k(kk),i(ii),j(jj)]==permA[i(ii),j(jj),k(kk)]
    end
    for ii ∈ 1:dim(i), jj ∈ 1:dim(j), kk ∈ 1:dim(k)
      @test A[k(kk),i(ii),j(jj)]==permA[i(ii),j(jj),k(kk)]
    end
    # TODO: I think this was doing slicing, but what is the output
    # of slicing an ITensor?
    #@testset "getindex and setindex with vector of IndexVals" begin
    #    k_inds = [k(kk) for kk ∈ 1:dim(k)]
    #    for ii ∈ 1:dim(i), jj ∈ 1:dim(j)
    #      @test A[k_inds,i(ii),j(jj)]==permA[i(ii),j(jj),k_inds...]
    #    end
    #    for ii ∈ 1:dim(i), jj ∈ 1:dim(j)
    #        A[k_inds,i(ii),j(jj)]=collect(1:length(k_inds))
    #    end
    #    permA = permute(A,k,j,i)
    #    for ii ∈ 1:dim(i), jj ∈ 1:dim(j)
    #      @test A[k_inds,i(ii),j(jj)]==permA[i(ii),j(jj),k_inds...]
    #    end
    #end
  end
  @testset "Set and get values with Ints" begin
    A = ITensor(SType,i,j,k)
    A = permute(A,k,i,j)
    for ii ∈ 1:dim(i), jj ∈ 1:dim(j), kk ∈ 1:dim(k)
      A[kk,ii,jj] = digits(SType,ii,jj,kk)
    end
    A = permute(A,i,j,k)
    for ii ∈ 1:dim(i), jj ∈ 1:dim(j), kk ∈ 1:dim(k)
      @test A[ii,jj,kk]==digits(SType,ii,jj,kk)
    end
  end
  @testset "Test scalar(ITensor)" begin
    x = SType(34)
    A = ITensor(x)
    @test x==scalar(A)
    A = ITensor(SType,i,j,k)
    @test_throws BoundsError scalar(A)
  end
  @testset "Test norm(ITensor)" begin
    A = randomITensor(SType,i,j,k)
    @test norm(A)≈sqrt(scalar(dag(A)*A))
  end
  @testset "Test add ITensors" begin
    A = randomITensor(SType,i,j,k)
    B = randomITensor(SType,k,i,j)
    C = A+B
    for ii ∈ 1:dim(i), jj ∈ 1:dim(j), kk ∈ 1:dim(k)
      @test C[i(ii),j(jj),k(kk)]==A[j(jj),i(ii),k(kk)]+B[i(ii),k(kk),j(jj)]
    end
    @test array(permute(C,i,j,k))==array(permute(A,i,j,k))+array(permute(B,i,j,k))
  end

  @testset "Test factorizations of an ITensor" begin

    A = randomITensor(SType,i,j,k,l)

    @testset "Test SVD of an ITensor" begin
      U,S,V,u,v = svd(A,(j,l))
      @test store(S) isa Diag{Vector{Float64}}
      @test A≈U*S*V
      @test U*dag(prime(U,u))≈δ(SType,u,u') atol=1e-14
      @test V*dag(prime(V,v))≈δ(SType,v,v') atol=1e-14
    end

    @testset "Test SVD truncation" begin
        ii = Index(4)
        jj = Index(4)
        T = randomITensor(ComplexF64,ii,jj)
        U,S,V = svd(T,ii;maxdim=2)
        u,s,v = svd(matrix(T))
        @test norm(U*S*V-T)≈sqrt(s[3]^2+s[4]^2)
    end

    @testset "Test QR decomposition of an ITensor" begin
      Q,R,q = qr(A,(i,l))
      q = commonindex(Q,R)
      @test A≈Q*R
      @test Q*dag(prime(Q,q))≈δ(SType,q,q') atol=1e-14
    end

    @testset "Test polar decomposition of an ITensor" begin
      U,P,u = polar(A,(k,l))
      @test A≈U*P
      #Note: this is only satisfied when left dimensions
      #are greater than right dimensions
      UUᵀ =  U*dag(prime(U,u))

      # TODO: use a combiner to combine the u indices to make
      # this test simpler
      for ii ∈ 1:dim(u[1]), jj ∈ 1:dim(u[2]), iip ∈ 1:dim(u[1]), jjp ∈ 1:dim(u[2])
        val = UUᵀ[u[1](ii),u[2](jj),u[1]'(iip),u[2]'(jjp)]
        if ii==iip && jj==jjp
          @test val≈one(SType) atol=1e-14
        else
          @test val≈zero(SType) atol=1e-14
        end
      end
    end

    @testset "Test Hermitian eigendecomposition of an ITensor" begin
      is = IndexSet(i,j)
      T = randomITensor(is...,prime(is)...)
      T = T + swapprime(dag(T),0,1)
      U,D,u = eigenHermitian(T)
      @test T ≈ U*D*prime(dag(U))
      UUᴴ =  U*prime(dag(U),u)
      @test UUᴴ ≈ δ(u,u') atol=1e-14
    end

  end # End ITensor factorization testset

end # End Dense storage test

end # End Dense ITensor basic functionality


########## Utility functions ##########

"""
Remove columns or rows containing nans.
"""
function remove_nan(X::AbstractVecOrMat; dims::Integer=1)
    if ndims(X) == 1
        idx = findall(.!isnan.(X))
        return view(X, idx), idx
    else
        idx = findall(vec(.!any(isnan.(X), dims=dims)))  # on a vector sum(X, dims=2) gives X
        return dims==1 ? view(X, :, idx) : view(X, idx, :), idx
    end
end

"""
    emb(x, idx, N, v=NaN)

Embedding of a vector `x` in a bigger vector of length `N` at the indices `idx`. The other indices are filled by `v`.
"""
function emb(x::AbstractVector{T}, idx::AbstractVector{<:Integer}, N::Integer, v::T=NaN) where {T<:Number}
    @assert length(x) == length(idx)
    @assert all(0 .< idx .<= N)
    y = zeros(eltype(x), N) * v
    y[idx] .= x
    return y
end

function emb(x::AbstractVector{<:Number}, mask::AbstractVector{Bool}, args...)
    idx = findall(mask)
    return emb(x, findall(mask), length(mask), args...)
    # @assert length(x) <= length(idx)
    # y = zeros(T, length(mask))
    # y[idx] = x
    # return y
end

"""
Return indexes of common elements of two vectors.
"""
function common_elements(tx::AbstractVector, ty::AbstractVector)
    tc = intersect(tx, ty)
    idx = vcat([findall(t.==tx) for t in tc]...)
    idy = vcat([findall(t.==ty) for t in tc]...)
    return idx, idy
end

"""
Spike train with logarithmic density.
"""
function logtrain(w::Integer, n::Integer)
    V = floor.(Int, w.^((0:n-1)./n))
    U = zeros(Bool, w)
    U[V] .= true
    return U
end

##### Algebra #####

function shrinkage_by_value(X0::AbstractArray{<:Number}, v::Real, mode::Symbol=:soft)
    X1 = fill(zero(T), size(X0))  # or zero(X0)
    idx = findall(abs.(X0) .> v)
    if mode == :soft
        X1[idx] .= sign.(X0[idx]) .* (abs.(X0[idx]) .- v)
    else
        X1[idx] .= X0[idx]
    end
    return X1
end

function shrinkage_by_number(X0::AbstractArray{<:Number}, n::Integer, mode::Symbol=:soft)
    Xv = ndims(X0)>1 ? vec(X0) : X0
    X1 = fill(zero(T), length(Xv))  # or zero(X0)
    idx = sortperm(abs.(Xv))[end-n+1:end]  # increasing order
    if mode == :soft
        X1[idx] .= sign.(Xv[idx]) .* (abs.(Xv[idx]) .- Xv[idx[1]])
    else
        X1[idx] .= Xv[idx]
    end
    return reshape(X1, size(X0))
end

function shrinkage_by_percentage(X0::AbstractArray{<:Number}, p::Real, mode::Symbol=:soft)
    @assert 0 <= p <= 1
    return shrinkage_by_number(X0, floor(Int,p*length(X0)), mode)
end

"""
Sigmoid function.
"""
sigmoid(α::Real) = exp(α)/(1+exp(α))

"""
Derivative of sigmoid function.
"""
diff_sigmoid(α::Real) = exp(α)/(1+2*exp(α)+exp(2α))

"""
Remove singular dimensions of an array.

# Note
This function is safe for array of size (1,1,..1).
"""
function squeezedims(X::AbstractArray{<:Real}; dims::Union{Int,AbstractVector{<:Integer}})
    dimx = tuple(intersect(tuple(findall(size(X).==1)...), dims)...)
    return if length(dimx) == 0
        X
    elseif length(dimx) == ndims(X)
        vec(X)
    else
        dropdims(X, dims=dimx)
    end
end

# function squeezedims(A::AbstractArray{T}) where {T<:Real}
#     dims = tuple(findall(size(A).==1)...)
#     return length(dims)>0 ? dropdims(A, dims=dims) : A
# end

ifloor(x::Int, y::Int) = floor(Int, x/y) * y
# ifloor(x::Int, y::Int) = x - (x%y)
iceil(x::Int, y::Int) = ceil(Int, x/y) * y

"""
Reshape a vector `x` to a matrix of `r` rows with truncation if `length(x)` does not divide `r`.
"""
function vec2mat(x::AbstractVector{<:Real}, r::Integer; keep::Symbol=:tail)
    @assert r<=length(x)  # number of row must be smaller than the size of x
    n = length(x) - (length(x)%r)
    return if keep == :head
        reshape(x[1:n], r, :)
    elseif keep == :tail
        reshape(x[end-n+1:end], r, :)
    else
        error("Unknown position: $(keep)")
    end
end

function norm(X::AbstractMatrix{<:Number}, p::Real=2; dims::Integer=1)
    if dims==1
        return [LinearAlgebra.norm(X[:,n],p) for n=1:size(X,2)]
    else
        return [LinearAlgebra.norm(X[n,:],p) for n=1:size(X,1)]
    end
end

"""
Compute A^-1 * B using the lsqr iterative method.
"""
function lsqr(A::AbstractMatrix{<:Real}, B::AbstractMatrix{<:Real}; kwargs...)
    # println("My lsqr")
    X = zeros(Float64, (size(A,2), size(B,2)))
    for n=1:size(B,2)
        X[:,n] = lsqr(A, B[:,n])
    end
    return X
end

"""
Diagonalization of a symmetric positive definite Toeplitz matrix using Levinson-Durbin (LD) method by providing `cseq`,
the covariance sequence of a stationary process.

# Returns
- `pseq::Vector{Vector{Float64}}`: linear prediction coefficients
- `sseq`: variances of residual
- `rseq`: partial correlation coefficients

# Explanation
`pseq` forms the lower triangular matrix diagonalizing the covariance matrix Γ, and `sseq` forms the resulting diagonal matrix. `rseq[n]` is just `pseq[n][n]`.
"""
function LevinsonDurbin(cseq::AbstractVector{<:Real})
    N = length(cseq)

    if N > 1
        # check that cseq is a validate covariance sequence
        @assert cseq[1] > 0
        @assert all(abs.(cseq[2:end]) .<= cseq[1])
        # @assert all(diff(abs.(cseq)) .<= 0)

        # initialization
        # pseq: linear prediction coefficients
        pseq = Vector{Vector{Float64}}(undef, N-1); pseq[1] = [cseq[2]/cseq[1]]
        # sseq: variances of residual
        sseq = zeros(N); sseq[1] = cseq[1]; sseq[2] = (1-pseq[1][1]^2) * sseq[1]
        # rseq: partial correlation coefficients
        rseq = zeros(N-1); rseq[1] = pseq[1][1]

        # recursive construction of the prediction coefficients and variances
        for n=2:N-1
            pseq[n] = zeros(n)
            pseq[n][n] = (cseq[n+1] - cseq[2:n]' * pseq[n-1][end:-1:1]) / sseq[n]
            pseq[n][1:n-1] = pseq[n-1] - pseq[n][n] * pseq[n-1][end:-1:1]
            sseq[n+1] = (1 - pseq[n][n]^2) * sseq[n]
            rseq[n] = pseq[n][n]
        end
    else
        pseq = Vector{Float64}[]
        sseq = copy(cseq)
        rseq = Float64[]
    end
    return pseq, sseq, rseq
end

LevinsonDurbin(p::StationaryProcess{T}, g::AbstractVector{<:T}) where T<:TimeStyle = LevinsonDurbin(covseq(p, g))

"""
Cholesky decomposition based on SVD.
"""
function chol_svd(W::AbstractMatrix{<:Real})
    Um, Sm, Vm = svd((W+W')/2)  # svd of forced symmetric matrix
    Ss = sqrt.(Sm[Sm.>0])  # truncation of negative singular values
    return Um*diagm(Ss)
end

"""
Vandermonde matrix.
"""
function vandermonde(dim::Tuple{Integer,Integer})
    nrow, ncol = dim
    V = zeros(Float64, dim)
    for c = 1:dim[2]
        V[:,c] = collect((1:dim[1]).^(c-1))
    end
    return V
end

vandermonde(nrow::Integer, ncol::Integer) = vandermonde((nrow, ncol))

function col_normalize(A::AbstractMatrix{<:Real}, p::Real=2)
    return A / diagm([norm(A[:,n], p) for n=1:size(A,2)])
end

function col_normalize!(A::AbstractMatrix{<:Real}, p::Real=2)
    for n=1:size(A,2)
        A[:,n] ./= norm(A[:,n], p)
    end
    return A
end

row_normalize(A) = col_normalize(transpose(A))
row_normalize!(A) = col_normalize!(transpose(A))

"""
    lagdiff(X, d; dims, mode)

Compute `d`-lag finite difference of a numerical array `X` along the dimension `dims` with the `mode` (:causal, :anticausal, :valid).
"""
function lagdiff(X::AbstractArray{<:Number}, d::Integer; dims::Integer=1, mode::Symbol=:causal)
    dX = fill(NaN * one(eltype(X)), size(X))
    s::Integer = if mode == :causal
        d
    elseif mode == :anticausal
        0
    else  # :valid, or :noncausal
        d÷2
    end
    # v = (1+s):(size(dX,dims)-d+s)
    selectdim(dX, dims, (1+s):(size(dX,dims)-d+s)) .= selectdim(collect(X), dims, d+1:size(X,dims)) .- selectdim(collect(X), dims, 1:size(X,dims)-d)

    return dX
    # old version
    # if ndims(X) == 1
    #     dX[d+1:end] = X[d+1:end] .- X[1:end-d]  # '-' or '.-' may raise "step cannot be zero" error on iterators like `X=1:100`, use `collect(X)`
    # else
    #     dX[d+1:end,:] = X[d+1:end,:] .- X[1:end-d,:]
    # end
    # return (mode==:causal) ? dX : circshift(dX, -d)
end

function lagdiff(X::AbstractVector{<:Number}, lags::AbstractVector{<:Integer};mode::Symbol=:causal)
    return hcat([lagdiff(X, d; mode=mode) for d in lags]...)
end

# function lagdiff(X::AbstractMatrix{<:Number}, dlags::AbstractVector{<:Integer}, mode::Symbol=:causal)
#     return reshape(hcat([lagdiff(X, d, mode) for d in dlags]...), (size(X,1),:,length(dlags)))
# end

###### Time series manipulation ######

function ffill!(X::AbstractVector{<:Number})
    for n=2:length(X)
        if ismissing(X[n]) || isnan(X[n])
            X[n] = X[n-1]
        end
    end
    return X
end

ffill(X) = ffill!(copy(X))

function bfill!(X::AbstractVector{<:Number})
    for n=length(X)-1:-1:1
        if ismissing(X[n]) || isnan(X[n])
            X[n] = X[n+1]
        end
    end
    return X
end
bfill(X) = bfill!(copy(X))

fbfill(X) = bfill(ffill(X))
bffill(X) = ffill(bfill(X))

function remove_outliers!(X::AbstractVector{<:Number})
    dX = lagdiff(X, 10, :causal)
    nidc = isnan.(dX)
    μ = median(dX[.!nidc])
    σ = median(abs.(dX[.!nidc] .- μ))
    # μ = mean(dX[.!nidc])
    # σ = std(dX[.!nidc])
    oidc = abs.(dX.-μ) .> 10σ
    X[oidc] .= NaN

    return X
end

"""
Split an object of `TimeArray` by applying truncation.

# Args
- data: input object of `TimeArray`
- T: unit of time period, e.g. `Dates.Day(1)`
- (wa,wb): relative starting and ending time w.r.t. to `T`, e.g., `Dates.Hour(9) + Dates.Minute(5)` for `09:05` and `Dates.Hour(17) + Dates.Minute(24)` for `17:24`.
- fillmode: fill mode for nan values, :o no fill, :f forward, :b backward, :fb forward-backward, :bf backward-forward
- endpoint: if true include the endpoint whenever possible

# Notes
"""
function window_split_timearray(data::TimeArray, T::AbstractTime, (wa,wb)::NTuple{2, Union{Nothing, AbstractTime}}=(nothing, nothing); ubase::Dates.TimePeriod=Dates.Minute(0), fillmode::Symbol=:fb, endpoint::Bool=true)
    stamp = TimeSeries.timestamp(data)  # time stamp
    time_begin, toto = Dates.floorceil(stamp[1], T)
    toto, time_end = Dates.floorceil(stamp[end], T)

    # time unit of data
    unit::Dates.TimePeriod = if ubase==Dates.Minute(0)
        minimum(diff(stamp)) # ÷ Dates.Millisecond(1000)
    else
        ubase
    end

    res = []
    for t in time_begin:T:time_end
        ta = (wa == nothing) ? t : t+wa
        tb = (wb == nothing) ? t+T : t+wb

        y = TimeSeries.to(TimeSeries.from(data,ta),tb)
        if length(y) > 0
            # x0 = (wa == nothing && wb == nothing) ? y : window_timearray(y, ta:unit:tb, fillmode)
            x0 = _window_timearray(y, ta:unit:tb, fillmode)
            x = (TimeSeries.timestamp(x0)[end] == tb && !endpoint) ? x0[1:end-1] : x0
            if length(x) > 0
                push!(res, x)
            end
        end
    end
    return res
end

function _window_timearray(A::TimeArray, tstp::AbstractVector{<:Dates.AbstractTime}, fillmode::Symbol)
    if fillmode == :o
        return A
    else
        cnames = TimeSeries.colnames(A)
        vstp = intersect(tstp, TimeSeries.timestamp(A))  # valid timestamps
        sidx = map(t->findall(isequal(t), tstp)[1], vstp)  # relative index of valid timestamp
        # shape = (ndims(TimeSeries.values(A)) == 1) ? (length(tstp),) : (length(tstp),length(cnames))

        x = fill(NaN, (length(tstp),length(cnames)))
        x[sidx,:] = TimeSeries.values(A[vstp])

        xf = if fillmode == :n
            x
        elseif fillmode==:f
            xf = hcat([ffill(x[:,n]) for n=1:size(x,2)]...)
        elseif fillmode==:b
            xf = hcat([bfill(x[:,n]) for n=1:size(x,2)]...)
        elseif fillmode==:fb
            xf = hcat([bfill(ffill(x[:,n])) for n=1:size(x,2)]...)
        elseif fillmode==:bf
            xf = hcat([ffill(bfill(x[:,n])) for n=1:size(x,2)]...)
        else
            error("Unknown symbol: $(fill)")
        end

        return TimeArray(tstp, (ndims(TimeSeries.values(A))==1) ? vec(xf) : xf, cnames)
    end
end

"""
Equalize the first point of the day with the last point of the previous day.

# Notes
- Useful in processing of stock price
"""
function equalize_daynight(sdata::AbstractVector)
    @assert typeof(sdata[1]) <: TimeArray
    N = ndims(sdata[1])
    @assert N<=2
    edata = [sdata[1]]
    func = (x,y) -> (N==2) ? x[[1],:]-y[[end],:] : x[[1]]-y[[end]]

    for n=2:length(sdata)
        d0 = TimeSeries.values(edata[n-1])
        d1 = TimeSeries.values(sdata[n])
        t1 = TimeSeries.timestamp(sdata[n])
        # push!(edata, TimeSeries.TimeArray(t1, d1.-(d1[1]-d0[end]), TimeSeries.colnames(sdata[1])))  # <- TODO: d0[end] or d1[1] is NaN? d1 is 2d array?
        dv = func(d1,d0)
        # println(typeof(dv))
        dv[isnan.(dv)] .= 0  # <- this makes the operation nan-safe
        push!(edata, TimeSeries.TimeArray(t1, d1.-dv, TimeSeries.colnames(sdata[1])))  # <- TODO: d0[end] or d1[1] is NaN? d1 is 2d array?
    end
    return edata
end

"""
Compute the daily average of a timearray.
"""
function daily_average_timearray(X::TimeArray)
    Xs = window_split_timearray(X, Dates.Day(1); ubase=Dates.Minute(1), fillmode=:o)

    days = [Dates.Date(TimeSeries.timestamp(A)[1]) for A in Xs]
    vals = zeros(length(Xs), size(X,2))

    for (n,A) in enumerate(Xs)
        v = TimeSeries.values(A)
        # v[isnan.(v)] .= 0
        vals[n,:] = mean(v, dims=1)
    end

    return TimeArray(days, vals, TimeSeries.colnames(X))
end


abstract type AbstractTimeIntegrationMethod end

"""
    RungeKuttaMethod(A, b, c=vec(sum(A, dims=2)))

Represent a Runge-Kutta method with Butcher coefficients `A`, `b`, and `c`.
If `c` is not provided, the usual "row sum" requirement of consistency with
autonomous problems is applied.
"""
struct RungeKuttaMethod{T, MatT<:AbstractMatrix{T}, VecT<:AbstractVector{T}} <: AbstractTimeIntegrationMethod
  A::MatT
  b::VecT
  c::VecT
end

function RungeKuttaMethod(A::AbstractMatrix, b::AbstractVector, c::AbstractVector=vec(sum(A, dims=2)))
  T = promote_type(eltype(A), eltype(b), eltype(c))
  _A = T.(A)
  _b = T.(b)
  _c = T.(c)
  return RungeKuttaMethod(_A, _b, _c)
end

Base.eltype(rk::RungeKuttaMethod{T}) where {T} = T

function Base.show(io::IO, rk::RungeKuttaMethod)
  print(io, "RungeKuttaMethod{", eltype(rk), "}")
  if get(io, :compact, false)
    print(io, "(")
    show(io, rk.A)
    print(io, ", ")
    show(io, rk.b)
    print(io, ", ")
    show(io, rk.c)
    print(io, ")")
  else
    print(io, " with\nA: ")
    show(io, MIME"text/plain"(), rk.A)
    print(io, "\nb: ")
    show(io, MIME"text/plain"(), rk.b)
    print(io, "\nc: ")
    show(io, MIME"text/plain"(), rk.c)
    print(io, "\n")
  end
end

"""
    elementary_weight(t::RootedTree, rk::RungeKuttaMethod)
    elementary_weight(t::RootedTree, A::AbstractMatrix, b::AbstractVector, c::AbstractVector)

Compute the elementary weight Φ(`t`) of the [`RungeKuttaMethod`](@ref) `rk`
with Butcher coefficients `A, b, c` for a rooted tree `t``.

Reference: Section 312 of
- Butcher, John Charles.
  Numerical methods for ordinary differential equations.
  John Wiley & Sons, 2008.
"""
function elementary_weight(t::RootedTree, rk::RungeKuttaMethod)
  dot(rk.b, derivative_weight(t, rk))
end

# TODO: Deprecate also this method?
function elementary_weight(t::RootedTree, A::AbstractMatrix, b::AbstractVector, c::AbstractVector)
  elementary_weight(t, RungeKuttaMethod(A, b, c))
end

"""
    derivative_weight(t::RootedTree, rk::RungeKuttaMethod)

Compute the derivative weight (ΦᵢD)(`t`) of the [`RungeKuttaMethod`](@ref) `rk`
with Butcher coefficients `A, b, c` for the rooted tree `t`.

Reference: Section 312 of
- Butcher, John Charles.
  Numerical methods for ordinary differential equations.
  John Wiley & Sons, 2008.
"""
function derivative_weight(t::RootedTree, rk::RungeKuttaMethod)
  A = rk.A
  c = rk.c

  # Initialize `result` with the identity element of pointwise multiplication `.*`
  result = zero(c) .+ one(eltype(c))

  # Iterate over all subtrees and update the `result` using recursion
  for subtree in SubtreeIterator(t)
    tmp = A * derivative_weight(subtree, rk)
    result = result .* tmp
  end

  return result
end

# TODO: Deprecations introduced in v2
@deprecate derivative_weight(t::RootedTree, A, b, c) derivative_weight(t, RungeKuttaMethod(A, b, c))

"""
    residual_order_condition(t::RootedTree, rk::RungeKuttaMethod)

The residual of the order condition
  `(Φ(t) - 1/γ(t)) / σ(t)`
with [`elementary_weight`](@ref) `Φ(t)`, [`density`](@ref) `γ(t)`, and
[`symmetry`](@ref) `σ(t)` of the [`RungeKuttaMethod`](@ref) `rk` with Butcher
coefficients `A, b, c` for the rooted tree `t`.

Reference: Section 315 of
- Butcher, John Charles.
  Numerical methods for ordinary differential equations.
  John Wiley & Sons, 2008.
"""
function residual_order_condition(t::RootedTree, rk::RungeKuttaMethod)
  ew = elementary_weight(t, rk)
  T = typeof(ew)

  (ew - one(T) / γ(t)) / σ(t)
end

# TODO: Deprecations introduced in v2
@deprecate residual_order_condition(t::RootedTree, A, b, c) residual_order_condition(t, RungeKuttaMethod(A, b, c))

"""
    AdditiveRungeKuttaMethod(rks)
    AdditiveRungeKuttaMethod(As, bs, cs=map(A -> vec(sum(A, dims=2)), As))

Represent an additive Runge-Kutta method with collections of Butcher
coefficients `As`, `bs`, and `cs`. Alternatively, you can pass a collection of
[`RungeKuttaMethod`](@ref)s to the constructor.
If the `cs` are not provided, the usual "row sum" requirement of consistency
with autonomous problems is applied.

An additive Runge-Kutta method applied to the ODE problem
```math
  u'(t) = \\sum_\\nu f^\\nu(t, u(t))
```
has the form
```math
\\begin{aligned}
  y^i &= u^n + \\Delta t \\sum_\\nu \\sum_j a^\\nu_{i,j} f^\\nu(y^i), \\\\
  u^{n+1} &= u^n + \\Delta t \\sum_\\nu \\sum_i b^\\nu_{i} f^\\nu(y^i).
\\end{aligned}
```

In particular, additive Runge-Kutta methods are a superset of partitioned RK
methods, which are applied to partitioned problems of the form
```math
  (u^1)'(t) = f^1(t, u^1, u^2),
  \\quad
  (u^2)'(t) = f^2(t, u^1, u^2).
```

# References

- A. L. Araujo, A. Murua, and J. M. Sanz-Serna.
  "Symplectic Methods Based on Decompositions".
  SIAM Journal on Numerical Analysis 34.5 (1997): 1926-1947.
  [DOI: 10.1137/S0036142995292128](https://doi.org/10.1137/S0036142995292128)
"""
struct AdditiveRungeKuttaMethod{T, RKs<:AbstractVector{<:RungeKuttaMethod{T}}} <: AbstractTimeIntegrationMethod
  rks::RKs
end

function AdditiveRungeKuttaMethod(rks) # if not all RK methods use the same eltype
  T = mapreduce(eltype, promote_type, rks)
  As = map(rk -> T.(rk.A), rks)
  bs = map(rk -> T.(rk.b), rks)
  cs = map(rk -> T.(rk.c), rks)
  AdditiveRungeKuttaMethod(As, bs, cs)
end

function AdditiveRungeKuttaMethod(As, bs, cs=map(A -> vec(sum(A, dims=2)), As))
  rks = map(RungeKuttaMethod, As, bs, cs)
  AdditiveRungeKuttaMethod(rks)
end

Base.eltype(ark::AdditiveRungeKuttaMethod{T}) where {T} = T

function Base.show(io::IO, ark::AdditiveRungeKuttaMethod)
  print(io, "AdditiveRungeKuttaMethod{", eltype(ark), "} with methods\n")
  for (idx, rk) in enumerate(ark.rks)
    print(io, idx, ". ")
    show(io, rk)
  end
end

# Colored trees are used for order conditions of additive Runge-Kutta methods.
# The function `color_to_index` maps a color to a one-based index of ARK
# coefficients.
# This is considered to be an internal interface but stable across minor/patch
# releases.
color_to_index(color::Integer) = Int(color)
color_to_index(color::Bool) = 1 + color

"""
    elementary_weight(t::ColoredRootedTree, ark::AdditiveRungeKuttaMethod)

Compute the elementary weight Φ(`t`) of the [`AdditiveRungeKuttaMethod`](@ref)
`ark` for a colored rooted tree `t``.

# References

- A. L. Araujo, A. Murua, and J. M. Sanz-Serna.
  "Symplectic Methods Based on Decompositions".
  SIAM Journal on Numerical Analysis 34.5 (1997): 1926–1947.
  [DOI: 10.1137/S0036142995292128](https://doi.org/10.1137/S0036142995292128)
- Butcher, John Charles.
  Numerical methods for ordinary differential equations.
  John Wiley & Sons, 2008.
  Section 312
"""
function elementary_weight(t::ColoredRootedTree, ark::AdditiveRungeKuttaMethod)
  # TODO: Check number of RK methods in `ark`?
  if isempty(t)
    return one(eltype(first(ark.rks).c))
  else
    color = first(t.color_sequence)
    dot(ark.rks[color_to_index(color)].b, derivative_weight(t, ark))
  end
end

"""
    derivative_weight(t::ColoredRootedTree, ark::AdditiveRungeKuttaMethod)

Compute the derivative weight (ΦᵢD)(`t`) of the [`AdditiveRungeKuttaMethod`](@ref)
`ark` for the colored rooted tree `t`.

# References

- A. L. Araujo, A. Murua, and J. M. Sanz-Serna.
  "Symplectic Methods Based on Decompositions".
  SIAM Journal on Numerical Analysis 34.5 (1997): 1926–1947.
  [DOI: 10.1137/S0036142995292128](https://doi.org/10.1137/S0036142995292128)
- Butcher, John Charles.
  Numerical methods for ordinary differential equations.
  John Wiley & Sons, 2008.
  Section 312
"""
function derivative_weight(t::ColoredRootedTree, ark::AdditiveRungeKuttaMethod)
  # Initialize `result` with the identity element of pointwise multiplication `.*`
  c = first(ark.rks).c
  result = zero(c) .+ one(eltype(c))

  # Iterate over all subtrees and update the `result` using recursion
  for subtree in SubtreeIterator(t)
    A = ark.rks[color_to_index(first(subtree.color_sequence))].A
    tmp = A * derivative_weight(subtree, ark)
    result = result .* tmp
  end

  return result
end

"""
    residual_order_condition(t::ColoredRootedTree, ark::AdditiveRungeKuttaMethod)

The residual of the order condition
  `(Φ(t) - 1/γ(t)) / σ(t)`
with [`elementary_weight`](@ref) `Φ(t)`, [`density`](@ref) `γ(t)`, and
[`symmetry`](@ref) `σ(t)` of the [`AdditiveRungeKuttaMethod`](@ref) `ark`
for the colored rooted tree `t`.

# References

- A. L. Araujo, A. Murua, and J. M. Sanz-Serna.
  "Symplectic Methods Based on Decompositions".
  SIAM Journal on Numerical Analysis 34.5 (1997): 1926–1947.
  [DOI: 10.1137/S0036142995292128](https://doi.org/10.1137/S0036142995292128)
- Butcher, John Charles.
  Numerical methods for ordinary differential equations.
  John Wiley & Sons, 2008.
  Section 312
"""
function residual_order_condition(t::ColoredRootedTree, ark::AdditiveRungeKuttaMethod)
  ew = elementary_weight(t, ark)
  T = typeof(ew)

  (ew - one(T) / γ(t)) / σ(t)
end


#######################################################################
#    Script for testing sensitivities calculated using Sundials.jl    #
#                                                                     #
#     Sensitivity equations are solved with a finite difference       #
#         approximation to the Jacobian computed by Sundials.         #
#                               content                               #
#                                                                     #
#     The example is chosen to allow calculation of analytical        #
#         sensitivities for comparison. These are implemented         #
#                              as well.                               #
#                                                                     #
#     Christian Haargaard Olsen <christianhaargaard@gmail.com>        #
#                                                                     #
#######################################################################

using Sundials

# Uncomment this to use QT interface for PyPlot
using PyCall
@pyimport(matplotlib)
pygui(:qt)
using PyPlot

# First we create the datatype used for parsing parameter values to
# the right-hand-side function rhs.
type UserData
  parameters :: Array{Float64, 1}
end

# using the function f(t) = a*sin(b*t), f(0) = 0, we create an ODE by differentiating
# with respect to t, to get fdot (t) = a*b*cos(b*t). We will then use this ODE to test
# the sensitivity capabilities of sundials, as the correct sensitivites are given by
# f_a = sin(b*t) and f_b = a*t*cos(b*t).

function rhs(
  t         :: Float64,
  x         :: Sundials.N_Vector,
  xdot      :: Sundials.N_Vector,
  data
  )

  x         = Sundials.asarray(x)                 :: Array{Float64,1}
  xdot      = Sundials.asarray(xdot)              :: Array{Float64,1}
  user_data = unsafe_pointer_to_objref(data)      :: UserData

  a = user_data.parameters[1]
  b = user_data.parameters[2]

  x1 = x[1]
  xdot[1] = a*b*cos(b*t)

  return Int32(0)
end

function true_solution(
  t                                  :: Array{Float64,1},
  data                               :: UserData
  )

  a             = data.parameters[1] :: Float64
  b             = data.parameters[2] :: Float64
  result        = a*sin(b*t)         :: Array{Float64,1}

  return result
end

function true_sens(
  t                                  :: Array{Float64,1},
  data                               :: UserData
  )

  neq = 1
  NS = length(data.parameters)

  sens = zeros(length(t),neq,NS)

  a             = data.parameters[1] :: Float64
  b             = data.parameters[2] :: Float64
  sens[:,1,1] =  sin(b*t)            :: Array{Float64,1}
  sens[:,1,2] =  t.*a.*cos(b*t)      :: Array{Float64,1}

  return sens
end

function load_states()
  return ([0.0], UserData([1.0,1.0,1.0]))
end

function calcsens_finitedif(parameters :: Array{Float64,1}; step = sqrt(1e-6))
  (x, data)                 = load_states()               :: Tuple{Array{Float64,1}, UserData}
  data.parameters           = parameters                  :: Array{Float64,1}
  xdot                      = zeros(x)                    :: Array{Float64,1}
  t                         = collect(linspace(0,10,101)) :: Array{Float64,1}
  (reltol, abstol,sens_tol) = (1e-8, 1e-6, 1e-6)          :: Tuple{Float64,Float64, Float64}
  neq                       = length(x)
  NS                        = length(parameters)

  sols = zeros(length(t),neq,NS) # No need to copy initial condition, they are already zero
  sens = zeros(length(t),neq,NS) # No need to copy initial condition, they are already zero

  (t0, y0, x) = runmodel(parameters)
  @printf("calculating finite difference sensitivities for %i parameters.\n",NS)
  for i in 1:NS
    newparameters = zeros(parameters)
    newparameters = copy(parameters)
    newparameters[i] = (1.0+step)*newparameters[i]
    (t0, sols[:,:,i], x) = runmodel(newparameters)
    sens[:,:,i] = (sols[:,:,i]-y0)/(step*parameters[i])
    @printf("Done for parameter %i.\n",i)
    end

  return (sens)
end

function runsens(parameters :: Array{Float64,1})
  (x, data)                 = load_states()               :: Tuple{Array{Float64,1}, UserData}
  data.parameters           = parameters                  :: Array{Float64,1}
  xdot                      = zeros(x)                    :: Array{Float64,1}
  t                         = collect(linspace(0,10,101)) :: Array{Float64,1}
  (reltol, abstol,sens_tol) = (1e-8, 1e-6, 1e-6)          :: Tuple{Float64,Float64, Float64}
  neq                       = length(x)
  NS                        = length(parameters)

  #######################################################################
  #                         ODE solver settings                         #
  #######################################################################

  #cvode_mem = Sundials.CVodeCreate(Sundials.CV_BDF, Sundials.CV_FUNCTIONAL)
  cvode_mem = Sundials.CVodeCreate(Sundials.CV_BDF, Sundials.CV_NEWTON)
  flag = Sundials.CVodeInit(cvode_mem,rhs, t[1], x)
  flag = Sundials.CVodeSetUserData(cvode_mem, data)
  flag = Sundials.CVodeSStolerances(cvode_mem, reltol, abstol)
  flag = Sundials.CVDense(cvode_mem, neq)

  # List of parameters for which to calculate sensitivities
  plist = round(Int32,1:NS)
  println(plist)

  # scaling of parameters
  pbar = copy(data.parameters)

  # Tolerances for solving sensitivities for each parameter
  #sens_tol_vec = [abstol :: Float64 for i in 1:NS]

  # IDEA: Maybe these should be scaled with parameter value?
  sens_tol_vec = [abstol*data.parameters[i] :: Float64 for i in 1:NS]

  # Sundials want a (pointer to a) vector where each component corresponds to
  # a parameter, and where the element is a pointer to an array with a place for
  # each state variable.
  ySJ    = Array(Array{Float64,1}, NS);
  for i in 1:NS
    ySJ[i] = zeros(Float64,neq)*1.0;
  end
  ySjulia = [Sundials.nvector(ySJ[i]) for i in 1:NS]
  yS = pointer(ySjulia)

  #######################################################################
  #                         Sensitivity settigs                         #
  #######################################################################

  sensi_meth = Int32(Sundials.CV_SIMULTANEOUS) # CV_SIMULTANEOUS, CV_STAGGERED or CV_STAGGERED1
  flag = Sundials.CVodeSensInit1(cvode_mem, NS, sensi_meth, 0, yS);
  flag = Sundials.CVodeSensEEtolerances(cvode_mem);
  flag = Sundials.CVodeSetSensErrCon(cvode_mem, 0);
  #flag = Sundials.CVodeSensSStolerances(cvode_mem, reltol, sens_tol_vec);
  flag = Sundials.CVodeSetSensDQMethod(cvode_mem, Sundials.CV_CENTERED, 0.0);
  flag = Sundials.CVodeSetSensParams(cvode_mem, data.parameters, 0, 0);

  # Placeholder for solution and sensitivities
  solution = zeros(length(t), neq)
  solution[1,:] = copy(x)
  sens = zeros(length(t),neq,NS) # No need to copy initial condition, they are already zero

  tout = [t[1]] # Set first output time

  # Loop through all the output times
  for k in 2:length(t)
    # Extract the solution to x, and the sensitivities to yS
    flag = Sundials.CVode(cvode_mem, t[k], x, tout, Sundials.CV_NORMAL)
    flag = Sundials.CVodeGetSens(cvode_mem, tout, yS)

    #Save the results
    solution[k,:] = x
    for i in 1:NS
      sens[k,:,i] = Sundials.asarray(unsafe_load(yS,i))
    end
  end

  return (solution,sens)
end

#######################################################################
#                          Plot results                               #
#######################################################################

t = collect(linspace(0,10,101));
(x, data) = load_states()

(solution, sens) = runsens(data.parameters)
senstrue = true_sens(t,data)

figure(1)
plot(t,sens[:,1,1])
plot(t,senstrue[:,1,1])
show()

figure(2)
plot(t,sens[:,1,2])
plot(t,senstrue[:,1,2])
show()


@testset "561.array-partition-i.jl" begin
    @test array_pair_sum([26, 26, 21, 28, -24, -11, 26, 19, -10, -22, -22, -29, 28, -11, -1,
                          -10, -23, 7, 16, 7, 0, 1, 16, 14, 22, 9, 29, -23, 21, -20, -10,
                          21, -23, 21, 18, 2, -26, 6, 19, 8, 11, -29, 9, -13, -11, -23, -24,
                          13, -21, -21, 9, -25, 8, -5, 17, -25, -17, 28, -12, -27, -1, -7,
                          1, -26, -28, -11, 29, -28, 2, 1, 4, -21, 23, 0, 18, -23, 18, 23,
                          3, -10, 12, 4, 28, 29, -27, -28, 23, -27, -16, 13]) == -33
    @test array_pair_sum([1, 4, 3, 2]) == 4
    @test array_pair_sum([6, 2, 6, 5, 1, 2]) == 9
end


function deleteat(M::Array, i1::Integer, i2::Integer)
    return [M[1:i1 - 1,1:i2 - 1] M[1:i1 - 1,i2 + 1:end];M[i1 + 1:end,1:i2 - 1] M[i1 + 1:end,i2 + 1:end]]
end

deleteat(M::Array,i::Integer) = deleteat(M, i, i)

function orthogonalrows(axis::AbstractVector)
    if norm(axis) > 0
        axis = normalize(axis)
    end
    A = svd(skew(axis)).Vt
    inds = SA[1; 2; 3]
    V1 = A[1,inds]'
    V2 = A[2,inds]'
    V3 = axis' # instead of A[3,:] for correct sign: abs(axis) = abs(A[3,:])

    return V1, V2, V3
end

function orthogonalcols(axis::AbstractVector)
    V1, V2, V3 = orthogonalrows(axis)
    return V1', V2', V3'
end

@inline function skewplusdiag(v::AbstractVector{T},w::T) where T
    SA[
         w    -v[3]  v[2]
         v[3]  w    -v[1]
        -v[2]  v[1]  w
    ]
end

function getfieldnumber(obj)
    i = 1
    while true
        !isdefined(obj, i) ? break : (i+=1)
    end
    return i-1
end

@inline offsetrange(offset, length) = (offset-1)*length+1:offset*length
@inline offsetrange(offset, length, totallength, inneroffset) = (offset-1)*totallength+(inneroffset-1)*length+1:(offset-1)*totallength+inneroffset*length