text: # Copyright 2019 Ondrej Skopek.
  #
  # 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 numpy as np
  import torch
  from torch import tensor as t
  
  import mt.mvae.ops.poincare as P
  from mt.mvae.ops.common import eps
  
  np.random.seed(42)
  random_nums = np.random.random_sample(100) * 100 + 1
  test_eps = 5e-6
  low_prec_test_eps = 1e-4
  radius = torch.tensor(2., dtype=torch.float32)
  
  
  def is_in_hyp_space(x: torch.Tensor, eps: float = eps, radius: torch.Tensor = radius) -> torch.Tensor:
      prod = x.dot(x)
      return prod <= (1 + eps) * torch.ones_like(prod) * radius * radius
  
  
  def is_in_tangent_space(x: torch.Tensor, at_point: torch.Tensor, radius: torch.Tensor = radius, eps: float = eps) -> \
          torch.Tensor:
      # TODO-LATER: This is most likely wrong, doesn't matter for VAE though, just for tests.
      assert is_in_hyp_space(at_point, radius=radius, eps=eps)
      prod = x.dot(at_point)
      return prod.allclose(torch.zeros_like(prod), atol=eps)
  
  
  def parallel_transport(x: torch.Tensor, src: torch.Tensor, dst: torch.Tensor, radius: torch.Tensor = radius) -> \
          torch.Tensor:
      return P.pm.parallel_transport(src, dst, x, c=P._c(radius))
  
  
  def inverse_parallel_transport(x: torch.Tensor, src: torch.Tensor, dst: torch.Tensor, radius: torch.Tensor = radius) \
          -> torch.Tensor:
      return P.pm.parallel_transport(dst, src, x, c=P._c(radius))
  
  
  def test_mu_0() -> None:
      res = P.mu_0((3, 3), dtype=torch.int64)
      expected = t([[0, 0, 0], [0, 0, 0], [0, 0, 0]])
      assert res.allclose(expected)
  
  
  def test_is_in_hyp_space() -> None:
      assert is_in_hyp_space(radius * t([1., 0, 0]))
      assert not is_in_hyp_space(radius * t([0.1, 2, 3]))
      assert not is_in_hyp_space(radius * t([1, 1, np.sqrt(2)]))
      assert not is_in_hyp_space(radius * t([1, 1, np.sqrt(2)]))
      assert not is_in_hyp_space(radius * t([0, 2, -np.sqrt(2)]))
  
  
  def test_is_in_tangent_space() -> None:
      assert is_in_tangent_space(t([0., 2, 3]), radius * t([1., 0, 0]))
      assert not is_in_tangent_space(t([0.1, 2, 3]), radius * t([1., 0, 0]))
      # -0*2 + 2*1 - 2 = 0
      assert is_in_tangent_space(t([0, 2, -np.sqrt(2)]), t([2, 1, np.sqrt(2)]) / radius, eps=test_eps)
  
  
  def test_poincare_distance() -> None:
      mu0 = t([0., 0, 0])
      mu = radius * t([2., 1., np.sqrt(2)])
      assert P.poincare_distance(mu0, mu0, radius=radius).allclose(t(0.), atol=5e-4)
      assert P.poincare_distance(mu, mu, radius=radius).allclose(t(0.), atol=5e-4)
      assert P.poincare_distance(mu0, mu, radius=radius) == P.poincare_distance(mu, mu0, radius=radius)
  
  
  def test_parallel_transport() -> None:
      mu1 = t([2., 1, np.sqrt(2)]).double() / radius
      mu2 = t([np.sqrt(5), 1, np.sqrt(3)]).double() / radius
      assert is_in_hyp_space(mu1)
      assert is_in_hyp_space(mu2)
  
      u = t([0, 2, -np.sqrt(2)]).double()
      # assert is_in_tangent_space(u, at_point=mu1, eps=test_eps)
  
      assert parallel_transport(u, src=mu1, dst=mu1).allclose(u, atol=5e-4)
  
      pt_u = parallel_transport(u, src=mu1, dst=mu2)
      # assert is_in_tangent_space(pt_u, at_point=mu2, eps=test_eps)
      u_ = parallel_transport(pt_u, src=mu2, dst=mu1)
      assert u.allclose(u_, atol=5e-4)
      u_inv = inverse_parallel_transport(pt_u, src=mu1, dst=mu2)
      assert u.allclose(u_inv)
  
  
  def test_parallel_transport_batch() -> None:
      mu1 = t([2., 1, np.sqrt(2)]) / radius
      mu2 = t([np.sqrt(5), 1, np.sqrt(3)]) / radius
      u = t([0, 2, -np.sqrt(2)])
      u2 = t([0, 4, -2 * np.sqrt(2)])
  
      U = torch.stack((u, u2), dim=0)
      res = parallel_transport(U, src=mu1, dst=mu2)
      U_ = inverse_parallel_transport(res, src=mu1, dst=mu2)
      assert U.allclose(U_, atol=test_eps)
  
  
  def test_parallel_transport_mu0() -> None:
      mu0 = t([0., 0, 0])
      mu2 = t([np.sqrt(5), 1, np.sqrt(3)]) / radius
      u = t([0, 2, -np.sqrt(2)])
  
      assert P.parallel_transport_mu0(u, dst=mu0, radius=radius).allclose(u)
  
      pt_u = P.parallel_transport_mu0(u, dst=mu2, radius=radius)
      assert parallel_transport(u, src=mu0, dst=mu2).allclose(pt_u)
  
      u_inv = P.inverse_parallel_transport_mu0(pt_u, src=mu2, radius=radius)
      assert u.allclose(u_inv)
  
  
  def test_parallel_transport_mu0_batch() -> None:
      mu2 = radius * t([np.sqrt(5), 1, np.sqrt(3)])
      u = t([0, 2, -np.sqrt(2)])
      u2 = t([0, 4, -2 * np.sqrt(2)])
  
      U = torch.stack((u, u2), dim=0)
      res = P.parallel_transport_mu0(U, dst=mu2, radius=radius)
      U_ = P.inverse_parallel_transport_mu0(res, src=mu2, radius=radius)
      assert U.allclose(U_)
  
  
  def test_exp_map() -> None:
      mu = t([2., 1, np.sqrt(2)]) / radius
      u = t([0, 2, -np.sqrt(2)])
  
      assert is_in_tangent_space(u, at_point=mu, eps=test_eps)
      u_mapped = P.exp_map(u, at_point=mu, radius=radius)
      u_ = P.inverse_exp_map(u_mapped, at_point=mu, radius=radius)
      assert u.allclose(u_, atol=test_eps)
  
      assert P.poincare_distance(mu, u_mapped,
                                 radius=radius).allclose(P.pm.lambda_x(mu, c=P._c(radius)) * torch.norm(u, p=2))
  
  
  def test_exp_map_mu0() -> None:
      mu0 = P.mu_0((3,))
      u = t([0, 2, -np.sqrt(2)])
  
      assert is_in_tangent_space(u, at_point=mu0, eps=test_eps)
      u_mapped = P.exp_map(u, at_point=mu0, radius=radius)
      u_mu0_mapped = P.exp_map_mu0(u, radius=radius)
      assert u_mapped.allclose(u_mu0_mapped)
      u_ = P.inverse_exp_map(u_mapped, at_point=mu0, radius=radius)
      u_mu0 = P.inverse_exp_map_mu0(u_mu0_mapped, radius=radius)
      assert u.allclose(u_, atol=test_eps)
      assert u.allclose(u_mu0, atol=test_eps)
  
      assert P.poincare_distance(mu0, u_mapped,
                                 radius=radius).allclose(P.pm.lambda_x(mu0, c=P._c(radius)) * torch.norm(u, p=2))
      assert P.poincare_distance(mu0, u_mu0_mapped,
                                 radius=radius).allclose(P.pm.lambda_x(mu0, c=P._c(radius)) * torch.norm(u, p=2))
  
  
  def test_exp_map_large() -> None:
      mu = t([2., 1, np.sqrt(2)]) / radius
      mu = mu.double()
      u = 2.5 * t([0, 2, -np.sqrt(2)])
      u = u.double()
      assert is_in_tangent_space(u, at_point=mu, eps=test_eps)  # This should hold.
      u_mapped = P.exp_map(u, at_point=mu, radius=radius)
      u_ = P.inverse_exp_map(u_mapped, at_point=mu, radius=radius)
      assert u.allclose(u_, atol=low_prec_test_eps)
  
  
  def test_exp_map_batch() -> None:
      mu = t([2., 1, np.sqrt(2)]).double() / radius
      u = t([0, 2, -np.sqrt(2)]).double()
      u2 = t([0, 4, -2 * np.sqrt(2)]).double()
      assert is_in_tangent_space(u, at_point=mu, eps=test_eps)
      assert is_in_tangent_space(u2, at_point=mu, eps=test_eps)
  
      U = torch.stack((u, u2), dim=0)
      U_mapped = P.exp_map(U, at_point=mu, radius=radius)
      U_ = P.inverse_exp_map(U_mapped, at_point=mu, radius=radius)
      assert U.allclose(U_, atol=1e-5)
  
  
  def test_sample_projection() -> None:
      v = t([0., 1, 2])
      assert not is_in_hyp_space(v)
  
      mu0 = t([0., 0, 0])
      assert is_in_hyp_space(mu0)
      assert is_in_tangent_space(v, at_point=mu0)
  
      mu = t([2., 1, np.sqrt(2)]) / radius
      assert is_in_hyp_space(mu)
  
      v_proj, _ = P.sample_projection_mu0(v, at_point=mu, radius=radius)
      assert is_in_hyp_space(v_proj, eps=low_prec_test_eps)
  
      _, v_ = P.inverse_sample_projection_mu0(v_proj, at_point=mu, radius=radius)
      assert v.allclose(v_, atol=test_eps)
      assert not is_in_hyp_space(v_)
      assert is_in_tangent_space(v_, at_point=mu0, eps=test_eps)
  