import numpy as np

from NFSFT import NSFT


def main():
    np.random.seed(1234)

    N = 6
    M = 6

    plan = NSFT(N, M)

    print("NSFT Test (equispaced spherical transform)")
    print("Using:", plan.lib_path)
    print("Number of Threads:", plan.num_threads)

    # pseudo-random spherical Fourier coefficients
    fhat = np.zeros((2 * plan.N + 2) ** 2, dtype=np.complex128)
    for k in range(plan.N + 1):
        for n in range(-k, k + 1):
            idx = plan.nfsft_index(k, n)
            fhat[idx] = (np.random.rand() - 0.5) + 1j * (np.random.rand() - 0.5)
    plan.fhat = np.ascontiguousarray(fhat)

    # compare direct and fast forward transform
    f_direct = plan.trafo_direct().copy()
    f_fast = plan.trafo().copy()

    err = f_direct - f_fast
    E2 = np.linalg.norm(err) / np.linalg.norm(f_direct)
    Einf = np.linalg.norm(err, np.inf) / np.linalg.norm(plan.fhat, 1)
    print("E_2 (trafo test):", E2)
    print("E_infty (trafo test):", Einf)

    if E2 >= 1e-9 or Einf >= 1e-9:
        raise RuntimeError("TRAFO TEST FAILED: Errors are too large.")

    # As in pyNFFT3 reference, we only check forward for equispaced variant.
    # Note: finalize() can crash for some AVX2 libnfsftjulia builds.
    # We intentionally omit explicit finalization in this example.


if __name__ == "__main__":
    main()
