# -*- coding: utf-8 -*-
import tensorflow as tf

from tensorflow.python.util import compat
from tensorflow.python.framework import ops
from tensorflow.python.framework import errors
from tensorflow.python.framework import importer
from tensorflow.python.framework import c_api_util
from tensorflow.python.framework import op_def_registry
from tensorflow.python import pywrap_tensorflow as c_api


def set_device(graph, device):
    """Set the device for the graph.
    Args:
      graph: Optimized TensorFlow Graph.
      device: string, the device name specified by user.

    Returns: 
      Optimized TensorFlow Graph, while partial op.device have been reset.
    """
    ops = graph.get_operations()
    for op in ops:
        if op.device == '/device:EPU:0':
            op._set_device(device)
    return graph


def import_graph(model_file):
    """Import optimized graph. 
    Args:
      model_file: Optimized TensorFlow Graph, e.g., resnet50.tflex.

    Returns: 
      Optimized TensorFlow Graph, while partial op.device have been reset.
    """
    input_arrays = []
    output_arrays = []
    with tf.gfile.GFile(model_file, 'rb') as fid:
        serialized_graph = fid.read()
        import_graph_def_from_serialized(serialized_graph, name='')
        graph = tf.get_default_graph()
        ops = graph.get_operations()
        for op in ops:
            try:
                op.get_attr("_is_input")
                input_arrays.append(op.name)
            except:
                pass
            try:
                op.get_attr("_is_output")
                output_arrays.append(op.name)
            except:
                pass

    return graph, input_arrays, output_arrays


def get_tensors(graph, tensor_names):
    """Gets the Tensors associated with the `tensor_names` in the provided graph.

    Args:
      graph: TensorFlow Graph.
      tensor_names: List of strings that represent names of tensors in the graph.

    Returns:
      A list of Tensor objects in the same order the names are provided.

    Raises:
      ValueError:
        tensor_names contains an invalid tensor name.
    """
    # Get the list of all of the tensors.
    tensor_name_to_tensor = {
        tensor_name(tensor): tensor for op in graph.get_operations()
        for tensor in op.values()
    }

    # Get the tensors associated with tensor_names.
    tensors = []
    invalid_tensors = []
    for name in tensor_names:
        tensor = tensor_name_to_tensor.get(name)
        if tensor is None:
            invalid_tensors.append(name)
        else:
            tensors.append(tensor)

    # Throw ValueError if any user input names are not valid tensors.
    if invalid_tensors:
        raise ValueError("Invalid tensors '{}' were found.".format(
            ",".join(invalid_tensors)))
    return tensors


def import_graph_def_from_serialized(serialized_graph, name=''):
    op_dict = op_def_registry.get_registered_ops()
    input_map = {}
    return_elements = None

    graph = ops.get_default_graph()
    with ops.name_scope(name, 'import', input_map.values()) as scope:
        # Save unique prefix generated by name_scope
        if scope:
            assert scope.endswith('/')
            prefix = scope[:-1]
        else:
            prefix = ''

    scoped_options = c_api_util.ScopedTFImportGraphDefOptions()
    options = scoped_options.options
    importer._PopulateTFImportGraphDefOptions(options, prefix, input_map, return_elements)

    with graph._lock:  # pylint: disable=protected-access
        with c_api_util.tf_buffer(serialized_graph) as serialized:
            try:
                results = c_api.TF_GraphImportGraphDefWithResults(graph._c_graph, serialized,
                                                                  options)  # pylint: disable=protected-access
                results = c_api_util.ScopedTFImportGraphDefResults(results)
            except errors.InvalidArgumentError as e:
                # Convert to ValueError for backwards compatibility.
                raise ValueError(str(e))

        importer._ProcessNewOps(graph)

    # Treat input mappings that don't appear in the graph as an error, because
    # they are likely to be due to a typo.
    missing_unused_input_keys = (c_api.TF_ImportGraphDefResultsMissingUnusedInputMappings_wrapper(results.results))
    if missing_unused_input_keys:
        missing_unused_input_keys = [compat.as_str(s) for s in missing_unused_input_keys]
        raise ValueError(
            'Attempted to map inputs that were not found in graph_def: [%s]' % ', '.join(missing_unused_input_keys))

    if return_elements is None:
        return None
    else:
        return importer._GatherReturnElements(return_elements, graph, results.results)


def tensor_name(x):
    return x.name.split(":")[0]
