text:
@doc raw"""
    Rotations{N} <: AbstractManifold{ℝ}

The special orthogonal manifold $\mathrm{SO}(n)$ represented by $n × n$
real-valued orthogonal matrices with determinant $+1$ is the manifold of `Rotations`,
since these matrices represent all rotations of points in $ℝ^n$.

# Constructor

    Rotations(n)

Generate the $\mathrm{SO}(n) \subset ℝ^{n × n}$
"""
struct Rotations{N} <: AbstractManifold{ℝ} end

Rotations(n::Int) = Rotations{n}()

"""
    NormalRotationDistribution(M::Rotations, d::Distribution, x::TResult)

Distribution that returns a random point on the manifold [`Rotations`](@ref)
`M`. Random point is generated using base distribution `d` and the type
of the result is adjusted to `TResult`.

See [`normal_rotation_distribution`](@ref) for details.
"""
struct NormalRotationDistribution{TResult,TM<:Rotations,TD<:Distribution} <:
       MPointDistribution{TM}
    manifold::TM
    distr::TD
end

function NormalRotationDistribution(
    M::Rotations,
    d::Distribution,
    x::TResult,
) where {TResult}
    return NormalRotationDistribution{TResult,typeof(M),typeof(d)}(M, d)
end

@doc raw"""
    angles_4d_skew_sym_matrix(A)

The Lie algebra of [`Rotations(4)`](@ref) in $ℝ^{4 × 4}$, $𝔰𝔬(4)$, consists of $4 × 4$
skew-symmetric matrices. The unique imaginary components of their eigenvalues are the
angles of the two plane rotations. This function computes these more efficiently than
`eigvals`.

By convention, the returned values are sorted in decreasing order
(corresponding to the same ordering of _angles_ as
[`cos_angles_4d_rotation_matrix`](@ref)).
"""
function angles_4d_skew_sym_matrix(A)
    @assert size(A) == (4, 4)
    @inbounds begin
        halfb = (A[1, 2]^2 + A[1, 3]^2 + A[2, 3]^2 + A[1, 4]^2 + A[2, 4]^2 + A[3, 4]^2) / 2
        c = (A[1, 2] * A[3, 4] - A[1, 3] * A[2, 4] + A[1, 4] * A[2, 3])^2
    end
    sqrtdisc = sqrt(halfb^2 - c)
    return sqrt(halfb + sqrtdisc), sqrt(halfb - sqrtdisc)
end

"""
    check_point(M, p; kwargs...)

Check whether `p` is a valid point on the [`Rotations`](@ref) `M`,
i.e. is an array of size [`manifold_dimension`](@ref)`(M)` and represents a
valid rotation.
The tolerance for the last test can be set using the `kwargs...`.
"""
function check_point(M::Rotations{N}, p; kwargs...) where {N}
    if size(p) != (N, N)
        return DomainError(
            size(p),
            "The point $(p) does not lie on $M, since its size is not $((N, N)).",
        )
    end
    if !isapprox(det(p), 1; kwargs...)
        return DomainError(det(p), "The determinant of $p has to be +1 but it is $(det(p))")
    end
    if !isapprox(transpose(p) * p, one(p); kwargs...)
        return DomainError(
            norm(transpose(p) * p - one(p)),
            "$p must be orthogonal but it's not at kwargs $kwargs",
        )
    end
    return nothing
end

"""
    check_vector(M, p, X; kwargs... )

Check whether `X` is a tangent vector to `p` on the [`Rotations`](@ref)
space `M`, i.e. after [`check_point`](@ref)`(M,p)`, `X` has to be of same
dimension and orthogonal to `p`.
The tolerance for the last test can be set using the `kwargs...`.
"""
function check_vector(M::Rotations{N}, p, X; kwargs...) where {N}
    return check_point(SkewSymmetricMatrices(N), X; kwargs...)
end

@doc raw"""
    embed(M::Rotations{N}, p, X)

Embed the tangent vector `X` at point `p` in `M` from
its Lie algebra representation (set of skew matrices) into the
Riemannian submanifold representation

The formula reads
```math
X_{\text{embedded}} = p * X
```
"""
embed(::Rotations, p, X)

function embed!(::Rotations, Y, p, X)
    return mul!(Y, p, X)
end

@doc raw"""
    cos_angles_4d_rotation_matrix(R)

4D rotations can be described by two orthogonal planes that are unchanged by
the action of the rotation (vectors within a plane rotate only within the
plane). The cosines of the two angles $α,β$ of rotation about these planes may be
obtained from the distinct real parts of the eigenvalues of the rotation
matrix. This function computes these more efficiently by solving the system

```math
\begin{aligned}
\cos α + \cos β &= \frac{1}{2} \operatorname{tr}(R)\\
\cos α + \cos β &= \frac{1}{8} \operatorname{tr}(R)^2
                 - \frac{1}{16} \operatorname{tr}((R - R^T)^2) - 1.
\end{aligned}
```

By convention, the returned values are sorted in increasing order. See
[`angles_4d_skew_sym_matrix`](@ref).
"""
function cos_angles_4d_rotation_matrix(R)
    trR = tr(R)
    a = trR / 4
    b = sqrt(clamp(tr((R .- transpose(R))^2) / 16 - a^2 + 1, 0, Inf))
    return (a + b, a - b)
end

@doc raw"""
    exp(M::Rotations, p, X)

Compute the exponential map on the [`Rotations`](@ref) from `p` into direction
`X`, i.e.

````math
\exp_p X = p \operatorname{Exp}(X),
````

where $\operatorname{Exp}(X)$  denotes the matrix exponential of $X$.

    exp(M::Rotations{4}, p, X)

Compute the exponential map of tangent vector `X` at point `p` from $\mathrm{SO}(4)$
manifold `M`.

The algorithm used is a more numerically stable form of those proposed in
[^Gallier2002] and [^Andrica2013].

[^Gallier2002]:
    > Gallier J.; Xu D.; Computing exponentials of skew-symmetric matrices
    > and logarithms of orthogonal matrices.
    > International Journal of Robotics and Automation (2002), 17(4), pp. 1-11.
    > [pdf](http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.35.3205).

[^Andrica2013]:
    > Andrica D.; Rohan R.-A.; Computing the Rodrigues coefficients of the
    > exponential map of the Lie groups of matrices.
    > Balkan Journal of Geometry and Its Applications (2013), 18(2), pp. 1-2.
    > [pdf](https://www.emis.de/journals/BJGA/v18n2/B18-2-an.pdf).
"""
exp(::Rotations, ::Any...)

exp!(M::Rotations, q, p, X) = copyto!(q, p * exp(X))
function exp!(M::Rotations{2}, q, p, X)
    @assert size(q) == (2, 2)
    θ = get_coordinates(M, p, X, DefaultOrthogonalBasis())[1]
    sinθ, cosθ = sincos(θ)
    return copyto!(q, p * SA[cosθ -sinθ; sinθ cosθ])
end
function exp!(M::Rotations{3}, q, p, X)
    θ = norm(M, p, X) / sqrt(2)
    if θ ≈ 0
        a = 1 - θ^2 / 6
        b = θ / 2
    else
        a = sin(θ) / θ
        b = (1 - cos(θ)) / θ^2
    end
    pinvq = I + a .* X .+ b .* (X^2)
    return copyto!(q, p * pinvq)
end
function exp!(M::Rotations{4}, q, p, X)
    T = eltype(X)
    α, β = angles_4d_skew_sym_matrix(X)
    sinα, cosα = sincos(α)
    sinβ, cosβ = sincos(β)
    α² = α^2
    β² = β^2
    Δ = β² - α²
    if !isapprox(Δ, 0; atol=1e-6)  # Case α > β ≥ 0
        sincα = sinα / α
        sincβ = β == 0 ? one(T) : sinβ / β
        a₀ = (β² * cosα - α² * cosβ) / Δ
        a₁ = (β² * sincα - α² * sincβ) / Δ
        a₂ = (cosα - cosβ) / Δ
        a₃ = (sincα - sincβ) / Δ
    elseif α == 0 # Case α = β = 0
        a₀ = one(T)
        a₁ = one(T)
        a₂ = T(1 / 2)
        a₃ = T(1 / 6)
    else  # Case α ⪆ β ≥ 0, α ≠ 0
        sincα = sinα / α
        r = β / α
        c = 1 / (1 + r)
        d = α * (α - β) / 2
        if α < 1e-2
            e = @evalpoly(α², T(1 / 3), T(-1 / 30), T(1 / 840), T(-1 / 45360))
        else
            e = (sincα - cosα) / α²
        end
        a₀ = (α * sinα + (1 + r - d) * cosα) * c
        a₁ = ((3 - d) * sincα - (2 - r) * cosα) * c
        a₂ = (sincα - (1 - r) / 2 * cosα) * c
        a₃ = (e + (1 - r) * (e - sincα / 2)) * c
    end

    X² = X * X
    X³ = X² * X
    pinvq = a₀ * I + a₁ .* X .+ a₂ .* X² .+ a₃ .* X³
    return copyto!(q, p * pinvq)
end

@doc raw"""
    get_coordinates(M::Rotations, p, X)

Extract the unique tangent vector components $X^i$ at point `p` on [`Rotations`](@ref)
$\mathrm{SO}(n)$ from the matrix representation `X` of the tangent
vector.

The basis on the Lie algebra $𝔰𝔬(n)$ is chosen such that
for $\mathrm{SO}(2)$, $X^1 = θ = X_{21}$ is the angle of rotation, and
for $\mathrm{SO}(3)$, $(X^1, X^2, X^3) = (X_{32}, X_{13}, X_{21}) = θ u$ is the
angular velocity and axis-angle representation, where $u$ is the unit vector
along the axis of rotation.

For $\mathrm{SO}(n)$ where $n ≥ 4$, the additional elements of $X^i$ are
$X^{j (j - 3)/2 + k + 1} = X_{jk}$, for $j ∈ [4,n], k ∈ [1,j)$.
"""
get_coordinates(::Rotations, ::Any...)
get_coordinates(::Rotations{2}, p, X, ::DefaultOrthogonalBasis{ℝ,TangentSpaceType}) = [X[2]]

function get_coordinates!(
    ::Rotations{2},
    Xⁱ,
    p,
    X,
    ::DefaultOrthogonalBasis{ℝ,TangentSpaceType},
)
    Xⁱ[1] = X[2]
    return Xⁱ
end
function get_coordinates!(
    M::Rotations{N},
    Xⁱ,
    p,
    X,
    ::DefaultOrthogonalBasis{ℝ,TangentSpaceType},
) where {N}
    @inbounds begin
        Xⁱ[1] = X[3, 2]
        Xⁱ[2] = X[1, 3]
        Xⁱ[3] = X[2, 1]

        k = 4
        for i in 4:N, j in 1:(i - 1)
            Xⁱ[k] = X[i, j]
            k += 1
        end
    end
    return Xⁱ
end
function get_coordinates!(
    M::Rotations{N},
    Xⁱ,
    p,
    X,
    ::DefaultOrthonormalBasis{ℝ,TangentSpaceType},
) where {N}
    T = Base.promote_eltype(p, X)
    get_coordinates!(M, Xⁱ, p, X, DefaultOrthogonalBasis())
    Xⁱ .*= sqrt(T(2))
    return Xⁱ
end

@doc raw"""
    get_vector(M::Rotations, p, Xⁱ, B::DefaultOrthogonalBasis)

Convert the unique tangent vector components `Xⁱ` at point `p` on [`Rotations`](@ref)
group $\mathrm{SO}(n)$ to the matrix representation $X$ of the tangent vector. See
[`get_coordinates`](@ref get_coordinates(::Rotations, ::Any...)) for the conventions used.
"""
get_vector(::Rotations, ::Any...)

function get_vector!(
    M::Rotations{2},
    X,
    p,
    Xⁱ,
    B::DefaultOrthogonalBasis{ℝ,TangentSpaceType},
)
    return get_vector!(M, X, p, Xⁱ[1], B)
end
function get_vector!(
    M::Rotations{2},
    X,
    p,
    Xⁱ::Real,
    ::DefaultOrthogonalBasis{ℝ,TangentSpaceType},
)
    @assert length(X) == 4
    @inbounds begin
        X[1] = 0
        X[2] = Xⁱ
        X[3] = -Xⁱ
        X[4] = 0
    end
    return X
end
function get_vector!(
    M::Rotations{N},
    X,
    p,
    Xⁱ,
    ::DefaultOrthogonalBasis{ℝ,TangentSpaceType},
) where {N}
    @assert size(X) == (N, N)
    @assert length(Xⁱ) == manifold_dimension(M)
    @inbounds begin
        X[1, 1] = 0
        X[1, 2] = -Xⁱ[3]
        X[1, 3] = Xⁱ[2]
        X[2, 1] = Xⁱ[3]
        X[2, 2] = 0
        X[2, 3] = -Xⁱ[1]
        X[3, 1] = -Xⁱ[2]
        X[3, 2] = Xⁱ[1]
        X[3, 3] = 0
        k = 4
        for i in 4:N
            for j in 1:(i - 1)
                X[i, j] = Xⁱ[k]
                X[j, i] = -Xⁱ[k]
                k += 1
            end
            X[i, i] = 0
        end
    end
    return X
end
function get_vector!(M::Rotations, X, p, Xⁱ, ::DefaultOrthonormalBasis{ℝ,TangentSpaceType})
    T = Base.promote_eltype(p, X)
    get_vector!(M, X, p, Xⁱ, DefaultOrthogonalBasis())
    X ./= sqrt(T(2))
    return X
end

@doc raw"""
    injectivity_radius(M::Rotations)
    injectivity_radius(M::Rotations, p)

Return the injectivity radius on the [`Rotations`](@ref) `M`, which is globally

````math
    \operatorname{inj}_{\mathrm{SO}(n)}(p) = π\sqrt{2}.
````

    injectivity_radius(M::Rotations, p, ::PolarRetraction)

Return the radius of injectivity for the [`PolarRetraction`](@ref) on the
[`Rotations`](@ref) `M` which is $\frac{π}{\sqrt{2}}$.
"""
injectivity_radius(::Rotations) = π * sqrt(2.0)
injectivity_radius(::Rotations, ::ExponentialRetraction) = π * sqrt(2.0)
eval(
    quote
        @invoke_maker 1 AbstractManifold injectivity_radius(
            M::Rotations,
            rm::AbstractRetractionMethod,
        )
    end,
)
injectivity_radius(::Rotations, ::Any) = π * sqrt(2.0)
injectivity_radius(::Rotations, ::Any, ::ExponentialRetraction) = π * sqrt(2.0)
injectivity_radius(::Rotations, ::PolarRetraction) = π / sqrt(2.0)
injectivity_radius(::Rotations, p, ::PolarRetraction) = π / sqrt(2.0)

@doc raw"""
    inner(M::Rotations, p, X, Y)

Compute the inner product of the two tangent vectors `X`, `Y` from the tangent
plane at `p` on the special orthogonal space `M=`$\mathrm{SO}(n)$ using the
restriction of the metric from the embedding, i.e.

````math
g_p(X, Y) = \operatorname{tr}(X^\mathrm{T} Y),
````

Tangent vectors are represented by matrices.
"""
inner(M::Rotations, p, X, Y) = dot(X, Y)

@doc raw"""
    inverse_retract(M, p, q, ::PolarInverseRetraction)

Compute a vector from the tangent space $T_p\mathrm{SO}(n)$
of the point `p` on the [`Rotations`](@ref) manifold `M`
with which the point `q` can be reached by the
[`PolarRetraction`](@ref) from the point `p` after time 1.

The formula reads
````math
\operatorname{retr}^{-1}_p(q)
= -\frac{1}{2}(p^{\mathrm{T}}qs - (p^{\mathrm{T}}qs)^{\mathrm{T}})
````

where $s$ is the solution to the Sylvester equation

$p^{\mathrm{T}}qs + s(p^{\mathrm{T}}q)^{\mathrm{T}} + 2I_n = 0.$
"""
inverse_retract(::Rotations, ::Any, ::Any, ::PolarInverseRetraction)

@doc raw"""
    inverse_retract(M::Rotations, p, q, ::QRInverseRetraction)

Compute a vector from the tangent space $T_p\mathrm{SO}(n)$ of the point `p` on the
[`Rotations`](@ref) manifold `M` with which the point `q` can be reached by the
[`QRRetraction`](@ref) from the point `q` after time 1.
"""
inverse_retract(::Rotations, ::Any, ::Any, ::QRInverseRetraction)

function inverse_retract!(M::Rotations, X, p, q, method::PolarInverseRetraction)
    A = transpose(p) * q
    Amat = A isa StaticMatrix ? A : convert(Matrix, A)
    H = copyto!(allocate(Amat), -2I)
    try
        B = lyap(A, H)
        mul!(X, A, B)
    catch e
        if isa(e, LinearAlgebra.LAPACKException)
            throw(OutOfInjectivityRadiusError())
        else
            rethrow()
        end
    end
    return project!(M, X, p, X)
end
function inverse_retract!(M::Rotations{N}, X, p, q, ::QRInverseRetraction) where {N}
    A = transpose(p) * q
    R = zero(X)
    for i in 1:N
        b = zeros(i)
        b[end] = 1
        b[1:(end - 1)] = -transpose(R[1:(i - 1), 1:(i - 1)]) * A[i, 1:(i - 1)]
        R[1:i, i] = A[1:i, 1:i] \ b
    end
    mul!(X, A, R)
    return project!(M, X, p, X)
end

@doc raw"""
    log(M::Rotations, p, q)

Compute the logarithmic map on the [`Rotations`](@ref) manifold
`M`$=\mathrm{SO}(n)$, which is given by
```math
\log_p q =
  \frac{1}{2} \bigl(\operatorname{Log}(p^{\mathrm{T}}q)
  - (\operatorname{Log}(p^{\mathrm{T}}q)^{\mathrm{T}}),
```
where $\operatorname{Log}$ denotes the matrix logarithm.

For antipodal rotations the function returns deterministically one of the tangent vectors
that point at `q`.
"""
log(::Rotations, ::Any...)

function log!(M::Rotations, X, p, q)
    U = transpose(p) * q
    X .= real(log_safe(U))
    return project!(M, X, p, X)
end
function log!(M::Rotations{2}, X, p, q)
    U = transpose(p) * q
    @assert size(U) == (2, 2)
    @inbounds θ = atan(U[2], U[1])
    return get_vector!(M, X, p, θ, DefaultOrthogonalBasis())
end
function log!(M::Rotations{3}, X, p, q)
    U = transpose(p) * q
    cosθ = (tr(U) - 1) / 2
    if cosθ ≈ -1
        eig = eigen_safe(U)
        ival = findfirst(λ -> isapprox(λ, 1), eig.values)
        inds = SVector{3}(1:3)
        ax = eig.vectors[inds, ival]
        return get_vector!(M, X, p, π * ax, DefaultOrthogonalBasis())
    end
    X .= U ./ usinc_from_cos(cosθ)
    return project!(M, X, p, X)
end
function log!(M::Rotations{4}, X, p, q)
    U = transpose(p) * q
    cosα, cosβ = cos_angles_4d_rotation_matrix(U)
    α = acos(clamp(cosα, -1, 1))
    β = acos(clamp(cosβ, -1, 1))
    if α ≈ π && β ≈ 0
        A² = Symmetric((U - I) ./ 2)
        P = eigvecs(A²)
        E = similar(U)
        fill!(E, 0)
        α = acos(clamp(cosα, -1, 1))
        @inbounds begin
            E[2, 1] = -α
            E[1, 2] = α
        end
        copyto!(X, P * E * transpose(P))
    else
        copyto!(X, real(log_safe(U)))
    end
    return project!(M, X, p, X)
end

@doc raw"""
    manifold_dimension(M::Rotations)

Return the dimension of the manifold $\mathrm{SO}(n)$, i.e.
```math
\dim_{\mathrm{SO}(n)} = \frac{n(n-1)}{2}.
```
"""
manifold_dimension(::Rotations{N}) where {N} = div(N * (N - 1), 2)

"""
    mean(
        M::Rotations,
        x::AbstractVector,
        [w::AbstractWeights,]
        method = GeodesicInterpolationWithinRadius(π/2/√2);
        kwargs...,
    )

Compute the Riemannian [`mean`](@ref mean(M::AbstractManifold, args...)) of `x` using
[`GeodesicInterpolationWithinRadius`](@ref).
"""
mean(::Rotations, ::Any)

function Statistics.mean!(M::Rotations, q, x::AbstractVector, w::AbstractVector; kwargs...)
    return mean!(M, q, x, w, GeodesicInterpolationWithinRadius(π / 2 / √2); kwargs...)
end

@doc raw"""
    norm(M::Rotations, p, X)

Compute the norm of a tangent vector `X` from the tangent space at `p` on the
[`Rotations`](@ref) `M`. The formula reads

````math
\lVert X \rVert_p = \lVert X \rVert,
````

i.e. the Frobenius norm of `X`, where tangent vectors are represented by
elements from the Lie algebra.
"""
LinearAlgebra.norm(M::Rotations, p, X) = norm(X)

@doc raw"""
    normal_rotation_distribution(M::Rotations, p, σ::Real)

Return a random point on the manifold [`Rotations`](@ref) `M`
by generating a (Gaussian) random orthogonal matrix with determinant $+1$. Let

$QR = A$

be the QR decomposition of a random matrix $A$, then the formula reads

$p = QD$

where $D$ is a diagonal matrix with the signs of the diagonal entries of $R$,
i.e.

$D_{ij}=\begin{cases} \operatorname{sgn}(R_{ij}) & \text{if} \; i=j \\ 0 & \, \text{otherwise} \end{cases}.$

It can happen that the matrix gets -1 as a determinant. In this case, the first
and second columns are swapped.

The argument `p` is used to determine the type of returned points.
"""
function normal_rotation_distribution(M::Rotations{N}, p, σ::Real) where {N}
    d = Distributions.MvNormal(zeros(N * N), σ)
    return NormalRotationDistribution(M, d, p)
end

@doc raw"""
    project(M::Rotations, p; check_det = true)

Project `p` to the nearest point on manifold `M`.

Given the singular value decomposition $p = U Σ V^\mathrm{T}$, with the
singular values sorted in descending order, the projection is

````math
\operatorname{proj}_{\mathrm{SO}(n)}(p) =
U\operatorname{diag}\left[1,1,…,\det(U V^\mathrm{T})\right] V^\mathrm{T}
````

The diagonal matrix ensures that the determinant of the result is $+1$.
If `p` is expected to be almost special orthogonal, then you may avoid this
check with `check_det = false`.
"""
project(::Rotations, ::Any)

function project!(::Rotations{N}, q, p; check_det=true) where {N}
    F = svd(p)
    mul!(q, F.U, F.Vt)
    if check_det && det(q) < 0
        d = similar(F.S)
        @inbounds fill!(view(d, 1:(N - 1)), 1)
        @inbounds d[N] = -1
        copyto!(q, F.U * Diagonal(d) * F.Vt)
    end
    return q
end

@doc raw"""
    project(M::Rotations, p, X)

Project the matrix `X` onto the tangent space by making `X` skew symmetric,

````math
\operatorname{proj}_p(X) = \frac{X-X^{\mathrm{T}}}{2},
````

where tangent vectors are represented by elements from the Lie group
"""
project(::Rotations, ::Any, ::Any)

project!(M::Rotations{N}, Y, p, X) where {N} = project!(SkewSymmetricMatrices(N), Y, X)

@doc raw"""
    representation_size(M::Rotations)

Return the `size()` of a point on the [`Rotations`](@ref) `M`, i.e. for the
$\mathrm{SO}(n)$ it's `(n,n)`.
"""
@generated representation_size(::Rotations{N}) where {N} = (N, N)

function Random.rand(
    rng::AbstractRNG,
    d::NormalRotationDistribution{TResult,Rotations{N}},
) where {TResult,N}
    return if N == 1
        convert(TResult, ones(1, 1))
    else
        A = reshape(rand(rng, d.distr), (N, N))
        convert(TResult, _fix_random_rotation(A))
    end
end

function Distributions._rand!(
    rng::AbstractRNG,
    d::NormalRotationDistribution{TResult,Rotations{N}},
    x::AbstractArray{<:Real},
) where {TResult,N}
    return copyto!(x, rand(rng, d))
end

function _fix_random_rotation(A::AbstractMatrix)
    s = diag(sign.(qr(A).R))
    D = Diagonal(s)
    C = qr(A).Q * D
    if det(C) < 0
        C[:, [1, 2]] = C[:, [2, 1]]
    end
    return C
end

@doc raw"""
    retract(M::Rotations, p, X, ::PolarRetraction)

Compute the SVD-based retraction on the [`Rotations`](@ref) `M` from `p` in direction `X`
(as an element of the Lie group) and is a second-order approximation of the exponential map.
Let

````math
USV = p + pX
````

be the singular value decomposition, then the formula reads

````math
\operatorname{retr}_p X = UV^\mathrm{T}.
````
"""
retract(::Rotations, ::Any, ::Any, ::PolarRetraction)

@doc raw"""
    retract(M, p, X, ::QRRetraction)

Compute the QR-based retraction on the [`Rotations`](@ref) `M` from `p` in direction `X`
(as an element of the Lie group), which is a first-order approximation of the exponential map.

This is also the default retraction on the [`Rotations`](@ref)
"""
retract(::Rotations, ::Any, ::Any, ::QRRetraction)

function retract!(M::Rotations, q::AbstractArray{T}, p, X, method::QRRetraction) where {T}
    A = p + p * X
    qr_decomp = qr(A)
    d = diag(qr_decomp.R)
    D = Diagonal(sign.(d .+ convert(T, 0.5)))
    return copyto!(q, qr_decomp.Q * D)
end
function retract!(M::Rotations, q, p, X, method::PolarRetraction)
    A = p + p * X
    return project!(M, q, A; check_det=false)
end

Base.show(io::IO, ::Rotations{N}) where {N} = print(io, "Rotations($(N))")

Distributions.support(d::NormalRotationDistribution) = MPointSupport(d.manifold)

@doc raw"""
    vector_transport_direction(M::Rotations, p, X, d)

Compute parallel transport of vector `X` tangent at `p` on the [`Rotations`](@ref)
manifold in the direction `d`. The formula, provided in [^Rentmeesters], reads:

```math
\mathcal P_{q\gets p}X = q^\mathrm{T}p \operatorname{Exp}(d/2) X \operatorname{Exp}(d/2)
```
where ``q=\exp_p d``.

The formula simplifies to identity for 2-D rotations.

[^Rentmeesters]:
    > Rentmeesters Q., “A gradient method for geodesic data fitting on some symmetric
    > Riemannian manifolds,” in 2011 50th IEEE Conference on Decision and Control and
    > European Control Conference, Dec. 2011, pp. 7141–7146. doi: 10.1109/CDC.2011.6161280.
"""
vector_transport_direction(M::Rotations, p, X, d)

function vector_transport_direction!(M::Rotations, Y, p, X, d, ::ParallelTransport)
    expdhalf = exp(d / 2)
    q = exp(M, p, d)
    return copyto!(Y, transpose(q) * p * expdhalf * X * expdhalf)
end
function vector_transport_direction!(M::Rotations{2}, Y, p, X, d, ::ParallelTransport)
    return copyto!(Y, X)
end

function vector_transport_to!(M::Rotations, Y, p, X, q, ::ParallelTransport)
    d = log(M, p, q)
    expdhalf = exp(d / 2)
    return copyto!(Y, transpose(q) * p * expdhalf * X * expdhalf)
end
function vector_transport_to!(M::Rotations{2}, Y, p, X, q, ::ParallelTransport)
    return copyto!(Y, X)
end

@doc raw"""
    zero_vector(M::Rotations, p)

Return the zero tangent vector from the tangent space art `p` on the [`Rotations`](@ref)
as an element of the Lie group, i.e. the zero matrix.
"""
zero_vector(M::Rotations, p) = zero(p)

zero_vector!(M::Rotations, X, p) = fill!(X, 0)


text:
using Mocking

create_cluster_1_patch = @patch function create_cluster(a...; b...)
    return Dict{String, Any}("cluster" => Dict{String, Any}("runningTasksCount" => 0, "status" => "ACTIVE", "activeServicesCount" => 0, "pendingTasksCount" => 0, "clusterArn" => "arn:aws:ecs:us-east-1:012345678910:cluster/my_cluster", "registeredContainerInstancesCount" => 0, "clusterName" => "my_cluster"))
end

create_service_1_patch = @patch function create_service(a...; b...)
    return Dict{String, Any}("service" => Dict{String, Any}("taskDefinition" => "arn:aws:ecs:us-east-1:012345678910:task-definition/hello_world:6", "deployments" => Any[Dict{String, Any}("runningCount" => 0, "status" => "PRIMARY", "taskDefinition" => "arn:aws:ecs:us-east-1:012345678910:task-definition/hello_world:6", "id" => "ecs-svc/9223370564342348388", "desiredCount" => 10, "updatedAt" => "2016-08-29T16:13:47.298Z", "pendingCount" => 0, "createdAt" => "2016-08-29T16:13:47.298Z"), Dict{String, Any}("runningCount" => 0, "status" => "ACTIVE", "taskDefinition" => "arn:aws:ecs:us-east-1:012345678910:task-definition/hello_world:6", "id" => "ecs-svc/9223370564343611322", "desiredCount" => 0, "updatedAt" => "2016-08-29T16:11:38.941Z", "pendingCount" => 0, "createdAt" => "2016-08-29T15:52:44.481Z")], "pendingCount" => 0, "createdAt" => "2016-08-29T16:13:47.298Z", "events" => Any[], "status" => "ACTIVE", "clusterArn" => "arn:aws:ecs:us-east-1:012345678910:cluster/default", "serviceArn" => "arn:aws:ecs:us-east-1:012345678910:service/ecs-simple-service", "deploymentConfiguration" => Dict{String, Any}("maximumPercent" => 200, "minimumHealthyPercent" => 100), "runningCount" => 0, "serviceName" => "ecs-simple-service", "loadBalancers" => Any[], "desiredCount" => 10))
end

create_service_2_patch = @patch function create_service(a...; b...)
    return Dict{String, Any}("service" => Dict{String, Any}("roleArn" => "arn:aws:iam::012345678910:role/ecsServiceRole", "taskDefinition" => "arn:aws:ecs:us-east-1:012345678910:task-definition/console-sample-app-static:6", "deployments" => Any[Dict{String, Any}("runningCount" => 0, "status" => "PRIMARY", "taskDefinition" => "arn:aws:ecs:us-east-1:012345678910:task-definition/console-sample-app-static:6", "id" => "ecs-svc/9223370564343000923", "desiredCount" => 10, "updatedAt" => "2016-08-29T16:02:54.884Z", "pendingCount" => 0, "createdAt" => "2016-08-29T16:02:54.884Z")], "pendingCount" => 0, "createdAt" => "2016-08-29T16:02:54.884Z", "events" => Any[], "status" => "ACTIVE", "clusterArn" => "arn:aws:ecs:us-east-1:012345678910:cluster/default", "serviceArn" => "arn:aws:ecs:us-east-1:012345678910:service/ecs-simple-service-elb", "deploymentConfiguration" => Dict{String, Any}("maximumPercent" => 200, "minimumHealthyPercent" => 100), "runningCount" => 0, "serviceName" => "ecs-simple-service-elb", "loadBalancers" => Any[Dict{String, Any}("loadBalancerName" => "EC2Contai-EcsElast-15DCDAURT3ZO2", "containerName" => "simple-app", "containerPort" => 80)], "desiredCount" => 10))
end

delete_account_setting_1_patch = @patch function delete_account_setting(a...; b...)
    return Dict{String, Any}("setting" => Dict{String, Any}("name" => "serviceLongArnFormat", "value" => "enabled", "principalArn" => "arn:aws:iam::<aws_account_id>:user/principalName"))
end

delete_account_setting_2_patch = @patch function delete_account_setting(a...; b...)
    return Dict{String, Any}("setting" => Dict{String, Any}("name" => "containerInstanceLongArnFormat", "value" => "enabled", "principalArn" => "arn:aws:iam::<aws_account_id>:user/principalName"))
end

delete_cluster_1_patch = @patch function delete_cluster(a...; b...)
    return Dict{String, Any}("cluster" => Dict{String, Any}("runningTasksCount" => 0, "status" => "INACTIVE", "activeServicesCount" => 0, "pendingTasksCount" => 0, "clusterArn" => "arn:aws:ecs:us-east-1:012345678910:cluster/my_cluster", "registeredContainerInstancesCount" => 0, "clusterName" => "my_cluster"))
end

delete_service_1_patch = @patch function delete_service(a...; b...)
    return Dict{String, Any}()
end

deregister_container_instance_1_patch = @patch function deregister_container_instance(a...; b...)
    return Dict{String, Any}()
end

describe_clusters_1_patch = @patch function describe_clusters(a...; b...)
    return Dict{String, Any}("clusters" => Any[Dict{String, Any}("status" => "ACTIVE", "clusterArn" => "arn:aws:ecs:us-east-1:aws_account_id:cluster/default", "clusterName" => "default")], "failures" => Any[])
end

describe_container_instances_1_patch = @patch function describe_container_instances(a...; b...)
    return Dict{String, Any}("containerInstances" => Any[Dict{String, Any}("runningTasksCount" => 1, "ec2InstanceId" => "i-807f3249", "remainingResources" => Any[Dict{String, Any}("name" => "CPU", "doubleValue" => 0.0, "integerValue" => 1948, "type" => "INTEGER", "longValue" => 0), Dict{String, Any}("name" => "MEMORY", "doubleValue" => 0.0, "integerValue" => 3668, "type" => "INTEGER", "longValue" => 0), Dict{String, Any}("name" => "PORTS", "doubleValue" => 0.0, "integerValue" => 0, "type" => "STRINGSET", "longValue" => 0, "stringSetValue" => Any["2376", "22", "80", "51678", "2375"])], "status" => "ACTIVE", "agentConnected" => true, "containerInstanceArn" => "arn:aws:ecs:us-east-1:012345678910:container-instance/f2756532-8f13-4d53-87c9-aed50dc94cd7", "pendingTasksCount" => 0, "registeredResources" => Any[Dict{String, Any}("name" => "CPU", "doubleValue" => 0.0, "integerValue" => 2048, "type" => "INTEGER", "longValue" => 0), Dict{String, Any}("name" => "MEMORY", "doubleValue" => 0.0, "integerValue" => 3768, "type" => "INTEGER", "longValue" => 0), Dict{String, Any}("name" => "PORTS", "doubleValue" => 0.0, "integerValue" => 0, "type" => "STRINGSET", "longValue" => 0, "stringSetValue" => Any["2376", "22", "51678", "2375"])])], "failures" => Any[])
end

describe_services_1_patch = @patch function describe_services(a...; b...)
    return Dict{String, Any}("services" => Any[Dict{String, Any}("taskDefinition" => "arn:aws:ecs:us-east-1:012345678910:task-definition/hello_world:6", "deployments" => Any[Dict{String, Any}("runningCount" => 0, "status" => "PRIMARY", "taskDefinition" => "arn:aws:ecs:us-east-1:012345678910:task-definition/hello_world:6", "id" => "ecs-svc/9223370564341623665", "desiredCount" => 1, "updatedAt" => "2016-08-29T16:25:52.130Z", "pendingCount" => 0, "createdAt" => "2016-08-29T16:25:52.130Z")], "pendingCount" => 0, "createdAt" => "2016-08-29T16:25:52.130Z", "events" => Any[Dict{String, Any}("message" => "(service ecs-simple-service) was unable to place a task because no container instance met all of its requirements. The closest matching (container-instance 3f4de1c5-ffdd-4954-af7e-75b4be0c8841) is already using a port required by your task. For more information, see the Troubleshooting section of the Amazon ECS Developer Guide.", "id" => "38c285e5-d335-4b68-8b15-e46dedc8e88d", "createdAt" => "2016-08-29T16:25:58.520Z")], "status" => "ACTIVE", "clusterArn" => "arn:aws:ecs:us-east-1:012345678910:cluster/default", "serviceArn" => "arn:aws:ecs:us-east-1:012345678910:service/ecs-simple-service", "deploymentConfiguration" => Dict{String, Any}("maximumPercent" => 200, "minimumHealthyPercent" => 100), "runningCount" => 0, "serviceName" => "ecs-simple-service", "loadBalancers" => Any[], "desiredCount" => 1)], "failures" => Any[])
end

describe_task_definition_1_patch = @patch function describe_task_definition(a...; b...)
    return Dict{String, Any}("taskDefinition" => Dict{String, Any}("containerDefinitions" => Any[Dict{String, Any}("essential" => true, "name" => "wordpress", "links" => Any["mysql"], "cpu" => 10, "image" => "wordpress", "mountPoints" => Any[], "portMappings" => Any[Dict{String, Any}("hostPort" => 80, "containerPort" => 80)], "memory" => 500, "environment" => Any[], "volumesFrom" => Any[]), Dict{String, Any}("essential" => true, "name" => "mysql", "cpu" => 10, "image" => "mysql", "mountPoints" => Any[], "portMappings" => Any[], "memory" => 500, "environment" => Any[Dict{String, Any}("name" => "MYSQL_ROOT_PASSWORD", "value" => "password")], "volumesFrom" => Any[])], "family" => "hello_world", "revision" => 8, "taskDefinitionArn" => "arn:aws:ecs:us-east-1:<aws_account_id>:task-definition/hello_world:8", "volumes" => Any[]))
end

describe_tasks_1_patch = @patch function describe_tasks(a...; b...)
    return Dict{String, Any}("tasks" => Any[Dict{String, Any}("overrides" => Dict{String, Any}("containerOverrides" => Any[Dict{String, Any}("name" => "ecs-demo")]), "desiredStatus" => "RUNNING", "taskDefinitionArn" => "arn:aws:ecs:<region>:<aws_account_id>:task-definition/amazon-ecs-sample:1", "lastStatus" => "RUNNING", "taskArn" => "arn:aws:ecs:<region>:<aws_account_id>:task/c5cba4eb-5dad-405e-96db-71ef8eefe6a8", "clusterArn" => "arn:aws:ecs:<region>:<aws_account_id>:cluster/default", "containerInstanceArn" => "arn:aws:ecs:<region>:<aws_account_id>:container-instance/18f9eda5-27d7-4c19-b133-45adc516e8fb", "containers" => Any[Dict{String, Any}("name" => "ecs-demo", "lastStatus" => "RUNNING", "taskArn" => "arn:aws:ecs:<region>:<aws_account_id>:task/c5cba4eb-5dad-405e-96db-71ef8eefe6a8", "containerArn" => "arn:aws:ecs:<region>:<aws_account_id>:container/7c01765b-c588-45b3-8290-4ba38bd6c5a6", "networkBindings" => Any[Dict{String, Any}("bindIP" => "0.0.0.0", "hostPort" => 80, "containerPort" => 80)])], "startedBy" => "ecs-svc/9223370608528463088")], "failures" => Any[])
end

list_account_settings_1_patch = @patch function list_account_settings(a...; b...)
    return Dict{String, Any}("settings" => Any[Dict{String, Any}("name" => "containerInstanceLongArnFormat", "value" => "disabled", "principalArn" => "arn:aws:iam::<aws_account_id>:user/principalName"), Dict{String, Any}("name" => "serviceLongArnFormat", "value" => "enabled", "principalArn" => "arn:aws:iam::<aws_account_id>:user/principalName"), Dict{String, Any}("name" => "taskLongArnFormat", "value" => "disabled", "principalArn" => "arn:aws:iam::<aws_account_id>:user/principalName")])
end

list_account_settings_2_patch = @patch function list_account_settings(a...; b...)
    return Dict{String, Any}("settings" => Any[Dict{String, Any}("name" => "containerInstanceLongArnFormat", "value" => "disabled", "principalArn" => "arn:aws:iam::<aws_account_id>:user/principalName"), Dict{String, Any}("name" => "serviceLongArnFormat", "value" => "enabled", "principalArn" => "arn:aws:iam::<aws_account_id>:user/principalName"), Dict{String, Any}("name" => "taskLongArnFormat", "value" => "disabled", "principalArn" => "arn:aws:iam::<aws_account_id>:user/principalName")])
end

list_clusters_1_patch = @patch function list_clusters(a...; b...)
    return Dict{String, Any}("clusterArns" => Any["arn:aws:ecs:us-east-1:<aws_account_id>:cluster/test", "arn:aws:ecs:us-east-1:<aws_account_id>:cluster/default"])
end

list_container_instances_1_patch = @patch function list_container_instances(a...; b...)
    return Dict{String, Any}("containerInstanceArns" => Any["arn:aws:ecs:us-east-1:<aws_account_id>:container-instance/f6bbb147-5370-4ace-8c73-c7181ded911f", "arn:aws:ecs:us-east-1:<aws_account_id>:container-instance/ffe3d344-77e2-476c-a4d0-bf560ad50acb"])
end

list_services_1_patch = @patch function list_services(a...; b...)
    return Dict{String, Any}("serviceArns" => Any["arn:aws:ecs:us-east-1:012345678910:service/my-http-service"])
end

list_tags_for_resource_1_patch = @patch function list_tags_for_resource(a...; b...)
    return Dict{String, Any}("tags" => Any[Dict{String, Any}("key" => "team", "value" => "dev")])
end

list_task_definition_families_1_patch = @patch function list_task_definition_families(a...; b...)
    return Dict{String, Any}("families" => Any["node-js-app", "web-timer", "hpcc", "hpcc-c4-8xlarge"])
end

list_task_definition_families_2_patch = @patch function list_task_definition_families(a...; b...)
    return Dict{String, Any}("families" => Any["hpcc", "hpcc-c4-8xlarge"])
end

list_task_definitions_1_patch = @patch function list_task_definitions(a...; b...)
    return Dict{String, Any}("taskDefinitionArns" => Any["arn:aws:ecs:us-east-1:<aws_account_id>:task-definition/sleep300:2", "arn:aws:ecs:us-east-1:<aws_account_id>:task-definition/sleep360:1", "arn:aws:ecs:us-east-1:<aws_account_id>:task-definition/wordpress:3", "arn:aws:ecs:us-east-1:<aws_account_id>:task-definition/wordpress:4", "arn:aws:ecs:us-east-1:<aws_account_id>:task-definition/wordpress:5", "arn:aws:ecs:us-east-1:<aws_account_id>:task-definition/wordpress:6"])
end

list_task_definitions_2_patch = @patch function list_task_definitions(a...; b...)
    return Dict{String, Any}("taskDefinitionArns" => Any["arn:aws:ecs:us-east-1:<aws_account_id>:task-definition/wordpress:3", "arn:aws:ecs:us-east-1:<aws_account_id>:task-definition/wordpress:4", "arn:aws:ecs:us-east-1:<aws_account_id>:task-definition/wordpress:5", "arn:aws:ecs:us-east-1:<aws_account_id>:task-definition/wordpress:6"])
end

list_tasks_1_patch = @patch function list_tasks(a...; b...)
    return Dict{String, Any}("taskArns" => Any["arn:aws:ecs:us-east-1:012345678910:task/0cc43cdb-3bee-4407-9c26-c0e6ea5bee84", "arn:aws:ecs:us-east-1:012345678910:task/6b809ef6-c67e-4467-921f-ee261c15a0a1"])
end

list_tasks_2_patch = @patch function list_tasks(a...; b...)
    return Dict{String, Any}("taskArns" => Any["arn:aws:ecs:us-east-1:012345678910:task/0cc43cdb-3bee-4407-9c26-c0e6ea5bee84"])
end

put_account_setting_1_patch = @patch function put_account_setting(a...; b...)
    return Dict{String, Any}("setting" => Dict{String, Any}("name" => "serviceLongArnFormat", "value" => "enabled", "principalArn" => "arn:aws:iam::<aws_account_id>:user/principalName"))
end

put_account_setting_2_patch = @patch function put_account_setting(a...; b...)
    return Dict{String, Any}("setting" => Dict{String, Any}("name" => "containerInstanceLongArnFormat", "value" => "enabled", "principalArn" => "arn:aws:iam::<aws_account_id>:user/principalName"))
end

put_account_setting_default_1_patch = @patch function put_account_setting_default(a...; b...)
    return Dict{String, Any}("setting" => Dict{String, Any}("name" => "serviceLongArnFormat", "value" => "enabled", "principalArn" => "arn:aws:iam::<aws_account_id>:root"))
end

register_task_definition_1_patch = @patch function register_task_definition(a...; b...)
    return Dict{String, Any}("taskDefinition" => Dict{String, Any}("containerDefinitions" => Any[Dict{String, Any}("essential" => true, "name" => "sleep", "cpu" => 10, "image" => "busybox", "mountPoints" => Any[], "command" => Any["sleep", "360"], "portMappings" => Any[], "memory" => 10, "environment" => Any[], "volumesFrom" => Any[])], "family" => "sleep360", "revision" => 1, "taskDefinitionArn" => "arn:aws:ecs:us-east-1:<aws_account_id>:task-definition/sleep360:19", "volumes" => Any[]))
end

run_task_1_patch = @patch function run_task(a...; b...)
    return Dict{String, Any}("tasks" => Any[Dict{String, Any}("overrides" => Dict{String, Any}("containerOverrides" => Any[Dict{String, Any}("name" => "sleep")]), "desiredStatus" => "RUNNING", "taskDefinitionArn" => "arn:aws:ecs:us-east-1:<aws_account_id>:task-definition/sleep360:1", "lastStatus" => "PENDING", "taskArn" => "arn:aws:ecs:us-east-1:<aws_account_id>:task/a9f21ea7-c9f5-44b1-b8e6-b31f50ed33c0", "containerInstanceArn" => "arn:aws:ecs:us-east-1:<aws_account_id>:container-instance/ffe3d344-77e2-476c-a4d0-bf560ad50acb", "containers" => Any[Dict{String, Any}("name" => "sleep", "lastStatus" => "PENDING", "taskArn" => "arn:aws:ecs:us-east-1:<aws_account_id>:task/a9f21ea7-c9f5-44b1-b8e6-b31f50ed33c0", "containerArn" => "arn:aws:ecs:us-east-1:<aws_account_id>:container/58591c8e-be29-4ddf-95aa-ee459d4c59fd")])])
end

tag_resource_1_patch = @patch function tag_resource(a...; b...)
    return Dict{String, Any}()
end

untag_resource_1_patch = @patch function untag_resource(a...; b...)
    return Dict{String, Any}()
end

update_service_1_patch = @patch function update_service(a...; b...)
    return Dict{String, Any}()
end

update_service_2_patch = @patch function update_service(a...; b...)
    return Dict{String, Any}()
end

ext: jl


text:
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/FEMBase.jl/blob/master/LICENSE

abstract type AbstractProblem end
abstract type FieldProblem<:AbstractProblem end
abstract type BoundaryProblem<:AbstractProblem end
abstract type MixedProblem<:AbstractProblem end

"""
General linearized problem to solve
    (K₁+K₂)Δu  +   C1'*Δλ = f₁+f₂
         C2Δu  +     D*Δλ = g
"""
mutable struct Assembly

    M :: SparseMatrixCOO  # mass matrix

    # for field assembly
    K :: SparseMatrixCOO   # stiffness matrix
    Kg :: SparseMatrixCOO  # geometric stiffness matrix
    f :: SparseMatrixCOO   # force vector
    fg :: SparseMatrixCOO  #

    # for boundary assembly
    C1 :: SparseMatrixCOO
    C2 :: SparseMatrixCOO
    D :: SparseMatrixCOO
    g :: SparseMatrixCOO
    c :: SparseMatrixCOO

    u :: Vector{Float64}  # solution vector u
    u_prev :: Vector{Float64}  # previous solution vector u
    u_norm_change :: Real  # change of norm in u

    la :: Vector{Float64}  # solution vector la
    la_prev :: Vector{Float64}  # previous solution vector u
    la_norm_change :: Real # change of norm in la

    removed_dofs :: Vector{Int64} # manually remove dofs from assembly
end

function Assembly()
    return Assembly(
        SparseMatrixCOO(),
        SparseMatrixCOO(),
        SparseMatrixCOO(),
        SparseMatrixCOO(),
        SparseMatrixCOO(),
        SparseMatrixCOO(),
        SparseMatrixCOO(),
        SparseMatrixCOO(),
        SparseMatrixCOO(),
        SparseMatrixCOO(),
        [], [], Inf,
        [], [], Inf,
        [])
end

function empty!(assembly::Assembly)
    empty!(assembly.K)
    empty!(assembly.Kg)
    empty!(assembly.f)
    empty!(assembly.fg)
    empty!(assembly.C1)
    empty!(assembly.C2)
    empty!(assembly.D)
    empty!(assembly.g)
    empty!(assembly.c)
end

function isempty(assembly::Assembly)
    T = isempty(assembly.K)
    T &= isempty(assembly.Kg)
    T &= isempty(assembly.f)
    T &= isempty(assembly.fg)
    T &= isempty(assembly.C1)
    T &= isempty(assembly.C2)
    T &= isempty(assembly.D)
    T &= isempty(assembly.g)
    T &= isempty(assembly.c)
    return T
end

"""
    Problem{P<:AbstractProblem}

Defines a new problem of type `P`, where `P` characterizes the physics of the
problem. `P` can be for example `Elasticity`, if the physics of the system is
described by Cauchy's stress equation ∇⋅σ + b = ̈ρu, or `Heat`, if the physics
of the problem is described by heat equation -∇⋅(k∇u) = f.

"""
mutable struct Problem{P<:AbstractProblem}
    name :: AbstractString           # descriptive name for the problem
    dimension :: Int                 # degrees of freedom per node
    parent_field_name :: AbstractString # (optional) name of the parent field e.g. "displacement"
    elements :: Vector{Element}
    dofmap :: Dict{Element, Vector{Int64}} # connects the element local dofs to the global dofs
    assembly :: Assembly
    fields :: Dict{String, AbstractField}
    postprocess_fields :: Vector{String}
    properties :: P
end

"""
    Problem(problem_type, problem_name, problem_dimension)

Construct a new field problem.

`problem_type` must be a subtype of `FieldProblem` (`Elasticity`, `Heat`, etc..).
`problem_dimensions` is the number of degrees of freedom each node is containing.

# Examples

To create vector-valued elasticity problem, having 3 dofs / node:
```julia
problem1 = Problem(Elasticity, "test problem", 3)
```

To create scalar-valued Poisson problem:
```julia
problem2 = Problem(Heat, "test problem 2", 1)
```

"""
function Problem(::Type{P}, name::AbstractString, dimension::Int64) where P<:FieldProblem
    parent_field_name = "none"
    elements = []
    dofmap = Dict()
    assembly = Assembly()
    fields = Dict()
    postprocess_fields = Vector()
    properties = P()
    problem = Problem{P}(name, dimension, parent_field_name, elements, dofmap,
                         assembly, fields, postprocess_fields, properties)
    @info("Creating a new problem of type $P, having name `$name` and " *
          "dimension $dimension dofs/node.")
    return problem
end

"""
    Problem(problem_type, problem_name, problem_dimension, parent_field_name)

Construct a new boundary problem.

`problem_type` must be a subtype of `BoundaryProblem` (`Dirichlet`, `Contact`,
etc..). `problem_dimensions` is the number of degrees of freedom each node is
containing. `parent_field_name` is describing the field, where the boundary
problem is affecting.

# Examples

To create a Dirichlet boundary condition for a vector-valued elasticity problem,
having 3 dofs / node:
```julia
bc1 = Problem(Dirichlet, "fix displacement on support", 3, "displacement")
```

To create a Dirichlet boundary condition for scalar-valued Poisson problem:
```julia
bc2 = Problem(Dirichlet, "fix surface temperature", 1, "temperature")
```
"""
function Problem(::Type{P}, name, dimension, parent_field_name) where P<:BoundaryProblem
    elements = []
    dofmap = Dict()
    assembly = Assembly()
    fields = Dict()
    postprocess_fields = Vector()
    properties = P()
    problem = Problem{P}(name, dimension, parent_field_name, elements, dofmap,
                         assembly, fields, postprocess_fields, properties)
    @info("Creating a new boundary problem of type $P, having name `$name` and " *
          "dimension $dimension dofs/node. This boundary problems fixes field " *
          "`$parent_field_name`.")
    return problem
end

function get_formulation_type(::Problem)
    return :incremental
end

"""
    get_unknown_field_dimension(problem)

Return the dimension of the unknown field of this problem.
"""
function get_unknown_field_dimension(problem::Problem)
    return problem.dimension
end

"""
    get_unknown_field_name(problem)

Default function if unknown field name is not defined for some problem.
"""
function get_unknown_field_name(::P) where P<:AbstractProblem
    @warn("The name of unknown field (e.g. displacement, temperature, ...) of the " *
          "problem type must be given by defining a function " *
          "`get_unknown_field_name(::$P)`")
    return "N/A"
end

""" Return the name of the unknown field of this problem. """
function get_unknown_field_name(problem::Problem{P}) where P
    return get_unknown_field_name(problem.properties)
end

""" Return the name of the parent field of this (boundary) problem. """
function get_parent_field_name(problem::Problem{P}) where P<:BoundaryProblem
    return problem.parent_field_name
end

function get_unknown_field_name(::P) where P<:BoundaryProblem
    return "lambda"
end

is_field_problem(::Problem) = false
is_field_problem(::Problem{P}) where {P<:FieldProblem} = true
is_boundary_problem(::Problem) = false
is_boundary_problem(::Problem{P}) where {P<:BoundaryProblem} = true

function get_elements(problem::Problem)
    return problem.elements
end

function update!(problem::P, attr::Pair{String, String}...) where P<:AbstractProblem
    for (name, value) in attr
        setfield!(problem, Meta.parse(name), Meta.parse(value))
    end
end

"""
    function initialize!(problem_type, element_name, time)

Initialize the element ready for calculation, where `problem_type` is the type
of the problem (Elasticity, Dirichlet, etc.), `element_name` is the name of a
constructed element (see Element(element_type, connectivity_vector)) and `time`
is the starting time of the initializing process.
"""
function initialize!(problem::Problem, element::Element, time::Float64)
    field_name = get_unknown_field_name(problem)
    field_dim = get_unknown_field_dimension(problem)
    nnodes = length(element)
    if field_dim == 1  # scalar field
        empty_field = tuple(zeros(nnodes)...)
    else  # vector field
        # FIXME: the most effective way to do
        # ([0.0,0.0], [0.0,0.0], ..., [0.0,0.0]) ?
        empty_field = tuple(map((x)->zeros(field_dim)*x, 1:nnodes)...)
    end

    # initialize primary field
    if !haskey(element, field_name)
        update!(element, field_name, time => empty_field)
    end

    # if a boundary problem, initialize also a field for the main problem
    is_boundary_problem(problem) || return
    field_name = get_parent_field_name(problem)
    if !haskey(element, field_name)
        update!(element, field_name, time => empty_field)
    end
end

function initialize!(problem::Problem, time::Float64=0.0)
    for element in get_elements(problem)
        initialize!(problem, element, time)
    end
end

function update!(problem::Problem, assembly::Assembly, u::Vector, la::Vector)

    # resize & fill with zeros vectors if length mismatch with current solution

    if length(u) != length(assembly.u)
        resize!(assembly.u, length(u))
        fill!(assembly.u, 0.0)
    end

    if length(la) != length(assembly.la)
        resize!(assembly.la, length(la))
        fill!(assembly.la, 0.0)
    end

    # copy current solutions to previous ones and add/replace new solution
    # TODO: here we have couple of options and they need to be clarified
    # for total formulation we are solving total quantity Ku = f while in
    # incremental formulation we solve KΔu = f and u = u + Δu
    assembly.u_prev = copy(assembly.u)
    assembly.la_prev = copy(assembly.la)

    if get_formulation_type(problem) == :total
        assembly.u = u
        assembly.la = la
    elseif get_formulation_type(problem) == :incremental
        assembly.u += u
        assembly.la = la
    elseif get_formulation_type(problem) == :forwarddiff
        assembly.u += u
        assembly.la += la
    else
        @info("$(problem.name): unknown formulation type, don't know what to do with results")
        error("serious failure with problem formulation: $(get_formulation_type(problem))")
    end

    # calculate change of norm
    assembly.u_norm_change = norm(assembly.u - assembly.u_prev)
    assembly.la_norm_change = norm(assembly.la - assembly.la_prev)
    return assembly.u, assembly.la
end

"""
    get_global_solution(problem, assembly)

Return a global solution (u, la) for a problem.

Notes
-----
If the length of solution vector != number of nodes, i.e. the field dimension is
something else than 1, reshape vectors so that their length matches to the
number of nodes. This helps to get nodal results easily.
"""
function get_global_solution(problem::Problem, assembly::Assembly)
    u = assembly.u
    la = assembly.la
    field_dim = get_unknown_field_dimension(problem)
    if field_dim == 1
        return u, la
    else
        nnodes = round(Int, length(u)/field_dim)
        u = reshape(u, field_dim, nnodes)
        u = Vector{Float64}[u[:,i] for i in 1:nnodes]
        la = reshape(la, field_dim, nnodes)
        la = Vector{Float64}[la[:,i] for i in 1:nnodes]
        return u, la
    end
end

function update!(problem::Problem{P}, assembly::Assembly, elements::Vector{Element}, time::Float64) where P<:FieldProblem
    u, la = get_global_solution(problem, assembly)
    field_name = get_unknown_field_name(problem)
    # update solution u for elements
    for element in elements
        connectivity = get_connectivity(element)
        update!(element, field_name, time => tuple(u[connectivity]...))
    end
end

function update!(problem::Problem{P}, assembly::Assembly, elements::Vector{Element}, time::Float64) where P<:BoundaryProblem
    u, la = get_global_solution(problem, assembly)
    parent_field_name = get_parent_field_name(problem) # displacement
    field_name = get_unknown_field_name(problem) # lambda
    # update solution and lagrange multipliers for boundary elements
    for element in elements
        connectivity = get_connectivity(element)
        update!(element, parent_field_name, time => tuple(u[connectivity]...))
        update!(element, field_name, time => tuple(la[connectivity]...))
    end
end

"""
    add_element!(problem, element1, element2, ...)

Add element(s) to the problem.
"""
function add_element!(problem, elements...)
    for element in elements
        push!(problem.elements, element)
    end
    return nothing
end

"""
    add_elements!(problem, element_set_1, element_set_2, ...)

Add vectors/tuples of element(s) to the problem.
"""
function add_elements!(problem, element_sets::Union{Vector,Tuple}...)
    for elements in element_sets
        nelements = length(elements)
        @info("Adding $nelements elements to problem `$(problem.name)`")
        add_element!(problem, elements...)
    end
    return nothing
end

add_elements!(problem, elements::Element...) = add_element!(problem, elements...)

function add_elements!(problem, elements_or_lists_of_elements...)
    for item in elements_or_lists_of_elements
        add_elements!(problem, item)
    end
end

get_assembly(problem::Problem) = problem.assembly
Base.length(problem::Problem) = length(problem.elements)

function update!(problem::Problem, field_name::AbstractString, data)
    #if haskey(problem.fields, field_name)
    #    update!(problem.fields[field_name], field_name::AbstractString, data)
    #else
    #    problem.fields[field_name] = Field(data)
    #end
    update!(problem.elements, field_name::AbstractString, data)
end

function haskey(problem::Problem, field_name::AbstractString)
    return haskey(problem.fields, field_name)
end

function getindex(problem::Problem, field_name::String)
    return problem.fields[field_name]
end

#""" Return field calculated to nodal points for elements in problem p. """
function (problem::Problem)(field_name::String, time::Float64)
    #if haskey(problem, field_name)
    #    return problem[field_name](time)
    #end
    f = Dict{Int64, Any}()
    for element in get_elements(problem)
        haskey(element, field_name) || continue
        for (c, v) in zip(get_connectivity(element), element(field_name, time))
            if haskey(f, c)
                if !isapprox(f[c], v)
                    @info("several values for single node when returning field $field_name")
                    @info("already have: $(f[c]), and trying to set $v")
                end
            else
                f[c] = v
            end
        end
    end
    #f == nothing && return f
    #update!(problem, field_name, time => f)
    return f
end

function push!(problem::Problem, elements...)
    push!(problem.elements, elements...)
end

function push!(problem::Problem, elements_::Vector...)
    for elements in elements_
        push!(problem.elements, elements...)
    end
end

"""
    set_gdofs!(problem, element)

Set element global degrees of freedom.
"""
function set_gdofs!(problem, element, dofs)
    problem.dofmap[element] = dofs
end

"""
    get_gdofs(problem, element)

Return the global degrees of freedom for element.

First make lookup from problem dofmap. If not defined there, make implicit
assumption that dofs follow formula `gdofs = [dim*(nid-1)+j for j=1:dim]`,
where `nid` is node id and `dim` is the dimension of problem. This formula
arranges dofs so that first comes all dofs of node 1, then node 2 and so on:
(u11, u12, u13, u21, u22, u23, ..., un1, un2, un3) for 3 dofs/node setting.
"""
function get_gdofs(problem::Problem, element::Element)
    if haskey(problem.dofmap, element)
        return problem.dofmap[element]
    end
    conn = get_connectivity(element)
    if length(conn) == 0
        error("element connectivity not defined, cannot determine global ",
              "degrees of freedom for element #: $(element.id)")
    end
    dim = get_unknown_field_dimension(problem)
    gdofs = [dim*(i-1)+j for i in conn for j=1:dim]
    return gdofs
end


text:
export MACLearner

using Flux

"""
    MACLearner(;kwargs...)
   Keyword arguments
- `approximator`::[`ActorCritic`](@ref)
- `γ::Float32`, reward discount rate
-  `bootstrap::bool`, if false then Q function is approximated using monte carlo returns.
"""

Base.@kwdef mutable struct MACLearner{A<:ActorCritic} <: AbstractLearner
    approximator::A
    γ::Float32
    max_grad_norm::Union{Nothing,Float32} = nothing
    norm::Float32 = 0.0f0
    actor_loss::Float32 = 0.0f0
    critic_loss::Float32 = 0.0f0
    loss::Float32 = 0.0f0
    bootstrap::Bool = true
end

function (learner::MACLearner)(env::MultiThreadEnv)
    learner.approximator.actor(send_to_device(
        device(learner.approximator),
        get_state(env),
    )) |> send_to_host
end

function (learner::MACLearner)(env)
    s = get_state(env)
    s = Flux.unsqueeze(s, ndims(s) + 1)
    s = send_to_device(device(learner.approximator), s)
    learner.approximator.actor(s) |> vec |> send_to_host
end

function RLBase.update!(learner::MACLearner, t::AbstractTrajectory)
    isfull(t) || return

    states = t[:state]
    actions = t[:action]
    rewards = t[:reward]
    terminals = t[:terminal]

    AC = learner.approximator
    γ = learner.γ
    D = device(AC)

    states = send_to_device(D, states)
    states_flattened = flatten_batch(states) # (state_size..., n_thread * update_step)

    actions = flatten_batch(actions)
    actions = CartesianIndex.(actions, 1:length(actions))

    if learner.bootstrap
        next_state = select_last_frame(t[:next_state])
        next_state = send_to_device(D, next_state)
        next_state_values = AC.critic(next_state)

        gains = discount_rewards(
            rewards,
            γ;
            dims = 2,
            init = send_to_host(next_state_values),
            terminal = terminals,
        )
        gains = send_to_device(D, gains)
    else
        next_state_flattened = flatten_batch(t[:next_state])
        next_state_flattened = send_to_device(D, next_state_flattened)
        rewards_flattened = flatten_batch(rewards)
        rewards_flattened = send_to_device(D, rewards_flattened)
    end

    action_values = AC.critic(states_flattened)

    ps1 = Flux.params(AC.actor)
    gs1 = gradient(ps1) do
        logits = AC.actor(states_flattened)
        probs = softmax(logits)
        actor_loss = -mean(sum((probs .* Zygote.dropgrad(action_values)), dims = 1))
        loss = actor_loss
        ignore() do
            learner.actor_loss = actor_loss
        end
        loss
    end
    if !isnothing(learner.max_grad_norm)
        learner.norm = clip_by_global_norm!(gs1, ps1, learner.max_grad_norm)
    end
    update!(AC.actor, gs1)

    ps2 = Flux.params(AC.critic)
    gs2 = gradient(ps2) do
        if learner.bootstrap
            critic_loss = mean((vec(gains) .- vec(action_values[actions])) .^ 2)
        else
            next_state_values = AC.critic(next_state_flattened)
            target_action_values =
                vec(rewards_flattened) .+
                γ * vec(Zygote.dropgrad(sum(
                    next_state_values .* softmax(AC.actor(next_state_flattened)),
                    dims = 1,
                )))
            critic_loss =
                mean((vec(target_action_values) .- vec(action_values[actions])) .^ 2)
        end

        loss = critic_loss
        ignore() do
            learner.critic_loss = critic_loss
        end
        loss
    end
    if !isnothing(learner.max_grad_norm)
        learner.norm = clip_by_global_norm!(gs2, ps2, learner.max_grad_norm)
    end
    update!(AC.critic, gs2)
end

function (agent::Agent{<:QBasedPolicy{<:MACLearner},<:CircularCompactSARTSATrajectory})(
    ::Training{PreActStage},
    env,
)
    action = agent.policy(env)
    state = get_state(env)
    push!(agent.trajectory; state = state, action = action)
    update!(agent.policy, agent.trajectory)

    # the main difference is we'd like to flush the buffer after each update!
    if isfull(agent.trajectory)
        empty!(agent.trajectory)
        push!(agent.trajectory; state = state, action = action)
    end

    action
end


text:
"""
    gmm = combine(gmmx::IsotropicGMM, gmmy::IsotropicGMM)
    mgmm = combine(mgmmx::MultiGMM, mgmmy::MultiGMM)
    gmm = combine(gmms::Union{AbstractVector{<:IsotropicGMM},AbstractVector{<:MultiGMM}})

Creates a new `IsotropicGMM` or `MultiGMM` by concatenating the vectors of `IsotroicGaussian`s in
the input GMMs.
"""
function combine(gmmx::AbstractSingleGMM, gmmy::AbstractSingleGMM)
    if dims(gmmx) != dims(gmmy)
        throw(ArgumentError("GMMs must have the same dimensionality"))
    end
    t = promote_type(typeof(gmmx), typeof(gmmy))
    return t(vcat(gmmx.gaussians, gmmy.gaussians))
end

function combine(mgmmx::AbstractMultiGMM, mgmmy::AbstractMultiGMM)
    if dims(gmmx) != dims(mgmmy)
        throw(ArgumentError("GMMs must have the same dimensionality"))
    end
    t = IsotropicGMM{dims(gmmx),promote_type(numbertype(mgmmx), numbertype(mgmmy))}
    gmms = Dict{Symbol, t}()
    xkeys, ykeys = keys(mgmmx.gmms), keys(mgmmy.gmms)
    for key in xkeys ∪ ykeys
        if key ∈ xkeys && key ∈ykeys
            push!(gmms, Pair(key, t(combine(mgmmx.gmms[key], mgmmy.gmms[key]))))
        elseif key ∈ xkeys
            push!(gmms, Pair(key, t(mgmmx.gmms[key])))
        else
            push!(gmms, Pair(key, t(mgmmy.gmms[key])))
        end
    end
    return promote_type(typeof(mgmmx), typeof(mgmmy))(gmms)
end

function combine(gmms::Union{AbstractVector{<:AbstractSingleGMM},AbstractVector{<:AbstractMultiGMM}})
    if length(gmms) > 1
        return combine([combine(gmms[1],gmms[2]), gmms[3:end]...])
    elseif length(gmms) == 1
        return gmms[1]
    else
        throw(ArgumentError("provided no GMMs to combine"))
    end
end


text:
include("INCLUDEME.jl")

using Yao, Yao.Blocks, JLD2
# using Yao, Circuit, UnicodePlots, GradOptim, Utils, ArgParse, JLD2, FileIO
import Kernels

# n = 6
# qcbm = QCBM{n, 10}(get_nn_pairs(n))

@load "data.jld" output

layer(::Val{:first}) = rollrepeat(chain(Rx(), Rz()))
layer(::Val{:last}) = rollrepeat(chain(Rz(), Rx()))
layer(::Val{:mid}) = rollrepeat(chain(Rz(), Rx(), Rz()))

c = kron(6, i=>chain(Rx(), Rz()) for i = 1:6)

# c = layer(Val(:first))(6)
dispatch!(c, [i for i=1:12])
c

out = statevec(apply!(register(bit"000000"), c))

# dispatch!(qcbm, params)
# out = statevec(qcbm())
@show out ≈ output
out

max_stars_repo_name: Roger-luo/QCBM

max_issues_repo_name: Roger-luo/QCBM

max_forks_repo_name: Roger-luo/QCBM


text:
using Test
using WorldBankData

country_data = Any[ Dict("total" => 4, "per_page" => "25000", "pages" => 1, "page" => 1),
                    Any[
                     Dict("latitude" => "42.5075", "id" => "AND", "iso2Code" => "AD", "incomeLevel" => Dict("id" => "NOC", "value" => "High income: nonOECD"), "adminregion" => Dict("id" => "", "value" => ""), "lendingType" => Dict("id" => "LNX", "value" => "Not classified"), "region" => Dict("id" => "ECS", "value" => "Europe & Central Asia (all income levels)"), "capitalCity" => "Andorra la Vella", "name" => "Andorra", "longitude" => "1.5218"),
                     Dict("latitude" => "", "id" => "ARB", "iso2Code" => "1A", "incomeLevel" => Dict("id" => "NA", "value" => "Aggregates"), "adminregion" => Dict("id" => "", "value" => ""), "lendingType" => Dict("id" => "", "value" => "Aggregates"), "region" => Dict("id" => "NA", "value" => "Aggregates"), "capitalCity" => "", "name" => "Arab World", "longitude" => ""),
                     Dict("latitude" => "24.4764", "id" => "ARE", "iso2Code" => "AE", "incomeLevel" => Dict("id" => "NOC", "value" => "High income: nonOECD"), "adminregion" => Dict("id" => "", "value" => ""), "lendingType" => Dict("id" => "LNX", "value" => "Not classified"), "region" => Dict("id" => "MEA", "value" => "Middle East & North Africa (all income levels)"), "capitalCity" => "Abu Dhabi", "name" => "United Arab Emirates", "longitude" => "54.3705"),
                     Dict("latitude" => "-34.6118", "id" => "ARG", "iso2Code" => "AR", "incomeLevel" => Dict("id" => "UMC", "value" => "Upper middle income"), "adminregion" => Dict("id" => "LAC", "value" => "Latin America & Caribbean (developing only)"), "lendingType" => Dict("id" => "IBD", "value" => "IBRD"), "region" => Dict("id" => "LCN", "value" => "Latin America & Caribbean (all income levels)"), "capitalCity" => "Buenos Aires", "name" => "Argentina", "longitude" => "-58.4173")
                    ]
                  ]

df_country = WorldBankData.parse_country(country_data)

indicator_data = Any[ Dict("total" => 5, "per_page" => "25000", "pages" => 1, "page" => 1),
                      Any[ Dict("sourceOrganization" => "World Bank and International Energy Agency (IEA Statistics © OECD/IEA, http://www.iea.org/stats/index.asp).  ", "id" => "12.1_TD.LOSSES", "topics" => [], "sourceNote" => "Transmission and distribution losses (%): Transmission and distribution (T&D) losses measure power lost in the transmission of (high-voltage) electricity from power generators to distributors and in the distribution of (medium- and low-voltage) electricity from distributors to end-users. T&D losses are represented as a percentage of gross electricity production. They include both technical and nontechnical (or commercial) losses. Included in the latter are unmetered, unbilled, and unpaid electricity, including theft, which could be significant in developing countries. Aggregate T&D system indicators may be dominated by factors other than losses. The location of primary energy resources (such as hydro lakes and coal seams) and large loads (cities and industries) may be more significant factors in T&D efficiency indicators than the losses or efficiency of the transmission system itself. Properly separating true losses (and hence the efficiency potential of transmission systems) from exogenous location and scale factors and nontechnical losses would require detailed studies of system-dynamic interactions and real operating requirements that are not practical for global tracking purposes.", "name" => "Transmission and distribution losses (%)", "source" => Dict("id" => "35", "value" => "Sustainable Energy for All")),
                       Dict("sourceOrganization" => "World Bank and International Energy Agency (IEA Statistics © OECD/IEA, http://www.iea.org/stats/index.asp).  ", "id" => "13.1_INDUSTRY.ENERGY.INTENSITY", "topics" => [], "sourceNote" => "Energy intensity of industrial sector (MJ/\$2005):  A ratio between energy consumption in industry (including energy industry own use) and industry sector value added measured at purchasing power parity. Energy intensity is an indication of how much energy is used to produce one unit of economic output. Lower ratio indicates that less energy is used to produce one unit of output. ", "name" => "Energy intensity of industrial sector (MJ/\$2005)", "source" => Dict("id" => "35", "value" => "Sustainable Energy for All")),
                       Dict("sourceOrganization" => "World Bank and International Energy Agency (IEA Statistics © OECD/IEA, http://www.iea.org/stats/index.asp).  ", "id" => "14.1_AGR.ENERGY.INTENSITY", "topics" => [], "sourceNote" => "Energy intensity of agricultural sector (MJ/\$2005):  A ratio between energy consumption in agricultural sector (including forestry and fishing) and agricultural sector value added measured at purchasing power parity. Energy intensity is an indication of how much energy is used to produce one unit of economic output. Lower ratio indicates that less energy is used to produce one unit of output. ", "name" => "Energy intensity of agricultural sector (MJ/\$2005)", "source" => Dict("id" => "35", "value" => "Sustainable Energy for All")),
                       Dict("sourceOrganization" => "World Bank and International Energy Agency (IEA Statistics © OECD/IEA, http://www.iea.org/stats/index.asp).  ", "id" => "15.1_OTHER.SECT.ENER.INTENS", "topics" => [], "sourceNote" => "Energy intensity of other sectors (MJ/\$2005): A ratio between energy consumption in other sectors (including services, residential and transport) and services sector value added measured at purchasing power parity. Energy intensity is an indication of how much energy is used to produce one unit of economic output. Lower ratio indicates that less energy is used to produce one unit of output. ", "name" => "Energy intensity of other sectors (MJ/\$2005)", "source" => Dict("id" => "35", "value" => "Sustainable Energy for All")),
                       Dict("sourceOrganization" => "World Bank and International Energy Agency (IEA Statistics © OECD/IEA, http://www.iea.org/stats/index.asp).  ", "id" => "16.1_DECOMP.EFFICIENCY.IND", "topics" => [], "sourceNote" => "The Index is created from the series of year to year Energy Intensity component of the Decomposition Analysis with the DIVISIA LMDI I method. The Significant Period is defined as the period for which the available data is consistent in the energy and value added databases such that all the sectors have energy and value added for those years. The CAGR is calculated over that Significant Period.", "name" => "Divisia Decomposition Analysis - Energy Intensity component Index", "source" => Dict("id" => "35", "value" => "Sustainable Energy for All"))
                      ]
                    ]

df_indicator = WorldBankData.parse_indicator(indicator_data)

WorldBankData.set_country_cache(df_country)

@test search_wdi("countries", "name", r"(Andorr|Arg)"i)[!, :capital] ==  String["Andorra la Vella", "Buenos Aires"]
@test search_wdi("countries", "capital", r"(Andorr|Abu)"i)[!, :iso2c] ==  String["AD", "AE"]

WorldBankData.set_indicator_cache(df_indicator)

@test search_wdi("indicators", "description", r"energy intensity of"i)[!, :name] == String["Energy intensity of industrial sector (MJ/\$2005)", "Energy intensity of agricultural sector (MJ/\$2005)", "Energy intensity of other sectors (MJ/\$2005)"]
@test search_wdi("indicators", "description", r"transmiss"i)[!, :indicator] == String["12.1_TD.LOSSES"]


text:
import InfrastructureSystems:
    increment_count!,
    LogEventSuppressionStats,
    LogSuppressionTracker,
    should_suppress!,
    get_log_events

TEST_MSG = "test log message"

@testset "Test LogEventTracker" begin
    levels = (Logging.Info, Logging.Warn, Logging.Error)
    tracker = IS.LogEventTracker(levels)

    events = (
        IS.LogEvent("file", 14, :id, TEST_MSG, Logging.Debug),
        IS.LogEvent("file", 14, :id, TEST_MSG, Logging.Info),
        IS.LogEvent("file", 14, :id, TEST_MSG, Logging.Warn),
        IS.LogEvent("file", 14, :id, TEST_MSG, Logging.Error),
    )

    for i in range(1, length = 2)
        for event in events
            increment_count!(tracker, event, false)
        end
    end

    @test length(IS.get_log_events(tracker, Logging.Debug)) == 0
    for level in levels
        test_events = collect(IS.get_log_events(tracker, level))
        @test length(test_events) == 1
        @test test_events[1].count == 2
    end

    text = IS.report_log_summary(tracker)
    @test !occursin("Debug", text)
    @test !occursin("suppressed", text)
    for level in ("Error", "Warn", "Info")
        @test occursin("1 $level event", text)
    end
end

@testset "Test MultiLogger with no event tracking" begin
    logger = IS.MultiLogger([
        TerminalLogger(devnull, Logging.Info),
        SimpleLogger(devnull, Logging.Debug),
    ])
    with_logger(logger) do
        @info TEST_MSG
    end

    @test_throws ErrorException IS.report_log_summary(logger)
end

@testset "Test MultiLogger with event tracking" begin
    levels = (Logging.Debug, Logging.Info, Logging.Warn, Logging.Error)
    logger = IS.MultiLogger(
        [TerminalLogger(devnull, Logging.Info), SimpleLogger(devnull, Logging.Debug)],
        IS.LogEventTracker(levels),
    )

    with_logger(logger) do
        for i in range(1, length = 2)
            @debug TEST_MSG
            @info TEST_MSG
            @warn TEST_MSG
            @error TEST_MSG maxlog = 1
        end
    end

    events = collect(IS.get_log_events(logger.tracker, Logging.Error))
    @test length(events) == 1
    events[1].suppressed == 1

    text = IS.report_log_summary(logger)
    for level in levels
        @test occursin("1 $level event", text)
    end
end

@testset "Test configure_logging" begin
    # Verify logging to a file.
    logfile = "testlog.txt"
    logger = IS.configure_logging(;
        file = true,
        filename = logfile,
        file_level = Logging.Info,
        set_global = false,
    )
    with_logger(logger) do
        @info TEST_MSG
    end

    close(logger)

    @test isfile(logfile)
    open(logfile) do io
        lines = readlines(io)
        @test length(lines) == 2  # two lines per message
        @test occursin(TEST_MSG, lines[1])
    end
    rm(logfile)

    # Verify logging with no file.
    logger = IS.configure_logging(;
        console = true,
        file = false,
        console_stream = devnull,
        filename = nothing,
        file_level = Logging.Info,
        set_global = false,
    )
    with_logger(logger) do
        @error TEST_MSG
    end

    events = collect(IS.get_log_events(logger.tracker, Logging.Error))
    @test length(events) == 1
    close(logger)

    @test !isfile(logfile)

    # Verify disabling of tracker.
    logger = IS.configure_logging(;
        console = true,
        file = false,
        console_stream = devnull,
        filename = logfile,
        file_level = Logging.Info,
        set_global = false,
        tracker = nothing,
    )
    with_logger(logger) do
        @error TEST_MSG
        @test isnothing(logger.tracker)
    end

    # Verify setting of global logger
    orig_logger = global_logger()
    logger = IS.configure_logging(;
        console = true,
        file = false,
        console_stream = devnull,
        filename = logfile,
        file_level = Logging.Info,
        set_global = true,
        tracker = nothing,
    )
    @error TEST_MSG
    @test orig_logger != global_logger()
    global_logger(orig_logger)
end

@testset "Test group level" begin
    levels = (Logging.Debug, Logging.Info, Logging.Warn, Logging.Error)
    tracker = IS.LogEventTracker(levels)
    logger = IS.MultiLogger([TerminalLogger(devnull, Logging.Debug)], tracker)
    IS.set_group_level!(logger, :test_group, Logging.Info)

    @test IS.get_group_level(logger, :test_group) == Logging.Info
    @test IS.get_group_level(logger, :not_stored) === nothing

    with_logger(logger) do
        @debug TEST_MSG _group = :test_group
    end
    @test length(IS.get_log_events(tracker, Logging.Debug)) == 0

    IS.empty_group_levels!(logger)
    with_logger(logger) do
        @debug TEST_MSG _group = :test_group
    end
    @test length(IS.get_log_events(tracker, Logging.Debug)) == 1
end

@testset "Test group levels" begin
    levels = (Logging.Debug, Logging.Info, Logging.Warn, Logging.Error)
    tracker = IS.LogEventTracker(levels)
    logger = IS.MultiLogger([TerminalLogger(devnull, Logging.Debug)], tracker)
    min_levels = Dict(:test_group => Logging.Warn)
    IS.set_group_levels!(logger, min_levels)
    @test IS.get_group_levels(logger) == min_levels
    with_logger(logger) do
        @info TEST_MSG _group = :test_group
    end
    @test length(IS.get_log_events(tracker, Logging.Info)) == 0

    IS.empty_group_levels!(logger)
    with_logger(logger) do
        @info TEST_MSG _group = :test_group
    end
    @test length(IS.get_log_events(tracker, Logging.Info)) == 1
end

@testset "Test make_logging_config_file" begin
    filename = joinpath(tempdir(), "logging_config.toml")
    try
        redirect_stdout(devnull) do
            IS.make_logging_config_file(filename)
            return IS.make_logging_config_file(filename, force = true)
        end
        @test IS.LoggingConfiguration(filename) isa IS.LoggingConfiguration
    finally
        isfile(filename) && rm(filename)
    end
end

@testset "Test log suppression for single event" begin
    event_id = :my_event
    cur_time = time()
    period = 10
    maxlog = 3
    stats = LogEventSuppressionStats(event_id, cur_time, period)
    @test !stats.is_tracking_active
    @test stats.period == period

    for i in 1:maxlog
        suppress, num = should_suppress!(stats, cur_time, maxlog)
        @test !suppress
        @test num == 0
        @test stats.is_tracking_active
    end

    suppress, num = should_suppress!(stats, cur_time, maxlog)
    @test suppress
    @test num == 1
    @test stats.is_suppression_enabled
    suppress, num = should_suppress!(stats, cur_time, maxlog)
    @test suppress
    @test num == 2

    # Fake a jump in time.
    stats.tracking_start_time -= period
    suppress, num = should_suppress!(stats, cur_time, maxlog)
    @test !suppress
    @test num == 2
    @test !stats.is_suppression_enabled
    @test stats.count == 1
    @test stats.num_suppressed == 0

    for i in 1:(maxlog - 1)
        suppress, num = should_suppress!(stats, cur_time, maxlog)
        @test !suppress
        @test num == 0
        @test stats.is_tracking_active
    end

    suppress, num = should_suppress!(stats, cur_time, maxlog)
    @test suppress
    @test num == 1
end

@testset "Test log suppression" begin
    levels = (Logging.Info, Logging.Warn, Logging.Error)
    tracker = IS.LogEventTracker(levels)
    logger = IS.MultiLogger(
        [TerminalLogger(devnull, Logging.Info), SimpleLogger(devnull, Logging.Debug)],
        tracker,
    )

    function run()
        @info TEST_MSG maxlog = 5 _suppression_period = 10
    end

    with_logger(logger) do
        for _ in 1:15
            run()
        end
    end

    log_events = values(get_log_events(tracker, Logging.Info))
    @test length(log_events) == 1
    log_event = first(log_events)
    @test log_event.count == 15
    @test log_event.suppressed == 10
end

@testset "Test log suppression logging of num_suppressed" begin
    io = IOBuffer()
    logger = IS.MultiLogger([TerminalLogger(io, Logging.Info)])

    function run()
        @info TEST_MSG maxlog = 1 _suppression_period = 5
    end

    with_logger(logger) do
        # Not suppressed
        run()
        # Suppressed
        run()
        # Fake a jump in time.
        for stats in values(logger.suppression_tracker.event_stats)
            stats.tracking_start_time -= 10
        end
        # Not suppressed
        return run()
    end

    count_occurrences(str, substr) = length(split(str, substr)) - 1

    output = String(take!(io))
    @test count_occurrences(output, TEST_MSG) == 2
    @test count_occurrences(output, "num_suppressed") == 1
    @test count_occurrences(output, "num_suppressed = 1") == 1
end

@testset "Test progress logger" begin
    io = IOBuffer()
    logger = IS.configure_logging(
        console_stream = io,
        console = false,
        file = false,
        progress = true,
        set_global = false,
    )
    with_logger(logger) do
        @progress for i in 1:5
            sleep(0.01)
        end
    end

    output = String(take!(io))
    @test occursin("Progress", output)
end

@testset "Test bad input" begin
    @test_throws ErrorException IS.configure_logging(
        console = false,
        file = false,
        progress = false,
    )
end


text:
############################## Lennard Jones ###################################
struct Morse <: EmpiricalPotential
    D
    α
    σ
    rcutoff
    species :: AbstractVector
end

get_parameters(m::Morse) = Parameter{ (:D,:α, :σ) }((m.D, m.α, m.σ))
set_parameters(p::Parameter{(:D,:α, :σ)}, m::Morse) = Morse(p.D, p.α, p.σ, m.rcutoff, m.species)

deserialize_parameters(p::Parameter{(:D,:α, :σ)}, m::Morse) = [p.ϵ, p.σ]
serialize_parameters(p::Vector, m::Morse) = Parameter{(:ϵ, :σ)}( (p[1], p[2]) )

get_hyperparameters(m::Morse) = Parameter{(:rcutoff,)}( (m.rcutoff,) )
set_hyperparameters(p::Parameter{(:rcutoff,)}, m::Morse) = Morse(m.D, m.α, m.σ, p.rcutoff, m.species)

deserialize_hyperparameters(p::Parameter{(:rcutoff,)}, m::Morse) = [p.rcutoff]
serialize_hyperparameters(p::Vector, m::Morse) = Parameter{(:rcutoff, )}( (p[1],) )

############################# Energies ##########################################

function potential_energy(R::AbstractFloat, p::Morse)
    p.D * (1.0 - exp(-p.α*(R - p.σ)) )^2
end

############################### Forces ##########################################

function force(R::AbstractFloat, r::SVector{3,<:AbstractFloat}, p::Morse)
    SVector( 2 * p.D * p.α * (1.0 - exp(-p.α*(R - p.σ)) ) * exp(-p.α*(R - p.σ)) .* r ./ R)
end


text:
## Step1EndGrid is an endogenous grid search to find the optimal policy function

#module Step1EndGrid

using LinearAlgebra, Random, Distributions, Statistics, Plots
using BenchmarkTools, Interpolations, Roots

"""
    Performs an endogenous search to find the optimal policy function.
    To do so, 3 functions are nested together:
    1. RHSEuler: computes the RHS of the Euler equation
    2. initialpolicy: for a guess value of consumption, finds a'
    3. EndGridSearch: solves for optimal consumption and a'
    Takes price q as given.
"""

function RHSEuler(A, E, Π, q, β, σ, e, avalue, aprime; tol = 0.01)

    # avalue is the actual value of assets
    # aprime is the actual value of the new assets

    """
    This function computes the RHS of the Euler Equation.
    Inputs:
        - A: assets today
        - E: exogenous shocks
        - Π: transtion matrix
        - q: price (guessed one)
        - β: discount factor
        - σ: preference parameter
        - e: position parameter for the shock (position, not value)
        - avalue: actual value of assets today (value, not position)
        - aprime: actual value of assets tomorrow
        - Aprimetol: provisional grid of future assets
    Output:
        - RHS: right hand side of the Euler Equation, for 2 exogenous state var.
    """

    a_size = length(A)
    e_size = size(E)[1]

    # Create a grid of A multiplied by the length of E
    Agrid = repeat(A, 1, e_size)

    Aprimetol = copy(Agrid) .+ 2*tol

    # Define the first derivative of utility function
    # (in the future, make it more general)
    lowerbound = 1E-7
    up(c::Real) = (max(c, lowerbound)).^(-σ)

    Uprime = zeros(e_size, 1)
    RHS = zeros(e_size, 1)

    aux = LinearInterpolation(Agrid[:, e], Aprimetol[:, e])
    adoubleprime = aux(aprime)
    Uprime = up(aprime .+ E[e] .- q .* adoubleprime)
    temp = Uprime .* ones(e_size, 1)
    temp2 = (Π[e, :]' * temp)[1,1]
    RHS = up(avalue + E[e] - aprime * q) - β * temp2

    return RHS

end

function initialpolicy(A, E, Π, β, σ, q; maxT = 600, tol = 0.01)

    """
    This function creates the first grid of assets in the following period.
    Inputs:
        - A: grid of assets today
        - E: exogenous shocks
        - Π: transtion matrix
        - β: discount factor
        - σ: preference parameter
        - q: price (guessed one)
    Output:
        - Aprime: grid of assets tomorrow
    """

    a_size = length(A)
    e_size = size(E)[1]

    # Create a grid of A multiplied by the length of E
    Agrid = repeat(A, 1, e_size)
    Aprime = copy(Agrid)
    Aprimetol = copy(Agrid) .+ 2*tol

    # Iteration
    normdiff = Inf
    t = 1

    while normdiff > tol && t <= maxT

        for a = 1:a_size
            for e = 1:e_size

                FOCfunction(aprime) = RHSEuler(A, E, Π, q, β, σ, e, Agrid[a, e], aprime)

                if FOCfunction(Agrid[1, e]) > 0.0
                    Aprime[a, e] = Agrid[1, e]
                elseif FOCfunction(Agrid[end, e]) < 0.0
                    Aprime[a, e] = Agrid[end, e]
                else
                    Aprime[a, e] = fzero(FOCfunction, Agrid[1, e], Agrid[end, e])
                end

            end
        end

        t += 1
        normdiff = maximum(abs.(Aprimetol - Aprime))
        Aprimetol = copy(Aprime)

    end

    return Aprime

end

function EndGridSearch(A, E, Π, β, σ, q; maxT = 600, tol = 0.10)

    """
    Function that finds the optimal policy function.
    Inputs:
        - A: assets today
        - E: exogenous shocks
        - Π: transition matrix
        - β: discount factor
        - σ: preference parameter
        - q: price (guess)
    Output:
        - pol_func: optimal policy function
    """

    # Check that parameters and price make sense
    @assert β < 1
    @assert σ > 1
    @assert q > β

    a_size = length(A)
    e_size = size(E)[1]

    # Check that transition matrix is well defined
    for e = 1:e_size
        @assert sum(Π[e,:]) == 1
    end

    # Define the first derivative of utility function and its inverse
    # (in the future, make it more general)
    lb = 1E-7
    lowerbound = lb .* ones(a_size, e_size)
    up(c) = max.(c, lowerbound).^(-σ)
    upinv(B) = B.^(-1/σ)

    # Compute Aprime: initial grid of next period's assets
    Aprime = initialpolicy(A, E, Π, β, σ, q)

    # Compute initial consumption (use it as guess for iteration)
    # Remember that it cannot be negative
    C = Aprime .+ E'
    for a = 1:a_size
        for e = 1:e_size
            C[a, e] = (C[a, e] <= 0) ? 0.0 : C[a,e]
        end
    end

    Cupdate = 0.0 .* similar(Aprime)
    Aupdate = similar(Aprime)

    # Iteration to find optimal consumption
    normdiff = Inf
    t = 1

    while normdiff > tol && t <= maxT

        # Find Ctilda from the Euler Equation (weakly positive)
        B = β .* up(C) * Π'
        Ctilda = upinv(B ./ q)
        for a = 1:a_size
            for e = 1:e_size
                Ctilda[a, e] = (Ctilda[a, e] <= 0) ? 0.0 : Ctilda[a,e]
            end
        end

        # Solve for A from the budget constraint
        A = Ctilda + Aprime .* q .- E'

        # Adjust for whether budget constraint binds (only low shock)
        binds = Aprime .< A[1,:]'
        Cupdate[binds] .= (Aprime .+ E' .- Aprime[1,:]' .* q)[binds]
        # Generate auxiliary matrix for when binding
        auxbinding = similar(A)
        for e = 1:e_size
            auxfunc = LinearInterpolation(A[:, e], Ctilda[:, e])
            auxbinding[:, e] = auxfunc.(Aprime[:, e])
        end
        # Fill Cupdate for binding
        Cupdate[.!binds] .= auxbinding[.!binds]

        # Replace negative consumption with 0, if exists
        for a = 1:a_size
            for e = 1:e_size
                Cupdate[a, e] = (Cupdate[a, e] <= 0) ? 0.0 : Cupdate[a,e]
            end
        end

        # values for next iteration
        t += 1
        normdiff = maximum(abs.(Cupdate - C))
        C = copy(Cupdate)

    end

    # Now find the optimal policy function in my original grid
    pol_func = zeros(a_size, e_size)
    for e = 1:e_size
        auxfunc = LinearInterpolation(A[:, e], Aprime[:, e])
        pol_func[:, e] = auxfunc.(Aprime[:, e])
    end

    return pol_func

end

#export RHSEuler, initialpolicy, EndGridSearch

#end


text:
#
# Script to generate the default data-files for the EMT
# Potential; from the ASE implementation
#
# this is essentially derived from the ASE code, and hence ought to be
# licensed by whatever ASE's license is.
#

# import JuLIP
# import JuLIP.Potentials.EMT : EMTParams
using HDF5
using PyCall

# get access to the atomic numbers
@pyimport ase.calculators.emt as ase_emt
@pyimport ase.data as ase_data

# import everything we need from ASE
const Bohr = ase_emt.Bohr
# get_atomic_numbers(id::AbstractString) = ase_data.atomic_numbers[id]
const ase_parameters = ase_emt.parameters

# convert ASE style Dict to a new Dict where parameters
# are stored in a type instead of tuple.
const params = Dict{String, Dict{String,Float64}}()

beta = 1.809
maxseq = maximum([par[1] for par in values(ase_parameters)]) * Bohr
rc = beta * maxseq * 0.5 * (sqrt(3) + sqrt(4))
rr = rc * 2 * sqrt(4) / (sqrt(3) + sqrt(4))
acut = log(9999.0) / (rr - rc)

for (key, p) in ase_parameters   # p is a tuple of parameters, key a species
   s0 = p[1] * Bohr
   eta2 = p[3] / Bohr
   kappa = p[4] / Bohr
   x = eta2 * beta * s0
   gamma1 = 0.0
   gamma2 = 0.0
   for (i, n) in enumerate([12; 6; 24])
      r = s0 * beta * sqrt(i + 1.0)
      x = n / (12.0 * (1.0 + exp(acut * (r - rc))))
      gamma1 += x * exp(-eta2 * (r - beta * s0))
      gamma2 += x * exp(-kappa / beta * (r - beta * s0))
      params[key] = Dict{String, Float64}(
         "E0" => p[1], "s0" => s0, "V0" => p[2], "eta2" => eta2,
         "kappa" => kappa, "lambda" => p[5] / Bohr, "n0" => p[6] / Bohr^3,
         "rc" => rc, "gamma1" => gamma1, "gamma2" => gamma2 )
   end
end

# store: params, rc, acut, beta
# save("emt_default_params.jld", "params", params, "beta", beta, "rc", rc, "acut", acut)
h5open("emtdefaultparams.h5", "w") do file
   write(file, "params", params)
   write(file, "beta", beta)
   write(file, "rc", rc)
   write(file, "acut", acut)
end


text:
@testset "util" begin
    using Swizzles: masktuple, imasktuple
    Is = [(),
          (1,),
          (nil,),
          (1, 2,),
          (nil, 2,),
          (2, nil,),
          (1, 2, 4,),
          (1, 2, nil, 4,),
          (1, nil, 3,)]
    for I in Is
        @test length(masktuple(+, -, I)) == length(I)
        @test masktuple(+, -, I) isa Tuple{Vararg{Int}}
        for j in 1:length(I)
            if I[j] isa Nil
                @test(masktuple(+, -, I)[j] == j)
            else
               @test(masktuple(+, -, I)[j] == -I[j])
            end
        end
        @test imasktuple(+, -, I, max(0, I...)) isa Tuple{Vararg{Int}}
        @test length(imasktuple(+, -, I, 0)) == 0
        @test length(imasktuple(+, -, I, 1)) == 1
        @test length(imasktuple(+, -, I, 2)) == 2
        for j in 1:length(I)
            if !(I[j] isa Nil)
                @test(imasktuple(+, -, I, max(0, I...))[I[j]] == -j)
            end
        end
        for j in 1:max(0, I...)
            if !(j in I)
                @test(imasktuple(+, -, I, max(0, I...))[j] == j)
            end
        end
    end
end


text:
using Base.Test

cd(dirname(@__FILE__)) do
    isdir("tmp") || mkpath("tmp")
    include("test_gdal_tutorials.jl")
    include("test_geometry.jl")
    include("test_types.jl")
    include("test_drivers.jl")
    include("test_feature.jl")
    include("test_featurelayer.jl")
    include("test_fielddefn.jl")
    include("test_styletable.jl")
    include("test_dataset.jl")
    include("test_rasterband.jl")
    include("test_rasterio.jl")
    include("test_rasterattrtable.jl")
    include("test_ospy_examples.jl")
    include("test_geos_operations.jl")
    include("test_cookbook_geometry.jl")
    include("test_cookbook_projection.jl")
end


text:
@testset "OLS scalar variate" begin
    for _ in 1:100
        k = rand(5:10)
        n = rand(90:100) + k
        X = randn(n, k)
        β = randn(k)
        Y = X * β + randn(n)
        data = OLSData(Y, X)
        @test repr(data) ==
            "OLS regression data, $(n) scalar observations, $(k) covariates"
        @test_throws ArgumentError OLSData(randn(n+1), X)
        params = MLE(OLSModel(), data)
        @test repr(params) == "OLS regression parameters, $(k) covariates"
        @test params.B ≈ (X \ Y)
        @test params.Σ ≈ sum(abs2, Y - X * (X \ Y)) / n
        for _ in 1:10
            β′ = randn(k)
            Σ′ = abs(randn()) + 1
            @test loglikelihood(OLSModel(), data, OLSParams(β′, Σ′)) ≈
                sum(logpdf.(Normal(0, √Σ′), Y - X * β′))
        end
    end
end

@testset "OLS matrix variate" begin
    m = rand(3:6)
    k = rand(5:10)
    n = rand(90:100) + k
    A = randn(m, m)
    X = randn(n, k)
    B = randn(k, m)
    Y = X * B + randn(n, m) * A
    data = OLSData(Y, X)
    @test repr(data) ==
        "OLS regression data, $(n) vector observations of length $(m), $(k) covariates"
    @test_throws ArgumentError OLSData(randn(n+1, ), X)
    params = MLE(OLSModel(), data)
    @test repr(params) ==
        "OLS regression parameters, vector observations of length $(m), $(k) covariates"
    @test params.B ≈ (X \ Y)
    E = Y - X * (X \ Y)
    @test params.Σ ≈ E'*E / n
    for _ in 1:10
        B′ = randn(k, m)
        A′ = randn(m, m)
        Σ′ = A′' * A′           # ensures PSD
        @test loglikelihood(OLSModel(), data, OLSParams(B′, Σ′)) ≈
            StatsBase.loglikelihood(MvNormal(zeros(m), Σ′), (Y - X * B′)')
    end
    # corner case: zero estimated variance matrix
    params2 = MLE(OLSModel(), data)
    params2.Σ .= 0
    @test loglikelihood(OLSModel(), data, params2) == -Inf
end

@testset "adding intercept" begin
    N = 11
    x = randn(N)
    X = randn(N, 3)
    x2 = randn(N)
    @test add_intercept(x) == hcat(ones(N), x)
    @test add_intercept(X) == hcat(ones(N), X)
    @test add_intercept(x, x2) == hcat(ones(N), x, x2)
    @test_throws DimensionMismatch add_intercept(x, ones(N+1))
end

@testset "vectorizing parameters (OLS)" begin
    b = randn(rand(3:10))
    σ² = abs(randn())
    @test vec_parameters(OLSParams(b, σ²)) == vcat(b, σ²)
    B = randn(10, 3)
    Σ = randn(3, 3) |> x->x*x'
    @test vec_parameters(OLSParams(B, Σ)) == vcat(vec_parameters(B),
                                                  vec_parameters(Σ))
end


text:
export DCDPTmpVars

""""
    ensure_diagonal(x)

Create (possibly sparse) matrix which ensures diagonal has all entries (even if zero).
"""
function ensure_diagonal(x)
    n = checksquare(x)
    A = x + SparseMatrixCSC(I, n, n)
    Arows = rowvals(A)
    Avals = nonzeros(A)
    for j = OneTo(n)
       for i in nzrange(A, j)
          j == Arows[i]  &&  (Avals[i] -= 1)
       end
    end
    check_rowvals_equal(A,x) || throw(error("rowvals are not the same..."))
    return A
end

check_rowvals_equal(A::SparseMatrixCSC, X::SparseMatrixCSC) = rowvals(A) == rowvals(X)
check_rowvals_equal(A, X) = true

"Temp vars for dynamic model"
struct DCDPTmpVars{
    T<:Real, SM<:AbstractMatrix{T}, AA3<:AbstractArray3{T}, AA3b<:AbstractArray3{T}, AA4<:AbstractArray4{T}
    } <: AbstractTmpVars
    ubVfull::AA3
    dubVfull::AA4
    dubVfullperm::AA4
    q::AA3b
    lse::Matrix{T}
    tmp::Matrix{T}
    tmp_cart::Matrix{CartesianIndex{3}}
    Πψtmp::Matrix{T}
    IminusTEVp::SM
    tmpEVj::Vector{T}

    ΠsumdubVj::Matrix{T} #  = view(lse(t), :, 1:nk) # Array{T}(nz,nθ)
    dev0tmpj::Matrix{T}  #  = view(tmp(t), :, 1:nk) # Array{T}(nz,nθ)

    function DCDPTmpVars(ubVfull::AA3, dubVfull::AA4, dubVfullperm::AA4, q::AA3b, lse, tmp, tmp_cart, Πψtmp, IminusTEVp::SM, tmpEVj, ΠsumdubVj, dev0tmpj) where {AA3, AA3b, AA4, SM}
        (nθt, nz, nψ, nd) = size(dubVfull)
        (     nz, nψ, nθt, nd) == size(dubVfullperm) || throw(DimensionMismatch())
        (nz, nψ, nd) == size(ubVfull) == size(q) || throw(DimensionMismatch())
        (nz,nψ) == size(lse) == size(tmp) == size(tmp_cart) || throw(DimensionMismatch())
        (nz,nθt) == size(dev0tmpj) == size(ΠsumdubVj) || throw(DimensionMismatch())
        nψ == checksquare(Πψtmp) || throw(DimensionMismatch())
        nz == checksquare(IminusTEVp) == length(tmpEVj)|| throw(DimensionMismatch())
        T = eltype(ubVfull)
        return new{T,SM,AA3,AA3b,AA4}(ubVfull, dubVfull, dubVfullperm, q, lse, tmp, tmp_cart, Πψtmp, IminusTEVp, tmpEVj, ΠsumdubVj, dev0tmpj)
    end
end

const dcdp_tmpvars = DCDPTmpVars
const DCDPTmpVarsArray{T,SM} = DCDPTmpVars{T,SM,Array{T,3},Array{T,3},Array{T,4}}
const DCDPTmpVarsView = DCDPTmpVars{T,SM,AA3} where {T,SM,AA3<:SubArray}

@getFieldFunction DCDPTmpVars ubVfull dubVfull dubVfullperm q lse tmp tmp_cart
@getFieldFunction DCDPTmpVars Πψtmp IminusTEVp tmpEVj ΠsumdubVj dev0tmpj

size(x::DCDPTmpVars) = size(dubVfullperm(x))

ubV(     x::DCDPTmpVars{T,SM,AA3}) where {T,SM,AA3<:SubArray} = ubVfull(x)
dubV(    x::DCDPTmpVars{T,SM,AA3}) where {T,SM,AA3<:SubArray} = dubVfull(x)
dubVperm(x::DCDPTmpVars{T,SM,AA3}) where {T,SM,AA3<:SubArray} = dubVfullperm(x)

function fill!(t::DCDPTmpVars, x)
    fill!(ubVfull(     t), x)
    fill!(dubVfull(    t), x)
    fill!(dubVfullperm(t), x)
    fill!(q(           t), x)
    fill!(lse(         t), x)
    fill!(tmp(         t), x)
    # fill!(tmp_cart(    t), x)
    fill!(Πψtmp(       t), x)
    fill!(tmpEVj(t), x)
end

function DCDPTmpVars(nθt, nz, nψ, nd, ztransition::AbstractMatrix{T}) where {T<:Real}
    dubVfull = Array{T,4}(undef, nθt, nz, nψ,      nd)
    dubVperm = Array{T,4}(undef,      nz, nψ, nθt, nd)
    ubVfull  = Array{T,3}(undef,      nz, nψ,      nd)
    q = similar(ubVfull)
    lse = Matrix{T}(undef, nz, nψ)
    tmp = similar(lse)
    tmp_cart = similar(lse, CartesianIndex{3})
    Πψtmp = Matrix{T}(undef, nψ, nψ)
    IminusTEVp = ensure_diagonal(ztransition)
    tmpEVj = Vector{T}(undef, nz)
    ΠsumdubVj = Matrix{T}(undef, nz, nθt)
    dev0tmpj = similar(ΠsumdubVj)

    return DCDPTmpVars(
        ubVfull, dubVfull, dubVperm, q, lse, tmp, tmp_cart,
        Πψtmp, IminusTEVp, tmpEVj, ΠsumdubVj, dev0tmpj
    )
end

function view(t::DCDPTmpVars, idxd::AbstractVector)
    first(idxd) == 1 || throw(DomainError())
    last(idxd) <= size(ubVfull(t),3) || throw(DomainError())
    @views ubV  = view(ubVfull(t), :,:,  idxd)
    @views dubV = view(dubVfull(t),:,:,:,idxd)
    @views dubvperm = view(dubVfullperm(t),:,:,:,idxd)
    @views qq   = view(q(t), :,:,idxd)
    return dcdp_tmpvars(ubV, dubV, dubvperm, qq, lse(t), tmp(t), tmp_cart(t),
        Πψtmp(t), IminusTEVp(t), tmpEVj(t), ΠsumdubVj(t), dev0tmpj(t))
end

# --------------------------------------------------------
# Fill reward matrices
# --------------------------------------------------------

function update_static_payoffs!(tmpv::DCDPTmpVars, ddm::AbstractDrillModel, θ::AbstractVector, sidx::Integer, ichars::Tuple, dograd::Bool)
    ψspace = psispace(ddm)
    zs = product(zspace(ddm)...)
    zψpdct = product(zs, ψspace)
    nk = _nparm(reward(ddm))
    nc = num_choices(statespace(ddm), sidx)

    sztup = (length(zs), length(ψspace), nk, nc)
    size(tmpv) == sztup || throw(DimensionMismatch("size tmvp = $(size(tmpv)) vs $sztup"))

    dubv  = reshape(dubVfull(tmpv),    nk,  :, nc)
    dubvp = reshape(dubVfullperm(tmpv), :, nk, nc)
    ubv   = reshape( ubVfull(tmpv),     :, nc)

    for dp1 in dp1space(statespace(ddm), sidx)
        @inbounds for (i, (z,ψ)) in enumerate(zψpdct)
            obs = ObservationDrill(ddm, ichars, z, dp1-1, sidx)
            grad = uview(dubv, :, i, dp1)
            ubv[i,dp1] = flow!(grad, obs, θ, ψ, dograd)
        end
    end
    dograd && permutedims!(dubvp, dubv, [2,1,3])
end

@deprecate flow!(t::DCDPTmpVars, ddm::AbstractDrillModel, θ::AbstractVector, sidx::Integer, ichars::Tuple, dograd::Bool) update_static_payoffs!(t,ddm,θ,sidx,ichars,dograd)


text:
module TestPackage1
using SpecialFunctions

export mySpecialFunction

greet() = print("Hello World!")

function mySpecialFunction(x)
    return x^2 * gamma(x)
end
end # module

max_stars_repo_name: JeffreySarnoff/techytok-examples

max_issues_repo_name: JeffreySarnoff/techytok-examples

max_forks_repo_name: JeffreySarnoff/techytok-examples


text:
# %%codeblock
# Atomic and isotopic information
struct element
    # Full name of element, such as `"hydrogen"` for hydrogen
    name::String
    # Symbol of element such as `:Ga` for gallium
    symbol::Symbol
    # Atomic number of element, such as `8` for oxygen
    atomic_number::Int
    # relative atomic mass, such as `14.007` for nitrogen
    mass::Float64
    # ionic radii for common ion in Angstroms
    radii::Float64
end

# %%codeblock
mutable struct atom
    # element that a specific atom is
    element::element
    # fractional position of atom
    r_frac::Array{Union{Float64, Int}, 1}
    # cartesian position of aotm
    r_cart::Array{Union{Float64, Int}, 1}
end

# %%codeblock
dummy = element("dummy", :dummy, 200, 0.000001, 0.000001)
oxygen = element("oxygen", :O, 8, 15.999, 1.12)
zinc = element("zinc", :Zn, 30, 65.38, 0.88)
gallium = element("gallium", :Ga, 31, 69.723, 0.76)
indium = element("indium", :In, 49, 114.818, 0.94)
aluminum = element("aluminum", :Al, 13, 26.982, 0.68)


text:
using Roots

# Asymptotic formula for zeros of Bessel J function of order nu
"""
    besselj_asymptotic_zero(nu,n)

Asymptotic formula for the `n`th zero fo the the Bessel J function of order `nu`.
`besselj_asymptotic_zero` is vectorized.
"""
besselj_asymptotic_zero(nu,n) = return pi * (n-1 + nu/2 + 3//4)

# Find nth zero of Bessel J. Use the asymptotic values as
# starting values. These find the correct zeros even
# for n = 1,2,...

"""
    besselj_zero(nu,n)

`n`th zero of the Bessel J function of order `nu`,
for `n` = `1,2,...`. `besselj_zero` is vectorized.
"""
function besselj_zero(nu,n)
    z = besselj_asymptotic_zero(nu,n)
    bf = (x) -> besselj(nu,x)
#    return fzero(bf,z,ftol=1e-15)
    return fzero(bf,z)
end

@mkapprule BesselJZero nargs => 2

@doap BesselJZero(n::Number,z::AbstractFloat) = besselj_zero(n,z)
@doap BesselJZero(n::AbstractFloat,z::Number) = besselj_zero(n,z)
@doap BesselJZero(n::AbstractFloat,z::AbstractFloat) = besselj_zero(n,z)