"""
G R A D I E N T - E N H A N C E D   N E U R A L   N E T W O R K S  (G E N N)

Author: Steven H. Berguin <steven.berguin@gtri.gatech.edu>

This package is distributed under New BSD license.
"""

from smt.surrogate_models.surrogate_model import SurrogateModel
from smt.utils.neural_net.model import Model

import numpy as np


# ------------------------------------ S U P P O R T   F U N C T I O N S -----------------------------------------------


def load_smt_data(model, xt, yt, dyt_dxt=None):
    """
    Utility function to load SMT data more easily

    :param model: SurrogateModel object for which to load training data
    :param xt: smt data points at which response is evaluated
    :param yt: response at xt
    :param dyt_dxt: gradient at xt
    """
    # Dimensionality
    if len(xt.shape) == 1:
        n_x = 1  # number of variables, x
        m = xt.size
    else:
        m, n_x = xt.shape

    if len(yt.shape) == 1:
        n_y = 1  # number of responses, y
    else:
        n_y = yt.shape[1]

    # Reshape arrays
    xt = xt.reshape((m, n_x))
    yt = yt.reshape((m, n_y))

    # Load values
    model.set_training_values(xt, yt)

    # Load partials
    if dyt_dxt is not None:
        dyt_dxt = dyt_dxt.reshape((m, n_x))
        for i in range(n_x):
            model.set_training_derivatives(xt, dyt_dxt[:, i].reshape((m, 1)), i)


def smt_to_genn(training_points):
    """
    Translate from SMT data structure to GENN data structure.

    Concretely, this neural net module works with numpy arrays in the form of (X, Y, J) as defined here-under. However,
    SMT uses a different format. Hence, we need a function that takes care of the translation.

    :param: training_points -- dict, training data in the format used by surrogate_model.py (see SMT API)

    Returns:
    :return X -- a numpy matrix of input features of shape (n_x, m) where n_x = no. of inputs, m = no. of train examples
    :return Y -- a numpy matrix of output labels of shape (n_y, m) where n_y = no. of outputs
    :return J -- a numpy array of size (n_y, n_x, m) representing the Jacobian of Y w.r.t. X:

        dY1/dX1 = J[0][0][:]
        dY1/dX2 = J[0][1][:]
        ...
        dY2/dX1 = J[1][0][:]
        dY2/dX2 = J[1][1][:]
        ...

        N.B. To retrieve the i^th example for dY2/dX1: J[1][0][i] for all i = 1,...,m

    """
    # Retrieve training data from SMT training_points
    xt, yt = training_points[None][
        0
    ]  # training_points[name][0] = [np.array(xt), np.array(yt)]

    # Deduce number of dimensions and training examples
    m, n_x = xt.shape
    _, n_y = yt.shape

    # Assign training data but transpose to match neural net implementation
    X = xt
    Y = yt

    # Loop to retrieve each partial derivative from SMT training_points
    J = np.zeros((m, n_x, n_y))
    for k in range(0, n_x):
        xt, dyt_dxt = training_points[None][k + 1]

        # assert that dimensions match
        assert xt.shape[0] == m
        assert xt.shape[1] == n_x
        assert dyt_dxt.shape[0] == m
        assert dyt_dxt.shape[1] == n_y

        # Assert that derivatives provided are for the same training points
        assert xt.all() == X.all()

        # Assign training derivative but transpose to match neural net implementation
        J[:, k, :] = dyt_dxt

    return X.T, Y.T, J.T


# ------------------------------------ C L A S S -----------------------------------------------------------------------


class GENN(SurrogateModel):
    def _initialize(self):
        """API function: set default values for user options"""
        declare = self.options.declare
        declare("alpha", 0.5, types=(int, float), desc="optimizer learning rate")
        declare(
            "beta1", 0.9, types=(int, float), desc="Adam optimizer tuning parameter"
        )
        declare(
            "beta2", 0.99, types=(int, float), desc="Adam optimizer tuning parameter"
        )
        declare("lambd", 0.1, types=(int, float), desc="regularization coefficient")
        declare(
            "gamma", 1.0, types=(int, float), desc="gradient-enhancement coefficient"
        )
        declare("deep", 2, types=int, desc="number of hidden layers")
        declare("wide", 2, types=int, desc="number of nodes per hidden layer")
        declare(
            "mini_batch_size",
            64,
            types=int,
            desc="split data into batches of specified size",
        )
        declare(
            "num_epochs", 10, types=int, desc="number of random passes through the data"
        )
        declare(
            "num_iterations",
            100,
            types=int,
            desc="number of optimizer iterations per mini-batch",
        )
        declare(
            "seed",
            None,
            types=int,
            desc="random seed to ensure repeatability of results when desired",
        )
        declare("is_print", True, types=bool, desc="print progress (or not)")

        self.supports["derivatives"] = True
        self.supports["training_derivatives"] = True
        self.name = "GENN"

        self.model = Model()

        self._is_trained = False