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

export angle_to_angleaxis, angle_to_angle, angle_to_dcm, angle_to_quat
export smallangle_to_dcm, smallangle_to_quat
export angle_to_rot, smallangle_to_rot
export inv

################################################################################
#                                  Operations
################################################################################

"""
    *(Θ₂::EulerAngles, Θ₁::EulerAngles)

Compute the composed rotation of `Θ₁ -> Θ₂`. Notice that the rotation will be
represented by Euler angles (see `EulerAngles`) with the same rotation sequence
as `Θ₂`.

"""
@inline function *(Θ₂::EulerAngles, Θ₁::EulerAngles)
    # Convert to quaternions, compute the composition, and convert back to Euler
    # angles.
    q₁ = angle_to_quat(Θ₁)
    q₂ = angle_to_quat(Θ₂)

    quat_to_angle(q₁*q₂, Θ₂.rot_seq)
end

"""
    inv(Θ::EulerAngles)

Return the Euler angles that represent the inverse rotation of `Θ`. Notice that
the rotation sequence of the result will be the inverse of the input. Hence, if
the input rotation sequence is, for example, `:XYZ`, then the result will be
represented using `:ZYX`.

"""
function inv(Θ::EulerAngles)
    # Check what will be the inverse rotation.
    if Θ.rot_seq == :XYZ
        inv_rot_seq = :ZYX
    elseif Θ.rot_seq == :XZY
        inv_rot_seq = :YZX
    elseif Θ.rot_seq == :YXZ
        inv_rot_seq = :ZXY
    elseif Θ.rot_seq == :YZX
        inv_rot_seq = :XZY
    elseif Θ.rot_seq == :ZXY
        inv_rot_seq = :YXZ
    elseif Θ.rot_seq == :ZYX
        inv_rot_seq = :XYZ
    else
        inv_rot_seq = Θ.rot_seq
    end

    # Return the Euler angle that represented the inverse rotation.
    EulerAngles(-Θ.a3, -Θ.a2, -Θ.a1, inv_rot_seq)
end

################################################################################
#                                      IO
################################################################################

"""
    show(io::IO, Θ::EulerAngles{T}) where T
    show(io::IO, mime::MIME"text/plain", Θ::EulerAngles{T}) where T

Print the Euler angles `Θ` to the IO `io`.

"""
function show(io::IO, Θ::EulerAngles{T}) where T
    str     = @sprintf "%8.4f %8.4f %8.4f rad" Θ.a1 Θ.a2 Θ.a3
    rot_seq = String(Θ.rot_seq)

    print(io, "EulerAngles{$T}:")
    print(io,   " R($rot_seq): " * str)

    nothing
end

function show(io::IO, mime::MIME"text/plain", Θ::EulerAngles{T}) where T
    # Check if the user wants colors.
    color = get(io, :color, false)

    d = (color) ? _d : ""
    g = (color) ? _g : ""
    y = (color) ? _y : ""
    u = (color) ? _u : ""

    str_θ₁  = @sprintf "%8.4f rad (%9.4f deg)" Θ.a1 rad2deg(Θ.a1)
    str_θ₂  = @sprintf "%8.4f rad (%9.4f deg)" Θ.a2 rad2deg(Θ.a2)
    str_θ₃  = @sprintf "%8.4f rad (%9.4f deg)" Θ.a3 rad2deg(Θ.a3)
    rot_seq = String(Θ.rot_seq)

    println(io, "EulerAngles{$T}:")
    println(io,   "$g  R($(rot_seq[1])): $d" * str_θ₁)
    println(io,   "$y  R($(rot_seq[2])): $d" * str_θ₂)
      print(io,   "$u  R($(rot_seq[3])): $d" * str_θ₃)

    nothing
end

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

# Euler Angle and Axis
# ==============================================================================

"""
    angle_to_angleaxis(θ₁::Number, θ₂::Number, θ₃::Number, rot_seq::Symbol = :ZYX)
    angle_to_angleaxis(Θ::EulerAngles)

Convert the Euler angles `θ₁`, `θ₂`, and `θ₃` [rad] with the rotation sequence
`rot_seq` to an Euler angle and axis representation.  Those values can also be
passed inside the structure `Θ` (see `EulerAngles`).

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> angle_to_angleaxis(1,0,0,:XYZ)
EulerAngleAxis{Float64}:
  Euler angle:   1.0000 rad ( 57.2958 deg)
   Euler axis: [  1.0000,   0.0000,   0.0000]

julia> Θ = EulerAngles(1,1,1,:XYZ);

julia> angle_to_angleaxis(Θ)
EulerAngleAxis{Float64}:
  Euler angle:   1.9391 rad (111.1015 deg)
   Euler axis: [  0.6924,   0.2031,   0.6924]

```
"""
@inline function angle_to_angleaxis(θ₁::Number, θ₂::Number, θ₃::Number,
                                    rot_seq::Symbol = :ZYX)
    # First convert to quaternion and then to Euler angle and axis.
    quat_to_angleaxis( angle_to_quat( θ₁, θ₂, θ₃, rot_seq ) )
end

@inline angle_to_angleaxis(Θ::EulerAngles) =
    angle_to_angleaxis(Θ.a1, Θ.a2, Θ.a3, Θ.rot_seq)

# Euler Angles
# ==============================================================================

"""
    angle_to_angle(θ₁::Number, θ₂::Number, θ₃::Number, rot_seq_orig::Symbol, rot_seq_dest::Symbol)
    angle_to_angle(Θ::EulerAngles, rot_seq_dest::Symbol)

Convert the Euler angles `θ₁`, `θ₂`, and `θ₃` [rad] with the rotation sequence
`rot_seq_orig` to a new set of Euler angles with rotation sequence
`rot_seq_dest`. The input values of the origin Euler angles can also be passed
inside the structure `Θ` (see `EulerAngles`).

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> angle_to_angle(-pi/2, -pi/3, -pi/4, :ZYX, :XYZ)
EulerAngles{Float64}:
  R(X):  -1.0472 rad ( -60.0000 deg)
  R(Y):   0.7854 rad (  45.0000 deg)
  R(Z):  -1.5708 rad ( -90.0000 deg)

julia> angle_to_angle(-pi/2, 0, 0, :ZYX, :XYZ)
EulerAngles{Float64}:
  R(X):   0.0000 rad (   0.0000 deg)
  R(Y):   0.0000 rad (   0.0000 deg)
  R(Z):  -1.5708 rad ( -90.0000 deg)

julia> Θ = EulerAngles(1,2,3,:XYX)
EulerAngles{Int64}:
  R(X):   1.0000 rad (  57.2958 deg)
  R(Y):   2.0000 rad ( 114.5916 deg)
  R(X):   3.0000 rad ( 171.8873 deg)

julia> angle_to_angle(Θ,:ZYZ)
EulerAngles{Float64}:
  R(Z):  -2.7024 rad (-154.8356 deg)
  R(Y):   1.4668 rad (  84.0393 deg)
  R(Z):  -1.0542 rad ( -60.3984 deg)
```
"""
@inline angle_to_angle(θ₁::Number, θ₂::Number, θ₃::Number, rot_seq_orig::Symbol,
                       rot_seq_dest::Symbol) =
    quat_to_angle(angle_to_quat(θ₁, θ₂, θ₃, rot_seq_orig), rot_seq_dest)

@inline angle_to_angle(Θ::EulerAngles, rot_seq_dest::Symbol) =
    angle_to_angle(Θ.a1, Θ.a2, Θ.a3, Θ.rot_seq, rot_seq_dest)

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

"""
    angle_to_dcm(θ₁::Number, θ₂::Number, θ₃::Number, rot_seq::Symbol = :ZYX)

Convert the Euler angles `θ₁`, `θ₂`, and `θ₃` [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 = angle_to_dcm(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 angle_to_dcm(θ₁::Number, θ₂::Number, θ₃::Number, rot_seq::Symbol = :ZYX)

    # Compute the sines and cosines.
    s₁, c₁ = sincos(θ₁)
    s₂, c₂ = sincos(θ₂)
    s₃, c₃ = sincos(θ₃)

    if rot_seq == :ZYX
        return DCM(       c₂*c₁    ,        c₂*s₁    , -s₂  ,
                   s₃*s₂*c₁ - c₃*s₁, s₃*s₂*s₁ + c₃*c₁, s₃*c₂,
                   c₃*s₂*c₁ + s₃*s₁, c₃*s₂*s₁ - s₃*c₁, c₃*c₂)'

    elseif rot_seq == :XYX
        return DCM( c₂  ,         s₁*s₂    ,       -c₁*s₂    ,
                   s₂*s₃, -s₁*c₂*s₃ + c₁*c₃, c₁*c₂*s₃ + s₁*c₃,
                   s₂*c₃, -s₁*c₃*c₂ - c₁*s₃, c₁*c₃*c₂ - s₁*s₃)'

    elseif rot_seq == :XYZ
        return DCM( c₂*c₃,  s₁*s₂*c₃ + c₁*s₃, -c₁*s₂*c₃ + s₁*s₃,
                   -c₂*s₃, -s₁*s₂*s₃ + c₁*c₃,  c₁*s₂*s₃ + s₁*c₃,
                      s₂ ,        -s₁*c₂    ,       c₁*c₂)'

    elseif rot_seq == :XZX
        return DCM(   c₂ ,         c₁*s₂    ,         s₁*s₂    ,
                   -s₂*c₃,  c₁*c₃*c₂ - s₁*s₃,  s₁*c₃*c₂ + c₁*s₃,
                    s₂*s₃, -c₁*c₂*s₃ - s₁*c₃, -s₁*c₂*s₃ + c₁*c₃)'

    elseif rot_seq == :XZY
        return DCM(c₃*c₂, c₁*c₃*s₂ + s₁*s₃, s₁*c₃*s₂ - c₁*s₃,
                    -s₂ ,        c₁*c₂    ,        s₁*c₂    ,
                   s₃*c₂, c₁*s₂*s₃ - s₁*c₃, s₁*s₂*s₃ + c₁*c₃)'

    elseif rot_seq == :YXY
        return DCM(-s₁*c₂*s₃ + c₁*c₃, s₂*s₃ , -c₁*c₂*s₃ - s₁*c₃,
                          s₁*s₂    ,   c₂  ,         c₁*s₂    ,
                   s₁*c₃*c₂ + c₁*s₃, -s₂*c₃,  c₁*c₃*c₂ - s₁*s₃)'

    elseif rot_seq == :YXZ
        return DCM( c₁*c₃ + s₂*s₁*s₃, c₂*s₃, -s₁*c₃ + s₂*c₁*s₃,
                   -c₁*s₃ + s₂*s₁*c₃, c₂*c₃,  s₁*s₃ + s₂*c₁*c₃,
                        s₁*c₂       ,  -s₂ ,      c₂*c₁       )'

    elseif rot_seq == :YZX
        return DCM(        c₁*c₂    ,    s₂ ,        -s₁*c₂    ,
                   -c₃*c₁*s₂ + s₃*s₁,  c₂*c₃,  c₃*s₁*s₂ + s₃*c₁,
                    s₃*c₁*s₂ + c₃*s₁, -s₃*c₂, -s₃*s₁*s₂ + c₃*c₁)'

    elseif rot_seq == :YZY
        return DCM(c₁*c₃*c₂ - s₁*s₃, s₂*c₃, -s₁*c₃*c₂ - c₁*s₃,
                         -c₁*s₂    ,   c₂ ,         s₁*s₂    ,
                   c₁*c₂*s₃ + s₁*c₃, s₂*s₃, -s₁*c₂*s₃ + c₁*c₃)'

    elseif rot_seq == :ZXY
        return DCM(c₃*c₁ - s₂*s₃*s₁, c₃*s₁ + s₂*s₃*c₁, -s₃*c₂,
                      -c₂*s₁       ,     c₂*c₁       ,    s₂ ,
                   s₃*c₁ + s₂*c₃*s₁, s₃*s₁ - s₂*c₃*c₁,  c₂*c₃)'

    elseif rot_seq == :ZXZ
        return DCM(-s₁*c₂*s₃ + c₁*c₃, c₁*c₂*s₃ + s₁*c₃, s₂*s₃,
                   -s₁*c₃*c₂ - c₁*s₃, c₁*c₃*c₂ - s₁*s₃, s₂*c₃,
                           s₁*s₂    ,       -c₁*s₂    ,  c₂)'

    elseif rot_seq == :ZYZ
        return DCM( c₁*c₃*c₂ - s₁*s₃,  s₁*c₃*c₂ + c₁*s₃, -s₂*c₃,
                   -c₁*c₂*s₃ - s₁*c₃, -s₁*c₂*s₃ + c₁*c₃,  s₂*s₃,
                           c₁*s₂    ,         s₁*s₂    ,    c₂)'
    else
        throw(ArgumentError("The rotation sequence :$rot_seq is not valid."))
    end
end

"""
    angle_to_dcm(Θ::EulerAngles)

Convert the Euler angles `Θ` (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> angle_to_dcm(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
```

"""
angle_to_dcm(Θ::EulerAngles) = angle_to_dcm(Θ.a1, Θ.a2, Θ.a3, Θ.rot_seq)

"""
    smallangle_to_dcm(θx::Number, θy::Number, θz::Number; normalize = true)

Create a direction cosine matrix from three small rotations of angles `θx`,
`θy`, and `θz` [rad] about the axes X, Y, and Z, respectively. If the keyword
`normalize` is `true`, then the matrix will be normalized using the function
`orthonormalize`.

# Example

```julia-repl
julia> smallangle_to_dcm(+0.01, -0.01, +0.01)
3×3 StaticArrays.SArray{Tuple{3,3},Float64,2,9}:
  0.9999     0.00989903  0.010098
 -0.009999   0.999901    0.00989802
 -0.009999  -0.009998    0.9999

julia> smallangle_to_dcm(+0.01, -0.01, +0.01; normalize = false)
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
```

"""
@inline function smallangle_to_dcm(θx::Number, θy::Number, θz::Number;
                                   normalize = true)
    D = DCM(  1, +θz, -θy,
            -θz,   1, +θx,
            +θy, -θx,   1)'

    normalize ? orthonormalize(D) : D
end

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

"""
    angle_to_quat(θ₁::T1, θ₂::T2, θ₃::T3, rot_seq::Symbol = :ZYX) where {T1<:Number, T2<:Number, T3<:Number}

Convert the Euler angles `θ₁`, `θ₂`, and `θ₃` [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> angle_to_quat(pi/2, pi/3, pi/4, :ZYX)
Quaternion{Float64}:
  + 0.7010573846499779 - 0.09229595564125714.i + 0.5609855267969309.j + 0.43045933457687935.k
```

"""
function angle_to_quat(θ₁::T1, θ₂::T2, θ₃::T3, rot_seq::Symbol = :ZYX) where
    {T1<:Number, T2<:Number, T3<:Number}

    T = promote_type(T1,T2,T3)

    # Compute the sines and cosines of half angle.
    s₁, c₁ = sincos(T(θ₁)/2)
    s₂, c₂ = sincos(T(θ₂)/2)
    s₃, c₃ = sincos(T(θ₃)/2)

    if rot_seq == :ZYX
        q0 = c₁*c₂*c₃ + s₁*s₂*s₃

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

        return Quaternion(s*q0,
                          s*(c₁*c₂*s₃ - s₁*s₂*c₃),
                          s*(c₁*s₂*c₃ + s₁*c₂*s₃),
                          s*(s₁*c₂*c₃ - c₁*s₂*s₃))
    elseif rot_seq == :XYX
        q0 = c₁*c₂*c₃ - s₁*c₂*s₃

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

        return Quaternion(s*q0,
                          s*(c₁*c₂*s₃ + s₁*c₂*c₃),
                          s*(c₁*s₂*c₃ + s₁*s₂*s₃),
                          s*(s₁*s₂*c₃ - c₁*s₂*s₃))
    elseif rot_seq == :XYZ
        q0 = c₁*c₂*c₃ - s₁*s₂*s₃

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

        return Quaternion(s*q0,
                          s*(s₁*c₂*c₃ + c₁*s₂*s₃),
                          s*(c₁*s₂*c₃ - s₁*c₂*s₃),
                          s*(c₁*c₂*s₃ + s₁*s₂*c₃))
    elseif rot_seq == :XZX
        q0 = c₁*c₂*c₃ - s₁*c₂*s₃

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

        return Quaternion(s*q0,
                          s*(c₁*c₂*s₃ + s₁*c₂*c₃),
                          s*(c₁*s₂*s₃ - s₁*s₂*c₃),
                          s*(c₁*s₂*c₃ + s₁*s₂*s₃))
    elseif rot_seq == :XZY
        q0 = c₁*c₂*c₃ + s₁*s₂*s₃

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

        return Quaternion(s*q0,
                          s*(s₁*c₂*c₃ - c₁*s₂*s₃),
                          s*(c₁*c₂*s₃ - s₁*s₂*c₃),
                          s*(c₁*s₂*c₃ + s₁*c₂*s₃))
    elseif rot_seq == :YXY
        q0 = c₁*c₂*c₃ - s₁*c₂*s₃

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

        return Quaternion(s*q0,
                          s*(c₁*s₂*c₃ + s₁*s₂*s₃),
                          s*(c₁*c₂*s₃ + s₁*c₂*c₃),
                          s*(c₁*s₂*s₃ - s₁*s₂*c₃))
    elseif rot_seq == :YXZ
        q0 = c₁*c₂*c₃ + s₁*s₂*s₃

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

        return Quaternion(s*q0,
                          s*(c₁*s₂*c₃ + s₁*c₂*s₃),
                          s*(s₁*c₂*c₃ - c₁*s₂*s₃),
                          s*(c₁*c₂*s₃ - s₁*s₂*c₃))
    elseif rot_seq == :YZX
        q0 = c₁*c₂*c₃ - s₁*s₂*s₃

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

        return Quaternion(s*q0,
                          s*(c₁*c₂*s₃ + s₁*s₂*c₃),
                          s*(s₁*c₂*c₃ + c₁*s₂*s₃),
                          s*(c₁*s₂*c₃ - s₁*c₂*s₃))
    elseif rot_seq == :YZY
        q0 = c₁*c₂*c₃ - s₁*c₂*s₃

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

        return Quaternion(s*q0,
                          s*(s₁*s₂*c₃ - c₁*s₂*s₃),
                          s*(c₁*c₂*s₃ + s₁*c₂*c₃),
                          s*(c₁*s₂*c₃ + s₁*s₂*s₃))
    elseif rot_seq == :ZXY
        q0 = c₁*c₂*c₃ - s₁*s₂*s₃

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

        return Quaternion(s*q0,
                          s*(c₁*s₂*c₃ - s₁*c₂*s₃),
                          s*(c₁*c₂*s₃ + s₁*s₂*c₃),
                          s*(s₁*c₂*c₃ + c₁*s₂*s₃))
    elseif rot_seq == :ZXZ
        q0 = c₁*c₂*c₃ - s₁*c₂*s₃

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

        return Quaternion(s*q0,
                          s*(c₁*s₂*c₃ + s₁*s₂*s₃),
                          s*(s₁*s₂*c₃ - c₁*s₂*s₃),
                          s*(c₁*c₂*s₃ + s₁*c₂*c₃))
    elseif rot_seq == :ZYZ
        q0 = c₁*c₂*c₃ - s₁*c₂*s₃

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

        return Quaternion(s*q0,
                          s*(c₁*s₂*s₃ - s₁*s₂*c₃),
                          s*(c₁*s₂*c₃ + s₁*s₂*s₃),
                          s*(c₁*c₂*s₃ + s₁*c₂*c₃))
    else
        throw(ArgumentError("The rotation sequence :$rot_seq is not valid."))
    end
end

"""
    angle_to_quat(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> angle_to_quat(pi/2, pi/3, pi/4, :ZYX)
Quaternion{Float64}:
  + 0.7010573846499779 - 0.09229595564125714.i + 0.5609855267969309.j + 0.43045933457687935.k
```

"""
angle_to_quat(Θ::EulerAngles) = angle_to_quat(Θ.a1, Θ.a2, Θ.a3, Θ.rot_seq)

"""
    smallangle_to_quat(θ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> smallangle_to_quat(+0.01, -0.01, +0.01)
Quaternion{Float64}:
  + 0.9999625021092433 + 0.004999812510546217.i - 0.004999812510546217.j + 0.004999812510546217.k
```

"""
function smallangle_to_quat(θ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
################################################################################

"""
    angle_to_rot([T,] θx::Number, θy::Number, θz::Number, rot_seq::Symbol)

Convert the Euler angles `Θx`, `Θy`, and `Θz` [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 = angle_to_rot(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   = angle_to_rot(Quaternion,pi/2, pi/3, pi/4, :ZYX)
Quaternion{Float64}:
  + 0.7010573846499779 - 0.09229595564125714.i + 0.5609855267969309.j + 0.43045933457687935.k
```

"""
@inline angle_to_rot(θx::Number, θy::Number, θz::Number, rot_seq::Symbol) =
    angle_to_dcm(θx, θy, θz, rot_seq)

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

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

"""
    angle_to_rot([T,] Θ::EulerAngles)

Convert the Euler angles `Θ` (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 = angle_to_rot(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   = angle_to_rot(Quaternion,EulerAngles(pi/2, pi/3, pi/4, :ZYX))
Quaternion{Float64}:
  + 0.7010573846499779 - 0.09229595564125714.i + 0.5609855267969309.j +
  0.43045933457687935.k
```

"""
@inline angle_to_rot(Θ::EulerAngles) = angle_to_dcm(Θ.a1, Θ.a2, Θ.a3, Θ.rot_seq)

@inline angle_to_rot(::Type{DCM}, Θ::EulerAngles) =
    angle_to_dcm(Θ.a1, Θ.a2, Θ.a3, Θ.rot_seq)

@inline angle_to_rot(::Type{Quaternion}, Θ::EulerAngles) =
    angle_to_quat(Θ.a1, Θ.a2, Θ.a3, Θ.rot_seq)

"""
    smallangle_to_rot([T,] θx::Number, θy::Number, θz::Number[; normalize = true])

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`.

If `T` is `DCM`, then the resulting matrix will be orthonormalized using the
`orthonormalize` function if the keyword `normalize` is `true`.

# Example

```julia-repl
julia> dcm = smallangle_to_rot(+0.01, -0.01, +0.01)
3×3 StaticArrays.SArray{Tuple{3,3},Float64,2,9}:
  0.9999     0.00989903  0.010098
 -0.009999   0.999901    0.00989802
 -0.009999  -0.009998    0.9999

julia> dcm = smallangle_to_rot(+0.01, -0.01, +0.01; normalize = false)
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   = smallangle_to_rot(Quaternion,+0.01, -0.01, +0.01)
Quaternion{Float64}:
  + 0.9999625021092433 + 0.004999812510546217.i - 0.004999812510546217.j + 0.004999812510546217.k
```
"""
@inline smallangle_to_rot(θx::Number, θy::Number, θz::Number; normalize = true) =
    smallangle_to_dcm(θx, θy, θz; normalize = normalize)

@inline smallangle_to_rot(::Type{DCM}, θx::Number, θy::Number, θz::Number;
                          normalize = true) =
    smallangle_to_dcm(θx, θy, θz; normalize = normalize)

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


"
expands the bond dimension by adding random unitary vectors
"
@with_kw struct RandExpand<:Algorithm
    trscheme::TruncationScheme = truncdim(1)
end

function changebonds(state::InfiniteMPS,alg::RandExpand)
    #determine optimal expansion spaces around bond i
    pexp = PeriodicArray(map(1:length(state)) do i
        ACAR = _permute_front(state.AC[i])*_permute_tail(state.AR[i+1])
        AC2 = TensorMap(rand,eltype(ACAR),codomain(ACAR),domain(ACAR));

        #Calculate nullspaces for AL and AR
        NL = leftnull(state.AL[i])
        NR = rightnull(_permute_tail(state.AR[i+1]))

        #Use this nullspaces and SVD decomposition to determine the optimal expansion space
        intermediate = adjoint(NL)*AC2*adjoint(NR)
        (U,S,V) = tsvd(intermediate,trunc=alg.trscheme,alg=SVD())

        (NL*U,V*NR)
    end)

    newstate = copy(state);

    #do the actual expansion
    for i in 1:length(state)
        al = _permute_tail(catdomain(newstate.AL[i],pexp[i][1]))
        lz = TensorMap(zeros,_lastspace(pexp[i-1][1])',domain(al))
        newstate.AL[i] = _permute_front(catcodomain(al,lz))

        ar = _permute_front(catcodomain(_permute_tail(newstate.AR[i+1]),pexp[i][2]))
        rz = TensorMap(zeros,codomain(ar),space(pexp[i+1][2],1))
        newstate.AR[i+1] = catdomain(ar,rz)

        l = TensorMap(zeros,codomain(newstate.CR[i]),space(pexp[i][2],1))
        newstate.CR[i] = catdomain(newstate.CR[i],l)
        r = TensorMap(zeros,_lastspace(pexp[i][1])',domain(newstate.CR[i]))
        newstate.CR[i] = catcodomain(newstate.CR[i],r)

        newstate.AC[i] = newstate.AL[i]*newstate.CR[i]
    end

    return newstate
end

function changebonds(state::MPSMultiline,alg::RandExpand)
    #=
        todo : merge this with the MPSCentergauged implementation
    =#
    #determine optimal expansion spaces around bond i
    pexp = PeriodicArray(map(Iterators.product(1:size(state,1),1:size(state,2))) do (i,j)

        ACAR = _permute_front(state.AC[i-1,j])*_permute_tail(state.AR[i-1,j+1])
        AC2 = TensorMap(rand,eltype(ACAR),codomain(ACAR),domain(ACAR));

        #Calculate nullspaces for AL and AR
        NL = leftnull(state.AL[i,j])
        NR = rightnull(_permute_tail(state.AR[i,j+1]))

        #Use this nullspaces and SVD decomposition to determine the optimal expansion space
        intermediate = adjoint(NL)*AC2*adjoint(NR)
        (U,S,V) = tsvd(intermediate,trunc=alg.trscheme,alg=SVD())

        (NL*U,V*NR)
    end)

    newstate = copy(state);

    #do the actual expansion
    for (i,j) in Iterators.product(1:size(state,1),1:size(state,2))
        al = _permute_tail(catdomain(newstate.AL[i,j],pexp[i,j][1]))
        lz = TensorMap(zeros,_lastspace(pexp[i,j-1][1])',domain(al))

        newstate.AL[i,j] = _permute_front(catcodomain(al,lz))

        ar = _permute_front(catcodomain(_permute_tail(newstate.AR[i,j+1]),pexp[i,j][2]))
        rz = TensorMap(zeros,codomain(ar),space(pexp[i,j+1][2],1))
        newstate.AR[i,j+1] = catdomain(ar,rz)

        l = TensorMap(zeros,codomain(newstate.CR[i,j]),space(pexp[i,j][2],1))
        newstate.CR[i,j] = catdomain(newstate.CR[i,j],l)
        r = TensorMap(zeros,_lastspace(pexp[i,j][1])',domain(newstate.CR[i,j]))
        newstate.CR[i,j] = catcodomain(newstate.CR[i,j],r)

        newstate.AC[i,j] = newstate.AL[i,j]*newstate.CR[i,j]
    end

    return newstate
end

changebonds(state::Union{FiniteMPS,MPSComoving}, alg::RandExpand) = changebonds!(copy(state),alg)
function changebonds!(state::Union{FiniteMPS,MPSComoving},alg::RandExpand)
    for i in 1:(length(state)-1)
        ACAR = _permute_front(state.AC[i])*_permute_tail(state.AR[i+1])
        AC2 = TensorMap(rand,eltype(ACAR),codomain(ACAR),domain(ACAR));

        #Calculate nullspaces for left and right
        NL = leftnull(state.AC[i])
        NR = rightnull(_permute_tail(state.AR[i+1]))

        #Use this nullspaces and SVD decomposition to determine the optimal expansion space
        intermediate = adjoint(NL) * AC2 * adjoint(NR);
        (U,S,V) = tsvd(intermediate,trunc=alg.trscheme,alg=SVD())

        ar_re = V*NR;
        ar_le = TensorMap(zeros,codomain(state.AC[i]),space(V,1))

        (nal,nc) = leftorth(catdomain(state.AC[i],ar_le),alg=QRpos())
        nar = _permute_front(catcodomain(_permute_tail(state.AR[i+1]),ar_re));

        state.AC[i] = (nal,nc)
        state.AC[i+1] = (nc,nar)
    end

    state
end


using QuantumCumulants
using Test

@testset "mixed-order" begin

# Hilbertspace
hf = FockSpace(:cavity)
ha = NLevelSpace(:atom, 2)
h = hf⊗ha

# Operators
a = Destroy(h,:a,1)
σ(i,j) = Transition(h,:σ,i,j)

@test isequal(average(a'*a, [2,1]), average(a'*a))
@test isequal(average(a'*σ(1,2), [2,1]; mix_choice=minimum), average(a')*average(σ(1,2)))

he = meanfield([a'*a,σ(2,2)],a'*a + σ(2,2) + a'*σ(1,2) + a*σ(2,1))
he_avg1 = cumulant_expansion(he,2)
he_avg2 = cumulant_expansion(he,[2,1])
he_avg3 = cumulant_expansion(he,[2,1];mix_choice=minimum)

@test isequal(he_avg1.equations, he_avg2.equations)
@test !isequal(he_avg1.equations, he_avg3.equations)

he = meanfield(a'*σ(1,2), a'*a + σ(2,2) + a*σ(2,1))

# N-atom laser
# Parameters
N = 2 #number of atoms
κ, g, Γ23, Γ13, Γ12, Ω, Δc, Δ3 = cnumbers("κ g Γ_{23} Γ_{13} Γ_{12} Ω Δ_c Δ_3")

# Hilbertspace
hf = FockSpace(:cavity)
ha = ⊗([NLevelSpace(Symbol(:atom,i),3) for i=1:N]...)
h = hf ⊗ ha

# Operators
a = Destroy(h,:a)
σ(i,j,k) = Transition(h,Symbol("σ_{$k}"),i,j,k+1)

# Hamiltonian
H = -Δc*a'a + sum(g*(a'*σ(1,2,i) + a*σ(2,1,i)) for i=1:N) + sum(Ω*(σ(3,1,i) + σ(1,3,i)) for i=1:N) - sum(Δ3*σ(3,3,i) for i=1:N)

# Jumps
J = [a;[σ(1,2,i) for i=1:N];[σ(1,3,i) for i=1:N];[σ(2,3,i) for i=1:N]]

# Rates
rates = [κ;[Γ12 for i=1:N];[Γ13 for i=1:N];[Γ23 for i=1:N]]

# list of operators
ops = [a'a, σ(2,2,1), σ(3,3,1)]

he = meanfield(ops,H,J; rates=rates)
he_avg_ = cumulant_expansion(he,[2,1,1]) #second order average

he_avg = complete(he_avg_;order=[2,1,1])
@test isempty(find_missing(he_avg))
@test isempty(findall(x -> (aon = acts_on(x); 2 in aon && 3 in aon), he_avg.states))

he_avg_ = cumulant_expansion(he,[2,1,1];mix_choice=minimum)
he_avg = complete(he_avg_;order=[2,1,1],mix_choice=minimum)
@test isempty(find_missing(he_avg))
@test isempty(findall(x -> (aon = acts_on(x); 2 in aon && 3 in aon), he_avg.states))
@test isempty(findall(x -> (aon = acts_on(x); 1 in aon && 3 in aon), he_avg.states))
@test isempty(findall(x -> (aon = acts_on(x); 1 in aon && 2 in aon), he_avg.states))

end # testset


"""
    DistanceMethod

Abstract supertype for methods of calculating distances and discretised dispersal
probabilities between cells in a grid.
between cells in a grid.

Distance calculation methods include:
- [`CentroidToCentroid`](@ref)
- [`AreaToArea`](@ref)
- [`AreaToCentroid`](@ref)

Which are adapted from: "On the approximation of continuous dispersal kernels
in discrete-space models, Joseph D. Chipperfield et al 2011"

The `CentroidToArea` method has not been implemented.
"""
abstract type DistanceMethod end

subsampling(method::DistanceMethod) = method.subsampling

"""
    CentroidToCentroid <: DistanceMethod

    CentroidToCentroid()

Calculates the discrete probability of dispersal between source and destination cell
centroids. This is the naive method, but it will not handle low grid resolution well
due to severe truncation.

See: "On the approximation of continuous dispersal kernels in discrete-space models,
Joseph D. Chipperfield et al 2011"
"""
struct CentroidToCentroid <: DistanceMethod end

dispersalprob(f, ::CentroidToCentroid, x, y, cellsize) = f(sqrt(x^2 + y^2) * cellsize)

"""
    AreaToCentroid <: DistanceMethod

    AreaToCentroid(subsampling)
    AreaToCentroid(; subsampling=10.0)

Calculates the discrete probability of dispersal between source cell area and destination
centroid.

See: "On the approximation of continuous dispersal kernels in discrete-space models,
Joseph D. Chipperfield et al 2011"
"""
Base.@kwdef struct AreaToCentroid{SS<:Number} <: DistanceMethod
    "Subsampling for numerical integration"
    subsampling::SS = Param(10.0; bounds=(2.0, 40.0))
end

@inline function dispersalprob(f, dm::AreaToCentroid, x, y, cellsize)
    prob = zero(cellsize)
    centerfirst = 1 / subsampling(dm) / 2 - 0.5
    centerlast = centerfirst * -1
    range = LinRange(centerfirst, centerlast, round(Int, subsampling(dm)))
    for j in range, i in range
        prob += sqrt((x + i)^2 + (y + j)^2) * cellsize |> f
    end
    prob / subsampling(dm)^2
end

"""
    AreaToArea <: DistanceMethod

    AreaToArea(subsampling)
    AreaToArea(; subsampling=10.0)

Calculates the discrete probability of dispersal between source and destination based on
cell areas.

See: "On the approximation of continuous dispersal kernels in discrete-space models,
Joseph D. Chipperfield et al 2011"
"""
Base.@kwdef struct AreaToArea{SS<:Number} <: DistanceMethod
    subsampling::SS = Param(10.0; bounds=(2.0, 40.0))
end

@inline function dispersalprob(f, dm::AreaToArea, x, y, cellsize)
    prob = zero(cellsize)
    # Get the center point of the first cell (for both dimensions)
    centerfirst = 1 / subsampling(dm) / 2 - 0.5
    centerlast = centerfirst * -1
    range = LinRange(centerfirst, centerlast, round(Int, subsampling(dm)))
    for i in range, j in range
        for a in range, b in range
            prob += sqrt((x + i + a)^2 + (y + j + b)^2) * cellsize |> f
        end
    end
    prob / subsampling(dm)^4
end


# Identity

"""
    tr_identity(i, x)
    tr_identity(x)
    tr_identity(x, X::AbstractVector)

Identity function. Already defined in Julia as `identity`, specialized for vectors.
When `X` is provided, the result is computed without allocations.
"""
tr_identity(i, x) = identity(x[i])
lazy(tr_identity)

# Count equalities

"""
    tr_count_eq(i, x)
    tr_count_eq(x)
    tr_count_eq(x, X::AbstractVector)

Count the number of elements equal to `x[i]`. Extended method to vector with sig `(x)` are generated.
When `X` is provided, the result is computed without allocations.
"""
tr_count_eq(i, x) = count(y -> x[i] == y, x) - 1

"""
    tr_count_eq_right(i, x)
    tr_count_eq_right(x)
    tr_count_eq_right(x, X::AbstractVector)

Count the number of elements to the right of and equal to `x[i]`. Extended method to vector with sig `(x)` are generated.
When `X` is provided, the result is computed without allocations.
"""
tr_count_eq_right(i, x) = tr_count_eq(1, @view x[i:end])

"""
    tr_count_eq_left(i, x)
    tr_count_eq_left(x)
    tr_count_eq_left(x, X::AbstractVector)

Count the number of elements to the left of and equal to `x[i]`. Extended method to vector with sig `(x)` are generated.
When `X` is provided, the result is computed without allocations.
"""
tr_count_eq_left(i, x) = tr_count_eq(i, @view x[1:i])

# Generating vetorized versions
lazy(tr_count_eq, tr_count_eq_left, tr_count_eq_right)

# Count greater/lesser

"""
    tr_count_greater(i, x)
    tr_count_greater(x)
    tr_count_greater(x, X::AbstractVector)

Count the number of elements greater than `x[i]`. Extended method to vector with sig `(x)` are generated.
When `X` is provided, the result is computed without allocations.
"""
tr_count_greater(i, x) = count(y -> x[i] < y, x)

"""
    tr_count_lesser(i, x)
    tr_count_lesser(x)
    tr_count_lesser(x, X::AbstractVector)

Count the number of elements lesser than `x[i]`. Extended method to vector with sig `(x)` are generated.
When `X` is provided, the result is computed without allocations.
"""
tr_count_lesser(i, x) = count(y -> x[i] > y, x)

"""
    tr_count_g_left(i, x)
    tr_count_g_left(x)
    tr_count_g_left(x, X::AbstractVector)

Count the number of elements to the left of and greater than `x[i]`. Extended method to vector with sig `(x)` are generated.
When `X` is provided, the result is computed without allocations.
"""
tr_count_g_left(i, x) = tr_count_greater(i, @view x[1:i])

"""
    tr_count_l_left(i, x)
    tr_count_l_left(x)
    tr_count_l_left(x, X::AbstractVector)

Count the number of elements to the left of and lesser than `x[i]`. Extended method to vector with sig `(x)` are generated.
When `X` is provided, the result is computed without allocations.
"""
tr_count_l_left(i, x) = tr_count_lesser(i, @view x[1:i])

"""
    tr_count_g_right(i, x)
    tr_count_g_right(x)
    tr_count_g_right(x, X::AbstractVector)

Count the number of elements to the right of and greater than `x[i]`. Extended method to vector with sig `(x)` are generated.
"""
tr_count_g_right(i, x) = tr_count_greater(1, @view x[i:end])

"""
    tr_count_l_right(i, x)
    tr_count_l_right(x)
    tr_count_l_right(x, X::AbstractVector)

Count the number of elements to the right of and lesser than `x[i]`. Extended method to vector with sig `(x)` are generated.
When `X` is provided, the result is computed without allocations.
"""
tr_count_l_right(i, x) = tr_count_lesser(1, @view x[i:end])

# Generating vetorized versions
lazy(tr_count_greater, tr_count_g_left, tr_count_g_right)
lazy(tr_count_lesser, tr_count_l_left, tr_count_l_right)

# Count param

"""
    tr_count_eq_param(i, x; param)
    tr_count_eq_param(x; param)
    tr_count_eq_param(x, X::AbstractVector; param)

Count the number of elements equal to `x[i] + param`. Extended method to vector with sig `(x, param)` are generated.
When `X` is provided, the result is computed without allocations.
"""
tr_count_eq_param(i, x; param) = count(y -> y == x[i] + param, x)

"""
    tr_count_l_param(i, x; param)
    tr_count_l_param(x; param)
    tr_count_l_param(x, X::AbstractVector; param)

Count the number of elements lesser than `x[i] + param`. Extended method to vector with sig `(x, param)` are generated.
When `X` is provided, the result is computed without allocations.
"""
tr_count_l_param(i, x; param) = count(y -> y < x[i] + param, x)

"""
    tr_count_g_param(i, x; param)
    tr_count_g_param(x; param)
    tr_count_g_param(x, X::AbstractVector; param)

Count the number of elements greater than `x[i] + param`. Extended method to vector with sig `(x, param)` are generated.
When `X` is provided, the result is computed without allocations.
"""
tr_count_g_param(i, x; param) = count(y -> y > x[i] + param, x)

# Generating vetorized versions
lazy_param(tr_count_eq_param, tr_count_l_param, tr_count_g_param)

# Bounding param

"""
    tr_count_bounding_param(i, x; param)
    tr_count_bounding_param(x; param)
    tr_count_bounding_param(x, X::AbstractVector; param)

Count the number of elements bounded (not strictly) by `x[i]` and `x[i] + param`. An extended method to vector with sig `(x, param)` is generated.
When `X` is provided, the result is computed without allocations.
"""
tr_count_bounding_param(i, x; param) = count(y -> x[i] ≤ y ≤ x[i] + param, x)

# Generating vetorized versions
lazy_param(tr_count_bounding_param)

# Val/param substractions

"""
    tr_val_minus_param(i, x; param)
    tr_val_minus_param(x; param)
    tr_val_minus_param(x, X::AbstractVector; param)

Return the difference `x[i] - param` if positive, `0.0` otherwise.  Extended method to vector with sig `(x, param)` are generated.
When `X` is provided, the result is computed without allocations.
"""
tr_val_minus_param(i, x; param) = max(0, x[i] - param)

"""
    tr_param_minus_val(i, x; param)
    tr_param_minus_val(x; param)
    tr_param_minus_val(x, X::AbstractVector; param)

Return the difference `param - x[i]` if positive, `0.0` otherwise.  Extended method to vector with sig `(x, param)` are generated.
When `X` is provided, the result is computed without allocations.
"""
tr_param_minus_val(i, x; param) = max(0, param - x[i])

# Generating vetorized versions
lazy_param(tr_val_minus_param, tr_param_minus_val)

# Continuous values substraction
"""
    tr_contiguous_vals_minus(i, x)
    tr_contiguous_vals_minus(x)
    tr_contiguous_vals_minus(x, X::AbstractVector)

Return the difference `x[i] - x[i + 1]` if positive, `0.0` otherwise. Extended method to vector with sig `(x)` are generated.
When `X` is provided, the result is computed without allocations.
"""
tr_contiguous_vals_minus(i, x; param=nothing) = length(x) == i ? 0 : tr_val_minus_param(i, x; param=x[i + 1])

"""
    tr_contiguous_vals_minus_rev(i, x)
    tr_contiguous_vals_minus_rev(x)
    tr_contiguous_vals_minus_rev(x, X::AbstractVector)

Return the difference `x[i + 1] - x[i]` if positive, `0.0` otherwise. Extended method to vector with sig `(x)` are generated.
When `X` is provided, the result is computed without allocations.
"""
function tr_contiguous_vals_minus_rev(i, x; param=nothing)
    return length(x) == i ? 0 : tr_param_minus_val(i, x; param=x[i + 1])
end

# Generating vetorized versions
lazy(tr_contiguous_vals_minus, tr_contiguous_vals_minus_rev)

"""
    transformation_layer(param = false)
Generate the layer of transformations functions of the ICN. Iff `param` value is true, also includes all the parametric transformations.
"""
function transformation_layer(param=false)
    transformations = LittleDict{Symbol,Function}(
        :identity => tr_identity,
        :count_eq => tr_count_eq,
        :count_eq_left => tr_count_eq_left,
        :count_eq_right => tr_count_eq_right,
        :count_greater => tr_count_greater,
        :count_lesser => tr_count_lesser,
        :count_g_left => tr_count_g_left,
        :count_l_left => tr_count_l_left,
        :count_g_right => tr_count_g_right,
        :count_l_right => tr_count_l_right,
        :contiguous_vals_minus => tr_contiguous_vals_minus,
        :contiguous_vals_minus_rev => tr_contiguous_vals_minus_rev,
    )

    if param
        transformations_param = LittleDict{Symbol,Function}(
            :count_eq_param => tr_count_eq_param,
            :count_l_param => tr_count_l_param,
            :count_g_param => tr_count_g_param,
            :count_bounding_param => tr_count_bounding_param,
            :val_minus_param => tr_val_minus_param,
            :param_minus_val => tr_param_minus_val,
        )
        transformations = LittleDict(union(transformations, transformations_param))
    end

    return Layer(transformations, false)
end


module UnconstrainedFESpacesTests

using Test
using Gridap.Arrays
using Gridap.TensorValues
using Gridap.ReferenceFEs
using Gridap.Geometry
using Gridap.FESpaces

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

order = 2
grid_topology = get_grid_topology(model)
polytopes = get_polytopes(grid_topology)
reffes = [LagrangianRefFE(Float64,p,order) for p in polytopes]

dirichlet_tags = ["tag_1","tag_6"]
V0 = GradConformingFESpace(reffes,model,dirichlet_tags)

f(x) = sin(4*pi*(x[1]-x[2]^2))+1

fh = interpolate(f,V0)
@test get_dirichlet_values(fh) == get_dirichlet_values(V0)

fh = interpolate_everywhere(f,V0)

fh = interpolate_dirichlet(f,V0)
@test get_free_values(fh) == zero_free_values(V0)

#trian = get_triangulation(model)
#
#using Gridap.Visualization
#
#writevtk(trian,"trian",nsubcells=10,cellfields=["fh" => fh])

end # module


module LAHMC

using LinearAlgebra

export LogLikelihood, GaussianLogLikelihood
include("loglikelihood.jl")

export Sampler, LAHMCSampler, HMCSampler
export set_x0!, sample!
include("sampler.jl")

export mcmc!, mcmc

function mcmc!(s, N, S::Sampler{T}; burnin = 0, thin = 1) where T
    if length(s) > 0
        set_x0!(S, s[end])
    end
    for k in 1:burnin
        sample!(S)
    end
    for k in 1:N
        sample!(S)
        if k % thin == 0
            push!(s, S.x)
        end
    end
    return s
end

function mcmc(x0, N, S0::Sampler{T}; kwargs...) where T
    S = copy(S0)
    set_x0!(S, x0)
    s = Array{T, 1}[]
    mcmc!(s, N, S; kwargs...)
    return s, S
end

end # module


# ------------------------------------------------------------------------------------------
# # Sequence to sequence model (S2S)
# (c) Deniz Yuret, 2018. Based on ([Sutskever et al.
# 2014](https://papers.nips.cc/paper/5346-sequence-to-sequence-learning-with-neural-
# networks.pdf)).
# S2S models learn to map input sequences to output sequences using an encoder and a decoder
# RNN. Note that this is an instructional example written in low-level Julia/Knet and it is
# slow to train. For a faster and high-level implementation please see `@doc RNN`.
# ------------------------------------------------------------------------------------------

using Pkg; haskey(Pkg.installed(),"Knet") || Pkg.add("Knet")
using Knet

# ------------------------------------------------------------------------------------------
# (image source)
# ------------------------------------------------------------------------------------------

# S2S model definition

function initmodel(H, V; atype=(gpu()>=0 ? KnetArray{Float32} : Array{Float32}))
    init(d...)=atype(xavier(d...))
    model = Dict{Symbol,Any}()
    model[:state0] = [ init(1,H), init(1,H) ]
    model[:embed1] = init(V,H)
    model[:encode] = [ init(2H,4H), init(1,4H) ]
    model[:embed2] = init(V,H)
    model[:decode] = [ init(2H,4H), init(1,4H) ]
    model[:output] = [ init(H,V), init(1,V) ]
    return model
end;

# S2S loss function and its gradient

function s2s(model, inputs, outputs)
    state = initstate(inputs[1], model[:state0])
    for input in inputs
        input = onehotrows(input, model[:embed1])
        input = input * model[:embed1]
        state = lstm(model[:encode], state, input)
    end
    EOS = eosmatrix(outputs[1], model[:embed2])
    input = EOS * model[:embed2]
    sumlogp = 0
    for output in outputs
        state = lstm(model[:decode], state, input)
        ypred = predict(model[:output], state[1])
        ygold = onehotrows(output, model[:embed2])
        sumlogp += sum(ygold .* logp(ypred,dims=2))
        input = ygold * model[:embed2]
    end
    state = lstm(model[:decode], state, input)
    ypred = predict(model[:output], state[1])
    sumlogp += sum(EOS .* logp(ypred,dims=2))
    return -sumlogp
end

s2sgrad = gradloss(s2s);

# ------------------------------------------------------------------------------------------
# (image source)
# ------------------------------------------------------------------------------------------

# A LSTM implementation in Knet

function lstm(param, state, input)
    weight,bias = param
    hidden,cell = state
    h       = size(hidden,2)
    gates   = hcat(input,hidden) * weight .+ bias
    forget  = sigm.(gates[:,1:h])
    ingate  = sigm.(gates[:,1+h:2h])
    outgate = sigm.(gates[:,1+2h:3h])
    change  = tanh.(gates[:,1+3h:4h])
    cell    = cell .* forget + ingate .* change
    hidden  = outgate .* tanh.(cell)
    return (hidden,cell)
end;

# S2S helper functions

function predict(param, input)
    input * param[1] .+ param[2]
end

function initstate(idx, state0)
    h,c = state0
    h = h .+ fill!(similar(value(h), length(idx), length(h)), 0)
    c = c .+ fill!(similar(value(c), length(idx), length(c)), 0)
    return (h,c)
end

function onehotrows(idx, embeddings)
    nrows,ncols = length(idx), size(embeddings,1)
    z = zeros(Float32,nrows,ncols)
    @inbounds for i=1:nrows
        z[i,idx[i]] = 1
    end
    oftype(value(embeddings),z)
end

let EOS=nothing; global eosmatrix
function eosmatrix(idx, embeddings)
    nrows,ncols = length(idx), size(embeddings,1)
    if EOS==nothing || size(EOS) != (nrows,ncols)
        EOS = zeros(Float32,nrows,ncols)
        EOS[:,1] .= 1
        EOS = oftype(value(embeddings), EOS)
    end
    return EOS
end
end;

# Use reversing English words as an example task
# This loads them from /usr/share/dict/words and converts each character to an int.

function readdata(file="words")
    isfile(file) || (file=download("http://people.csail.mit.edu/deniz/models/tutorial/words","words"))
    global strings = map(chomp,readlines(file))
    global tok2int = Dict{Char,Int}()
    global int2tok = Vector{Char}()
    push!(int2tok,'\n'); tok2int['\n']=1 # We use '\n'=>1 as the EOS token
    sequences = Vector{Vector{Int}}()
    for w in strings
        s = Vector{Int}()
        for c in collect(w)
            if !haskey(tok2int,c)
                push!(int2tok,c)
                tok2int[c] = length(int2tok)
            end
            push!(s, tok2int[c])
        end
        push!(sequences, s)
    end
    return sequences
end;

sequences = readdata();
for x in (sequences, strings, int2tok, tok2int); println(summary(x)); end
for x in strings[501:505]; println(x); end

# Minibatch sequences putting equal length sequences together:

function minibatch(sequences, batchsize)
    table = Dict{Int,Vector{Vector{Int}}}()
    data = Any[]
    for s in sequences
        n = length(s)
        nsequences = get!(table, n, Any[])
        push!(nsequences, s)
        if length(nsequences) == batchsize
            push!(data, [[ nsequences[i][j] for i in 1:batchsize] for j in 1:n ])
            empty!(nsequences)
        end
    end
    return data
end

batchsize, statesize, vocabsize = 128, 128, length(int2tok)
data = minibatch(sequences,batchsize)
summary(data)

# Training loop

function train(model, data, opts)
    sumloss = cntloss = 0
    for sequence in data
        grads,loss = s2sgrad(model, sequence, reverse(sequence))
        update!(model, grads, opts)
        sumloss += loss
        cntloss += (1+length(sequence)) * length(sequence[1])
    end
    return sumloss/cntloss
end

file = "rnnreverse113.jld2"; model = opts = nothing; Knet.gc() # clean memory from previous run
if (print("Train from scratch? ");readline()[1]=='y')
    # Initialize model and optimization parameters
    model = initmodel(statesize,vocabsize)
    opts = optimizers(model,Adam)
    @time for epoch=1:10
        @time loss = train(model,data,opts) # ~17 sec/epoch
        println((epoch,loss))
    end
    Knet.save(file,"model",model)
else
    isfile(file) || download("http://people.csail.mit.edu/deniz/models/tutorial/$file",file)
    model = Knet.load(file,"model")
end
summary(model)

# Test on some examples:

function translate(model, str)
    state = model[:state0]
    for c in collect(str)
        input = onehotrows(tok2int[c], model[:embed1])
        input = input * model[:embed1]
        state = lstm(model[:encode], state, input)
    end
    input = eosmatrix(1, model[:embed2]) * model[:embed2]
    output = Char[]
    for i=1:100 #while true
        state = lstm(model[:decode], state, input)
        pred = predict(model[:output], state[1])
        i = argmax(vec(Array(pred)))
        i == 1 && break
        push!(output, int2tok[i])
        input = onehotrows(i, model[:embed2]) * model[:embed2]
    end
    String(output)
end;

translate(model,"capricorn")


using Test, OMEinsumContractionOrders, OMEinsum

@testset "save load" begin
    for code in [
        EinCode([[1,2], [2,3], [3,4]], [1,4]),
        EinCode([['a','b'], ['b','c'], ['c','d']], ['a','d'])
    ]
        for optcode in [optimize_code(code, uniformsize(code, 2), GreedyMethod()),
            optimize_code(code, uniformsize(code, 2), TreeSA(nslices=1))]
            filename = tempname()
            writejson(filename, optcode)
            code2 = readjson(filename)
            @test optcode == code2
        end
    end
end


function run(s)
    # Your code here
    joltages = readlines(IOBuffer(s))
    joltages = [parse(Int, j) for j in joltages]
    sort!(joltages)
    pushfirst!(joltages, 0)
    cache = zeros(Int, size(joltages))
    cache[1] = 1
    for i in 1:length(joltages)
        if (i - 1 > 0) && (joltages[i] - joltages[i - 1] <= 3)
            cache[i] += cache[i - 1]
        end
        if (i - 2 > 0) && (joltages[i] - joltages[i - 2] <= 3)
            cache[i] += cache[i - 2]
        end
        if (i - 3 > 0) && (joltages[i] - joltages[i - 3] <= 3)
            cache[i] += cache[i - 3]
        end
    end
    return cache[end]
end

#########################################

function main()
    res, time, memory = @timed run(ARGS[1])
    println("_duration:$(time * 1000)")
    println(res)
end

main()


r=1.0;s=1.0;t=1.0;

Relaxation_success = true

# psi(x::Any,r::Float64,s::Float64,t::Float64)
@test norm(MPCCSolver.psi(s,r,s,t)-s) <= eps(Float64)*10
# x = linspace(-0.5,10,1000); y = MPCCSolver.psi(x,r,s,t)
# PyPlot.plot(x, y, color="red", linewidth=2.0, linestyle="-")
# x = linspace(-0.5,10,1000); y = MPCCSolver.psi(x,r,s,t)
# PyPlot.plot(y, x, color="blue", linewidth=2.0, linestyle="-")

#invpsi(x::Any,r::Float64,s::Float64,t::Float64)
@test norm(MPCCSolver.invpsi(s,r,s,t)-s) <= eps(Float64)*10

# x = linspace(-0.5,10,1000); y = MPCCSolver.psi(x,r,s,t)
# PyPlot.plot(x, y, color="red", linewidth=2.0, linestyle="-")
# x = linspace(-0.5,10,1000); y = MPCCSolver.psi(x,r,s,t)
# PyPlot.plot(y, x, color="blue", linewidth=2.0, linestyle="-")

#dpsi(x::Any,r::Float64,s::Float64,t::Float64)
#phi(yg::Any,yh::Any,r::Float64,s::Float64,t::Float64)
#phi(x::Any,nb_comp::Int64,r::Float64,s::Float64,t::Float64)
#dphi(yg::Any,yh::Any,r::Float64,s::Float64,t::Float64)
#dphi(x::Any,nb_comp::Int64,r::Float64,s::Float64,t::Float64)

#alpha_max(x::Float64,dx::Float64,y::Float64,dy::Float64,r::Float64,s::Float64,t::Float64)
#alpha_max : résoud l'équation en alpha où psi(x,r,s,t)=s+t*theta(x,r)
r=1.0;s=1.0;t=0.0 #résoud l'éq.

x=1.0;y=2.0;dx=1.0;dy=-1.0;
#doit rendre (0,1)
@test isequal(MPCCSolver.alpha_max(x,dx,y,dy,r,s,t),[0.0,1.0])

x=0.0;y=2.0;dx=1.0;dy=-2.0;
#solution : (1,0.5)
@test isequal(MPCCSolver.alpha_max(x,dx,y,dy,r,s,t),[1.0,0.5])

r=1.0;s=1.0;t=1.0
#on tape la contrainte :
x=1.0;y=2.0;dx=1.0;dy=0.0;

@test minimum(MPCCSolver.alpha_max(x,dx,y,dy,r,s,t))>=0.0
@test maximum(MPCCSolver.alpha_max(x,dx,y,dy,r,s,t))==Inf

x=2.0;y=1.0;dx=-1.0;dy=2.0;

@test minimum(MPCCSolver.alpha_max(x,dx,y,dy,r,s,t)) >= 0.0
@test maximum(MPCCSolver.alpha_max(x,dx,y,dy,r,s,t)) <  Inf

#on est sous (s,s) et donc on ne tape pas la contrainte virtuelle
x=0.0;y=0.0;dx=0.0;dy=-2.0;

@test isequal(maximum(MPCCSolver.alpha_max(x,dx,y,dy,r,s,t)),Inf)

#Conclusion:
@test Relaxation_success


############################################################################################
"""
$(SIGNATURES)
Compute ESS as a function of proposed temperature `λ` against current temperature `λₜ₋₁`.
Note that weights are in the original (non-log) space, and should have temperature == 1.
# Examples
```julia
```

"""
function weightedESS(weights::AbstractVector{T}, λ::T, λₜ₋₁::T) where {T<:AbstractFloat}
    ArgCheck.@argcheck 0 < λₜ₋₁ <= λ <= 1.0
    num = sum(weights[iter]^(λ - λₜ₋₁) for iter in eachindex(weights))^2
    denom = sum(weights[iter]^(λ - λₜ₋₁)^2 for iter in eachindex(weights))
    return num/denom
end

############################################################################################
"""
$(SIGNATURES)
Closure to compute new temperature `λ` as a function of weights `weights` and current temperature `λₜ₋₁`.
`ESSTarget` is the benchmark to optimize against.

# Examples
```julia
```

"""
function get_ESSDifference(weights::AbstractVector{T}, λₜ₋₁::T, ESSTarget::T) where {T<:AbstractFloat}
    ArgCheck.@argcheck 0 < ESSTarget
    function ESSDifference(λ)
        return ESSTarget - weightedESS(weights, λ, λₜ₋₁)
    end
end

############################################################################################
"""
$(SIGNATURES)
Temperature adaption that takes into account updating schedule, see Chopin, Papaspiliopoulos (2020).
# Examples
```julia
```

"""
function update(weights::AbstractVector{T}, λₜ₋₁::T, ESSTarget::T) where {T<:AbstractFloat}
    # Assing ESS closure
    ESSDifference = get_ESSDifference(weights, λₜ₋₁, ESSTarget)
    # Check if ESS > ESSTarget; !NOTE: Compute ESS, not ESSDifference!
    if weightedESS(weights, T(1.0), λₜ₋₁) >= ESSTarget
        return T(1.0)
    end
    # Else tune for new stepsize
    λ = bisect(ESSDifference, λₜ₋₁, T(1.0))
    return λ
end

"Inverse CDF of Univariate Normal Mixture Model, where prob can be vectorized as a bivariate array of observations."
function bisect(f, min, max, atol = 10e-4, increasing = sign(f(max)), iter_max = 1000)
    c = Statistics.middle(min, max)
    z = f(c) * increasing
    iter = 0
    if z > 0 #
        maxi = c
        mini = typeof(maxi)(min)
    else
        mini = c
        maxi = typeof(mini)(max)
    end
    @inbounds while abs(mini - maxi) > atol && iter < iter_max
        iter += 1
        c = Statistics.middle(mini,maxi)
        if f(c) * increasing > 0 #
            maxi = c
        else
            mini = c
        end
    end
    #return output
    if iter < iter_max
        return Statistics.middle(mini,maxi)
    else
        return -Inf
    end
end

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

Tuning container for joint uptdates of temperature.

# Fields
$(TYPEDFIELDS)
"""
struct JointTempering{A<:UpdateBool, T<:AbstractFloat} <: TemperingMethod
    "Checks if temperature will be updated after new iterations"
    adaption::A
    "current temperature."
    val::ValueHolder{T}
    "Target Effective Sample Size."
    ESSTarget::ValueHolder{T}
    "Buffer for weights used to adapt temperature"
    weights::Vector{T}
    function JointTempering(adaption::A, val::ValueHolder{T}, ESSTarget::ValueHolder{T}, weights::Vector{T}
    ) where {A<:UpdateBool,T<:AbstractFloat}
        @argcheck 0.0 < val.current <= 1.0
        @argcheck 0.0 < ESSTarget.current <= length(weights)
        return new{A,T}(adaption, val, ESSTarget, weights)
    end
end

function JointTempering(::Type{T}, adaption::B, val::F, ESSTarget::S, chains::Integer
) where {T<:AbstractFloat, B<:UpdateBool, F<:AbstractFloat, S<:AbstractFloat}
    # Return tuning struct
    return JointTempering(adaption, ValueHolder(T(val)), ValueHolder(T(ESSTarget)), zeros(T, chains))
end

############################################################################################
function update!(tempering::JointTempering, adaption::UpdateTrue, ℓweights::AbstractVector)
    # Update weights
    for iter in eachindex(tempering.weights)
        tempering.weights[iter] = ℓweights[iter]
    end
    #Normalize numerically stable
    tempering.weights .-= logsumexp(tempering.weights)
    tempering.weights .= exp.(tempering.weights)
    # Compute new temperature
    tempering.val.current = update(tempering.weights, tempering.val.current, tempering.ESSTarget.current)
    return tempering.val.current
end

############################################################################################
function update!(tempering::JointTempering, adaption::UpdateFalse, weights::AbstractVector)
    return tempering.val.current
end
function update!(tempering::JointTempering, weights::AbstractVector)
    return update!(tempering, tempering.adaption, weights)
end

############################################################################################
"""
$(SIGNATURES)
Return initial temperature of tuning container.

# Examples
```julia
```

"""
function initial(tempertune::JointTempering)
    return tempertune.val.current
end

############################################################################################
# Export

export
    JointTempering,
    initial,
    update,
    update!


include("../utils.jl")

@testset "Cholesky Space" begin
    M = Manifolds.CholeskySpace(3)
    @test repr(M) == "CholeskySpace(3)"

    types = [Matrix{Float64}]
    TEST_FLOAT32 && push!(types, Matrix{Float32})
    TEST_STATIC_SIZED && push!(types, MMatrix{3,3,Float64,9})

    for T in types
        A(α) = [1.0 0.0 0.0; 0.0 cos(α) sin(α); 0.0 -sin(α) cos(α)]
        ptsF = [#
            cholesky([1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1]).L,
            cholesky([2.0 0.0 0.0; 0.0 2.0 0.0; 0.0 0.0 1]).L,
            cholesky(
                A(π / 6) * [1.0 0.0 0.0; 0.0 2.0 0.0; 0.0 0.0 1] * transpose(A(π / 6)),
            ).L,
        ]
        pts = [convert(T, a) for a in ptsF]
        test_manifold(
            M,
            pts;
            test_injectivity_radius=false,
            test_default_vector_transport=true,
            vector_transport_methods=[
                ParallelTransport(),
                SchildsLadderTransport(),
                PoleLadderTransport(),
            ],
            test_vee_hat=false,
            exp_log_atol_multiplier=8.0,
            test_inplace=true,
        )
    end
    @testset "Test Error cases in is_point and is_vector" begin
        pt1f = zeros(2, 3) # wrong size
        pt2f = [1.0 0.0 0.0; 0.0 -1.0 0.0; 0.0 0.0 1.0] # nonpos diag
        pt3f = [2.0 0.0 1.0; 0.0 1.0 0.0; 0.0 0.0 4.0] # no lower and nonsym
        pt4 = [2.0 0.0 0.0; 1.0 2.0 0.0; 0.0 0.0 4.0]
        @test !is_point(M, pt1f)
        @test_throws DomainError is_point(M, pt1f, true)
        @test !is_point(M, pt2f)
        @test_throws DomainError is_point(M, pt2f, true)
        @test !is_point(M, pt3f)
        @test_throws DomainError is_point(M, pt3f, true)
        @test is_point(M, pt4)
        @test !is_vector(M, pt3f, pt1f)
        @test_throws DomainError is_vector(M, pt3f, pt1f, true)
        @test !is_vector(M, pt4, pt1f)
        @test_throws DomainError is_vector(M, pt4, pt1f, true)
        @test !is_vector(M, pt4, pt3f)
        @test_throws DomainError is_vector(M, pt4, pt3f, true)
        @test is_vector(M, pt4, pt2f)
    end
end


using Mosek, Base.Test

function test_lo1()
    printstream(msg::String) = nothing # print(msg)

    bkc = [MSK_BK_FX, MSK_BK_LO, MSK_BK_UP]
    blc = [30.0, 15.0, -Inf]
    buc = [30.0, +Inf, 25.0]
    bkx = [ MSK_BK_LO, MSK_BK_RA, MSK_BK_LO, MSK_BK_LO ]
    blx = [   0.0,  0.0,    0.0,    0.0]
    bux = [+Inf, 10.0, +Inf, +Inf]
    numvar = length(bkx)
    numcon = length(bkc)
    c = [ 3.0, 1.0, 5.0, 1.0 ]
    A = sparse([1, 2, 1, 2, 3, 1, 2, 2, 3],
               [1, 1, 2, 2, 2, 3, 3, 4, 4],
               [3.0, 2.0, 1.0, 1.0, 2.0, 2.0, 3.0, 1.0, 3.0 ],
               numcon,numvar)

    task = maketask()
    putstreamfunc(task,MSK_STREAM_LOG,printstream)
    putobjname(task,"lo1")
    appendcons(task,numcon)
    for i=1:numcon
        putconname(task,i,@sprintf("c%02d",i))
    end
    appendvars(task,numvar)
    for j=1:numvar
        putvarname(task,j,@sprintf("x%02d",j))
    end
    putclist(task,[1,2,3,4], c)

    putacolslice(task,
                 1, numvar+1,
                 A.colptr[1:numvar], A.colptr[2:numvar+1],
                 A.rowval,
                 A.nzval)
    putvarboundslice(task, 1, numvar+1, bkx,blx,bux)
    putconboundslice(task,1,numcon+1,bkc,blc,buc)
    putobjsense(task,MSK_OBJECTIVE_SENSE_MAXIMIZE)
    optimize(task)
    solutionsummary(task,MSK_STREAM_MSG)

    solsta = getsolsta(task,MSK_SOL_BAS)
    prosta = getprosta(task,MSK_SOL_BAS)

    @test getsolsta(task,MSK_SOL_ITR) in [ MSK_SOL_STA_OPTIMAL, MSK_SOL_STA_NEAR_OPTIMAL ]
    @test solsta in [ MSK_SOL_STA_OPTIMAL, MSK_SOL_STA_NEAR_OPTIMAL ]
    @test prosta in [MSK_PRO_STA_PRIM_AND_DUAL_FEAS,MSK_PRO_STA_NEAR_PRIM_AND_DUAL_FEAS]

    if solsta in     [ MSK_SOL_STA_OPTIMAL,
                       MSK_SOL_STA_NEAR_OPTIMAL ]
        xx = getxx(task,MSK_SOL_BAS)
        # check feasibility and optimality of solution
    end
end

function test_qo1()
    printstream(msg::String) = nothing # print(msg)
    bkc   = [ MSK_BK_LO ]
    blc   = [ 1.0 ]
    buc   = [ Inf ]
    bkx   = [ MSK_BK_LO, MSK_BK_LO, MSK_BK_LO ]
    blx   = [ 0.0,  0.0, 0.0 ]
    bux   = [ Inf,  Inf, Inf ]
    numvar = length(bkx)
    numcon = length(bkc)
    c     = [ 0.0, -1.0, 0.0 ]
    A     = sparse( [ 1, 1, 1 ],
                    [ 1, 2, 3 ],
                    [ 1.0, 1.0, 1.0 ],
                    numcon, numvar )
    task = maketask()
    putstreamfunc(task,MSK_STREAM_LOG,printstream)
    appendcons(task,numcon)
    appendvars(task,numvar)
    putclist(task,[1:numvar;],c)
    putvarboundslice(task,1,numvar+1,bkx,blx,bux)
    putacolslice(task,1,numvar+1,
                 A.colptr[1:numvar],A.colptr[2:numvar+1],
                 A.rowval,A.nzval)
    qsubi = [ 1,   2,    3,   3   ]
    qsubj = [ 1,   2,    1,   3   ]
    qval  = [ 2.0, 0.2, -1.0, 2.0 ]
    putqobj(task,qsubi,qsubj,qval)
    putobjsense(task,MSK_OBJECTIVE_SENSE_MINIMIZE)
    optimize(task)
    solutionsummary(task,MSK_STREAM_MSG)

    prosta = getprosta(task,MSK_SOL_ITR)
    solsta = getsolsta(task,MSK_SOL_ITR)

    @test solsta in (MSK_SOL_STA_OPTIMAL, MSK_SOL_STA_NEAR_OPTIMAL)
    @test prosta in (MSK_PRO_STA_PRIM_AND_DUAL_FEAS,MSK_PRO_STA_NEAR_PRIM_AND_DUAL_FEAS)

    if solsta in     [ MSK_SOL_STA_OPTIMAL,
                       MSK_SOL_STA_NEAR_OPTIMAL ]
        xx = getxx(task,MSK_SOL_ITR)
        # check feasibility and optimality of solution
    end
end

function test_qcqo1()
    printstream(msg::String) = nothing # print(msg)
    task = maketask()
    putstreamfunc(task,MSK_STREAM_LOG,m -> nothing)
    bkc   = [ MSK_BK_LO ]
    blc   = [ 1.0 ]
    buc   = [ Inf ]
    bkx   = [ MSK_BK_LO, MSK_BK_LO, MSK_BK_LO ]
    blx   = [ 0.0,  0.0, 0.0 ]
    bux   = [ Inf,  Inf, Inf ]
    c     = [ 0.0, -1.0, 0.0 ]
    asub  = [ 1 ,2, 3 ]
    aval  = [ 1.0, 1.0, 1.0 ]
    numvar = length(bkx)
    numcon = length(bkc)
    appendcons(task,numcon)
    appendvars(task,numvar)
    putcfix(task,0.0)
    putclist(task,[1:numvar;],c)
    putvarboundslice(task,1,numvar+1,bkx,blx,bux)
    putarow(task,1,asub,aval)
    putconbound(task,1,bkc[1],blc[1],buc[1])
    qsubi = [ 1,   2,    3,   3   ]
    qsubj = [ 1,   2,    1,   3   ]
    qval  = [ 2.0, 0.2, -1.0, 2.0 ]
    putqobj(task,qsubi,qsubj,qval)
    qsubi = [  1,    2,    3,   3   ]
    qsubj = [  1,    2,    3,   1   ]
    qval  = [ -2.0, -2.0, -0.2, 0.2 ]
    putqconk(task,1, qsubi,qsubj, qval)
    putobjsense(task,MSK_OBJECTIVE_SENSE_MINIMIZE)
    optimize(task)
    solutionsummary(task,MSK_STREAM_MSG)

    prosta = getprosta(task,MSK_SOL_ITR)
    solsta = getsolsta(task,MSK_SOL_ITR)

    @test solsta in (MSK_SOL_STA_OPTIMAL, MSK_SOL_STA_NEAR_OPTIMAL)
    @test prosta in (MSK_PRO_STA_PRIM_AND_DUAL_FEAS,MSK_PRO_STA_NEAR_PRIM_AND_DUAL_FEAS)

    if solsta in     [ MSK_SOL_STA_OPTIMAL,
                       MSK_SOL_STA_NEAR_OPTIMAL ]
        xx = getxx(task,MSK_SOL_ITR)
        # check feasibility and optimality of solution
    end
end

function test_milo1()
    printstream(msg::String) = print(msg)  #print(msg)
    task = maketask()
    putstreamfunc(task,MSK_STREAM_LOG,printstream)
    bkc = [ MSK_BK_UP, MSK_BK_LO  ]
    blc = [      -Inf,      -4.0  ]
    buc = [     250.0,       Inf  ]
    bkx = [ MSK_BK_LO, MSK_BK_LO  ]
    blx = [       0.0,       0.0  ]
    bux = [       Inf,       Inf  ]
    c   = [       1.0,      0.64 ]
    A    = sparse( [ 1, 1, 2, 2],
                   [ 1, 2, 1, 2],
                   [ 50.0, 3.0, 31.0, -2.0])
    numvar = length(bkx)
    numcon = length(bkc)
    appendcons(task,numcon)
    appendvars(task,numvar)
    putclist(task,[1:numvar;],c)
    putvarboundslice(task,1,numvar+1,bkx,blx,bux)
    putacolslice(task,1,numvar+1, A.colptr[1:numvar],A.colptr[2:numvar+1],A.rowval,A.nzval)
    putconboundslice(task,1,numcon+1,bkc,blc,buc)
    putobjsense(task,MSK_OBJECTIVE_SENSE_MAXIMIZE)
    putvartypelist(task,[ 1, 2 ], [ MSK_VAR_TYPE_INT, MSK_VAR_TYPE_INT ])
    optimize(task)
    solutionsummary(task,MSK_STREAM_MSG)
    prosta = getprosta(task,MSK_SOL_ITG)
    solsta = getsolsta(task,MSK_SOL_ITG)

    @test solsta in (MSK_SOL_STA_INTEGER_OPTIMAL, MSK_SOL_STA_NEAR_INTEGER_OPTIMAL)
    @test prosta in (MSK_PRO_STA_PRIM_FEAS,MSK_PRO_STA_NEAR_PRIM_FEAS)

    if solsta in     [ MSK_SOL_STA_INTEGER_OPTIMAL,
                       MSK_SOL_STA_NEAR_INTEGER_OPTIMAL ]
        xx = getxx(task,MSK_SOL_ITG)
    end
end

function test_cqo()
        printstream(msg::String) = nothing # print(msg)
    callback(where,dinf,iinf,liinf) = 0
    bkc = [ MSK_BK_FX ]
    blc = [ 1.0 ]
    buc = [ 1.0 ]
    c   = [               0.0,              0.0,              0.0,
                          1.0,              1.0,              1.0 ]
    bkx = [ MSK_BK_LO,MSK_BK_LO,MSK_BK_LO,
            MSK_BK_FR,MSK_BK_FR,MSK_BK_FR ]
    blx = [               0.0,              0.0,              0.0,
                          -Inf,             -Inf,             -Inf ]
    bux = [               Inf,              Inf,              Inf,
                          Inf,              Inf,              Inf ]
    asub  = [ 1 ,2, 3 ]
    aval  = [ 1.0, 1.0, 1.0 ]
    numvar = length(bkx)
    numcon = length(bkc)
    task = maketask()
    putstreamfunc(task,MSK_STREAM_LOG,printstream)
    putcallbackfunc(task,callback)
    appendcons(task,numcon)
    appendvars(task,numvar)
    putclist(task,[1:6;],c)
    putvarboundslice(task,1,numvar+1,bkx,blx,bux)
    putarow(task,1,asub,aval)
    putconbound(task,1,bkc[1],blc[1],buc[1])
    appendcone(task,MSK_CT_QUAD, 0.0, [ 4, 1, 2 ])
    appendcone(task,MSK_CT_RQUAD, 0.0, [ 5, 6, 3 ])
    putobjsense(task,MSK_OBJECTIVE_SENSE_MINIMIZE)
    optimize(task)
    solutionsummary(task,MSK_STREAM_MSG)
    prosta = getprosta(task,MSK_SOL_ITR)
    solsta = getsolsta(task,MSK_SOL_ITR)

    @test solsta in (MSK_SOL_STA_OPTIMAL, MSK_SOL_STA_NEAR_OPTIMAL)
    @test prosta in (MSK_PRO_STA_PRIM_AND_DUAL_FEAS,MSK_PRO_STA_NEAR_PRIM_AND_DUAL_FEAS)

    if solsta in     [ MSK_SOL_STA_OPTIMAL,
                       MSK_SOL_STA_NEAR_OPTIMAL ]
        xx = getxx(task,MSK_SOL_ITR)
        # check feasibility and optimality of solution
    end
end

function test_sdo1()
    task = maketask()
    putstreamfunc(task,MSK_STREAM_LOG,m -> nothing)
    bkc = [MSK_BK_FX,
           MSK_BK_FX]
    blc = [1.0, 0.5]
    buc = [1.0, 0.5]
    A = sparse( [1,2,2],[1,2,3],[1.0, 1.0, 1.0])
    conesub = [1, 2, 3]
    barci = [1, 2, 2, 3, 3]
    barcj = [1, 1, 2, 2, 3]
    barcval = [2.0, 1.0, 2.0, 1.0, 2.0]
    barai   = Any[ [1, 2, 3],
                   [1, 2, 3, 2, 3, 3] ]
    baraj   = Any[ [1, 2, 3],
                   [1, 1, 1, 2, 2, 3] ]
    baraval = Any[ [1.0, 1.0, 1.0],
                   [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] ]
    numvar = 3
    numcon = length(bkc)
    barvardim = [3]
    appendvars(task,numvar)
    appendcons(task,numcon)
    appendbarvars(task,barvardim)
    putcj(task, 1, 1.0)
    putvarboundslice(task,1,numvar+1,
                     [ MSK_BK_FR for i in 1:numvar ],
                     [ -Inf      for i in 1:numvar ],
                     [ +Inf      for i in 1:numvar ])
    putconboundslice(task,1,numcon+1, bkc,blc,buc)
    putacolslice(task,1,numvar+1,
                 A.colptr[1:numvar], A.colptr[2:numvar+1],
                 A.rowval,A.nzval)
    appendcone(task,MSK_CT_QUAD, 0.0, conesub)
    symc  = appendsparsesymmat(task,barvardim[1],
                               barci,
                               barcj,
                               barcval)

    syma0 = appendsparsesymmat(task,
                               barvardim[1],
                               barai[1],
                               baraj[1],+
                               baraval[1])
    syma1 = appendsparsesymmat(task,barvardim[1],
                               barai[2],
                               baraj[2],
                               baraval[2])
    putbarcj(task,1, [symc], [1.0])
    putbaraij(task,1, 1, [syma0], [1.0])
    putbaraij(task,2, 1, [syma1], [1.0])
    putobjsense(task,MSK_OBJECTIVE_SENSE_MINIMIZE)
    optimize(task)
    solutionsummary(task,MSK_STREAM_MSG)
    prosta = getprosta(task,MSK_SOL_ITR)
    solsta = getsolsta(task,MSK_SOL_ITR)

    @test solsta in (MSK_SOL_STA_OPTIMAL, MSK_SOL_STA_NEAR_OPTIMAL)
    @test prosta in (MSK_PRO_STA_PRIM_AND_DUAL_FEAS,MSK_PRO_STA_NEAR_PRIM_AND_DUAL_FEAS)

    if solsta in     [ MSK_SOL_STA_OPTIMAL,
                       MSK_SOL_STA_NEAR_OPTIMAL ]
        xx = getxx(task,MSK_SOL_ITR)
        # check feasibility and optimality of solution
    end
end

function test_nlo1()
    t = maketask()
    putstreamfunc(t,MSK_STREAM_LOG,m -> nothing)
    appendvars(t,3)
    appendcons(t,1)
    putvarbound(t,1,MSK_BK_LO,0.0, +Inf)  # x0
    putvarbound(t,2,MSK_BK_LO,0.0, +Inf)  # x1
    putvarbound(t,3,MSK_BK_LO,0.0, +Inf)  # x2
    putarow(t,1,[1,2,3],[1.0,1.0,1.0])   # x0 + x1 + x2
    putconbound(t,1,MSK_BK_FX,1.0, 1.0)  # = 1.0
    putclist(t,[1], [-1.0])

    function indexof(v,a)
        for i in 1:length(a)
            if a[i] == v return i
            end
        end
        return 0
    end
    evalobj(x :: Vector{Float64}) = log(x[2]+x[3])
    evalconi(x:: Vector{Float64},i:: Int32) = 0.0
    function grdobj(x :: Vector{Float64},sub:: Vector{Int32}, val:: Vector{Float64})
        for j in 1:length(sub)
            if sub[j] == 2 || sub[j] == 3
                val[j] = -1.0 / (x[2]+x[3])
            end
        end
    end
    function grdconi(x  :: Vector{Float64},
                     i  :: Int32,
                     sub:: Vector{Int32},
                     val:: Vector{Float64})
        none
    end
    function grdlag(x ::   Vector{Float64},
                    yo::   Float64,
                    yc::   Vector{Float64},
                    subi:: Vector{Int32},
                    val::  Vector{Float64})
        val[2] = - yo * 1.0 / (x[2]+x[3])
        val[3] = - yo * 1.0 / (x[2]+x[3])
    end
    function heslag(x ::      Vector{Float64},
                    yo::      Float64,
                    yc::      Vector{Float64},
                    subi::    Vector{Int32},
                    hessubi:: Vector{Int32},
                    hessubj:: Vector{Int32},
                    hesval::  Vector{Float64})

        hessubi[1] = 2; hessubj[1] = 2; hesval[1] = (x[2]+x[3])^(-2)
        hessubi[2] = 3; hessubj[2] = 2; hesval[2] = (x[2]+x[3])^(-2)
        hessubi[3] = 3; hessubj[3] = 3; hesval[3] = (x[2]+x[3])^(-2)
    end

    putnlcallbacks(t,
                   [2,3], # subscripts of non-zeros in the gradient of the objective
                   Int[], # subscripts of non-zeros in the gradient of the constraints
                   [1,1], # rowptr for subscripts of non-zeros in the gradient of the constraints
                   [2,3,3], # hessubi
                   [2,2,3], # hessubj
                   evalobj,
                   evalconi,
                   grdlag,
                   grdobj,
                   grdconi,
                   heslag)

    optimize(t)

    solsta = getsolsta(t,MSK_SOL_ITR)
    prosta = getprosta(t,MSK_SOL_ITR)

    @test solsta in (MSK_SOL_STA_OPTIMAL, MSK_SOL_STA_NEAR_OPTIMAL)
    @test prosta in (MSK_PRO_STA_PRIM_AND_DUAL_FEAS,MSK_PRO_STA_NEAR_PRIM_AND_DUAL_FEAS)
end

@testset "[apitest]" begin
    @testset "lo1" begin
        test_lo1()
    end

    @testset "qo1" begin
        test_qo1()
    end

    @testset "qcqo1" begin
        test_qcqo1()
    end

    @testset "milo1" begin
        test_milo1()
    end

    @testset "cqo" begin
        test_cqo()
    end

    @testset "sdo1" begin
        test_sdo1()
    end

    @testset "nlo1" begin
        test_nlo1()
    end
end


"""
	qpvc : example from Bazinga.jl

	Quadratic program with vanishing constraints, from [KPB11].

	Original formulations:

	minimize      1/2 x' Q x + x' q
	subject to    x[i] ≥ 0                             ∀ i ∈ [1:nvc]
                  x[i] (G[i,:] x - g[i]) ≥ 0           ∀ i ∈ [1:nvc]

	Reformulation as a constrained structured problem in the form

	minimize     f(x)
	subject to   c(x) in S

	where
	    f(x) = 1/2 x' Q x + x' q
	    c(x) = [ x[1:nvc] ]
       	       [ G x - g  ]
	    S    = { (a,b) | ∀i ∈ [1:nvc] a_i = 0  or  a_i, b_i ≥ 0 }

    References:
    [KPB11]     Kirches, Potschka, Bock,, Sager, "A parametric active set method
                for quadratic programs with vanishing constraints" (2011).
                Pacific Journal of Optimization
"""

using Bazinga, OptiMo
using Random, LinearAlgebra
using DataFrames, Query, CSV
using Printf, Plots

###################################################################################
# problem definition
###################################################################################
mutable struct QPVC <: AbstractOptiModel
    meta::OptiModelMeta
    Q::Matrix
    q::Vector
    G::Matrix
    g::Vector
    nvc::Int
end

function QPVC(; nx::Int = 100, nvc::Int = 20, x0::Vector = zeros(Float64, nx))
    @assert nx >= 2 * nvc
    name = "qpvc"
    ncon = 2 * nvc
    R = Float64
    Q = randn(R, nx, nx)
    Q .= (Q' * Q)
    q = randn(R, nx)
    G = randn(R, nvc, nx)
    g = randn(R, nvc)
    meta = OptiModelMeta(nx, ncon, x0 = x0, name = name)
    return QPVC(meta, Q, q, G, g, nvc)
end

# necessary methods:
# obj, grad!: cons!, jprod!, jtprod!, proj!, prox!, objprox!
function OptiMo.obj(prob::QPVC, x::AbstractVector)
    OptiMo.@lencheck prob.meta.nvar x
    return 0.5 * x' * prob.Q * x + x' * prob.q
end

function OptiMo.grad!(prob::QPVC, x::AbstractVector, dfx::AbstractVector)
    OptiMo.@lencheck prob.meta.nvar x
    dfx .= prob.Q * x .+ prob.q
    return nothing
end

function OptiMo.cons!(prob::QPVC, x::AbstractVector, cx::AbstractVector)
    OptiMo.@lencheck prob.meta.nvar x
    OptiMo.@lencheck prob.meta.ncon cx
    cx .= [
        x[1:prob.nvc]
        prob.G * x .- prob.g
    ]
    return nothing
end

function OptiMo.jprod!(prob::QPVC, x::AbstractVector, v::AbstractVector, Jv::AbstractVector)
    OptiMo.@lencheck prob.meta.nvar x v
    OptiMo.@lencheck prob.meta.ncon Jv
    Jv .= [
        v[1:prob.nvc]
        prob.G * v
    ]
    return nothing
end

function OptiMo.jtprod!(
    prob::QPVC,
    x::AbstractVector,
    v::AbstractVector,
    Jtv::AbstractVector,
)
    OptiMo.@lencheck prob.meta.nvar x Jtv
    OptiMo.@lencheck prob.meta.ncon v
    Jtv .= prob.G' * v[prob.nvc+1:prob.meta.ncon]
    Jtv[1:prob.nvc] .+= v[1:prob.nvc]
    return nothing
end

function OptiMo.prox!(prob::QPVC, x::AbstractVector, a::Real, z::AbstractVector)
    OptiMo.@lencheck prob.meta.nvar x z
    z .= x
    z[1:prob.nvc] .= max.(x[1:prob.nvc], 0)
    return nothing
end

function OptiMo.objprox!(prob::QPVC, x::AbstractVector, a::Real, z::AbstractVector)
    OptiMo.@lencheck prob.meta.nvar x z
    z .= x
    z[1:prob.nvc] .= max.(x[1:prob.nvc], 0)
    return 0.0
end

function OptiMo.proj!(prob::QPVC, cx::AbstractVector, px::AbstractVector)
    OptiMo.@lencheck prob.meta.ncon cx px
    for i = 1:prob.nvc
        a = cx[i]
        b = cx[i+prob.nvc]
        a = max(a, 0)
        if a + b < 0
            # if a + b ≤ 0 ...
            a = 0
        else
            b = max(0, b)
        end
        px[i] = a
        px[i+prob.nvc] = b
    end
    return nothing
end

foldername = "/home/alberto/Documents/Bazinga.jl/demo/data/"
filename = "qpvc"

# problem build
problem = QPVC()

# solver build
solver = Bazinga.ALPX(max_sub_iter = 1000, verbose = false)
# solver warm-up
out = solver(problem)

data = DataFrame()
ntests = 1000

for i = 1:ntests
    local p_nx = rand(10:250)
    local p_nvc = Int(ceil(p_nx / 5))
    local problem = QPVC(nx = p_nx, nvc = p_nvc)
    local out = solver(problem)
    #print(out)
    push!(
        data,
        (
            id = i,
            nx = p_nx,
            nvc = p_nvc,
            time = out.time,
            iters = out.iterations,
            subiters = out.solver[:sub_iterations],
            cviol = out.cviolation,
            optim = out.optimality,
            cslack = out.solver[:cslackness],
            solved = out.status == :first_order ? 1 : 0,
        ),
    )
    @printf "."
    if mod(i, 50) == 0
        @printf "\n"
    end
end
@printf "\n"

# write
CSV.write(
    foldername * filename * ".csv",
    data,
    header = false,
)

max_cviol = maximum(data[!, 5])
max_optim = maximum(data[!, 6])
max_cslack = maximum(data[!, 7])
datatmp = data |> @filter(_.solved == 1) |> DataFrame
n_first_order = size(datatmp, 1)

pyplot()

scatter( data[!,2], data[!,4],
    color = :blue,
    marker = :circle,
    markerstrokewidth = 0,
    legend = false,
    yaxis= :log, xaxis=:log,
    xlabel="Problem size nx",
    ylabel="Run time [s]"
)

savefig(foldername * filename * ".pdf")


using MLDatasets, CUDA, ChowLiuTrees

function mnist_cpu()
    train_int = transpose(reshape(MNIST.traintensor(UInt8), 28*28, :));
    test_int = transpose(reshape(MNIST.testtensor(UInt8), 28*28, :));

    function bitsfeatures(data_int)
        data_bits = zeros(Bool, size(data_int,1), 28*28*8)
        for ex = 1:size(data_int,1), pix = 1:size(data_int,2)
            x = data_int[ex,pix]
            for b = 0:7
                if (x & (one(UInt8) << b)) != zero(UInt8)
                    data_bits[ex, (pix-1)*8+b+1] = true
                end
            end
        end
        data_bits
    end

    train_cpu = bitsfeatures(train_int);
    test_cpu = bitsfeatures(test_int);

    train_cpu, test_cpu
end

function mnist_gpu()
    cu.(mnist_cpu())
end

# Mnist GPU data
train_gpu, test_gpu = mnist_gpu()

# Learn a Chow Liu Tree Structure
@time learn_chow_liu_tree(train_gpu; Float=Float64, pseudocount=1.0)
@time learn_chow_liu_tree(train_gpu; Float=Float64, pseudocount=1.0)

@time learn_chow_liu_tree(train_gpu; Float=Float32, pseudocount=1.0) # faster
@time learn_chow_liu_tree(train_gpu; Float=Float32, pseudocount=1.0) # faster


module Sim
include("./name2idx/parameters.jl")
include("./name2idx/species.jl")
include("./set_model.jl")

using .C
using .V
using Printf
using DelayDiffEq

const dt = 0.01

function solvedde(
        diffeq::Function,
        u0::Vector{Float64},
        history::Vector{Float64},
        tspan::Tuple{Float64,Float64},
        p::Vector{Float64},
        tau::Float64)
    h(p,t) = history
    lags = [tau]
    prob = DDEProblem(diffeq,u0,h,tspan,p;constant_lags=lags)
    alg = MethodOfSteps(BS3())
    sol = solve(prob,alg,saveat=dt)
    return sol
end

const tspan = (0.0,48.0)
const t = collect(tspan[1]:dt:tspan[end])

const n_condition = 4

p53_conc = Matrix{Float64}(undef, length(t), n_condition)

for i in 1:n_condition
    local p::Vector{Float64}
    local u0::Vector{Float64}

    if i == 1
        p = param_values_U2OS()
        u0 = initialValues_U2OS()
        p[C.Eto] = 1
    elseif i == 2
        p = param_values_U2OS()
        u0 = initialValues_U2OS()
        p[C.Eto] = 100
    elseif i == 3
        p = param_values_MCF7()
        u0 = initialValues_MCF7()
        p[C.Eto] = 5
    elseif i == 4
        p = param_values_MCF7()
        u0 = initialValues_MCF7()
        p[C.Eto] = 200
    end

    sol = solvedde(diffeq,u0,u0,tspan,p,2.1)
    for j in eachindex(t)
        p53_conc[j,i] = sol.u[j][V.p53]
    end
end

end  # module


@testset "Maximum flow" begin
    #### Graphs for testing
    graphs = [
    # Graph with 8 vertices
    (8,
    [
    (1, 2, 10), (1, 3, 5), (1, 4, 15), (2, 3, 4), (2, 5, 9),
    (2, 6, 15), (3, 4, 4), (3, 6, 8), (4, 7, 16), (5, 6, 15),
    (5, 8, 10), (6, 7, 15), (6, 8, 10), (7, 3, 6), (7, 8, 10)
    ],
    1, 8, # source/target
    3,    # answer for default capacity
    28,   # answer for custom capacity
    15, 5 # answer for restricted capacity/restriction
    ),

    # Graph with 6 vertices
    (6,
    [
    (1, 2, 9), (1, 3, 9), (2, 3, 10), (2, 4, 8), (3, 4, 1),
    (3, 5, 3), (5, 4, 8), (4, 6, 10), (5, 6, 7)
    ],
    1, 6, # source/target
    2,   # answer for default capacity
    12,  # answer for custom capacity
    8, 5  # answer for restricted capacity/restriction
    )
    ]

    for (nvertices, flow_edges, s, t, fdefault, fcustom, frestrict, caprestrict) in graphs
        flow_graph = Graphs.DiGraph(nvertices)
        for g in testdigraphs(flow_graph)
            capacity_matrix = zeros(Int, nvertices, nvertices)
            for e in flow_edges
                u, v, f = e
                Graphs.add_edge!(g, u, v)
                capacity_matrix[u, v] = f
            end

            # Test DefaultCapacity
            d = @inferred(GraphsFlows.DefaultCapacity(g))
            T = eltype(d)
            @test typeof(d) <: AbstractMatrix{T}
            @test d[s, t] == 0
            @test size(d) == (nvertices, nvertices)
            @test typeof(transpose(d)) <: GraphsFlows.DefaultCapacity
            @test typeof(adjoint(d)) <: GraphsFlows.DefaultCapacity

            # Test all algorithms - type instability in PushRelabel #553
            for ALGO in [EdmondsKarpAlgorithm, DinicAlgorithm, BoykovKolmogorovAlgorithm, PushRelabelAlgorithm]
                @test maximum_flow(g, s, t, algorithm=ALGO())[1] == fdefault
                @test maximum_flow(g, s, t, capacity_matrix, algorithm=ALGO())[1] == fcustom
                @test maximum_flow(g, s, t, capacity_matrix, algorithm=ALGO(), restriction=caprestrict)[1] == frestrict
            end
        end
    end
end


struct QNVal
  name::SmallString
  val::Int
  modulus::Int

  function QNVal(name, v::Int, m::Int=1)
    am = abs(m)
    if am > 1
      return new(SmallString(name), mod(v, am), m)
    end
    return new(SmallString(name), v, m)
  end
end

QNVal(v::Int, m::Int=1) = QNVal("", v, m)
QNVal() = QNVal("", 0, 0)

name(qv::QNVal) = qv.name
val(qv::QNVal) = qv.val
modulus(qv::QNVal) = qv.modulus
isactive(qv::QNVal) = modulus(qv) != 0
(qv1::QNVal < qv2::QNVal) = (name(qv1) < name(qv2))

function qn_mod(val::Int, modulus::Int)
  amod = abs(modulus)
  amod <= 1 && return val
  return mod(val, amod)
end

function -(qv::QNVal)
  return QNVal(name(qv), qn_mod(-val(qv), modulus(qv)), modulus(qv))
end

zero(qv::QNVal) = QNVal(name(qv), 0, modulus(qv))

(dir::Arrow * qv::QNVal) = QNVal(name(qv), Int(dir) * val(qv), modulus(qv))

(qv::QNVal * dir::Arrow) = (dir * qv)

function pm(qv1::QNVal, qv2::QNVal, fac::Int)
  if name(qv1) != name(qv2)
    error("Cannot add QNVals with different names \"$(name(qv1))\", \"$(name(qv2))\"")
  end
  if modulus(qv1) != modulus(qv2)
    error(
      "QNVals with matching name \"$(name(qv1))\" cannot have different modulus values "
    )
  end
  m1 = modulus(qv1)
  if m1 == 1 || m1 == -1
    return QNVal(name(qv1), val(qv1) + fac * val(qv2), m1)
  end
  return QNVal(name(qv1), Base.mod(val(qv1) + fac * val(qv2), abs(m1)), m1)
end

(qv1::QNVal + qv2::QNVal) = pm(qv1, qv2, +1)
(qv1::QNVal - qv2::QNVal) = pm(qv1, qv2, -1)

const ZeroVal = QNVal()

const maxQNs = 4
const QNStorage = SVector{maxQNs,QNVal}
const MQNStorage = MVector{maxQNs,QNVal}

"""
A QN object stores a collection of up to four
named values such as ("Sz",1) or ("N",0).
These values can include a third integer "m"
which makes them obey addition modulo m, for
example ("P",1,2) for a value obeying addition mod 2.
(The default is regular integer addition).

Adding or subtracting pairs of QN objects performs
addition and subtraction element-wise on each of
the named values. If a name is missing from the
collection, its value is treated as zero.
"""
struct QN
  data::QNStorage

  function QN()
    s = QNStorage(ntuple(_ -> ZeroVal, Val(maxQNs)))
    return new(s)
  end

  QN(s::QNStorage) = new(s)
end

QN(mqn::MQNStorage) = QN(QNStorage(mqn))
QN(mqn::NTuple{N,QNVal}) where {N} = QN(QNStorage(mqn))

function hash(obj::QN, h::UInt)
  # TODO: use an MVector or SVector
  # for performance here; put non-zero QNVals
  # to front and slice when passing to hash
  nzqv = QNVal[]
  for qv in obj.data
    if val(qv) != 0
      push!(nzqv, qv)
    end
  end
  return hash(nzqv, h)
end

"""
    QN(qvs...)

Construct a QN from a set of up to four
named value tuples.

Examples
```julia
q = QN(("Sz",1))
q = QN(("N",1),("Sz",-1))
q = QN(("P",0,2),("Sz",0)).
```
"""
function QN(qvs...)
  m = MQNStorage(ntuple(_ -> ZeroVal, Val(maxQNs)))
  for (n, qv) in enumerate(qvs)
    m[n] = QNVal(qv...)
  end
  Nvals = length(qvs)
  sort!(@view m[1:Nvals]; by=name, alg=InsertionSort)
  for n in 1:(length(qvs) - 1)
    if name(m[n]) == name(m[n + 1])
      error("Duplicate name \"$(name(m[n]))\" in QN")
    end
  end
  return QN(QNStorage(m))
end

"""
    QN(name,val::Int,modulus::Int=1)

Construct a QN with a single named value
by providing the name, value, and optional
modulus.
"""
QN(name, val::Int, modulus::Int=1) = QN((name, val, modulus))

"""
    QN(val::Int,modulus::Int=1)

Construct a QN with a single unnamed value
(equivalent to the name being the empty string)
with optional modulus.
"""
QN(val::Int, modulus::Int=1) = QN(("", val, modulus))

data(qn::QN) = qn.data

getindex(q::QN, n::Int) = getindex(data(q), n)

length(qn::QN) = length(data(qn))

lastindex(qn::QN) = length(qn)

isactive(qn::QN) = isactive(qn[1])

function nactive(q::QN)
  for n in 1:maxQNs
    !isactive(q[n]) && (return n - 1)
  end
  return maxQNs
end

function iterate(qn::QN, state::Int=1)
  (state > length(qn)) && return nothing
  return (qn[state], state + 1)
end

keys(qn::QN) = keys(data(qn))

"""
    val(q::QN,name)

Get the value within the QN q
corresponding to the string `name`
"""
function val(q::QN, name_)
  sname = SmallString(name_)
  for n in 1:maxQNs
    name(q[n]) == sname && return val(q[n])
  end
  return 0
end

"""
    modulus(q::QN,name)

Get the modulus within the QN q
corresponding to the string `name`
"""
function modulus(q::QN, name_)
  sname = SmallString(name_)
  for n in 1:maxQNs
    name(q[n]) == sname && return modulus(q[n])
  end
  return 0
end

"""
    zero(q::QN)

Returns a QN object containing
the same names as q, but with
all values set to zero.
"""
function zero(qn::QN)
  mqn = MQNStorage(undef)
  for i in 1:length(mqn)
    mqn[i] = zero(qn[i])
  end
  return QN(mqn)
end

function (dir::Arrow * qn::QN)
  mqn = MQNStorage(undef)
  for i in 1:length(mqn)
    mqn[i] = dir * qn[i]
  end
  return QN(mqn)
end

(qn::QN * dir::Arrow) = (dir * qn)

function -(qn::QN)
  mqn = MQNStorage(undef)
  for i in 1:length(mqn)
    mqn[i] = -qn[i]
  end
  return QN(mqn)
end

function (a::QN + b::QN)
  !isactive(b[1]) && return a

  ma = MQNStorage(data(a))
  for nb in 1:maxQNs
    !isactive(b[nb]) && break
    bname = name(b[nb])
    for na in 1:maxQNs
      aname = name(a[na])
      if !isactive(ma[na])
        ma[na] = b[nb]
        break
      elseif name(ma[na]) == bname
        ma[na] = ma[na] + b[nb]
        break
      elseif (bname < aname) && (na == 1 || bname > name(ma[na - 1]))
        for j in maxQNs:-1:(na + 1)
          ma[j] = ma[j - 1]
        end
        ma[na] = b[nb]
        break
      end
    end
  end
  return QN(QNStorage(ma))
end

(a::QN - b::QN) = (a + (-b))

function hasname(qn::QN, qv_find::QNVal)
  for qv in qn
    name(qv) == name(qv_find) && return true
  end
  return false
end

# Does not perform checks on if QN is already full, drops
# the last QNVal
# Rename insert?
function NDTensors.insertafter(qn::QN, qv::QNVal, pos::Int)
  return QN(NDTensors.insertafter(Tuple(qn), qv, pos)[1:length(qn)])
end

function addqnval(qn::QN, qv_add::QNVal)
  isactive(qn[end]) &&
    error("Cannot add QNVal, QN already contains maximum number of QNVals")
  for (pos, qv) in enumerate(qn)
    if qv_add < qv || !isactive(qv)
      return NDTensors.insertafter(qn, qv_add, pos - 1)
    end
  end
end

# Fills in the qns of qn1 that qn2 has but
# qn1 doesn't
function fillqns_from(qn1::QN, qn2::QN)
  # If qn1 has no non-trivial qns, fill
  # with qn2
  !isactive(qn1) && return zero(qn2)
  for qv2 in qn2
    if !hasname(qn1, qv2)
      qn1 = addqnval(qn1, zero(qv2))
    end
  end
  return qn1
end

# Make sure qn1 and qn2 have all of the same qns
function fillqns(qn1::QN, qn2::QN)
  qn1_filled = fillqns_from(qn1, qn2)
  qn2_filled = fillqns_from(qn2, qn1)
  return qn1_filled, qn2_filled
end

function isequal_assume_filled(qn1::QN, qn2::QN)
  for (qv1, qv2) in zip(qn1, qn2)
    modulus(qv1) != modulus(qv2) && error("QNVals must have same modulus to compare")
    qv1 != qv2 && return false
  end
  return true
end

function ==(qn1::QN, qn2::QN; assume_filled=false)
  if !assume_filled
    qn1, qn2 = fillqns(qn1, qn2)
  end
  return isequal_assume_filled(qn1, qn2)
end

function isless_assume_filled(qn1::QN, qn2::QN)
  for n in 1:length(qn1)
    val1 = val(qn1[n])
    val2 = val(qn2[n])
    val1 != val2 && return val1 < val2
  end
  return false
end

function isless(qn1::QN, qn2::QN; assume_filled=false)
  return <(qn1, qn2; assume_filled=assume_filled)
end

function <(qn1::QN, qn2::QN; assume_filled=false)
  if !assume_filled
    qn1, qn2 = fillqns(qn1, qn2)
  end
  return isless_assume_filled(qn1, qn2)
end

function have_same_qns(qn1::QN, qn2::QN)
  for n in 1:length(qn1)
    name(qn1[n]) != name(qn2[n]) && return false
  end
  return true
end

function have_same_mods(qn1::QN, qn2::QN)
  for n in 1:length(qn1)
    modulus(qn1[n]) != modulus(qn2[n]) && return false
  end
  return true
end

function show(io::IO, q::QN)
  print(io, "QN(")
  Na = nactive(q)
  for n in 1:Na
    v = q[n]
    n > 1 && print(io, ",")
    Na > 1 && print(io, "(")
    if name(v) != SmallString("")
      print(io, "\"$(name(v))\",")
    end
    print(io, "$(val(v))")
    if modulus(v) != 1
      print(io, ",$(modulus(v))")
    end
    Na > 1 && print(io, ")")
  end
  return print(io, ")")
end

function write(parent::Union{HDF5.File,HDF5.Group}, gname::AbstractString, q::QN)
  g = create_group(parent, gname)
  attributes(g)["type"] = "QN"
  attributes(g)["version"] = 1
  names = [String(name(q[n])) for n in 1:maxQNs]
  vals = [val(q[n]) for n in 1:maxQNs]
  mods = [modulus(q[n]) for n in 1:maxQNs]
  write(g, "names", names)
  write(g, "vals", vals)
  return write(g, "mods", mods)
end

function read(parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ::Type{QN})
  g = open_group(parent, name)
  if read(attributes(g)["type"]) != "QN"
    error("HDF5 group or file does not contain QN data")
  end
  names = read(g, "names")
  vals = read(g, "vals")
  mods = read(g, "mods")
  mqn = ntuple(n -> QNVal(names[n], vals[n], mods[n]), maxQNs)
  return QN(mqn)
end


export Hermite,GaussWeight

#TODO: Add general lines
immutable Hermite{T} <: PolynomialSpace{Line{T}}
    L::T
end
Hermite()=Hermite(1.0)

domain(::Hermite)=Line()
canonicalspace(H::Hermite)=Hermite()
spacescompatible(::Hermite,::Hermite)=true #TODO:L
canonicaldomain(H::Hermite)=Line()
tocanonical(H::Hermite,x)=x

#####
# recα/β/γ are given by
#       x p_{n-1} =γ_n p_{n-2} + α_n p_{n-1} +  p_n β_n
#####

recα(::Type,::Hermite,k)=0;recβ(::Type,::Hermite,k)=0.5;recγ(::Type,::Hermite,k)=k-1
recA(::Type,::Hermite,k)=2;recB(::Type,::Hermite,k)=0;recC(::Type,::Hermite,k)=2k

Derivative(H::Hermite,order)=ConcreteDerivative(H,order)

bandinds{H<:Hermite}(D::ConcreteDerivative{H})=0,D.order
rangespace{H<:Hermite}(D::ConcreteDerivative{H})=domainspace(D)
getindex{H<:Hermite}(D::ConcreteDerivative{H},k::Integer,j::Integer) =
        j==k+D.order?one(eltype(D))*2^m*pochhammer(k,D.order):zero(eltype(D))

function hermitep(r::Range,x::Number)
    n=r[end]+1
    T = typeof(x)
    H = Hermite()
    if n≤2
        v=[1.,2x]
    else
        v=Array(promote_type(Float64,typeof(x)),n)  # x may be complex
        v[1]=1.
        v[2]=2x

        for k=2:n-1
            v[k+1]=((x-recα(T,H,k))*v[k] - recγ(T,H,k)*v[k-1])/recβ(T,H,k)
        end
    end
    v[r+1]
end
hermitep(n::Integer,v::Number)=hermitep(n:n,v)[1]
hermitep(n::Range,v::Vector)=transpose(hcat(map(x->hermitep(n,x),v)...))
hermitep(n::Integer,v::Vector)=map(x->hermitep(n,x),v)

identity_fun(sp::Hermite)=Fun([0.,0.5],sp)

# exp(-Lx^2)
immutable GaussWeight{S,T} <: WeightSpace{S,RealBasis,Line{Float64},1}
    space::S
    L::T
end

GaussWeight(H::Hermite)=GaussWeight(H,H.L)
GaussWeight()=GaussWeight(Hermite())

identity_fun(sp::GaussWeight)=identity_fun(sp.space)

spacescompatible(a::GaussWeight,b::GaussWeight)=spacescompatible(a.space,b.space)&&isapprox(a.L,b.L)

function Derivative(sp::GaussWeight,k::Integer)
   if k==1
        x=Multiplication(Fun(identity,sp.space),sp.space)
        D=Derivative(sp.space)
        D2=D-(2sp.L)x
        DerivativeWrapper(SpaceOperator(D2,sp,GaussWeight(rangespace(D2),sp.L)),1)
    else
        D=Derivative(sp)
        DerivativeWrapper(TimesOperator(Derivative(rangespace(D),k-1).op,D.op),k)
    end
end

weight(H::GaussWeight,x)=exp(-H.L*x.^2)

function Base.sum{H<:Hermite,T}(f::Fun{GaussWeight{H,T}})
    @assert space(f).space.L==space(f).L  # only implemented with matching weight
    f.coefficients[1]*sqrt(π)/sqrt(space(f).L)
end

include("hermitetransform.jl")

function Multiplication{H<:Hermite}(f::Fun{H},S::GaussWeight{H})
    M=Multiplication(f,S.space)
    rs=rangespace(M)
    MultiplicationWrapper(f,SpaceOperator(M,S,GaussWeight(rs,rs.L)))
end

function Multiplication{H<:Hermite,T}(f::Fun{GaussWeight{H,T}},S::Hermite)
    M=Multiplication(Fun(f.coefficients,space(f).space),S)
    rs=rangespace(M)
    MultiplicationWrapper(f,SpaceOperator(M,S,GaussWeight(rs,rs.L)))
end


using ScenTrees
using Test
using Statistics: std
using CSV, DataFrames, XLSX

@testset "ScenTrees.jl" begin
    @testset "Predefined tree - Tree 402" begin
        a = Tree(402)
        @test typeof(a) == Tree
        @test length(a.parent) == 15
        @test length(a.state) == length(a.probability) == length(a.parent) == 15
        @test sum(a.probability) == 8.0
        @test length(a.children) == 8
        @test length(root(a)) == 1
    end
    @testset "Initial Trees" begin
        init = Tree([1,2,2,2],1)
        @test typeof(init) == Tree
        @test length(init.parent) == 15
        @test length(init.state) == length(init.probability) == length(init.parent) == 15
        @test length(init.children) == 8
        @test length(stage(init)) == 15
        @test height(init) == 3
        @test length(leaves(init)) == 3
        @test nodes(init) == 1:15
        @test length(nodes(init)) == 15
        @test length(root(init)) == 1
    end
    @testset "A sample of a Scenario Tree 1D" begin
        x = Tree([1,3,3,3,3],1)
        @test typeof(x) == Tree
        @test length(x.parent) == 121
        @test length(x.state) == length(x.probability) == length(x.parent) == 121
        @test sum(x.probability) ≈ 41.0
        @test length(x.children) == 41
        @test length(root(x)) == 1
        @test length(leaves(x)) == 3
    end
    @testset "A sample of a Scenario Tree 2D" begin
        y = Tree([1,3,3,3,3],2)
        @test typeof(y) == Tree
        @test length(y.parent) == 121
        @test length(y.probability) == length(y.parent) == 121
        @test length(y.state) == length(y.parent)*2 == 242
        @test sum(y.probability) ≈ 41.0
        @test length(y.children) == 41
    end
    @testset "Sample stochastic functions" begin
        a = gaussian_path1D()
        b = running_maximum1D()
        c = path()
        d = gaussian_path2D()
        e = running_maximum2D()
        @test length(a) == 4
        @test length(b) == 4
        @test length(c) == 4
        @test size(d) == (4,2)
        @test size(e) == (4,2)
    end
    @testset "ScenTrees.jl - Tree Approximation 1D" begin
        paths = [gaussian_path1D,running_maximum1D]
        trees = [Tree([1,2,2,2]),Tree([1,3,3,3])]
        samplesize = 100000
        p = 2
        r = 2

        for path in paths
            for newtree in trees
                tree_approximation!(newtree,path,samplesize,p,r)
                @test length(newtree.parent) == length(newtree.state)
                @test length(newtree.parent) == length(newtree.probability)
                @test length(stage(newtree)) == length(newtree.parent)
                @test height(newtree) == maximum(stage(newtree))
                @test round(sum(leaves(newtree)[3]),digits=1) == 1.0   #sum of unconditional probabilities of the leaves
                @test length(root(newtree)) == 1
            end
        end
    end
    @testset "ScenTrees.jl - Tree Approximation 2D" begin
        twoD = tree_approximation!(Tree([1,3,3,3],2),gaussian_path2D,100000,2,2)
        @test size(twoD.state,2) == 2
        @test size(twoD.state,1) == length(twoD.parent) == length(twoD.probability)
    end

    @testset "ScenTrees.jl - Lattice Approximation" begin
        tstLat = lattice_approximation([1,2,3,4],gaussian_path1D,500000,2,1)
        @test length(tstLat.state) == length(tstLat.probability)
        @test round.(sum.(tstLat.probability), digits = 1)  == [1.0, 1.0, 1.0, 1.0] #sum of probs at every stage
    end

    @testset "ScenTrees.jl - Lattice Approximation 2D" begin
        lat2 = lattice_approximation([1,2,3,4],gaussian_path2D,500000,2,2)
        @test length(lat2) == 2 # resultant lattices are 2
        @test length(lat2[1].state) == length(lat2[1].probability)
        @test round.(sum.(lat2[1].probability), digits = 1)  == [1.0, 1.0, 1.0, 1.0]
        @test round.(sum.(lat2[2].probability), digits = 1)  == [1.0, 1.0, 1.0, 1.0]
    end
    @testset "ScenTrees.jl - Test Example Data" begin
        data = CSV.read("data5.csv",DataFrame) # check the new way of loading data
        gsData = Matrix(data)
        df = DataFrame(XLSX.readtable("Mappe1.xlsx", "Sheet1")...)
        df1 = Matrix(df)
        df1 = convert(Array{Float64,2},df1)
        df2 = DataFrame(XLSX.readtable("Mappe1.xlsx", "Sheet2")...)
        df22 = Matrix(df2)
        df22 = convert(Array{Float64,2},df22)
        RandomWalkData = CSV.read("RandomDataWalk.csv",DataFrame)
        RWData = Matrix(RandomWalkData)

        @test size(gsData) ==(100,10)
        @test size(df1) == (1000,7)
        @test size(df22) == (1000,4)
        @test size(RWData) == (1000,5)

        sd1 = zeros(5)
        sd2 = zeros(7)

        for t = 1 : size(RWData,2)
            sd1[t] = std(RWData[:,t])
        end

        for t = 1:size(df1,2)
            sd2[t] = std(df1[:,t])
        end
        @test (sd1 .< 5) == Bool[true, true, true, true, true]
        @test (sd2 .< 10) == Bool[true, true, true, true, true, true, true]

        LatFromKernel = lattice_approximation([1,3,4,5,6],kernel_scenarios(RWData),100000,2,1)
        @test round.(sum.(LatFromKernel.probability),digits=1) == [1.0, 1.0, 1.0, 1.0, 1.0]
        @test length(LatFromKernel.state) == length(LatFromKernel.probability)
    end
end