import math


def dgauss_normalizing_constant(sigma_sq):
  """Compute the normalizing constant of the discrete Gaussian.

  Source: https://arxiv.org/pdf/2004.00010.pdf.

  Args:
    sigma_sq: Variance proxy, squared scale, sigma^2.

  Returns:
    The normalizing constant.
  """
  original = None
  poisson = None
  if sigma_sq <= 1:
    original = 0
    x = 1000
    while x > 0:
      original = original + math.exp(-x * x / (2.0 * sigma_sq))
      x = x - 1
    original = 2 * original + 1

  if sigma_sq * 100 >= 1:
    poisson = 0
    y = 1000
    while y > 0:
      poisson = poisson + math.exp(-math.pi * math.pi * sigma_sq * 2 * y * y)
      y = y - 1
    poisson = math.sqrt(2 * math.pi * sigma_sq) * (1 + 2 * poisson)

  if poisson is None:
    return original
  if original is None:
    return poisson

  scale = max(1, math.sqrt(2 * math.pi * sigma_sq))
  precision = 1e-15
  assert -precision * scale <= original - poisson <= precision * scale
  return (original + poisson) / 2