import numpy as np
from pint import UnitRegistry

ureg = UnitRegistry()


def is_pint_unit(unit: str) -> bool:
    """Check if unit is available in pint

    Args:
        unit (str): unit to test

    Returns: (bool) True if unit is a pint unit, False otherwise

    """
    try:
        res = unit in ureg
    except (ValueError, AttributeError):
        res = False
    return res


def can_convert(from_unit: str, to_unit: str) -> bool:
    """Check if two units are compatible for conversion

    Args:
        from_unit: (str) from unit
        to_unit: (str) to unit

    Returns:(bool) True if conversion is possible, False otherwise

    """
    res = False
    if is_pint_unit(from_unit):
        res = ureg(from_unit).is_compatible_with(to_unit)
    return res


def convert_from_base_unit(unit: str, array: np.array) -> np.array:
    """Convert data from base unit values to wanted unit

    Args:
        unit:(str) wanted unit
        array: (np.array) data to convert

    Returns: (np.array) : converted data if unit is a pint unit, otherwise initial data

    """
    if is_pint_unit(unit):
        pintu = ureg(unit)
        # replace None by nan because pint dont support None
        array = np.where(array == None, np.nan, array)
        quantity = ureg.Quantity(array, pintu.to_base_units()).to(pintu)
        result = quantity.magnitude
        result = np.where(np.isnan(result.astype(float)), None, result)
        return result
    else:
        return array


def convert_to_base_unit(unit: str, array: np.array) -> np.array:
    """Convert data to base unit from initial unit

    Args:
        unit: (str) initial unit
        array: (np.array) data to convert

    Returns: (np.array) converted data if unit is a pint unit, otherwise initial data

    """
    if is_pint_unit(unit):
        # replace None by nan because pint dont support None
        array = np.where(array == None, np.nan, array)
        quantity = ureg.Quantity(array, unit).to_base_units()
        result = quantity.magnitude
        result = np.where(np.isnan(result.astype(float)), None, result)
        return result
    else:
        return array


def unit_conversion(from_unit: str, to_unit: str, array: np.array) -> np.array:
    """Conversion of data from a unit to a new one

    Args:
        from_unit: (str) input unit
        to_unit: (str) output unit
        array: (np.array) data to convert

    Returns: (np.array) converted data if both units are pint unit, otherwise initial data

    """

    if is_pint_unit(from_unit) and is_pint_unit(to_unit) and array is not None:
        array = np.where(array == None, np.nan, array)
        quantity = ureg.Quantity(array, from_unit).to(to_unit)
        result = quantity.magnitude
        result = np.where(np.isnan(result.astype(float)), None, result).astype(float)
        return result
    else:
        return array
