#############
# Rotations #
#############

export Rx
"""
Rotation matrix, for a rotation about the x-axis.

Arguments:
- `angle::Real`: Counter-clockwise angle of rotation as viewed looking back along the postive direction of the rotation axis.
- `use_degrees:Bool`: If `true` interpret input as being in degrees.

Returns:
- `r::Array{<:Real, 2}`: Rotation matrix

References:
1. O. Montenbruck, and E. Gill, _Satellite Orbits: Models, Methods and Applications_, 2012, p.27.
"""
function Rx(angle::Real ; use_degrees::Bool=false)

    if use_degrees
        angle *= pi/180.0
    end

    c = cos(angle)
    s = sin(angle)

    return [ +1.0  0.0  0.0;
              0.0 +c   +s;
              0.0 -s   +c]
end

export Ry
"""
Rotation matrix, for a rotation about the y-axis.

Arguments:
- `angle::Real`: Counter-clockwise angle of rotation as viewed looking back along the postive direction of the rotation axis.
- `use_degrees:Bool`: If `true` interpret input as being in degrees.

Returns:
- `r::Array{<:Real, 2}`: Rotation matrix

References:
1. O. Montenbruck, and E. Gill, _Satellite Orbits: Models, Methods and Applications_, 2012, p.27.
"""
function Ry(angle::Real ; use_degrees::Bool=false)

    if use_degrees
        angle *= pi/180.0
    end

    c = cos(angle)
    s = sin(angle)

    return [ +c    0.0 -s;
              0.0 +1.0  0.0;
             +s    0.0 +c]
end

export Rz
"""
Rotation matrix, for a rotation about the z-axis.

Arguments:
- `angle::Real`: Counter-clockwise angle of rotation as viewed looking back along the postive direction of the rotation axis.
- `use_degrees:Bool`: If `true` interpret input as being in degrees.

Returns:
- `r::Array{<:Real, 2}`: Rotation matrix

References:
1. O. Montenbruck, and E. Gill, _Satellite Orbits: Models, Methods and Applications_, 2012, p.27.
"""
function Rz(angle::Real ; use_degrees::Bool=false)

    if use_degrees
        angle *= pi/180.0
    end

    c = cos(angle)
    s = sin(angle)

    return [ +c    +s   0.0;
             -s    +c   0.0;
              0.0  0.0 +1.0]
end

###################
# Data Structures #
###################

export Quaternion
"""
The `Quaternion` type defines a _scalar-first_ quaternion for representation
of attitude.

Data members:
- `q0::Float64`: Scalar part of quaternion
- `q1::Float64`: First vector component of quaternion
- `q2::Float64`: Second vector component of quaternion
- `q3::Float64`: Third vector component of quaternion

References:
1. J. Diebel, _Representing attitude: Euler angles, unit quaternions, and rotation vectors._ Matrix 58(15-16) (2006).
"""
mutable struct Quaternion
    q0::Float64
    q1::Float64
    q2::Float64
    q3::Float64

    function Quaternion(q0::Real, q1::Real, q2::Real, q3::Real)
        n = sqrt(q0^2 + q1^2 + q2^2 + q3^2)

        return new(q0/n, q1/n, q2/n, q3/n)
    end
end

export EulerAngle
"""
The `EulerAngle` type provides a represenation of EulerAngles for storing attitude
information.

Valid sequences are: `121, 123, 131, 132, 212, 213, 231, 232, 312, 313, 321, 323`.

Data members:
- `seq::Integer`: Order of application of angles with respect to body axis.
- `phi::Float64`: First Euler angle
- `theta::Float64`: Second Euler angle
- `psi::Float64`: Third Euler angle

References:
1. J. Diebel, _Representing attitude: Euler angles, unit quaternions, and rotation vectors._ Matrix 58(15-16) (2006).
"""
mutable struct EulerAngle
    seq::Int
    phi::Float64
    theta::Float64
    psi::Float64

    function EulerAngle(seq::Integer, phi::Real, theta::Real, psi::Real)
        if !(seq in [121, 123, 131, 132, 212, 213, 231, 232, 312, 313, 321, 323])
            throw(ArgumentError("Invalid EulerAngle sequence: $seq"))
        end

        return new(seq, phi, theta, psi)
    end
end

export EulerAxis
"""
The `EulerAxis` type provides a representation of the Euler angle-and-axis attitude
representation.

Data members:
- `theta::Float64`: Angle of rotation
- `vec::Array{Float64, 1}`: Axis of rotation

References:
1. J. Diebel, _Representing attitude: Euler angles, unit quaternions, and rotation vectors._ Matrix 58(15-16) (2006).
"""
mutable struct EulerAxis
    angle::Float64
    axis::Array{Float64, 1}

    function EulerAxis(angle::Real, axis::Array{<:Real, 1})
        if length(axis) != 3
            throw(ArgumentError("Invalid array for EulerAxis initialization. Input size: $(size(axis)), Required size: (3,)"))
        end

        return new(angle, axis)
    end
end

##############
# Quaternion #
##############

# Quaternion Constructors
function Quaternion(vec::Array{<:Real, 1})
    if length(vec) != 4
        throw(ArgumentError("Invalid array for Quaternion initialization. Input length: $(length(vec)), Required length: 4"))
    end

    Quaternion(vec...)
end

function Quaternion(mat::Array{<:Real, 2})
    if size(mat) == (1,4)
        # Actually vector initialization. so it and return early
        return Quaternion(mat...)
    elseif size(mat) != (3,3)
        throw(ArgumentError("Invalid array for Quaternion initialization. Input size: $(size(mat)), Required size: (3,3)"))
    end

    temp = zeros(Float64, 4)
    temp[1] = 1 + mat[1, 1] + mat[2, 2] + mat[3, 3]
    temp[2] = 1 + mat[1, 1] - mat[2, 2] - mat[3, 3]
    temp[3] = 1 - mat[1, 1] + mat[2, 2] - mat[3, 3]
    temp[4] = 1 - mat[1, 1] - mat[2, 2] + mat[3, 3]

    # Get the maximum value and its index
    den, ind = findmax(temp)
    den      = sqrt(den) # Short-cut equivalence to reuse test information

    # Select optimal inverse mapping
    q0, q1, q2, q3 = 0.0, 0.0, 0.0, 0.0
    if ind == 1
        q0 = 0.5 * den
        q1 = 0.5 * (mat[2, 3] - mat[3, 2]) / den
        q2 = 0.5 * (mat[3, 1] - mat[1, 3]) / den
        q3 = 0.5 * (mat[1, 2] - mat[2, 1]) / den
    elseif ind == 2
        q0 = 0.5 * (mat[2, 3] - mat[3, 2]) / den
        q1 = 0.5 * den
        q2 = 0.5 * (mat[1, 2] + mat[2, 1]) / den
        q3 = 0.5 * (mat[3, 1] + mat[1, 3]) / den
    elseif ind == 3
        q0 = 0.5 * (mat[3, 1] - mat[1, 3]) / den
        q1 = 0.5 * (mat[1, 2] + mat[2, 1]) / den
        q2 = 0.5 * den
        q3 = 0.5 * (mat[2, 3] + mat[3, 2]) / den
    elseif ind == 4
        q0 = 0.5 * (mat[1, 2] - mat[2, 1]) / den
        q1 = 0.5 * (mat[3, 1] + mat[1, 3]) / den
        q2 = 0.5 * (mat[2, 3] + mat[3, 2]) / den
        q3 = 0.5 * den
    end

    return Quaternion(q0, q1, q2, q3)
end

function Quaternion(e::EulerAngle)
    # Extract Quaternion components
    q0, q1, q2, q3 = 0.0, 0.0, 0.0, 0.0

    # Compute sine and cosine values for angles

    # Get Quaternion components depending on Euler Angle sequence
    # reduce number of trig calls
    c1 = cos(e.phi/2.0)
    c2 = cos(e.theta/2.0)
    c3 = cos(e.psi/2.0)
    s1 = sin(e.phi/2.0)
    s2 = sin(e.theta/2.0)
    s3 = sin(e.psi/2.0)

    # Populate the quaternion
    if e.seq == 121
        q0 = c1*c2*c3 - s1*c2*s3
        q1 = c1*c2*s3 + c2*c3*s1
        q2 = c1*c3*s2 + s1*s2*s3
        q3 = c1*s2*s3 - s1*c3*s2

    elseif e.seq == 123
        q0 =  c1*c2*c3 + s1*s2*s3
        q1 = -c1*s2*s3 + c2*c3*s1
        q2 =  c1*c3*s2 + s1*c2*s3
        q3 =  c1*c2*s3 - s1*c3*s2

    elseif e.seq == 131
        q0 =  c1*c2*c3 - s1*c2*s3
        q1 =  c1*c2*s3 + c2*c3*s1
        q2 = -c1*s2*s3 + s1*c3*s2
        q3 =  c1*c3*s2 + s1*s2*s3

    elseif e.seq == 132
        q0 =  c1*c2*c3 - s1*s2*s3
        q1 =  c1*s2*s3 + c2*c3*s1
        q2 =  c1*c2*s3 + s1*c3*s2
        q3 =  c1*c3*s2 - s1*c2*s3

    elseif e.seq == 212
        q0 =  c1*c2*c3 - s1*c2*s3
        q1 =  c1*c3*s2 + s1*s2*s3
        q2 =  c1*c2*s3 + c2*c3*s1
        q3 = -c1*s2*s3 + s1*c3*s2

    elseif e.seq == 213
        q0 =  c1*c2*c3 - s1*s2*s3
        q1 =  c1*c3*s2 - s1*c2*s3
        q2 =  c1*s2*s3 + c2*c3*s1
        q3 =  c1*c2*s3 + s1*c3*s2

    elseif e.seq == 231
        q0 =  c1*c2*c3 + s1*s2*s3
        q1 =  c1*c2*s3 - s1*c3*s2
        q2 = -c1*s2*s3 + c2*c3*s1
        q3 =  c1*c3*s2 + s1*c2*s3

    elseif e.seq == 232
        q0 =  c1*c2*c3 - s1*c2*s3
        q1 =  c1*s2*s3 - s1*c3*s2
        q2 =  c1*c2*s3 + c2*c3*s1
        q3 =  c1*c3*s2 + s1*s2*s3

    elseif e.seq == 312
        q0 =  c1*c2*c3 + s1*s2*s3
        q1 =  c1*c3*s2 + s1*c2*s3
        q2 =  c1*c2*s3 - s1*c3*s2
        q3 = -c1*s2*s3 + c2*c3*s1

    elseif e.seq == 313
        q0 =  c1*c2*c3 - s1*c2*s3
        q1 =  c1*c3*s2 + s1*s2*s3
        q2 =  c1*s2*s3 - s1*c3*s2
        q3 =  c1*c2*s3 + c2*c3*s1

    elseif e.seq == 321
        q0 =  c1*c2*c3 - s1*s2*s3
        q1 =  c1*c2*s3 + s1*c3*s2
        q2 =  c1*c3*s2 - s1*c2*s3
        q3 =  c1*s2*s3 + c2*c3*s1

    elseif e.seq == 323
        q0 =  c1*c2*c3 - s1*c2*s3
        q1 = -c1*s2*s3 + s1*c3*s2
        q2 =  c1*c3*s2 + s1*s2*s3
        q3 =  c1*c2*s3 + c2*c3*s1
    else
        # Should get an invalid sequence, but it is possible if a user
        # Directly sets the sequence number
        throw(ArgumentError("Invalid EulerAngle sequence: $e.seq"))
    end

    return Quaternion(q0, q1, q2, q3)
end

function Quaternion(e::EulerAxis)
    # Extract Quaternion components
    q0 = cos(e.angle/2.0)
    q1 = e.axis[1]*sin(e.angle/2.0)
    q2 = e.axis[2]*sin(e.angle/2.0)
    q3 = e.axis[3]*sin(e.angle/2.0)

    return Quaternion(q0, q1, q2, q3)
end

#########################
# Quaternion Operations #
#########################

function Base.getindex(q::Quaternion, I::UnitRange{<:Integer})
    # Allocate vector once
    vec = as_vector(q)

    # Return selected index or range
    return [vec[i] for i in I]
end

function Base.getindex(q::Quaternion, I::Integer)
    if I == 1
        return q.q0
    elseif I == 2
        return q.q1
    elseif I == 3
        return q.q2
    elseif I == 4
        return q.q3
    else
        throw(BoundsError())
    end
end

Base.getindex(q::Quaternion, ::Colon) = [q.q0, q.q1, q.q2, q.q3]

# Return quaternion as a vector
export as_vector
"""
Return quaternion as a vector.

Equivalent to q[:]

Arguments:
- `q::Quaternion`: Quaternion

Returns:
- `vec::Array{Float64, 1}`: Quaternion as a (4,) vector
"""
function as_vector(q::Quaternion)
    return q[:]
end

# Return quaternion as a matrix
export as_matrix
"""
Return the rotation matrix representation of a Quaternion.

Arguments:
- `q::Quaternion`: Quaternion

Returns:
- `mat::Array{Float64, 2}`: Rotation Matrix on SO(3).
"""
function as_matrix(q::Quaternion)
    # initialize Empty Matrix
    mat = zeros(Float64, 3, 3)

    # Construct matrix from Quaternion
    mat[1, 1] = q.q0*q.q0 + q.q1*q.q1 - q.q2*q.q2 - q.q3*q.q3
    mat[1, 2] = 2*q.q1*q.q2 + 2*q.q0*q.q3
    mat[1, 3] = 2*q.q1*q.q3 - 2*q.q0*q.q2
    mat[2, 1] = 2*q.q1*q.q2 - 2*q.q0*q.q3
    mat[2, 2] = q.q0*q.q0 - q.q1*q.q1 + q.q2*q.q2 - q.q3*q.q3
    mat[2, 3] = 2*q.q2*q.q3 + 2*q.q0*q.q1
    mat[3, 1] = 2*q.q1*q.q3 + 2*q.q0*q.q2
    mat[3, 2] = 2*q.q2*q.q3 - 2*q.q0*q.q1
    mat[3, 3] = q.q0*q.q0 - q.q1*q.q1 - q.q2*q.q2 + q.q3*q.q3

    return mat
end

function Base.copy(q::Quaternion)
    return Quaternion(q.q0, q.q1, q.q2, q.q3)
end

function Base.deepcopy(q::Quaternion)
    return Quaternion(q.q0, q.q1, q.q2, q.q3)
end

"""
Compute the norm of a Quaternion.

Equivalent to `sqrt(q0^2 + q1^2 + q2^2 + q3^2)`

Arguments:
- `q::Quaternion`: Quaternion

Returns:
- `q_norm::Float64`: Norm of quaternion.
"""
function LinearAlgebra.norm(q::Quaternion)
    return sqrt(q.q0^2 + q.q1^2 + q.q2^2 + q.q3^2)
end

"""
Normalize a Quaternion in-place.

Equivalent to q=q/norm(q)

Arguments:
- `q::Quaternion`: Quaternion

Returns:
- `q_norm::Float64`: Norm of quaternion.
"""
function LinearAlgebra.normalize(q::Quaternion)
    # Get Quaternion norm
    q_norm = norm(q)

    # Normalize q in-place
    q.q0 = q.q0/q_norm
    q.q1 = q.q1/q_norm
    q.q2 = q.q2/q_norm
    q.q3 = q.q3/q_norm

    # Ensure return value is nothing
    nothing
end

"""
Get conjugate Quaternion.

Arguments:
- `q::Quaternion`: Input Quaternion

Returns:
- `q_conj::Quaternion`: Conjugate Quaternion of input
"""
function Base.conj(q::Quaternion)
    return Quaternion(q.q0, -q.q1, -q.q2, -q.q3)
end

"""
Get Quaternion inverse.

Arguments:
- `q::Quaternion`: Input Quaternion

Returns:
- `q_inv::Quaternion`: Inverse Quaternion of input
"""
function Base.inv(q::Quaternion)
    # Same as Quaternion conjugate since all quaternions are normalized to have
    # unit norm on construction
    return conj(q)
end

function Base.:-(q::Quaternion)
    return Quaternion(-q.q0, -q.q1, -q.q2, -q.q3)
end

function Base.:-(qa::Quaternion, qb::Quaternion)
    return [qa.q0 - qb.q0
            qa.q1 - qb.q1
            qa.q2 - qb.q2
            qa.q3 - qb.q3]
end

function Base.:+(qa::Quaternion, qb::Quaternion)
    return [qa.q0 + qb.q0
            qa.q1 + qb.q1
            qa.q2 + qb.q2
            qa.q3 + qb.q3]
end

function Base.:+(q::Quaternion, n::Real)
    return [q.q0 + n
            q.q1 + n
            q.q2 + n
            q.q3 + n]
end

function Base.:+(n::Real, q::Quaternion)
    return q+n
end

function Base.:*(qa::Quaternion, qb::Quaternion)
    # # Quaternion Multiplication
    # qcos = self.data[0]*other.data[0] - np.dot(self.data[1:4], other.data[1:4])
    # qvec = self.data[0]*other.data[1:4] + other.data[0]*self.data[1:4] + np.cross(self.data[1:4], other.data[1:4])

    # Quaternion Multiplication
    qcos = qa.q0*qb.q0 - dot(qa[2:4], qb[2:4])
    qvec = qa.q0*qb[2:4] + qb.q0*qa[2:4] + cross(qa[2:4], qb[2:4])

    return Quaternion(qcos, qvec...)
end

function Base.:*(q::Quaternion, n::Real)
    return q[:]*n
end

function Base.:*(n::Real, q::Quaternion)
    return q*n
end

export slerp
"""
Perform spherical linear interpolation (SLERP) on two quaternions. Interpolatles
from quaternion, `q1`, to quaternion, `q2`, at normalized interpolation time, `t`.

Interpolation time must be in the range `[0, 1]` a value of `0` will return `q1`,
while a value of `1` will return `q2`.

Arguments:
- `q1::Quaternion`: Starting Quaternion
- `q2::Quaternion`: Ending Quaternion
- `t::Real`: Normalized interpolation time. [0, 1]

Returns:
- `q:Quaternion`: Quaternion attitude interpolation from q1 toward q2 at time t.
"""
function slerp(q0::Quaternion, q1::Quaternion, t::Real)
    # Check Range on t
    if t < 0.0 || t > 1.0
        throw(ArgumentError("Invalid interpolation time $t. t must be in the range [0, 1]."))
    end

    # Extract vectors and normalize
    q0 = copy(q0)[:]
    q1 = copy(q1)[:]

    # Compute cosine of the angle between the two vectors
    dp = dot(q0, q1)

    # If the dot product is negative, the quaternions have opposite handed-ness
    # and slerp won't take the shortest path. Fix by reversing one quaternion.
    if dp < 0.0
        q1  = -q1
        dp = -dp
    end

    # If the inputs are too close we use linear interpolation instead
    if dp > 0.9995
        return Quaternion(q0 + (q1 - q0)*t)
    end

    theta0 = acos(dp) # Angle between input vectors
    theta  = theta0*t  # Angle between q0 and result quaternion

    s0 = cos(theta) - dp*sin(theta)/sin(theta0)
    s1 = sin(theta) / sin(theta0)

    return Quaternion((s0 * q0) + (s1 * q1))
end

##############
# EulerAngle #
##############

function EulerAngle(seq::Integer, vec::Array{<:Real, 1})
    if length(vec) != 3
        throw(ArgumentError("Invalid array for EulerAngle initialization. Input length: $(length(vec)), Required length: 3"))
    end

    EulerAngle(seq, vec...)
end

function EulerAngle(seq::Integer, mat::Array{<:Real, 2})
    if size(mat) != (3,3)
        throw(ArgumentError("Invalid array for Quaternion initialization. Input size: $(size(mat)), Required size: (3,3)"))
    end

    # Extract elements out of rotation matrix
    r11 = mat[1, 1]
    r12 = mat[1, 2]
    r13 = mat[1, 3]
    r21 = mat[2, 1]
    r22 = mat[2, 2]
    r23 = mat[2, 3]
    r31 = mat[3, 1]
    r32 = mat[3, 2]
    r33 = mat[3, 3]

    # Select euler angle sequence
    phi, theta, psi = 0.0, 0.0, 0.0
    if seq == 121
        phi   = atan(r21, r31)
        theta = acos(r11)
        psi   = atan(r12, -r13)
    elseif seq == 123
        phi   = atan(r23, r33)
        theta = -asin(r13)
        psi   = atan(r12, r11)
    elseif seq == 131
        phi   = atan(r31, -r21)
        theta = acos(r11)
        psi   = atan(r13, r12)
    elseif seq == 132
        phi   = atan(-r32, r22)
        theta = asin(r12)
        psi   = atan(-r13, r11)
    elseif seq == 212
        phi   = atan(r12, -r32)
        theta = acos(r22)
        psi   = atan(r21, r23)
    elseif seq == 213
        phi   = atan(-r13, r33)
        theta = asin(r23)
        psi   = atan(-r21, r22)
    elseif seq == 231
        phi   = atan(r31, r11)
        theta = -asin(r21)
        psi   = atan(r23, r22)
    elseif seq == 232
        phi   = atan(r32, r12)
        theta = acos(r22)
        psi   = atan(r23, -r21)
    elseif seq == 312
        phi   = atan(r12, r22)
        theta = -asin(r32)
        psi   = atan(r31, r33)
    elseif seq == 313
        phi   = atan(r13, r23)
        theta = acos(r33)
        psi   = atan(r31, -r32)
    elseif seq == 321
        phi   = atan(-r21, r11)
        theta = asin(r31)
        psi   = atan(-r32, r33)
    elseif seq == 323
        phi   = atan(r23, -r13)
        theta = acos(r33)
        psi   = atan(r32, r31)
    else
        throw(ArgumentError("Invalid EulerAngle sequence: $seq"))
    end

    return EulerAngle(seq, phi, theta, psi)
end

function EulerAngle(seq::Integer, q::Quaternion)
    # Construct angle from Quaternion by going through a rotation matrix
    return EulerAngle(seq::Integer, as_matrix(q))
end

function EulerAngle(seq::Integer, ea::EulerAxis)
    # Construct angle from EulerAxis by going through a rotation matrix
    return EulerAngle(seq::Integer, as_matrix(ea))
end

# Access Operators
function Base.getindex(e::EulerAngle, I::UnitRange{<:Integer})
    # Allocate vector once
    vec = as_vector(e)

    # Return selected index or range
    return [vec[i] for i in I]
end

function Base.getindex(e::EulerAngle, I::Integer)
    if I == 1
        return e.phi
    elseif I == 2
        return e.theta
    elseif I == 3
        return e.psi
    else
        throw(BoundsError())
    end
end

Base.getindex(e::EulerAngle, ::Colon) = [e.phi, e.theta, e.psi]

"""
Return Euler angles as a vector.

Equivalent to: `[e.phi, e.theta, e.psi]` for `EulerAngle` `e`

Arguments:
- `e::EulerAngle` Euler Angle

Returns:
- `evec::Array{Float64, 1}` Euler angles components in vector form.
"""
function as_vector(e::EulerAngle)
    return [e.phi, e.theta, e.psi]
end

function as_matrix(e::EulerAngle)
    # Get EulerAngle as matrix by going through Quaternions
    return as_matrix(Quaternion(e))
end

function Base.copy(e::EulerAngle)
    return EulerAngle(e.seq, e.phi, e.theta, e.psi)
end

function Base.deepcopy(e::EulerAngle)
    return EulerAngle(e.seq, e.phi, e.theta, e.psi)
end

#############
# EulerAxis #
#############

function EulerAxis(angle::Real, v1::Real, v2::Real, v3::Real)
    return EulerAxis(angle, [v1, v2, v3])
end

function EulerAxis(vec::Array{<:Real, 1})
    if length(vec) != 4
        throw(ArgumentError("Invalid array for EulerAxis initialization. Input size: $(size(vec)), Required size: (4,)"))
    end

    return EulerAxis(vec[1], [vec[2], vec[3], vec[4]])
end

function EulerAxis(mat::Array{<:Real, 2})
    if size(mat) != (3,3)
        throw(ArgumentError("Invalid array for EulerAxis initialization. Input size: $(size(mat)), Required size: (3,3)"))
    end

    return EulerAxis(Quaternion(mat))
end

function EulerAxis(q::Quaternion)
    # Extract quaternion vector and normalize
    qv = as_vector(q)
    q  = qv/norm(qv)

    # Ensure first element is positive
    if q[1] < 0
        q = -q
    end

    # Compute Euler Angle
    angle     = 2*acos(q[1])
    qvec_norm = norm(q[2:4])
    vec       = [0.0, 0.0, 0.0]

    if qvec_norm > 1e-15
        vec = q[2:4]/qvec_norm
    end

    return EulerAxis(angle, vec)
end

function EulerAxis(e::EulerAngle)
    # If input is an EulerAngle first compute a quaternion then get the
    # EulerAxis form
    EulerAxis(Quaternion(e))
end

# Access operators

function Base.getindex(e::EulerAxis, I::UnitRange{<:Integer})
    # Allocate vector once
    vec = as_vector(e)

    # Return selected index or range
    return [vec[i] for i in I]
end

function Base.getindex(e::EulerAxis, I::Integer)
    if I == 1
        return e.angle
    elseif I == 2
        return e.axis[1]
    elseif I == 3
        return e.axis[2]
    elseif I == 4
        return e.axis[3]
    else
        throw(BoundsError())
    end
end

Base.getindex(e::EulerAxis, ::Colon) = [e.angle, e.axis[1], e.axis[2], e.axis[3]]

function as_vector(e::EulerAxis)
    return e[:]
end

function as_matrix(e::EulerAxis)
    # Get matrix form from Quaternion for ease
    return as_matrix(Quaternion(e))
end

function Base.copy(e::EulerAxis)
    return EulerAxis(e.angle, e.axis)
end

function Base.deepcopy(e::EulerAxis)
    return EulerAxis(e.angle, e.axis)
end

####################
# Type Conversions #
####################


"""
    sample_entropy(x, m, r, δ)

Compute the Sample Entropy[^1] of `x`.

Compute Sample Entropy with template size `m` and tolerance factor `r`.
If timestep `δ` is provided, downsample signal `x` with timestep `δ`.

# Examples
```julia
julia> sample_entropy([1, 2, 3, 1, 2, 3], m=2, r=0.2, δ=1)
0.6931471805599453
```
"""
function sample_entropy(x::AbstractArray, m=2, r=0.2, δ=1)
    #=
    x -> normalized signal (1d vector)
    m -> embedding dimension (must be > length of signal)
    r -> tolerance factor
    δ -> delay =#
    N = length(x)
    σ = std(x)
    tolerance = σ * r

    # Create template vectors (of length m + 1)
    matches = zeros(m + 1, N)
    for i in 1:m
        matches[i, 1:(N + 1 - i)] = x[i:end]
    end

    matches[m + 1, 1:(N + 1 - m - δ)] = x[m + δ:end]
    matches = matches[:, 1:N + 1 - m - δ]

    # Calculate pairwise distances for templates of length m
    dist_m = zeros(N - m, N - m)
    pairwise!(dist_m, Chebyshev(), matches[1:m,:], dims=2)
    # Extract upper triangle of distance matrix
    dist_m = dist_m[tril!(trues(size(dist_m)), -1)]
    # Count pairs of template that are within tolerance
    B = count(x -> x <= tolerance, dist_m)

    if B == 0
        return Inf, 0, 0
    end

    # Repeat for templates m+1
    dist_m1 = zeros(N - m, N - m)
    pairwise!(dist_m1, Chebyshev(), matches[1:m + 1,:], dims=2)
    dist_m1 = dist_m1[tril!(trues(size(dist_m1)), -1)]
    A = count(x -> x <= tolerance, dist_m1)

    -log(A / B), A, B

end

function multiscale_entropy(x::AbstractArray, τ=1, m=2, r=0.2, δ=1)

    if τ == 1
        sample_entropy(x, m, r, δ)
    end

    N = length(x)
    x = x[1:end - (N % τ)]
    N = length(x)
    # TODO ordre de dims à vérifier
    dims = (τ, Int(N / τ))
    x = reshape(x, dims)
    x_mse = mean(x, dims=1)

    sample_entropy(x_mse, m, r, δ)

end


using SparseArrays

function poisson_linear_system(ndofs)
        operator = spdiagm(
                -1 => 1.0 * ones(ndofs - 1),
                0 => -2.0 * ones(ndofs),
                1 => 1.0 * ones(ndofs - 1),
        )
        operator[1,1] = 1.0
        operator[1,2] = 0.0
        operator[ndofs,ndofs-1] = 0.0
        operator[ndofs,ndofs] = 1.0
        return operator
end

function solve_direct(operator, rhs, bcleft, bcright, stepsize)
    numpts = length(rhs)
    copyrhs = copy(rhs)
    copyrhs[1] = bcleft / stepsize^2
    copyrhs[numpts] = bcright / stepsize^2
    sol = stepsize^2 * (operator \ copyrhs)
    return sol
end

function solve_direct_all_rhs(operator, bcleft, bcright, stepsize)
    ndofs, ndofs = size(operator)
    rhs = diagm(ones(ndofs))[:, 2:ndofs-1]
    rhs[1, :] .= bcleft / stepsize^2
    rhs[ndofs, :] .= bcright / stepsize^2
    sol = stepsize^2 * (operator \ rhs)
    return sol
end


"""
## beam_gm

This subroutine forms the beam geometric matrix for stability analysis.

### Method
```julia
beam_gm(ell::Float64)
```

### Arguments
```julia
* ell::Float64                   : Element length
```

### Return value
```julia
* gm::::Matrix{Float64}(4,4)     : Geometric matrix for beam element
```
"""
function beam_gm(ell::Float64)
  #
  # This subroutine forms the beam geometric matrix for stability analysis.
  #
  # Updated gm
  #
  gm = zeros(4, 4)
  gm[1,1] = 1.2/ell
  gm[1,2] = 0.1
  gm[2,1] = 0.1
  gm[1,3] = -1.2/ell
  gm[3,1] = -1.2/ell
  gm[1,4] = 0.1
  gm[4,1] = 0.1
  gm[2,2] = 2.0*ell/15.0
  gm[2,3] = -0.1
  gm[3,2] = -0.1
  gm[2,4] = -ell/30.0
  gm[4,2] = -ell/30.0
  gm[3,3] = 1.2/ell
  gm[3,4] = -0.1
  gm[4,3] = -0.1
  gm[4,4] = 2.0*ell/15.0
  gm
end


using splines

println("1D interpolation")
orders = [100]
smin = [0.0]
smax = [1.0]
N_tries = 10

f(x) = sin(x*10)
# 1d test
a = linspace(0,1,orders[1])
b = f(a)
b
di = 1
M = length(a)

coefs = filter_coeffs(smin,smax,orders,b)

#s = linspace(-0.1,1.1,100000)
s = linspace(0.01,0.99,100000)

s = copy(s'')

interp_values = eval_UC_spline(smin,smax,orders,coefs,s)
val, grad = eval_UC_spline_G(smin,smax,orders,coefs,s)

interp_values = eval_UC_spline(smin,smax,orders,coefs,s)
val, grad = eval_UC_spline_G(smin,smax,orders,coefs,s)

fun = interpolant_cspline(smin,smax,orders,b)
vals = fun(s)

println("Error: ", maximum(maximum(abs(interp_values - vals )  )))

tic()
for i = 1:N_tries
    interp_values = eval_UC_spline(smin,smax,orders,coefs,s)
end

true_values = f(s)

println("Maximum interpolation error : ", norm(true_values-interp_values,Inf))
toc()

# Test multi splines:

CC = zeros(2,size(coefs,1))
CC[1,:] = coefs   # first spline
CC[2,:] = coefs/2   # second spline
interp_multi_values = eval_UC_multi_spline(smin,smax,orders,CC,s)

#using PyPlot
#plot(s,sol)
#plot(s,true_values)
#plot(a,b,"o")
#side effect of pyplot: 0.0 not recognized as a constant

# 2d test
println("2D interpolation")
smin = [0.0,0.0]
smax = [1.0,1.0]
orders = [10,20]
a1 = linspace(0,1,orders[1])
a2 = linspace(0,1,orders[2])
b = [sin(i+j) for i=a1, j=a2]
b = convert(Array{Float64},b)
#b = reshape(b, orders[1], orders[2])
#di = [1.0,1.0]
coefs = filter_coeffs(smin,smax,orders,b)

forders = [1000,1000]
N = prod(forders)
mgrid = [ [x y] for x=linspace(0,1,forders[1]), y=linspace(0,1,forders[2])]

vgrid = vcat(mgrid...)
interp_vals = eval_UC_spline(smin,smax,orders,coefs,vgrid)

tic()
for i = 1:10
interp_vals = eval_UC_spline(smin,smax,orders,coefs,vgrid)
end
toc()
true_vals = [sin(x+y) for x=linspace(0,1,forders[1]), y=linspace(0,1,forders[2])]
true_vals = convert(Array{Float64},true_vals[:])

println("Maximum interpolation error : ", norm(true_vals-interp_vals,Inf))

# 3d test
println("3D interpolation")

a = [0.0, 0.0, 0.0]
b = [1.0, 1.0, 1.0]
orders = [20,20,20]

a1 = linspace(0,1,orders[1])
a2 = linspace(0,1,orders[2])
a3 = linspace(0,1,orders[3])

mat = [(i+j+k)^2 for i=a1, j=a2, k=a3]
mat = convert(Array{Float64},mat)

coefs = filter_coeffs(a,b,orders,mat)

forders = [100,100,100]
N = prod(forders)

mgrid = [ [x y z] for x=linspace(0,1,forders[1]), y=linspace(0,1,forders[2]), z=linspace(0,1,forders[3])]

vgrid = vcat(mgrid...)
interp_vals = eval_UC_spline(a,b,orders,coefs,vgrid)

tic()
for i = 1:10
interp_vals = eval_UC_spline(a,b,orders,coefs,vgrid)
end
toc()
true_vals = [(x+y+z)^2 for x=linspace(0,1,forders[1]), y=linspace(0,1,forders[2]),  z=linspace(0,1,forders[3])]
true_vals = convert(Array{Float64},true_vals[:])

println("Maximum interpolation error : ", norm(true_vals-interp_vals,Inf))


using Jecco
using LinearAlgebra
using SparseArrays
using Plots

#= solving equation

 f_xx + f_yy + x y f_xy + x f_x + y f_y + (x^2 + y^2) f
  + (f_x)^2 + (f_y)^2 + x y f_x f_y = Source(x,y)

for the source below, the solution should be given by

 f = exp(-x^2 - y^2)
=#

source(x,y) = exp(-x^2 - y^2) * (-4 + 3 * (x^2 + y^2) + 4 * x^2 * y^2) +
    4 * exp(-2*x^2 - 2*y^2) * (x^2 + y^2 + x^2 * y^2)

function compute_residual!(res::Array, f::Array, x::Vector, y::Vector,
                           Dxx_op::Jecco.FiniteDiffDeriv, Dyy_op::Jecco.FiniteDiffDeriv,
                           Dx_op::Jecco.FiniteDiffDeriv, Dy_op::Jecco.FiniteDiffDeriv)
    Nx = length(x)
    Ny = length(y)

    f_xx = Dxx_op * f
    f_yy = Dyy_op * f
    f_x  = Dx_op  * f
    f_y  = Dy_op  * f
    f_xy = Dx_op  * f_y

    for j in 1:Ny, i in 1:Nx
        x1 = x[i]
        y1 = y[j]

        res[i,j] = f_xx[i,j] + f_yy[i,j] + x1 * y1 * f_xy[i,j] +
            x1 * f_x[i,j] + y1 * f_y[i,j] +
            (x1*x1 + y1*y1) * f[i,j] + f_x[i,j]^2 + f_y[i,j]^2 +
            x1 * y1 * f_x[i,j] * f_y[i,j] - source(x1,y1)
    end

    res
end

function deriv_operators(hx, hy, Nx::Int, Ny::Int, ord::Int)
    Dx_op  = CenteredDiff{1}(1, ord, hx, Nx)
    Dxx_op = CenteredDiff{1}(2, ord, hx, Nx)

    Dy_op  = CenteredDiff{2}(1, ord, hy, Ny)
    Dyy_op = CenteredDiff{2}(2, ord, hy, Ny)

    Dxx_op, Dyy_op, Dx_op, Dy_op
end

#=
use the Kronecker product (kron) to build the 2-dimensional derivation matrices
from the 1-dimensional ones. see for instance:

  https://en.wikipedia.org/wiki/Kronecker_product

  https://arxiv.org/pdf/1801.01483.pdf (section 5)
=#
function deriv_matrices(Dxx_op::Jecco.FiniteDiffDeriv, Dyy_op::Jecco.FiniteDiffDeriv,
                        Dx_op::Jecco.FiniteDiffDeriv, Dy_op::Jecco.FiniteDiffDeriv)
    Dx  = kron(I(Ny), SparseMatrixCSC(Dx_op))
    Dxx = kron(I(Ny), SparseMatrixCSC(Dxx_op))
    Dy  = kron(SparseMatrixCSC(Dy_op), I(Nx))
    Dyy = kron(SparseMatrixCSC(Dyy_op), I(Nx))

    Dxx, Dyy, Dx * Dy, Dx, Dy
end

# returns axx Dxx + ayy Dyy + axy Dxy + bx Dx + by Dy + cc. note that this
# function overwrites the input matrices to save memory
function build_jacobian(Dxx::SparseMatrixCSC, Dyy::SparseMatrixCSC, Dxy::SparseMatrixCSC,
                        Dx::SparseMatrixCSC, Dy::SparseMatrixCSC,
                        axx::Diagonal, ayy::Diagonal, axy::Diagonal,
                        bx::Diagonal, by::Diagonal, cc::Diagonal)
    Dxx_ = similar(Dxx)
    Dyy_ = similar(Dyy)
    Dxy_ = similar(Dxy)
    Dx_  = similar(Dx)
    Dy_  = similar(Dy)

    mul!(Dxx_, axx, Dxx)
    mul!(Dyy_, ayy, Dyy)
    mul!(Dxy_, axy, Dxy)
    mul!(Dx_, bx, Dx)
    mul!(Dy_, by, Dy)

    Dxx_ + Dyy_ + Dxy_ + Dx_ + Dy_ + cc
end

x_min    = -5.0
x_max    =  5.0
x_nodes  =  128
# x_nodes  =  256
y_min    = -5.0
y_max    =  5.0
y_nodes  =  64
# y_nodes  =  256

ord = 4
# ord = 2

xcoord  = Cartesian{1}("x", x_min, x_max, x_nodes, endpoint=false)
ycoord  = Cartesian{2}("y", y_min, y_max, y_nodes, endpoint=false)

hx = Jecco.delta(xcoord)
hy = Jecco.delta(ycoord)

Nx = xcoord.nodes
Ny = ycoord.nodes

f_exact = [exp(-xcoord[i]^2 - ycoord[j]^2) for i in 1:Nx, j in 1:Ny]

Dxx_op, Dyy_op, Dx_op, Dy_op = deriv_operators(hx, hy, Nx, Ny, ord)
Dxx, Dyy, Dxy, Dx, Dy = deriv_matrices(Dxx_op, Dyy_op, Dx_op, Dy_op)

# initial guess
# fsol   = zeros(Nx,Ny)
fsol   = copy(f_exact)

res    = zeros(Nx,Ny)
ind2D  = LinearIndices(fsol)

M = Nx * Ny

b_vec   = zeros(M)
axx     = ones(M)
ayy     = ones(M)
axy     = zeros(M)
bx      = zeros(M)
by      = zeros(M)
cc      = zeros(M)
f0      = zeros(M)

itmax   = 8
epsilon = 1e-12

for it in 1:itmax

    compute_residual!(res, fsol, xcoord[:], ycoord[:], Dxx_op, Dyy_op, Dx_op, Dy_op)
    max_res = maximum(abs.(res))

    @show it, max_res
    if max_res < epsilon
        break
    end

    # build (linearized) operator
    # A = Dxx + Dyy + x y Dxy + (x + 2 f0_x + x y f0_y) Dx + (y + 2 f0_y + x y f0_x) Dy + (x^2 + y^2)

    @inbounds for idx in eachindex(f0)
        f0[idx] = fsol[idx]
    end
    f0_x = Dx * f0
    f0_y = Dy * f0

    for j in 1:Ny, i in 1:Nx
        idx = ind2D[i,j]

        x1  = xcoord[i]
        y1  = ycoord[j]

        axy[idx] = x1 * y1
        bx[idx]  = x1 + 2 * f0_x[idx] + x1 * y1 * f0_y[idx]
        by[idx]  = y1 + 2 * f0_y[idx] + x1 * y1 * f0_x[idx]
        cc[idx]  = x1^2 + y1^2

        b_vec[idx] = -res[idx]
    end

    A_mat = build_jacobian(Dxx, Dyy, Dxy, Dx, Dy,
                           Diagonal(axx), Diagonal(ayy), Diagonal(axy),
                           Diagonal(bx), Diagonal(by), Diagonal(cc))

    A_fact = factorize(A_mat)
    ldiv!(f0, A_fact, b_vec)

    # update solution
    @inbounds for idx in eachindex(fsol)
        fsol[idx] += f0[idx]
    end

end

j_slice = div(Ny,2) + 1
x = xcoord[:]

plot(x, f_exact[:,j_slice])
scatter!(x, fsol[:,j_slice])


## activate package environment
cd(@__DIR__)
using Pkg
Pkg.activate("Project.toml")

## percentage x of Al, wavelengths λ and index of refraction n from publication Aspnes 1986
x = [0.491 0.590]
λ = [0.520]
n = [3.8264 + 0.19880im;
     3.7477 + 0.15232im]
N = real(n)
k = imag(n)

## interpolation of n for wanted x
x_wanted = [0.509]
l = length(λ)
m = length(x_wanted)
N_wanted = zeros(m,l)
k_wanted = zeros(m,l)
for i = 1:l
    for j = 0:m-1
        local coeffs_N = [x[j+1,:] ones(2)]\N[(1+j*2):(1+j*2)+1,i]
        local coeffs_k = [x[j+1,:] ones(2)]\k[(1+j*2):(1+j*2)+1,i]
        N_wanted[i,j+1] = coeffs_N[1]*x_wanted[j+1] + coeffs_N[2]
        k_wanted[i,j+1] = coeffs_k[1]*x_wanted[j+1] + coeffs_k[2]
    end
end
N_wanted = round.(N_wanted,sigdigits=4)
k_wanted = round.(k_wanted,sigdigits=4)

## output solution
println("")
for j = 1:m
    for i = 1:l
        println("n(x = $(x_wanted[j]), λ = $(lambda[i])) = $(N_wanted[i,j]) + $(k_wanted[i,j])i")
    end
end


function square_distance(coords1, coords2)
    """Compute the square distances between coords1 and coords2.

    Note: Depending on the input value and the dtype, negative values may
    be mixed in the return value.

    Args:
        coords1: coordinates (N, C)
        coords2: coordinates (M, C)

    Returns:
        square distances:
            square distances between coords1 and coords2 (N, M)
    """
    dot_product = -2 * coords1 * coords2'
    column = repeat(sum(coords1 .* coords1, dims=2), 1, size(coords2, 1))
    row = repeat(sum(coords2 .* coords2, dims=2), 1, size(coords1, 1))'
    square_dist = column + dot_product + row
    return square_dist
end

if abspath(PROGRAM_FILE) == @__FILE__
    arr1 = reshape(range(0, 23, step=1), 8, 3)
    arr2 = reshape(range(24, 47, step=1), 8, 3)
    @time square_distance(arr1, arr2)
    @time square_distance(arr1, arr2)
end


# Fixed Talbot method
#
# Abate, J. and Valkó, P.P.
# Multi-precision Laplace transform inversion
# International Journal for Numerical Methods in Engineering, Vol. 60 (Iss. 5-7)  2004  pp 979–993

"""
    talbot(func::Function, t::AbstractFloat, M::Integer=talbot_default_num_terms)

Evaluate the inverse Laplace transform of `func` at the point `t`. Use `M` terms in the algorithm.
For `typeof(t)` is `Float64`, the default for `M` is `32`. For `BigFloat` the default is `64`.

If `BigFloat` precision is larger than default, try increasing `M`.

# Example

```jldoctest
julia> InverseLaplace.talbot(s -> 1 / s^3, 3)
4.50000000000153
```

!!! note
    This function uses the fixed Talbot method. It evaluates `func` for complex arguments.
"""
function talbot(func, t, M)
    bM = convert(typeof(t), M)
    r = (2 * bM) / (5 * t)
    term = (1//2) * exp(r * t) * func(r)
    for i in 1:M-1
        theta = i * (pi/bM)
        s = r * theta * (complex(cot(theta), one(theta)))
        sigma = theta + (theta * cot(theta) - 1) * cot(theta)
        term += real(exp(t * s) * complex(one(t), sigma) * func(s))
    end
    return term * 2 / (5 * t)
end

talbot(func, t) = talbot(func, t, talbot_default_num_terms)
const talbot_BigFloat_default_num_terms = 64
talbot(func, t::BigFloat) = talbot(func, t, talbot_BigFloat_default_num_terms)
talbot(func, t::Integer, args...) = talbot(func, BigFloat(t), args...)
# Hmm, at some point, one of these routines actually gave a Rational result. Don't recall how.
# But, it can't be talbot.
talbot(func, t::Rational, args...) = talbot(func, BigFloat(t), args...)

# Operate on an array of values of t. A single function evaluation
# f(s) is used for all t
# This gives more inaccurate results the further values of
# t are from tmax
"""
    talbotarr(func, ta::AbstractArray, M)

Compute the inverse Laplace transform for each element in `ta`. Each evaluation
of `func(s)` is used for all elements of `ta`. This may be faster than a
broadcast application of `talbot` (i.e. `talbot.(...`) , but is in general, less
accurate.  `talbotarr` uses the "fixed" Talbot method.
"""
function talbotarr(func, t::AbstractArray, M)
    tt = typeof(t[1])
    bM = convert(tt, M)
    terms = similar(t)
    tmax = maximum(t)
    r = (2 * bM) / (5 * tmax)
    fr = (1//2) * func(r)
    terms .= exp.(r .* t) .* fr
    for i in 1:M-1
        theta = i * (pi/bM)
        s = r * theta * (complex(cot(theta), one(theta)))
        sigma = theta + (theta * cot(theta) - 1) * cot(theta)
        fs = complex(one(tt), sigma) * func(s)
        for j in 1:length(terms)
            terms[j] += real(exp(t[j] * s) * fs)
        end
    end
    terms .= terms .* 2 / (5 * tmax)
    return terms
end

talbotarr(func, t::AbstractArray) = talbotarr(func, t, talbot_default_num_terms)
talbotarr(func, t::Vector{BigFloat}) = talbotarr(func, t, talbot_BigFloat_default_num_terms)

#  LocalWords:  talbotarr func talbot


# # Unconstrained optimization

#=
In this example, we show how to differentiate through the solution of the following unconstrained optimization problem:
```math
\hat{y}(x) = \min_{y \in \mathbb{R}^m} f(x, y)
```
The optimality conditions are given by gradient stationarity:
```math
F(x, \hat{y}(x)) = 0 \quad \text{with} \quad F(x,y) = \nabla_2 f(x, y) = 0
```

=#

using ImplicitDifferentiation
using Optim: optimize, minimizer, LBFGS
using Zygote

using ChainRulesTestUtils  #src
using Test  #src

# ## Implicit function wrapper

#=
To make verification easy, we minimize a quadratic objective
```math
f(x, y) = \lVert y - x \rVert^2
```
In this case, the optimization algorithm is very simple, but still we can implement it as a black box to show that it doesn't change the result.
=#

function forward(x)
    f(y) = sum(abs2, y-x)
    y0 = zero(x)
    res = optimize(f, y0, LBFGS(); autodiff=:forward)
    y = minimizer(res)
    return y
end;

#=
On the other hand, optimality conditions should be provided explicitly whenever possible, so as to avoid nesting automatic differentiation calls.
=#

conditions(x, y) = 2(y - x);

# We now have all the ingredients to construct our implicit function.

implicit = ImplicitFunction(forward, conditions);

# ## Testing

x = rand(5)

# Let's start by taking a look at the forward pass, which should be the identity function.

implicit(x)

# We now check whether the behavior of our `ImplicitFunction` wrapper is coherent with the theoretical derivatives.

Zygote.jacobian(implicit, x)[1]

# As expected, we recover the identity matrix as Jacobian.

# The following tests are not included in the docs.  #src

@testset verbose = true "ChainRules" begin  #src
    test_frule(implicit, x; check_inferred=false)  #src
    test_rrule(implicit, x; check_inferred=false)  #src
end  #src


"""
    expv(t, A, B; kwargs...)

Compute `exp(t*A) * B` without computing `t * A` or the matrix exponential.
This is significantly faster than the matrix exponential when the second dimension of `B` is
much smaller than the first one. The "time" `t` may be real or complex.
The algorithm is described in [^AlMohyHigham2011].

# Keywords

  - `shift=true`: Expand the Taylor series of `exp(t*A)` about ``A-μI=0`` instead of
    ``A=0``, where ``μ = \\operatorname{tr}(A) / n`` to speed up convergence. See
    §3.1 of [^AlMohyHigham2011].
  - `tol`: The tolerance at which to compute the result. Defaults to the tolerance of the
    eltype of the result.

[^AlMohyHigham2011]: Al-Mohy, Awad H. and Higham, Nicholas J. (2011) Computing the Action of the Matrix
    Exponential, with an Application to Exponential Integrators. SIAM Journal on Scientific
    Computing, 33 (2). pp. 488-511. ISSN 1064-8275
    doi: [10.1137/100788860](https://doi.org/10.1137/100788860)
    eprint: [eprints.maths.manchester.ac.uk/id/eprint/1591](http://eprints.maths.manchester.ac.uk/id/eprint/1591)
"""
function expv(t, A, B; shift=true, tol=eps(float(real(Base.promote_eltype(t, A, B)))))
    n = LinearAlgebra.checksquare(A)
    # §3: “Our experience indicates that p_max = 8 and m_max = 55 are appropriate choices.”
    p_max = 8
    m_max = 55
    n0 = size(B, 2)
    if shift
        μ = tr(A) / n
        A -= μ * I
    else
        μ = zero(float(eltype(A)))
    end
    params = parameters(t, A, n0, m_max, p_max, tol)
    η = exp(t * μ / params.s)
    F = one(η) * B
    for i in 1:(params.s)
        c1 = _opnormInf(B)
        for j in 1:(params.m)
            B = (A * B) * (t / (params.s * j))
            c2 = _opnormInf(B)
            F += B
            c1 + c2 ≤ tol * _opnormInf(F) && break
            c1 = c2
        end
        F *= η
        B = F
    end
    return F
end
expv(t, A::Diagonal, B; kwargs...) = exp.(t .* A.diag) .* B


"""
`Energy` - Compute the energy

E = -0.5 * J ∑ xᵢxⱼ

over a periodic N×N square lattice.  Note that this counts each bond twice,
hence the factor of 0.5

### Fields
* `x` - State of the lattice
### Optional Fields
* 'J=1' - Coupling constant>0
"""
function Energy(x::Tx; J::TF=1.0) where {TI<:Integer, Tx<:AbstractArray{TI}, TF<:AbstractFloat}
    E = 0.0;

    L= size(x)[1];
    for i in 1:L, j in 1:L
        # sum of energy contributions
        k, l = mod1(i+1,L), j;
        E += - 0.5 * J * x[i,j] * x[k,l];

        k,l = mod1(i-1,L), j;
        E += - 0.5 * J * x[i,j] * x[k,l];

        k, l = i, mod1(j+1,L);;
        E += - 0.5 * J * x[i,j] * x[k,l];

        k, l = i, mod1(j-1,L);
        E += - 0.5 * J * x[i,j] * x[k,l];
    end
    return E
end

"""
`Magnetization` - Compute the magnetization

M = ∑ xᵢ

over a periodic N×N square lattice.

### Fields
* `x` - State of the lattice
"""
function Magnetization(x::Tx) where {TI<:Integer, Tx<:AbstractArray{TI}}
    return sum(x)
end

"""
`ClusterNeighbors!` - Mark on the cluster map, `c`, where the current cluster
continues.

### Fields
* `c` - Cluster map
* `x` - State of the lattice
* `i,j` - Indices of the current site
* `D` - Direction of the current cluster (+/-1)
* `id` - Cluster id
* `L` - Lattice size, L × L
"""
function ClusterNeighbors!(c, x, i, j, D, id, L)

    k,l = mod1(i+1,L), j;
    if (x[k,l] == D && c[k,l]==0)
        c[k,l] = id;
        ClusterNeighbors!(c, x, k, l, D, id, L)
    end

    k,l = mod1(i-1,L), j;
    if (x[k,l]==D && c[k,l]==0)
        c[k,l] = id;
        ClusterNeighbors!(c, x, k, l, D, id, L)
    end

    k,l = i, mod1(j+1,L);
    if (x[k,l]==D && c[k,l]==0)
        c[k,l] = id;
        ClusterNeighbors!(c, x, k, l, D, id, L)
    end

    k,l = i, mod1(j-1,L);
    if (x[k,l]==D && c[k,l]==0)
        c[k,l] = id;
        ClusterNeighbors!(c, x, k, l, D, id, L)
    end
    c
end

"""
`Cluster` - Find all the clusters in the lattice.  Note that the numbering is
arbitrary.

### Fields
* `x` - State of the lattice
"""
function Clusters(x::Tx) where {TI<:Integer, Tx<:AbstractArray{TI}}
    L= size(x)[1];
    id = 0;
    c = zeros(Int, L, L);

    for i in 1:L, j in 1:L
        if(c[i,j]==0)
            id+=1;
            c[i,j] = id;
            D = x[i,j];
            ClusterNeighbors!(c, x, i, j, D, id, L)
        end
    end

    return c
end


# This file is a part of Julia. License is MIT: https://julialang.org/license

## number-theoretic functions ##

"""
    gcd(x,y)

Greatest common (positive) divisor (or zero if `x` and `y` are both zero).

# Examples
```jldoctest
julia> gcd(6,9)
3

julia> gcd(6,-9)
3
```
"""
function gcd(a::T, b::T) where T<:Integer
    while b != 0
        t = b
        b = rem(a, b)
        a = t
    end
    checked_abs(a)
end

# binary GCD (aka Stein's) algorithm
# about 1.7x (2.1x) faster for random Int64s (Int128s)
function gcd(a::T, b::T) where T<:Union{Int8,UInt8,Int16,UInt16,Int32,UInt32,Int64,UInt64,Int128,UInt128}
    a == 0 && return abs(b)
    b == 0 && return abs(a)
    za = trailing_zeros(a)
    zb = trailing_zeros(b)
    k = min(za, zb)
    u = unsigned(abs(a >> za))
    v = unsigned(abs(b >> zb))
    while u != v
        if u > v
            u, v = v, u
        end
        v -= u
        v >>= trailing_zeros(v)
    end
    r = u << k
    # T(r) would throw InexactError; we want OverflowError instead
    r > typemax(T) && __throw_gcd_overflow(a, b)
    r % T
end
@noinline __throw_gcd_overflow(a, b) = throw(OverflowError("gcd($a, $b) overflows"))

"""
    lcm(x,y)

Least common (non-negative) multiple.

# Examples
```jldoctest
julia> lcm(2,3)
6

julia> lcm(-2,3)
6
```
"""
function lcm(a::T, b::T) where T<:Integer
    # explicit a==0 test is to handle case of lcm(0,0) correctly
    if a == 0
        return a
    else
        return checked_abs(a * div(b, gcd(b,a)))
    end
end

gcd(a::Integer) = a
lcm(a::Integer) = a
gcd(a::Integer, b::Integer) = gcd(promote(a,b)...)
lcm(a::Integer, b::Integer) = lcm(promote(a,b)...)
gcd(a::Integer, b::Integer...) = gcd(a, gcd(b...))
lcm(a::Integer, b::Integer...) = lcm(a, lcm(b...))

lcm(abc::AbstractArray{<:Integer}) = reduce(lcm, abc; init=one(eltype(abc)))

function gcd(abc::AbstractArray{<:Integer})
    a = zero(eltype(abc))
    for b in abc
        a = gcd(a,b)
        if a == 1
            return a
        end
    end
    return a
end

# return (gcd(a,b),x,y) such that ax+by == gcd(a,b)
"""
    gcdx(x,y)

Computes the greatest common (positive) divisor of `x` and `y` and their Bézout
coefficients, i.e. the integer coefficients `u` and `v` that satisfy
``ux+vy = d = gcd(x,y)``. ``gcdx(x,y)`` returns ``(d,u,v)``.

# Examples
```jldoctest
julia> gcdx(12, 42)
(6, -3, 1)

julia> gcdx(240, 46)
(2, -9, 47)
```

!!! note
    Bézout coefficients are *not* uniquely defined. `gcdx` returns the minimal
    Bézout coefficients that are computed by the extended Euclidean algorithm.
    (Ref: D. Knuth, TAoCP, 2/e, p. 325, Algorithm X.)
    For signed integers, these coefficients `u` and `v` are minimal in
    the sense that ``|u| < |y/d|`` and ``|v| < |x/d|``. Furthermore,
    the signs of `u` and `v` are chosen so that `d` is positive.
    For unsigned integers, the coefficients `u` and `v` might be near
    their `typemax`, and the identity then holds only via the unsigned
    integers' modulo arithmetic.
"""
function gcdx(a::T, b::T) where T<:Integer
    # a0, b0 = a, b
    s0, s1 = oneunit(T), zero(T)
    t0, t1 = s1, s0
    # The loop invariant is: s0*a0 + t0*b0 == a
    while b != 0
        q = div(a, b)
        a, b = b, rem(a, b)
        s0, s1 = s1, s0 - q*s1
        t0, t1 = t1, t0 - q*t1
    end
    a < 0 ? (-a, -s0, -t0) : (a, s0, t0)
end
gcdx(a::Integer, b::Integer) = gcdx(promote(a,b)...)

# multiplicative inverse of n mod m, error if none

"""
    invmod(x,m)

Take the inverse of `x` modulo `m`: `y` such that ``x y = 1 \\pmod m``,
with ``div(x,y) = 0``. This is undefined for ``m = 0``, or if
``gcd(x,m) \\neq 1``.

# Examples
```jldoctest
julia> invmod(2,5)
3

julia> invmod(2,3)
2

julia> invmod(5,6)
5
```
"""
function invmod(n::T, m::T) where T<:Integer
    g, x, y = gcdx(n, m)
    g != 1 && throw(DomainError((n, m), "Greatest common divisor is $g."))
    m == 0 && throw(DomainError(m, "`m` must not be 0."))
    # Note that m might be negative here.
    # For unsigned T, x might be close to typemax; add m to force a wrap-around.
    r = mod(x + m, m)
    # The postcondition is: mod(r * n, m) == mod(T(1), m) && div(r, m) == 0
    r
end
invmod(n::Integer, m::Integer) = invmod(promote(n,m)...)

# ^ for any x supporting *
to_power_type(x) = convert(Base._return_type(*, Tuple{typeof(x), typeof(x)}), x)
@noinline throw_domerr_powbysq(::Any, p) = throw(DomainError(p,
    string("Cannot raise an integer x to a negative power ", p, '.',
           "\nConvert input to float.")))
@noinline throw_domerr_powbysq(::Integer, p) = throw(DomainError(p,
   string("Cannot raise an integer x to a negative power ", p, '.',
          "\nMake x a float by adding a zero decimal (e.g., 2.0^$p instead ",
          "of 2^$p), or write 1/x^$(-p), float(x)^$p, or (x//1)^$p")))
@noinline throw_domerr_powbysq(::AbstractMatrix, p) = throw(DomainError(p,
   string("Cannot raise an integer matrix x to a negative power ", p, '.',
          "\nMake x a float matrix by adding a zero decimal ",
          "(e.g., [2.0 1.0;1.0 0.0]^$p instead ",
          "of [2 1;1 0]^$p), or write float(x)^$p or Rational.(x)^$p")))
function power_by_squaring(x_, p::Integer)
    x = to_power_type(x_)
    if p == 1
        return copy(x)
    elseif p == 0
        return one(x)
    elseif p == 2
        return x*x
    elseif p < 0
        isone(x) && return copy(x)
        isone(-x) && return iseven(p) ? one(x) : copy(x)
        throw_domerr_powbysq(x, p)
    end
    t = trailing_zeros(p) + 1
    p >>= t
    while (t -= 1) > 0
        x *= x
    end
    y = x
    while p > 0
        t = trailing_zeros(p) + 1
        p >>= t
        while (t -= 1) >= 0
            x *= x
        end
        y *= x
    end
    return y
end
power_by_squaring(x::Bool, p::Unsigned) = ((p==0) | x)
function power_by_squaring(x::Bool, p::Integer)
    p < 0 && !x && throw_domerr_powbysq(x, p)
    return (p==0) | x
end

^(x::T, p::T) where {T<:Integer} = power_by_squaring(x,p)
^(x::Number, p::Integer)  = power_by_squaring(x,p)

# x^p for any literal integer p is lowered to Base.literal_pow(^, x, Val(p))
# to enable compile-time optimizations specialized to p.
# However, we still need a fallback that calls the function ^ which may either
# mean Base.^ or something else, depending on context.
# We mark these @inline since if the target is marked @inline,
# we want to make sure that gets propagated,
# even if it is over the inlining threshold.
@inline literal_pow(f, x, ::Val{p}) where {p} = f(x,p)

# Restrict inlining to hardware-supported arithmetic types, which
# are fast enough to benefit from inlining.
const HWReal = Union{Int8,Int16,Int32,Int64,UInt8,UInt16,UInt32,UInt64,Float32,Float64}
const HWNumber = Union{HWReal, Complex{<:HWReal}, Rational{<:HWReal}}

# Core.Compiler has complicated logic to inline x^2 and x^3 for
# numeric types.  In terms of Val we can do it much more simply.
# (The first argument prevents unexpected behavior if a function ^
# is defined that is not equal to Base.^)
@inline literal_pow(::typeof(^), x::HWNumber, ::Val{0}) = one(x)
@inline literal_pow(::typeof(^), x::HWNumber, ::Val{1}) = x
@inline literal_pow(::typeof(^), x::HWNumber, ::Val{2}) = x*x
@inline literal_pow(::typeof(^), x::HWNumber, ::Val{3}) = x*x*x

# don't use the inv(x) transformation here since float^p is slightly more accurate
@inline literal_pow(::typeof(^), x::AbstractFloat, ::Val{p}) where {p} = x^p
@inline literal_pow(::typeof(^), x::AbstractFloat, ::Val{-1}) = inv(x)

# for other types, define x^-n as inv(x)^n so that negative literal powers can
# be computed in a type-stable way even for e.g. integers.
@inline @generated function literal_pow(f::typeof(^), x, ::Val{p}) where {p}
    if p < 0
        :(literal_pow(^, inv(x), $(Val{-p}())))
    else
        :(f(x,$p))
    end
end

# note: it is tempting to add optimized literal_pow(::typeof(^), x, ::Val{n})
#       methods here for various n, but this easily leads to method ambiguities
#       if anyone has defined literal_pow(::typeof(^), x::T, ::Val).

# b^p mod m

"""
    powermod(x::Integer, p::Integer, m)

Compute ``x^p \\pmod m``.

# Examples
```jldoctest
julia> powermod(2, 6, 5)
4

julia> mod(2^6, 5)
4

julia> powermod(5, 2, 20)
5

julia> powermod(5, 2, 19)
6

julia> powermod(5, 3, 19)
11
```
"""
function powermod(x::Integer, p::Integer, m::T) where T<:Integer
    p < 0 && return powermod(invmod(x, m), -p, m)
    p == 0 && return mod(one(m),m)
    (m == 1 || m == -1) && return zero(m)
    b = oftype(m,mod(x,m))  # this also checks for divide by zero

    t = prevpow(2, p)
    r::T = 1
    while true
        if p >= t
            r = mod(widemul(r,b),m)
            p -= t
        end
        t >>>= 1
        t <= 0 && break
        r = mod(widemul(r,r),m)
    end
    return r
end

# optimization: promote the modulus m to BigInt only once (cf. widemul in generic powermod above)
powermod(x::Integer, p::Integer, m::Union{Int128,UInt128}) = oftype(m, powermod(x, p, big(m)))

_nextpow2(x::Unsigned) = oneunit(x)<<((sizeof(x)<<3)-leading_zeros(x-oneunit(x)))
_nextpow2(x::Integer) = reinterpret(typeof(x),x < 0 ? -_nextpow2(unsigned(-x)) : _nextpow2(unsigned(x)))
_prevpow2(x::Unsigned) = one(x) << unsigned((sizeof(x)<<3)-leading_zeros(x)-1)
_prevpow2(x::Integer) = reinterpret(typeof(x),x < 0 ? -_prevpow2(unsigned(-x)) : _prevpow2(unsigned(x)))

"""
    ispow2(n::Integer) -> Bool

Test whether `n` is a power of two.

# Examples
```jldoctest
julia> ispow2(4)
true

julia> ispow2(5)
false
```
"""
ispow2(x::Integer) = x > 0 && count_ones(x) == 1

"""
    nextpow(a, x)

The smallest `a^n` not less than `x`, where `n` is a non-negative integer. `a` must be
greater than 1, and `x` must be greater than 0.

# Examples
```jldoctest
julia> nextpow(2, 7)
8

julia> nextpow(2, 9)
16

julia> nextpow(5, 20)
25

julia> nextpow(4, 16)
16
```

See also [`prevpow`](@ref).
"""
function nextpow(a::Real, x::Real)
    x <= 0 && throw(DomainError(x, "`x` must be positive."))
    # Special case fast path for x::Integer, a == 2.
    # This is a very common case. Constant prop will make sure that a call site
    # specified as `nextpow(2, x)` will get this special case inlined.
    a == 2 && isa(x, Integer) && return _nextpow2(x)
    a <= 1 && throw(DomainError(a, "`a` must be greater than 1."))
    x <= 1 && return one(a)
    n = ceil(Integer,log(a, x))
    p = a^(n-1)
    # guard against roundoff error, e.g., with a=5 and x=125
    p >= x ? p : a^n
end

"""
    prevpow(a, x)

The largest `a^n` not greater than `x`, where `n` is a non-negative integer.
`a` must be greater than 1, and `x` must not be less than 1.

# Examples
```jldoctest
julia> prevpow(2, 7)
4

julia> prevpow(2, 9)
8

julia> prevpow(5, 20)
5

julia> prevpow(4, 16)
16
```
See also [`nextpow`](@ref).
"""
function prevpow(a::Real, x::Real)
    x < 1 && throw(DomainError(x, "`x` must be ≥ 1."))
    # See comment in nextpos() for a == special case.
    a == 2 && isa(x, Integer) && return _prevpow2(x)
    a <= 1 && throw(DomainError(a, "`a` must be greater than 1."))
    n = floor(Integer,log(a, x))
    p = a^(n+1)
    p <= x ? p : a^n
end

## ndigits (number of digits) in base 10 ##

# decimal digits in an unsigned integer
const powers_of_ten = [
    0x0000000000000001, 0x000000000000000a, 0x0000000000000064, 0x00000000000003e8,
    0x0000000000002710, 0x00000000000186a0, 0x00000000000f4240, 0x0000000000989680,
    0x0000000005f5e100, 0x000000003b9aca00, 0x00000002540be400, 0x000000174876e800,
    0x000000e8d4a51000, 0x000009184e72a000, 0x00005af3107a4000, 0x00038d7ea4c68000,
    0x002386f26fc10000, 0x016345785d8a0000, 0x0de0b6b3a7640000, 0x8ac7230489e80000,
]
function bit_ndigits0z(x::Base.BitUnsigned64)
    lz = (sizeof(x)<<3)-leading_zeros(x)
    nd = (1233*lz)>>12+1
    nd -= x < powers_of_ten[nd]
end
function bit_ndigits0z(x::UInt128)
    n = 0
    while x > 0x8ac7230489e80000
        x = div(x,0x8ac7230489e80000)
        n += 19
    end
    return n + ndigits0z(UInt64(x))
end

ndigits0z(x::BitSigned) = bit_ndigits0z(unsigned(abs(x)))
ndigits0z(x::BitUnsigned) = bit_ndigits0z(x)
ndigits0z(x::Integer) = ndigits0zpb(x, 10)

## ndigits with specified base ##

# The suffix "nb" stands for "negative base"
function ndigits0znb(x::Integer, b::Integer)
    d = 0
    if x isa Unsigned
        d += (x != 0)::Bool
        x = -signed(fld(x, -b))
    end
    # precondition: b < -1 && !(typeof(x) <: Unsigned)
    while x != 0
        x = cld(x,b)
        d += 1
    end
    return d
end

# do first division before conversion with signed here, which can otherwise overflow
ndigits0znb(x::Bool, b::Integer) = x % Int

# The suffix "pb" stands for "positive base"
function ndigits0zpb(x::Integer, b::Integer)
    # precondition: b > 1
    x == 0 && return 0
    b = Int(b)
    x = abs(x)
    if x isa Base.BitInteger
        x = unsigned(x)::Unsigned
        b == 2  && return sizeof(x)<<3 - leading_zeros(x)
        b == 8  && return (sizeof(x)<<3 - leading_zeros(x) + 2) ÷ 3
        b == 16 && return sizeof(x)<<1 - leading_zeros(x)>>2
        b == 10 && return bit_ndigits0z(x)
    end

    d = 0
    while x > typemax(Int)
        x = div(x,b)
        d += 1
    end
    x = div(x,b)
    d += 1

    m = 1
    while m <= x
        m *= b
        d += 1
    end
    return d
end

ndigits0zpb(x::Bool, b::Integer) = x % Int

# The suffix "0z" means that the output is 0 on input zero (cf. #16841)
"""
    ndigits0z(n::Integer, b::Integer=10)

Return 0 if `n == 0`, otherwise compute the number of digits in
integer `n` written in base `b` (i.e. equal to `ndigits(n, base=b)`
in this case).
The base `b` must not be in `[-1, 0, 1]`.

# Examples
```jldoctest
julia> Base.ndigits0z(0, 16)
0

julia> Base.ndigits(0, base=16)
1

julia> Base.ndigits0z(0)
0

julia> Base.ndigits0z(10, 2)
4

julia> Base.ndigits0z(10)
2
```

See also [`ndigits`](@ref).
"""
function ndigits0z(x::Integer, b::Integer)
    if b < -1
        ndigits0znb(x, b)
    elseif b > 1
        ndigits0zpb(x, b)
    else
        throw(DomainError(b, "The base must not be in `[-1, 0, 1]`."))
    end
end

"""
    ndigits(n::Integer; base::Integer=10, pad::Integer=1)

Compute the number of digits in integer `n` written in base `base`
(`base` must not be in `[-1, 0, 1]`), optionally padded with zeros
to a specified size (the result will never be less than `pad`).

# Examples
```jldoctest
julia> ndigits(12345)
5

julia> ndigits(1022, base=16)
3

julia> string(1022, base=16)
"3fe"

julia> ndigits(123, pad=5)
5
```
"""
ndigits(x::Integer; base::Integer=10, pad::Integer=1) = max(pad, ndigits0z(x, base))

## integer to string functions ##

function bin(x::Unsigned, pad::Integer, neg::Bool)
    i = neg + max(pad,sizeof(x)<<3-leading_zeros(x))
    a = StringVector(i)
    while i > neg
        @inbounds a[i] = 48+(x&0x1)
        x >>= 1
        i -= 1
    end
    if neg; @inbounds a[1]=0x2d; end
    String(a)
end

function oct(x::Unsigned, pad::Integer, neg::Bool)
    i = neg + max(pad,div((sizeof(x)<<3)-leading_zeros(x)+2,3))
    a = StringVector(i)
    while i > neg
        @inbounds a[i] = 48+(x&0x7)
        x >>= 3
        i -= 1
    end
    if neg; @inbounds a[1]=0x2d; end
    String(a)
end

function dec(x::Unsigned, pad::Integer, neg::Bool)
    i = neg + ndigits(x, base=10, pad=pad)
    a = StringVector(i)
    while i > neg
        @inbounds a[i] = 48+rem(x,10)
        x = oftype(x,div(x,10))
        i -= 1
    end
    if neg; @inbounds a[1]=0x2d; end
    String(a)
end

function hex(x::Unsigned, pad::Integer, neg::Bool)
    i = neg + max(pad,(sizeof(x)<<1)-(leading_zeros(x)>>2))
    a = StringVector(i)
    while i > neg
        d = x & 0xf
        @inbounds a[i] = 48+d+39*(d>9)
        x >>= 4
        i -= 1
    end
    if neg; @inbounds a[1]=0x2d; end
    String(a)
end

const base36digits = ['0':'9';'a':'z']
const base62digits = ['0':'9';'A':'Z';'a':'z']

function _base(b::Integer, x::Integer, pad::Integer, neg::Bool)
    (x >= 0) | (b < 0) || throw(DomainError(x, "For negative `x`, `b` must be negative."))
    2 <= abs(b) <= 62 || throw(DomainError(b, "base must satisfy 2 ≤ abs(base) ≤ 62"))
    digits = abs(b) <= 36 ? base36digits : base62digits
    i = neg + ndigits(x, base=b, pad=pad)
    a = StringVector(i)
    @inbounds while i > neg
        if b > 0
            a[i] = digits[1+rem(x,b)]
            x = div(x,b)
        else
            a[i] = digits[1+mod(x,-b)]
            x = cld(x,b)
        end
        i -= 1
    end
    if neg; a[1]='-'; end
    String(a)
end

split_sign(n::Integer) = unsigned(abs(n)), n < 0
split_sign(n::Unsigned) = n, false

"""
    string(n::Integer; base::Integer = 10, pad::Integer = 1)

Convert an integer `n` to a string in the given `base`,
optionally specifying a number of digits to pad to.

```jldoctest
julia> string(5, base = 13, pad = 4)
"0005"

julia> string(13, base = 5, pad = 4)
"0023"
```
"""
function string(n::Integer; base::Integer = 10, pad::Integer = 1)
    if base == 2
        (n_positive, neg) = split_sign(n)
        bin(n_positive, pad, neg)
    elseif base == 8
        (n_positive, neg) = split_sign(n)
        oct(n_positive, pad, neg)
    elseif base == 10
        (n_positive, neg) = split_sign(n)
        dec(n_positive, pad, neg)
    elseif base == 16
        (n_positive, neg) = split_sign(n)
        hex(n_positive, pad, neg)
    else
        _base(base, base > 0 ? unsigned(abs(n)) : convert(Signed, n), pad, (base>0) & (n<0))
    end
end

string(b::Bool) = b ? "true" : "false"

"""
    bitstring(n)

A string giving the literal bit representation of a number.

# Examples
```jldoctest
julia> bitstring(4)
"0000000000000000000000000000000000000000000000000000000000000100"

julia> bitstring(2.2)
"0100000000000001100110011001100110011001100110011001100110011010"
```
"""
function bitstring end

bitstring(x::Union{Bool,Int8,UInt8})           = string(reinterpret(UInt8,x), pad = 8, base = 2)
bitstring(x::Union{Int16,UInt16,Float16})      = string(reinterpret(UInt16,x), pad = 16, base = 2)
bitstring(x::Union{Char,Int32,UInt32,Float32}) = string(reinterpret(UInt32,x), pad = 32, base = 2)
bitstring(x::Union{Int64,UInt64,Float64})      = string(reinterpret(UInt64,x), pad = 64, base = 2)
bitstring(x::Union{Int128,UInt128})            = string(reinterpret(UInt128,x), pad = 128, base = 2)

"""
    digits([T<:Integer], n::Integer; base::T = 10, pad::Integer = 1)

Return an array with element type `T` (default `Int`) of the digits of `n` in the given
base, optionally padded with zeros to a specified size. More significant digits are at
higher indices, such that `n == sum([digits[k]*base^(k-1) for k=1:length(digits)])`.

# Examples
```jldoctest
julia> digits(10, base = 10)
2-element Array{Int64,1}:
 0
 1

julia> digits(10, base = 2)
4-element Array{Int64,1}:
 0
 1
 0
 1

julia> digits(10, base = 2, pad = 6)
6-element Array{Int64,1}:
 0
 1
 0
 1
 0
 0
```
"""
digits(n::Integer; base::Integer = 10, pad::Integer = 1) =
    digits(typeof(base), n, base = base, pad = pad)

function digits(T::Type{<:Integer}, n::Integer; base::Integer = 10, pad::Integer = 1)
    digits!(zeros(T, ndigits(n, base=base, pad=pad)), n, base=base)
end

"""
    hastypemax(T::Type) -> Bool

Return `true` if and only if `typemax(T)` is defined.
"""
hastypemax(::Base.BitIntegerType) = true
hastypemax(::Type{T}) where {T} = applicable(typemax, T)

"""
    digits!(array, n::Integer; base::Integer = 10)

Fills an array of the digits of `n` in the given base. More significant digits are at higher
indices. If the array length is insufficient, the least significant digits are filled up to
the array length. If the array length is excessive, the excess portion is filled with zeros.

# Examples
```jldoctest
julia> digits!([2,2,2,2], 10, base = 2)
4-element Array{Int64,1}:
 0
 1
 0
 1

julia> digits!([2,2,2,2,2,2], 10, base = 2)
6-element Array{Int64,1}:
 0
 1
 0
 1
 0
 0
```
"""
function digits!(a::AbstractVector{T}, n::Integer; base::Integer = 10) where T<:Integer
    2 <= abs(base) || throw(DomainError(base, "base must be ≥ 2 or ≤ -2"))
    hastypemax(T) && abs(base) - 1 > typemax(T) &&
        throw(ArgumentError("type $T too small for base $base"))
    isempty(a) && return a

    if base > 0
        for i in eachindex(a)
            n, d = divrem(n, base)
            a[i] = d
        end
    else
        # manually peel one loop iteration for type stability
        n, d = fldmod(n, -base)
        a[firstindex(a)] = d
        n = -signed(n)
        for i in firstindex(a)+1:lastindex(a)
            n, d = fldmod(n, -base)
            a[i] = d
            n = -n
        end
    end
    return a
end

"""
    isqrt(n::Integer)

Integer square root: the largest integer `m` such that `m*m <= n`.

```jldoctest
julia> isqrt(5)
2
```
"""
isqrt(x::Integer) = oftype(x, trunc(sqrt(x)))

function isqrt(x::Union{Int64,UInt64,Int128,UInt128})
    x==0 && return x
    s = oftype(x, trunc(sqrt(x)))
    # fix with a Newton iteration, since conversion to float discards
    # too many bits.
    s = (s + div(x,s)) >> 1
    s*s > x ? s-1 : s
end

"""
    factorial(n::Integer)

Factorial of `n`. If `n` is an [`Integer`](@ref), the factorial is computed as an
integer (promoted to at least 64 bits). Note that this may overflow if `n` is not small,
but you can use `factorial(big(n))` to compute the result exactly in arbitrary precision.

# Examples
```jldoctest
julia> factorial(6)
720

julia> factorial(21)
ERROR: OverflowError: 21 is too large to look up in the table; consider using `factorial(big(21))` instead
Stacktrace:
[...]

julia> factorial(big(21))
51090942171709440000
```

# See also
* [`binomial`](@ref)

# External links
* [Factorial](https://en.wikipedia.org/wiki/Factorial) on Wikipedia.
"""
function factorial(n::Integer)
    n < 0 && throw(DomainError(n, "`n` must be nonnegative."))
    f::typeof(n*n) = 1
    for i::typeof(n*n) = 2:n
        f *= i
    end
    return f
end

"""
    binomial(n::Integer, k::Integer)

The _binomial coefficient_ ``\\binom{n}{k}``, being the coefficient of the ``k``th term in
the polynomial expansion of ``(1+x)^n``.

If ``n`` is non-negative, then it is the number of ways to choose `k` out of `n` items:
```math
\\binom{n}{k} = \\frac{n!}{k! (n-k)!}
```
where ``n!`` is the [`factorial`](@ref) function.

If ``n`` is negative, then it is defined in terms of the identity
```math
\\binom{n}{k} = (-1)^k \\binom{k-n-1}{k}
```

# Examples
```jldoctest
julia> binomial(5, 3)
10

julia> factorial(5) ÷ (factorial(5-3) * factorial(3))
10

julia> binomial(-5, 3)
-35
```

# See also
* [`factorial`](@ref)

# External links
* [Binomial coeffient](https://en.wikipedia.org/wiki/Binomial_coefficient) on Wikipedia.
"""
function binomial(n::T, k::T) where T<:Integer
    n0, k0 = n, k
    k < 0 && return zero(T)
    sgn = one(T)
    if n < 0
        n = -n + k -1
        if isodd(k)
            sgn = -sgn
        end
    end
    k > n && return zero(T)
    (k == 0 || k == n) && return sgn
    k == 1 && return sgn*n
    if k > (n>>1)
        k = (n - k)
    end
    x::T = nn = n - k + 1
    nn += 1
    rr = 2
    while rr <= k
        xt = div(widemul(x, nn), rr)
        x = xt % T
        x == xt || throw(OverflowError("binomial($n0, $k0) overflows"))
        rr += 1
        nn += 1
    end
    convert(T, copysign(x, sgn))
end


# # Spectrahedron
#
# This example shows an optimization problem over the spectraplex:
# ```math
# S = \{X \in \mathbb{S}_+^n, Tr(X) = 1\}
# ```
# with $\mathbb{S}_+^n$ the set of positive semidefinite matrices.
# Linear optimization with symmetric objective $D$ over the spetraplex consists in computing the
# leading eigenvector of $D$.
#
# The package also exposes `UnitSpectrahedronLMO` which corresponds to the feasible set:
# ```math
# S_u = \{X \in \mathbb{S}_+^n, Tr(X) \leq 1\}
# ```

using FrankWolfe
using LinearAlgebra
using Random
using SparseArrays

# The objective function will be the symmetric squared distance to a set of known or observed entries $Y_{ij}$ of the matrix.
# ```math
# f(X) = \sum_{(i,j) \in L} 1/2 (X_{ij} - Y_{ij})^2
# ```

# ## Setting up the input data, objective, and gradient

# Dimension, number of iterations and number of known entries:
n = 500
k = 10000
n_entries = 50

Random.seed!(41)

const entry_indices = unique!([minmax(rand(1:n, 2)...) for _ in 1:n_entries])
const entry_values = randn(length(entry_indices))

function f(X)
    r = zero(eltype(X))
    for (idx, (i, j)) in enumerate(entry_indices)
        r += 1/2 * (X[i,j] - entry_values[idx])^2
        r += 1/2 * (X[j,i] - entry_values[idx])^2
    end
    return r
end

function grad!(storage, X)
    storage .= 0
    for (idx, (i, j)) in enumerate(entry_indices)
        storage[i,j] += (X[i,j] - entry_values[idx])
        storage[j,i] += (X[j,i] - entry_values[idx])
    end
end

# Note that the `ensure_symmetry = false` argument to `SpectraplexLMO`.
# It skips an additional step making the used direction symmetric.
# It is not necessary when the gradient is a `LinearAlgebra.Symmetric` (or more rarely a `LinearAlgebra.Diagonal` or `LinearAlgebra.UniformScaling`).

const lmo = FrankWolfe.SpectraplexLMO(1.0, n, false)
const x0 = FrankWolfe.compute_extreme_point(lmo, spzeros(n, n))

target_tolerance = 1e-6;

#src the following two lines are used only to precompile the functions
FrankWolfe.frank_wolfe(f, grad!, lmo, x0, max_iteration=2, line_search=FrankWolfe.MonotonousStepSize()) #src
FrankWolfe.lazified_conditional_gradient(f, grad!, lmo, x0, max_iteration=2, line_search=FrankWolfe.MonotonousStepSize()) #src

# ## Running standard and lazified Frank-Wolfe

Xfinal, Vfinal, primal, dual_gap, trajectory = FrankWolfe.frank_wolfe(
    f,
    grad!,
    lmo,
    x0,
    max_iteration=k,
    line_search=FrankWolfe.MonotonousStepSize(),
    print_iter=k / 10,
    emphasis=FrankWolfe.memory,
    verbose=true,
    trajectory=true,
    epsilon=target_tolerance,
)

Xfinal, Vfinal, primal, dual_gap, trajectory_lazy = FrankWolfe.lazified_conditional_gradient(
    f,
    grad!,
    lmo,
    x0,
    max_iteration=k,
    line_search=FrankWolfe.MonotonousStepSize(),
    print_iter=k / 10,
    emphasis=FrankWolfe.memory,
    verbose=true,
    trajectory=true,
    epsilon=target_tolerance,
)

# ## Plotting the resulting trajectories

data = [trajectory, trajectory_lazy]
label = ["FW", "LCG"]
FrankWolfe.plot_trajectories(data, label, xscalelog=true)


# Experimental functions for linear extensions

export linear_extension, linear_extension_empty!, all_linear_extensions, random_linear_extension

"""
`linear_extension(P)` returns a linear extension of the poset `P`.
"""
function linear_extension(P::SimplePoset{T}) where T
    PP = deepcopy(P)
    result = linear_extension_empty!(PP)
    return result
end

"""
`linear_extension_empty!(P)` returns a linear extension of the poset `P`,
but also empties the poset `P`. This is a destructive operation.
"""
function linear_extension_empty!(PP::SimplePoset{T}) where T
    result = T[]
    while card(PP)>0
        M = minimals(PP)
        append!(result, M)
        for x in M
            delete!(PP,x)
        end
    end
    return result
end

"""
`random_linear_extension(P)` generates a linear extension of `P`
at random. This is slower than `linear_extension(P)` if one just
wants some linear extension.

Note that linear extensions are not generated uniformly
at random by this function. Rather we proceed recursively: A
minimal element `x` of `P` is chosen uniformly at random from among
all minimal elements, and that becomes the first (lowest) element in
the linear extension. Then `x` is deleted and the same process is
repeated until all elements have been processed.
"""
function random_linear_extension(P::SimplePoset{T}) where T
    result = T[]
    PP = deepcopy(P)
    while card(PP) > 0
        x = rand(minimals(PP))  # choose a min el't at random
        push!(result,x)
        delete!(PP,x)
    end
    return result
end

_LX_table = Dict{SimplePoset, Set}()
export clear_LX_table

"""
`clear_LX_table()` releases cached results computed by
`all_linear_extensions`.
"""
function clear_LX_table()
    global _LX_table = Dict{SimplePoset, Set}()
    nothing
end

"""
`all_linear_extensions(P)` returns the `Set` of all linear extensions
of `P`. This can take a very long time and eat up a lot of memory
(which can be freed using `clear_LX_table`).
"""
function all_linear_extensions(P::SimplePoset)::Set
    # see if we already have an answer
    global _LX_table
    if haskey(_LX_table,P)
        return _LX_table[P]
    end

    T = eltype(P)
    result::Set = Set{Array{T,1}}()
    if card(P) == 0
        return result
    end
    if card(P) == 1
        L = elements(P)
        push!(result, L)
        return result
    end

    M = maximals(P)
    for x in M
        PP = deepcopy(P)
        delete!(PP,x)
        PP_exts = all_linear_extensions(PP)
        for L in PP_exts
            LL = deepcopy(L)
            append!(LL,[x])
            push!(result, LL)
        end
    end
    _LX_table[P] = result   # save this answer
    return result
end


const cardpbk = 14012298 # Card's public key
const doorpbk = 74241  # Door's public key
# Puzze input

const subjectnum = 7
# Subject number

const transfnum = 20201227
# Number used in the operaration that transform a subject number

# The card transforms the subject number of 7 according to the card's
# secret loop size. The result is called the card's public key (cardpbk).

function findloopzise(pbk)
    val = 1
    for i ∈ 1:typemax(Int)
        # To transform a subject number (subjectnum), start with the value 1.
        # Then, a number of times called the loop size, perform the following steps:

        val *= subjectnum
        # 1. Set the value to itself multiplied by the subject number.

        val %= transfnum
        # 2. Set the value to the remainder after dividing the value
        # by 20201227 (transfnum)

        if val == pbk
            # The card transforms the subject number of 7 according to
            # the card's secret loop size. The result is called the card's public key.
            # Then when val is equal to out public ket we found our secret loop size

            return i # Secret loop size
        end
    end
end

function findencriptkey(cardpbk, doorpbk)
    cardloopsize = findloopzise(cardpbk)
    doorloopsize = findloopzise(doorpbk)

    # You can use either device's loop size with the other device's public key
    # to calculate the encryption key

    subjectnum = doorpbk
    encriptkey₁ = 1
    for i ∈ 1:cardloopsize
        # The card transforms the subject number of the door's public key
        # according to the card's loop size. The result is the encryption key.

        encriptkey₁ *= subjectnum
        encriptkey₁ %= transfnum
    end

    subjectnum = cardpbk
    encriptkey₂ = 1
    for i ∈ 1:doorloopsize
        # Transforming the subject number of cardpbk (the card's public key)
        # with a loop size of doorloopsize produces the encryption key

        encriptkey₂ *= subjectnum
        encriptkey₂ %= transfnum
    end
    return encriptkey₁, encriptkey₂
end

findencriptkey(cardpbk, doorpbk)


#' ---
#' title: Getting Started with JuMP
#' ---

#' **Originally Contributed by**: Arpit Bhatia

#' This tutorial is aimed at providing a quick introduction to writing JuMP code. It assumes familiar with basic optimization and
#' the notion of an [AML](https://en.wikipedia.org/wiki/Algebraic_modeling_language).

#' ## What is JuMP?
#' JuMP ("Julia for Mathematical Programming") is an open-source modeling language that is embedded in Julia. It allows users to
#' users formulate various classes of optimization problems (linear, mixed-integer, quadratic, conic quadratic, semidefinite,
#' and nonlinear) with easy-to-read code. These problems can then be solved using state-of-the-art open-source and commercial solvers.
#' JuMP also makes advanced optimization techniques easily accessible from a high-level language.

#' ## Installing JuMP
#' JuMP is a package for Julia. From Julia, JuMP is installed by using the built-in package manager.
#+ tangle = false

import Pkg
Pkg.add("JuMP")

#' ## A Complete Example
#' Let's try to solve the following linear programming problem by using JuMP and GLPK (a linear and mixed integer programming
#' solver). We will first look at the complete code to solve the problem and then go through it step by step.

#' $$
#' \begin{align*}
#' & \min & 12x + 20y \\
#' & \;\;\text{s.t.} & 6x + 8y \geq 100 \\
#' & & 7x + 12y \geq 120 \\
#' & & x \geq 0 \\
#' & & y \geq 0 \\
#' \end{align*}
#' $$

using JuMP
using GLPK

model = Model(with_optimizer(GLPK.Optimizer))
@variable(model, x >= 0)
@variable(model, y >= 0)
@constraint(model, 6x + 8y >= 100)
@constraint(model, 7x + 12y >= 120)
@objective(model, Min, 12x + 20y)

optimize!(model)

@show value(x);
@show value(y);
@show objective_value(model);

#' ## Step by Step JuMP Code
#' Once JuMP is installed, to use JuMP in your programs, we just need to write-

using JuMP

#' We also need to include a Julia package which provides an appropriate solver. We want to use GLPK.Optimizer here which is
#' provided by the GLPK.jl package.

using GLPK

#' A model object is a container for variables, constraints, solver options, etc. Models are created with the Model() function.
#' The `with_optimizer` syntax is used to specify the optimizer to be used which is GLPK in this case.

model = Model(with_optimizer(GLPK.Optimizer));

#' A variable is modelled using `@variable(name of the model object, variable name and bound, variable type)`. The bound can be a
#' lower bound, an upper bound or both. If no variable type is defined, then it is treated as real.

@variable(model, x >= 0)
@variable(model, y >= 0);

#' A constraint is modelled using `@constraint(name of the model object, constraint)`.

@constraint(model, 6x + 8y >= 100)
@constraint(model, 7x + 12y >= 120);

#' The objective is set in a similar manner using `@objective(name of the model object, Min/Max, function to be optimized)`

@objective(model, Min, 12x + 20y);

#' To solve the optimization problem, we call the optimize function.

optimize!(model)

#' Let's now check the value of objective and variables.

@show value(x);
@show value(y);
@show objective_value(model);


using Plots
using LinearAlgebra

theta = [pi/15, -pi/9, -pi/20]
Sigma = diagm([3.0, 1.0, 0.5]) # scale x, then y, then z

Rx = [1 0 0; 0 cos(theta[1]) -sin(theta[1]);  0 sin(theta[1]) cos(theta[1])] # rotate about x-axis
Ry = [cos(theta[2]) 0 sin(theta[2]); 0 1 0; -sin(theta[2]) 0 cos(theta[2])] # rotate about y-axis
Rz = [cos(theta[3]) -sin(theta[3]) 0; sin(theta[3]) cos(theta[3]) 0; 0 0 1] # rotate about z-axis

X = Rz*Ry*Rx*Sigma

# Plot sphere
n = 100
u = range(0, stop = 2pi, length = n)
v = range(0, stop = pi, length = n)

x = cos.(u) * sin.(v)'
y = sin.(u) * sin.(v)'
z = ones(n) * cos.(v)'

p1 = plot(x, y, z, legend = false)

xR = zeros(n, n)
yR = zeros(n, n)
zR = zeros(n, n)

for i = 1:n
    for j = 1:n
        vecR = X*[x[i,j]; y[i,j]; z[i,j]]
        xR[i,j] = vecR[1]
        yR[i,j] = vecR[2]
        zR[i,j] = vecR[3]
    end
end

p2 = plot(xR, yR, zR, legend = false)


using LinearAlgebra
using Makie

μ1 = μ2 = 0.0

𝛍 = [
    μ1
    μ2
]

σ11 = σ22 = 1.0
σ12 = σ21 = 0.0

Σ = [
    σ11 σ12
    σ21 σ22
]

function 𝒩(x1, x2; 𝛍 = 𝛍, Σ = Σ)
    𝐱 = [x1, x2]
    Λ = inv(Σ)
    D = length(𝛍)
    return √(det(Λ)) * exp(-0.5 * (𝐱 .- 𝛍)' * Λ * (𝐱 .- 𝛍)) / √(2π)^D
end

x1min, x1max = -1.0, 1.0
x2min, x2max = -1.0, 1.0
x1 = range(x1min, x1max, length = 50)
x2 = range(x2min, x2max, length = 50)
z = 𝒩.(x1', x2)
wireframe(x1, x2, z, colormap = :Spectral)


### TRIANGLE
area(v1::SVector{3,T}, v2::SVector{3,T}, v3::SVector{3,T}) where {T} = norm(vector_area(v1, v2, v3))
centroid(v1::SVector{3,T}, v2::SVector{3,T}, v3::SVector{3,T}) where {T} = (v1 + v2 + v3) * Float64(1/3)  # 4 times faster than dividing by 3
vector_area(v1::SVector{3,T}, v2::SVector{3,T}, v3::SVector{3,T}) where {T} = cross(v2 - v1, v3 - v2) * 0.5
function triangle_area(v123::NTuple{3,SVector{3,T}}, n̂::SVector{3,T}) where {T}
	vec_area = vector_area(v123[1], v123[2], v123[3])
	return dot(n̂, vec_area)
end
triangleNormal(v1::SVector{3,T}, v2::SVector{3,T}, v3::SVector{3,T}) where {T} = normalize(vector_area(v1, v2, v3))

for funName in (:area, :centroid, :triangleNormal, :triangleCross)
    @eval begin
        function $funName(sv::SVector{3,SVector{3,T}}) where {T}
            return $funName(sv[1], sv[2], sv[3])
        end
    end
end

### TETRAHEDRON
function centroid(v1::SVector{3,T}, v2::SVector{3,T}, v3::SVector{3,T}, v4::SVector{3,T}) where {T}
    return (v1 + v2 + v3 + v4) * 0.25
end

function volume(v1::SVector{3,T}, v2::SVector{3,T}, v3::SVector{3,T}, v4::SVector{3,T}) where {T}
    # NOTE: This is an algebraic refactorization of a symbolic answer
    a1, a2, a3 = v1[1], v1[2], v1[3]
    b1, b2, b3 = v2[1], v2[2], v2[3]
    c1, c2, c3 = v3[1], v3[2], v3[3]
    d1, d2, d3 = v4[1], v4[2], v4[3]
    V =       (b1 - a1) * (c2 * d3 - c3 * d2)
    V = muladd(b2 - a2,    c3 * d1 - c1 * d3, V)
    V = muladd(b3 - a3,    c1 * d2 - c2 * d1, V)
    V = muladd(c1 - d1,    a3 * b2 - a2 * b3, V)
    V = muladd(c2 - d2,    a1 * b3 - a3 * b1, V)
    V = muladd(c3 - d3,    a2 * b1 - a1 * b2, V)
    return V * Float64(1/6)
end

for funName in (:centroid, :volume)
    @eval begin
        function $funName(sv::SVector{4,SVector{3,T}}) where {T}
            return $funName(sv[1], sv[2], sv[3], sv[4])
        end
    end
end


using Plots

max_n = 30
n_arr = collect(5:max_n)
probs = []
probs_corner = []
probs_line = []

for n in n_arr
    N_m = n^2 - 4*(n-1)
    prob = N_m / n^2 * (N_m-1) / (n^2 - 1)
    push!(probs, prob)

    # Prob at least one node is a corner node
    prob_corner = 4/n^2 + (n^2 - 4)/(n^2) * (4/(n^2-1))
    push!(probs_corner, prob_corner)

    # Prob that users end on same row or column
    prob_line = 2*(n-1)/(n^2 - 1)
    push!(probs_line, prob_line)
end

plot(n_arr, 1 .- probs, label="probs not middle")
plot!(n_arr, probs_corner, label="probs corner")
plot!(n_arr, probs_line, label="probs line")


# Statistical model for analyzing Radial Velocity observations
# Assumes existance of data in arrays times, obs, sigma_obs
# Assumes existance of functions: loglikelihood(theta), logprior_noise(theta), and is_valid_noise(theta)

export logprior_planets, logprior, logtarget, negative_logtarget
export set_times, set_obs, set_sigma_obs, num_noise_param

# Observational data to compare model to
global times # ::Array{Float64,1}
global obs # ::Array{Float64,1}
global sigma_obs #::Array{Float64,1}

# Functions to set global data within module
function set_times(t::Array{Float64,1}) global times = t   end
function set_obs(o::Array{Float64,1})  global obs = o   end
function set_sigma_obs(so::Array{Float64,1}) global sigma_obs = so  end

function is_valid(p::Vector)
   is_valid_noise(p) && is_valid_planets(p)
end

function logprior_period(P::Real)
  const min_period = 1.0
  const max_period = 10000.0
  const norm = -log(max_period/min_period)
  if !(min_period<=P<=max_period)
     return -Inf
  end
  return norm-log(P)
end

function logprior_amplitude(K::Real)
  const min_K = 0.0
  const max_K = 999.0
  const K0 = 1.0
  const norm = -log(K0*log1p(maxK/min_K))
  if !(min_K<=K<=max_K)
     return -Inf
  end
  return norm-log1p(K/K0)
end

function logprior_planets(p::Vector)
  if !is_valid_planets(p) return -Inf end  # prempt model evaluation
  num_pl = num_planets(p)
  logp = zero(eltype(p))
  if num_pl <= 0 return logp end
  if num_param_per_planet == 3
     logp -= num_pl*log(2pi)
  elseif num_param_per_planet == 5
     logp -= 2*num_pl*log(2pi)
  end
  for plid in 1:num_pl
    P::eltype(p) = extract_period(p,plid=plid)
    K::eltype(p) = extract_amplitude(p,plid=plid)
    logp += logprior_period(P) + logprior_amplitude(K)
  end
  return logp::eltype(p)
end

function logprior_offset(C::Real)
  const max_C = 1000.0
  if !(-max_C<=C<=max_C)
     return -Inf
  end
  return -log(2*max_C)
end

logprior_offset(p::Vector) = logprior_offset(extract_rvoffset(p))

function logprior(p::Vector)
  logprior_noise(p) + logprior_offset(p) + logprior_planets(p)
end

function logtarget(p::Vector)
  val = logprior(p)
  if val==-Inf
     return val
  end
  val += loglikelihood(p)
  return val
end

negative_logtarget(p::Vector) = -logtarget(p)


using JuMP, Ipopt, Plots, LinearAlgebra, SparseArrays

#A JuMP model is initialized
mod = Model(optimizer_with_attributes( Ipopt.Optimizer ,  "max_iter" => 100,
            "mumps_mem_percent" => 500))

#The parameters of the 1-D rod
L = 0.1                                 #Length of rod
λ = 45.0                                # Conductivity
c = 460.0                               # Specific heat capacitivity
ρ = 7800.0                              # Material density
α = λ / (c*ρ)                           # Diffusivity

#Number of spatial discretization points and length of each
N = 11                                 #Discretization points / nodes
Δx = L / (N-1)                         # Δx = x[i+1] - x[i]

#Integration from t0 = 0 to tf
tf = 10000.0;                          #Final time
Δt = 10^(-1);                          #Sampling period
k = round(Int,tf / Δt);                #k = number of steps

#System's matrix: θ' = Aθ + Bu
A = spdiagm(-1 => ones( N-1), 0 => -2*ones(N), 1 => ones(N-1));
A[2,1] = 2.0;
A[end - 1,end] = 2.0;
A = (α/( Δx^2)) * A;

#input matrix
B = spzeros(N);
B[1] = 1;
B = (2/ (c * ρ * Δx)) * B;

#output matrix
C = spzeros(1,N)
C[1,end] = 1.0;

#Initial and reference temperatures in kelvin
θinit = 273.0
θref = 500.0
#bounded input values[u_min,u_max]
u_min = 0.0
u_max = 15000.0

#Decision variables of the optimization problem
@variable(mod,θ[1:N,1:k]);                              # temperature at x[1:N] at time step k
@variable(mod,u_min <= u[1:k - 1] <= u_max);            # bounded constrained Input u
@variable(mod, y[1:k]);                                 # Output y
#Initial values
@constraint(mod, θ[:,1] .== θinit)
@constraint(mod, y[1] .== θinit)
#System dynamics
for j in 1:k - 1,i in 1:N
    @constraint(mod,θ[i,j + 1] - (θ[i,j] + Δt * ((A[:,i]' * θ[:,j]) + B[i] * u[j])).== 0.0)
end

for j in 1:k - 1
    @constraint(mod,y[j + 1] .== (C * θ[:,j + 1]))
end
#Weighing matrices
Q = C' * C * 1000.0;
R = 0.0001;

#err = sum(e' Q e)
err= @NLexpression(mod,(sum(Q[i,i] * (θref - θ[i,j]) ^ 2
                        for j in 1:k-1,i in 1:N)))
#in_err = sum(u'Ru)
in_err = @NLexpression(mod,(sum(R * u[j] ^ 2 for j in 1:k-1)))

J = @NLexpression(mod,0.5 * Δt * (err + in_err))
#defining the cost function
@NLobjective(mod, Min, J)
#optimizes the model
optimize!(mod)

#plotting values
states_θstart = zeros(k)
states_θend = zeros(k)
input_u = zeros(k - 1)
for i = 1 : k-1
    states_θstart[i] = JuMP.value(θ[1,i])
    states_θend[i] = JuMP.value(θ[end,i])
    input_u[i] = JuMP.value(u[i])
end

states_θend[end] = JuMP.value(θ[end,end])
states_θstart[end] = JuMP.value(θ[1,end])

plot(states_θstart,label = "θ1")
plot!(states_θend,label="y = Cθ")
savefig("lq-optimal_1D_Q1000_R0001_output-bounded.png")
p3 = plot(input_u,label="input,Q = 1000,R = 0.001")
savefig("lq-optimal_1D_Q1000_R0001-input-bounded.png")


# Michelle Medina
# Finite Differencing
# DATE

export Params, Params2D
export dX, dX_2D, dY_2D, functionOnGrid, Deriv1D, DerivX_2D, DerivY_2D

struct Params
#structure contains # of points N and Domain Space L
    N::Int
    L::Float64
end

struct Params2D
    #contains number of points Nx and Ny and the length of the
    #2D grid Lx and Ly
    Nx::Int
    Ny::Int
    Lx::Float64
    Ly::Float64
end

# Defining the Grid Space
""" computing grid space h """
function dX(grid::Params)
    return (grid.L/(grid.N-1))
end

""" computing grid space hx """
function dX_2D(grid::Params2D)
    return (grid.Lx/(grid.Nx-1))
end

""" computing grid space hy """
function dY_2D(grid::Params2D)
    #computes grid space hy
    return (grid.Ly/(grid.Ny-1))
end

""" 1D Function is computed on the grid """
function functionOnGrid(u::Function, grid::Params)::Array{Float64,1}
    ux=zeros(grid.N)
    h=dX(grid)
    for i in 1:grid.N
        x=(i-1)*h
        ux[i]=u(x)
    end
    return ux
end

""" 2D Function is computed on the grid """
function functionOnGrid(u::Function, grid::Params2D)::Array{Float64,2}
    uxy=zeros(grid.Nx,grid.Ny)
    hx=dX_2D(grid)
    hy=dY_2D(grid)
    for i in 1:grid.Nx, j in 1:grid.Ny
        x = (i-1)*hx
        y= (j-1)*hy
        uxy[i,j]=u(x,y)
    end
    return uxy
end

""" Derivative of a function is computed on a 1D grid """
function Deriv1D(u::Array{Float64,1},grid::Params)::Array{Float64}
    dudx = zeros(grid.N)
    h = dX(grid)
    for i in 1:grid.N
        if i == 1
            dudx[i] = (u[i+1] - u[i])/h
        elseif i == grid.N
            dudx[i] = (u[i]-u[i-1])/h
        else
            dudx[i] = (u[i+1]-u[i-1])/2h
        end
    end
    return dudx
end

function DerivX_2D(u::Array{Float64,2}, grid::Params2D)::Array{Float64}
    dudx= zeros(grid.Nx,grid.Ny)
    hx = dX_2D(grid)

    for i in 1:grid.Nx, j in 1:grid.Ny
        if i==1
            dudx[i,j] = (u[i+1,j] - u[i,j])/hx
        elseif i==grid.Nx
            dudx[i,j] = (u[i,j] - u[i-1,j])/hx
        else
            dudx[i,j] = (u[i+1,j]-u[i-1,j])/2hx
        end
    end
    return dudx
end

function DerivY_2D(u::Array{Float64,2},grid::Params2D)::Array{Float64}
    dudy= zeros(grid.Nx,grid.Ny)
    hy = dY_2D(grid)

    for i in 1:grid.Nx, j in 1:grid.Ny
        if j == 1
            dudy[i,j] = (u[i,j+1] - u[i,j])/hy
        elseif j== grid.Ny
            dudy[i,j] = (u[i,j]-u[i,j-1])/hy
        else
            dudy[i,j] = (u[i,j+1]-u[i,j-1])/2hy
        end
    end
    return dudy
end