flag_spikes.py

import os
from datetime import datetime

import typer
from typing import Optional

from ..metadata import Metadata, MetadataFactory
from ..algorithm.SpikesAlgorithm import SpikesAlgorithm


def flag_spikes(
    observatory: str = typer.Option(..., help="Observatory code"),
    channels: str = typer.Option(
        default=["F"], help="Channels to despike, default is F. Example input: HEZF"
    ),
    starttime: datetime = typer.Option(
        default=None,
        help="Start time for metadata, default is None. CLI example usage: --starttime '2024-01-01' ",
    ),
    endtime: datetime = typer.Option(
        default=None,
        help="End time for metadata, default is None. CLI example usage: --endtime '2024-01-01' ",
    ),
    window_size: Optional[int] = typer.Option(
        default=1000,
        help="Window size for despike algorithm",
    ),
    threshold: Optional[int] = typer.Option(
        default=6,
        help="Threshold for spike size",
    ),
    metadata_token: str = typer.Option(
        default=os.getenv("GITLAB_API_TOKEN"),
        help="Token for metadata web service.",
        metavar="TOKEN",
        show_default="environment variable GITLAB_API_TOKEN",
    ),
    metadata_url: str = typer.Option(
        default="https://staging-geomag.cr.usgs.gov/ws/secure/metadata",
        help="URL to metadata web service",
        metavar="URL",
    ),
    force: bool = typer.Option(
        default=False,
        help="Force skip the checks",
    ),
):
    """Flag spikes and create spike metadata in the metadata webservice."""

    # TODO: Consider adding below variables as options
    # run spike detection algorithm on data and return spike metadata
    spike_algorithm = SpikesAlgorithm(
        starttime=starttime,
        endtime=endtime,
        observatory=observatory,
        window_size=window_size,
        threshold=threshold,
        channels=channels,
        type="variation",
        interval="second",
        metadata_token=metadata_token,
        metadata_url=metadata_url,
    )

    spikes = spike_algorithm.run()

    # confirm whether or not to create spike metadata
    if not force:
        typer.confirm(f"Are you sure you want to create flag metadata?", abort=True)
        print("Creating flag metadata")

    # write flag metadata to metadata service
    metadata_factory = MetadataFactory(token=metadata_token, url=metadata_url)
    with typer.progressbar(
        iterable=spikes, label="Uploading to metadata service"
    ) as progressbar:
        for spike in progressbar:
            upload_spike_metadata(algorithm=spike_algorithm, spike=spike)


def main() -> None:
    """Entrypoint for flag spikes method."""
    # for one command, can use typer.run
    typer.run(flag_spikes)


def upload_spike_metadata(algorithm: SpikesAlgorithm, spike: Metadata) -> Metadata:
    """Upload new or existing spike metadata to metadata service.

    Parameters
    ----------
    algorithm: SpikesAlgorithm
        algorithm to update/load spike metadata
    spike:  Metadata
        spike metadata to be uploaded.

    Returns
    -------
    Metadata
        created metadata object.
    """

    # check if metadata already exists for period before uploading
    prior_metadata_ID = SpikesAlgorithm.check_existing_metadata(
        algorithm, metadata_obj=spike
    )
    if prior_metadata_ID:
        # TODO: Confirm whether or not to add force or simply update automatically
        typer.confirm(
            f"Spikes already exist for this period, would you like to update this metadata?",
            abort=True,
        )

        # set spike.id to prior_metadata.id to update that metadata object in the database
        spike.id = prior_metadata_ID
        return SpikesAlgorithm.update_metadata(algorithm, spike)
    else:
        return SpikesAlgorithm.create_new_metadata(algorithm, spike)


if __name__ == "__main__":
    main()