import numpy as np
import math

class RoboticArmIK:
    def __init__(self):
        # DH parameters for 6-DOF arm
        self.a = [0, 150, 120, 0, 0, 0]        # Link lengths (mm)
        self.d = [100, 0, 0, 95, 0, 60]        # Link offsets (mm)
        self.alpha = [90, 0, 0, 90, -90, 0]    # Link twists (degrees)
        
        # Convert degrees to radians
        self.alpha = [math.radians(a) for a in self.alpha]
        
    def forward_kinematics(self, joint_angles):
        """
        Calculate forward kinematics using DH parameters
        """
        T = np.eye(4)
        
        for i in range(6):
            theta = joint_angles[i]
            
            # DH transformation matrix
            ct = math.cos(theta)
            st = math.sin(theta)
            ca = math.cos(self.alpha[i])
            sa = math.sin(self.alpha[i])
            
            T_i = np.array([
                [ct, -st*ca,  st*sa, self.a[i]*ct],
                [st,  ct*ca, -ct*sa, self.a[i]*st],
                [0,   sa,     ca,    self.d[i]],
                [0,   0,      0,     1]
            ])
            
            T = np.dot(T, T_i)
        
        return T
        
    def inverse_kinematics(self, target_pos, target_orient):
        """
        Calculate inverse kinematics using geometric approach
        """
        x, y, z = target_pos
        
        # Calculate joint 1 (base rotation)
        theta1 = math.atan2(y, x)
        
        # Calculate wrist center position
        r = math.sqrt(x**2 + y**2)
        wrist_z = z - self.d[5]
        wrist_r = r
        
        # Calculate joint 3 (elbow)
        D = (wrist_r**2 + (wrist_z - self.d[0])**2 - 
             self.a[1]**2 - self.a[2]**2) / (2 * self.a[1] * self.a[2])
        
        if abs(D) > 1:
            return None  # No solution exists
            
        theta3 = math.atan2(math.sqrt(1 - D**2), D)
        
        # Calculate joint 2 (shoulder)
        s3 = math.sin(theta3)
        c3 = math.cos(theta3)
        
        k1 = self.a[1] + self.a[2] * c3
        k2 = self.a[2] * s3
        
        theta2 = math.atan2(wrist_z - self.d[0], wrist_r) - math.atan2(k2, k1)
        
        # Calculate remaining joints based on orientation
        # This is simplified - full implementation would include orientation
        theta4 = 0  # Wrist pitch
        theta5 = 0  # Wrist roll
        theta6 = 0  # End-effector rotation
        
        return [theta1, theta2, theta3, theta4, theta5, theta6]
        
    def check_joint_limits(self, joint_angles):
        """
        Check if joint angles are within limits
        """
        limits = [
            (-180, 180),  # Base
            (-90, 90),    # Shoulder
            (-180, 0),    # Elbow
            (-180, 180),  # Wrist 1
            (-90, 90),    # Wrist 2
            (-180, 180)   # Wrist 3
        ]
        
        for i, (angle, (min_limit, max_limit)) in enumerate(zip(joint_angles, limits)):
            angle_deg = math.degrees(angle)
            if angle_deg < min_limit or angle_deg > max_limit:
                return False, f"Joint {i+1} out of range: {angle_deg:.1f}°"
        
        return True, "All joints within limits"

# Example usage
if __name__ == "__main__":
    arm = RoboticArmIK()
    
    # Test forward kinematics
    joint_angles = [0, math.radians(45), math.radians(-90), 0, 0, 0]
    end_effector_pose = arm.forward_kinematics(joint_angles)
    
    print("End effector position:")
    print(f"X: {end_effector_pose[0,3]:.1f} mm")
    print(f"Y: {end_effector_pose[1,3]:.1f} mm")
    print(f"Z: {end_effector_pose[2,3]:.1f} mm")
    
    # Test inverse kinematics
    target_pos = [200, 100, 150]
    target_orient = [0, 0, 0]  # Simplified
    
    solution = arm.inverse_kinematics(target_pos, target_orient)
    if solution:
        print("\nInverse kinematics solution:")
        for i, angle in enumerate(solution):
            print(f"Joint {i+1}: {math.degrees(angle):.1f}°")
    else:
        print("\nNo valid solution found")

import cv2
import numpy as np
import torch
from ultralytics import YOLO

class ObjectDetectionSystem:
    def __init__(self, model_path='yolov8n.pt'):
        self.model = YOLO(model_path)
        self.camera = cv2.VideoCapture(0)
        
        # Camera calibration parameters (example values)
        self.camera_matrix = np.array([
            [800, 0, 320],
            [0, 800, 240],
            [0, 0, 1]
        ], dtype=np.float32)
        
        self.dist_coeffs = np.array([0.1, -0.2, 0, 0, 0], dtype=np.float32)
        
        # Object classes we're interested in
        self.target_classes = ['bottle', 'cup', 'book', 'cell phone']
        
    def capture_frame(self):
        """Capture a frame from the camera"""
        ret, frame = self.camera.read()
        if ret:
            return cv2.undistort(frame, self.camera_matrix, self.dist_coeffs)
        return None
        
    def detect_objects(self, frame):
        """Detect objects in the frame using YOLO"""
        results = self.model(frame)
        detections = []
        
        for result in results:
            boxes = result.boxes
            if boxes is not None:
                for box in boxes:
                    # Get bounding box coordinates
                    x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                    confidence = box.conf[0].cpu().numpy()
                    class_id = int(box.cls[0].cpu().numpy())
                    class_name = self.model.names[class_id]
                    
                    if class_name in self.target_classes and confidence > 0.5:
                        # Calculate center point
                        center_x = int((x1 + x2) / 2)
                        center_y = int((y1 + y2) / 2)
                        
                        detection = {
                            'class': class_name,
                            'confidence': confidence,
                            'bbox': (int(x1), int(y1), int(x2), int(y2)),
                            'center': (center_x, center_y),
                            '3d_position': self.estimate_3d_position(center_x, center_y)
                        }
                        detections.append(detection)
        
        return detections
        
    def estimate_3d_position(self, pixel_x, pixel_y, estimated_depth=500):
        """
        Convert 2D pixel coordinates to 3D world coordinates
        This is simplified - real implementation would use stereo vision or depth sensor
        """
        # Convert pixel coordinates to camera coordinates
        fx = self.camera_matrix[0, 0]
        fy = self.camera_matrix[1, 1]
        cx = self.camera_matrix[0, 2]
        cy = self.camera_matrix[1, 2]
        
        # Calculate 3D coordinates (relative to camera)
        x_cam = (pixel_x - cx) * estimated_depth / fx
        y_cam = (pixel_y - cy) * estimated_depth / fy
        z_cam = estimated_depth
        
        # Transform to robot base coordinates
        # This transformation depends on camera mounting position
        # Example transformation (camera mounted above workspace)
        x_robot = x_cam + 0      # Camera X offset
        y_robot = y_cam + 200    # Camera Y offset  
        z_robot = 100 - z_cam    # Camera height - object depth
        
        return (x_robot, y_robot, z_robot)
        
    def draw_detections(self, frame, detections):
        """Draw detection results on the frame"""
        for detection in detections:
            x1, y1, x2, y2 = detection['bbox']
            center_x, center_y = detection['center']
            
            # Draw bounding box
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
            
            # Draw center point
            cv2.circle(frame, (center_x, center_y), 5, (0, 0, 255), -1)
            
            # Draw label
            label = f"{detection['class']}: {detection['confidence']:.2f}"
            cv2.putText(frame, label, (x1, y1-10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
            
            # Draw 3D coordinates
            pos_3d = detection['3d_position']
            coord_text = f"({pos_3d[0]:.0f}, {pos_3d[1]:.0f}, {pos_3d[2]:.0f})"
            cv2.putText(frame, coord_text, (x1, y2+20), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)
        
        return frame
        
    def get_pickup_targets(self):
        """Get list of objects suitable for pickup"""
        frame = self.capture_frame()
        if frame is None:
            return []
            
        detections = self.detect_objects(frame)
        
        # Filter detections based on position and accessibility
        pickup_targets = []
        for detection in detections:
            x, y, z = detection['3d_position']
            
            # Check if object is within robot's reach
            distance = np.sqrt(x**2 + y**2)
            if distance < 300 and z > 0:  # Within 300mm reach and above table
                pickup_targets.append(detection)
        
        return pickup_targets
        
    def release(self):
        """Clean up resources"""
        self.camera.release()
        cv2.destroyAllWindows()

# Example usage
if __name__ == "__main__":
    detector = ObjectDetectionSystem()
    
    try:
        while True:
            frame = detector.capture_frame()
            if frame is None:
                break
                
            detections = detector.detect_objects(frame)
            frame_with_detections = detector.draw_detections(frame, detections)
            
            cv2.imshow('Object Detection', frame_with_detections)
            
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
                
    finally:
        detector.release()