# This is essentially a hack to fix two problems: 
# 1. The infamous Julia "slow closure" bug.
# 2. Closures cannot be reliably serialized and deserialized with JLD2.

# f
struct NeuralVectorField{S,E,R}
    system::S
    experiment_version::E
    restructure::R
end

function (nvf::NeuralVectorField)(u, θ, t)
    (; system, experiment_version, restructure) = nvf
    return rhs_neural(u, θ, t, restructure, system, Val(experiment_version))
end

# g
struct ConstraintsFunction{S,E,I}
    system::S
    experiment_version::E
    constraints0::I
end

function ConstraintsFunction(system, experiment_version, u0, t0)
    constraints0 = constraints(u0, t0, system, Val(experiment_version))
    return ConstraintsFunction(system, experiment_version, constraints0)
end

function (constraints_function::ConstraintsFunction)(u, t)
    (; system, experiment_version, constraints0) = constraints_function
    return constraints(u, t, system, Val(experiment_version)) .- constraints0
end

# F
abstract type AbstractStabilizationMatrix end

struct ConstraintsPseudoinverse{S,E} <: AbstractStabilizationMatrix
    system::S
    experiment_version::E
end

function (constraints_inv::ConstraintsPseudoinverse)(u, t)
    (; system, experiment_version) = constraints_inv
    J = constraints_jacobian(u, t, system, Val(experiment_version))
    return J' * inv(J * J')  # Turns out to be faster and well-conditioned always
end
