import numpy as np
import pyopencl as cl
from boxtree import TreeBuilder
from boxtree.traversal import FMMTraversalBuilder
from boxtree.tools import make_normal_particle_array as p_normal
from boxtree.tools import particle_array_to_host
from boxtree.pyfmmlib_integration import FMMLibExpansionWrangler
from boxtree.fmm import drive_fmm
from pyopencl.clrandom import PhiloxGenerator
import numpy.linalg as la

# parameters
dims = 3
nsources = 50000
ntargets = 50000
dtype = np.float64

# Set up PyOpenCL runtime
ctx = cl.create_some_context()
queue = cl.CommandQueue(ctx)
print(queue.context.devices)

# Generate particles
sources = p_normal(queue, nsources, dims, dtype, seed=15)
targets = p_normal(queue, ntargets, dims, dtype, seed=18)
sources_host = particle_array_to_host(sources)
targets_host = particle_array_to_host(targets)
rng = PhiloxGenerator(queue.context, seed=20)
sources_weights = rng.uniform(queue, nsources, dtype=np.float64).get()
rng = PhiloxGenerator(queue.context, seed=22)
target_radii = rng.uniform(queue, ntargets, a=0, b=0.05, dtype=np.float64).get()

# Calculate potentials using direct evaluation
# distances = la.norm(sources_host.reshape(1, nsources, 2) -
#                    targets_host.reshape(ntargets, 1, 2),
#                    ord=2, axis=2)
# pot_naive = np.sum(-np.log(distances)*sources_weights, axis=1)/2/np.pi

# Build the tree and interaction lists
tb = TreeBuilder(ctx)
# tree, _ = tb(queue, sources, targets=targets, target_radii=target_radii,
#             stick_out_factor=0.25, max_particles_in_box=30, debug=True)
tree, _ = tb(queue, sources, targets=targets, max_particles_in_box=30, debug=True)
tg = FMMTraversalBuilder(ctx)
trav, _ = tg(queue, tree, debug=True)
trav = trav.get(queue=queue)


# Get pyfmmlib expansion wrangler
def fmm_level_to_nterms(tree, level):
    return 20


wrangler = FMMLibExpansionWrangler(
    trav.tree, 0, fmm_level_to_nterms=fmm_level_to_nterms)

# Compute FMM using shared memory parallelism
pot_fmm = drive_fmm(trav, wrangler, sources_weights)

# Print accuracy
# print(la.norm(pot_fmm - pot_naive, ord=np.inf)/la.norm(pot_naive, ord=np.inf))