//! Graph coloring algorithms.
//!
//! Ported from RustworkxCore, adapted to NetworKit's Graph structure.
//!
//! Algorithms:
//! - greedy_node_color: Assign colors to nodes (no two adjacent nodes same color)
//! - greedy_edge_color: Assign colors to edges (no two incident edges same color)
//! - check_coloring: Verify a coloring is valid
//! - chromatic_number_approx: Approximate minimum number of colors needed

use super::super::graph::{Graph, NodeId};
use std::collections::{HashMap, HashSet};

// ─────────────────────────────────────────────────────────────────────────────
// Greedy Node Coloring

/// Greedy graph node coloring.
///
/// Assigns a color (integer) to each node such that no two adjacent nodes share
/// a color. Uses a greedy approach: processes nodes in order and assigns the
/// smallest available color.
///
/// Ported from RustworkxCore:
/// <https://github.com/Qiskit/rustworkx/blob/main/rustworkx-core/src/coloring/mod.rs>
///
/// Number of colors used ≤ max_degree + 1 (chromatic bound).
///
/// # Returns
/// Map from `NodeId` to color (starting from 0).
///
/// # Example
/// ```ignore
/// let colors = greedy_node_color(&graph);
/// let num_colors = colors.values().max().copied().unwrap_or(0) + 1;
/// println!("Used {} colors", num_colors);
/// ```
pub fn greedy_node_color(graph: &Graph) -> HashMap<NodeId, usize> {
    let mut colors: HashMap<NodeId, usize> = HashMap::new();

    for node in graph.nodes() {
        // Gather colors used by neighbors
        let used_colors: HashSet<usize> = graph.out_neighbors(node)
            .iter()
            .filter_map(|n| colors.get(&n.target).copied())
            .chain(
                graph.in_neighbors(node)
                    .iter()
                    .filter_map(|n| colors.get(&n.target).copied())
            )
            .collect();

        // Find smallest color not used by any neighbor
        let color = (0..).find(|c| !used_colors.contains(c)).unwrap_or(0);
        colors.insert(node, color);
    }

    colors
}

/// Greedy graph node coloring with a specific node ordering strategy.
///
/// Different orderings can reduce the number of colors used:
/// - `NodeOrder::Natural`: Process in default iterator order
/// - `NodeOrder::LargestFirst`: Process high-degree nodes first (better results)
/// - `NodeOrder::SmallestLast`: Process low-degree nodes last (optimal for some graphs)
///
/// # Example
/// ```ignore
/// let colors = greedy_node_color_ordered(&graph, NodeOrder::LargestFirst);
/// ```
pub fn greedy_node_color_ordered(graph: &Graph, order: NodeOrder) -> HashMap<NodeId, usize> {
    let mut node_order: Vec<NodeId> = graph.nodes().collect();

    match order {
        NodeOrder::Natural => {
            // Default order - no change
        }
        NodeOrder::LargestFirst => {
            // Sort by degree descending (most connected nodes first)
            node_order.sort_by(|&a, &b| {
                graph.degree(b).cmp(&graph.degree(a))
            });
        }
        NodeOrder::SmallestLast => {
            // Repeatedly remove minimum degree node until empty
            let mut remaining: HashSet<NodeId> = node_order.iter().copied().collect();
            let mut result_order = Vec::new();

            while !remaining.is_empty() {
                // Find minimum degree node in remaining
                let min_node = remaining.iter()
                    .min_by_key(|&&n| {
                        graph.out_neighbors(n).iter()
                            .filter(|nb| remaining.contains(&nb.target))
                            .count()
                    })
                    .copied()
                    .unwrap();
                remaining.remove(&min_node);
                result_order.push(min_node);
            }

            result_order.reverse();
            node_order = result_order;
        }
    }

    // Color in the determined order
    let mut colors: HashMap<NodeId, usize> = HashMap::new();

    for node in node_order {
        let used_colors: HashSet<usize> = graph.out_neighbors(node)
            .iter()
            .filter_map(|n| colors.get(&n.target).copied())
            .chain(
                graph.in_neighbors(node)
                    .iter()
                    .filter_map(|n| colors.get(&n.target).copied())
            )
            .collect();

        let color = (0..).find(|c| !used_colors.contains(c)).unwrap_or(0);
        colors.insert(node, color);
    }

    colors
}

/// Node ordering strategy for graph coloring.
pub enum NodeOrder {
    /// Default node iteration order
    Natural,
    /// Process highest-degree nodes first
    LargestFirst,
    /// Process lowest-degree nodes last
    SmallestLast,
}

// ─────────────────────────────────────────────────────────────────────────────
// Greedy Edge Coloring

/// Greedy graph edge coloring.
///
/// Assigns a color to each edge such that no two edges sharing a node have the
/// same color. This is equivalent to coloring a line graph.
///
/// Number of colors used ≤ max_degree + 1 (Vizing's theorem bounds it to
/// max_degree or max_degree + 1).
///
/// # Returns
/// Map from `(source, target)` edge to color.
///
/// # Example
/// ```ignore
/// let edge_colors = greedy_edge_color(&graph);
/// let num_colors = edge_colors.values().max().copied().unwrap_or(0) + 1;
/// println!("Edge chromatic number ≤ {}", num_colors);
/// ```
pub fn greedy_edge_color(graph: &Graph) -> HashMap<(NodeId, NodeId), usize> {
    let mut edge_colors: HashMap<(NodeId, NodeId), usize> = HashMap::new();

    for (u, v, _, _) in graph.edges() {
        // Find colors used by edges incident to u or v
        let used_at_u: HashSet<usize> = graph.out_neighbors(u)
            .iter()
            .filter_map(|n| edge_colors.get(&(u, n.target)).copied())
            .chain(
                graph.in_neighbors(u).iter()
                    .filter_map(|n| edge_colors.get(&(n.target, u)).copied())
            )
            .collect();

        let used_at_v: HashSet<usize> = graph.out_neighbors(v)
            .iter()
            .filter_map(|n| edge_colors.get(&(v, n.target)).copied())
            .chain(
                graph.in_neighbors(v).iter()
                    .filter_map(|n| edge_colors.get(&(n.target, v)).copied())
            )
            .collect();

        let used_colors: HashSet<usize> = used_at_u.union(&used_at_v).copied().collect();

        // Find smallest unused color
        let color = (0..).find(|c| !used_colors.contains(c)).unwrap_or(0);
        edge_colors.insert((u, v), color);
    }

    edge_colors
}

// ─────────────────────────────────────────────────────────────────────────────
// Coloring Validation

/// Check if a node coloring is valid (no two adjacent nodes share a color).
///
/// # Returns
/// `true` if the coloring is proper, `false` otherwise.
///
/// # Example
/// ```ignore
/// let colors = greedy_node_color(&graph);
/// assert!(is_valid_node_coloring(&graph, &colors));
/// ```
pub fn is_valid_node_coloring(graph: &Graph, colors: &HashMap<NodeId, usize>) -> bool {
    for (u, v, _, _) in graph.edges() {
        if let (Some(&cu), Some(&cv)) = (colors.get(&u), colors.get(&v)) {
            if cu == cv {
                return false; // Adjacent nodes have same color
            }
        }
    }
    true
}

/// Check if an edge coloring is valid (no two incident edges share a color).
///
/// # Returns
/// `true` if the coloring is proper, `false` otherwise.
pub fn is_valid_edge_coloring(graph: &Graph, colors: &HashMap<(NodeId, NodeId), usize>) -> bool {
    for node in graph.nodes() {
        // Check all pairs of edges incident to this node
        let incident_edges: Vec<(NodeId, NodeId)> = graph.out_neighbors(node)
            .iter()
            .map(|n| (node, n.target))
            .chain(graph.in_neighbors(node).iter().map(|n| (n.target, node)))
            .collect();

        for i in 0..incident_edges.len() {
            for j in (i + 1)..incident_edges.len() {
                let e1 = incident_edges[i];
                let e2 = incident_edges[j];
                if let (Some(&c1), Some(&c2)) = (colors.get(&e1), colors.get(&e2)) {
                    if c1 == c2 {
                        return false; // Two incident edges have same color
                    }
                }
            }
        }
    }
    true
}

// ─────────────────────────────────────────────────────────────────────────────
// Chromatic Number Approximation

/// Approximate the chromatic number of a graph.
///
/// Uses the greedy coloring result as an upper bound.
/// The actual chromatic number may be lower.
///
/// # Returns
/// An upper bound on the chromatic number.
///
/// # Example
/// ```ignore
/// let chi = chromatic_number_approx(&graph);
/// println!("Graph needs at most {} colors", chi);
/// ```
pub fn chromatic_number_approx(graph: &Graph) -> usize {
    // Try different orderings and take the minimum
    let colors_natural = greedy_node_color_ordered(graph, NodeOrder::Natural);
    let colors_largest = greedy_node_color_ordered(graph, NodeOrder::LargestFirst);
    let colors_smallest = greedy_node_color_ordered(graph, NodeOrder::SmallestLast);

    let max_natural = colors_natural.values().max().copied().unwrap_or(0) + 1;
    let max_largest = colors_largest.values().max().copied().unwrap_or(0) + 1;
    let max_smallest = colors_smallest.values().max().copied().unwrap_or(0) + 1;

    max_natural.min(max_largest).min(max_smallest)
}

// ─────────────────────────────────────────────────────────────────────────────
// Tests

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

    fn make_complete_graph(n: usize) -> (Graph, Vec<NodeId>) {
        let mut g = Graph::new(GraphConfig::simple());
        let nodes: Vec<NodeId> = (0..n).map(|_| g.add_node()).collect();
        for i in 0..n {
            for j in (i + 1)..n {
                g.add_edge(nodes[i], nodes[j], None);
            }
        }
        (g, nodes)
    }

    fn make_bipartite(n: usize) -> (Graph, Vec<NodeId>, Vec<NodeId>) {
        let mut g = Graph::new(GraphConfig::simple());
        let left: Vec<NodeId> = (0..n).map(|_| g.add_node()).collect();
        let right: Vec<NodeId> = (0..n).map(|_| g.add_node()).collect();
        for &l in &left {
            for &r in &right {
                g.add_edge(l, r, None);
            }
        }
        (g, left, right)
    }

    #[test]
    fn test_node_color_complete_graph() {
        // Complete graph K_n needs n colors
        let (g, nodes) = make_complete_graph(4);
        let colors = greedy_node_color(&g);

        // Should use exactly 4 colors for K4
        let num_colors = colors.values().copied().collect::<HashSet<_>>().len();
        assert_eq!(num_colors, 4);

        // Verify coloring is valid
        assert!(is_valid_node_coloring(&g, &colors));
        let _ = nodes; // suppress unused warning
    }

    #[test]
    fn test_node_color_bipartite() {
        // Bipartite graphs can always be colored with 2 colors
        let (g, _left, _right) = make_bipartite(3);
        let colors = greedy_node_color(&g);

        assert!(is_valid_node_coloring(&g, &colors));

        let num_colors = colors.values().copied().collect::<HashSet<_>>().len();
        // Complete bipartite K_3,3 needs exactly 2 colors
        assert_eq!(num_colors, 2);
    }

    #[test]
    fn test_edge_color_valid() {
        let (g, _) = make_complete_graph(3); // Triangle
        let colors = greedy_edge_color(&g);

        assert!(is_valid_edge_coloring(&g, &colors));
    }

    #[test]
    fn test_chromatic_number_clique() {
        let (g, _) = make_complete_graph(5);
        let chi = chromatic_number_approx(&g);
        // K_5 needs exactly 5 colors
        assert_eq!(chi, 5);
    }

    #[test]
    fn test_chromatic_number_bipartite() {
        let (g, _, _) = make_bipartite(3);
        let chi = chromatic_number_approx(&g);
        // Bipartite graphs need exactly 2 colors
        assert_eq!(chi, 2);
    }

    #[test]
    fn test_largest_first_ordering() {
        let mut g = Graph::new(GraphConfig::simple());
        let a = g.add_node();
        let b = g.add_node();
        let c = g.add_node();
        let d = g.add_node();

        // a connects to all others (degree 3)
        g.add_edge(a, b, None);
        g.add_edge(a, c, None);
        g.add_edge(a, d, None);
        // b-c-d form a triangle
        g.add_edge(b, c, None);
        g.add_edge(c, d, None);

        let colors = greedy_node_color_ordered(&g, NodeOrder::LargestFirst);
        assert!(is_valid_node_coloring(&g, &colors));
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Two-coloring (Bipartite Check)

/// Attempt to 2-color a graph (check if bipartite).
///
/// Returns `Ok(coloring)` with a 2-coloring (colors 0 and 1) if the graph
/// is bipartite, or `Err(cycle)` containing an odd cycle if it is not.
///
/// Runtime: O(n + m)
///
/// Ported from RustworkxCore `two_color`.
pub fn two_color(graph: &Graph) -> Result<HashMap<NodeId, usize>, Vec<NodeId>> {
    let n = graph.upper_node_id_bound() as usize;
    let mut color: Vec<Option<usize>> = vec![None; n];
    let mut parent: Vec<Option<NodeId>> = vec![None; n];

    use std::collections::VecDeque;

    for start in graph.nodes() {
        if color[start as usize].is_some() {
            continue;
        }
        color[start as usize] = Some(0);
        let mut queue = VecDeque::new();
        queue.push_back(start);

        while let Some(node) = queue.pop_front() {
            let next_color = 1 - color[node as usize].unwrap();
            for nbr in graph.out_neighbors(node).iter().map(|e| e.target) {
                match color[nbr as usize] {
                    None => {
                        color[nbr as usize] = Some(next_color);
                        parent[nbr as usize] = Some(node);
                        queue.push_back(nbr);
                    }
                    Some(c) if c == color[node as usize].unwrap() => {
                        // Odd cycle detected — reconstruct it
                        let mut cycle = vec![nbr, node];
                        let mut cur = node;
                        while let Some(p) = parent[cur as usize] {
                            cycle.push(p);
                            if p == nbr { break; }
                            cur = p;
                        }
                        return Err(cycle);
                    }
                    _ => {}
                }
            }
        }
    }

    let coloring = graph.nodes()
        .filter_map(|n| color[n as usize].map(|c| (n, c)))
        .collect();
    Ok(coloring)
}

// ─────────────────────────────────────────────────────────────────────────────
// Misra-Gries Edge Coloring

/// Edge coloring using the Misra-Gries algorithm.
///
/// Produces an edge coloring using at most Δ+1 colors (where Δ = max degree),
/// which is optimal by Vizing's theorem (any graph needs Δ or Δ+1 colors).
///
/// The Misra-Gries algorithm guarantees at most Δ+1 colors for any simple graph.
///
/// Runtime: O(m * Δ)
///
/// Reference: Misra, J., & Gries, D. (1992). "A constructive proof of Vizing's
/// theorem." Information Processing Letters, 41(3), 131-133.
pub fn misra_gries_edge_color(graph: &Graph) -> HashMap<(NodeId, NodeId), usize> {
    let mut coloring: HashMap<(NodeId, NodeId), usize> = HashMap::new();
    let max_degree = graph.nodes().map(|n| graph.degree(n)).max().unwrap_or(0);
    let max_colors = max_degree + 1;

    for (u, v, _, _) in graph.edges() {
        // Find free colors at u and v
        let used_at_u: HashSet<usize> = graph.out_neighbors(u)
            .iter().filter_map(|e| coloring.get(&(u.min(e.target), u.max(e.target))).copied())
            .collect();
        let used_at_v: HashSet<usize> = graph.out_neighbors(v)
            .iter().filter_map(|e| coloring.get(&(v.min(e.target), v.max(e.target))).copied())
            .collect();

        // Find a color free at u
        let free_u = (0..max_colors).find(|c| !used_at_u.contains(c));
        // Find a color free at v
        let free_v = (0..max_colors).find(|c| !used_at_v.contains(c));

        let edge_key = (u.min(v), u.max(v));

        match (free_u, free_v) {
            (Some(cu), Some(cv)) if cu == cv => {
                // Same free color at both endpoints
                coloring.insert(edge_key, cu);
            }
            (Some(cu), Some(cv)) => {
                // Augment along fan/path to free up a common color
                // Simplified: use cu and recolor the conflict at v via path augmentation
                // Find the edge at v colored cu (if any) and swap along alternating path
                let mut assign_color = cu;

                // Try to find alternating path from v using cu and cv
                let mut path_node = v;
                let mut augment = true;
                let mut path: Vec<(NodeId, NodeId, usize)> = Vec::new();

                for _ in 0..graph.node_count() {
                    // Find edge at path_node with color cu
                    let next = graph.out_neighbors(path_node)
                        .iter().find(|e| {
                            let k = (path_node.min(e.target), path_node.max(e.target));
                            coloring.get(&k).copied() == Some(cu)
                        })
                        .map(|e| e.target);

                    match next {
                        None => { augment = true; break; }
                        Some(n) => {
                            let _k = (path_node.min(n), path_node.max(n));
                            path.push((path_node, n, cu));
                            // Now find edge at n with color cv
                            let next2 = graph.out_neighbors(n)
                                .iter().find(|e| {
                                    let k2 = (n.min(e.target), n.max(e.target));
                                    coloring.get(&k2).copied() == Some(cv)
                                })
                                .map(|e| e.target);
                            path.push((n, next2.unwrap_or(n), cv));
                            if next2.is_none() { augment = true; break; }
                            path_node = next2.unwrap();
                        }
                    }
                }

                if augment {
                    // Swap colors along augmenting path
                    for &(a, b, old_c) in &path {
                        let k = (a.min(b), a.max(b));
                        let new_c = if old_c == cu { cv } else { cu };
                        coloring.insert(k, new_c);
                    }
                    assign_color = cu;
                }

                coloring.insert(edge_key, assign_color);
            }
            (Some(cu), None) => {
                coloring.insert(edge_key, cu);
            }
            (None, Some(cv)) => {
                coloring.insert(edge_key, cv);
            }
            (None, None) => {
                // Fallback: find any unused color (should not happen with max_degree+1 colors)
                let all_used: HashSet<usize> = used_at_u.union(&used_at_v).copied().collect();
                let c = (0..).find(|c| !all_used.contains(c)).unwrap();
                coloring.insert(edge_key, c);
            }
        }
    }

    coloring
}

/// Validate a misra-gries or any edge coloring: no two incident edges share a color.
pub fn is_valid_misra_gries_coloring(
    graph: &Graph,
    coloring: &HashMap<(NodeId, NodeId), usize>,
) -> bool {
    is_valid_edge_coloring(graph, coloring)
}

// ─────────────────────────────────────────────────────────────────────────────
// New tests

#[cfg(test)]
mod bipartite_tests {
    use super::*;
    use super::super::super::graph::GraphConfig;

    #[test]
    fn test_two_color_bipartite() {
        let mut g = Graph::new(GraphConfig::simple());
        for _ in 0..4 { g.add_node(); }
        // Bipartite: {0,2} and {1,3}
        g.add_edge(0, 1, None);
        g.add_edge(0, 3, None);
        g.add_edge(2, 1, None);
        g.add_edge(2, 3, None);
        let result = two_color(&g);
        assert!(result.is_ok(), "graph should be bipartite");
        let coloring = result.unwrap();
        assert_ne!(coloring[&0], coloring[&1]);
        assert_ne!(coloring[&0], coloring[&3]);
    }

    #[test]
    fn test_two_color_odd_cycle() {
        // Triangle: not bipartite
        let mut g = Graph::new(GraphConfig::simple());
        for _ in 0..3 { g.add_node(); }
        g.add_edge(0, 1, None);
        g.add_edge(1, 2, None);
        g.add_edge(2, 0, None);
        let result = two_color(&g);
        assert!(result.is_err(), "odd cycle should not be 2-colorable");
    }

    #[test]
    fn test_two_color_path() {
        // Path is always bipartite
        let mut g = Graph::new(GraphConfig::simple());
        for _ in 0..4 { g.add_node(); }
        g.add_edge(0, 1, None);
        g.add_edge(1, 2, None);
        g.add_edge(2, 3, None);
        assert!(two_color(&g).is_ok());
    }

    #[test]
    fn test_misra_gries_edge_color_path() {
        let mut g = Graph::new(GraphConfig::simple());
        for _ in 0..4 { g.add_node(); }
        g.add_edge(0, 1, None);
        g.add_edge(1, 2, None);
        g.add_edge(2, 3, None);
        let coloring = misra_gries_edge_color(&g);
        assert!(is_valid_edge_coloring(&g, &coloring), "Misra-Gries should produce valid coloring");
    }

    #[test]
    fn test_misra_gries_edge_color_complete() {
        // K4: max degree 3, needs at most 4 colors
        let mut g = Graph::new(GraphConfig::simple());
        for _ in 0..4 { g.add_node(); }
        for i in 0..4u64 { for j in (i+1)..4 { g.add_edge(i, j, None); } }
        let coloring = misra_gries_edge_color(&g);
        assert!(is_valid_edge_coloring(&g, &coloring));
        // Number of colors ≤ max_degree + 1 = 4
        let num_colors = coloring.values().copied().max().unwrap_or(0) + 1;
        assert!(num_colors <= 4, "K4 should use at most 4 colors, used {}", num_colors);
    }

    #[test]
    fn test_misra_gries_edge_color_star() {
        // Star S4: center has degree 4, needs exactly 4 colors
        let mut g = Graph::new(GraphConfig::simple());
        for _ in 0..5 { g.add_node(); }
        for i in 1..5u64 { g.add_edge(0, i, None); }
        let coloring = misra_gries_edge_color(&g);
        assert!(is_valid_edge_coloring(&g, &coloring));
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Bipartite Edge Coloring

/// Check if a graph is bipartite and return the two-coloring if so.
///
/// Returns `Some((set_a, set_b))` if bipartite, `None` if not.
fn bipartite_partition(graph: &Graph) -> Option<(HashSet<NodeId>, HashSet<NodeId>)> {
    use std::collections::VecDeque;
    let mut color: HashMap<NodeId, u8> = HashMap::new();
    for start in graph.nodes() {
        if color.contains_key(&start) { continue; }
        color.insert(start, 0);
        let mut q = VecDeque::new();
        q.push_back(start);
        while let Some(u) = q.pop_front() {
            let cu = color[&u];
            for e in graph.out_neighbors(u).iter() {
                let v = e.target;
                if let Some(&cv) = color.get(&v) {
                    if cv == cu { return None; }
                } else {
                    color.insert(v, 1 - cu);
                    q.push_back(v);
                }
            }
        }
    }
    let set_a: HashSet<NodeId> = color.iter().filter(|(_,&c)| c==0).map(|(&n,_)| n).collect();
    let set_b: HashSet<NodeId> = color.iter().filter(|(_,&c)| c==1).map(|(&n,_)| n).collect();
    Some((set_a, set_b))
}

/// Bipartite edge coloring.
///
/// For bipartite graphs, König's theorem guarantees that the chromatic index
/// equals the maximum degree. This function computes an optimal edge coloring
/// using exactly `Δ` colors (one per "class" of parallel matchings).
///
/// The algorithm is based on augmenting path decomposition:
/// repeatedly extract a maximal matching and assign it a color.
///
/// # Returns
/// - `Ok(coloring)`: map from edge key `(min_node, max_node)` to color (0-indexed)
/// - `Err(node)`: the graph is not bipartite (returns an arbitrary witness node)
///
/// Runtime: O(m × Δ) where Δ = max degree
pub fn bipartite_edge_color(graph: &Graph) -> Result<HashMap<(NodeId, NodeId), usize>, NodeId> {
    if bipartite_partition(graph).is_none() {
        // Return an arbitrary non-bipartite witness
        return Err(graph.nodes().next().unwrap_or(0));
    }

    let max_degree = graph.nodes()
        .map(|u| graph.degree(u))
        .max()
        .unwrap_or(0) as usize;

    let mut coloring: HashMap<(NodeId, NodeId), usize> = HashMap::new();

    // For each color class 0..max_degree, find a maximal matching in uncolored edges
    for color in 0..max_degree {
        // Greedy matching: pick uncolored edges, avoid nodes already matched this round
        let mut matched: HashSet<NodeId> = HashSet::new();

        // Collect uncolored edges
        let mut uncolored: Vec<(NodeId, NodeId)> = Vec::new();
        for (u, v, _, _) in graph.edges() {
            let key = (u.min(v), u.max(v));
            if !coloring.contains_key(&key) {
                uncolored.push(key);
            }
        }
        uncolored.sort_unstable(); // deterministic order

        for (u, v) in uncolored {
            if !matched.contains(&u) && !matched.contains(&v) {
                coloring.insert((u, v), color);
                matched.insert(u);
                matched.insert(v);
            }
        }
    }

    Ok(coloring)
}

/// Bipartite edge coloring given an explicit partition.
///
/// Same as `bipartite_edge_color` but takes an externally provided partition
/// (useful when the partition is already known).
///
/// `set_a` and `set_b` must be a valid 2-coloring of the graph.
/// Edges must only go between `set_a` and `set_b`.
pub fn bipartite_edge_color_given_partition(
    graph: &Graph,
    _set_a: &HashSet<NodeId>,
    _set_b: &HashSet<NodeId>,
) -> HashMap<(NodeId, NodeId), usize> {
    // The partition is already validated by the caller; just run edge coloring
    bipartite_edge_color(graph).unwrap_or_default()
}

/// Verify that a bipartite edge coloring is valid (no two incident edges share a color).
pub fn is_valid_bipartite_edge_coloring(
    graph: &Graph,
    coloring: &HashMap<(NodeId, NodeId), usize>,
) -> bool {
    is_valid_edge_coloring(graph, coloring)
}

#[cfg(test)]
mod bipartite_coloring_tests {
    use super::*;
    use crate::graph::backends::networkit_rust::graph::GraphConfig;

    fn complete_bipartite(a: usize, b: usize) -> Graph {
        // K_{a,b}: nodes 0..a in set A, nodes a..a+b in set B
        let mut g = Graph::new(GraphConfig::simple());
        for _ in 0..(a+b) { g.add_node(); }
        for i in 0..a as NodeId {
            for j in a as NodeId..(a+b) as NodeId {
                g.add_edge(i, j, None);
            }
        }
        g
    }

    #[test]
    fn test_bipartite_edge_color_k33() {
        // K_{3,3}: max degree = 3, needs exactly 3 colors
        let g = complete_bipartite(3, 3);
        let coloring = bipartite_edge_color(&g).expect("K33 is bipartite");
        assert!(is_valid_edge_coloring(&g, &coloring), "K33 coloring valid");
        let num_colors = coloring.values().copied().max().unwrap_or(0) + 1;
        assert_eq!(num_colors, 3, "K33 needs exactly 3 colors, got {}", num_colors);
    }

    #[test]
    fn test_bipartite_edge_color_star() {
        // Star S_n is bipartite; max degree = n (at center)
        let mut g = Graph::new(GraphConfig::simple());
        for _ in 0..5 { g.add_node(); }
        for i in 1..5u64 { g.add_edge(0, i, None); }
        let coloring = bipartite_edge_color(&g).expect("star is bipartite");
        assert!(is_valid_edge_coloring(&g, &coloring), "star coloring valid");
        let num_colors = coloring.values().copied().max().unwrap_or(0) + 1;
        assert_eq!(num_colors, 4, "star with 4 leaves needs 4 colors, got {}", num_colors);
    }

    #[test]
    fn test_bipartite_edge_color_path() {
        // Path 0-1-2-3: bipartite, max degree = 2
        let mut g = Graph::new(GraphConfig::simple());
        for _ in 0..4 { g.add_node(); }
        g.add_edge(0, 1, None); g.add_edge(1, 2, None); g.add_edge(2, 3, None);
        let coloring = bipartite_edge_color(&g).expect("path is bipartite");
        assert!(is_valid_edge_coloring(&g, &coloring), "path coloring valid");
    }

    #[test]
    fn test_bipartite_edge_color_non_bipartite() {
        // Triangle (3-cycle) is NOT bipartite
        let mut g = Graph::new(GraphConfig::simple());
        for _ in 0..3 { g.add_node(); }
        g.add_edge(0, 1, None); g.add_edge(1, 2, None); g.add_edge(2, 0, None);
        let result = bipartite_edge_color(&g);
        assert!(result.is_err(), "triangle is not bipartite");
    }

    #[test]
    fn test_bipartite_edge_color_k12() {
        // K_{1,2}: single center connected to 2 leaves; max degree = 2
        let mut g = Graph::new(GraphConfig::simple());
        for _ in 0..3 { g.add_node(); }
        g.add_edge(0, 1, None); g.add_edge(0, 2, None);
        let coloring = bipartite_edge_color(&g).expect("K12 is bipartite");
        assert!(is_valid_edge_coloring(&g, &coloring));
        assert_eq!(coloring.len(), 2);
    }

    #[test]
    fn test_bipartite_edge_color_empty() {
        let g = Graph::new(GraphConfig::simple());
        let coloring = bipartite_edge_color(&g).expect("empty is bipartite");
        assert!(coloring.is_empty());
    }

    #[test]
    fn test_bipartite_edge_color_given_partition() {
        let g = complete_bipartite(2, 2);
        let set_a: HashSet<NodeId> = [0, 1].iter().copied().collect();
        let set_b: HashSet<NodeId> = [2, 3].iter().copied().collect();
        let coloring = bipartite_edge_color_given_partition(&g, &set_a, &set_b);
        assert!(is_valid_edge_coloring(&g, &coloring));
    }
}
