# coding=utf-8
# Copyright 2021 Google LLC.
#
# 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.

# Lint as: python3
"""Several cost/norm functions for relevant vector types."""
import abc

import jax
from jax.lib import xla_bridge
import jax.numpy as jnp
from ott.geometry import matrix_square_root


def dot(x: jnp.ndarray, y: jnp.ndarray):
  """Accelerator dependent dot. Implemented to avoid OOMs with online mode."""
  platform = xla_bridge.get_backend().platform
  return jnp.sum(x * y) if platform == 'gpu' else jnp.vdot(x, y)


@jax.tree_util.register_pytree_node_class
class CostFn(abc.ABC):
  """A generic cost function, taking two vectors as input.

  Cost functions evaluate a function on a pair of inputs. For convenience,
  that function is split into two norms -- evaluated on each input separately --
  followed by a pairwise cost that involves both inputs, as in

  c(x,y) = norm(x) + norm(y) + pairwise(x,y)

  If the norm function is not implemented, that value is handled as a 0.
  """

  norm = None  #  no norm function created by default.

  @abc.abstractmethod
  def pairwise(self, x, y):
    pass

  def __call__(self, x, y):
    return self.pairwise(x, y) + (
        0 if self.norm is None else self.norm(x) + self.norm(y))  # pylint: disable=not-callable

  def all_pairs(self, x: jnp.ndarray, y: jnp.ndarray):
    """Computes matrix of all costs (including norms) for vectors in x / y.

    Args:
      x: [num_a, d] jnp.ndarray
      y: [num_b, d] jnp.ndarray
    Returns:
      [num_a, num_b] matrix of cost evaluations.
    """
    return jax.vmap(lambda x_: jax.vmap(lambda y_: self(x_, y_))(y))(x)

  def all_pairs_pairwise(self, x: jnp.ndarray, y: jnp.ndarray):
    """Computes matrix of all pairwise-costs (no norms) for vectors in x / y.

    Args:
      x: [num_a, d] jnp.ndarray
      y: [num_b, d] jnp.ndarray
    Returns:
      [num_a, num_b] matrix of pairwise cost evaluations.
    """
    return jax.vmap(lambda x_: jax.vmap(lambda y_: self.pairwise(x_, y_))(y))(x)

  def tree_flatten(self):
    return (), None

  @classmethod
  def tree_unflatten(cls, aux_data, children):
    del aux_data
    return cls(*children)


@jax.tree_util.register_pytree_node_class
class Euclidean(CostFn):
  """Squared Euclidean distance CostFn."""

  def norm(self, x):
    return jnp.sum(x ** 2, axis=-1)

  def pairwise(self, x, y):
    return -2 * dot(x, y)


@jax.tree_util.register_pytree_node_class
class Cosine(CostFn):
  """Cosine distance CostFn."""

  def __init__(self, ridge=1e-8):
    super().__init__()
    self._ridge = ridge

  def pairwise(self, x, y):
    ridge = self._ridge
    x_norm = jnp.linalg.norm(x, axis=-1)
    y_norm = jnp.linalg.norm(y, axis=-1)
    cosine_similarity = dot(x, y) / (x_norm * y_norm + ridge)
    cosine_distance = 1.0 - cosine_similarity
    return cosine_distance


@jax.tree_util.register_pytree_node_class
class Bures(CostFn):
  """Bures distance between a pair of (mean, cov matrix) raveled as vectors."""

  def __init__(self, dimension, **kwargs):
    super().__init__()
    self._dimension = dimension
    self._sqrtm_kw = kwargs

  def norm(self, x):
    norm = jnp.sum(x[..., 0:self._dimension]**2, axis=-1)
    x_mat = jnp.reshape(x[..., self._dimension:],
                        (-1, self._dimension, self._dimension))

    norm += jnp.trace(x_mat, axis1=-2, axis2=-1)
    return norm

  def pairwise(self, x, y):
    mean_dot_prod = dot(x[0:self._dimension], y[0:self._dimension])
    x_mat = jnp.reshape(x[self._dimension:],
                        (self._dimension, self._dimension))
    y_mat = jnp.reshape(y[self._dimension:],
                        (self._dimension, self._dimension))

    sq_x = matrix_square_root.sqrtm(x_mat, self._dimension,
                                    **self._sqrtm_kw)[0]
    sq_x_y_sq_x = jnp.matmul(sq_x, jnp.matmul(y_mat, sq_x))
    sq__sq_x_y_sq_x = matrix_square_root.sqrtm(sq_x_y_sq_x, self._dimension,
                                               **self._sqrtm_kw)[0]
    return -2 * (
        mean_dot_prod + jnp.trace(sq__sq_x_y_sq_x, axis1=-2, axis2=-1))

  def tree_flatten(self):
    return (), (self._dimension, self._sqrtm_kw)

  @classmethod
  def tree_unflatten(cls, aux_data, children):
    del children
    return cls(aux_data[0], **aux_data[1])


@jax.tree_util.register_pytree_node_class
class UnbalancedBures(CostFn):
  """Regularized, unbalanced Bures distance between triplets.

  This cost implements the value defined in https://arxiv.org/pdf/2006.02572.pdf
  Equation 37, 39, 40. We follow their notations. It is assumed inputs are given
  as triplets (mass, mean, covariance) raveled as vectors, in that order.
  """

  def __init__(self,
               dimension: int,
               gamma: float = 1,
               sigma: float = 1,
               **kwargs):
    super().__init__()
    self._dimension = dimension
    self._gamma = gamma
    self._sigma2 = sigma ** 2
    self._sqrtm_kw = kwargs

  def norm(self, x):
    return self._gamma * x[0]

  def pairwise(self, x, y):
    # Sets a few constants
    gam = self._gamma
    sig2 = self._sigma2
    lam = sig2 + gam / 2
    tau = gam / (2 * lam)

    # Extracts mass, mean vector, covariance matrices
    mass_x, mass_y = x[0], y[0]
    diff_means = x[1:1+self._dimension] - y[1:1+self._dimension]
    x_mat = jnp.reshape(x[1+self._dimension:],
                        (self._dimension, self._dimension))
    y_mat = jnp.reshape(y[1+self._dimension:],
                        (self._dimension, self._dimension))

    # Identity matrix of suitable size
    iden = jnp.eye(self._dimension)

    # Creates matrices needed in the computation
    tilde_a = 0.5 * gam * (iden - lam * jnp.linalg.inv(x_mat + lam * iden))
    tilde_b = 0.5 * gam * (iden - lam * jnp.linalg.inv(y_mat + lam * iden))

    tilde_a_b = jnp.matmul(tilde_a, tilde_b)
    c_mat = matrix_square_root.sqrtm(
        1 / tau * tilde_a_b + 0.25 * (sig2 ** 2) * iden, **self._sqrtm_kw)[0]
    c_mat -= 0.5 * sig2 * iden

    # Computes log determinants (their sign should be >0).
    sldet_c, ldet_c = jnp.linalg.slogdet(c_mat)
    sldet_t_ab, ldet_t_ab = jnp.linalg.slogdet(tilde_a_b)
    sldet_ab, ldet_ab = jnp.linalg.slogdet(jnp.matmul(x_mat, y_mat))
    sldet_c_ab, ldet_c_ab = jnp.linalg.slogdet(c_mat - 2 * tilde_a_b / gam)

    # Gathers all these results to compute log total mass of transport
    log_m_pi = (0.5 * self._dimension * sig2 / (gam + sig2)) * jnp.log(sig2)

    log_m_pi += (1 / (tau + 1)) * (jnp.log(mass_x) + jnp.log(mass_y) + ldet_c
                                   + 0.5 * (tau * ldet_t_ab - ldet_ab))

    log_m_pi += -jnp.sum(diff_means * jnp.linalg.solve(
        x_mat + y_mat + lam * iden, diff_means)) / (2 * (tau + 1))

    log_m_pi += - 0.5 * ldet_c_ab

    # If all logdet signs are 1, output value, nan otherwise.
    return jnp.where(
        sldet_c == 1 and sldet_c_ab == 1 and sldet_ab == 1 and sldet_t_ab == 1,
        2 * sig2 * mass_x * mass_y - 2 * (sig2 + gam) * jnp.exp(log_m_pi),
        jnp.nan)

  def tree_flatten(self):
    return (), (self._dimension, self._gamma, self._sigma2, self._sqrtm_kw)

  @classmethod
  def tree_unflatten(cls, aux_data, children):
    del children
    return cls(aux_data[0], aux_data[1], aux_data[2], **aux_data[3])
