use typed_generational_arena::{Arena, Index};

use crate::traits::{GraphTopology, GraphTopologyDeletion};

type Vertex = Index<VertexIncidenceHeader, usize, usize>;

type Edge = Index<IncidenceEntry, usize, usize>;

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

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

pub struct VertexIncidenceHeader {
    incidence_count: usize,
    first_incidence: Option<IncidenceSlot>,
}

pub struct IncidenceEntry {
    next: Option<IncidenceSlot>,
    neighbor: VertexSlot,
}

struct VertexNeighborIterator<'a> {
    graph: &'a Graph,
    incidence: Option<IncidenceSlot>,
}

impl<'a> Iterator for VertexNeighborIterator<'a> {
    type Item = Vertex;

    fn next(&mut self) -> Option<Self::Item> {
        // TODO: Benchmark storing full Index (one read, larger entries) vs. slot + get_idx() (two reads, smaller entries).
        let incidence = self.incidence?;
        let index = self.graph.incidences.get_idx(incidence.0)?;
        let entry = &self.graph.incidences[index];
        self.incidence = entry.next;
        self.graph.vertices.get_idx(entry.neighbor.0)
    }
}

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 {
    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(),
            neighbor: VertexSlot(v2.arr_idx()),
        });
        self.vertices[v1].incidence_count += 1;
        self.vertices[v1].first_incidence = Some(IncidenceSlot(edge.arr_idx()));
        edge
    }
}

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

impl GraphTopology for Graph {
    type Vertex = Vertex;

    type Edge = Edge;

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

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

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

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

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

    fn neighbors(&self, v: Self::Vertex) -> impl Iterator<Item = Self::Vertex> {
        VertexNeighborIterator {
            graph: self,
            incidence: self.vertices[v].first_incidence,
        }
    }

    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) {
        todo!()
    }

    fn delete_edge(&mut self, e: Self::Edge) {
        todo!()
    }
}

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

    #[test]
    fn add_vertex() {
        let mut graph = Graph::new();
        let v = graph.add_vertex();
        assert_ne!(graph.add_vertex(), v, "unexpected duplicate vertex");
    }

    #[test]
    fn vertex_count_empty() {
        let graph = Graph::new();
        assert_eq!(graph.vertex_count(), 0, "unexpected vertex count");
    }

    #[test]
    fn vertex_count() {
        let (graph, _) = make_test_graph();
        assert_eq!(graph.vertex_count(), 10, "unexpected vertex count");
    }

    #[test]
    fn add_edge() {
        let mut graph = Graph::new();
        let v1 = graph.add_vertex();
        let v2 = graph.add_vertex();
        let e = graph.add_edge(v1, v2);
        assert_ne!(graph.add_edge(v1, v2), e, "unexpected duplicate edge");
    }

    #[test]
    fn edge_count_empty() {
        let graph = Graph::new();
        assert_eq!(graph.edge_count(), 0, "unexpected edge count");
    }

    #[test]
    fn edge_count() {
        let (graph, _) = make_test_graph();
        assert_eq!(graph.edge_count(), 14, "unexpected edge count");
    }

    #[test]
    fn degree_zero() {
        let mut graph = Graph::new();
        let v = graph.add_vertex();
        assert_eq!(graph.degree(v), 0, "unexpected non-zero degree");
    }

    #[test]
    fn degree() {
        let (graph, vertices) = make_test_graph();
        let expected_degrees = [1, 4, 4, 2, 4, 2, 3, 3, 2, 3];
        for i in 0..graph.vertex_count() {
            assert_eq!(
                graph.degree(vertices[i]),
                expected_degrees[i],
                "unexpected degree of {:?}",
                vertices[i]
            );
        }
    }

    #[test]
    fn are_adjacent_vertex_self() {
        let mut graph = Graph::new();
        let v = graph.add_vertex();
        assert!(
            !graph.are_adjacent(v, v),
            "should not be adjacent to itself"
        );
    }

    #[test]
    fn are_adjacent_single_edge() {
        let mut graph = Graph::new();
        let v1 = graph.add_vertex();
        let v2 = graph.add_vertex();
        assert!(!graph.are_adjacent(v1, v2), "should not be adjacent");
        assert!(!graph.are_adjacent(v2, v1), "should not be adjacent");
        graph.add_edge(v1, v2);
        assert!(graph.are_adjacent(v1, v2), "should be adjacent");
        assert!(graph.are_adjacent(v2, v1), "should be adjacent");
    }

    #[test]
    fn are_adjacent() {
        let (graph, vertices) = make_test_graph();
        assert!(
            graph.are_adjacent(vertices[0], vertices[1]),
            "expected {:?} and {:?} to be adjacent",
            vertices[0],
            vertices[1]
        );
        assert!(
            graph.are_adjacent(vertices[9], vertices[5]),
            "expected {:?} and {:?} to be adjacent",
            vertices[9],
            vertices[5]
        );
        assert!(
            !graph.are_adjacent(vertices[9], vertices[3]),
            "unexpected adjacency of {:?} and {:?}",
            vertices[9],
            vertices[3]
        );

        for i in 0..graph.vertex_count() {
            let exp = match i {
                2 => continue,
                1 | 4 | 5 | 6 => true,
                _ => false,
            };
            assert_eq!(
                graph.are_adjacent(vertices[2], vertices[i]),
                exp,
                "unexpected adjacency of {:?} and {:?}",
                vertices[2],
                vertices[i]
            );
            assert_eq!(
                graph.are_adjacent(vertices[i], vertices[2]),
                exp,
                "unexpected adjacency of {:?} and {:?}",
                vertices[i],
                vertices[2]
            );
        }
    }

    #[test]
    fn vertices_empty() {
        let graph = Graph::new();
        assert_eq!(
            graph.vertices().count(),
            0,
            "vertex iterator of empty graph should have no elements"
        );
    }

    #[test]
    fn vertices() {
        let (graph, vertices) = make_test_graph();
        assert_eq!(graph.vertices().count(), 10, "unexpected vertex count");

        // Expects each vertex to appear exactly once.
        for v in graph.vertices() {
            assert_eq!(
                vertices.iter().filter(|&x| *x == v).count(),
                1,
                "unexpected vertex {v:?} from the iterator"
            );
        }
    }

    #[test]
    fn neighbors_empty() {
        let mut graph = Graph::new();
        let vertex = graph.add_vertex();
        assert_eq!(
            graph.neighbors(vertex).count(),
            0,
            "neighbor iterator of vertex with degree 0 should have no elements"
        );
    }

    #[test]
    fn neighbors() {
        let (graph, vertices) = make_test_graph();
        // Checks neighbors of vertex 4.
        assert_eq!(
            graph.neighbors(vertices[4]).count(),
            4,
            "unexpected neighbor count"
        );

        // Expects each neighbor to appear exactly once. This will not work if there are multiple edges.
        let neighbors = vec![vertices[1], vertices[2], vertices[7], vertices[8]];
        for v in graph.neighbors(vertices[4]) {
            assert_eq!(
                neighbors.iter().filter(|&x| *x == v).count(),
                1,
                "unexpected neighbor {v:?} of {:?} from the iterator",
                vertices[4]
            );
        }
    }

    #[test]
    fn loop_edge() {
        let mut graph = Graph::new();
        let v = graph.add_vertex();
        graph.add_edge(v, v);
        assert_eq!(graph.vertex_count(), 1, "unexpected vertex count");
        assert_eq!(graph.edge_count(), 1, "unexpected edge count");
        assert_eq!(graph.degree(v), 2, "unexpected degree");
        assert!(
            graph.are_adjacent(v, v),
            "vertex with loop edge should be adjacent to itself"
        );
        let mut neighbors = graph.neighbors(v);
        assert_eq!(
            neighbors.next(),
            Some(v),
            "vertex should be neighbor of itself"
        );
        assert_eq!(
            neighbors.next(),
            Some(v),
            "vertex should be neighbor of itself twice"
        );
        assert_eq!(neighbors.next(), None, "too many neighbors from iterator");
    }

    #[test]
    fn multiple_edges() {
        let k = 3;
        let mut graph = Graph::new();
        let vertices = [graph.add_vertex(), graph.add_vertex()];
        for _ in 0..k {
            graph.add_edge(vertices[0], vertices[1]);
        }
        assert_eq!(
            graph.vertex_count(),
            vertices.len(),
            "unexpected vertex count"
        );
        assert_eq!(graph.edge_count(), k, "unexpected edge count");
        for v in vertices {
            assert_eq!(graph.degree(v), k, "unexpected degree of {v:?}");
        }
        assert!(
            graph.are_adjacent(vertices[0], vertices[1]),
            "should be adjacent"
        );
        for i in 0..2 {
            let mut neighbors = graph.neighbors(vertices[i]);
            for j in 0..k {
                assert_eq!(
                    neighbors.next(),
                    Some(vertices[1 - i]),
                    "neighbor {j} of vertex {:?} should be {:?}",
                    vertices[i],
                    vertices[1 - i]
                );
            }
            assert_eq!(
                neighbors.next(),
                None,
                "too many neighbors of {:?} from iterator",
                vertices[i]
            );
        }
    }

    #[test]
    fn delete_vertex() {
        let mut graph = Graph::new();
        let v = graph.add_vertex();
        assert_eq!(
            graph.vertex_count(),
            1,
            "unexpected vertex count before delete"
        );
        graph.delete_vertex(v);
        assert_eq!(
            graph.vertex_count(),
            0,
            "unexpected vertex count after delete"
        );
        assert_ne!(
            graph.add_vertex(),
            v,
            "unexpected duplicate vertex after delete"
        );
    }

    #[test]
    fn delete_vertex_invalid_index() {
        let mut graph = Graph::new();
        let v = graph.add_vertex();
        graph.delete_vertex(v);
        let result = std::panic::catch_unwind(move || graph.delete_vertex(v));
        assert!(result.is_err(), "second deletion should panic");
    }

    #[test]
    fn delete_vertex_connected() {
        let (mut graph, vertices) = make_test_graph();
        assert_eq!(
            graph.vertex_count(),
            10,
            "unexpected vertex count before delete"
        );
        assert_eq!(
            graph.edge_count(),
            14,
            "unexpected edge count before delete"
        );
        graph.delete_vertex(vertices[2]);
        assert_eq!(
            graph.vertex_count(),
            9,
            "unexpected vertex count after delete"
        );
        assert_eq!(graph.edge_count(), 10, "unexpected edge count after delete");
    }

    #[test]
    fn delete_edge() {
        let mut graph = Graph::new();
        let v1 = graph.add_vertex();
        let v2 = graph.add_vertex();
        let e = graph.add_edge(v1, v2);
        assert_eq!(
            graph.vertex_count(),
            2,
            "unexpected vertex count before delete"
        );
        assert_eq!(graph.edge_count(), 1, "unexpected edge count before delete");
        graph.delete_edge(e);
        assert_eq!(
            graph.vertex_count(),
            2,
            "unexpected vertex count after delete"
        );
        assert_eq!(graph.edge_count(), 0, "unexpected edge count after delete");
        assert_ne!(
            graph.add_edge(v1, v2),
            e,
            "unexpected duplicate edge after delete"
        );
    }

    #[test]
    fn delete_edge_invalid_index() {
        let mut graph = Graph::new();
        let v1 = graph.add_vertex();
        let v2 = graph.add_vertex();
        let e = graph.add_edge(v1, v2);
        graph.delete_edge(e);
        let result = std::panic::catch_unwind(move || graph.delete_edge(e));
        assert!(result.is_err(), "second deletion should panic");
    }

    fn make_test_graph() -> (Graph, [Vertex; 10]) {
        let mut graph = Graph::new();
        let vertices: [Vertex; 10] = core::array::from_fn(|_| graph.add_vertex());
        graph.add_edge(vertices[0], vertices[1]);
        graph.add_edge(vertices[1], vertices[2]);
        graph.add_edge(vertices[1], vertices[3]);
        graph.add_edge(vertices[1], vertices[4]);
        graph.add_edge(vertices[2], vertices[4]);
        graph.add_edge(vertices[2], vertices[5]);
        graph.add_edge(vertices[2], vertices[6]);
        graph.add_edge(vertices[3], vertices[6]);
        graph.add_edge(vertices[4], vertices[7]);
        graph.add_edge(vertices[4], vertices[8]);
        graph.add_edge(vertices[5], vertices[9]);
        graph.add_edge(vertices[6], vertices[9]);
        graph.add_edge(vertices[7], vertices[8]);
        graph.add_edge(vertices[7], vertices[9]);
        (graph, vertices)
    }
}