# Copyright 2018 The TensorFlow Probability 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.
# ============================================================================
"""The GaussianProcessRegressionModel distribution class."""

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

# Dependency imports
import tensorflow.compat.v2 as tf

from tensorflow_probability.python import positive_semidefinite_kernels as tfpk
from tensorflow_probability.python.distributions import gaussian_process
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import tensorshape_util

__all__ = [
    'GaussianProcessRegressionModel',
]

def _add_diagonal_shift(matrix, shift):
  diag_plus_shift = tf.linalg.diag_part(matrix) + shift
  return tf.linalg.set_diag(matrix, diag_plus_shift)

def _is_empty_observation_data(
    feature_ndims, observation_index_points, observations):
  """Returns `True` if given observation data is empty.

  Emptiness means either
    1. Both `observation_index_points` and `observations` are `None`, or
    2. the "number of observations" shape is 0. The shape of
    `observation_index_points` is `[..., N, f1, ..., fF]`, where `N` is the
    number of observations and the `f`s are feature dims. Thus, we look at the
    shape element just to the left of the leftmost feature dim. If that shape is
    zero, we consider the data empty.

  We don't check the shape of observations; validations are checked elsewhere in
  the calling code, to ensure these shapes are consistent.

  Args:
    feature_ndims: the number of feature dims, as reported by the GP kernel.
    observation_index_points: the observation data locations in the index set.
    observations: the observation data.

  Returns:
    is_empty: True if the data were deemed to be empty.
  """
  # If both input locations and observations are `None`, we consider this
  # "empty" observation data.
  if observation_index_points is None and observations is None:
    return True
  num_obs = tf.compat.dimension_value(
      observation_index_points.shape[-(feature_ndims + 1)])
  if num_obs is not None and num_obs == 0:
    return True
  return False

def _validate_observation_data(
    kernel, observation_index_points, observations):
  """Ensure that observation data and locations have consistent shapes.

  This basically means that the batch shapes are broadcastable. We can only
  ensure this when those shapes are fully statically defined.

  Args:
    kernel: The GP kernel.
    observation_index_points: the observation data locations in the index set.
    observations: the observation data.

  Raises:
    ValueError: if the observations' batch shapes are not broadcastable.
  """
  # Check that observation index points and observation counts broadcast.
  ndims = kernel.feature_ndims
  if (tensorshape_util.is_fully_defined(
      observation_index_points.shape[:-ndims]) and
      tensorshape_util.is_fully_defined(observations.shape)):
    index_point_count = observation_index_points.shape[:-ndims]
    observation_count = observations.shape
    try:
      tf.broadcast_static_shape(index_point_count, observation_count)
    except ValueError:
      # Re-raise with our own more contextual error message.
      raise ValueError(
          'Observation index point and observation counts are not '
          'broadcastable: {} and {}, respectively.'.format(
              index_point_count, observation_count))

class GaussianProcessRegressionModel(gaussian_process.GaussianProcess):
  """Posterior predictive distribution in a conjugate GP regression model.

  This class represents the distribution over function values at a set of points
  in some index set, conditioned on noisy observations at some other set of
  points. More specifically, we assume a Gaussian process prior, `f ~ GP(m, k)`
  with IID normal noise on observations of function values. In this model
  posterior inference can be done analytically. This `Distribution` is
  parameterized by

    * the mean and covariance functions of the GP prior,
    * the set of (noisy) observations and index points to which they correspond,
    * the set of index points at which the resulting posterior predictive
      distribution over function values is defined,
    * the observation noise variance,
    * jitter, to compensate for numerical instability of Cholesky decomposition,

  in addition to the usual params like `validate_args` and `allow_nan_stats`.

  #### Mathematical Details

  Gaussian process regression (GPR) assumes a Gaussian process (GP) prior and a
  normal likelihood as a generative model for data. Given GP mean function `m`,
  covariance kernel `k`, and observation noise variance `v`, we have

  ```none
    f ~ GP(m, k)

                     iid
    (y[i] | f, x[i])  ~  Normal(f(x[i]), v),   i = 1, ... , N
  ```

  where `y[i]` are the noisy observations of function values at points `x[i]`.

  In practice, `f` is an infinite object (eg, a function over `R^n`) which can't
  be realized on a finite machine, but fortunately the marginal distribution
  over function values at a finite set of points is just a multivariate normal
  with mean and covariance given by the mean and covariance functions applied at
  our finite set of points (see [Rasmussen and Williams, 2006][1] for a more
  extensive discussion of these facts).

  We spell out the generative model in detail below, but first, a digression on
  notation. In what follows we drop the indices on vectorial objects such as
  `x[i]`, it being implied that we are generally considering finite collections
  of index points and corresponding function values and noisy observations
  thereof. Thus `x` should be considered to stand for a collection of index
  points (indeed, themselves often vectorial). Furthermore:

    * `f(x)` refers to the collection of function values at the index points in
      the collection `x`",
    * `m(t)` refers to the collection of values of the mean function at the
      index points in the collection `t`, and
    * `k(x, t)` refers to the *matrix* whose entries are values of the kernel
      function `k` at all pairs of index points from `x` and `t`.

  With these conventions in place, we may write

  ```none
    (f(x) | x) ~ MVN(m(x), k(x, x))

    (y | f(x), x) ~ Normal(f(x), v)
  ```

  When we condition on observed data `y` at the points `x`, we can derive the
  posterior distribution over function values `f(x)` at those points. We can
  then compute the posterior predictive distribution over function values `f(t)`
  at a new set of points `t`, conditional on those observed data.

  ```none
    (f(t) | t, x, f(x)) ~ MVN(loc, cov)

    where

    loc = k(t, x) @ inv(k(x, x) + v * I) @ (y - loc)
    cov = k(t, t) - k(t, x) @ inv(k(x, x) + v * I) @ k(x, t)
  ```

  where `I` is the identity matrix of appropriate dimension. Finally, the
  distribution over noisy observations at the new set of points `t` is obtained
  by adding IID noise from `Normal(0., observation_noise_variance)`.

  #### Examples

  ##### Draw joint samples from the posterior predictive distribution in a GP
  regression model

  ```python
  import numpy as np
  import tensorflow as tf
  import tensorflow_probability as tfp

  tfd = tfp.distributions
  psd_kernels = tfp.positive_semidefinite_kernels

  # Generate noisy observations from a known function at some random points.
  observation_noise_variance = .5
  f = lambda x: np.sin(10*x[..., 0]) * np.exp(-x[..., 0]**2)
  observation_index_points = np.random.uniform(-1., 1., 50)[..., np.newaxis]
  observations = (f(observation_index_points) +
                  np.random.normal(0., np.sqrt(observation_noise_variance)))

  index_points = np.linspace(-1., 1., 100)[..., np.newaxis]

  kernel = psd_kernels.MaternFiveHalves()

  gprm = tfd.GaussianProcessRegressionModel(
      kernel=kernel,
      index_points=index_points,
      observation_index_points=observation_index_points,
      observations=observations,
      observation_noise_variance=observation_noise_variance)

  samples = gprm.sample(10)
  # ==> 10 independently drawn, joint samples at `index_points`.
  ```

  Above, we have used the kernel with default parameters, which are unlikely to
  be good. Instead, we can train the kernel hyperparameters on the data, as in
  the next example.

  ##### Optimize model parameters via maximum marginal likelihood

  Here we learn the kernel parameters as well as the observation noise variance
  using gradient descent on the maximum marginal likelihood.

  ```python
  # Suppose we have some data from a known function. Note the index points in
  # general have shape `[b1, ..., bB, f1, ..., fF]` (here we assume `F == 1`),
  # so we need to explicitly consume the feature dimensions (just the last one
  # here).
  f = lambda x: np.sin(10*x[..., 0]) * np.exp(-x[..., 0]**2)

  observation_index_points = np.random.uniform(-1., 1., 50)[..., np.newaxis]
  observations = f(observation_index_points) + np.random.normal(0., .05, 50)

  # Define a kernel with trainable parameters. Note we transform the trainable
  # variables to apply a positivity constraint.
  amplitude = tf.exp(tf.Variable(np.float64(0)), name='amplitude')
  length_scale = tf.exp(tf.Variable(np.float64(0)), name='length_scale')
  kernel = psd_kernels.ExponentiatedQuadratic(amplitude, length_scale)

  observation_noise_variance = tf.exp(
      tf.Variable(np.float64(-5)), name='observation_noise_variance')

  # We'll use an unconditioned GP to train the kernel parameters.
  gp = tfd.GaussianProcess(
      kernel=kernel,
      index_points=observation_index_points,
      observation_noise_variance=observation_noise_variance)
  neg_log_likelihood = -gp.log_prob(observations)

  optimizer = tf.train.AdamOptimizer(learning_rate=.05, beta1=.5, beta2=.99)
  optimize = optimizer.minimize(neg_log_likelihood)

  # We can construct the posterior at a new set of `index_points` using the same
  # kernel (with the same parameters, which we'll optimize below).
  index_points = np.linspace(-1., 1., 100)[..., np.newaxis]
  gprm = tfd.GaussianProcessRegressionModel(
      kernel=kernel,
      index_points=index_points,
      observation_index_points=observation_index_points,
      observations=observations,
      observation_noise_variance=observation_noise_variance)

  samples = gprm.sample(10)
  # ==> 10 independently drawn, joint samples at `index_points`.

  # Now execute the above ops in a Session, first training the model
  # parameters, then drawing and plotting posterior samples.
  with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())

    for i in range(1000):
      _, neg_log_likelihood_ = sess.run([optimize, neg_log_likelihood])
      if i % 100 == 0:
        print("Step {}: NLL = {}".format(i, neg_log_likelihood_))

    print("Final NLL = {}".format(neg_log_likelihood_))
    samples_ = sess.run(samples)

    plt.scatter(np.squeeze(observation_index_points), observations)
    plt.plot(np.stack([index_points[:, 0]]*10).T, samples_.T, c='r', alpha=.2)
  ```

  ##### Marginalization of model hyperparameters

  Here we use TensorFlow Probability's MCMC functionality to perform
  marginalization of the model hyperparameters: kernel params as well as
  observation noise variance.

  ```python
  f = lambda x: np.sin(10*x[..., 0]) * np.exp(-x[..., 0]**2)
  observation_index_points = np.random.uniform(-1., 1., 25)[..., np.newaxis]
  observations = np.random.normal(f(observation_index_points), .05)

  def joint_log_prob(
      index_points, observations, amplitude, length_scale, noise_variance):

    # Hyperparameter Distributions.
    rv_amplitude = tfd.LogNormal(np.float64(0.), np.float64(1))
    rv_length_scale = tfd.LogNormal(np.float64(0.), np.float64(1))
    rv_noise_variance = tfd.LogNormal(np.float64(0.), np.float64(1))

    gp = tfd.GaussianProcess(
        kernel=psd_kernels.ExponentiatedQuadratic(amplitude, length_scale),
        index_points=index_points,
        observation_noise_variance=noise_variance)

    return (
        rv_amplitude.log_prob(amplitude) +
        rv_length_scale.log_prob(length_scale) +
        rv_noise_variance.log_prob(noise_variance) +
        gp.log_prob(observations)
    )

  initial_chain_states = [
      1e-1 * tf.ones([], dtype=np.float64, name='init_amplitude'),
      1e-1 * tf.ones([], dtype=np.float64, name='init_length_scale'),
      1e-1 * tf.ones([], dtype=np.float64, name='init_obs_noise_variance')
  ]

  # Since HMC operates over unconstrained space, we need to transform the
  # samples so they live in real-space.
  unconstraining_bijectors = [
      tfp.bijectors.Softplus(),
      tfp.bijectors.Softplus(),
      tfp.bijectors.Softplus(),
  ]

  def unnormalized_log_posterior(amplitude, length_scale, noise_variance):
    return joint_log_prob(
        observation_index_points, observations, amplitude, length_scale,
        noise_variance)

  num_results = 200
  [
      amplitudes,
      length_scales,
      observation_noise_variances
  ], kernel_results = tfp.mcmc.sample_chain(
      num_results=num_results,
      num_burnin_steps=500,
      num_steps_between_results=3,
      current_state=initial_chain_states,
      kernel=tfp.mcmc.TransformedTransitionKernel(
          inner_kernel = tfp.mcmc.HamiltonianMonteCarlo(
              target_log_prob_fn=unnormalized_log_posterior,
              step_size=[np.float64(.15)],
              num_leapfrog_steps=3),
          bijector=unconstraining_bijectors))

  # Now we can sample from the posterior predictive distribution at a new set
  # of index points.
  index_points = np.linspace(-1., 1., 200)[..., np.newaxis]
  gprm = tfd.GaussianProcessRegressionModel(
      # Batch of `num_results` kernels parameterized by the MCMC samples.
      kernel=psd_kernels.ExponentiatedQuadratic(amplitudes, length_scales),
      index_points=index_points,
      observation_index_points=observation_index_points,
      observations=observations,
      # We reshape this to align batch dimensions.
      observation_noise_variance=observation_noise_variances[..., np.newaxis])
  samples = gprm.sample()

  with tf.Session() as sess:
    kernel_results_, samples_ = sess.run([kernel_results, samples])

    print("Acceptance rate: {}".format(
        np.mean(kernel_results_.inner_results.is_accepted)))

    # Plot posterior samples and their mean, target function, and observations.
    plt.plot(np.stack([index_points[:, 0]]*num_results).T,
             samples_.T,
             c='r',
             alpha=.01)
    plt.plot(index_points[:, 0], np.mean(samples_, axis=0), c='k')
    plt.plot(index_points[:, 0], f(index_points))
    plt.scatter(observation_index_points[:, 0], observations)
  ```

  #### References
  [1]: Carl Rasmussen, Chris Williams. Gaussian Processes For Machine Learning,
       2006.
  """

  def __init__(self,
               kernel,
               index_points=None,
               observation_index_points=None,
               observations=None,
               observation_noise_variance=0.,
               predictive_noise_variance=None,
               mean_fn=None,
               jitter=1e-6,
               validate_args=False,
               allow_nan_stats=False,
               name='GaussianProcessRegressionModel'):
    """Construct a GaussianProcessRegressionModel instance.

    Args:
      kernel: `PositiveSemidefiniteKernel`-like instance representing the
        GP's covariance function.
      index_points: `float` `Tensor` representing finite collection, or batch of
        collections, of points in the index set over which the GP is defined.
        Shape has the form `[b1, ..., bB, e, f1, ..., fF]` where `F` is the
        number of feature dimensions and must equal `kernel.feature_ndims` and
        `e` is the number (size) of index points in each batch. Ultimately this
        distribution corresponds to an `e`-dimensional multivariate normal. The
        batch shape must be broadcastable with `kernel.batch_shape` and any
        batch dims yielded by `mean_fn`.
      observation_index_points: `float` `Tensor` representing finite collection,
        or batch of collections, of points in the index set for which some data
        has been observed. Shape has the form `[b1, ..., bB, e, f1, ..., fF]`
        where `F` is the number of feature dimensions and must equal
        `kernel.feature_ndims`, and `e` is the number (size) of index points in
        each batch. `[b1, ..., bB, e]` must be broadcastable with the shape of
        `observations`, and `[b1, ..., bB]` must be broadcastable with the
        shapes of all other batched parameters (`kernel.batch_shape`,
        `index_points`, etc). The default value is `None`, which corresponds to
        the empty set of observations, and simply results in the prior
        predictive model (a GP with noise of variance
        `predictive_noise_variance`).
      observations: `float` `Tensor` representing collection, or batch of
        collections, of observations corresponding to
        `observation_index_points`. Shape has the form `[b1, ..., bB, e]`, which
        must be brodcastable with the batch and example shapes of
        `observation_index_points`. The batch shape `[b1, ..., bB]` must be
        broadcastable with the shapes of all other batched parameters
        (`kernel.batch_shape`, `index_points`, etc.). The default value is
        `None`, which corresponds to the empty set of observations, and simply
        results in the prior predictive model (a GP with noise of variance
        `predictive_noise_variance`).
      observation_noise_variance: `float` `Tensor` representing the variance
        of the noise in the Normal likelihood distribution of the model. May be
        batched, in which case the batch shape must be broadcastable with the
        shapes of all other batched parameters (`kernel.batch_shape`,
        `index_points`, etc.).
        Default value: `0.`
      predictive_noise_variance: `float` `Tensor` representing the variance in
        the posterior predictive model. If `None`, we simply re-use
        `observation_noise_variance` for the posterior predictive noise. If set
        explicitly, however, we use this value. This allows us, for example, to
        omit predictive noise variance (by setting this to zero) to obtain
        noiseless posterior predictions of function values, conditioned on noisy
        observations.
      mean_fn: Python `callable` that acts on `index_points` to produce a
        collection, or batch of collections, of mean values at `index_points`.
        Takes a `Tensor` of shape `[b1, ..., bB, f1, ..., fF]` and returns a
        `Tensor` whose shape is broadcastable with `[b1, ..., bB]`.
        Default value: `None` implies the constant zero function.
      jitter: `float` scalar `Tensor` added to the diagonal of the covariance
        matrix to ensure positive definiteness of the covariance matrix.
        Default value: `1e-6`.
      validate_args: Python `bool`, default `False`. When `True` distribution
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may silently render incorrect
        outputs.
        Default value: `False`.
      allow_nan_stats: Python `bool`, default `True`. When `True`,
        statistics (e.g., mean, mode, variance) use the value `NaN` to
        indicate the result is undefined. When `False`, an exception is raised
        if one or more of the statistic's batch members are undefined.
        Default value: `False`.
      name: Python `str` name prefixed to Ops created by this class.
        Default value: 'GaussianProcessRegressionModel'.

    Raises:
      ValueError: if either
        - only one of `observations` and `observation_index_points` is given, or
        - `mean_fn` is not `None` and not callable.
    """
    parameters = dict(locals())
    with tf.name_scope(name) as name:
      dtype = dtype_util.common_dtype([
          index_points, observation_index_points, observations,
          observation_noise_variance, predictive_noise_variance, jitter
      ], tf.float32)
      if index_points is not None:
        index_points = tf.convert_to_tensor(
            value=index_points, dtype=dtype, name='index_points')
      observation_index_points = (None if observation_index_points is None else
                                  tf.convert_to_tensor(
                                      value=observation_index_points,
                                      dtype=dtype,
                                      name='observation_index_points'))
      observations = (None if observations is None else tf.convert_to_tensor(
          value=observations, dtype=dtype, name='observations'))
      observation_noise_variance = tf.convert_to_tensor(
          value=observation_noise_variance,
          dtype=dtype,
          name='observation_noise_variance')
      predictive_noise_variance = (
          observation_noise_variance
          if predictive_noise_variance is None else tf.convert_to_tensor(
              value=predictive_noise_variance,
              dtype=dtype,
              name='predictive_noise_variance'))
      jitter = tf.convert_to_tensor(value=jitter, dtype=dtype, name='jitter')

      if (observation_index_points is None) != (observations is None):
        raise ValueError(
            '`observations` and `observation_index_points` must both be given '
            'or None. Got {} and {}, respectively.'.format(
                observations, observation_index_points))
      # Default to a constant zero function, borrowing the dtype from
      # index_points to ensure consistency.
      if mean_fn is None:
        mean_fn = lambda x: tf.zeros([1], dtype=dtype)
      else:
        if not callable(mean_fn):
          raise ValueError('`mean_fn` must be a Python callable')

      self._name = name
      self._observation_index_points = observation_index_points
      self._observations = observations
      self._observation_noise_variance = observation_noise_variance
      self._predictive_noise_variance = predictive_noise_variance
      self._jitter = jitter
      self._validate_args = validate_args

      with tf.name_scope('init'):
        conditional_kernel = tfpk.SchurComplement(
            base_kernel=kernel,
            fixed_inputs=observation_index_points,
            diag_shift=jitter + observation_noise_variance)

        # Special logic for mean_fn only; SchurComplement already handles the
        # case of empty observations (ie, falls back to base_kernel).
        if _is_empty_observation_data(
            feature_ndims=kernel.feature_ndims,
            observation_index_points=observation_index_points,
            observations=observations):
          conditional_mean_fn = mean_fn
        else:
          _validate_observation_data(
              kernel=kernel,
              observation_index_points=observation_index_points,
              observations=observations)

          def conditional_mean_fn(x):
            k_x_obs_linop = tf.linalg.LinearOperatorFullMatrix(
                kernel.matrix(x, observation_index_points))
            chol_linop = tf.linalg.LinearOperatorLowerTriangular(
                conditional_kernel.divisor_matrix_cholesky)

            diff = observations - mean_fn(observation_index_points)
            return mean_fn(x) + k_x_obs_linop.matvec(
                chol_linop.solvevec(chol_linop.solvevec(diff), adjoint=True))

        graph_parents = [observation_noise_variance, jitter]
        def _maybe_append(x):
          if x is not None:
            graph_parents.append(x)
        _maybe_append(index_points)
        _maybe_append(observation_index_points)
        _maybe_append(observations)

        super(GaussianProcessRegressionModel, self).__init__(
            kernel=conditional_kernel,
            mean_fn=conditional_mean_fn,
            index_points=index_points,
            jitter=jitter,
            # What the GP super class calls "observation noise variance" we call
            # here the "predictive noise variance". We use the observation noise
            # variance for the fit/solve process above, and predictive for
            # downstream computations like sampling.
            observation_noise_variance=predictive_noise_variance,
            validate_args=validate_args,
            allow_nan_stats=allow_nan_stats, name=name)
        self._parameters = parameters
        self._graph_parents = graph_parents

  @property
  def observation_index_points(self):
    return self._observation_index_points

  @property
  def observations(self):
    return self._observations

  @property
  def predictive_noise_variance(self):
    return self._predictive_noise_variance


# Copyright 2018 The TensorFlow Probability 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.
# ============================================================================
"""The Gamma distribution class."""

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

# Dependency imports
import numpy as np
import tensorflow.compat.v2 as tf

from tensorflow_probability.python.distributions import distribution
from tensorflow_probability.python.distributions import kullback_leibler
from tensorflow_probability.python.internal import assert_util
from tensorflow_probability.python.internal import distribution_util
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import reparameterization

__all__ = [
    "Gamma",
]

class Gamma(distribution.Distribution):
  """Gamma distribution.

  The Gamma distribution is defined over positive real numbers using
  parameters `concentration` (aka "alpha") and `rate` (aka "beta").

  #### Mathematical Details

  The probability density function (pdf) is,

  ```none
  pdf(x; alpha, beta, x > 0) = x**(alpha - 1) exp(-x beta) / Z
  Z = Gamma(alpha) beta**(-alpha)
  ```

  where:

  * `concentration = alpha`, `alpha > 0`,
  * `rate = beta`, `beta > 0`,
  * `Z` is the normalizing constant, and,
  * `Gamma` is the [gamma function](
    https://en.wikipedia.org/wiki/Gamma_function).

  The cumulative density function (cdf) is,

  ```none
  cdf(x; alpha, beta, x > 0) = GammaInc(alpha, beta x) / Gamma(alpha)
  ```

  where `GammaInc` is the [lower incomplete Gamma function](
  https://en.wikipedia.org/wiki/Incomplete_gamma_function).

  The parameters can be intuited via their relationship to mean and stddev,

  ```none
  concentration = alpha = (mean / stddev)**2
  rate = beta = mean / stddev**2 = concentration / mean
  ```

  Distribution parameters are automatically broadcast in all functions; see
  examples for details.

  Warning: The samples of this distribution are always non-negative. However,
  the samples that are smaller than `np.finfo(dtype).tiny` are rounded
  to this value, so it appears more often than it should.
  This should only be noticeable when the `concentration` is very small, or the
  `rate` is very large. See note in `tf.random_gamma` docstring.

  Samples of this distribution are reparameterized (pathwise differentiable).
  The derivatives are computed using the approach described in the paper

  [Michael Figurnov, Shakir Mohamed, Andriy Mnih.
  Implicit Reparameterization Gradients, 2018](https://arxiv.org/abs/1805.08498)

  #### Examples

  ```python
  import tensorflow_probability as tfp
  tfd = tfp.distributions

  dist = tfd.Gamma(concentration=3.0, rate=2.0)
  dist2 = tfd.Gamma(concentration=[3.0, 4.0], rate=[2.0, 3.0])
  ```

  Compute the gradients of samples w.r.t. the parameters:

  ```python
  concentration = tf.constant(3.0)
  rate = tf.constant(2.0)
  dist = tfd.Gamma(concentration, rate)
  samples = dist.sample(5)  # Shape [5]
  loss = tf.reduce_mean(tf.square(samples))  # Arbitrary loss function
  # Unbiased stochastic gradients of the loss function
  grads = tf.gradients(loss, [concentration, rate])
  ```

  """

  def __init__(self,
               concentration,
               rate,
               validate_args=False,
               allow_nan_stats=True,
               name="Gamma"):
    """Construct Gamma with `concentration` and `rate` parameters.

    The parameters `concentration` and `rate` must be shaped in a way that
    supports broadcasting (e.g. `concentration + rate` is a valid operation).

    Args:
      concentration: Floating point tensor, the concentration params of the
        distribution(s). Must contain only positive values.
      rate: Floating point tensor, the inverse scale params of the
        distribution(s). Must contain only positive values.
      validate_args: Python `bool`, default `False`. When `True` distribution
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may silently render incorrect
        outputs.
      allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
        (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
        result is undefined. When `False`, an exception is raised if one or
        more of the statistic's batch members are undefined.
      name: Python `str` name prefixed to Ops created by this class.

    Raises:
      TypeError: if `concentration` and `rate` are different dtypes.
    """
    parameters = dict(locals())
    with tf.name_scope(name) as name:
      dtype = dtype_util.common_dtype([concentration, rate], tf.float32)
      concentration = tf.convert_to_tensor(
          value=concentration, name="concentration", dtype=dtype)
      rate = tf.convert_to_tensor(value=rate, name="rate", dtype=dtype)
      with tf.control_dependencies([
          assert_util.assert_positive(concentration),
          assert_util.assert_positive(rate),
      ] if validate_args else []):
        self._concentration = tf.identity(concentration)
        self._rate = tf.identity(rate)
        dtype_util.assert_same_float_dtype([self._concentration, self._rate])
    super(Gamma, self).__init__(
        dtype=dtype,
        validate_args=validate_args,
        allow_nan_stats=allow_nan_stats,
        reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,
        parameters=parameters,
        graph_parents=[self._concentration, self._rate],
        name=name)

  @staticmethod
  def _param_shapes(sample_shape):
    return dict(
        zip(("concentration", "rate"),
            ([tf.convert_to_tensor(value=sample_shape, dtype=tf.int32)] * 2)))

  @classmethod
  def _params_event_ndims(cls):
    return dict(concentration=0, rate=0)

  @property
  def concentration(self):
    """Concentration parameter."""
    return self._concentration

  @property
  def rate(self):
    """Rate parameter."""
    return self._rate

  def _batch_shape_tensor(self):
    return tf.broadcast_dynamic_shape(
        tf.shape(input=self.concentration), tf.shape(input=self.rate))

  def _batch_shape(self):
    return tf.broadcast_static_shape(
        self.concentration.shape,
        self.rate.shape)

  def _event_shape_tensor(self):
    return tf.constant([], dtype=tf.int32)

  def _event_shape(self):
    return tf.TensorShape([])

  @distribution_util.AppendDocstring(
      """Note: See `tf.random_gamma` docstring for sampling details and
      caveats.""")
  def _sample_n(self, n, seed=None):
    return tf.random.gamma(
        shape=[n],
        alpha=self.concentration,
        beta=self.rate,
        dtype=self.dtype,
        seed=seed)

  def _log_prob(self, x):
    return self._log_unnormalized_prob(x) - self._log_normalization()

  def _cdf(self, x):
    x = self._maybe_assert_valid_sample(x)
    # Note that igamma returns the regularized incomplete gamma function,
    # which is what we want for the CDF.
    return tf.math.igamma(self.concentration, self.rate * x)

  def _log_unnormalized_prob(self, x):
    x = self._maybe_assert_valid_sample(x)
    return tf.math.xlogy(self.concentration - 1., x) - self.rate * x

  def _log_normalization(self):
    return (tf.math.lgamma(self.concentration) -
            self.concentration * tf.math.log(self.rate))

  def _entropy(self):
    return (self.concentration - tf.math.log(self.rate) +
            tf.math.lgamma(self.concentration) +
            ((1. - self.concentration) * tf.math.digamma(self.concentration)))

  def _mean(self):
    return self.concentration / self.rate

  def _variance(self):
    return self.concentration / tf.square(self.rate)

  def _stddev(self):
    return tf.sqrt(self.concentration) / self.rate

  @distribution_util.AppendDocstring(
      """The mode of a gamma distribution is `(shape - 1) / rate` when
      `shape > 1`, and `NaN` otherwise. If `self.allow_nan_stats` is `False`,
      an exception will be raised rather than returning `NaN`.""")
  def _mode(self):
    mode = (self.concentration - 1.) / self.rate
    if self.allow_nan_stats:
      nan = tf.fill(
          self.batch_shape_tensor(),
          dtype_util.as_numpy_dtype(self.dtype)(np.nan),
          name="nan")
      return tf.where(self.concentration > 1., mode, nan)
    else:
      return distribution_util.with_dependencies([
          assert_util.assert_less(
              tf.ones([], self.dtype),
              self.concentration,
              message="mode not defined when any concentration <= 1"),
      ], mode)

  def _maybe_assert_valid_sample(self, x):
    dtype_util.assert_same_float_dtype(tensors=[x], dtype=self.dtype)
    if not self.validate_args:
      return x
    return distribution_util.with_dependencies([
        assert_util.assert_positive(x),
    ], x)

@kullback_leibler.RegisterKL(Gamma, Gamma)
def _kl_gamma_gamma(g0, g1, name=None):
  """Calculate the batched KL divergence KL(g0 || g1) with g0 and g1 Gamma.

  Args:
    g0: instance of a Gamma distribution object.
    g1: instance of a Gamma distribution object.
    name: (optional) Name to use for created operations. Default is
      "kl_gamma_gamma".

  Returns:
    kl_gamma_gamma: `Tensor`. The batchwise KL(g0 || g1).
  """
  with tf.name_scope(name or "kl_gamma_gamma"):
    # Result from:
    #   http://www.fil.ion.ucl.ac.uk/~wpenny/publications/densities.ps
    # For derivation see:
    #   http://stats.stackexchange.com/questions/11646/kullback-leibler-divergence-between-two-gamma-distributions   pylint: disable=line-too-long
    return (((g0.concentration - g1.concentration) *
             tf.math.digamma(g0.concentration)) +
            tf.math.lgamma(g1.concentration) -
            tf.math.lgamma(g0.concentration) +
            g1.concentration * tf.math.log(g0.rate) -
            g1.concentration * tf.math.log(g1.rate) + g0.concentration *
            (g1.rate / g0.rate - 1.))


#
# Copyright 2016 The BigDL 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.
#

import multiprocessing
import os
import sys
from distutils.dir_util import mkpath

from py4j.java_gateway import JavaObject
from pyspark.rdd import RDD

from bigdl.util.common import DOUBLEMAX
from bigdl.util.common import JTensor
from bigdl.util.common import JavaValue
from bigdl.util.common import callBigDlFunc
from bigdl.util.common import callJavaFunc
from bigdl.util.common import get_node_and_core_number
from bigdl.util.common import init_engine
from bigdl.util.common import to_list
from bigdl.dataset.dataset import *

if sys.version >= '3':
    long = int
    unicode = str

class Top1Accuracy(JavaValue):
    """
    Caculate the percentage that output's max probability index equals target.

    >>> top1 = Top1Accuracy()
    creating: createTop1Accuracy
    """
    def __init__(self, bigdl_type="float"):
        JavaValue.__init__(self, None, bigdl_type)

class TreeNNAccuracy(JavaValue):
    """
    Caculate the percentage that output's max probability index equals target.

    >>> top1 = TreeNNAccuracy()
    creating: createTreeNNAccuracy
    """
    def __init__(self, bigdl_type="float"):
        JavaValue.__init__(self, None, bigdl_type)

class Top5Accuracy(JavaValue):
    """
    Caculate the percentage that output's max probability index equals target.

    >>> top5 = Top5Accuracy()
    creating: createTop5Accuracy
    """
    def __init__(self, bigdl_type="float"):
        JavaValue.__init__(self, None, bigdl_type)

class Loss(JavaValue):

    """
    This evaluation method is calculate loss of output with respect to target
    >>> from bigdl.nn.criterion import ClassNLLCriterion
    >>> loss = Loss()
    creating: createClassNLLCriterion
    creating: createLoss

    >>> loss = Loss(ClassNLLCriterion())
    creating: createClassNLLCriterion
    creating: createLoss
    """
    def __init__(self, cri=None, bigdl_type="float"):
        from bigdl.nn.criterion import ClassNLLCriterion
        if cri is None:
            cri = ClassNLLCriterion()
        JavaValue.__init__(self, None, bigdl_type, cri)

class HitRatio(JavaValue):
    """
    Hit Ratio(HR) used in recommandation application.
    HR intuitively measures whether the test item is present on the top-k list.

    >>> hr10 = HitRatio(k = 10)
    creating: createHitRatio
    """
    def __init__(self, k = 10, neg_num = 100, bigdl_type="float"):
        """
        Create hit ratio validation method.

        :param k: top k
        :param neg_num: number of negative items.
        """
        JavaValue.__init__(self, None, bigdl_type, k, neg_num)

class NDCG(JavaValue):
    """
    Normalized Discounted Cumulative Gain(NDCG).
    NDCG accounts for the position of the hit by assigning higher scores to hits at top ranks.

    >>> ndcg = NDCG(k = 10)
    creating: createNDCG
    """
    def __init__(self, k = 10, neg_num = 100, bigdl_type="float"):
        """
        Create NDCG validation method.

        :param k: top k
        :param neg_num: number of negative items.
        """
        JavaValue.__init__(self, None, bigdl_type, k, neg_num)

class MAE(JavaValue):
    """
    This evaluation method calculates the mean absolute error of output with respect to target.

    >>> mae = MAE()
    creating: createMAE
    """
    def __init__(self, bigdl_type="float"):
        JavaValue.__init__(self, None, bigdl_type)

class MaxIteration(JavaValue):
    """
    A trigger specifies a timespot or several timespots during training,
    and a corresponding action will be taken when the timespot(s) is reached.
    MaxIteration is a trigger that triggers an action when training reaches
    the number of iterations specified by "max".
    Usually used as end_trigger when creating an Optimizer.

    >>> maxIteration = MaxIteration(20)
    creating: createMaxIteration
    """
    def __init__(self, max, bigdl_type="float"):
        """
        Create a MaxIteration trigger.

        :param max: max
        """
        JavaValue.__init__(self, None, bigdl_type, max)

class MaxEpoch(JavaValue):
    """
    A trigger specifies a timespot or several timespots during training,
    and a corresponding action will be taken when the timespot(s) is reached.
    MaxEpoch is a trigger that triggers an action when training reaches
    the number of epochs specified by "max_epoch".
    Usually used as end_trigger when creating an Optimizer.

    >>> maxEpoch = MaxEpoch(2)
    creating: createMaxEpoch
    """
    def __init__(self, max_epoch, bigdl_type="float"):
        """
        Create a MaxEpoch trigger.

        :param max_epoch: max_epoch
        """
        JavaValue.__init__(self, None, bigdl_type, max_epoch)

class EveryEpoch(JavaValue):
    """
    A trigger specifies a timespot or several timespots during training,
    and a corresponding action will be taken when the timespot(s) is reached.
    EveryEpoch is a trigger that triggers an action when each epoch finishs.
    Could be used as trigger in setvalidation and setcheckpoint in Optimizer,
    and also in TrainSummary.set_summary_trigger.

    >>> everyEpoch = EveryEpoch()
    creating: createEveryEpoch
    """
    def __init__(self, bigdl_type="float"):
        """
        Create a EveryEpoch trigger.
        """
        JavaValue.__init__(self, None, bigdl_type)

class SeveralIteration(JavaValue):
    """
    A trigger specifies a timespot or several timespots during training,
    and a corresponding action will be taken when the timespot(s) is reached.
    SeveralIteration is a trigger that triggers an action every "n"
    iterations.
    Could be used as trigger in setvalidation and setcheckpoint in Optimizer,
    and also in TrainSummary.set_summary_trigger.

    >>> serveralIteration = SeveralIteration(2)
    creating: createSeveralIteration
    """
    def __init__(self, interval, bigdl_type="float"):
        """
        Create a SeveralIteration trigger.

        :param interval: interval is the "n" where an action is triggeredevery "n" iterations
        """
        JavaValue.__init__(self, None, bigdl_type, interval)

class MaxScore(JavaValue):
    """
    A trigger that triggers an action when validation score larger than "max" score

    >>> maxScore = MaxScore(0.4)
    creating: createMaxScore
    """
    def __init__(self, max, bigdl_type="float"):
        """
        Create a MaxScore trigger.

        :param max: max score
        """
        JavaValue.__init__(self, None, bigdl_type, max)

class MinLoss(JavaValue):
    """
    A trigger that triggers an action when training loss less than "min" loss

    >>> minLoss = MinLoss(0.1)
    creating: createMinLoss
    """
    def __init__(self, min, bigdl_type="float"):
        """
        Create a MinLoss trigger.

        :param min: min loss
        """
        JavaValue.__init__(self, None, bigdl_type, min)

class TriggerAnd(JavaValue):
    """
    A trigger contains other triggers and triggers when all of them trigger (logical AND)

    >>> a = TriggerAnd(MinLoss(0.1), MaxEpoch(2))
    creating: createMinLoss
    creating: createMaxEpoch
    creating: createTriggerAnd
    """
    def __init__(self, first, *other):
        """
        Create a And trigger.

        :param first: first Trigger
        :param other: other Trigger
        """
        JavaValue.__init__(self, None, "float", first, list(other))

class TriggerOr(JavaValue):
    """
    A trigger contains other triggers and triggers when any of them trigger (logical OR)

    >>> o = TriggerOr(MinLoss(0.1), MaxEpoch(2))
    creating: createMinLoss
    creating: createMaxEpoch
    creating: createTriggerOr
    """
    def __init__(self, first, *other):
        """
        Create a Or trigger.

        :param first: first Trigger
        :param other: other Trigger
        """
        JavaValue.__init__(self, None, "float", first, list(other))

class Poly(JavaValue):
    """
    A learning rate decay policy, where the effective learning rate
    follows a polynomial decay, to be zero by the max_iteration.
    Calculation: base_lr (1 - iter/max_iteration) ^ (power)

    :param power: coeffient of decay, refer to calculation formula
    :param max_iteration: max iteration when lr becomes zero

    >>> poly = Poly(0.5, 2)
    creating: createPoly
    """
    def __init__(self, power, max_iteration, bigdl_type="float"):
            JavaValue.__init__(self, None, bigdl_type, power, max_iteration)

class Exponential(JavaValue):
    """
    [[Exponential]] is a learning rate schedule, which rescale the learning rate by
    lr_{n + 1} = lr * decayRate `^` (iter / decayStep)
    :param decay_step the inteval for lr decay
    :param decay_rate decay rate
    :param stair_case if true, iter / decayStep is an integer division
                     and the decayed learning rate follows a staircase function.

    >>> exponential = Exponential(100, 0.1)
    creating: createExponential
    """
    def __init__(self, decay_step, decay_rate, stair_case=False, bigdl_type="float"):
        JavaValue.__init__(self, None, bigdl_type, decay_step, decay_rate, stair_case)

class Step(JavaValue):
    """
    A learning rate decay policy, where the effective learning rate is
    calculated as base_lr * gamma ^ (floor(iter / step_size))

    :param step_size:
    :param gamma:

    >>> step = Step(2, 0.3)
    creating: createStep
    """
    def __init__(self, step_size, gamma, bigdl_type="float"):
            JavaValue.__init__(self, None, bigdl_type, step_size, gamma)

class Default(JavaValue):
    """
    A learning rate decay policy, where the effective learning rate is
    calculated as base_lr * gamma ^ (floor(iter / step_size))

    :param step_size
    :param gamma

    >>> step = Default()
    creating: createDefault
    """
    def __init__(self, bigdl_type="float"):
        JavaValue.__init__(self, None, bigdl_type)

class Plateau(JavaValue):
    """
    Plateau is the learning rate schedule when a metric has stopped improving.
    Models often benefit from reducing the learning rate by a factor of 2-10
    once learning stagnates. It monitors a quantity and if no improvement
    is seen for a 'patience' number of epochs, the learning rate is reduced.

    :param monitor quantity to be monitored, can be Loss or score
    :param factor factor by which the learning rate will be reduced. new_lr = lr * factor
    :param patience number of epochs with no improvement after which learning rate will be reduced.
    :param mode one of {min, max}.
                In min mode, lr will be reduced when the quantity monitored has stopped decreasing;
                in max mode it will be reduced when the quantity monitored has stopped increasing
    :param epsilon threshold for measuring the new optimum, to only focus on significant changes.
    :param cooldown number of epochs to wait before resuming normal operation
                    after lr has been reduced.
    :param min_lr lower bound on the learning rate.

    >>> plateau = Plateau("score")
    creating: createPlateau
    """
    def __init__(self,
                 monitor,
                 factor=0.1,
                 patience=10,
                 mode="min",
                 epsilon=1e-4,
                 cooldown=0,
                 min_lr=0.0,
                 bigdl_type="float"):
        JavaValue.__init__(self, None, bigdl_type, monitor, factor, patience, mode, epsilon,
                           cooldown, min_lr)

class Warmup(JavaValue):
    """
    A learning rate gradual increase policy, where the effective learning rate
    increase delta after each iteration.
    Calculation: base_lr + delta * iteration

    :param delta: increase amount after each iteration

    >>> warmup = Warmup(0.05)
    creating: createWarmup
    """
    def __init__(self, delta, bigdl_type="float"):
        JavaValue.__init__(self, None, bigdl_type, delta)

class SequentialSchedule(JavaValue):
    """
    Stack several learning rate schedulers.

    :param iterationPerEpoch: iteration numbers per epoch

    >>> sequentialSchedule = SequentialSchedule(5)
    creating: createSequentialSchedule
    >>> poly = Poly(0.5, 2)
    creating: createPoly
    >>> test = sequentialSchedule.add(poly, 5)

    """
    def __init__(self, iteration_per_epoch, bigdl_type="float"):
        JavaValue.__init__(self, None, bigdl_type, iteration_per_epoch)

    def add(self, scheduler, max_iteration, bigdl_type="float"):
        """
        Add a learning rate scheduler to the contained `schedules`

        :param scheduler: learning rate scheduler to be add
        :param max_iteration: iteration numbers this scheduler will run
        """
        return callBigDlFunc(bigdl_type, "addScheduler", self.value, scheduler, max_iteration)

class OptimMethod(JavaValue):

    def __init__(self, jvalue, bigdl_type, *args):
        if (jvalue):
            assert(type(jvalue) == JavaObject)
            self.value = jvalue
        else:
            self.value = callBigDlFunc(
                bigdl_type, JavaValue.jvm_class_constructor(self), *args)
        self.bigdl_type = bigdl_type

    @staticmethod
    def load(path, bigdl_type="float"):
        """
        load optim method
        :param path: file path
        """
        return callBigDlFunc(bigdl_type, "loadOptimMethod", path)

    def save(self, path, overWrite):
        """
        save OptimMethod
        :param path      path
        :param overWrite whether to overwrite
        """
        method=self.value
        return callBigDlFunc(self.bigdl_type, "saveOptimMethod", method, path, overWrite)

class SGD(OptimMethod):
    """
    A plain implementation of SGD

    :param learningrate learning rate
    :param learningrate_decay learning rate decay
    :param weightdecay weight decay
    :param momentum momentum
    :param dampening dampening for momentum
    :param nesterov enables Nesterov momentum
    :param learningrates 1D tensor of individual learning rates
    :param weightdecays 1D tensor of individual weight decays
    >>> sgd = SGD()
    creating: createDefault
    creating: createSGD
    """
    def __init__(self,
                 learningrate=1e-3,
                 learningrate_decay=0.0,
                 weightdecay=0.0,
                 momentum=0.0,
                 dampening=DOUBLEMAX,
                 nesterov=False,
                 leaningrate_schedule=None,
                 learningrates=None,
                 weightdecays=None,
                 bigdl_type="float"):
        super(SGD, self).__init__(None, bigdl_type, learningrate, learningrate_decay, weightdecay,
                           momentum, dampening, nesterov,
                           leaningrate_schedule if (leaningrate_schedule) else Default(),
                           JTensor.from_ndarray(learningrates), JTensor.from_ndarray(weightdecays))

class Adagrad(OptimMethod):
    """
    An implementation of Adagrad. See the original paper:
    http://jmlr.org/papers/volume12/duchi11a/duchi11a.pdf

    :param learningrate learning rate
    :param learningrate_decay learning rate decay
    :param weightdecay weight decay
    >>> adagrad = Adagrad()
    creating: createAdagrad
    """
    def __init__(self,
                 learningrate=1e-3,
                 learningrate_decay=0.0,
                 weightdecay=0.0,
                 bigdl_type="float"):
        super(Adagrad, self).__init__(None, bigdl_type, learningrate, learningrate_decay, weightdecay)

class LBFGS(OptimMethod):
    """
    This implementation of L-BFGS relies on a user-provided line
    search function (state.lineSearch). If this function is not
    provided, then a simple learningRate is used to produce fixed
    size steps. Fixed size steps are much less costly than line
    searches, and can be useful for stochastic problems.
    The learning rate is used even when a line search is provided.
    This is also useful for large-scale stochastic problems, where
    opfunc is a noisy approximation of f(x). In that case, the learning
    rate allows a reduction of confidence in the step size.

    :param max_iter Maximum number of iterations allowed
    :param max_eval Maximum number of function evaluations
    :param tolfun Termination tolerance on the first-order optimality
    :param tolx Termination tol on progress in terms of func/param changes
    :param ncorrection
    :param learningrate
    :param verbose
    :param linesearch A line search function
    :param linesearch_options If no line search provided, then a fixed step size is used
    >>> lbfgs = LBFGS()
    creating: createLBFGS
    """
    def __init__(self,
                 max_iter=20,
                 max_eval=DOUBLEMAX,
                 tolfun=1e-5,
                 tolx=1e-9,
                 ncorrection=100,
                 learningrate=1.0,
                 verbose=False,
                 linesearch=None,
                 linesearch_options=None,
                 bigdl_type="float"):
        if linesearch or linesearch_options:
            raise ValueError('linesearch and linesearch_options must be None in LBFGS')
        super(LBFGS, self).__init__(None, bigdl_type, max_iter, max_eval, tolfun, tolx,
                       ncorrection, learningrate, verbose, linesearch, linesearch_options)

class Adadelta(OptimMethod):
    """
    Adadelta implementation for SGD: http://arxiv.org/abs/1212.5701

    :param decayrate interpolation parameter rho
    :param epsilon for numerical stability
    >>> adagrad = Adadelta()
    creating: createAdadelta
    """
    def __init__(self,
                 decayrate = 0.9,
                 epsilon = 1e-10,
                 bigdl_type="float"):
        super(Adadelta, self).__init__(None, bigdl_type, decayrate, epsilon)

class Adam(OptimMethod):
    """
    An implementation of Adam http://arxiv.org/pdf/1412.6980.pdf
    :param learningrate learning rate
    :param learningrate_decay learning rate decay
    :param beta1 first moment coefficient
    :param beta2 second moment coefficient
    :param epsilon for numerical stability
    >>> adam = Adam()
    creating: createAdam
    """
    def __init__(self,
                 learningrate = 1e-3,
                 learningrate_decay = 0.0,
                 beta1 = 0.9,
                 beta2 = 0.999,
                 epsilon = 1e-8,
                 bigdl_type="float"):
        super(Adam, self).__init__(None, bigdl_type, learningrate, learningrate_decay,
                           beta1, beta2, epsilon)

class ParallelAdam(OptimMethod):
    """
    An implementation of Adam http://arxiv.org/pdf/1412.6980.pdf
    :param learningrate learning rate
    :param learningrate_decay learning rate decay
    :param beta1 first moment coefficient
    :param beta2 second moment coefficient
    :param epsilon for numerical stability
    >>> init_engine()
    >>> pAdam = ParallelAdam()
    creating: createParallelAdam
    """
    def __init__(self,
                 learningrate = 1e-3,
                 learningrate_decay = 0.0,
                 beta1 = 0.9,
                 beta2 = 0.999,
                 epsilon = 1e-8,
                 parallel_num = -1,
                 bigdl_type="float"):
        if parallel_num == -1:
            parallel_num = get_node_and_core_number()[1]
        super(ParallelAdam, self).__init__(None, bigdl_type, learningrate, learningrate_decay,
                                   beta1, beta2, epsilon, parallel_num)

class Ftrl(OptimMethod):
    """
    An implementation of Ftrl https://www.eecs.tufts.edu/~dsculley/papers/ad-click-prediction.pdf.
    Support L1 penalty, L2 penalty and shrinkage-type L2 penalty.

    :param learningrate learning rate
    :param learningrate_power double, must be less or equal to zero. Default is -0.5.
    :param initial_accumulator_value double, the starting value for accumulators,
        require zero or positive values.
    :param l1_regularization_strength double, must be greater or equal to zero. Default is zero.
    :param l2_regularization_strength double, must be greater or equal to zero. Default is zero.
    :param l2_shrinkage_regularization_strength double, must be greater or equal to zero.
        Default is zero. This differs from l2RegularizationStrength above. L2 above is a
        stabilization penalty, whereas this one is a magnitude penalty.
    >>> ftrl = Ftrl()
    creating: createFtrl
    >>> ftrl2 = Ftrl(1e-2, -0.1, 0.2, 0.3, 0.4, 0.5)
    creating: createFtrl
    """
    def __init__(self,
                 learningrate = 1e-3,
                 learningrate_power = -0.5,
                 initial_accumulator_value = 0.1,
                 l1_regularization_strength = 0.0,
                 l2_regularization_strength = 0.0,
                 l2_shrinkage_regularization_strength = 0.0,
                 bigdl_type="float"):
        super(Ftrl, self).__init__(None, bigdl_type, learningrate, learningrate_power,
                                   initial_accumulator_value,
                                   l1_regularization_strength,
                                   l2_regularization_strength,
                                   l2_shrinkage_regularization_strength)

class Adamax(OptimMethod):
    """
    An implementation of Adamax http://arxiv.org/pdf/1412.6980.pdf
    :param learningrate learning rate
    :param beta1 first moment coefficient
    :param beta2 second moment coefficient
    :param epsilon for numerical stability
    >>> adagrad = Adamax()
    creating: createAdamax
    """
    def __init__(self,
                 learningrate = 0.002,
                 beta1 = 0.9,
                 beta2 = 0.999,
                 epsilon = 1e-38,
                 bigdl_type="float"):
        super(Adamax, self).__init__(None, bigdl_type, learningrate, beta1, beta2, epsilon)

class RMSprop(OptimMethod):
    """
    An implementation of RMSprop
    :param learningrate learning rate
    :param learningrate_decay learning rate decay
    :param decayrate decay rate, also called rho
    :param epsilon for numerical stability
    >>> adagrad = RMSprop()
    creating: createRMSprop
    """
    def __init__(self,
                 learningrate = 1e-2,
                 learningrate_decay = 0.0,
                 decayrate = 0.99,
                 epsilon = 1e-8,
                 bigdl_type="float"):
        super(RMSprop, self).__init__(None, bigdl_type, learningrate, learningrate_decay, decayrate, epsilon)

class MultiStep(JavaValue):
    """
    similar to step but it allows non uniform steps defined by stepSizes

    :param step_size: the series of step sizes used for lr decay
    :param gamma: coefficient of decay

    >>> step = MultiStep([2, 5], 0.3)
    creating: createMultiStep
    """
    def __init__(self, step_sizes, gamma, bigdl_type="float"):
        JavaValue.__init__(self, None, bigdl_type, step_sizes, gamma)

class BaseOptimizer(JavaValue):

    def set_model(self, model):
        """
        Set model.

        :param model: new model
        """
        self.value.setModel(model.value)

    def set_criterion(self, criterion):
        """
        set new criterion, for optimizer reuse

        :param criterion: new criterion
        :return:
        """
        callBigDlFunc(self.bigdl_type, "setCriterion", self.value,
                      criterion)

    def set_checkpoint(self, checkpoint_trigger,
                       checkpoint_path, isOverWrite=True):
        """
        Configure checkpoint settings.

        :param checkpoint_trigger: the interval to write snapshots
        :param checkpoint_path: the path to write snapshots into
        :param isOverWrite: whether to overwrite existing snapshots in path.default is True
        """
        if not os.path.exists(checkpoint_path):
            mkpath(checkpoint_path)
        callBigDlFunc(self.bigdl_type, "setCheckPoint", self.value,
                      checkpoint_trigger, checkpoint_path, isOverWrite)

    def set_gradclip_const(self, min_value, max_value):
        """
        Configure constant clipping settings.

        :param min_value: the minimum value to clip by
        :param max_value: the maxmimum value to clip by
        """
        callBigDlFunc(self.bigdl_type, "setConstantClip", self.value, min_value, max_value)

    def set_gradclip_l2norm(self, clip_norm):
        """
        Configure L2 norm clipping settings.

        :param clip_norm: gradient L2-Norm threshold
        """
        callBigDlFunc(self.bigdl_type, "setL2NormClip", self.value, clip_norm)

    def disable_gradclip(self):
        """
        disable clipping.
        """
        callBigDlFunc(self.bigdl_type, "disableClip", self.value)

    # return a module
    def optimize(self):
        """
        Do an optimization.
        """
        jmodel = callJavaFunc(self.value.optimize)
        from bigdl.nn.layer import Layer
        return Layer.of(jmodel)

    def set_train_summary(self, summary):
        """
        Set train summary. A TrainSummary object contains information
        necessary for the optimizer to know how often the logs are recorded,
        where to store the logs and how to retrieve them, etc. For details,
        refer to the docs of TrainSummary.

        :param summary: a TrainSummary object
        """
        callBigDlFunc(self.bigdl_type, "setTrainSummary", self.value,
                      summary)
        return self

    def set_val_summary(self, summary):
        """
        Set validation summary. A ValidationSummary object contains information
        necessary for the optimizer to know how often the logs are recorded,
        where to store the logs and how to retrieve them, etc. For details,
        refer to the docs of ValidationSummary.

        :param summary: a ValidationSummary object

        """
        callBigDlFunc(self.bigdl_type, "setValSummary", self.value,
                      summary)
        return self

    def prepare_input(self):
        """
        Load input. Notebook user can call this method to seprate load data and
        create optimizer time
        """
        print("Loading input ...")
        self.value.prepareInput()

    def set_end_when(self, end_when):
        """
        When to stop, passed in a [[Trigger]]
        """
        self.value.setEndWhen(end_when.value)
        return self

class Optimizer(BaseOptimizer):

    # NOTE: This is a deprecated method, you should use `create` method instead.
    def __init__(self,
                 model,
                 training_rdd,
                 criterion,
                 end_trigger,
                 batch_size,
                 optim_method=None,
                 bigdl_type="float"):
        """
        Create a distributed optimizer.

        :param model: the neural net model
        :param training_rdd: the training dataset
        :param criterion: the loss function
        :param optim_method: the algorithm to use for optimization,
           e.g. SGD, Adagrad, etc. If optim_method is None, the default algorithm is SGD.
        :param end_trigger: when to end the optimization
        :param batch_size: training batch size
        """
        self.pvalue = DistriOptimizer(model,
                                      training_rdd,
                                      criterion,
                                      end_trigger,
                                      batch_size,
                                      optim_method,
                                      bigdl_type)
        self.value = self.pvalue.value
        self.bigdl_type = self.pvalue.bigdl_type

    @staticmethod
    def create(model,
               training_set,
               criterion,
               end_trigger=None,
               batch_size=32,
               optim_method=None,
               cores=None,
               bigdl_type="float"):
        """
        Create an optimizer.
        Depend on the input type, the returning optimizer can be a local optimizer \
        or a distributed optimizer.

        :param model: the neural net model
        :param training_set: (features, label) for local mode. RDD[Sample] for distributed mode.
        :param criterion: the loss function
        :param optim_method: the algorithm to use for optimization,
           e.g. SGD, Adagrad, etc. If optim_method is None, the default algorithm is SGD.
        :param end_trigger: when to end the optimization. default value is MapEpoch(1)
        :param batch_size: training batch size
        :param cores: This is for local optimizer only and use total physical cores as the default value
        """
        if not end_trigger:
            end_trigger = MaxEpoch(1)
        if not optim_method:
            optim_method = SGD()
        if isinstance(training_set, RDD) or isinstance(training_set, DataSet):
            return DistriOptimizer(model=model,
                                   training_rdd=training_set,
                                   criterion=criterion,
                                   end_trigger=end_trigger,
                                   batch_size=batch_size,
                                   optim_method=optim_method,
                                   bigdl_type=bigdl_type)
        elif isinstance(training_set, tuple) and len(training_set) == 2:
            x, y = training_set
            return LocalOptimizer(X=x,
                                  Y=y,
                                  model=model,
                                  criterion=criterion,
                                  end_trigger=end_trigger,
                                  batch_size=batch_size,
                                  optim_method=optim_method,
                                  cores=cores,
                                  bigdl_type="float")
        else:
            raise Exception("Not supported training set: %s" % type(training_set))

    def set_validation(self, batch_size, val_rdd, trigger, val_method=None):
        """
        Configure validation settings.

        :param batch_size: validation batch size
        :param val_rdd: validation dataset
        :param trigger: validation interval
        :param val_method: the ValidationMethod to use,e.g. "Top1Accuracy", "Top5Accuracy", "Loss"
        """
        if val_method is None:
            val_method = [Top1Accuracy()]
        func_name = "setValidation"
        if isinstance(val_rdd, DataSet):
            func_name = "setValidationFromDataSet"
        callBigDlFunc(self.bigdl_type, func_name, self.value, batch_size,
                      trigger, val_rdd, to_list(val_method))

    def set_traindata(self, training_rdd, batch_size):
        """
        Set new training dataset, for optimizer reuse

        :param training_rdd: the training dataset
        :param batch_size: training batch size
        :return:
        """
        callBigDlFunc(self.bigdl_type, "setTrainData", self.value,
                     training_rdd, batch_size)

class DistriOptimizer(Optimizer):
    def __init__(self,
                 model,
                 training_rdd,
                 criterion,
                 end_trigger,
                 batch_size,
                 optim_method=None,
                 bigdl_type="float"):
        """
        Create an optimizer.

        :param model: the neural net model
        :param training_data: the training dataset
        :param criterion: the loss function
        :param optim_method: the algorithm to use for optimization,
           e.g. SGD, Adagrad, etc. If optim_method is None, the default algorithm is SGD.
        :param end_trigger: when to end the optimization
        :param batch_size: training batch size
        """
        if not optim_method:
            optim_methods = {model.name(): SGD()}
        elif isinstance(optim_method, OptimMethod):
            optim_methods = {model.name(): optim_method}
        elif isinstance(optim_method, JavaObject):
            optim_methods = {model.name(): OptimMethod(optim_method, bigdl_type)}
        else:
            optim_methods = optim_method
        if isinstance(training_rdd, RDD):
            JavaValue.__init__(self, None, bigdl_type, model.value,
                               training_rdd, criterion,
                               optim_methods, end_trigger, batch_size)
        elif isinstance(training_rdd, DataSet):
            self.bigdl_type = bigdl_type
            self.value = callBigDlFunc(self.bigdl_type, "createDistriOptimizerFromDataSet",
                                       model.value, training_rdd, criterion,
                                       optim_methods, end_trigger, batch_size)

class LocalOptimizer(BaseOptimizer):
    """
    Create an optimizer.

    :param model: the neural net model
    :param X: the training features which is an ndarray or list of ndarray
    :param Y: the training label which is an ndarray
    :param criterion: the loss function
    :param optim_method: the algorithm to use for optimization,
       e.g. SGD, Adagrad, etc. If optim_method is None, the default algorithm is SGD.
    :param end_trigger: when to end the optimization
    :param batch_size: training batch size
    :param cores: by default is the total physical cores.
    """
    def __init__(self,
                 X,
                 Y,
                 model,
                 criterion,
                 end_trigger,
                 batch_size,
                 optim_method=None,
                 cores=None,
                 bigdl_type="float"):
        if not optim_method:
            optim_methods = {model.name(): SGD()}
        elif isinstance(optim_method, OptimMethod):
            optim_methods = {model.name(): optim_method}
        elif isinstance(optim_method, JavaObject):
            optim_methods = {model.name(): OptimMethod(optim_method, bigdl_type)}
        else:
            optim_methods = optim_method
        if cores is None:
            cores = multiprocessing.cpu_count()
        JavaValue.__init__(self, None, bigdl_type,
                           [JTensor.from_ndarray(X) for X in to_list(X)],
                           JTensor.from_ndarray(Y),
                           model.value,
                           criterion,
                           optim_methods, end_trigger, batch_size, cores)

    def set_validation(self, batch_size, X_val, Y_val, trigger, val_method=None):
        """
        Configure validation settings.

        :param batch_size: validation batch size
        :param X_val: features of validation dataset
        :param Y_val: label of validation dataset
        :param trigger: validation interval
        :param val_method: the ValidationMethod to use,e.g. "Top1Accuracy", "Top5Accuracy", "Loss"
        """
        if val_method is None:
            val_method = [Top1Accuracy()]
        callBigDlFunc(self.bigdl_type, "setValidation", self.value, batch_size,
                      trigger, [JTensor.from_ndarray(X) for X in to_list(X_val)],
                      JTensor.from_ndarray(Y_val), to_list(val_method))

class TrainSummary(JavaValue, ):
    """
    A logging facility which allows user to trace how indicators (e.g.
    learning rate, training loss, throughput, etc.) change with iterations/time
    in an optimization process. TrainSummary is for training indicators only
    (check ValidationSummary for validation indicators).  It contains necessary
    information for the optimizer to know where to store the logs, how to
    retrieve the logs, and so on. - The logs are written in tensorflow-compatible
    format so that they can be visualized directly using tensorboard. Also the
    logs can be retrieved as ndarrays and visualized using python libraries
    such as matplotlib (in notebook, etc.).

    Use optimizer.setTrainSummary to enable train logger.
    """
    def __init__(self, log_dir, app_name, bigdl_type="float"):
        """
        Create a TrainSummary. Logs will be saved to log_dir/app_name/train.

        :param log_dir: the root dir to store the logs
        :param app_name: the application name
        """
        JavaValue.__init__(self, None, bigdl_type, log_dir, app_name)

    def read_scalar(self, tag):
        """
        Retrieve train logs by type. Return an array of records in the format
        (step,value,wallClockTime). - "Step" is the iteration count by default.

        :param tag: the type of the logs, Supported tags are: "LearningRate","Loss", "Throughput"
        """
        return callBigDlFunc(self.bigdl_type, "summaryReadScalar", self.value,
                             tag)

    def set_summary_trigger(self, name, trigger):
        """
        Set the interval of recording for each indicator.

        :param tag: tag name. Supported tag names are "LearningRate", "Loss","Throughput", "Parameters". "Parameters" is an umbrella tag thatincludes weight, bias, gradWeight, gradBias, and some running status(eg. runningMean and runningVar in BatchNormalization). If youdidn't set any triggers, we will by default record Loss and Throughputin each iteration, while *NOT* recording LearningRate and Parameters,as recording parameters may introduce substantial overhead when themodel is very big, LearningRate is not a public attribute for allOptimMethod.
        :param trigger: trigger
        """
        return callBigDlFunc(self.bigdl_type, "summarySetTrigger", self.value,
                             name, trigger)

class ValidationSummary(JavaValue):
    """
     A logging facility which allows user to trace how indicators (e.g.
     validation loss, top1 accuray, top5 accuracy etc.) change with
     iterations/time in an optimization process. ValidationSummary is for
     validation indicators only (check TrainSummary for train indicators).
     It contains necessary information for the optimizer to know where to
     store the logs, how to retrieve the logs, and so on. - The logs are
     written in tensorflow-compatible format so that they can be visualized
     directly using tensorboard. Also the logs can be retrieved as ndarrays
     and visualized using python libraries such as matplotlib
     (in notebook, etc.).

     Use optimizer.setValidationSummary to enable validation logger.
    """
    def __init__(self, log_dir, app_name, bigdl_type="float"):
        """
        Create a ValidationSummary. Logs will be saved to
        log_dir/app_name/train. By default, all ValidationMethod set into
        optimizer will be recorded and the recording interval is the same
        as trigger of ValidationMethod in the optimizer.

        :param log_dir: the root dir to store the logs
        :param app_name: the application name
        """
        JavaValue.__init__(self, None, bigdl_type, log_dir, app_name)

    def read_scalar(self, tag):
        """
        Retrieve validation logs by type. Return an array of records in the
        format (step,value,wallClockTime). - "Step" is the iteration count
        by default.

        :param tag: the type of the logs. The tag should match the name ofthe ValidationMethod set into the optimizer. e.g."Top1AccuracyLoss","Top1Accuracy" or "Top5Accuracy".
        """
        return callBigDlFunc(self.bigdl_type, "summaryReadScalar", self.value,
                             tag)

class L1L2Regularizer(JavaValue):
    """
    Apply both L1 and L2 regularization

    :param l1 l1 regularization rate
    :param l2 l2 regularization rate

    """
    def __init__(self, l1, l2, bigdl_type="float"):
        JavaValue.__init__(self, None, bigdl_type, l1, l2)

class ActivityRegularization(JavaValue):
    """
    Apply both L1 and L2 regularization

    :param l1 l1 regularization rate
    :param l2 l2 regularization rate

    """
    def __init__(self, l1, l2, bigdl_type="float"):
        JavaValue.__init__(self, None, bigdl_type, l1, l2)

class L1Regularizer(JavaValue):
    """
    Apply L1 regularization

    :param l1 l1 regularization rate

    """
    def __init__(self, l1, bigdl_type="float"):
        JavaValue.__init__(self, None, bigdl_type, l1)

class L2Regularizer(JavaValue):
    """
    Apply L2 regularization

    :param l2 l2 regularization rate

    """
    def __init__(self, l2, bigdl_type="float"):
        JavaValue.__init__(self, None, bigdl_type, l2)

def _test():
    import doctest
    from pyspark import SparkContext
    from bigdl.optim import optimizer
    from bigdl.util.common import init_engine
    from bigdl.util.common import create_spark_conf
    globs = optimizer.__dict__.copy()
    sc = SparkContext(master="local[4]", appName="test optimizer",
                      conf=create_spark_conf())
    init_engine()
    globs['sc'] = sc
    (failure_count, test_count) = doctest.testmod(globs=globs,
                                                  optionflags=doctest.ELLIPSIS)
    if failure_count:
        exit(-1)

if __name__ == "__main__":
    _test()


#
# Copyright 2016 The BigDL 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.
#
# Part of the code originally from Tensorflow

import gzip

import numpy

from bigdl.dataset import base
from bigdl.dataset.transformer import *

SOURCE_URL = 'http://yann.lecun.com/exdb/mnist/'

TRAIN_MEAN = 0.13066047740239506 * 255
TRAIN_STD = 0.3081078 * 255
TEST_MEAN = 0.13251460696903547 * 255
TEST_STD = 0.31048024 * 255

def _read32(bytestream):
    dt = numpy.dtype(numpy.uint32).newbyteorder('>')
    return numpy.frombuffer(bytestream.read(4), dtype=dt)[0]

def extract_images(f):
    """Extract the images into a 4D uint8 numpy array [index, y, x, depth].

    :param: f: A file object that can be passed into a gzip reader.
    :return: data: A 4D unit8 numpy array [index, y, x, depth].
    :raise: ValueError: If the bytestream does not start with 2051.

    """
    print('Extracting', f.name)
    with gzip.GzipFile(fileobj=f) as bytestream:
        magic = _read32(bytestream)
        if magic != 2051:
            raise ValueError(
                'Invalid magic number %d in MNIST image file: %s' %
                (magic, f.name))
        num_images = _read32(bytestream)
        rows = _read32(bytestream)
        cols = _read32(bytestream)
        buf = bytestream.read(rows * cols * num_images)
        data = numpy.frombuffer(buf, dtype=numpy.uint8)
        data = data.reshape(num_images, rows, cols, 1)
        return data

def extract_labels(f):
    print('Extracting', f.name)
    with gzip.GzipFile(fileobj=f) as bytestream:
        magic = _read32(bytestream)
        if magic != 2049:
            raise ValueError(
                'Invalid magic number %d in MNIST label file: %s' %
                (magic, f.name))
        num_items = _read32(bytestream)
        buf = bytestream.read(num_items)
        labels = numpy.frombuffer(buf, dtype=numpy.uint8)
        return labels

def read_data_sets(train_dir, data_type="train"):
    """
    Parse or download mnist data if train_dir is empty.

    :param: train_dir: The directory storing the mnist data

    :param: data_type: Reading training set or testing set.It can be either "train" or "test"

    :return:

    ```
    (ndarray, ndarray) representing (features, labels)
    features is a 4D unit8 numpy array [index, y, x, depth] representing each pixel valued from 0 to 255.
    labels is 1D unit8 nunpy array representing the label valued from 0 to 9.
    ```

    """
    TRAIN_IMAGES = 'train-images-idx3-ubyte.gz'
    TRAIN_LABELS = 'train-labels-idx1-ubyte.gz'
    TEST_IMAGES = 't10k-images-idx3-ubyte.gz'
    TEST_LABELS = 't10k-labels-idx1-ubyte.gz'

    if data_type == "train":
        local_file = base.maybe_download(TRAIN_IMAGES, train_dir,
                                         SOURCE_URL + TRAIN_IMAGES)
        with open(local_file, 'rb') as f:
            train_images = extract_images(f)

        local_file = base.maybe_download(TRAIN_LABELS, train_dir,
                                         SOURCE_URL + TRAIN_LABELS)
        with open(local_file, 'rb') as f:
            train_labels = extract_labels(f)
        return train_images, train_labels

    else:
        local_file = base.maybe_download(TEST_IMAGES, train_dir,
                                         SOURCE_URL + TEST_IMAGES)
        with open(local_file, 'rb') as f:
            test_images = extract_images(f)

        local_file = base.maybe_download(TEST_LABELS, train_dir,
                                         SOURCE_URL + TEST_LABELS)
        with open(local_file, 'rb') as f:
            test_labels = extract_labels(f)
        return test_images, test_labels

def load_data(location="/tmp/mnist"):
    (train_images, train_labels) = read_data_sets(location, "train")
    (test_images, test_labels) = read_data_sets(location, "test")
    X_train = normalizer(train_images, TRAIN_MEAN, TRAIN_STD)
    X_test = normalizer(test_images, TRAIN_MEAN, TRAIN_STD)
    Y_train = train_labels + 1
    Y_test = test_labels + 1
    return (X_train, Y_train), (X_test, Y_test)

if __name__ == "__main__":
    train, _ = read_data_sets("/tmp/mnist/", "train")
    test, _ = read_data_sets("/tmp/mnist", "test")
    assert numpy.abs(numpy.mean(train) - TRAIN_MEAN) / TRAIN_MEAN < 1e-7
    assert numpy.abs(numpy.std(train) - TRAIN_STD) / TRAIN_STD < 1e-7
    assert numpy.abs(numpy.mean(test) - TEST_MEAN) / TEST_MEAN < 1e-7
    assert numpy.abs(numpy.std(test) - TEST_STD) / TEST_STD < 1e-7


#
# Copyright 2016 The BigDL 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.
#

import tarfile
from bigdl.dataset import base
import os
import sys

NEWS20_URL = 'http://qwone.com/~jason/20Newsgroups/20news-18828.tar.gz'  # noqa
GLOVE_URL = 'http://nlp.stanford.edu/data/glove.6B.zip'  # noqa

CLASS_NUM = 20

def download_news20(dest_dir):
    file_name = "20news-18828.tar.gz"
    file_abs_path = base.maybe_download(file_name, dest_dir, NEWS20_URL)
    tar = tarfile.open(file_abs_path, "r:gz")
    extracted_to = os.path.join(dest_dir, "20news-18828")
    if not os.path.exists(extracted_to):
        print("Extracting %s to %s" % (file_abs_path, extracted_to))
        tar.extractall(dest_dir)
        tar.close()
    return extracted_to

def download_glove_w2v(dest_dir):
    file_name = "glove.6B.zip"
    file_abs_path = base.maybe_download(file_name, dest_dir, GLOVE_URL)
    import zipfile
    zip_ref = zipfile.ZipFile(file_abs_path, 'r')
    extracted_to = os.path.join(dest_dir, "glove.6B")
    if not os.path.exists(extracted_to):
        print("Extracting %s to %s" % (file_abs_path, extracted_to))
        zip_ref.extractall(extracted_to)
        zip_ref.close()
    return extracted_to

def get_news20(source_dir="./data/news20/"):
    """
    Parse or download news20 if source_dir is empty.

    :param source_dir: The directory storing news data.
    :return: A list of (tokens, label)
    """
    news_dir = download_news20(source_dir)
    texts = []  # list of text samples
    label_id = 0
    for name in sorted(os.listdir(news_dir)):
        path = os.path.join(news_dir, name)
        label_id += 1
        if os.path.isdir(path):
            for fname in sorted(os.listdir(path)):
                if fname.isdigit():
                    fpath = os.path.join(path, fname)
                    if sys.version_info < (3,):
                        f = open(fpath)
                    else:
                        f = open(fpath, encoding='latin-1')
                    content = f.read()
                    texts.append((content, label_id))
                    f.close()

    print('Found %s texts.' % len(texts))
    return texts

def get_glove_w2v(source_dir="./data/news20/", dim=100):
    """
    Parse or download the pre-trained glove word2vec if source_dir is empty.

    :param source_dir: The directory storing the pre-trained word2vec
    :param dim: The dimension of a vector
    :return: A dict mapping from word to vector
    """
    w2v_dir = download_glove_w2v(source_dir)
    w2v_path = os.path.join(w2v_dir, "glove.6B.%sd.txt" % dim)
    if sys.version_info < (3,):
        w2v_f = open(w2v_path)
    else:
        w2v_f = open(w2v_path, encoding='latin-1')
    pre_w2v = {}
    for line in w2v_f.readlines():
        items = line.split(" ")
        pre_w2v[items[0]] = [float(i) for i in items[1:]]
    w2v_f.close()
    return pre_w2v

if __name__ == "__main__":
    get_news20("./data/news20/")
    get_glove_w2v("./data/news20/")


#
# Copyright 2016 The BigDL 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.
#

from bigdl.nn.keras.layer import KerasLayer
from bigdl.optim.optimizer import *
from bigdl.nn.criterion import *
import multiprocessing

from bigdl.nn.layer import SharedStaticUtils, Container

class KerasModel(KerasLayer, Container, SharedStaticUtils):
    def __convert_optim_method(self, optimizer):
        optimizer = optimizer.lower()
        if optimizer == "adagrad":
            return Adagrad(learningrate=0.01)
        elif optimizer == "sgd":
            return SGD(learningrate=0.01)
        elif optimizer == "adam":
            return Adam()
        elif optimizer == "rmsprop":
            return RMSprop(learningrate=0.001, decayrate=0.9)
        elif optimizer == "adadelta":
            return Adadelta(decayrate=0.95, epsilon=1e-8)
        elif optimizer == "adamax":
            return Adamax(epsilon=1e-8)
        else:
            raise TypeError("Unsupported optimizer: %s" % optimizer)

    def __convert_criterion(self, criterion):
        criterion = criterion.lower()
        if criterion == "categorical_crossentropy":
            return CategoricalCrossEntropy()
        elif criterion == "mse" or criterion == "mean_squared_error":
            return MSECriterion()
        elif criterion == "binary_crossentropy":
            return BCECriterion()
        elif criterion == "mae" or criterion == "mean_absolute_error":
            return AbsCriterion()
        elif criterion == "hinge":
            return MarginCriterion()
        elif criterion == "mean_absolute_percentage_error" or criterion == "mape":
            return MeanAbsolutePercentageCriterion()
        elif criterion == "mean_squared_logarithmic_error" or criterion == "msle":
            return MeanSquaredLogarithmicCriterion()
        elif criterion == "squared_hinge":
            return MarginCriterion(squared=True)
        elif criterion == "sparse_categorical_crossentropy":
            return ClassNLLCriterion(logProbAsInput=False)
        elif criterion == "kullback_leibler_divergence" or criterion == "kld":
            return KullbackLeiblerDivergenceCriterion()
        elif criterion == "poisson":
            return PoissonCriterion()
        elif criterion == "cosine_proximity" or criterion == "cosine":
            return CosineProximityCriterion()
        else:
            raise TypeError("Unsupported loss: %s" % criterion)

    def __convert_metrics(self, metrics):
        metrics = to_list(metrics)
        bmetrics = []
        for metric in metrics:
            if metric.lower() == "accuracy":
                bmetrics.append(Top1Accuracy())
            else:
                raise TypeError("Unsupported metrics: %s" % metric)
        return bmetrics

    def compile(self, optimizer, loss, metrics=None):
        """
        Configures the learning process. Must be called before fit or evaluate.

        # Arguments
        optimizer: Optimization method to be used. One can alternatively pass in the corresponding
                   string representation, such as 'sgd'.
        loss: Criterion to be used. One can alternatively pass in the corresponding string
              representation, such as 'mse'.
        metrics: List of validation methods to be used. Default is None. One can alternatively use ['accuracy'].
        """
        if isinstance(optimizer, six.string_types):
            optimizer = self.__convert_optim_method(optimizer)
        if isinstance(loss, six.string_types):
            loss = self.__convert_criterion(loss)
        if all(isinstance(metric, six.string_types) for metric in metrics):
            metrics = self.__convert_metrics(metrics)
        callBigDlFunc(self.bigdl_type, "compile",
                      self.value,
                      optimizer,
                      loss,
                      metrics)

    def fit(self, x, y=None, batch_size=32, nb_epoch=10, validation_data=None, distributed=True):
        """
        Train a model for a fixed number of epochs on a dataset.

        # Arguments
        x: Input data. A Numpy array or RDD of Sample or Image DataSet.
        y: Labels. A Numpy array. Default is None if x is already RDD of Sample or Image DataSet.
        batch_size: Number of samples per gradient update.
        nb_epoch: Number of iterations to train.
        validation_data: Tuple (x_val, y_val) where x_val and y_val are both Numpy arrays.
                         Or RDD of Sample. Default is None if no validation is involved.
        distributed: Boolean. Whether to train the model in distributed mode or local mode.
                     Default is True. In local mode, x and y must both be Numpy arrays.
        """
        if distributed:
            if isinstance(x, np.ndarray) and isinstance(y, np.ndarray):
                training_data = to_sample_rdd(x, y)
                if validation_data:
                    validation_data = to_sample_rdd(*validation_data)
            elif (isinstance(x, RDD) or isinstance(x, DataSet)) and not y:
                training_data = x
            else:
                raise TypeError("Unsupported training data type: %s" % type(x))
            callBigDlFunc(self.bigdl_type, "fit",
                          self.value,
                          training_data,
                          batch_size,
                          nb_epoch,
                          validation_data)
        else:
            if validation_data:
                val_x = [JTensor.from_ndarray(x) for x in to_list(validation_data[0])]
                val_y = JTensor.from_ndarray(validation_data[1])
            else:
                val_x, val_y = None, None
            callBigDlFunc(self.bigdl_type, "fit",
                          self.value,
                          [JTensor.from_ndarray(x) for x in to_list(x)],
                          JTensor.from_ndarray(y),
                          batch_size,
                          nb_epoch,
                          val_x,
                          val_y,
                          multiprocessing.cpu_count())

    def evaluate(self, x, y=None, batch_size=32):
        """
        Evaluate a model on a given dataset in distributed mode.

        # Arguments
        x: Input data. A Numpy array or RDD of Sample.
        y: Labels. A Numpy array. Default is None if x is already RDD of Sample.
        batch_size: Number of samples per gradient update.
        """
        if isinstance(x, np.ndarray) and isinstance(y, np.ndarray):
            evaluation_data = to_sample_rdd(x, y)
        elif isinstance(x, RDD) and not y:
            evaluation_data = x
        else:
            raise TypeError("Unsupported evaluation data type: %s" % type(x))
        return callBigDlFunc(self.bigdl_type, "evaluate",
                             self.value,
                             evaluation_data,
                             batch_size)

    def predict(self, x, distributed=True):
        """
        Use a model to do prediction.

        # Arguments
        x: Input data. A Numpy array or RDD of Sample.
        distributed: Boolean. Whether to do prediction in distributed mode or local mode.
                     Default is True. In local mode, x must be a Numpy array.
        """
        if is_distributed:
            if isinstance(x, np.ndarray):
                features = to_sample_rdd(x, np.zeros([x.shape[0]]))
            elif isinstance(x, RDD):
                features = x
            else:
                raise TypeError("Unsupported prediction data type: %s" % type(x))
            return self.predict_distributed(features)
        else:
            if isinstance(x, np.ndarray):
                return self.predict_local(x)
            else:
                raise TypeError("Unsupported prediction data type: %s" % type(x))

class Sequential(KerasModel):
    """
    Container for a sequential model.

    # Arguments
    name: String to specify the name of the sequential model. Default is None.

    >>> sequential = Sequential(name="seq1")
    creating: createKerasSequential
    """
    def __init__(self, jvalue=None, **kwargs):
        super(Sequential, self).__init__(jvalue, **kwargs)

    @staticmethod
    def from_jvalue(jvalue, bigdl_type="float"):
        """
        Create a Python Model base on the given java value
        :param jvalue: Java object create by Py4j
        :return: A Python Model
        """
        model = Sequential(jvalue=jvalue)
        model.value = jvalue
        return model

    def add(self, model):
        self.value.add(model.value)
        return self

class Model(KerasModel):
    """
    Container for a graph model.

    # Arguments
    input: An input node or a list of input nodes.
    output: An output node or a list of output nodes.
    name: String to specify the name of the graph model. Default is None.
    """
    def __init__(self, input, output, jvalue=None,  **kwargs):
        super(Model, self).__init__(jvalue,
                                    to_list(input),
                                    to_list(output),
                                    **kwargs)
    @staticmethod
    def from_jvalue(jvalue, bigdl_type="float"):
        """
        Create a Python Model base on the given java value
        :param jvalue: Java object create by Py4j
        :return: A Python Model
        """
        model = Model([], [], jvalue=jvalue)
        model.value = jvalue
        return model


#
# Copyright 2016 The BigDL 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.
#

from bigdl.dataset import mnist
from bigdl.dataset.transformer import *
from bigdl.optim.optimizer import *

def get_mnist(sc, data_type="train", location="/tmp/mnist"):
    """
    Get mnist dataset and parallelize into RDDs.
    Data would be downloaded automatically if it doesn't present at the specific location.

    :param sc: SparkContext.
    :param data_type: "train" for training data and "test" for testing data.
    :param location: Location to store mnist dataset.
    :return: RDD of (features: ndarray, label: ndarray).
    """
    (images, labels) = mnist.read_data_sets(location, data_type)
    images = sc.parallelize(images)
    labels = sc.parallelize(labels + 1)  # Target start from 1 in BigDL
    record = images.zip(labels)
    return record

def preprocess_mnist(sc, options):
    """
    Preprocess mnist dataset.
    Normalize and transform into Sample of RDDs.
    """
    train_data = get_mnist(sc, "train", options.dataPath)\
        .map(lambda rec_tuple: (normalizer(rec_tuple[0], mnist.TRAIN_MEAN, mnist.TRAIN_STD),
                                rec_tuple[1]))\
        .map(lambda t: Sample.from_ndarray(t[0], t[1]))
    test_data = get_mnist(sc, "test", options.dataPath)\
        .map(lambda rec_tuple: (normalizer(rec_tuple[0], mnist.TEST_MEAN, mnist.TEST_STD),
                                rec_tuple[1]))\
        .map(lambda t: Sample.from_ndarray(t[0], t[1]))
    return train_data, test_data

def get_end_trigger(options):
    """
    When to end the optimization based on input option.
    """
    if options.endTriggerType.lower() == "epoch":
        return MaxEpoch(options.endTriggerNum)
    else:
        return MaxIteration(options.endTriggerNum)

def validate_optimizer(optimizer, test_data, options):
    """
    Set validation and checkpoint for distributed optimizer.
    """
    optimizer.set_validation(
        batch_size=options.batchSize,
        val_rdd=test_data,
        trigger=EveryEpoch(),
        val_method=[Top1Accuracy()]
    )
    optimizer.set_checkpoint(EveryEpoch(), options.checkpointPath)