import jax
import jax.numpy as jnp
import haiku as hk

from e3nn_jax import IrrepsData, radius_graph, spherical_harmonics
from e3nn_jax.experimental.point_convolution import Convolution


def test_point_convolution(keys):
    @hk.without_apply_rng
    @hk.transform
    def model(pos, src, dst, x):
        edge_attr = spherical_harmonics("0e + 1e + 2e", pos[dst] - pos[src], True)

        return Convolution(
            irreps_node_attr=None,
            irreps_node_output="8x0e + 8x0o",
            fc_neurons=[],
            num_neighbors=2.0,
        )(x, src, dst, edge_attr)

    apply = jax.jit(model.apply)

    pos = jnp.array([
        [0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0],
        [1.0, 1.0, 0.0],
    ])
    x = IrrepsData.randn("16x0e + 1o", next(keys), (pos.shape[0],))
    src, dst = radius_graph(pos, 2.0)

    w = model.init(next(keys), pos, src, dst, x)
    apply(w, pos, src, dst, x)

    # TODO test equivariance

@@ -420,62 +420,8 @@
     
-irreps_node_attr=None,%0A            irreps_node_output=
 %228x0

fix test