struct RobotArm{T} <: AbstractDynamicalSystem{T} end

function endpoint(θ)
    θ₁, θ₂, θ₃ = θ
    return [cos(θ₁) + cos(θ₂) + cos(θ₃), sin(θ₁) + sin(θ₂) + sin(θ₃)]
end

function pose(t)
    return [-sin(2π * t) / 2π, 0.0]
end

function dpose(t)
    return [-cos(2π * t), 0.0]
end

function (ra::RobotArm)(du, u, p, t)
    θ₁, θ₂, θ₃ = u
    G = [
        -sin(θ₁) -sin(θ₂) -sin(θ₃)
        cos(θ₁) cos(θ₂) cos(θ₃)
    ]
    du .= G' * inv(G * G') * dpose(t)
    return nothing
end

function initial_conditions(::RobotArm{T}) where {T}
    θ₀ = π / 4 + π / 8 * rand()  # [2π/8, 3π/8]
    return T[θ₀, -θ₀, θ₀]
end

function get_trajectories(
    system::RobotArm{T},
    experiment_version,
    seconds,
    dt,
    transient_seconds,
    solver,
    reltol,
    abstol,
    N,
    steps,
    stabilization_param,
    θ,
    restructure,
) where {T}
    f = NeuralVectorField(system, experiment_version, restructure)
    F = ConstraintsPseudoinverse(system, experiment_version)
    γ = T(stabilization_param)

    systemBF = RobotArm{BigFloat}()
    trajectories = []
    for _ = 1:N
        time_series = generate_data(
            systemBF;
            seconds,
            dt,
            transient_seconds,
            solver,
            reltol,
            abstol,
            u0 = initial_conditions(systemBF),
            NF = T,
        )
        u0 = time_series.trajectory[:, 1]
        t0 = time_series.times[1]

        # Set up the SNDE
        if γ == 0
            rhs = f
        else
            g = ConstraintsFunction(system, experiment_version, u0, t0)
            rhs = StabilizedNDE(f, g, F, γ)
        end
        prob =
            ODEProblem{false,SciMLBase.FullSpecialize}(rhs, zeros(T), (zero(T), one(T)), θ)
        data_ms = multiple_shooting(prob, time_series; steps)
        push!(trajectories, data_ms)
    end
    return vcat(trajectories...)
end

function get_mlp(hidden_layers, hidden_width, activation, ::RobotArm{T}, ::Val{1}) where {T}
    return get_mlp(8 => 3, hidden_layers, hidden_width, activation, T)
end

function constraints(u, t, ::RobotArm{T}, ::Val{1}) where {T}
    return endpoint(u) .- pose(t)
end

function constraints_jacobian(u::AbstractVector{T}, t::T, ::RobotArm{T}, ::Val{1}) where {T}
    θ₁, θ₂, θ₃ = u
    return [
        -sin(θ₁) -sin(θ₂) -sin(θ₃)
        cos(θ₁) cos(θ₂) cos(θ₃)
    ]
end

function rhs_neural(u, θ, t, re::Optimisers.Restructure, ::RobotArm, ::Val{1})
    θ₁, θ₂, θ₃ = u
    x = vcat(cos(θ₁), cos(θ₂), cos(θ₃), sin(θ₁), sin(θ₂), sin(θ₃), dpose(t))
    return re(θ)(x)
end
