alg-grid 0.1.1

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
use core::ops::Add;
use core::cmp::{max, min};
use heapless::{
    Vec,
    binary_heap::{BinaryHeap, Max},
    spsc::Queue,
    consts::*
};
use map_vec::Map;
use num_rational::BigRational;

/// Representation of a point in three dimensions
#[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd)]
pub struct Point3D {
    pub x: i32,
    pub y: i32,
    pub z: i32,
}

impl Default for Point3D {
    fn default() -> Self {
        Self {
            x: 0,
            y: 0,
            z: 0,
        }
    }
}

impl Add for Point3D {
    type Output = Self;
    fn add(self, other: Self) -> Self {
        Self {
            x: self.x + other.x,
            y: self.y + other.y,
            z: self.z + other.z,
        }
    }
}

/// Trait for implementing a Grid in three dimensions
pub trait Grid3D {
    /// The dimensions of the Grid. The only method that must be defined.
    fn dimensions(&self) -> Point3D;

    /// The lower bound. Defaults to (0, 0, 0).
    fn lower_bound(&self) -> Point3D {
        Default::default()
    }

    /// The upper bound. Defaults to the dimensions itself.
    fn upper_bound(&self) -> Point3D {
        let dim = self.dimensions();
        let low = self.lower_bound();
        Point3D {
            x: low.x + dim.x,
            y: low.y + dim.y,
            z: low.z + dim.z
        }
    }

    /// Check if a point is in the bounds of the grid.
    fn in_bounds(&self, point: Point3D) -> bool {
        let low = self.lower_bound();
        let upp = self.upper_bound();
        low.x <= point.x && point.x < upp.x &&
        low.y <= point.y && point.y < upp.y
    }

    /// Convert a point to an index.
    /// 
    /// Useful if you store the grid in a one-dimensional array.
    fn point3d_to_index(&self, point: Point3D) -> usize {
        let dim = self.dimensions();
        (
            point.z * dim.y * dim.x +
            point.y * dim.x +
            point.x
        ) as usize
    }

    /// Convert an index to a point.
    /// 
    /// Useful if you store the grid in a one-dimensional array.
    fn index_to_point3d(&self, index: usize) -> Point3D {
        let dim = self.dimensions();
        let mut idx = index as i32;
        let z = idx / (dim.x * dim.y);
        idx -= z * dim.x * dim.y;
        let y = idx / dim.x;
        let x = idx % dim.x;
        Point3D {
            x,
            y,
            z,
        }
    }

    /// Check if a point is traversable.
    /// 
    /// Defaults to always, so you may want to implement
    /// this.
    #[allow(unused_variables)]
    fn is_opaque(&self, point: Point3D) -> bool {
        true
    }

    /// Get all possible neighbors of the point, regardless if the
    /// point or its neighbors is in bounds, opaque, or neither.
    fn get_possible_neighbors(&self, point: Point3D) -> [Point3D; 26] {
        [
            // centers
            point + Point3D {
                x: 0,
                y: 1,
                z: 0,
            },
            point + Point3D {
                x: 0,
                y: -1,
                z: 0,
            },
            point + Point3D {
                x: 1,
                y: 0,
                z: 0,
            },
            point + Point3D {
                x: -1,
                y: 0,
                z: 0,
            },
            point + Point3D {
                x: 0,
                y: 0,
                z: 1,
            },
            point + Point3D {
                x: 0,
                y: 0,
                z: -1,
            },
            // sides
            point + Point3D {
                x: 1,
                y: 1,
                z: 0,
            },
            point + Point3D {
                x: 1,
                y: -1,
                z: 0,
            },
            point + Point3D {
                x: -1,
                y: -1,
                z: 0,
            },
            point + Point3D {
                x: -1,
                y: 0,
                z: 0,
            },
            point + Point3D {
                x: 0,
                y: 1,
                z: 1,
            },
            point + Point3D {
                x: 0,
                y: 1,
                z: -1,
            },
            point + Point3D {
                x: 0,
                y: -1,
                z: -1,
            },
            point + Point3D {
                x: 0,
                y: -1,
                z: 1,
            },
            point + Point3D {
                x: 1,
                y: 0,
                z: 1,
            },
            point + Point3D {
                x: 1,
                y: 0,
                z: -1,
            },
            point + Point3D {
                x: -1,
                y: 0,
                z: -1,
            },
            point + Point3D {
                x: -1,
                y: 0,
                z: 1,
            },
            // corners
            point + Point3D {
                x: 1,
                y: 1,
                z: 1,
            },
            point + Point3D {
                x: 1,
                y: -1,
                z: 1,
            },
            point + Point3D {
                x: -1,
                y: -1,
                z: 1,
            },
            point + Point3D {
                x: -1,
                y: 1,
                z: 1,
            },
            point + Point3D {
                x: 1,
                y: 1,
                z: -1,
            },
            point + Point3D {
                x: 1,
                y: -1,
                z: -1,
            },
            point + Point3D {
                x: -1,
                y: -1,
                z: -1,
            },
            point + Point3D {
                x: -1,
                y: 1,
                z: -1,
            },
        ]
    }

    /// Check if two points are possible neighbors.
    /// 
    /// Does not check if either points are inbounds or non-opaque.
    fn is_possible_neighbor(&self, p1: Point3D, p2: Point3D) -> bool {
        self.get_possible_neighbors(p1)
            .iter()
            .any(|&x| x == p2)
    }

    /// Get the neighbors that is in bounds and not opaque.
    fn get_neighbors(&self, point: Point3D) -> [Option<Point3D>; 26] {
        let mut arr: [Option<Point3D>; 26] = [None; 26];
        let possible_neighbors = self.get_possible_neighbors(point);
        for (i, n) in possible_neighbors.iter().enumerate() {
            if !self.is_opaque(*n) && self.in_bounds(*n) {
                arr[i] = Some(*n);
            }
        }
        arr
    }

    /// Check if two points are neighbors.
    /// 
    /// Checks if either points are inbounds or non-opaque.
    fn is_neighbor(&self, p1: Point3D, p2: Point3D) -> bool {
        self.get_neighbors(p1)
            .iter()
            .any(|&x| x == Some(p2))
    }

    /// Get the neighbors with the associated cost.
    /// 
    /// Defaults to all eight neighbors having a cost of 1.0
    /// if the neighbor is valid.
    /// 
    /// If you want the diagonals to cost sqrt(2), reimplement
    /// this method yourself.
    fn get_neighbors_with_cost(&self, point: Point3D) -> [(Option<Point3D>, BigRational); 26] {
        let mut arr: [(Option<Point3D>, BigRational); 26] = [
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
            (None, BigRational::from_float(0.0).unwrap()),
        ];
        let neighbors = self.get_neighbors(point);
        for (i, n) in neighbors.iter().enumerate() {
            let one = BigRational::from_float(1.0).unwrap();
            if n != &None {
                arr[i] = (*n, one);
            }
        }
        arr
    }
}

/// Algorithms to get the distance of two points.
pub enum Distance2D {
    Pythagoras,
    Manhattan,
    Chebyshev,
}

impl Distance2D {
    pub fn distance(&self, start: Point3D, end: Point3D) -> f64 {
        use Distance2D::*;
        match self {
            Pythagoras => {
                let dx = (max(start.x, end.x) - min(start.x, end.x)) as f64;
                let dy = (max(start.y, end.y) - min(start.y, end.y)) as f64;
                let dz = (max(start.z, end.z) - min(start.z, end.z)) as f64;
                libm::sqrt((dx * dx) + (dy * dy) + (dz * dz))
            }
            Manhattan => {
                let dx = (max(start.x, end.x) - min(start.x, end.x)) as f64;
                let dy = (max(start.y, end.y) - min(start.y, end.y)) as f64;
                let dz = (max(start.z, end.z) - min(start.z, end.z)) as f64;
                dx + dy + dz
            }
            Chebyshev => {
                let dx = (max(start.x, end.x) - min(start.x, end.x)) as f64;
                let dy = (max(start.y, end.y) - min(start.y, end.y)) as f64;
                let dz = (max(start.z, end.z) - min(start.z, end.z)) as f64;
                libm::sqrt((dx * dx) + (dy * dy) + (dz * dz))
            }
        }
    }
}

pub fn breadth_first_search(graph: &dyn Grid3D, start: Point3D, goal: Point3D) -> Map<Point3D, Option<Point3D>> {
    let mut frontier: Queue<Point3D, U16> = Queue::new();
    frontier.enqueue(start).expect("Failed to enqueue");
    let mut came_from = Map::<Point3D, Option<Point3D>>::new();
    came_from.insert(start, None);
    while !frontier.is_empty() {
        let current = frontier.dequeue().unwrap();
        if current == goal { break }
        for next in &graph.get_neighbors(current) {
            let next = next.unwrap();
            if !came_from.contains_key(&next) {
                frontier.enqueue(next).expect("Failed to enqueue");
                came_from.insert(next, Some(current));
            }
        }
    }
    came_from
}

pub fn dijkstra_search(
    graph: &dyn Grid3D,
    start: Point3D,
    goal: Point3D
) -> (
    Map<Point3D, Option<Point3D>>,
    Map<Point3D, BigRational>,
) {
    let mut frontier: BinaryHeap<(BigRational, Point3D), U16, Max> = BinaryHeap::new();
    
    let zero1 = BigRational::from_float(0.0).unwrap();
    let zero2 = BigRational::from_float(0.0).unwrap();
    
    frontier.push((zero1, start)).expect("fail to push to heap");
    let mut came_from = Map::<Point3D, Option<Point3D>>::new();
    let mut cost_so_far = Map::<Point3D, BigRational>::new();
    
    came_from.insert(start, None);
    cost_so_far.insert(start, zero2);
    
    while !frontier.is_empty() {
        let current = frontier.pop().unwrap();
        let point = current.1;
        
        if current.1 == goal { break }
        
        for (next, cost) in &graph.get_neighbors_with_cost(point) {
            let next = next.unwrap();
            let new_cost1 = cost_so_far.get(&point).unwrap() + cost;
            let new_cost2 = cost_so_far.get(&point).unwrap() + cost;
            
            if !cost_so_far.contains_key(&next) || new_cost1 < *cost_so_far.get(&next).unwrap() {
                cost_so_far.insert(next, new_cost1);
                
                let priority = new_cost2;
                
                frontier.push((priority, next)).expect("fail to push to heap");
                came_from.insert(next, Some(point));
            }
        }
    }
    (came_from, cost_so_far)
}

pub fn a_star_search(
    graph: &dyn Grid3D,
    start: Point3D,
    goal: Point3D
) -> (
    Map<Point3D, Option<Point3D>>,
    Map<Point3D, BigRational>,
) {
    use num_bigint::ToBigInt;
    
    let mut frontier: BinaryHeap<(BigRational, Point3D), U16, Max> = BinaryHeap::new();
    
    let zero1 = BigRational::from_float(0.0).unwrap();
    let zero2 = BigRational::from_float(0.0).unwrap();
    
    frontier.push((zero1, start)).expect("fail to push to heap");
    
    let mut came_from = Map::<Point3D, Option<Point3D>>::new();
    let mut cost_so_far = Map::<Point3D, BigRational>::new();
    
    came_from.insert(start, None);
    cost_so_far.insert(start, zero2);
    
    while !frontier.is_empty() {
        let current = frontier.pop().unwrap();
        let point = current.1;
        
        if current.1 == goal { break }
        
        for (next, cost) in &graph.get_neighbors_with_cost(point) {
            let next = next.unwrap();
            
            let new_cost1 = cost_so_far.get(&point).unwrap() + cost;
            let new_cost2 = cost_so_far.get(&point).unwrap() + cost;
            
            if !cost_so_far.contains_key(&next) || new_cost1 < *cost_so_far.get(&next).unwrap() {
                cost_so_far.insert(next, new_cost1);
                
                let h = BigRational::from_integer(
                    heuristic(goal, next).to_bigint().unwrap()
                );
                let priority = new_cost2 + h;
                
                frontier.push((priority, next)).expect("fail to push to heap");
                came_from.insert(next, Some(point));
            }
        }
    }
    (came_from, cost_so_far)
}


fn heuristic(p1: Point3D, p2: Point3D) -> i32 {
    (p1.x - p2.x).abs() + (p1.y - p2.y).abs()
}

pub fn reconstruct_path(
    came_from: Map<Point3D, Option<Point3D>>,
    start: Point3D,
    goal: Point3D
) -> Vec<Point3D, U16> {
    let mut current = goal;
    let mut path = Vec::<Point3D, U16>::new();
    while current != start {
        path.push(current).expect("Cannot push to vector");
        current = came_from.get(&current).unwrap().unwrap();
    }
    path.push(start).expect("Cannot push to vector");
    #[cfg(feature = "reverse_path")]
    let path = path.iter()
        .cloned()
        .rev()
        .collect::<Vec<_, U16>>();
    path
}

// TODO: Create some tests
#[cfg(test)]
mod tests {
    #[test]
    fn it_works() {
        assert_eq!(2 + 2, 4);
    }
}