from math import pi

 

import cozmo

 

from .kine import *

from cozmo_fsm import geometry

from .geometry import tprint

from .rrt_shapes import *

 

 

wheelbase = 45 # millimeters

center_of_rotation_offset = -19.7 # millimeters

 

 

class CozmoKinematics(Kinematics):

    def __init__(self,robot):

        base_frame = Joint('base',

                           description='Base frame: the root of the kinematic tree')

 

        # cor is center of rotation

        cor_frame = Joint('cor', parent=base_frame,

                          description='Center of rotation',

                          r=-19.,

                          collision_model=Rectangle(geometry.point(),

                                                    dimensions=(95,60)))

 

        # Use link instead of joint for world_frame

        world_frame = Joint('world', parent=base_frame, type='world', getter=self.get_world,

                            description='World origin in base frame coordinates',

                            qmin=None, qmax=None)

 

        front_axle_frame = Joint('front_axle', parent=base_frame,

                                 description='Center of the front axle',

                                 alpha=pi/2)

       

        left_front_wheel_frame = Joint('left_front_wheel', parent = front_axle_frame,

                                      description = 'Axis of left wheel rotation - \

                                     z points to the right', d = -27.5)

 

        right_front_wheel_frame = Joint('right_front_wheel', parent = front_axle_frame,

                                      description = 'Axis of right wheel rotation - \

                                     z points to the right', d = 27.5)

 

        back_axle_frame = Joint('back_axle', parent=base_frame, r=-46., alpha=pi/2)

 

        # This frame is on the midline.  Could add separate left and right shoulders.

        # Positive angle is up, so z must point to the right.

        # x is forward, y points up.

        shoulder_frame = Joint('shoulder', parent=base_frame,

                               type='revolute', getter=self.get_shoulder,

                               description='Rotation axis of the lift; z points to the right',

                               qmin=cozmo.robot.MIN_LIFT_ANGLE.radians,

                               qmax=cozmo.robot.MAX_LIFT_ANGLE.radians,

                               d=21., r=-39., alpha=pi/2)

 

        lift_attach_frame = \

            Joint('lift_attach', parent=shoulder_frame, type='revolute',

                  description='Tip of the lift, where cubes attach; distal end of four-bar linkage',

                  getter=self.get_lift_attach, r=66.,

                  qmax = - cozmo.robot.MIN_LIFT_ANGLE.radians,

                  qmin = - cozmo.robot.MAX_LIFT_ANGLE.radians,

                  #collision_model=Circle(geometry.point(), radius=10))

            )

 

        #Go get exact measurements

 

        left_hook_dummy = Joint('left_hook_dummy', parent= lift_attach_frame,

                            description = 'Dummy joint at horizontal position of the hooks - \

                               z points up', d = -17.5, alpha = -pi/2)

 

        left_hook_frame = Joint('left_hook', parent = left_hook_dummy,

                                description = 'frame with origin in the center of the left hooks \

                                    attachment to the robot - z points upwards', d = -10)

 

        right_hook_dummy = Joint('left_hook_dummy', parent= lift_attach_frame,

                            description = 'Dummy joint at horizontal position of the hooks - \

                               z points up', d = 17.5, alpha = -pi/2)

 

        right_hook_frame = Joint('left_hook', parent = right_hook_dummy,

                                description = 'frame with origin in the center of the right hooks \

                                    attachment to the robot - z points upwards', d = -10)

 

        # Positive head angle is up, so z must point to the right.

        # With x pointing forward, y must point up.

        head_frame = Joint('head', parent=base_frame, type='revolute',

                           getter=self.get_head,

                           description='Axis of head rotation; z points to the right',

                           qmin=cozmo.robot.MIN_HEAD_ANGLE.radians,

                           qmax=cozmo.robot.MAX_HEAD_ANGLE.radians,

                           d=35., r=-10., alpha=pi/2)

 

        # Dummy joint located below head joint at level of the camera frame,

        # and x axis points down, z points forward, y points left

        camera_dummy = Joint('camera_dummy', parent=head_frame,

                             description='Dummy joint below the head, at the level of the camera frame',

                             theta=-pi/2, r=7.5, alpha=-pi/2)

        # x axis points right, y points down, z points forward

        camera_frame = Joint('camera', parent=camera_dummy,

                             description='Camera reference frame; y is down and z is outward',

                             d=15., theta=-pi/2)

 

        joints = [base_frame, world_frame, cor_frame,

                  front_axle_frame, back_axle_frame,

                  shoulder_frame, lift_attach_frame,

                  head_frame, camera_dummy, camera_frame,

                  left_front_wheel_frame, right_front_wheel_frame,

                  left_hook_dummy, right_hook_dummy,

                  left_hook_frame, right_hook_frame]

 

        super().__init__(joints,robot)

 

    def get_head(self):

        return self.robot.head_angle.radians

 

    def get_shoulder(self):

        # Formula supplied by Mark Wesley at Anki

        # Check SDK documentation for new lift-related calls that might replace this

        return math.asin( (self.robot.lift_height.distance_mm-45.0) / 66.0)

 

    def get_lift_attach(self):

        return -self.get_shoulder()

 

    def get_world(self):

        return self.robot.world.particle_filter.pose_estimate()