# 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)

# Build the controller (we will iteratively
# update it at every step of the simulation)
brain = opensim.PrescribedController()
controllers = []

# Add actuators to the controller
state = model.initSystem() # we need to initialize the system (?)
muscleSet = model.getMuscles()
forceSet = model.getForceSet()

for j in range(muscleSet.getSize()):
    func = opensim.Constant(1.0)
    controllers.append(func)
    brain.addActuator(muscleSet.get(j))
    brain.prescribeControlForActuator(j, func)

model.addController(brain)

blockos = opensim.Body('blockos', 0.0001 , opensim.Vec3(0), opensim.Inertia(1,1,.0001,0,0,0) );
pj = opensim.PinJoint("pinblock",
                         model.getGround(), # PhysicalFrame
                         opensim.Vec3(-0.5, 0, 0),
                         opensim.Vec3(0, 0, 0),
                         blockos, # PhysicalFrame
                         opensim.Vec3(0, 0, 0),
                         opensim.Vec3(0, 0, 0))

bodyGeometry = opensim.Ellipsoid(0.1, 0.1, 0.1)
bodyGeometry.setColor(opensim.Gray)
blockos.attachGeometry(bodyGeometry)

model.addComponent(pj)
model.addComponent(blockos)
#block = opensim.ContactMesh('block.obj',opensim.Vec3(0,0,0), opensim.Vec3(0,0,0), blockos);
block = opensim.ContactSphere(0.4, opensim.Vec3(0,0,0), blockos);
model.addContactGeometry(block)

# Reinitialize the system with the new controller
state0 = model.initSystem()
state = opensim.State(state0)

# Change max force
muscleSet.get(0).setMaxIsometricForce(100000.0)

# Get ligaments
ligamentSet = []
for j in range(20, 26):
    ligamentSet.append(opensim.CoordinateLimitForce.safeDownCast(forceSet.get(j)))

state.setTime(0)
manager = opensim.Manager(model)
manager.setIntegratorAccuracy(5e-4)
manager.initialize(state)

for i in range(20):
    # Set some excitation values
    for j in range(muscleSet.getSize()):
        controllers[j].setValue( ((i + j) % 10) * 0.1)

    # Integrate
    t = (i+1) * stepsize
    manager.integrate(t)

    # Report activations and excitations
    model.realizeDynamics(state)
    print("%f %f" % (t, muscleSet.get(0).getActivation(state)))
    print("%f %f" % (t, muscleSet.get(0).getExcitation(state)))

    # Ligaments
    for lig in ligamentSet:
        print(lig.calcLimitForce(state))

    # Restart the model every 100 frames
    if (i + 1) % 100 == 0:
        state = opensim.State(state0)