class Loss:
  """A loss function for use in training models."""

  def _compute_tf_loss(self, output, labels):
    """Compute the loss function for TensorFlow tensors.

    The inputs are tensors containing the model's outputs and the labels for a
    batch.  The return value should be a tensor of shape (batch_size) or
    (batch_size, tasks) containing the value of the loss function on each
    sample or sample/task.

    Parameters
    ----------
    output: tensor
      the output of the model
    labels: tensor
      the expected output

    Returns
    -------
    The value of the loss function on each sample or sample/task pair
    """
    raise NotImplementedError("Subclasses must implement this")

  def _create_pytorch_loss(self):
    """Create a PyTorch loss function."""
    raise NotImplementedError("Subclasses must implement this")


class L1Loss(Loss):
  """The absolute difference between the true and predicted values."""

  def _compute_tf_loss(self, output, labels):
    import tensorflow as tf
    output, labels = _make_tf_shapes_consistent(output, labels)
    output, labels = _ensure_float(output, labels)
    return tf.abs(output - labels)

  def _create_pytorch_loss(self):
    import torch

    def loss(output, labels):
      output, labels = _make_pytorch_shapes_consistent(output, labels)
      return torch.nn.functional.l1_loss(output, labels, reduction='none')

    return loss


class HuberLoss(Loss):
  """Modified version of L1 Loss, also known as Smooth L1 loss.
  Less sensitive to small errors, linear for larger errors.
  Huber loss is generally better for cases where are are both large outliers as well as small, as compared to the L1 loss.
  By default, Delta = 1.0 and reduction = 'none'.
  """

  def _compute_tf_loss(self, output, labels):
    import tensorflow as tf
    output, labels = _make_tf_shapes_consistent(output, labels)
    return tf.keras.losses.Huber(reduction='none')(output, labels)

  def _create_pytorch_loss(self):
    import torch

    def loss(output, labels):
      output, labels = _make_pytorch_shapes_consistent(output, labels)
      return torch.nn.functional.smooth_l1_loss(
          output, labels, reduction='none')

    return loss


class L2Loss(Loss):
  """The squared difference between the true and predicted values."""

  def _compute_tf_loss(self, output, labels):
    import tensorflow as tf
    output, labels = _make_tf_shapes_consistent(output, labels)
    output, labels = _ensure_float(output, labels)
    return tf.square(output - labels)

  def _create_pytorch_loss(self):
    import torch

    def loss(output, labels):
      output, labels = _make_pytorch_shapes_consistent(output, labels)
      return torch.nn.functional.mse_loss(output, labels, reduction='none')

    return loss


class HingeLoss(Loss):
  """The hinge loss function.

  The 'output' argument should contain logits, and all elements of 'labels'
  should equal 0 or 1.
  """

  def _compute_tf_loss(self, output, labels):
    import tensorflow as tf
    output, labels = _make_tf_shapes_consistent(output, labels)
    return tf.keras.losses.hinge(labels, output)

  def _create_pytorch_loss(self):
    import torch

    def loss(output, labels):
      output, labels = _make_pytorch_shapes_consistent(output, labels)
      return torch.mean(torch.clamp(1 - labels * output, min=0), dim=-1)

    return loss


class SquaredHingeLoss(Loss):
  """The Squared Hinge loss function.
  
  Defined as the square of the hinge loss between y_true and y_pred. The Squared Hinge Loss is differentiable.
  """

  def _compute_tf_loss(self, output, labels):
    import tensorflow as tf
    output, labels = _make_tf_shapes_consistent(output, labels)
    return tf.keras.losses.SquaredHinge(reduction='none')(labels, output)

  def _create_pytorch_loss(self):
    import torch

    def loss(output, labels):
      output, labels = _make_pytorch_shapes_consistent(output, labels)
      return torch.mean(
          torch.pow(
              torch.max(1 - torch.mul(labels, output), torch.tensor(0.0)), 2),
          dim=-1)

    return loss


class PoissonLoss(Loss):
  """The Poisson loss function is defined as the mean of the elements of y_pred - (y_true * log(y_pred) for an input of (y_true, y_pred).
  Poisson loss is generally used for regression tasks where the data follows the poisson
  """

  def _compute_tf_loss(self, output, labels):
    import tensorflow as tf
    output, labels = _make_tf_shapes_consistent(output, labels)
    loss = tf.keras.losses.Poisson(reduction='auto')
    return loss(labels, output)

  def _create_pytorch_loss(self):
    import torch

    def loss(output, labels):
      output, labels = _make_pytorch_shapes_consistent(output, labels)
      return torch.mean(output - labels * torch.log(output))

    return loss


class BinaryCrossEntropy(Loss):
  """The cross entropy between pairs of probabilities.

  The arguments should each have shape (batch_size) or (batch_size, tasks) and
  contain probabilities.
  """

  def _compute_tf_loss(self, output, labels):
    import tensorflow as tf
    output, labels = _make_tf_shapes_consistent(output, labels)
    output, labels = _ensure_float(output, labels)
    return tf.keras.losses.binary_crossentropy(labels, output)

  def _create_pytorch_loss(self):
    import torch
    bce = torch.nn.BCELoss(reduction='none')

    def loss(output, labels):
      output, labels = _make_pytorch_shapes_consistent(output, labels)
      return torch.mean(bce(output, labels), dim=-1)

    return loss


class CategoricalCrossEntropy(Loss):
  """The cross entropy between two probability distributions.

  The arguments should each have shape (batch_size, classes) or
  (batch_size, tasks, classes), and represent a probability distribution over
  classes.
  """

  def _compute_tf_loss(self, output, labels):
    import tensorflow as tf
    output, labels = _make_tf_shapes_consistent(output, labels)
    output, labels = _ensure_float(output, labels)
    return tf.keras.losses.categorical_crossentropy(labels, output)

  def _create_pytorch_loss(self):
    import torch

    def loss(output, labels):
      output, labels = _make_pytorch_shapes_consistent(output, labels)
      return -torch.sum(labels * torch.log(output), dim=-1)

    return loss


class SigmoidCrossEntropy(Loss):
  """The cross entropy between pairs of probabilities.

  The arguments should each have shape (batch_size) or (batch_size, tasks).  The
  labels should be probabilities, while the outputs should be logits that are
  converted to probabilities using a sigmoid function.
  """

  def _compute_tf_loss(self, output, labels):
    import tensorflow as tf
    output, labels = _make_tf_shapes_consistent(output, labels)
    output, labels = _ensure_float(output, labels)
    return tf.nn.sigmoid_cross_entropy_with_logits(labels, output)

  def _create_pytorch_loss(self):
    import torch
    bce = torch.nn.BCEWithLogitsLoss(reduction='none')

    def loss(output, labels):
      output, labels = _make_pytorch_shapes_consistent(output, labels)
      return bce(output, labels)

    return loss


class SoftmaxCrossEntropy(Loss):
  """The cross entropy between two probability distributions.

  The arguments should each have shape (batch_size, classes) or
  (batch_size, tasks, classes).  The labels should be probabilities, while the
  outputs should be logits that are converted to probabilities using a softmax
  function.
  """

  def _compute_tf_loss(self, output, labels):
    import tensorflow as tf
    output, labels = _make_tf_shapes_consistent(output, labels)
    output, labels = _ensure_float(output, labels)
    return tf.nn.softmax_cross_entropy_with_logits(labels, output)

  def _create_pytorch_loss(self):
    import torch
    ls = torch.nn.LogSoftmax(dim=-1)

    def loss(output, labels):
      output, labels = _make_pytorch_shapes_consistent(output, labels)
      return -torch.sum(labels * ls(output), dim=-1)

    return loss


class SparseSoftmaxCrossEntropy(Loss):
  """The cross entropy between two probability distributions.

  The labels should have shape (batch_size) or (batch_size, tasks), and be
  integer class labels.  The outputs have shape (batch_size, classes) or
  (batch_size, tasks, classes) and be logits that are converted to probabilities
  using a softmax function.
  """

  def _compute_tf_loss(self, output, labels):
    import tensorflow as tf

    if len(labels.shape) == len(output.shape):
      labels = tf.squeeze(labels, axis=-1)

    labels = tf.cast(labels, tf.int32)

    return tf.nn.sparse_softmax_cross_entropy_with_logits(labels, output)

  def _create_pytorch_loss(self):
    import torch
    ce_loss = torch.nn.CrossEntropyLoss(reduction='none')

    def loss(output, labels):
      # Convert (batch_size, tasks, classes) to (batch_size, classes, tasks)
      # CrossEntropyLoss only supports (batch_size, classes, tasks)
      # This is for API consistency
      if len(output.shape) == 3:
        output = output.permute(0, 2, 1)

      if len(labels.shape) == len(output.shape):
        labels = labels.squeeze(-1)
      return ce_loss(output, labels.long())

    return loss


class VAE_ELBO(Loss):
  """The Variational AutoEncoder loss, KL Divergence Regularize + marginal log-likelihood.
  
  This losses based on _[1].
  ELBO(Evidence lower bound) lexically replaced Variational lower bound. 
  BCE means marginal log-likelihood, and KLD means KL divergence with normal distribution.
  Added hyper parameter 'kl_scale' for KLD.
  
  The logvar and mu should have shape (batch_size, hidden_space).
  The x and reconstruction_x should have (batch_size, attribute). 
  The kl_scale should be float.
  
  Examples
  --------
  Examples for calculating loss using constant tensor.
  
  batch_size = 2,
  hidden_space = 2,
  num of original attribute = 3
  >>> import numpy as np
  >>> import torch
  >>> import tensorflow as tf
  >>> logvar = np.array([[1.0,1.3],[0.6,1.2]])
  >>> mu = np.array([[0.2,0.7],[1.2,0.4]])
  >>> x = np.array([[0.9,0.4,0.8],[0.3,0,1]])
  >>> reconstruction_x = np.array([[0.8,0.3,0.7],[0.2,0,0.9]])
  
  Case tensorflow
  >>> VAE_ELBO()._compute_tf_loss(tf.constant(logvar), tf.constant(mu), tf.constant(x), tf.constant(reconstruction_x))
  <tf.Tensor: shape=(2,), dtype=float64, numpy=array([0.70165154, 0.76238271])>
  
  Case pytorch
  >>> (VAE_ELBO()._create_pytorch_loss())(torch.tensor(logvar), torch.tensor(mu), torch.tensor(x), torch.tensor(reconstruction_x))
  tensor([0.7017, 0.7624], dtype=torch.float64)
  
  
  References
  ----------
  .. [1] Kingma, Diederik P., and Max Welling. "Auto-encoding variational bayes." arXiv preprint arXiv:1312.6114 (2013).
  
  """

  def _compute_tf_loss(self, logvar, mu, x, reconstruction_x, kl_scale=1):
    import tensorflow as tf
    x, reconstruction_x = _make_tf_shapes_consistent(x, reconstruction_x)
    x, reconstruction_x = _ensure_float(x, reconstruction_x)
    BCE = tf.keras.losses.binary_crossentropy(x, reconstruction_x)
    KLD = VAE_KLDivergence()._compute_tf_loss(logvar, mu)
    return BCE + kl_scale * KLD

  def _create_pytorch_loss(self):
    import torch
    bce = torch.nn.BCELoss(reduction='none')

    def loss(logvar, mu, x, reconstruction_x, kl_scale=1):
      x, reconstruction_x = _make_pytorch_shapes_consistent(x, reconstruction_x)
      BCE = torch.mean(bce(reconstruction_x, x), dim=-1)
      KLD = (VAE_KLDivergence()._create_pytorch_loss())(logvar, mu)
      return BCE + kl_scale * KLD

    return loss


class VAE_KLDivergence(Loss):
  """The KL_divergence between hidden distribution and normal distribution.
  
  This loss represents KL divergence losses between normal distribution(using parameter of distribution)
  based on  _[1].
  
  The logvar should have shape (batch_size, hidden_space) and each term represents
  standard deviation of hidden distribution. The mean shuold have 
  (batch_size, hidden_space) and each term represents mean of hidden distribtuon.
  
  Examples
  --------
  Examples for calculating loss using constant tensor.
  
  batch_size = 2,
  hidden_space = 2,
  >>> import numpy as np
  >>> import torch
  >>> import tensorflow as tf
  >>> logvar = np.array([[1.0,1.3],[0.6,1.2]])
  >>> mu = np.array([[0.2,0.7],[1.2,0.4]])
  
  Case tensorflow
  >>> VAE_KLDivergence()._compute_tf_loss(tf.constant(logvar), tf.constant(mu))
  <tf.Tensor: shape=(2,), dtype=float64, numpy=array([0.17381787, 0.51425203])>
  
  Case pytorch
  >>> (VAE_KLDivergence()._create_pytorch_loss())(torch.tensor(logvar), torch.tensor(mu))
  tensor([0.1738, 0.5143], dtype=torch.float64)
  
  References
  ----------
  .. [1] Kingma, Diederik P., and Max Welling. "Auto-encoding variational bayes." arXiv preprint arXiv:1312.6114 (2013).
  
  """

  def _compute_tf_loss(self, logvar, mu):
    import tensorflow as tf
    logvar, mu = _make_tf_shapes_consistent(logvar, mu)
    logvar, mu = _ensure_float(logvar, mu)
    return 0.5 * tf.reduce_mean(
        tf.square(mu) + tf.square(logvar) -
        tf.math.log(1e-20 + tf.square(logvar)) - 1, -1)

  def _create_pytorch_loss(self):
    import torch

    def loss(logvar, mu):
      logvar, mu = _make_pytorch_shapes_consistent(logvar, mu)
      return 0.5 * torch.mean(
          torch.square(mu) + torch.square(logvar) -
          torch.log(1e-20 + torch.square(logvar)) - 1, -1)

    return loss


class ShannonEntropy(Loss):
  """The ShannonEntropy of discrete-distribution.
  
  This loss represents shannon entropy based on _[1].
  
  The inputs should have shape (batch size, num of variable) and represents
  probabilites distribution.
  
  Examples
  --------
  Examples for calculating loss using constant tensor.
  
  batch_size = 2,
  num_of variable = variable,
  >>> import numpy as np
  >>> import torch
  >>> import tensorflow as tf
  >>> inputs = np.array([[0.7,0.3],[0.9,0.1]])
  
  Case tensorflow
  >>> ShannonEntropy()._compute_tf_loss(tf.constant(inputs))
  <tf.Tensor: shape=(2,), dtype=float64, numpy=array([0.30543215, 0.16254149])>
  
  Case pytorch
  >>> (ShannonEntropy()._create_pytorch_loss())(torch.tensor(inputs))
  tensor([0.3054, 0.1625], dtype=torch.float64)
  
  References
  ----------
  .. [1] Chen, Ricky Xiaofeng. "A Brief Introduction to Shannon’s Information Theory." arXiv preprint arXiv:1612.09316 (2016).
  
  """

  def _compute_tf_loss(self, inputs):
    import tensorflow as tf
    #extended one of probabilites to binary distribution
    if inputs.shape[-1] == 1:
      inputs = tf.concat([inputs, 1 - inputs], axis=-1)
    return tf.reduce_mean(-inputs * tf.math.log(1e-20 + inputs), -1)

  def _create_pytorch_loss(self):
    import torch

    def loss(inputs):
      #extended one of probabilites to binary distribution
      if inputs.shape[-1] == 1:
        inputs = torch.cat((inputs, 1 - inputs), dim=-1)
      return torch.mean(-inputs * torch.log(1e-20 + inputs), -1)

    return loss


def _make_tf_shapes_consistent(output, labels):
  """Try to make inputs have the same shape by adding dimensions of size 1."""
  import tensorflow as tf
  shape1 = output.shape
  shape2 = labels.shape
  len1 = len(shape1)
  len2 = len(shape2)
  if len1 == len2:
    return (output, labels)
  if isinstance(shape1, tf.TensorShape):
    shape1 = tuple(shape1.as_list())
  if isinstance(shape2, tf.TensorShape):
    shape2 = tuple(shape2.as_list())
  if len1 > len2 and all(i == 1 for i in shape1[len2:]):
    for i in range(len1 - len2):
      labels = tf.expand_dims(labels, -1)
    return (output, labels)
  if len2 > len1 and all(i == 1 for i in shape2[len1:]):
    for i in range(len2 - len1):
      output = tf.expand_dims(output, -1)
    return (output, labels)
  raise ValueError("Incompatible shapes for outputs and labels: %s versus %s" %
                   (str(shape1), str(shape2)))


def _make_pytorch_shapes_consistent(output, labels):
  """Try to make inputs have the same shape by adding dimensions of size 1."""
  import torch
  shape1 = output.shape
  shape2 = labels.shape
  len1 = len(shape1)
  len2 = len(shape2)
  if len1 == len2:
    return (output, labels)
  shape1 = tuple(shape1)
  shape2 = tuple(shape2)
  if len1 > len2 and all(i == 1 for i in shape1[len2:]):
    for i in range(len1 - len2):
      labels = torch.unsqueeze(labels, -1)
    return (output, labels)
  if len2 > len1 and all(i == 1 for i in shape2[len1:]):
    for i in range(len2 - len1):
      output = torch.unsqueeze(output, -1)
    return (output, labels)
  raise ValueError("Incompatible shapes for outputs and labels: %s versus %s" %
                   (str(shape1), str(shape2)))


def _ensure_float(output, labels):
  """Make sure the outputs and labels are both floating point types."""
  import tensorflow as tf
  if output.dtype not in (tf.float32, tf.float64):
    output = tf.cast(output, tf.float32)
  if labels.dtype not in (tf.float32, tf.float64):
    labels = tf.cast(labels, tf.float32)
  return (output, labels)
