from typing import Dict, Tuple, Union, List, Set
import networkx as nx

def exporter_graphe_filtre(
    G: nx.Graph,
    liens_chemins: List[Tuple[Union[str, int], Union[str, int]]]
) -> nx.Graph:
    """Gère l'export du graphe filtré au format DOT.

    Args:
        G (nx.Graph): Le graphe d'origine à exporter.
        liens_chemins (list): Liste des tuples contenant les paires de nœuds reliées
            par un chemin dans le graphe, avec leurs attributs associés.

    Returns:
        tuple: Un tuple contenant le graphe exporté sous forme de DiGraph
            et le dictionnaire des attributs du graphe exporté.
    """

    G_export = nx.DiGraph()
    for u, v in liens_chemins:
        G_export.add_node(u, **G.nodes[u])
        G_export.add_node(v, **G.nodes[v])
        data = G.get_edge_data(u, v)
        if isinstance(data, dict) and all(isinstance(k, int) for k in data):
            G_export.add_edge(u, v, **data[0])
        elif isinstance(data, dict):
            G_export.add_edge(u, v, **data)
        else:
            G_export.add_edge(u, v)

    return(G_export)

def extraire_liens_filtres(
    chemins: List[List[Union[str, int]]],
    niveaux: Dict[str | int, int],
    niveau_depart: int,
    niveau_arrivee: int,
    niveaux_speciaux: list[int]
) -> List[Tuple[Union[str, int], Union[str, int]]]:
    """Extrait les liens des chemins en respectant les niveaux.

    Args:
        chemins (list): Liste des chemins dans le graphe.
        niveaux (dict): Dictionnaire associant chaque nœud au niveau correspondant.
        niveau_depart (int): Niveau de départ pour la sélection des liens.
        niveau_arrivee (int): Niveau d'arrivée pour la sélection des liens.
        niveaux_speciaux (set): Ensemble des niveaux spéciaux à considérer dans le filtrage.

    Returns:
        set: Ensemble des paires de nœuds constituant les liens du graphe filtré
            en respectant les niveaux et les niveaux spéciaux demandés.
    """
    liens = set()
    for chemin in chemins:
        for i in range(len(chemin) - 1):
            u, v = chemin[i], chemin[i + 1]
            niveau_u = niveaux.get(u, 999)
            niveau_v = niveaux.get(v, 999)
            if (
                (niveau_depart <= niveau_u <= niveau_arrivee or niveau_u in niveaux_speciaux)
                and (niveau_depart <= niveau_v <= niveau_arrivee or niveau_v in niveaux_speciaux)
            ):
                liens.add((u, v))
    return liens