//! Structural graph contraction and elimination tree construction.
//!
//! This module performs the topology-only contraction phase. It determines which shortcut
//! edges need to be added to the graph based on the node ordering.
//!
//! It features a standard parallelized `contract` method, as well as a highly optimized
//! `recontract` method. The `recontract` method allows the engine to update the hierarchy
//! when small structural changes occur (e.g., a road closure) without needing to recompute
//! the entire graph from scratch.

use std::ops::Range;

use rayon::prelude::*;

use crate::cch::{Cch, CchGraph};

#[derive(Clone)]
pub struct Topology {
    pub first_out: Vec<u32>,
    pub head: Vec<u32>,
    pub etree: Vec<Option<u32>>,
}

impl Topology {
    fn from_nodes(nodes: Vec<ContractionNode>, etree: Vec<Option<u32>>) -> Self {
        let mut first_out = Vec::with_capacity(nodes.len() + 1);
        let mut head = Vec::with_capacity(nodes.iter().map(|n| n.edges.len()).sum());
        first_out.push(0);
        for node in nodes {
            head.extend_from_slice(&node.edges);
            first_out.push(head.len() as u32);
        }
        Self { first_out, head, etree }
    }
    pub fn build_scheduler(&self) -> Vec<Vec<u32>> {
        let n = self.etree.len();
        let mut depths = vec![0; n];
        let mut max_depth = 0;
        for u in (0..n).rev() {
            if let Some(p) = self.etree[u] {
                depths[u] = depths[p as usize] + 1;
                max_depth = max_depth.max(depths[u]);
            }
        }
        let mut levels = vec![Vec::new(); max_depth + 1];
        for u in 0..n {
            levels[max_depth - depths[u]].push(u as u32);
        }
        levels
    }
    /// Finds the edge ID between two nodes.
    ///
    /// # Safety
    /// The caller must ensure that `from` and `to` are valid rank indices
    /// within the bounds of the hierarchy's topology.
    #[inline(always)]
    pub unsafe fn find_edge_id(&self, from: u32, to: u32) -> Option<u32> {
        let r = unsafe { self.get_range(from as usize) };
        self.head[r.clone()].binary_search(&to).ok().map(|idx| (r.start + idx) as u32)
    }
    /// Returns the edge range for a given node.
    ///
    /// # Safety
    /// `u` must be a valid node index within the range `0..num_nodes`.
    #[inline(always)]
    pub unsafe fn get_range(&self, u: usize) -> Range<usize> { unsafe { *self.first_out.get_unchecked(u) as usize..*self.first_out.get_unchecked(u + 1) as usize } }
    /// Returns the head node of a specific edge.
    ///
    /// # Safety
    /// `edge_id` must be a valid edge index within the range `0..num_edges`.
    #[inline(always)]
    pub unsafe fn get_head(&self, edge_id: usize) -> usize { unsafe { *self.head.get_unchecked(edge_id) as usize } }
}

#[derive(Debug, Default)]
struct ContractionNode {
    edges: Vec<u32>,
}

impl Cch {
    fn internal_contract(mut nodes: Vec<ContractionNode>, mut etree: Vec<Option<u32>>, start_rank: usize) -> Topology {
        for i in start_rank..nodes.len() {
            nodes[i].edges.sort_unstable();
            nodes[i].edges.dedup();

            if let Some(&lowest) = nodes[i].edges.first() {
                etree[i] = Some(lowest);
                let suffix = nodes[i].edges[1..].to_vec();
                nodes[lowest as usize].edges.extend_from_slice(&suffix);
            }
        }
        Topology::from_nodes(nodes, etree)
    }

    pub fn contract<G: CchGraph + Sync>(ranks: &[u32], graph: &G) -> Topology {
        let n = graph.num_nodes();

        let mut up_edges = vec![vec![]; n];
        for u in 0..n {
            let u_r = ranks[u];
            for v in graph.neighbors(u as u32) {
                let v_r = ranks[v as usize];
                if u_r < v_r {
                    up_edges[u_r as usize].push(v_r);
                } else if v_r < u_r {
                    up_edges[v_r as usize].push(u_r);
                }
            }
        }

        let nodes: Vec<ContractionNode> = up_edges.into_par_iter().map(|edges| ContractionNode { edges }).collect();
        let etree = vec![None; n];
        Self::internal_contract(nodes, etree, 0)
    }

    pub fn recontract<G: CchGraph + Sync>(&self, graph: &G, start_rank: u32, old: &Topology) -> Topology {
        let n = graph.num_nodes();
        let start_r = start_rank as usize;

        let mut new_up_edges = vec![vec![]; n - start_r];

        for u in 0..n {
            let u_r = self.ranks[u] as usize;
            for v in graph.neighbors(u as u32) {
                let v_r = self.ranks[v as usize] as usize;
                if u_r < v_r && u_r >= start_r {
                    new_up_edges[u_r - start_r].push(v_r as u32);
                } else if v_r < u_r && v_r >= start_r {
                    new_up_edges[v_r - start_r].push(u_r as u32);
                }
            }
        }
        let mut nodes: Vec<ContractionNode> = (0..start_r)
            .into_par_iter()
            .map(|r| ContractionNode {
                edges: old.head[unsafe { old.get_range(r) }].to_vec(),
            })
            .chain(new_up_edges.into_par_iter().map(|edges| ContractionNode { edges }))
            .collect();
        let etree = old.etree.clone();
        let (low, high) = nodes.split_at_mut(start_r);

        for (r, node) in low.iter().enumerate() {
            if let Some(parent) = etree[r].filter(|&p| p >= start_rank) {
                let high_idx = (parent - start_rank) as usize;
                if let Some((_, suffix)) = node.edges.split_first() {
                    high[high_idx].edges.extend_from_slice(suffix);
                }
            }
        }
        Self::internal_contract(nodes, etree, start_r)
    }
}
