array

class Array[T](CGType):
    """A homogeneous sequence.

    ``Array[T]`` is storable when ``T`` is (so ``Array[Path[V]]`` is a
    non-storable view); a bare ``Array`` is the schemaless broadcast view a
    projection lands in. The type parameter types the element surface
    (indexing, iteration, ``append``); derived views — broadcasts,
    projections, ``map`` — carry ``object`` elements, narrowed by the caller.
    """

    __slots__ = ("_items",)
    _items: list[T]
    _element: ClassVar[type[CGType] | None] = None

    def __init__(self, items: Iterable[T] = ()) -> None:
        """Wrap ``items`` (any iterable, copied into a plain list) as an Array."""
        object.__setattr__(self, "_items", list(items))

    def __class_getitem__(cls, element: object) -> type[Array] | GenericAlias:
        """Build ``Array[element]`` — a storable subclass, or a typing-only alias.

        A CGType element makes ``Array[T]`` whose storability *follows* the
        element: ``Array[Image]`` is storable, while ``Array[Path[V]]`` is a
        non-storable fan-out collection (an array of positions, never written to
        disk). A non-CGType element is a typing-only alias, so loaded values can
        be annotated without a schema.
        """
        if not (isinstance(element, type) and issubclass(element, CGType)):
            return GenericAlias(cls, element)
        check_cgtype(element, "Array element")
        return type(
            f"Array[{element.__name__}]",
            (cls,),
            {"__slots__": (), "_element": element, "storable": element.storable},
        )

    # --- storage (a schema'd Array[T] only) ---------------------------------

    @classmethod
    def schema(cls) -> dict[str, JSONValue]:
        """This array's wire form: its element's schema, tagged ``"array"``."""
        if cls._element is None:
            raise TypeError("a schemaless Array has no wire form")
        return {"type": "array", "element": cls._element.schema()}

    @classmethod
    def load(cls, base: Path) -> Self:
        """Read the zero-padded index entries under ``base`` as this array type.

        Typed by the annotation naming this class — ``Array[Image].load(...)``
        is an ``Array[Image]``. The schema supplies the runtime ``_element``;
        the cast below is the one boundary where the type-erased registry meets
        the static element type the schema declares.
        """
        if cls._element is None:
            raise TypeError(f"{cls.__name__} cannot load without an element type")
        if not base.exists():
            return cls()
        loaded = (cls._element.load(base / stem) for stem in entry_stems(base))
        return cls(cast("Iterable[T]", loaded))

    @classmethod
    def store(cls, base: Path, value: object) -> None:
        """Write ``value`` (an iterable) under ``base`` as zero-padded index entries."""
        if cls._element is None:
            raise TypeError(f"{cls.__name__} cannot store without an element type")
        if not isinstance(value, Iterable):
            raise TypeError(f"{cls.__name__}.store needs an iterable, got {value!r}")
        base.mkdir(parents=True, exist_ok=True)
        for index, item in enumerate(value):
            cls._element.store(base / index_name(index), item)

    # --- sequence (any Array instance) --------------------------------------

    def __iter__(self) -> Iterator[T]:
        """Iterate the wrapped elements, in order."""
        return iter(self._items)

    def __len__(self) -> int:
        """The number of elements."""
        return len(self._items)

    def __bool__(self) -> bool:
        """Falsy when empty, like a plain list."""
        return bool(self._items)

    def append(self, item: T) -> None:
        """Add ``item`` to the end of the array, in place."""
        self._items.append(item)

    def __eq__(self, other: object) -> bool:
        """Equal to another ``Array`` or a plain ``list`` with the same elements."""
        if isinstance(other, Array):
            return self._items == other._items
        if isinstance(other, list):
            return self._items == other
        return NotImplemented

    __hash__ = None  # type: ignore[assignment]

    def __repr__(self) -> str:
        """The wrapped elements' own ``repr``."""
        return repr(self._items)

    def __add__(self, other: Iterable[T]) -> Array[T]:
        """Concatenate with another sequence into a new ``Array``."""
        return Array([*self._items, *other])

    def __radd__(self, other: Iterable[T]) -> Array[T]:
        """Concatenate a sequence with this array into a new ``Array``."""
        return Array([*other, *self._items])

    # --- broadcast attribute access -----------------------------------------

    def __getattr__(self, name: str) -> Array[object]:
        """Broadcast attribute access over the elements.

        A name no element has raises (a typo), a name some lack chains as
        ``Missing``. Underscore names use normal lookup.
        """
        if name.startswith("_"):
            raise AttributeError(name)
        return Array(broadcast_attr(self._items, name))

    def __setattr__(self, name: str, value: object) -> None:
        """Broadcast attribute assignment over the elements.

        An empty array adopts the length of ``value``, creating a plain ``dict``
        per entry; a non-empty array assigns element-wise (a same-length
        sequence), broadcasts a single-element sequence or a scalar, and skips
        ``Missing`` elements. Each element is written by attribute when it
        supports one, else by key (a ``dict`` row).
        """
        if name.startswith("_"):
            object.__setattr__(self, name, value)
            return

        if isinstance(value, (list, tuple, Array)):
            values = list(value)
            if not self._items:
                _backing(self).extend({name: item} for item in values)
                return
            if len(values) == len(self._items):
                pairs = zip(self._items, values, strict=True)
            elif len(values) == 1:
                pairs = ((item, values[0]) for item in self._items)
            else:
                raise ValueError(
                    f"Cannot assign sequence of length {len(values)} to Array of "
                    f"length {len(self._items)}; lengths must match, or use a "
                    f"single-element sequence / scalar to broadcast."
                )
            for item, item_value in pairs:
                if not is_missing(item):
                    _set(item, name, item_value)
            return
        for item in self._items:
            if not is_missing(item):
                _set(item, name, value)

    # --- indexing / projection ----------------------------------------------

    @overload
    def __getitem__(self, key: SupportsIndex) -> T: ...
    @overload
    def __getitem__(self, key: slice) -> Array[T]: ...
    @overload
    def __getitem__(self, key: str | tuple[str, ...]) -> Array[object]: ...
    def __getitem__(self, key: SupportsIndex | slice | str | tuple[str, ...]) -> object:
        """Index by position (``int``/``slice``) or project by ``str``/``tuple``.

        ``numbers.Integral`` is accepted alongside plain ``int`` so a NumPy
        index (``np.argmax`` etc.) works like one — this Array converts to and
        from NumPy for analysis. An out-of-range index chains as ``Missing``
        (typed as the element).
        """
        if isinstance(key, numbers.Integral):
            try:
                return _broadcast(self._items[int(key)])
            except IndexError as error:
                return _Missing(error)
        if isinstance(key, slice):
            return Array(_broadcast(item) for item in self._items[key])
        if isinstance(key, str):
            return Array(self._project(key))
        if isinstance(key, tuple):
            columns = {name: self._project(name) for name in key}
            return Array(
                {name: columns[name][index] for name in key} for index in range(len(self._items))
            )
        raise TypeError(f"Array indices must be int, slice, str, or tuple[str, ...]; got {key!r}")

    def _project(self, name: str) -> list[object]:
        """``access(item, name)`` over the elements, guarded so a typo raises."""
        return guard_broadcast(
            [access(item, name) for item in self._items],
            f"key {name!r}",
        )

    # --- function application -----------------------------------------------

    def map(self, function: Callable[[T], object]) -> Array[object]:
        """Apply ``function`` to each element, collecting failures as ``Missing``.

        Raises if the call fails on every element for a non-data reason (a broken
        lambda fails loudly); a sparse data-absence failure chains as ``Missing``.
        """
        if not callable(function):
            raise TypeError(f"Array.map needs a callable, got {function!r}")
        results = [_apply(function, item) for item in self._items]
        return Array(guard_broadcast(results, "map"))

    def pmap(
        self,
        function: Callable[[T], object],
        workers: int | None = None,
    ) -> Array[object]:
        """Like ``map``, but applies ``function`` across a thread pool (I/O-bound)."""
        from concurrent.futures import ThreadPoolExecutor

        if not callable(function):
            raise TypeError(f"Array.pmap needs a callable, got {function!r}")
        if not self._items:
            return Array()
        with ThreadPoolExecutor(max_workers=workers or min(32, len(self._items))) as pool:
            results = list(pool.map(lambda item: _apply(function, item), self._items))
        return Array(guard_broadcast(results, "pmap"))

    def map_leaves(self, function: Callable[[object], object]) -> Array[object]:
        """Apply ``function`` to every leaf, descending into nested Arrays.

        Raises if the call fails on every leaf for a non-data reason (a broken
        lambda fails loudly); a sparse data-absence failure chains as ``Missing``.
        """
        result = Array(_map_leaves(item, function) for item in self._items)
        guard_broadcast(_leaves(result), "map_leaves")
        return result

    def __call__(self, *args: object, **kwargs: object) -> Array[object]:
        """Call each element with the same arguments (the elements are callables).

        Raises if the call fails on every element for a non-data reason (every
        element non-callable is a typo); a sparse failure chains as ``Missing``.
        """

        def called(item: object) -> object:
            if not callable(item):
                raise TypeError(f"Array element {item!r} is not callable")
            return item(*args, **kwargs)

        results = [_apply(called, item) for item in self._items]
        return Array(guard_broadcast(results, "call"))

    def __or__(self, default: object) -> Array[object]:
        """Fill ``Missing``/``None`` leaves with ``default``.

        A parallel sequence fills element-wise; this recurses into nested Arrays.
        """
        if isinstance(default, (Array, list)) and len(default) == len(self._items):
            return Array(
                _fill_missing(item, fill) for item, fill in zip(self._items, default, strict=True)
            )
        return Array(_fill_missing(item, default) for item in self._items)

    # --- inspection ---------------------------------------------------------

    @property
    def dtype(self) -> type | None:
        """The type of the first present leaf (descends nested arrays), or ``None``."""
        return _first_type(self._items)

    @property
    def dkeys(self) -> Array[str]:
        """The union of field/key names across all present leaves."""
        keys: set[str] = set()
        _collect_keys(self._items, keys)
        return Array(keys)

    def __dir__(self) -> list[str]:
        """Standard members plus the broadcastable element field names.

        Surfaces ``photos.location`` in REPL/notebook autocomplete; the
        broadcast itself stays dynamically typed.
        """
        return [*super().__dir__(), *self.dkeys]

    # --- conversion ---------------------------------------------------------

    def to_numpy(self) -> np.ndarray:
        """Convert to a NumPy array, mapping ``Missing``/``None`` to ``nan``."""
        import numpy as np

        return np.array(_unwrap(self._items))

    def __array__(
        self,
        dtype: DTypeLike | None = None,
        copy: bool = True,
    ) -> np.ndarray:
        """NumPy's array-conversion protocol hook — see ``to_numpy``."""
        array = self.to_numpy()
        return array if dtype is None else array.astype(dtype)

    def to_dict(
        self,
        keys: tuple[str, ...] | list[str] | None = None,
    ) -> dict[str, list[object]]:
        """Column-orient as ``{name: [values]}`` for a DataFrame.

        With no ``keys``, the union of leaf field names is used; ``Missing`` is
        written as ``None``.
        """
        if keys is None:
            keys = list(self.dkeys)
        if not self._items:
            return {key: [] for key in keys}
        return {key: [_unwrap(item) for item in self[key]] for key in keys}