# Copyright 2015 The TensorFlow Authors. 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.
# ==============================================================================
"""Utilities borrowed from the tf.keras codebase."""
import collections
import weakref

from tensorflow.python.framework import ops

# A global dictionary mapping graph objects to an index of counters used
# for various layer/optimizer names in each graph.
# Allows to give unique autogenerated names to layers, in a graph-specific way.
PER_GRAPH_OBJECT_NAME_UIDS = weakref.WeakKeyDictionary()


def _get_default_graph_uid_map():
    graph = ops.get_default_graph()
    name_uid_map = PER_GRAPH_OBJECT_NAME_UIDS.get(graph, None)
    if name_uid_map is None:
        name_uid_map = collections.defaultdict(int)
        PER_GRAPH_OBJECT_NAME_UIDS[graph] = name_uid_map
    return name_uid_map


def unique_object_name(
    name,
    name_uid_map=None,
    avoid_names=None,
    namespace="",
    zero_based=False,
):
    """Makes a object name (or arbitrary string) unique within a TensorFlow graph.
    Arguments:
      name: String name to make unique.
      name_uid_map: An optional defaultdict(int) to use when creating unique
        names. If None (default), uses a per-Graph dictionary.
      avoid_names: An optional set or dict with names which should not be used. If
        None (default) does not avoid any names.
      namespace: Gets a name which is unique within the (graph, namespace). Layers
        which are not Networks use a blank namespace and so get graph-global
        names.
      zero_based: If True, name sequences start with no suffix (e.g. "dense",
        "dense_1"). If False, naming is one-based ("dense_1", "dense_2").
    Returns:
      Unique string name.
    Example:
    ```python
    _unique_layer_name('dense')  # dense_1
    _unique_layer_name('dense')  # dense_2
    ```
    """
    if name_uid_map is None:
        name_uid_map = _get_default_graph_uid_map()
    if avoid_names is None:
        avoid_names = set()
    proposed_name = None
    while proposed_name is None or proposed_name in avoid_names:
        name_key = (namespace, name)
        if zero_based:
            number = name_uid_map[name_key]
            if number:
                proposed_name = name + "_" + str(number)
            else:
                proposed_name = name
            name_uid_map[name_key] += 1
        else:
            name_uid_map[name_key] += 1
            proposed_name = name + "_" + str(name_uid_map[name_key])
    return proposed_name
