#  Copyright (c) 2020, Apple Inc. All rights reserved.
#
#  Use of this source code is governed by a BSD-3-clause license that can be
#  found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause

from collections import OrderedDict
from itertools import islice

import torch


def _make_ssa_name(name):
    """Converts a symbol name (string) into an SSA name, by prepending '%'.
    Only used for pretty printing the graph.
    """
    return "%" + name


def _ssa_name_list(names):
    """Take a list of symbol names (strings) and return them as SSA names. Only
    used for pretty printing the graph.
    """
    return [_make_ssa_name(x) for x in names]


def _find_new_name(old_name, node_names):
    """Disambiguate a node's name from a list of existing node names by adding
    successively larger integers.
    """
    count = 0
    new_name = old_name + "." + str(count)
    while new_name in node_names:
        count += 1
        new_name = old_name + "." + str(count)
    return new_name


def _replace_in_list(ls, old_val, new_val):
    """Helper function to replace a value in a list."""
    try:
        idx = ls.index(old_val)
    except ValueError:
        pass
    else:
        ls[idx] = new_val


class InternalTorchIRBlock:
    """CoreML internal representation of a torch IR block.

        Arguments:
            raw_block: The torch._C.Block to convert, or None.
            parent: The InternalTorchIRNode this block belongs to.
            nodes: If @raw_block is None, the list of InternalTorchIRNodes in the block
            inputs: If @raw_block is None, the list of input symbols.
            outputs: If @raw_block is None, the list of output symbols.
    """

    def __init__(self, raw_block=None, parent=None, nodes=None, inputs=None, outputs=None):
        self.nodes = []
        node_names = set()
        self.inputs = []
        self.outputs = []
        self.parent = parent

        if raw_block:
            # Add nodes
            for raw_node in raw_block.nodes():
                new_node = InternalTorchIRNode(raw_node, parent=self)
                if new_node.name == new_node.kind:
                    new_node.name = _find_new_name(new_node.name, node_names)
                self.nodes.append(new_node)
                node_names.add(new_node.name)

            # Add inputs
            for inp in raw_block.inputs():
                self.inputs.append(inp.debugName())

            # Add outputs
            for outp in raw_block.outputs():
                self.outputs.append(outp.debugName())
        else:
            self.nodes = nodes
            self.inputs = inputs
            self.outputs = outputs

    def __str__(self, indent=2):
        indent_str = " " * indent
        graph_str = "{}block({}):\n".format(
            indent_str, ", ".join(_ssa_name_list(self.inputs))
        )
        graph_str += "{}\n".format(indent_str).join(
            [x.__str__(indent=indent + 2) for x in self.nodes]
        )
        graph_str += "\n{}return ({})".format(
            indent_str, ", ".join(_ssa_name_list(self.outputs))
        )
        return graph_str

    def __repr__(self):
        return str(self)

    def replace_name(self, old_name, new_name):
        """Replaces all instances of @old_name with @new_name in @self."""

        # Replace graph inputs/outputs
        _replace_in_list(self.inputs, old_name, new_name)
        _replace_in_list(self.outputs, old_name, new_name)

        for node in self.nodes:
            node.replace_name(old_name, new_name)


class InternalTorchIRNode:
    """CoreML internal representation of a torch IR node.
    Can construct itself from a provided torchIR node or manually constructed with
    args for testing.

    See InternalTorchIRGraph for the motivation behind this structure.

        Arguments:
            node: The torch._C.Node to convert, or None.
            parent: The InternalTorchIRGraph/Block this node belongs to.
            attr: If @node is not specified, the dict of named attributes.
            inputs: If @node is not specified, the list of input symbols.
            outputs: If @node is not specified, the list of output symbols.
            kind: If @node is not specified, the kind (op) of the node.
            blocks: If @node is not specified, the list of InternalTorchIRBlock.
    """

    def __init__(
        self, node=None, parent=None, attr=None, inputs=None, outputs=None, kind=None, blocks=None,
    ):
        self.parent = parent
        if node is not None:
            self.inputs = [_input.debugName() for _input in node.inputs()]
            self.outputs = [output.debugName() for output in node.outputs()]
            self.kind = node.kind().split("::")[-1].lower()
            self.blocks = [InternalTorchIRBlock(raw_block=b, parent=self) for b in node.blocks()]
            self.attr = {
                name: getattr(node, node.kindOf(name))(name)
                for name in node.attributeNames()
            }
            if "value" not in self.attr:
                self.attr["value"] = None
            # If the output is boolean, explicitly cast it so type inference
            # will work correctly.
            if len(self.outputs) == 1 and next(node.outputs()).type().str() == "bool":
                self.attr["value"] = bool(self.attr["value"])
        else:
            self.inputs = inputs
            self.outputs = outputs
            self.kind = kind
            self.blocks = blocks if blocks is not None else []
            self.attr = attr if attr is not None else {"value": None}
        # On rare occassions, a node has no outputs. In that case, the node's
        # name will be its kind. However, this no longer guarantees the node's
        # name is unique. It will be up to the graph constructing the node to
        # make sure names are unique.
        self.name = self.outputs[0] if len(self.outputs) > 0 else self.kind

    def __str__(self, indent=2):
        node_str = " " * indent + "{} = {}".format(
            ", ".join(_ssa_name_list(self.outputs)), self.kind
        )
        node_str += "[{}]".format(
            ", ".join(
                ["{}={}".format(n, v) for n, v in self.attr.items() if v is not None]
            )
        )
        node_str += "({})".format(", ".join(_ssa_name_list(self.inputs)))
        for b in self.blocks:
            node_str += "\n" + b.__str__(indent=indent + 2)
        return node_str

    def __repr__(self):
        return str(self)

    def replace_name(self, old_name, new_name):
        """Replaces all instances of @old_name with @new_name in @self."""

        _replace_in_list(self.inputs, old_name, new_name)
        _replace_in_list(self.outputs, old_name, new_name)

        if self.name == old_name:
            self.name = new_name
        for block in self.blocks:
            block.replace_name(old_name, new_name)


class InternalTorchIRGraph:
    """CoreML internal representation of a torch IR graph. A torch._C.Graph
    object is not an ideal structure to use in converting to CoreML. Conversion
    to an InternalTorchIRGraph is inserted between the original graph and the
    final CoreML model to address several issues:
        1. A torch._C.graph is hard to work with. For example, its .inputs()
          and .outputs() functions return iterators, so the only way to
          determine the number of inputs/outputs is by counting to the end.
          There are other examples of why the torch structure is hard to work
          with, and this structure alleviates those isses.
        2. torch._C.graph is an internal API and so we can't count on its
          stability. By inserting a layer in between, we can handle any changes
          to torch._C.graph here and isolate the ops code that processes the
          graph.
        3. torch._C.graph does not expose a Python constructor. This makes
          it impossible to write unit tests that isolate specific ops since
          they have to come from actually converting a PyTorch graph. With an
          internal structure, we can directly build the test cases we need for
          unit testing.

        Arguments:
            raw_graph: raw_graph: The torch._C.Graph to convert, or None.
            params_dict: A dictionary mapping graph parameter names to tensors.
                Must be given if @raw_graph is not None.
            input_values: A list of inputs to the graph. Must be given is
                @raw_graph if not None.
            cut_at_symbols: The list of desired outputs from the graph. Symbols
                must be present in the graph. For debugging use only. Can only
                be given if @raw_graph is not None.
            nodes: If @raw_graph is None, the list of InternalTorchIRNodes in
                the graph.
            params: If @raw_graph is None, the dict mapping parameter names to
                their numpy value.
            inputs: If @raw_graph is None, the OrderedDict mapping input names
                to their example values.
            outputs: If @raw_graph is None, the list of outputs from the graph.
    """

    def __init__(
        self, raw_graph=None, params_dict=None, input_values=None, cut_at_symbols=None, nodes=None, params=None, inputs=None, outputs=None,
    ):
        self.nodes = []
        node_names = set()
        self.params = {}
        self.inputs = OrderedDict()
        self.outputs = []

        if raw_graph is not None:
            # Add nodes
            for raw_node in raw_graph.nodes():
                new_node = InternalTorchIRNode(raw_node, parent=self)
                if new_node.name == new_node.kind:
                    new_node.name = _find_new_name(new_node.name, node_names)
                self.nodes.append(new_node)
                node_names.add(new_node.name)

            # Add params
            for name, param in params_dict.items():
                value = param.detach().cpu().numpy()
                self.params[name] = value

            # Add inputs
            for index, _input in enumerate(islice(raw_graph.inputs(), len(input_values))):
                name = _input.debugName()
                value = input_values[index]
                self.inputs[name] = value

            # Add outputs, cutting if @cut_at_symbols is set
            output_names = cut_at_symbols
            if output_names is None:
                output_names = [x.debugName() for x in raw_graph.outputs()]
            for output in output_names:
                self.outputs.append(output)
        else:
            self.nodes = nodes
            self.params = params
            self.inputs = inputs
            self.outputs = outputs

    def __str__(self):
        graph_str = "graph(\n"
        graph_str += self._format_inputs(self.inputs, unpack=True)
        graph_str += self._format_inputs(self.params)
        graph_str += "):\n"
        graph_str += "\n".join([str(x) for x in self.nodes]) + "\n"
        graph_str += "return ({})".format(", ".join(_ssa_name_list(self.outputs)))
        return graph_str

    def _format_inputs(self, inputs, unpack=False):
        def tensor_str(x):
            return "Tensor{}".format(
                tuple(list(x.shape.shape if unpack else x.shape) + [str(x.dtype)])
            )

        inp_str = ""
        for k, v in inputs.items():
            if isinstance(v, (tuple, list)):
                shape_str = "({})".format(", ".join([tensor_str(x) for x in v]))
            else:
                shape_str = tensor_str(v)
            inp_str += "    {} : {},\n".format(_make_ssa_name(k), shape_str)
        return inp_str

    def __repr__(self):
        return str(self)

    def replace_name(self, old_name, new_name):
        """Replaces all instances of @old_name with @new_name in @self."""

        # Replace graph inputs/outputs
        _replace_in_list(self.inputs, old_name, new_name)
        _replace_in_list(self.outputs, old_name, new_name)

        for node in self.nodes:
            node.replace_name(old_name, new_name)
