import torch
from ultralytics import YOLO

# Load a pre-trained YOLOv8 model
model = YOLO('yolov8n.pt')

# Fine-tune the model on our custom PVC pipe dataset
results = model.train(
    data='pvc_pipe_dataset.yaml',
    epochs=100,
    imgsz=640,
    batch=16,
    device=0  # Use GPU for training
)

import cv2

def detect_pvc_pipes(frame):
    """
    Detect PVC pipes in the input frame using our fine-tuned YOLOv8 model.
    
    Args:
        frame: Input image frame from the robot's camera
        
    Returns:
        List of detected pipe bounding boxes in format [x_min, y_min, x_max, y_max, confidence]
    """
    # Run YOLOv8 inference on the frame
    results = model(frame)
    
    # Process the results to extract pipe detections
    detections = []
    for result in results:
        boxes = result.boxes
        for box in boxes:
            # Filter by class (0 = PVC pipe in our model) and confidence
            if box.cls == 0 and box.conf > 0.7:
                x_min, y_min, x_max, y_max = box.xyxy[0].tolist()
                confidence = box.conf.item()
                detections.append([x_min, y_min, x_max, y_max, confidence])
    
    return detections

import numpy as np

def convert_2d_to_3d(detections, camera_matrix, depth_map=None):
    """
    Convert 2D bounding box detections to 3D world coordinates.
    
    Args:
        detections: List of [x_min, y_min, x_max, y_max, confidence]
        camera_matrix: 3x3 camera intrinsic matrix
        depth_map: Optional depth map from stereo vision or depth sensor
        
    Returns:
        List of pipe centers in 3D coordinates [x, y, z, diameter]
    """
    pipe_3d_coords = []
    
    # Extract camera parameters
    fx = camera_matrix[0, 0]
    fy = camera_matrix[1, 1]
    cx = camera_matrix[0, 2]
    cy = camera_matrix[1, 2]
    
    for detection in detections:
        x_min, y_min, x_max, y_max, conf = detection
        
        # Calculate center of the bounding box
        center_x = (x_min + x_max) / 2
        center_y = (y_min + y_max) / 2
        
        # Estimate pipe diameter in pixels
        diameter_pixels = max(x_max - x_min, y_max - y_min)
        
        # Get depth at the center point (either from depth map or using default)
        if depth_map is not None:
            depth = depth_map[int(center_y), int(center_x)]
        else:
            # Fallback to estimation based on known pipe diameter
            depth = 500  # mm
        
        # Convert from image to world coordinates
        x = (center_x - cx) * depth / fx
        y = (center_y - cy) * depth / fy
        z = depth
        
        # Estimate actual pipe diameter in mm
        diameter_mm = diameter_pixels * depth / fx
        
        pipe_3d_coords.append([x, y, z, diameter_mm])
    
    return pipe_3d_coords

import numpy as np
from scipy.optimize import minimize

class AdaptiveGripForceOptimizer:
    def __init__(self, mu=0.5, ext_force=10, torque_limit=5):
        """
        Initialize the gripper force optimizer.
        
        Args:
            mu: Friction coefficient between gripper and PVC pipe
            ext_force: External force (primarily gravity) in Newtons
            torque_limit: Maximum allowable joint torque in Nm
        """
        self.mu = mu
        self.ext_force = ext_force
        self.torque_limit = torque_limit
        self.p = 0.1  # Distance from center to contact point in meters
        
    def objective(self, F):
        """Minimize the maximum applied force."""
        return max(abs(F))

    def force_balance(self, F):
        """Force balance constraint: normal forces must balance external force."""
        return F[2] + F[5] - self.ext_force
        
    def friction_cone_1(self, F):
        """Friction cone constraint for first contact point."""
        return self.mu * F[2] - np.sqrt(F[0]**2 + F[1]**2)
        
    def friction_cone_2(self, F):
        """Friction cone constraint for second contact point."""
        return self.mu * F[5] - np.sqrt(F[3]**2 + F[4]**2)
        
    def torque_limit_1(self, F):
        """Torque limit constraint for first finger."""
        return self.torque_limit - abs(self.p * F[0])
        
    def torque_limit_2(self, F):
        """Torque limit constraint for second finger."""
        return self.torque_limit - abs(self.p * F[3])

    def optimize(self, pipe_diameter=None):
        """
        Optimize gripping forces for the given pipe diameter.
        
        Args:
            pipe_diameter: Diameter of the pipe in mm (optional)
            
        Returns:
            Optimized forces for each contact point
        """
        # Adjust external force based on pipe diameter if provided
        if pipe_diameter is not None:
            # Calculate pipe mass based on diameter (simplified calculation)
            pipe_volume = np.pi * (pipe_diameter/2000)**2 * 0.1  # 10cm length
            pipe_density = 1400  # kg/m³ for PVC
            pipe_mass = pipe_volume * pipe_density
            self.ext_force = pipe_mass * 9.81  # Force in Newtons
        
        # Define constraints for SLSQP optimization
        constraints = [
            {'type': 'eq', 'fun': self.force_balance},
            {'type': 'ineq', 'fun': self.friction_cone_1},
            {'type': 'ineq', 'fun': self.friction_cone_2},
            {'type': 'ineq', 'fun': self.torque_limit_1},
            {'type': 'ineq', 'fun': self.torque_limit_2}
        ]
        
        # Initial guess for forces [f_x1, f_y1, f_z1, f_x2, f_y2, f_z2]
        initial_guess = [1, 1, self.ext_force/2, 1, 1, self.ext_force/2]
        
        # Solve optimization using SLSQP
        result = minimize(self.objective, initial_guess, method='SLSQP', 
                         constraints=constraints)
                         
        return result.x

import numpy as np

class MinimumJerkTrajectoryPlanner:
    """Generate smooth robot trajectories using minimum-jerk optimization."""
    
    def generate_trajectory(self, x0, xf, duration, num_points=100):
        """
        Generate a minimum-jerk trajectory between start and end positions.
        
        Args:
            x0: Starting position as [x, y, z]
            xf: Ending position as [x, y, z]
            duration: Time to complete the movement in seconds
            num_points: Number of points in the trajectory
            
        Returns:
            Dictionary containing positions, velocities, accelerations,
            and jerks for each time point
        """
        # Create time vector
        t = np.linspace(0, duration, num_points)
        
        # Initialize arrays for trajectory data
        positions = np.zeros((num_points, 3))
        velocities = np.zeros((num_points, 3))
        accelerations = np.zeros((num_points, 3))
        jerks = np.zeros((num_points, 3))
        
        # For each dimension (x, y, z)
        for i in range(3):
            # Get the normalized time vector
            tau = t / duration
            
            # Apply minimum jerk formula
            positions[:, i] = x0[i] + (xf[i] - x0[i]) * (10*tau**3 - 15*tau**4 + 6*tau**5)
            
            # Calculate derivatives
            v_scale = (xf[i] - x0[i]) / duration
            velocities[:, i] = v_scale * (30*tau**2 - 60*tau**3 + 30*tau**4)
            
            a_scale = (xf[i] - x0[i]) / (duration**2)
            accelerations[:, i] = a_scale * (60*tau - 180*tau**2 + 120*tau**3)
            
            j_scale = (xf[i] - x0[i]) / (duration**3)
            jerks[:, i] = j_scale * (60 - 360*tau + 360*tau**2)
        
        return {
            'time': t,
            'position': positions,
            'velocity': velocities,
            'acceleration': accelerations,
            'jerk': jerks
        }

    def generate_multi_waypoint_trajectory(self, waypoints, durations):
        """
        Generate a trajectory that passes through multiple waypoints.
        
        Args:
            waypoints: List of positions [[x1,y1,z1], [x2,y2,z2], ...]
            durations: List of durations for each segment
            
        Returns:
            Combined trajectory through all waypoints
        """
        if len(waypoints) < 2:
            raise ValueError("At least 2 waypoints required")
            
        if len(waypoints) != len(durations) + 1:
            raise ValueError("Number of durations must be one less than number of waypoints")
            
        # Convert lists to numpy arrays for consistency
        waypoints = np.array(waypoints)
        durations = np.array(durations)
        
        # Generate segment trajectories
        segment_trajectories = []
        for i in range(len(waypoints) - 1):
            start_pos = waypoints[i]
            end_pos = waypoints[i+1]
            duration = durations[i]
            
            trajectory = self.generate_trajectory(start_pos, end_pos, duration)
            segment_trajectories.append(trajectory)
            
        # Concatenate segment trajectories
        total_points = sum(len(traj['time']) for traj in segment_trajectories)
        combined_time = np.zeros(total_points)
        combined_position = np.zeros((total_points, 3))
        combined_velocity = np.zeros((total_points, 3))
        combined_acceleration = np.zeros((total_points, 3))
        
        # Fill in the combined trajectory
        idx = 0
        time_offset = 0
        for traj in segment_trajectories:
            n_points = len(traj['time'])
            combined_time[idx:idx+n_points] = traj['time'] + time_offset
            combined_position[idx:idx+n_points] = traj['position']
            combined_velocity[idx:idx+n_points] = traj['velocity']
            combined_acceleration[idx:idx+n_points] = traj['acceleration']
            
            idx += n_points
            time_offset += traj['time'][-1]
            
        return {
            'time': combined_time,
            'position': combined_position,
            'velocity': combined_velocity,
            'acceleration': combined_acceleration
        }

    def plan_pvc_coupling_path(self, pickup_pos, water_bath_pos, mold_pos, home_pos):
        """
        Plan a complete path for the PVC coupling process:
        1. Pickup position → Water bath
        2. Water bath → Mold position
        3. Mold position → Home position
        
        Args:
            pickup_pos: Position to pick up the PVC pipe [x,y,z]
            water_bath_pos: Position of the water bath [x,y,z]
            mold_pos: Position of the mold for placement [x,y,z]
            home_pos: Home position of the robot [x,y,z]
            
        Returns:
            Complete trajectory for the entire process
        """
        # Define waypoints
        waypoints = [
            pickup_pos,
            water_bath_pos,
            mold_pos,
            home_pos
        ]
        
        # Define appropriate durations for each segment
        durations = [
            3.0,  # Pickup to water bath (slower for careful pickup)
            2.0,  # Water bath to mold (faster while pipe is heated)
            3.0   # Mold to home (medium speed)
        ]
        
        return self.generate_multi_waypoint_trajectory(waypoints, durations)

import cv2
import numpy as np
import matplotlib.pyplot as plt

# Function to calculate the center of a bounding box
def get_bounding_box_center_2d(x_min, y_min, x_max, y_max):
    center_x = (x_min + x_max) / 2
    center_y = (y_min + y_max) / 2
    return center_x, center_y

# Function to simulate depth estimation (this could be replaced with actual depth sensing logic)
def depth_estimation(x, y, depth_map=None):
    # For simplicity, return a constant depth value for now
    # In production, this would use a depth camera or stereo vision
    return 10.0  # Placeholder value for depth

# Function to calculate the 3D center from YOLO detection
def get_bounding_box_center_3d(x_min, y_min, x_max, y_max, depth_map=None):
    # Get 2D center
    center_x, center_y = get_bounding_box_center_2d(x_min, y_min, x_max, y_max)
    
    # Get the depth value for the center (z-coordinate)
    center_z = depth_estimation(center_x, center_y, depth_map)
    
    return center_x, center_y, center_z

# Example usage: Process each detected pipe
def process_pipe_detection(frame, depth_map=None):
    # Example of detected bounding boxes from YOLO
    detections = [(50, 30, 150, 130), (200, 80, 350, 220)]
    
    for (x_min, y_min, x_max, y_max) in detections:
        # Get 3D center of the detected bounding box
        center_x, center_y, center_z = get_bounding_box_center_3d(
            x_min, y_min, x_max, y_max, depth_map
        )
        
        print(f"3D center of pipe: ({center_x:.2f}, {center_y:.2f}, {center_z:.2f})")
        
        # Draw the bounding box and center on the image for visualization
        cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
        cv2.circle(frame, (int(center_x), int(center_y)), 5, (0, 0, 255), -1)
        cv2.putText(
            frame, 
            f"3D: ({center_x:.1f}, {center_y:.1f}, {center_z:.1f})", 
            (x_min, y_min - 10), 
            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 0), 1
        )

import numpy as np
from scipy.optimize import minimize

# Define given parameters
mu = 0.5  # Friction coefficient
external_force = 10  # External force (e.g., gravity in Newtons)
torque_limit = 5  # Maximum allowable torque
p = 0.1  # Distance from the center to contact points (meters)

# Objective function: Minimize max(|f_i|) to reduce excessive gripping force
def objective(F):
    return max(abs(F))  # Minimize the largest applied force

# Constraint 1: Force balance equation
def force_balance(F):
    return F[2] + F[5] - external_force  # Normal forces must balance external force

# Constraint 2: Friction cone (Prevent slipping)
def friction_cone_1(F):
    return mu * F[2] - np.sqrt(F[0]**2 + F[1]**2)  # mu*f_z1 >= sqrt(f_x1^2 + f_y1^2)

def friction_cone_2(F):
    return mu * F[5] - np.sqrt(F[3]**2 + F[4]**2)  # mu*f_z2 >= sqrt(f_x2^2 + f_y2^2)

# Constraint 3: Torque limits (Prevent excessive joint load)
def torque_limit_1(F):
    return torque_limit - abs(p * F[0])  # Ensure torque does not exceed limit for finger 1

def torque_limit_2(F):
    return torque_limit - abs(p * F[3])  # Ensure torque does not exceed limit for finger 2

# Define constraints for SLSQP optimization
constraints = [
    {'type': 'eq', 'fun': force_balance},
    {'type': 'ineq', 'fun': friction_cone_1},
    {'type': 'ineq', 'fun': friction_cone_2},
    {'type': 'ineq', 'fun': torque_limit_1},
    {'type': 'ineq', 'fun': torque_limit_2}
]

# Initial guess for forces [f_x1, f_y1, f_z1, f_x2, f_y2, f_z2]
initial_guess = [1, 1, 5, 1, 1, 5]

# Solve the optimization using SLSQP
result = minimize(objective, initial_guess, method='SLSQP', constraints=constraints)

import numpy as np

class HierarchicalIKSolver:
    """
    A class for solving Hierarchical Inverse Kinematics (IK) problems with Nullspace Projection.
    This solver handles primary and secondary tasks using damped least squares for stability.
    """

    def __init__(self, damping=0.01):
        """
        Initialize the IK solver.

        Args:
            damping: Damping factor for stability (default=0.01).
        """
        self.damping = damping

    def damped_least_squares(self, J, error):
        """
        Perform the damped least squares solution for a given Jacobian and error vector.

        Args:
            J: Jacobian matrix (m x n).
            error: Error vector (m x 1) - difference between current and target position.

        Returns:
            Joint angle adjustments (delta_theta).
        """
        JT = J.T
        identity = np.eye(J.shape[0])
        damped_term = JT @ np.linalg.inv(J @ JT + (self.damping ** 2) * identity)
        delta_theta = damped_term @ error
        return delta_theta

    def solve_hierarchical_ik(self, J_tasks, errors):
        """
        Solve the hierarchical IK problem with nullspace projection.

        Args:
            J_tasks: List of Jacobian matrices for tasks (ordered by priority).
            errors: List of error vectors for each task (ordered by priority).

        Returns:
            Joint angle adjustments (delta_theta) respecting task priorities.
        """
        num_joints = J_tasks[0].shape[1]
        delta_theta = np.zeros(num_joints)
        nullspace_projector = np.eye(num_joints)  # Initialize nullspace projector as identity

        for task_idx, (J, error) in enumerate(zip(J_tasks, errors)):
            # Project the current task into the available nullspace
            J_effective = J @ nullspace_projector

            # Check if `J_effective` has valid dimensions
            if J_effective.shape[0] == 0 or J_effective.shape[1] == 0:
                print(f"Skipping task {task_idx} due to invalid Jacobian dimensions")
                continue

            # Solve for delta_theta for the current task
            delta_theta_task = self.damped_least_squares(J_effective, error)
            delta_theta += delta_theta_task

            # Update the nullspace projector
            try:
                # Compute pseudo-inverse of effective Jacobian
                J_effective_pseudo_inverse = np.linalg.pinv(J_effective)

                # Debug prints for each task
                print(f"[Task {task_idx}] J_effective shape: {J_effective.shape}")
                print(f"[Task {task_idx}] J_effective_pseudo_inverse shape: {J_effective_pseudo_inverse.shape}")

                # Update nullspace projector
                nullspace_update = J_effective_pseudo_inverse @ J_effective
                nullspace_projector = nullspace_projector - nullspace_update
            except ValueError as e:
                print(f"Error updating nullspace projector at task {task_idx}: {e}")
                continue

        return delta_theta

# Example usage
if __name__ == "__main__":
    # Example setup for a 2-task IK problem
    solver = HierarchicalIKSolver(damping=0.01)

    # Define Jacobians for primary and secondary tasks
    J_primary = np.array([
        [-1.0, 0.0, 0.0],
        [0.0, 1.0, 0.0]
    ])  # Shape: (2, 3)

    J_secondary = np.array([
        [0.0, 0.0, 1.0]
    ])  # Shape: (1, 3)

    # Define error vectors for primary and secondary tasks
    error_primary = np.array([0.1, -0.2])  # Example target position error (x, y)
    error_secondary = np.array([0.05])     # Example secondary task error

    # Solve hierarchical IK
    delta_theta = solver.solve_hierarchical_ik(
        J_tasks=[J_primary, J_secondary],
        errors=[error_primary, error_secondary]
    )

    print("Computed joint angle adjustments:", delta_theta)

import numpy as np
import matplotlib.pyplot as plt

class MinimumJerkTrajectoryPlanner:
    """
    A class for generating minimum-jerk trajectories for smooth robotic arm movement
    when handling PVC pipes in the hot water coupling automation process.
    """
    
    def generate_trajectory(self, start_pos, end_pos, duration, num_points=100):
        """Generate a minimum-jerk trajectory between start and end positions."""
        # Create the time vector
        t = np.linspace(0, duration, num_points)
        
        # Initialize arrays to store the trajectory data
        positions = np.zeros((num_points, 3))
        velocities = np.zeros((num_points, 3))
        accelerations = np.zeros((num_points, 3))
        jerks = np.zeros((num_points, 3))
        
        # For each dimension (x, y, z)
        for i in range(3):
            # Calculate the minimum-jerk trajectory coefficients
            x0 = start_pos[i]
            xf = end_pos[i]
            
            # The 5th order polynomial coefficients for minimum jerk trajectory
            a0 = x0
            a1 = 0  # zero initial velocity
            a2 = 0  # zero initial acceleration
            a3 = 10 * (xf - x0) / (duration**3)
            a4 = -15 * (xf - x0) / (duration**4)
            a5 = 6 * (xf - x0) / (duration**5)
            
            # Calculate position, velocity, acceleration, and jerk for each time point
            for j, tj in enumerate(t):
                # Position (5th order polynomial)
                positions[j, i] = a0 + a1*tj + a2*tj**2 + a3*tj**3 + a4*tj**4 + a5*tj**5
                
                # Velocity (derivative of position)
                velocities[j, i] = a1 + 2*a2*tj + 3*a3*tj**2 + 4*a4*tj**3 + 5*a5*tj**4
                
                # Acceleration (derivative of velocity)
                accelerations[j, i] = 2*a2 + 6*a3*tj + 12*a4*tj**2 + 20*a5*tj**3
                
                # Jerk (derivative of acceleration)
                jerks[j, i] = 6*a3 + 24*a4*tj + 60*a5*tj**2
        
        return {
            'time': t,
            'position': positions,
            'velocity': velocities,
            'acceleration': accelerations,
            'jerk': jerks
        }
    
    def generate_multi_waypoint_trajectory(self, waypoints, durations, num_points_per_segment=50):
        """Generate a minimum-jerk trajectory that passes through multiple waypoints."""
        if len(waypoints) < 2:
            raise ValueError("At least 2 waypoints are required")
        
        if len(waypoints) != len(durations) + 1:
            raise ValueError("Number of durations must be one less than number of waypoints")
        
        # Initialize arrays for the complete trajectory
        total_points = num_points_per_segment * len(durations)
        combined_time = np.zeros(total_points)
        combined_positions = np.zeros((total_points, 3))
        combined_velocities = np.zeros((total_points, 3))
        combined_accelerations = np.zeros((total_points, 3))
        combined_jerks = np.zeros((total_points, 3))
        
        time_offset = 0
        
        # Generate trajectory for each segment
        for i in range(len(durations)):
            start_pos = waypoints[i]
            end_pos = waypoints[i+1]
            duration = durations[i]
            
            segment_traj = self.generate_trajectory(start_pos, end_pos, duration, num_points_per_segment)
            
            # Calculate indices for this segment
            start_idx = i * num_points_per_segment
            end_idx = (i + 1) * num_points_per_segment
            
            # Store the segment trajectory data with appropriate time offset
            combined_time[start_idx:end_idx] = segment_traj['time'] + time_offset
            combined_positions[start_idx:end_idx] = segment_traj['position']
            combined_velocities[start_idx:end_idx] = segment_traj['velocity']
            combined_accelerations[start_idx:end_idx] = segment_traj['acceleration']
            combined_jerks[start_idx:end_idx] = segment_traj['jerk']
            
            # Update time offset for the next segment
            time_offset += duration
        
        return {
            'time': combined_time,
            'position': combined_positions,
            'velocity': combined_velocities,
            'acceleration': combined_accelerations,
            'jerk': combined_jerks
        }

def plan_pvc_pipe_coupling_path(self, pickup_position, waterbath_position, 
                               mold_position, home_position,
                               segment_durations=None):
    """
    Plan a complete path for the PVC pipe coupling process with four key positions:
    1. Pickup position (from pallet)
    2. Water bath position (for heating)
    3. Mold position (for final placement)
    4. Home position (to restart cycle)
    """
    waypoints = [pickup_position, waterbath_position, mold_position, home_position]
    
    # Default durations if none provided
    if segment_durations is None:
        segment_durations = [3.0, 2.0, 3.0]  # 3 durations for 4 waypoints
        
    # Verify that we have the correct number of durations
    if len(segment_durations) != len(waypoints) - 1:
        raise ValueError(f"Need {len(waypoints) - 1} durations for {len(waypoints)} waypoints")
        
    trajectory = self.generate_multi_waypoint_trajectory(waypoints, segment_durations)
    
    return trajectory