text:
module ModernRoboticsBook

import LinearAlgebra
const linalg = LinearAlgebra

export NearZero,
       Normalize,
       RotInv,
       VecToso3,
       so3ToVec,
       AxisAng3,
       MatrixExp3,
       MatrixLog3,
       RpToTrans,
       TransToRp,
       TransInv,
       VecTose3,
       se3ToVec,
       Adjoint,
       ScrewToAxis,
       AxisAng6,
       MatrixExp6,
       MatrixLog6,
       ProjectToSO3,
       ProjectToSE3,
       DistanceToSO3,
       DistanceToSE3,
       TestIfSO3,
       TestIfSE3,
       FKinBody,
       FKinSpace,
       JacobianBody,
       JacobianSpace,
       IKinBody,
       IKinSpace,
       ad,
       InverseDynamics,
       MassMatrix,
       VelQuadraticForces,
       GravityForces,
       EndEffectorForces,
       ForwardDynamics,
       EulerStep,
       InverseDynamicsTrajectory,
       ForwardDynamicsTrajectory,
       CubicTimeScaling,
       QuinticTimeScaling,
       JointTrajectory,
       ScrewTrajectory,
       CartesianTrajectory,
       ComputedTorque,
       SimulateControl

# """
# *** BASIC HELPER FUNCTIONS ***
# """

"""
    NearZero(z)

Determines whether a scalar is small enough to be treated as zero.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> NearZero(-1e-7)
true
```
"""
NearZero(z::Number) = abs(z) < 1e-6

"""
    Normalize(V)

Normalizes a vector.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> Normalize([1, 2, 3])
3-element Array{Float64,1}:
 0.2672612419124244
 0.5345224838248488
 0.8017837257372732
```
"""
Normalize(V::Array) = V / linalg.norm(V)

# """
# *** CHAPTER 3: RIGID-BODY MOTIONS ***
# """

"""
    RotInv(R)

Inverts a rotation matrix.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> RotInv([0 0 1; 1 0 0; 0 1 0])
3×3 LinearAlgebra.Adjoint{Int64,Array{Int64,2}}:
 0  1  0
 0  0  1
 1  0  0
```
"""
RotInv(R::Array) = R'

"""
    VecToso3(ω)

Converts a 3-vector to an so(3) representation.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> VecToso3([1 2 3])
3×3 Array{Int64,2}:
  0  -3   2
  3   0  -1
 -2   1   0
```
"""
function VecToso3(ω::Array)
    [ 0    -ω[3]  ω[2];
      ω[3]  0    -ω[1];
     -ω[2]  ω[1]  0   ]
end

"""
    so3ToVec(so3mat)

Converts an so(3) representation to a 3-vector.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> so3ToVec([0 -3 2; 3 0 -1; -2 1 0])
3-element Array{Int64,1}:
 1
 2
 3
```
"""
function so3ToVec(so3mat::Array)
    [so3mat[3, 2], so3mat[1, 3], so3mat[2, 1]]
end

"""
    AxisAng3(expc3)

Converts a 3-vector of exponential coordinates for rotation into axis-angle form.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> AxisAng3([1, 2, 3])
([0.267261, 0.534522, 0.801784], 3.7416573867739413)
```
"""
AxisAng3(expc3::Array) = Normalize(expc3), linalg.norm(expc3)

"""
    MatrixExp3(so3mat)

Computes the matrix exponential of a matrix in so(3).

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> MatrixExp3([0 -3 2; 3 0 -1; -2 1 0])
3×3 Array{Float64,2}:
 -0.694921   0.713521  0.0892929
 -0.192007  -0.303785  0.933192
  0.692978   0.63135   0.348107
```
"""
function MatrixExp3(so3mat::Array)
    omgtheta = so3ToVec(so3mat)
    if NearZero(linalg.norm(omgtheta))
        return linalg.I
    else
        θ = AxisAng3(omgtheta)[2]
        omgmat = so3mat / θ
        return linalg.I + sin(θ) * omgmat + (1 - cos(θ)) * omgmat * omgmat
    end
end

"""
    MatrixLog3(R)

Computes the matrix logarithm of a rotation matrix.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> MatrixLog3([0 0 1; 1 0 0; 0 1 0])
3×3 Array{Float64,2}:
  0.0     -1.2092   1.2092
  1.2092   0.0     -1.2092
 -1.2092   1.2092   0.0
```
"""
function MatrixLog3(R::Array)
    acosinput = (linalg.tr(R) - 1) / 2
    if acosinput >= 1
        return zeros(3, 3)
    elseif acosinput <= -1
        if !NearZero(1 + R[3, 3])
            omg = (1 / √(2 * (1 + R[3, 3]))) * [R[1, 3], R[2, 3], 1 + R[3, 3]]
        elseif !NearZero(1 + R[2, 2])
            omg = (1 / √(2 * (1 + R[2, 2]))) * [R[1, 2], 1 + R[2, 2], R[3, 2]]
        else
            omg = (1 / √(2 * (1 + R[1, 1]))) * [1 + R[1, 1], R[2, 1], R[3, 1]]
        end
        return VecToso3(π * omg)
    else
        θ = acos(acosinput)
        return θ / 2 / sin(θ) * (R - R')
    end
end

"""
    RpToTrans(R, p)

Converts a rotation matrix and a position vector into homogeneous transformation matrix.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> RpToTrans([1 0 0; 0 0 -1; 0 1 0], [1, 2, 5])
4×4 Array{Int64,2}:
 1  0   0  1
 0  0  -1  2
 0  1   0  5
 0  0   0  1
```
"""
RpToTrans(R::Array, p::Array) = vcat(hcat(R, p), [0 0 0 1])

"""
    TransToRp(T)

Converts a homogeneous transformation matrix into a rotation matrix and position vector.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> TransToRp([1 0 0 0; 0 0 -1 0; 0 1 0 3; 0 0 0 1])
([1 0 0; 0 0 -1; 0 1 0], [0, 0, 3])
```
"""
TransToRp(T::Array) = T[1:3, 1:3], T[1:3, 4]

"""
    TransInv(T)

Inverts a homogeneous transformation matrix.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> TransInv([1 0 0 0; 0 0 -1 0; 0 1 0 3; 0 0 0 1])
4×4 Array{Int64,2}:
 1   0  0   0
 0   0  1  -3
 0  -1  0   0
 0   0  0   1
```
"""
function TransInv(T::Array)
    R, p = TransToRp(T)
    vcat(hcat(R', -R' * p), [0 0 0 1])
end

"""
    VecTose3(V)

Converts a spatial velocity vector into a 4x4 matrix in se3.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> VecTose3([1 2 3 4 5 6])
4×4 Array{Float64,2}:
  0.0  -3.0   2.0  4.0
  3.0   0.0  -1.0  5.0
 -2.0   1.0   0.0  6.0
  0.0   0.0   0.0  0.0
```
"""
VecTose3(V::Array) = vcat(hcat(VecToso3(V[1:3]), V[4:6]), zeros(1, 4))

"""
    se3ToVec(se3mat)

Converts an se3 matrix into a spatial velocity vector.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> se3ToVec([0 -3 2 4; 3 0 -1 5; -2 1 0 6; 0 0 0 0])
6-element Array{Int64,1}:
 1
 2
 3
 4
 5
 6
```
"""
se3ToVec(se3mat::Array) = vcat([se3mat[3, 2],
                                se3mat[1, 3],
                                se3mat[2, 1]], se3mat[1:3, 4])

"""
    Adjoint(T)

Computes the adjoint representation of a homogeneous transformation matrix.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> Adjoint([1 0 0 0; 0 0 -1 0; 0 1 0 3; 0 0 0 1])
6×6 Array{Float64,2}:
 1.0  0.0   0.0  0.0  0.0   0.0
 0.0  0.0  -1.0  0.0  0.0   0.0
 0.0  1.0   0.0  0.0  0.0   0.0
 0.0  0.0   3.0  1.0  0.0   0.0
 3.0  0.0   0.0  0.0  0.0  -1.0
 0.0  0.0   0.0  0.0  1.0   0.0
```
"""
function Adjoint(T::Array)
    R, p = TransToRp(T)
    vcat(hcat(R, zeros(3, 3)), hcat(VecToso3(p) * R, R))
end

"""
    ScrewToAxis(q, s, h)

Takes a parametric description of a screw axis and converts it to a normalized screw axis.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> ScrewToAxis([3; 0; 0], [0; 0; 1], 2)
6-element Array{Int64,1}:
  0
  0
  1
  0
 -3
  2
```
"""
ScrewToAxis(q::Array, s::Array, h::Number) = vcat(s, linalg.cross(q, s) + h * s)

"""
    AxisAng6(expc6)

Converts a 6-vector of exponential coordinates into screw axis-angle form.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> AxisAng6([1, 0, 0, 1, 2, 3])
([1.0, 0.0, 0.0, 1.0, 2.0, 3.0], 1.0)
```
"""
function AxisAng6(expc6::Array)
    θ = linalg.norm(expc6[1:3])
    if NearZero(θ)
        θ = linalg.norm(expc6[3:6])
    end
    expc6 / θ, θ
end

"""
    MatrixExp6(se3mat)

Computes the matrix exponential of an se3 representation of exponential coordinates.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> MatrixExp6([0 0 0 0; 0 0 -1.57079632 2.35619449; 0 1.57079632 0 2.35619449; 0 0 0 0])
4×4 Array{Float64,2}:
 1.0  0.0         0.0        0.0
 0.0  6.7949e-9  -1.0        1.01923e-8
 0.0  1.0         6.7949e-9  3.0
 0.0  0.0         0.0        1.0
```
"""
function MatrixExp6(se3mat::Array)
    omgtheta = so3ToVec(se3mat[1:3, 1:3])
    if NearZero(linalg.norm(omgtheta))
        return vcat(hcat(linalg.I, se3mat[1:3, 4]), [0 0 0 1])
    else
        θ = AxisAng3(omgtheta)[2]
        omgmat = se3mat[1:3, 1:3] / θ
        return vcat(hcat(MatrixExp3(se3mat[1:3, 1:3]),
                         (linalg.I * θ +
                          (1 - cos(θ)) * omgmat +
                          (θ - sin(θ)) * omgmat * omgmat) *
                         se3mat[1:3, 4] / θ),
                    [0 0 0 1])
    end
end

"""
    MatrixLog6(T)

Computes the matrix logarithm of a homogeneous transformation matrix.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> MatrixLog6([1 0 0 0; 0 0 -1 0; 0 1 0 3; 0 0 0 1])
4×4 Array{Float64,2}:
 0.0  0.0      0.0     0.0
 0.0  0.0     -1.5708  2.35619
 0.0  1.5708   0.0     2.35619
 0.0  0.0      0.0     0.0
```
"""
function MatrixLog6(T::Array)
    R, p = TransToRp(T)
    omgmat = MatrixLog3(R)
    if omgmat == zeros(3, 3)
        return vcat(hcat(zeros(3, 3), T[1:3, 4]), [0 0 0 0])
    else
        θ = acos((linalg.tr(R) - 1) / 2)
        return vcat(hcat(omgmat,
                         (linalg.I - omgmat / 2 +
                          (1 / θ - 1 / tan(θ / 2) / 2) *
                          omgmat * omgmat / θ) * T[1:3, 4]),
                    [0 0 0 0])
    end
end

"""
    ProjectToSO3(mat)

Returns a projection of mat into SO(3).

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> ProjectToSO3([0.675 0.150  0.720; 0.370 0.771 -0.511; -0.630 0.619  0.472])
3×3 Array{Float64,2}:
  0.679011  0.148945   0.718859
  0.373207  0.773196  -0.512723
 -0.632187  0.616428   0.469421
```
"""
function ProjectToSO3(mat::Array)
    F  = linalg.svd(mat)
    R = F.U * F.Vt
    if linalg.det(R) < 0
        # In this case the result may be far from mat.
        # Hmm, I think this needs to be double-checked...
        R[:, Int(F.S[3])] = -R[:, Int(F.S[3])]
    end
    return R
end

"""
    ProjectToSE3(mat)

Returns a projection of mat into SE(3).

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> ProjectToSE3([0.675 0.150 0.720 1.2; 0.370 0.771 -0.511 5.4; -0.630 0.619 0.472 3.6; 0.003 0.002 0.010 0.9])
4×4 Array{Float64,2}:
  0.679011  0.148945   0.718859  1.2
  0.373207  0.773196  -0.512723  5.4
 -0.632187  0.616428   0.469421  3.6
  0.0       0.0        0.0       1.0
```
"""
ProjectToSE3(mat::Array) = RpToTrans(ProjectToSO3(mat[1:3, 1:3]), mat[1:3, 4])

"""
    DistanceToSO3(mat)

Returns the Frobenius norm to describe the distance of mat from the SO(3) manifold.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> DistanceToSO3([1.0 0.0 0.0; 0.0 0.1 -0.95; 0.0 1.0 0.1])
0.08835298523536149
```
"""
DistanceToSO3(mat::Array) = linalg.det(mat) > 0 ? linalg.norm(mat'mat - linalg.I) : 1e+9

"""
    DistanceToSE3(mat)

Returns the Frobenius norm to describe the distance of mat from the SE(3) manifold.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> DistanceToSE3([1.0 0.0 0.0 1.2; 0.0 0.1 -0.95 1.5; 0.0 1.0 0.1 -0.9; 0.0 0.0 0.1 0.98])
0.13493053768513638
```
"""
function DistanceToSE3(mat::Array)
    matR = mat[1:3, 1:3]
    if linalg.det(matR) > 0
        linalg.norm(hcat(vcat(matR'matR, zeros(1, 3)), mat[4, :]) - linalg.I)
    else
        1e+9
    end
end

"""
    TestIfSO3(mat)

Returns true if mat is close to or on the manifold SO(3).

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> TestIfSO3([1.0 0.0 0.0; 0.0 0.1 -0.95; 0.0 1.0 0.1])
false
```
"""
TestIfSO3(mat::Array) = abs(DistanceToSO3(mat)) < 1e-3

"""
    TestIfSE3(mat)

Returns true if mat is close to or on the manifold SE(3).

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> TestIfSE3([1.0 0.0 0.0 1.2; 0.0 0.1 -0.95 1.5; 0.0 1.0 0.1 -0.9; 0.0 0.0 0.1 0.98])
false
```
"""
TestIfSE3(mat::Array) = abs(DistanceToSE3(mat)) < 1e-3

# """
# *** CHAPTER 4: FORWARD KINEMATICS ***
# """

"""
    FKinBody(M, Blist, thetalist)

Computes forward kinematics in the body frame for an open chain robot.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> M = [ -1  0  0  0 ;
              0  1  0  6 ;
              0  0 -1  2 ;
              0  0  0  1 ];

julia> Blist = [  0  0 -1  2  0  0   ;
                  0  0  0  0  1  0   ;
                  0  0  1  0  0  0.1 ]';

julia> thetalist = [ π/2, 3, π ];

julia> FKinBody(M, Blist, thetalist)
4×4 Array{Float64,2}:
 -1.14424e-17  1.0           0.0  -5.0
  1.0          1.14424e-17   0.0   4.0
  0.0          0.0          -1.0   1.68584
  0.0          0.0           0.0   1.0
```
"""
function FKinBody(M::AbstractMatrix, Blist::AbstractMatrix, thetalist::Array)
    for i = 1:length(thetalist)
        M *= MatrixExp6(VecTose3(Blist[:, i] * thetalist[i]))
    end
    M
end

"""
    FKinSpace(M, Slist, thetalist)

Computes forward kinematics in the space frame for an open chain robot.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> M = [ -1  0  0  0 ;
              0  1  0  6 ;
              0  0 -1  2 ;
              0  0  0  1 ];

julia> Slist = [  0  0  1  4  0  0   ;
                  0  0  0  0  1  0   ;
                  0  0 -1 -6  0 -0.1 ]';

julia> thetalist = [ π/2, 3, π ];

julia> FKinSpace(M, Slist, thetalist)
4×4 Array{Float64,2}:
 -1.14424e-17  1.0           0.0  -5.0
  1.0          1.14424e-17   0.0   4.0
  0.0          0.0          -1.0   1.68584
  0.0          0.0           0.0   1.0
```
"""
function FKinSpace(M::AbstractMatrix, Slist::AbstractMatrix, thetalist::Array)
    for i = length(thetalist):-1:1
        M = MatrixExp6(VecTose3(Slist[:, i] * thetalist[i])) * M
    end
    M
end

# """
# *** CHAPTER 5: VELOCITY KINEMATICS AND STATICS ***
# """

"""
    JacobianBody(Blist, thetalist)

Computes the body Jacobian for an open chain robot.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> Blist = [0 0 1   0 0.2 0.2;
                1 0 0   2   0   3;
                0 1 0   0   2   1;
                1 0 0 0.2 0.3 0.4]';

julia> thetalist = [0.2, 1.1, 0.1, 1.2];

julia> JacobianBody(Blist, thetalist)
6×4 Array{Float64,2}:
 -0.0452841  0.995004    0.0       1.0
  0.743593   0.0930486   0.362358  0.0
 -0.667097   0.0361754  -0.932039  0.0
  2.32586    1.66809     0.564108  0.2
 -1.44321    2.94561     1.43307   0.3
 -2.0664     1.82882    -1.58869   0.4
```
"""
function JacobianBody(Blist::AbstractMatrix, thetalist::Array)
    T = linalg.I
    Jb = copy(Blist)
    for i = length(thetalist)-1:-1:1
        T *= MatrixExp6(VecTose3(Blist[:, i+1] * -thetalist[i+1]))
        Jb[:, i] = Adjoint(T) * Blist[:, i]
    end
    Jb
end

"""
    JacobianSpace(Slist, thetalist)

Computes the space Jacobian for an open chain robot.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> Slist = [0 0 1   0 0.2 0.2;
                1 0 0   2   0   3;
                0 1 0   0   2   1;
                1 0 0 0.2 0.3 0.4]';

julia> thetalist = [0.2, 1.1, 0.1, 1.2];

julia> JacobianSpace(Slist, thetalist)
6×4 Array{Float64,2}:
 0.0  0.980067  -0.0901156   0.957494
 0.0  0.198669   0.444554    0.284876
 1.0  0.0        0.891207   -0.0452841
 0.0  1.95219   -2.21635    -0.511615
 0.2  0.436541  -2.43713     2.77536
 0.2  2.96027    3.23573     2.22512
```
"""
function JacobianSpace(Slist::AbstractMatrix, thetalist::Array)
    T = linalg.I
    Js = copy(Slist)
    for i = 2:length(thetalist)
        T *= MatrixExp6(VecTose3(Slist[:, i - 1] * thetalist[i - 1]))
        Js[:, i] = Adjoint(T) * Slist[:, i]
    end
    Js
end

# """
# *** CHAPTER 6: INVERSE KINEMATICS ***
# """

"""
    IKinBody(Blist, M, T, thetalist0, eomg, ev)

Computes inverse kinematics in the body frame for an open chain robot.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> Blist = [  0  0 -1  2  0  0   ;
                  0  0  0  0  1  0   ;
                  0  0  1  0  0  0.1 ]';

julia> M = [ -1  0  0  0 ;
              0  1  0  6 ;
              0  0 -1  2 ;
              0  0  0  1 ];

julia> T = [ 0  1  0     -5 ;
             1  0  0      4 ;
             0  0 -1 1.6858 ;
             0  0  0      1 ];

julia> thetalist0 = [1.5, 2.5, 3];

julia> eomg, ev = 0.01, 0.001;

julia> IKinBody(Blist, M, T, thetalist0, eomg, ev)
([1.57074, 2.99967, 3.14154], true)
```
"""
function IKinBody(Blist::AbstractMatrix,
                      M::AbstractMatrix,
                      T::AbstractMatrix,
             thetalist0::Array,
                   eomg::Number,
                     ev::Number)
    thetalist = copy(thetalist0)
    i = 0
    maxiterations = 20
    Vb = se3ToVec(MatrixLog6(TransInv(FKinBody(M, Blist, thetalist)) * T))
    err = linalg.norm(Vb[1:3]) > eomg || linalg.norm(Vb[4:6]) > ev
    while err && i < maxiterations
        thetalist += linalg.pinv(JacobianBody(Blist, thetalist)) * Vb
        i += 1
        Vb = se3ToVec(MatrixLog6(TransInv(FKinBody(M, Blist, thetalist)) * T))
        err = linalg.norm(Vb[1:3]) > eomg || linalg.norm(Vb[4:6]) > ev
    end
    return thetalist, !err
end

"""
    IKinSpace(Slist, M, T, thetalist0, eomg, ev)

Computes inverse kinematics in the space frame for an open chain robot.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> Slist = [  0  0  1  4  0  0   ;
                  0  0  0  0  1  0   ;
                  0  0 -1 -6  0 -0.1 ]';

julia> M = [ -1  0  0  0 ;
              0  1  0  6 ;
              0  0 -1  2 ;
              0  0  0  1 ];

julia> T = [ 0  1  0     -5 ;
             1  0  0      4 ;
             0  0 -1 1.6858 ;
             0  0  0      1 ];

julia> thetalist0 = [1.5, 2.5, 3];

julia> eomg, ev = 0.01, 0.001;

julia> IKinSpace(Slist, M, T, thetalist0, eomg, ev)
([1.57074, 2.99966, 3.14153], true)
```
"""
function IKinSpace(Slist::AbstractMatrix,
                       M::AbstractMatrix,
                       T::AbstractMatrix,
              thetalist0::Array,
                    eomg::Number,
                      ev::Number)
    thetalist = copy(thetalist0)
    i = 0
    maxiterations = 20
    Tsb = FKinSpace(M, Slist, thetalist)
    Vs = Adjoint(Tsb) * se3ToVec(MatrixLog6(TransInv(Tsb) * T))
    err = linalg.norm(Vs[1:3]) > eomg || linalg.norm(Vs[4:6]) > ev
    while err && i < maxiterations
        thetalist += linalg.pinv(JacobianSpace(Slist, thetalist)) * Vs
        i += 1
        Tsb = FKinSpace(M, Slist, thetalist)
        Vs = Adjoint(Tsb) * se3ToVec(MatrixLog6(TransInv(Tsb) * T))
        err = linalg.norm(Vs[1:3]) > eomg || linalg.norm(Vs[4:6]) > ev
    end
    thetalist, !err
end

# """
# *** CHAPTER 8: DYNAMICS OF OPEN CHAINS ***
# """

"""
    ad(V)

Calculate the 6x6 matrix [adV] of the given 6-vector.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> ad([1, 2, 3, 4, 5, 6])
6×6 Array{Float64,2}:
  0.0  -3.0   2.0   0.0   0.0   0.0
  3.0   0.0  -1.0   0.0   0.0   0.0
 -2.0   1.0   0.0   0.0   0.0   0.0
  0.0  -6.0   5.0   0.0  -3.0   2.0
  6.0   0.0  -4.0   3.0   0.0  -1.0
 -5.0   4.0   0.0  -2.0   1.0   0.0
```
"""
function ad(V::Array)
    omgmat = VecToso3(V[1:3])
    vcat(hcat(omgmat, zeros(3, 3)),
         hcat(VecToso3(V[4:6]), omgmat))
end

"""
    InverseDynamics(thetalist, dthetalist, ddthetalist, g, Ftip, Mlist, Glist, Slist)

Computes inverse dynamics in the space frame for an open chain robot.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> InverseDynamics(thetalist, dthetalist, ddthetalist, g, Ftip, Mlist, Glist, Slist)
3-element Array{Float64,1}:
  74.69616155287451
 -33.06766015851458
  -3.230573137901424
```
"""
function InverseDynamics(thetalist::Array,
                        dthetalist::Array,
                       ddthetalist::Array,
                                 g::Array,
                              Ftip::Array,
                             Mlist::Array,
                             Glist::Array,
                             Slist::AbstractMatrix)
    n = length(thetalist)
    Mi = linalg.I
    Ai = zeros(eltype(thetalist), 6, n)
    AdTi = Array{Array{eltype(thetalist), 2}}(undef, n + 1)
    Vi = zeros(eltype(thetalist), 6, n + 1)
    Vdi = zeros(eltype(thetalist), 6, n + 1)
    Vdi[:, 1] = vcat(zeros(eltype(thetalist), 3), -g)
    AdTi[n+1] = Adjoint(TransInv(Mlist[n+1]))
    Fi = copy(Ftip)
    taulist = zeros(eltype(thetalist), n)

    for i = 1:n
        Mi *= Mlist[i]
        Ai[:, i] = Adjoint(TransInv(Mi)) * Slist[:, i]
        AdTi[i] = Adjoint(MatrixExp6(VecTose3(Ai[:, i] * -thetalist[i])) *
                          TransInv(Mlist[i]))
        Vi[:, i + 1] = AdTi[i] * Vi[:,i] + Ai[:, i] * dthetalist[i]
        Vdi[:, i + 1] = AdTi[i] * Vdi[:, i] + Ai[:, i] * ddthetalist[i] +
                        ad(Vi[:, i + 1]) * Ai[:, i] * dthetalist[i]
    end

    for i = n:-1:1
        Fi = AdTi[i + 1]' * Fi + Glist[i] * Vdi[:, i + 1] -
             ad(Vi[:, i + 1])' * Glist[i] * Vi[:, i + 1]
        taulist[i] = Fi' * Ai[:, i]
    end

    return taulist
end

"""
    MassMatrix(thetalist, Mlist, Glist, Slist)

Computes the mass matrix of an open chain robot based on the given configuration.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> MassMatrix(thetalist, Mlist, Glist, Slist)
3×3 Array{Float64,2}:
 22.5433      -0.307147  -0.00718426
 -0.307147     1.96851    0.432157
 -0.00718426   0.432157   0.191631
```
"""
function MassMatrix(thetalist::Array,
                        Mlist::Array,
                        Glist::Array,
                        Slist::AbstractMatrix)
    n = length(thetalist)
    M = zeros(n, n)

    for i = 1:n
        ddthetalist = zeros(n)
        ddthetalist[i] = 1
        M[:, i] = InverseDynamics(thetalist, zeros(n), ddthetalist, zeros(3),
                                  zeros(6), Mlist, Glist, Slist)
    end

    return M
end

"""
    VelQuadraticForces(thetalist, dthetalist, Mlist, Glist, Slist)

Computes the Coriolis and centripetal terms in the inverse dynamics of an open chain robot.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> VelQuadraticForces(thetalist, dthetalist, Mlist, Glist, Slist)
3-element Array{Float64,1}:
  0.26453118054501235
 -0.0550515682891655
 -0.006891320068248911
```
"""
function VelQuadraticForces(thetalist::Array,
                           dthetalist::Array,
                                Mlist::Array,
                                Glist::Array,
                                Slist::AbstractMatrix)
    InverseDynamics(thetalist, dthetalist, zeros(length(thetalist)),
                    zeros(3), zeros(6), Mlist, Glist, Slist)
end

"""
    GravityForces(thetalist, g, Mlist, Glist, Slist)

Computes the joint forces/torques an open chain robot requires to overcome gravity at its configuration.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> GravityForces(thetalist, g, Mlist, Glist, Slist)
3-element Array{Float64,1}:
  28.40331261821983
 -37.64094817177068
  -5.4415891999683605
```
"""
function GravityForces(thetalist::Array,
                               g::Array,
                           Mlist::Array,
                           Glist::Array,
                           Slist::AbstractMatrix)
    n = length(thetalist)
    InverseDynamics(thetalist, zeros(n), zeros(n), g, zeros(6), Mlist, Glist, Slist)
end

"""
    EndEffectorForces(thetalist, Ftip, Mlist, Glist, Slist)

Computes the joint forces/torques an open chain robot requires only to create the end-effector force `Ftip`.

# Arguments
- `thetalist`: the ``n``-vector of joint variables.
- `Ftip`: the spatial force applied by the end-effector expressed in frame `{n+1}`.
- `Mlist`: the list of link frames `i` relative to `i-1` at the home position.
- `Glist`: the spatial inertia matrices `Gi` of the links.
- `Slist`: the screw axes `Si` of the joints in a space frame, in the format of a matrix with axes as the columns.

Returns the joint forces and torques required only to create the end-effector force `Ftip`.
This function calls InverseDynamics with `g = 0`, `dthetalist = 0`, and `ddthetalist = 0`.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> import LinearAlgebra

julia> const linalg = LinearAlgebra;

julia> thetalist = [0.1, 0.1, 0.1]
3-element Array{Float64,1}:
 0.1
 0.1
 0.1

julia> Ftip = [1, 1, 1, 1, 1, 1]
6-element Array{Int64,1}:
 1
 1
 1
 1
 1
 1

julia> M01 = [1 0 0        0;
              0 1 0        0;
              0 0 1 0.089159;
              0 0 0        1]
4×4 Array{Float64,2}:
 1.0  0.0  0.0  0.0
 0.0  1.0  0.0  0.0
 0.0  0.0  1.0  0.089159
 0.0  0.0  0.0  1.0

julia> M12 = [ 0 0 1    0.28;
               0 1 0 0.13585;
              -1 0 0       0;
               0 0 0       1]
4×4 Array{Float64,2}:
  0.0  0.0  1.0  0.28
  0.0  1.0  0.0  0.13585
 -1.0  0.0  0.0  0.0
  0.0  0.0  0.0  1.0

julia> M23 = [1 0 0       0;
              0 1 0 -0.1197;
              0 0 1   0.395;
              0 0 0       1]
4×4 Array{Float64,2}:
 1.0  0.0  0.0   0.0
 0.0  1.0  0.0  -0.1197
 0.0  0.0  1.0   0.395
 0.0  0.0  0.0   1.0

julia> M34 = [1 0 0       0;
              0 1 0       0;
              0 0 1 0.14225;
              0 0 0       1]
4×4 Array{Float64,2}:
 1.0  0.0  0.0  0.0
 0.0  1.0  0.0  0.0
 0.0  0.0  1.0  0.14225
 0.0  0.0  0.0  1.0

julia> Mlist = [M01, M12, M23, M34]
4-element Array{Array{Float64,2},1}:
 [1.0 0.0 0.0 0.0; 0.0 1.0 0.0 0.0; 0.0 0.0 1.0 0.089159; 0.0 0.0 0.0 1.0]
 [0.0 0.0 1.0 0.28; 0.0 1.0 0.0 0.13585; -1.0 0.0 0.0 0.0; 0.0 0.0 0.0 1.0]
 [1.0 0.0 0.0 0.0; 0.0 1.0 0.0 -0.1197; 0.0 0.0 1.0 0.395; 0.0 0.0 0.0 1.0]
 [1.0 0.0 0.0 0.0; 0.0 1.0 0.0 0.0; 0.0 0.0 1.0 0.14225; 0.0 0.0 0.0 1.0]

julia> G1 = linalg.Diagonal([0.010267, 0.010267, 0.00666, 3.7, 3.7, 3.7])
6×6 LinearAlgebra.Diagonal{Float64,Array{Float64,1}}:
 0.010267   ⋅         ⋅        ⋅    ⋅    ⋅
  ⋅        0.010267   ⋅        ⋅    ⋅    ⋅
  ⋅         ⋅        0.00666   ⋅    ⋅    ⋅
  ⋅         ⋅         ⋅       3.7   ⋅    ⋅
  ⋅         ⋅         ⋅        ⋅   3.7   ⋅
  ⋅         ⋅         ⋅        ⋅    ⋅   3.7

julia> G2 = linalg.Diagonal([0.22689, 0.22689, 0.0151074, 8.393, 8.393, 8.393])
6×6 LinearAlgebra.Diagonal{Float64,Array{Float64,1}}:
 0.22689   ⋅        ⋅          ⋅      ⋅      ⋅
  ⋅       0.22689   ⋅          ⋅      ⋅      ⋅
  ⋅        ⋅       0.0151074   ⋅      ⋅      ⋅
  ⋅        ⋅        ⋅         8.393   ⋅      ⋅
  ⋅        ⋅        ⋅          ⋅     8.393   ⋅
  ⋅        ⋅        ⋅          ⋅      ⋅     8.393

julia> G3 = linalg.Diagonal([0.0494433, 0.0494433, 0.004095, 2.275, 2.275, 2.275])
6×6 LinearAlgebra.Diagonal{Float64,Array{Float64,1}}:
 0.0494433   ⋅          ⋅         ⋅      ⋅      ⋅
  ⋅         0.0494433   ⋅         ⋅      ⋅      ⋅
  ⋅          ⋅         0.004095   ⋅      ⋅      ⋅
  ⋅          ⋅          ⋅        2.275   ⋅      ⋅
  ⋅          ⋅          ⋅         ⋅     2.275   ⋅
  ⋅          ⋅          ⋅         ⋅      ⋅     2.275

julia> Glist = [G1, G2, G3]
3-element Array{LinearAlgebra.Diagonal{Float64,Array{Float64,1}},1}:
 [0.010267 0.0 … 0.0 0.0; 0.0 0.010267 … 0.0 0.0; … ; 0.0 0.0 … 3.7 0.0; 0.0 0.0 … 0.0 3.7]
 [0.22689 0.0 … 0.0 0.0; 0.0 0.22689 … 0.0 0.0; … ; 0.0 0.0 … 8.393 0.0; 0.0 0.0 … 0.0 8.393]
 [0.0494433 0.0 … 0.0 0.0; 0.0 0.0494433 … 0.0 0.0; … ; 0.0 0.0 … 2.275 0.0; 0.0 0.0 … 0.0 2.275]

julia> Slist = [ 1  0  1      0  1      0;
                 0  1  0 -0.089  0      0;
                 0  1  0 -0.089  0  0.425]'
6×3 LinearAlgebra.Adjoint{Float64,Array{Float64,2}}:
 1.0   0.0     0.0
 0.0   1.0     1.0
 1.0   0.0     0.0
 0.0  -0.089  -0.089
 1.0   0.0     0.0
 0.0   0.0     0.425

julia> EndEffectorForces(thetalist, Ftip, Mlist, Glist, Slist)
3-element Array{Float64,1}:
 1.4095460782639782
 1.8577149723180628
 1.392409
```
"""
function EndEffectorForces(thetalist::Array,
                                Ftip::Array,
                               Mlist::Array,
                               Glist::Array,
                               Slist::AbstractMatrix)
    n = length(thetalist)
    InverseDynamics(thetalist, zeros(n), zeros(n), zeros(3), Ftip, Mlist, Glist, Slist)
end

"""
    ForwardDynamics(thetalist, dthetalist, taulist, g, Ftip, Mlist, Glist, Slist)

Computes forward dynamics in the space frame for an open chain robot.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> ForwardDynamics(thetalist, dthetalist, taulist, g, Ftip, Mlist, Glist, Slist)
3-element Array{Float64,1}:
  -0.9739290670855626
  25.584667840340558
 -32.91499212478149
```
"""
function ForwardDynamics(thetalist::Array,
                        dthetalist::Array,
                           taulist::Array,
                                 g::Array,
                              Ftip::Array,
                             Mlist::Array,
                             Glist::Array,
                             Slist::AbstractMatrix)
    linalg.inv(MassMatrix(thetalist, Mlist, Glist, Slist)) *
    (taulist - VelQuadraticForces(thetalist, dthetalist, Mlist, Glist, Slist)
             - GravityForces(thetalist, g, Mlist, Glist, Slist)
             - EndEffectorForces(thetalist, Ftip, Mlist, Glist, Slist))
end

"""
    EulerStep(thetalist, dthetalist, ddthetalist, dt)

Compute the joint angles and velocities at the next timestep using first order Euler integration.

# Arguments
- `thetalist`: the ``n``-vector of joint variables.
- `dthetalist`: the ``n``-vector of joint rates.
- `ddthetalist`: the ``n``-vector of joint accelerations.
- `dt`: the timestep delta t.

# Return
- `thetalistNext`: the vector of joint variables after `dt` from first order Euler integration.
- `dthetalistNext`: the vector of joint rates after `dt` from first order Euler integration.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> EulerStep([0.1, 0.1, 0.1], [0.1, 0.2, 0.3], [2, 1.5, 1], 0.1)
([0.11, 0.12, 0.13], [0.3, 0.35, 0.4])
```
"""
function EulerStep(thetalist::Array, dthetalist::Array, ddthetalist::Array, dt::Number)
    thetalist + dt * dthetalist, dthetalist + dt * ddthetalist
end

"""
    InverseDynamicsTrajectory(thetamat, dthetamat, ddthetamat, g, Ftipmat, Mlist, Glist, Slist)

Calculates the joint forces/torques required to move the serial chain along the given trajectory using inverse dynamics.
"""
function InverseDynamicsTrajectory(thetamat::Array,
                                  dthetamat::Array,
                                 ddthetamat::Array,
                                          g::Array,
                                    Ftipmat::Array,
                                      Mlist::Array,
                                      Glist::Array,
                                      Slist::AbstractMatrix)
    thetamat = thetamat'
    dthetamat = dthetamat'
    ddthetamat = ddthetamat'
    Ftipmat = Ftipmat'
    taumat = copy(thetamat)

    for i = 1:size(thetamat, 2)
        taumat[:, i] = InverseDynamics(thetamat[:, i], dthetamat[:, i],
                                       ddthetamat[:, i], g, Ftipmat[:, i],
                                       Mlist, Glist, Slist)
    end

    taumat'
end

"""
    ForwardDynamicsTrajectory(thetalist, dthetalist, taumat, g, Ftipmat, Mlist, Glist, Slist, dt, intRes)

Simulates the motion of a serial chain given an open-loop history of joint forces/torques.
"""
function ForwardDynamicsTrajectory(thetalist::Array,
                                  dthetalist::Array,
                                      taumat::AbstractMatrix,
                                           g::Array,
                                     Ftipmat::Array,
                                       Mlist::Array,
                                       Glist::Array,
                                       Slist::AbstractMatrix,
                                          dt::Number,
                                      intRes::Number)
    taumat = taumat'
    Ftipmat = Ftipmat'
    thetamat = copy(taumat)
    thetamat[:, 1] = thetalist
    dthetamat = copy(taumat)
    dthetamat[:, 1] = dthetalist

    for i = 1:size(taumat, 2)-1
        for j = 1:intRes
            ddthetalist = ForwardDynamics(thetalist, dthetalist, taumat[:, i], g, Ftipmat[:, i], Mlist, Glist, Slist)
            thetalist, dthetalist = EulerStep(thetalist, dthetalist, ddthetalist, 1.0 * dt / intRes)
        end

        thetamat[:, i + 1] = thetalist
        dthetamat[:, i + 1] = dthetalist
    end

    thetamat = thetamat'
    dthetamat = dthetamat'
    return thetamat, dthetamat
end

# """
# *** CHAPTER 9: TRAJECTORY GENERATION ***
# """

"""
    CubicTimeScaling(Tf, t)

Computes s(t) for a cubic time scaling.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> CubicTimeScaling(2, 0.6)
0.21600000000000003
```
"""
CubicTimeScaling(Tf::Number, t::Number) = 3(t / Tf)^2 - 2(t / Tf)^3

"""
    QuinticTimeScaling(Tf, t)

Computes s(t) for a quintic time scaling.

# Examples
```jldoctest; setup = :(using ModernRoboticsBook)
julia> QuinticTimeScaling(2, 0.6)
0.16308
```
"""
QuinticTimeScaling(Tf::Number, t::Number) = 10(t / Tf)^3 - 15(t / Tf)^4 + 6(t / Tf)^5

"""
    JointTrajectory(thetastart, thetaend, Tf, N, method)

Computes a straight-line trajectory in joint space.
"""
function JointTrajectory(thetastart::Array, thetaend::Array, Tf::Number, N::Integer, method::Integer)
    timegap = Tf / (N - 1)
    traj = zeros(length(thetastart), N)

    for i = 1:N
        if method == 3
            s = CubicTimeScaling(Tf, timegap * (i - 1))
        else
            s = QuinticTimeScaling(Tf, timegap * (i - 1))
        end

        traj[:, i] = s * thetaend + (1 - s) * thetastart
    end

    traj'
end

"""
    ScrewTrajectory(Xstart, Xend, Tf, N, method)

Computes a trajectory as a list of N SE(3) matrices corresponding to the screw motion about a space screw axis.
"""
function ScrewTrajectory(Xstart::Array, Xend::Array, Tf::Number, N::Integer, method::Integer)
    timegap = Tf / (N - 1)
    traj = Array{Array{Float64, 2}}(undef, N)

    for i = 1:N
        if method == 3
            s = CubicTimeScaling(Tf, timegap * (i - 1))
        else
            s = QuinticTimeScaling(Tf, timegap * (i - 1))
        end

        traj[i] = Xstart * MatrixExp6(MatrixLog6(TransInv(Xstart) * Xend) * s)
    end

    return traj
end

"""
    CartesianTrajectory(Xstart, Xend, Tf, N, method)

Computes a trajectory as a list of N SE(3) matrices corresponding to the origin of the end-effector frame following a straight line.
"""
function CartesianTrajectory(Xstart::Array, Xend::Array, Tf::Number, N::Integer, method::Integer)
    timegap = Tf / (N - 1)
    traj = Array{Array{Float64, 2}}(undef, N)

    Rstart, pstart = TransToRp(Xstart)
    Rend, pend = TransToRp(Xend)

    for i = 1:N
        if method == 3
            s = CubicTimeScaling(Tf, timegap * (i - 1))
        else
            s = QuinticTimeScaling(Tf, timegap * (i - 1))
        end

        traj[i] = vcat(hcat(Rstart * MatrixExp3(MatrixLog3(Rstart' * Rend) * s), s * pend + (1 - s) * pstart), [0 0 0 1])
    end

    return traj
end

# """
# *** CHAPTER 11: ROBOT CONTROL ***
# """

"""
    ComputedTorque(thetalist, dthetalist, eint, g, Mlist, Glist, Slist, thetalistd, dthetalistd, ddthetalistd, Kp, Ki, Kd)

Computes the joint control torques at a particular time instant.
"""
function ComputedTorque(thetalist::Array,
                       dthetalist::Array,
                             eint::Array,
                                g::Array,
                            Mlist::Array,
                            Glist::Array,
                            Slist::AbstractMatrix,
                       thetalistd::Array,
                      dthetalistd::Array,
                     ddthetalistd::Array,
                               Kp::Number,
                               Ki::Number,
                               Kd::Number)
    e = thetalistd - thetalist
    MassMatrix(thetalist, Mlist, Glist, Slist) *
    (Kp * e + Ki * (eint + e) + Kd * (dthetalistd - dthetalist)) +
    InverseDynamics(thetalist, dthetalist, ddthetalistd, g,
                    zeros(6), Mlist, Glist, Slist)
end

"""
    SimulateControl(thetalist, dthetalist, g, Ftipmat, Mlist, Glist,
                    Slist, thetamatd, dthetamatd, ddthetamatd, gtilde,
                    Mtildelist, Gtildelist, Kp, Ki, Kd, dt, intRes)

Simulates the computed torque controller over a given desired trajectory.
"""
function SimulateControl(thetalist::Array,
                        dthetalist::Array,
                                 g::Array,
                           Ftipmat::Array,
                             Mlist::Array,
                             Glist::Array,
                             Slist::AbstractMatrix,
                         thetamatd::Array,
                        dthetamatd::Array,
                       ddthetamatd::Array,
                            gtilde::Array,
                        Mtildelist::Array,
                        Gtildelist::Array,
                                Kp::Number,
                                Ki::Number,
                                Kd::Number,
                                dt::Number,
                            intRes::Number)
    Ftipmat = Ftipmat'
    thetamatd = thetamatd'
    dthetamatd = dthetamatd'
    ddthetamatd = ddthetamatd'

    m, n = size(thetamatd)

    thetacurrent = copy(thetalist)
    dthetacurrent = copy(dthetalist)

    eint = reshape(zeros(m, 1), (m,))
    taumat = zeros(size(thetamatd))
    thetamat = zeros(size(thetamatd))

    for i = 1:n
        taulist = ComputedTorque(thetacurrent, dthetacurrent, eint, gtilde, Mtildelist, Gtildelist, Slist, thetamatd[:, i], dthetamatd[:, i], ddthetamatd[:, i], Kp, Ki, Kd)

        for j = 1:intRes
            ddthetalist = ForwardDynamics(thetacurrent, dthetacurrent, taulist, g, Ftipmat[:, i], Mlist, Glist, Slist)
            thetacurrent, dthetacurrent = EulerStep(thetacurrent, dthetacurrent, ddthetalist, dt / intRes)
        end

        taumat[:, i] = taulist
        thetamat[:, i] = thetacurrent

        eint += dt * (thetamatd[:, i] - thetacurrent)
    end

    taumat', thetamat'
end

end # module

max_stars_repo_name: veysiadn/AdvancedRoboticsJulia

max_issues_repo_name: veysiadn/AdvancedRoboticsJulia

max_forks_repo_name: veysiadn/AdvancedRoboticsJulia


text:
@testset "using" begin
    @test_broken isempty(lintstr("""
    module TmpTestBase
        export foobar
        foobar(x) = x
    end

    module Test
        using ..TmpTestBase
        g(x) = foobar(x)
    end
    """))
end

ext: jl


text:
itergeneration(itr, gen) = itr

"""
    RepeatPartitionIterator
    RepeatPartitionIterator(base, nrep)

Iteratates over iterators of a subset of size `nrep` elements in `base`.

Generally useful for training all models in a population with the same data in each evolution epoch.

Tailored for situations where iterating over models is more expensive than iterating over data, for
example if candidates are stored in host RAM or on disk and needs to be transferred to the GPU for training.

Example training loop:
```julia

for partiter in iter

    for model in population
        train!(model, partiter)
    end

    evolvepopulation(population)
end
```

"""
struct RepeatPartitionIterator{T, VS}
    base::Iterators.Stateful{T, VS}
    ntake::Int
end
RepeatPartitionIterator(base, nrep) = RepeatPartitionIterator(Iterators.Stateful(base), nrep)
RepeatPartitionIterator(base::Iterators.Stateful, nrep) = RepeatPartitionIterator(base, nrep)

function Base.iterate(itr::RepeatPartitionIterator, reset=true)
    if reset
        Iterators.reset!(itr.base, itr.base.itr)
    end
    Base.IteratorSize(itr) !== Base.IsInfinite() && length(itr) == 0 && return nothing
    return Iterators.take(RepeatStatefulIterator(itr.base), itr.ntake), false
end

Base.length(itr::RepeatPartitionIterator) = cld(length(itr.base), itr.ntake)
Base.eltype(itr::RepeatPartitionIterator) = eltype(itr.base)
Base.size(itr::RepeatPartitionIterator) = size(itr.base.itr)

Base.IteratorSize(itr::RepeatPartitionIterator) = Base.IteratorSize(itr.base.itr)
Base.IteratorEltype(itr::RepeatPartitionIterator) = Base.IteratorEltype(itr.base.itr)

struct RepeatStatefulIterator{T, VS}
    base::Iterators.Stateful{T, VS}
    start::VS
    taken::Int
end
RepeatStatefulIterator(base) = RepeatStatefulIterator(base, base.nextvalstate, base.taken)

function Base.iterate(itr::RepeatStatefulIterator, reset=true)
    if reset
        itr.base.nextvalstate = itr.start
        itr.base.taken = itr.taken
    end
    val, state = IterTools.@ifsomething iterate(itr.base)
    return val, false
end

Base.length(itr::RepeatStatefulIterator) = length(itr.base.itr) - itr.taken
Base.eltype(itr::RepeatStatefulIterator) = eltype(itr.base)
Base.size(itr::RepeatStatefulIterator) = size(itr.base.itr)

Base.IteratorSize(::Type{RepeatStatefulIterator{T,VS}}) where {T, VS} = Base.IteratorSize(Iterators.Stateful{T,VS})
Base.IteratorEltype(::Type{RepeatStatefulIterator{T,VS}}) where {T, VS} = Base.IteratorEltype(Iterators.Stateful{T,VS})

"""
    StatefulGenerationIter{T, VS}

Uses a `RepeatPartitionIterator` to ensure that the same iterator is returned for the same generation number.
"""
struct StatefulGenerationIter{I, T, VS}
    currgen::Ref{Int}
    curriter::Ref{I}
    iter::RepeatPartitionIterator{T, VS}
end
# TODO : This is a bit of cludge-on-cludge. Try to refactor someday to a more straighforward design, perhaps use LearnBase.getobs
StatefulGenerationIter(iter::RepeatPartitionIterator, gen=0) = StatefulGenerationIter(Ref(gen), Ref(first(iterate(iter))), iter)

function itergeneration(itr::StatefulGenerationIter, gen)
    if gen != itr.currgen[]
        itr.currgen[] = gen
        newiterstate = iterate(itr.iter, false)
        if newiterstate === nothing
            newiterstate = iterate(itr.iter)
        end
        itr.curriter[] = first(newiterstate)
    end
    return itr.curriter[]
end

"""
    SeedIterator
    SeedIterator(base; rng=rng_default, seed=rand(rng, UInt32))

Iterator which has the random seed of an `AbstractRNG` as state.

Calls `Random.seed!(rng, seed)` every iteration so that wrapped iterators which depend on `rng` will produce the same sequence.

Useful in conjunction with [`RepeatPartitionIterator`](@ref) and [`BatchIterator`](@ref) and/or random data augmentation so that all candidates in a generation are trained with identical data.
"""
struct SeedIterator{R <: AbstractRNG,T}
    rng::R
    seed::UInt32
    base::T
end
SeedIterator(base; rng=rng_default, seed=rand(rng, UInt32)) = SeedIterator(rng, UInt32(seed), base)

function Base.iterate(itr::SeedIterator)
    Random.seed!(itr.rng, itr.seed)
    val, state = IterTools.@ifsomething iterate(itr.base)
    return val, (itr.seed+1, state)
end

function Base.iterate(itr::SeedIterator, state)
    seed,basestate = state
    Random.seed!(itr.rng, seed)
    val, state = IterTools.@ifsomething iterate(itr.base, basestate)
    return val, (seed+1, state)
end

Base.length(itr::SeedIterator) = length(itr.base)
Base.eltype(itr::SeedIterator) = eltype(itr.base)
Base.size(itr::SeedIterator) = size(itr.base)

Base.IteratorSize(::Type{SeedIterator{R, T}}) where {R,T}= Base.IteratorSize(T)
Base.IteratorEltype(::Type{SeedIterator{R, T}}) where {R,T} = Base.IteratorEltype(T)

"""
    GpuIterator(itr)

Return an iterator which sends values from `itr` to the GPU.
"""
GpuIterator(itr) = Iterators.map(gpuitr, itr)
gpuitr(a) = Flux.gpu(a)
gpuitr(a::SubArray) = gpuitr(collect(a))
gpuitr(a::Tuple) = gpuitr.(a)

struct NoShuffle end
Random.shuffle(::NoShuffle, x) = x

"""
    BatchIterator{R, D}
    BatchIterator(data, batchsize; shuffle)

Return an iterator which iterates `batchsize` samples along the last dimension of `data`
or all elements of `data` if `data` is a `Tuple` (e.g `(features, labels)`).

Will shuffle examples if `shuffle` is `true` or an `AbstractRNG`. Shuffling will be
different each time iteration starts (subject to implementation of shuffle(rng,...)).
"""
struct BatchIterator{R, D}
    nobs::Int
    batchsize::Int
    rng::R
    data::D
end
BatchIterator(data::Union{AbstractArray, Singleton}, bs::Int; kwargs...) = BatchIterator(size(data)[end], bs, data; kwargs...)
function BatchIterator(data::Tuple, bs; kwargs...)
    @assert all(x -> size(x)[end] === size(data[1])[end], data) "Mistmatched batch dimensions! Got sizes $(size.(data))"
    BatchIterator(size(data[1])[end], bs, data; kwargs...)
end
function BatchIterator(nobs::Integer, batchsize::Integer, data; shuffle=false)
    return BatchIterator(nobs, batchsize, shufflerng(shuffle), data)
end
shufflerng(b::Bool) = b ? rng_default : NoShuffle()
shufflerng(rng) = rng

ndata(itr::BatchIterator) = ndata(itr.data)
ndata(data::Tuple) = length(data)
ndata(data::AbstractArray) = 1

function Base.iterate(itr::BatchIterator, inds = shuffle(itr.rng, 1:itr.nobs))
    isempty(inds) && return nothing
    return batch(itr.data, @view(inds[1:min(end, itr.batchsize)])), @view(inds[itr.batchsize+1:end])
end

Base.length(itr::BatchIterator) = cld(itr.nobs, itr.batchsize)
Base.size(itr::BatchIterator) = tuple(length(itr))

Base.IteratorEltype(::Type{BatchIterator{R,D}}) where {R, D} = Base.IteratorEltype(D)

batch(s::Singleton, inds) = batch(val(s), inds)
batch(b::Tuple, inds) = batch.(b, Ref(inds))
batch(a::AbstractArray{T,N}, inds) where {T,N} = selectdim(a, N, inds)

Base.show(io::IO, itr::BatchIterator{R}) where R = print(io, "BatchIterator(size=$(size(itr.data)), batchsize=$(itr.batchsize), shuffle=$(R !== NoShuffle))")
Base.show(io::IO, itr::BatchIterator{R, <:Tuple}) where R = print(io, "BatchIterator(size=$(size.(itr.data)), batchsize=$(itr.batchsize), shuffle=$(R !== NoShuffle))")

"""
    TimedIterator{F,A,I}
    TimedIterator(;timelimit, patience, timeoutaction, accumulate_timeouts, base)

Measures time between iterations and calls `timeoutaction()` if this time is longer than `timelimit` `patience` number of times.

Intended use is to quickly abort training of models which take very long time to train, typically also assigning them a very low fitness in case of a timeout.

By default, calling `timeoutaction()` will not stop the iteration as this would break otherwise convenient functions like `length`, `collect` and `map`. Let `timeoutaction()` return `TimedIteratorStop` to stop iteration.

If `accumulate_timeouts` is `false` then counting will reset when time between iterations is shorter than `timelimit`, otherwise it will not.
"""
struct TimedIterator{F,A,I}
    timelimit::F
    patience::Int
    timeoutaction::A
    accumulate_timeouts::Bool
    base::I
end
TimedIterator(;timelimit, patience=5, timeoutaction, accumulate_timeouts=false, base) = TimedIterator(timelimit, patience, timeoutaction, accumulate_timeouts, base)
"""
    TimedIteratorStop

Special type for stopping a `TimedIterator` if returned from `timeoutaction()`.
"""
struct TimedIteratorStop
    TimedIteratorStop() = throw(DomainError("You typically do not want to instantiate this as `TimedIterator` expects the type and not a value, i.e `return TimedIteratorStop` instead of `return TimedIteratorStop()`. Call TimedIteratorStop(true) if you really want to instantiate"))
    TimedIteratorStop(::Bool) = new()
end

Base.length(itr::TimedIterator) = length(itr.base)
Base.size(itr::TimedIterator) = size(itr.base)

Base.IteratorSize(::Type{TimedIterator{F,A,I}}) where {F,A,I} = Base.IteratorSize(I)
Base.IteratorEltype(::Type{TimedIterator{F,A,I}}) where {F,A,I}  = Base.IteratorEltype(I)

function Base.iterate(itr::TimedIterator)
    val, bstate = IterTools.@ifsomething iterate(itr.base)
    tstamp = NaN #Always skip first due to compilation times etc
    return val, (tstamp, 0, bstate)
end

function Base.iterate(itr::TimedIterator, (tstamp, ntimeout, bstate)::Tuple)
    duration = time() - tstamp
    # Note reiles on NaN > x === false for all x
    ntimeout = if duration > itr.timelimit
        ntimeout + 1
    else
         itr.accumulate_timeouts ? ntimeout : 0
    end

    if ntimeout >= itr.patience
        stop = itr.timeoutaction()
        # Note: we allow users to return e.g. nothing here to continue
        stop === TimedIteratorStop && return nothing
    end

    val, bstate = IterTools.@ifsomething iterate(itr.base, bstate)
    tstamp = time()
    return val, (tstamp, ntimeout, bstate)
end

# itergeneration on timeoutaction just in case user e.g. wants to deepcopy it to avoid some shared state.
itergeneration(itr::TimedIterator, gen) = TimedIterator(itr.timelimit, itr.patience, itergeneration(itr.timeoutaction, gen), itr.accumulate_timeouts, itergeneration(itr.base, gen))


text:
export ode45, rk4, αscheme, αscheme_time,
    αscheme_atime, TR_BDF2, ExplicitNewmark, Newmark

function runge_kutta_one_step(f::Function, t::PyObject, y::PyObject, Δt::PyObject, θ::Union{PyObject, Missing})
    k1 = Δt*f(t, y, θ)
    k2 = Δt*f(t+Δt/2, y+k1/2, θ)
    k3 = Δt*f(t+Δt/2, y+k2/2, θ)
    k4 = Δt*f(t+Δt, y+k3, θ)
    y = y + k1/6 + k2/3 + k3/3 + k4/6
end

function ode45_one_step(f::Function, t::PyObject, y::PyObject, h::PyObject, θ::Union{PyObject, Missing})
    k1 = h * f(t, y, θ);
    k2 = h * f(t + (1/5)*h, y + (1/5)*k1, θ);
    k3 = h * f(t + (3/10)*h, y + (3/40)*k1 + (9/40)*k2, θ);
    k4 = h * f(t + (4/5)*h, y + (44/45)*k1 + (-56/15)*k2 + (32/9)*k3, θ);
    k5 = h * f(t + (8/9)*h, y + (19372/6561)*k1 + (-25360/2187)*k2 + (64448/6561)*k3 + (-212/729)*k4, θ);
    k6 = h * f(t + h, y + (9017/3168)*k1 + (-355/33)*k2 + (46732/5247 )*k3 + (49/176)*k4 + (-5103/18656)*k5, θ);
    k7 = h * f(t + h, y + (35/384)*k1 + (500/1113)*k3 + (125/192)*k4 + (-2187/6784)*k5 + (11/84)*k6, θ);

    y_new = y + (35/384)*k1 + (500/1113)*k3 + (125/192)*k4 + (-2187/6784)*k5 + (11/84)*k6;
end

@doc raw"""
    runge_kutta(f::Function, T::Union{PyObject, Float64},
                NT::Union{PyObject,Int64}, y::Union{PyObject, Float64, Array{Float64}}, θ::Union{PyObject, Missing}=missing; method::String="rk4")

Solves
```math
\frac{dy}{dt} = f(y, t, \theta)
```
with Runge-Kutta method.

For example, the default solver, `RK4`, has the following numerical scheme per time step
```math
\begin{aligned}
k_1 &= \Delta t f(t_n, y_n, \theta)\\
k_2 &= \Delta t f(t_n+\Delta t/2, y_n + k_1/2, \theta)\\
k_3 &= \Delta t f(t_n+\Delta t/2, y_n + k_2/2, \theta)\\
k_4 &= \Delta t f(t_n+\Delta t, y_n + k_3, \theta)\\
y_{n+1} &= y_n + \frac{k_1}{6} +\frac{k_2}{3} +\frac{k_3}{3} +\frac{k_4}{6}
\end{aligned}
```
"""
function runge_kutta(f::Function, T::Union{PyObject, Float64},
                NT::Union{PyObject,Int64}, y::Union{PyObject, Float64, Array{Float64}}, θ::Union{PyObject, Missing}=missing; method::String="rk4")

    local one_step
    if lowercase(method)=="rk4"
        one_step = runge_kutta_one_step
    elseif lowercase(method)=="rk45"
        one_step = ode45_one_step
    else
        error("Method $method not implemented yet")
    end
    y = convert_to_tensor(y)
    Δt = convert_to_tensor(T/NT)

    ta = TensorArray(NT+1) # storing y
    function condition(i, ta)
        i <= NT+1
    end
    function body(i, ta)
        y = read(ta, i-1)
        y_ = one_step(f, (cast(eltype(Δt), i)-1)*Δt, y, Δt, θ)
        ta = write(ta, i, y_)
        i+1, ta
    end
    ta = write(ta, 1, y)
    i = constant(2, dtype=Int32)
    _, out = while_loop(condition, body, [i, ta])
    res = stack(out)
end

@doc raw"""
    rk4(y::Union{PyObject, Float64, Array{Float64}}, T::Union{PyObject, Float64},
                NT::Union{PyObject,Int64}, f::Function, θ::Union{PyObject, Missing}=missing)

Solves
```math
\frac{dy}{dt} = f(y, t, \theta)
```
with Runge-Kutta (order 4) method.
"""
rk4(args...;kwargs...) = runge_kutta(args...;method="rk4", kwargs...)

@doc raw"""
    ode45(y::Union{PyObject, Float64, Array{Float64}}, T::Union{PyObject, Float64},
                NT::Union{PyObject,Int64}, f::Function, θ::Union{PyObject, Missing}=missing)

Solves
```math
\frac{dy}{dt} = f(y, t, \theta)
```
with six-stage, fifth-order, Runge-Kutta method.
"""
ode45(args...;kwargs...) = runge_kutta(args...;method="rk45", kwargs...)

@doc raw"""
    αscheme(M::Union{SparseTensor, SparseMatrixCSC},
        C::Union{SparseTensor, SparseMatrixCSC},
        K::Union{SparseTensor, SparseMatrixCSC},
        Force::Union{Array{Float64}, PyObject},
        d0::Union{Array{Float64, 1}, PyObject},
        v0::Union{Array{Float64, 1}, PyObject},
        a0::Union{Array{Float64, 1}, PyObject},
        Δt::Array{Float64};
        solve::Union{Missing, Function} = missing,
        extsolve::Union{Missing, Function} = missing,
        ρ::Float64 = 1.0)

Generalized α-scheme.
$$M u_{tt} + C u_{t} + K u = F$$

`Force` must be an array of size `n`×`p`, where `d0`, `v0`, and `a0` have a size `p`
`Δt` is an array (variable time step).

The generalized α scheme solves the equation by the time stepping
```math
\begin{aligned}
\bf d_{n+1} &= \bf d_n + h\bf v_n + h^2 \left(\left(\frac{1}{2}-\beta_2 \right)\bf a_n + \beta_2 \bf a_{n+1}  \right)\\
\bf v_{n+1} &= \bf v_n + h((1-\gamma_2)\bf a_n + \gamma_2 \bf a_{n+1})\\
\bf F(t_{n+1-\alpha_{f_2}}) &= M \bf a _{n+1-\alpha_{m_2}} + C \bf v_{n+1-\alpha_{f_2}} + K \bf{d}_{n+1-\alpha_{f_2}}
\end{aligned}
```
where
```math
\begin{aligned}
\bf d_{n+1-\alpha_{f_2}} &= (1-\alpha_{f_2})\bf d_{n+1} + \alpha_{f_2} \bf d_n\\
\bf v_{n+1-\alpha_{f_2}} &= (1-\alpha_{f_2}) \bf v_{n+1} + \alpha_{f_2} \bf v_n \\
\bf a_{n+1-\alpha_{m_2} } &= (1-\alpha_{m_2}) \bf a_{n+1} + \alpha_{m_2} \bf a_n\\
t_{n+1-\alpha_{f_2}} & = (1-\alpha_{f_2}) t_{n+1 + \alpha_{f_2}} + \alpha_{f_2}t_n
\end{aligned}
```

Here the parameters are computed using
```math
\begin{aligned}
\gamma_2 &= \frac{1}{2} - \alpha_{m_2} + \alpha_{f_2}\\
\beta_2 &= \frac{1}{4} (1-\alpha_{m_2}+\alpha_{f_2})^2 \\
\alpha_{m_2} &= \frac{2\rho_\infty-1}{\rho_\infty+1}\\
\alpha_{f_2} &= \frac{\rho_\infty}{\rho_\infty+1}
\end{aligned}
```

∘ `solve`: users can provide a solver function, `solve(A, rhs)` for solving `Ax = rhs`
∘ `extsolve`: similar to `solve`, but the signature has the form
```julia
extsolve(A, rhs, i)
```
This provides the users with more control, e.g., (time-dependent) Dirichlet boundary conditions.
See [Generalized α Scheme](https://kailaix.github.io/ADCME.jl/dev/alphascheme/) for details.

!!! note
    In the case $u$ has a nonzero essential boundary condition $u_b$, we let $\tilde u=u-u_b$, then
    $$M \tilde u_{tt} + C \tilde u_t + K u = F - K u_b - C \dot u_b$$
"""
function αscheme(M::Union{SparseTensor, SparseMatrixCSC},
                      C::Union{SparseTensor, SparseMatrixCSC},
                      K::Union{SparseTensor, SparseMatrixCSC},
                      Force::Union{Array{Float64, 2}, PyObject},
                      d0::Union{Array{Float64, 1}, PyObject},
                      v0::Union{Array{Float64, 1}, PyObject},
                      a0::Union{Array{Float64, 1}, PyObject},
                      Δt::Array{Float64, 1};
                      solve::Union{Missing, Function} = missing,
                      extsolve::Union{Missing, Function} = missing,
                      ρ::Float64 = 1.0)
    if !ismissing(solve) && !ismissing(extsolve)
        error("You cannot provide `solve` and `extsolve` at the same time.")
    end
    nt = length(Δt)
    αm = (2ρ-1)/(ρ+1)
    αf = ρ/(1+ρ)
    γ = 1/2-αm+αf
    β = 0.25*(1-αm+αf)^2
    d = length(d0)

    M = isa(M, SparseMatrixCSC) ? constant(M) : M
    C = isa(C, SparseMatrixCSC) ? constant(C) : C
    K = isa(K, SparseMatrixCSC) ? constant(K) : K
    Force, d0, v0, a0, Δt = convert_to_tensor([Force, d0, v0, a0, Δt], [Float64, Float64, Float64, Float64, Float64])

    function equ(dc, vc, ac, dt, Force, i)
        dn = dc + dt*vc + dt^2/2*(1-2β)*ac
        vn = vc + dt*((1-γ)*ac)

        df = (1-αf)*dn + αf*dc
        vf = (1-αf)*vn + αf*vc
        am = αm*ac

        rhs = Force - (M*am + C*vf + K*df)
        A = (1-αm)*M + (1-αf)*C*dt*γ + (1-αf)*K*β*dt^2

        if !ismissing(solve)
            return solve(A, rhs)
        elseif !ismissing(extsolve)
            return extsolve(A, rhs, i)
        else
            return A\rhs
        end
    end

    function condition(i, tas...)
        return i<=nt
    end
    function body(i, tas...)
        dc_arr, vc_arr, ac_arr = tas
        dc = read(dc_arr, i)
        vc = read(vc_arr, i)
        ac = read(ac_arr, i)
        y = equ(dc, vc, ac, Δt[i], Force[i], i)
        dn = dc + Δt[i]*vc + Δt[i]^2/2*((1-2β)*ac+2β*y)
        vn = vc + Δt[i]*((1-γ)*ac+γ*y)
        i+1, write(dc_arr, i+1, dn), write(vc_arr, i+1, vn), write(ac_arr, i+1, y)
    end

    dM = TensorArray(nt+1); vM = TensorArray(nt+1); aM = TensorArray(nt+1)
    dM = write(dM, 1, d0)
    vM = write(vM, 1, v0)
    aM = write(aM, 1, a0)
    i = constant(1, dtype=Int32)
    _, d, v, a = while_loop(condition, body, [i,dM, vM, aM])
    set_shape(stack(d), (nt+1, length(d0))), set_shape(stack(v), (nt+1, length(v0))), set_shape(stack(a), (nt+1, length(a0)))
end

@doc raw"""
    αscheme_time(Δt::Array{Float64}; ρ::Float64 = 1.0)

Returns the integration time $t_{i+1-\alpha_{f_2}}$ between $[t_i, t_{i+1}]$ using the alpha scheme.
If $\Delta t$ has length $n$, the output will also have length $n$.
"""
function αscheme_time(Δt::Array{Float64}; ρ::Float64 = 1.0)
    n = length(Δt)
    αm = (2ρ-1)/(ρ+1)
    αf = ρ/(1+ρ)
    γ = 1/2-αm+αf
    β = 0.25*(1-αm+αf)^2
    function equ(tc, dt)
        tf1 = (1-αf)*(tc+dt) + αf*tc
        return tf1
    end

    tcf = Float64[]
    tc = 0.0
    for i = 1:n
        t1 = equ(tc, Δt[i])
        push!(tcf, t1)
        tc += Δt[i]
    end
    return tcf
end

@doc raw"""
    TR_BDF2(D0::Union{SparseTensor, SparseMatrixCSC},
        D1::Union{SparseTensor, SparseMatrixCSC},
        Δt::Float64)

Constructs a TR-BDF2 (the Trapezoidal Rule with Second Order Backward Difference Formula) handler for
the DAE

$$D_1 \dot y + D_0 y = f$$

The struct is a functor, which performs one step simulation

```
(tr::TR_BDF2)(y::Union{PyObject, Array{Float64, 1}},
    f1::Union{PyObject, Array{Float64, 1}},
    f2::Union{PyObject, Array{Float64, 1}},
    f3::Union{PyObject, Array{Float64, 1}})
```
Here `f1`, `f2`, and `f3` correspond to the right hand side at time step $n$, $n+\frac12$, and $n+1$.

Or we can pass a batched `F` defined as a `(2NT+1) × DOF` array

```
(tr::TR_BDF2)(y0::Union{PyObject, Array{Float64, 1}},
    F::Union{PyObject, Array{Float64, 2}})
```

The output will be the entire solution of size `(NT+1) × DOF`.

!!! info
    The scheme takes the following form for n = 0, 1, ...
    $$\begin{aligned} D_1(y^{n+\frac12}-y^n) = \frac12\frac{\Delta t}{2}\left(f^{n+\frac12} + f^n - D_0 \left(y^{n+\frac12} + y^n\right)\right)\\ \left(\frac{\Delta t}{2}\right)^{-1} D_1 \left(\frac32y^{n+1} - 2y^{n+\frac12} + \frac12 y^n\right) + D_0 y^{n+1} = f^{n+1}\end{aligned}$$
"""
mutable struct TR_BDF2
    D0::Union{SparseTensor, SparseMatrixCSC}
    D1::Union{SparseTensor, SparseMatrixCSC}
    Δt::Float64
    _D0::Union{SparseTensor, SparseMatrixCSC}
    _D1::Union{SparseTensor, SparseMatrixCSC}
    symbolic::Bool
    function TR_BDF2(D0::Union{SparseTensor, SparseMatrixCSC},
            D1::Union{SparseTensor, SparseMatrixCSC},
            Δt::Float64)
        if isa(D0, SparseTensor) || isa(D1, SparseTensor)
            symbolic = true
            D0 = constant(D0)
            D1 = constant(D1)
        else
            symbolic = false
        end
        @assert size(D0, 1)==size(D0,2)==size(D1,1)==size(D1,2)
        _D0 = D1 + Δt/4 * D0
        _D1 = 1/(Δt/2) * 3/2 * D1  + D0
        new(D0, D1, Δt, _D0, _D1, symbolic)
    end
end

"""
    constant(tr::TR_BDF2)

Converts `tr` to a symbolic solver.
"""
function constant(tr::TR_BDF2)
    TR_BDF2(constant(tr.D0), constant(tr.D1), tr.Δt)
end

function Base.:show(io::IO, tr::TR_BDF2)
    print("""TR_BDF2 (DOF = $(size(tr.D0, 1)), Δt = $(tr.Δt)$(tr.symbolic ? ", symbolic" : ""))""")
end

function (tr::TR_BDF2)(y::Union{PyObject, Array{Float64, 1}},
    f1::Union{PyObject, Array{Float64, 1}},
    f2::Union{PyObject, Array{Float64, 1}},
    f3::Union{PyObject, Array{Float64, 1}})
    y, f1, f2, f3 = convert_to_tensor([y, f1, f2, f3], [Float64, Float64, Float64, Float64])
    r1 = tr.Δt/4 * (f2 + f1) - tr.Δt/4 * (tr.D0 * y) + tr.D1 * y
    yn = tr._D0\r1
    r2 = f3 + 1/(tr.Δt/2)*(tr.D1* (2*yn - 0.5*y))
    tr._D1\r2
end

function (tr::TR_BDF2)(y0::Union{PyObject, Array{Float64, 1}},
    F::Union{PyObject, Array{Float64, 2}})
    @assert size(F, 2)==size(tr.D0, 1) && mod(size(F, 1), 2)==1
    y0, F = convert_to_tensor([y0, F], [Float64, Float64])
    NT = size(F, 1)÷2
    y_arr = TensorArray(NT+1)
    y_arr = write(y_arr, 1, y0)
    function condition(i, y_arr)
        i<=NT
    end
    function body(i, y_arr)
        y = read(y_arr, i)
        f1, f2, f3 = F[2*i-1], F[2*i], F[2*i+1]
        yn = tr(y, f1, f2, f3)
        i+1, write(y_arr, i+1, yn)
    end
    i = constant(1, dtype = Int32)
    _, ya = while_loop(condition, body, [i, y_arr])
    set_shape(stack(ya), (NT+1, length(y0)))
end

function (tr::TR_BDF2)(y::Array{Float64, 1},
    f1::Array{Float64, 1},
    f2::Array{Float64, 1},
    f3::Array{Float64, 1})
    if tr.symbolic
        return tr(constant(y), f1, f2, f3)
    end
    r1 = tr.Δt/4 * (f2 + f1) - tr.Δt/4 * (tr.D0 * y) + tr.D1 * y
    yn = tr._D0\r1
    r2 = f3 + 1/(tr.Δt/2)*(tr.D1* (2*yn - 0.5*y))
    tr._D1\r2
end

function (tr::TR_BDF2)(y0::Array{Float64,1}, F::Array{Float64, 2})
    if tr.symbolic
        return tr(constant(y0), F)
    end
    @assert size(F, 2)==size(tr.D0, 1) && mod(size(F, 1), 2)==1
    y = zeros(size(F, 1)÷2+1, 2)
    y[1,:] = y0
    for i = 1:size(F, 1)÷2
        y[i+1,:] = tr(y[i,:], F[2*i-1,:], F[2*i,:], F[2*i+1,:])
    end
    y
end

# design principle: proceed one step
# (states) u0, u1, ..., (coefficients) M0, M1, ..., (step size) Δt

@doc raw"""
    ExplicitNewmark(M::Union{SparseTensor, SparseMatrixCSC}, Z1::Union{Missing, SparseTensor, SparseMatrixCSC}, Z2::Union{Missing, SparseTensor, SparseMatrixCSC}, Δt::Float64)

An explicit Newmark integrator for

$$M \ddot{\mathbf{d}} + Z_1 \dot{\mathbf{d}} + Z_2 \mathbf{d} + f = 0$$

The numerical scheme is

$$\left(\frac{1}{\Delta t^2} M + \frac{1}{2\Delta t}Z_1\right)d^{n+1} = \left(\frac{2}{\Delta t^2} M - Z_2\right)d^n - \left(\frac{1}{\Delta t^2} M - \frac{1}{2\Delta t}Z_1\right) d^{n-1} - f$$

To use this integrator,

```julia
en = ExplicitNewmark(M, Z1, Z2, Δt)
d2 = step(en, d0, d1, f)
```
"""
mutable struct ExplicitNewmark
    A::Union{PyObject, Tuple{SparseTensor, PyObject}, SparseTensor}
    B::Union{PyObject, SparseTensor}
    C::Union{PyObject, SparseTensor}
    function ExplicitNewmark(M::Union{SparseTensor, PyObject, Array{Float64, 2}, SparseMatrixCSC},
            Z1::Union{Missing, PyObject, Array{Float64, 2}, SparseTensor, SparseMatrixCSC},
            Z2::Union{Missing, PyObject, Array{Float64, 2}, SparseTensor, SparseMatrixCSC}, Δt::Float64;
            factorize_A::Bool = true)
        M = constant(M)
        if ismissing(Z1)
            A = 1/Δt^2 * M
            C = -1/Δt^2 * M
        else
            Z1 = constant(Z1)
            A = (1/Δt^2 * M + 1/(2Δt) * Z1)
            C = -(1/Δt^2 * M - 1/(2Δt) * Z1)
        end
        if isa(A, SparseTensor) && factorize_A
            A = factorize(A)
        end
        if ismissing(Z2)
            B = 2/Δt^2 * M
        else
            Z2 = constant(Z2)
            B = (2/Δt^2 * M - Z2)
        end
        new(A, B, C)
    end
end

function Base.:show(io::IO, en::ExplicitNewmark)
    print("ExplicitNewmark(DOF=$(size(en.B, 1)))")
end

function Base.:step(en::ExplicitNewmark,
    d0::Union{Array{Float64, 1}, PyObject},
    d1::Union{Array{Float64, 1}, PyObject},
    f::Union{Array{Float64, 1}, PyObject})
    d0, d1, f = convert_to_tensor([d0, d1, f], [Float64, Float64, Float64])
    en.A \ (en.B * d1 + en.C * d0 - f)
end

@doc raw"""
    Newmark(M::Union{SparseTensor, PyObject, Array{Float64, 2}, SparseMatrixCSC},
                        C::Union{SparseTensor, PyObject, Array{Float64, 2}, SparseMatrixCSC},
                        K::Union{SparseTensor, PyObject, Array{Float64, 2}, SparseMatrixCSC},
                        Δt::Float64; β::Float64 = 1/4, γ::Float64 = 1/2, factorize_S::Bool = true)

Creates a Newmark algorithm object. This function uses "a-form" implementation for

$$M\ddot \mathbf{d} + C \dot\mathbf{d} + K \mathbf{d} + f = 0$$

The numerical scheme is

$$\begin{aligned}
\tilde d_{n+1} &= d_n + \Delta t v_n + \Delta^2/2(1-2\beta)a_n \\
\tilde v_{n+1} &= v_n + (1-\gamma)\Delta a_n \\
(M+\gamma\Delta C+\beta \Delta^2 K)a_{n+1} &= F_{n+1} - C\tilde v_{n+1} - K\tilde d_{n+1}\\
d_{n+1} &= \tilde d_{n+1} + \beta \Delta^2 a_{n+1}\\
v_{n+1} &= \tilde v_{n+1} + \gamma \Delta ta_n
\end{aligned}$$

For unconditional stability
$$2\beta \geq \gamma \geq \frac12$$

For an explicit form, see [`ExplicitNewmark`](@ref), which corresponds to $\beta = 0, \gamma = \frac12$.

A common choice for Newmark algorithm is $\beta = \frac14, \gamma = \frac12$ (a.k.a., average acceleration, or
Trapezoidal rule), which is unconditionally stable,
has second order of accuracy.

!!! info
    Internally, $M+\gamma\Delta C+\beta \Delta^2 K$ is represented by `S`. Users can elect to factorize $S$ (default $=$ true).

To perform time marching, use
```
step(nm::Newmark, d::PyObject, v::PyObject, a::PyObject,
        F::Union{Array{Float64, 1}, PyObject})
```
"""
mutable struct Newmark
    M::SparseTensor
    C::SparseTensor
    K::SparseTensor
    S::Union{Tuple{SparseTensor, PyCall.PyObject}, SparseTensor}
    β::Float64
    γ::Float64
    Δt::Float64
    function Newmark(M::Union{SparseTensor, PyObject, Array{Float64, 2}, SparseMatrixCSC},
            C::Union{SparseTensor, PyObject, Array{Float64, 2}, SparseMatrixCSC},
            K::Union{SparseTensor, PyObject, Array{Float64, 2}, SparseMatrixCSC},
            Δt::Float64; β::Float64 = 1/4, γ::Float64 = 1/2, factorize_S::Bool = true)
        M = SparseTensor(M)
        C = SparseTensor(C)
        K = SparseTensor(K)
        S = M + γ*Δt*C + β*Δt^2*K
        if factorize_S
            S = factorize(S)
        end
        new(M, C, K, S, β, γ, Δt)
    end
end

function Base.:show(io::IO, en::Newmark)
    print("Newmark(DOF=$(size(en.S, 1)), γ=$(en.γ), β=$(en.β))")
end

function Base.:step(nm::Newmark, d::PyObject, v::PyObject, a::PyObject,
        F::Union{Array{Float64, 1}, PyObject})
    β, γ, Δt = nm.β, nm.γ, nm.Δt
    td = d + Δt * v + Δt^2/2*(1-2β)*a
    tv = v + (1-γ)*Δt*a
    a_ = nm.S\(-F - nm.C*tv - nm.K*td)
    d_ = td + β*Δt^2*a_
    v_ = tv + γ*Δt*a_
    d_, v_, a_
end


text:
include("schema.jl")

function estimate_aide(gold_standard_sampler::Sampler, target_sampler::Sampler, num_replicates::Int)
    aide_estimates = aide(gold_standard_sampler, target_sampler, 1, 1, num_replicates, num_replicates)
    estimate = mean(aide_estimates)
    stderr = std(aide_estimates) / sqrt(num_replicates)
    return (estimate, stderr)
end

function estimate_expected_num_clusters(target_sampler::Sampler, num_replicates::Int)
    target_samples = pmap((i) -> simulate(target_sampler)[1], 1:num_replicates)
    num_clusters = map((s) -> length(clusters(s.crp)), target_samples)
    num_clusters = convert(Vector{Float64}, num_clusters)
    estimate = mean(num_clusters)
    stderr = std(num_clusters) / sqrt(length(num_clusters))
    return (estimate, stderr)
end

function generate_sampler_data(data::Vector{Float64}, hypers::Hyperparameters,
                               num_aide_replicates::Int, num_diagnostic_replicates::Int,
                               gold_standard_sampler::Sampler)
    proposal_name_column = String[]
    num_sweeps_column = Int[]
    num_particles_column = Int[]
    aide_estimate_column = Float64[]
    aide_stderr_column = Float64[]
    num_clusters_estimate_column = Float64[]
    num_clusters_stderr_column = Float64[]

    for num_sweeps in num_sweeps_list
        for num_particles in num_particles_list
            prior_target_sampler = make_prior_proposal_smc_sampler(data, hypers, num_sweeps, num_particles)
            optimal_target_sampler = make_optimal_proposal_smc_sampler(data, hypers, num_sweeps, num_particles)

            target_samplers = [(prior_target_sampler, PRIOR_PROPOSAL_NAME),
                               (optimal_target_sampler, OPTIMAL_PROPOSAL_NAME)]

            for (target_sampler, proposal_name) in target_samplers
                println("$proposal_name, num_sweeps=$num_sweeps, num_particles=$num_particles")
                push!(proposal_name_column, proposal_name)
                push!(num_sweeps_column, num_sweeps)
                push!(num_particles_column, num_particles)

                # AIDE estimate
                # NOTE: it is possible to reuse the gold-standard samplers when evaluating different target samplers
                # We do not do this here.
                (aide_estimate, aide_stderr) = estimate_aide(gold_standard_sampler, target_sampler, num_aide_replicates)
                push!(aide_estimate_column, aide_estimate)
                push!(aide_stderr_column, aide_stderr)

                # estimate the expected number of clusters
                (num_clusters_estimate, num_clusters_stderr) = estimate_expected_num_clusters(target_sampler, num_diagnostic_replicates)
                push!(num_clusters_estimate_column, num_clusters_estimate)
                push!(num_clusters_stderr_column, num_clusters_stderr)
            end
        end
    end

    data = DataFrame()
    data[COL_PROPOSAL_NAME] = proposal_name_column
    data[COL_NUM_PARTICLES] = num_particles_column
    data[COL_NUM_SWEEPS] = num_sweeps_column
    data[COL_AIDE_ESTIMATE] = aide_estimate_column
    data[COL_AIDE_STDERR] = aide_stderr_column
    data[COL_NUM_CLUSTERS_ESTIMATE] = num_clusters_estimate_column
    data[COL_NUM_CLUSTERS_STDERR] = num_clusters_stderr_column
    return data
end


text:
# Functions used to give workers tasks for distributed processing

# Imports for distributed processors
#using OPFLearn

"""
Produces samples from the given polytope and stores them in the sample channel.
Adds any new polytope constraints that are put in the polytope channel to the
polytope defined by Ax < b. When a new constraint is added all samples in the
sample channel are replaced.

Runs until the given termination criteria are met.
"""
function sample_producer(A, b, sampler, sampler_opts::Dict, base_load_feasible,
					     K, U, S, V, max_iter, T, num_procs, results,
						 discard, variance, reset_level,
						 save_while, save_infeasible, stat_track, save_certs,
						 net_name, dual_vars, save_order, replace_samples,
						 save_path, sample_ch, polytope_ch, result_ch, final_ch,
						 print_level)
	now_str = Dates.format(Dates.now(), "mm-dd-yyy_HH.MM.SS")

	AC_inputs = results["inputs"]
    AC_outputs = results["outputs"]
    duals = results["duals"]

	save_infeasible && (infeasible_AC_inputs = results["infeasible_inputs"])
	save_certs && (store_cert(results, A, b))
	stat_track > 0 && (stats = results["stats"])

	# Gather initial chebychev center for sampler
	center, radius = chebyshev_center(A, b)
	x0 = center'

	#TEST: Use base_load_feasible in starting point
	x0 = base_load_feasible * 0.9 + center' * 0.1

	# Initialize load sample, x, containing [PG; QG]
    new_samples = sampler(A, b, x0, 4*num_procs; sampler_opts...)
	for sample in eachcol(new_samples)
		put!(sample_ch, sample)
	end

	k = 0  # Count of feasible samples collected
    i = 0  # Count of samples generated
	u = 0  # Count of feasible samples since last unique active set found
	v = 0  # Count of feasible samples since last increase in variance seen
	s = 0  # Count of feasible samples since last saved (not discarded) sample
	(print_level > 0) && println("Starting sample production...")
	start_time = time()
	while (k < K) & (u < (1 / U)) & (s < (1 / S)) & (v < 1 / V)	 &
		  (i < max_iter) & ((time() - start_time) < T)
		n_certs = length(b)
		while isready(polytope_ch)
			println("ADDING SLICE")
			x_star, x = take!(polytope_ch)

			A, b = add_infeasibility_certificate(A, b, x_star, x)
			save_certs && (store_cert(results, A, b))
			if save_while
				save_cert(A, b, net_name*"_"*now_str, dir=save_path)
			end

			x0 = x_star  #TASK: Should work since x_star always r-AC-OPF feasible
		end

		if length(b) > n_certs
			# Fully replace values in sample_chnl
			center, radius = chebyshev_center(A, b)
			if reset_level > 1
				x0 = (base_load_feasible * 0.1 + x0 * 0.9) * 0.9 + center' * 0.1
			elseif reset_level > 0
				x0 = x0 * 0.9 + center' * 0.1
			else
				x0 = center'
			end

			# Replace values in sample queue
			#TEST: Sometimes this seems to block and make the process halt
			if replace_samples
				num_new_samples = 0
				while isready(sample_ch)
					take!(sample_ch)  #TASK: Figure out if this can be blocking
					num_new_samples += 1  # and can get stuck waiting on a sample when there are none
				end

				new_samples = sampler(A, b, x0, num_new_samples; sampler_opts...)
				for sample in eachcol(new_samples)
					put!(sample_ch, sample)
				end
			end
		end

		# Only takes a sample once per iteration which can have two results
		#TEST: Will it still work with while instead of if? Would likely improve speeds as procs could be waiting on samples?
		# What happens if results are being returned faster than this loop? Put a limit on it? 10 iters per while?
		num_new_samples = 0
		while isready(result_ch) & ((k < K) & (u < (1 / U)) & (s < (1 / S)) & (v < 1 / V))	 &
								   (i < max_iter) & ((time() - start_time) < T)
			# Get sample result from result channel
			x, result, feasible, new_cert, iter_elapsed_time = take!(result_ch)

			# Set stats to defaults
			iter_stats = Dict(:new_cert => new_cert,
							  :new_set => false,
							  :var_inc => false,
							  :added => false,
							  :feasible => false,
							  :iter_time => iter_elapsed_time,
							  :active_set => []
							  )

			#TASK: Determine if there is a way to continue sampling off of old load
			# Can only tell if the result is feasible
			#if feasible
			#	x0 = x #TASK: Set this up to be when SOC feasible or ACOPF feasible
			#end

			num_new_samples += 1
			i += 1  # Increment iterations

			if feasible
				s, u, v, k, iter_stats = store_feasible_sample(s, u, v, k, i, K, iter_stats,
											  AC_inputs, AC_outputs, duals, dual_vars,
											  x, result, discard, variance, net_name,
											  now_str, save_path, save_while, save_order,
											  print_level)

			elseif !isnothing(result)  # Infeasible
				save_infeasible && store_infeasible_sample(infeasible_AC_inputs, x, result,
								save_while, net_name, now_str, save_order, dual_vars, save_path)
			end

			print_level > 0 && println("Samples: $(k) / $(K),\t Iter: $(i)")

			if stat_track > 0
				update_stats!(stats, duals, iter_stats, save_level=stat_track)
				save_while && (save_stats(iter_stats, net_name*"_"*now_str*"_stats", dir=save_path))
			end
		end

		if num_new_samples > 0
			new_samples = sampler(A, b, x0, num_new_samples; sampler_opts...)
			for sample in eachcol(new_samples)
				put!(sample_ch, sample)
			end
			x0 = new_samples[:, end]  # Set sampler origin to last sampled sample
		end
	end
	# Put results objects into the final channel pipeline to get returned
	put!(final_ch, "DONE FLAG")

	return results
end

"""
Adds a new constraint to a polytope defined by Ax < b.
The constraint is constructed from a point, p1, and a
normal vector, p2 - p1.
"""
function add_infeasibility_certificate(A, b, p1, p2)
	normal = (p2 - p1)'

	# Add infeasibility certificate
	A = vcat(A, normal)
	b = vcat(b, normal * p1)
	return A, b
end

"""
Tests feasibility of generated samples from the sample channel and if
infeasible finds the nearest feasible point. Adds this nearest feasible
point and original sample to the polytope channel to add a infeasibility
certificate in the polytope sample space. Runs until a value is added to
the done_ch.
"""
function sample_processor(net, net_r, r_solver, opf_solver,
						  sample_ch, done_ch, poly_ch, result_ch,
						  print_level=0, model_type=PM.QCLSPowerModel)

	(print_level > 0) && println("Create FNFP Model...")
	pm = PowerModels.instantiate_model(net_r, model_type, build_opf_var_load)
	fnfp_model = find_nearest_feasible_model(pm, print_level=print_level, solver=r_solver)

	(print_level > 0) && println("Starting sample processing...")
    i = 1  # Count of iterations
    while !isready(done_ch)  #TASK: Determine the best way to terminate this loop
        iter_start_time = time()
		println("Iter: $(i)") #TEST: Printing samples every iteration

		# Gather sample from sample channel
		x = take!(sample_ch)

        # Set network loads to sampled values
        set_network_load(net, x, scale_load=false)

        # Solve OPF for the load sample
        result, feasible = run_ac_opf(net, solver=opf_solver)
        println("OPF SUCCESS: " * string(feasible))

		new_cert = false  # Default for iteration stat tracking

		if !feasible
			println("FINDING NEAREST FEASIBLE")
            px = x[1:Integer(length(x)/2)]
            qx = x[Integer(length(x)/2) + 1:end]

            r, pd, qd, solved = find_nearest_feasible(fnfp_model, px, qx, print_level=print_level)

			r = sum((pd .- px).^2 + (qd .- qx).^2)
			println("R:", r)

            if (r > R_TOLERANCE) & solved
				new_cert = true  # For iteration stat tracking
                # Solve OPF for the random sample:
                # try to solve the non-convex AC OPF at this solution
                # note that: while this solution is feasible for the relaxation
                # it can be infeasible for the original problem.

                # Get infeasibility certificate
                xp_ = pd .- INFEASIBILITY_CERT_SHIFT
                xq_ = qd .- INFEASIBILITY_CERT_SHIFT
				xp_[xp_ .< 0] .= 0 # Was used when substracting 1e-2 here
				xq_[xq_ .< 0] .= 0 # Ensure loads do not go negative when shifted

                x_star = vcat(xp_, xq_)

				# Puts new infeasibility certificate points to be added to the polytope
				put!(poly_ch, (x_star, x))

				xp_ = pd .- FNFP_SHIFT
                xq_ = qd .- FNFP_SHIFT
				xp_[xp_ .< 0] .= 0
				xq_[xq_ .< 0] .= 0
                x = vcat(xp_, xq_)

                set_network_load(net, x, scale_load=false)

                # Solve OPF for the relaxation feasible sample
                result, feasible = run_ac_opf(net, solver=opf_solver)
                println("FNFP OPF SUCCESS: " * string(feasible))
			end
		end
		# Puts OPF results to the results channel to be processed by main thread
		iter_elapsed_time = time() - iter_start_time

		put!(result_ch, (x, result, feasible, new_cert, iter_elapsed_time))
        i = i + 1
	end
end


text:
module Figures
using AverageShiftedHistograms, Distributions
using Plots
pyplot()

o = ash(rand(Gamma(5, 1), 0), 0:.05:20, m = 15)
sz = 20
anim = @animate for i in 1:50
    fit!(o, rand(Gamma(5, 1), sz); warnout = false)
    plot(o, title = "Nobs = $(10 * i)", ylim = (0, .5))
end
gif(anim, "animation.gif", fps = 10)

end


text:
hover_test(line, char) = LanguageServer.textDocument_hover_request(LanguageServer.TextDocumentPositionParams(LanguageServer.TextDocumentIdentifier("testdoc"), LanguageServer.Position(line, char)), server, server.jr_endpoint)

settestdoc("""
1234
Base
+
vari = 1234
\"\"\"
    Text
\"\"\"
function func(arg) end
func() = nothing
module M end
struct T end
mutable struct T end
for i = 1:1 end
while true end
begin end
sin()
struct S
    a
    b
    c
    d
    e
    f
    g
end
S(a,b,c,d,e,f,g)
""")

@test hover_test(0, 1) === nothing
@test hover_test(1, 1) !== nothing
@test hover_test(2, 1) !== nothing
@test hover_test(3, 1) !== nothing
@test hover_test(7, 11) !== nothing
@test hover_test(8, 2) !== nothing
@test hover_test(7, 16) !== nothing
@test hover_test(7, 20) !== nothing
@test hover_test(9, 11) !== nothing
@test hover_test(10, 11) !== nothing
@test hover_test(11, 18) !== nothing
@test hover_test(12, 14) !== nothing
@test hover_test(13, 13) !== nothing
@test hover_test(14, 7) !== nothing
@test hover_test(15, 2) !== nothing
@test hover_test(15, 5) === nothing
@test hover_test(25, 15) !== nothing


text:
###### readtimearray ############

function readtimearray(source; delim::Char = ',', meta = nothing,
                       format::AbstractString = "", header::Bool = true)
    cfile = readdlm(source, delim, header = header)
    if header
        cfile, hd = cfile
    end

    # remove empty lines if any
    inoempty = findall(s -> length(s) > 2, cfile[:, 1])
    cfile = cfile[inoempty, :]

    # create a DataFormat instance to improve performance
    df = isempty(format) ? nothing : DateFormat(format)

    time = cfile[1:end, 1]
    if length(time[1]) < 11
        # assuming Date not DateTime
        tstamps = df === nothing ?
            Date[Date(t) for t in time] :
            Date[Date(t, df) for t in time]
    else
        tstamps = df === nothing ?
            DateTime[DateTime(t) for t in time] :
            DateTime[DateTime(t, df) for t in time]
    end

    vals   = insertNaN(cfile[1:end, 2:end])
    cnames = header ? Symbol.(hd[2:end]) : gen_colnames(size(cfile, 2) - 1)
    TimeArray(tstamps, vals, cnames, meta)
end  # readtimearray

function insertNaN(aa::Array{Any, N}) where {N}
    for i in 1:size(aa, 1)
        for j in 1:size(aa, 2)
            if !isa(aa[i, j], Real)
                aa[i, j] = NaN
            end
        end
    end
    convert(Array{Float64, N}, aa)
end

function writetimearray(ta::TimeArray, fname::AbstractString;
                        delim::Char = ',', format::AbstractString = "",
                        header::Bool = true)
    open(fname, "w") do io
        if header
            strvals = join(colnames(ta), delim)
            write(io, string("Timestamp", delim, strvals, "\n"))
        end

        for i in eachindex(timestamp(ta))
            strvals = replace(join(values(ta)[i, :], delim), "NaN" => "")
            strtstamp = isempty(format) ?
                string(timestamp(ta)[i]) :
                Dates.format(timestamp(ta)[i], format)
            write(io, string(strtstamp, delim, strvals, "\n"))
        end
    end
end  # writetimearray


text:
"""
The constructor takes following keyword arguments:
- `populationSize`: The size of the population
- `F`: the differentiation (mutation) scale factor (default: 0.9). It's usually defined in range ``F \\in (0, 1+]``
- `n`: the number of differences used in the perturbation (default: 1)
- `selection`: the selection strategy function (default: [`random`](@ref))
- `recombination`: the recombination functions (default: [`BINX(0.5)`](@ref))
- `K`: the recombination scale factor (default: 0.5*(F+1))
- `metrics` is a collection of convergence metrics.
"""
Base.@kwdef struct DE <: AbstractOptimizer
    populationSize::Integer = 50
    F::Real = 0.9
    n::Integer = 1
    K::Real = 0.5*(F+1)
    selection = random
    recombination = BINX(0.5)
    metrics::ConvergenceMetrics = [AbsDiff(1e-12)]
end

population_size(method::DE) = method.populationSize
default_options(method::DE) = (iterations=1000,)
summary(m::DE) = "DE/$(m.selection)/$(m.n)/$(m.recombination)"
show(io::IO,m::DE) = print(io, summary(m))

mutable struct DEState <: AbstractOptimizerState
    iteration
    start_time
    stop_time
    metrics::ConvergenceMetrics

    N::Int
    population
    fittest
end
value(s::DEState) = objectives(s.fittest)
minimizer(s::DEState) = variables(s.fittest)
copy(s::DEState) = DEState(s.iteration, s.start_time, s.stop_time, copy(s.metrics), s.N, copy(s.population), copy(s.fittest))

function initial_state(method::DE, objfun, population, options)
    T = typeof(value(objfun))
    individual = first(population)
    N = length(variables(individual)) # 决策变量维度
    value!(objfun, population)
    _, i = findmin(objectives(population))
    fittest = population[i]

    return DEState(0, 0, 0, copy(method.metrics), N, copy(population), copy(fittest))
end

function update_state!(
    method::DE,
    state,
    objective,
    constraints,
    options
)
    offspring = copy(state.population)

    targets = method.selection(state.population, method.populationSize)

    for (i, target) in enumerate(targets)

        # mutation
        idxs = randexcl(options.rng, 1:method.populationSize, [i], 3)
        donors = copy(state.population[idxs])
        donor = donors[1]
        differentiation!(donors[1], donors[2:end]; F=method.F)

        # crossover
        trial, _ = method.recombination(donor, target, rng=options.rng)
        offspring[i] = trial
    end

    value!(objective, offspring)

    for i in 1:method.populationSize
        # o = apply!(constraints, offspring[i])
        v = objectives(offspring[i])
        if (v <= objectives(state.population[i]))
            state.population[i] = offspring[i]
            if (v < objectives(state.fittest))
                state.fittest = copy(offspring[i])
            end
        end
    end

    return false
end


text:
module UtilitiesForMRI

using LinearAlgebra, SparseArrays, CUDA, AbstractLinearOperators, FINUFFT, Flux, FFTW, Interpolations

const RealOrComplex{T<:Real} = Union{T,Complex{T}}

include("./abstract_types.jl")
include("./spatial_sampling.jl")
include("./kspace_sampling.jl")
include("./translations.jl")
include("./rotations.jl")
include("./nfft.jl")
include("./motion_parameter_utils.jl")
include("./optimization_utils.jl")
include("./imresize_utils.jl")

end

max_stars_repo_name: grizzuti/UtilitiesForMRI

max_issues_repo_name: grizzuti/UtilitiesForMRI

max_forks_repo_name: grizzuti/UtilitiesForMRI


text:
using SpecialFunctions

mutable struct GammaDist
    _a::Float64
    _b::Float64
    function GammaDist(a::Float64, b::Float64)
        new(a, b)
    end
end

function pdf(gamma::GammaDist, x::Float64)
    a = gamma._a
    b = gamma._b
    return (b^a) * (x^(a - 1)) * exp(-b * x) / gamma(a)
end

function pdf(dist::GammaDist, X::AbstractArray{Float64,1})
    a = dist._a
    b = dist._b
    return (b^a) .* (x .^ (a - 1)) .* exp.(-b .* X) ./ gamma(a)
end

max_stars_repo_name: soblin/prml-julia

max_issues_repo_name: soblin/prml-julia

max_forks_repo_name: soblin/prml-julia


text:
###########################
# Function Cell Interface #
###########################
"""
	kpca([;kwargs])

Constructs an untrained cell that when piped data inside, calculates the Kernel-PCA
projection matrix of the input data. We assume that `d` is the dimesionality
of the input dataset, `n` the number of samples

# Keyword arguments (same as in `MultivariateStats`)
  * `kernel` the kernel; a function that accepts two vector arguments `x` and `y`, and returns a scalar value (default `(x,y)->x'y`)
  * `solver` is the solver. Can be `:eig` uses `eigfact`, `:eigs` uses `eigs` (default `:eig`, for sparse data, `:eigs` is always used)
  smaller than the number of samples, `svd` otherwise)
  * `maxoutdim` is the number of output dimensions, default `min(<number of dimensions>,<number of samples>)`
  * `inverse` specified whether to perform calculation for inverse transform for non-precomputed kernelsi (default `false`)
  * `β` is a hyperparameter of the ridge regression that learns the inverse transform, when inverse is true (default `1.0`)
  * `tol` convergence tolerance for eigs solver (default `0.0`)
  * `maxiter` maximum number of iterations for eigs solver (default `300`)

# Kernels
List of the commonly used kernels:
    function 				description
    `(x,y)->x'y` 			 `Linear`
    `(x,y)->(x'y+c)^d` 			 `Polynomial`
    `(x,y)->exp(-γ*norm(x-y)^2.0)` 	 `Radial basis function (RBF)`

Read the `MultivariateStats.jl` documentation for more information.
"""
kpca(;kwargs...) = FunctionCell(kpca, (), ModelProperties(), kwtitle("Kernel-PCA", kwargs);kwargs...)

############################
# DataCell/Array Interface #
############################
"""
	kpca(x, [;kwargs])

Trains a function cell that when piped data into will compute the principal components of
the observations based on the .
"""
# Training
kpca(x::T where T<:CellData; kwargs...) = kpca(getx!(x);kwargs...)
kpca(x::T where T<:AbstractVector; kwargs...) = kpca(mat(x, LearnBase.ObsDim.Constant{2}());kwargs...)
kpca(x::T where T<:AbstractMatrix; kwargs...)  = begin

	kpcadata = fit(MultivariateStats.KernelPCA, getobs(x); kwargs...)

	# Build model properties
	modelprops = ModelProperties(MultivariateStats.indim(kpcadata),
				     MultivariateStats.outdim(kpcadata))

	# Returned trained cell
	FunctionCell(kpca, Model(kpcadata, modelprops), kwtitle("Kernel-PCA",kwargs))
end

# Execution
kpca(x::T where T<:CellData, model::Model{<:MultivariateStats.KernelPCA}) =
	datacell(kpca(getx!(x), model), gety(x))
kpca(x::T where T<:AbstractVector, model::Model{<:MultivariateStats.KernelPCA}) =
	kpca(mat(x, LearnBase.ObsDim.Constant{2}()), model)
kpca(x::T where T<:AbstractMatrix, model::Model{<:MultivariateStats.KernelPCA}) =
	MultivariateStats.transform(model.data, getobs(x))

"""
	kpcar(M::Cell), where `M=kpca(x)`

Trains a cell based on a previously trained Kernel-PCA function cell. When piped data into,
the returned output will be the reconstruction of the observations in the original
space. Obviously, the input data (e.g. data to be reconstructed) is expected to be
principal components extracted with `M`.

Read the `MultivariateStats.jl` documentation for more information.
```
"""
# Training for the reconstruction tranform (fixed cell)
kpcar(x::T where T<:CellFunT{<:Model{<:MultivariateStats.KernelPCA}}) = begin

	# Build model properties
	modelprops = ModelProperties(getx(x).properties.odim, getx(x).properties.idim)

	FunctionCell(kpcar, Model(getx(x).data, modelprops),"Kernel-PCA: reconstruct")
end

# Execution
kpcar(x::T where T<:CellData, model::Model{<:MultivariateStats.KernelPCA}) = datacell(kpcar(getx!(x), model), gety(x))
kpcar(x::T where T<:AbstractVector, model::Model{<:MultivariateStats.KernelPCA}) = kpcar(mat(x, LearnBase.ObsDim.Constant{2}()), model)
kpcar(x::T where T<:AbstractMatrix, model::Model{<:MultivariateStats.KernelPCA})::Matrix{Float64} = MultivariateStats.reconstruct(model.data, getobs(x))