## Implementation of a mesh based full solver with boundary conditions etc.
import numpy as np
from . import mesh_forces
import logging
logger = logging.getLogger(__name__)


def mesh_solver(
        particles: np.ndarray,
        G: float,
        mapping: callable,
        n_grid: int,
        bounds: tuple = (-1, 1),
        boundary: str = "vanishing",
    ) -> np.ndarray:
    """
    Computes the gravitational force acting on a set of particles using a mesh-based approach. The mesh is of fixed size: n_grid x n_grid x n_grid spanning the given bounds. Particles reaching the boundary are treated according to the boundary condition.
    Args:
    - particles: np.ndarray, shape=(n, 4). Assumes that the particles array has the following columns: x, y, z, m.
    - G: float, the gravitational constant.
    - mapping: callable, the mapping function to use.
    - n_grid: int, the number of grid points in each direction.
    - bounds: tuple, the bounds of the mesh.
    - boundary: str, the boundary condition to apply.
    """
    if particles.shape[1] != 4:
        raise ValueError("Particles array must have 4 columns: x, y, z, m")

    logger.debug(f"Computing forces for {particles.shape[0]} particles using mesh [mapping={mapping.__name__}, {n_grid=}]")

    # the function is fine, let's abuse it somewhat
    
    axis = np.linspace(bounds[0], bounds[1], n_grid)
    mesh = np.zeros((n_grid, n_grid, n_grid))
    spacing = np.abs(axis[1] - axis[0])
    logger.debug(f"Using mesh spacing: {spacing}")


    # Check that the boundary condition is fullfilled
    if boundary == "periodic":
        raise NotImplementedError("Periodic boundary conditions are not implemented yet")
    elif boundary == "vanishing":
        # remove the particles that are outside the mesh
        outlier_mask = particles[:, :3] < bounds[0] | particles[:, :3] > bounds[1]

        if np.any(outlier_mask):
            logger.warning(f"Removing {np.sum(outlier_mask)} particles that are outside the mesh")
            particles = particles[~outlier_mask]
            logger.debug(f"New particles shape: {particles.shape}")
    else:
        raise ValueError(f"Unknown boundary condition: {boundary}")


    if logger.isEnabledFor(logging.DEBUG):
        mesh_forces.show_mesh_information(mesh, "Density mesh")

    # compute the potential and its gradient
    phi_grad = mesh_forces.mesh_poisson(rho, G, spacing)

    if logger.isEnabledFor(logging.DEBUG):
        logger.debug(f"Got phi_grad with: {phi_grad.shape}, {np.max(phi_grad)}")
        mesh_forces.show_mesh_information(phi_grad[0], "Potential gradient (x-direction)")

    # compute the particle forces from the mesh potential
    forces = np.zeros_like(particles[:, :3])
    for i, p in enumerate(particles):
        ijk = np.digitize(p, axis) - 1
        logger.debug(f"Particle {p} maps to cell {ijk}")
        # this gives 4 entries since p[3] the mass is digitized as well -> this is meaningless and we discard it
        # logger.debug(f"Particle {p} maps to cell {ijk}")
        forces[i] = - p[3] * phi_grad[..., ijk[0], ijk[1], ijk[2]]

    return forces


def particles_to_mesh(particles: np.ndarray, mesh: np.ndarray, axis: np.ndarray, mapping: callable) -> None:
    """
    Maps a list of particles to an existing mesh, filling it inplace
    """
    if particles.shape[1] < 4:
        raise ValueError("Particles array must have at least 4 columns: x, y, z, m")

    # axis provide an easy way to map the particles to the mesh
    for p in particles:
        m = p[-1]
        # spread the star onto cells through the shape function, taking into account the mass
        ijks, weights = mapping(p, axis)
        for ijk, weight in zip(ijks, weights):
            mesh[ijk[0], ijk[1], ijk[2]] += weight * m