Vector stereographicToEuclidean(const Vector& ζ) {
  unsigned N = ζ.size() + 1;
  Vector z(N);
  Scalar a = ζ.dot(ζ);
  std::complex<double> b = 2.0 * sqrt(N) / (1.0 + a);

  for (unsigned i = 0; i < N - 1; i++) {
    z(i) = b * ζ(i);
  }

  z(N - 1) = b * (a - 1.0) / 2.0;

  return z;
}

Vector euclideanToStereographic(const Vector& z) {
  unsigned N = z.size();
  Vector ζ(N - 1);

  for (unsigned i = 0; i < N - 1; i++) {
    ζ(i) = z(i) / (sqrt(N) - z(N - 1));
  }

  return ζ;
}

Matrix stereographicJacobian(const Vector& ζ) {
  unsigned N = ζ.size() + 1;
  Matrix J(N - 1, N);

  Scalar b = ζ.dot(ζ);

  for (unsigned i = 0; i < N - 1; i++) {
    for (unsigned j = 0; j < N - 1; j++) {
      J(i, j) = -4.0 * ζ(i) * ζ(j);
      if (i == j) {
        J(i, j) += 2.0 * (1.0 + b);
      }
      J(i, j) *= sqrt(N) / pow(1.0 + b, 2);
    }

    J(i, N - 1) = 4.0 * sqrt(N) * ζ(i) / pow(1.0 + b, 2);
  }

  return J;
}

Matrix stereographicMetric(const Vector& ζ) {
  unsigned N = ζ.size();
  Matrix g(N, N);

  double a = ζ.cwiseAbs2().sum();
  Scalar b = ζ.dot(ζ);

  for (unsigned i = 0; i < N; i++) {
    for (unsigned j = 0; j < N; j++) {
      g(i, j) = 16.0 * (std::conj(ζ(i)) * ζ(j) * (1.0 + a) -
                        std::real(std::conj(ζ(i) * ζ(j)) * (1.0 + b)));
      if (i == j) {
        g(i, j) += 4.0 * std::abs(1.0 + b);
      }
      g(i, j) *= (N + 1) / std::norm(pow(b, 2));
    }
  }

  return g;
}

std::tuple<std::complex<double>, Vector, Matrix> stereographicHamGradHess(const Tensor3& J,
                                                                          const Vector& ζ) {
  Vector grad;
  Matrix hess;

  Matrix jacobian = stereographicJacobian(ζ);
  Vector z = stereographicToEuclidean(ζ);

  std::complex<double> hamiltonian;
  Vector gradZ;
  Matrix hessZ;
  std::tie(hamiltonian, gradZ, hessZ) = hamGradHess(J, z);

  Matrix g = stereographicMetric(ζ);
  Matrix gj = g.llt().solve(jacobian);

  grad = gj * gradZ;
  hess = gj * hessZ * gj.transpose();

  return {hamiltonian, grad, hess};
}