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


M_SCALE: int = None
R_SCALE: int = None

def seed_scales(r_scale: u.Quantity, m_scale: u.Quantity):
    """
    Set the scales for the given simulation.
    Parameters:
    - r_scale: astropy.units.Quantity with units of length - the characteristic length scale of the simulation. Particle positions are expressed in units of this scale.
    - m_scale: astropy.units.Quantity with units of mass - the characteristic mass scale of the simulation. Particle masses are expressed in units of this scale.
    """
    global M_SCALE, R_SCALE
    M_SCALE = m_scale
    R_SCALE = r_scale
    logger.info(f"Set scales: M_SCALE = {M_SCALE}, R_SCALE = {R_SCALE}")


def apply_units(columns: np.array, quantity: str):
    if quantity == "mass":
        return columns * M_SCALE
    elif quantity == "position":
        return columns * R_SCALE
    elif quantity == "velocity":
        return columns * R_SCALE / u.s
    elif quantity == "time":
        return columns * u.s
    elif quantity == "acceleration":
        return columns * R_SCALE / u.s**2
    elif quantity == "force":
        return columns * M_SCALE * R_SCALE / u.s**2
    elif quantity == "volume":
        return columns * R_SCALE**3
    elif quantity == "density":
        return columns * M_SCALE / R_SCALE**3
    else:
        raise ValueError(f"Unknown quantity: {quantity}")