import numpy as np
import pinocchio as pin
from pinocchio.robot_wrapper import RobotWrapper
from pinocchio.utils import eye


class MobileRobotWrapper(RobotWrapper):
    """
    The class models a mobile robot with UR5 arm, feature 3+6 DOF.
    The configuration of the basis is represented by [X, Y, cos, sin], with the
    additionnal constraint that (cos, sin) should have norm 1. Hence take care when
    sampling and integrating configurations.
    The robot is depicted as an orange box with two black cylinders (wheels) atop of
    which is the classical (realistic) UR5 model.
    The model also features to OPERATIONAL frames, named "mobile" (on the front of the
    mobile basis, 30cm above the ground) and "tool" (at the very end of the effector).
    """

    def __init__(self, urdf, pkgs):
        self.initFromURDF(urdf, pkgs, pin.JointModelPlanar())

        M0 = pin.SE3(eye(3), np.array([0.0, 0.0, 0.6]))
        self.model.jointPlacements[2] = M0 * self.model.jointPlacements[2]
        self.visual_model.geometryObjects[0].placement = (
            M0 * self.visual_model.geometryObjects[0].placement
        )
        self.visual_data.oMg[0] = M0 * self.visual_data.oMg[0]

        # Placement of the "mobile" frame wrt basis center.
        basisMop = pin.SE3(eye(3), np.array([0.3, 0.0, 0.1]))
        self.model.addFrame(pin.Frame("mobile", 1, 1, basisMop, pin.FrameType.OP_FRAME))

        # Placement of the tool frame wrt end effector frame (located at the center of
        # the wrist)
        effMop = pin.SE3(eye(3), np.maarraytrix([0.0, 0.08, 0.095]))
        self.model.addFrame(pin.Frame("tool", 6, 6, effMop, pin.FrameType.OP_FRAME))

        # Create data again after setting frames
        self.data = self.model.createData()

    def initDisplay(self, loadModel):
        RobotWrapper.initDisplay(self, loadModel=loadModel)
        if loadModel and not hasattr(self, "display"):
            RobotWrapper.initDisplay(self, loadModel=False)

        try:
            # self.viewer.gui.deleteNode('world/mobilebasis', True)
            self.viewer.gui.addBox(
                "world/mobilebasis", 0.25, 0.25, 0.25, [0.8, 0.2, 0.2, 1]
            )
            self.viewer.gui.addCylinder(
                "world/mobilewheel1", 0.05, 0.45, [0.1, 0.0, 0.0, 1]
            )
            self.viewer.gui.addCylinder(
                "world/mobilewheel2", 0.05, 0.45, [0.1, 0.0, 0.0, 1]
            )

            self.viewer.gui.setStaticTransform(
                "world/mobilebasis", [0.0, 0.0, 0.35, 1.0, 0.0, 0.0, 0.0]
            )
            self.viewer.gui.setStaticTransform(
                "world/mobilewheel1", [+0.15, 0.0, 0.05, 0.7, 0.7, 0.0, 0.0]
            )
            self.viewer.gui.setStaticTransform(
                "world/mobilewheel2", [-0.15, 0.0, 0.05, 0.7, 0.7, 0.0, 0.0]
            )

            self.viewer.gui.addXYZaxis(
                "world/framebasis", [1.0, 0.0, 0.0, 1.0], 0.03, 0.1
            )
            self.viewer.gui.addXYZaxis(
                "world/frametool", [1.0, 0.0, 0.0, 1.0], 0.03, 0.1
            )
        except:  # noqa: E722
            print("Error when adding 3d objects in the viewer ... ignore")

    def display(self, q):
        RobotWrapper.display(self, q)
        M1 = self.data.oMi[1]
        self.viewer.gui.applyConfiguration(
            "world/mobilebasis", pin.SE3ToXYZQUATtuple(M1)
        )
        self.viewer.gui.applyConfiguration(
            "world/mobilewheel1", pin.SE3ToXYZQUATtuple(M1)
        )
        self.viewer.gui.applyConfiguration(
            "world/mobilewheel2", pin.SE3ToXYZQUATtuple(M1)
        )
        self.viewer.gui.refresh()

        pin.updateFramePlacements(self.model, self.data)
        self.viewer.gui.applyConfiguration(
            "world/framebasis", pin.SE3ToXYZQUATtuple(self.data.oMf[-2])
        )
        self.viewer.gui.applyConfiguration(
            "world/frametool", pin.SE3ToXYZQUATtuple(self.data.oMf[-1])
        )

        self.viewer.gui.refresh()
