"""Module to compute distance between nodes or cells on topological domains."""

from collections.abc import Hashable, Iterable

import networkx as nx

from toponetx.classes.cell import Cell
from toponetx.classes.cell_complex import CellComplex
from toponetx.classes.colored_hypergraph import ColoredHyperGraph
from toponetx.classes.combinatorial_complex import CombinatorialComplex
from toponetx.classes.hyperedge import HyperEdge
from toponetx.exception import TopoNetXNoPath

__all__ = ["cell_distance", "distance"]

# In this module, only cell complexes, combinatorial complexes and colored
# hypergraphs are supported.
ComplexType = CellComplex | CombinatorialComplex | ColoredHyperGraph


def distance(
    domain: ComplexType, source: Hashable, target: Hashable, s: int = 1
) -> int:
    """Return shortest s-walk distance between two nodes in the cell complex.

    Parameters
    ----------
    domain : CellComplex or CombinaorialComplex or ColoredHyperGraph
        The domain on which to compute the s-walk distance between source and target.
    source : Hashable
        A node in the input complex.
    target : Hashable
        A node in the input complex.
    s : int, optional
        The number of intersections between pairwise consecutive cells.

    Returns
    -------
    int
        Shortest s-walk distance between two nodes in the cell complex.

    Raises
    ------
    TopoNetXNoPath
        If no s-walk path exists between source and target nodes.

    See Also
    --------
    cell_distance
        Method implemented in the library that returns
        the shortest s-walk distance between two cells in the cell complex.

    Notes
    -----
    The s-distance is the shortest s-walk length between the nodes.
    An s-walk between nodes is a sequence of nodes that pairwise share
    at least s cells. The length of the shortest s-walk is 1 less than
    the number of nodes in the path sequence.

    Uses the networkx shortest_path_length method on the graph
    generated by the s-adjacency matrix.

    Examples
    --------
    >>> CC = tnx.CellComplex()
    >>> CC.add_cell([2, 3, 4], rank=2)
    >>> CC.add_cell([5, 6, 7], rank=2)
    >>> list(tnx.node_diameters(CC))
    >>> CCC = CC.to_combinatorial_complex()
    >>> list(tnx.node_diameters(CCC))
    >>> CHG = CC.to_colored_hypergraph()
    >>> list(tnx.node_diameters(CHG))
    """
    if not isinstance(domain, CellComplex | CombinatorialComplex | ColoredHyperGraph):
        raise TypeError(f"Input complex {domain} is not supported.")
    if isinstance(source, Cell):
        source = source.elements
    if isinstance(target, Cell):
        target = target.elements
    if isinstance(source, Iterable):
        source = tuple(source)
    if isinstance(target, Iterable):
        target = tuple(target)
    rowdict, A = domain.node_to_all_cell_adjacnecy_matrix(s=s, index=True)
    G = nx.from_scipy_sparse_array(A)
    try:
        return nx.shortest_path_length(G, rowdict[source], rowdict[target])
    except nx.NetworkXNoPath as exc:
        raise TopoNetXNoPath from exc


def cell_distance(
    domain: ComplexType,
    source: Iterable | HyperEdge | Cell,
    target: Iterable | HyperEdge | Cell,
    s: int = 1,
) -> int:
    """Return the shortest s-walk distance between two cells in the cell complex.

    Parameters
    ----------
    domain : CellComplex or CombinatorialComplex or ColoredHyperGraph
        The domain on which to compute the s-walk distance between source and target cells.
    source : Iterable or HyperEdge or Cell
        An Iterable representing a cell in the input complex cell complex.
    target : Iterable or HyperEdge or Cell
        An Iterable representing a cell in the input complex cell complex.
    s : int
        The number of intersections between pairwise consecutive cells.

    Returns
    -------
    int
        Shortest s-walk cell distance between `source` and `target`. A shortest s-walk
        is computed as a sequence of cells, the s-walk distance is the number of cells
        in the sequence minus 1.

    Raises
    ------
    TopoNetXNoPath
        If no s-walk path exists between source and target cells.

    See Also
    --------
    distance
        Method implemented in the library that returns
        shortest s-walk distance between two nodes in the cell complex.

    Notes
    -----
    The s-distance is the shortest s-walk length between the cells.
    An s-walk between cells is a sequence of cells such that consecutive pairwise
    cells intersect in at least s nodes. The length of the shortest s-walk is 1 less than
    the number of cells in the path sequence.

    Uses the networkx shortest_path_length method on the graph
    generated by the s-cell_adjacency matrix.

    Examples
    --------
    >>> CC = tnx.CellComplex()
    >>> CC.add_cell([2, 3, 4], rank=2)
    >>> CC.add_cell([5, 6, 7], rank=2)
    >>> CC.add_cell([5, 2], rank=1)
    >>> tnx.cell_distance(CC, [2, 3], [6, 7])
    >>> CHG = CC.to_colored_hypergraph()
    >>> tnx.cell_distance(CHG, (frozenset({2, 3}), 0), (frozenset({6, 7}), 0))
    >>> CCC = CC.to_combinatorial_complex()
    >>> tnx.cell_distance(CCC, frozenset({2, 3}), frozenset({6, 7}))
    """
    if not isinstance(domain, CellComplex | CombinatorialComplex | ColoredHyperGraph):
        raise TypeError(f"Input complex {domain} is not supported.")
    if isinstance(source, Cell | HyperEdge):
        source = source.elements
    if isinstance(target, Cell | HyperEdge):
        target = target.elements

    if isinstance(domain, CellComplex):
        if isinstance(source, Iterable):
            source = tuple(source)
        if isinstance(target, Iterable):
            target = tuple(target)
    if isinstance(domain, ColoredHyperGraph):
        if isinstance(source, Iterable):
            source = frozenset(source)
        if isinstance(target, Iterable):
            target = frozenset(target)

    cell_dict, A = domain.all_cell_to_node_coadjacency_matrix(s=s, index=True)
    print(cell_dict)
    G = nx.from_scipy_sparse_array(A)
    try:
        return nx.shortest_path_length(G, cell_dict[source], cell_dict[target])
    except nx.NetworkXNoPath as exc:
        raise TopoNetXNoPath from exc
