Enhanced ICP

"""
A 3D visualization program that simulates an articulated cylinder that can be deformed by applying forces.
The program uses Open3D for visualization and allows interactive manipulation of the cylinder through keyboard controls.
"""

import open3d as o3d
import numpy as np
import copy
import time
from dataclasses import dataclass

@dataclass
class Point3D:
    """
    Represents a point in 3D space with x, y, z coordinates.
    Implements basic vector operations like addition and normalization.
    """
    x: float
    y: float
    z: float
    
    def __add__(self, other):
        """Vector addition operator for Point3D objects."""
        return Point3D(
            self.x + other.x,
            self.y + other.y,
            self.z + other.z
        )
    
    def normalize(self):
        """
        Returns a normalized version of the point vector (unit length).
        If the vector has zero length, returns a copy of the original point.
        """
        length = np.sqrt(self.x**2 + self.y**2 + self.z**2)
        if length > 0:
            return Point3D(
                self.x / length,
                self.y / length,
                self.z / length
            )
        return copy.deepcopy(self)
    
    def to_numpy(self):
        """Converts the point to a numpy array format."""
        return np.array([self.x, self.y, self.z])
    
    @staticmethod
    def from_numpy(arr):
        """Creates a Point3D object from a numpy array."""
        return Point3D(arr[0], arr[1], arr[2])

class ArticulatedCylinder:
    """
    Represents a deformable cylinder composed of multiple circular sections.
    The cylinder can be deformed by applying forces at different points.
    """
    def __init__(self, height=4.0, radius=0.5, n_sections=20, points_per_section=30):
        """
        Initialize the cylinder with specified dimensions and resolution.
        
        Args:
            height (float): Total height of the cylinder
            radius (float): Radius of the cylinder
            n_sections (int): Number of vertical sections
            points_per_section (int): Number of points in each circular section
        """
        self.radius = radius
        self.height = height
        self.n_sections = n_sections
        self.points_per_section = points_per_section
        self.threshold = height / (n_sections * 2)  # Section height threshold
        
        # Generate initial points and center line
        self.points = self.generate_points()
        self.center_line = self.generate_center_line()
        self.sections = []
        
        # Create sections
        self.create_sections()
        
    def generate_points(self):
        """
        Generate initial cylinder points arranged in circular sections.
        Returns a list of Point3D objects representing the cylinder surface.
        """
        points = []
        for h_idx in range(self.n_sections):
            h = (h_idx / (self.n_sections - 1)) * self.height
            for theta_idx in range(self.points_per_section):
                theta = (theta_idx / self.points_per_section) * 2 * np.pi
                x = self.radius * np.cos(theta)
                y = self.radius * np.sin(theta)
                points.append(Point3D(x, y, h))
        return points
    
    def generate_center_line(self):
        """
        Generate points along the central axis of the cylinder.
        Returns a list of Point3D objects representing the center line.
        """
        return [Point3D(0, 0, h) for h in np.linspace(0, self.height, self.n_sections)]
    
    def create_sections(self):
        """
        Group points into vertical sections based on their height.
        Each section contains points at approximately the same height.
        """
        self.sections = [[] for _ in range(self.n_sections)]
        
        for point in self.points:
            section_idx = int((point.z / self.height) * (self.n_sections - 1) + 0.5)
            section_idx = max(0, min(section_idx, self.n_sections - 1))
            self.sections[section_idx].append(point)
    
    def find_nearest_section(self, force_point):
        """
        Find the section index closest to the force application point.
        
        Args:
            force_point (Point3D): Point where force is being applied
            
        Returns:
            int: Index of the nearest section
        """
        force_height = force_point.z
        section_idx = int((force_height / self.height) * (self.n_sections - 1) + 0.5)
        return max(0, min(section_idx, self.n_sections - 1))
    
    def calculate_section_center(self, section):
        """
        Calculate the center point of a given section.
        
        Args:
            section (list): List of Point3D objects in the section
            
        Returns:
            Point3D: Center point of the section
        """
        if not section:
            return Point3D(0, 0, 0)
        x_avg = sum(p.x for p in section) / len(section)
        y_avg = sum(p.y for p in section) / len(section)
        z_avg = sum(p.z for p in section) / len(section)
        return Point3D(x_avg, y_avg, z_avg)
    
    def update_deformation(self, force_point, force_direction):
        """
        Update cylinder shape based on applied force.
        The deformation effect decreases with distance from the force application point.
        
        Args:
            force_point (Point3D): Point where force is applied
            force_direction (Point3D): Direction and magnitude of the force
        """
        impact_section = self.find_nearest_section(force_point)
        
        # Calculate deformation influence on each section
        for i in range(len(self.sections)):
            distance_factor = 1 / (1 + abs(i - impact_section))
            deformation = Point3D(
                force_direction.x * distance_factor,
                force_direction.y * distance_factor,
                force_direction.z * distance_factor
            )
            
            # Update points in section
            for point in self.sections[i]:
                new_x = point.x + deformation.x
                new_y = point.y + deformation.y
                new_z = point.z + deformation.z
                
                # Maintain radial constraint to preserve cylinder shape
                center = self.calculate_section_center(self.sections[i])
                dx, dy = new_x - center.x, new_y - center.y
                dist = np.sqrt(dx**2 + dy**2)
                
                if dist > 0:
                    scale = self.radius / dist
                    point.x = center.x + dx * scale
                    point.y = center.y + dy * scale
                    point.z = new_z

class InteractiveCylinder:
    """
    Handles the interactive visualization and manipulation of the articulated cylinder.
    Provides keyboard controls for applying forces to deform the cylinder.
    """
    def __init__(self):
        """Initialize the visualization window and create the cylinder object."""
        self.vis = o3d.visualization.VisualizerWithKeyCallback()
        self.vis.create_window()
        
        # Create articulated cylinder
        self.cylinder = ArticulatedCylinder()
        
        # Create visualization objects
        self.point_cloud = o3d.geometry.PointCloud()
        self.update_visualization()
        self.vis.add_geometry(self.point_cloud)
        
        # Create target point (force application point)
        self.target_point = Point3D(0.0, 0.0, self.cylinder.height/2)
        self.target_sphere = self.create_target_sphere()
        self.vis.add_geometry(self.target_sphere)
        
        # Movement parameters
        self.move_speed = 0.1
        self.force_strength = 0.1
        
        # Register keyboard callbacks
        self.register_key_callbacks()
    
    def create_target_sphere(self):
        """
        Create a red sphere to represent the force application point.
        Returns an Open3D sphere mesh.
        """
        sphere = o3d.geometry.TriangleMesh.create_sphere(radius=0.1)
        sphere.translate(self.target_point.to_numpy())
        sphere.paint_uniform_color([1, 0, 0])
        return sphere
    
    def update_visualization(self):
        """
        Update the point cloud visualization with current cylinder state.
        Colors points based on their height for better visual feedback.
        """
        # Convert points to numpy array
        points = np.array([p.to_numpy() for p in self.cylinder.points])
        
        # Update point cloud
        self.point_cloud.points = o3d.utility.Vector3dVector(points)
        
        # Update colors based on height
        colors = np.zeros((len(points), 3))
        heights = points[:, 2]
        normalized_heights = (heights - np.min(heights)) / (np.max(heights) - np.min(heights))
        colors[:, 0] = 0.7  # Red
        colors[:, 1] = 0.7  # Green
        colors[:, 2] = normalized_heights * 0.3 + 0.4  # Blue varies with height
        self.point_cloud.colors = o3d.utility.Vector3dVector(colors)
    
    def move_target(self, direction):
        """
        Move the target point and apply deformation force to the cylinder.
        
        Args:
            direction (str): Direction of movement ('LEFT', 'RIGHT', 'UP', 'DOWN', 'FORWARD', 'BACKWARD')
        """
        movement = np.zeros(3)
        if direction == 'LEFT':
            movement[0] = -self.move_speed
        elif direction == 'RIGHT':
            movement[0] = self.move_speed
        elif direction == 'UP':
            movement[2] = self.move_speed
        elif direction == 'DOWN':
            movement[2] = -self.move_speed
        elif direction == 'FORWARD':
            movement[1] = self.move_speed
        elif direction == 'BACKWARD':
            movement[1] = -self.move_speed
        
        # Update target position while keeping it within cylinder bounds
        new_position = np.array([self.target_point.x, self.target_point.y, self.target_point.z]) + movement
        if 0.1 <= new_position[2] <= (self.cylinder.height - 0.1):
            self.target_point = Point3D.from_numpy(new_position)
            self.target_sphere.translate(movement)
            
            # Apply force to cylinder
            force_direction = Point3D(movement[0], movement[1], movement[2])
            self.cylinder.update_deformation(self.target_point, force_direction)
    
    def register_key_callbacks(self):
        """Register keyboard controls for interactive manipulation."""
        def move_callback(direction):
            def callback(vis):
                self.move_target(direction)
                return True
            return callback
        
        # Map arrow keys and W/S keys to movement directions
        self.vis.register_key_callback(262, move_callback('RIGHT'))  # Right arrow
        self.vis.register_key_callback(263, move_callback('LEFT'))   # Left arrow
        self.vis.register_key_callback(265, move_callback('UP'))     # Up arrow
        self.vis.register_key_callback(264, move_callback('DOWN'))   # Down arrow
        self.vis.register_key_callback(87, move_callback('FORWARD')) # W key
        self.vis.register_key_callback(83, move_callback('BACKWARD'))# S key
    
    def run(self):
        """
        Main run loop for the interactive visualization.
        Updates the visualization and handles user input continuously.
        """
        print("Controls:")
        print("Arrow keys: Move target point horizontally and vertically")
        print("W/S: Move target point forward/backward")
        print("Press Ctrl+C to exit")
        
        while True:
            self.update_visualization()
            self.vis.update_geometry(self.point_cloud)
            self.vis.update_geometry(self.target_sphere)
            self.vis.poll_events()
            self.vis.update_renderer()
            time.sleep(0.01)

if __name__ == "__main__":
    cylinder = InteractiveCylinder()
    try:
        cylinder.run()
    except KeyboardInterrupt:
        cylinder.vis.destroy_window()