import numpy as np

from NFFT import Plan


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

    print("2d NFFT Test")

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

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

    X = np.ascontiguousarray(rng.random((M, len(N))) - 0.5, dtype=np.float64)
    plan.x = X

    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]
            for k0 in range(-N[0] // 2, N[0] // 2)
            for k1 in range(-N[1] // 2, N[1] // 2)
        ],
        dtype=np.int32,
    )

    F = np.exp(-2j * np.pi * (X @ I.T))
    f_direct = F @ fhat

    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(fhat, 1)

    print("Forward E2:", E2)
    print("Forward Einf:", Einf)

    plan.f = np.ascontiguousarray(f_fast, dtype=np.complex128)
    fhat_fast = plan.adjoint().copy()

    fhat_direct = F.conj().T @ f_fast

    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(fhat, 1)

    print("Adjoint E2:", E2_adj)
    print("Adjoint Einf:", Einf_adj)

    if E2 >= 1e-8 or Einf >= 1e-8:
        raise RuntimeError("Forward errors are too large.")
    if E2_adj >= 1e-8 or Einf_adj >= 1e-8:
        raise RuntimeError("Adjoint errors are too large.")

    plan.finalize()


if __name__ == "__main__":
    main()
