from pinocchio.utils import *
from pinocchio.explog import exp,log
from numpy.linalg import pinv,norm
from pinocchio import forwardKinematics, Inertia, JointModelRX, Model, SE3
import gepetto.corbaserver
from display import Display
import pinocchio as pin
import eigenpy
import hppfcl
from pinocchio.visualize import GepettoVisualizer, MeshcatVisualizer

eigenpy.switchToNumpyArray()

color   = [red,green,blue,transparency] = [1,1,0.78,1.0]
red = [1.0,0.0,0.0,1.0]

class Joint:
    def __init__(self,name,parent=None,rootid=0, **kargs):
        self.name=name
        self.joint_models = kwargs.get("joint_models", [pin.JointModelRX()])
        self.placement = kwargs.get("placement", self.__calculate_placement(parent))
        self.ids = []
        self.geo_ids = []
        self.inertias = kwargs.get("inertias")
        
        #check if parent =0
        
        if not self.inertias:
            self.inertias = [None] * len(self.joint_models)
            for i in range(len(self.joint_models)):
                self.inertias[i] = pin.Inertia.Random()
        for i, (joint_model, inertia) in enumerate(
            zip(self.joint_models, self.inertias)
        ):
            current_parent = None
            current_placement = None
            if i == 0:
                current_parent = parent.id if parent else root_id
                current_placement = self.placement
            else:
                current_parent = self.ids[-1]
                current_placement = pin.SE3.Identity()
            self.ids.append(
                model.addJoint(
                    current_parent,
                    joint_model,
                    current_placement,
                    self.name,
                )
            model.appendBodyToJoint(self.ids[-1], inertia, pin.SE3.Identity())
        self.id = self.ids[-1]
        self.sphere_radius = kwargs.get("sphere_radius", 0.1)
        self.sphere = None
        if self.sphere_radius:
            self.sphere_name = kwargs.get("sphere_name", self.name)
            self.sphere_color = kwargs.get("sphere_color", red)
            self.sphere_placement = kwargs.get("sphere_placement", pin.SE3.Identity())
        self.sphere = self.__create_geo_obj(
                self.sphere_name,
                self.sphere_placement,
                self.sphere_color,
                hppfcl.Sphere,
                self.sphere_radius,
            )
            self.geo_ids.append(geo_model.addGeometryObject(self.sphere))
        self.box_x = kwargs.get("box_x", 0.1)
        self.box_y = kwargs.get("box_y", 0.1)
        self.box_z = kwargs.get("box_z", 0)
        self.box = None
        if self.box_x and self.box_y and self.box_z:
            self.box_name = kwargs.get("box_name", f"box_{name}")
            self.box_color = kwargs.get("box_color", WHITE)
            # Placement of the geometry with respect to the joint frame
            self.box_placement = kwargs.get(
                "box_placement", pin.SE3(np.eye(3), np.array([0, 0, self.box_z / 2]))
            )
            self.box = self.__create_geo_obj(
                self.box_name,
                self.box_placement,
                self.box_color,
                hppfcl.Box,
                self.box_x,
                self.box_y,
                self.box_z,
            )
            self.geo_ids.append(geo_model.addGeometryObject(self.box))
        def __create_geo_obj(self, name, placement, color, factory, *nargs, **kwargs):
        return pin.GeometryObject(
            name=name,
            parent_joint=self.id,
            collision_geometry=factory(*nargs, **kwargs),
            placement=placement,
            mesh_path="",
            mesh_scale=np.ones(3),
            override_material=False,
            mesh_color=color,
        )
        def __calculate_placement(self, parent=None):
        if parent is None:
            return pin.SE3.Identity()
        else:
            parent_z = 0
            if parent.box_z != 0:
                parent_z = parent.box_z
        return pin.SE3(np.eye(3), np.array([0, 0, parent_z]))

class Robot:
    def __init__(self, root_id=0, show_origin=True, add_plane=True):
        self.viewer = Display()
        self.visuals = []
        self.model = Model()
        collision_model = pin.GeometryModel()
        visual_model = pin.GeometryModel()