blackout/src/pathfinding.rs

use bevy::{
    ecs::entity::Entities,
    prelude::*,
    tasks::{prelude::*, Task},
    utils::HashMap,
};
use bevy_rapier2d::{
    na::{Isometry2, Vector2},
    prelude::*,
    rapier::prelude::{ColliderHandle, ColliderSet, QueryPipeline, RigidBodySet},
};
use futures_lite::future;
use leafwing_input_manager::{axislike::DualAxisData, plugin::InputManagerSystem, prelude::*};
use pathfinding::prelude::*;

use crate::{
    core::{PointLike, TransformExt},
    map::{Map, MapObstruction},
    navigation::{NavigationAction, RotationSpeed},
};

#[derive(PartialEq, Eq, Clone, Copy, Hash, Debug)]
pub struct NegotiatePathAction;

impl Actionlike for NegotiatePathAction {
    fn n_variants() -> usize {
        1
    }

    fn get_at(index: usize) -> Option<Self> {
        if index == 0 {
            Some(Self)
        } else {
            None
        }
    }

    fn index(&self) -> usize {
        0
    }
}

#[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, Deref, DerefMut)]
#[reflect(Component)]
pub struct CostMap(pub HashMap<(i32, i32), f32>);

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 p.0 >= 0 && p.1 >= 0 {
                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.get(&(tile.0 as i32, tile.1 as i32))
                                {
                                    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(
    initiator: ColliderHandle,
    start: Transform,
    destination: Destination,
    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![];
            let x = p.0;
            let y = p.1;
            let exits = vec![
                ((x - 1, y), 1.),
                ((x + 1, y), 1.),
                ((x, y - 1), 1.),
                ((x, y + 1), 1.),
                ((x - 1, y - 1), 1.5),
                ((x + 1, y - 1), 1.5),
                ((x - 1, y + 1), 1.5),
                ((x + 1, y + 1), 1.5),
            ];
            let pos = Vector2::new(x as f32, y as f32);
            let shape_pos = Isometry2::new(pos, 0.);
            for exit in &exits {
                let mut should_push = true;
                let dest = Vector2::new(exit.0 .0 as f32, exit.0 .1 as f32);
                let shape_vel = dest - pos;
                let max_toi = 1.;
                if query_pipeline
                    .cast_shape(
                        &rigid_body_set,
                        &collider_set,
                        &shape_pos,
                        &shape_vel,
                        &*shape.raw,
                        max_toi,
                        true,
                        bevy_rapier2d::rapier::pipeline::QueryFilter::new()
                            .predicate(&|h, _c| h != initiator),
                    )
                    .is_some()
                {
                    should_push = false;
                }
                if should_push {
                    let mut cost = exit.1 * 100.;
                    if let Some(cost_map) = cost_map {
                        if let Some(modifier) = cost_map.get(&(exit.0 .0, exit.0 .1)) {
                            cost *= modifier;
                        }
                    }
                    successors.push((exit.0, 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(
    mut commands: Commands,
    rapier_context: Res<RapierContext>,
    obstructions: Query<&RapierColliderHandle, With<MapObstruction>>,
    query: Query<
        (
            Entity,
            &RapierColliderHandle,
            &Destination,
            &Transform,
            &Collider,
            Option<&CostMap>,
        ),
        Changed<Destination>,
    >,
) {
    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>();
            continue;
        }
        trace!(
            "Calculating path from {:?}",
            coordinates.translation.truncate().i32()
        );
        let coordinates_clone = *coordinates;
        let destination_clone = *destination;
        let query_pipeline = rapier_context.query_pipeline.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();
        trace!(
            "{entity:?}: path: calculating from {:?} to {destination:?}",
            coordinates.i32(),
        );
        let pool = AsyncComputeTaskPool::get();
        let task = pool.spawn(async move {
            find_path_for_shape(
                handle_clone.0,
                coordinates_clone,
                destination_clone,
                &cost_map_clone,
                query_pipeline,
                collider_set,
                bodies,
                shape_clone,
            )
        });
        commands
            .entity(entity)
            .insert(Calculating(task))
            .remove::<Path>()
            .remove::<NoPath>();
    }
}

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

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

fn negotiate_path(
    mut commands: Commands,
    mut query: Query<(
        Entity,
        &mut ActionState<NavigationAction>,
        &mut Path,
        &mut Transform,
        &Collider,
        Option<&RotationSpeed>,
    )>,
    rapier_context: Res<RapierContext>,
    obstructions: Query<&MapObstruction>,
) {
    for (entity, mut actions, mut path, mut transform, collider, rotation_speed) in &mut query {
        let start_i32 = transform.translation.truncate().i32();
        trace!(
            "{entity:?}: start pathfinding from {start_i32:?} to {:?}: {:?}",
            path.last(),
            transform.translation.truncate()
        );
        if path.len() > 0 {
            if path.len() > 1 && path[0] != start_i32 && path[1] != start_i32 {
                trace!("Off path");
            } else if path.len() > 1 && path[1] == start_i32 {
                trace!("Next in path is start of path, shouldn't happen");
            }
            if path[0] == start_i32 {
                trace!("At start, removing");
                path.remove(0);
            }
        }
        if let Some(next) = path.first() {
            trace!("{entity:?}: path: moving from {start_i32:?} to {next:?}");
            let start = transform.translation.truncate();
            let mut next = Vec2::new(next.0 as f32, next.1 as f32);
            if next.x > 0. {
                next.x += 0.5;
            } else {
                next -= 0.5;
            }
            if next.y > 0. {
                next.y += 0.5;
            } else {
                next.y -= 0.5;
            }
            let mut direction = next - start;
            direction = direction.normalize();
            trace!(
                "{entity:?}: Direction: {direction:?}, Distance: {}",
                (next - start).length()
            );
            if let Some((_hit, _toi)) = rapier_context.cast_shape(
                start,
                transform.yaw().radians(),
                direction,
                collider,
                rapier_context.integration_parameters.dt,
                QueryFilter::new()
                    .predicate(&|entity| obstructions.get(entity).is_ok())
                    .exclude_sensors()
                    .exclude_collider(entity),
            ) {
                // println!("{entity:?} is stuck, hit: {hit:?}, TOI: {toi:?}");
                // TODO: Remove when we have an actual character controller.
                transform.translation = next.extend(0.);
                continue;
            }
            trace!("{entity:?}: path: direction: {direction:?}");
            actions.press(NavigationAction::Move);
            actions.action_data_mut(NavigationAction::Move).axis_pair =
                Some(DualAxisData::from_xy(Vec2::new(-direction.y, direction.x)));
            if rotation_speed.is_some() {
                let angle = direction.y.atan2(direction.x);
                transform.rotation = Quat::from_rotation_z(angle);
            }
        } else {
            trace!("{entity:?}: empty path, cleaning");
            commands
                .entity(entity)
                .remove::<Path>()
                .remove::<NoPath>()
                .remove::<Destination>();
            actions.release(NavigationAction::Move);
            trace!("{entity:?}: pathfinding: cleaned up");
        }
    }
}

fn actions(
    mut commands: Commands,
    mut query: Query<(
        Entity,
        &mut ActionState<NegotiatePathAction>,
        &mut ActionState<NavigationAction>,
        Option<&mut Destination>,
    )>,
) {
    for (entity, mut actions, mut navigation_action, destination) in &mut query {
        if actions.pressed(NegotiatePathAction) {
            if let Some(pair) = actions.axis_pair(NegotiatePathAction) {
                trace!("Negotiating path to {pair:?}");
                let dest = Destination((pair.x() as i32, pair.y() as i32));
                if let Some(mut current_dest) = destination {
                    if *current_dest != dest {
                        trace!("{entity:?}: New destination, zeroing velocity");
                        navigation_action.press(NavigationAction::SetLinearVelocity);
                        navigation_action
                            .action_data_mut(NavigationAction::SetLinearVelocity)
                            .axis_pair = Some(DualAxisData::from_xy(Vec2::ZERO));
                        *current_dest = dest;
                    }
                } else {
                    trace!("{entity:?}: Adding destination, zeroing velocity");
                    navigation_action.press(NavigationAction::SetLinearVelocity);
                    navigation_action
                        .action_data_mut(NavigationAction::SetLinearVelocity)
                        .axis_pair = Some(DualAxisData::from_xy(Vec2::ZERO));
                    commands.entity(entity).insert(dest);
                }
            } else if destination.is_some() {
                commands
                    .entity(entity)
                    .remove::<Destination>()
                    .remove::<NoPath>();
            }
            actions.release(NegotiatePathAction);
            actions.action_data_mut(NegotiatePathAction).axis_pair = None;
        }
    }
}

pub struct PathfindingPlugin;

impl Plugin for PathfindingPlugin {
    fn build(&self, app: &mut App) {
        app.add_plugin(InputManagerPlugin::<NegotiatePathAction>::default())
            .register_type::<Destination>()
            .register_type::<NoPath>()
            .register_type::<Path>()
            .register_type::<CostMap>()
            .add_systems((negotiate_path, poll_tasks).in_base_set(CoreSet::PreUpdate))
            .add_systems(
                (actions, calculate_path)
                    .chain()
                    .in_base_set(CoreSet::PreUpdate)
                    .after(InputManagerSystem::Tick),
            )
            .add_system(remove_destination.in_base_set(CoreSet::PostUpdate));
    }
}