blackout/src/pathfinding.rs

use std::collections::HashMap;

use bevy::{
    ecs::entity::Entities,
    prelude::*,
    tasks::{prelude::*, Task},
};
use bevy_rapier2d::{
    na::UnitComplex,
    prelude::*,
    rapier::data::{ComponentSet, ComponentSetOption, Index},
};
use derive_more::{Deref, DerefMut};
use futures_lite::future;
use mapgen::Tile;
use pathfinding::prelude::*;

use crate::{
    core::{Coordinates, 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);

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

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

struct StaticColliderComponentsSet(
    HashMap<Entity, (ColliderPosition, ColliderShape, ColliderFlags)>,
);

impl<'world, 'state>
    From<(
        &Query<
            'world,
            'state,
            (
                Entity,
                &ColliderPositionComponent,
                &ColliderShapeComponent,
                &ColliderFlagsComponent,
            ),
        >,
        &Query<'world, 'state, &MapObstruction>,
    )> for StaticColliderComponentsSet
{
    fn from(
        query: (
            &Query<(
                Entity,
                &ColliderPositionComponent,
                &ColliderShapeComponent,
                &ColliderFlagsComponent,
            )>,
            &Query<&MapObstruction>,
        ),
    ) -> Self {
        let entries = query
            .0
            .iter()
            .filter(|(a, _, _, _)| query.1.get(*a).is_ok())
            .map(|(a, b, c, d)| (a, **b, (*c).clone(), **d))
            .collect::<Vec<(Entity, ColliderPosition, ColliderShape, ColliderFlags)>>();
        let mut m = HashMap::new();
        for (e, a, b, c) in entries {
            m.insert(e, (a, b, c));
        }
        Self(m)
    }
}

impl ComponentSet<ColliderShape> for StaticColliderComponentsSet {
    fn size_hint(&self) -> usize {
        self.0.len()
    }

    fn for_each(&self, _f: impl FnMut(Index, &ColliderShape)) {
        unimplemented!()
    }
}

impl ComponentSetOption<ColliderShape> for StaticColliderComponentsSet {
    fn get(&self, index: Index) -> Option<&ColliderShape> {
        self.0.get(&index.entity()).map(|v| &v.1)
    }
}

impl ComponentSet<ColliderFlags> for StaticColliderComponentsSet {
    fn size_hint(&self) -> usize {
        self.0.len()
    }

    fn for_each(&self, _f: impl FnMut(Index, &ColliderFlags)) {
        unimplemented!()
    }
}

impl ComponentSetOption<ColliderFlags> for StaticColliderComponentsSet {
    fn get(&self, index: Index) -> Option<&ColliderFlags> {
        self.0.get(&index.entity()).map(|v| &v.2)
    }
}

impl ComponentSet<ColliderPosition> for StaticColliderComponentsSet {
    fn size_hint(&self) -> usize {
        self.0.len()
    }

    fn for_each(&self, _f: impl FnMut(Index, &ColliderPosition)) {
        unimplemented!()
    }
}

impl ComponentSetOption<ColliderPosition> for StaticColliderComponentsSet {
    fn get(&self, index: Index) -> Option<&ColliderPosition> {
        let v = self.0.get(&index.entity()).map(|v| &v.0);
        v
    }
}

fn find_path_for_shape(
    initiator: Entity,
    start: Coordinates,
    destination: Destination,
    query_pipeline: QueryPipeline,
    map: Map,
    collider_set: StaticColliderComponentsSet,
    shape: ColliderShape,
) -> Option<Path> {
    let path = astar(
        &start.i32(),
        |p| {
            let mut successors: Vec<((i32, i32), u32)> = vec![];
            if map.at(p.0 as usize, p.1 as usize).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 = Isometry::new(vector![tile.0 as f32, tile.1 as f32], 0.);
                    query_pipeline.intersections_with_shape(
                        &collider_set,
                        &shape_pos,
                        &*shape,
                        InteractionGroups::all(),
                        Some(&|v| v.entity() != initiator),
                        |_handle| {
                            should_push = false;
                            false
                        },
                    );
                    if should_push {
                        successors.push(((tile.0 as i32, tile.1 as i32), (tile.2 * 100.) 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,
    pool: Res<AsyncComputeTaskPool>,
    query_pipeline: Res<QueryPipeline>,
    obstructions: Query<&MapObstruction>,
    collider_query: QueryPipelineColliderComponentsQuery,
    query: Query<
        (Entity, &Destination, &Coordinates, &ColliderShapeComponent),
        Changed<Destination>,
    >,
    map: Query<&Map>,
) {
    for (entity, destination, coordinates, shape) in query.iter() {
        if coordinates.i32() == **destination {
            commands
                .entity(entity)
                .remove::<Path>()
                .remove::<NoPath>()
                .remove::<Calculating>()
                .remove::<Destination>()
                .remove::<Speed>();
            continue;
        }
        let map = map.single();
        let coordinates_clone = *coordinates;
        let destination_clone = *destination;
        let query_pipeline_clone = query_pipeline.clone();
        let map_clone = map.clone();
        let shape_clone = (*shape).clone();
        let collider_set: StaticColliderComponentsSet = (&collider_query, &obstructions).into();
        let task = pool.spawn(async move {
            find_path_for_shape(
                entity,
                coordinates_clone,
                destination_clone,
                query_pipeline_clone,
                map_clone,
                collider_set,
                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 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 RigidBodyPositionComponent,
        &mut RigidBodyVelocityComponent,
        &Speed,
        Option<&RotationSpeed>,
    )>,
) {
    for (entity, mut path, mut position, mut velocity, speed, rotation_speed) in query.iter_mut() {
        let start_i32 = (
            position.position.translation.x,
            position.position.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(
                position.position.translation.x,
                position.position.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.into();
            if rotation_speed.is_some() {
                let v = next - start;
                let angle = v.y.atan2(v.x);
                position.position.rotation = UnitComplex::new(angle);
            }
            continue;
        } else {
            commands
                .entity(entity)
                .remove::<Path>()
                .remove::<NoPath>()
                .remove::<Destination>()
                .remove::<Speed>();
            velocity.linvel = Vec2::ZERO.into();
        }
    }
}

fn remove_calculating(
    mut commands: Commands,
    query: Query<Entity, (Changed<Destination>, With<Calculating>)>,
) {
    for entity in query.iter() {
        commands.entity(entity).remove::<Calculating>();
    }
}

pub struct PathfindingPlugin;

impl Plugin for PathfindingPlugin {
    fn build(&self, app: &mut App) {
        app.add_system(calculate_path)
            .add_system_to_stage(CoreStage::PostUpdate, remove_destination)
            .add_system(poll_tasks)
            .add_system(negotiate_path)
            .add_system_to_stage(CoreStage::PostUpdate, remove_calculating);
    }
}