paiagram_core/

graph.rs

//! Definitions for the graph.

use std::collections::HashMap;

pub mod arrange;

use bevy::ecs::entity::{EntityHash, EntityHashMap, EntityHashSet};
use bevy::prelude::*;
use bevy::tasks::futures_lite::future::poll_once;
use bevy::tasks::{AsyncComputeTaskPool, Task, block_on};
use moonshine_core::kind::{Instance, SpawnInstance};
use moonshine_core::prelude::{MapEntities, ReflectMapEntities};
use petgraph::algo::astar;
use petgraph::prelude::DiGraphMap;
use petgraph::visit::EdgeRef;
use rstar::{AABB, PointDistance, RTree, RTreeObject};
use serde::{Deserialize, Serialize};
use smallvec::SmallVec;

use crate::entry::EntryStop;
use crate::interval::{Interval, IntervalQuery};
use crate::route::Route;
use crate::station::{Platforms, Station};
use crate::units::distance::Distance;

/// The graph plugin. Remember to register this plugin.
pub struct GraphPlugin;
impl Plugin for GraphPlugin {
    fn build(&self, app: &mut App) {
        app.init_resource::<Graph>()
            .init_resource::<GraphSpatialIndex>()
            .init_resource::<GraphSpatialIndexState>()
            .init_resource::<GraphIntervalSpatialIndex>()
            .init_resource::<GraphIntervalSpatialIndexState>()
            .add_systems(Update, arrange::apply_graph_layout_task)
            .add_systems(
                Update,
                (
                    mark_graph_spatial_index_dirty,
                    start_graph_spatial_index_rebuild,
                    apply_graph_spatial_index_task,
                )
                    .chain(),
            )
            .add_systems(
                Update,
                (
                    mark_graph_interval_spatial_index_dirty,
                    start_graph_interval_spatial_index_rebuild,
                    apply_graph_interval_spatial_index_task,
                )
                    .chain(),
            )
            .add_observer(update_graph_on_station_removal)
            .add_observer(update_graph_on_interval_removal)
            .add_observer(add_interval_pair);
        #[cfg(debug_assertions)]
        {
            use bevy::time::common_conditions::on_real_timer;
            app.add_systems(
                PostUpdate,
                check_stations_in_graph.run_if(on_real_timer(std::time::Duration::from_secs(10))),
            );
        }
    }
}

/// The graph. Graph stores the node using an entity and the edge also using an entity entity in a
/// [`DiGraphMap`] using [`EntityHash`].
#[derive(Reflect, Clone, Resource, Serialize, Deserialize, Default, Deref, DerefMut)]
#[reflect(Resource, opaque, Serialize, Deserialize, MapEntities)]
pub struct Graph {
    /// The actual map
    pub map: DiGraphMap<Entity, Entity, EntityHash>,
}

/// Reconstructs the map when mapping entities
impl MapEntities for Graph {
    fn map_entities<E: EntityMapper>(&mut self, entity_mapper: &mut E) {
        // construct a new graph instead
        let (nodes, edges) = self.capacity();
        let mut new_graph = DiGraphMap::with_capacity(nodes, edges);
        for mut node in self.nodes() {
            node.map_entities(entity_mapper);
            new_graph.add_node(node);
        }
        for (mut source, mut target, weight) in self.all_edges() {
            let mut weight = *weight;
            source.map_entities(entity_mapper);
            target.map_entities(entity_mapper);
            weight.map_entities(entity_mapper);
            new_graph.add_edge(source, target, weight);
        }
        self.map = new_graph;
    }
}

impl Graph {
    /// Find a route between the source stop and the target stop.
    /// Returns [`None`] if no valid route exists.
    /// Returns the total length in i32 and the stations on the route if a valid route is found.
    pub fn route_between(
        &self,
        source: Entity,
        target: Entity,
        interval_q: &Query<IntervalQuery>,
    ) -> Option<(i32, Vec<Entity>)> {
        astar(
            &self.map,
            source,
            |f| f == target,
            |e| {
                let Ok(i) = interval_q.get(*e.weight()) else {
                    return i32::MAX;
                };
                i.distance().0
            },
            |_| 0,
        )
    }
    /// Find a route given a set of stations that must be on the route.
    /// Returns [`None`] if no valid route exists.
    /// Returns the total length in i32 and the stations on the route if a valid route is found.
    pub fn route_between_source_waypoint_target(
        &self,
        mut points: impl Iterator<Item = Entity>,
        interval_q: &Query<IntervalQuery>,
    ) -> Option<(i32, Vec<Entity>)> {
        let mut prev = points.next()?;
        let mut total_length = 0;
        let mut passes = vec![prev];
        for curr in points {
            let (leg_length, leg_points) = astar(
                &self.map,
                prev,
                |f| f == curr,
                |e| {
                    let Ok(i) = interval_q.get(*e.weight()) else {
                        return i32::MAX;
                    };
                    i.distance().0
                },
                |_| 0,
            )?;
            total_length += leg_length;
            passes.extend_from_slice(&leg_points[1..]);
            prev = curr;
        }
        Some((total_length, passes))
    }

    /// Transforms the graph into a normal [`petgraph::Graph`], which could be useful for mapping to
    /// e.g. graphviz graphs.
    pub fn into_graph(self) -> petgraph::Graph<Entity, Entity> {
        self.map.into_graph()
    }
}

#[derive(Clone, Copy, Debug)]
struct SpatialIndexedEntity {
    entity: Entity,
    point: [f64; 2],
}

#[derive(Clone, Copy, Debug)]
struct IntervalSpatialIndexedEntity {
    interval: Entity,
    p0: [f64; 2],
    p1: [f64; 2],
}

impl RTreeObject for SpatialIndexedEntity {
    type Envelope = AABB<[f64; 2]>;

    fn envelope(&self) -> Self::Envelope {
        AABB::from_point(self.point)
    }
}

impl PointDistance for SpatialIndexedEntity {
    fn distance_2(&self, point: &[f64; 2]) -> f64 {
        let dx = self.point[0] - point[0];
        let dy = self.point[1] - point[1];
        dx * dx + dy * dy
    }
}

impl RTreeObject for IntervalSpatialIndexedEntity {
    type Envelope = AABB<[f64; 2]>;

    fn envelope(&self) -> Self::Envelope {
        AABB::from_corners(
            [self.p0[0].min(self.p1[0]), self.p0[1].min(self.p1[1])],
            [self.p0[0].max(self.p1[0]), self.p0[1].max(self.p1[1])],
        )
    }
}

#[derive(Resource, Default)]
pub struct GraphSpatialIndex {
    tree: RTree<SpatialIndexedEntity>,
}

#[derive(Clone, Copy, Debug)]
pub struct GraphIntervalSpatialSample {
    pub interval: Entity,
    pub p0: [f64; 2],
    pub p1: [f64; 2],
}

#[derive(Resource, Default)]
pub struct GraphIntervalSpatialIndex {
    tree: RTree<IntervalSpatialIndexedEntity>,
}

impl GraphSpatialIndex {
    pub fn is_empty(&self) -> bool {
        self.tree.size() == 0
    }

    pub fn clear(&mut self) {
        self.tree = RTree::new();
    }

    pub fn insert_xy(&mut self, entity: Entity, x: f64, y: f64) {
        self.tree.insert(SpatialIndexedEntity {
            entity,
            point: [x, y],
        });
    }

    pub fn insert_lon_lat(&mut self, entity: Entity, lon: f64, lat: f64) {
        let (x, y) = lon_lat_to_xy(lon, lat);
        self.insert_xy(entity, x, y);
    }

    pub fn entities_in_xy_aabb(
        &self,
        min_x: f64,
        min_y: f64,
        max_x: f64,
        max_y: f64,
    ) -> Vec<Entity> {
        let envelope = AABB::from_corners(
            [min_x.min(max_x), min_y.min(max_y)],
            [min_x.max(max_x), min_y.max(max_y)],
        );
        self.tree
            .locate_in_envelope_intersecting(&envelope)
            .map(|entry| entry.entity)
            .collect()
    }

    pub fn entities_in_lon_lat_aabb(
        &self,
        min_lon: f64,
        min_lat: f64,
        max_lon: f64,
        max_lat: f64,
    ) -> Vec<Entity> {
        let (x0, y0) = lon_lat_to_xy(min_lon, min_lat);
        let (x1, y1) = lon_lat_to_xy(max_lon, max_lat);
        self.entities_in_xy_aabb(x0, y0, x1, y1)
    }

    pub fn nearest_in_xy(&self, x: f64, y: f64) -> Option<Entity> {
        self.tree
            .nearest_neighbor(&[x, y])
            .map(|entry| entry.entity)
    }

    pub fn nearest_in_lon_lat(&self, lon: f64, lat: f64) -> Option<Entity> {
        let (x, y) = lon_lat_to_xy(lon, lat);
        self.nearest_in_xy(x, y)
    }

    fn replace_tree(&mut self, tree: RTree<SpatialIndexedEntity>) {
        self.tree = tree;
    }
}

impl GraphIntervalSpatialIndex {
    pub fn is_empty(&self) -> bool {
        self.tree.size() == 0
    }

    pub fn query_xy_aabb(
        &self,
        min_x: f64,
        min_y: f64,
        max_x: f64,
        max_y: f64,
    ) -> Vec<GraphIntervalSpatialSample> {
        if self.is_empty() {
            return Vec::new();
        }

        let envelope = AABB::from_corners(
            [min_x.min(max_x), min_y.min(max_y)],
            [min_x.max(max_x), min_y.max(max_y)],
        );

        self.tree
            .locate_in_envelope_intersecting(&envelope)
            .map(|item| GraphIntervalSpatialSample {
                interval: item.interval,
                p0: item.p0,
                p1: item.p1,
            })
            .collect()
    }

    fn replace_tree(&mut self, tree: RTree<IntervalSpatialIndexedEntity>) {
        self.tree = tree;
    }
}

#[derive(Resource)]
struct GraphSpatialIndexState {
    dirty: bool,
    task: Option<Task<RTree<SpatialIndexedEntity>>>,
}

#[derive(Resource)]
struct GraphIntervalSpatialIndexState {
    dirty: bool,
    task: Option<Task<RTree<IntervalSpatialIndexedEntity>>>,
}

impl Default for GraphSpatialIndexState {
    fn default() -> Self {
        Self {
            dirty: true,
            task: None,
        }
    }
}

impl Default for GraphIntervalSpatialIndexState {
    fn default() -> Self {
        Self {
            dirty: true,
            task: None,
        }
    }
}

// EPSG:3857
const EARTH_RADIUS_METERS: f64 = 6_378_137.0;
const WEB_MERCATOR_MAX_LAT: f64 = 85.051_128_78;

pub fn lon_lat_to_xy(lon: f64, lat: f64) -> (f64, f64) {
    let x = EARTH_RADIUS_METERS * lon.to_radians();
    let lat = lat.clamp(-WEB_MERCATOR_MAX_LAT, WEB_MERCATOR_MAX_LAT);
    let lat_rad = lat.to_radians();
    let y = -EARTH_RADIUS_METERS * (std::f64::consts::FRAC_PI_4 + lat_rad / 2.0).tan().ln();
    (x, y)
}

pub fn xy_to_lon_lat(x: f64, y: f64) -> (f64, f64) {
    let lon = (x / EARTH_RADIUS_METERS).to_degrees();
    let lat =
        (2.0 * (-y / EARTH_RADIUS_METERS).exp().atan() - std::f64::consts::FRAC_PI_2).to_degrees();
    (lon, lat)
}

// TODO: partial update
fn mark_graph_spatial_index_dirty(
    mut state: ResMut<GraphSpatialIndexState>,
    changed_nodes: Query<(), Or<(Added<Node>, Changed<Node>)>>,
    mut removed_nodes: RemovedComponents<Node>,
) {
    if !changed_nodes.is_empty() || removed_nodes.read().next().is_some() {
        state.dirty = true;
    }
}

fn mark_graph_interval_spatial_index_dirty(
    mut state: ResMut<GraphIntervalSpatialIndexState>,
    graph: Res<Graph>,
    changed_nodes: Query<(), Or<(Added<Node>, Changed<Node>)>>,
    changed_intervals: Query<(), Or<(Added<Interval>, Changed<Interval>)>>,
    mut removed_nodes: RemovedComponents<Node>,
    mut removed_intervals: RemovedComponents<Interval>,
) {
    if graph.is_added()
        || graph.is_changed()
        || !changed_nodes.is_empty()
        || !changed_intervals.is_empty()
        || removed_nodes.read().next().is_some()
        || removed_intervals.read().next().is_some()
    {
        state.dirty = true;
    }
}

fn start_graph_spatial_index_rebuild(
    mut state: ResMut<GraphSpatialIndexState>,
    nodes: Query<(Entity, &Node)>,
) {
    if !state.dirty || state.task.is_some() {
        return;
    }
    state.dirty = false;

    let snapshot: Vec<(Entity, [f64; 2])> = nodes
        .iter()
        .map(|(entity, node)| (entity, node.coor.to_xy_arr()))
        .collect();
    state.task = Some(AsyncComputeTaskPool::get().spawn(async move {
        let entries: Vec<SpatialIndexedEntity> = snapshot
            .into_iter()
            .map(|(entity, point)| SpatialIndexedEntity { entity, point })
            .collect();
        RTree::bulk_load(entries)
    }));
}

fn start_graph_interval_spatial_index_rebuild(
    mut state: ResMut<GraphIntervalSpatialIndexState>,
    graph: Res<Graph>,
    nodes: Query<&Node>,
) {
    if !state.dirty || state.task.is_some() {
        return;
    }
    state.dirty = false;

    let mut snapshot = Vec::<IntervalSpatialIndexedEntity>::new();
    for (source, target, interval) in graph.all_edges() {
        let Ok(source_node) = nodes.get(source) else {
            continue;
        };
        let Ok(target_node) = nodes.get(target) else {
            continue;
        };
        snapshot.push(IntervalSpatialIndexedEntity {
            interval: *interval,
            p0: source_node.coor.to_xy_arr(),
            p1: target_node.coor.to_xy_arr(),
        });
    }

    state.task = Some(AsyncComputeTaskPool::get().spawn(async move { RTree::bulk_load(snapshot) }));
}

fn apply_graph_spatial_index_task(
    mut state: ResMut<GraphSpatialIndexState>,
    mut index: ResMut<GraphSpatialIndex>,
) {
    let Some(task) = state.task.as_mut() else {
        return;
    };
    let Some(tree) = block_on(poll_once(task)) else {
        return;
    };
    index.replace_tree(tree);
    state.task = None;
}

fn apply_graph_interval_spatial_index_task(
    mut state: ResMut<GraphIntervalSpatialIndexState>,
    mut index: ResMut<GraphIntervalSpatialIndex>,
) {
    let Some(task) = state.task.as_mut() else {
        return;
    };
    let Some(tree) = block_on(poll_once(task)) else {
        return;
    };
    index.replace_tree(tree);
    state.task = None;
}

/// The position of the node.
///
/// This stores longitude and latitude values only.
#[derive(Reflect, Clone, Copy, Debug, PartialEq)]
pub struct NodeCoor {
    pub lon: f64,
    pub lat: f64,
}

impl std::fmt::Display for NodeCoor {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let lat_dir = if self.lat < 0.0 { 'S' } else { 'N' };
        let lon_dir = if self.lon < 0.0 { 'W' } else { 'E' };
        write!(
            f,
            "{:.4}°{}, {:.4}°{}",
            self.lat.abs(),
            lat_dir,
            self.lon.abs(),
            lon_dir
        )
    }
}

impl Default for NodeCoor {
    fn default() -> Self {
        Self::new(0.0, 0.0)
    }
}

impl NodeCoor {
    pub fn new(lon: f64, lat: f64) -> Self {
        Self { lon, lat }
    }
    pub fn from_xy(x: f64, y: f64) -> Self {
        let (lon, lat) = xy_to_lon_lat(x, y);
        Self::new(lon, lat)
    }
    pub fn to_xy(&self) -> (f64, f64) {
        lon_lat_to_xy(self.lon, self.lat)
    }
    pub fn to_xy_arr(&self) -> [f64; 2] {
        let (x, y) = self.to_xy();
        [x, y]
    }
    /// Shift the node on the canvas by x and y
    pub fn shift(&mut self, dx: f64, dy: f64) {
        self.lon += dx;
        self.lat += dy;
    }
    /// Linearly interpolates between `self` and `other` by fraction `t`.
    /// `t` is typically between 0.0 and 1.0.
    pub fn lerp(&self, other: &Self, t: f64) -> Self {
        let end_lon = other.lon;
        let end_lat = other.lat;

        Self {
            lon: self.lon + (end_lon - self.lon) * t,
            lat: self.lat + (end_lat - self.lat) * t,
        }
    }
}

#[derive(Default, Reflect, Component, Debug)]
#[reflect(Component)]
pub struct Node {
    pub coor: NodeCoor,
}

fn update_graph_on_station_removal(
    removed_station: On<Remove, Station>,
    mut commands: Commands,
    mut graph: ResMut<Graph>,
) {
    let s = removed_station.entity;
    for e in graph
        .neighbors_directed(s, petgraph::Direction::Incoming)
        .chain(graph.neighbors_directed(s, petgraph::Direction::Outgoing))
    {
        commands.entity(e).despawn();
    }
    graph.remove_node(s);
}

fn update_graph_on_interval_removal(
    removed_interval: On<Remove, Interval>,
    mut graph: ResMut<Graph>,
) {
    let i = removed_interval.entity;
    let mut source = None;
    let mut target = None;
    for (s, t, weight) in graph.all_edges() {
        if i != *weight {
            continue;
        }
        source = Some(s);
        target = Some(t);
        break;
    }
    let (Some(s), Some(t)) = (source, target) else {
        return;
    };
    graph.remove_edge(s, t);
}

#[cfg(debug_assertions)]
fn check_stations_in_graph(
    graph: Res<Graph>,
    stations: Populated<Entity, With<Station>>,
    intervals: Populated<Entity, With<Interval>>,
    names: Query<&Name>,
) {
    let queried_station_set: EntityHashSet = stations.iter().collect();
    let queried_interval_set: EntityHashSet = intervals.iter().collect();
    let mut graphed_station_set = EntityHashSet::new();
    let mut graphed_interval_set = EntityHashSet::new();
    for (_, _, w) in graph.all_edges() {
        graphed_interval_set.insert(*w);
    }
    for node in graph.nodes() {
        graphed_station_set.insert(node);
    }
    if queried_station_set != graphed_station_set {
        debug_graph_set_diff(
            "station",
            &queried_station_set,
            &graphed_station_set,
            &names,
        );
    }
    if queried_interval_set != graphed_interval_set {
        debug_graph_set_diff(
            "interval",
            &queried_interval_set,
            &graphed_interval_set,
            &names,
        );
    }
    debug_assert_eq!(queried_station_set, graphed_station_set);
    debug_assert_eq!(queried_interval_set, graphed_interval_set);
}

#[cfg(debug_assertions)]
fn debug_graph_set_diff(
    label: &str,
    queried: &EntityHashSet,
    graphed: &EntityHashSet,
    names: &Query<&Name>,
) {
    let intersection: EntityHashSet = queried.intersection(graphed).copied().collect();
    let only_queried: EntityHashSet = queried.difference(graphed).copied().collect();
    let only_graphed: EntityHashSet = graphed.difference(queried).copied().collect();

    let list_with_names = |set: &EntityHashSet| -> Vec<String> {
        let mut out: Vec<String> = set
            .iter()
            .map(|e| match names.get(*e) {
                Ok(name) => format!("{} ({})", name.as_str(), e.index()),
                Err(_) => format!("<??> ({})", e.index()),
            })
            .collect();
        out.sort_unstable();
        out
    };

    warn!(
        "Graph {label} set mismatch: intersection={:#?} | only_queried={:#?} | only_graphed={:#?}",
        list_with_names(&intersection),
        list_with_names(&only_queried),
        list_with_names(&only_graphed)
    );
}

// TODO: instead of merging them, make stations platforms instead
pub fn merge_station_by_name(
    mut commands: Commands,
    mut graph: ResMut<Graph>,
    stations: Query<(Entity, &Name, &Platforms), With<Station>>,
    entry_stops: Query<(Entity, &EntryStop)>,
    mut routes: Query<&mut Route>,
) {
    let mut name_map: HashMap<&str, SmallVec<[Entity; 1]>> = HashMap::new();
    for (entity, name, _) in &stations {
        let v = name_map.entry(name.as_str()).or_default();
        v.push(entity);
    }

    let mut remap: EntityHashMap<Entity> = EntityHashMap::default();

    for (_name, mut entities) in name_map.into_iter().filter(|(_, v)| v.len() > 1) {
        entities.sort_unstable_by_key(|entity| entity.index());
        let keep = entities[0];

        for duplicate in entities.into_iter().skip(1) {
            if let Ok((_, _, platforms)) = stations.get(duplicate) {
                let to_move: SmallVec<[Entity; 8]> = platforms.iter().collect();
                if !to_move.is_empty() {
                    commands.entity(keep).add_children(&to_move);
                }
            }
            remap.insert(duplicate, keep);
        }
    }

    if remap.is_empty() {
        return;
    }

    let (nodes, edges) = graph.capacity();
    let mut new_graph = DiGraphMap::with_capacity(nodes, edges);
    let mut removed_intervals = EntityHashSet::default();
    for (source, target, weight) in graph.all_edges() {
        let source = remap.get(&source).copied().unwrap_or(source);
        let target = remap.get(&target).copied().unwrap_or(target);
        if let Some(existing_weight) = new_graph.edge_weight(source, target) {
            if *existing_weight != *weight {
                removed_intervals.insert(*weight);
            }
            continue;
        }
        new_graph.add_edge(source, target, *weight);
    }
    for node in graph.nodes() {
        let node = remap.get(&node).copied().unwrap_or(node);
        new_graph.add_node(node);
    }
    graph.map = new_graph;

    for (entry, stop) in &entry_stops {
        if let Some(new_stop) = remap.get(&stop.0) {
            commands.entity(entry).insert(EntryStop(*new_stop));
        }
    }

    for mut route in &mut routes {
        for stop in &mut route.stops {
            if let Some(new_stop) = remap.get(stop) {
                *stop = *new_stop;
            }
        }
    }

    for interval in removed_intervals {
        commands.entity(interval).despawn();
    }

    for duplicate in remap.into_keys() {
        commands.entity(duplicate).despawn();
    }
}

#[derive(Event, Clone, Copy)]
pub struct AddIntervalPair {
    pub source: Entity,
    pub target: Entity,
    pub length: Distance,
}

fn add_interval_pair(msg: On<AddIntervalPair>, mut graph: ResMut<Graph>, mut commands: Commands) {
    if !graph.contains_edge(msg.source, msg.target) {
        let e1: Instance<Interval> = commands
            .spawn_instance(Interval { length: msg.length })
            .into();
        graph.add_edge(msg.source, msg.target, e1.entity());
    }
    if !graph.contains_edge(msg.target, msg.source) {
        let e2: Instance<Interval> = commands
            .spawn_instance(Interval { length: msg.length })
            .into();
        graph.add_edge(msg.target, msg.source, e2.entity());
    }
}