# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
This module defines structured units and quantities.
"""

# Standard library
import operator
from collections.abc import Collection
from functools import cached_property
from typing import Self

import numpy as np

from .core import UNITY, Unit, UnitBase

__all__ = ["StructuredUnit"]


DTYPE_OBJECT = np.dtype("O")


def _names_from_dtype(dtype):
    """Recursively extract field names from a dtype."""
    names = []
    for name in dtype.names:
        subdtype = dtype.fields[name][0].base
        if subdtype.names:
            names.append([name, _names_from_dtype(subdtype)])
        else:
            names.append(name)
    return tuple(names)


def _normalize_names(names):
    """Recursively normalize, inferring upper level names for unadorned tuples.

    Generally, we want the field names to be organized like dtypes, as in
    ``(['pv', ('p', 'v')], 't')``.  But we automatically infer upper
    field names if the list is absent from items like ``(('p', 'v'), 't')``,
    by concatenating the names inside the tuple.
    """
    result = []
    for name in names:
        if isinstance(name, str) and len(name) > 0:
            result.append(name)
        elif (
            isinstance(name, list)
            and len(name) == 2
            and isinstance(name[0], str)
            and len(name[0]) > 0
            and isinstance(name[1], tuple)
            and len(name[1]) > 0
        ):
            result.append([name[0], _normalize_names(name[1])])
        elif isinstance(name, tuple) and len(name) > 0:
            new_tuple = _normalize_names(name)
            name = "".join([(i[0] if isinstance(i, list) else i) for i in new_tuple])
            result.append([name, new_tuple])
        else:
            raise ValueError(
                f"invalid entry {name!r}. Should be a name, "
                "tuple of names, or 2-element list of the "
                "form [name, tuple of names]."
            )

    return tuple(result)


class StructuredUnit:
    """Container for units for a structured Quantity.

    Parameters
    ----------
    units : unit-like, tuple of unit-like, or `~astropy.units.StructuredUnit`
        Tuples can be nested.  If a `~astropy.units.StructuredUnit` is passed
        in, it will be returned unchanged unless different names are requested.
    names : tuple of str, tuple or list; `~numpy.dtype`; or `~astropy.units.StructuredUnit`, optional
        Field names for the units, possibly nested. Can be inferred from a
        structured `~numpy.dtype` or another `~astropy.units.StructuredUnit`.
        For nested tuples, by default the name of the upper entry will be the
        concatenation of the names of the lower levels.  One can pass in a
        list with the upper-level name and a tuple of lower-level names to
        avoid this.  For tuples, not all levels have to be given; for any level
        not passed in, default field names of 'f0', 'f1', etc., will be used.

    Notes
    -----
    It is recommended to initialize the class indirectly, using
    `~astropy.units.Unit`.  E.g., ``u.Unit('AU,AU/day')``.

    When combined with a structured array to produce a structured
    `~astropy.units.Quantity`, array field names will take precedence.
    Generally, passing in ``names`` is needed only if the unit is used
    unattached to a `~astropy.units.Quantity` and one needs to access its
    fields.

    Examples
    --------
    Various ways to initialize a `~astropy.units.StructuredUnit`::

        >>> import astropy.units as u
        >>> su = u.Unit('(AU,AU/day),yr')
        >>> su
        Unit("((AU, AU / d), yr)")
        >>> su.field_names
        (['f0', ('f0', 'f1')], 'f1')
        >>> su['f1']
        Unit("yr")
        >>> su2 = u.StructuredUnit(((u.AU, u.AU/u.day), u.yr), names=(('p', 'v'), 't'))
        >>> su2 == su
        True
        >>> su2.field_names
        (['pv', ('p', 'v')], 't')
        >>> su3 = u.StructuredUnit((su2['pv'], u.day), names=(['p_v', ('p', 'v')], 't'))
        >>> su3.field_names
        (['p_v', ('p', 'v')], 't')
        >>> su3.keys()
        ('p_v', 't')
        >>> su3.values()
        (Unit("(AU, AU / d)"), Unit("d"))

    Structured units share most methods with regular units::

        >>> su.physical_type
        astropy.units.structured.Structure((astropy.units.structured.Structure((PhysicalType('length'), PhysicalType({'speed', 'velocity'})), dtype=[('f0', 'O'), ('f1', 'O')]), PhysicalType('time')), dtype=[('f0', 'O'), ('f1', 'O')])
        >>> su.si
        Unit("((1.49598e+11 m, 1.73146e+06 m / s), 3.15576e+07 s)")

    """

    def __new__(cls, units, names=None):
        dtype = None
        if names is not None:
            if isinstance(names, StructuredUnit):
                dtype = names._units.dtype
                names = names.field_names
            elif isinstance(names, np.dtype):
                if not names.fields:
                    raise ValueError("dtype should be structured, with fields.")
                dtype = np.dtype([(name, DTYPE_OBJECT) for name in names.names])
                names = _names_from_dtype(names)
            else:
                if not isinstance(names, tuple):
                    names = (names,)
                names = _normalize_names(names)

        if not isinstance(units, tuple):
            units = Unit(units)
            if isinstance(units, StructuredUnit):
                # Avoid constructing a new StructuredUnit if no field names
                # are given, or if all field names are the same already anyway.
                if names is None or units.field_names == names:
                    return units

                # Otherwise, turn (the upper level) into a tuple, for renaming.
                units = units.values()
            else:
                # Single regular unit: make a tuple for iteration below.
                units = (units,)

        if names is None:
            names = tuple(f"f{i}" for i in range(len(units)))

        elif len(units) != len(names):
            raise ValueError("lengths of units and field names must match.")

        converted = []
        for unit, name in zip(units, names):
            if isinstance(name, list):
                # For list, the first item is the name of our level,
                # and the second another tuple of names, i.e., we recurse.
                unit = cls(unit, name[1])
                name = name[0]
            else:
                # We are at the lowest level.  Check unit.
                unit = Unit(unit)
                if dtype is not None and isinstance(unit, StructuredUnit):
                    raise ValueError(
                        "units do not match in depth with field "
                        "names from dtype or structured unit."
                    )

            converted.append(unit)

        self = super().__new__(cls)
        if dtype is None:
            dtype = np.dtype(
                [
                    ((name[0] if isinstance(name, list) else name), DTYPE_OBJECT)
                    for name in names
                ]
            )
        # Decay array to void so we can access by field name and number.
        self._units = np.array(tuple(converted), dtype)[()]
        return self

    def __getnewargs__(self):
        """When de-serializing, e.g. pickle, start with a blank structure."""
        return (), None

    @property
    def field_names(self):
        """Possibly nested tuple of the field names of the parts."""
        return tuple(
            ([name, unit.field_names] if isinstance(unit, StructuredUnit) else name)
            for name, unit in self.items()
        )

    # Allow StructuredUnit to be treated as an (ordered) mapping.
    def __len__(self):
        return len(self._units.dtype.names)

    def __getitem__(self, item):
        # Since we are based on np.void, indexing by field number works too.
        return self._units[item]

    def values(self):
        return self._units.item()

    def keys(self):
        return self._units.dtype.names

    def items(self):
        return tuple(zip(self._units.dtype.names, self._units.item()))

    def __iter__(self):
        yield from self._units.dtype.names

    # Helpers for methods below.
    def _recursively_apply(self, func, cls=None):
        """Apply func recursively.

        Parameters
        ----------
        func : callable
            Function to apply to all parts of the structured unit,
            recursing as needed.
        cls : type, optional
            If given, should be a subclass of `~numpy.void`. By default,
            will return a new `~astropy.units.StructuredUnit` instance.
        """
        results = np.void(tuple(map(func, self.values())), self._units.dtype)
        if cls is not None:
            return results.view((cls, results.dtype))

        # Short-cut; no need to interpret field names, etc.
        result = super().__new__(self.__class__)
        result._units = results
        return result

    def _recursively_get_dtype(self, value, enter_lists=True):
        """Get structured dtype according to value, using our field names.

        This is useful since ``np.array(value)`` would treat tuples as lower
        levels of the array, rather than as elements of a structured array.
        The routine does presume that the type of the first tuple is
        representative of the rest.  Used in ``get_converter``.

        For the special value of ``UNITY``, all fields are assumed to be 1.0,
        and hence this will return an all-float dtype.

        """
        if enter_lists:
            while isinstance(value, list):
                value = value[0]
        if value is UNITY:
            value = (UNITY,) * len(self)
        elif not isinstance(value, tuple) or len(self) != len(value):
            raise ValueError(f"cannot interpret value {value} for unit {self}.")
        descr = []
        for (name, unit), part in zip(self.items(), value):
            if isinstance(unit, StructuredUnit):
                descr.append(
                    (name, unit._recursively_get_dtype(part, enter_lists=False))
                )
            else:
                # Got a part associated with a regular unit. Gets its dtype.
                # Like for Quantity, we cast integers to float.
                part = np.array(part)
                part_dtype = part.dtype
                if part_dtype.kind in "iu":
                    part_dtype = np.dtype(float)
                descr.append((name, part_dtype, part.shape))
        return np.dtype(descr)

    @property
    def si(self):
        """The `StructuredUnit` instance in SI units."""
        return self._recursively_apply(operator.attrgetter("si"))

    @property
    def cgs(self):
        """The `StructuredUnit` instance in cgs units."""
        return self._recursively_apply(operator.attrgetter("cgs"))

    # Needed to pass through Unit initializer, so might as well use it.
    @cached_property
    def _physical_type_id(self):
        return self._recursively_apply(
            operator.attrgetter("_physical_type_id"), cls=Structure
        )

    @property
    def physical_type(self):
        """Physical types of all the fields."""
        return self._recursively_apply(
            operator.attrgetter("physical_type"), cls=Structure
        )

    def decompose(self, bases: Collection[UnitBase] = ()) -> Self:
        """The `StructuredUnit` composed of only irreducible units.

        Parameters
        ----------
        bases : sequence of `~astropy.units.UnitBase`, optional
            The bases to decompose into.  When not provided,
            decomposes down to any irreducible units.  When provided,
            the decomposed result will only contain the given units.
            This will raises a `UnitsError` if it's not possible
            to do so.

        Returns
        -------
        `~astropy.units.StructuredUnit`
            With the unit for each field containing only irreducible units.
        """
        return self._recursively_apply(operator.methodcaller("decompose", bases=bases))

    def is_equivalent(self, other, equivalencies=[]):
        """`True` if all fields are equivalent to the other's fields.

        Parameters
        ----------
        other : `~astropy.units.StructuredUnit`
            The structured unit to compare with, or what can initialize one.
        equivalencies : list of tuple, optional
            A list of equivalence pairs to try if the units are not
            directly convertible.  See :ref:`unit_equivalencies`.
            The list will be applied to all fields.

        Returns
        -------
        bool
        """
        try:
            other = StructuredUnit(other)
        except Exception:
            return False

        if len(self) != len(other):
            return False

        for self_part, other_part in zip(self.values(), other.values()):
            if not self_part.is_equivalent(other_part, equivalencies=equivalencies):
                return False

        return True

    def get_converter(self, other, equivalencies=[]):
        if not isinstance(other, type(self)):
            other = self.__class__(other, names=self)

        converters = [
            self_part.get_converter(other_part, equivalencies=equivalencies)
            for (self_part, other_part) in zip(self.values(), other.values())
        ]

        def converter(value):
            if not hasattr(value, "dtype"):
                value = np.array(value, self._recursively_get_dtype(value))
            result = np.empty_like(value)
            for name, converter_ in zip(result.dtype.names, converters):
                result[name] = converter_(value[name])
            # Index with empty tuple to decay array scalars to numpy void.
            return result if result.shape else result[()]

        return converter

    get_converter.__doc__ = UnitBase.get_converter.__doc__

    def to(self, other, value=np._NoValue, equivalencies=[]):
        """Return values converted to the specified unit.

        Parameters
        ----------
        other : `~astropy.units.StructuredUnit`
            The unit to convert to.  If necessary, will be converted to
            a `~astropy.units.StructuredUnit` using the dtype of ``value``.
        value : array-like, optional
            Value(s) in the current unit to be converted to the
            specified unit.  If a sequence, the first element must have
            entries of the correct type to represent all elements (i.e.,
            not have, e.g., a ``float`` where other elements have ``complex``).
            If not given, assumed to have 1. in all fields.
        equivalencies : list of tuple, optional
            A list of equivalence pairs to try if the units are not
            directly convertible.  See :ref:`unit_equivalencies`.
            This list is in addition to possible global defaults set by, e.g.,
            `set_enabled_equivalencies`.
            Use `None` to turn off all equivalencies.

        Returns
        -------
        values : scalar or array
            Converted value(s).

        Raises
        ------
        UnitsError
            If units are inconsistent
        """
        if value is np._NoValue:
            # We do not have UNITY as a default, since then the docstring
            # would list 1.0 as default, yet one could not pass that in.
            value = UNITY
        return self.get_converter(other, equivalencies=equivalencies)(value)

    def to_string(self, format="generic"):
        """Output the unit in the given format as a string.

        Units are separated by commas.

        Parameters
        ----------
        format : `astropy.units.format.Base` subclass or str
            The name of a format or a formatter class.  If not
            provided, defaults to the generic format.

        Notes
        -----
        Structured units can be written to all formats, but can be
        re-read only with 'generic'.

        """
        parts = [part.to_string(format) for part in self.values()]
        out_fmt = "({})" if len(self) > 1 else "({},)"
        if format.startswith("latex"):
            # Strip $ from parts and add them on the outside.
            parts = [part[1:-1] for part in parts]
            out_fmt = "$" + out_fmt + "$"
        return out_fmt.format(", ".join(parts))

    def _repr_latex_(self):
        return self.to_string("latex")

    __array_ufunc__ = None

    def __mul__(self, other):
        if isinstance(other, str):
            try:
                other = Unit(other, parse_strict="silent")
            except Exception:
                return NotImplemented
        if isinstance(other, UnitBase):
            new_units = tuple(part * other for part in self.values())
            return self.__class__(new_units, names=self)
        if isinstance(other, StructuredUnit):
            return NotImplemented

        # Anything not like a unit, try initialising as a structured quantity.
        try:
            from .quantity import Quantity

            return Quantity(other, unit=self)
        except Exception:
            return NotImplemented

    def __rmul__(self, other):
        return self.__mul__(other)

    def __truediv__(self, other):
        if isinstance(other, str):
            try:
                other = Unit(other, parse_strict="silent")
            except Exception:
                return NotImplemented

        if isinstance(other, UnitBase):
            new_units = tuple(part / other for part in self.values())
            return self.__class__(new_units, names=self)
        return NotImplemented

    def __rlshift__(self, m):
        try:
            from .quantity import Quantity

            return Quantity(m, self, copy=False, subok=True)
        except Exception:
            return NotImplemented

    def __str__(self):
        return self.to_string()

    def __repr__(self):
        return f'Unit("{self.to_string()}")'

    def __hash__(self):
        return hash(self.values())

    def __eq__(self, other):
        try:
            other = StructuredUnit(other)
        except Exception:
            return NotImplemented

        return self.values() == other.values()

    def __ne__(self, other):
        if not isinstance(other, type(self)):
            try:
                other = StructuredUnit(other)
            except Exception:
                return NotImplemented

        return self.values() != other.values()


class Structure(np.void):
    """Single element structure for physical type IDs, etc.

    Behaves like a `~numpy.void` and thus mostly like a tuple which can also
    be indexed with field names, but overrides ``__eq__`` and ``__ne__`` to
    compare only the contents, not the field names.  Furthermore, this way no
    `FutureWarning` about comparisons is given.

    """

    # Note that it is important for physical type IDs to not be stored in a
    # tuple, since then the physical types would be treated as alternatives in
    # :meth:`~astropy.units.UnitBase.is_equivalent`.  (Of course, in that
    # case, they could also not be indexed by name.)

    def __eq__(self, other):
        if isinstance(other, np.void):
            other = other.item()

        return self.item() == other

    def __ne__(self, other):
        if isinstance(other, np.void):
            other = other.item()

        return self.item() != other


def _structured_unit_like_dtype(
    unit: UnitBase | StructuredUnit, dtype: np.dtype
) -> StructuredUnit:
    """Make a `StructuredUnit` of one unit, with the structure of a `numpy.dtype`.

    Parameters
    ----------
    unit : UnitBase
        The unit that will be filled into the structure.
    dtype : `numpy.dtype`
        The structure for the StructuredUnit.

    Returns
    -------
    StructuredUnit
    """
    if isinstance(unit, StructuredUnit):
        # If unit is structured, it should match the dtype. This function is
        # only used in Quantity, which performs this check, so it's fine to
        # return as is.
        return unit

    # Make a structured unit
    units = []
    for name in dtype.names:
        subdtype = dtype.fields[name][0]
        if subdtype.names is not None:
            units.append(_structured_unit_like_dtype(unit, subdtype))
        else:
            units.append(unit)
    return StructuredUnit(tuple(units), names=dtype.names)
