import numpy as np

from NFST import NFSTplan


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

    print("3d NFST Test")

    N = np.array([32, 16, 8], dtype=np.int32)
    M = 1000

    p = NFSTplan(N, M)

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

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

    fhat = np.ascontiguousarray(rng.random(p.Ns), dtype=np.float64)
    p.fhat = fhat

    f2 = p.trafo().copy()

    I = np.array(
        [
            [j, i, k]
            for j in range(1, int(N[0]))
            for i in range(1, int(N[1]))
            for k in range(1, int(N[2]))
        ],
        dtype=np.int32,
    )

    F = (
        np.sin(2.0 * np.pi * (A[:, 0:1] * I[:, 0]))
        * np.sin(2.0 * np.pi * (A[:, 1:2] * I[:, 1]))
        * np.sin(2.0 * np.pi * (A[:, 2:3] * I[:, 2]))
    )

    f1 = F @ fhat

    error_vector = f1 - f2
    E_2 = np.linalg.norm(error_vector) / np.linalg.norm(f1)
    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.")

    f2 = p.adjoint().copy()
    f1 = F.T @ p.f

    error_vector = f1 - f2
    E_2 = np.linalg.norm(error_vector) / np.linalg.norm(f1)
    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.")

    p.finalize()


if __name__ == "__main__":
    main()
