aimbat.lib

def string_to_uuid(
    session: Session,
    id: str,
    aimbat_class: type[
        AimbatDataSource
        | AimbatStation
        | AimbatEvent
        | AimbatEventParameters
        | AimbatSeismogram
        | AimbatSeismogramParameters
        | AimbatSnapshot
        | AimbatEventParametersSnapshot
        | AimbatSeismogramParametersSnapshot
    ],
    custom_error: str | None = None,
) -> UUID:
    """Determine a UUID from a string containing the first few characters.

    Parameters:
        session: Database session.
        id: Input string to find UUID for.
        aimbat_class: Aimbat class to use to find UUID.
        custom_error: Overrides the default error message.

    Returns:
        The full UUID.

    Raises:
        ValueError: If the UUID could not be determined.
    """
    uuid_set = {
        u for u in session.exec(select(aimbat_class.id)).all() if str(u).startswith(id)
    }
    if len(uuid_set) == 1:
        return uuid_set.pop()
    if len(uuid_set) == 0:
        raise ValueError(
            custom_error or f"Unable to find {aimbat_class.__name__} using id: {id}."
        )
    raise ValueError(f"Found more than one {aimbat_class.__name__} using id: {id}")

def add_files_to_project(
    datasources: Sequence[str | os.PathLike],
    datatype: DataType,
    disable_progress_bar: bool = True,
) -> None:
    """Add files to the AIMBAT database.

    Parameters:
        datasources: List of data sources to add.
        datatype: Type of data.
        disable_progress_bar: Do not display progress bar.
    """

    logger.info(f"Adding {len(datasources)} {datatype} files to project.")

    with Session(engine) as session:
        for datasource in track(
            sequence=datasources,
            description="Adding files ...",
            disable=disable_progress_bar,
        ):
            aimbat_station = _create_station(session, datasource, datatype)
            aimbat_event = _create_event(session, datasource, datatype)
            aimbat_seismogram = _create_seismogram(session, datasource, datatype)

            # TODO: perhaps adding potentially updated station and event information should be optional?
            aimbat_seismogram.station = aimbat_station
            aimbat_seismogram.event = aimbat_event

            # Create AimbatDataSource instance with relationship to AimbatSeismogram
            select_aimbat_data_source = select(AimbatDataSource).where(
                AimbatDataSource.sourcename == str(datasource)
            )
            aimbat_data_source = session.exec(select_aimbat_data_source).one_or_none()
            if aimbat_data_source is None:
                logger.debug(f"Adding data source {datasource} to project.")
                aimbat_data_source_create = AimbatDataSourceCreate(
                    sourcename=str(datasource), datatype=datatype
                )
                aimbat_data_source = AimbatDataSource.model_validate(
                    aimbat_data_source_create, update={"seismogram": aimbat_seismogram}
                )

            else:
                logger.debug(
                    f"Using existing data source {datasource} instead of adding new one."
                )
                aimbat_data_source.seismogram = aimbat_seismogram
            session.add(aimbat_data_source)

        session.commit()

def dump_data_table() -> None:
    """Dump the table data to json."""

    logger.info("Dumping AIMBAT datasources table to json.")

    with Session(engine) as session:
        aimbat_data_sources = session.exec(select(AimbatDataSource)).all()
        dump_to_json(aimbat_data_sources)

def get_data_for_active_event(session: Session) -> Sequence[AimbatDataSource]:
    """Returns the AimbatFiles belonging to the active event.

    Parameters:
        session: Database session.

    Returns:
        List of AimbatFiles.
    """

    logger.info("Getting aimbatfiles in active event.")

    select_files = (
        select(AimbatDataSource)
        .join(AimbatSeismogram)
        .join(AimbatEvent)
        .where(AimbatEvent.active == 1)
    )
    return session.exec(select_files).all()

def print_data_table(short: bool, all_events: bool = False) -> None:
    """Print a pretty table with AIMBAT data.

    Parameters:
        short: Shorten UUIDs and format data.
        all_events: Print all files instead of limiting to the active event.
    """

    logger.info("Printing AIMBAT data table.")

    with Session(engine) as session:
        if all_events:
            aimbat_data_sources = session.exec(select(AimbatDataSource)).all()
            title = "AIMBAT data for all events"
        else:
            active_event = get_active_event(session)
            aimbat_data_sources = get_data_for_active_event(session)
            time = (
                active_event.time.strftime("%Y-%m-%d %H:%M:%S")
                if short
                else active_event.time
            )
            id = uuid_shortener(session, active_event) if short else active_event.id
            title = f"AIMBAT data for event {time} (ID={id})"

        logger.debug(f"Found {len(aimbat_data_sources)} files in total.")

        rows = [
            [
                uuid_shortener(session, a) if short else str(a.id),
                str(a.datatype),
                str(a.sourcename),
                (
                    uuid_shortener(session, a.seismogram)
                    if short
                    else str(a.seismogram.id)
                ),
            ]
            for a in aimbat_data_sources
        ]

        table = make_table(title=title)

        table.add_column(
            "ID (shortened)" if short else "ID",
            justify="center",
            style=TABLE_STYLING.id,
            no_wrap=True,
        )
        table.add_column("Datatype", justify="center", style=TABLE_STYLING.mine)
        table.add_column(
            "Filename", justify="left", style=TABLE_STYLING.mine, no_wrap=True
        )
        table.add_column(
            "Seismogram ID", justify="center", style=TABLE_STYLING.linked, no_wrap=True
        )

        for row in rows:
            table.add_row(*row)

        console = Console()
        console.print(table)

def delete_event(session: Session, event: AimbatEvent) -> None:
    """Delete an AimbatEvent from the database.

    Parameters:
        session: Database session.
        event: Event to delete.
    """

    logger.info(f"Deleting event {event.id}.")

    session.delete(event)
    session.commit()

def delete_event_by_id(session: Session, event_id: UUID) -> None:
    """Delete an AimbatEvent from the database by ID.

    Parameters:
        session: Database session.
        event_id: Event ID.

    Raises:
        NoResultFound: If no AimbatEvent is found with the given ID.
    """

    logger.debug(f"Getting event with id={event_id}.")

    event = session.get(AimbatEvent, event_id)
    if event is None:
        raise NoResultFound(
            f"Unable to find event using id: {event_id}. {HINTS.LIST_EVENTS}"
        )
    delete_event(session, event)

def dump_event_table() -> None:
    """Dump the table data to json."""

    logger.info("Dumping AIMBAT event table to json.")

    with Session(engine) as session:
        aimbat_events = session.exec(select(AimbatEvent)).all()
        dump_to_json(aimbat_events)

def get_active_event(session: Session) -> AimbatEvent:
    """
    Return the currently active event (i.e. the one being processed).

    Parameters:
        session: SQL session.

    Returns:
        Active Event

    Raises
        NoResultFound: When no event is active.
    """

    logger.debug("Attempting to determine active event.")

    select_active_event = select(AimbatEvent).where(AimbatEvent.active == 1)

    # NOTE: While there technically can be no active event in the database,
    # we typically don't really want to go beyond this point when that is the
    # case. Hence we call `one` rather than `one_or_none`.
    try:
        active_event = session.exec(select_active_event).one()
    except NoResultFound:
        raise NoResultFound(f"No active event found. {HINTS.ACTIVATE_EVENT}")

    logger.debug(f"Active event: {active_event.id}")

    return active_event

def get_completed_events(session: Session) -> Sequence[AimbatEvent]:
    """Get the events marked as completed.

    Parameters:
        session: SQL session.
    """

    select_completed_events = (
        select(AimbatEvent)
        .join(AimbatEventParameters)
        .where(AimbatEventParameters.completed == 1)
    )

    return session.exec(select_completed_events).all()

def get_event_parameter(
    session: Session, name: EventParameter
) -> timedelta | bool | float:
    """Get event parameter value for the active event.

    Parameters:
        session: Database session.
        name: Name of the parameter.
    """

    active_event = get_active_event(session)

    logger.info(f"Getting {name=} value for {active_event=}.")

    return getattr(active_event.parameters, name)

def get_events_using_station(
    session: Session, station: AimbatStation
) -> Sequence[AimbatEvent]:
    """Get all events that use a particular station.

    Parameters:
        session: Database session.
        station: Station to return events for.

    Returns: Events that use the station.
    """

    logger.info(f"Getting events for station: {station.id}.")

    select_events = (
        select(AimbatEvent)
        .join(AimbatSeismogram)
        .join(AimbatStation)
        .where(AimbatStation.id == station.id)
    )

    events = session.exec(select_events).all()

    logger.debug(f"Found {len(events)}.")

    return events

def print_event_table(short: bool = True) -> None:
    """Print a pretty table with AIMBAT events.

    Parameters:
        short: Shorten and format the output to be more human-readable.
    """

    logger.info("Printing AIMBAT events table.")

    table = make_table(title="AIMBAT Events")
    table.add_column(
        "ID (shortened)" if short else "ID",
        justify="center",
        style=TABLE_STYLING.id,
        no_wrap=True,
    )
    table.add_column("Active", justify="center", style=TABLE_STYLING.mine, no_wrap=True)
    table.add_column(
        "Date & Time", justify="center", style=TABLE_STYLING.mine, no_wrap=True
    )
    table.add_column("Latitude", justify="center", style=TABLE_STYLING.mine)
    table.add_column("Longitude", justify="center", style=TABLE_STYLING.mine)
    table.add_column("Depth", justify="center", style=TABLE_STYLING.mine)
    table.add_column("Completed", justify="center", style=TABLE_STYLING.parameters)
    table.add_column("# Seismograms", justify="center", style=TABLE_STYLING.linked)
    table.add_column("# Stations", justify="center", style=TABLE_STYLING.linked)

    with Session(engine) as session:
        for event in session.exec(select(AimbatEvent)).all():
            logger.debug(f"Adding event with id={event.id} to the table.")
            table.add_row(
                uuid_shortener(session, event) if short else str(event.id),
                TABLE_STYLING.bool_formatter(event.active),
                TABLE_STYLING.datetime_formatter(event.time, short),
                f"{event.latitude:.3f}" if short else str(event.latitude),
                f"{event.longitude:.3f}" if short else str(event.longitude),
                f"{event.depth:.0f}" if short else str(event.depth),
                TABLE_STYLING.bool_formatter(event.parameters.completed),
                str(len(event.seismograms)),
                str(len(station.get_stations_in_event(session, event))),
            )

    console = Console()
    console.print(table)

def set_active_event(session: Session, event: AimbatEvent) -> None:
    """
    Set the active event (i.e. the one being processed).

    Parameters:
        session: SQL session.
        event: AIMBAT Event to set as active.
    """

    logger.info(f"Activating {event=}")

    event.active = True
    session.add(event)
    session.commit()

def set_active_event_by_id(session: Session, event_id: UUID) -> None:
    """
    Set the currently selected event (i.e. the one being processed) by its ID.

    Parameters:
        session: SQL session.
        event_id: ID of AIMBAT Event to set as active one.

    Raises:
        ValueError: If no event with the given ID is found.
    """
    logger.info(f"Setting active event to event with id={event_id}.")

    if event_id not in session.exec(select(AimbatEvent.id)).all():
        raise ValueError(
            f"No AimbatEvent found with id: {event_id}. {HINTS.LIST_EVENTS}"
        )

    aimbat_event = session.exec(
        select(AimbatEvent).where(AimbatEvent.id == event_id)
    ).one()
    set_active_event(session, aimbat_event)

def set_event_parameter(
    session: Session, name: EventParameter, value: timedelta | bool | float | str
) -> None:
    """Set event parameter value for the active event.

    Parameters:
        session: Database session.
        name: Name of the parameter.
        value: Value to set.
    """

    active_event = get_active_event(session)

    logger.info(f"Setting {name=} to {value} for {active_event=}.")

    parameters = AimbatEventParametersBase.model_validate(
        active_event.parameters, update={name: value}
    )
    setattr(active_event.parameters, name, getattr(parameters, name))
    session.add(active_event)
    session.commit()

def create_iccs_instance(session: Session) -> ICCS:
    """Create an ICCS instance for the active event.

    Parameters:
        session: Database session.

    Returns:
        ICCS instance.
    """

    logger.info("Creating ICCS instance for active event.")

    active_event = event.get_active_event(session)

    return ICCS(
        seismograms=active_event.seismograms,
        window_pre=active_event.parameters.window_pre,
        window_post=active_event.parameters.window_post,
        min_ccnorm=active_event.parameters.min_ccnorm,
        context_width=settings.context_width,
    )

def plot_seismograms(iccs: ICCS, context: bool, all: bool) -> None:
    """Plot the ICCS seismograms as an image.

    Parameters:
        iccs: ICCS instance.
        context: Whether to use seismograms with extra context.
        all: Whether to plot all seismograms.
    """

    logger.info("Plotting ICCS seismograms for active event.")

    _plot_seismograms(iccs, context, all)

def plot_stack(iccs: ICCS, context: bool, all: bool) -> None:
    """Plot the ICCS stack.

    Parameters:
        iccs: ICCS instance.
        context: Whether to use seismograms with extra context.
        all: Whether to plot all seismograms.
    """

    logger.info("Plotting ICCS stack for active event.")
    _plot_stack(iccs, context, all)

def run_iccs(session: Session, iccs: ICCS, autoflip: bool, autoselect: bool) -> None:
    """Run ICCS algorithm.

    Parameters:
        session: Database session.
        iccs: ICCS instance.
        autoflip: Whether to automatically flip seismograms.
        autoselect: Whether to automatically select seismograms.
    """

    logger.info(f"Running ICCS with {autoflip=}, {autoselect=}.")

    results = iccs(autoflip=autoflip, autoselect=autoselect)
    logger.info(f"ICCS {results = }")
    session.commit()

def update_min_ccnorm(session: Session, iccs: ICCS, context: bool, all: bool) -> None:
    """Update the minimum cross correlation coefficient for the active event.

    Parameters:
        iccs: ICCS instance.
        context: Whether to use seismograms with extra context.
        all: Whether to plot all seismograms.
    """

    logger.info("Updating minimum cross correlation coefficient for active event.")

    logger.debug(f"Current {iccs.min_ccnorm = }.")
    _update_min_ccnorm(iccs, context, all)
    logger.debug(f"Updated {iccs.min_ccnorm = }.")

    active_event = event.get_active_event(session)
    active_event.parameters.min_ccnorm = float(iccs.min_ccnorm)
    session.commit()

def update_pick(
    session: Session, iccs: ICCS, context: bool, all: bool, use_seismogram_image: bool
) -> None:
    """Update the pick for the active event.

    Parameters:
        iccs: ICCS instance.
        context: Whether to use seismograms with extra context.
        all: Whether to plot all seismograms.
        use_seismogram_image: Whether to use the seismogram image to update pick.
    """

    logger.info("Updating pick for active event.")

    _update_pick(iccs, context, all, use_seismogram_image)
    session.commit()

def update_timewindow(
    session: Session, iccs: ICCS, context: bool, all: bool, use_seismogram_image: bool
) -> None:
    """Update the time window for the active event.

    Parameters:
        iccs: ICCS instance.
        context: Whether to use seismograms with extra context.
        all: Whether to plot all seismograms.
        use_seismogram_image: Whether to use the seismogram image to update pick.
    """

    logger.info("Updating time window for active event.")

    logger.debug(f"Current {iccs.window_pre = }, {iccs.window_post = }.")
    _update_timewindow(iccs, context, all, use_seismogram_image)
    logger.debug(f"Updated {iccs.window_pre = }, {iccs.window_post = }.")

    active_event = event.get_active_event(session)
    active_event.parameters.window_pre = iccs.window_pre
    active_event.parameters.window_post = iccs.window_post
    session.commit()

def create_event(datasource: str | PathLike, datatype: DataType) -> AimbatEvent:
    """Read event data from a data source and create an AimbatEvent.

    Parameters:
        datasource: Name of the data source.
        datatype: AIMBAT compatible datatype.

    Returns:
        AimbatEvent instance.

    Raises:
        NotImplementedError: If the datatype is not supported.
    """

    logger.debug(f"Creating AimbatEvent from {datasource}.")

    event_creator_fn = event_creator.get(datatype)
    if event_creator_fn is None:
        raise NotImplementedError(
            f"I don't know how to create an AimbatEvent from {datatype}."
        )
    return event_creator_fn(datasource)

def create_seismogram(
    datasource: str | PathLike, datatype: DataType
) -> AimbatSeismogram:
    """Read seismogram data from a data source and create an AimbatSeismogram.

    Parameters:
        datasource: Name of the data source.
        datatype: AIMBAT compatible datatype.

    Returns:
        AimbatSeismogram instance.

    Raises:
        NotImplementedError: If the datatype is not supported.
    """

    logger.debug(f"Creating AimbatSeismogram from {datasource}.")

    station_creator_fn = seismogram_creator.get(datatype)
    if station_creator_fn is None:
        raise NotImplementedError(
            f"I don't know how to create an AimbatSeismgoram from {datatype}."
        )
    return station_creator_fn(datasource)