#!/usr/bin/env python
# ******************************************************************************
# Copyright 2022 Brainchip Holdings Ltd.
#
# 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.
# ******************************************************************************
import keras
import tensorflow as tf
import numpy as np
from typing import Callable
import warnings

_GLOBAL_CUSTOM_OBJECTS = {}


class StaticHashTable(keras.layers.Layer):
    """Reference object that extends the :obj:`tf.lookup.StaticHashTable` declaration,
    in order to store the input data in independent variables and **build the static table at
    time of building the layer**, looking to manipulate the layer as an general object.

    Args:
        keys (:obj:`np.ndarray`): the keys tensor in the static table.
        values (:obj:`np.ndarray`): the values tensor in the static table.
        default_value (:obj:`np.ndarray`): the default value.
        name (str, optional): the name of the operation. Defaults to None.
    """

    def __init__(self, keys, values, default_value, name=None, **kwargs):
        super().__init__(name=name, **kwargs)
        if not (isinstance(keys, np.ndarray) and isinstance(values, np.ndarray)):
            raise ValueError("StaticHashTable does not support Tensor. Please use numpy values.")
        self._keys, self._values = keys, values
        self._kwargs = {'default_value': default_value}

    def build(self, _):
        self.hash_table = tf.lookup.StaticHashTable(tf.lookup.KeyValueTensorInitializer(
            self._keys, self._values, self._keys.dtype, self._values.dtype), **self._kwargs)

    def call(self, x):
        return self.hash_table.lookup(tf.cast(x, self.hash_table.key_dtype))


class RegisterCustomHashTable:
    """Register values generated by a function in a global variable, as a static hash table.

    Given a generator function, :obj:`RegisterCustomHashTable` will store the data generated by it
    as a global variable, identified with a table name.

    The following code is an example of the common uses:

    .. code-block:: python
        with RegisterCustomHashTable(table_name='my_table', gen_func=my_func) as register:
            table = register.get_table()

    Args:
        table_name (str): the name of the hash table to get.
        gen_func (Callable): the generator function. Expected to return a tuple of three elements:
            - keys (:obj:`np.ndarray`): the input tensor.
            - values (:obj:`np.ndarray`): the output tensor.
            - default_value (:obj:`np.ndarray`): the default value.
        drop_table (bool, optional): if True, the hash table will be dropped after the context is
            exited. Defaults to False.
        *args, **kwargs: additional arguments used in the call of generator function.

    Note:
        The generator function only will be called if the hash table is not registered.

    Raises:
        ValueError: if the generator function does not return a tuple of three elements.
    """

    def __init__(self, table_name: str, gen_func: Callable, *args, drop_table=False, **kwargs):
        self.table_name = table_name
        self.gen_func = gen_func
        self.drop_table = drop_table
        self._args = args
        self._kwargs = kwargs

    def get_table(self):
        """Return the hash table given a table name setting in context.

        Returns:
            :obj:`tf.lookup.StaticHashTable`: the static hash table.
        """
        return self.get_static_hash_table(self.table_name)

    @staticmethod
    def get_static_hash_table(table_name: str):
        """Return the hash table registered previously.

        Args:
            table_name (str): the name of the hash table.

        Returns:
            :obj:`tf.lookup.StaticHashTable`: the static hash table.
        """
        # Retrieves hash static table name
        if table_name not in _GLOBAL_CUSTOM_OBJECTS:
            raise ValueError(f'Hash table {table_name} not found')
        return _GLOBAL_CUSTOM_OBJECTS[table_name]

    @staticmethod
    def drop_hash_table(table_name: str):
        """Drop a hash table given its name.

        Args:
            table_name (str): the name of the hash table.
        """
        if table_name not in _GLOBAL_CUSTOM_OBJECTS:
            raise ValueError(f'Hash table {table_name} not found')
        _GLOBAL_CUSTOM_OBJECTS.pop(table_name)

    def _register_static_hash_table(self):
        """ Register a static hash table, where the data is generated by :attr:`gen_func`.
        """
        if self.table_name in _GLOBAL_CUSTOM_OBJECTS:
            warnings.warn(f'Hash table {self.table_name} already exists. Overwriting...')

        # Register hash table
        keys, values, default_value = self.gen_func(*self._args, **self._kwargs)
        _GLOBAL_CUSTOM_OBJECTS[self.table_name] = StaticHashTable(
            keys, values, name=self.table_name, default_value=default_value)

    def _check_registered_hash_table(self):
        """Check if a hash table is registered given its name.

        Returns:
            bool: True if the hash table is registered, False otherwise.
        """
        return self.table_name in _GLOBAL_CUSTOM_OBJECTS

    def __enter__(self):
        """ Builtin method to enter of custom context scope ('with' statement),
        registering a hash table if it is not in the context.
        """
        if not self._check_registered_hash_table():
            self._register_static_hash_table()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """ Builtin method to exit of custom context scope ('with' statement) and
        drop the hash table if :attr:`drop_table` is True.
        """
        if self.drop_table and self._check_registered_hash_table():
            self.drop_hash_table(self.table_name)
