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``), and ``map``/``pmap`` carry the
    function's return type. A broadcast or projection names a field no type
    system can resolve, so it carries ``Any`` elements — the caller's next
    call types them.
    """

    __slots__ = ("_items", "_cg_base", "_cg_pin")
    _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")
        key = (cls, element)
        if key not in _SUBSCRIPTS:
            _SUBSCRIPTS[key] = type(
                f"Array[{element.__name__}]",
                (cls,),
                {"__slots__": (), "_element": element, "storable": element.storable},
            )
        return _SUBSCRIPTS[key]

    # --- 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 _children(cls) -> list[tuple[str, type[CGType]]]:
        """The element type, labelled ``[*]``; a schemaless array declares none."""
        return [("[*]", cls._element)] if cls._element is not None else []

    @classmethod
    def _descend(cls, component: str) -> type[CGType]:
        """The element type, reached by any canonical index name.

        Raises:
            SchemaPathError: If this array is schemaless, or ``component`` is not
                a canonical index name.
        """
        if cls._element is None:
            raise SchemaPathError(f"{cls.__name__} is a schemaless array")
        try:
            parse_index_name(component)
        except ValueError as error:
            raise SchemaPathError(str(error)) from error
        return cls._element

    @classmethod
    def load(cls, base: Path) -> Self:
        """The array stored at ``base``, elements in index order; an unwritten tree is empty.

        Raises:
            TypeError: If this array is schemaless.
        """
        if cls._element is None:
            raise TypeError(f"{cls.__name__} cannot load without an element type")
        if not base.is_dir():
            return cls()
        items = [cls._element.load(base / stem) for stem in entry_stems(base)]
        array = cls(cast("list[T]", items))
        remember_base(array, base)
        return array

    @classmethod
    def store(cls, base: Path, value: object) -> None:
        """Write ``value`` under ``base``, one zero-padded index entry per element.

        Raises:
            TypeError: If this array is schemaless, ``value`` is not iterable, or
                an element is ``Missing`` — omitting it would renumber the rest.
        """
        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):
            if is_missing(item):
                raise TypeError(f"{cls.__name__}.store: element {index_name(index)} is Missing")
            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."""
        forget_base(self)
        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  # pyright: ignore[reportAssignmentType] # a mutable container is unhashable

    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[Any]:
        """Broadcast attribute access over the elements.

        A name no element has raises (a typo), a name some lack chains as
        ``Missing``; this array's own machinery is not an element field
        (``broadcastable``).
        """
        if not broadcastable(self, name):
            raise AttributeError(name)
        return Array(broadcast(self._items, f"attribute {name!r}", _accessor(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 not broadcastable(self, name):
            object.__setattr__(self, name, value)
            return

        forget_base(self)
        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[Any]: ...
    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 _wrapped(self._items[int(key)])
            except IndexError as error:
                return _Missing(error)
        if isinstance(key, slice):
            return Array(_wrapped(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[Any]:
        """``access(item, name)`` over the elements, guarded so a typo raises."""
        return broadcast(self._items, f"key {name!r}", _accessor(name))

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

    def map[R](self, function: Callable[[T], R], label: str | None = None) -> Array[R]:
        """Apply ``function`` to each element, collecting failures as ``Missing``.

        The walk is a progress step counting the elements, shown as ``label``
        or ``function``'s name. 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}")
        return Array(map_broadcast(self._items, function, label))

    def pmap[R](self, function: Callable[[T], R], label: str | None = None) -> Array[R]:
        """Like ``map``, across the process's one thread pool, sized to the CPU count.

        Parallel only where ``function`` releases the GIL: decoding an image,
        NumPy, a subprocess. ``cg.parallelism(workers)`` caps how many elements
        are in flight at once. Each worker runs in a copy of the caller's
        context, so the active scratch, workspace and progress step carry over;
        a ``pmap`` inside a worker runs sequentially, the outer pool being the
        parallelism.

        ```python
        thumbnails = sessions.photo.pmap(session_thumbnail)
        ```
        """
        if not callable(function):
            raise TypeError(f"Array.pmap needs a callable, got {function!r}")
        return Array(pbroadcast(self._items, "pmap", function, label))

    def map_leaves(self, function: Callable[[Any], object]) -> Array[Any]:
        """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``.
        """

        def leaf(value: object) -> object:
            if is_missing(value):
                return value if isinstance(value, _Missing) else Missing
            return isolated(function, value)

        result = Array(_rebuilt(item, leaf) for item in self._items)
        guard_broadcast(list(_leaves(result)), "map_leaves")
        return result

    def __call__(self, *args: object, **kwargs: object) -> Array[Any]:
        """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)

        return Array(broadcast(self._items, "call", called))

    def __or__(self, default: object) -> Array[T]:
        """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):
            filled = (
                _rebuilt(item, _filler(fill))
                for item, fill in zip(self._items, default, strict=True)
            )
        else:
            filled = (_rebuilt(item, _filler(default)) for item in self._items)
        return cast("Array[T]", Array(filled))

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

    @property
    def dtype(self) -> type | None:
        """The type of the first present row (descends nested arrays), or ``None``."""
        return next((type(row) for row in _rows(self._items) if not is_missing(row)), None)

    @property
    def dkeys(self) -> Array[str]:
        """The union of field/key names across all present rows."""
        keys: set[str] = set()
        for row in _rows(self._items):
            keys.update(_field_names(row))
        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(_unwrapped(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: [_unwrapped(item) for item in self[key]] for key in keys}