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

class ActivationsTest(unittest.TestCase):
    def test_activations(self):
        env = L2RunEnv(visualize=False)
        observation = env.reset()

        newact = np.array([0.0] * 18)
        env.osim_model.set_activations(newact)
        
        current = np.array(env.osim_model.get_activations())
        dist = np.linalg.norm(newact - current)
        self.assertTrue(dist < 0.05)
        
        newact = np.array([1.0] * 18)
        env.osim_model.set_activations(newact)

        current = np.array(env.osim_model.get_activations())
        dist = np.linalg.norm(newact - current)
        self.assertTrue(dist < 0.05)

    def test_activations_changes(self):
        env = L2RunEnv(visualize=False)

        # Do not set new activations
        newAct = [0.9] * 18
        observation = env.reset()
        env.osim_model.set_activations(newAct)
        for i in range(5): 
            withoutAct = env.osim_model.get_activations()
            observation, reward, done, info = env.step([0.5]*18)

        # Set new activations
        newAct = [0.1] * 18
        observation = env.reset()
        env.osim_model.set_activations(newAct)
        for i in range(5): 
            withAct = env.osim_model.get_activations()
            observation, reward, done, info = env.step([0.5]*18)

        dist = np.linalg.norm(np.array(withAct) - np.array(withoutAct))

        self.assertFalse(dist < 1e-2,"Activations after 5 steps haven't changed (despite different initial conditions)")

        
if __name__ == '__main__':
    unittest.main()