serogen.core¶

class ScalarView(Protocol):
    r"""Formal structural contract for immutable scalar views into a [`BatchedContainer`][serogen.core.collections.BatchedContainer].

    All scalar views must be `@dataclass(slots=True, frozen=True)` records representing a
    single element extracted from a Structure-of-Arrays (SoA) batch container.
    """

    __dataclass_fields__: ClassVar[dict[str, Field[Any]]]

class BatchedContainer[T: ScalarView, S: BatchedContainer[Any, Any]](Protocol):
    r"""A formal contract for Structure-of-Arrays (SoA) and batched containers.

    This protocol ensures that all high-performance vectorized collections provide a standard
    interface for instantiation, iteration, concatenation, and indexing without relying on slow
    dynamic mixins or runtime type introspection.

    Type Parameters:
        T: The [`ScalarView`][serogen.core.collections.ScalarView] dataclass returned when accessing a single index.
        S: The batched container type returned when indexing with slices or array masks.
    """

    def __len__(self) -> int:
        r"""Return the number of records in the batch.

        Returns:
            int: The total count of items in the container.
        """
        ...

    @overload
    def __getitem__(self, item: int | np.integer[Any]) -> T: ...

    @overload
    def __getitem__(self, item: slice | np.ndarray[tuple[int], Any] | list[Any]) -> S: ...

    @overload
    def __getitem__(self, item: int | np.integer[Any] | slice | np.ndarray[tuple[int], Any] | list[Any]) -> T | S: ...

    def __getitem__(self, item: int | np.integer[Any] | slice | np.ndarray[tuple[int], Any] | list[Any]) -> T | S:
        r"""Access records by index, slice, or boolean/integer array mask.

        Args:
            item (int | slice | npt.NDArray | list): An integer index, slice, or mask/indices array.

        Returns:
            T | S: A single scalar record (`T`) if an integer index is provided, or a new
                batched collection (`S`) if a slice or mask is provided.

        Raises:
            IndexError: If an integer index is out of bounds.
        """
        ...

    def __iter__(self) -> Iterator[T]:
        r"""Iterate over the container, yielding each scalar view `T` in order."""
        for i in range(len(self)):
            yield self[i]

    def _normalize_indices(
        self, item: slice | np.ndarray[Any, Any] | list[Any]
    ) -> np.ndarray[tuple[int], np.dtype[np.int32]]:
        r"""Resolve a slice, boolean mask, or index list into a 1D `int32` index array."""
        return normalize_indices(item, len(self))

    @classmethod
    def empty(cls) -> Self:
        r"""Create an empty, 0-length collection with correctly typed arrays.

        Returns:
            S: An empty instance of the container.
        """
        ...

    @classmethod
    def concat(cls, batches: Iterable[Self]) -> Self:
        r"""Concatenate multiple collections into a single, larger collection.

        Args:
            batches (Iterable[S]): An iterable of collections of the same type.

        Returns:
            S: A new, combined collection.

        Raises:
            ValueError: If the input iterable is empty or contains incompatible batches.
        """
        ...

class RaggedArrayContainer[T: ScalarView, S: BatchedContainer[Any, Any]](BatchedContainer[T, S], Protocol):
    r"""A formal contract for ragged Structure-of-Arrays (SoA) containers.

    Ragged containers store variable-length data (such as sequences, cluster members, or CIGAR
    operations) in a flat, contiguous memory layout, managing element boundaries using
    `offsets` and `lengths` arrays.
    """

    @property
    def offsets(self) -> np.ndarray[tuple[int], np.dtype[np.int32]]:
        r"""Starting indices into the flat data array for each record."""
        ...

    @property
    def lengths(self) -> np.ndarray[tuple[int], np.dtype[np.int32]]:
        r"""Lengths of each record."""
        ...

    @property
    def indptr(self) -> np.ndarray[tuple[int], np.dtype[np.int32]]:
        r"""CSR-compatible $(N + 1)$-length index pointer array (`[0, *cumsum(lengths)]`)."""
        n = len(self.offsets)
        if n == 0:
            return np.zeros(1, dtype=np.int32)
        out = np.empty(n + 1, dtype=np.int32)
        out[:n] = self.offsets
        out[n] = self.offsets[-1] + self.lengths[-1]
        return out

    @staticmethod
    def lengths_to_offsets(
        lengths: np.ndarray[tuple[int], np.dtype[np.int32]],
    ) -> np.ndarray[tuple[int], np.dtype[np.int32]]:
        r"""Compute 0-based starting offsets from a 1D `lengths` array."""
        return lengths_to_offsets(lengths)

    def _slice_ragged(
        self,
        item: slice | np.ndarray[Any, Any] | list[Any],
        *arrays: np.ndarray[tuple[int], Any],
    ) -> tuple[
        np.ndarray[tuple[int], np.dtype[np.int32]],
        np.ndarray[tuple[int], np.dtype[np.int32]],
        np.ndarray[tuple[int], np.dtype[np.int32]],
        tuple[np.ndarray[tuple[int], Any], ...],
    ]:
        r"""Slice or gather one or more parallel ragged data arrays while maintaining contiguous packing.

        Uses an $O(1)$ zero-copy view slice on `arrays` when `item` is a contiguous `slice(start, stop, 1)`,
        and falls back to the parallel Numba `_gather_ragged_kernel` for stepped slices, boolean masks,
        or arbitrary index arrays.

        Args:
            item: Slice, boolean mask, or integer index list/array.
            *arrays: Parallel 1D flat data arrays indexed by `self.offsets` and `self.lengths`.

        Returns:
            tuple: `(indices, new_offsets, new_lengths, sliced_arrays)`.
        """
        n = len(self)
        if isinstance(item, slice) and item.step in (None, 1):
            start, stop, _ = item.indices(n)
            if stop <= start:
                empty_i32 = np.empty(0, dtype=np.int32)
                return (
                    empty_i32,
                    empty_i32,
                    empty_i32,
                    tuple(np.empty(0, dtype=a.dtype) for a in arrays),
                )
            indices = np.arange(start, stop, dtype=np.int32)
            new_lengths = self.lengths[start:stop]
            base_offset = int(self.offsets[start])
            end_offset = int(self.offsets[stop - 1]) + int(self.lengths[stop - 1])
            new_offsets = self.offsets[start:stop] - base_offset
            return (
                indices,
                new_offsets,
                new_lengths,
                tuple(a[base_offset:end_offset] for a in arrays),
            )

        indices = self._normalize_indices(item)
        if len(indices) == 0:
            empty_i32 = np.empty(0, dtype=np.int32)
            return (
                indices,
                empty_i32,
                empty_i32,
                tuple(np.empty(0, dtype=a.dtype) for a in arrays),
            )

        new_lengths = self.lengths[indices].astype(np.int32, copy=False)
        new_offsets = lengths_to_offsets(new_lengths)
        total_len = int(new_offsets[-1] + new_lengths[-1])
        if total_len == 0:
            return (
                indices,
                new_offsets,
                new_lengths,
                tuple(np.empty(0, dtype=a.dtype) for a in arrays),
            )

        sliced = tuple(
            _gather_ragged_kernel(a, self.offsets, new_offsets, new_lengths, indices, total_len) for a in arrays
        )
        return indices, new_offsets, new_lengths, sliced

class Strand(IntEnum):
    r"""Integer representation of genomic strand orientation.

    Supports conversion from common string formats (`'+'`, `'-'`, `'1'`, `'-1'`).

    Attributes:
        FORWARD (int): Forward strand (`+1`).
        REVERSE (int): Reverse strand (`-1`).
        UNSTRANDED (int): Unstranded or unknown orientation (`0`).
    """

    FORWARD = 1
    REVERSE = -1
    UNSTRANDED = 0

    @classmethod
    def _missing_(cls, value: object) -> Strand:
        r"""Coerce strings, bytes, or integers to a valid [`Strand`][serogen.core.interval.Strand] instance.

        Args:
            value (object): Value to coerce (e.g. `'+'`, `'-'`, `b'+'`, `1`, `-1`).

        Returns:
            Strand: Corresponding [`Strand`][serogen.core.interval.Strand] enum member.
        """
        if isinstance(value, bytes):
            value = value.decode("ascii")
        if isinstance(value, str):
            if value in ("+", "1", "+1"):
                return Strand.FORWARD
            elif value in ("-", "-1"):
                return Strand.REVERSE
        return Strand.UNSTRANDED

    def __str__(self) -> str:
        r"""Return the string representation of the strand (`'+'`, `'-'`, or `'.'`).

        Returns:
            str: Single-character string representation.
        """
        if self == Strand.FORWARD:
            return "+"
        if self == Strand.REVERSE:
            return "-"
        return "."

@dataclass(frozen=True, slots=True)
class Interval:
    r"""A single 0-based, half-open genomic interval with strand orientation.

    Represents a discrete sequence segment defined by `start` (inclusive), `end` (exclusive),
    and `strand` orientation.

    Attributes:
        start (int): 0-based starting position (inclusive).
        end (int): 0-based ending position (exclusive).
        strand (Strand): Strand orientation [`Strand`][serogen.core.interval.Strand].
    """

    start: int
    end: int
    strand: Strand = Strand.UNSTRANDED

    def __len__(self) -> int:
        r"""Calculate the length of the interval in base pairs (`end - start`).

        Returns:
            int: Interval length.
        """
        return self.end - self.start

@dataclass(frozen=True, slots=True)
class Intervals(BatchedContainer[Interval, "Intervals"]):
    r"""A high-performance Structure-of-Arrays (SoA) collection of genomic intervals.

    Stores starts, ends, and strands as parallel 1D NumPy arrays
    for memory-efficient vectorized operations and Numba spatial algorithms.

    Attributes:
        starts (np.ndarray[tuple[int], np.dtype[np.int32]]): 1D array of 0-based start coordinates.
        ends (np.ndarray[tuple[int], np.dtype[np.int32]]): 1D array of 0-based end coordinates.
        strands (np.ndarray[tuple[int], np.dtype[np.int8]]): 1D array of strand values (+1, -1, 0).
    """

    starts: np.ndarray[tuple[int], np.dtype[np.int32]]
    ends: np.ndarray[tuple[int], np.dtype[np.int32]]
    strands: np.ndarray[tuple[int], np.dtype[np.int8]]

    @classmethod
    def empty(cls) -> Self:
        r"""Create an empty [`Intervals`][serogen.core.interval.Intervals] object with zero-length arrays.

        Returns:
            Intervals: Empty [`Intervals`][serogen.core.interval.Intervals] collection.
        """
        return cls(
            np.empty(0, dtype=np.int32),
            np.empty(0, dtype=np.int32),
            np.empty(0, dtype=np.int8),
        )

    def __len__(self) -> int:
        r"""Return the number of intervals in the batch.

        Returns:
            int: Total interval count.
        """
        return len(self.starts)

    @overload
    def __getitem__(self, item: int | np.integer[Any]) -> Interval: ...

    @overload
    def __getitem__(self, item: slice | np.ndarray[tuple[int], Any] | list[Any]) -> Intervals: ...

    @overload
    def __getitem__(
        self, item: int | np.integer[Any] | slice | np.ndarray[tuple[int], Any] | list[Any]
    ) -> Interval | Intervals: ...

    def __getitem__(
        self, item: int | np.integer[Any] | slice | np.ndarray[tuple[int], Any] | list[Any]
    ) -> Interval | Intervals:
        r"""Access intervals by integer index, slice, or boolean/integer NumPy array.

        Args:
            item (int | slice | np.ndarray | list): Integer index, slice, or NumPy mask array.

        Returns:
            Interval | Intervals: A single [`Interval`][serogen.core.interval.Interval] or a
                sliced [`Intervals`][serogen.core.interval.Intervals] collection.

        Raises:
            IndexError: If integer index is out of bounds.
        """
        if isinstance(item, (int, np.integer)):
            item_idx = int(item)
            if item_idx < 0:
                item_idx += len(self)
            if item_idx < 0 or item_idx >= len(self):
                raise IndexError("Batch index out of range")
            return Interval(self.starts[item_idx], self.ends[item_idx], self.strands[item_idx])

        return Intervals(
            self.starts[item],
            self.ends[item],
            self.strands[item],
        )

    @classmethod
    def concat(cls, batches: Iterable[Self]) -> Self:
        r"""Concatenate multiple [`Intervals`][serogen.core.interval.Intervals] collections into a single batch.

        Args:
            batches (Iterable[Intervals]): Iterable of [`Intervals`][serogen.core.interval.Intervals] objects to merge.

        Returns:
            Intervals: Concatenated [`Intervals`][serogen.core.interval.Intervals] collection.

        Raises:
            ValueError: If `batches` is empty.
        """
        batches_list = list(batches)
        if not batches_list:
            raise ValueError("Cannot concatenate empty list of batches")
        return cls(
            np.concatenate([b.starts for b in batches_list]),
            np.concatenate([b.ends for b in batches_list]),
            np.concatenate([b.strands for b in batches_list]),
        )

    @property
    def lengths(self) -> np.ndarray[tuple[int], np.dtype[np.int32]]:
        r"""Compute the lengths of all intervals in base pairs (`ends - starts`).

        Returns:
            np.ndarray[tuple[int], np.dtype[np.int32]]: 1D array of interval lengths.
        """
        return self.ends - self.starts

@dataclass(slots=True, frozen=True)
class GenomeAssembly(Sequences):
    r"""Container for a genome assembly with support for transparent decompression.

    Stores contig sequence data as a contiguous memory block using [`Sequences`][serogen.core.seq.seq.Sequences]
    and supports thread-safe, lazy-cached index creation for high-performance sequence alignment.

    Attributes:
        id (bytes): Unique genome assembly identifier.
    """

    _SEQUENCE_FILE_REGEX: ClassVar[Pattern[str]] = re_compile(
        r"\.(?P<ext>f(asta|a|na|fn|as))(\.(?P<compression>gz|bz2|xz|zst|zstd))?$"
    )
    _OPENERS: ClassVar[dict[str, Callable[..., Any]]] = dict(
        gz=gzip.open, bz2=bz2.open, xz=lzma.open, zst=zstd.open, zstd=zstd.open
    )

    id: bytes = b"unknown"

    @classmethod
    def ensure(cls, genome: Self | str | Path | IO[bytes]) -> Self:
        r"""Ensure the input object is coerced into a [`GenomeAssembly`][serogen.core.genome.GenomeAssembly].

        Args:
            genome (Self | str | Path | IO[bytes]): Existing assembly object, file path,
                or binary FASTA stream.

        Returns:
            GenomeAssembly: Validated or loaded [`GenomeAssembly`][serogen.core.genome.GenomeAssembly] instance.
        """
        if isinstance(genome, cls):
            return genome
        elif isinstance(genome, (str, Path)):
            return cls.from_file(genome)
        return cls.from_stream(genome)  # type: ignore

    @classmethod
    def from_fasta(
        cls,
        data: bytes | bytearray | memoryview,
        id_: bytes | str = b"unknown",
        *,
        truncate_id: bool = True,
        uppercase: bool = True,
    ) -> Self:
        r"""Construct a GenomeAssembly directly from raw binary FASTA data.

        Args:
            data (bytes | bytearray | memoryview): Raw binary FASTA contents.
            id_ (bytes | str): Genome assembly identifier.
            truncate_id (bool): If `True`, truncates header IDs at first whitespace.
            uppercase (bool): If `True`, normalizes sequence bytes to uppercase ASCII.

        Returns:
            GenomeAssembly: Constructed [`GenomeAssembly`][serogen.core.genome.GenomeAssembly] instance.
        """
        seqs = FastaReader.from_bytes(data, truncate_id=truncate_id, uppercase=uppercase).load()
        id_bytes = id_.encode("utf-8") if isinstance(id_, str) else bytes(id_)
        return cls(seqs.ids, seqs.seqs, seqs.offsets, seqs.lengths, id=id_bytes)

    @classmethod
    def from_file(cls, filepath: str | Path) -> Self:
        r"""Load a genome assembly from a FASTA file path.

        Supports plain `.fasta` / `.fa` / `.fna` files as well as compressed
        `.gz`, `.bz2`, and `.xz` files.

        Args:
            filepath (str | Path): Path to the target FASTA file.

        Returns:
            GenomeAssembly: Loaded [`GenomeAssembly`][serogen.core.genome.GenomeAssembly] instance.

        Raises:
            NotImplementedError: If the file format or compression extension is unsupported.
        """
        filepath = Path(filepath)
        if not (m := cls._SEQUENCE_FILE_REGEX.search(filepath.name)):
            raise NotImplementedError(f"Unsupported format: {filepath}")

        with cls._OPENERS.get(m.group("compression"), open)(filepath, mode="rb") as handle:
            stem = filepath.name.removesuffix(m.group()).encode("utf-8")
            return cls.from_stream(handle, id_=stem)

    @classmethod
    def from_stream(cls, handle: IO[bytes], id_: bytes | str | None = None) -> Self:
        r"""Load a genome assembly from an open binary stream.

        Args:
            handle (IO[bytes]): Open binary FASTA stream handle.
            id_ (bytes | str | None): Custom assembly identifier. If None, derived from
                handle name or defaults to `b'unknown'`.

        Returns:
            GenomeAssembly: Loaded [`GenomeAssembly`][serogen.core.genome.GenomeAssembly] instance.
        """
        if id_ is None:
            raw_name = getattr(handle, "name", "unknown")
            id_bytes = raw_name.encode("utf-8") if isinstance(raw_name, str) else bytes(raw_name)
        elif isinstance(id_, str):
            id_bytes = id_.encode("utf-8")
        else:
            id_bytes = bytes(id_)

        seqs = FastaReader(handle).load()
        return cls(seqs.ids, seqs.seqs, seqs.offsets, seqs.lengths, id=id_bytes)

    @classmethod
    def from_records(
        cls,
        records: list[Sequence] | tuple[Sequence, ...] | Iterable[Sequence],
        id_: bytes | str = b"unknown",
    ) -> Self:
        r"""Construct a [`GenomeAssembly`][serogen.core.genome.GenomeAssembly] from sequence records.

        Args:
            records (Iterable[Sequence]): Iterable of [`Sequence`][serogen.core.seq.seq.Sequence] objects.
            id_ (bytes | str): Genome assembly identifier.

        Returns:
            GenomeAssembly: Constructed [`GenomeAssembly`][serogen.core.genome.GenomeAssembly] instance.
        """
        seqs = Sequences.from_records(records)
        id_bytes = id_.encode("utf-8") if isinstance(id_, str) else bytes(id_)
        return cls(seqs.ids, seqs.seqs, seqs.offsets, seqs.lengths, id=id_bytes)