import warnings

import numpy as np
from numpy.linalg import norm

from .. import pinocchio_pywrap_default as pin
from ..utils import npToTuple
from . import BaseVisualizer

try:
    import coal

    WITH_COAL = True
except:  # noqa: E722
    WITH_COAL = False


class GepettoVisualizer(BaseVisualizer):
    """A Pinocchio display using Gepetto Viewer"""

    def getViewerNodeName(self, geometry_object, geometry_type):
        """Return the name of the geometry object inside the viewer"""
        if geometry_type is pin.GeometryType.VISUAL:
            return self.viewerVisualGroupName + "/" + geometry_object.name
        elif geometry_type is pin.GeometryType.COLLISION:
            return self.viewerCollisionGroupName + "/" + geometry_object.name

    def initViewer(
        self,
        viewer=None,
        windowName="python-pinocchio",
        sceneName="world",
        loadModel=False,
    ):
        """Init GepettoViewer by loading the gui and creating a window."""

        try:
            import gepetto.corbaserver
        except ImportError:
            import warnings

            msg = (
                "Error while importing the viewer client.\n"
                "Check whether gepetto-gui is properly installed"
            )
            warnings.warn(msg, category=UserWarning, stacklevel=2)

        try:
            self.viewer = gepetto.corbaserver.Client() if viewer is None else viewer
            gui = self.viewer.gui

            # Create window
            window_l = gui.getWindowList()
            if windowName not in window_l:
                self.windowID = self.viewer.gui.createWindow(windowName)
            else:
                self.windowID = self.viewer.gui.getWindowID(windowName)

            # Create scene if needed
            scene_l = gui.getSceneList()
            if sceneName not in scene_l:
                gui.createScene(sceneName)
            self.sceneName = sceneName
            gui.addSceneToWindow(sceneName, self.windowID)

            if loadModel:
                self.loadViewerModel()
        except:  # noqa: E722
            import warnings

            msg = (
                "Error while starting the viewer client.\n"
                "Check whether gepetto-viewer is properly started"
            )
            warnings.warn(msg, category=UserWarning, stacklevel=2)

    def loadPrimitive(self, meshName, geometry_object):
        gui = self.viewer.gui

        meshColor = geometry_object.meshColor

        geom = geometry_object.geometry
        if isinstance(geom, coal.Capsule):
            return gui.addCapsule(
                meshName, geom.radius, 2.0 * geom.halfLength, npToTuple(meshColor)
            )
        elif isinstance(geom, coal.Cylinder):
            return gui.addCylinder(
                meshName, geom.radius, 2.0 * geom.halfLength, npToTuple(meshColor)
            )
        elif isinstance(geom, coal.Box):
            w, h, d = npToTuple(2.0 * geom.halfSide)
            return gui.addBox(meshName, w, h, d, npToTuple(meshColor))
        elif isinstance(geom, coal.Sphere):
            return gui.addSphere(meshName, geom.radius, npToTuple(meshColor))
        elif isinstance(geom, coal.Cone):
            return gui.addCone(
                meshName, geom.radius, 2.0 * geom.halfLength, npToTuple(meshColor)
            )
        elif isinstance(geom, coal.Plane) or isinstance(geom, coal.Halfspace):
            res = gui.createGroup(meshName)
            if not res:
                return False
            planeName = meshName + "/plane"
            res = gui.addFloor(planeName)
            if not res:
                return False
            normal = geom.n
            rot = pin.Quaternion.FromTwoVectors(normal, pin.ZAxis)
            alpha = geom.d / norm(normal, 2) ** 2
            trans = alpha * normal
            plane_offset = pin.SE3(rot, trans)
            gui.applyConfiguration(planeName, pin.SE3ToXYZQUATtuple(plane_offset))
        elif isinstance(geom, coal.Convex):
            pts = [
                npToTuple(geom.points(geom.polygons(f)[i]))
                for f in range(geom.num_polygons)
                for i in range(3)
            ]
            gui.addCurve(meshName, pts, npToTuple(meshColor))
            gui.setCurveMode(meshName, "TRIANGLES")
            gui.setLightingMode(meshName, "ON")
            gui.setBoolProperty(meshName, "BackfaceDrawing", True)
            return True
        elif isinstance(geom, coal.ConvexBase):
            pts = [npToTuple(geom.points(i)) for i in range(geom.num_points)]
            gui.addCurve(meshName, pts, npToTuple(meshColor))
            gui.setCurveMode(meshName, "POINTS")
            gui.setLightingMode(meshName, "OFF")
            return True
        else:
            msg = f"Unsupported geometry type for {geometry_object.name} ({type(geom)})"
            warnings.warn(msg, category=UserWarning, stacklevel=2)
            return False

    def loadViewerGeometryObject(self, geometry_object, geometry_type):
        """Load a single geometry object"""

        gui = self.viewer.gui

        meshName = self.getViewerNodeName(geometry_object, geometry_type)
        meshPath = geometry_object.meshPath
        meshTexturePath = geometry_object.meshTexturePath
        meshScale = geometry_object.meshScale
        meshColor = geometry_object.meshColor

        try:
            if WITH_COAL and isinstance(geometry_object.geometry, coal.ShapeBase):
                success = self.loadPrimitive(meshName, geometry_object)
            else:
                if meshName == "":
                    msg = (
                        "Display of geometric primitives is supported only if "
                        "pinocchio is build with coal bindings."
                    )
                    warnings.warn(msg, category=UserWarning, stacklevel=2)
                    return
                success = gui.addMesh(meshName, meshPath)
            if not success:
                return
        except Exception as e:
            msg = (
                "Error while loading geometry object: "
                f"{geometry_object.name}\nError message:\n{e}"
            )
            warnings.warn(msg, category=UserWarning, stacklevel=2)
            return

        gui.setScale(meshName, npToTuple(meshScale))
        if geometry_object.overrideMaterial:
            gui.setColor(meshName, npToTuple(meshColor))
            if meshTexturePath != "":
                gui.setTexture(meshName, meshTexturePath)

    def loadViewerModel(self, rootNodeName="pinocchio"):
        """Create the scene displaying the robot meshes in gepetto-viewer"""

        # Start a new "scene" in this window, named "world", with just a floor.
        gui = self.viewer.gui
        self.viewerRootNodeName = self.sceneName + "/" + rootNodeName

        if not gui.nodeExists(self.viewerRootNodeName):
            gui.createGroup(self.viewerRootNodeName)

        self.viewerCollisionGroupName = self.viewerRootNodeName + "/" + "collisions"
        if not gui.nodeExists(self.viewerCollisionGroupName):
            gui.createGroup(self.viewerCollisionGroupName)

        self.viewerVisualGroupName = self.viewerRootNodeName + "/" + "visuals"
        if not gui.nodeExists(self.viewerVisualGroupName):
            gui.createGroup(self.viewerVisualGroupName)

        # iterate over visuals and create the meshes in the viewer
        if self.collision_model is not None:
            for collision in self.collision_model.geometryObjects:
                self.loadViewerGeometryObject(collision, pin.GeometryType.COLLISION)
        # Display collision if we have them and there is no visual
        self.displayCollisions(
            self.collision_model is not None and self.visual_model is None
        )

        if self.visual_model is not None:
            for visual in self.visual_model.geometryObjects:
                self.loadViewerGeometryObject(visual, pin.GeometryType.VISUAL)
        self.displayVisuals(self.visual_model is not None)

        # Finally, refresh the layout to obtain your first rendering.
        gui.refresh()

    def display(self, q=None):
        """
        Display the robot at configuration q in the viewer by placing all the bodies.
        """
        if "viewer" not in self.__dict__:
            return

        gui = self.viewer.gui
        # Update the robot kinematics and geometry.
        if q is not None:
            pin.forwardKinematics(self.model, self.data, q)

        if self.display_collisions:
            pin.updateGeometryPlacements(
                self.model, self.data, self.collision_model, self.collision_data
            )
            gui.applyConfigurations(
                [
                    self.getViewerNodeName(collision, pin.GeometryType.COLLISION)
                    for collision in self.collision_model.geometryObjects
                ],
                [
                    pin.SE3ToXYZQUATtuple(
                        self.collision_data.oMg[
                            self.collision_model.getGeometryId(collision.name)
                        ]
                    )
                    for collision in self.collision_model.geometryObjects
                ],
            )

        if self.display_visuals:
            pin.updateGeometryPlacements(
                self.model, self.data, self.visual_model, self.visual_data
            )
            gui.applyConfigurations(
                [
                    self.getViewerNodeName(visual, pin.GeometryType.VISUAL)
                    for visual in self.visual_model.geometryObjects
                ],
                [
                    pin.SE3ToXYZQUATtuple(
                        self.visual_data.oMg[
                            self.visual_model.getGeometryId(visual.name)
                        ]
                    )
                    for visual in self.visual_model.geometryObjects
                ],
            )

        gui.refresh()

    def displayCollisions(self, visibility):
        """Set whether to display collision objects or not"""
        gui = self.viewer.gui
        self.display_collisions = visibility
        if self.collision_model is None:
            return

        if visibility:
            visibility_mode = "ON"
        else:
            visibility_mode = "OFF"

        for collision in self.collision_model.geometryObjects:
            nodeName = self.getViewerNodeName(collision, pin.GeometryType.COLLISION)
            gui.setVisibility(nodeName, visibility_mode)

    def displayVisuals(self, visibility):
        """Set whether to display visual objects or not"""
        gui = self.viewer.gui
        self.display_visuals = visibility
        if self.visual_model is None:
            return

        if visibility:
            visibility_mode = "ON"
        else:
            visibility_mode = "OFF"

        for visual in self.visual_model.geometryObjects:
            nodeName = self.getViewerNodeName(visual, pin.GeometryType.VISUAL)
            gui.setVisibility(nodeName, visibility_mode)

    def setBackgroundColor(self):
        raise NotImplementedError()

    def setCameraTarget(self, target):
        raise NotImplementedError()

    def setCameraPosition(self, position: np.ndarray):
        raise NotImplementedError()

    def setCameraZoom(self, zoom: float):
        raise NotImplementedError()

    def setCameraPose(self, pose: np.ndarray):
        raise NotImplementedError()

    def captureImage(self, w=None, h=None):
        raise NotImplementedError()

    def disableCameraControl(self):
        raise NotImplementedError()

    def enableCameraControl(self):
        raise NotImplementedError()

    def drawFrameVelocities(self, *args, **kwargs):
        raise NotImplementedError()


__all__ = ["GepettoVisualizer"]
