from osim.env import ProstheticsEnv
import numpy as np
import unittest

class SetStateTest(unittest.TestCase):
    def test_activations(self):
        env = ProstheticsEnv(visualize=False, integrator_accuracy=1e-1)  # we quickly want to see what happens
        env.reset()
        state_checkpoint = env.osim_model.get_state()  # store state

        for i in range(5):
            env.step(env.action_space.high)  # execute step with static action
        obs1 = env.get_observation()

        env.osim_model.set_state(state_checkpoint)  # restore state
        for i in range(5):
            env.step(env.action_space.high)
        obs2 = env.get_observation()

        dist = np.sum((np.array(obs1) - np.array(obs2))**2)
        self.assertTrue(dist < 0.05)
        
if __name__ == '__main__':
    unittest.main()