blackout/src/pathfinding.rs

use std::marker::PhantomData;

use bevy::{
    ecs::entity::Entities,
    prelude::*,
    tasks::{prelude::*, Task},
};
use bevy_rapier2d::{
    na::{Isometry2, Vector2},
    prelude::*,
    rapier::prelude::{ColliderHandle, ColliderSet, QueryPipeline, RigidBodySet},
};
use futures_lite::future;

use pathfinding::prelude::*;

use crate::{
    commands::RunIfExistsExt,
    core::PointLike,
    map::{Map, MapObstruction},
    navigation::{RotationSpeed, Speed},
};

#[derive(Component, Debug, Deref, DerefMut)]
struct Calculating(Task<Option<Path>>);

#[derive(Component, Clone, Copy, Debug, Default, Deref, DerefMut, Eq, Hash, PartialEq, Reflect)]
#[reflect(Component)]
pub struct Destination(pub (i32, i32));

impl_pointlike_for_tuple_component!(Destination);
impl_pointlike_for_tuple_component!(&Destination);

#[derive(Component, Clone, Debug, Default, Reflect)]
#[reflect(Component)]
pub struct NoPath;

#[derive(Component, Clone, Debug, Default, Deref, DerefMut, Reflect)]
#[reflect(Component)]
pub struct Path(pub Vec<(i32, i32)>);

#[derive(Component, Clone, Debug, Default, Reflect)]
#[reflect(Component)]
pub struct CostMap {
    width: usize,
    costs: Vec<f32>,
}

impl CostMap {
    pub fn new(width: usize, height: usize) -> Self {
        let count = width * height;
        Self {
            width,
            costs: vec![1.; count],
        }
    }

    pub fn modifier(&self, x: usize, y: usize) -> Option<&f32> {
        let idx = (y as usize) * self.width + (x as usize);
        self.costs.get(idx)
    }

    pub fn set_modifier(&mut self, x: usize, y: usize, cost: f32) {
        let idx = (y as usize) * self.width + (x as usize);
        self.costs[idx] = cost;
    }
}

pub fn find_path<D: 'static + Clone + Default + Send + Sync>(
    start: &dyn PointLike,
    destination: &dyn PointLike,
    map: &Map<D>,
    cost_map: &Option<CostMap>,
) -> Option<(Vec<(i32, i32)>, u32)> {
    astar(
        &start.into(),
        |p| {
            let mut successors: Vec<((i32, i32), u32)> = vec![];
            if let Some(tile) = map.at(p.0 as usize, p.1 as usize) {
                if tile.is_walkable() {
                    for tile in map.get_available_exits(p.0 as usize, p.1 as usize) {
                        let mut cost = tile.2 * 100.;
                        if let Some(cost_map) = cost_map {
                            if let Some(modifier) = cost_map.modifier(tile.0, tile.1) {
                                cost *= modifier;
                            }
                        }
                        successors.push(((tile.0 as i32, tile.1 as i32), cost as u32));
                    }
                }
            }
            successors
        },
        |p| (p.distance_squared(destination) * 100.) as u32,
        |p| *p == destination.into(),
    )
}

fn find_path_for_shape<D: 'static + Clone + Default + Send + Sync>(
    initiator: ColliderHandle,
    start: Transform,
    destination: Destination,
    map: Map<D>,
    cost_map: &Option<CostMap>,
    query_pipeline: QueryPipeline,
    collider_set: ColliderSet,
    rigid_body_set: RigidBodySet,
    shape: Collider,
) -> Option<Path> {
    let path = astar(
        &start.i32(),
        |p| {
            let mut successors: Vec<((i32, i32), u32)> = vec![];
            if let Some(tile) = map.at(p.0 as usize, p.1 as usize) {
                if tile.is_walkable() {
                    for tile in map.get_available_exits(p.0 as usize, p.1 as usize) {
                        let mut should_push = true;
                        let shape_pos = Isometry2::new(
                            Vector2::new(tile.0 as f32 + 0.5, tile.1 as f32 + 0.5),
                            0.,
                        );
                        query_pipeline.intersections_with_shape(
                            &rigid_body_set,
                            &collider_set,
                            &shape_pos,
                            &*shape.raw,
                            bevy_rapier2d::rapier::pipeline::QueryFilter::new()
                                .predicate(&|h, _c| h != initiator),
                            |_handle| {
                                should_push = false;
                                false
                            },
                        );
                        if should_push {
                            let mut cost = tile.2 * 100.;
                            if let Some(cost_map) = cost_map {
                                if let Some(modifier) = cost_map.modifier(tile.0, tile.1) {
                                    cost *= modifier;
                                }
                            }
                            successors
                                .push(((tile.0 as i32, tile.1 as i32), cost as u32));
                        }
                    }
                }
            }
            successors
        },
        |p| (p.distance_squared(&destination) * 100.) as u32,
        |p| *p == destination.i32(),
    );
    if let Some(path) = path {
        Some(Path(path.0))
    } else {
        None
    }
}

fn calculate_path<D: 'static + Clone + Default + Send + Sync>(
    mut commands: Commands,
    rapier_context: Res<RapierContext>,
    obstructions: Query<&RapierColliderHandle, With<MapObstruction>>,
    query: Query<
        (
            Entity,
            &RapierColliderHandle,
            &Destination,
            &Transform,
            &Collider,
            Option<&CostMap>,
        ),
        Changed<Destination>,
    >,
    map: Query<&Map<D>>,
) {
    for (entity, handle, destination, coordinates, shape, cost_map) in &query {
        if coordinates.i32() == **destination {
            commands
                .entity(entity)
                .remove::<Path>()
                .remove::<NoPath>()
                .remove::<Calculating>()
                .remove::<Destination>()
                .remove::<Speed>();
            continue;
        }
        if let Ok(map) = map.get_single() {
            let coordinates_clone = *coordinates;
            let destination_clone = *destination;
            let query_pipeline = rapier_context.query_pipeline.clone();
            let map_clone = map.clone();
            let cost_map_clone = cost_map.cloned();
            let handle_clone = *handle;
            let mut collider_set = rapier_context.colliders.clone();
            let mut to_remove = vec![];
            for handle in collider_set.iter() {
                if !obstructions.iter().map(|v| v.0).any(|x| x == handle.0) {
                    to_remove.push(handle.0);
                }
            }
            let mut bodies = rapier_context.bodies.clone();
            if !to_remove.is_empty() {
                let mut islands = rapier_context.islands.clone();
                for handle in to_remove {
                    collider_set.remove(handle, &mut islands, &mut bodies, false);
                }
            }
            let shape_clone = (*shape).clone();
            let pool = AsyncComputeTaskPool::get();
            let task = pool.spawn(async move {
                find_path_for_shape(
                    handle_clone.0,
                    coordinates_clone,
                    destination_clone,
                    map_clone,
                    &cost_map_clone,
                    query_pipeline,
                    collider_set,
                    bodies,
                    shape_clone,
                )
            });
            commands.run_if_exists(entity, |mut entity| {
                entity.insert(Calculating(task));
                entity.remove::<Path>();
                entity.remove::<NoPath>();
            });
        }
    }
}

fn poll_tasks(mut commands: Commands, mut query: Query<(Entity, &mut Calculating)>) {
    for (entity, mut calculating) in query.iter_mut() {
        if let Some(result) = future::block_on(future::poll_once(&mut **calculating)) {
            if let Some(path) = result {
                commands.entity(entity).insert(path);
            } else {
                commands.entity(entity).insert(NoPath);
            }
            commands.entity(entity).remove::<Calculating>();
        }
    }
}

fn remove_destination(
    mut commands: Commands,
    entities: &Entities,
    removed: RemovedComponents<Destination>,
) {
    for entity in removed.iter() {
        if entities.contains(entity) {
            commands.entity(entity).remove::<Calculating>();
        }
    }
}

fn negotiate_path(
    mut commands: Commands,
    mut query: Query<(
        Entity,
        &mut Path,
        &mut Transform,
        &mut Velocity,
        &Speed,
        Option<&RotationSpeed>,
    )>,
) {
    for (entity, mut path, mut transform, mut velocity, speed, rotation_speed) in query.iter_mut() {
        let start_i32 = (transform.translation.x, transform.translation.y).i32();
        let mut new_path = path
            .iter()
            .cloned()
            .skip_while(|v| *v == start_i32)
            .collect::<Vec<(i32, i32)>>();
        new_path.retain(|v| *v != start_i32);
        **path = new_path;
        if let Some(next) = path.first() {
            let start = Vec2::new(transform.translation.x, transform.translation.y);
            let next = Vec2::new(next.0 as f32 + 0.5, next.1 as f32 + 0.5);
            let mut direction = next - start;
            direction = direction.normalize();
            direction *= **speed;
            velocity.linvel = direction;
            if rotation_speed.is_some() {
                let v = next - start;
                let angle = v.y.atan2(v.x);
                transform.rotation = Quat::from_rotation_z(angle);
            }
            continue;
        } else {
            commands
                .entity(entity)
                .remove::<Path>()
                .remove::<NoPath>()
                .remove::<Destination>()
                .remove::<Speed>();
            velocity.linvel = Vec2::ZERO;
        }
    }
}

pub struct PathfindingPlugin<D: 'static + Clone + Default + Send + Sync>(PhantomData<D>);

impl<D: 'static + Clone + Default + Send + Sync> Default for PathfindingPlugin<D> {
    fn default() -> Self {
        Self(Default::default())
    }
}

impl<D: 'static + Clone + Default + Send + Sync> Plugin for PathfindingPlugin<D> {
    fn build(&self, app: &mut App) {
        app.add_system(calculate_path::<D>)
            .add_system_to_stage(CoreStage::PostUpdate, remove_destination)
            .add_system(poll_tasks)
            .add_system(negotiate_path.after(crate::navigation::limit_speed));
    }
}