Source code for satorbis_kit.raster.generate_ndvi

"""
Generate NDVI using OpenEO and create patches from the result.
"""

from typing import Optional
import openeo

from satorbis_kit.constants import Constants




def generate_ndvi(
    collection_name: str,
    spatial_extent: dict,
    temporal_extent: list,
    job_title: Optional[str] = None,
    wait_for_completion: bool = False,
) -> openeo.rest.job.BatchJob:
    """
    Generate NDVI from OpenEO collection using the backend.

    This function loads a satellite data collection, calculates NDVI using
    the NIR (B08) and Red (B04) bands, and processes the data through the
    OpenEO backend.

    Args:
        collection_name (str): OpenEO collection name to load.
            Example: "TCX8021"
        spatial_extent (dict): Spatial bounding box with keys: west, south, east, north.
            Example: {"west": 77.0, "south": 28.5, "east": 77.5, "north": 29.0}
        temporal_extent (list): Start and end dates for temporal filtering.
            Example: ["2021-02-01", "2021-02-03"]
        job_title (str, optional): Custom job title. If None, defaults to
            "NDVI Patch Generation Job".
        wait_for_completion (bool, optional): If True, blocks until job completes.
            If False, returns immediately after starting. Defaults to False.

    Returns:
        openeo.rest.job.BatchJob: OpenEO batch job object that can be monitored
            and managed. Use job.status() to check status.

    Raises:
        ConnectionError: If unable to connect to OpenEO backend.
        ValueError: If spatial/temporal extent is invalid or collection doesn't exist.

    Examples:
        Basic NDVI generation:

        >>> from satorbis_kit.raster import generate_ndvi
        >>> job = generate_ndvi(
        ...     collection_name="TCX8021",
        ...     spatial_extent={"west": 77.0, "south": 28.5, "east": 77.5, "north": 29.0},
        ...     temporal_extent=["2021-02-01", "2021-02-03"]
        ... )
        >>> print(job.status())

        Wait for completion:

        >>> job = generate_ndvi(
        ...     collection_name="TCX8021",
        ...     spatial_extent={"west": 77.0, "south": 28.5, "east": 77.5, "north": 29.0},
        ...     temporal_extent=["2021-02-01", "2021-02-03"],
        ...     wait_for_completion=True
        ... )
        >>> print(f"Generation complete: {job.status()}")

        Custom job title:

        >>> job = generate_ndvi(
        ...     collection_name="TCX8021",
        ...     spatial_extent={"west": 77.0, "south": 28.5, "east": 77.5, "north": 29.0},
        ...     temporal_extent=["2021-02-01", "2021-02-03"],
        ...     job_title="My NDVI Analysis",
        ...     wait_for_completion=True
        ... )

    Note:
        - Requires AWS credentials configured for S3 access
        - OpenEO backend accessible via Constants.DEV_URL
        - Collection must have NIR (B08) and Red (B04) bands for NDVI calculation
        - Results can be downloaded using job.download_result() after completion
    """
    # Connect to OpenEO
    con = openeo.connect(
        Constants.DEV_URL,
    )

    # Load collection
    cube = con.load_collection(
        collection_name,
        spatial_extent=spatial_extent,
        temporal_extent=temporal_extent,
    )

    # Create job with descriptive title
    if job_title is None:
        job_title = "NDVI Patch Generation Job"

    job = cube.create_job(title=job_title)
    if wait_for_completion:
        job.start_and_wait()
    else:
        job.start()

    return job