from typing import Optional, Union

from anndata import AnnData
from lndb_storage.object._lazy_field import LazySelector
from lnschema_core import File

File.__doc__ = """Files: serialized data objects.

- Guide: :doc:`/guide/track`
- FAQ: :doc:`/faq/ingest`

Args:
   data: Union[PathLike, DataLike] = None - A file path or an in-memory data
      object to serialize. Can be a cloud path.
   key: Optional[str] = None - A storage key, a relative filepath within the
      storage location, e.g., an S3 or GCP bucket.
   name: Optional[str] = None - A name. Defaults to a file name for a file.
   run: Optional[:class:`~lamindb.Run`] = None - The generating run.
   features: List[:class:`~lamindb.Features`] = None - A feature set record.
   id: Optional[str] = None - The id of the file. Auto-generated if not passed.
   input_of: List[:class:`~lamindb.Run`] = None - Runs for which the file
      is as an input.

Often, files represent atomic datasets in object storage:
jointly measured observations of features (:class:`~lamindb.Features`).
They are generated by running code (:class:`~lamindb.Transform`),
instances of :class:`~lamindb.Run`.

Data objects often have canonical on-disk and in-memory representations. LaminDB
makes some configurable default choices (e.g., serialize a `DataFrame` as a
`.parquet` file).

Some datasets do not have a canonical in-memory representation, for
instance, `.fastq`, `.vcf`, or files describing QC of datasets.

.. note:: Examples for storage ⟷ memory correspondence:

   - Table: `.csv`, `.tsv`, `.parquet`, `.ipc`, `.feather`
     ⟷ `pd.DataFrame`, `polars.DataFrame`
   - Annotated matrix: `.h5ad`, `.h5mu`, `.zarr` ⟷ `anndata.AnnData`, `mudata.MuData`
   - Image: `.jpg`, `.png` ⟷ `np.ndarray`, ...
   - Tensor: zarr directory, TileDB store ⟷ zarr loader, TileDB loader
   - Fastq: `.fastq` ⟷ /
   - VCF: `.vcf` ⟷ /
   - QC: `.html` ⟷ /

"""


def subset(
    self: File,
    query_obs: Optional[Union[str, LazySelector]] = None,
    query_var: Optional[Union[str, LazySelector]] = None,
) -> AnnData:
    """Subset the AnnData File and stream the result into memory.

    Args:
        query_obs: The pandas query to evaluate on `.obs` of the
            underlying `AnnData` object.
        query_var: The pandas query to evaluate on `.var` of the
            underlying `AnnData` object.
    """
    from lndb_storage.object import _subset_anndata_file

    if query_obs is None and query_var is None:
        raise ValueError("Please specify at least one of query_obs or query_var")

    if self.suffix not in (".h5ad", ".zarr"):
        raise ValueError("File should have an AnnData object as the underlying data")

    return _subset_anndata_file(self, query_obs, query_var)


File.subset = subset
