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 (and its magnitude)
        r_vec = x_vec - x_current
        r = np.linalg.norm(r_vec, axis=1)

        # add softening to the denominator
        r_adjusted = r**2 + softening**2
        # usually with a square root: r' = sqrt(r^2 + softening^2) and then cubed, but we combine that below

        # the numerator is tricky:
        # m is a list of scalars and r_vec is a list of vectors (2D array)
        # we only want row_wise multiplication
        num = G * (masses * r_vec.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!= 0 and i % 5000 == 0:
            logger.debug(f"Particle {i} done")

    return forces


def n_body_forces_basic(particles: np.ndarray, G: float, softening: float = 0):
    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))
    for i in range(n):
        for j in range(n):
            if i == j:
                continue # keep the value at zero
            r_vec = x_vec[j] - x_vec[i]
            r = np.linalg.norm(r_vec)
            f = - G * masses[i] * masses[j] * r_vec / (r**3 + softening**3)
            forces[i] += f

    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