import functools, networkx
from typing import List, Optional, Union
from context.proto.context_pb2 import EndPoint
from context.service.database.EndPointModel import EndPointModel
from .Enums import EdgeTypeEnum, LayerTypeEnum
#from .Tools import get_endpoint

def select_copper_edges(topology, n1, n2):
    selected_edges = {EdgeTypeEnum.COPPER, EdgeTypeEnum.INTERNAL, EdgeTypeEnum.OTHER}
    return topology[n1][n2].get('edge_type', EdgeTypeEnum.OTHER) in selected_edges

def select_optical_edges(topology, n1, n2):
    selected_edges = {EdgeTypeEnum.OPTICAL, EdgeTypeEnum.INTERNAL, EdgeTypeEnum.OTHER}
    return topology[n1][n2].get('edge_type', EdgeTypeEnum.OTHER) in selected_edges

def select_microwave_edges(topology, n1, n2):
    selected_edges = {EdgeTypeEnum.MICROWAVE, EdgeTypeEnum.INTERNAL, EdgeTypeEnum.OTHER}
    return topology[n1][n2].get('edge_type', EdgeTypeEnum.OTHER) in selected_edges

SELECT_LAYER = {
    LayerTypeEnum.COPPER    : select_copper_edges,
    LayerTypeEnum.OPTICAL   : select_optical_edges,
    LayerTypeEnum.MICROWAVE : select_microwave_edges,
}

def get_layer(topology : networkx.Graph, layer : LayerTypeEnum):
    filter_edge = functools.partial(SELECT_LAYER[layer], topology)
    return networkx.subgraph_view(topology, filter_edge=filter_edge)

def select_layer_from_endpoints(endpoints : List[Union[EndPointModel, EndPoint]]) -> Optional[LayerTypeEnum]:
    endpoint_types = set()
    for endpoint in endpoints: endpoint_types.add(endpoint.endpoint_type)
    if len(endpoint_types) != 1: return None
    # pylint: disable=no-member,protected-access
    return LayerTypeEnum._value2member_map_.get(endpoint_types.pop())

#a_endpoint = get_endpoint(self.__topology, a_endpoint_key)
#z_endpoint = get_endpoint(self.__topology, z_endpoint_key)
#endpoints = [a_endpoint, z_endpoint]
#layer_type = select_layer_from_endpoints(endpoints)
#topology = self.__topology if layer_type is None else get_layer(self.__topology, layer_type)
#write_dot(topology, '../data/layer-{:s}.dot'.format('all' if layer_type is None else str(layer_type.value)))

