import argparse

import numpy as np

from FASTSUM import FASTSUM


def sample_points_in_disk(count, radius, rng):
    r = np.sqrt(rng.random(count)) * radius
    phi = rng.random(count) * (2.0 * np.pi)
    return np.column_stack((r * np.cos(phi), r * np.sin(phi)))


def main():
    parser = argparse.ArgumentParser(description="fastsum Python example")
    parser.add_argument("--N", type=int, default=20000, help="number of source nodes")
    parser.add_argument("--M", type=int, default=20000, help="number of target nodes")
    parser.add_argument("--seed", type=int, default=1234)
    args = parser.parse_args()

    rng = np.random.default_rng(args.seed)

    print("fastsum test")

    d = 2
    N = args.N
    M = args.M
    kernel = "multiquadric"
    c = 1.0 / np.sqrt(N)
    p = 8
    m = p
    n = 256
    eps_I = p / n
    eps_B = max(1.0 / 16.0, p / n)
    nn = 2 * n

    plan = FASTSUM(d, N, M, kernel, c, n=n, p=p, eps_I=eps_I, eps_B=eps_B, nn=nn, m=m)

    print("Using:", plan.lib_path)

    radius = 0.25 - eps_B / 2.0

    # source nodes
    X = sample_points_in_disk(N, radius, rng)
    plan.x = X

    # source coefficients
    alpha = rng.random(N) + 1.0j * rng.random(N)
    plan.alpha = np.ascontiguousarray(alpha, dtype=np.complex128)

    # target nodes
    Y = sample_points_in_disk(M, radius, rng)
    plan.y = Y

    # fast transform
    f_fast = plan.trafo().copy()

    # exact transform
    f_exact = plan.trafo_exact().copy()

    err = f_fast - f_exact
    E_2 = np.linalg.norm(err) / np.linalg.norm(f_fast)
    E_infty = np.linalg.norm(err, np.inf) / np.linalg.norm(plan.alpha, 1)

    print("E_2 error:")
    print(E_2)
    print("E_infty error:")
    print(E_infty)

    # Keep same threshold as pyNFFT3 fastsum_test.py
    if E_2 >= 1e-5 or E_infty >= 1e-5:
        raise RuntimeError("Errors are too large.")

    plan.finalize()


if __name__ == "__main__":
    main()
