//! [`Graph`], an undirected graph topology supporting addition and deletion.

use typed_generational_arena::{Arena, Index};

use crate::maps::EntityMap;
use crate::traits::{
    GraphTopology, GraphTopologyAddition, GraphTopologyDeletion, Incidence, IncidenceCursor,
};

/// An opaque handle identifying a vertex in a [`Graph`].
///
/// Handles remain valid until the vertex is explicitly deleted. Obtain via graph methods like
/// [`Graph::add_vertex`] and [`Graph::vertices`].
pub type Vertex = Index<VertexIncidenceHeader, usize, usize>;

/// An opaque handle identifying an edge in a [`Graph`].
///
/// Handles remain valid until the edge is explicitly deleted, or one of its endpoint vertices is
/// deleted. Obtain via graph methods like [`Graph::add_edge`] and [`Graph::edges`].
pub type Edge = Index<IncidenceEntry, usize, usize>;

/// An [`EntityMap`] for [`Graph`] vertices.
pub type GraphVertexMap<T> = EntityMap<Vertex, T>;

/// An [`EntityMap`] for [`Graph`] edges.
pub type GraphEdgeMap<T> = EntityMap<Edge, T>;

/// An [`Incidence`] for [`Graph`].
pub type GraphIncidence = Incidence<Vertex, Edge>;

#[derive(Copy, Clone, PartialEq, Eq, Debug)]
struct VertexSlot(usize);

#[derive(Copy, Clone, PartialEq, Eq, Debug)]
struct IncidenceSlot(usize);

// TODO: Check if VertexIncidenceHeader and IncidenceEntry can be made smaller. Currently they both take 24 bytes (on 64bit), see https://stackoverflow.com/a/79653173
/// `pub` because [`Vertex`] references it as a type parameter of the underlying arena. Not
/// intended for direct external use.
#[doc(hidden)]
pub struct VertexIncidenceHeader {
    incidence_count: usize,
    first_incidence: Option<IncidenceSlot>,
}

/// `pub` because [`Edge`] references it as a type parameter of the underlying arena. Not intended
/// for direct external use.
#[doc(hidden)]
#[derive(Copy, Clone)]
pub struct IncidenceEntry {
    next: Option<IncidenceSlot>,
    adjacent: VertexSlot,
}

struct IncidentEdgeCursor {
    incidence: Option<IncidenceSlot>,
}

impl IncidentEdgeCursor {
    fn next(&mut self, graph: &Graph) -> Option<Edge> {
        graph.step_incidence(&mut self.incidence).map(|(_, e)| e)
    }
}

/// A resumable cursor over the incidences of a single vertex in a [`Graph`].
///
/// Obtain via [`Graph::incidence_cursor`]. See [`IncidenceCursor`] on usage guidance.
#[derive(Copy, Clone)]
pub struct GraphIncidenceCursor {
    incidence: Option<IncidenceSlot>,
}

impl IncidenceCursor<Graph> for GraphIncidenceCursor {
    fn next(&mut self, graph: &Graph) -> Option<GraphIncidence> {
        graph
            .step_incidence(&mut self.incidence)
            .map(|(vs, e)| Incidence {
                vertex: graph.vertices.get_idx(vs.0).unwrap(),
                edge: graph.normalize_edge(e),
            })
    }
}

/// An undirected graph that supports adding vertices and edges, and deleting them.
///
/// `Graph` is suited for workloads that need to modify the graph structure over time, i.e. adding
/// amd removing vertices and edges. [`Vertex`] and [`Edge`] handles remain valid until the vertex
/// or edge they identify is explicitly deleted. Use [`AppendGraph`] instead if you only need to
/// append vertices and edges.
///
/// Incidences are stored as interleaved adjacency lists in a generational [`Arena`]. In general,
/// vertex neighborhood traversals result in scattered memory accesses.
///
/// # Examples
///
/// ```
/// use grapherity::prelude::*;
/// use grapherity::models::Graph;
///
/// let mut graph = Graph::new();
/// let v1 = graph.add_vertex();
/// let v2 = graph.add_vertex();
/// let _e = graph.add_edge(v1, v2);
/// assert!(graph.are_adjacent(v1, v2));
/// graph.delete_vertex(v1);
/// assert_eq!(graph.degree(v2), 0);
/// ```
///
/// # Time and space complexity
///
/// Both [`add_vertex`] and [`add_edge`] run in amortised *O(1)* time. [`degree`] runs in *O(1)*
/// time since vertex degrees are stored. [`delete_vertex`] runs in *O(degree(v))* time.
/// [`delete_edge`] runs in *O(degree(u) + degree(v))* time, where *u* and *v* are the edge's
/// endpoints.
///
/// Space complexity is *O(|V| + |E|)*, but freed slots for vertices and edges are not compacted and
/// their allocation is never reclaimed. Capacity only grows. If additions and deletions are both
/// planned, performing deletions first allows freed slots to be reused by subsequent additions,
/// potentially avoiding reallocation.
///
/// [`add_vertex`]: Self::add_vertex
/// [`add_edge`]: Self::add_edge
/// [`degree`]: Self::degree
/// [`delete_vertex`]: Self::delete_vertex
/// [`delete_edge`]: Self::delete_edge
/// [`AppendGraph`]: crate::models::AppendGraph
pub struct Graph {
    // TODO: Arena index and generation types could be externalized to Graph.
    vertices: Arena<VertexIncidenceHeader, usize, usize>,
    incidences: Arena<IncidenceEntry, usize, usize>,
}

impl Graph {
    /// Creates an empty graph instance with no vertices or edges.
    pub fn new() -> Self {
        Self {
            vertices: Arena::new(),
            incidences: Arena::new(),
        }
    }

    /// Adds a single incidence of an edge, which is composed by two such incidences, to the
    /// incidences arena, and returns its index.
    fn add_incidence(&mut self, v1: Vertex, v2: Vertex) -> Edge {
        let edge = self.incidences.insert(IncidenceEntry {
            next: self.vertices[v1].first_incidence.take(),
            adjacent: VertexSlot(v2.arr_idx()),
        });
        self.vertices[v1].incidence_count += 1;
        self.vertices[v1].first_incidence = Some(IncidenceSlot(edge.arr_idx()));
        edge
    }

    fn remove_incidence_pair(&mut self, e: Edge) -> (Edge, IncidenceEntry, IncidenceEntry) {
        let f = self
            .incidences
            .get_idx(e.arr_idx() ^ 1)
            .expect("invalid paired incidence index, corrupt internal data state");
        let e_entry = self
            .incidences
            .remove(e)
            .expect("attempt to delete an invalid edge");
        let f_entry = self
            .incidences
            .remove(f)
            .expect("cannot read paired incidence, corrupt internal data state");
        (f, e_entry, f_entry)
    }

    /// Updates the `source` vertex incidence list after the incidence `e` was deleted from the
    /// incidence arena. `next` is the next incidence after `e` in the `source` vertex incidence
    /// list.
    fn update_incidence_list(
        &mut self,
        e: Edge,
        source: VertexSlot,
        next: Option<IncidenceSlot>,
        is_loop: bool,
    ) {
        let source_vertex = self
            .vertices
            .get_idx(source.0)
            .expect("missing incident vertex, corrupt internal data state");
        let vertex_header = &mut self.vertices[source_vertex];
        vertex_header.incidence_count -= if is_loop { 2 } else { 1 };
        let first = vertex_header
            .first_incidence
            .expect("incident vertex without incidences, corrupt internal data state");
        if first.0 == e.arr_idx() {
            vertex_header.first_incidence = next;
        } else {
            let graph: &Graph = self;
            let (_, previous) = self
                .raw_incidences(source_vertex)
                .find(|(_, f)| {
                    graph.incidences[*f]
                        .next
                        .is_some_and(|i| i.0 == e.arr_idx())
                })
                .expect("cannot find previous incidence, corrupt internal data state");
            self.incidences[previous].next = next;
        }
    }

    fn raw_incidences(&self, v: Vertex) -> impl Iterator<Item = (VertexSlot, Edge)> {
        let mut incidence = self.vertices[v].first_incidence;
        std::iter::from_fn(move || self.step_incidence(&mut incidence))
    }

    fn step_incidence(&self, incidence: &mut Option<IncidenceSlot>) -> Option<(VertexSlot, Edge)> {
        // TODO: Benchmark storing full Index (one read, larger entries) vs. slot + get_idx() (two reads, smaller entries).
        let current = (*incidence)?;
        let e = self.incidences.get_idx(current.0).unwrap();
        let entry = self.incidences[e];
        *incidence = entry.next;
        Some((entry.adjacent, e))
    }

    fn normalize_edge(&self, e: Edge) -> Edge {
        let i = e.arr_idx();
        if i & 1 == 0 {
            e
        } else {
            self.incidences.get_idx(i ^ 1).unwrap()
        }
    }
}

impl Default for Graph {
    fn default() -> Self {
        Self::new()
    }
}

impl GraphTopology for Graph {
    type Vertex = Vertex;
    type Edge = Edge;
    type IncidenceCursor = GraphIncidenceCursor;

    fn vertex_count(&self) -> usize {
        self.vertices.len()
    }

    fn vertex_map<T: Clone>(&self, default: T) -> EntityMap<Self::Vertex, T> {
        EntityMap::new(default, |v| v.arr_idx(), self.vertex_capacity())
    }

    fn edge_count(&self) -> usize {
        self.incidences.len() / 2
    }

    fn edge_map<T: Clone>(&self, default: T) -> EntityMap<Self::Edge, T> {
        EntityMap::new(default, |e| e.arr_idx() / 2, self.edge_capacity())
    }

    fn degree(&self, v: Self::Vertex) -> usize {
        self.vertices[v].incidence_count
    }

    fn are_adjacent(&self, v1: Self::Vertex, v2: Self::Vertex) -> bool {
        self.adjacent_vertices(v1).any(|x| x == v2)
    }

    fn vertices(&self) -> impl Iterator<Item = Self::Vertex> {
        self.vertices.iter().map(|(i, _)| i)
    }

    fn adjacent_vertices(&self, v: Self::Vertex) -> impl Iterator<Item = Self::Vertex> {
        self.raw_incidences(v)
            .map(|(vs, _)| self.vertices.get_idx(vs.0).unwrap())
    }

    fn incident_vertices(&self, e: Self::Edge) -> (Self::Vertex, Self::Vertex) {
        let v2 = self
            .vertices
            .get_idx(self.incidences[e].adjacent.0)
            .unwrap();
        let f = self.incidences.get_idx(e.arr_idx() ^ 1).unwrap();
        let v1 = self
            .vertices
            .get_idx(self.incidences[f].adjacent.0)
            .unwrap();
        (v1, v2)
    }

    fn edges(&self) -> impl Iterator<Item = Self::Edge> {
        self.incidences
            .iter()
            .filter(|(i, _)| i.arr_idx() % 2 == 0)
            .map(|(i, _)| i)
    }

    fn incident_edges(&self, v: Self::Vertex) -> impl Iterator<Item = Self::Edge> {
        self.raw_incidences(v).map(|(_, e)| self.normalize_edge(e))
    }

    fn incidences(&self, v: Self::Vertex) -> impl Iterator<Item = GraphIncidence> {
        self.raw_incidences(v).map(|(vs, e)| Incidence {
            vertex: self.vertices.get_idx(vs.0).unwrap(),
            edge: self.normalize_edge(e),
        })
    }

    fn incidence_cursor(&self, v: Self::Vertex) -> Self::IncidenceCursor {
        GraphIncidenceCursor {
            incidence: self.vertices[v].first_incidence,
        }
    }
}

impl GraphTopologyAddition for Graph {
    fn vertex_capacity(&self) -> usize {
        self.vertices.capacity()
    }

    fn edge_capacity(&self) -> usize {
        self.incidences.capacity() / 2
    }

    fn add_vertex(&mut self) -> Self::Vertex {
        self.vertices.insert(VertexIncidenceHeader {
            incidence_count: 0,
            first_incidence: None,
        })
    }

    fn add_edge(&mut self, v1: Self::Vertex, v2: Self::Vertex) -> Self::Edge {
        let first = self.add_incidence(v1, v2);
        self.add_incidence(v2, v1);
        first
    }
}

// TODO: Benchmark delete with storing "previous" in O(1) vs. linear lookup in O(degree).
impl GraphTopologyDeletion for Graph {
    fn delete_vertex(&mut self, v: Self::Vertex) {
        let v_header = self
            .vertices
            .remove(v)
            .expect("attempt to delete an invalid vertex");
        let mut cursor = IncidentEdgeCursor {
            incidence: v_header.first_incidence,
        };
        while let Some(e) = cursor.next(self) {
            // Since v is being deleted, there are no update_incidence_list() calls for e, no need
            // to fix v's incidence list.
            let (f, e_entry, f_entry) = self.remove_incidence_pair(e);
            if e_entry.adjacent != f_entry.adjacent {
                self.update_incidence_list(f, e_entry.adjacent, f_entry.next, false);
            } else {
                cursor.incidence = f_entry.next;
            }
        }
    }

    fn delete_edge(&mut self, e: Self::Edge) {
        // The incidence entries are removed before patching the linked lists. This is safe because
        // `update_incidence_list` searches by raw slot index (IncidenceSlot.0) and only
        // dereferences the predecessor, never the removed entries themselves.
        let (f, e_entry, f_entry) = self.remove_incidence_pair(e);
        if e_entry.adjacent != f_entry.adjacent {
            self.update_incidence_list(e, f_entry.adjacent, e_entry.next, false);
            self.update_incidence_list(f, e_entry.adjacent, f_entry.next, false);
        } else if f_entry.next.is_some_and(|i| e.arr_idx() == i.0) {
            self.update_incidence_list(f, e_entry.adjacent, e_entry.next, true);
        } else {
            self.update_incidence_list(e, e_entry.adjacent, f_entry.next, true);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    crate::graph_topology_test_fixtures!(Graph);
    crate::graph_topology_tests!(Graph);
    crate::graph_topology_deletion_tests!(Graph);

    #[test]
    fn incident_vertices_paired_index() {
        let mut graph = Graph::new();
        let v1 = graph.add_vertex();
        let v2 = graph.add_vertex();
        let e = graph.add_edge(v1, v2);
        let f = graph
            .incidences
            .get_idx(e.arr_idx() + 1)
            .expect("paired index should be valid");
        let (u1, u2) = graph.incident_vertices(f);
        assert!(
            (u1 == v1 && u2 == v2) || (u1 == v2 && u2 == v1),
            "unexpected incident vertices {u1:?} and {u2:?} for edge {f:?}"
        );
    }

    #[test]
    fn delete_edge_paired_index() {
        let mut graph = Graph::new();
        let v1 = graph.add_vertex();
        let v2 = graph.add_vertex();
        let e = graph.add_edge(v1, v2);
        let f = graph
            .incidences
            .get_idx(e.arr_idx() + 1)
            .expect("paired index should be valid");
        graph.delete_edge(f);
        assert_eq!(graph.edge_count(), 0, "unexpected edge count after delete");
    }

    #[test]
    fn delete_edge_loop_paired_index() {
        let mut graph = Graph::new();
        let v = graph.add_vertex();
        let e = graph.add_edge(v, v);
        let f = graph
            .incidences
            .get_idx(e.arr_idx() + 1)
            .expect("paired index should be valid");
        graph.delete_edge(f);
        assert_eq!(graph.edge_count(), 0, "unexpected edge count after delete");
    }
}