"""Representation of coordinates in different systems."""

__all__ = ["Space"]

import math
from collections.abc import ItemsView, Iterable, KeysView, Mapping, ValuesView
from textwrap import indent
from typing import Any, final

import equinox as eqx
import jax
import wadler_lindig as wl
from astropy.units import PhysicalType as Dimension
from plum import dispatch

import quaxed.numpy as jnp
import unxt as u
from xmmutablemap import ImmutableMap

from .base import AbstractVectors
from .utils import DimensionLike, _get_dimension_name, can_broadcast_shapes
from coordinax._src.vectors.api import vector
from coordinax._src.vectors.base import AbstractVector


# TODO: figure out how to make the keys into Dimension objects, not str. This is
#       running afoul of Jax's tree flattening, where ImmutableMap and
#       eqx.Module differ.
@final
class Space(AbstractVectors, ImmutableMap[Dimension, AbstractVector]):  # type: ignore[misc]
    """A collection of vectors that acts like the primary vector.

    Parameters
    ----------
    *args, **kwargs
        See input to `dict` for the input data.

    Examples
    --------
    >>> import coordinax as cx

    >>> x = cx.CartesianPos3D.from_([1, 2, 3], "km")
    >>> v = cx.CartesianVel3D.from_([4, 5, 6], "km/s")
    >>> a = cx.vecs.CartesianAcc3D.from_([7, 8, 9], "km/s^2")

    All the vectors can be brought together into a space:

    >>> space = cx.Space(length=x, speed=v, acceleration=a)
    >>> print(space)
    Space({
       'length': <CartesianPos3D (x[km], y[km], z[km])
           [1 2 3]>,
       'speed': <CartesianVel3D (x[km / s], y[km / s], z[km / s])
           [4 5 6]>,
       'acceleration': <CartesianAcc3D (x[km / s2], y[km / s2], z[km / s2])
           [7 8 9]>
    })

    The vectors can initialized from `unxt.Quantity` objects and can have
    (brodcastable) batch shapes:

    >>> w = cx.Space(
    ...     length=u.Quantity([[8.5, 0, 0], [10, 0, 0]], "kpc"),
    ...     speed=u.Quantity([0, 200, 0], "km/s"))
    >>> print(w)
    Space({
       'length': <CartesianPos3D (x[kpc], y[kpc], z[kpc])
           [[ 8.5  0.   0. ]
            [10.   0.   0. ]]>,
       'speed': <CartesianVel3D (x[km / s], y[km / s], z[km / s])
           [  0 200   0]>
    })

    The vectors can then be accessed by key:

    >>> space["length"]
    CartesianPos3D( ... )

    The space can also be converted to different representations. If the
    conversion is on the primary vector, the other vectors will be
    correspondingly converted.

    >>> space.vconvert(cx.SphericalPos)
    Space({
        'length': SphericalPos( ... ),
        'speed': SphericalVel( ... ),
        'acceleration': SphericalAcc( ... )
    })

    >>> cx.vconvert(cx.SphericalPos, space)
    Space({
        'length': SphericalPos( ... ),
        'speed': SphericalVel( ... ),
        'acceleration': SphericalAcc( ... )
    })

    Actions on the space are done on the contained vectors.

    >>> w = cx.Space(
    ...     length=cx.CartesianPos3D.from_([[[1, 2, 3], [4, 5, 6]]], "m"),
    ...     speed=cx.CartesianVel3D.from_([[[1, 2, 3], [4, 5, 6]]], "m/s")
    ... )

    >>> w.ndim
    2

    >>> w.shape
    (1, 2)

    >>> w.shapes
    mappingproxy({'length': (1, 2), 'speed': (1, 2)})

    >>> w.mT.shapes
    mappingproxy({'length': (2, 1), 'speed': (2, 1)})

    There are convenience ways to initialize the vectors in the space:

    >>> space = cx.Space.from_({"length": u.Quantity([1, 2, 3], "km"),
    ...                         "speed": u.Quantity([4, 5, 6], "km/s")})
    >>> print(space)
    Space({
       'length': <CartesianPos3D (x[km], y[km], z[km])
           [1 2 3]>,
       'speed': <CartesianVel3D (x[km / s], y[km / s], z[km / s])
           [4 5 6]>
    })



    """

    _data: dict[str, AbstractVector] = eqx.field(init=False)

    def __init__(  # pylint: disable=super-init-not-called  # TODO: resolve this
        self,
        /,
        *args: Mapping[DimensionLike, Any]
        | tuple[DimensionLike, Any]
        | Iterable[tuple[DimensionLike, Any]],
        **kwargs: Any,
    ) -> None:
        # Consolidate the inputs into a single dict, then process keys & values.
        raw = dict(*args, **kwargs)  # process the input data

        # Process the keys
        dims = tuple(u.dimension(k) for k in raw)
        keys = tuple(_get_dimension_name(dim) for dim in dims)
        # Convert the values to vectors
        values = tuple(vector(v) for v in raw.values())

        # TODO: check the dimension makes sense for the value

        # Check that the shapes are broadcastable
        values = eqx.error_if(
            values,
            not can_broadcast_shapes([v.shape for v in values]),
            "vector shapes are not broadcastable.",
        )

        ImmutableMap.__init__(self, dict(zip(keys, values, strict=True)))

    @classmethod
    def _dimensionality(cls) -> int:
        """Dimensionality of the Space.

        Examples
        --------
        >>> import coordinax as cx

        >>> w = cx.Space.from_(cx.CartesianPos3D.from_([1, 2, 3], "kpc"))
        >>> try: w._dimensionality()
        ... except NotImplementedError as e: print("not implemented")
        not implemented

        """
        raise NotImplementedError  # TODO: implement this

    # ===============================================================
    # Mapping API

    @dispatch
    def __getitem__(self, key: Any) -> Any:
        """Get the vector by key.

        Examples
        --------
        >>> import coordinax as cx
        >>> w = cx.Space(length=cx.CartesianPos3D.from_([[[1, 2, 3], [4, 5, 6]]], "m"),
        ...              speed=cx.CartesianVel3D.from_([[[1, 2, 3], [4, 5, 6]]], "m/s"))

        By number:

        >>> w[0]
        Space({
            'length': CartesianPos3D(
                x=Quantity[...](value=i32[2], unit=Unit("m")),
                y=Quantity[...](value=i32[2], unit=Unit("m")),
                z=Quantity[...](value=i32[2], unit=Unit("m")) ),
            'speed': CartesianVel3D(
                x=Quantity[...]( value=i32[2], unit=Unit("m / s") ),
                y=Quantity[...]( value=i32[2], unit=Unit("m / s") ),
                z=Quantity[...]( value=i32[2], unit=Unit("m / s") ) )
        })

        By slice:

        >>> w[1:]
        Space({
            'length': CartesianPos3D(
                x=Quantity[...](value=i32[0,2], unit=Unit("m")),
                y=Quantity[...](value=i32[0,2], unit=Unit("m")),
                z=Quantity[...](value=i32[0,2], unit=Unit("m")) ),
            'speed': CartesianVel3D(
                x=Quantity[...]( value=i32[0,2], unit=Unit("m / s") ),
                y=Quantity[...]( value=i32[0,2], unit=Unit("m / s") ),
                z=Quantity[...]( value=i32[0,2], unit=Unit("m / s") ) )
        })

        By Ellipsis:

        >>> w[...]
        Space({
            'length': CartesianPos3D(
                x=Quantity[...](value=i32[1,2], unit=Unit("m")),
                y=Quantity[...](value=i32[1,2], unit=Unit("m")),
                z=Quantity[...](value=i32[1,2], unit=Unit("m")) ),
            'speed': CartesianVel3D(
                x=Quantity[...]( value=i32[1,2], unit=Unit("m / s") ),
                y=Quantity[...]( value=i32[1,2], unit=Unit("m / s") ),
                z=Quantity[...]( value=i32[1,2], unit=Unit("m / s") ) )
        })

        By tuple[int, ...]:

        >>> w[(0, 1)]
        Space({
            'length': CartesianPos3D(
                x=Quantity[...](value=i32[], unit=Unit("m")),
                y=Quantity[...](value=i32[], unit=Unit("m")),
                z=Quantity[...](value=i32[], unit=Unit("m")) ),
            'speed': CartesianVel3D(
                x=Quantity[...]( value=i32[], unit=Unit("m / s") ),
                y=Quantity[...]( value=i32[], unit=Unit("m / s") ),
                z=Quantity[...]( value=i32[], unit=Unit("m / s") ) )
        })

        This also supports numpy index arrays. But this example section
        highlights core python indexing.

        """
        return Space(**{k: v[key] for k, v in self.items()})

    @dispatch
    def __getitem__(self, key: str | Dimension) -> Any:
        """Get the vector by key.

        Examples
        --------
        >>> import coordinax as cx
        >>> w = cx.Space(length=cx.CartesianPos3D.from_([[[1, 2, 3], [4, 5, 6]]], "m"),
        ...              speed=cx.CartesianVel3D.from_([[[1, 2, 3], [4, 5, 6]]], "m/s"))

        By string key (which is the dimension name):

        >>> print(w["length"])
        <CartesianPos3D (x[m], y[m], z[m])
            [[[1 2 3]
            [4 5 6]]]>

        By the actual dimension object:

        >>> import unxt as u
        >>> print(w[u.dimension("length")])
        <CartesianPos3D (x[m], y[m], z[m])
            [[[1 2 3]
            [4 5 6]]]>

        """
        if isinstance(key, Dimension):
            key = _get_dimension_name(key)

        return ImmutableMap.__getitem__(self, key)

    def keys(self) -> KeysView[str]:
        return self._data.keys()

    def values(self) -> ValuesView[AbstractVector]:
        return self._data.values()

    def items(self) -> ItemsView[str, AbstractVector]:
        return self._data.items()

    # ===============================================================
    # Quax API

    # TODO: unify this with AvalMixin
    def aval(self) -> jax.core.ShapedArray:
        """Return the vector as a JAX array.

        Examples
        --------
        >>> import coordinax as cx
        >>> w = cx.Space(
        ...     length=cx.CartesianPos3D.from_([[[1, 2, 3], [4, 5, 6]]], "m"),
        ...     speed=cx.CartesianVel3D.from_([7, 8, 9], "m/s")
        ... )
        >>> w.aval()
        ShapedArray(int32[1,2,6])

        """
        avals = tuple(v.aval() for v in self.values())
        shapes = [a.shape for a in avals]
        shape = (
            *jnp.broadcast_shapes(*[s[:-1] for s in shapes]),
            sum(s[-1] for s in shapes),
        )
        dtype = jnp.result_type(*map(jnp.dtype, avals))
        return jax.core.ShapedArray(shape, dtype)

    # ===============================================================
    # Array API

    # ---------------------------------------------------------------
    # Attributes

    @property
    def size(self) -> int:
        """Total number of elements in the vector.

        Examples
        --------
        >>> import coordinax as cx

        >>> w = cx.Space(
        ...     length=cx.CartesianPos3D.from_([[[1, 2, 3], [4, 5, 6]]], "m"),
        ...     speed=cx.CartesianVel3D.from_([[[1, 2, 3], [4, 5, 6]]], "m/s") )

        >>> w.size
        2

        """
        return math.prod(self.shape)

    @property
    def T(self) -> "Self":  # noqa: N802
        """Transpose each vector in the space.

        Examples
        --------
        >>> import coordinax as cx

        >>> w = cx.Space(
        ...     length=cx.CartesianPos3D.from_([[[1, 2, 3], [4, 5, 6]]], "m"),
        ...     speed=cx.CartesianVel3D.from_([[[1, 2, 3], [4, 5, 6]]], "m/s")
        ... )

        >>> w.T.shapes
        mappingproxy({'length': (2, 1), 'speed': (2, 1)})

        """
        return super().T

    # ---------------------------------------------------------------
    # Methods

    def __pdoc__(self, **kwargs: Any) -> wl.AbstractDoc:
        """Return the Wadler-Lindig representation.

        Examples
        --------
        >>> import coordinax as cx

        >>> q = cx.CartesianPos3D.from_([1, 2, 3], "m")
        >>> p = cx.CartesianVel3D.from_([1, 2, 3], "m/s")
        >>> w = cx.Space(length=q, speed=p)
        >>> w
        Space({
            'length': CartesianPos3D(
                x=Quantity[...](value=i32[], unit=Unit("m")),
                y=Quantity[...](value=i32[], unit=Unit("m")),
                z=Quantity[...](value=i32[], unit=Unit("m"))
            ),
            'speed': CartesianVel3D(
                x=Quantity[...]( value=i32[], unit=Unit("m / s") ),
                y=Quantity[...]( value=i32[], unit=Unit("m / s") ),
                z=Quantity[...]( value=i32[], unit=Unit("m / s") )
            ) })

        """
        return (
            wl.TextDoc(self.__class__.__name__)
            + wl.TextDoc("(")
            + wl.pdoc(self._data, **kwargs)
            + wl.TextDoc(")")
        )

    def __str__(self) -> str:  # TODO: update using wadler-lindig
        """Return the string representation.

        Examples
        --------
        >>> import coordinax as cx

        >>> q = cx.CartesianPos3D.from_([1, 2, 3], "m")
        >>> p = cx.CartesianVel3D.from_([4, 5, 6], "m/s")
        >>> w = cx.Space(length=q, speed=p)
        >>> print(w)
        Space({
            'length': <CartesianPos3D (x[m], y[m], z[m])
                [1 2 3]>,
            'speed': <CartesianVel3D (x[m / s], y[m / s], z[m / s])
                [4 5 6]>
        })

        """
        cls_name = self.__class__.__name__
        kv = (f"{k!r}: {v!s}" for k, v in self._data.items())
        data = "{\n" + indent(",\n".join(kv), "   ") + "\n}"
        return cls_name + "(" + data + ")"
