import numpy as np

from NFSFT import NFSFT


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

    # The AVX2 nfsft build is stable with moderate sizes.
    N = 20
    M = 20

    plan = NFSFT(N, M)

    print("NFSFT Test")
    print("Using:", plan.lib_path)
    print("Number of Threads:", plan.num_threads)

    # pseudo-random nodes: x[0] in [-0.5,0.5], x[1] in [0,0.5]
    X = np.random.rand(2, M)
    X[0, :] -= 0.5
    X[1, :] *= 0.5
    plan.x = np.ascontiguousarray(X, dtype=np.float64)

    # 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)

    # direct vs fast forward
    f_direct = plan.trafo_direct().copy()
    f_fast = plan.trafo().copy()

    # direct vs fast adjoint
    fhat_direct_all = plan.adjoint_direct().copy()
    fhat_direct = np.zeros_like(fhat_direct_all)
    for k in range(plan.N + 1):
        for n in range(-k, k + 1):
            idx = plan.nfsft_index(k, n)
            fhat_direct[idx] = fhat_direct_all[idx]

    fhat_fast_all = plan.adjoint().copy()
    fhat_fast = np.zeros_like(fhat_fast_all)
    for k in range(plan.N + 1):
        for n in range(-k, k + 1):
            idx = plan.nfsft_index(k, n)
            fhat_fast[idx] = fhat_fast_all[idx]

    # errors
    err_traf = f_direct - f_fast
    E2_traf = np.linalg.norm(err_traf) / np.linalg.norm(f_direct)
    Einf_traf = np.linalg.norm(err_traf, np.inf) / np.linalg.norm(plan.fhat, 1)
    print("E_2 (trafo test):", E2_traf)
    print("E_infty (trafo test):", Einf_traf)

    err_adj = fhat_direct - fhat_fast
    E2_adj = np.linalg.norm(err_adj) / np.linalg.norm(fhat_direct)
    Einf_adj = np.linalg.norm(err_adj, np.inf) / np.linalg.norm(plan.f, 1)
    print("E_2 (adjoint test):", E2_adj)
    print("E_infty (adjoint test):", Einf_adj)

    if E2_traf >= 1e-9 or Einf_traf >= 1e-9:
        raise RuntimeError("TRAFO TEST FAILED: Errors are too large.")
    if E2_adj >= 1e-9 or Einf_adj >= 1e-9:
        raise RuntimeError("ADJOINT TEST FAILED: Errors are too large.")

    # Note: finalize() can crash for some AVX2 libnfsftjulia builds.
    # We intentionally omit explicit finalization in this example.


if __name__ == "__main__":
    main()
