import numpy as np
import logging
logger = logging.getLogger(__name__)


def n_body_forces(particles: np.ndarray, G: float, softening: float = 0):
    """
    Computes the gravitational forces between a set of particles.
    Assumes that the particles array has the following columns: x, y, z, m.
    """
    if particles.shape[1] != 4:
        raise ValueError("Particles array must have 4 columns: x, y, z, m")

    x_vec = particles[:, 0:3]
    masses = particles[:, 3]

    n = particles.shape[0]
    forces = np.zeros((n, 3))
    logger.debug(f"Computing forces for {n} particles using n^2 algorithm (using {softening=:.2g})")

    for i in range(n):
        # the current particle is at x_current
        x_current = x_vec[i, :]
        m_current = masses[i]

        # first compute the displacement to all other particles
        displacements = x_vec - x_current
        # and its magnitude
        r = np.linalg.norm(displacements, axis=1)
        # add softening to the denominator
        r_adjusted = r**2 + softening**2
        # the numerator is tricky:
        # m is a list of scalars and displacements is a list of vectors (2D array)
        # we only want row_wise multiplication
        num = G * (masses * displacements.T).T
        
        # a zero value is expected where we have the same particle
        r_adjusted[i] = 1
        num[i] = 0
        
        f = np.sum((num.T / r_adjusted**1.5).T, axis=0) * m_current
        forces[i] = -f

        if i % 5000 == 0:
            logger.debug(f"Particle {i} done")

    return forces



def analytical_forces(particles: np.ndarray):
    """
    Computes the interparticle forces without computing the n^2 interactions.
    This is done by using newton's second theorem for a spherical mass distribution.
    The force on a particle at radius r is simply the force exerted by a point mass with the enclosed mass.
    Assumes that the particles array has the following columns: x, y, z, m.
    """
    n = particles.shape[0]
    forces = np.zeros((n, 3))

    logger.debug(f"Computing forces for {n} particles using spherical approximation")

    r_particles = np.linalg.norm(particles[:, :3], axis=1)
    for i in range(n):
        r_current = np.linalg.norm(particles[i, 0:3])
        m_current = particles[i, 3]

        m_enclosed = np.sum(particles[r_particles < r_current, 3])

        # the force is the same as the force exerted by a point mass at the center
        f = - m_current * m_enclosed / r_current**2
        forces[i] = f

        if i % 5000 == 0:
            logger.debug(f"Particle {i} done")

    return forces