Source code for coincident.search.stac

"""
STAC Search Functions
"""

from __future__ import annotations

# maxar-platorm still hasn't renamed and expects MAXAR_API_KEY, so make sure we have both environment variables set before importing!
import os
import warnings
from typing import Any

import arro3.core
import geopandas as gpd
import pystac
import pystac_client
import rustac
from pystac_client.stac_api_io import StacApiIO

if "MAXAR_API_KEY" not in os.environ and "VANTOR_API_KEY" in os.environ:
    os.environ["MAXAR_API_KEY"] = os.environ["VANTOR_API_KEY"]

try:
    import maxar_platform
except ImportError:
    msg_notfound = "'maxar-platform' package not found. Install for vantor functionality: https://pypi.org/project/maxar-platform/"
    warnings.warn(msg_notfound, stacklevel=2)

try:
    import maxar_platform.discovery
except maxar_platform.session.NoSessionCredentials:
    msg_noauth = "Unable to authenticate with Vantor API. Please set VANTOR_API_KEY environment variable."
    warnings.warn(msg_noauth, stacklevel=2)
except maxar_platform.exceptions.UnAuthorizedException:
    msg_noauth = (
        "Unable to authenticate with Vantor API. Please check VANTOR_API_KEY is valid."
    )
    warnings.warn(msg_noauth, stacklevel=2)


def _filter_assets(assets: gpd.GeoDataFrame) -> dict[str, str]:
    """Remove key:None pairs from assets"""
    keep_keys = []
    for k, v in assets.items():
        if v is not None:
            keep_keys.append(k)

    return {key: assets[key] for key in keep_keys}




def to_geopandas(
    collection: list[pystac.Item]
    | pystac.item_collection.ItemCollection
    | arro3.core.Table,
) -> gpd.GeoDataFrame:
    """
    Convert a STAC ItemCollection to a GeoDataFrame.
    It also adds an additional column 'dayofyear' for convenience.

    Parameters
    ----------
    collection
        The STAC ItemCollection to be converted.

    Returns
    -------
        A GeoDataFrame containing the data from the STAC ItemCollection.

    Raises
    ------
    ValueError
        If the provided ItemCollection is empty.
    """
    # Catch if no items are passed
    if collection is None or (hasattr(collection, "__len__") and len(collection) == 0):
        message = "ItemCollection is empty, cannot convert to GeoDataFrame"
        raise ValueError(message)

    if isinstance(collection, pystac.item_collection.ItemCollection):
        collection = rustac.to_arrow(collection.to_dict())

    # NOTE: I believe all STAC footprints are expected as EPSG:4326 (proj:geometry is optional)
    gf = gpd.GeoDataFrame.from_arrow(collection).set_crs(epsg=4326)

    # Workaround stac-geoparquet limitation https://github.com/stac-utils/stac-geoparquet/issues/82
    gf["assets"] = gf["assets"].apply(_filter_assets)

    # Additional columns for convenience
    # NOTE: these become entries under STAC properties
    gf["dayofyear"] = gf["datetime"].dt.dayofyear

    return gf





def to_pystac_items(gf: gpd.GeoDataFrame) -> list[pystac.Item]:
    """
    Converts a GeoDataFrame to a list of PySTAC Items.

    Parameters
    ----------
    gf
        The GeoDataFrame to be converted.

    Returns
    -------
        A list of PySTAC Items created from the GeoDataFrame.
    """

    feature_collection = rustac.from_arrow(gf.to_arrow())
    item_collection = pystac.item_collection.ItemCollection.from_dict(
        feature_collection
    )
    return item_collection.items  # type: ignore[no-any-return]



def search(
    client: pystac_client.client.Client, **kwargs: dict[str, Any] | None
) -> pystac_client.item_search.ItemSearch:
    """Search any STAC API (e.g. https://github.com/nasa/cmr-stac)"""
    # NOTE: add logging for kwargs?
    # print(kwargs)
    results = client.search(
        **kwargs,
    )
    return results.item_collection()  # actually run the query


def configure_vantor_client(
    area_based_calc: bool = True,
) -> pystac_client.client.Client:
    # automatically checks authentication

    client = maxar_platform.discovery.open_catalog(catalog="imagery")
    # Custom Vantor setting
    client.area_based_calc = area_based_calc

    return client


def configure_stac_client(url: str) -> pystac_client.client.Client:
    return pystac_client.Client.open(
        url=url, stac_io=StacApiIO(timeout=60, max_retries=5)
    )