from metagraph import translator
from metagraph.plugins import has_networkx, has_scipy


if has_networkx and has_scipy:
    import networkx as nx
    import numpy as np
    from .types import NetworkXGraph
    from ..scipy.types import ScipyGraph
    from ..numpy.types import NumpyNodeSet, NumpyNodeMap

    @translator
    def graph_from_scipy(x: ScipyGraph, **props) -> NetworkXGraph:
        from ..python.types import dtype_casting

        aprops = ScipyGraph.Type.compute_abstract_properties(
            x, {"is_directed", "edge_type", "edge_dtype"}
        )

        nx_graph = nx.from_scipy_sparse_matrix(
            x.edges.value,
            create_using=nx.DiGraph if aprops["is_directed"] else nx.Graph,
            edge_attribute="weight",
        )

        if aprops["edge_type"] == "set":
            # Remove weight attribute
            for _, _, attr in nx_graph.edges(data=True):
                del attr["weight"]
        else:
            caster = dtype_casting[aprops["edge_dtype"]]
            for _, _, attr in nx_graph.edges(data=True):
                attr["weight"] = caster(attr["weight"])

        if x.edges.node_list is not None:
            pos2id = dict(enumerate(x.edges.node_list))
            nx.relabel_nodes(nx_graph, pos2id, False)

        if x.nodes is not None:
            if isinstance(x.nodes, NumpyNodeSet):
                nx_graph.add_nodes_from(x.nodes)
            elif isinstance(x.nodes, NumpyNodeMap):
                # TODO make __iter__ a required method for NodeMap implementations or making __getitem__ handle sets of ids to simplify this sort of code
                make_weight_dict = lambda weight: {"weight": weight}
                if x.nodes.mask is not None:
                    ids = np.flatnonzero(x.nodes.mask)
                    attrs = map(make_weight_dict, x.nodes.value[x.nodes.mask])
                    id2attr = dict(zip(ids, attrs))
                elif x.nodes.id2pos is not None:
                    id2attr = {
                        node_id: make_weight_dict(x.nodes.value[pos])
                        for node_id, pos in x.nodes.id2pos.items()
                    }
                else:
                    id2attr = dict(enumerate(map(make_weight_dict, x.nodes.value)))
                nx.set_node_attributes(nx_graph, id2attr, name="weight")

        return NetworkXGraph(nx_graph)
