import math
import time
from robodk.robolink import (
    ITEM_TYPE_ROBOT,
    ROBOTCOM_READY,
    RUNMODE_RUN_ROBOT,
    RUNMODE_SIMULATE,
    Robolink,
)
from robodk.robomath import rotx, transl, xyzrpw_2_pose, roty


class Robot:
    """Run one motion-only pass around the outside of the part."""

    ROBOT_SIMULATION = False
    PASS_X = 100
    PASS_RADIUS = 32.5
    NORMAL_SPEED = (250, 125)  # mm/s, deg/s
    SCAN_SPEED = (250, 125)  # mm/s, deg/s
    START_ANGLE = -90
    ARC_ANGLES = ((0, 90),)
    # ARC_ANGLES = ((0, 90), (180, 270))

    def __init__(self):
        self.object_frame = xyzrpw_2_pose([850.0, 1.3556660681488173, -230.2612160099226, 0.0, -1.273, -0.7020000000000002]) # Object frame pose (pipe frame)
        self.tool_pose = xyzrpw_2_pose([-84.196000, 16.500000, 314.251000, 0.000000, -90.000000, 180.000000])

    def initialize(self):
        self.rdk = Robolink()
        self.robot = self.rdk.ItemUserPick("Select a robot", ITEM_TYPE_ROBOT)
        if not self.robot.Valid():
            raise RuntimeError("No robot selected or available")

        if self.ROBOT_SIMULATION:
            self.rdk.setRunMode(RUNMODE_SIMULATE)
            return

        if not self.robot.Connect():
            raise RuntimeError("Could not start the robot connection")

        status, status_message = self.robot.ConnectedState()
        if status != ROBOTCOM_READY:
            raise RuntimeError(f"Failed to connect: {status_message}")

        self.rdk.setRunMode(RUNMODE_RUN_ROBOT)

    def target_at(self, angle_degrees):
        angle = math.radians(angle_degrees)
        y = self.PASS_RADIUS * math.sin(angle)
        z = self.PASS_RADIUS * math.cos(angle)
        return transl(self.PASS_X, y, z) * rotx(-angle) * roty(math.radians(-5))

    def centered_start_joints(self):
        """Choose the valid start IK solution closest to joint zero."""
        solutions = self.robot.SolveIK_All(
            self.target_at(self.START_ANGLE),
            self.tool_pose,
            self.object_frame,
        )
        lower_matrix, upper_matrix, _ = self.robot.JointLimits()
        lower = lower_matrix.list()
        upper = upper_matrix.list()
        joint_count = len(lower)

        candidates = []
        for column in range(solutions.size()[1]):
            joints = [solutions[joint, column] for joint in range(joint_count)]
            if all(
                lower[joint] <= joints[joint] <= upper[joint]
                for joint in range(joint_count)
            ):
                candidates.append(joints)

        if not candidates:
            raise RuntimeError("The scan start target has no valid IK solution")

        return min(candidates, key=lambda joints: sum(abs(value) for value in joints))

    def run_one_pass(self):
        print(
            f"Running one pass at X={self.PASS_X} mm and "
            f"radius={self.PASS_RADIUS} mm"
        )
        self.robot.setPoseTool(self.tool_pose)
        self.robot.setPoseFrame(self.object_frame)
        self.robot.setSpeed(*self.NORMAL_SPEED)
        self.robot.setRounding(100)

        print(f"Moving to the {self.START_ANGLE:g}-degree start position")
        self.robot.MoveJ(self.centered_start_joints())

        start_time = time.time()
        self.robot.setSpeed(*self.SCAN_SPEED)
        for index, (via_angle, end_angle) in enumerate(self.ARC_ANGLES, start=1):
            print(
                f"Arc {index}/{len(self.ARC_ANGLES)}: "
                f"via {via_angle:g} to {end_angle:g} degrees"
            )
            self.robot.MoveC(
                self.target_at(via_angle),
                self.target_at(end_angle),
                blocking=True,
            )
        print(f"Pass completed in {time.time() - start_time:.2f} seconds")

    def main(self):
        self.initialize()
        self.run_one_pass()
        print("Motion pass complete")


if __name__ == "__main__":
    Robot().main()