import time
import cv2
from ultralytics import YOLO

from robolink import *
from robodk import *

# ==========================================================
# ROBODK CONNECTION
# ==========================================================
RDK = Robolink()

# ==========================================================
# ROBOTS
# ==========================================================
robotA = RDK.Item('JAKA Zu7', ITEM_TYPE_ROBOT)
robotB = RDK.Item('JAKA Zu12', ITEM_TYPE_ROBOT)

# ==========================================================
# GRIPPERS
# ==========================================================
gripperA = RDK.Item('RG6_Gripper_Zu7', ITEM_TYPE_ROBOT)
gripperB = RDK.Item('2FG7_Gripper_Zu12', ITEM_TYPE_ROBOT)

needle = RDK.Item('Needle', ITEM_TYPE_OBJECT)
needle_holder = RDK.Item('Needle Holder', ITEM_TYPE_OBJECT)

GRIPPER_OPEN = [80]      # Fully open
GRIPPER_GRASP = [13]     # Needle grasp position

# ==========================================================
# TOOLS
# ==========================================================
toolA = RDK.Item('Tool 1')
toolB = RDK.Item('Tool 2')

# ==========================================================
# TARGETS - ZU7 / ROBOT A
# ==========================================================
A_Home = RDK.Item('A_Home', ITEM_TYPE_TARGET)
A_Approach_Needle_Pickup = RDK.Item('A_Approach_Needle_Pickup', ITEM_TYPE_TARGET)
A_Needle_Pickup = RDK.Item('A_Needle_Pickup', ITEM_TYPE_TARGET)
A_Needle_Drop = RDK.Item('A_Needle_Drop', ITEM_TYPE_TARGET)
A_Receive_Approach = RDK.Item('A_Receive_Approach_new', ITEM_TYPE_TARGET)
A_Receive = RDK.Item('A_Receive_new', ITEM_TYPE_TARGET)
A_Intermediate = RDK.Item('A_Intermediate', ITEM_TYPE_TARGET)
A_Intermediate2 = RDK.Item('A_Intermediate2', ITEM_TYPE_TARGET)
A_Handover_Approach = RDK.Item('A_Handover_Approach', ITEM_TYPE_TARGET)
A_Handover = RDK.Item('A_Handover', ITEM_TYPE_TARGET)

# ==========================================================
# TARGETS - ZU12 / ROBOT B
# ==========================================================
B_Home = RDK.Item('B_Home', ITEM_TYPE_TARGET)
B_Intermediate = RDK.Item('B_Intermediate', ITEM_TYPE_TARGET)
B_Receive_Approach = RDK.Item('B_Receive_Approach', ITEM_TYPE_TARGET)
B_Receive = RDK.Item('B_Receive', ITEM_TYPE_TARGET)
B_Intermediate2 = RDK.Item('B_Intermediate2', ITEM_TYPE_TARGET)
B_Handover_Approach = RDK.Item('B_Handover_Approach', ITEM_TYPE_TARGET)
B_Handover = RDK.Item('B_Handover', ITEM_TYPE_TARGET)

# ==========================================================
# ROBOT SETTINGS
# ==========================================================
robotA.setSpeed(100, 50)
robotA.setPoseTool(toolA)

robotB.setSpeed(100, 50)
robotB.setPoseTool(toolB)

# ==========================================================
# YOLO MODEL
# ==========================================================
MODEL_PATH = (r"D:\MSc Robotics and Automation\Individual Research Project\Machine Vision\VS_Environment\needle_yolov11n_gripping_colab.pt")

model = YOLO(MODEL_PATH)

# ==========================================================
# WEBCAM
# ==========================================================
CAMERA_INDEX = 1

cap = cv2.VideoCapture(CAMERA_INDEX)

if not cap.isOpened():
    raise Exception("ERROR: Could not open webcam.")

cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)

# ==========================================================
# YOLO SETTINGS
# ==========================================================
CONFIDENCE = 0.25

# Number of consecutive frames required
# before accepting a detection.
STABLE_FRAMES = 5

# ==========================================================
# CLASS DEFINITIONS
# ==========================================================
GRIPPER_A_CLOSE = 0
GRIPPER_A_OPEN = 1

GRIPPER_B_CLOSE = 2
GRIPPER_B_OPEN = 3

NEEDLE = 4

# =====================================================
# GRIPPER FUNCTIONS
# =====================================================
def open_gripper(gripper):
    gripper.MoveJ(GRIPPER_OPEN)
    gripper.WaitMove()
    pause(0.2)

def grasp_needle(gripper):
    gripper.MoveJ(GRIPPER_GRASP)
    gripper.WaitMove()
    pause(0.2)

# ==========================================================
# VISION FUNCTION
# ==========================================================
def get_detected_classes(frame):
    """
    Run YOLO on one frame and return the detected class IDs.
    """

    results = model(
        frame,
        imgsz=640,
        conf=CONFIDENCE,
        verbose=False
    )

    detected_classes = set()

    result = results[0]

    if result.boxes is not None:

        for box in result.boxes:

            class_id = int(box.cls[0])
            confidence = float(box.conf[0])

            if confidence >= CONFIDENCE:
                detected_classes.add(class_id)

    return detected_classes, result

# ==========================================================
# WAIT FOR VISION CONDITION
# ==========================================================
def wait_for_vision(required_classes, description, stable_frames=STABLE_FRAMES):
    """
    Wait until all required classes are detected
    for several consecutive frames.

    Example:

    wait_for_vision(
        {GRIPPER_A_CLOSE, NEEDLE},
        "Needle held by Gripper A"
    )
    """

    print()
    print("==========================================")
    print("VISION CHECK")
    print(description)
    print("==========================================")

    consecutive_frames = 0

    while True:

        ret, frame = cap.read()

        if not ret:

            print("ERROR: Could not read webcam frame.")

            time.sleep(0.1)
            continue

        detected_classes, result = get_detected_classes(frame)

        # --------------------------------------------------
        # Check whether all required classes are detected
        # --------------------------------------------------
        condition_met = required_classes.issubset(
            detected_classes
        )

        if condition_met:

            consecutive_frames += 1

            print(
                f"Vision condition detected "
                f"({consecutive_frames}/{stable_frames})"
            )

        else:

            consecutive_frames = 0

        # --------------------------------------------------
        # Display detection
        # --------------------------------------------------
        annotated_frame = result.plot()

        cv2.imshow(
            "YOLO11n Vision - Robot Handover",
            annotated_frame
        )

        # --------------------------------------------------
        # Confirm stable detection
        # --------------------------------------------------
        if consecutive_frames >= stable_frames:

            print()
            print("VISION CONDITION CONFIRMED")
            print(description)

            cv2.waitKey(1)

            return True

        # --------------------------------------------------
        # Q = stop program
        # --------------------------------------------------
        if cv2.waitKey(1) & 0xFF == ord('q'):

            print("Vision stopped by user.")

            raise KeyboardInterrupt

        time.sleep(0.01)

# ==========================================================
# MAIN ROBOT SEQUENCE
# ==========================================================
try:

    while True:

        # ==================================================
        # 1. ROBOT A - ZU7
        # PICK NEEDLE
        # ==================================================
        robotA.setDO(1, 0)
        time.sleep(2.0)

        print("Moving Zu7 to A_Home...")
        robotA.MoveJ(A_Home)
        robotA.WaitMove()

        print("Moving Zu7 to A_Approach_Needle_Pickup...")
        robotA.MoveJ(A_Approach_Needle_Pickup)
        robotA.WaitMove()

        print("Moving Zu7 to A_Needle_Pickup...")
        robotA.MoveL(A_Needle_Pickup)
        robotA.WaitMove()

        robotA.setDO(1, 1)
        time.sleep(2.0)

        grasp_needle(gripperA)
        needle.setParentStatic(toolA)

        print("Moving Zu7 back to A_Approach_Needle_Pickup...")
        robotA.MoveL(A_Approach_Needle_Pickup)
        robotA.WaitMove()

        print("Moving Zu7 to A_Receive_Approach...")
        robotA.MoveJ(A_Receive_Approach)
        robotA.WaitMove()

        print("Moving Zu7 to A_Receive...")
        robotA.MoveL(A_Receive)
        robotA.WaitMove()

        robotB.setDO(1, 0)
        time.sleep(2.0)

        # open_gripper(gripperB)

        print("Moving Zu12 to B_Home...")
        robotB.MoveJ(B_Home)
        robotB.WaitMove()

        print("Moving Zu12 to B_Intermediate...")
        robotB.MoveL(B_Intermediate)
        robotB.WaitMove()

        print("Moving Zu12 to B_Receive_Approach...")
        robotB.MoveJ(B_Receive_Approach)
        robotB.WaitMove()

        # ==================================================
        # 3. VISION CHECK
        #
        # Robot B is now at B_Receive_Approach.
        #
        # Camera on Zu12 checks:
        #
        #   Class 0 = gripper_a_close
        #   Class 4 = needle
        #
        # Both must be visible.
        # ==================================================
        wait_for_vision(
            {
                GRIPPER_A_CLOSE,
                NEEDLE
            },
            "Confirming needle is held by Gripper A"
        )

        # ==================================================
        # 4. ROBOT B MOVES TO RECEIVE
        # ==================================================
        print("Needle confirmed in Gripper A.")

        print("Moving Zu12 to B_Receive...")
        robotB.MoveL(B_Receive)
        robotB.WaitMove()

        robotB.setDO(1, 1)
        time.sleep(2.0)

        grasp_needle(gripperB)
        needle.setParentStatic(toolB)

        # ==================================================
        # 5. VISION CHECK
        #
        # Robot B has closed its gripper.
        #
        # Class 2 = gripper_b_close
        #
        # Wait until YOLO confirms Gripper B is closed.
        # ==================================================
        wait_for_vision(
            {
                GRIPPER_B_CLOSE
            },
            "Confirming Gripper B is closed"
        )

        # ==================================================
        # 6. ROBOT A RELEASES NEEDLE
        # ==================================================
        print("Gripper B confirmed closed.")

        print("Opening Gripper A...")

        robotA.setDO(1, 0)
        time.sleep(2.0)

        open_gripper(gripperA)

        # ==================================================
        # 7. VISION CHECK
        #
        # Confirm Gripper A is open.
        #
        # Class 1 = gripper_a_open
        # ==================================================
        wait_for_vision(
            {
                GRIPPER_A_OPEN
            },
            "Confirming Gripper A is open"
        )

        # ==================================================
        # 8. ROBOT A MOVES AWAY
        # ==================================================
        print("Moving Zu7 back to A_Receive_Approach...")
        robotA.MoveL(A_Receive_Approach)
        robotA.WaitMove()

        print("Moving Zu7 to A_Approach_Needle_Pickup...")
        robotA.MoveJ(A_Approach_Needle_Pickup)
        robotA.WaitMove()

        print("Moving Zu7 to A_Home...")
        robotA.MoveL(A_Home)
        robotA.WaitMove()

        # ==================================================
        # 9. ROBOT B NOW HOLDS THE NEEDLE
        #
        # Camera is no longer needed for the next movements
        # because Robot A has moved away from the camera view.
        # ==================================================
        print("Needle is now held by Gripper B.")

        print("Moving Zu12 to B_Receive_Approach...")
        robotB.MoveL(B_Receive_Approach)
        robotB.WaitMove()

        print("Moving Zu12 to B_Intermediate...")
        robotB.MoveL(B_Intermediate)
        robotB.WaitMove()

        print("Moving Zu12 to B_Home...")
        robotB.MoveJ(B_Home)
        robotB.WaitMove()

        # ==================================================
        # 10. ROBOT B MOVES TO SECOND HANDOVER
        # ==================================================
        print("Moving Zu12 to B_Intermediate2...")
        robotB.MoveJ(B_Intermediate2)
        robotB.WaitMove()

        print("Moving Zu12 to B_Handover_Approach...")
        robotB.MoveJ(B_Handover_Approach)
        robotB.WaitMove()

        print("Moving Zu12 to B_Handover...")
        robotB.MoveL(B_Handover)
        robotB.WaitMove()

        # ==================================================
        # 11. ROBOT A MOVES TO SECOND HANDOVER
        # ==================================================
        print("Moving Zu7 to A_Intermediate...")
        robotA.MoveJ(A_Intermediate)
        robotA.WaitMove()

        print("Moving Zu7 to A_Intermediate2...")
        robotA.MoveJ(A_Intermediate2)
        robotA.WaitMove()

        print("Moving Zu7 to A_Handover_Approach...")
        robotA.MoveL(A_Handover_Approach)
        robotA.WaitMove()

        print("Moving Zu7 to A_Handover...")
        robotA.MoveJ(A_Handover)
        robotA.WaitMove()

        # ==================================================
        # 12. ROBOT A CLOSES ITS GRIPPER
        # ==================================================
        robotA.setDO(1, 1)
        time.sleep(2.0)

        grasp_needle(gripperA)
        needle.setParentStatic(toolA)

        # ==================================================
        # 13. VISION CHECK
        #
        # Confirm Gripper A has closed.
        #
        # Class 0 = gripper_a_close
        #
        # We can also require needle = class 4 to make sure
        # the needle is still visible during the handover.
        # ==================================================
        wait_for_vision(
            {
                GRIPPER_A_CLOSE,
                NEEDLE
            },
            "Confirming Gripper A has closed around the needle"
        )

        # ==================================================
        # 14. ROBOT B RELEASES NEEDLE
        # ==================================================
        print("Gripper A confirmed closed.")

        print("Opening Gripper B...")

        robotB.setDO(1, 0)
        time.sleep(2.0)

        # open_gripper(gripperB)

        # ==================================================
        # 15. VISION CHECK
        #
        # Confirm Gripper B is open.
        #
        # Class 3 = gripper_b_open
        # ==================================================
        wait_for_vision(
            {
                GRIPPER_B_OPEN
            },
            "Confirming Gripper B is open"
        )

        # ==================================================
        # 16. ROBOT B MOVES AWAY
        # ==================================================
        print("Moving Zu12 to B_Handover_Approach...")
        robotB.MoveL(B_Handover_Approach)
        robotB.WaitMove()

        print("Moving Zu12 to B_Intermediate2...")
        robotB.MoveL(B_Intermediate2)
        robotB.WaitMove()

        print("Moving Zu12 to B_Home...")
        robotB.MoveJ(B_Home)
        robotB.WaitMove()

        # ==================================================
        # 17. ROBOT A RETURNS
        # ==================================================
        print("Moving Zu7 to A_Handover_Approach...")
        robotA.MoveJ(A_Handover_Approach)
        robotA.WaitMove()

        print("Moving Zu7 to A_Intermediate2...")
        robotA.MoveL(A_Intermediate2)
        robotA.WaitMove()

        print("Moving Zu7 to A_Intermediate...")
        robotA.MoveJ(A_Intermediate)
        robotA.WaitMove()

        print("Moving Zu7 to A_Home...")
        robotA.MoveJ(A_Home)
        robotA.WaitMove()

        # ==================================================
        # 18. ROBOT A RETURNS NEEDLE TO HOLDER
        # ==================================================
        print("Moving Zu7 to A_Approach_Needle_Pickup...")
        robotA.MoveJ(A_Approach_Needle_Pickup)
        robotA.WaitMove()

        print("Moving Zu7 to A_Needle_Drop...")
        robotA.MoveL(A_Needle_Drop)
        robotA.WaitMove()

        robotA.setDO(1, 0)
        time.sleep(2.0)

        open_gripper(gripperA)
        needle.setParentStatic(needle_holder)

        print("Moving Zu7 back to A_Approach_Needle_Pickup...")
        robotA.MoveL(A_Approach_Needle_Pickup)
        robotA.WaitMove()

        # ==================================================
        # COMPLETE CYCLE
        # ==================================================
        print()
        print("==========================================")
        print("HANDOVER CYCLE COMPLETED")
        print("==========================================")
        print()

except KeyboardInterrupt:

    print()
    print("Program stopped by user.")

finally:
    # ======================================================
    # CLEANUP CAMERA
    # ======================================================
    cap.release()
    cv2.destroyAllWindows()

    print("Camera released.")