"""
Deep Dreaming for time series in Keras.
This code is roughly based on https://github.com/keras-team/keras/blob/master/examples/deep_dream.py,
but has been heavily modified by <NAME> (LoadingByte).

In this file:
- u_timeser: A single univariate time series in this format: [a0,a1,a2]
- m_timeser: A single multivariate time series in this format: [[a0,b0], [a1,b1], [a2,b2]]
- b_timeser: A batch of multivariate time series in this format: [[[a0,b0],[a1,b1]], [[c0,d0],[c1,d1]]]

Process:
- Load the base time series.
- Define a number of processing scales (i.e. time series lengths), from smallest to largest.
- Resize the base time series to the smallest scale.
- For every scale, starting with the smallest (i.e. current one):
    - Run gradient ascent.
    - Upscale time series to the next scale.
    - Reinject the detail that was lost at upscaling time.
- Stop when we are back to the original size.

To obtain the detail lost during upscaling, we simply take the original time series, shrink it down, upscale it,
and compare the result to the (resized) base time series.
"""

from collections import Mapping
from numbers import Number

import numpy as np
from keras import backend as K

DEFAULT_HYPERPARAMS = {
    "step": 0.01,
    "num_octave": 3,
    "octave_scale": 1.4,
    "iterations": 20,
    "max_loss": 10.0
}


def dream(model, loss_layers, base_umb_timeser,
          step=DEFAULT_HYPERPARAMS["step"],
          num_octave=DEFAULT_HYPERPARAMS["num_octave"],
          octave_scale=DEFAULT_HYPERPARAMS["octave_scale"],
          iterations=DEFAULT_HYPERPARAMS["iterations"],
          max_loss=DEFAULT_HYPERPARAMS["max_loss"],
          intermediates=False,
          fetch_loss_and_grads=None, backend_kwargs=None,
          verbose=False):
    """
    Lets a network manipulate the given base time series such that L2 activation in the given loss layers is maximized.

    Manipulation of the base series is done through gradient ascent with respect to a loss function, which is defined
    by the given loss layers.

    Beware of memory leaks! If and only if fetch_loss_and_grads is not provided, the underlying graph of the model
    will be modified by this function. Thus, when dreaming multiple times, it is recommended to generate
    fetch_loss_and_grads in advance using instrument_model() and provide the result to this function each time you
    dream. Otherwise, multiple instrumentations will be generated, bloating the graph unnecessarily.

    Args:
        model: The Keras neural network model that shall dream.
        loss_layers: A dict containing the names of the layers (key) whose L2 activation we try to maximize,
            as well as the layer output's weight (value) in the final "loss" we try to maximize.
            The weight value can be one of these:
            - In case of a scalar, the output of each neuron in the layer is weighed with the scalar divided by
              the number of neurons in the layer.
            - In case of a list, each neuron is weighed with the element of the list at the corresponding index.
            - In case of a dict, the dict is converted to a list such that the keys from the dict are indices into
              the list. We then continue with the list case.
        base_umb_timeser: The base time series batch which the model manipulates to maximize L2 activation;
            must be given in one of the following formats:
            - Univariate - [a0,a1,a2]
            - Multivariate - [[a0,b0], [a1,b1], [a2,b2]]
            - Batch - [[[a0,b0],[a1,b1]], [[c0,d0],[c1,d1]]]
            Note that the dream time series will be returned in exactly the same format.
        step: Gradient ascent step size.
        num_octave: Number of consecutive scales at which to run gradient ascent. We first scale the base series
            down to the lowest scale and run gradient ascent there. Next, we scale the resulting series up to the
            next higher scale and run gradient ascent again. We continue until we reach the highest scale, which is
            just the original size of the time series.
        octave_scale: Size ratio between scales.
        iterations: Number of gradient ascent steps per scale.
        max_loss: When, during a scale, we reach a loss this high, the rest of the scale is skipped and
            we proceed with the subsequent scale.
        intermediates: When true, a list of all intermediate dreams is returned, instead of only the final dream.
        fetch_loss_and_grads: Optional instrument function previously generated by instrument_model().
            If provided, the dreaming process will not need to alter the model graph in any way.
        backend_kwargs: Passed to keras.backend.function and thus to tf.Session.run.
            If fetch_loss_and_grads is provided, this parameter is ignored!
        verbose: Whether to print what the algorithm is doing.

    Returns:
        The dream time series, in the same format (univariate/multivariate/batch) as the base time series.

    """

    K.set_learning_phase(0)

    # If not already provided, integrate fetch_loss_and_grads into the model.
    if fetch_loss_and_grads is None:
        fetch_loss_and_grads = instrument_model(model, backend_kwargs)

    # Convert the loss layers dict provided by the user to weight tensors that can be used by fetch_loss_and_grads.
    weights_tensors = _get_weight_tensors(model, loss_layers)

    # Determine the base format of the time series and
    # convert univariate and multivariate time series to batch time series.
    if not 1 <= np.ndim(base_umb_timeser) <= 3:
        raise Exception("Invalid time series dimension; only 1 (univariate), 2 (multivariate), "
                        f"or 3 (multivariate batch) are allowed, but not: {base_umb_timeser}")
    base_timeser_format = {1: "u", 2: "m", 3: "b"}[np.ndim(base_umb_timeser)]
    if base_timeser_format == "u":
        base_b_timeser = np.reshape(base_umb_timeser, (1, -1, 1))
    elif base_timeser_format == "m":
        base_b_timeser = np.expand_dims(base_umb_timeser, axis=0)
    elif base_timeser_format == "b":
        base_b_timeser = np.array(base_umb_timeser)

    original_b_timeser = np.copy(base_b_timeser)
    dream_b_timesers = [np.copy(base_b_timeser)]

    # Make a list of shapes (= time series sizes).
    # The algorithm will start dreaming on the smallest shape and
    # then gradually work its way upwards to larger shapes.
    original_shape = original_b_timeser.shape[1]
    successive_shapes = [original_shape]
    for i in range(1, num_octave):
        shape = int(original_shape / (octave_scale ** i))
        successive_shapes.append(shape)
    successive_shapes = successive_shapes[::-1]  # reverse array

    # This is where the actual dreaming process begins.
    shrunk_original_b_timeser = _resize_b_timeser(original_b_timeser, successive_shapes[0])
    for shape in successive_shapes:
        if verbose:
            print("Processing time series shape", shape)

        # Upscale the dream time series to the current shape.
        dream_b_timesers.append(_resize_b_timeser(dream_b_timesers[-1], shape))

        # Run gradient ascent on the dream time series, doing the actual dreaming.
        _gradient_ascent(fetch_loss_and_grads, weights_tensors, iterations, step, max_loss, verbose, dream_b_timesers)

        # Calculate the detail lost in downscaling.
        upscaled_shrunk_original_b_timeser = _resize_b_timeser(shrunk_original_b_timeser, shape)
        same_size_original = _resize_b_timeser(original_b_timeser, shape)
        lost_detail = same_size_original - upscaled_shrunk_original_b_timeser

        # Add that detail back to the upscaled dream time series.
        dream_b_timesers.append(dream_b_timesers[-1] + lost_detail)
        shrunk_original_b_timeser = _resize_b_timeser(original_b_timeser, shape)

    # Convert the batch dream time series back to the format of the base time series.
    if base_timeser_format == "u":
        dream_umb_timesers = [t[0].reshape(-1, order="F") for t in dream_b_timesers]
    elif base_timeser_format == "m":
        dream_umb_timesers = [t[0] for t in dream_b_timesers]
    elif base_timeser_format == "b":
        dream_umb_timesers = dream_b_timesers

    return dream_umb_timesers if intermediates else dream_umb_timesers[-1]


def instrument_model(model, backend_kwargs=None):
    """
    Extends the graph of the given model with a custom function that computes a single gradient ascent step.

    The custom function takes the current dream time series and one tensor of weights for each layer in the model.
    It then computes the current loss and a gradient ascent step aiming at maximizing the L2 activation of neurons
    as defined by the weight tensors.

    Args:
        model: The Keras neural network model that shall be instrumented.
        backend_kwargs: Passed to keras.backend.function and thus to tf.Session.run.

    Returns:
        The fetch_loss_and_grads function that is required to do the actual dreaming.
        It has to be passed to dream() later on.

    """
    # The model's input time series is what we want to manipulate.
    dream = model.input

    # For each layer, declare one tensor that holds the weights for each feature (= output neuron) of the layer.
    weights_tensors = [K.placeholder(_weights_shape(layer), name=layer.name + "_dream_weights")
                       for layer in model.layers]

    # Define the loss.
    loss = K.variable(0.0)
    for layer, weights in zip(model.layers, weights_tensors):
        x = layer.output

        # Expand weights into "None" (= unknown) dimensions
        true_shape = tuple(1 if dim is None else dim for dim in K.int_shape(x))
        true_shape_weights = K.reshape(weights, true_shape)

        # Quick'n'dirtily scale down the weights tensor by just slicing it to the layer's shape ("gathering").
        # This is required for octaves to work because they obviously change the size of the layers.
        # When looking at this, it becomes clear that you should not use single-neuron weights on
        # any layer except those which do not resize (mainly the last one).
        effective_weights = K.slice(true_shape_weights, start=np.zeros(len(true_shape), dtype="int32"), size=K.shape(x))

        # Add the L2-norm of the features of a layer to the loss.
        loss = loss + K.sum(effective_weights * K.square(x))

    # Compute the gradients of the dream with respect to the loss.
    grads = K.gradients(loss, dream)[0]
    # Normalize the gradients.
    grads /= K.maximum(K.mean(K.abs(grads)), K.epsilon())

    # Set up a function to retrieve the value of the loss as well as gradients given an input time series.
    inputs = [dream, *weights_tensors]
    outputs = [loss, grads]
    if backend_kwargs is None:
        fetch_loss_and_grads = K.function(inputs, outputs)
    else:
        fetch_loss_and_grads = K.function(inputs, outputs, **backend_kwargs)

    return fetch_loss_and_grads


def _get_weight_tensors(model, loss_layers):
    def _get_single_weight_tensor(layer):
        weights_shape = _weights_shape(layer)

        if layer.name not in loss_layers:
            return np.zeros(weights_shape)
        user_weights = loss_layers[layer.name]

        if isinstance(user_weights, Number):
            scaling = np.prod(weights_shape)
            uniform_weight = user_weights / scaling
            return np.full(weights_shape, uniform_weight)
        elif isinstance(user_weights, Mapping):
            def _weight_for_coord(coord):
                if tuple(coord) in user_weights:
                    return user_weights[tuple(coord)]
                elif len(coord) == 1 and coord[0] in user_weights:
                    return user_weights[coord[0]]
                else:
                    return 0

            def _func(*dims):
                return np.apply_along_axis(_weight_for_coord, axis=0, arr=np.stack(dims, axis=0))

            return np.fromfunction(_func, weights_shape, dtype=int)
        elif np.iterable(user_weights):
            if np.shape(user_weights) != weights_shape:
                raise ValueError(f"For layer {layer.name}: Shape of weight array {np.shape(user_weights)} must match"
                                 f"output shape of layer {weights_shape}")
            return np.array(user_weights)
        else:
            raise ValueError(f"Weight of layer {layer.name} must be a number, a dict, or an iterable,"
                             f"not {type(user_weights)}")

    return [_get_single_weight_tensor(layer) for layer in model.layers]


def _weights_shape(layer):
    return tuple(comp for comp in layer.output_shape if comp is not None)


# Hint: "size" refers to the number of points in time.
def _resize_b_timeser(b_timeser, new_size):
    old_size = b_timeser.shape[1]

    def resize_u_timeser(u_timeser):
        return np.interp(np.linspace(0, old_size - 1, num=new_size), np.arange(old_size), u_timeser)

    def resize_m_timeser(m_timeser):
        # Convert from [[a0,b0],[a1,b1],[a2,b2]] to [[a0,a1,a2],[b0,b1,b2]]
        reshaped_m_timeser = m_timeser.reshape(-1, order="F").reshape((-1, old_size))
        # Resize each univariate time series individually.
        resized_reshaped_m_timeser = np.array([resize_u_timeser(u_timeser) for u_timeser in reshaped_m_timeser])
        # Reverse the conversion made in the first statement.
        resized_m_timeser = resized_reshaped_m_timeser.reshape(-1).reshape((new_size, -1), order="F")
        return resized_m_timeser

    return np.array([resize_m_timeser(m_timeser) for m_timeser in b_timeser])


def _gradient_ascent(fetch_loss_and_grads, weights_tensors, iterations, step, max_loss, verbose,
                     dream_b_timesers):
    for i in range(iterations):
        loss_value, grad_values = fetch_loss_and_grads([dream_b_timesers[-1], *weights_tensors])

        if max_loss is not None and loss_value > max_loss:
            break
        if verbose:
            print(f"..Loss value at iteration {i}: {loss_value}")

        dream_b_timesers.append(dream_b_timesers[-1] + (step * grad_values))
