# This script goes through OpenSim funcionalties

# required for OpenSim-RL

import opensim

# Settings

stepsize = 0.01

# Load existing model

model_path = "../osim/models/gait9dof18musc.osim"

model = opensim.Model(model_path)

model.setUseVisualizer(True)

# Create the ball

r = 0.000001

ballBody = opensim.Body('ball', 0.0001 , opensim.Vec3(0), opensim.Inertia(1,1,.0001,0,0,0) );

ballGeometry = opensim.Ellipsoid(r, r, r)

ballGeometry.setColor(opensim.Gray)

ballBody.attachGeometry(ballGeometry)

# Attach ball to the model

ballJoint = opensim.FreeJoint("weldball",

                         model.getGround(), # PhysicalFrame

                         opensim.Vec3(0, 0, 0),

                         opensim.Vec3(0, 0, 0),

                         ballBody, # PhysicalFrame

                         opensim.Vec3(0, 0, 0),

                         opensim.Vec3(0, 0, 0))

model.addComponent(ballJoint)

model.addComponent(ballBody)

# Add contact

ballContact = opensim.ContactSphere(r, opensim.Vec3(0,0,0), ballBody);

model.addContactGeometry(ballContact)

# Reinitialize the system with the new controller

state = model.initSystem()

for i in range(6):

    ballJoint.getCoordinate(i).setLocked(state, True)

state.setTime(0)

manager = opensim.Manager(model)

manager.setIntegratorAccuracy(5e-4)

manager.initialize(state)

# Simulate

for i in range(100):

    manager.integrate(i + stepsize)

    # Restart the model every 10 frames, with the new position of the ball

    if (i + 1) % 10 == 0:

        newloc = opensim.Vec3(float(i) / 5, 0, 0)

        opensim.PhysicalOffsetFrame.safeDownCast(ballJoint.getChildFrame()).set_translation(newloc)

        r = i * 0.005

        #ballGeometry.set_radii(opensim.Vec3(r,r,r))

        #ballBody.scale(opensim.Vec3(r, r, r))

        ballContact.setRadius(r)

        state = model.initializeState()

        ballJoint.getCoordinate(3).setValue(state, i / 100.0)

        for i in range(6):

            ballJoint.getCoordinate(i).setLocked(state, True)