# Copyright 2016 James Hensman, Mark van der Wilk, Valentine Svensson, alexggmatthews, PabloLeon, fujiisoup
# Copyright 2017 ST John
#
# 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 __future__ import absolute_import
import types
from functools import wraps
import tensorflow as tf
from ._settings import settings
from . import session
float_type = settings.dtypes.float_type


class AutoFlow:
    """
    This decorator-class is designed for use on methods of the Parameterized class
    (below).

    The idea is that methods that compute relevant quantities (such as
    predictions) can define a tf graph which we automatically run when the
    decorated function is called.

    The syntax looks like:

    >>> class MyClass(Parameterized):
    >>>
    >>>   @AutoFlow((tf.float64), (tf.float64))
    >>>   def my_method(self, x, y):
    >>>       #compute something, returning a tf graph.
    >>>       return tf.foo(self.baz, x, y)

    >>> m = MyClass()
    >>> x = np.random.randn(3,1)
    >>> y = np.random.randn(3,1)
    >>> result = my_method(x, y)

    Now the output of the method call is the _result_ of the computation,
    equivalent to

    >>> m = MyModel()
    >>> x = np.random.randn(3,1)
    >>> y = np.random.randn(3,1)
    >>> x_tf = tf.placeholder(tf.float64)
    >>> y_tf = tf.placeholder(tf.float64)
    >>> with m.tf_mode():
    >>>     graph = tf.foo(m.baz, x_tf, y_tf)
    >>> result = m._session.run(graph,
                                feed_dict={x_tf:x,
                                           y_tf:y,
                                           m._free_vars:m.get_free_state()})

    Not only is the syntax cleaner, but multiple calls to the method will
    result in the graph being constructed only once.
    """

    def __init__(self, *tf_arg_tuples):
        # NB. TF arg_tuples is a list of tuples, each of which can be used to
        # construct a tf placeholder.
        self.tf_arg_tuples = tf_arg_tuples

    def __call__(self, tf_method):
        @wraps(tf_method)
        def runnable(instance, *np_args):
            storage_name = '_' + tf_method.__name__ + '_AF_storage'
            if hasattr(instance, storage_name):
                # the method has been compiled already, get things out of storage
                storage = getattr(instance, storage_name)
            else:
                # the method needs to be compiled.
                storage = {}  # an empty dict to keep things in
                setattr(instance, storage_name, storage)
                storage['graph'] = tf.Graph()
                # storage['session'] = tf.Session(graph=storage['graph'])
                storage['session'] = session.get_session(
                    graph=storage['graph'],
                    output_file_name=settings.profiling.output_file_name + "_" + tf_method.__name__,
                    output_directory=settings.profiling.output_directory,
                    each_time=settings.profiling.each_time
                )
                with storage['graph'].as_default():
                    storage['tf_args'] = [tf.placeholder(*a) for a in self.tf_arg_tuples]
                    storage['free_vars'] = tf.placeholder(float_type, [None])
                    instance.make_tf_array(storage['free_vars'])
                    with instance.tf_mode():
                        storage['tf_result'] = tf_method(instance, *storage['tf_args'])
                    storage['feed_dict_keys'] = instance.get_feed_dict_keys()
                    feed_dict = {}
                    instance.update_feed_dict(storage['feed_dict_keys'], feed_dict)
                    storage['session'].run(tf.global_variables_initializer(), feed_dict=feed_dict)
            feed_dict = dict(zip(storage['tf_args'], np_args))
            feed_dict[storage['free_vars']] = instance.get_free_state()
            instance.update_feed_dict(storage['feed_dict_keys'], feed_dict)
            return storage['session'].run(storage['tf_result'], feed_dict=feed_dict)

        return runnable


def AutoFlowify(*args):
    """
    Decorator that saves its arguments in the `_autoflow_args` field in the
    function so that they can be passed to AutoFlow when constructing a
    compute_method in `make_compute_method`.
    """
    def decorator(func):
        func._autoflow_args = args
        return func
    return decorator


def make_compute_method(compute_name, build_method, autoflow_args):
    """
    Returns a new method with name `compute_name` that calls `build_method`
    but is wrapped by the AutoFlow decorator with the arguments in
    `autoflow_args`.

    `compute_name` should correspond to the name given to the new method in
    the model class.
    """
    @wraps(build_method)
    def compute_method(self, *build_args):
        return build_method(self, *build_args)

    # AutoFlow uses <method>.__name__ for its graph-caching, so we need to
    # keep the original method name, and we need to apply the AutoFlow
    # decorator *after* changing the method's __name__ attribute
    compute_method.__name__ = compute_name
    compute_method = AutoFlow(*autoflow_args)(compute_method)

    return compute_method


class MetaAutoFlow(type):
    """
    Metaclass for GPflow models that automatically constructs an @AutoFlow-
    wrapped compute_<method> for each build_<method> that is decorated with
    @AutoFlowify.

    E.g., when adding this as a metaclass to a model,

    >>> @AutoFlowify((float_type, [None,None]))
    >>> def build_foo(self, a):
    >>>     return tf.reduce_sum(a)

    is equivalent to the more verbose

    >>> @AutoFlow((float_type, [None,None]))
    >>> def compute_foo(self, a):
    >>>     return self.build_foo(a)
    >>>
    >>> def build_foo(self, a):
    >>>     return tf.reduce_sum(a)

    Methods that do not start with 'build_' or that are not decorated by
    @AutoFlowify will be left alone.  If the class defines a compute_ method
    itself, that will take precedence over the metaclass-autogenerated one.
    """

    def __new__(meta, name, bases, class_dict):
        new_class_dict = {}
        for key, value in class_dict.items():
            if (key.startswith('build_') and
                    isinstance(value, types.FunctionType) and
                    hasattr(value, '_autoflow_args')):
                build_name, build_method = key, value

                compute_name = build_name.replace('build_', 'compute_', 1)
                compute_method = make_compute_method(compute_name, build_method, build_method._autoflow_args)
                new_class_dict[compute_name] = compute_method

        new_class_dict.update(class_dict)  # add all other methods, overwriting compute_<> if already exists

        cls = type.__new__(meta, name, bases, new_class_dict)
        return cls