"""
Dimensions of physical quantities


"""

from functools import wraps
from itertools import chain

from sympy import Rational, Symbol, sympify

from unyt._deprecation import warn_deprecated

#: mass
mass = Symbol("(mass)", positive=True)
#: length
length = Symbol("(length)", positive=True)
#: time
time = Symbol("(time)", positive=True)
#: temperature
temperature = Symbol("(temperature)", positive=True)
#: angle
angle = Symbol("(angle)", positive=True)
#: current_mks
current_mks = Symbol("(current_mks)", positive=True)
#: luminous_intensity
luminous_intensity = Symbol("(luminous_intensity)", positive=True)
#: dimensionless
dimensionless = sympify(1)
#: logarithmic
logarithmic = Symbol("(logarithmic)", positive=True)

#: A list of all of the base dimensions
base_dimensions = [
    mass,
    length,
    time,
    temperature,
    angle,
    current_mks,
    dimensionless,
    luminous_intensity,
    logarithmic,
]

#
# Derived dimensions
#

# rate
rate = 1 / time
# frequency (alias for rate)
frequency = rate
# angular frequency
angular_frequency = angle * rate

# spatial frequency
spatial_frequency = 1 / length

#: solid_angle
solid_angle = angle * angle
#: velocity
velocity = length / time
#: acceleration
acceleration = length / time**2
#: jerk
jerk = length / time**3
#: snap
snap = length / time**4
#: crackle
crackle = length / time**5
#: pop
pop = length / time**6

#: area
area = length * length
#: volume
volume = area * length
#: momentum
momentum = mass * velocity
#: force
force = mass * acceleration
#: surface tension
tension = force / length
#: pressure
pressure = force / area
#: energy
energy = force * length
#: power
power = energy / time
#: flux
flux = power / area
#: specific_flux
specific_flux = flux / rate
#: number_density
number_density = 1 / (length * length * length)
#: density
density = mass * number_density
#: angular_momentum
angular_momentum = mass * length * velocity
#: specific_angular_momentum
specific_angular_momentum = angular_momentum / mass
#: specific_energy
specific_energy = energy / mass
#: count_flux
count_flux = 1 / (area * time)
#: count_intensity
count_intensity = count_flux / solid_angle
#: luminous_flux
luminous_flux = luminous_intensity * solid_angle
#: luminance
luminance = luminous_intensity / area

# Gaussian electromagnetic units
#: charge_cgs
charge_cgs = (energy * length) ** Rational(1, 2)  # proper 1/2 power
#: current_cgs
current_cgs = charge_cgs / time
#: electric_field_cgs
electric_field_cgs = charge_cgs / length**2
#: magnetic_field_cgs
magnetic_field_cgs = electric_field_cgs
#: electric_potential_cgs
electric_potential_cgs = energy / charge_cgs
#: resistance_cgs
resistance_cgs = electric_potential_cgs / current_cgs
#: magnetic_flux_cgs
magnetic_flux_cgs = magnetic_field_cgs * area

# SI electromagnetic units
#: charge
charge = charge_mks = current_mks * time
#: electric_field
electric_field = electric_field_mks = force / charge_mks
#: magnetic_field
magnetic_field = magnetic_field_mks = electric_field_mks / velocity
#: electric_potential
electric_potential = electric_potential_mks = energy / charge_mks
#: resistance
resistance = resistance_mks = electric_potential_mks / current_mks
#: capacitance
capacitance = capacitance_mks = charge / electric_potential
#: magnetic_flux
magnetic_flux = magnetic_flux_mks = magnetic_field_mks * area
#: inductance
inductance = inductance_mks = magnetic_flux_mks / current_mks

#: a list containing all derived_dimensions
derived_dimensions = [
    rate,
    velocity,
    acceleration,
    jerk,
    snap,
    crackle,
    pop,
    momentum,
    force,
    energy,
    power,
    charge_cgs,
    electric_field_cgs,
    magnetic_field_cgs,
    solid_angle,
    flux,
    specific_flux,
    volume,
    luminous_flux,
    area,
    current_cgs,
    charge_mks,
    electric_field_mks,
    magnetic_field_mks,
    electric_potential_cgs,
    electric_potential_mks,
    resistance_cgs,
    resistance_mks,
    magnetic_flux_mks,
    magnetic_flux_cgs,
    luminance,
    spatial_frequency,
    angular_frequency,
]


#: a list containing all dimensions
dimensions = base_dimensions + derived_dimensions

#: a dict containing a bidirectional mapping from
#: mks dimension to cgs dimension
em_dimensions = {
    magnetic_field_mks: magnetic_field_cgs,
    magnetic_flux_mks: magnetic_flux_cgs,
    charge_mks: charge_cgs,
    current_mks: current_cgs,
    electric_potential_mks: electric_potential_cgs,
    resistance_mks: resistance_cgs,
}

for k, v in list(em_dimensions.items()):
    em_dimensions[v] = k


def accepts(**arg_units):
    """Decorator for checking units of function arguments.

    Parameters
    ----------
    arg_units: dict
        Mapping of function arguments to dimensions, of the form 'arg1'=dimension1 etc
        where ``'arg1'`` etc are the function arguments and ``dimension1`` etc
        are SI base units (or combination of units), eg. length/time.

    Notes
    -----
    Keyword args are not dimensionally check, being directly passed to the
    decorated function.

    Function arguments that don't have attached units can be skipped can bypass
    dimensionality checking by not being passed to the decorator. See ``baz`` in
    the examples, where ``a`` has no units.

    Examples
    --------
    >>> import unyt as u
    >>> from unyt.dimensions import length, time
    >>> @accepts(a=time, v=length/time)
    ... def foo(a, v):
    ...     return a * v
    ...
    >>> res = foo(a= 2 * u.s, v = 3 * u.m/u.s)
    >>> print(res)
    6 m
    >>> @accepts(a=length, v=length/time)
    ... def bar(a, v):
    ...     return a * v
    ...
    >>> bar(a= 2 * u.s, v = 3 * u.m/u.s)
    Traceback (most recent call last):
    ...
    TypeError: arg 'a=2 s' does not match (length)
    >>> @accepts(v=length/time)
    ... def baz(a, v):
    ...     return a * v
    ...
    >>> res = baz(a= 2, v = 3 * u.m/u.s)
    >>> print(res)
    6 m/s

    """

    def check_accepts(f):
        """Decorates original function.

        Parameters
        ----------
        f : function
            Function being decorated.

        Returns
        -------
        new_f: function
            Decorated function.

        """
        names_of_args = f.__code__.co_varnames

        @wraps(f)
        def new_f(*args, **kwargs):
            """The new function being returned from the decorator.

            Check units of `args` and `kwargs`, then run original function.

            Raises
            ------
            TypeError
                If the units do not match.

            """
            for arg_name, arg_value in chain(zip(names_of_args, args), kwargs.items()):
                if arg_name in arg_units:  # function argument needs to be checked
                    dimension = arg_units[arg_name]
                    if not _has_dimensions(arg_value, dimension):
                        raise TypeError(
                            f"arg '{arg_name}={arg_value}' does not match {dimension}"
                        )
            return f(*args, **kwargs)

        return new_f

    return check_accepts


def returns(*r_units, r_unit=None):
    """Decorator for checking function return units.

    Parameters
    ----------
    *r_units: :py:class:`sympy.core.symbol.Symbol`
        SI base unit (or combination of units), eg. length/time
        of the value(s) returned by the original function
    r_unit: :py:class:`sympy.core.symbol.Symbol`
        Deprecated version of `r_units` which supports only one named return value.

    Examples
    --------
    >>> import unyt as u
    >>> from unyt.dimensions import length, time
    >>> @returns(length)
    ... def f(a, v):
    ...     return a * v
    ...
    >>> res = f(a= 2 * u.s, v = 3 * u.m/u.s)
    >>> print(res)
    6 m
    >>> @returns(length/time)
    ... def f(a, v):
    ...     return a * v
    ...
    >>> f(a= 2 * u.s, v = 3 * u.m/u.s)
    Traceback (most recent call last):
    ...
    TypeError: result '6 m' does not match (length)/(time)
    >>> @returns(length, length/time**2)
    ... def f(a, v):
    ...     return a * v, v / a
    ...
    >>> res = f(a= 2 * u.s, v = 3 * u.m/u.s)
    >>> print(*res)
    6 m 1.5 m/s**2
    """

    # Convert deprecated arguments into current ones where possible.
    if r_unit is not None:
        if len(r_units) > 0:
            raise ValueError(
                "Cannot specify `r_unit` and other return values simultaneously"
            )
        else:
            warn_deprecated(
                "@unyt.returns(r_unit=...)",
                replacement="use @unyt.returns(...)",
                since_version="3.0",
            )
            r_units = (r_unit,)

    def check_returns(f):
        """Decorates original function.

        Parameters
        ----------
        f : function
            Function being decorated.

        Returns
        -------
        new_f: function
            Decorated function.

        """

        @wraps(f)
        def new_f(*args, **kwargs):
            """The decorated function, which checks the return units.

            Raises
            ------
            TypeError
                If the units do not match.

            """
            results = f(*args, **kwargs)

            # Make results a tuple so we can treat single and multiple return values the
            # same way.
            if isinstance(results, tuple):
                result_tuple = results
            else:
                result_tuple = (results,)

            for result, dimension in zip(result_tuple, r_units):
                if not _has_dimensions(result, dimension):
                    raise TypeError(f"result '{result}' does not match {dimension}")
            return results

        return new_f

    return check_returns


def _has_dimensions(quant, dim):
    """Checks the argument has the right dimensionality.

    Parameters
    ----------
    quant : :py:class:`unyt.array.unyt_quantity`
        Quantity whose dimensionality we want to check.
    dim : :py:class:`sympy.core.symbol.Symbol`
        SI base unit (or combination of units), eg. length/time

    Returns
    -------
    bool
        True if check successful.

    Examples
    --------
    >>> import unyt as u
    >>> from unyt.dimensions import length, time
    >>> _has_dimensions(3 * u.m/u.s, length/time)
    True
    >>> _has_dimensions(3, length)
    False
    """
    try:
        arg_dim = quant.units.dimensions
    except AttributeError:
        arg_dim = dimensionless
    return arg_dim == dim