import simbios.simtk as simtk
import sys

class BounceHandler(simtk.TriggeredEventHandler):
    def __init__(self):
        super(BounceHandler, self).__init__(simtk.Stage.Position)
        self.getTriggerInfo().setTriggerOnRisingSignTransition(False)

    def getValue(self, state):
        return state.q[0]
    
    def handleEvent(self, state, accuracy, yWeights, ooConstraintTols, lowestModified, shouldTerminate):
        state.u[0] *= -1
        lowestModified = simtk.Stage.Velocity
        return lowestModified
    
# Create the system.
system = simtk.MultibodySystem()
matter = simtk.SimbodyMatterSubsystem(system)
forces = simtk.GeneralForceSubsystem(system)
gravity = simtk.Force.UniformGravity(forces, matter, (0, -9.8, 0))
pendulumBody = simtk.Body.Rigid(simtk.MassProperties(1.0, (0,0,0), simtk.Inertia(1)))
pendulumBody.addDecoration(simtk.Transform(), simtk.DecorativeSphere(0.1))
pendulum = simtk.MobilizedBody.Pin(matter.ground, simtk.Transform((0,0,0)), pendulumBody, simtk.Transform((0, 1, 0)))
system.defaultSubsystem.addEventReporter(simtk.VTKEventReporter(system, 0.01))
handler = BounceHandler()
system.defaultSubsystem.addEventHandler(handler)

# Initialize the system and state.
system.realizeTopology()
state = system.defaultState
pendulum.setOneU(state, 0, 1.0)

# Simulate it.
integ = simtk.RungeKuttaMersonIntegrator(system)
ts = simtk.TimeStepper(system, integ)
ts.initialize(state)
ts.stepTo(4.0)