# coding=utf-8
# Copyright 2019 The Edward2 Authors.
#
# 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.

"""Parses real and synthetic datasets.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import google_type_annotations
from __future__ import print_function

import collections
import tensorflow as tf

NPRegressionDescription = collections.namedtuple(
    "NPRegressionDescription",
    ("context_x", "context_y", "target_x", "target_y"))


class GPCurvesReader(object):
  """Generates curves using a Gaussian Process (GP).

  Supports vector inputs (x) and vector outputs (y). Kernel is
  mean-squared exponential, using the x-value l2 coordinate distance scaled by
  some factor chosen randomly in a range. Outputs are independent gaussian
  processes.
  """

  def __init__(self,
               batch_size,
               max_num_context,
               x_size=1,
               y_size=1,
               l1_scale=0.6,
               sigma_scale=1.0,
               random_kernel_parameters=False,
               testing=False):
    """Creates a regression dataset of functions sampled from a GP.

    Args:
      batch_size: An integer.
      max_num_context: The max number of observations in the context.
      x_size: Integer >= 1 for length of "x values" vector.
      y_size: Integer >= 1 for length of "y values" vector.
      l1_scale: Float; typical scale for kernel distance function.
      sigma_scale: Float; typical scale for variance.
      random_kernel_parameters: If `True`, the kernel parameters (l1 and sigma)
          are sampled uniformly within [0.1, l1_scale] and [0.1, sigma_scale].
      testing: Boolean that indicates whether we are testing. If so there are
          more targets for visualization.
    """
    self._batch_size = batch_size
    self._max_num_context = max_num_context
    self._x_size = x_size
    self._y_size = y_size
    self._l1_scale = l1_scale
    self._sigma_scale = sigma_scale
    self._random_kernel_parameters = random_kernel_parameters
    self._testing = testing

  def _gaussian_kernel(self, xdata, l1, sigma_f, sigma_noise=2e-2):
    """Applies the Gaussian kernel to generate curve data.

    Args:
      xdata: Tensor of shape [B, num_total_points, x_size] with
          the values of the x-axis data.
      l1: Tensor of shape [B, y_size, x_size], the scale
          parameter of the Gaussian kernel.
      sigma_f: Tensor of shape [B, y_size], the magnitude
          of the std.
      sigma_noise: Float, std of the noise that we add for stability.

    Returns:
      The kernel, a float tensor of shape
      [B, y_size, num_total_points, num_total_points].
    """
    num_total_points = tf.shape(xdata)[1]

    # Expand and take the difference
    xdata1 = tf.expand_dims(xdata, axis=1)  # [B, 1, num_total_points, x_size]
    xdata2 = tf.expand_dims(xdata, axis=2)  # [B, num_total_points, 1, x_size]
    diff = xdata1 - xdata2  # [B, num_total_points, num_total_points, x_size]

    # [B, y_size, num_total_points, num_total_points, x_size]
    norm = tf.square(diff[:, None, :, :, :] / l1[:, :, None, None, :])

    norm = tf.reduce_sum(
        norm, -1)  # [B, data_size, num_total_points, num_total_points]

    # [B, y_size, num_total_points, num_total_points]
    kernel = tf.square(sigma_f)[:, :, None, None] * tf.exp(-0.5 * norm)

    # Add some noise to the diagonal to make the cholesky work.
    kernel += (sigma_noise**2) * tf.eye(num_total_points)

    return kernel

  def generate_curves(self, num_context=None):
    """Builds the op delivering the data.

    Generated functions are `float32` with x values between -2 and 2.

    Args:
      num_context: Number of context points. If None, chosen randomly.

    Returns:
      A `CNPRegressionDescription` namedtuple.
    """
    if num_context is None:
      num_context = tf.random_uniform(
          shape=[], minval=3, maxval=self._max_num_context, dtype=tf.int32)

    # If we are testing we want to have more targets and have them evenly
    # distributed in order to plot the function.
    if self._testing:
      num_target = 400
      num_total_points = num_target
      x_values = tf.tile(
          tf.expand_dims(tf.range(-2., 2., 1. / 100, dtype=tf.float32), axis=0),
          [self._batch_size, 1])
      x_values = tf.expand_dims(x_values, axis=-1)
    # During training the number of target points and their x-positions are
    # selected at random
    else:
      num_target = tf.random_uniform(shape=(), minval=0,
                                     maxval=self._max_num_context - num_context,
                                     dtype=tf.int32)
      num_total_points = num_context + num_target
      x_values = tf.random_uniform(
          [self._batch_size, num_total_points, self._x_size], -2, 2)

    # Set kernel parameters
    # Either choose a set of random parameters for the mini-batch
    if self._random_kernel_parameters:
      l1 = tf.random_uniform([self._batch_size, self._y_size,
                              self._x_size], 0.1, self._l1_scale)
      sigma_f = tf.random_uniform([self._batch_size, self._y_size],
                                  0.1, self._sigma_scale)
    # Or use the same fixed parameters for all mini-batches
    else:
      l1 = tf.ones(shape=[self._batch_size, self._y_size,
                          self._x_size]) * self._l1_scale
      sigma_f = tf.ones(shape=[self._batch_size,
                               self._y_size]) * self._sigma_scale

    # Pass the x_values through the Gaussian kernel
    # [batch_size, y_size, num_total_points, num_total_points]
    kernel = self._gaussian_kernel(x_values, l1, sigma_f)

    # Calculate Cholesky, using double precision for better stability:
    cholesky = tf.cast(tf.cholesky(tf.cast(kernel, tf.float64)), tf.float32)

    # Sample a curve
    # [batch_size, y_size, num_total_points, 1]
    y_values = tf.matmul(
        cholesky,
        tf.random_normal([self._batch_size, self._y_size, num_total_points, 1]))

    # [batch_size, num_total_points, y_size]
    y_values = tf.transpose(tf.squeeze(y_values, 3), [0, 2, 1])

    if self._testing:
      # Select the targets
      target_x = x_values
      target_y = y_values

      # Select the observations
      idx = tf.random_shuffle(tf.range(num_target))
      context_x = tf.gather(x_values, idx[:num_context], axis=1)
      context_y = tf.gather(y_values, idx[:num_context], axis=1)

    else:
      # Select the targets which will consist of the context points as well as
      # some new target points
      target_x = x_values[:, :num_target + num_context, :]
      target_y = y_values[:, :num_target + num_context, :]

      # Select the observations
      context_x = x_values[:, :num_context, :]
      context_y = y_values[:, :num_context, :]

    return NPRegressionDescription(
        context_x=context_x,
        context_y=context_y,
        target_x=target_x,
        target_y=target_y)

[(27, 'collections.namedtuple', 'collections.namedtuple', False, 'import collections\n'), (92, 'tensorflow.expand_dims', 'tf.expand_dims', True, 'import tensorflow as tf\n'), (93, 'tensorflow.expand_dims', 'tf.expand_dims', True, 'import tensorflow as tf\n'), (97, 'tensorflow.square', 'tf.square', True, 'import tensorflow as tf\n'), (99, 'tensorflow.reduce_sum', 'tf.reduce_sum', True, 'import tensorflow as tf\n'), (89, 'tensorflow.shape', 'tf.shape', True, 'import tensorflow as tf\n'), (103, 'tensorflow.exp', 'tf.exp', True, 'import tensorflow as tf\n'), (106, 'tensorflow.eye', 'tf.eye', True, 'import tensorflow as tf\n'), (122, 'tensorflow.random_uniform', 'tf.random_uniform', True, 'import tensorflow as tf\n'), (133, 'tensorflow.expand_dims', 'tf.expand_dims', True, 'import tensorflow as tf\n'), (137, 'tensorflow.random_uniform', 'tf.random_uniform', True, 'import tensorflow as tf\n'), (141, 'tensorflow.random_uniform', 'tf.random_uniform', True, 'import tensorflow as tf\n'), (147, 'tensorflow.random_uniform', 'tf.random_uniform', True, 'import tensorflow as tf\n'), (149, 'tensorflow.random_uniform', 'tf.random_uniform', True, 'import tensorflow as tf\n'), (169, 'tensorflow.random_normal', 'tf.random_normal', True, 'import tensorflow as tf\n'), (172, 'tensorflow.squeeze', 'tf.squeeze', True, 'import tensorflow as tf\n'), (181, 'tensorflow.gather', 'tf.gather', True, 'import tensorflow as tf\n'), (182, 'tensorflow.gather', 'tf.gather', True, 'import tensorflow as tf\n'), (103, 'tensorflow.square', 'tf.square', True, 'import tensorflow as tf\n'), (153, 'tensorflow.ones', 'tf.ones', True, 'import tensorflow as tf\n'), (155, 'tensorflow.ones', 'tf.ones', True, 'import tensorflow as tf\n'), (163, 'tensorflow.cast', 'tf.cast', True, 'import tensorflow as tf\n'), (180, 'tensorflow.range', 'tf.range', True, 'import tensorflow as tf\n'), (131, 'tensorflow.range', 'tf.range', True, 'import tensorflow as tf\n')]