network_graph 0.1.2

use std::collections::{HashMap, HashSet, hash_map::Keys};

use egui::epaint::CircleShape;
pub use egui::{Color32, Pos2};
use egui::{Frame, Id, Sense, Shape, Stroke, Ui, Vec2, vec2};
#[cfg(feature = "log_gui")]
use log::trace;

use crate::graph::{
    Edge, Graph, Identifier, InvalidEdgeEndError, Node, RemainingConnectedEdgesError,
};

pub mod helpers;

/// Style attributes for a graph node.
#[derive(Debug, Clone)]
pub struct NodeStyle {
    pub radius: f32,
    pub max_radius: f32,
    pub min_radius: f32,
    pub fill_color: Color32,
    pub border_width: f32,
    pub border_color: Color32,
    /// Disables the node from being draggable.
    pub fixed: bool,
    pub hidden: bool,
}

impl NodeStyle {
    fn shape(&self, center: Pos2, scale: f32) -> Shape {
        let mut radius = self.radius * scale;
        radius = radius.min(self.max_radius);
        radius = radius.max(self.min_radius);

        Shape::Circle(CircleShape {
            center,
            radius,
            fill: self.fill_color,
            stroke: Stroke {
                width: self.border_width,
                color: self.border_color,
            },
        })
    }
}

impl Default for NodeStyle {
    fn default() -> Self {
        return Self {
            radius: 10.0,
            max_radius: 20.0,
            min_radius: 5.0,
            fill_color: Color32::WHITE,
            border_width: 2.0,
            border_color: Color32::DARK_GRAY,
            fixed: false,
            hidden: false,
        };
    }
}

/// Style attributes for a graph edge.
#[derive(Debug, Clone)]
pub struct EdgeStyle {
    pub line_width: f32,
    pub line_color: Color32,
    pub hidden: bool,
}

impl EdgeStyle {
    fn shape(&self, end_1: Pos2, end_2: Pos2) -> Shape {
        return Shape::LineSegment {
            points: [end_1, end_2],
            stroke: Stroke {
                width: self.line_width,
                color: self.line_color,
            },
        };
    }
}

impl Default for EdgeStyle {
    fn default() -> Self {
        return Self {
            line_width: 2.0,
            line_color: Color32::LIGHT_GRAY,
            hidden: false,
        };
    }
}

/// Entry for a node within a graph.
#[derive(Debug)]
pub struct NodeEntry<NodeIdentifier>
where
    NodeIdentifier: Identifier,
{
    /// Unique identifier for the node within the graph.
    pub id: NodeIdentifier,
    /// Absolute position within the graph.
    pub pos: Pos2,
    /// Appearance of the node.
    pub style: NodeStyle,
}

impl<NodeIdentifier> Node for NodeEntry<NodeIdentifier>
where
    NodeIdentifier: Identifier,
{
    type Identifier = NodeIdentifier;

    fn identifier(&self) -> &Self::Identifier {
        return &self.id;
    }
}

/// Entry for an edge within a graph.
#[derive(Debug)]
pub struct EdgeEntry<EdgeIdentifier, NodeIdentifier>
where
    NodeIdentifier: Identifier,
    EdgeIdentifier: Identifier,
{
    /// Unique identifier for the edge within the graph.
    pub id: EdgeIdentifier,
    /// Appearane of the edge.
    pub style: EdgeStyle,
    /// Which node identifiers the edge attaches to.
    pub ends: [NodeIdentifier; 2],
}

impl<NodeIdentifier, EdgeIdentifier> Edge<NodeEntry<NodeIdentifier>>
    for EdgeEntry<EdgeIdentifier, NodeIdentifier>
where
    NodeIdentifier: Identifier,
    EdgeIdentifier: Identifier,
{
    type Identifier = EdgeIdentifier;

    fn identifier(&self) -> &Self::Identifier {
        return &self.id;
    }

    fn ends(&self) -> [<NodeEntry<NodeIdentifier> as Node>::Identifier; 2] {
        return self.ends.clone();
    }
}

/// Configuration for the current position and zoom within the graph.
#[derive(Debug)]
pub struct CanvasState {
    /// Scalar to zoom in/out.
    zoom_scale: f32,
    /// Coordinate at center of graph in original (non-mapped) units.
    pos: Vec2,
}

impl CanvasState {
    /// Helper function to adjuct canvas position.
    pub fn drag(&mut self, delta: Vec2) {
        self.pos -= delta / self.zoom_scale;
    }

    /// Helper function to adjust zoom multiplier.
    pub fn zoom(&mut self, scale: f32) {
        self.zoom_scale *= scale;
    }
}

impl Default for CanvasState {
    fn default() -> Self {
        return Self {
            zoom_scale: 1.0,
            pos: vec2(0.0, 0.0),
        };
    }
}

#[derive(Debug, Default)]
pub struct NodeInteraction {
    /// Distance the node was dragged by, pre-adjusted for zoom.
    pub drag: Option<Vec2>,
}

#[derive(Debug, Default)]
pub struct CanvasInteraction {
    /// Vector the canvas was dragged by, pre-adjusted for zoom.
    pub drag: Option<Vec2>,
    /// Multiplier the canvas was zoomed by.
    pub zoom: Option<f32>,
}

/// Utility for modifying nodes while within the graph.
pub struct NodeEditContext<'a, NodeIdentifier, EdgeIdentifier>
where
    NodeIdentifier: Identifier,
    EdgeIdentifier: Identifier,
{
    /// Unique identifier for the node within the graph.
    pub id: &'a NodeIdentifier,
    /// Absolute position within the graph.
    pub pos: &'a mut Pos2,
    /// Appearance of the node.
    pub style: &'a mut NodeStyle,
    /// Edges that the node is connected to.
    pub edges: &'a HashSet<EdgeIdentifier>,
    /// Interaction from the last widget render.
    pub interaction: Option<NodeInteraction>,
}

impl<'a, NodeIdentifier, EdgeIdentifier> NodeEditContext<'a, NodeIdentifier, EdgeIdentifier>
where
    NodeIdentifier: Identifier,
    EdgeIdentifier: Identifier,
{
    /// Helper function to create new edit contexts.
    fn new(
        node: &'a mut NodeEntry<NodeIdentifier>,
        edges: &'a HashSet<EdgeIdentifier>,
        interaction: Option<NodeInteraction>,
    ) -> Self {
        return Self {
            id: &node.id,
            pos: &mut node.pos,
            style: &mut node.style,
            edges,
            interaction,
        };
    }
    /// Applies trivial node interaction.
    pub fn default_interact(&mut self) {
        if let Some(drag) = self.interaction.as_ref().map(|i| i.drag).flatten() {
            *self.pos += drag;
        }
    }
}

/// Utility for modifying edges while within the graph.
pub struct EdgeEditContext<'a, EdgeIdentifier, NodeIdentifier>
where
    EdgeIdentifier: Identifier,
    NodeIdentifier: Identifier,
{
    /// Unique identifier for the edge within the graph.
    pub id: &'a EdgeIdentifier,
    /// Appearane of the edge.
    pub style: &'a mut EdgeStyle,
    /// Which node identifiers the edge attaches to.
    pub ends: &'a [NodeIdentifier; 2],
}

/// Utility for modying the canvas.
pub struct CanvasEditContext<'a> {
    /// Canvas state.
    pub canvas: &'a mut CanvasState,
    /// Interaction from the last widget render.
    pub interaction: Option<CanvasInteraction>,
}

impl<'a> CanvasEditContext<'a> {
    /// Applies trivial canvas interaction.
    pub fn default_interact(&mut self) {
        if let Some(scale) = self.interaction.as_ref().map(|i| i.zoom).flatten() {
            self.canvas.zoom(scale.clone());
        }
        if let Some(delta) = &self.interaction.as_ref().map(|i| i.drag).flatten() {
            self.canvas.drag(*delta);
        }
    }
}

/// Utility for mapping locations of points based on canvas position and zoom.
#[derive(Debug)]
struct CanvasMapper<'a> {
    pos: &'a CanvasState,
    bl: Vec2,
}

impl<'a> CanvasMapper<'a> {
    fn compile(screen_dims: Vec2, canvas: &'a CanvasState) -> Self {
        // Range of canvas in pre-transformed canvas
        let range = screen_dims / canvas.zoom_scale;
        // Bottom left in pre-transformed canvas when centered at origin
        let bl = range / -2.0;
        return Self { pos: canvas, bl };
    }

    pub fn map_pos(&self, pos: Pos2) -> Pos2 {
        return (pos - self.pos.pos - self.bl) * self.pos.zoom_scale;
    }
}

/// Raw widget for rendering node graphs.
///
/// This widget by itself is very dumb, and should be wrapped with a controller.
#[derive(Debug, Default)]
pub struct GraphWidget<NodeIdentifier, EdgeIdentifier>
where
    NodeIdentifier: Identifier,
    EdgeIdentifier: Identifier,
{
    nodes: HashMap<NodeIdentifier, NodeEntry<NodeIdentifier>>,
    nodes_adjacency: HashMap<NodeIdentifier, HashSet<EdgeIdentifier>>,
    edges: HashMap<EdgeIdentifier, EdgeEntry<EdgeIdentifier, NodeIdentifier>>,
    canvas: CanvasState,
    canvas_interaction: Option<CanvasInteraction>,
    node_interactions: HashMap<NodeIdentifier, NodeInteraction>,
}

impl<NodeIdentifier, EdgeIdentifier> GraphWidget<NodeIdentifier, EdgeIdentifier>
where
    NodeIdentifier: Identifier,
    EdgeIdentifier: Identifier,
{
    /// Inserts an edge into the graph.
    ///
    /// If the node ends are not present at the time the widget is rendered, it will be removed.
    pub fn insert_edge_lazy(
        &mut self,
        edge: EdgeEntry<EdgeIdentifier, NodeIdentifier>,
    ) -> Option<EdgeEntry<EdgeIdentifier, NodeIdentifier>> {
        return self.edges.insert(edge.id.clone(), edge);
    }

    /// Removes a node from the graph.
    ///
    /// If there are edges that depend on the node after the node is removed, they will be removed
    /// at the time the widget is rendered.
    pub fn remove_node_lazy(
        &mut self,
        identifier: &NodeIdentifier,
    ) -> Option<NodeEntry<NodeIdentifier>> {
        self.nodes_adjacency.remove(identifier);
        return self.nodes.remove(identifier);
    }

    /// Recieve interaction from the canvas and transform the canvas' state.
    ///
    /// This takes and returns the canvas interaction.
    pub fn modify_canvas<'a>(&'a mut self) -> CanvasEditContext<'a> {
        return CanvasEditContext {
            canvas: &mut self.canvas,
            interaction: self.canvas_interaction.take(),
        };
    }

    /// Recieve interaction from a node and transform the node's state.
    pub fn modify_node<'a>(
        &'a mut self,
        node_identifier: NodeIdentifier,
    ) -> Option<NodeEditContext<'a, NodeIdentifier, EdgeIdentifier>> {
        let node_ref = self.nodes.get_mut(&node_identifier)?;
        return Some(NodeEditContext::new(
            node_ref,
            self.nodes_adjacency.get(&node_identifier).unwrap(),
            self.node_interactions.remove(&node_identifier),
        ));
    }

    /// Return iterator of nodes that have interactions.
    pub fn interacted_nodes<'a>(&'a self) -> Keys<'a, NodeIdentifier, NodeInteraction> {
        return self.node_interactions.keys();
    }

    /// Return iterator of node interaction contexts.
    pub fn modify_interacted_nodes<F, R>(&mut self, apply: F) -> Vec<R>
    where
        F: Fn(NodeEditContext<NodeIdentifier, EdgeIdentifier>) -> R,
    {
        let mut res = Vec::new();
        for (id, interaction) in self.node_interactions.drain() {
            let Some(node) = self.nodes.get_mut(&id) else {
                continue;
            };
            let edges = self.nodes_adjacency.get(&id).unwrap();
            res.push(apply(NodeEditContext::new(node, edges, Some(interaction))));
        }
        return res;
    }

    /// Transform an edge's style.
    pub fn modify_edge<'a>(
        &'a mut self,
        edge_identifier: EdgeIdentifier,
    ) -> Option<EdgeEditContext<'a, EdgeIdentifier, NodeIdentifier>> {
        let edge_ref = self.edges.get_mut(&edge_identifier)?;
        return Some(EdgeEditContext {
            id: &edge_ref.id,
            style: &mut edge_ref.style,
            ends: &edge_ref.ends,
        });
    }
}

impl<NodeIdentifier, EdgeIdentifier>
    Graph<NodeEntry<NodeIdentifier>, EdgeEntry<EdgeIdentifier, NodeIdentifier>>
    for GraphWidget<NodeIdentifier, EdgeIdentifier>
where
    NodeIdentifier: Identifier,
    EdgeIdentifier: Identifier,
{
    fn insert_edge(
        &mut self,
        edge: EdgeEntry<EdgeIdentifier, NodeIdentifier>,
    ) -> Result<
        Option<EdgeEntry<EdgeIdentifier, NodeIdentifier>>,
        InvalidEdgeEndError<NodeIdentifier>,
    > {
        let [id_1_valid, id_2_valid] = edge.ends.clone().map(|e| {
            if self.nodes.contains_key(&e) {
                None
            } else {
                Some(e)
            }
        });
        if id_1_valid.is_none() && id_2_valid.is_none() {
            return Ok(self.insert_edge_lazy(edge));
        }
        return Err(InvalidEdgeEndError([id_1_valid, id_2_valid]));
    }

    fn remove_edge(
        &mut self,
        identifier: &EdgeIdentifier,
    ) -> Option<EdgeEntry<EdgeIdentifier, NodeIdentifier>> {
        return self.edges.remove(identifier);
    }

    fn insert_node(
        &mut self,
        node: NodeEntry<NodeIdentifier>,
    ) -> Option<NodeEntry<NodeIdentifier>> {
        self.nodes_adjacency.insert(node.id.clone(), HashSet::new());
        return self.nodes.insert(node.id.clone(), node);
    }

    fn remove_node(
        &mut self,
        identifier: &NodeIdentifier,
    ) -> Result<Option<NodeEntry<NodeIdentifier>>, RemainingConnectedEdgesError<EdgeIdentifier>>
    {
        let Some(adjacency) = self.nodes_adjacency.get(identifier) else {
            return Ok(None);
        };
        if adjacency.is_empty() {
            return Ok(self.remove_node_lazy(identifier));
        }
        return Err(RemainingConnectedEdgesError(adjacency.clone()));
    }
}

impl<NodeIdentifier, EdgeIdentifier> egui::Widget
    for &mut GraphWidget<NodeIdentifier, EdgeIdentifier>
where
    NodeIdentifier: Identifier,
    EdgeIdentifier: Identifier,
{
    fn ui(self, ui: &mut Ui) -> egui::Response {
        #[cfg(feature = "log_gui")]
        trace!("Re-rendering graph canvas");
        // Reset interaction values.
        if let Some(_) = self.canvas_interaction.take() {
            #[cfg(feature = "log_gui")]
            trace!("Dumping canvas interaction");
        }
        let node_interactions = self.node_interactions.drain().count();
        if node_interactions != 0 {
            #[cfg(feature = "log_gui")]
            trace!("Dumping {} node interactions", node_interactions);
        }

        return Frame::canvas(ui.style())
            .show(ui, |ui| {
                // Allocate space and set clipping boundary
                let screen_dims = vec2(ui.available_width(), ui.available_height());
                let (_id, rect) = ui.allocate_space(screen_dims);
                ui.set_clip_rect(rect);

                // Capture canvas drag.
                let response = ui.interact(rect, Id::new("frame_dragged"), Sense::click_and_drag());
                if response.dragged() {
                    let mut canvas_interact = self.canvas_interaction.take().unwrap_or_default();
                    canvas_interact.drag = Some(response.drag_delta());
                    self.canvas_interaction = Some(canvas_interact);
                }

                // Capture canvas scroll
                let y_scroll = ui.input(|i| i.smooth_scroll_delta.y);
                if y_scroll != 0.0 {
                    let mut canvas_interact = self.canvas_interaction.take().unwrap_or_default();
                    canvas_interact.zoom = Some((1.1 as f32).powf(y_scroll));
                    self.canvas_interaction = Some(canvas_interact);
                }

                // Compile canvas mapper for this round
                let point_mapper = CanvasMapper::compile(screen_dims, &self.canvas);

                // Build graph

                // Compute nodes
                let mut node_circles: Vec<(NodeIdentifier, Shape)> = Vec::new();
                for node in self.nodes.values_mut() {
                    let center = point_mapper.map_pos(node.pos);
                    let shape = node.style.shape(center, point_mapper.pos.zoom_scale);

                    let mut node_interaction: Option<NodeInteraction> = None;

                    // Compute node interactions.
                    let rect = shape.visual_bounding_rect();
                    let response = ui.interact(rect, Id::new(&node.id), Sense::click_and_drag());
                    if response.dragged() {
                        node_interaction = Some({
                            let mut node_interaction = node_interaction.take().unwrap_or_default();
                            node_interaction.drag =
                                Some(response.drag_delta() / self.canvas.zoom_scale);
                            node_interaction
                        });
                    }

                    if let Some(node_interaction) = node_interaction {
                        self.node_interactions
                            .insert(node.id.clone(), node_interaction);
                    }

                    node_circles.push((node.id.clone(), shape));
                }
                node_circles.sort_by(|(uuid1, _), (uuid2, _)| uuid1.cmp(uuid2));

                // Compute connections
                let mut edge_lines: Vec<(EdgeIdentifier, Shape)> = Vec::new();
                let mut cleanup = Vec::new();
                for (id, edge) in self.edges.iter() {
                    let [Some((end_1, end_1_hidden)), Some((end_2, end_2_hidden))]: [Option<(
                        Pos2,
                        bool,
                    )>;
                        2] = edge.ends.clone().map(|e| {
                        self.nodes
                            .get(&e)
                            .map(|n| (point_mapper.map_pos(n.pos), n.style.hidden))
                    }) else {
                        // Put edge into cleanup list and continue to next edge.
                        cleanup.push(id.clone());
                        continue;
                    };

                    // Skip if the edge, or either of the ends, are hidden.
                    if end_1_hidden || end_2_hidden || edge.style.hidden {
                        continue;
                    }

                    let shape = edge.style.shape(end_1, end_2);

                    edge_lines.push((edge.id.clone(), shape));
                }

                // Cleanup dangling edges
                for cleanup in cleanup {
                    self.edges.remove(&cleanup);
                }

                edge_lines.sort_by(|(uuid1, _), (uuid2, _)| uuid1.cmp(uuid2));

                let painter = ui.painter();
                painter.extend(edge_lines.into_iter().map(|s| s.1));
                painter.extend(node_circles.into_iter().map(|s| s.1));
            })
            .response;
    }
}