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


def density_distribution(r_bins: np.ndarray, particles: np.ndarray, ret_error: bool = False):
    """
    Computes the radial density distribution of a set of particles.
    Assumes that the particles array has the following columns: x, y, z, m.
    If ret_error is True, it will return the absolute error of the density.
    """
    if particles.shape[1] != 4:
        raise ValueError("Particles array must have 4 columns: x, y, z, m")

    m = particles[:, 3]
    r = np.linalg.norm(particles[:, :3], axis=1)
    
    m_shells = np.zeros_like(r_bins)
    v_shells = np.zeros_like(r_bins)
    error_relative = np.zeros_like(r_bins)
    r_bins = np.insert(r_bins, 0, 0)

    for i in range(len(r_bins) - 1):
        mask = (r >= r_bins[i]) & (r < r_bins[i + 1])
        m_shells[i] = np.sum(m[mask])
        v_shells[i] = 4/3 * np.pi * (r_bins[i + 1]**3 - r_bins[i]**3)
        if ret_error:
            count = np.count_nonzero(mask)
            if count > 0:
                error_relative[i] = 1 / np.sqrt(count)
            else:
                error_relative[i] = 0

    density = m_shells / v_shells

    if ret_error:
        return density, density * error_relative
    else:
        return density



def r_distribution(particles: np.ndarray):
    """
    Computes the distribution of distances (to the origin) of a set of particles.
    Assumes that the particles array has the following columns: x, y, z ...
    """
    if particles.shape[1] < 3:
        raise ValueError("Particles array must have at least 3 columns: x, y, z")

    r = np.linalg.norm(particles[:, :3], axis=1)
    return r



def remove_outliers(particles: np.ndarray, std_threshold: float = 3):
    """
    Removes outliers from a set of particles.
    Assumes that the particles array has the following columns: x, y, z ...
    """
    if particles.shape[1] < 3:
        raise ValueError("Particles array must have at least 3 columns: x, y, z")

    r = np.linalg.norm(particles[:, :3], axis=1)
    r_std = np.std(r)
    r_mean = np.mean(r)
    mask = np.abs(r - r_mean) < std_threshold * r_std
    return particles[mask]



def mean_interparticle_distance(particles: np.ndarray):
    """
    Computes the mean interparticle distance of a set of particles.
    Assumes that the particles array has the following columns: x, y, z ...
    """
    if particles.shape[1] < 3:
        raise ValueError("Particles array must have at least 3 columns: x, y, z")
    

    r_half_mass = half_mass_radius(particles)
    r = np.linalg.norm(particles[:, :3], axis=1)

    n_half_mass = np.sum(r < r_half_mass)
    logger.debug(f"Number of particles within half mass radius: {n_half_mass} of {particles.shape[0]}")

    rho = n_half_mass / (4/3 * np.pi * r_half_mass**3)
    # the mean distance between particles is the inverse of the density

    epsilon = (1 / rho)**(1/3)
    logger.info(f"Found mean interparticle distance: {epsilon}")
    return epsilon
    # TODO: check if this is correct



def half_mass_radius(particles: np.ndarray):
    """
    Computes the half mass radius of a set of particles.
    Assumes that the particles array has the following columns: x, y, z ...
    """
    if particles.shape[1] < 3:
        raise ValueError("Particles array must have at least 3 columns: x, y, z")
    
    # even though in the simple example, all the masses are the same, we will consider the general case 
    total_mass = np.sum(particles[:, 3])
    half_mass = total_mass / 2
    
    # sort the particles by distance
    r = np.linalg.norm(particles[:, :3], axis=1)
    indices = np.argsort(r)
    r = r[indices]
    masses = particles[indices, 3]
    masses_cumsum = np.cumsum(masses)

    i = np.argmin(np.abs(masses_cumsum - half_mass))
    logger.debug(f"Half mass radius: {r[i]} for {i}th particle of {particles.shape[0]}")
    r_hm = r[i]

    return r_hm



def total_energy(particles: np.ndarray):
    """
    Computes the total energy of a set of particles.
    Assumes that the particles array has the following columns: x, y, z, vx, vy, vz, m.
    Uses the approximation that the particles are in a central potential as computed in analytical.py
    """
    if particles.shape[1] != 7:
        raise ValueError("Particles array must have 7 columns: x, y, z, vx, vy, vz, m")

    # compute the kinetic energy
    v = particles[:, 3:6]
    m = particles[:, 6]
    ke = 0.5 * np.sum(m * np.linalg.norm(v, axis=1)**2)

    # # compute the potential energy
    # forces = forces_basic.analytical_forces(particles)
    # r = np.linalg.norm(particles[:, :3], axis=1)
    # pe_particles = -forces[:, 0] * particles[:, 0] - forces[:, 1] * particles[:, 1] - forces[:, 2] * particles[:, 2]
    # pe = np.sum(pe_particles)
    # # TODO: i am pretty sure this is wrong
    pe = 0
    return ke + pe