import jax
import jax.lax as lax
import jax.numpy as jnp

import jaxpm
import jaxpm.ops
from jaxpm._src.spmd_config import pm_operators
from jaxpm.kernels import cic_compensation


def cic_paint(particle_mesh, positions, halo_size=0):
    return pm_operators.cic_paint(particle_mesh,
                                  positions,
                                  halo_size=halo_size)


def cic_read(mesh, positions, halo_size=0):
    return pm_operators.cic_read(mesh, positions, halo_size=halo_size)


def cic_paint_dx(displacements, halo_size=0):
    return pm_operators.cic_paint_dx(displacements, halo_size=halo_size)


def cic_read_dx(particle_mesh, halo_size=0):
    return pm_operators.cic_read_dx(particle_mesh, halo_size=halo_size)


def cic_paint_2d(mesh, positions, weight):
    """ Paints positions onto a 2d mesh
    mesh: [nx, ny]
    positions: [npart, 2]
    weight: [npart]
    """
    positions = jnp.expand_dims(positions, 1)
    floor = jnp.floor(positions)
    connection = jnp.array([[0, 0], [1., 0], [0., 1], [1., 1]])

    neighboor_coords = floor + connection
    kernel = 1. - jnp.abs(positions - neighboor_coords)
    kernel = kernel[..., 0] * kernel[..., 1]
    if weight is not None:
        kernel = kernel * weight[..., jnp.newaxis]

    neighboor_coords = jnp.mod(
        neighboor_coords.reshape([-1, 4, 2]).astype('int32'),
        jnp.array(mesh.shape))

    dnums = jax.lax.ScatterDimensionNumbers(update_window_dims=(),
                                            inserted_window_dims=(0, 1),
                                            scatter_dims_to_operand_dims=(0,
                                                                          1))
    mesh = lax.scatter_add(mesh, neighboor_coords, kernel.reshape([-1, 4]),
                           dnums)
    return mesh


def compensate_cic(field):
    """
    Compensate for CiC painting
    Args:
      field: input 3D cic-painted field
    Returns:
      compensated_field
    """
    nc = field.shape
    kvec = jaxpm.ops.fftk(nc)

    delta_k = jaxpm.ops.fftn(field)
    delta_k = cic_compensation(kvec) * delta_k
    return jaxpm.ops.ifftn(delta_k).real