# Copyright 2019 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.
# ============================================================================
"""Experimental Numpy backend."""

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

import contextlib
import importlib
import logging
import types

import numpy as onp
import jax.numpy as np


from tensorflow_probability.python.internal.backend.jax import _utils as utils


__all__ = [
    'parameterized_truncated_normal',
    'prevent_gradient',
    'remove_undocumented',
]


JAX_MODE = True


if JAX_MODE:
  from jax import random  # pylint: disable=g-import-not-at-top


def _parameterized_truncated_normal(shape, means=0., stddevs=1.,
                                    minvals=-2, maxvals=2., dtype=np.float32,
                                    seed=None, name=None):  # pylint: disable=unused-argument
  """Implementation of ops.random_ops.parameterized_truncated_normal."""

  # NOTE: The docstring for parameterized_truncated_normal is wrong.  The actual
  # requirements are:
  #  - Parameter `shape` must represent a shape with rank at least 1.
  #  - Parameters `means`, `stddevs`, `minvals`, `maxvals` can be rank 0 or
  #    rank 1.  If rank 1, then they must have shape `[shape[0]]`.

  def _right_expand(x, shape):
    if (np.ndim(x) == 0) or (np.ndim(x) >= len(shape)):
      return x
    # Add size-1 dimensions on the right to `x` so its rank is `len(shape)`.
    return np.reshape(x, np.shape(x) + (1,) * (len(shape) - 1))

  if JAX_MODE:
    min_z = (minvals - means) / stddevs
    max_z = (maxvals - means) / stddevs

    min_z = _right_expand(min_z, shape)
    max_z = _right_expand(max_z, shape)
    means = _right_expand(means, shape)
    stddevs = _right_expand(stddevs, shape)

    z = random.truncated_normal(
        seed, lower=min_z, upper=max_z, shape=shape, dtype=dtype)
    return z * stddevs + means

  raise NotImplementedError


def remove_undocumented(module_name, allowed_exception_list=None,  # pylint: disable=unused-argument
                        doc_string_modules=None):  # pylint: disable=unused-argument
  pass


def GraphOrParentsInXlaContext(*args, **kwargs):  # pylint: disable=invalid-name
  del args, kwargs
  return True


def numpy_text(tensor, is_repr=False):
  if is_repr:
    return repr(tensor)
  return str(tensor)


@contextlib.contextmanager
def eager_mode():
  yield


class LazyLoader(types.ModuleType):
  """Reimplementation of TF's LazyLoader."""

  def __init__(self, local_name, parent_module_globals, name, warning=None):  # pylint: disable=super-on-old-class
    self._local_name = local_name
    self._parent_module_globals = parent_module_globals
    self._warning = warning

    super(LazyLoader, self).__init__(name)

  def _load(self):
    """Load the module and insert it into the parent's globals."""
    # Import the target module and insert it into the parent's namespace
    module = importlib.import_module(self.__name__)
    self._parent_module_globals[self._local_name] = module

    # Emit a warning if one was specified
    if self._warning:
      logging.warning(self._warning)
      # Make sure to only warn once.
      self._warning = None

    # Update this object's dict so that if someone keeps a reference to the
    #   LazyLoader, lookups are efficient (__getattr__ is only called on lookups
    #   that fail).
    self.__dict__.update(module.__dict__)

    return module

  def __getattr__(self, item):
    module = self._load()
    return getattr(module, item)

  def __dir__(self):
    module = self._load()
    return dir(module)


def max_error(*args, **kwargs):
  del args, kwargs
  raise NotImplementedError

parameterized_truncated_normal = utils.copy_docstring(
    'random_ops.parameterized_truncated_normal',
    _parameterized_truncated_normal)


def _prevent_gradient(input, message='', name=None):  # pylint: disable=unused-argument,redefined-builtin
  raise NotImplementedError

prevent_gradient = utils.copy_docstring(
    'array_ops.prevent_gradient',
    _prevent_gradient)

