blackout/src/pathfinding.rs

use avian2d::prelude::*;
use bevy::{prelude::*, utils::HashMap};
use leafwing_input_manager::{plugin::InputManagerSystem, prelude::*};
use pathfinding::prelude::*;

use crate::{
    core::GlobalTransformExt,
    map::{Map, MapObstruction},
    navigation::{NavigationAction, RotationSpeed, Speed},
};

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

impl Actionlike for NegotiatePathAction {
    fn input_control_kind(&self) -> InputControlKind {
        InputControlKind::DualAxis
    }
}

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

#[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<IVec2>);

#[derive(Component, Clone, Debug, Default, Reflect, Deref, DerefMut)]
#[reflect(Component)]
pub struct CostMap(pub HashMap<IVec2, f32>);

#[derive(Bundle, Deref, DerefMut)]
pub struct PathfindingControlsBundle(pub ActionState<NegotiatePathAction>);

impl Default for PathfindingControlsBundle {
    fn default() -> Self {
        let mut input: ActionState<NegotiatePathAction> = Default::default();
        input.disable();
        Self(input)
    }
}

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

fn calculate_path(
    mut commands: Commands,
    spatial_query: SpatialQuery,
    mut query: Query<
        (
            Entity,
            &Destination,
            &GlobalTransform,
            &Collider,
            Option<&CostMap>,
            Option<&mut ActionState<NegotiatePathAction>>,
        ),
        Changed<Destination>,
    >,
    obstructions: Query<Entity, With<MapObstruction>>,
    sensors: Query<Entity, With<Sensor>>,
) {
    for (entity, destination, transform, collider, cost_map, actions) in &mut query {
        trace!("{entity}: destination: {destination:?}");
        if transform.translation().truncate().as_ivec2() == **destination {
            commands
                .entity(entity)
                .remove::<Path>()
                .remove::<NoPath>()
                .remove::<Destination>();
            continue;
        }
        commands.entity(entity).remove::<Path>().remove::<NoPath>();
        let path = astar(
            &transform.translation().truncate().as_ivec2(),
            |p| {
                let mut successors: Vec<(IVec2, u32)> = vec![];
                let x = p.x;
                let y = p.y;
                let exits = vec![
                    (IVec2::new(x - 1, y), 1.),
                    (IVec2::new(x + 1, y), 1.),
                    (IVec2::new(x, y - 1), 1.),
                    (IVec2::new(x, y + 1), 1.),
                    (IVec2::new(x - 1, y - 1), 1.5),
                    (IVec2::new(x + 1, y - 1), 1.5),
                    (IVec2::new(x - 1, y + 1), 1.5),
                    (IVec2::new(x + 1, y + 1), 1.5),
                ];
                for exit in &exits {
                    let mut should_push = true;
                    let mut check = exit.0.as_vec2();
                    if check.x >= 0. {
                        check.x += 0.5;
                    } else {
                        check.x -= 0.5;
                    }
                    if check.y >= 0. {
                        check.y += 0.5;
                    } else {
                        check.y -= 0.5;
                    }
                    let hits = spatial_query.shape_intersections(
                        collider,
                        check,
                        transform.yaw().as_radians(),
                        SpatialQueryFilter::from_excluded_entities(&sensors),
                    );
                    for entity in &hits {
                        if obstructions.contains(*entity) {
                            should_push = false;
                            break;
                        }
                    }
                    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) {
                                cost *= modifier;
                            }
                        }
                        successors.push((exit.0, cost as u32));
                    }
                }
                successors
            },
            |p| (p.distance_squared(**destination) * 100) as u32,
            |p| *p == **destination,
        );
        if let Some(path) = path {
            commands.entity(entity).insert(Path(path.0));
        } else {
            commands.entity(entity).insert(NoPath);
            if let Some(mut actions) = actions {
                trace!("{entity:?}: Disabling and resetting because no path");
                actions.disable();
                actions.reset_all();
            }
        }
    }
}

fn remove_destination(mut commands: Commands, mut removed: RemovedComponents<Destination>) {
    for entity in removed.read() {
        if let Some(mut commands) = commands.get_entity(entity) {
            commands.remove::<Path>().remove::<NoPath>();
        }
    }
}

fn negotiate_path(
    mut commands: Commands,
    time: Res<Time>,
    physics_time: Res<Time<Physics>>,
    spatial_query: SpatialQuery,
    mut query: Query<(
        Entity,
        Option<&mut ActionState<NegotiatePathAction>>,
        &mut ActionState<NavigationAction>,
        &mut Path,
        &mut Transform,
        &GlobalTransform,
        &Collider,
        &Speed,
        Option<&RotationSpeed>,
    )>,
    obstructions: Query<Entity, With<MapObstruction>>,
    sensors: Query<Entity, With<Sensor>>,
) {
    if physics_time.is_paused() {
        return;
    }
    for (
        entity,
        actions,
        mut navigation_actions,
        mut path,
        mut transform,
        global_transform,
        collider,
        speed,
        rotation_speed,
    ) in &mut query
    {
        trace!("{entity:?}: negotiating path");
        let start = transform.translation.truncate().as_ivec2();
        if path.len() > 0 && path[0] == start {
            trace!("At start, removing");
            path.remove(0);
        }
        if let Some(next) = path.first() {
            let start = global_transform.translation().truncate();
            let mut next = next.as_vec2();
            if next.x >= 0. {
                next.x += 0.5;
            } else {
                next.x -= 0.5;
            }
            if next.y >= 0. {
                next.y += 0.5;
            } else {
                next.y -= 0.5;
            }
            let mut direction = next - start;
            direction = direction.normalize();
            direction *= **speed;
            let mut hits = vec![];
            spatial_query.shape_hits_callback(
                collider,
                start,
                global_transform.yaw().as_radians(),
                Dir2::new_unchecked(direction.normalize()),
                time.delta_seconds(),
                true,
                SpatialQueryFilter::from_excluded_entities(&sensors),
                |hit| {
                    if obstructions.contains(hit.entity) {
                        hits.push(hit.entity);
                    }
                    true
                },
            );
            if !hits.is_empty() {
                // TODO: Remove when we have an actual character controller.
                next.x = next.x.trunc();
                next.y = next.y.trunc();
                transform.translation = next.extend(0.);
            } else {
                let delta = direction * time.delta_seconds();
                navigation_actions.set_axis_pair(&NavigationAction::Translate, delta);
            }
            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");
            if let Some(mut actions) = actions {
                trace!("{entity:?}: Disabling pathfind because at destination");
                actions.reset_all();
                actions.disable();
            }
            commands
                .entity(entity)
                .remove::<Path>()
                .remove::<NoPath>()
                .remove::<Destination>();
            trace!("{entity:?}: pathfinding: cleaned up");
        }
    }
}

fn actions(
    mut commands: Commands,
    mut query: Query<(
        Entity,
        &ActionState<NegotiatePathAction>,
        &mut ActionState<NavigationAction>,
        Option<&mut Destination>,
    )>,
) {
    for (entity, actions, mut navigation_action, destination) in &mut query {
        if actions.action_disabled(&NegotiatePathAction) {
            if destination.is_some() {
                commands.entity(entity).remove::<Destination>();
            }
        } else {
            let pair = actions.axis_pair(&NegotiatePathAction);
            trace!("{entity:?}: Negotiating path to {pair:?}");
            let dest = Destination(pair.as_ivec2());
            if let Some(mut current_dest) = destination {
                trace!("Got a destination");
                if *current_dest != dest {
                    trace
!("{entity:?}: New destination {dest:?} differs from {current_dest:?}, zeroing velocity");
                    navigation_action
                        .set_axis_pair(&NavigationAction::SetLinearVelocity, Vec2::ZERO);
                    *current_dest = dest;
                }
            } else {
                trace!("{entity:?}: Adding destination, zeroing velocity");
                navigation_action.set_axis_pair(&NavigationAction::SetLinearVelocity, Vec2::ZERO);
                commands.entity(entity).insert(dest);
            }
        }
    }
}

pub struct PathfindingPlugin;

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