# 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)

[((871, 976), 'collections.namedtuple', 'collections.namedtuple', (['"""NPRegressionDescription"""', "('context_x', 'context_y', 'target_x', 'target_y')"], {}), "('NPRegressionDescription', ('context_x', 'context_y',\n    'target_x', 'target_y'))\n", (893, 976), False, 'import collections\n'), ((3274, 3303), 'tensorflow.expand_dims', 'tf.expand_dims', (['xdata'], {'axis': '(1)'}), '(xdata, axis=1)\n', (3288, 3303), True, 'import tensorflow as tf\n'), ((3353, 3382), 'tensorflow.expand_dims', 'tf.expand_dims', (['xdata'], {'axis': '(2)'}), '(xdata, axis=2)\n', (3367, 3382), True, 'import tensorflow as tf\n'), ((3571, 3636), 'tensorflow.square', 'tf.square', (['(diff[:, (None), :, :, :] / l1[:, :, (None), (None), :])'], {}), '(diff[:, (None), :, :, :] / l1[:, :, (None), (None), :])\n', (3580, 3636), True, 'import tensorflow as tf\n'), ((3643, 3666), 'tensorflow.reduce_sum', 'tf.reduce_sum', (['norm', '(-1)'], {}), '(norm, -1)\n', (3656, 3666), True, 'import tensorflow as tf\n'), ((3204, 3219), 'tensorflow.shape', 'tf.shape', (['xdata'], {}), '(xdata)\n', (3212, 3219), True, 'import tensorflow as tf\n'), ((3837, 3856), 'tensorflow.exp', 'tf.exp', (['(-0.5 * norm)'], {}), '(-0.5 * norm)\n', (3843, 3856), True, 'import tensorflow as tf\n'), ((3955, 3979), 'tensorflow.eye', 'tf.eye', (['num_total_points'], {}), '(num_total_points)\n', (3961, 3979), True, 'import tensorflow as tf\n'), ((4359, 4447), 'tensorflow.random_uniform', 'tf.random_uniform', ([], {'shape': '[]', 'minval': '(3)', 'maxval': 'self._max_num_context', 'dtype': 'tf.int32'}), '(shape=[], minval=3, maxval=self._max_num_context, dtype=\n    tf.int32)\n', (4376, 4447), True, 'import tensorflow as tf\n'), ((4816, 4849), 'tensorflow.expand_dims', 'tf.expand_dims', (['x_values'], {'axis': '(-1)'}), '(x_values, axis=-1)\n', (4830, 4849), True, 'import tensorflow as tf\n'), ((4980, 5081), 'tensorflow.random_uniform', 'tf.random_uniform', ([], {'shape': '()', 'minval': '(0)', 'maxval': '(self._max_num_context - num_context)', 'dtype': 'tf.int32'}), '(shape=(), minval=0, maxval=self._max_num_context -\n    num_context, dtype=tf.int32)\n', (4997, 5081), True, 'import tensorflow as tf\n'), ((5219, 5295), 'tensorflow.random_uniform', 'tf.random_uniform', (['[self._batch_size, num_total_points, self._x_size]', '(-2)', '(2)'], {}), '([self._batch_size, num_total_points, self._x_size], -2, 2)\n', (5236, 5295), True, 'import tensorflow as tf\n'), ((5452, 5543), 'tensorflow.random_uniform', 'tf.random_uniform', (['[self._batch_size, self._y_size, self._x_size]', '(0.1)', 'self._l1_scale'], {}), '([self._batch_size, self._y_size, self._x_size], 0.1, self\n    ._l1_scale)\n', (5469, 5543), True, 'import tensorflow as tf\n'), ((5585, 5660), 'tensorflow.random_uniform', 'tf.random_uniform', (['[self._batch_size, self._y_size]', '(0.1)', 'self._sigma_scale'], {}), '([self._batch_size, self._y_size], 0.1, self._sigma_scale)\n', (5602, 5660), True, 'import tensorflow as tf\n'), ((6441, 6512), 'tensorflow.random_normal', 'tf.random_normal', (['[self._batch_size, self._y_size, num_total_points, 1]'], {}), '([self._batch_size, self._y_size, num_total_points, 1])\n', (6457, 6512), True, 'import tensorflow as tf\n'), ((6588, 6611), 'tensorflow.squeeze', 'tf.squeeze', (['y_values', '(3)'], {}), '(y_values, 3)\n', (6598, 6611), True, 'import tensorflow as tf\n'), ((6829, 6875), 'tensorflow.gather', 'tf.gather', (['x_values', 'idx[:num_context]'], {'axis': '(1)'}), '(x_values, idx[:num_context], axis=1)\n', (6838, 6875), True, 'import tensorflow as tf\n'), ((6894, 6940), 'tensorflow.gather', 'tf.gather', (['y_values', 'idx[:num_context]'], {'axis': '(1)'}), '(y_values, idx[:num_context], axis=1)\n', (6903, 6940), True, 'import tensorflow as tf\n'), ((3798, 3816), 'tensorflow.square', 'tf.square', (['sigma_f'], {}), '(sigma_f)\n', (3807, 3816), True, 'import tensorflow as tf\n'), ((5776, 5837), 'tensorflow.ones', 'tf.ones', ([], {'shape': '[self._batch_size, self._y_size, self._x_size]'}), '(shape=[self._batch_size, self._y_size, self._x_size])\n', (5783, 5837), True, 'import tensorflow as tf\n'), ((5897, 5944), 'tensorflow.ones', 'tf.ones', ([], {'shape': '[self._batch_size, self._y_size]'}), '(shape=[self._batch_size, self._y_size])\n', (5904, 5944), True, 'import tensorflow as tf\n'), ((6277, 6304), 'tensorflow.cast', 'tf.cast', (['kernel', 'tf.float64'], {}), '(kernel, tf.float64)\n', (6284, 6304), True, 'import tensorflow as tf\n'), ((6789, 6809), 'tensorflow.range', 'tf.range', (['num_target'], {}), '(num_target)\n', (6797, 6809), True, 'import tensorflow as tf\n'), ((4710, 4758), 'tensorflow.range', 'tf.range', (['(-2.0)', '(2.0)', '(1.0 / 100)'], {'dtype': 'tf.float32'}), '(-2.0, 2.0, 1.0 / 100, dtype=tf.float32)\n', (4718, 4758), True, 'import tensorflow as tf\n')]