#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Distributed under terms of the MIT license.

"""An example of using tfp.optimizer.lbfgs_minimize to optimize a TensorFlow model.

This code shows a naive way to wrap a tf.keras.Model and optimize it with the L-BFGS
optimizer from TensorFlow Probability.

Python interpreter version: 3.6.9
TensorFlow version: 2.0.0
TensorFlow Probability version: 0.8.0
NumPy version: 1.17.2
Matplotlib version: 3.1.1
"""
import numpy
import tensorflow as tf
import tensorflow_probability as tfp
from matplotlib import pyplot
from . import complex_math 

__all__ = ['function_factory']

def function_factory(model, loss, dataset):
    """A factory to create a function required by tfp.optimizer.lbfgs_minimize.

    Args:
        model [in]: an instance of `tf.keras.Model` or its subclasses.
        loss [in]: a function with signature loss_value = loss(pred_y, true_y).
        train_x [in]: the input part of training data.
        train_y [in]: the output part of training data.

    Returns:
        A function that has a signature of:
            loss_value, gradients = f(model_parameters).
    """

    # obtain the shapes of all trainable parameters in the model
    shapes = tf.shape_n(model.trainable_variables)
    n_tensors = len(shapes)

    # we'll use tf.dynamic_stitch and tf.dynamic_partition later, so we need to
    # prepare required information first
    count = 0
    idx = [] # stitch indices
    part = [] # partition indices

    for i, shape in enumerate(shapes):
        n = numpy.product(shape)
        idx.append(tf.reshape(tf.range(count, count+n, dtype=tf.int32), shape))
        part.extend([i]*n)
        count += n

    part = tf.constant(part)

    @tf.function
    @tf.autograph.experimental.do_not_convert
    def assign_new_model_parameters(params_1d):
        """A function updating the model's parameters with a 1D tf.Tensor.

        Args:
            params_1d [in]: a 1D tf.Tensor representing the model's trainable parameters.
        """

        params = tf.dynamic_partition(params_1d, part, n_tensors)
        for i, (shape, param) in enumerate(zip(shapes, params)):
            model.trainable_variables[i].assign(tf.reshape(param, shape))
            #tf.print(model.trainable_variables[i])

    @tf.function
    def volume_form(x, Omega_Omegabar, mass, restriction):

        kahler_metric = complex_math.complex_hessian(tf.math.real(model(x)), x)
        volume_form = tf.math.real(tf.linalg.det(tf.matmul(restriction, tf.matmul(kahler_metric, restriction, adjoint_b=True))))
        weights = mass / tf.reduce_sum(mass)
        factor = tf.reduce_sum(weights * volume_form / Omega_Omegabar)
        #factor = tf.constant(35.1774, dtype=tf.complex64)
        return volume_form / factor


    # now create a function that will be returned by this factory
    def f(params_1d):
        """A function that can be used by tfp.optimizer.lbfgs_minimize.

        This function is created by function_factory.

        Args:
           params_1d [in]: a 1D tf.Tensor.

        Returns:
            A scalar loss and the gradients w.r.t. the `params_1d`.
        """

        # use GradientTape so that we can calculate the gradient of loss w.r.t. parameters
        for step, (points, Omega_Omegabar, mass, restriction) in enumerate(dataset):
            with tf.GradientTape() as tape:
                # update the parameters in the model
                assign_new_model_parameters(params_1d)
                # calculate the loss
                det_omega = volume_form(points, Omega_Omegabar, mass, restriction)
                loss_value = loss(Omega_Omegabar, det_omega, mass)

            # calculate gradients and convert to 1D tf.Tensor
            grads = tape.gradient(loss_value, model.trainable_variables)
            grads = tf.dynamic_stitch(idx, grads)
          
            # reweight the loss and grads 
            mass_sum = tf.reduce_sum(mass) 
            try:
                total_loss += loss_value * mass_sum 
                total_grads += grads * mass_sum
                total_mass += mass_sum
            except NameError:
                total_loss = loss_value * mass_sum 
                total_grads = grads * mass_sum
                total_mass = mass_sum

        total_loss = total_loss / total_mass
        total_grads = total_grads / total_mass

        # print out iteration & loss
        f.iter.assign_add(1)
        tf.print("Iter:", f.iter, "loss:", total_loss)

        # store loss value so we can retrieve later
        tf.py_function(f.history.append, inp=[total_loss], Tout=[])

        return total_loss, total_grads

    # store these information as members so we can use them outside the scope
    f.iter = tf.Variable(0)
    f.idx = idx
    f.part = part
    f.shapes = shapes
    f.assign_new_model_parameters = assign_new_model_parameters
    f.history = []

    return f
