# Copyright 2021 Sony Semiconductor Israel, Inc. All rights reserved.
#
# 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.
# ==============================================================================


from typing import Any, Dict

from model_compression_toolkit.core.common.graph.base_node import BaseNode

# Edge attributes:
EDGE_SOURCE_INDEX = 'source_index'
EDGE_SINK_INDEX = 'sink_index'


class Edge(object):
    """
    Edge in a directed graph.
    """

    def __init__(self,
                 source_node: BaseNode,
                 sink_node: BaseNode,
                 source_index: int,
                 sink_index: int):
        """
        Construct an edge by two nodes and indices.
        Args:
            source_node: Source node of the edge.
            sink_node: Sink node of the edge.
            source_index: Index of the source node of the edge.
            sink_index: Index of the sink node of the edge.
        """
        self.source_node = source_node
        self.sink_node = sink_node
        self.source_index = source_index
        self.sink_index = sink_index

    def get_attributes(self) -> Dict[str, Any]:
        """
        Get edge's attributes as a dictionary mapping an attribute name to its edge value.
        Returns: Edge's attributes.
        """

        return {EDGE_SOURCE_INDEX: self.source_index,
                EDGE_SINK_INDEX: self.sink_index}

    def __eq__(self, other: Any) -> bool:
        """
        Check if this edge is identical to the edge that was passed.
        Args:
            other: An object to compare to this edge.

        Returns:
            Whether the passed object is identical to this edge.
        """

        if isinstance(other, Edge):
            return other.sink_node == self.sink_node and \
                   other.source_node == self.source_node and \
                   other.source_index == self.source_index and \
                   other.sink_index == self.source_index

        return False

    def __repr__(self) -> str:
        """
        A string representing this edge for display.
        """

        return f'{self.source_node.name}:{self.source_index} -> ' \
               f'{self.sink_node.name}:{self.sink_index}'


def convert_to_edge(edge: Any) -> Edge:
    """
    Get an edge in the graph either as an Edge object (source node, destination node,
    source index and destination index) or as a networkx edge with edge data (triplet
    tuple with nodes and edge data) and return it as an Edge object.

    Args:
        edge: Edge in a graph to convert.

    Returns:
        An Edge object containing source node, destination node,
        source index and destination index.
    """

    if isinstance(edge, tuple) and len(edge) == 3:  # networkx edge representation with edge data
        src_node = edge[0]
        dst_node = edge[1]
        edge_data = edge[2]
        return Edge(src_node,
                    dst_node,
                    edge_data[EDGE_SOURCE_INDEX],
                    edge_data[EDGE_SINK_INDEX])

    elif isinstance(edge, Edge):  # it's already an Edge and no change need to be done
        return edge

    raise Exception('Edges list contains an object that is not a known edge format.')
