text:
@doc raw"""
    Tucker{N, R, D, 𝔽} <: AbstractManifold{𝔽}

The manifold of ``N_1 \times \dots \times N_D`` real-valued or complex-valued tensors of
fixed multilinear rank ``(R_1, \dots, R_D)`` . If ``R_1 = \dots = R_D = 1``, this is the
Segre manifold, i.e., the set of rank-1 tensors.

# Representation in HOSVD format

Let ``\mathbb{F}`` be the real or complex numbers.
Any tensor ``p`` on the Tucker manifold can be represented as a multilinear product in HOSVD
[^DeLathauwer2000] form
```math
p = (U_1,\dots,U_D) \cdot \mathcal{C}
```
where ``\mathcal C \in \mathbb{F}^{R_1 \times \dots \times R_D}`` and, for ``d=1,\dots,D``,
the matrix ``U_d \in \mathbb{F}^{N_d \times R_d}`` contains the singular vectors of the
``d``th unfolding of ``\mathcal{A}``

# Tangent space

The tangent space to the Tucker manifold at
``p = (U_1,\dots,U_D) \cdot \mathcal{C}`` is [^Koch2010]
```math
T_p \mathcal{M} =
\bigl\{
(U_1,\dots,U_D) \cdot \mathcal{C}^\prime
+ \sum_{d=1}^D \bigl(
    (U_1, \dots, U_{d-1}, U_d^\prime, U_{d+1}, \dots, U_D)
    \cdot \mathcal{C}
\bigr)
\bigr\}
```
where ``\mathcal{C}^\prime`` is arbitrary, ``U_d^{\mathrm{H}}`` is the Hermitian adjoint of
``U_d``, and ``U_d^{\mathrm{H}} U_d^\prime = 0`` for all ``d``.

# Constructor
    Tucker(N::NTuple{D, Int}, R::NTuple{D, Int}[, field = ℝ])

Generate the manifold of `field`-valued tensors of dimensions  `N[1] × … × N[D]` and
multilinear rank `R = (R[1], …, R[D])`.

[^DeLathauwer2000]:
    > Lieven De Lathauwer, Bart De Moor, Joos Vandewalle: "A multilinear singular value decomposition"
    > SIAM Journal on Matrix Analysis and Applications, 21(4), pp. 1253-1278, 2000
    > doi: [10.1137/S0895479896305696](https://doi.org/10.1137/S0895479896305696)

[^Koch2010]:
    > Othmar Koch, Christian Lubic, "Dynamical Tensor approximation"
    > SIAM Journal on Matrix Analysis and Applications, 31(5), pp. 2360-2375, 2010
    > doi: [10.1137/09076578X](https://doi.org/10.1137/09076578X)
"""
struct Tucker{N,R,D,𝔽} <: AbstractManifold{𝔽} end
function Tucker(n⃗::NTuple{D,Int}, r⃗::NTuple{D,Int}, field::AbstractNumbers=ℝ) where {D}
    @assert is_valid_mlrank(n⃗, r⃗)
    return Tucker{n⃗,r⃗,D,field}()
end

#=
HOSVD{T, D}

Higher-order singular value decomposition of an order D tensor with eltype T
fields:
* U: singular vectors of the unfoldings
* core: core tensor
* σ : singular values of the unfoldings
=#
struct HOSVD{T,D}
    U::NTuple{D,Matrix{T}}
    core::Array{T,D}
    σ::NTuple{D,Vector{T}}
end

@doc raw"""
    TuckerPoint{T,D}

An order `D` tensor of fixed multilinear rank and entries of type `T`, which makes it a
point on the [`Tucker`](@ref) manifold. The tensor is represented in HOSVD form.

# Constructors:
    TuckerPoint(core::AbstractArray{T,D}, factors::Vararg{<:AbstractMatrix{T},D}) where {T,D}

Construct an order `D` tensor of element type `T` that can be represented as the
multilinear product `(factors[1], …, factors[D]) ⋅ core`.
It is assumed that the dimensions of the core are the multilinear
rank of the tensor and that the matrices `factors` each have full rank. No further
assumptions are made.

    TuckerPoint(p::AbstractArray{T,D}, mlrank::NTuple{D,Int}) where {T,D}

The low-multilinear rank tensor arising from the sequentially truncated the higher-order
singular value decomposition of the `D`-dimensional array `p` of type `T`. The singular
values are truncated to get a multilinear rank `mlrank`
[^Vannieuwenhoven2012].

[^Vannieuwenhoven2012]:
    > Nick Vannieuwenhoven, Raf Vandebril, Karl Meerbergen: "A new truncation strategy for the higher-order singular value decomposition"
    > SIAM Journal on Scientific Computing, 34(2), pp. 1027-1052, 2012
    > doi: [10.1137/110836067](https://doi.org/10.1137/110836067)

"""
struct TuckerPoint{T,D} <: AbstractManifoldPoint
    hosvd::HOSVD{T,D}
end
function TuckerPoint(
    core::AbstractArray{T,D},
    factors::Vararg{MtxT,D},
) where {T,D,MtxT<:AbstractMatrix{T}}
    # Take the QR decompositions of the factors and multiply the R factors into the core
    qrfacs = qr.(factors)
    Q = map(qrfac -> qrfac.Q, qrfacs)
    R = map(qrfac -> qrfac.R, qrfacs)
    core′ = reshape(Kronecker.:⊗(reverse(R)...) * vec(core), size(core))

    # Convert to HOSVD format by taking the HOSVD of the core
    decomp = st_hosvd(core′)
    factors′ = Q .* decomp.U
    return TuckerPoint(HOSVD{T,D}(factors′, decomp.core, decomp.σ))
end
function TuckerPoint(A::AbstractArray{T,D}, mlrank::NTuple{D,Int}) where {T,D}
    @assert is_valid_mlrank(size(A), mlrank)
    return TuckerPoint(st_hosvd(A, mlrank))
end

@doc raw"""
    TuckerTVector{T, D} <: TVector

Tangent vector to the `D`-th order [`Tucker`](@ref) manifold at
``p = (U_1,\dots,U_D) ⋅ \mathcal{C}``. The numbers are of type `T` and the vector is
represented as
````math
X =
(U_1,\dots,U_D) \cdot \mathcal{C}^\prime +
\sum_{d=1}^D (U_1,\dots,U_{d-1},U_d^\prime,U_{d+1},\dots,U_D) \cdot \mathcal{C}
````
where ``U_d^\mathrm{H} U_d^\prime = 0``.

# Constructor
    TuckerTVector(C′::Array{T,D}, U′::NTuple{D,Matrix{T}}) where {T,D}

Constructs a `D`th order [`TuckerTVector`](@ref) of number type `T` with ``C^\prime``
and ``U^\prime``, so that, together with a
[`TuckerPoint`](@ref) ``p`` as above, the tangent vector can be represented as ``X`` in the
above expression.
"""
struct TuckerTVector{T,D} <: TVector
    Ċ::Array{T,D}
    U̇::NTuple{D,Matrix{T}}
end

# An implicitly stored basis of the tangent space to the Tucker manifold. This is the basis
# from [Dewaele2021] and acts as the default orthonormal basis.
struct HOSVDBasis{T,D}
    point::TuckerPoint{T,D}
    U⊥::NTuple{D,Matrix{T}}
end
CachedHOSVDBasis{𝔽,T,D} =
    CachedBasis{𝔽,DefaultOrthonormalBasis{𝔽,TangentSpaceType},HOSVDBasis{T,D}}

⊗ᴿ(a...) = Kronecker.:⊗(reverse(a)...)

Base.:*(s::Number, x::TuckerTVector) = TuckerTVector(s * x.Ċ, s .* x.U̇)
Base.:*(x::TuckerTVector, s::Number) = TuckerTVector(x.Ċ * s, x.U̇ .* s)
Base.:/(x::TuckerTVector, s::Number) = TuckerTVector(x.Ċ / s, x.U̇ ./ s)
Base.:\(s::Number, x::TuckerTVector) = TuckerTVector(s \ x.Ċ, s .\ x.U̇)
Base.:+(x::TuckerTVector, y::TuckerTVector) = TuckerTVector(x.Ċ + y.Ċ, x.U̇ .+ y.U̇)
Base.:-(x::TuckerTVector, y::TuckerTVector) = TuckerTVector(x.Ċ - y.Ċ, x.U̇ .- y.U̇)
Base.:-(x::TuckerTVector) = TuckerTVector(-x.Ċ, map(-, x.U̇))
Base.:+(x::TuckerTVector) = TuckerTVector(x.Ċ, x.U̇)
Base.:(==)(x::TuckerTVector, y::TuckerTVector) = (x.Ċ == y.Ċ) && all(x.U̇ .== y.U̇)

allocate(p::TuckerPoint) = allocate(p, number_eltype(p))
function allocate(p::TuckerPoint{Tp,D}, ::Type{T}) where {T,Tp,D}
    @assert promote_type(Tp, T) == T
    return TuckerPoint(
        HOSVD(allocate(p.hosvd.U, T), allocate(p.hosvd.core, T), allocate(p.hosvd.σ, T)),
    )
end
allocate(x::TuckerTVector) = allocate(x, number_eltype(x))
function allocate(x::TuckerTVector, ::Type{T}) where {T}
    return TuckerTVector(allocate(x.Ċ, T), allocate(x.U̇, T))
end

# Tuple-like broadcasting of TuckerTVector
Base.axes(::TuckerTVector) = ()

function Broadcast.BroadcastStyle(::Type{TuckerTVector{T,D}}) where {T,D}
    return Broadcast.Style{TuckerTVector{Any,D}}()
end
function Broadcast.BroadcastStyle(
    ::Broadcast.AbstractArrayStyle{0},
    b::Broadcast.Style{<:TuckerTVector},
)
    return b
end

function Broadcast.instantiate(
    bc::Broadcast.Broadcasted{Broadcast.Style{TuckerTVector{Any,D}},Nothing},
) where {D}
    return bc
end
function Broadcast.instantiate(
    bc::Broadcast.Broadcasted{Broadcast.Style{TuckerTVector{Any,D}}},
) where {D}
    Broadcast.check_broadcast_axes(bc.axes, bc.args...)
    return bc
end

Broadcast.broadcastable(v::TuckerTVector) = v

Base.@propagate_inbounds function Broadcast._broadcast_getindex(
    v::TuckerTVector,
    ::Val{I},
) where {I}
    if I isa Symbol
        return getfield(v, I)
    else
        return getfield(v, I[1])[I[2]]
    end
end

####

@doc raw"""
    check_point(M::Tucker{N,R,D}, p; kwargs...) where {N,R,D}

Check whether the multidimensional array or [`TuckerPoint`](@ref) `p` is a point on the
[`Tucker`](@ref) manifold, i.e. it is a `D`th order `N[1] × … × N[D]` tensor of multilinear
rank `(R[1], …, R[D])`. The keyword arguments are passed to the matrix rank function applied
to the unfoldings.
For a [`TuckerPoint`](@ref) it is checked that the point is in correct HOSVD form.
"""
function check_point(M::Tucker{N,R,D}, x; kwargs...) where {N,R,D}
    s = "The point $(x) does not lie on $(M), "
    size(x) == N || return DomainError(size(x), s * "since its size is not $(N).")
    x_buffer = similar(x)
    for d in 1:ndims(x)
        r = rank(tensor_unfold!(x_buffer, x, d); kwargs...)
        r == R[d] || return DomainError(size(x), s * "since its rank is not $(R).")
    end
    return nothing
end
function check_point(M::Tucker{N,R,D}, x::TuckerPoint; kwargs...) where {N,R,D}
    s = "The point $(x) does not lie on $(M), "
    U = x.hosvd.U
    ℭ = x.hosvd.core
    if size(ℭ) ≠ R
        return DomainError(
            size(x.hosvd.core),
            s * "since the size of the core is not $(R).",
        )
    end
    if size(x) ≠ N
        return DomainError(size(x), s * "since its dimensions are not $(N).")
    end
    for u in U
        if u' * u ≉ LinearAlgebra.I
            return DomainError(
                norm(u' * u - LinearAlgebra.I),
                s * "since its factor matrices are not unitary.",
            )
        end
    end
    ℭ_buffer = similar(ℭ)
    for d in 1:ndims(x.hosvd.core)
        ℭ⁽ᵈ⁾ = tensor_unfold!(ℭ_buffer, ℭ, d)
        gram = ℭ⁽ᵈ⁾ * ℭ⁽ᵈ⁾'
        if gram ≉ Diagonal(x.hosvd.σ[d])^2
            return DomainError(
                norm(gram - Diagonal(x.hosvd.σ[d])^2),
                s *
                "since the unfoldings of the core are not diagonalised by" *
                "the singular values.",
            )
        end
        if rank(Diagonal(x.hosvd.σ[d]); kwargs...) ≠ R[d]
            return DomainError(
                minimum(x.hosvd.σ[d]),
                s * "since the core does not have full multilinear rank.",
            )
        end
    end
    return nothing
end

@doc raw"""
    check_vector(M::Tucker{N,R,D}, p::TuckerPoint{T,D}, X::TuckerTVector) where {N,R,T,D}

Check whether a [`TuckerTVector`](@ref) `X` is is in the tangent space to
the `D`th order [`Tucker`](@ref) manifold `M` at the `D`th order [`TuckerPoint`](@ref) `p`.
This
is the case when the dimensions of the factors in `X` agree with those of `p` and the factor
matrices of `X` are in the orthogonal complement of the HOSVD factors of `p`.
"""
function check_vector(
    M::Tucker{N,R,D},
    p::TuckerPoint{T,D},
    v::TuckerTVector,
) where {N,R,T,D}
    s = "The tangent vector $(v) is not a tangent vector to $(p) on $(M), "
    if size(p.hosvd.core) ≠ size(v.Ċ) || any(size.(v.U̇) .≠ size.(p.hosvd.U))
        return DomainError(
            size(v.Ċ),
            s * "since the array dimensons of $(p) and $(v)" * "do not agree.",
        )
    end
    for (U, U̇) in zip(p.hosvd.U, v.U̇)
        if norm(U' * U̇) ≥ √eps(eltype(U)) * √length(U)
            return DomainError(
                norm(U' * U̇),
                s *
                "since the columns of x.hosvd.U are not" *
                "orthogonal to those of v.U̇.",
            )
        end
    end
    return nothing
end

"""
    Base.convert(::Type{Matrix{T}}, basis::CachedBasis{𝔽,DefaultOrthonormalBasis{𝔽, TangentSpaceType},HOSVDBasis{T, D}}) where {𝔽, T, D}
    Base.convert(::Type{Matrix}, basis::CachedBasis{𝔽,DefaultOrthonormalBasis{𝔽, TangentSpaceType},HOSVDBasis{T, D}}) where {𝔽, T, D}

Convert a HOSVD-derived cached basis from [^Dewaele2021] of the `D`th order
[`Tucker`](@ref) manifold with number type `T` to a matrix.
The columns of this matrix are the vectorisations of the
[`embed`](@ref)dings of the basis vectors.
"""
function Base.convert(::Type{Matrix{T}}, ℬ::CachedHOSVDBasis{𝔽,T,D}) where {𝔽,T,D}
    𝔄 = ℬ.data.point
    r⃗ = size(𝔄.hosvd.core)
    n⃗ = size(𝔄)
    ℳ = Tucker(n⃗, r⃗)

    J = Matrix{T}(undef, prod(n⃗), manifold_dimension(ℳ))
    # compute all possible ∂𝔄╱∂ℭ (in one go is quicker than one vector at a time)
    J[:, 1:prod(r⃗)] = ⊗ᴿ(𝔄.hosvd.U...)
    # compute all possible ∂𝔄╱∂U[d] for d = 1,...,D
    function fill_column!(i, vᵢ)
        Jᵢ_tensor = reshape(view(J, :, i), n⃗) # changes to this apply to J as well
        return embed!(ℳ, Jᵢ_tensor, 𝔄, vᵢ)
    end
    foreach(fill_column!, ℳ, 𝔄, ℬ, (prod(r⃗) + 1):manifold_dimension(ℳ))
    return J
end
function Base.convert(::Type{Matrix}, basis::CachedHOSVDBasis{𝔽,T,D}) where {𝔽,T,D}
    return convert(Matrix{T}, basis)
end

@inline function Base.copy(
    bc::Broadcast.Broadcasted{Broadcast.Style{TuckerTVector{Any,D}}},
) where {D}
    return TuckerTVector(
        @inbounds(Broadcast._broadcast_getindex(bc, Val(:Ċ))),
        ntuple(i -> @inbounds(Broadcast._broadcast_getindex(bc, Val((:U̇, i)))), Val(D)),
    )
end
Base.copy(x::TuckerTVector) = TuckerTVector(copy(x.Ċ), map(copy, x.U̇))

function Base.copyto!(q::TuckerPoint, p::TuckerPoint)
    for d in 1:ndims(q)
        copyto!(q.hosvd.U[d], p.hosvd.U[d])
        copyto!(q.hosvd.σ[d], p.hosvd.σ[d])
    end
    copyto!(q.hosvd.core, p.hosvd.core)
    return q
end
function Base.copyto!(y::TuckerTVector, x::TuckerTVector)
    for d in 1:ndims(y.Ċ)
        copyto!(y.U̇[d], x.U̇[d])
    end
    copyto!(y.Ċ, x.Ċ)
    return y
end
@inline function Base.copyto!(
    dest::TuckerTVector,
    bc::Broadcast.Broadcasted{Broadcast.Style{TuckerTVector{Any,D}}},
) where {D}
    # Performance optimization: broadcast!(identity, dest, A) is equivalent to copyto!(dest, A) if indices match
    if bc.f === identity && bc.args isa Tuple{TuckerTVector} # only a single input argument to broadcast!
        A = bc.args[1]
        return copyto!(dest, A)
    end
    bc′ = Broadcast.preprocess(dest, bc)
    copyto!(dest.Ċ, Broadcast._broadcast_getindex(bc′, Val(:Ċ)))
    for i in 1:D
        copyto!(dest.U̇[i], Broadcast._broadcast_getindex(bc, Val((:U̇, i))))
    end
    return dest
end

@doc raw"""
    embed(::Tucker{N,R,D}, p::TuckerPoint) where {N,R,D}

Convert a [`TuckerPoint`](@ref) `p` on the rank `R` [`Tucker`](@ref) manifold to a full
`N[1] × … × N[D]`-array by evaluating the Tucker decomposition.

    embed(::Tucker{N,R,D}, p::TuckerPoint, X::TuckerTVector) where {N,R,D}

Convert a tangent vector `X` with base point `p` on the rank `R` [`Tucker`](@ref)
manifold to a full tensor, represented as an `N[1] × … × N[D]`-array.
"""
embed(::Tucker, ::Any, ::TuckerPoint)

function embed!(::Tucker, q, p::TuckerPoint)
    return copyto!(q, reshape(⊗ᴿ(p.hosvd.U...) * vec(p.hosvd.core), size(p)))
end
function embed!(ℳ::Tucker, Y, 𝔄::TuckerPoint{T,D}, X::TuckerTVector) where {T,D}
    mul!(vec(Y), ⊗ᴿ(𝔄.hosvd.U...), vec(X.Ċ))
    𝔄_embedded = embed(ℳ, 𝔄)
    buffer = similar(𝔄_embedded)
    for k in 1:D
        U̇ₖUₖᵀ𝔄₍ₖ₎ = X.U̇[k] * (𝔄.hosvd.U[k]' * tensor_unfold!(buffer, 𝔄_embedded, k))
        Y .= Y + tensor_fold!(buffer, U̇ₖUₖᵀ𝔄₍ₖ₎, k)
    end
    return Y
end

@doc raw"""
    Base.foreach(f, M::Tucker, p::TuckerPoint, basis::AbstractBasis, indices=1:manifold_dimension(M))

Let `basis` be and [`AbstractBasis`](https://juliamanifolds.github.io/ManifoldsBase.jl/stable/bases.html#ManifoldsBase.AbstractBasis) at a point `p` on `M`. Suppose `f` is a function
that takes an index and a vector as an argument.
This function applies `f` to `i` and the `i`th basis vector sequentially for each `i` in
`indices`.
Using a [`CachedBasis`](https://juliamanifolds.github.io/ManifoldsBase.jl/stable/bases.html#ManifoldsBase.CachedBasis) may speed up the computation.

**NOTE**: The i'th basis vector is overwritten in each iteration. If any information about
the vector is to be stored, `f` must make a copy.
"""
function Base.foreach(
    f,
    M::Tucker,
    p::TuckerPoint,
    basis::AbstractBasis,
    indices=1:manifold_dimension(M),
)
    # Use mutating variants to avoid superfluous allocation
    bᵢ = zero_vector(M, p)
    eᵢ = zeros(number_eltype(p), manifold_dimension(M))
    for i in indices
        eᵢ[i] = one(eltype(eᵢ))
        get_vector!(M, bᵢ, p, eᵢ, basis)
        eᵢ[i] = zero(eltype(eᵢ))
        f(i, bᵢ)
    end
end

@doc raw"""
    get_basis(:: Tucker, p::TuckerPoint, basisType::DefaultOrthonormalBasis{𝔽, TangentSpaceType}) where 𝔽

An implicitly stored basis of the tangent space to the Tucker manifold.
Assume ``p = (U_1,\dots,U_D) \cdot \mathcal{C}`` is in HOSVD format and that, for
``d=1,\dots,D``, the singular values of the
``d``'th unfolding are ``\sigma_{dj}``, with ``j = 1,\dots,R_d``.
The basis of the tangent space is as follows: [^Dewaele2021]

````math
\bigl\{
(U_1,\dots,U_D) e_i
\bigr\} \cup \bigl\{
(U_1,\dots, \sigma_{dj}^{-1} U_d^{\perp} e_i e_j^T,\dots,U_D) \cdot \mathcal{C}
\bigr\}
````

for all ``d = 1,\dots,D`` and all canonical basis vectors ``e_i`` and ``e_j``.
Every ``U_d^\perp`` is such that ``[U_d \quad U_d^{\perp}]`` forms an orthonormal basis
of ``\mathbb{R}^{N_d}``.

[^Dewaele2021]:
    > Nick Dewaele, Paul Breiding, Nick Vannieuwenhoven, "The condition number of many tensor decompositions is invariant under Tucker compression"
    > arxiv: [2106.13034](https://arxiv.org/abs/2106.13034)
"""
function get_basis(
    ::Tucker,
    𝔄::TuckerPoint,
    basisType::DefaultOrthonormalBasis{𝔽,TangentSpaceType}=DefaultOrthonormalBasis(),
) where {𝔽}
    D = ndims(𝔄)
    n⃗ = size(𝔄)
    r⃗ = size(𝔄.hosvd.core)

    U = 𝔄.hosvd.U
    U⊥ = ntuple(d -> Matrix(qr(I - U[d] * U[d]', Val(true)).Q)[:, 1:(n⃗[d] - r⃗[d])], D)

    basis = HOSVDBasis(𝔄, U⊥)
    return CachedBasis(basisType, basis)
end

#=
get_coordinates(::Tucker, A, X::TuckerTVector, b)

The coordinates of a tangent vector X at point A on the Tucker manifold with respect to the
basis b.
=#
function get_coordinates(::Tucker, 𝔄, X::TuckerTVector, ℬ::CachedHOSVDBasis)
    coords = vec(X.Ċ)
    for d in 1:length(X.U̇)
        coord_mtx = (ℬ.data.U⊥[d] \ X.U̇[d]) * Diagonal(𝔄.hosvd.σ[d])
        coords = vcat(coords, vec(coord_mtx'))
    end
    return coords
end
function get_coordinates(
    M::Tucker,
    𝔄,
    X,
    ℬ::DefaultOrthonormalBasis{𝔽,TangentSpaceType},
) where {𝔽}
    return get_coordinates(M, 𝔄, X, get_basis(M, 𝔄, ℬ))
end

#=
get_vector(::Tucker, A, x, b)

The tangent vector at a point A whose coordinates with respect to the basis b are x.
=#
function get_vector!(
    ::Tucker,
    y,
    𝔄::TuckerPoint,
    x::AbstractVector{T},
    ℬ::CachedHOSVDBasis,
) where {T}
    ξ = convert(Vector{promote_type(number_eltype(𝔄), eltype(x))}, x)
    ℭ = 𝔄.hosvd.core
    σ = 𝔄.hosvd.σ
    U⊥ = ℬ.data.U⊥
    D = ndims(ℭ)
    r⃗ = size(ℭ)
    n⃗ = size(𝔄)

    # split ξ into ξ_core and ξU so that vcat(ξ_core, ξU...) == ξ, but avoid copying
    ξ_core = view(ξ, 1:length(ℭ))
    ξU = Vector{typeof(ξ_core)}(undef, D)
    nextcolumn = length(ℭ) + 1
    for d in 1:D
        numcols = r⃗[d] * (n⃗[d] - r⃗[d])
        ξU[d] = view(ξ, nextcolumn:(nextcolumn + numcols - 1))
        nextcolumn += numcols
    end

    # Construct ∂U[d] by plugging in the definition of the orthonormal basis [Dewaele2021]
    # ∂U[d] = ∑ᵢⱼ { ξU[d]ᵢⱼ (σ[d]ⱼ)⁻¹ U⊥[d] 𝐞ᵢ 𝐞ⱼᵀ }
    #       = U⊥[d] * ∑ⱼ (σ[d]ⱼ)⁻¹ (∑ᵢ ξU[d]ᵢⱼ  𝐞ᵢ) 𝐞ⱼᵀ
    # ξU[d] = [ξ₁₁, ..., ξ₁ⱼ, ..., ξᵢ₁, ..., ξᵢⱼ, ..., ]
    # => turn these i and j into matrix indices and do matrix operations
    for d in 1:D
        grid = transpose(reshape(ξU[d], r⃗[d], n⃗[d] - r⃗[d]))
        mul!(y.U̇[d], U⊥[d], grid * Diagonal(1 ./ σ[d]))
    end

    y.Ċ .= reshape(ξ_core, size(y.Ċ))
    return y
end
function get_vector!(
    ℳ::Tucker,
    y,
    𝔄::TuckerPoint,
    x,
    ℬ::DefaultOrthonormalBasis{𝔽,TangentSpaceType},
) where {𝔽}
    return get_vector!(ℳ, y, 𝔄, x, get_basis(ℳ, 𝔄, ℬ))
end

function get_vectors(ℳ::Tucker, 𝔄::TuckerPoint{T,D}, ℬ::CachedHOSVDBasis) where {T,D}
    vectors = Vector{TuckerTVector{T,D}}(undef, manifold_dimension(ℳ))
    foreach((i, vᵢ) -> setindex!(vectors, copy(vᵢ), i), ℳ, 𝔄, ℬ)
    return vectors
end
function get_vectors(ℳ::Tucker, 𝔄::TuckerPoint, ℬ::DefaultOrthonormalBasis)
    return get_vectors(ℳ, 𝔄, get_basis(ℳ, 𝔄, ℬ))
end

"""
    inner(M::Tucker, p::TuckerPoint, X::TuckerTVector, Y::TuckerTVector)

The Euclidean inner product between tangent vectors `X` and `X` at the point `p` on
the Tucker manifold. This is equal to `embed(M, p, X) ⋅ embed(M, p, Y)`.

    inner(::Tucker, A::TuckerPoint, X::TuckerTVector, Y)
    inner(::Tucker, A::TuckerPoint, X, Y::TuckerTVector)

The Euclidean inner product between `X` and `Y` where `X` is a vector tangent to the Tucker
manifold at `p` and `Y` is a vector in the ambient space or vice versa. The vector in the
ambient space is represented as a full tensor, i.e., a multidimensional array.
"""
function inner(::Tucker, 𝔄::TuckerPoint, x::TuckerTVector, y::TuckerTVector)
    ℭ = 𝔄.hosvd.core
    dotprod = dot(x.Ċ, y.Ċ)
    ℭ_buffer = similar(ℭ)
    for k in 1:ndims(𝔄)
        ℭ₍ₖ₎ = tensor_unfold!(ℭ_buffer, ℭ, k)
        dotprod += dot(x.U̇[k] * ℭ₍ₖ₎, y.U̇[k] * ℭ₍ₖ₎)
    end
    return dotprod
end
inner(M::Tucker, 𝔄::TuckerPoint, x::TuckerTVector, y) = dot(embed(M, 𝔄, x), y)
inner(M::Tucker, 𝔄::TuckerPoint, x, y::TuckerTVector) = dot(x, embed(M, 𝔄, y))

"""
    inverse_retract(M::Tucker, p::TuckerPoint, q::TuckerPoint, ::ProjectionInverseRetraction)

The projection inverse retraction on the Tucker manifold interprets `q` as a point in the
ambient Euclidean space (see [`embed`](@ref)) and projects it onto the tangent space at
to `M` at `p`.
"""
inverse_retract(
    ::Tucker,
    ::Any,
    ::TuckerPoint,
    ::TuckerPoint,
    ::ProjectionInverseRetraction,
)

function inverse_retract_project!(ℳ::Tucker, X, 𝔄::TuckerPoint, 𝔅::TuckerPoint)
    diffVector = embed(ℳ, 𝔅) - embed(ℳ, 𝔄)
    return project!(ℳ, X, 𝔄, diffVector)
end

function isapprox(p::TuckerPoint, q::TuckerPoint; kwargs...)
    ℳ = Tucker(size(p), size(p.hosvd.core))
    return isapprox(embed(ℳ, p), embed(ℳ, q); kwargs...)
end
isapprox(::Tucker, p::TuckerPoint, q::TuckerPoint; kwargs...) = isapprox(p, q; kwargs...)
function isapprox(M::Tucker, p::TuckerPoint, x::TuckerTVector, y::TuckerTVector; kwargs...)
    return isapprox(embed(M, p, x), embed(M, p, y); kwargs...)
end

#=
Determines whether there are tensors of dimensions n⃗ with multilinear rank r⃗
=#
function is_valid_mlrank(n⃗, r⃗)
    return all(r⃗ .≥ 1) &&
           all(r⃗ .≤ n⃗) &&
           all(ntuple(i -> r⃗[i] ≤ prod(r⃗) ÷ r⃗[i], length(r⃗)))
end

@doc raw"""
    manifold_dimension(::Tucker{N,R,D}) where {N,R,D}

The dimension of the manifold of ``N_1 \times \dots \times N_D`` tensors of multilinear
rank ``(R_1, \dots, R_D)``, i.e.
````math
\mathrm{dim}(\mathcal{M}) = \prod_{d=1}^D R_d + \sum_{d=1}^D R_d (N_d - R_d).
````
"""
manifold_dimension(::Tucker{n⃗,r⃗}) where {n⃗,r⃗} = prod(r⃗) + sum(r⃗ .* (n⃗ .- r⃗))

@doc raw"""
    Base.ndims(p::TuckerPoint{T,D}) where {T,D}

The order of the tensor corresponding to the [`TuckerPoint`](@ref) `p`, i.e., `D`.
"""
Base.ndims(::TuckerPoint{T,D}) where {T,D} = D

number_eltype(::TuckerPoint{T,D}) where {T,D} = T
number_eltype(::TuckerTVector{T,D}) where {T,D} = T

"""
    project(M::Tucker, p::TuckerPoint, X)

The least-squares projection of a dense tensor `X` onto the tangent space to `M` at `p`.
"""
project(::Tucker, ::Any, ::TuckerPoint, ::Any)

function project!(ℳ::Tucker, Y, 𝔄::TuckerPoint, X)
    ℬ = get_basis(ℳ, 𝔄, DefaultOrthonormalBasis())
    coords = Vector{number_eltype(𝔄)}(undef, manifold_dimension(ℳ))
    f!(i, ℬᵢ) = setindex!(coords, inner(ℳ, 𝔄, ℬᵢ, X), i)
    foreach(f!, ℳ, 𝔄, ℬ)
    return get_vector!(ℳ, Y, 𝔄, coords, ℬ)
end

@doc raw"""
    retract(::Tucker, p::TuckerPoint, X::TuckerTVector, ::PolarRetraction)

The truncated HOSVD-based retraction [^Kressner2014] to the Tucker manifold, i.e.
the result is the sequentially tuncated HOSVD approximation of ``p + X``.

In the exceptional case that the multilinear rank of ``p + X`` is lower than that of ``p``, this
retraction produces a boundary point, which is outside the manifold.

[^Kressner2014]:
    > Daniel Kressner, Michael Steinlechner, Bart Vandereycken: "Low-rank tensor completion by Riemannian optimization"
    > BIT Numerical Mathematics, 54(2), pp. 447-468, 2014
    > doi: [10.1007/s10543-013-0455-z](https://doi.org/10.1007/s10543-013-0455-z)

"""
retract(::Tucker, ::Any, ::Any, ::PolarRetraction)

function retract_polar!(
    ::Tucker,
    q::TuckerPoint,
    p::TuckerPoint{T,D},
    x::TuckerTVector,
) where {T,D}
    U = p.hosvd.U
    V = x.U̇
    ℭ = p.hosvd.core
    𝔊 = x.Ċ
    r⃗ = size(ℭ)

    # Build the core tensor S and the factors [Uᵈ  Vᵈ]
    S = zeros(T, 2 .* size(ℭ))
    S[CartesianIndices(ℭ)] = ℭ + 𝔊
    UQ = Matrix{T}[]
    buffer = similar(ℭ)
    for k in 1:D
        # We make the following adaptation to Kressner2014:
        # Fix the i'th term of the sum and replace Vᵢ by Qᵢ Rᵢ.
        # We can absorb the R factor into the core by replacing Vᵢ by Qᵢ
        # and C (in the i'th term of the sum) by C ×ᵢ Rᵢ
        Q, R = qr(V[k])
        idxOffset = CartesianIndex(ntuple(i -> i == k ? r⃗[k] : 0, D))
        ℭ⨉ₖR = tensor_fold!(buffer, R * tensor_unfold!(buffer, ℭ, k), k)
        S[CartesianIndices(ℭ) .+ idxOffset] = ℭ⨉ₖR
        push!(UQ, hcat(U[k], Matrix(Q)))
    end

    #Convert to truncated HOSVD of p + x
    hosvd_S = st_hosvd(S, r⃗)
    factors = UQ .* hosvd_S.U
    for i in 1:D
        q.hosvd.U[i] .= factors[i]
        q.hosvd.σ[i] .= hosvd_S.σ[i]
    end
    q.hosvd.core .= hosvd_S.core
    return q
end

function Base.show(io::IO, ::MIME"text/plain", 𝒯::Tucker{N,R,D,𝔽}) where {N,R,D,𝔽}
    return print(io, "Tucker(", N, ", ", R, ", ", 𝔽, ")")
end
function Base.show(io::IO, ::MIME"text/plain", 𝔄::TuckerPoint)
    pre = " "
    summary(io, 𝔄)
    for d in eachindex(𝔄.hosvd.U)
        println(io, string("\nU factor ", d, ":"))
        su = sprint(show, "text/plain", 𝔄.hosvd.U[d]; context=io, sizehint=0)
        su = replace(su, '\n' => "\n$(pre)")
        println(io, pre, su)
    end
    println(io, "\nCore:")
    su = sprint(show, "text/plain", 𝔄.hosvd.core; context=io, sizehint=0)
    su = replace(su, '\n' => "\n$(pre)")
    return print(io, pre, su)
end
function Base.show(io::IO, ::MIME"text/plain", x::TuckerTVector)
    pre = " "
    summary(io, x)
    for d in eachindex(x.U̇)
        println(io, string("\nU̇ factor ", d, ":"))
        su = sprint(show, "text/plain", x.U̇[d]; context=io, sizehint=0)
        su = replace(su, '\n' => "\n$(pre)")
        println(io, pre, su)
    end
    println(io, "\nĊ factor:")
    su = sprint(show, "text/plain", x.Ċ; context=io, sizehint=0)
    su = replace(su, '\n' => "\n$(pre)")
    return print(io, pre, su)
end
function Base.show(io::IO, ::MIME"text/plain", ℬ::CachedHOSVDBasis{𝔽,T,D}) where {𝔽,T,D}
    summary(io, ℬ)
    print(io, " ≅")
    su = sprint(show, "text/plain", convert(Matrix{T}, ℬ); context=io, sizehint=0)
    su = replace(su, '\n' => "\n ")
    return println(io, " ", su)
end

"""
    Base.size(p::TuckerPoint)

The dimensions of a [`TuckerPoint`](@ref) `p`, when regarded as a full tensor
(see [`embed`](@ref)).
"""
Base.size(𝔄::TuckerPoint) = map(u -> size(u, 1), 𝔄.hosvd.U)

#=
Modification of the ST-HOSVD from [Vannieuwenhoven2012]
This is the HOSVD of an approximation of 𝔄, i.e. the core of this decomposition
is also in HOSVD format.
=#
function st_hosvd(𝔄, mlrank=size(𝔄))
    T = eltype(𝔄)
    D = ndims(𝔄)
    n⃗ = size(𝔄)
    # Add type assertions to U and σ for type stability
    U::NTuple{D,Matrix{T}} = ntuple(d -> Matrix{T}(undef, n⃗[d], mlrank[d]), D)
    σ::NTuple{D,Vector{T}} = ntuple(d -> Vector{T}(undef, mlrank[d]), D)
    # Initialise arrays to store successive truncations (𝔄′) and unfoldings (buffer)
    # so that the type remains constant at every truncation
    𝔄′ = reshape(view(𝔄, 1:length(𝔄)), n⃗)
    fold_buffer = reshape(view(similar(𝔄), 1:length(𝔄)), n⃗)
    unfold_buffer = view(similar(𝔄), 1:length(𝔄))

    for k in 1:D
        rₖ = mlrank[k]
        𝔄′₍ₖ₎ = tensor_unfold!(unfold_buffer, 𝔄′, k)
        # truncated SVD + incremental construction of the core
        UΣVᵀ = svd(𝔄′₍ₖ₎)
        U[k] .= UΣVᵀ.U[:, 1:rₖ]
        σ[k] .= UΣVᵀ.S[1:rₖ]
        𝔄′₍ₖ₎_trunc = Diagonal(σ[k]) * UΣVᵀ.Vt[1:rₖ, :]
        size𝔄′ = ntuple(i -> i ≤ k ? mlrank[i] : n⃗[i], D)
        fold_buffer = reshape(view(fold_buffer, 1:prod(size𝔄′)), size𝔄′)
        unfold_buffer = view(unfold_buffer, 1:prod(size𝔄′))
        𝔄′ = tensor_fold!(fold_buffer, 𝔄′₍ₖ₎_trunc, k)
    end
    core = Array(𝔄′)

    # Make sure the truncated core is in "all-orthogonal" HOSVD format
    if mlrank ≠ n⃗
        hosvd_core = st_hosvd(core, mlrank)
        U = U .* hosvd_core.U
        core = hosvd_core.core
        σ = hosvd_core.σ
    end

    return HOSVD{T,D}(U, core, σ)
end

# In-place inverse of the k'th unfolding of a size n₁ × ... × n_D tensor.
# The size of the reshaped tensor is determined by the size of 𝔄.
# The result is stored in 𝔄. The returned value uses the same address space as 𝔄.
function tensor_fold!(𝔄::AbstractArray{T,D}, 𝔄₍ₖ₎::AbstractMatrix{T}, k) where {T,D}
    @assert length(𝔄₍ₖ₎) == length(𝔄) && size(𝔄₍ₖ₎, 1) == size(𝔄, k)
    @assert pointer(𝔄) !== pointer(𝔄₍ₖ₎)
    # Caution: tuple operations can be type unstable if used incorrectly
    σ(i) = i == 1 ? k : i ≤ k ? i - 1 : i
    σ⁻¹(i) = i < k ? i + 1 : i == k ? 1 : i
    permuted_size = ntuple(i -> size(𝔄, σ(i)), D)
    return permutedims!(𝔄, reshape(𝔄₍ₖ₎, permuted_size), ntuple(σ⁻¹, D))
end

# In-place mode-k unfolding of the array 𝔄 of order D ≥ k.
# The argument buffer is an array of arbitrary dimensions of the same length as 𝔄.
# The returned value uses the same address space as the buffer.
function tensor_unfold!(buffer, 𝔄::AbstractArray{T,D}, k) where {T,D}
    @assert length(buffer) == length(𝔄)
    @assert pointer(𝔄) !== pointer(buffer)
    𝔄₍ₖ₎ = reshape(buffer, size(𝔄, k), :)
    # Caution: tuple operations can be type unstable if used incorrectly
    σ(i) = i == 1 ? k : i ≤ k ? i - 1 : i
    permuted_size = ntuple(i -> size(𝔄, σ(i)), D)
    permutedims!(reshape(𝔄₍ₖ₎, permuted_size), 𝔄, ntuple(σ, D))
    return 𝔄₍ₖ₎
end

@doc raw"""
    zero_vector(::Tucker, p::TuckerPoint)

The zero element in the tangent space to `p` on the [`Tucker`](@ref) manifold, represented
as a [`TuckerTVector`](@ref).
"""
zero_vector(::Tucker, ::TuckerPoint)

function zero_vector!(::Tucker, X::TuckerTVector, ::TuckerPoint)
    for U̇ in X.U̇
        fill!(U̇, zero(eltype(U̇)))
    end
    fill!(X.Ċ, zero(eltype(X.Ċ)))
    return X
end

# The standard implementation of allocate_result on vector-valued functions gives an element
# of the same type as the manifold point. We want a vector instead.
for fun in [:get_vector, :inverse_retract, :project, :zero_vector]
    @eval function ManifoldsBase.allocate_result(
        ::Tucker,
        ::typeof($(fun)),
        p::TuckerPoint,
        args...,
    )
        return TuckerTVector(allocate(p.hosvd.core), allocate(p.hosvd.U))
    end
end

function ManifoldsBase.allocate_result(::Tucker{N}, f::typeof(embed), p, args...) where {N}
    dims = N
    return Array{number_eltype(p),length(dims)}(undef, dims)
end


text: abstract type Plotter end

ext: jl


text:
module MendelIHT

    import Distances: euclidean, chebyshev, sqeuclidean
    import Base.show
    import GLM: glmvar, mueta, fit, linkinv, Link, GeneralizedLinearModel, devresid, checky, canonicallink
    import SpecialFunctions: digamma, trigamma
    import Pkg
    import StatsBase: sample, aweights

    using GLM
    using SnpArrays
    using DataFrames
    using LinearAlgebra
    using Distributions
    using Random
    using DelimitedFiles
    using ProgressMeter
    using Reexport
    using ThreadPools
    using VCFTools
    using BGEN

    @reexport using Distributions
    @reexport using SnpArrays

    export iht, cross_validate
    export fit_iht, cv_iht, cv_iht_distribute_fold, iht_run_many_models
    export loglikelihood, deviance, score!, mle_for_r
    export project_k!, project_group_sparse!
    export IHTVariable, make_snparray, standardize!, maf_weights
    export save_prev!, naive_impute, initialize_beta
    export simulate_random_snparray, simulate_correlated_snparray
    export make_bim_fam_files, simulate_random_response, adhoc_add_correlation
    export random_covariance_matrix
    export pve

    # IHT will only work on single/double precision floats!
    const Float = Union{Float64,Float32}

    include("data_structures.jl")
    include("utilities.jl")
    include("simulate_utilities.jl")
    include("fit.jl")
    include("cross_validation.jl")
    include("multivariate.jl")
    include("wrapper.jl")
    include("pve.jl")

    # test data directory
    datadir(parts...) = joinpath(@__DIR__, "..", "data", parts...)

end # end module


text:
######################################
"""
    x1, x2, y1, y2 = getlims(tiles::VisClaw.AMRGrid)

get min/max range of a tile
"""
function getlims(tile::VisClaw.AMRGrid)
    return tile.xlow, tile.xlow+(tile.mx-1)*tile.dx, tile.ylow, tile.ylow+(tile.my-1)*tile.dy
end
######################################
"""
    x1, x2, y1, y2 = getlims(tiles::Vector{VisClaw.AMRGrid})

get min/max range of tiles
"""
function getlims(tiles::Vector{VisClaw.AMRGrid})
    x1 = minimum(getfield.(tiles, :xlow))
    y1 = minimum(getfield.(tiles, :ylow))
    x2 = maximum(round.(getfield.(tiles, :xlow) .+ getfield.(tiles, :mx).*getfield.(tiles, :dx), digits=4))
    y2 = maximum(round.(getfield.(tiles, :ylow) .+ getfield.(tiles, :my).*getfield.(tiles, :dy), digits=4))
    return x1, x2, y1, y2
end
######################################
"""
    xmesh, ymesh = meshtile(tile::VisClaw.AMRGrid)

generate meshgrids of tile
"""
function meshtile(tile::VisClaw.AMRGrid)
    ## set the boundary
    x = [tile.xlow, tile.xlow+tile.dx*tile.mx]
    y = [tile.ylow, tile.ylow+tile.dy*tile.my]
    ## grid info
    xline = collect(Float64, x[1]+0.5tile.dx:tile.dx:x[2]-0.5tile.dx+1e-4)
    yline = collect(Float64, y[1]+0.5tile.dy:tile.dy:y[2]-0.5tile.dy+1e-4)
    xmesh = repeat(xline', outer=(tile.my,1))
    ymesh = repeat(yline,  outer=(1,tile.mx))

    ## return values
    return xmesh, ymesh
end
######################################

"""
    var = keytile(tile::VisClaw.AMRGrid)

Get the main property name from VisClaw.AMRGrid
"""
function keytile(tile::VisClaw.AMRGrid)
    # check
    !isa(tile, VisClaw.AMRGrid) && error("Invalid input argument. It must be a type of VisClaw.AMRGrid")
    # assign
    varset = [:eta, :vel, :slp]
    ind = map(T -> isa(tile, T), [VisClaw.SurfaceHeight, VisClaw.Velocity, VisClaw.Storm])
    # return value
    return varset[ind][1]
end
##########################################################

##########################################################
"""
    xvec, yvec, val = tilezmargin(tile::VisClaw.AMRGrid, var::Symbol; digits=4)

Get Z-values of cells including their margins
"""
function tilezmargin(tile::VisClaw.AMRGrid, var::Symbol; digits=4)
    ## set the boundary
    x = [tile.xlow, round(tile.xlow+tile.dx*tile.mx, digits=digits)]
    y = [tile.ylow, round(tile.ylow+tile.dy*tile.my, digits=digits)]

    ## grid info
    xvec = collect(LinRange(x[1]-0.5tile.dx, x[2]+0.5tile.dx, tile.mx+2));
    yvec = collect(LinRange(y[1]-0.5tile.dy, y[2]+0.5tile.dy, tile.my+2));
    xvec = round.(xvec, digits=digits)
    yvec = round.(yvec, digits=digits)
    ## adjust data
    val = zeros(tile.my+2,tile.mx+2)
    val[2:end-1,2:end-1] = getfield(tile, var)
    val[2:end-1,1] = val[2:end-1,2]
    val[2:end-1,end] = val[2:end-1,end-1]
    val[1,:] = val[2,:]
    val[end,:] = val[end-1,:]

    # return val
    return xvec, yvec, val
end
##########################################################

##########################################################
"""
    xvec, yvec, val = tilez(tile::VisClaw.AMRGrid, var::Symbol; digits=4)

Get Z-values of cells at the grid lines
"""
function tilez(tile::VisClaw.AMRGrid, var::Symbol; digits=4)
    xvec, yvec, val = VisClaw.tilezmargin(tile, var, digits=digits)
    itp = Interpolations.interpolate((yvec, xvec), val, Interpolations.Gridded(Interpolations.Linear()))

    ## set the boundary
    x = [tile.xlow, round(tile.xlow+tile.dx*tile.mx, digits=digits)]
    y = [tile.ylow, round(tile.ylow+tile.dy*tile.my, digits=digits)]

    xvec = collect(LinRange(x[1], x[2], tile.mx+1));
    yvec = collect(LinRange(y[1], y[2], tile.my+1));
    xvec = round.(xvec, digits=digits)
    yvec = round.(yvec, digits=digits)

    val = itp(yvec,xvec);

    # return val
    return xvec, yvec, val
end
############################################################


text:
module StaggeredRK4
"""
Staggered Time Integrators for Wave Equations
Michelle Ghrist, Bengt Fornberg, and Tobin A. Driscoll
https://doi.org/10.1137/S0036142999351777
"""

function staggered_rk4(
                f::Function,
                g::Function,
               u::AbstractArray,
               v::AbstractArray,
               t::Float64, dt::Float64)

        uout = step(f, g, u, v, t, dt)
        vout = step(g, f, v, uout, t + 0.5 * dt, dt)

        return uout, vout

end

function step(
                f::Function,
                g::Function,
               u::AbstractArray,
               v::AbstractArray,
               t::Float64, dt::Float64)

        z1 = dt * g(u, t)
        z1 = v - z1
        z1 = dt * f(z1, t + 0.5 * dt - dt)
        z2 = dt * f(v, t + 0.5 * dt)
        z3 = u + z2
        z3 = dt * g(z3, t + dt)
        z3 = v + z3
        z3 = dt * f(z3, t + 0.5 * dt + dt)
        u = u + z1 / 24 + 11 / 12 * z2  + z3 / 24

        #z1 = dt * f(v - dt * g(u, t), t + 0.5 * dt - dt)
        #z2 = dt * f(v, t + 0.5 * dt)
        #z3 = dt * f(v + dt * g(u + z2, t + dt) , t + 0.5 * dt + dt)
        #u = u + z1 / 24 + 11 / 12 * z2  + z3 / 24

        return u

end

function staggered_rk4_linear(
                              Ap, Av, p, v, dt::Float64)
        p = step_linear(Ap, Av, p, v, dt)
        v = step_linear(Av, Ap, v, p, dt)
        return p, v

end

function step_linear(Ap, Av, p, v, dt)
        r1 = dt * Ap * v
        p = p + r1
        r2 = dt * Av * r1
        r1 = dt * Ap * r2
        p = p + r1 / 24
        return p
end

end

max_stars_repo_name: ooreilly/sbpjl

max_issues_repo_name: ooreilly/sbpjl

max_forks_repo_name: ooreilly/sbpjl


text:
export funLS

"""
f,g = funLS(r::AbstractArray; params::Dict)\n

Normalize r using the 2 norm. This function uses no params\n

Returns: f: 2-Norm of r\n
         g: Normalized r\n
"""
function funLS(r::Tr,params::Dict{String,Any}) where {ETr<:Real, Tr<:Union{AbstractArray{ETr}, AbstractArray{Complex{ETr}}}}

    f = norm(r, 2)
    g = r/f

    return f, g
end


text:
using Pipe

function parse_line(l)
    lhs, rhs = split(l, " -> ")
    x₁, y₁ = 1 .+ parse.(Int, split(lhs, ","))
    x₂, y₂ = 1 .+ parse.(Int, split(rhs, ","))
    @pipe ((x₁,x₂), (y₁,y₂)) .|> make_range(_...)
end

make_range(u, v) = u ≤ v ? (u:v) : (u:-1:v)

const data = parse_line.(readlines())

grid = zeros(Int, (1000, 1000))
for (r_x, r_y) in data
    (length(r_x) > 1) ⊻ (length(r_y) > 1) && (grid[r_x, r_y] .+= 1)

    # Comment for Part 1
    length(r_y) == length(r_x) && for (xx, yy) in zip(r_x, r_y)
        grid[xx, yy] += 1
    end
end

@pipe grid |> count(≥(2), _) |> println

max_stars_repo_name: pLOPeGG/JuliaOfCode

max_issues_repo_name: pLOPeGG/JuliaOfCode

max_forks_repo_name: pLOPeGG/JuliaOfCode


text:
module ModelBasedCF

# package code goes here
using Persa
using ProgressMeter
using Statistics
using LinearAlgebra: norm
using Random: shuffle

abstract type MatrixFactorization{T} <: Persa.Model{T}
end

include("irsvd.jl")
include("rsvd.jl")
include("train.jl")

include("baseline.jl")
include("random.jl")

end # module


text:
# COMBINE FUNCTIONS

"""
    combine(code1::SimpleCode, code2::SimpleCode) -> SimpleCode
    combine(code1::TensorNetworkCode, code2::TensorNetworkCode) -> TensorNetworkCode

Return a new code that is the tensor product of the given codes. Physically equivalent to
preparing two codes on different sets of physical qubits.

# Examples
```jldoctest
julia> code = combine(five_qubit_code(), steane_code());  # combine 5 and 7 qubit codes

julia> num_qubits(code)
12
```
"""
function combine(code1::SimpleCode, code2::SimpleCode)
    n1 = num_qubits(code1)
    n2 = num_qubits(code2)

    stabilizers = Vector{Int}[]
    logicals = Vector{Int}[]
    pure_errors = Vector{Int}[]

    for stabilizer in code1.stabilizers
        push!(stabilizers, vcat(stabilizer, zeros(Int, n2)))
    end
    for stabilizer in code2.stabilizers
        push!(stabilizers, vcat(zeros(Int, n1), stabilizer))
    end

    for logical in code1.logicals
        push!(logicals, vcat(logical, zeros(Int, n2)))
    end
    for logical in code2.logicals
        push!(logicals, vcat(zeros(Int, n1), logical))
    end

    for pure_error in code1.pure_errors
        push!(pure_errors, vcat(pure_error, zeros(Int, n2)))
    end
    for pure_error in code2.pure_errors
        push!(pure_errors, vcat(zeros(Int, n1), pure_error))
    end

    return SimpleCode("", stabilizers, logicals, pure_errors)
end
function combine(code1::TensorNetworkCode, code2::TensorNetworkCode)
    # code2 = (code1 === code2) ? deepcopy(code2) : code2  # copy if codes identical

    # copy & change indices of one of the codes.  necessary to avoid
    # duplicate ITensor indices
    if num_qubits(code1) < num_qubits(code2)
        code1 = new_indices(code1)
    else
        code2 = new_indices(code2)
    end

    new_code = combine(SimpleCode(code1), SimpleCode(code2))
    new_code_graph = _combine(code1.code_graph, code2.code_graph)
    new_seed_codes = merge(code1.seed_codes, code2.seed_codes)

    return TensorNetworkCode(new_code, new_code_graph, new_seed_codes)
end

"""
    _combine(graph1,graph2)

Combines two `CodeGraphs` with physical and virtual nodes and
preserves their metadata (`type`, `indices`, `coords` and 'qubit'
labels).
"""
function _combine(code_graph1::CodeGraph, code_graph2::CodeGraph)
    n1 = num_nodes(code_graph1)
    Labels = nodes(code_graph1)
    n_qubits1 = length(filter(x -> x > 0, Labels))

    new_code_graph2 = _shift_keys(code_graph2, n_qubits1, n1)

    coords = merge(code_graph1.coords, new_code_graph2.coords)
    node_types = merge(code_graph1.node_types, new_code_graph2.node_types)
    edge_types = merge(code_graph1.edge_types, new_code_graph2.edge_types)
    node_indices = merge(code_graph1.node_indices, new_code_graph2.node_indices)
    edge_indices = merge(code_graph1.edge_indices, new_code_graph2.edge_indices)

    return CodeGraph(coords, node_types, edge_types, node_indices, edge_indices)
end

"""
    _shift_keys(code_graph,n_qubits,n_vert)

When merging two `Graphs`, the second `Graph` has all its edges
relabelled.  So to merge two `CodeGraphs`, it is necessary to preemptively
shift the keys of all the dictionaries in the second `CodeGraph` to
allow for this.

`code_graph` is the second of the two `CodeGraphs` to be merged.  `n_vert`
and `n` are the number of vertices and qubits respectively of the first
`CodeGraph` to be merged.
"""
function _shift_keys(code_graph::CodeGraph, n_qubits::Int64, n_vert::Int64)

    new_coords = Dict{Int64,Vector{Float64}}()
    new_node_types = Dict{Int64,String}()
    new_edge_types = Dict{Set{Int64},String}()
    new_node_indices = Dict{Int64,Vector{Index{Int64}}}()
    new_edge_indices = Dict{Set{Int64},Vector{Index{Int64}}}()

    n_virtual = n_vert - n_qubits

    # Update node data
    for (key, value) in coords(code_graph)
        if key > 0
            new_coords[key + n_qubits] = value
        else
            new_coords[key - n_virtual] = value
        end
    end
    for (key, value) in node_types(code_graph)
        if key > 0
            new_node_types[key + n_qubits] = value
        else
            new_node_types[key - n_virtual] = value
        end
    end
    for (key, value) in node_indices(code_graph)
        if key < 0
            new_node_indices[key - n_virtual] = value
        end
    end

    # Edge data is a bit trickier
    for old_edge in edges(code_graph)
        new_edge = old_edge .+ 0 # converts to vector

        for α in 1:2
            if new_edge[α] > 0
                new_edge[α] = new_edge[α] + n_qubits
            else
                new_edge[α] = new_edge[α] - n_virtual
            end
        end

        new_edge = Set(new_edge)

        new_edge_types[new_edge] = edge_types(code_graph, old_edge)
        new_edge_indices[new_edge] = edge_indices(code_graph, old_edge)
    end

    return CodeGraph(
        new_coords,
        new_node_types,
        new_edge_types,
        new_node_indices,
        new_edge_indices
    )
end

# FUSION FUNCTIONS

"""
    fusion(code::SimpleCode, qubit_pair::AbstractVector{Int}) -> SimpleCode
    fusion(code::SimpleCode, qubit_pairs) -> SimpleCode

    fusion(code::TensorNetworkCode, qubit_pair::AbstractVector{Int}) -> TensorNetworkCode
    fusion(code::TensorNetworkCode, qubit_pairs) -> TensorNetworkCode

Return a new code that results from fusing the physical qubits with labels given by each
qubit pair. The new code has two fewer physical qubits, for each qubit pair, but the same
number of logical qubits. Physically equivalent to updating stabilizers, logicals and pure
errors after measuring ``XX`` and ``ZZ`` on each pair of qubits.

The versions that take `qubit_pairs` take iterables of `AbstractVector{Int}`. An
`ErrorException` is thrown if the fusion is not possible, i.e., if the logical
degrees of freedom would not be preserved by the measurements.

# Examples
```jldoctest
julia> code = steane_code();

julia> num_qubits(code), length(code.logicals) ÷ 2
(7, 1)

julia> fused_code = fusion(code, [1, 2]);

julia> num_qubits(fused_code), length(fused_code.logicals) ÷ 2
(5, 1)
```
"""
function fusion(code::SimpleCode, qubit_pair::AbstractVector{Int})
    stabilizers = deepcopy(code.stabilizers)
    pure_errors = deepcopy(code.pure_errors)
    logicals = deepcopy(code.logicals)

    annoying_ops = _annoying_stabilizers(stabilizers, qubit_pair)
    if length(annoying_ops) != 0
        @debug "There were annoying operators!"
    end
    for logical in logicals
        if !pauli_are_commuting(vcat(annoying_ops, [logical]))
            error("Fusion is not possible between selected qubit pair!")
        end
    end

    useful_inds = _find_useful_stabilizer_indices(stabilizers, qubit_pair)
    if length(useful_inds) != 0
        sort!(useful_inds)
        useful_ops = [stabilizers[ind] for ind in useful_inds]
        deleteat!(stabilizers, sort!(useful_inds))
        deleteat!(pure_errors, sort!(useful_inds))

        stabilizers = _make_ready(stabilizers, useful_ops, qubit_pair)
        logicals = _make_ready(logicals, useful_ops, qubit_pair)
        pure_errors = _make_ready(pure_errors, useful_ops, qubit_pair)
    end

    _remove_qubits!(stabilizers, qubit_pair)
    if length(logicals) != 0
        _remove_qubits!(logicals, qubit_pair)
    end

    if length(useful_inds) != 2
        new_stabilizers = Array{Int64,1}[]
        for α in 1:length(stabilizers)
            if pauli_are_independent(vcat(new_stabilizers, [stabilizers[α]]))
                push!(new_stabilizers, stabilizers[α])
            end
        end
        stabilizers = new_stabilizers
    end

    _remove_qubits!(pure_errors, qubit_pair)

    return SimpleCode(" ", stabilizers, logicals, pure_errors)
end
function fusion(code::SimpleCode, qubit_pairs)
    output_code = deepcopy(code)
    qubit_pair_list = collect(qubit_pairs) # convert interable to list
    # qubit labels change after each fusion, so must account for this:
    _update_qubit_pairs!(qubit_pair_list)

    for qubit_pair in qubit_pair_list
        if num_qubits(output_code)  == 0
            return SimpleCode()
        end
        output_code = fusion(output_code, qubit_pair)
    end

    return output_code
end
function fusion(code::TensorNetworkCode, qubit_pairs)
    new_code = fusion(SimpleCode(code), qubit_pairs)
    new_code_graph = _fusion(code.code_graph, qubit_pairs)
    return TensorNetworkCode(new_code, new_code_graph, code.seed_codes)
end
function fusion(code::TensorNetworkCode, qubit_pair::AbstractVector{Int})
    return fusion(code, [qubit_pair])
end

"""
    _fusion(code_graph,qubit_pair)

Returns a `CodeGraph` after fusion is performed on the qubits in
`qubit_pair`.
"""
function _fusion(code_graph::CodeGraph, qubit_pair::AbstractVector{Int})
    new_edge_types = deepcopy(code_graph.edge_types)
    new_node_indices = deepcopy(code_graph.node_indices)
    new_edge_indices = deepcopy(code_graph.edge_indices)

    # Find new edge
    edges_to_remove = edges(code_graph)
    edges_to_remove =
    filter(x -> length(intersect(qubit_pair, x)) != 0, edges_to_remove)

    virtual1 = setdiff(edges_to_remove[1], qubit_pair)
    virtual2 = setdiff(edges_to_remove[2], qubit_pair)
    new_edge = union(virtual1, virtual2)

    if length(new_edge) == 1
        println("Self contraction occurred.  Contraction algorithms may not work!")
        @goto skip_new_edge_stuff
    end

    # Add edge data
    new_edge_types[new_edge] = "bond"
    new_index = Index(4, "bond")
    new_edge_indices[new_edge] = [new_index]

    # replace node indices with new bond index
    indices_to_replace = [new_edge_indices[edge][1] for edge in edges_to_remove]
    for label in new_edge # this is cool: new_edge is a Set
        index_list = new_node_indices[label]
        for α in 1:length(index_list)
            if index_list[α] in indices_to_replace
                index_list[α] = new_index
            end
        end
    end

    @label skip_new_edge_stuff

    output = CodeGraph(
        code_graph.coords,
        code_graph.node_types,
        new_edge_types,
        new_node_indices,
        new_edge_indices
    )

    sort!(qubit_pair)
    qubit_pair[2] = qubit_pair[2] - 1
    for qubit in qubit_pair
        output = _remove_node(output, qubit)
    end

    return output
end
function _fusion(code_graph::CodeGraph, qubit_pairs::Array{Array{Int64,1},1})
    output_graph = code_graph
    for qubit_pair in qubit_pairs
        output_graph = _fusion(output_graph, qubit_pair)
    end

    return output_graph
end

function _annoying_stabilizers(
    stabilizers::AbstractVector{<:AbstractVector{Int}},
    qubit_pair::AbstractVector{Int}
)
    i, j = qubit_pair
    n = length(stabilizers[1])
    output = Array{Array{Int64,1},1}()

    for σ in [1,2,3]

        for α in 1:length(stabilizers)
            if !pauli_are_commuting([[σ,σ],[stabilizers[α][i],stabilizers[α][j]]])
                @goto not_this_one
            end
        end

        operator = zeros(Int64, n)
        operator[i] = σ
        operator[j] = σ
        push!(output, operator)

        if length(output) == 2
            break
        end

        @label not_this_one
    end

    return output
end

"""
    _find_useful_stabilizer_indices(stabilizers,qubit_pair)

Finds two `stabilizers` that can be used to make other operators
ready for fusion.  Basically it returns two operators that are independent
from each other and XX, YY, ZZ on the qubits in `qubit_pair`.
"""
function _find_useful_stabilizer_indices(
    stabilizers::AbstractVector{<:AbstractVector{Int}},
    qubit_pair::AbstractVector{Int}
)
    indices = Int64[]

    for α in 1:length(stabilizers)
        operator = stabilizers[α]
        if operator[qubit_pair[1]] == operator[qubit_pair[2]]
            continue
        end

        if length(indices) == 0
            push!(indices, α)
            continue
        end

        if length(indices) == 1
            useful_operator = stabilizers[indices[1]]
            new_operator = pauli_product.(operator, useful_operator)

            if new_operator[qubit_pair[1]] != new_operator[qubit_pair[2]]
                push!(indices, α)
                return indices
            end
        end
    end

    return indices
end

"""
    _make_ready(operator,useful_ops,qubit_pair)

Makes an `operator` ready for fusion on `qubit_pair`, i.e.,
multiplies a combination of operators in `useful_ops` until the
resulting operator has the same Paulis on both qubits in `qubit_pair`.
"""
function _make_ready(
    operator::AbstractVector{Int},
    useful_ops::AbstractVector{<:AbstractVector{Int}},
    qubit_pair::AbstractVector{Int}
)
    if _is_ready(operator, qubit_pair)
        return operator
    end

    for useful_op in useful_ops
        new_operator = pauli_product.(operator, useful_op)
        if _is_ready(new_operator, qubit_pair)
            return new_operator
        end
    end

    if length(useful_ops) == 2
        operator2 = pauli_product.(useful_ops[1], useful_ops[2])
        new_operator = pauli_product.(operator, operator2)
        if _is_ready(new_operator, qubit_pair)
            return new_operator
        end
    end
end
_make_ready(
    operators::AbstractVector{<:AbstractVector{Int}},
    useful_ops::Array{Array{Int64,1}},
    qubit_pair::AbstractVector{Int}
) = _make_ready.(operators, Ref(useful_ops), Ref(qubit_pair))

"""
    _is_ready(operator,qubit_pair)

Checks if `operator` is ready for fusion on `qubit_pair`, i.e.,
does `operator` have the same Paulis on both qubits.
"""
function _is_ready(operator::AbstractVector{Int}, qubit_pair::AbstractVector{Int})
    return (operator[qubit_pair[1]] == operator[qubit_pair[2]])
end

"""
    _remove_qubits!(operators, qubits_to_remove)

Given a set of operators, remove entries corresponding to qubits being removed.
"""
function _remove_qubits!(
    operators::AbstractVector{<:AbstractVector{Int}},
    qubits_to_remove::AbstractVector{Int}
)
    for α in 1:length(operators)
        deleteat!(operators[α], sort!(qubits_to_remove))
    end
end

function _update_qubit_pairs!(qubit_pairs::AbstractVector{<:AbstractVector{Int}})
    num_pairs = length(qubit_pairs)
    for α in 1:num_pairs
        for x in 1:2
            for β in 1:α - 1
                if qubit_pairs[β][1] <= qubit_pairs[α][x]
                    qubit_pairs[α][x] -= 1
                end
                if qubit_pairs[β][2] <= qubit_pairs[α][x]
                    qubit_pairs[α][x] -= 1
                end
            end
        end
    end
end

"""
    _remove_node(code_graph,label)

Removes a node from a `CodeGraph`; returns a new `CodeGraph`.
"""
function _remove_node(code_graph::CodeGraph, label::Int)

    # Remove graph vertex and update values in `nodes`
#     # A little annoying because of how rem_vertex! works
#     # in LightGraphs.jl
#     g = code_graph.graph
#     end_label = nodes_to_labels(code_graph)[end]
#     removed_node = nodes(code_graph,label)

#     new_g = deepcopy(g)
#     rem_vertex!(new_g,removed_node)

    # Update dictionaries with nodes as keys
    old_labels = nodes(code_graph)
    new_labels = Dict{Int64,Int64}()
    for old_label in old_labels
        if old_label < label
            new_labels[old_label] = old_label
        elseif old_label > label
           new_labels[old_label] = old_label - 1
        end
    end

    new_coords = _update_keys(code_graph.coords, new_labels)
    new_node_types = _update_keys(code_graph.node_types, new_labels)
    filter!(x -> x.first < 0, new_labels)
    new_node_indices = _update_keys(code_graph.node_indices, new_labels)

    # Update dictionaries with edges as keys
    old_edges = edges(code_graph)
    new_edges = Dict{Set{Int64},Set{Int64}}()
    for old_edge in old_edges

        new_edge = old_edge .+ 0 # converts to vector

        if label in new_edge
            continue
        end
        for α in 1:2
            if new_edge[α] > label
                new_edge[α] = new_edge[α] - 1
            end
        end

        new_edge = Set(new_edge)
        new_edges[old_edge] = new_edge
    end

    new_edge_types = _update_keys(code_graph.edge_types, new_edges)
    new_edge_indices = _update_keys(code_graph.edge_indices, new_edges)

    return CodeGraph(
        new_coords,
        new_node_types,
        new_edge_types,
        new_node_indices,
        new_edge_indices
    )
end

# maybe don't bother including types
function _update_keys(
    dict::Dict{Int,T}, new_keys::Dict{Int,Int}
) where T <: Union{Int,Vector{Float64},String,Vector{Index{Int}}}

    value_type = typeof(dict[findfirst(x -> true, dict)])  # this is lame
    new_dict = Dict{Int64,value_type}()

    for old_key in keys(new_keys)
        new_key = new_keys[old_key]
        new_dict[new_key] = dict[old_key]
    end

    return new_dict
end
function _update_keys(
    dict::Dict{Set{Int},T}, new_keys::Dict{Set{Int},Set{Int}}
) where T <: Union{Int,Vector{Float64},String,Vector{Index{Int}}}

    value_type = typeof(dict[findfirst(x -> true, dict)])  # this is lame
    new_dict = Dict{Set{Int64},value_type}()

    for old_key in keys(new_keys)
        new_key = new_keys[old_key]
        new_dict[new_key] = dict[old_key]
    end

    return new_dict
end

# CONTRACT FUNCTIONS

"""
    contract(
        code1::TensorNetworkCode,
        code2::TensorNetworkCode,
        qubit_pair::AbstractVector{Int}
    ) -> TensorNetworkCode

    contract(
        code1::TensorNetworkCode,
        code2::TensorNetworkCode,
        qubit_pairs
    ) -> TensorNetworkCode

Return a new code that results from combining the codes and fusing physical qubits
identified by the qubit pairs. The first and second elements of a qubit pair refer to qubit
labels from the first and second codes, respectively.

This is equivalent to [`combine`](@ref) followed by [`fusion`](@ref), with the qubit pair
labels referring to the code qubit labels before combining. The version that takes
`qubit_pairs` takes iterables of `AbstractVector{Int}`. An `ErrorException` is thrown if the
fusion is not possible.

See also: [`combine`](@ref), [`fusion`](@ref).

# Examples
```jldoctest
julia> code1 = TensorNetworkCode(five_qubit_code());

julia> code2 = TensorNetworkCode(steane_code());

julia> contracted_code = contract(code1, code2, [[1, 2], [2, 7]]);

julia> num_qubits(contracted_code), length(contracted_code.logicals) ÷ 2
(8, 2)

julia> fusion(combine(code1, code2), [[1, 7], [2, 12]]);  # equivalent code
```
"""
function contract(code1::TensorNetworkCode, code2::TensorNetworkCode, qubit_pairs)
    code2 = (code1 === code2) ? deepcopy(code2) : code2  # copy if codes identical

    n = num_qubits(code1)
    qubit_pair_list = [[pair[1], pair[2] + n] for pair in qubit_pairs]

    output_code = combine(code1, code2)
    output_code = fusion(output_code, qubit_pair_list)
    return output_code
end
function contract(
    code1::TensorNetworkCode,
    code2::TensorNetworkCode,
    qubit_pair::AbstractVector{Int}
)
    return contract(code1, code2, [qubit_pair])
end

"""
    contract_by_coords(code1::TensorNetworkCode,code2::TensorNetworkCode)
        -> TensorNetworkCode

Return a new code that results from combining the codes and fusing physical qubits with
coincident coordinates.

This is equivalent to [`combine`](@ref) followed by [`fusion`](@ref), with the fusion qubit
pairs being those with coincident coordinates. An `ErrorException` is thrown if the fusion
is not possible.

See also: [`combine`](@ref), [`fusion`](@ref).

# Examples
```jldoctest
julia> code1 = TensorNetworkCode(five_qubit_code());

julia> code2 = TensorNetworkCode(steane_code());

julia> set_coords!(code1, 1, [5, 5]); set_coords!(code2, 2, [5, 5]); # qubits 1 & 2 coincide

julia> set_coords!(code1, 2, [6, 6]); set_coords!(code2, 7, [6, 6]); # qubits 2 & 7 coincide

julia> contracted_code = contract_by_coords(code1, code2);

julia> num_qubits(contracted_code), length(contracted_code.logicals) ÷ 2
(8, 2)

julia> contract(code1, code2, [[1, 2], [2, 7]]);  # equivalent code
```
"""
function contract_by_coords(code1::TensorNetworkCode, code2::TensorNetworkCode)
    graph1 = code1.code_graph
    graph2 = code2.code_graph
    qubit_pairs = Vector{Int}[]

    for label1 in nodes(graph1), label2 in nodes(graph2)
        if (graph1.node_types[label1] != "physical" ||
            graph2.node_types[label2] != "physical")
            continue
        end

        coord1 = coords(graph1, label1)
        coord2 = coords(graph2, label2)
        if coord1 ≈ coord2
            push!(qubit_pairs, [label1,label2])
        end
    end

    if length(qubit_pairs) == 0
        return combine(code1, code2)
    end

    return contract(code1, code2, qubit_pairs)
end


text:
using CLIMA.Mesh.Elements
using GaussQuadrature
using Test

@testset "GaussQuadrature" begin
  for T ∈ (Float32, Float64, BigFloat)
    let
      x, w = GaussQuadrature.legendre(T, 1)
      @test iszero(x)
      @test w ≈ [2one(T)]
    end

    let
      endpt = GaussQuadrature.left
      x, w = GaussQuadrature.legendre(T, 1, endpt)
      @test x ≈ [-one(T)]
      @test w ≈ [2one(T)]
    end

    let
      endpt = GaussQuadrature.right
      x, w = GaussQuadrature.legendre(T, 1, endpt)
      @test x ≈ [one(T)]
      @test w ≈ [2one(T)]
    end

    let
      endpt = GaussQuadrature.left
      x, w = GaussQuadrature.legendre(T, 2, endpt)
      @test x ≈ [-one(T); T(1//3)]
      @test w ≈ [T(1//2); T(3//2)]
    end

    let
      endpt = GaussQuadrature.right
      x, w = GaussQuadrature.legendre(T, 2, endpt)
      @test x ≈ [T(-1//3); one(T)]
      @test w ≈ [T(3//2); T(1//2)]
    end
  end

  let
    err = ErrorException("Must have at least two points for both ends.")
    endpt = GaussQuadrature.both
    @test_throws err GaussQuadrature.legendre(1, endpt)
  end

  let
    T = Float64
    n = 100
    endpt = GaussQuadrature.both

    a, b = GaussQuadrature.legendre_coefs(T, n)

    err = ErrorException("No convergence after 1 iterations " *
                         "(try increasing maxits)")

    @test_throws err GaussQuadrature.custom_gauss_rule(-one(T), one(T),
                                                              a, b, endpt, 1)
  end
end

@testset "Operators" begin
  P5(r::AbstractVector{T}) where T = T(1)/T(8) * (T(15) * r - T(70) * r.^3 +
                                                  T(63) * r.^5)

  P6(r::AbstractVector{T}) where T =
    T(1)/T(16) * (-T(5) .+ T(105) * r.^2 - T(315) * r.^4 + T(231) * r.^6)
  DP6(r::AbstractVector{T}) where T =
    T(1)/T(16) * (T(2 * 105) * r - T(4 * 315) * r.^3 + T(6 * 231) * r.^5)

  IPN(::Type{T}, N) where T = T(2)/ T(2 * N + 1)

  N = 6
  for test_type ∈ (Float32, Float64, BigFloat)
    r, w = Elements.lglpoints(test_type, N)
    D = Elements.spectralderivative(r)
    x = LinRange{test_type}(-1, 1, 101)
    I = Elements.interpolationmatrix(r, x)

    @test sum(P5(r).^2 .* w) ≈ IPN(test_type, 5)
    @test D * P6(r) ≈ DP6(r)
    @test I * P6(r) ≈ P6(x)
  end

  for test_type ∈ (Float32, Float64, BigFloat)
    r, w = Elements.lgpoints(test_type, N)
    D = Elements.spectralderivative(r)

    @test sum(P5(r).^2 .* w) ≈ IPN(test_type, 5)
    @test sum(P6(r).^2 .* w) ≈ IPN(test_type, 6)
    @test D * P6(r) ≈ DP6(r)
  end
end

max_stars_repo_name: akshaysridhar/CLIMA

max_issues_repo_name: akshaysridhar/CLIMA

max_forks_repo_name: akshaysridhar/CLIMA


text:
robotsEndpoint = "api/$(curApiVersion)/users/{1}/robots"
robotEndpoint = "api/$(curApiVersion)/users/{1}/robots/{2}"

"""
$(SIGNATURES)
Gets all robots managed by the specified user.
Return: A vector of robots for a given user.
"""
function getRobots()::Vector{RobotResponse}
    config = getGraffConfig()
    if config == nothing
        error("Graff config is not set, please call setGraffConfig with a valid configuration.")
    end

    url = "$(config.apiEndpoint)/$(format(robotsEndpoint, config.userId))"
    response = @mock _sendRestRequest(config, HTTP.get, url)
    if(response.status != 200)
        error("Error getting robots, received $(response.status) with body '$(String(response.body))'.")
    end
    # Some manual effort done here because it's a vector response.
    robots = JSON2.read(String(response.body), Vector{RobotResponse})
    return robots
end

"""
$(SIGNATURES)
Return: Returns true if the robot exists already.
"""
function isRobotExisting(robotId::String)::Bool
    config = getGraffConfig()
    if config == nothing
        error("Graff config is not set, please call setGraffConfig with a valid configuration.")
    end

    url = "$(config.apiEndpoint)/$(format(robotEndpoint, config.userId, robotId))/exists"
    response = @mock _sendRestRequest(config, HTTP.get, url)
    body = String(response.body)
    if(response.status != 200)
        error("Error getting robot existence, received $(response.status) with body '$body'.")
    end
    return lowercase(body) == "true"
end

"""
$(SIGNATURES)
Return: Returns true if the robot in the config exists already.
"""
function isRobotExisting()::Bool
    config = getGraffConfig()
    if config == nothing
        error("Graff config is not set, please call setGraffConfig with a valid configuration.")
    end

    if config.robotId == ""
        error("Your config doesn't have a robot specified, please attach your config to a valid robot by setting the robotId field. Robot = $(config.robotId)")
    end

    return isRobotExisting(config.robotId)
end

"""
$(SIGNATURES)
Get a specific robot given a user ID and robot ID. Will retrieve config.robotId by default.
Return: The robot for the provided user ID and robot ID.
"""
function getRobot(robotId::String)::RobotResponse
    config = getGraffConfig()
    if config == nothing
        error("Graff config is not set, please call setGraffConfig with a valid configuration.")
    end

    url = "$(config.apiEndpoint)/$(format(robotEndpoint, config.userId, robotId))"
    response = @mock _sendRestRequest(config, HTTP.get, url)
    body = String(response.body)
    if(response.status != 200)
        error("Error getting robot, received $(response.status) with body '$body'.")
    end
    return JSON2.read(body, RobotResponse)
end

"""
$(SIGNATURES)
Get a specific robot given a user ID and robot ID. Will retrieve config.robotId by default.
Return: The robot for the provided user ID and robot ID.
"""
function getRobot()::RobotResponse
    config = getGraffConfig()
    if config == nothing
        error("Graff config is not set, please call setGraffConfig with a valid configuration.")
    end

    if config.robotId == ""
        error("Your config doesn't have a robot specified, please attach your config to a valid robot by setting the robotId field. Robot = $(config.robotId)")
    end

    return getRobot(config.robotId)
end

"""
$(SIGNATURES)
Create a robot in Synchrony and associate it with the given user.
Return: Returns the created robot.
"""
function addRobot(robot::RobotRequest)::RobotResponse
    config = getGraffConfig()
    if config == nothing
        error("Graff config is not set, please call setGraffConfig with a valid configuration.")
    end
    url = "$(config.apiEndpoint)/$(format(robotEndpoint, config.userId, robot.id))"
    response = @mock _sendRestRequest(config, HTTP.post, url, data=JSON.json(robot), headers=Dict{String, String}("Content-Type" => "application/json"))
    if(response.status != 200)
        error("Error creating robot, received $(response.status) with body '$(String(response.body))'.")
    end
    body = String(response.body)
    return JSON2.read(body, RobotResponse)
end

"""
$(SIGNATURES)
Update a robot.
Return: The updated robot from the service.
"""
function updateRobot(robot::RobotRequest)::RobotResponse
    config = getGraffConfig()
    if config == nothing
        error("Graff config is not set, please call setGraffConfig with a valid configuration.")
    end
    url = "$(config.apiEndpoint)/$(format(robotEndpoint, config.userId, robot.id))"
    response = @mock _sendRestRequest(config, HTTP.put, url, data=JSON.json(robot), headers=Dict{String, String}("Content-Type" => "application/json"))
    if(response.status != 200)
        error("Error updating robot, received $(response.status) with body '$(String(response.body))'.")
    end
    body = String(response.body)
    return JSON2.read(body, RobotResponse)
end

"""
$(SIGNATURES)
Delete a robot given a robot ID.
Return: The deleted robot.
"""
function deleteRobot(robotId::String)::RobotResponse
    config = getGraffConfig()
    if config == nothing
        error("Graff config is not set, please call setGraffConfig with a valid configuration.")
    end
    url = "$(config.apiEndpoint)/$(format(robotEndpoint, config.userId, robotId))"
    response = @mock _sendRestRequest(config, HTTP.delete, url)
    if(response.status != 200)
        error("Error deleting robot, received $(response.status) with body '$(String(response.body))'.")
    end
    body = String(response.body)
    return JSON2.read(body, RobotResponse)
end

"""
$(SIGNATURES)
Will retrieve the robot configuration (user settings) for the given robot ID.
Return: The robot config for the provided user ID and robot ID.
"""
function getRobotConfig(robotId::String)::String
    config = getGraffConfig()
    if config == nothing
        error("Graff config is not set, please call setGraffConfig with a valid configuration.")
    end
    url = "$(config.apiEndpoint)/$(format(robotEndpoint, config.userId, robotId))/config"
    response = @mock _sendRestRequest(config, HTTP.get, url)
    if(response.status != 200 && response.status != 204)
        error("Error getting robot, received $(response.status) with body '$(String(response.body))'.")
    end
    return String(response.body)
end

"""
$(SIGNATURES)
Will retrieve the robot configuration (user settings) for the default robot ID.
Return: The robot config for the provided user ID and robot ID.
"""
function getRobotConfig()::String
    config = getGraffConfig()
    if config == nothing
        error("Graff config is not set, please call setGraffConfig with a valid configuration.")
    end

    if config.robotId == ""
        error("Your config doesn't have a robot specified, please attach your config to a valid robot by setting the robotId field. Robot = $(config.robotId)")
    end

    return getRobotConfig(config.robotId)
end

"""
$(SIGNATURES)
Update a robot configuration.
Return: The updated robot configuration from the service.
"""
function updateRobotConfig(robotId::String, robotConfig::String)::String
    config = getGraffConfig()
    if config == nothing
        error("Graff config is not set, please call setGraffConfig with a valid configuration.")
    end
    url = "$(config.apiEndpoint)/$(format(robotEndpoint, config.userId, robotId))/config"
    response = @mock _sendRestRequest(config, HTTP.put, url, data=JSON.json(robotConfig))
    if(response.status != 200)
        error("Error updating robot config, received $(response.status) with body '$(String(response.body))'.")
    end
    return String(response.body)
end

"""
$(SIGNATURES)
Update a robot configuration.
Return: The updated robot configuration from the service.
"""
function updateRobotConfig(robotConfig::String)::String
    config = getGraffConfig()
    if config == nothing
        error("Graff config is not set, please call setGraffConfig with a valid configuration.")
    end
    if config.robotId == ""
        error("Your config doesn't have a robot specified, please attach your config to a valid robot by setting the robotId field. Robot = $(config.robotId)")
    end

    return updateRobotConfig(config.robotId, robotConfig)
end


text:
using JLD
include("Transformation.jl")
include("AbstractSystem.jl")
include("Tree.jl")
include("Evaluation.jl")
include("BackPropogation.jl")
include("Facility.jl")
import Base.push!

type Axiom
    tree1 :: Tree
    tree2 :: Tree
    index :: Dict
end

function Base.show(io :: IO, m :: Axiom)
    print(io, "Axiom[")
    print(io, m.tree1)
    print(io, ", ")
    print(io, m.tree2)
    print(io, "]")
end

function beautify(m :: Axiom)
    string("Axiom:\n", beautify(m.tree1, 1), "\n", beautify(m.tree2, 1), "\n")
end

Axiom(tree1 :: Tree, tree2 :: Tree) = Axiom(tree1, tree2, add!(index(tree1), index(tree2)))

Axiom(skeleton1, skeleton2) = Axiom(toTree(skeleton1), toTree(skeleton2))

function push!(index :: Dict, ops :: Array)
    n = length(ops)
    for i in 1:n
        for t in index[i]
            t.op = ops[i]
        end
    end
end

function push!(index :: Dict, ops :: Dict)
    for key in keys(ops)
        if haskey(index, key)
            ts = index[key]
            for t in ts
                t.op = ops[key]
            end
        end
    end
end

function push!(axiom :: Axiom, ops)
    push!(axiom.index, variables)
end

function init_axiom!(axiom :: Axiom, variables)
    push!(axiom.index, variables)
    init_tree!(axiom.tree1)
    init_tree!(axiom.tree2)
end

init_axioms! = distribute(init_axiom!)

## add some loss function to deal with degenerating problem?

function loss(a, b)
    b * (1. - a * b)
end

function train_axiom!(axiom :: Axiom, variables, n = 1, randomize = identity)
    push!(axiom.index, variables)
    d1 :: Array{Float64, 1} = axiom.tree1.value[:d]
    d2 :: Array{Float64, 1} = axiom.tree2.value[:d]
    v1 :: Array{Float64, 1} = axiom.tree1.value[:value]
    v2 :: Array{Float64, 1} = axiom.tree2.value[:value]
    for i in 1:n
        randomize(variables)
        push!(axiom.index, variables)
        eval_tree!(axiom.tree2)
        eval_tree!(axiom.tree1)
        for j in 1:length(d1)
            d1[j] = loss(v1[j], v2[j])
            d2[j] = loss(v2[j], v1[j])
        end
        bp_tree!(axiom.tree1)
        bp_tree!(axiom.tree2)
    end
end

train_axioms! = distribute(train_axiom!)

## to prevent degeneration problem, we use anti-traing to deal with the problem.

function anti_train_axiom!(axiom :: Axiom, n = 1, randomize = identity)
    d1 :: Array{Float64, 1} = axiom.tree1.value[:d]
    d2 :: Array{Float64, 1} = axiom.tree2.value[:d]
    v1 :: Array{Float64, 1} = axiom.tree1.value[:value]
    v2 :: Array{Float64, 1} = axiom.tree2.value[:value]
    for i in 1:n
        randomize(axiom.tree1)
        randomize(axiom.tree2)
        eval_tree!(axiom.tree2)
        eval_tree!(axiom.tree1)
        for j in 1:length(d1)
            d1[j] = - loss(v1[j], v2[j])
            d2[j] = - loss(v2[j], v1[j])
        end
        bp_tree!(axiom.tree1)
        bp_tree!(axiom.tree2)
    end
end

anti_train_axioms! = distribute(anti_train_axiom!)

include("DataStructs.jl")

@load "infinity.jld"

## The Repl facility, note that for the program structure parsing, we rely on the julia parser;
## and we also need to read the variable name from string.
include("Repl.jl")
r = REPL(Env.empty, Seq.empty)

a1 = repl(r, "function(x) x end").value;
a2 = repl(r, "var y = function(x) x end").value;
a3 = repl(r, "y").value;
a4 = repl(r, "function(x) x(x) end").value;
a5 = repl(r, "begin var y = function(x) x end; y end").value;
as = [a1, a2, a3, a4, a5];

pair_computation(function(x, y) sum(abs(as[x] - as[y])) end, 1:5)

max_stars_repo_name: Non-Contradiction/NextGeneration

max_issues_repo_name: Non-Contradiction/NextGeneration

max_forks_repo_name: Non-Contradiction/NextGeneration


text:
const FILES = String["CLITest.jl",
                     "DecodeTest.jl",
                     "EncodeTest.jl"
                    ]

function run_test(test_module::Module)::Nothing
    Base.invokelatest(test_module.run)

    return nothing
end

function run_tests(test_modules::Vector{Module})::Nothing
    for test_module in test_modules
        run_test(test_module)
    end

    return nothing
end

run_tests([include(test_module_path) for test_module_path in FILES])


text:
function get_image(data, xmax, ymax, xmin, ymin)
    imm = zeros(Int64, xmax - xmin + 1, ymax - ymin + 1)
    for i in 1:length(data[:, 1])
        point = data[i, :]
        x = point[1] - xmin + 1
        y = point[2] - ymin + 1
        @assert 0 < x <= xmax - xmin + 1
        @assert 0 < y <= ymax - ymin + 1
        imm[x, y] = 1
    end
    return imm
end

function prettyprint(imm)
    for i in 1:length(imm[1, :])
        dt = -11 .* imm[:, i] .+ 46
        println(join(convert.(Char, dt)))
    end
end

function foo()
    lines = readlines(open("data/10.txt", "r"))

    points = zeros(Int, length(lines), 2)
    velocity = zeros(Int, length(lines), 2)
    for (i, line) in enumerate(lines)
        start = findfirst(isequal('<'), line) + 1
        middle = findfirst(isequal(','), line)
        stop = findfirst(isequal('>'), line) - 1
        x = parse(Int, line[start:middle-1])
        y = parse(Int, line[middle+1:stop])
        points[i, :] = [x, y]

        start = findnext(isequal('<'), line, stop+2) + 1
        middle = findnext(isequal(','), line, stop+2)
        stop = findnext(isequal('>'), line, stop+2) - 1
        x = parse(Int, line[start:middle-1])
        y = parse(Int, line[middle+1:stop])
        velocity[i, :] = [x, y]
    end

    xdist = 100000
    ydist = 100000
    xtime = 0
    ytime = 0
    for time in 10000:12000
        data = points + (time .* velocity)

        ys = maximum(data[:, 2]) - minimum(data[:, 2])
        xs = maximum(data[:, 1]) - minimum(data[:, 1])
        if xs < xdist
            xdist = xs
            xtime = time
        end
        #= colsum = sum(imm, dims=1) =#
        if ys < ydist
            ydist = ys
            ytime = time
        end
    end

    @assert ytime == xtime
    data = points + (xtime .* velocity)
    xmax = maximum(data[:, 1])
    xmin = minimum(data[:, 1])
    ymax = maximum(data[:, 2])
    ymin = minimum(data[:, 2])

    imm = get_image(data, xmax, ymax, xmin, ymin)

    prettyprint(imm)
end

foo()


text:
function make_progressbar(progress_every)
    id = uuid4()
    @info ProgressLogging.Progress(id)
    next_threshold = progress_every
    function update(; fraction)
        if fraction >= next_threshold
            @info ProgressLogging.Progress(id, fraction)
            next_threshold += progress_every
        end
    end
    close() = @info ProgressLogging.Progress(id, done=true)
    return update, close
end

max_stars_repo_name: nathanaelbosch/capos

max_issues_repo_name: nathanaelbosch/capos

max_forks_repo_name: nathanaelbosch/capos


text:
using Plots
gr(fmt=:p)
function plt(x_vals,y_vals)
    plot(x_vals,y_vals, arrow=true,color=:blue,
    legend=:none,  framestyle=:origin)
end

max_stars_repo_name: bkgsur/QuantitativeEconomics

max_issues_repo_name: bkgsur/QuantitativeEconomics

max_forks_repo_name: bkgsur/QuantitativeEconomics


text:
using .GalacticOptim

function GalacticOptim.OptimizationFunction(post::Posterior, args...; kwargs...)
    tpost = asflat(post)
    ℓ(x,p) = -logdensity(tpost, x)
    return OptimizationFunction(ℓ, args...; kwargs...), tpost.transform
end

function GalacticOptim.OptimizationFunction(post::TransformedPosterior, args...; kwargs...)
    ℓ(x,p) = -logdensity(post, x)
    return OptimizationFunction(ℓ, args...; kwargs...), post.transform
end

"""
    solve(prob, opt, args...; transform=nothing, kwargs...)
Find the optimum of the problem.

Optional keyword argument `transform``
will transform the solution to parameter space, and the return will be
a tuple with the galactic optim solution in the first element and the optimum
location in model space in the second argument
"""
function GalacticOptim.solve(prob::OptimizationProblem, opt, transform::Union{Nothing, HypercubeTransform.AbstractHypercubeTransform, HypercubeTransform.TransformVariables.AbstractTransform}, args...; kwargs...)
    sol = solve(prob, opt, args...; kwargs...)
    if isnothing(transform)
        return sol
    else
        p = HypercubeTransform.transform(transform, sol.u)
        return sol, p
    end
end


text:
function ldlfor!(a::Matrix{Float64}, b::Vector{Float64})
  n = size(a, 1)
  for i in 1:n
    total = b[i]
    if i > 1
      for j in 1:(i-1)
        total -= a[j, i] * b[j]
      end
      #(i=i, total=total, ad=a[i, i], newb=total / a[i, i]) |> display
    end
    b[i] = total / a[i, i]
    #i == 1 && (i=i, total=total, ad=a[i, i], newb=total / a[i, i]) |> display
  end
end

function ldlfor(a::Matrix{Float64}, b::Vector{Float64})
  bt = deepcopy(b)
  n = size(a, 1)
  for i in 1:n
    total = bt[i]
    if i > 1
      for j in 1:(i-1)
        total -= a[j, i] * bt[j]
      end
      #(i=i, total=total, ad=a[i, i], newb=total / a[i, i]) |> display
    end
    bt[i] = total / a[i, i]
    #i == 1 && (i=i, total=total, ad=a[i, i], newb=total / a[i, i]) |> display
  end
  bt
end

export
  ldlfor!,
  ldlfor


text:
using Clang
using Clang.LibClang
using Test

@testset "macro" begin
    trans_unit = parse_header(joinpath(@__DIR__, "c", "macro.h"),
                              flags = CXTranslationUnit_DetailedPreprocessingRecord)
    ctx = DefaultContext()
    push!(ctx.trans_units, trans_unit)
    root_cursor = getcursor(trans_unit)
    cursors = children(root_cursor)
    g_uri_xxx = search(cursors, x->name(x)=="G_URI_RESERVED_CHARS_SUBCOMPONENT_DELIMITERS")[1]
    wrap!(ctx, g_uri_xxx)
    expr = :(const G_URI_RESERVED_CHARS_SUBCOMPONENT_DELIMITERS = "!\$&'()*+,;=")
    @test ctx.common_buffer[:G_URI_RESERVED_CHARS_SUBCOMPONENT_DELIMITERS].items[1] == expr

    # test tokenize for macro instantiation
    func = search(cursors, x->kind(x)==CXCursor_FunctionDecl)[1]
    body = children(func)[1]
    op = children(body)[3]
    subop = children(op)[2]
    @test mapreduce(x->x.text, *, tokenize(subop)) == "FOO(foo,1,x)"
    wrap!(ctx, func)
    @test !isempty(ctx.api_buffer)

    # Issue #270 Support '==' in macros
    equals_macro = search(cursors, x->name(x)=="EQUALS_A_B")[1]
    wrap!(ctx, equals_macro)
    @test ctx.common_buffer[:EQUALS_A_B].items[1] == :(const EQUALS_A_B = A == B)
end


text:
#>  (==)(x::T,y::T), (!=)(x::T, y::T)

function (==)(x::Mag, y::Mag)
    0 != ccall(@libarb(mag_equal), Cint, (Ref{Mag}, Ref{Mag}), x, y)
end
function (!=)(x::Mag, y::Mag)
    0 == ccall(@libarb(mag_equal), Cint, (Ref{Mag}, Ref{Mag}), x, y)
end

function (==)(x::ArbFloat{P}, y::ArbFloat{P})  where {P}
    0 != ccall(@libarb(arf_equal), Cint, (Ref{ArbFloat}, Ref{ArbFloat}), x, y)
end
function (!=)(x::ArbFloat{P}, y::ArbFloat{P})  where {P}
    0 == ccall(@libarb(arf_equal), Cint, (Ref{ArbFloat}, Ref{ArbFloat}), x, y)
end

function (==)(x::ArbReal{P}, y::ArbReal{P})  where {P}
    0 != ccall(@libarb(arb_equal), Cint, (Ref{ArbReal}, Ref{ArbReal}), x, y)
end
function (!=)(x::ArbReal{P}, y::ArbReal{P})  where {P}
    0 == ccall(@libarb(arb_equal), Cint, (Ref{ArbReal}, Ref{ArbReal}), x, y)
end

function (==)(x::ArbComplex{P}, y::ArbComplex{P})  where {P}
    0 != ccall(@libarb(acb_equal), Cint, (Ref{ArbComplex}, Ref{ArbComplex}), x, y)
end
function (!=)(x::ArbComplex{P}, y::ArbComplex{P})  where {P}
    0 == ccall(@libarb(acb_equal), Cint, (Ref{ArbComplex}, Ref{ArbComplex}), x, y)
end

#>  F(x::T,y::T), for F as {<, >, <=, >=}

#=
        x < y     x == y     x > y
          -         0          +
=#
function cmp(x::Mag, y::Mag)
    ccall(@libarb(mag_cmp), Cint, (Ref{Mag}, Ref{Mag}), x, y)
end
function cmp(x::ArbFloat{P}, y::ArbFloat{P})  where {P}
    ccall(@libarb(arf_cmp), Cint, (Ref{ArbFloat}, Ref{ArbFloat}), x, y)
end

function (<)(x::Mag, y::Mag)
    signbit(cmp(x, y))
end
function (>)(x::Mag, y::Mag)
    signbit(cmp(y, x))
end
function (<=)(x::Mag, y::Mag)
    (x < y) || (x == y)
end
function (>=)(x::Mag, y::Mag)
    (x > y) || (x == y)
end

function (<)(x::ArbFloat{P}, y::ArbFloat{P})  where {P}
    signbit(cmp(x, y))
end
function (>)(x::ArbFloat{P}, y::ArbFloat{P})  where {P}
    signbit(cmp(y, x))
end
function (<=)(x::ArbFloat{P}, y::ArbFloat{P})  where {P}
    (x < y) || (x == y)
end
function (>=)(x::ArbFloat{P}, y::ArbFloat{P})  where {P}
    (x > y) || (x == y)
end

function (<)(x::ArbReal{P}, y::ArbReal{P})  where {P}
    x = upperbound(x)
    y = lowerbound(y)
    0 != ccall(@libarb(arb_lt), Cint, (Ref{ArbReal}, Ref{ArbReal}), x, y)
end

function (>)(x::ArbReal{P}, y::ArbReal{P})  where {P}
    x = lowerbound(x)
    y = upperbound(y)
    0 != ccall(@libarb(arb_gt), Cint, (Ref{ArbReal}, Ref{ArbReal}), x, y)
end

function (<=)(x::ArbReal{P}, y::ArbReal{P})  where {P}
    x < y ||
    0 != ccall(@libarb(arb_contains), Cint, (Ref{ArbReal}, Ref{ArbReal}), x, y)
end

function (>=)(x::ArbReal{P}, y::ArbReal{P})  where {P}
    x > y ||
    0 != ccall(@libarb(arb_contains), Cint, (Ref{ArbReal}, Ref{ArbReal}), x, y)
end

# ArbComplex comparisons < > <= >=

for F in (:(==), :(!=), :(<), :(<=), :(>=), :(>), :isequal, :isless)
    @eval begin
        $F(x::ArbFloat{P}, y::T) where {P, T<:Integer} =
            $F(promote(x, y)...,)
        $F(x::T, y::ArbFloat{P}) where {P, T<:Integer} =
            $F(promote(x, y)...,)
        $F(x::ArbReal{P}, y::T) where {P, T<:Integer} =
            $F(promote(x, y)...,)
        $F(x::T, y::ArbReal{P}) where {P, T<:Integer} =
            $F(promote(x, y)...,)

        $F(x::ArbFloat{P}, y::T) where {P, T<:IEEEFloat} =
            $F(promote(x, y)...,)
        $F(x::T, y::ArbFloat{P}) where {P, T<:IEEEFloat} =
            $F(promote(x, y)...,)
        $F(x::ArbReal{P}, y::T) where {P, T<:IEEEFloat} =
            $F(promote(x, y)...,)
        $F(x::T, y::ArbReal{P}) where {P, T<:IEEEFloat} =
            $F(promote(x, y)...,)

        $F(x::ArbFloat{P}, y::Rational) where {P} =
            $F(promote(x, y)...,)
        $F(x::Rational, y::ArbFloat{P}) where {P} =
            $F(promote(x, y)...,)
        $F(x::ArbReal{P}, y::Rational) where {P} =
            $F(promote(x, y)...,)
        $F(x::Rational, y::ArbReal{P}) where {P} =
            $F(promote(x, y)...,)
    end
end

#=
       x < 2^y   x == 2^y   x > 2^y
          -         0          +
=#
function cmp2pow(x::Mag, y::Signed)
    ccall(@libarb(mag_cmp_2exp_si), Cint, (Ref{Mag}, Clong), x, y)
end
function cmp2pow(x::ArbFloat{P}, y::Signed) where {P}
    ccall(@libarb(mag_cmp_2exp_si), Cint, (Ref{Mag}, Clong), x, y)
end

#>  min(x::T,y::T), max(x::T, y::T)

function min(x::Mag, y::Mag)
    z = Mag()
    ccall(@libarb(mag_min), Cvoid, (Ref{Mag}, Ref{Mag}, Ref{Mag}), z, x, y)
    return z
end

function max(x::Mag, y::Mag)
    z = Mag()
    ccall(@libarb(mag_max), Cvoid, (Ref{Mag}, Ref{Mag}, Ref{Mag}), z, x, y)
    return z
end

minmax(x::Mag, y::Mag) = (x <= y ? (x, y) : (y, x))

function min(x::ArbFloat{P}, y::ArbFloat{P}) where {P}
    z = ArbFloat{P}()
    ccall(@libarb(arf_min), Cvoid, (Ref{ArbFloat}, Ref{ArbFloat}, Ref{ArbFloat}), z, x, y)
    return z
end

function max(x::ArbFloat{P}, y::ArbFloat{P}) where {P}
    z = ArbFloat{P}()
    ccall(@libarb(arf_max), Cvoid, (Ref{ArbFloat}, Ref{ArbFloat}, Ref{ArbFloat}), z, x, y)
    return z
end

minmax(x::ArbFloat{P}, y::ArbFloat{P}) where {P} = (x <= y ? (x, y) : (y, x))

function min(x::ArbReal{P}, y::ArbReal{P}) where {P}
    z = ArbReal{P}()
    ccall(@libarb(arb_min), Cvoid, (Ref{ArbReal}, Ref{ArbReal}, Ref{ArbReal}), z, x, y)
    return z
end

function max(x::ArbReal{P}, y::ArbReal{P}) where {P}
    z = ArbReal{P}()
    ccall(@libarb(arb_max), Cvoid, (Ref{ArbReal}, Ref{ArbReal}, Ref{ArbReal}), z, x, y)
    return z
end

minmax(x::ArbReal{P}, y::ArbReal{P}) where {P} = (x <= y ? (x, y) : (y, x))

function min(x::ArbComplex{P}, y::ArbComplex{P}, noNaNs::Bool=false) where {P}
    z = ArbComplex{P}()
    analytic = noNaNs ? 0 : 1
    ccall(@libarb(acb_real_min), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Cint, Clong), z, x, y, analytic, P)
    return z
end

function max(x::ArbComplex{P}, y::ArbComplex{P}; noNaNs::Bool=true) where {P}
    z = ArbComplex{P}()
    analytic = noNaNs ? 0 : 1
    ccall(@libarb(acb_real_max), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Cint, Clong), z, x, y, analytic, P)
    return z
end

minmax(x::ArbComplex{P}, y::ArbComplex{P}) where {P} = min(x, y), max(x, y)

isapprox(x::ArbFloat{P}, y::ArbFloat{P}) where {P} = isapprox(x, y, rtol=√eps(max(abs(x),abs(y))))
isapprox(x::ArbReal{P}, y::ArbReal{P}) where {P} = isapprox(midpoint(x), midpoint(y))
isapprox(x::ArbComplex{P}, y::ArbComplex{P}) where {P} = isapprox(midpoint(real(x)), midpoint(real(y))) && isapprox(midpoint(imag(x)), midpoint(imag(y)))

isapprox(x::ArbFloat{P}, y::ArbReal{P})  where {P} = isapprox(x, midpoint(y))
isapprox(x::ArbReal{P}, y::ArbFloat{P})  where {P} = isapprox(midpoint(x), y)
isapprox(x::ArbFloat{P}, y::ArbComplex{P}) where {P} = isapprox(x, midpoint(real(y))) && isapprox(zero(ArbFloat{P}), midpoint(imag(y)))
isapprox(x::ArbComplex{P}, y::ArbFloat{P}) where {P} = isapprox(midpoint(real(x)), y) && isapprox(zero(ArbFloat{P}), midpoint(imag(x)))
isapprox(x::ArbReal{P}, y::ArbComplex{P}) where {P} = isapprox(midpoint(x), midpoint(real(y))) && isapprox(zero(ArbFloat{P}), midpoint(imag(y)))
isapprox(x::ArbComplex{P}, y::ArbReal{P}) where {P} = isapprox(midpoint(real(x)), midpoint(y)) && isapprox(zero(ArbFloat{P}), midpoint(imag(x)))

isapprox(x::ArbFloat{P}, y::F) where {P, F<:IEEEFloat} = isapprox(F(x), y)
isapprox(x::F, y::ArbFloat{P}) where {P, F<:IEEEFloat} = isapprox(x, F(y))
isapprox(x::ArbReal{P}, y::F) where {P, F<:IEEEFloat} = isapprox(F(x), y)
isapprox(x::F, y::ArbReal{P}) where {P, F<:IEEEFloat} = isapprox(x, F(y))
isapprox(x::ArbComplex{P}, y::F) where {P, F<:IEEEFloat} = isapprox(F(real(x)), y)
isapprox(x::F, y::ArbComplex{P}) where {P, F<:IEEEFloat} = isapprox(x, F(real(y)))

isapprox(x::ArbFloat{P}, y::Real) where {P} = isapprox(x, ArbFloat{P}(y))
isapprox(x::Real, y::ArbFloat{P}) where {P} = isapprox(ArbFloat{P}(x), y)
isapprox(x::ArbReal{P}, y::Real) where {P} = isapprox(x, ArbReal{P}(y))
isapprox(x::Number, y::ArbReal{P}) where {P} = isapprox(ArbReal{P}(x), y)
isapprox(x::ArbComplex{P}, y::Number) where {P} = isapprox(x, ArbComplex{P}(y))
isapprox(x::Number, y::ArbComplex{P}) where {P} = isapprox(ArbComplex{P}(x), y)


text:
# indicator of a hyperslab

export IndHyperslab

"""
**Indicator of a hyperslab**

    IndHyperslab(low, a, upp)

For an array `a` and scalars `low` and `upp`, returns the indicator of set
```math
S = \\{x : low \\leq \\langle a,x \\rangle \\leq upp \\}.
```
"""
struct IndHyperslab{R <: Real, T <: AbstractArray{R}} <: ProximableFunction
  low::R
  a::T
  upp::R
  norm_a::R
  function IndHyperslab{R, T}(low::R, a::T, upp::R) where {R <: Real, T <: AbstractArray{R}}
    norm_a = norm(a)
    if (norm_a == 0 && (upp < 0 || low > 0)) || upp < low
        error("function is improper")
    end
    new(low, a, upp, norm_a)
  end
end

IndHyperslab(low::R, a::T, upp::R) where {R <: Real, T <: AbstractArray{R}} = IndHyperslab{R, T}(low, a, upp)

is_convex(f::IndHyperslab) = true
is_set(f::IndHyperslab) = true
is_cone(f::IndHyperslab{R}) where R =
    (f.low == f.upp == 0) ||
    (f.low == 0 && f.upp == Inf) ||
    (f.low == -Inf && f.upp == 0) ||
    (f.low == -Inf && f.upp == Inf)

function (f::IndHyperslab{R})(x::AbstractArray{R}) where R
  s = dot(f.a, x)
  # if f.low <= s <= f.upp
  # tol = (R(1) + abs(s))*eps(R)
  if f.low - s <= eps(R)*f.norm_a*(1 + abs(f.low)) && s - f.upp <= eps(R)*f.norm_a*(1 + abs(f.upp))
    return zero(R)
  end
  return R(Inf)
end

function prox!(y::AbstractArray{R}, f::IndHyperslab{R}, x::AbstractArray{R}, gamma::R=one(R)) where R
  s = dot(f.a, x)
  if s < f.low && f.norm_a > 0
    y .= x .- ((s - f.low)/f.norm_a^2) .* f.a
  elseif s > f.upp && f.norm_a > 0
    y .= x .- ((s - f.upp)/f.norm_a^2) .* f.a
  else
    copyto!(y, x)
  end
  return zero(R)
end

function prox_naive(f::IndHyperslab{R}, x::AbstractArray{R}, gamma::R=one(R)) where R
  s = dot(f.a, x)
  if s < f.low && f.norm_a > 0
    return x - ((s - f.low)/norm(f.a)^2) * f.a, 0
  elseif s > f.upp && f.norm_a > 0
    return x - ((s - f.upp)/norm(f.a)^2) * f.a, 0
  else
    return x, 0
  end
end


text:
using SpecialFunctions, StatsFuns, Base.Iterators

"""
    log_int_prod_2_Gammas(α1, β1, α2, β2)

Compute the log of ``\\int_{\\mathbb{R}^+} Gamma(α1, β1)Gammas(α2, β2)`` exactly. This formula is only valid for ``α1 + α2 > 1``.

# Examples
```julia-repl
julia> log_int_prod_2_Gammas(1, 1, 1, 1) #== -log(2)
-0.6931471805599453
julia> log_int_prod_2_Gammas(2, 0.5, 1, 0.5) #==  log(0.5^2*0.5 )
-0.6931471805599453
```
"""
function log_int_prod_2_Gammas(α1, β1, α2, β2)
    @assert α1 + α2 > 1 "The formula for the L2 distance is only valid for α1 + α2 > 1"
    return α1 * log(β1) + α2 * log(β2) + lgamma_local(α1 + α2 - 1) - lgamma_local(α1) - lgamma_local(α2) - (α1 + α2 - 1) * log(β1 + β2)
end

"""
    int_prod_2_Gammas(α1, β1, α2, β2)

Compute ``\\int_{\\mathbb{R}^+} Gamma(α1, β1)Gammas(α2, β2)`` exactly. This formula is only valid for ``α1 + α2 > 1``.

# Examples
```julia-repl
julia> int_prod_2_Gammas(1, 1, 1, 1)
0.5
julia> int_prod_2_Gammas(2, 0.5, 1, 0.5) #==  0.5^2*0.5
-0.6931471805599453
```
"""
int_prod_2_Gammas(α1, β1, α2, β2) = log_int_prod_2_Gammas(α1, β1, α2, β2) |> exp

function log_L2_dist_Gamma_mixtures(logw1, αlist_1, βlist_1, logw2, αlist_2, βlist_2)
    I = length(logw1)
    J = length(logw2)
    @assert I == length(αlist_1)
    @assert I == length(βlist_1)
    @assert J == length(αlist_2)
    @assert J == length(βlist_2)
    # println("I=$I, J=$J")

    first_logterm = logsumexp(logw1[i] + logw1[j] + log_int_prod_2_Gammas(αlist_1[i], βlist_1[i], αlist_1[j], βlist_1[j]) for (i,j) in Base.Iterators.product(1:I, 1:I))

    second_logterm =  logsumexp(logw2[i] + logw2[j] + log_int_prod_2_Gammas(αlist_2[i], βlist_2[i], αlist_2[j], βlist_2[j]) for (i,j) in Base.Iterators.product(1:J, 1:J))

    third_log_term = logsumexp(log(2) + logw1[i] + logw2[j] + log_int_prod_2_Gammas(αlist_1[i], βlist_1[i], αlist_2[j], βlist_2[j]) for (i,j) in Base.Iterators.product(1:I, 1:J))

    # println([first_logterm, second_logterm, third_log_term])
    # println(exp.([first_logterm, second_logterm, third_log_term]))
    # println(sum(exp.([first_logterm, second_logterm, third_log_term]) .* [1,1,-1]))
    #
    # lx = [first_logterm, second_logterm, third_log_term]
    # signs = [1,1,-1]
    # m = maximum(lx)
    #
    # println("m= $m")
    # scaled_sum = sum(signs .* exp.(lx .- m))
    # println("scaled_sum= $scaled_sum")
    #
    # println("abs(scaled_sum) <= eps(Float64): $(abs(scaled_sum) <= eps(Float64))")

    log_squared_L2_dist = ExactWrightFisher.signed_logsumexp([first_logterm, second_logterm, third_log_term], [1,1,-1])[2] #The result will always be positive
    return 0.5*log_squared_L2_dist
end

"""
    log_int_prod_2_Dir(α1, α2)

Compute the log of ``\\int_{\\nabla_K} Dir(α1)Dir(α2)`` exactly. This formula is only valid for ``\\forall j \\in \\{1,\\ldots, K\\}, \\alpha_{1,j} + \\alpha_{2,j} > 1``.

# Examples
```julia-repl
julia> DualOptimalFiltering.log_int_prod_2_Dir([1,1], [1,1])
0.0
julia> log_int_prod_2_Gammas(2, 0.5, 1, 0.5) #==  log(0.5^2*0.5 )
-0.6931471805599453
```
"""
function log_int_prod_2_Dir(α1, α2)
    @assert all(α1 .+ α2 .> 1) "The formula for the L2 distance is only valid for α1 + α2 > 1 for each component of α"
    return lmvbeta(α1 .+ α2 .- 1) - lmvbeta(α1) - lmvbeta(α2)
end

"""
    lmvbeta(α::AbstractArray{T,1}) for T <: Real

Compute the log of the [multivariate beta function](https://en.wikipedia.org/wiki/Beta_function#Multivariate_beta_function): ``\\Beta(\\alpha_1,\\alpha_2,\\ldots\\alpha_n) = \\frac{\\Gamma(\\alpha_1)\\Gamma(\\alpha_2) \\cdots \\Gamma(\\alpha_n)}{\\Gamma(\\alpha_1 + \\alpha_2 + \\cdots + \\alpha_n)}``.

# Examples
```julia-repl
julia> lmvbeta([1,1])
0.0
```
"""
function lmvbeta(α::AbstractArray{T,1}) where T <: Real
    return sum(lgamma_local.(α)) - lgamma_local(sum(α))
end

# function log_L2_dist_Dirichlet_mixtures(logw1::RealVector, αlist_1::Array{RealVector, 1}, logw2::RealVector, αlist_2::Array{RealVector, 1})
function log_L2_dist_Dirichlet_mixtures(logw1::RealVector, αlist_1::Array{T, 1}, logw2::RealVector, αlist_2::Array{T, 1}) where T <: RealVector
    I = length(logw1)
    J = length(logw2)
    first_logterm = logsumexp(logw1[i] + logw1[j] + log_int_prod_2_Dir(αlist_1[i], αlist_1[j]) for (i,j) in Base.Iterators.product(1:I, 1:I))
    second_logterm =  logsumexp(logw2[i] + logw2[j] + log_int_prod_2_Dir(αlist_2[i], αlist_2[j]) for (i,j) in Base.Iterators.product(1:J, 1:J))
    third_log_term = logsumexp(log(2) + logw1[i] + logw2[j] + log_int_prod_2_Dir(αlist_1[i], αlist_2[j]) for (i,j) in Base.Iterators.product(1:I, 1:J))

    # println([first_logterm, second_logterm, third_log_term])
    # println(exp.([first_logterm, second_logterm, third_log_term]))
    # println(sum(exp.([first_logterm, second_logterm, third_log_term]) .* [1,1,-1]))
    #
    # lx = [first_logterm, second_logterm, third_log_term]
    # signs = [1,1,-1]
    # m = maximum(lx)
    #
    # println("m= $m")
    # scaled_sum = sum(signs .* exp.(lx .- m))
    # println("scaled_sum= $scaled_sum")
    #
    # println("abs(scaled_sum) <= eps(Float64): $(abs(scaled_sum) <= eps(Float64))")

    log_squared_L2_dist = ExactWrightFisher.signed_logsumexp([first_logterm, second_logterm, third_log_term], [1,1,-1])[2] #The result will always be positive
    return 0.5*log_squared_L2_dist
end

# DualOptimalFiltering.log_L2_dist_Dirichlet_mixtures(log.([0.2,0.3,0.5]), [[1,2,3],[1,2,4],[5,3,2]], log.([0.6,0.4]), [[1,2,1],[4,2,3]])
# #
# DualOptimalFiltering.log_L2_dist_Dirichlet_mixtures(log.([0.2,0.3,0.5]), [[1,2,3],[1,2,4],[5,3,2]], log.([0.2,0.3,0.5]), [[1,2,3],[1,2,4],[5,3,2]])
#
#
# DualOptimalFiltering.log_L2_dist_Dirichlet_mixtures(log.([0.2,0.3,0.5]), [[1,2,3],[1,2,4],[5,3,2]], log.([0.200001,0.3,0.5] |> DualOptimalFiltering.normalise), [[1,2,3],[1,2,4],[5,3,2]])
#
# DualOptimalFiltering.log_L2_dist_Dirichlet_mixtures(log.([0.2,0.3,0.5]), [[1,2,3],[1,2,4],[5,3,2]], log.([0.2000001,0.3,0.5] |> DualOptimalFiltering.normalise), [[1,2,3],[1,2,4],[5,3,2]])
#
# DualOptimalFiltering.log_L2_dist_Dirichlet_mixtures(log.([0.2,0.3,0.5]), [[1,2,3],[1,2,4],[5,3,2]], log.([0.20000001,0.3,0.5] |> DualOptimalFiltering.normalise), [[1,2,3],[1,2,4],[5,3,2]])
#
# DualOptimalFiltering.log_L2_dist_Dirichlet_mixtures(log.([0.2,0.3,0.5]), [[1,2,3],[1,2,4],[5,3,2]], log.([0.2000000001,0.3,0.5] |> DualOptimalFiltering.normalise), [[1,2,3],[1,2,4],[5,3,2]])
#
# import DualOptimalFiltering.RR
#
# ExactWrightFisher.signed_logsumexp([log(RR(3.00000000000001)), log(RR(3)), log(RR(6))],[1,1,-1])
#
# ExactWrightFisher.signed_logsumexp([log(3), log(3), log(6)],[1,1,-1])

max_stars_repo_name: konkam/DualOptimalFiltering

max_issues_repo_name: konkam/DualOptimalFiltering

max_forks_repo_name: konkam/DualOptimalFiltering


text:
module AWG
using Libdl # for NI dlls
include("../INSTR_TYPES.jl")
include("../lib_locations.jl")
include("../visa/VISA.jl")
#
## include all AWG files
include("./Agilent/AWG_Agilent_common.jl")
include("./Agilent/Agilent33500B.jl")
include("./Agilent/Agilent33250A.jl")
# NI
include("NIFgen/NIFgen.jl")
include("NIFgen/NIFgen_Misc.jl")
# Rigol
include("Rigol/AWG_Rigol_common.jl")
include("Rigol/Rigol_DG1000Z.jl")

## Common functions
# Type Alias for SCPI devices
AWG_SCPI = Union{DG1000Z, F335x2, F332x1}
include("AWG_SCPI_common.jl")

"""
# Generate a sine Waveform
`function gen_sin(obj::T; ch="1", freq::Real=1e6, amp::Real=1, offset::Signed=0, phase::Real=0) where {T <: AWG_SCPI}`
- ch = {"1", "2"}
- freq = Frequency in Hertz
- amp = Amplitude in Volts
- offset = Offsett in Volts
- phase = Phase in Deg
"""
function gen_sin(obj::T; ch="1", freq::Real=1e6, amp::Real=1, offset::Signed=0, phase::Real=0) where {T <: AWG_SCPI}
    set_freq(obj; ch, fr=freq);
    set_wfm(obj; ch, func="SIN");
    set_amplit(obj; ch, volt=amp);
    set_offs(obj; ch, volt=offset);
    set_phase(obj; ch, phase);
end

"""
# Generate a square Waveform
`function gen_square(obj::T; ch="1", freq::Real=1e6, amp::Real=1, offset::Signed=0, phase::Real=0, duty::Real=50) where {T <: AWG_SCPI}`
- ch = {"1", "2"}
- freq = Frequency in Hertz
- amp = Amplitude in Volts
- offset = Offsett in Volts
- phase = Phase in Deg
- duty = Duty cycle in %
"""
function gen_square(obj::T; ch="1", freq::Real=1e6, amp::Real=1, offset::Signed=0, phase::Real=0, duty::Real=50) where {T <: AWG_SCPI}
    set_freq(obj; ch, fr=freq);
    set_wfm(obj; ch, func="SQU");
    set_amplit(obj; ch, volt=amp);
    set_offs(obj; ch, volt=offset);
    set_phase(obj; ch, phase);
    set_duty(obj; ch, duty);
end

end #endmodule