// Licensed under the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations
// under the License.

use crate::GraphNotBipartite;
use crate::{digraph, graph, EdgeIndex, NodeIndex};

use pyo3::prelude::*;
use pyo3::types::PyDict;
use pyo3::Python;

use std::convert::Infallible;

use rustworkx_core::bipartite_coloring::bipartite_edge_color;
use rustworkx_core::coloring::{
    greedy_edge_color_with_coloring_strategy, greedy_node_color_with_coloring_strategy,
    misra_gries_edge_color, two_color,
};

pub use rustworkx_core::coloring::ColoringStrategy as ColoringStrategyCore;

/// Greedy coloring strategies available for `graph_greedy_color`
///
/// .. list-table:: Strategy description
///     :header-rows: 1
///
///     * - Strategy
///       - Reference
///     * - Degree
///       - `Largest-first` strategy in [1] (section 1.2.2.2)
///     * - Saturation
///       - `DSATUR` strategy in [1] (section 1.2.2.8)
///     * - IndependentSet
///       - `GIS` strategy in [1] (section 1.2.2.9)
///
/// [1] Adrian Kosowski, and Krzysztof Manuszewski, Classical Coloring of Graphs, Graph Colorings, 2-19, 2004. ISBN 0-8218-3458-4.
#[pyclass(module = "rustworkx", eq, eq_int)]
#[derive(Clone, PartialEq)]
pub enum ColoringStrategy {
    Degree,
    Saturation,
    IndependentSet,
}

/// Color a :class:`~.PyGraph` object using a greedy graph coloring algorithm.
///
/// This function uses one of several greedy strategies described in: [1]_.
/// The `Degree` (aka `largest-first`) strategy colors the nodes with higher degree
/// first. The `Saturation` (aka `DSATUR` and `SLF`) strategy dynamically
/// chooses the vertex that has the largest number of different colors already
/// assigned to its neighbors, and, in case of a tie, the vertex that has the
/// largest number of uncolored neighbors. The `IndependentSet` strategy finds
/// independent subsets of the graph and assigns a different color to each of these
/// subsets.
///
/// .. note::
///
///     The coloring problem is NP-hard and this is a heuristic algorithm which
///     may not return an optimal solution.
///
/// :param PyGraph: The input PyGraph object to color.
/// :param preset_color_fn: An optional callback function that is used to manually
///     specify a color to use for particular nodes in the graph. If specified
///     this takes a callable that will be passed a node index and is expected to
///     either return an integer representing a color or ``None`` to indicate there
///     is no preset. Note if you do use a callable there is no validation that
///     the preset values are valid colors. You can generate an invalid coloring
///     if you the specified function returned invalid colors for any nodes.
/// :param strategy: The strategy used by the algorithm. When the
///     strategy is not explicitly specified, the `Degree` strategy is used by
///     default.
///
/// :returns: A dictionary where keys are node indices and the value is
///     the color
/// :rtype: dict
///
/// .. jupyter-execute::
///
///     import rustworkx as rx
///     from rustworkx.visualization import mpl_draw
///
///     graph = rx.generators.generalized_petersen_graph(5, 2)
///     coloring = rx.graph_greedy_color(graph)
///     colors = [coloring[node] for node in graph.node_indices()]
///
///     # Draw colored graph
///     layout = rx.shell_layout(graph, nlist=[[0, 1, 2, 3, 4],[6, 7, 8, 9, 5]])
///     mpl_draw(graph, node_color=colors, pos=layout)
///
///
/// .. [1] Adrian Kosowski, and Krzysztof Manuszewski, Classical Coloring of Graphs,
///     Graph Colorings, 2-19, 2004. ISBN 0-8218-3458-4.
#[pyfunction]
#[pyo3(text_signature = "(graph, /, preset_color_fn=None, strategy=ColoringStrategy::Degree)")]
#[pyo3(signature=(graph, /, preset_color_fn=None, strategy=ColoringStrategy::Degree))]
pub fn graph_greedy_color(
    py: Python,
    graph: &graph::PyGraph,
    preset_color_fn: Option<PyObject>,
    strategy: ColoringStrategy,
) -> PyResult<PyObject> {
    let inner_strategy = match strategy {
        ColoringStrategy::Saturation => ColoringStrategyCore::Saturation,
        ColoringStrategy::Degree => ColoringStrategyCore::Degree,
        ColoringStrategy::IndependentSet => ColoringStrategyCore::IndependentSet,
    };

    let colors = match preset_color_fn {
        Some(preset_color_fn) => {
            let callback = |node_idx: NodeIndex| -> PyResult<Option<usize>> {
                preset_color_fn
                    .call1(py, (node_idx.index(),))
                    .map(|x| x.extract(py).ok())
            };
            greedy_node_color_with_coloring_strategy(&graph.graph, callback, inner_strategy)?
        }
        None => {
            let callback = |_: NodeIndex| -> Result<Option<usize>, Infallible> { Ok(None) };
            greedy_node_color_with_coloring_strategy(&graph.graph, callback, inner_strategy)?
        }
    };
    let out_dict = PyDict::new(py);
    for (node, color) in colors {
        out_dict.set_item(node.index(), color)?;
    }
    Ok(out_dict.into())
}

/// Color edges of a :class:`~.PyGraph` object using a greedy approach.
///
/// This function works by greedily coloring the line graph of the given graph.
///
/// This function uses one of several greedy strategies described in: [1]_.
/// The `Degree` (aka `largest-first`) strategy colors the nodes with higher degree
/// first. The `Saturation` (aka `DSATUR` and `SLF`) strategy dynamically
/// chooses the vertex that has the largest number of different colors already
/// assigned to its neighbors, and, in case of a tie, the vertex that has the
/// largest number of uncolored neighbors. The `IndependentSet` strategy finds
/// independent subsets of the graph and assigns a different color to each of these
/// subsets.
///
/// :param PyGraph: The input PyGraph object to edge-color.
/// :param preset_color_fn: An optional callback function that is used to manually
///     specify a color to use for particular edges in the graph. If specified
///     this takes a callable that will be passed an edge index and is expected to
///     either return an integer representing a color or ``None`` to indicate there
///     is no preset. Note if you do use a callable there is no validation that
///     the preset values are valid colors. You can generate an invalid coloring
///     if you the specified function returned invalid colors for any edges.
/// :param strategy: The greedy strategy used by the algorithm. When the
///     strategy is not explicitly specified, the `Degree` strategy is used by
///     default.
///
/// :returns: A dictionary where keys are edge indices and the value is the color
/// :rtype: dict
///
/// .. jupyter-execute::
///
///     import rustworkx as rx
///
///     graph = rx.generators.cycle_graph(7)
///     edge_colors = rx.graph_greedy_edge_color(graph)
///     assert edge_colors == {0: 0, 1: 1, 2: 0, 3: 1, 4: 0, 5: 1, 6: 2}
///
///
/// .. [1] Adrian Kosowski, and Krzysztof Manuszewski, Classical Coloring of Graphs,
///     Graph Colorings, 2-19, 2004. ISBN 0-8218-3458-4.
#[pyfunction]
#[pyo3(text_signature = "(graph, /, preset_color_fn=None, strategy=ColoringStrategy::Degree)")]
#[pyo3(signature=(graph, /, preset_color_fn=None, strategy=ColoringStrategy::Degree))]
pub fn graph_greedy_edge_color(
    py: Python,
    graph: &graph::PyGraph,
    preset_color_fn: Option<PyObject>,
    strategy: ColoringStrategy,
) -> PyResult<PyObject> {
    let inner_strategy = match strategy {
        ColoringStrategy::Saturation => ColoringStrategyCore::Saturation,
        ColoringStrategy::Degree => ColoringStrategyCore::Degree,
        ColoringStrategy::IndependentSet => ColoringStrategyCore::IndependentSet,
    };

    let colors = match preset_color_fn {
        Some(preset_color_fn) => {
            let callback = |edge_idx: EdgeIndex| -> PyResult<Option<usize>> {
                preset_color_fn
                    .call1(py, (edge_idx.index(),))
                    .map(|x| x.extract(py).ok())
            };
            greedy_edge_color_with_coloring_strategy(&graph.graph, callback, inner_strategy)?
        }
        None => {
            let callback = |_: EdgeIndex| -> Result<Option<usize>, Infallible> { Ok(None) };
            greedy_edge_color_with_coloring_strategy(&graph.graph, callback, inner_strategy)?
        }
    };

    let out_dict = PyDict::new(py);
    for (node, color) in colors {
        out_dict.set_item(node.index(), color)?;
    }
    Ok(out_dict.into())
}

/// Color edges of a :class:`~.PyGraph` object using the Misra-Gries edge
/// coloring algorithm..
///
/// Based on the paper: "A constructive proof of Vizing's theorem" by
/// Misra and Gries, 1992.
/// <https://www.cs.utexas.edu/users/misra/psp.dir/vizing.pdf>
///
/// The coloring produces at most d + 1 colors where d is the maximum degree
/// of the graph.
///
/// :param PyGraph: The input PyGraph object to edge-color
///
/// :returns: A dictionary where keys are edge indices and the value is the color
/// :rtype: dict
///
/// .. jupyter-execute::
///
///     import rustworkx as rx
///
///     graph = rx.generators.cycle_graph(7)
///     edge_colors = rx.graph_misra_gries_edge_color(graph)
///     assert edge_colors == {0: 0, 1: 1, 2: 2, 3: 0, 4: 1, 5: 0, 6: 2}
///
#[pyfunction]
#[pyo3(text_signature = "(graph, /)")]
pub fn graph_misra_gries_edge_color(py: Python, graph: &graph::PyGraph) -> PyResult<PyObject> {
    let colors = misra_gries_edge_color(&graph.graph);
    let out_dict = PyDict::new(py);
    for (node, color) in colors {
        out_dict.set_item(node.index(), color)?;
    }
    Ok(out_dict.into())
}

/// Compute a two-coloring of a graph
///
/// If a two coloring is not possible for the input graph (meaning it is not
/// bipartite), ``None`` is returned.
///
/// :param PyGraph graph: The graph to find the coloring for
///
/// :returns: If a coloring is possible return a dictionary of node indices to the color as an
/// integer (0 or 1)
/// :rtype: dict
#[pyfunction]
pub fn graph_two_color(py: Python, graph: &graph::PyGraph) -> PyResult<Option<PyObject>> {
    match two_color(&graph.graph) {
        Some(colors) => {
            let out_dict = PyDict::new(py);
            for (node, color) in colors {
                out_dict.set_item(node.index(), color)?;
            }
            Ok(Some(out_dict.into()))
        }
        None => Ok(None),
    }
}

/// Compute a two-coloring of a directed graph
///
/// If a two coloring is not possible for the input graph (meaning it is not
/// bipartite), ``None`` is returned.
///
/// :param PyDiGraph graph: The graph to find the coloring for
///
/// :returns: If a coloring is possible return a dictionary of node indices to the color as an
/// integer (0 or 1)
/// :rtype: dict
#[pyfunction]
pub fn digraph_two_color(py: Python, graph: &digraph::PyDiGraph) -> PyResult<Option<PyObject>> {
    match two_color(&graph.graph) {
        Some(colors) => {
            let out_dict = PyDict::new(py);
            for (node, color) in colors {
                out_dict.set_item(node.index(), color)?;
            }
            Ok(Some(out_dict.into()))
        }
        None => Ok(None),
    }
}

/// Color edges of a graph by checking whether the graph is bipartite,
/// and if so, calling the algorithm for edge-coloring bipartite graphs.
///
/// If the input graph is not bipartite, ``None`` is returned.
///
/// The implementation is based on the following paper:
///
/// Noga Alon. "A simple algorithm for edge-coloring bipartite multigraphs".
/// Inf. Process. Lett. 85(6), (2003).
/// <https://www.tau.ac.il/~nogaa/PDFS/lex2.pdf>
///
/// The algorithm runs in time `O (n + m log m)`, where `n` is the number of
/// vertices and `m` is the number of edges of the graph.
///
/// :param PyGraph graph: The graph to find the coloring for
///
/// :returns: A dictionary where keys are edge indices and the value is the color
///  (provided that the graph is bipartite)
/// :rtype: dict
#[pyfunction]
#[pyo3(text_signature = "(graph, /)")]
pub fn graph_bipartite_edge_color(py: Python, graph: &graph::PyGraph) -> PyResult<PyObject> {
    let colors = match bipartite_edge_color(&graph.graph) {
        Ok(colors) => colors,
        Err(_) => return Err(GraphNotBipartite::new_err("Graph is not bipartite")),
    };
    let out_dict = PyDict::new(py);
    for (node, color) in colors {
        out_dict.set_item(node.index(), color)?;
    }
    Ok(out_dict.into())
}
