import numpy as np

from NFFT import Plan


def main():
    rng = np.random.default_rng(1234)

    print("3d NFFT Test")

    N = np.array([16, 8, 4], dtype=np.int32)
    M = 10000

    plan = Plan(N, M)

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

    A = np.ascontiguousarray(rng.random((M, 3)) - 0.5, dtype=np.float64)
    plan.x = A

    fhat = np.ascontiguousarray(rng.random(plan.Ns) + 1.0j * rng.random(plan.Ns), dtype=np.complex128)
    plan.fhat = fhat

    f_fast = plan.trafo().copy()

    I = np.array(
        [
            [k0, k1, k2]
            for k0 in range(-N[0] // 2, N[0] // 2)
            for k1 in range(-N[1] // 2, N[1] // 2)
            for k2 in range(-N[2] // 2, N[2] // 2)
        ],
        dtype=np.int32,
    )
    F = np.exp(-2j * np.pi * (A @ I.T))
    f_direct = F @ fhat

    error_vector = f_direct - f_fast
    E_2 = np.linalg.norm(error_vector) / np.linalg.norm(f_direct)
    E_infty = np.linalg.norm(error_vector, np.inf) / np.linalg.norm(fhat, 1)
    print("E_2 error:")
    print(E_2)
    print("E_infty error:")
    print(E_infty)

    if E_2 >= 1e-8 or E_infty >= 1e-8:
        raise RuntimeError("Errors are too large.")

    plan.f = np.ascontiguousarray(f_fast, dtype=np.complex128)
    fhat_fast = plan.adjoint().copy()
    fhat_direct = F.conj().T @ plan.f

    error_vector = fhat_direct - fhat_fast
    E_2 = np.linalg.norm(error_vector) / np.linalg.norm(fhat_direct)
    E_infty = np.linalg.norm(error_vector, np.inf) / np.linalg.norm(fhat, 1)
    print("E_2 error:")
    print(E_2)
    print("E_infty error:")
    print(E_infty)

    if E_2 >= 1e-8 or E_infty >= 1e-8:
        raise RuntimeError("Errors are too large.")

    plan.finalize()


if __name__ == "__main__":
    main()
