#
# Copyright (c) 2016 CNRS
#
# This file is part of Pinocchio
# Pinocchio is free software: you can redistribute it
# and/or modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation, either version
# 3 of the License, or (at your option) any later version.
#
# Pinocchio is distributed in the hope that it will be
# useful, but WITHOUT ANY WARRANTY; without even the implied warranty
# of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# General Lesser Public License for more details. You should have
# received a copy of the GNU Lesser General Public License along with
# Pinocchio If not, see
# <http://www.gnu.org/licenses/>.

import pinocchio as se3
from math import pi
from pinocchio.utils import np, rotate, XYZQUATToViewerConfiguration, se3ToXYZQUAT

DENSITY = 1


def placement(x=0, y=0, z=0, rx=0, ry=0, rz=0):
    m = se3.SE3.Identity()
    m.translation = np.matrix([[float(i)] for i in [x, y, z]])
    m.rotation *= rotate('x', rx)
    m.rotation *= rotate('y', ry)
    m.rotation *= rotate('z', rz)
    return m


def color(body_number=1):
    return [int(i) for i in '%03d' % int(bin(body_number % 8)[2:])] + [1]


class ModelWrapper(object):
    def __init__(self, model, name=None, display=False):
        self.visuals = [('universe', se3.SE3.Identity(), se3.SE3.Identity().translation)]
        self.name = self.__class__.__name__ if name is None else name
        self.model = model
        self.display = display
        self.add_joints()

    def add_joints(self):
        for joint in self.joints:
            self.add_joint(**joint)

    def add_joint(self, joint_name, joint_model=None, joint_placement=None, lever=None, shape="box",
                  dimensions=1, mass=None, body_color=1, parent=0):
        if joint_model is None:
            joint_model = se3.JointModelFreeFlyer()
        elif isinstance(joint_model, str):
            joint_model = se3.__dict__['JointModel' + joint_model]()
        if joint_placement is None:
            joint_placement = se3.SE3.Identity()
        elif isinstance(joint_placement, dict):
            joint_placement = placement(**joint_placement)
        if lever is None:
            lever = se3.SE3.Identity()
        elif isinstance(lever, dict):
            lever = placement(**lever)
        joint_name, body_name = ('world/%s_%s_%s' % (self.name, joint_name, i) for i in ('joint', 'body'))
        body_inertia = se3.Inertia.Random()
        if shape == "box":
            w, h, d = (float(i) for i in dimensions) if isinstance(dimensions, tuple) else [float(dimensions)] * 3
            if mass is None:
                mass = w * h * d * DENSITY
            body_inertia = se3.Inertia.FromBox(mass, w, h, d)
            if self.display:
                self.display.viewer.gui.addBox(body_name, w, h, d, color(body_color))
        elif shape == "cylinder":
            r, h = dimensions
            if mass is None:
                mass = pi * r ** 2 * h * DENSITY
            body_inertia = se3.Inertia.FromCylinder(mass, r, h)
            if self.display:
                self.display.viewer.gui.addCylinder(body_name, r, h, color(body_color))
        elif shape == "sphere":
            w, h, d = (float(i) for i in dimensions) if isinstance(dimensions, tuple) else [float(dimensions)] * 3
            if mass is None:
                mass = 4. / 3. * pi * w * h * d * DENSITY
            body_inertia = se3.Inertia.FromEllipsoid(mass, w, h, d)
            if self.display:
                self.display.viewer.gui.addSphere(body_name, dimensions, color(body_color))
        body_inertia.lever = lever.translation
        self.model.addJointAndBody(parent, joint_model, joint_placement, body_inertia, joint_name, body_name)
        self.visuals.append((body_name, joint_placement, lever))
        self.data = self.model.createData()
        if self.display:
            self.place()

    def place(self):
        for i, (name, placement, lever) in enumerate(self.visuals):
            if i == 0:
                continue
            self.display.place(name, self.data.oMi[i] * placement * lever)
        self.display.viewer.gui.refresh()


class RobotDisplay(object):
    def __init__(self, window_name="pinocchio"):
        from gepetto import corbaserver

        self.viewer = corbaserver.Client()
        try:
            window_id = self.viewer.gui.getWindowID(window_name)
        except:
            window_id = self.viewer.gui.createWindow(window_name)
            self.viewer.gui.createSceneWithFloor("world")
            self.viewer.gui.addSceneToWindow("world", window_id)
        self.viewer.gui.setLightingMode('world', 'OFF')
        self.viewer.gui.setVisibility('world/floor', "OFF")
        self.viewer.gui.refresh()

    def place(self, obj_name, m):
        self.viewer.gui.applyConfiguration(obj_name, XYZQUATToViewerConfiguration(se3ToXYZQUAT(m)))


class SimplestWalker(ModelWrapper):
    joints = [
        {'joint_name': "pelvis",
         'dimensions': (.1, .2, .1),
         'mass': .5,
         },
        {'joint_name': "left_leg",
         'joint_model': "RY",
         'joint_placement': {'y': -.15},
         'lever': {'z': -.45},
         'shape': "cylinder",
         'dimensions': (.1, .9),
         'mass': 20,
         'body_color': 2,
         'parent': 1,
         },
        {'joint_name': "right_leg",
         'joint_model': "RY",
         'joint_placement': {'y': .15},
         'lever': {'z': -.45},
         'shape': "cylinder",
         'dimensions': (.1, .9),
         'mass': 20,
         'body_color': 3,
         'parent': 1,
         },
    ]

SimplestWalker(model)