import unittest

import numpy as np
import pinocchio as pin

from core.model_contract import (
    MASTER_FRAMES,
    MASTER_JOINT_NAMES,
    SLAVE_FRAMES,
    SLAVE_JOINT_NAMES,
    clip_configuration,
    finite_joint_limits,
    joint_q_indices,
    load_models,
    require_frame,
    safe_configuration,
)


class ModelContractTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.models = load_models()

    def test_canonical_models_and_simulated_tcp_load(self):
        self.assertEqual(self.models.master.nq, 7)
        self.assertEqual(self.models.slave.nq, 7)
        for name in MASTER_FRAMES.values():
            require_frame(self.models.master, name)
        for name in SLAVE_FRAMES.values():
            require_frame(self.models.slave, name)

    def test_safe_slave_configuration_is_strictly_inside_limits(self):
        q = safe_configuration(self.models.slave, SLAVE_JOINT_NAMES)
        qidx = joint_q_indices(self.models.slave, SLAVE_JOINT_NAMES)
        lower, upper = finite_joint_limits(self.models.slave, SLAVE_JOINT_NAMES)
        self.assertTrue(np.all(q[qidx] > lower))
        self.assertTrue(np.all(q[qidx] < upper))

    def test_clipping_reports_limit_event(self):
        q = pin.neutral(self.models.slave)
        qidx = joint_q_indices(self.models.slave, SLAVE_JOINT_NAMES)
        lower, upper = finite_joint_limits(self.models.slave, SLAVE_JOINT_NAMES)
        q[qidx] = upper + 0.1
        clipped, event = clip_configuration(
            self.models.slave, q, SLAVE_JOINT_NAMES, margin=1e-4
        )
        self.assertTrue(event)
        self.assertTrue(np.all(clipped[qidx] <= upper - 1e-4 + 1e-12))
        self.assertTrue(np.all(clipped[qidx] >= lower + 1e-4 - 1e-12))


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