text: r"""Transformation between two representations of a signal on the sphere.
  
  .. math:: f: S^2 \longrightarrow \mathbb{R}
  
  is a signal on the sphere.
  
  One representation that we like to call "spherical tensor" is
  
  .. math:: f(x) = \sum_{l=0}^{l_{\mathit{max}}} F^l \cdot Y^l(x)
  
  it is made of :math:`(l_{\mathit{max}} + 1)^2` real numbers represented in the above formula by the familly of vectors :math:`F^l \in \mathbb{R}^{2l+1}`.
  
  Another representation is the discretization around the sphere. For this representation we chose a particular grid of size :math:`(N, M)`
  
  .. math::
  
      x_{ij} &= (\sin(\beta_i) \sin(\alpha_j), \cos(\beta_i), \sin(\beta_i) \cos(\alpha_j))
  
      \beta_i &= \pi (i + 0.5) / N
  
      \alpha_j &= 2 \pi j / M
  
  In the code, :math:`N` is called ``res_beta`` and :math:`M` is ``res_alpha``.
  
  The discrete representation is therefore
  
  .. math:: \{ h_{ij} = f(x_{ij}) \}_{ij}
  """
  import math
  
  import torch
  import torch.fft
  from e3nn import o3
  from e3nn.util import explicit_default_types
  from e3nn.util.jit import compile_mode
  
  
  def _quadrature_weights(b, dtype=None, device=None):
      """
      function copied from ``lie_learn.spaces.S3``
  
      Compute quadrature weights for the grid used by Kostelec & Rockmore [1, 2].
      """
      k = torch.arange(b, device=device)
      w = torch.tensor([
          (
              (2. / b) * torch.sin(math.pi * (2. * j + 1.) / (4. * b)) * (
                  (1. / (2 * k + 1)) * torch.sin((2 * j + 1) * (2 * k + 1) * math.pi / (4. * b))
              ).sum()
          )
          for j in torch.arange(2 * b, device=device)
      ], dtype=dtype, device=device)
  
      w /= 2. * ((2 * b) ** 2)
      return w
  
  
  def s2_grid(res_beta, res_alpha, dtype=None, device=None):
      r"""grid on the sphere
  
      Parameters
      ----------
      res_beta : int
          :math:`N`
  
      res_alpha : int
          :math:`M`
  
      dtype : torch.dtype or None
          ``dtype`` of the returned tensors. If ``None`` then set to ``torch.get_default_dtype()``.
  
      device : torch.device or None
          ``device`` of the returned tensors. If ``None`` then set to the default device of the current context.
  
      Returns
      -------
      betas : `torch.Tensor`
          tensor of shape ``(res_beta)``
  
      alphas : `torch.Tensor`
          tensor of shape ``(res_alpha)``
      """
      dtype, device = explicit_default_types(dtype, device)
  
      i = torch.arange(res_beta, dtype=dtype, device=device)
      betas = (i + 0.5) / res_beta * math.pi
  
      i = torch.arange(res_alpha, dtype=dtype, device=device)
      alphas = i / res_alpha * 2 * math.pi
      return betas, alphas
  
  
  def spherical_harmonics_s2_grid(lmax, res_beta, res_alpha, dtype=None, device=None):
      r"""spherical harmonics evaluated on the grid on the sphere
  
      .. math::
  
          f(x) = \sum_{l=0}^{l_{\mathit{max}}} F^l \cdot Y^l(x)
  
          f(\beta, \alpha) = \sum_{l=0}^{l_{\mathit{max}}} F^l \cdot S^l(\alpha) P^l(\cos(\beta))
  
      Parameters
      ----------
      lmax : int
          :math:`l_{\mathit{max}}`
  
      res_beta : int
          :math:`N`
  
      res_alpha : int
          :math:`M`
  
      Returns
      -------
      betas : `torch.Tensor`
          tensor of shape ``(res_beta)``
  
      alphas : `torch.Tensor`
          tensor of shape ``(res_alpha)``
  
      shb : `torch.Tensor`
          tensor of shape ``(res_beta, (lmax + 1)**2)``
  
      sha : `torch.Tensor`
          tensor of shape ``(res_alpha, 2 lmax + 1)``
      """
      betas, alphas = s2_grid(res_beta, res_alpha, dtype=dtype, device=device)
      shb = o3.Legendre(list(range(lmax + 1)))(betas.cos(), betas.sin().abs())  # [b, l * m]
      sha = o3.spherical_harmonics_alpha(lmax, alphas)  # [a, m]
      return betas, alphas, shb, sha
  
  
  def _complete_lmax_res(lmax, res_beta, res_alpha):
      """
      try to use FFT
      i.e. 2 * lmax + 1 == res_alpha
      """
      if res_beta is None:
          if lmax is not None:
              res_beta = 2 * (lmax + 1)  # minimum req. to go on sphere and back
          elif res_alpha is not None:
              res_beta = 2 * ((res_alpha + 1) // 2)
          else:
              raise ValueError("All the entries are None")
  
      if res_alpha is None:
          if lmax is not None:
              if res_beta is not None:
                  res_alpha = max(2 * lmax + 1, res_beta - 1)
              else:
                  res_alpha = 2 * lmax + 1  # minimum req. to go on sphere and back
          elif res_beta is not None:
              res_alpha = res_beta - 1
  
      if lmax is None:
          lmax = min(res_beta // 2 - 1, (res_alpha - 1) // 2)  # maximum possible to go on sphere and back
          # see tests -------------------------------^
  
      assert res_beta % 2 == 0
      assert lmax + 1 <= res_beta // 2
  
      return lmax, res_beta, res_alpha
  
  
  def _expand_matrix(ls, like=None, dtype=None, device=None):
      """
      convertion matrix between a flatten vector (L, m) like that
      (0, 0) (1, -1) (1, 0) (1, 1) (2, -2) (2, -1) (2, 0) (2, 1) (2, 2)
  
      and a bidimensional matrix representation like that
                      (0, 0)
              (1, -1) (1, 0) (1, 1)
      (2, -2) (2, -1) (2, 0) (2, 1) (2, 2)
  
      :return: tensor [l, m, l * m]
      """
      lmax = max(ls)
      if like is None:
          m = torch.zeros(len(ls),
                          2 * lmax + 1,
                          sum(2 * l + 1 for l in ls),
                          dtype=dtype,
                          device=device)
      else:
          m = like.new_zeros((len(ls), 2 * lmax + 1, sum(2 * l + 1 for l in ls)),
                             dtype=dtype,
                             device=device)
      i = 0
      for j, l in enumerate(ls):
          m[j, lmax - l:lmax + l + 1, i:i + 2 * l + 1] = torch.eye(2 * l + 1, dtype=dtype, device=device)
          i += 2 * l + 1
      return m
  
  
  def rfft(x, l):
      r"""Real fourier transform
  
      Parameters
      ----------
      x : `torch.Tensor`
          tensor of shape ``(..., 2 l + 1)``
  
      res : int
          output resolution, has to be an odd number
  
      Returns
      -------
      `torch.Tensor`
          tensor of shape ``(..., res)``
  
      Examples
      --------
  
      >>> lmax = 8
      >>> res = 101
      >>> _betas, _alphas, _shb, sha = spherical_harmonics_s2_grid(lmax, res, res)
      >>> x = torch.randn(res)
      >>> (rfft(x, lmax) - x @ sha).abs().max().item() < 1e-4
      True
      """
      *size, res = x.shape
      x = x.reshape(-1, res)
      x = torch.fft.rfft(x, dim=1)
      x = torch.cat([
          x[:, 1:l + 1].imag.flip(1).mul(-math.sqrt(2)),
          x[:, :1].real,
          x[:, 1:l + 1].real.mul(math.sqrt(2)),
      ], dim=1)
      return x.reshape(*size, 2 * l + 1)
  
  
  def irfft(x, res):
      r"""Inverse of the real fourier transform
  
      Parameters
      ----------
      x : `torch.Tensor`
          tensor of shape ``(..., 2 l + 1)``
  
      res : int
          output resolution, has to be an odd number
  
      Returns
      -------
      `torch.Tensor`
          positions on the sphere, tensor of shape ``(..., res, 3)``
  
      Examples
      --------
  
      >>> lmax = 8
      >>> res = 101
      >>> _betas, _alphas, _shb, sha = spherical_harmonics_s2_grid(lmax, res, res)
      >>> x = torch.randn(2 * lmax + 1)
      >>> (irfft(x, res) - sha @ x).abs().max().item() < 1e-4
      True
      """
      assert res % 2 == 1
      *size, sm = x.shape
      x = x.reshape(-1, sm)
      x = torch.cat([
          x.new_zeros((x.shape[0], (res - sm) // 2)),
          x,
          x.new_zeros((x.shape[0], (res - sm) // 2)),
      ], dim=-1)
      assert x.shape[1] == res
      l = res // 2
      x = torch.complex(
          torch.cat([
              x[:, l:l + 1],
              x[:, l + 1:].div(math.sqrt(2))
          ], dim=1),
          torch.cat([
              torch.zeros_like(x[:, :1]),
              x[:, :l].flip(-1).div(-math.sqrt(2)),
          ], dim=1),
      )
      x = torch.fft.irfft(x, n=res, dim=1) * res
      return x.reshape(*size, res)
  
  
  @compile_mode('trace')
  class ToS2Grid(torch.nn.Module):
      r"""Transform spherical tensor into signal on the sphere
  
      The inverse transformation of `FromS2Grid`
  
      Parameters
      ----------
      lmax : int
      res : int, tuple of int
          resolution in ``beta`` and in ``alpha``
  
      normalization : {'norm', 'component', 'integral'}
      dtype : torch.dtype or None, optional
      device : torch.device or None, optional
  
      Examples
      --------
  
      >>> m = ToS2Grid(6, (100, 101))
      >>> x = torch.randn(3, 49)
      >>> m(x).shape
      torch.Size([3, 100, 101])
  
  
      `ToS2Grid` and `FromS2Grid` are inverse of each other
  
      >>> m = ToS2Grid(6, (100, 101))
      >>> k = FromS2Grid((100, 101), 6)
      >>> x = torch.randn(3, 49)
      >>> y = k(m(x))
      >>> (x - y).abs().max().item() < 1e-4
      True
  
      Attributes
      ----------
      grid : `torch.Tensor`
          positions on the sphere, tensor of shape ``(res_beta, res_alpha, 3)``
      """
  
      def __init__(self, lmax=None, res=None, normalization='component', dtype=None, device=None):
          super().__init__()
  
          assert normalization in ['norm', 'component', 'integral'] or torch.is_tensor(normalization), "normalization needs to be 'norm', 'component' or 'integral'"
  
          if isinstance(res, int) or res is None:
              lmax, res_beta, res_alpha = _complete_lmax_res(lmax, res, None)
          else:
              lmax, res_beta, res_alpha = _complete_lmax_res(lmax, *res)
  
          betas, alphas, shb, sha = spherical_harmonics_s2_grid(lmax, res_beta, res_alpha, dtype=dtype, device=device)
  
          n = None
          if normalization == 'component':
              # normalize such that all l has the same variance on the sphere
              # given that all componant has mean 0 and variance 1
              n = math.sqrt(4 * math.pi) * betas.new_tensor([
                  1 / math.sqrt(2 * l + 1)
                  for l in range(lmax + 1)
              ]) / math.sqrt(lmax + 1)
          if normalization == 'norm':
              # normalize such that all l has the same variance on the sphere
              # given that all componant has mean 0 and variance 1/(2L+1)
              n = math.sqrt(4 * math.pi) * betas.new_ones(lmax + 1) / math.sqrt(lmax + 1)
          if normalization == 'integral':
              n = betas.new_ones(lmax + 1)
          if torch.is_tensor(normalization):
              n = normalization
          m = _expand_matrix(range(lmax + 1), dtype=dtype, device=device)  # [l, m, i]
          shb = torch.einsum('lmj,bj,lmi,l->mbi', m, shb, m, n)  # [m, b, i]
  
          self.lmax, self.res_beta, self.res_alpha = lmax, res_beta, res_alpha
          self.register_buffer('alphas', alphas)
          self.register_buffer('betas', betas)
          self.register_buffer('sha', sha)
          self.register_buffer('shb', shb)
  
      def __repr__(self):
          return f"{self.__class__.__name__}(lmax={self.lmax} res={self.res_beta}x{self.res_alpha} (beta x alpha))"
  
      @property
      def grid(self):
          beta, alpha = torch.meshgrid(self.betas, self.alphas)
          return o3.angles_to_xyz(alpha, beta)
  
      def forward(self, x):
          r"""Evaluate
  
          Parameters
          ----------
          x : `torch.Tensor`
              tensor of shape ``(..., (l+1)^2)``
  
          Returns
          -------
          `torch.Tensor`
              tensor of shape ``[..., beta, alpha]``
          """
          size = x.shape[:-1]
          x = x.reshape(-1, x.shape[-1])
  
          x = torch.einsum('mbi,zi->zbm', self.shb, x)  # [batch, beta, m]
  
          sa, sm = self.sha.shape
          if sa >= sm and sa % 2 == 1:
              x = irfft(x, sa)
          else:
              x = torch.einsum('am,zbm->zba', self.sha, x)
          return x.reshape(*size, *x.shape[1:])
  
      def _make_tracing_inputs(self, n: int):
          return [{'forward': (torch.randn(self.lmax**2),)} for _ in range(n)]
  
  
  @compile_mode('trace')
  class FromS2Grid(torch.nn.Module):
      r"""Transform signal on the sphere into spherical tensor
  
      The inverse transformation of `ToS2Grid`
  
      Parameters
      ----------
      res : int, tuple of int
          resolution in ``beta`` and in ``alpha``
  
      lmax : int
      normalization : {'norm', 'component', 'integral'}
      lmax_in : int, optional
      dtype : torch.dtype or None, optional
      device : torch.device or None, optional
  
      Examples
      --------
  
      >>> m = FromS2Grid((100, 101), 6)
      >>> x = torch.randn(3, 100, 101)
      >>> m(x).shape
      torch.Size([3, 49])
  
  
      `ToS2Grid` and `FromS2Grid` are inverse of each other
  
      >>> m = FromS2Grid((100, 101), 6)
      >>> k = ToS2Grid(6, (100, 101))
      >>> x = torch.randn(3, 100, 101)
      >>> x = k(m(x))  # remove high frequencies
      >>> y = k(m(x))
      >>> (x - y).abs().max().item() < 1e-4
      True
  
      Attributes
      ----------
      grid : `torch.Tensor`
          positions on the sphere, tensor of shape ``(res_beta, res_alpha, 3)``
  
      """
  
      def __init__(self, res=None, lmax=None, normalization='component', lmax_in=None, dtype=None, device=None):
          super().__init__()
  
          assert normalization in ['norm', 'component', 'integral'] or torch.is_tensor(normalization), "normalization needs to be 'norm', 'component' or 'integral'"
  
          if isinstance(res, int) or res is None:
              lmax, res_beta, res_alpha = _complete_lmax_res(lmax, res, None)
          else:
              lmax, res_beta, res_alpha = _complete_lmax_res(lmax, *res)
  
          if lmax_in is None:
              lmax_in = lmax
  
          betas, alphas, shb, sha = spherical_harmonics_s2_grid(lmax, res_beta, res_alpha, dtype=dtype, device=device)
  
          # normalize such that it is the inverse of ToS2Grid
          n = None
          if normalization == 'component':
              n = math.sqrt(4 * math.pi) * betas.new_tensor([
                  math.sqrt(2 * l + 1)
                  for l in range(lmax + 1)
              ]) * math.sqrt(lmax_in + 1)
          if normalization == 'norm':
              n = math.sqrt(4 * math.pi) * betas.new_ones(lmax + 1) * math.sqrt(lmax_in + 1)
          if normalization == 'integral':
              n = 4 * math.pi * betas.new_ones(lmax + 1)
          if torch.is_tensor(normalization):
              n = normalization
          m = _expand_matrix(range(lmax + 1), dtype=dtype, device=device)  # [l, m, i]
          assert res_beta % 2 == 0
          qw = _quadrature_weights(res_beta // 2, dtype=dtype, device=device) * res_beta**2 / res_alpha  # [b]
          shb = torch.einsum('lmj,bj,lmi,l,b->mbi', m, shb, m, n, qw)  # [m, b, i]
  
          self.lmax, self.res_beta, self.res_alpha = lmax, res_beta, res_alpha
          self.register_buffer('alphas', alphas)
          self.register_buffer('betas', betas)
          self.register_buffer('sha', sha)
          self.register_buffer('shb', shb)
  
      def __repr__(self):
          return f"{self.__class__.__name__}(lmax={self.lmax} res={self.res_beta}x{self.res_alpha} (beta x alpha))"
  
      @property
      def grid(self):
          beta, alpha = torch.meshgrid(self.betas, self.alphas)
          return o3.angles_to_xyz(alpha, beta)
  
      def forward(self, x):
          r"""Evaluate
  
          Parameters
          ----------
          x : `torch.Tensor`
              tensor of shape ``[..., beta, alpha]``
  
          Returns
          -------
          `torch.Tensor`
              tensor of shape ``(..., (l+1)^2)``
          """
          size = x.shape[:-2]
          res_beta, res_alpha = x.shape[-2:]
          x = x.reshape(-1, res_beta, res_alpha)
  
          sa, sm = self.sha.shape
          if sm <= sa and sa % 2 == 1:
              x = rfft(x, sm // 2)
          else:
              x = torch.einsum('am,zba->zbm', self.sha, x)
          x = torch.einsum('mbi,zbm->zi', self.shb, x)
          return x.reshape(*size, x.shape[1])
  
      def _make_tracing_inputs(self, n: int):
          return [{'forward': (torch.randn(self.res_beta, self.res_alpha),)} for _ in range(n)]
  