using RigidBodyDynamics
using LinearAlgebra
using StaticArrays
using RigidBodyDynamics.OdeIntegrators: step, process
using RigidBodyDynamics.OdeIntegrators
using Zygote;

function init_RigidBody()    
    world_body = RigidBody{Float64}("world");
    rigid_body_mechanism = Mechanism(world_body; gravity = SVector(0, 0, -9.81));

    # Body
    frame = CartesianFrame3D("frame"); # the reference frame in which the spatial inertia will be expressed
    inertia = SpatialInertia( frame,
                              moment = Matrix(0.01I, 3, 3), # moment of inertia,
                              com = SVector(0, 0, 0), # center of mass location with respect to joint axis
                              mass = 1.0);
    body = RigidBody(inertia)

    # The joint is used by the sim_callback! therefore it must be available globally.
    global joint = Joint("joint", QuaternionFloating{Float64}());
    to_world = one(Transform3D, frame_before(joint), default_frame(world_body))
    attach!(rigid_body_mechanism, world_body, body, joint, joint_pose = to_world);

    mechanism_state = MechanismState(rigid_body_mechanism);


    # Set the initial mechanism_state configuration
    set_configuration!(mechanism_state, joint, [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) # in World Frame
    set_velocity!(mechanism_state, joint, [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) # In Body frame
    
    return mechanism_state, joint
end;

function sim_callback!(torques::AbstractVector{T}, t, state::MechanismState{T})  where T
    # takes the Global Torque, Force and Joint Value
    torques[velocity_range(state, joint)] = [torque_B; force_B]
end;

function init_integrator(mechanism_state)
    T = Float64;
    result = DynamicsResult{T}(mechanism_state.mechanism)
    control_torques::AbstractVector{T} = [0.0; 0.0; 0.0; 0.0; 0.0; 0.0];


    #a callable dynamics!(vd, t, state) that updates the joint acceleration vector vd at time t and in state state;
    closed_loop_dynamics! = function (v̇::AbstractArray, ṡ::AbstractArray, t, state)
                sim_callback!(control_torques, t, state)
                dynamics!(result, state, control_torques)
                copyto!(v̇, result.v̇)
                copyto!(ṡ, result.ṡ) 
            end


    tableau = runge_kutta_4(T) # Return the Butcher tableau for the standard fourth order Runge-Kutta integrator.
    storage = RingBufferStorage{T}(mechanism_state, 1); # Ringbuffer with one place as OdeResultsSink
    
    integrator = MuntheKaasIntegrator(mechanism_state, closed_loop_dynamics!, tableau, storage);
    
    return integrator, storage
end;

function rigid_body_simulation_step(integrator, t, Δt, storage)
    step(integrator, t, Δt)

    # Returns the position of the body
    x_W = configuration(integrator.state)[7]
    return x_W
end;

function set_next_torque_force!(torque_B_::AbstractVector{T}, force_B_::AbstractVector{T}) where T
    # Since the MuntheKaasIntegrator calls the dynamic update, the new force values must be passed globally.
    global torque_B = torque_B_;
    global force_B = force_B_;
end;

function f(x)
    t = 0.0;
    Δt = 1e-2;
    loss = 0.0;
    
    # The force acting on the body along the gravity can be varied by x.
    force_B = [0.0; 0.0; x];
    torque_B = [0.0; 0.0; 0.0];
    

    for i in 1:10
        set_next_torque_force!(torque_B, force_B)
        x_W = rigid_body_simulation_step(integrator, t, Δt, storage)
        loss += norm(x_W)
        t = Δt;
    end
    
    # The loss is the sum of all positions deviating from the origin.
    return loss
end;

mechanism_state, joint = init_RigidBody();
integrator, storage = init_integrator(mechanism_state);