blackout/src/pathfinding.rs

use std::collections::HashMap;

use bevy::{prelude::*, tasks::prelude::*};
use bevy_rapier2d::{
    na::UnitComplex,
    prelude::*,
    rapier::data::{ComponentSet, ComponentSetOption, Index},
};
use crossbeam_channel::{unbounded, Receiver, Sender};
use derive_more::{Deref, DerefMut};
use pathfinding::prelude::*;

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

#[derive(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(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![];
            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
    From<(
        &Query<'_, (Entity, &ColliderPosition, &SharedShape, &ColliderFlags)>,
        &Query<'_, &MapObstruction>,
    )> for StaticColliderComponentsSet
{
    fn from(
        query: (
            &Query<(Entity, &ColliderPosition, &SharedShape, &ColliderFlags)>,
            &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<SharedShape> for StaticColliderComponentsSet {
    fn size_hint(&self) -> usize {
        self.0.len()
    }

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

impl ComponentSetOption<SharedShape> for StaticColliderComponentsSet {
    fn get(&self, index: Index) -> Option<&SharedShape> {
        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(
    pool: &AsyncComputeTaskPool,
    query_pipeline: QueryPipeline,
    initiator: Entity,
    start: &dyn PointLike,
    destination: &dyn PointLike,
    map: &Map,
    collider_query: &QueryPipelineColliderComponentsQuery,
    obstructions: &Query<&MapObstruction>,
    shape: &SharedShape,
    channel: &Sender<Option<Path>>,
) {
    let collider_set: StaticColliderComponentsSet = (collider_query, obstructions).into();
    let start = start.i32();
    let destination = destination.i32();
    let map_clone = map.clone();
    let shape_clone = shape.clone();
    let channel_clone = channel.clone();
    pool.spawn(async move {
        let path = astar(
            &start,
            |p| {
                let mut successors: Vec<((i32, i32), u32)> = vec![];
                for tile in map_clone.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_clone,
                        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,
        );
        channel_clone
            .send(if let Some(path) = path {
                Some(Path(path.0))
            } else {
                None
            })
            .ok();
    })
    .detach();
}

fn calculate_path(
    mut commands: Commands,
    pool: Res<AsyncComputeTaskPool>,
    query_pipeline: Res<QueryPipeline>,
    obstructions: Query<&MapObstruction>,
    collider_query: QueryPipelineColliderComponentsQuery,
    mut calculating: Local<HashMap<Entity, Receiver<Option<Path>>>>,
    query: Query<(Entity, &Destination, &Coordinates, &ColliderShape), Changed<Destination>>,
    destinations: Query<&Destination>,
    map: Query<&Map>,
) {
    let calculating_clone = calculating.clone();
    for (entity, rx) in calculating_clone.iter() {
        if destinations.get(*entity).is_ok() {
            if let Ok(msg) = rx.try_recv() {
                if let Some(path) = msg {
                    commands.entity(*entity).insert(path);
                } else {
                    commands.entity(*entity).remove::<Destination>();
                }
                calculating.remove(&entity);
            }
        } else {
            calculating.remove(&entity);
        }
    }
    for (entity, destination, coordinates, shape) in query.iter() {
        if calculating.contains_key(&entity) {
            continue;
        }
        if coordinates.i32() == **destination {
            commands.entity(entity).remove::<Path>();
            commands.entity(entity).remove::<Destination>();
            continue;
        }
        let (tx, rx) = unbounded();
        calculating.insert(entity, rx);
        for map in map.iter() {
            find_path_for_shape(
                &*pool,
                query_pipeline.clone(),
                entity,
                coordinates,
                destination,
                &map,
                &collider_query,
                &obstructions,
                shape,
                &tx,
            );
        }
    }
}

fn negotiate_path(
    mut commands: Commands,
    mut query: Query<(
        Entity,
        &mut Path,
        &mut RigidBodyPosition,
        &mut RigidBodyVelocity,
        &Speed,
        Option<&RotationSpeed>,
    )>,
) {
    for (entity, mut path, mut position, mut velocity, speed, rotation_speed) in query.iter_mut() {
        let mut new_path = path.0.clone();
        let start_i32 = (
            position.position.translation.x,
            position.position.translation.y,
        )
            .i32();
        let new_path_clone = new_path.clone();
        let mut iter = new_path_clone.split(|p| *p == start_i32);
        if iter.next().is_some() {
            if let Some(upcoming) = iter.next() {
                new_path = vec![start_i32];
                new_path.append(&mut upcoming.to_vec());
            } else {
                let start = Vec2::new(
                    position.position.translation.x,
                    position.position.translation.y,
                );
                let next = path[1];
                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();
                continue;
            }
        } else {
            commands.entity(entity).remove::<Path>();
            velocity.linvel = Vec2::ZERO.into();
        }
        **path = new_path;
        if path.len() >= 2 {
            let start = Vec2::new(
                position.position.translation.x,
                position.position.translation.y,
            );
            let next = path[1];
            let next = Vec2::new(next.0 as f32 + 0.5, next.1 as f32 + 0.5);
            if rotation_speed.is_some() {
                let v = next - start;
                let angle = v.y.atan2(v.x);
                position.position.rotation = UnitComplex::new(angle);
            }
            let mut direction = next - start;
            direction = direction.normalize();
            direction *= **speed;
            velocity.linvel = direction.into();
        } else {
            commands.entity(entity).remove::<Path>();
            velocity.linvel = Vec2::ZERO.into();
            commands.entity(entity).remove::<Destination>();
        }
    }
}

pub struct PathfindingPlugin;

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