from __future__ import print_function

import numpy as np
import struct

import solvers
import pid
from util import *

MOTORSPEED = 0.9
MOTORMARGIN = 1
MOTORSLOPE = 30

ERRORLIM = 5.0

class ArmConfig:
    """Holds an arm's proportions, limits and other configuration data"""
    def __init__(self,
                 main_length = 148.4,
                 forearm_length = 160,
                 linkage_length = 155,
                 lower_actuator_length = 65,
                 upper_actuator_length = 54.4,
                 wrist_length = 90.52,
                 shoulder_offset = [-9.7, 18.71]):
        self.main_length = main_length
        self.forearm_length = forearm_length
        self.linkage_length = linkage_length
        self.lower_actuator_length = lower_actuator_length
        self.upper_actuator_length = upper_actuator_length
        self.wrist_length = wrist_length;
        self.shoulder_offset = shoulder_offset


class ArmPose:
    """
    Defines a physical configuration of a LiteArm robot arm.
    Internal angles are relative to vertical (elevator/actuator) or straight
    forward (swing), and are stored in radians. Extracted servo angles range
    0-300 and are measured in degrees.

    Provides methods for:
        - finding the required servo angles to reach the pose
        - checking the validity of the pose
    """
    structFormat = 'fffff'

    @staticmethod
    def calcElevatorAngle(servoAngle):
        return radians(178.21 - servoAngle)

    @staticmethod
    def calcSwingAngle(servoAngle):
        return radians(150.0 - servoAngle)

    @staticmethod
    def calcActuatorAngle(servoAngle):
        return radians(servoAngle - 204.78)

    @staticmethod
    def calcWristXAngle(servoAngle):
        return radians(150.0 - servoAngle)

    @staticmethod
    def calcWristYAngle(servoAngle):
        return radians(servoAngle - 147.0)

    def __init__(self,
                 arm_config,
                 swing_angle,
                 shoulder_angle,
                 actuator_angle,
                 elbow_angle,
                 elbow2D,
                 wrist2D,
                 effector2D,
                 effector,
                 wrist_x,
                 wrist_y):
        self.cfg = arm_config
        self.swing_angle = swing_angle
        self.shoulder_angle = shoulder_angle
        self.actuator_angle = actuator_angle
        self.elbow_angle = elbow_angle
        # Joints in the arm
        shoulder = rotate(self.cfg.shoulder_offset, swing_angle)
        self.shoulder2D = [self.cfg.shoulder_offset[1], 0]
        self.shoulder = [shoulder[0], 0, shoulder[1]]
        self.wrist2D = wrist2D
        self.effector2D = effector2D
        self.effector = effector
        # Construct the 3D elbow & wrist positions from the 2D (planar) IK
        # solution
        arm_vec = effector - self.shoulder
        arm_vec[1] = 0
        self.elbow2D = elbow2D
        self.elbow = self.shoulder + normalize(arm_vec)*elbow2D[0]
        self.elbow[1] = elbow2D[1]
        self.wrist = self.effector - normalize(arm_vec)*arm_config.wrist_length
        # Wrist pose
        self.wristXAngle = wrist_x
        self.wristYAngle = wrist_y

    def getServoElevator(self):
        return 178.21 - degrees(self.shoulder_angle)

    def getServoActuator(self):
        return degrees(self.actuator_angle) + 204.78

    def getServoSwing(self):
        return 150 - degrees(self.swing_angle)

    def getServoWristX(self):
        return 150 - degrees(self.wristXAngle)

    def getServoWristY(self):
        return 147 + degrees(self.wristYAngle)

    def armDiffAngle(self):
        return degrees(self.shoulder_angle - self.actuator_angle)

    def checkActuator(self):
        angle = self.getServoActuator()
        return angle >= 95 and angle <= 250

    def checkDiff(self):
        angle = self.armDiffAngle()
        return angle >= 44 and angle <= 175

    def checkElevator(self):
        angle = self.getServoElevator()
        return angle >= 60 and angle <= 210

    def checkForearm(self):
        angle = degrees(self.elbow_angle + self.shoulder_angle)
        return angle < 200 and angle > 80

    def checkSwing(self):
        angle = self.getServoSwing()
        return angle >= 60 and angle <= 240

    def checkWristX(self):
        angle = self.getServoWristX()
        return angle >= 60 and angle <= 240

    def checkWristY(self):
        angle = self.getServoWristY()
        return angle >= 60 and angle <= 160

    def checkPositioning(self):
        # When Y>0 Forearm always faces outwards
        if self.wrist2D[1] > 0 and self.wrist2D[0] < self.elbow2D[0]:
            return False
        # No valid positions X<=0
        if self.wrist2D[0] <= 0:
            return False
        # Effector height range
        if self.effector[1] > 180 or self.effector[1] < -200:
            return False
        return True

    def checkClearance(self):
        return (self.checkDiff() and self.checkActuator() and
                self.checkElevator() and self.checkSwing() and
                self.checkWristX() and self.checkWristY() and
                self.checkPositioning() and self.checkForearm())

    def serialize(self):
        """Returns a packed struct holding the pose information"""
        return struct.pack(
            ArmPose.structFormat,
            self.swing_angle,
            self.shoulder_angle,
            self.elbow_angle,
            self.wristXAngle,
            self.wristYAngle
        )

class ArmController:
    def __init__(self,
                 servo_swing,
                 servo_shoulder,
                 servo_elbow,
                 servo_wrist_x,
                 servo_wrist_y,
                 arm_config,
                 motion_enable = False):
        # Solvers are responsible for calculating the target servo positions to
        # reach a given goal position
        self.ik = solvers.IKSolver(
            arm_config.main_length,
            arm_config.forearm_length,
            arm_config.wrist_length,
            arm_config.shoulder_offset)
        self.physsolver = solvers.PhysicalSolver(
            arm_config.main_length,
            arm_config.linkage_length,
            arm_config.lower_actuator_length,
            arm_config.upper_actuator_length)
        # Servos
        self.servos = {}
        self.servos["swing"] = servo_swing
        self.servos["shoulder"] = servo_shoulder
        self.servos["elbow"] = servo_elbow
        self.servos["wrist_x"] = servo_wrist_x
        self.servos["wrist_y"] = servo_wrist_y
        for key, servo in self.servos.iteritems():
            if servo is None:
                print ("Warning: {0} servo not connected".format(key))
            else:
                # Initialise a PID controller for the servo
                if servo.protocol == 1:
                    servo.setGoalSpeed(-MOTORSPEED)
                    servo.data['pid'] = pid.PIDControl(2.4, 0, 0.4)
                else:
                    servo.setGoalSpeed(0)
                servo.data['error'] = 0.0
                # Make sure the goal speed is set
                servo.setTorqueEnable(1)
                if servo.protocol == 1:
                    print("Setting slope")
                    servo.setCWMargin(MOTORMARGIN)
                    servo.setCCWMargin(MOTORMARGIN)
                    servo.setCWSlope(MOTORSLOPE)
                    servo.setCCWSlope(MOTORSLOPE)
        # Store parameters
        self.motion_enable = True
        self.enableMovement(False)
        self.cfg = arm_config
        # Dirty flags for stored poses
        self.ik_pose = None
        self.ik_dirty = True
        self.real_pose = None
        self.real_dirty = True
        # Current target pose
        self.target_pose = None

    def enableMovement(self, enable):
        changed = False
        if enable and not self.motion_enable:
            print ("Warning: Arm enabled")
            self.motion_enable = True
            changed = True
        elif not enable:
            self.motion_enable = False
            changed = True
        if changed:
            # Set servos on/off
            if self.servos['swing'] is not None:
                self.servos['swing'].setTorqueEnable(self.motion_enable)
            if self.servos['shoulder'] is not None:
                self.servos['shoulder'].setTorqueEnable(self.motion_enable)
            if self.servos['elbow'] is not None:
                self.servos['elbow'].setTorqueEnable(self.motion_enable)
            if self.servos['wrist_x'] is not None:
                self.servos['wrist_x'].setTorqueEnable(self.motion_enable)
            if self.servos['wrist_y'] is not None:
                self.servos['wrist_y'].setTorqueEnable(self.motion_enable)

    def setWristGoalPosition(self, pos):
        self.ik.setGoal(pos)
        self.ik_dirty = True

    def setWristGoalDirection(self, normal):
        self.ik.setWristDir(normal)
        self.ik_dirty = True

    def getIKPose(self):
        if self.ik_dirty and self.ik.valid:
            # Construct geometry of arm from IK state
            main_arm = self.ik.elbow - self.ik.originpl
            arm_vert_angle = sigangle(main_arm, vertical)
            forearm = self.ik.wristpl - self.ik.elbow
            elbow_angle = angle_between(main_arm, forearm)
            # Solve actuator angle for given elbow angle
            # Base angle is between the main arm and actuator
            base_angle = self.physsolver.inverse_forearm(elbow_angle)
            actuator_angle = arm_vert_angle - base_angle

            self.ik_pose = ArmPose(
                self.cfg,
                swing_angle = self.ik.swing,
                # angles from vertical
                shoulder_angle = arm_vert_angle,
                actuator_angle = actuator_angle,
                # angle between the main arm and forearm
                elbow_angle = elbow_angle,
                elbow2D = self.ik.elbow,
                wrist2D = self.ik.wristpl,
                effector2D = self.ik.goalpl,
                effector = self.ik.goal,
                wrist_x = self.ik.wrist_x,
                wrist_y = self.ik.wrist_y
            )
        return self.ik_pose

    def pollServos(self):
        """Poll the real-world servo positions"""
        for servo in self.servos.itervalues():
            if servo is not None:
                newPos = servo.getPosition()
                if type(newPos) is float:
                    servo.data['pos'] = newPos

    def clearPositionError(self):
        """Clears the servo's position-error accumulators"""
        for servo in self.servos.itervalues():
            if servo is not None and servo.protocol == 1:
                servo.data['error'] = 0.0

    def getRealPose(self):
        """Retrieve the real-world arm pose, or None if not all servos are
        connected.
        """
        if any([servo is None for servo in self.servos.itervalues()]):
            return None

        # This whole function is essentially just FK based on the known servo
        # angles
        swing_servo = self.servos['swing'].data['pos']
        elevator_servo = self.servos['shoulder'].data['pos']
        actuator_servo = self.servos['elbow'].data['pos']
        wrist_x_servo = self.servos['wrist_x'].data['pos']
        wrist_y_servo = self.servos['wrist_y'].data['pos']

        # Find the internal arm-pose angles for the given servo positions
        swing_angle = ArmPose.calcSwingAngle(swing_servo)
        elevator_angle = ArmPose.calcElevatorAngle(elevator_servo)
        actuator_angle = ArmPose.calcActuatorAngle(actuator_servo)
        wrist_x_angle = ArmPose.calcWristXAngle(wrist_x_servo)
        wrist_y_angle = ArmPose.calcWristYAngle(wrist_y_servo)
        # Solve elbow angle for given actuator and elevator angles
        # (this is the angle from the elevator arm's direction to the forearm's)
        elbow_angle = self.physsolver.solve_forearm(elevator_angle, actuator_angle)

        # FK positions from config and angles
        offset = self.cfg.shoulder_offset
        shoulder2D = np.array([offset[1], 0])
        elbow2D = shoulder2D + rotate(vertical, elevator_angle)*self.cfg.main_length
        wrist2D = elbow2D + rotate(vertical, elevator_angle + elbow_angle)*self.cfg.forearm_length
        effector2D = wrist2D + [self.cfg.wrist_length, 0]
        # 3D Effector calculation is a little more involved
        td = rotate([offset[0], effector2D[0]], swing_angle)
        effector = np.array([td[0], effector2D[1], td[1]])

        pose = ArmPose(
            self.cfg,
            swing_angle, elevator_angle, actuator_angle,
            elbow_angle, elbow2D, wrist2D, effector2D,
            effector, wrist_x_angle, wrist_y_angle)
        return pose

    def setTargetPose(self, new_pose):
        self.target_pose = new_pose

    def tick(self):
        if self.target_pose is not None:
            if self.motion_enable:
                # Drive servos
                gain = 0.1
                if self.servos['swing'] is not None:
                    s = self.servos['swing']
                    pos = s.data['pos']
                    target = self.target_pose.getServoSwing()
                    # err = min(10, pos-target)
                    # s.data['error'] += err*gain
                    s.setGoalPosition(target)
                if self.servos['shoulder'] is not None:
                    s = self.servos['shoulder']
                    # cumulative error
                    pos = s.data['pos']
                    target = self.target_pose.getServoElevator()
                    err = min(10, pos-target)
                    s.data['error'] += err*gain
                    s.data['error'] = np.clip(s.data['error'], -ERRORLIM, ERRORLIM)
                    s.setGoalPosition(target - s.data['error'])
                if self.servos['elbow'] is not None:
                    s = self.servos['elbow']
                    pos = s.data['pos']
                    target = self.target_pose.getServoActuator()
                    err = min(10, pos-target)
                    s.data['error'] += err*gain
                    s.data['error'] = np.clip(s.data['error'], -ERRORLIM, ERRORLIM)
                    s.setGoalPosition(target - s.data['error'])

                if self.servos['wrist_x'] is not None:
                    self.servos['wrist_x'].setGoalPosition(self.target_pose.getServoWristX())
                if self.servos['wrist_y'] is not None:
                    self.servos['wrist_y'].setGoalPosition(self.target_pose.getServoWristY())

[(170, 'struct.pack', 'struct.pack', False, 'import struct\n'), (190, 'solvers.IKSolver', 'solvers.IKSolver', False, 'import solvers\n'), (195, 'solvers.PhysicalSolver', 'solvers.PhysicalSolver', False, 'import solvers\n'), (338, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (344, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (214, 'pid.PIDControl', 'pid.PIDControl', False, 'import pid\n'), (375, 'numpy.clip', 'np.clip', True, 'import numpy as np\n'), (383, 'numpy.clip', 'np.clip', True, 'import numpy as np\n')]