mecha10-planning 0.6.3

//! A* (A-Star) path planning algorithm
//!
//! This module implements the A* algorithm for optimal grid-based path planning.

use crate::types::{Obstacle, Path, Point2D};
use priority_queue::PriorityQueue;
use std::cmp::Ordering;
use std::collections::{HashMap, HashSet};

/// A* path planner
#[derive(Debug, Clone)]
pub struct AStarPlanner {
    start: Point2D,
    goal: Point2D,
    resolution: f32,
    obstacles: Vec<Obstacle>,
    diagonal_movement: bool,
}

impl AStarPlanner {
    /// Create a new A* planner
    ///
    /// # Arguments
    ///
    /// * `start` - Starting position
    /// * `goal` - Goal position
    /// * `resolution` - Grid resolution (smaller = more precise, slower)
    pub fn new(start: impl Into<Point2D>, goal: impl Into<Point2D>, resolution: f32) -> Self {
        Self {
            start: start.into(),
            goal: goal.into(),
            resolution,
            obstacles: Vec::new(),
            diagonal_movement: true,
        }
    }

    /// Add an obstacle to the environment
    pub fn add_obstacle(&mut self, obstacle: Obstacle) {
        self.obstacles.push(obstacle);
    }

    /// Set whether diagonal movement is allowed
    pub fn with_diagonal_movement(mut self, allow: bool) -> Self {
        self.diagonal_movement = allow;
        self
    }

    /// Plan a path from start to goal
    pub fn plan(&self) -> Option<Path> {
        let start_cell = self.to_grid_cell(&self.start);
        let goal_cell = self.to_grid_cell(&self.goal);

        let mut open_set = PriorityQueue::new();
        open_set.push(start_cell, FloatOrd(-self.heuristic(&self.start)));

        let mut came_from: HashMap<GridCell, GridCell> = HashMap::new();
        let mut g_score: HashMap<GridCell, f32> = HashMap::new();
        g_score.insert(start_cell, 0.0);

        let mut closed_set = HashSet::new();

        while let Some((current, _)) = open_set.pop() {
            if current == goal_cell {
                return Some(self.reconstruct_path(&came_from, current));
            }

            closed_set.insert(current);

            let current_point = self.to_world_point(&current);

            for neighbor in self.get_neighbors(&current) {
                if closed_set.contains(&neighbor) {
                    continue;
                }

                let neighbor_point = self.to_world_point(&neighbor);

                if self.is_in_obstacle(&neighbor_point) {
                    continue;
                }

                let tentative_g_score = g_score[&current] + current_point.distance_to(&neighbor_point);

                if tentative_g_score < *g_score.get(&neighbor).unwrap_or(&f32::INFINITY) {
                    came_from.insert(neighbor, current);
                    g_score.insert(neighbor, tentative_g_score);

                    let f_score = tentative_g_score + self.heuristic(&neighbor_point);
                    open_set.push(neighbor, FloatOrd(-f_score));
                }
            }
        }

        None
    }

    fn to_grid_cell(&self, point: &Point2D) -> GridCell {
        GridCell {
            x: (point.x / self.resolution).round() as i32,
            y: (point.y / self.resolution).round() as i32,
        }
    }

    fn to_world_point(&self, cell: &GridCell) -> Point2D {
        Point2D::new(cell.x as f32 * self.resolution, cell.y as f32 * self.resolution)
    }

    fn heuristic(&self, point: &Point2D) -> f32 {
        point.distance_to(&self.goal)
    }

    fn is_in_obstacle(&self, point: &Point2D) -> bool {
        self.obstacles.iter().any(|obs| obs.contains(point))
    }

    fn get_neighbors(&self, cell: &GridCell) -> Vec<GridCell> {
        let mut neighbors = Vec::new();

        for (dx, dy) in &[(0, 1), (1, 0), (0, -1), (-1, 0)] {
            neighbors.push(GridCell {
                x: cell.x + dx,
                y: cell.y + dy,
            });
        }

        if self.diagonal_movement {
            for (dx, dy) in &[(1, 1), (1, -1), (-1, 1), (-1, -1)] {
                neighbors.push(GridCell {
                    x: cell.x + dx,
                    y: cell.y + dy,
                });
            }
        }

        neighbors
    }

    fn reconstruct_path(&self, came_from: &HashMap<GridCell, GridCell>, mut current: GridCell) -> Path {
        let mut waypoints = vec![self.to_world_point(&current)];

        while let Some(&prev) = came_from.get(&current) {
            current = prev;
            waypoints.push(self.to_world_point(&current));
        }

        waypoints.reverse();
        Path::new(waypoints)
    }
}

/// Grid cell coordinate
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
struct GridCell {
    x: i32,
    y: i32,
}

/// Wrapper for f32 to implement Ord for priority queue
#[derive(Debug, Clone, Copy, PartialEq)]
struct FloatOrd(f32);

impl Eq for FloatOrd {}

impl PartialOrd for FloatOrd {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for FloatOrd {
    fn cmp(&self, other: &Self) -> Ordering {
        self.0.partial_cmp(&other.0).unwrap_or(Ordering::Equal)
    }
}