text:
using FiniteDifferences, ForwardDiff
using LinearAlgebra: I
using StatsBase: AbstractWeights, pweights
import Manifolds: mean!, median!, InducedBasis, induced_basis, get_chart_index

include("utils.jl")

struct TestEuclidean{N} <: AbstractManifold{ℝ} end
struct TestEuclideanMetric <: AbstractMetric end

Manifolds.manifold_dimension(::TestEuclidean{N}) where {N} = N
function Manifolds.local_metric(
    M::MetricManifold{ℝ,<:TestEuclidean,<:TestEuclideanMetric},
    ::Any,
    ::InducedBasis,
)
    return Diagonal(1.0:manifold_dimension(M))
end
function Manifolds.local_metric(
    M::MetricManifold{ℝ,<:TestEuclidean,<:TestEuclideanMetric},
    ::Any,
    ::InducedBasis,
)
    return Diagonal(1.0:manifold_dimension(M))
end

struct TestSphere{N,T} <: AbstractManifold{ℝ}
    r::T
end

struct TestSphericalMetric <: AbstractMetric end

Manifolds.manifold_dimension(::TestSphere{N}) where {N} = N
function Manifolds.local_metric(
    M::MetricManifold{ℝ,<:TestSphere,<:TestSphericalMetric},
    p,
    ::InducedBasis,
)
    r = base_manifold(M).r
    d = allocate(p)
    d[1] = r^2
    d[2] = d[1] * sin(p[1])^2
    return Diagonal(d)
end
sph_to_cart(θ, ϕ) = [cos(ϕ) * sin(θ), sin(ϕ) * sin(θ), cos(θ)]

struct BaseManifold{N} <: AbstractManifold{ℝ} end
struct BaseManifoldMetric{M} <: AbstractMetric end
struct DefaultBaseManifoldMetric <: AbstractMetric end
struct NotImplementedMetric <: AbstractMetric end

Manifolds.manifold_dimension(::BaseManifold{N}) where {N} = N
Manifolds.inner(::BaseManifold, p, X, Y) = 2 * dot(X, Y)
Manifolds.exp!(::BaseManifold, q, p, X) = q .= p + 2 * X
Manifolds.log!(::BaseManifold, Y, p, q) = Y .= (q - p) / 2
Manifolds.project!(::BaseManifold, Y, p, X) = Y .= 2 .* X
Manifolds.project!(::BaseManifold, q, p) = (q .= p)
Manifolds.injectivity_radius(::BaseManifold) = Inf
Manifolds.injectivity_radius(::BaseManifold, ::Any) = Inf
Manifolds.injectivity_radius(::BaseManifold, ::AbstractRetractionMethod) = Inf
Manifolds.injectivity_radius(::BaseManifold, ::ExponentialRetraction) = Inf
Manifolds.injectivity_radius(::BaseManifold, ::Any, ::AbstractRetractionMethod) = Inf
Manifolds.injectivity_radius(::BaseManifold, ::Any, ::ExponentialRetraction) = Inf
function Manifolds.local_metric(
    ::MetricManifold{ℝ,BaseManifold{N},BaseManifoldMetric{N}},
    p,
    ::InducedBasis,
) where {N}
    return 2 * one(p * p')
end
function Manifolds.exp!(
    M::MetricManifold{ℝ,BaseManifold{N},BaseManifoldMetric{N}},
    q,
    p,
    X,
) where {N}
    return exp!(base_manifold(M), q, p, X)
end
function Manifolds.vector_transport_to!(::BaseManifold, Y, p, X, q, ::ParallelTransport)
    return (Y .= X)
end
function Manifolds.get_basis(
    ::BaseManifold{N},
    p,
    B::DefaultOrthonormalBasis{<:Any,ManifoldsBase.TangentSpaceType},
) where {N}
    return CachedBasis(B, [(Matrix{eltype(p)}(I, N, N)[:, i]) for i in 1:N])
end
function Manifolds.get_coordinates!(
    ::BaseManifold,
    Y,
    p,
    X,
    ::DefaultOrthonormalBasis{<:Any,ManifoldsBase.TangentSpaceType},
)
    return Y .= X
end
function Manifolds.get_vector!(
    ::BaseManifold,
    Y,
    p,
    X,
    ::DefaultOrthonormalBasis{<:Any,ManifoldsBase.TangentSpaceType},
)
    return Y .= X
end
Manifolds.default_metric_dispatch(::BaseManifold, ::DefaultBaseManifoldMetric) = Val(true)
function Manifolds.projected_distribution(M::BaseManifold, d)
    return ProjectedPointDistribution(M, d, project!, rand(d))
end
function Manifolds.projected_distribution(M::BaseManifold, d, p)
    return ProjectedPointDistribution(M, d, project!, p)
end
function Manifolds.mean!(::BaseManifold, y, x::AbstractVector, w::AbstractVector; kwargs...)
    return fill!(y, 1)
end
function Manifolds.median!(
    ::BaseManifold,
    y,
    x::AbstractVector,
    w::AbstractVector;
    kwargs...,
)
    return fill!(y, 2)
end
function Manifolds.mean!(
    ::MetricManifold{ℝ,BaseManifold{N},BaseManifoldMetric{N}},
    y,
    x::AbstractVector,
    w::AbstractVector;
    kwargs...,
) where {N}
    return fill!(y, 3)
end
function Manifolds.median!(
    ::MetricManifold{ℝ,BaseManifold{N},BaseManifoldMetric{N}},
    y,
    x::AbstractVector,
    w::AbstractVector;
    kwargs...,
) where {N}
    return fill!(y, 4)
end

function Manifolds.flat!(
    ::BaseManifold,
    v::FVector{Manifolds.CotangentSpaceType},
    x,
    w::FVector{Manifolds.TangentSpaceType},
)
    v.data .= 2 .* w.data
    return v
end
function Manifolds.sharp!(
    ::BaseManifold,
    v::FVector{Manifolds.TangentSpaceType},
    x,
    w::FVector{Manifolds.CotangentSpaceType},
)
    v.data .= w.data ./ 2
    return v
end

@testset "Metrics" begin
    # some tests failed due to insufficient accuracy for a particularly bad RNG state
    Random.seed!(42)
    @testset "Metric Basics" begin
        #one for MetricManifold, one for AbstractManifold & Metric
        @test length(methods(is_default_metric)) == 2
    end

    @testset "solve_exp_ode error message" begin
        E = TestEuclidean{3}()
        g = TestEuclideanMetric()
        M = MetricManifold(E, g)

        p = [1.0, 2.0, 3.0]
        X = [2.0, 3.0, 4.0]
        @test_throws ErrorException exp(M, p, X)
        using OrdinaryDiffEq
        exp(M, p, X)
    end
    @testset "Local Metric Error message" begin
        M = MetricManifold(BaseManifold{2}(), NotImplementedMetric())
        A = Manifolds.get_default_atlas(M)
        p = [3, 4]
        i = get_chart_index(M, A, p)

        B = induced_basis(M, A, i, TangentSpace)
        @test_throws MethodError local_metric(M, p, B)
    end
    @testset "scaled Euclidean metric" begin
        n = 3
        E = TestEuclidean{n}()
        g = TestEuclideanMetric()
        M = MetricManifold(E, g)
        A = Manifolds.get_default_atlas(M)
        @test repr(M) == "MetricManifold(TestEuclidean{3}(), TestEuclideanMetric())"

        @test TestEuclideanMetric()(E) === M
        @test TestEuclideanMetric(E) === M

        G = Diagonal(1.0:n)
        invG = inv(G)
        @test manifold_dimension(M) == n
        @test base_manifold(M) === E
        @test metric(M) === g

        i_zeros = get_chart_index(M, A, zeros(3))
        B_i_zeros = induced_basis(M, A, i_zeros, TangentSpace)
        @test_throws MethodError local_metric_jacobian(E, zeros(3), B_i_zeros)
        @test_throws MethodError christoffel_symbols_second_jacobian(E, zeros(3), B_i_zeros)

        for vtype in (Vector, MVector{n})
            p, X, Y = vtype(randn(n)), vtype(randn(n)), vtype(randn(n))

            chart_p = get_chart_index(M, A, p)
            B_chart_p = induced_basis(M, A, chart_p, TangentSpace)

            @test check_point(M, p) == check_point(E, p)
            @test check_vector(M, p, X) == check_vector(E, p, X)

            @test local_metric(M, p, B_chart_p) ≈ G
            @test inverse_local_metric(M, p, B_chart_p) ≈ invG
            @test det_local_metric(M, p, B_chart_p) ≈ *(1.0:n...)
            @test log_local_metric_density(M, p, B_chart_p) ≈ sum(log.(1.0:n)) / 2

            fX = ManifoldsBase.TFVector(X, B_chart_p)
            fY = ManifoldsBase.TFVector(Y, B_chart_p)
            @test inner(M, p, fX, fY) ≈ dot(X, G * Y) atol = 1e-6
            @test norm(M, p, fX) ≈ sqrt(dot(X, G * X)) atol = 1e-6

            if VERSION ≥ v"1.1"
                T = 0:0.5:10
                @test geodesic(M, p, X, T) ≈ [p + t * X for t in T] atol = 1e-6
            end

            @test christoffel_symbols_first(M, p, B_chart_p) ≈ zeros(n, n, n) atol = 1e-6
            @test christoffel_symbols_second(M, p, B_chart_p) ≈ zeros(n, n, n) atol = 1e-6
            @test riemann_tensor(M, p, B_chart_p) ≈ zeros(n, n, n, n) atol = 1e-6
            @test ricci_tensor(M, p, B_chart_p) ≈ zeros(n, n) atol = 1e-6
            @test ricci_curvature(M, p, B_chart_p) ≈ 0 atol = 1e-6
            @test gaussian_curvature(M, p, B_chart_p) ≈ 0 atol = 1e-6
            @test einstein_tensor(M, p, B_chart_p) ≈ zeros(n, n) atol = 1e-6

            fdm = FiniteDifferencesBackend(forward_fdm(2, 1))
            @test christoffel_symbols_first(M, p, B_chart_p; backend=fdm) ≈ zeros(n, n, n) atol =
                1e-6
            @test christoffel_symbols_second(M, p, B_chart_p; backend=fdm) ≈ zeros(n, n, n) atol =
                1e-6
            @test riemann_tensor(M, p, B_chart_p; backend=fdm) ≈ zeros(n, n, n, n) atol =
                1e-6
            @test ricci_tensor(M, p, B_chart_p; backend=fdm) ≈ zeros(n, n) atol = 1e-6
            @test ricci_curvature(M, p, B_chart_p; backend=fdm) ≈ 0 atol = 1e-6
            @test gaussian_curvature(M, p, B_chart_p; backend=fdm) ≈ 0 atol = 1e-6
            @test einstein_tensor(M, p, B_chart_p; backend=fdm) ≈ zeros(n, n) atol = 1e-6

            fwd_diff = Manifolds.ForwardDiffBackend()
            @test christoffel_symbols_first(M, p, B_chart_p; backend=fwd_diff) ≈
                  zeros(n, n, n) atol = 1e-6
            @test christoffel_symbols_second(M, p, B_chart_p; backend=fwd_diff) ≈
                  zeros(n, n, n) atol = 1e-6
            @test riemann_tensor(M, p, B_chart_p; backend=fwd_diff) ≈ zeros(n, n, n, n) atol =
                1e-6
            @test ricci_tensor(M, p, B_chart_p; backend=fwd_diff) ≈ zeros(n, n) atol = 1e-6
            @test ricci_curvature(M, p, B_chart_p; backend=fwd_diff) ≈ 0 atol = 1e-6
            @test gaussian_curvature(M, p, B_chart_p; backend=fwd_diff) ≈ 0 atol = 1e-6
            @test einstein_tensor(M, p, B_chart_p; backend=fwd_diff) ≈ zeros(n, n) atol =
                1e-6
        end
    end

    @testset "scaled Sphere metric" begin
        n = 2
        r = 10 * rand()
        θ, ϕ = π * rand(), 2π * rand()
        Sr = TestSphere{n,Float64}(r)
        S = Manifolds.Sphere(n)
        g = TestSphericalMetric()
        M = MetricManifold(Sr, g)
        A = Manifolds.get_default_atlas(M)

        @test manifold_dimension(M) == n
        @test base_manifold(M) === Sr
        @test metric(M) === g

        for vtype in (Vector, MVector{n})
            p = vtype([θ, ϕ])
            chart_p = get_chart_index(M, A, p)
            B_p = induced_basis(M, A, chart_p, TangentSpace)
            G = Diagonal(vtype([1, sin(θ)^2])) .* r^2
            invG = Diagonal(vtype([1, 1 / sin(θ)^2])) ./ r^2
            X, Y = normalize(randn(n)), normalize(randn(n))

            @test local_metric(M, p, B_p) ≈ G atol = 1e-6
            @test inverse_local_metric(M, p, B_p) ≈ invG atol = 1e-6
            @test det_local_metric(M, p, B_p) ≈ r^4 * sin(θ)^2 atol = 1e-6
            @test log_local_metric_density(M, p, B_p) ≈ 2 * log(r) + log(sin(θ)) atol = 1e-6
            fX = ManifoldsBase.TFVector(X, B_p)
            fY = ManifoldsBase.TFVector(Y, B_p)
            @test inner(M, p, fX, fY) ≈ dot(X, G * Y) atol = 1e-6
            @test norm(M, p, fX) ≈ sqrt(dot(X, G * X)) atol = 1e-6

            pcart = sph_to_cart(θ, ϕ)
            Xcart = [
                cos(ϕ)*cos(θ) -sin(ϕ)*sin(θ)
                sin(ϕ)*cos(θ) cos(ϕ)*sin(θ)
                -sin(θ) 0
            ] * X

            if !Sys.iswindows() || Sys.ARCH == :x86_64
                @testset "numerically integrated geodesics for $vtype" begin
                    T = 0:0.1:1
                    @test isapprox(
                        [sph_to_cart(yi...) for yi in geodesic(M, p, X, T)],
                        geodesic(S, pcart, Xcart, T);
                        atol=1e-3,
                        rtol=1e-3,
                    )
                end
            end

            Γ₁ = christoffel_symbols_first(M, p, B_p)
            for i in 1:n, j in 1:n, k in 1:n
                if (i, j, k) == (1, 2, 2) || (i, j, k) == (2, 1, 2)
                    @test Γ₁[i, j, k] ≈ r^2 * cos(θ) * sin(θ) atol = 1e-6
                elseif (i, j, k) == (2, 2, 1)
                    @test Γ₁[i, j, k] ≈ -r^2 * cos(θ) * sin(θ) atol = 1e-6
                else
                    @test Γ₁[i, j, k] ≈ 0 atol = 1e-6
                end
            end

            Γ₂ = christoffel_symbols_second(M, p, B_p)
            for l in 1:n, i in 1:n, j in 1:n
                if (l, i, j) == (1, 2, 2)
                    @test Γ₂[l, i, j] ≈ -cos(θ) * sin(θ) atol = 1e-6
                elseif (l, i, j) == (2, 1, 2) || (l, i, j) == (2, 2, 1)
                    @test Γ₂[l, i, j] ≈ cot(θ) atol = 1e-6
                else
                    @test Γ₂[l, i, j] ≈ 0 atol = 1e-6
                end
            end

            R = riemann_tensor(M, p, B_p)
            for l in 1:n, i in 1:n, j in 1:n, k in 1:n
                if (l, i, j, k) == (2, 1, 1, 2)
                    @test R[l, i, j, k] ≈ -1 atol = 2e-6
                elseif (l, i, j, k) == (2, 1, 2, 1)
                    @test R[l, i, j, k] ≈ 1 atol = 2e-6
                elseif (l, i, j, k) == (1, 2, 1, 2)
                    @test R[l, i, j, k] ≈ sin(θ)^2 atol = 1e-6
                elseif (l, i, j, k) == (1, 2, 2, 1)
                    @test R[l, i, j, k] ≈ -sin(θ)^2 atol = 1e-6
                else
                    @test R[l, i, j, k] ≈ 0 atol = 1e-6
                end
            end

            @test ricci_tensor(M, p, B_p) ≈ G ./ r^2 atol = 2e-6
            @test ricci_curvature(M, p, B_p) ≈ 2 / r^2 atol = 2e-6
            @test gaussian_curvature(M, p, B_p) ≈ 1 / r^2 atol = 2e-6
            @test einstein_tensor(M, p, B_p) ≈
                  ricci_tensor(M, p, B_p) - gaussian_curvature(M, p, B_p) .* G atol = 1e-6
        end
    end

    @testset "Metric decorator" begin
        M = BaseManifold{3}()
        g = BaseManifoldMetric{3}()
        MM = MetricManifold(M, g)

        @test DefaultBaseManifoldMetric(BaseManifold{3}()) ===
              MetricManifold(BaseManifold{3}(), DefaultBaseManifoldMetric())
        MT = DefaultBaseManifoldMetric()
        @test MT(BaseManifold{3}()) ===
              MetricManifold(BaseManifold{3}(), DefaultBaseManifoldMetric())

        g2 = DefaultBaseManifoldMetric()
        MM2 = MetricManifold(M, g2)
        A = Manifolds.get_default_atlas(M)

        @test (@inferred Manifolds.default_metric_dispatch(MM)) ===
              (@inferred Manifolds.default_metric_dispatch(base_manifold(MM), metric(MM)))
        @test (@inferred Manifolds.default_metric_dispatch(MM2)) ===
              (@inferred Manifolds.default_metric_dispatch(base_manifold(MM2), metric(MM2)))
        @test (@inferred Manifolds.default_metric_dispatch(MM2)) === Val(true)
        @test is_default_metric(MM) == is_default_metric(base_manifold(MM), metric(MM))
        @test is_default_metric(MM2) == is_default_metric(base_manifold(MM2), metric(MM2))
        @test is_default_metric(MM2)
        @test Manifolds.default_decorator_dispatch(MM) ===
              Manifolds.default_metric_dispatch(MM)
        @test Manifolds.default_decorator_dispatch(MM2) ===
              Manifolds.default_metric_dispatch(MM2)

        @test convert(typeof(MM2), M) == MM2
        @test_throws ErrorException convert(typeof(MM), M)
        p = [0.1, 0.2, 0.4]
        X = [0.5, 0.7, 0.11]
        Y = [0.13, 0.17, 0.19]
        q = allocate(p)

        p2 = allocate(p)
        copyto!(MM, p2, p)
        p3 = allocate(p)
        copyto!(M, p3, p)
        @test p2 == p3
        X = zero_vector(MM, p)
        Y = allocate(X)
        copyto!(MM, Y, p, X)
        Y2 = allocate(X)
        copyto!(M, Y2, p, X)
        @test Y == Y2

        chart_p = get_chart_index(M, A, p)
        B_p = induced_basis(M, A, chart_p, TangentSpace)
        fX = ManifoldsBase.TFVector(X, B_p)
        fY = ManifoldsBase.TFVector(Y, B_p)

        @test inner(M, p, X, Y) == 2 * dot(X, Y)
        @test inner(MM, p, fX, fY) === inner(M, p, X, Y)
        @test norm(MM, p, fX) === norm(M, p, X)
        @test exp(M, p, X) == p + 2 * X
        @test exp(MM2, p, X) == exp(M, p, X)
        @test exp!(MM, q, p, X) === exp!(M, q, p, X)
        @test retract!(MM, q, p, X) === retract!(M, q, p, X)
        @test retract!(MM, q, p, X, 1) === retract!(M, q, p, X, 1)
        # without a definition for the metric from the embedding, no projection possible
        @test_throws ErrorException log!(MM, Y, p, q) === project!(M, Y, p, q)
        @test_throws ErrorException project!(MM, Y, p, X) === project!(M, Y, p, X)
        @test_throws ErrorException project!(MM, q, p) === project!(M, q, p)
        @test_throws ErrorException vector_transport_to!(MM, Y, p, X, q) ===
                                    vector_transport_to!(M, Y, p, X, q)
        # without DiffEq, these error
        # @test_throws ErrorException exp(MM,x, X, 1:3)
        # @test_throws ErrorException exp!(MM, q, p, X)
        # these always fall back anyways.
        @test zero_vector!(MM, X, p) === zero_vector!(M, X, p)

        @test injectivity_radius(MM, p) === injectivity_radius(M, p)
        @test injectivity_radius(MM) === injectivity_radius(M)
        @test injectivity_radius(MM, ProjectionRetraction()) ===
              injectivity_radius(M, ProjectionRetraction())
        @test injectivity_radius(MM, ExponentialRetraction()) ===
              injectivity_radius(M, ExponentialRetraction())
        @test injectivity_radius(MM) === injectivity_radius(M)

        @test is_point(MM, p) === is_point(M, p)
        @test is_vector(MM, p, X) === is_vector(M, p, X)

        A = Manifolds.get_default_atlas(MM2)
        chart_p = get_chart_index(MM2, A, p)
        B_p = induced_basis(MM2, A, chart_p, TangentSpace)
        @test_throws MethodError local_metric(MM2, p, B_p)
        @test_throws MethodError local_metric_jacobian(MM2, p, B_p)
        @test_throws MethodError christoffel_symbols_second_jacobian(MM2, p, B_p)
        # MM falls back to nondefault error
        @test_throws MethodError projected_distribution(MM, 1, p)
        @test_throws MethodError projected_distribution(MM, 1)

        @test inner(MM2, p, X, Y) === inner(M, p, X, Y)
        @test norm(MM2, p, X) === norm(M, p, X)
        @test distance(MM2, p, q) === distance(M, p, q)
        @test exp!(MM2, q, p, X) === exp!(M, q, p, X)
        @test exp(MM2, p, X) == exp(M, p, X)
        @test log!(MM2, X, p, q) === log!(M, X, p, q)
        @test log(MM2, p, q) == log(M, p, q)
        @test retract!(MM2, q, p, X) === retract!(M, q, p, X)
        @test retract!(MM2, q, p, X, 1) === retract!(M, q, p, X, 1)

        @test project!(MM2, q, p) === project!(M, q, p)
        @test project!(MM2, Y, p, X) === project!(M, Y, p, X)
        @test vector_transport_to!(MM2, Y, p, X, q) == vector_transport_to!(M, Y, p, X, q)
        @test zero_vector!(MM2, X, p) === zero_vector!(M, X, p)
        @test injectivity_radius(MM2, p) === injectivity_radius(M, p)
        @test injectivity_radius(MM2) === injectivity_radius(M)
        @test injectivity_radius(MM2, p, ExponentialRetraction()) ===
              injectivity_radius(M, p, ExponentialRetraction())
        @test injectivity_radius(MM2, ExponentialRetraction()) ===
              injectivity_radius(M, ExponentialRetraction())
        @test injectivity_radius(MM2, p, ProjectionRetraction()) ===
              injectivity_radius(M, p, ProjectionRetraction())
        @test injectivity_radius(MM2, ProjectionRetraction()) ===
              injectivity_radius(M, ProjectionRetraction())
        @test is_point(MM2, p) === is_point(M, p)
        @test is_vector(MM2, p, X) === is_vector(M, p, X)

        a = Manifolds.projected_distribution(M, Distributions.MvNormal(zero(zeros(3)), 1.0))
        b = Manifolds.projected_distribution(
            MM2,
            Distributions.MvNormal(zero(zeros(3)), 1.0),
        )
        @test isapprox(Matrix(a.distribution.Σ), Matrix(b.distribution.Σ))
        @test isapprox(a.distribution.μ, b.distribution.μ)
        @test get_basis(M, p, DefaultOrthonormalBasis()).data ==
              get_basis(MM2, p, DefaultOrthonormalBasis()).data
        @test_throws ErrorException get_basis(MM, p, DefaultOrthonormalBasis())

        fX = ManifoldsBase.TFVector(X, B_p)
        fY = ManifoldsBase.TFVector(Y, B_p)
        coX = flat(M, p, X)
        coY = flat(M, p, Y)
        cofX = flat(M, p, fX)
        cofY = flat(M, p, fY)
        @test coX(X) ≈ norm(M, p, X)^2
        @test coY(X) ≈ inner(M, p, X, Y)
        cotspace = CotangentBundleFibers(M)
        cotspace2 = CotangentBundleFibers(MM)
        @test coX.X ≈ X
        @test inner(M, p, X, Y) ≈ inner(cotspace, p, coX, coY)
        @test inner(MM, p, fX, fY) ≈ inner(cotspace, p, coX, coY)

        @test inner(MM, p, fX, fY) ≈ inner(cotspace2, p, cofX, cofY)
        @test sharp(M, p, coX) ≈ X

        coMMfX = flat(MM, p, fX)
        coMMfY = flat(MM, p, fY)
        @test inner(MM, p, fX, fY) ≈ inner(cotspace2, p, coMMfX, coMMfY)
        @test isapprox(sharp(MM, p, coMMfX).data, fX.data)

        @testset "Mutating flat/sharp" begin
            cofX2 = allocate(cofX)
            flat!(M, cofX2, p, fX)
            @test isapprox(cofX2.data, cofX.data)

            fX2 = allocate(fX)
            sharp!(M, fX2, p, cofX2)
            @test isapprox(fX2.data, fX.data)

            cofX2 = allocate(cofX)
            flat!(MM, cofX2, p, fX)
            @test isapprox(cofX2.data, cofX.data)

            fX2 = allocate(fX)
            sharp!(MM, fX2, p, cofX2)
            @test isapprox(fX2.data, fX.data)
        end

        psample = [[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]]
        Y = pweights([0.5, 0.5])
        # test despatch with results from above
        @test mean(M, psample, Y) ≈ ones(3)
        @test mean(MM2, psample, Y) ≈ ones(3)
        @test mean(MM, psample, Y) ≈ 3 .* ones(3)

        @test median(M, psample, Y) ≈ 2 .* ones(3)
        @test median(MM2, psample, Y) ≈ 2 * ones(3)
        @test median(MM, psample, Y) ≈ 4 .* ones(3)
    end

    @testset "Metric decorator dispatches" begin
        M = BaseManifold{3}()
        g = BaseManifoldMetric{3}()
        MM = MetricManifold(M, g)
        x = [1, 2, 3]
        # nonmutating always go to parent for allocation
        for f in [exp, flat, inverse_retract, log, mean, median, project]
            @test Manifolds.decorator_transparent_dispatch(f, MM) === Val{:parent}()
        end
        for f in [sharp, retract, get_vector, get_coordinates]
            @test Manifolds.decorator_transparent_dispatch(f, MM) === Val{:parent}()
        end
        for f in [vector_transport_along, vector_transport_direction, vector_transport_to]
            @test Manifolds.decorator_transparent_dispatch(f, MM) === Val{:parent}()
        end
        for f in [get_basis, inner]
            @test Manifolds.decorator_transparent_dispatch(f, MM) === Val{:intransparent}()
        end
        for f in [get_coordinates!, get_vector!]
            @test Manifolds.decorator_transparent_dispatch(f, MM) === Val{:intransparent}()
        end

        # mirroring ones are mostly intransparent despite for a few cases - e.g. dispatch/default last variables
        for f in [exp!, flat!, inverse_retract!, log!, mean!, median!]
            @test Manifolds.decorator_transparent_dispatch(f, MM) === Val{:intransparent}()
        end
        for f in [norm, project!, sharp!, retract!]
            @test Manifolds.decorator_transparent_dispatch(f, MM) === Val{:intransparent}()
        end
        for f in [vector_transport_along!, vector_transport_to!]
            @test Manifolds.decorator_transparent_dispatch(f, MM) === Val{:intransparent}()
        end
        @test Manifolds.decorator_transparent_dispatch(vector_transport_direction!, MM) ===
              Val{:parent}()

        @test Manifolds.decorator_transparent_dispatch(exp!, MM, x, x, x, x) ===
              Val{:parent}()
        @test Manifolds.decorator_transparent_dispatch(
            inverse_retract!,
            MM,
            x,
            x,
            x,
            LogarithmicInverseRetraction(),
        ) === Val{:parent}()
        @test Manifolds.decorator_transparent_dispatch(
            retract!,
            MM,
            x,
            x,
            x,
            ExponentialRetraction(),
        ) === Val{:parent}()
    end
end


text:
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.8.0-DEV
#     language: julia
#     name: julia-1.8
# ---

# %%
using Plots
h=0.05
y, x = range(-5,  stop=5+h, step=h), range(-5, stop=10+h, step=h)
X = [xx for yy in y, xx in x]
Y = [yy for yy in y, xx in x]
z = (sin.(X)) .^10 + cos.(10 .+ Y .* X) + cos.(X) + 0.2*Y + 0.1*X
heatmap(x, y, z, c=:viridis)

# %%
using Plots
h=0.05
y, x = range(-5,  stop=5+h, step=h), range(-5, stop=10+h, step=h)
X = [xx for yy in y, xx in x]
Y = [yy for yy in y, xx in x]
z = (sin.(X)) .^10 + cos.(10 .+ Y .* X) + cos.(X) + 0.2*Y + 0.1*X
heatmap(x, y, z, c=reverse(cgrad(:viridis)))

# %%
using Plots
h=0.05
y, x = range(-5,  stop=5+h, step=h), range(-5, stop=10+h, step=h)
X = [xx for yy in y, xx in x]
Y = [yy for yy in y, xx in x]
z = (sin.(X)) .^10 + cos.(10 .+ Y .* X) + cos.(X) + 0.2*Y + 0.1*X
heatmap(x, y, z, c=cgrad(:viridis, rev=true))

# %%
using Plots
h=0.05
x, y = -5:h:10+h, -5:h:5+h
f(x, y) = sin(x) ^10 + cos(10 + x * y) + cos(x) + 0.2y + 0.1x
heatmap(x, y, f, c=cgrad(:viridis, rev=true))

# %% [markdown]
# You can simply do
#
# ```julia
# heatmap(x, y, z, c=cgrad(:viridis, rev=true))
# ```
#
# or
#
# ```julia
# heatmap(x, y, z, c=reverse(cgrad(:viridis)))
# ```
#
# ref. [ColorGradient](http://docs.juliaplots.org/latest/colorschemes/#ColorGradient)
#
# P.S. More simply,
#
# ```julia
# using Plots
# h=0.05
# x, y = -5:h:10+h, -5:h:5+h
# f(x, y) = sin(x) ^10 + cos(10 + x * y) + cos(x) + 0.2y + 0.1x
# heatmap(x, y, f, c=cgrad(:viridis, rev=true))
# ```
#
# A lot of dots in one line can often cause bugs

# %%

ext: jl

max_stars_repo_path: 0005/cgrad(name, rev=true).jl

max_stars_repo_name: genkuroki/public

max_issues_repo_path: 0005/cgrad(name, rev=true).jl

max_issues_repo_name: genkuroki/public

max_forks_repo_path: 0005/cgrad(name, rev=true).jl

max_forks_repo_name: genkuroki/public


text:
"""
    CareerOneStop::Module

API wraper for the Career One Stop API.
"""
module CareerOneStop
    using HTTP: URI, request
    using JSON3: JSON3
    for (root, dirs, files) in walkdir(joinpath(@__DIR__))
        for file in files
            if occursin("assets", root) || isequal("CareerOneStop.jl", file)
            else
                include(joinpath(root, file))
            end
        end
    end
    export certificationfinder, skillgap
end


text:
# parameters for the model
mutable struct modelParams{t1 <: Real, t2 <: Real, t3 <: Real, t4 <: Real, t5 <: Real, t6 <: Real}
    J::t1 # number of sectors
    N::t2 # number of regions
    R::t3 # number of regions with labor reallocation (may need R_1, R_2 for separate areas)
    T::t4 # number of years
    beta::t5 # discount factor
    ν::t6 # variance of idiosyncratic shock to households
end

# parameters for the algorithm
mutable struct algoParams{t1 <: Real, t2 <: Real, t3 <: Real, t4 <: Real, t5 <: Real, t6 <: Real, t7 <: Real, t8 <: Real}
    tol_out::t1 # solution tolerance for labor loop
    tol_in::t2 # solution tolerance for inner price loop
    maxit::t3 # maximum number of outer loop iterations
    Hmax::t4 # maximum difference between udot guess and actual udot
    iter::t5 # current iteration
    damp::t6 # damping parameter
    howard_start::t7 # iteration to start howard improvement
    howard_its::t8 # number of howard iterations
end

max_stars_repo_name: rafserqui/class-repo

max_issues_repo_name: rafserqui/class-repo

max_forks_repo_name: rafserqui/class-repo


text:
using ELMNIST
using Test

@testset "ELMNIST.jl" begin
    # Write your own tests here.
end


text:
using StaticArrays

struct Sphere{T}
    midpoint::SVector{3,T}
    r::T
end

struct SphereGenerator{Fr}
    r::Fr
end

function (g::SphereGenerator)(midpoint, rot)
    Sphere(midpoint, g.r())
end

volume(shape::Sphere) = 4/3 * π * shape.r^3

function set_midpoint(shape::Sphere, midpoint)
    return Sphere(midpoint, t.r)
end

function bounding_box(shape::Sphere)
    lower = shape.midpoint .- shape.r
    upper = shape.midpoint .+ shape.r
    return lower, upper
end

function inside(shape::Sphere, p0)
    p = p0 - shape.midpoint
    return shape.r^2 < sum(p.^2)
end


text:
#This file is simply a small script to load the enviroment files/types in the correct order

#import needed external libraries
import SPICE: furnsh, pxform, sxform, spkezr, spkpos

#Load Kernels function
LoadKernels(kernels) = furnsh(kernels)

#load PlanetEnviroment and CelestialBodies first
include("PlanetEnvironment.jl")
include("CelestialBodies.jl")
include("StarsAndPlanets.jl")

#Then, load the transformation Quaternions, which requires as input the current body
include("Transformations.jl")
include("Utils.jl")

#Then, load the models (they require the types defined in CelestialBodies.jl)
include("EnvironmentModels/GravityModels/CentralGravity.jl")
include("EnvironmentModels/PlanetModels/SphericalPlanet.jl")

#Load the Instances of the Star, Planets, and Moons
include("CelestialBodies/SolarSystemBarycenter.jl")
include("CelestialBodies/Sun.jl")
include("CelestialBodies/Mercury/MercuryBarycenter.jl")
include("CelestialBodies/Mercury/Mercury.jl")
include("CelestialBodies/Venus/VenusBarycenter.jl")
include("CelestialBodies/Venus/Venus.jl")
include("CelestialBodies/Earth/EarthMoonBarycenter.jl")
include("CelestialBodies/Earth/Earth.jl")
include("CelestialBodies/Earth/Moon.jl")
include("CelestialBodies/Mars/MarsBarycenter.jl")
include("CelestialBodies/Mars/Mars.jl")
include("CelestialBodies/Jupiter/JupiterBarycenter.jl")
include("CelestialBodies/Jupiter/Jupiter.jl")
include("CelestialBodies/Saturn/SaturnBarycenter.jl")
include("CelestialBodies/Saturn/Saturn.jl")
include("CelestialBodies/Uranus/UranusBarycenter.jl")
include("CelestialBodies/Uranus/Uranus.jl")
include("CelestialBodies/Neptune/NeptuneBarycenter.jl")
include("CelestialBodies/Neptune/Neptune.jl")
include("CelestialBodies/Pluto/PlutoBarycenter.jl")
include("CelestialBodies/Pluto/Pluto.jl")

#Load the planetary System type
include("PlanetarySystem.jl")

#Load GravitationalAccelerations
include("GravitationalAcceleration.jl")


text:
module MyLearning

end

max_stars_repo_name: Plant-Form-Function/juliant-template

max_issues_repo_name: Plant-Form-Function/juliant-template

max_forks_repo_name: Plant-Form-Function/juliant-template


text:
@recipe function plot(p::AbstractParticles)
    seriestype --> :histogram
    @series p.particles
end

# @recipe f(::Type{<:AbstractParticles}, p::AbstractParticles) = p.particles # Does not seem to be needed

const RealOrTuple = Union{Real, Tuple}
handle_args(y::AbstractVecOrMat{<:AbstractParticles}, q::RealOrTuple=0.025) = 1:size(y,1), y, q
handle_args(x::AbstractVector, y::AbstractVecOrMat{<:AbstractParticles}, q::RealOrTuple=0.025) = x, y, q
handle_args(p) = handle_args(p.args...)
handle_args(args...) = throw(ArgumentError("The plot function should be called with the signature plotfun([x=1:length(y)], y::Vector{Particles}, [q=0.025])"))

function quantiles(y,q::Number)
    m = vec(mean.(y))
    q > 0.5 && (q = 1-q)
    lower = reshape(-(quantile.(vec(y),q)-m), size(y))
    upper = reshape(quantile.(vec(y),1-q)-m, size(y))
    lower,upper
end

function quantiles(y,q)
    m = vec(mean.(y))
    lower = reshape(-(quantile.(vec(y),q[1])-m), size(y))
    upper = reshape(quantile.(vec(y),q[2])-m, size(y))
    lower,upper
end

@userplot Errorbarplot
@recipe function plt(p::Errorbarplot)
    x,y,q = handle_args(p)
    m = mean.(y)
    label --> "Mean with $q quantile"
    Q = quantiles(y, q)
    if y isa AbstractMatrix
        for c in 1:size(y,2)
            @series begin
                yerror := (Q[1][:,c], Q[2][:,c])
                x,m[:,c]
            end
        end
    else
        yerror := Q
        @series x,m
    end
end

"This is a helper function to make multiple series into one series separated by `Inf`. This makes plotting vastly more efficient."
function to1series(x,y)
    r,c = size(y)
    y2 = vec([y; fill(Inf, 1, c)])
    x2 = repeat([x; Inf], c)
    x2,y2
end

to1series(y) = to1series(1:size(y,1),y)

@userplot MCplot
@recipe function plt(p::MCplot)
    x,y,q = handle_args(p)
    N = nparticles(y)
    selected = q > 1 ? randperm(N)[1:q] : 1:N
    N = length(selected)
    label --> ""
    seriesalpha --> 1/log(N)
    if y isa AbstractMatrix
        for c in 1:size(y,2)
            m = Matrix(y[:,c])'
            @series to1series(x, m[:, selected])
        end
    else
        m = Matrix(y)'
        @series to1series(x, m[:, selected])
    end
end

@userplot Ribbonplot
@recipe function plt(p::Ribbonplot; N=false)
    x,y,q = handle_args(p)
    if N > 0
        for col = 1:size(y,2)
            yc = y[:,col]
            m = mean.(yc)
            @series begin
                label --> "Mean with $q quantile"
                ribbon := quantiles(yc, q)
                x,m
            end
            @series begin
                ribbon := quantiles(yc, q)
                m
            end
            @series begin
                M = Matrix(yc)
                np,ny = size(M)
                primary := false
                nc = N > 1 ? N : min(np, 50)
                seriesalpha --> max(1/sqrt(nc), 0.1)
                chosen = randperm(np)[1:nc]
                to1series(M[chosen, :]')
            end
        end
    else
        @series begin
            label --> "Mean with $q quantile"
            m = mean.(y)
            ribbon := quantiles(y, q)
            x,m
        end
    end
end

"""
    errorbarplot(x,y,[q=0.025])

Plots a vector of particles with error bars at quantile `q`.
If `q::Tuple`, then you can specify both lower and upper quantile, e.g., `(0.01, 0.99)`.
"""
errorbarplot

"""
    mcplot(x,y,[N=0])

Plots all trajectories represented by a vector of particles. `N > 1` controls the number of trajectories to plot.
"""
mcplot

"""
    ribbonplot(x,y,[q=0.025]; N=true)

Plots a vector of particles with a ribbon covering quantiles `q, 1-q`.
If `q::Tuple`, then you can specify both lower and upper quantile, e.g., `(0.01, 0.99)`.

If a positive number `N` is provided, `N` sample trajectories will be plotted on top of the ribbon.
"""
ribbonplot

@recipe function plt(y::Union{MvParticles,AbstractMatrix{<:AbstractParticles}}, q=0.025; N=true)
    label --> "Mean with ($q, $(1-q)) quantiles"
    if N > 0
        for col = 1:size(y,2)
            yc = y[:,col]
            @series begin
                ribbon := quantiles(yc, q)
                mean.(yc)
            end
            @series begin
                M = Matrix(yc)
                np,ny = size(M)
                primary := false
                nc = N > 1 ? N : min(np, 50)
                seriesalpha --> max(1/sqrt(nc), 0.1)
                chosen = randperm(np)[1:nc]
                M[chosen, :]'
            end
        end
    end
end

@recipe function plt(func::Function, x::Union{MvParticles,AbstractMatrix{<:AbstractParticles}}, q=0.025)
    y = func.(x)
    label --> "Mean with ($q, $(1-q)) quantiles"
    xerror := quantiles(x, q)
    yerror := quantiles(y, q)
    mean.(x), mean.(y)
end

@recipe function plt(x::Union{MvParticles,AbstractMatrix{<:AbstractParticles}}, y::Union{MvParticles,AbstractMatrix{<:AbstractParticles}}, q=0.025; points=false)
    my = mean.(y)
    mx = mean.(x)
    if points
        @series begin
            seriestype --> :scatter
            primary := false
            seriesalpha --> 0.1
            Matrix(x), Matrix(y)
        end
        @series begin
            mx, my
        end
    else
        @series begin
            yerror := quantiles(y, q)
            xerror := quantiles(x, q)
            label --> "Mean with $q quantile"
            mx, my
        end
    end
end

@recipe function plt(x::Union{MvParticles,AbstractMatrix{<:AbstractParticles}}, y::AbstractArray, q=0.025)
    mx = mean.(x)
    lower,upper = quantiles(x, q)
    xerror := (lower,upper)
    mx, y
end

@recipe function plt(x::AbstractArray, y::Union{MvParticles,AbstractMatrix{<:AbstractParticles}}, q=0.025; N=true)
    if N > 0
        for col = 1:size(y,2)
            yc = y[:,col]
            @series begin
                ribbon := quantiles(yc, q)
                label --> "Mean with ($q, $(1-q)) quantiles"
                x, mean.(yc)
            end
            @series begin
                M = Matrix(yc)
                np,ny = size(M)
                primary := false
                nc = N > 1 ? N : min(np, 50)
                seriesalpha --> max(1/sqrt(nc), 0.1)
                chosen = randperm(np)[1:nc]
                to1series(x, M[chosen, :]')
            end
        end
    else
        @series begin
            ribbon := quantiles(y, q)
            label --> "Mean with ($q, $(1-q)) quantiles"
            x, mean.(y)
        end
    end
end


text:
facts("Timing functions") do

    context("fpswhen") do
        b = Signal(false)
        t = fpswhen(b, 10)
        acc = foldp((x, y) -> x+1, 0, t)
        sleep(0.14)

        @fact queue_size() --> 0
        push!(b, true)

        dt = @elapsed Reactive.run(11) # the first one starts the timer
        push!(b, false)
        Reactive.run(1)

        sleep(0.11) # no more updates
        @fact queue_size() --> 0

        @fact dt --> roughly(1, atol=0.2) # mac OSX needs a lot of tolerence here
        @fact value(acc) --> 10

    end

    context("fps") do
        t = fps(10)
        acc = foldp(push!, Float64[], t)
        Reactive.run(11) # Starts with 0
        log = copy(value(acc))

        @fact log[1] --> roughly(0.1, atol=0.1) # First one's always crappy
        log = log[2:end]

        @fact sum(log) --> roughly(1.0, atol=0.05)
        @fact [1/10 for i=1:10] --> roughly(log, atol=0.03)

        @fact queue_size() --> 0
        sleep(0.11)
        @fact queue_size() --> 1
        sleep(0.22)
        @fact queue_size() --> 1

        close(acc)
        close(t)

        Reactive.run_till_now()
    end

    context("every") do
        t = every(0.1)
        acc = foldp(push!, Float64[], t)
        Reactive.run(11)
        end_t = time()
        log = copy(value(acc))

        @fact log[end-1] --> roughly(end_t, atol=0.001)

        close(acc)
        close(t)
        Reactive.run_till_now()

        @fact [0.1 for i=1:10] --> roughly(diff(log), atol=0.03)

        sleep(0.2)
        # make sure close actually also closed the timer
        @fact queue_size() --> 0
    end

    context("throttle") do
        x = Signal(0)
        y = throttle(0.1, x)
        y′ = throttle(0.2, x, push!, Int[], x->Int[]) # collect intermediate updates
        z = foldp((acc, x) -> acc+1, 0, y)
        z′ = foldp((acc, x) -> acc+1, 0, y′)

        push!(x, 1)
        step()

        push!(x, 2)
        step()

        push!(x, 3)
        t0=time()
        step()

        @fact value(y) --> 0
        @fact value(z) --> 0
        @fact queue_size() --> 0

        sleep(0.1)

        @fact queue_size() --> 1
        step()
        @fact value(y) --> 3
        @fact value(z) --> 1
        sleep(0.1)

        @fact queue_size() --> 1
        step()
        @fact value(z′) --> 1
        @fact value(y′) --> Int[1,2,3]

        push!(x, 3)
        step()

        push!(x, 2)
        step()

        push!(x, 1)
        step()
        sleep(0.2)

        @fact queue_size() --> 2
        step()
        step()
        @fact value(y) --> 1
        @fact value(z′) --> 2
        @fact value(y′) --> Int[3,2,1]

        # type safety
        s1 = Signal(3)
        s2 = Signal(rand(2,2))
        m = merge(s1, s2)
        t = throttle(1/60, m; typ=Any)
        r = rand(3,3)
        push!(s2, r)
        Reactive.run(1)
        sleep(0.05)
        Reactive.run(1)
        @fact value(t) --> r
    end
end


text:
export PARAMETER_DEFAULT, PARAMETER_UNSTABLE_BUT_FAST, PARAMETER_STABLE_BUT_SLOW

using MathOptInterface
const MOI = MathOptInterface
const AFF = MOI.ScalarAffineFunction{Cdouble}
const EQ = MOI.EqualTo{Cdouble}
const AFFEQ = MOI.ConstraintIndex{AFF,EQ}

mutable struct Optimizer <: MOI.AbstractOptimizer
    objective_constant::Cdouble
    objective_sign::Int
    blockdims::Vector{Int}
    varmap::Vector{Tuple{Int, Int, Int}} # Variable Index vi -> blk, i, j
    b::Vector{Cdouble}
    problem::Union{Nothing, SDPAProblem}
    optimized::Bool
    solve_time::Float64
    silent::Bool
    options::Dict{Symbol, Any}
    function Optimizer(; kwargs...)
		optimizer = new(
            zero(Cdouble), 1, Int[], Tuple{Int, Int, Int}[], Cdouble[],
            nothing, false, NaN, false, Dict{Symbol, Any}())
        if !isempty(kwargs)
            @warn("""Passing optimizer attributes as keyword arguments to
            SDPA.Optimizer is deprecated. Use
                MOI.set(model, MOI.RawOptimizerAttribute("key"), value)
            or
                JuMP.set_optimizer_attribute(model, "key", value)
            instead.
            """)
        end
		for (key, value) in kwargs
			MOI.set(optimizer, MOI.RawOptimizerAttribute(key), value)
		end
		return optimizer
    end
end

varmap(optimizer::Optimizer, vi::MOI.VariableIndex) = optimizer.varmap[vi.value]

function MOI.supports(optimizer::Optimizer, param::MOI.RawOptimizerAttribute)
	return param.name in keys(SET_PARAM)
end
function MOI.set(optimizer::Optimizer, param::MOI.RawOptimizerAttribute, value)
	if !MOI.supports(optimizer, param)
		throw(MOI.UnsupportedAttribute(param))
	end
	optimizer.options[param.name] = value
end
function MOI.get(optimizer::Optimizer, param::MOI.RawOptimizerAttribute)
	# TODO: This gives a poor error message if the name of the parameter is invalid.
	return optimizer.options[param.name]
end

MOI.supports(::Optimizer, ::MOI.Silent) = true
function MOI.set(optimizer::Optimizer, ::MOI.Silent, value::Bool)
	optimizer.silent = value
end
MOI.get(optimizer::Optimizer, ::MOI.Silent) = optimizer.silent

MOI.get(::Optimizer, ::MOI.SolverName) = "SDPA"

# FIXME `hash` needs this, this should be fixed upstream in CxxWrap
Base.hash(p::PhaseType, u::UInt64) = hash(convert(Int32, p), u)

# See https://www.researchgate.net/publication/247456489_SDPA_SemiDefinite_Programming_Algorithm_User's_Manual_-_Version_600
# "SDPA (SemiDefinite Programming Algorithm) User's Manual — Version 6.00" Section 6.2
const RAW_STATUS = Dict(
    noINFO     => "The iteration has exceeded the maxIteration and stopped with no informationon the primal feasibility and the dual feasibility.",
    pdOPT      => "The normal termination yielding both primal and dual approximate optimal solutions.",
    pFEAS      => "The primal problem got feasible but the iteration has exceeded the maxIteration and stopped.",
    dFEAS      => "The dual problem got feasible but the iteration has exceeded the maxIteration and stopped.",
    pdFEAS     => "Both primal problem and the dual problem got feasible, but the iterationhas exceeded the maxIteration and stopped.",
    pdINF      => "At least one of the primal problem and the dual problem is expected to be infeasible.",
    pFEAS_dINF => "The primal problem has become feasible but the dual problem is expected to be infeasible.",
    pINF_dFEAS => "The dual problem has become feasible but the primal problem is expected to be infeasible.",
    pUNBD      => "The primal problem is expected to be unbounded.",
    dUNBD      => "The dual problem is expected to be unbounded.")

function MOI.get(optimizer::Optimizer, ::MOI.RawStatusString)
    if optimizer.problem === nothing
        return "`MOI.optimize!` not called."
    end
	return RAW_STATUS[getPhaseValue(optimizer.problem)]
end
function MOI.get(optimizer::Optimizer, ::MOI.SolveTimeSec)
	return optimizer.solve_time
end

function MOI.is_empty(optimizer::Optimizer)
    return iszero(optimizer.objective_constant) &&
        optimizer.objective_sign == 1 &&
        isempty(optimizer.blockdims) &&
        isempty(optimizer.varmap) &&
        isempty(optimizer.b)
end
function MOI.empty!(optimizer::Optimizer)
    optimizer.objective_constant = zero(Cdouble)
    optimizer.objective_sign = 1
    empty!(optimizer.blockdims)
    empty!(optimizer.varmap)
    empty!(optimizer.b)
    optimizer.problem = nothing
    optimizer.optimized = false
end

function MOI.supports(
    optimizer::Optimizer,
    ::Union{MOI.ObjectiveSense,
            MOI.ObjectiveFunction{AFF}})
    return true
end

MOI.supports_add_constrained_variables(::Optimizer, ::Type{MOI.Reals}) = false

const SupportedSets = Union{MOI.Nonnegatives, MOI.PositiveSemidefiniteConeTriangle}
MOI.supports_add_constrained_variables(::Optimizer, ::Type{<:SupportedSets}) = true
function MOI.supports_constraint(
    ::Optimizer, ::Type{AFF},
    ::Type{EQ})
    return true
end

function _new_block(optimizer::Optimizer, set::MOI.Nonnegatives)
    push!(optimizer.blockdims, -MOI.dimension(set))
    blk = length(optimizer.blockdims)
    for i in 1:MOI.dimension(set)
        push!(optimizer.varmap, (blk, i, i))
    end
end

function _new_block(optimizer::Optimizer, set::MOI.PositiveSemidefiniteConeTriangle)
    push!(optimizer.blockdims, set.side_dimension)
    blk = length(optimizer.blockdims)
    for i in 1:set.side_dimension
        for j in 1:i
            push!(optimizer.varmap, (blk, i, j))
        end
    end
end

function _add_constrained_variables(optimizer::Optimizer,
                                             set::SupportedSets)
    offset = length(optimizer.varmap)
    _new_block(optimizer, set)
    ci = MOI.ConstraintIndex{MOI.VectorOfVariables, typeof(set)}(offset + 1)
    return [MOI.VariableIndex(i) for i in offset .+ (1:MOI.dimension(set))], ci
end

function _error(start, stop)
    error(start, ". Use `MOI.instantiate(SDPA.Optimizer, with_bridge_type = Float64)` ", stop)
end

function constrain_variables_on_creation(
    dest::MOI.ModelLike,
    src::MOI.ModelLike,
    index_map::MOI.Utilities.IndexMap,
    ::Type{S},
) where {S<:MOI.AbstractVectorSet}
    for ci_src in
        MOI.get(src, MOI.ListOfConstraintIndices{MOI.VectorOfVariables,S}())
        f_src = MOI.get(src, MOI.ConstraintFunction(), ci_src)
        if !allunique(f_src.variables)
            _error("Cannot copy constraint `$(ci_src)` as variables constrained on creation because there are duplicate variables in the function `$(f_src)`",
                   "to bridge this by creating slack variables.")
        elseif any(vi -> haskey(index_map, vi), f_src.variables)
            _error("Cannot copy constraint `$(ci_src)` as variables constrained on creation because some variables of the function `$(f_src)` are in another constraint as well.",
                   "to bridge constraints having the same variables by creating slack variables.")
        else
            set = MOI.get(src, MOI.ConstraintSet(), ci_src)::S
            vis_dest, ci_dest = _add_constrained_variables(dest, set)
            index_map[ci_src] = ci_dest
            for (vi_src, vi_dest) in zip(f_src.variables, vis_dest)
                index_map[vi_src] = vi_dest
            end
        end
    end
end

# Loads objective coefficient α * vi
function load_objective_term!(optimizer::Optimizer, α, vi::MOI.VariableIndex)
    blk, i, j = varmap(optimizer, vi)
    coef = optimizer.objective_sign * α
    if i != j
        coef /= 2
    end
    # in SDP format, it is max and in MPB Conic format it is min
    inputElement(optimizer.problem, 0, blk, i, j, float(coef), false)
end

function MOI.copy_to(dest::Optimizer, src::MOI.ModelLike)
    MOI.empty!(dest)
    index_map = MOI.Utilities.IndexMap()

    # Step 1) Compute the dimensions of what needs to be allocated
    constrain_variables_on_creation(
        dest,
        src,
        index_map,
        MOI.Nonnegatives,
    )
    constrain_variables_on_creation(
        dest,
        src,
        index_map,
        MOI.PositiveSemidefiniteConeTriangle,
    )
    vis_src = MOI.get(src, MOI.ListOfVariableIndices())
    if length(vis_src) != length(index_map.var_map)
        _error("Free variables are not supported by SDPA",
               "to bridge free variables into `x - y` where `x` and `y` are nonnegative.")
    end
    cis_src = MOI.get(src, MOI.ListOfConstraintIndices{AFF,EQ}())
    dest.b = Vector{Cdouble}(undef, length(cis_src))
    funcs = Vector{AFF}(undef, length(cis_src))
    for (k, ci_src) in enumerate(cis_src)
        funcs[k] = MOI.get(src, MOI.CanonicalConstraintFunction(), ci_src)
        set = MOI.get(src, MOI.ConstraintSet(), ci_src)
        if !iszero(MOI.constant(funcs[k]))
            throw(MOI.ScalarFunctionConstantNotZero{
                Cdouble, AFF, EQ}(
                    MOI.constant(funcs[k])))
        end
        dest.b[k] = MOI.constant(set)
        index_map[ci_src] = AFFEQ(k)
    end

    # Step 2) Allocate SDPA datastructures
    dummy = isempty(dest.b)
    if dummy
        dest.b = [one(Cdouble)]
        dest.blockdims = [dest.blockdims; -1]
    end
    dest.problem = SDPAProblem()
    setParameterType(dest.problem, PARAMETER_DEFAULT)
	# TODO Take `silent` into account here
    setparameters!(dest.problem, dest.options)
    inputConstraintNumber(dest.problem, length(dest.b))
    inputBlockNumber(dest.problem, length(dest.blockdims))
    for (i, blkdim) in enumerate(dest.blockdims)
        inputBlockSize(dest.problem, i, blkdim)
        inputBlockType(dest.problem, i, blkdim < 0 ? LP : SDP)
    end
    initializeUpperTriangleSpace(dest.problem)
    for i in eachindex(dest.b)
        inputCVec(dest.problem, i, dest.b[i])
    end
    if dummy
        inputElement(dest.problem, 1, length(dest.blockdims), 1, 1, one(Cdouble), false)
    end

    # Step 3) Load data in the datastructures
    for k in eachindex(funcs)
        for term in funcs[k].terms
            if !iszero(term.coefficient)
                blk, i, j = varmap(dest, index_map[term.variable])
                coef = term.coefficient
                if i != j
                    coef /= 2
                end
                inputElement(dest.problem, k, blk, i, j, float(coef), false)
            end
        end
    end

    MOI.Utilities.pass_attributes(dest, src, index_map, vis_src)
    # Throw error for constraint attributes
    MOI.Utilities.pass_attributes(dest, src, index_map, cis_src)

    # Pass objective attributes and throw error for other ones
    model_attributes = MOI.get(src, MOI.ListOfModelAttributesSet())
    for attr in model_attributes
        if attr != MOI.ObjectiveSense() && attr != MOI.ObjectiveFunction{AFF}()
            throw(MOI.UnsupportedAttribute(attr))
        end
    end
    # We make sure to set `objective_sign` first before setting the objective
    if MOI.ObjectiveSense() in model_attributes
        sense = MOI.get(src, MOI.ObjectiveSense())
        dest.objective_sign = sense == MOI.MIN_SENSE ? -1 : 1
    end
    if MOI.ObjectiveFunction{AFF}() in model_attributes
        func = MOI.get(src, MOI.ObjectiveFunction{AFF}())
        obj = MOI.Utilities.canonical(func)
        dest.objective_constant = obj.constant
        for term in obj.terms
            if !iszero(term.coefficient)
                load_objective_term!(dest, term.coefficient, index_map[term.variable])
            end
        end
    end
    return index_map
end

function MOI.optimize!(m::Optimizer)
    # SDPA segfault if `optimize!` is called twice
    if !m.optimized
        start_time = time()
        SDPA.initializeUpperTriangle(m.problem, false)
        SDPA.initializeSolve(m.problem)
        SDPA.solve(m.problem)
        m.solve_time = time() - start_time
        m.optimized = true
    end
end

function MOI.get(m::Optimizer, ::MOI.TerminationStatus)
    if m.problem === nothing
        return MOI.OPTIMIZE_NOT_CALLED
    end
    status = getPhaseValue(m.problem)
    if status == noINFO
        return MOI.ITERATION_LIMIT
    elseif status == pFEAS
        return MOI.SLOW_PROGRESS
    elseif status == dFEAS
        return MOI.SLOW_PROGRESS
    elseif status == pdFEAS
        return MOI.OPTIMAL
    elseif status == pdINF
        return MOI.INFEASIBLE_OR_UNBOUNDED
    elseif status == pFEAS_dINF
        return MOI.DUAL_INFEASIBLE
    elseif status == pINF_dFEAS
        return MOI.INFEASIBLE
    elseif status == pdOPT
        return MOI.OPTIMAL
    elseif status == pUNBD
        return MOI.DUAL_INFEASIBLE
    elseif status == dUNBD
        return MOI.INFEASIBLE
    end
end

function MOI.get(m::Optimizer, attr::MOI.PrimalStatus)
    if attr.result_index > MOI.get(m, MOI.ResultCount())
        return MOI.NO_SOLUTION
    end
    status = getPhaseValue(m.problem)
    if status == noINFO
        return MOI.UNKNOWN_RESULT_STATUS
    elseif status == pFEAS
        return MOI.FEASIBLE_POINT
    elseif status == dFEAS
        return MOI.UNKNOWN_RESULT_STATUS
    elseif status == pdFEAS
        return MOI.FEASIBLE_POINT
    elseif status == pdINF
        return MOI.UNKNOWN_RESULT_STATUS
    elseif status == pFEAS_dINF
        return MOI.INFEASIBILITY_CERTIFICATE
    elseif status == pINF_dFEAS
        return MOI.INFEASIBLE_POINT
    elseif status == pdOPT
        return MOI.FEASIBLE_POINT
    elseif status == pUNBD
        return MOI.INFEASIBILITY_CERTIFICATE
    elseif status == dUNBD
        return MOI.INFEASIBLE_POINT
    end
end

function MOI.get(m::Optimizer, attr::MOI.DualStatus)
    if attr.result_index > MOI.get(m, MOI.ResultCount())
        return MOI.NO_SOLUTION
    end
    status = getPhaseValue(m.problem)
    if status == noINFO
        return MOI.UNKNOWN_RESULT_STATUS
    elseif status == pFEAS
        return MOI.UNKNOWN_RESULT_STATUS
    elseif status == dFEAS
        return MOI.FEASIBLE_POINT
    elseif status == pdFEAS
        return MOI.FEASIBLE_POINT
    elseif status == pdINF
        return MOI.UNKNOWN_RESULT_STATUS
    elseif status == pFEAS_dINF
        return MOI.INFEASIBLE_POINT
    elseif status == pINF_dFEAS
        return MOI.INFEASIBILITY_CERTIFICATE
    elseif status == pdOPT
        return MOI.FEASIBLE_POINT
    elseif status == pUNBD
        return MOI.INFEASIBLE_POINT
    elseif status == dUNBD
        return MOI.INFEASIBILITY_CERTIFICATE
    end
end

MOI.get(m::Optimizer, ::MOI.ResultCount) = m.problem === nothing ? 0 : 1
function MOI.get(m::Optimizer, attr::MOI.ObjectiveValue)
    MOI.check_result_index_bounds(m, attr)
    return m.objective_sign * getPrimalObj(m.problem) + m.objective_constant
end
function MOI.get(m::Optimizer, attr::MOI.DualObjectiveValue)
    MOI.check_result_index_bounds(m, attr)
    return m.objective_sign * getDualObj(m.problem) + m.objective_constant
end
struct PrimalSolutionMatrix <: MOI.AbstractModelAttribute end
MOI.is_set_by_optimize(::PrimalSolutionMatrix) = true
MOI.get(optimizer::Optimizer, ::PrimalSolutionMatrix) = PrimalSolution(optimizer.problem)

struct DualSolutionVector <: MOI.AbstractModelAttribute end
MOI.is_set_by_optimize(::DualSolutionVector) = true
function MOI.get(optimizer::Optimizer, ::DualSolutionVector)
    return unsafe_wrap(Array, getResultXVec(optimizer.problem).cpp_object, getConstraintNumber(optimizer.problem))
end

struct DualSlackMatrix <: MOI.AbstractModelAttribute end
MOI.is_set_by_optimize(::DualSlackMatrix) = true
MOI.get(optimizer::Optimizer, ::DualSlackMatrix) = VarDualSolution(optimizer.problem)

function block(optimizer::Optimizer, ci::MOI.ConstraintIndex{MOI.VectorOfVariables})
    return optimizer.varmap[ci.value][1]
end
function dimension(optimizer::Optimizer, ci::MOI.ConstraintIndex{MOI.VectorOfVariables})
    blockdim = optimizer.blockdims[block(optimizer, ci)]
    if blockdim < 0
        return -blockdim
    else
        return MOI.dimension(MOI.PositiveSemidefiniteConeTriangle(blockdim))
    end
end
function vectorize_block(M, blk::Integer, s::Type{MOI.Nonnegatives})
    return diag(block(M, blk))
end
function vectorize_block(M::AbstractMatrix{Cdouble}, blk::Integer, s::Type{MOI.PositiveSemidefiniteConeTriangle}) where T
    B = block(M, blk)
    d = LinearAlgebra.checksquare(B)
    n = MOI.dimension(MOI.PositiveSemidefiniteConeTriangle(d))
    v = Vector{Cdouble}(undef, n)
    k = 0
    for j in 1:d
        for i in 1:j
            k += 1
            v[k] = B[i, j]
        end
    end
    @assert k == n
    return v
end

function MOI.get(optimizer::Optimizer, attr::MOI.VariablePrimal, vi::MOI.VariableIndex)
    MOI.check_result_index_bounds(optimizer, attr)
    blk, i, j = varmap(optimizer, vi)
    return block(MOI.get(optimizer, PrimalSolutionMatrix()), blk)[i, j]
end

function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintPrimal,
                 ci::MOI.ConstraintIndex{MOI.VectorOfVariables, S}) where S<:SupportedSets
    MOI.check_result_index_bounds(optimizer, attr)
    return vectorize_block(MOI.get(optimizer, PrimalSolutionMatrix()), block(optimizer, ci), S)
end
function MOI.get(m::Optimizer, attr::MOI.ConstraintPrimal, ci::AFFEQ)
    MOI.check_result_index_bounds(m, attr)
    return m.b[ci.value]
end

function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintDual,
                 ci::MOI.ConstraintIndex{MOI.VectorOfVariables, S}) where S<:SupportedSets
    MOI.check_result_index_bounds(optimizer, attr)
    return vectorize_block(MOI.get(optimizer, DualSlackMatrix()), block(optimizer, ci), S)
end
function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintDual, ci::AFFEQ)
    MOI.check_result_index_bounds(optimizer, attr)
    return -MOI.get(optimizer, DualSolutionVector())[ci.value]
end


text:
module Day19
__revise_mode__ = :eval

# day 2 not solved

using Advent: save_input

function validstr(line, dict, rule)
    isempty(line) && return (true, "")
    rule isa Char && return (line[1] == rule ? (true, line[2:end]) : (false, line))
    res = line
    for r in dict[rule]
        res = line
        k = true
        for pos in r
            (ok, res) = validstr(res, dict, pos)
            k = ok
            res == line && (ok = false; break)
        end
        k == true && return (true, res)
    end
    return (false, line)
end

function solve_a(filename)
    open(filename) do io
        rules = Dict{Int, Vector{Vector{Union{Char,Int}}}}()
        for line in eachline(io)
            line == "" && break
            num, vals = split(line, ": ")
            vals = split(vals, " | ")
            rules[parse(Int, num)] = [((v[1][1] == '"') ? map(x -> x[2], v) : parse.(Int, v)) for v in split.(vals)]
        end

        count(eachline(io)) do line
            empty!(rcounts)
            x, y = validstr(line, rules, 0)
            x && isempty(y)
        end
    end
end

function solve_b(filename)
end

solve(filename = save_input(2020, 19)) = solve_a(filename), solve_b(filename)

end

max_stars_repo_name: AnAverageHuman/competitive

max_issues_repo_name: AnAverageHuman/competitive

max_forks_repo_name: AnAverageHuman/competitive


text:
using StructIO
using Test

# First, exercise the `@io` macro a bit, to ensure it can handle different
# kinds of type declarations
@io struct TwoUInts
    x::UInt
    y::UInt
end

abstract type AbstractType end
@io struct ConcreteType <: AbstractType
    A::UInt32
    B::UInt16
    C::UInt128
    D::UInt8
end align_packed

@io struct PackedNestedType
    A::ConcreteType
    B::ConcreteType
end align_packed

@io struct DefaultNestedType
    A::ConcreteType
    B::ConcreteType
end

@io struct PackedParametricType{T}
    x::T
    y::T
end align_packed

# Also test documenting a type
"""
This is a docstring
"""
@io struct ParametricType{S,T}
    A::S
    B::T
    C::T
end

@testset "unpack()" begin
    # Test native unpacking
    buf = IOBuffer()
    write(buf, UInt(1))
    write(buf, UInt(2))
    seekstart(buf)
    @test unpack(buf, TwoUInts) == TwoUInts(1,2)

    # Now, test explicitly setting the endianness
    for endian in [:LittleEndian, :BigEndian]
        buf = IOBuffer()
        write(buf, fix_endian(UInt32(1), endian))
        write(buf, fix_endian(UInt16(2), endian))
        write(buf, fix_endian(UInt128(3), endian))
        write(buf, fix_endian(UInt8(4),  endian))
        seekstart(buf)

        @test unpack(buf, ConcreteType, endian) == ConcreteType(1, 2, 3, 4)
    end

    # Test packed nested types across endianness
    for endian in [:LittleEndian, :BigEndian]
        buf = IOBuffer()
        write(buf, fix_endian(UInt32(1), endian))
        write(buf, fix_endian(UInt16(2), endian))
        write(buf, fix_endian(UInt128(3), endian))
        write(buf, fix_endian(UInt8(4), endian))
        write(buf, fix_endian(UInt32(5), endian))
        write(buf, fix_endian(UInt16(6), endian))
        write(buf, fix_endian(UInt128(7), endian))
        write(buf, fix_endian(UInt8(8), endian))
        seekstart(buf)

        x = PackedNestedType(ConcreteType(1,2,3,4), ConcreteType(5,6,7,8))
        @test unpack(buf, PackedNestedType, endian) == x
    end

    # Test mixed Packed/Default nested types across endianness
    for endian in [:BigEndian, :LittleEndian]
        # Helper function to write a value, then write zeros afterward to build
        # a stream that mocks up a `Default` packing strategy memory layout
        function write_skip(buf, x, field_idx)
            n_written = write(buf, fix_endian(x, endian))

            n_size = Int32(StructIO.fieldsize(ConcreteType, field_idx))
            write(buf, zeros(UInt8, n_size - n_written))
        end

        buf = IOBuffer()
        write_skip(buf, UInt32(1), 1)
        write_skip(buf, UInt16(2), 2)
        write_skip(buf, UInt128(3), 3)
        write_skip(buf, UInt8(4), 4)

        write_skip(buf, UInt32(5), 1)
        write_skip(buf, UInt16(6), 2)
        write_skip(buf, UInt128(7), 3)
        write_skip(buf, UInt8(8), 4)
        seekstart(buf)

        x = DefaultNestedType(ConcreteType(1,2,3,4), ConcreteType(5,6,7,8))
        @test unpack(buf, DefaultNestedType, endian) == x
    end
end

@testset "pack()" begin
    # Pack simple types
    for endian in [:BigEndian, :LittleEndian]
        buf = IOBuffer()
        pack(buf, UInt8(1), endian)
        pack(buf, Int16(2), endian)
        pack(buf, UInt32(4), endian)
        pack(buf, Int64(8), endian)
        pack(buf, UInt128(16), endian)
        @test position(buf) == 1 + 2 + 4 + 8 + 16
        seekstart(buf)
        @test unpack(buf, UInt8, endian) === UInt8(1)
        @test unpack(buf, Int16, endian) === Int16(2)
        @test unpack(buf, UInt32, endian) === UInt32(4)
        @test unpack(buf, Int64, endian) === Int64(8)
        @test unpack(buf, UInt128, endian) === UInt128(16)
    end

    # Pack a simple object
    buf = IOBuffer()
    tu = TwoUInts(2, 3)
    pack(buf, tu)

    # Test that the stream looks reasonable
    @test position(buf) == sizeof(TwoUInts)
    seekstart(buf)
    @test read(buf, UInt) == 2
    @test read(buf, UInt) == 3

    # Test that we can unpack a packed stream
    buf = IOBuffer()
    pack(buf, tu)
    seekstart(buf)
    @test unpack(buf, TwoUInts) == TwoUInts(2, 3)

    # Test packed/default nested types across endianness
    for NT in [PackedNestedType, DefaultNestedType]
        for endian in [:LittleEndian, :BigEndian]
            buf = IOBuffer()
            nt = NT(ConcreteType(1,2,3,4), ConcreteType(5,6,7,8))
            pack(buf, nt)
            seekstart(buf)
            @test unpack(buf, NT) == nt
        end
    end
end

@testset "packed_sizeof()" begin
    @test packed_sizeof(TwoUInts) == 2*sizeof(UInt)
    @test packed_sizeof(ConcreteType) == 1 + 2 + 4 + 16
    @test packed_sizeof(PackedNestedType) == 2*packed_sizeof(ConcreteType)
    @test packed_sizeof(PackedParametricType{UInt8}) == 2
    @test packed_sizeof(PackedParametricType{UInt32}) == 8
    psCT = packed_sizeof(ConcreteType)
    @test packed_sizeof(PackedParametricType{ConcreteType}) == 2*psCT
end

@testset "Documentation" begin
    @test string(@doc ParametricType) == "This is a docstring\n"
end


text:
function clean_data(file)
    df = CSV.read(file)

    districts = unique(df.district)
    dname = format_string.(districts)

    dataset = Dict{String, Dict{String, Union{Dict{String, Vector{Float64}}, Nothing}}}()

    for (d, dn) in zip(districts, dname)
        if occursin("other", dn) || occursin("foreign", dn)
            continue
        end
        dsplit = df[df.district .== d, :]
        states = unique(dsplit.state)
        for state in states
            dsplit_2 = dsplit[dsplit.state .== state, :]

            state = format_string(state)
            if state ∉ keys(dataset)
                dataset[state] = Dict{String, Union{Dict{String, Vector{Float64}}, Nothing}}()
            end

            infected = Array(dsplit_2[dsplit_2.category .== "confirmed", :])[4:end]
            recovered = Array(dsplit_2[dsplit_2.category .== "recovered", :])[4:end]
            deceased = Array(dsplit_2[dsplit_2.category .== "deceased", :])[4:end]

            infected = infected .- recovered .- deceased

            idx = findfirst((infected .> 10) .& (.!ismissing.(infected)))
            if isnothing(idx)
                dataset[state][dn] = nothing
            else
                dataset[state][dn] = Dict{String, Vector{Float64}}()
                time = []
                infect = []
                recov = []
                dead = []
                test_time = Float64.([length(infected) - idx + 1 + i for i in 1:7])
                infected = infected[idx:end]
                recovered = recovered[idx:end]
                deceased = deceased[idx:end]
                for (i, (inf, rec, dec)) in enumerate(zip(infected, recovered, deceased))
                    if all(.!ismissing.([inf, rec, dec]))
                        append!(time, Float64(i))
                        append!(infect, inf)
                        append!(recov, rec)
                        append!(dead, dec)
                    end
                end
                dataset[state][dn]["infected"] = Float64.(infect)
                dataset[state][dn]["recovered"] = Float64.(recov)
                dataset[state][dn]["deceased"] = Float64.(dead)
                dataset[state][dn]["time_train"] = Float64.(time)
                dataset[state][dn]["time_test"] = Float64.(test_time)
            end
        end
    end

    return dataset
end


text:
#Store details of the optical system
mutable struct SAIMOptics
	#Required parameters for standard SAIM experiments with single excitation
	nB::Float64		        	#The refractive index of the ambient buffer / cytoplasm
	dOx::Float64		        #The thickness of the SiO2 layer in units of nm
	λ_Ex_1::Float64		        #The wavelength of excitation light #1 in units of nm
	nOx_1::Float64		        #The refractive index of SiO2 for excitation #1
	nSi_1::Complex{Float64}	    #The complex refractive index of Si for excitation #1

	#Additional parameters for two-color and three-color SAIM; these parameters are
	#not used for single excitation experiments
	λ_Ex_2::Float64		        #The wavelength of excitation light #2 in units of nm
	nOx_2::Float64		        #The refractive index of SiO2 for excitation #2
	nSi_2::Complex{Float64}	    #The complex refractive index of Si for excitation #2
	λ_Ex_3::Float64		        #The wavelength of excitation light #3 in units of nm
	nOx_3::Float64		        #The refractive index of SiO2 for excitation #3
	nSi_3::Complex{Float64}	    #The complex refractive index of Si for excitation #3

	#Constructor for the structure with all variaable initialized to zero
	function SAIMOptics()
		this = new(0.0, 0.0, 0.0, 0.0, 0.0+0.0im, 0.0, 0.0, 0.0+0.0im, 0.0, 0.0, 0.0+0.0im)
	end
end


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

const Stack = Main.TestGCExt.Stack

function make()
    ccall(:stk_make, Stack, ())
end

function push(stack :: Stack, val :: String)
    ccall(:stk_push, Nothing, (Stack, String), stack, val)
end

function top(stack :: Stack)
    return ccall(:stk_top, String, (Stack,), stack)
end

function pop(stack :: Stack)
    return ccall(:stk_pop, String, (Stack,), stack)
end

function size(stack :: Stack)
    return ccall(:stk_size, UInt, (Stack,), stack)
end

function empty(stack :: Stack)
    return size(stack) == 0
end

function blob(stack :: Stack)
    return ccall(:stk_blob, Any, (Stack,), stack)
end

function gc_counter_full()
    return ccall(:get_gc_counter, UInt, (Cint,), 1)
end

function gc_counter_inc()
    return ccall(:get_gc_counter, UInt, (Cint,), 0)
end

function gc_counter()
    return gc_counter_full() + gc_counter_inc()
end

function num_obj_sweeps()
    return ccall(:get_obj_sweeps, UInt, ())
end

function get_aux_root(n :: Int)
    return ccall(:get_aux_root, String, (UInt,), n)
end

function set_aux_root(n :: Int, x :: String)
    return ccall(:set_aux_root, Nothing, (UInt, String), n, x)
end

function internal_obj_scan(p :: Any)
    if ccall(:internal_obj_scan, Cint, (Any,), p) == 0
        global internal_obj_scan_failures += 1
    end
end

global internal_obj_scan_failures = 0

for i in 0:1000
    set_aux_root(i, string(i))
end

function test()
    local stack = make()
    for i in 1:100000
        push(stack, string(i, base=2))
        internal_obj_scan(top(stack))
    end
    for i in 1:1000
        local stack2 = make()
        internal_obj_scan(stack2)
        internal_obj_scan(blob(stack2))
        while !empty(stack)
            push(stack2, pop(stack))
        end
        stack = stack2
        if i % 100 == 0
            GC.gc()
        end
    end
end

@time test()

global corrupted_roots = 0
for i in 0:1000
    if get_aux_root(i) != string(i)
        global corrupted_roots += 1
    end
end

print(gc_counter_full(), " full collections.\n")
print(gc_counter_inc(), " partial collections.\n")
print(num_obj_sweeps(), " object sweeps.\n")
print(internal_obj_scan_failures, " internal object scan failures.\n")
print(corrupted_roots, " corrupted auxiliary roots.\n")

max_stars_repo_name: greimel/julia

max_issues_repo_name: greimel/julia

max_forks_repo_name: greimel/julia


text:
module ElasticPDMats
import Base: size, getindex, setindex!, append!, show, view, deleteat!
import LinearAlgebra: mul!, ldiv!
using LinearAlgebra, MacroTools
using PDMats
import PDMats: dim, Matrix, diag, pdadd!, *, \, inv, logdet, eigmax, eigmin, whiten!, unwhiten!, quad, quad!, invquad, invquad!, X_A_Xt, Xt_A_X, X_invA_Xt, Xt_invA_X
export ElasticPDMat, AllElasticArray, ElasticSymmetricMatrix, ElasticCholesky, setcapacity!, setstepsize!, setdimension!

mutable struct AllElasticArray{T, N} <: AbstractArray{T, N}
    dims::Tuple{Vararg{Int, N}}
    capacity::Tuple{Vararg{Int, N}}
    stepsize::Tuple{Vararg{Int, N}}
    data::Array{T, N}
end
function AllElasticArray(m::AbstractArray{T, N}; dims = size(m),
                         capacity = fill(10^2, N),
                         stepsize = fill(10^2, N)) where {T, N}
    AllElasticArray(dims, Tuple(max.(dims, capacity)), Tuple(max.(dims, stepsize)), m)
end
AllElasticArray(N; capacity = tuple(fill(10^2, N)...), stepsize = tuple(fill(10^2, N)...)) =
AllElasticArray(zeros(capacity...), dims = tuple(fill(0, N)...), capacity = capacity,
                    stepsize = stepsize)

size(m::AllElasticArray) = m.dims
for i in 2:3
    vars = [Symbol(:i, j) for j in 1:i]
    @eval begin
        getindex(m::AllElasticArray{T, $i}, $(vars...)) where {T} = getindex(m.data, $(vars...))
        setindex!(m::AllElasticArray{T, $i}, v::T, $(vars...)) where {T} = setindex!(m.data,v, $(vars...))
        view(m::AllElasticArray{T, $i}, $(vars...)) where {T} = view(m.data, $(vars...))
        view(m::AllElasticArray{T, $i}) where {T} = view(m, $([:(UnitRange(1, m.dims[$j])) for j in 1:i]...))
        setdimension!(m::AllElasticArray{T, $i}, v::Int, k::Int) where {T} = m.dims = tuple($([:(k == $j ? v : m.dims[$j]) for j in 1:i]...))
    end
end
getindex(m::AllElasticArray{T, N}, i, j, k, l, I...) where {T, N} = getindex(m.data, i, j, k, l, I...)
view(m::AllElasticArray{T, N}) where {T, N} = view(m, (UnitRange.(1, m.dims))...)
setindex!(m::AllElasticArray{T, N}, I::Vararg{Int, N}) where {T, N} = setindex!(m.data, I...)
setdimension!(m::AllElasticArray{T, N}, v::Int, k::Int) where {T, N} = m.dims = tuple([k == j ? v : m.dims[j] for j in 1:N]...)
setdimension!(m::AllElasticArray, v::Int, k::AbstractArray{Int, 1}) = for i in k setdimension!(m, v, i) end
function grow!(obj::AllElasticArray)
    obj.capacity = obj.capacity .+ obj.stepsize
    resize!(obj)
end
resize!(obj::AllElasticArray) = resize!(obj, obj.data)
function resize!(obj::AllElasticArray, data::AbstractArray{T, N}) where {T, N}
    tmp = zeros(T, obj.capacity...)
    ind = CartesianIndices(Tuple(UnitRange.(1, obj.dims)))
    copyto!(tmp, ind, data, ind)
    obj.data = tmp
    obj.capacity
end

mul!(Y::AbstractArray{T, 1}, M::AllElasticArray{T, 2}, V::AbstractArray{T, 1}) where {T} = mul!(Y, view(M), V)
mul!(Y::AbstractArray{T, 2}, M::AllElasticArray{T, 2}, V::AbstractArray{T, 2}) where {T} = mul!(Y, view(M), V)
mul!(Y::AbstractArray{T, 2}, M1::AllElasticArray{T, 2}, M2::AllElasticArray{T, 2}) where {T} = mul!(Y, view(M1), view(M2))

struct ElasticSymmetricMatrix{T} <: AbstractArray{T, 2}
    m::AllElasticArray{T, 2}
end
function ElasticSymmetricMatrix(m::AbstractArray{T, 2}; N = size(m, 1), capacity = 10^3, stepsize = 10^3) where {T}
    !issymmetric(m) && error("Data is not symmetric.")
    data = zeros(T, capacity, capacity)
    ind = CartesianIndices((1:N, 1:N))
    copyto!(data, ind, m, ind)
    ElasticSymmetricMatrix(AllElasticArray((N, N), (capacity, capacity), (stepsize, stepsize), data))
end
ElasticSymmetricMatrix(; capacity = 10^3, stepsize = 10^3) = ElasticSymmetricMatrix(AllElasticArray((0, 0), (capacity, capacity), (stepsize, stepsize), zeros(capacity, capacity)))

# from Lazy.jl
macro forward(ex, fs)
  @capture(ex, T_.field_) || error("Syntax: @forward T.x f, g, h")
  T = esc(T)
  fs = isexpr(fs, :tuple) ? map(esc, fs.args) : [esc(fs)]
  :($([:($f(x::$T, args...) = (Base.@_inline_meta; $f(x.$field, args...)))
       for f in fs]...);
    nothing)
end
@forward ElasticSymmetricMatrix.m size, getindex, setindex!, view, mul!, grow!, resize!, setnewdata!
setdimension!(e::ElasticSymmetricMatrix, N::Int) = setdimension!(e.m, N, 1:2)
function append!(g::ElasticSymmetricMatrix{T}, data::AbstractArray{T, 2}) where {T}
    n, m = size(data)
    oldm = size(g, 1)
    oldm + m > g.m.capacity[1] && grow!(g)
    gd = g.m.data
    @inbounds @simd for j in 1:m
        for i in 1:min(oldm + j, n)
            gd[i, j + oldm] = gd[j + oldm, i] = data[i, j]
        end
    end
    setdimension!(g, oldm + m)
    g
end
function deleteat!(g::ElasticSymmetricMatrix, i::Int)
    N = size(g, 1)
    gd = g.m.data
    copyto!(gd, CartesianIndices((1:N, i:N - 1)),
            gd, CartesianIndices((1:N, i+1:N)))
    copyto!(gd, CartesianIndices((i:N-1, 1:N)),
            gd, CartesianIndices((i+1:N, 1:N)))
    setdimension!(g, N - 1)
    g
end
function setcapacity!(x::ElasticSymmetricMatrix, c::Int)
    x.m.capacity = (c, c)
    resize!(x)
end
setstepsize!(x::ElasticSymmetricMatrix, c::Int) = x.m.stepsize = (c, c)

mutable struct ElasticCholesky{T, A} <: Factorization{T}
    N::Int
    capacity::Int
    stepsize::Int
    c::Cholesky{T, A}
end
function ElasticCholesky(c::Cholesky{T, A}; capacity = 10^3, stepsize = 10^3) where {T, A}
    N = size(c, 1)
    data = zeros(T, capacity, capacity)
    ind = CartesianIndices((1:N, 1:N))
    copyto!(data, ind, c.factors, ind)
    ElasticCholesky(N, capacity, stepsize, Cholesky(data, 'U', LinearAlgebra.BlasInt(0)))
end
ElasticCholesky(; capacity = 10^3, stepsize = 10^3) = ElasticCholesky(0, capacity, stepsize, Cholesky(zeros(capacity, capacity), 'U', LinearAlgebra.BlasInt(0)))

function setcapacity!(x::ElasticCholesky, c::Int)
    x.capacity = c
    resize!(x)
end
setstepsize!(x::ElasticCholesky, c::Int) = x.stepsize = c

view(c::ElasticCholesky, i, j) = Cholesky(view(c.c.factors, i, j), c.c.uplo, c.c.info)
view(c::ElasticCholesky) = view(c, 1:c.N, 1:c.N)
show(io::IO, m::MIME{Symbol("text/plain")}, c::ElasticCholesky) = show(io, m, view(c))
size(c::ElasticCholesky) = (c.N, c.N)
size(c::ElasticCholesky, i::Int) = c.N

ldiv!(c::ElasticCholesky, x) = ldiv!(view(c), x)

function grow!(c::ElasticCholesky)
    c.capacity += c.stepsize
    resize!(c)
end
resize!(c::ElasticCholesky) = resize!(c, c.c.factors)
function resize!(obj::ElasticCholesky, data::AbstractArray{T, N}) where {T, N}
    tmp = zeros(T, obj.capacity, obj.capacity)
    ind = CartesianIndices((1:obj.N, 1:obj.N))
    copyto!(tmp, ind, data, ind)
    obj.c = Cholesky(tmp, 'U', LinearAlgebra.BlasInt(0))
    obj.capacity
end

append!(c::Union{ElasticCholesky{T, A}, ElasticSymmetricMatrix{T}}, data::Vector{T}) where {T, A} = append!(c, reshape(data, :, 1))
function append!(c::ElasticCholesky{T,A}, data::A) where {T, A}
    n, m = size(data)
    c.N + m > c.capacity && grow!(c)
    s = data[1:c.N, 1:m]
    LAPACK.trtrs!('U', 'C', 'N', view(c).factors, s)
    colrange = c.N + 1:c.N + m
    copyto!(c.c.factors, CartesianIndices((1:c.N, colrange)), s, CartesianIndices((1:c.N, 1:m)))
    copyto!(c.c.factors, CartesianIndices((colrange, colrange)), view(data, colrange,1:m), CartesianIndices((1:m, 1:m)))
    BLAS.syrk!('U', 'T', -1., s, 1., view(c.c.factors, colrange, colrange))
    LinearAlgebra._chol!(view(c.c.factors, colrange, colrange), UpperTriangular)
    c.N += m
    c
end
# TODO: Check if this can be optimized.
function deleteat!(c::ElasticCholesky, i::Int)
    R = view(c.c.factors, i, i+1:c.N) * view(c.c.factors, i, i+1:c.N)'
    R += view(c, i+1:c.N, i+1:c.N).U' * view(c, i+1:c.N, i+1:c.N).U
    cholesky!(Symmetric(R))
    copyto!(c.c.factors, CartesianIndices((1:i-1, i:c.N-1)),
            c.c.factors, CartesianIndices((1:i-1, i+1:c.N)))
    copyto!(c.c.factors, CartesianIndices((i:c.N-1, i:c.N-1)),
            R, CartesianIndices((1:c.N-i, 1:c.N-i)))
    c.N -= 1
    c
end

struct ElasticPDMat{T, A} <: AbstractPDMat{T}
    mat::ElasticSymmetricMatrix{T}
    chol::ElasticCholesky{T, A}
end
"""
    ElasticPDMat([m [, chol]]; capacity = 10^3, stepsize = 10^3)

Creates an elastic positive definite matrix with initial `capacity = 10^3` and
`stepsize = 10^3`. The optional argument `m` is a positive definite, symmetric
matrix and `chol` its cholesky decomposition. Use `append!` and `deleteat!` to
change an ElasticPDMat.
"""
ElasticPDMat(; kwargs...) = ElasticPDMat(ElasticSymmetricMatrix(; kwargs...), ElasticCholesky(; kwargs...))
ElasticPDMat(m; kwargs...) = ElasticPDMat(m, cholesky(m); kwargs...)
function ElasticPDMat(m, chol; kwargs...)
    ElasticPDMat(ElasticSymmetricMatrix(m; kwargs...),
                 ElasticCholesky(chol; kwargs...))
end

function setcapacity!(x::ElasticPDMat, c::Int)
    setcapacity!(x.mat, c)
    setcapacity!(x.chol, c)
end
function setstepsize!(x::ElasticPDMat, c::Int)
    setstepsize!(x.mat, c)
    setstepsize!(x.chol, c)
end

function append!(a::ElasticPDMat, data)
    append!(a.mat, data)
    append!(a.chol, data)
    a
end
function deleteat!(a::ElasticPDMat, i::Int)
    deleteat!(a.mat, i)
    deleteat!(a.chol, i)
    a
end
# TODO: more efficient blockwise deletion
function deleteat!(g::Union{ElasticSymmetricMatrix, ElasticCholesky, ElasticPDMat}, idxs::AbstractArray{Int, 1})
    map(i -> deleteat!(g, i), sort(idxs, rev = true))
end

dim(a::ElasticPDMat) = size(a.mat, 1)
Base.Matrix(a::ElasticPDMat) = Matrix(view(a.mat))
LinearAlgebra.diag(a::ElasticPDMat) = diag(view(a.mat))
function pdadd!(r::Matrix, a::Matrix, gb::ElasticPDMat, c::Real)
    b = view(gb.mat)
    PDMats.@check_argdims size(r) == size(a) == size(b)
    # PDMats._addscal!(r, m, view(a.mat), c) doesn't work because _addscal! does
    # not accept views. Below is copy-paste of PDMats
    if c == one(c)
        for i = 1:length(b)
            @inbounds r[i] = a[i] + b[i]
        end
    else
        for i = 1:length(b)
            @inbounds r[i] = a[i] + b[i] * c
        end
    end
    return r
end

*(a::ElasticPDMat, c::Real) = ElasticPDMat(c * Matrix(a), capacity = a.chol.capacity, stepsize = a.chol.stepsize)
*(a::ElasticPDMat, x::AbstractArray) = a.mat * x
\(a::ElasticPDMat, x::AbstractArray) = a.chol \ x

inv(a::ElasticPDMat) = ElasticPDMat(inv(a.chol), capacity = a.chol.capacity, stepsize = a.chol.stepsize)
logdet(a::ElasticPDMat) = logdet(view(a.chol))
eigmax(a::ElasticPDMat) = eigmax(view(a.mat))
eigmin(a::ElasticPDMat) = eigmin(view(a.mat))

function whiten!(r::DenseVecOrMat, a::ElasticPDMat, x::DenseVecOrMat)
    cf = view(a.chol).UL
    v = PDMats._rcopy!(r, x)
    istriu(cf) ? ldiv!(transpose(cf), v) : ldiv!(cf, v)
end

function unwhiten!(r::DenseVecOrMat, a::ElasticPDMat, x::DenseVecOrMat)
    cf = view(a.chol).UL
    v = PDMats._rcopy!(r, x)
    istriu(cf) ? lmul!(transpose(cf), v) : lmul!(cf, v)
end

quad(a::ElasticPDMat{T, A}, x::AbstractArray{T, 1}) where {T, A} = dot(x, a * x)
quad!(r::AbstractArray, a::ElasticPDMat, x::DenseMatrix) = PDMats.colwise_dot!(r, x, a.mat * x)
invquad(a::ElasticPDMat{T, A}, x::AbstractArray{T, 1}) where {T, A} = dot(x, a \ x)
invquad!(r::AbstractArray, a::ElasticPDMat, x::DenseMatrix) = PDMats.colwise_dot!(r, x, a.mat \ x)

function X_A_Xt(a::ElasticPDMat, x::DenseMatrix)
    z = copy(x)
    cf = view(a.chol).UL
    rmul!(z, istriu(cf) ? transpose(cf) : cf)
    z * transpose(z)
end

function Xt_A_X(a::ElasticPDMat, x::DenseMatrix)
    cf = view(a.chol).UL
    z = lmul!(istriu(cf) ? cf : transpose(cf), copy(x))
    transpose(z) * z
end

function X_invA_Xt(a::ElasticPDMat, x::DenseMatrix)
    cf = view(a.chol).UL
    z = rdiv!(copy(x), istriu(cf) ? cf : transpose(cf))
    z * transpose(z)
end

function Xt_invA_X(a::ElasticPDMat, x::DenseMatrix)
    cf = view(a.chol).UL
    z = ldiv!(istriu(cf) ? transpose(cf) : cf, copy(x))
    transpose(z) * z
end
end # module


text:
# __precompile__(true)  "Now the default"
module Symata

import MacroTools
import SpecialFunctions
import REPL
import Markdown
import InteractiveUtils
import Base: setindex!, getindex, replace
import Dates
import Statistics
import Formatting

export @symExpr, @extomx
export Mxpr, getsymata, isympy, margs, mhead, mxpr, mxpra, newargs, pytosj,
       setsymata, sjtopy, @sym, symeval, symparseeval, symparsestring, symprintln, symtranseval,
       unpacktoList
export mmul, mplus, mpow, mabs, mminus, symatamath
export debugmxpr

"""
    devimport()

import some symbols from `Symata` into `Main` that are useful for development.
"""
function devimport()
    Core.eval(Main, Meta.parse("""
       import Symata: @canonexpr, doeval, getpart, infseval, meval, mpmath, mxpr, pytypeof, setpart!,
                      setsymval, symata, symname, sympy, symval!, testex_wrap
"""))
end

export devimport

# FIXME: why export these *and* have devimport ?
export  @canonexpr, doeval, getpart, infseval, meval, mpmath, mxpr, pytypeof, setpart!,
        setsymval, symata, symname, sympy, symval!, testex_wrap

## Set const debugging parametres at compile-time in debug.jl
include("debug.jl") # must use include here, because @inc is defined in debug.jl
@inc("LambertW.jl")  # remove this when Pkg.jl is better developed.
@inc("version.jl")
@inc("util.jl")
@inc("sjcompat.jl")
@inc("mxpr.jl")
@inc("wrapout.jl") # move up once more !
@inc("attributes.jl")
@inc("mxpr_util.jl")
#@inc("wrapout.jl") #
@inc("exceptions.jl")
@inc("level_specification.jl")
@inc("sequence_specification.jl")
@inc("sjiterator.jl")
@inc("sortpattern.jl")
@inc("AST_translation_tables.jl")
@inc("julia_level.jl")
@inc("arithmetic.jl")
@inc("apprules_core.jl")
@inc("constants.jl")
@inc("autoload.jl")
@inc("namedparts.jl")
@inc("AST_translation.jl")
@inc("alteval.jl")
@inc("doc.jl")
@inc("misc_doc.jl")
@inc("symbols.jl")
@inc("random.jl")
@inc("numcomb.jl")
@inc("kernelstate.jl")
@inc("evaluation.jl")
@inc("comparison_logic.jl")
@inc("test.jl")
#@inc("wrapout.jl")
#@inc("formatting/Formatting.jl")
@inc("predicates.jl")
@inc("symata_julia_interface.jl")
@inc("measurements.jl")
@inc("flowcontrol.jl")
@inc("Output.jl")
@inc("system.jl")
@inc("IO.jl")
@inc("latex.jl")
@inc("pattern.jl")
## experimental, eventually replace pattern with this.
#@inc("pattern2.jl")
@inc("parts.jl")
@inc("lists.jl")
@inc("stats.jl")
#@inc("table.jl")
@inc("table_dyn.jl")
@inc("expressions.jl")
@inc("algebra.jl")
@inc("expanda.jl")
@inc("flatten.jl")
@inc("lexcmp.jl")
@inc("sortorderless.jl")
@inc("module.jl")
@inc("strings.jl")
@inc("wrappers.jl")
@inc("keyword_translation.jl")
@inc("sympy.jl")
@inc("math_functions.jl")
# @inc("mittleff_glue.jl") moved mittag-leffler to external dependence. But, it's not registered.
@inc("functionzeros.jl")
@inc("sympy_application.jl")
@inc("alternate_syntax.jl")
@inc("matrix.jl")
@inc("numerical.jl")
@inc("protected_symbols.jl")
@inc("REPL_symata.jl")
@inc("client_symata.jl")
@inc("Jupyter/Jupyter.jl")
@inc("plot.jl")
@inc("docautoloaded.jl")
@inc("function.jl")

import .Jupyter
import .Jupyter: isymata
import .Jupyter: is_isymata_mode
# These are for IJulia. We could probably import insymata in the interface code instead.
export isymata, insymata

function __init__()
    do_init()
end

is_loaded_IJulia() = isdefined(Main, :IJulia)

function do_init()
    init_sympy()
    sjimportall(:System, :Main) # These are Symata "contexts"
    set_current_context(:Main)
    if is_loaded_IJulia() # Use the Jupyter interface
        Jupyter.isymata()
    elseif isinteractive() # Started perhaps from the Julia REPL
        if isdefined(Base, :active_repl)
            RunSymataREPL(Base.active_repl)
        else
            nothing
        end
    else
        nothing
    end
end

end # module Symata

nothing


text:
using Test

tests = ["immctm", "mmctm", "ilda", "lda", "common"]

for test in tests
    include(test * ".jl")
end


text:
 function control_constraints!(model, qs, u, control_limits, vals)
    n_modes = vals.n_modes
    N = vals.N
    k_c = vals.k_c

    S_u = k_c + n_modes * (N - k_c) - 1

    ω_max = control_limits[1]
    ω_min = control_limits[2]
    a_max = control_limits[3]
    a_min = control_limits[4]
    vs_max = control_limits[5]
    vs_min = control_limits[6]

    tan_param = Parameter(() -> tand(30.), model)
    L_param = Parameter(() -> 4., model)

    # TODO: revise the rotation control constraint and address hardcoding
    for k in 1:S_u
        @constraint(model, L_param * u[1, k] - tan_param * qs[4, k]  <= 0) #omega <= abs((v/L)*tan(d))
        @constraint(model, L_param  * u[1, k] + tan_param * qs[4, k] >= 0)
        @constraint(model, u[1, k]  <= ω_max)
        @constraint(model, u[1, k] >= ω_min)
        @constraint(model, u[2, k] <= a_max)
        @constraint(model, u[2, k] >= a_min)
        @constraint(model, u[3, k] <= vs_max)
        @constraint(model, u[3, k] >= vs_min)
    end
end

function dynamics_constraints!(model, q, u, qs, us, vals)
    # TODO: create a function that performs the linearization relevant steps externally

    n_modes = vals.n_modes
    N = vals.N
    k_c = vals.k_c
    dt = vals.dt
    n = vals.n
    m = vals.m

    As = []
    Bs = []
    gs = []

    # constraints for transitions 1 to 2, 2 to 3, ..., k_c - 1 to k_c.
    for k = 1:k_c-1 # TODO: check this indexing
        Ak, Bk = linearize_dynamics(discrete_dynamics, qs[1:n-1, k], us[1:m-1, k], dt)
        vals.As[k] = Ak
        vals.Bs[k] = Bk
        vals.gs[k] = discrete_dynamics(qs[1:n-1, k], us[1:m-1, k], dt) - Ak * qs[1:n-1, k] - Bk * us[1:m-1, k]

        push!(As, Parameter(() -> vals.As[k], model))
        push!(Bs, Parameter(() -> vals.Bs[k], model))
        push!(gs, Parameter(() -> vals.gs[k], model))

        @constraint(model, As[k] * q[1:n-1, k] + Bs[k] * u[1:m-1, k] + gs[k] == q[1:n-1, k+1])
        @constraint(model, dt * u[m, k] + q[n, k] == q[n, k+1])
    end

    # constraints for transitions from k_c to subsequent time step for each mode
    # (i.e., k_c to k_c + 1, k_c to k_c + (j - 1) * N + 1, etc., with j mode index)
    # only need linearization about state at k_c for these transitions
    Ak_c, Bk_c = linearize_dynamics(discrete_dynamics, qs[1:n-1, k_c], us[1:m-1, k_c], dt)
    vals.As[k_c] = Ak_c
    vals.Bs[k_c] = Bk_c
    vals.gs[k_c] = discrete_dynamics(qs[1:n-1, k_c], us[1:m-1, k_c], dt) - Ak_c * qs[1:n-1, k_c] - Bk_c * us[1:m-1, k_c]
    push!(As, Parameter(() -> vals.As[k_c], model))
    push!(Bs, Parameter(() -> vals.Bs[k_c], model))
    push!(gs, Parameter(() -> vals.gs[k_c], model))

    for j = 1:n_modes
        next_state_idx = k_c + (j - 1) * (N - k_c) + 1

        @constraint(model, As[k_c] * q[1:n-1, k_c] + Bs[k_c] * u[1:m-1, k_c] + gs[k_c] == q[1:n-1, next_state_idx])
        @constraint(model, dt * u[m, k_c] + q[n, k_c] == q[n, next_state_idx])
    end

    # constraints for transitions for each parallel horizon "tail" #TODO: check this
    for j = 1:n_modes
        idx_os = (j - 1) * (N - k_c) # index offset
        start_idx = k_c + 1 + idx_os
        end_idx = start_idx + N - k_c - 2

        for k = start_idx:end_idx
            Ak, Bk = linearize_dynamics(discrete_dynamics, qs[1:n-1, k], us[1:m-1, k-(j-1)], dt)
            vals.As[k - (j-1)] = Ak
            vals.Bs[k - (j-1)] = Bk
            vals.gs[k - (j-1)] = discrete_dynamics(qs[1:n-1, k], us[1:m-1, k-(j-1)], dt) - Ak * qs[1:n-1, k] - Bk * us[1:m-1, k-(j-1)]

            push!(As, Parameter(() -> vals.As[k-(j-1)], model))
            push!(Bs, Parameter(() -> vals.Bs[k-(j-1)], model))
            push!(gs, Parameter(() -> vals.gs[k-(j-1)], model))

            @constraint(model, As[k - (j-1)] * q[1:n-1, k] + Bs[k-(j-1)] * u[1:m-1, k-(j-1)] + gs[k-(j-1)] == q[1:n-1, k+1])
            @constraint(model, dt * u[m, k-(j-1)] + q[n, k] == q[n, k+1])
        end
    end

    return As, Bs, gs
end

#TODO: This function can be neater.
function state_constraints!(model, q, q0, vals)
    # initial state constraint
    @constraint(model, q0 == q[:,1])

    # max/min velocity constraint #TODO: this is a hard coded constraint
    for k = 1:vals.N
        @constraint(model, q[4, k] <= 12) # 5 cm/s ;)
        @constraint(model, q[4, k] >= 0.0) # 5 cm/s ;)
    end
end

"""
Creates constraints for static obstacles. Obstacles should be of form of list
position vectors (i.e., [[x_1, y_1], [x_2, y_2], ...]).
"""
function obstacle_constraints!(model, q, qs, vals, scene)
    n_modes = vals.n_modes
    N = vals.N
    k_c = vals.k_c
    n_obs = vals.n_obs
    N_obs = vals.obs_horizon
    node_ids = scene.node_ids
    S_q = k_c + n_modes * (N - k_c)

    ps = []
    p_obs = []
    obs_on = [Parameter(() -> vals.p_obs[node_id].active, model) for node_id in node_ids]

    #TODO: These are hard-coded distances. Change to chance-constraints.
    for k = 1:S_q
        vals.ps[k] = qs[1:2, k]
        push!(ps, Parameter(() -> vals.ps[k], model))
    end

    for j = 1:n_modes
        for k = 1:k_c
            for (i, node_id) in enumerate(node_ids)
                push!(p_obs, Parameter(() -> vals.p_obs[node_id].ps[j][k], model))
                a = @expression (ps[k] - p_obs[length(p_obs)]) / norm(ps[k] - p_obs[length(p_obs)])
                b = 3.
                @constraint(model, obs_on[i]*transpose(a)*(q[1:2, k] - p_obs[length(p_obs)]) >= obs_on[i]*(b))
            end
        end
    end

    for j = 1:n_modes
        for k = k_c+1:N_obs
            for (i, node_id) in enumerate(node_ids)
                idx_os = (j - 1) * (N - k_c) # index offset
                k_os = k + idx_os

                push!(p_obs, Parameter(() -> vals.p_obs[node_id].ps[j][k], model))
                a = @expression (ps[k_os] - p_obs[length(p_obs)]) / norm(ps[k_os] - p_obs[length(p_obs)])
                b = 3.
                @constraint(model, obs_on[i]*transpose(a)*(q[1:2, k_os] - p_obs[length(p_obs)]) >= obs_on[i]*(b))
            end
        end
    end

    return ps, p_obs, obs_on
end

max_stars_repo_name: StanfordASL/MATS

max_issues_repo_name: StanfordASL/MATS

max_forks_repo_name: StanfordASL/MATS


text:
using PSCOPF

using .PSCOPFFixtures

using Test
using JuMP
using Dates
using DataStructures
using Printf

@testset verbose=true "test_bilevel_pilotables_costing" begin

    TS = [DateTime("2015-01-01T11:00:00")]
    ech = DateTime("2015-01-01T07:00:00")
    next_ech = DateTime("2015-01-01T07:30:00")
    function create_instance(load_1, load_2,
                             prop_cost_1, prop_cost_2,
                            limit::Float64=35.,
                            logs=nothing)
        network = PSCOPFFixtures.network_2buses(limit=limit)
        # Pilotables
        PSCOPF.Networks.add_new_generator_to_bus!(network, "bus_1", "prod_1_1", PSCOPF.Networks.PILOTABLE,
                                                0., 200.,
                                                0., prop_cost_1,
                                                Dates.Second(0), Dates.Second(0))
        PSCOPF.Networks.add_new_generator_to_bus!(network, "bus_2", "prod_2_1", PSCOPF.Networks.PILOTABLE,
                                                0., 200.,
                                                0., prop_cost_2,
                                                Dates.Second(0), Dates.Second(0))

        # Uncertainties
        uncertainties = PSCOPF.Uncertainties()
        PSCOPF.add_uncertainty!(uncertainties, ech, "bus_1", TS[1], "S1", load_1)
        PSCOPF.add_uncertainty!(uncertainties, ech, "bus_2", TS[1], "S1", load_2)

        context = PSCOPF.PSCOPFContext(network, TS, PSCOPF.PSCOPF_MODE_1,
                                    SortedDict{String,PSCOPF.GeneratorState}(), #gen_initial_state
                                    uncertainties, nothing, logs)

        return context
    end

    #=
    TS: [11h, 11h15]
    S: [S1,S2]
                        bus 1                   bus 2
                        |                      |
    (limitable) prod_1_1|       "1_2"          |prod_2_1
    Pmin=0, Pmax=200    |                      |Pmin=0, Pmax=200
    Csta=0, Cprop=10    |                      |Csta=0, Cprop=50
                        |----------------------|
                        |         35           |
                        |                      |
                        |                      |
           load(bus_1)  |                      |load(bus_2)
      S1: 10            |                      | S1: 100

    prod_1_1 is cheaper than prod_2_1
     but using only prod_1_1, like a market would, will cause RSO problems
    => The TSO needs to oblige
            prod_1_1 to produce at most 10+35=45 => P_1_1 in [0,45] => cost : 200-45 = 1*155
            prod_2_1 to produce at least 100-35=65 => P_2_1 in [65,200] => cost : 65 = 1*65
        Imposing one of the preceding conditions guarantees RSO constraints
        => TSO will impose one of the two preceding conditions, the cheaper one => option 2

    =#
    @testset "configured_same_cost" begin

        context = create_instance(10., 100.,
                                10., 50.,
                                35.)

        tso = PSCOPF.TSOBilevel(PSCOPF.TSOBilevelConfigs(USE_UNITS_PROP_COST_AS_TSO_BOUNDING_COST=false))
        firmness = PSCOPF.compute_firmness(tso,
                                            ech, next_ech,
                                            TS, context)
        result = PSCOPF.run(tso, ech, firmness,
                    PSCOPF.get_target_timepoints(context),
                    context)

        # Solution is optimal
        @test PSCOPF.get_status(result) == PSCOPF.pscopf_OPTIMAL

        #TSO RSO constraints are OK
        @test value(result.upper.limitable_model.p_global_capping[TS[1],"S1"]) < 1e-09
        @test value(result.upper.lol_model.p_global_loss_of_load[TS[1],"S1"]) < 1e-09

        #Market EOD constraints are OK
        @test value(result.lower.limitable_model.p_global_capping[TS[1],"S1"]) < 1e-09
        @test value(result.lower.lol_model.p_global_loss_of_load[TS[1],"S1"]) < 1e-09

        #TSO sets the bounds for pilotable production respecting units' pmin and pmax
        #prod_1_1 is not bound
        @test value(result.upper.pilotable_model.p_imposition_min["prod_1_1",TS[1],"S1"]) < 1e-09
        @test 200. ≈ value(result.upper.pilotable_model.p_imposition_max["prod_1_1",TS[1],"S1"])
        #prod_2_1 bounding + EOD, guarantees RSO constraints
        @test 65. ≈ value(result.upper.pilotable_model.p_imposition_min["prod_2_1",TS[1],"S1"])
        @test 200. ≈ value(result.upper.pilotable_model.p_imposition_max["prod_2_1",TS[1],"S1"])

        #Market chooses the levels sets the bounds for pilotable production
        @test 45. ≈ value(result.lower.pilotable_model.p_injected["prod_1_1",TS[1],"S1"])
        @test 65. ≈ value(result.lower.pilotable_model.p_injected["prod_2_1",TS[1],"S1"])

        @test objective_value(result.upper.model) ≈ ( 65 )

    end

    #=
    TS: [11h, 11h15]
    S: [S1,S2]
                        bus 1                   bus 2
                        |                      |
    (limitable) prod_1_1|       "1_2"          |prod_2_1
    Pmin=0, Pmax=200    |                      |Pmin=0, Pmax=200
    Csta=0, Cprop=10    |                      |Csta=0, Cprop=50
                        |----------------------|
                        |         35           |
                        |                      |
                        |                      |
           load(bus_1)  |                      |load(bus_2)
      S1: 10            |                      | S1: 100

    TSOBilevelConfigs.USE_UNITS_PROP_COST_AS_TSO_BOUNDING_COST allows using the units' prop cost as
     coeff in the objective function for limiting pilotables.

    prod_1_1 is cheaper than prod_2_1
     but using only prod_1_1, like a market would, will cause RSO problems
    => The TSO needs to oblige
            bus1 to produce at most 10+35=45 => P_1_1 in [0,45] => cost : 200-45 = 10*155 = 1550
            bus2 to produce at least 100-35=65 => P_2_1 in [65,200] => cost : 65 = 50*65 = 3250
        Imposing one of the preceding conditions guarantees RSO constraints
        => TSO will impose one of the two preceding conditions, the cheaper one => option 1

    =#
    @testset "use_prop_costs" begin

        context = create_instance(10., 100.,
                                10., 50.,
                                35.)

        tso = PSCOPF.TSOBilevel(PSCOPF.TSOBilevelConfigs(USE_UNITS_PROP_COST_AS_TSO_BOUNDING_COST=true))

        firmness = PSCOPF.compute_firmness(tso,
                                            ech, next_ech,
                                            TS, context)
        result = PSCOPF.run(tso, ech, firmness,
                    PSCOPF.get_target_timepoints(context),
                    context)

        # Solution is optimal
        @test PSCOPF.get_status(result) == PSCOPF.pscopf_OPTIMAL

        #TSO RSO constraints are OK
        @test value(result.upper.limitable_model.p_global_capping[TS[1],"S1"]) < 1e-09
        @test value(result.upper.lol_model.p_global_loss_of_load[TS[1],"S1"]) < 1e-09

        #Market EOD constraints are OK
        @test value(result.lower.limitable_model.p_global_capping[TS[1],"S1"]) < 1e-09
        @test value(result.lower.lol_model.p_global_loss_of_load[TS[1],"S1"]) < 1e-09

        #TSO sets the bounds for pilotable production respecting units' pmin and pmax
        #prod_1_1 is not bound
        @test value(result.upper.pilotable_model.p_imposition_min["prod_1_1",TS[1],"S1"]) < 1e-09
        @test 45. ≈ value(result.upper.pilotable_model.p_imposition_max["prod_1_1",TS[1],"S1"])
        #prod_2_1 bounding + EOD, guarantees RSO constraints
        @test value(result.upper.pilotable_model.p_imposition_min["prod_2_1",TS[1],"S1"]) < 1e-09
        @test 200. ≈ value(result.upper.pilotable_model.p_imposition_max["prod_2_1",TS[1],"S1"])

        #Market chooses the levels sets the bounds for pilotable production
        @test 45. ≈ value(result.lower.pilotable_model.p_injected["prod_1_1",TS[1],"S1"])
        @test 65. ≈ value(result.lower.pilotable_model.p_injected["prod_2_1",TS[1],"S1"])

        @test objective_value(result.upper.model) ≈ 1550.

    end

end

max_stars_repo_name: klorel/pscopf

max_issues_repo_name: klorel/pscopf

max_forks_repo_name: klorel/pscopf