aimbat.plot

def plot_matrix_image(
    iccs: ICCS, context: bool, all_seismograms: bool, return_fig: bool
) -> tuple["Figure", "Axes"] | None:
    """Plot the ICCS seismograms as a matrix image.

    The matrix is assembled from individual waveforms, with each row representing
    a different seismogram.

    Args:
        iccs: ICCS instance.
        context: If True, plot waveforms with extra context around the taper window.
        all_seismograms: If True, include deselected seismograms in the plot.
        return_fig: If True, return the figure and axes objects instead of showing the plot.

    Returns:
        A tuple of (Figure, Axes) if return_fig is True, otherwise None.
    """

    logger.info("Plotting matrix image.")

    return _plot_matrix_image(iccs, context, all_seismograms, return_fig=return_fig)  # type: ignore[call-overload]

def plot_seismograms(
    plot_for: AimbatEvent | AimbatStation, return_fig: bool
) -> tuple[plt.Figure, plt.Axes] | None:
    """Plot all seismograms for a particular event or station.

    Args:
        plot_for: What to plot the seismograms for (Event or Station).
        return_fig: Whether to return the figure and axes objects instead of showing the plot.

    Returns:
        figure and axes objects if return_fig is True, otherwise None.

    Raises:
        RuntimeError: If no seismograms are found for the event or station.

    Note:
        The seismograms use the filter settings specified in the event parameters.
        `plot_for` is read while the figure is built, so it must still be attached
        to a session at that point; nothing touches it afterwards. Pass
        `return_fig=True` and call `show_figure` separately to close the session
        before the window opens.
    """
    logger.info(f"Plotting seismograms for {type(plot_for).__name__}: {plot_for.id}.")

    fig, ax = _plot_seis(plot_for)
    _add_scroll_pan(ax)

    if return_fig:
        return fig, ax
    show_figure(fig)
    return None

def plot_stack(
    iccs: ICCS, context: bool, all_seismograms: bool, return_fig: bool
) -> tuple["Figure", "Axes"] | None:
    """Plot the ICCS stack.

    Args:
        iccs: ICCS instance.
        context: If True, plot waveforms with extra context around the taper window.
        all_seismograms: If True, include deselected seismograms in the plot.
        return_fig: If True, return the figure and axes objects instead of showing the plot.

    Returns:
        A tuple of (Figure, Axes) if return_fig is True, otherwise None.
    """

    logger.info("Plotting ICCS stack.")
    return _plot_stack(iccs, context, all_seismograms, return_fig=return_fig)  # type: ignore[call-overload]

def show_figure(fig: plt.Figure) -> None:
    """Show a figure and block until the user closes its window.

    Args:
        fig: Figure to show.
    """
    plt.show()
    plt.close(fig)

def update_bandpass(
    event_id: UUID,
    iccs: ICCS,
    context: bool,
    all_seismograms: bool,
    use_matrix_image: bool,
    return_fig: bool,
    engine: Engine | None = None,
) -> tuple["Figure", "Axes", Any] | None:
    """Update the bandpass filter parameters for an event.

    Args:
        event_id: UUID of the event being updated.
        iccs: ICCS instance.
        context: If True, plot waveforms with extra context around the taper window.
        all_seismograms: If True, include deselected seismograms in the plot.
        use_matrix_image: If True, pick from the matrix image; otherwise pick from the stack plot.
        return_fig: If True, return the figure, axes and widget objects instead of showing the plot.
        engine: Engine to persist the change with. Defaults to the project engine.

    Returns:
        A tuple of (Figure, Axes, widgets) if return_fig is True, otherwise None.
    """

    logger.info(f"Updating bandpass filter parameters for event {event_id}.")

    result = _update_bandpass(  # type: ignore[call-overload]
        iccs, context, all_seismograms, use_matrix_image, return_fig=return_fig
    )

    if not return_fig:
        logger.debug(
            f"Saving new bandpass filter parameters for event {event_id}: apply="
            + f"{iccs.bandpass_apply}, fmin={iccs.bandpass_fmin}, fmax="
            + f"{iccs.bandpass_fmax}, corners={iccs.corners}"
        )
        with _evict_cached_iccs_on_failure(event_id), _write_session(engine) as session:
            set_event_parameters(
                session,
                event_id,
                {
                    EventParameter.BANDPASS_APPLY: iccs.bandpass_apply,
                    EventParameter.BANDPASS_FMIN: iccs.bandpass_fmin,
                    EventParameter.BANDPASS_FMAX: iccs.bandpass_fmax,
                    EventParameter.CORNERS: iccs.corners,
                },
            )
        return None

    logger.debug(_RETURN_FIG_WARNING)
    return result

def update_min_cc(
    event_id: UUID,
    iccs: ICCS,
    context: bool,
    all_seismograms: bool,
    causal: bool,
    return_fig: bool,
    engine: Engine | None = None,
) -> tuple["Figure", "Axes", Any] | None:
    """Update the minimum cross-correlation threshold for the given event.

    Args:
        event_id: UUID of the event being updated.
        iccs: ICCS instance.
        context: If True, plot waveforms with extra context around the taper window.
        all_seismograms: If True, include deselected seismograms in the plot.
        causal: If True, use causal (single-pass) instead of zero-phase filtering.
        return_fig: If True, return the figure, axes and widget objects instead of showing the plot.
        engine: Engine to persist the change with. Defaults to the project engine.

    Returns:
        A tuple of (Figure, Axes, widgets) if return_fig is True, otherwise None.
    """

    logger.info(f"Updating minimum cross-correlation threshold for event {event_id}.")

    result = _update_min_cc(  # type: ignore[call-overload]
        iccs, context, all_seismograms, causal, return_fig=return_fig
    )

    if not return_fig:
        logger.debug(
            f"Saving new minimum cross-correlation threshold for event {event_id}:"
            + f" {iccs.min_cc}"
        )
        with _evict_cached_iccs_on_failure(event_id), _write_session(engine) as session:
            set_event_parameter(
                session, event_id, EventParameter.MIN_CC, float(iccs.min_cc)
            )
        return None

    logger.debug(_RETURN_FIG_WARNING)
    return result

def update_pick(
    event_id: UUID,
    iccs: ICCS,
    context: bool,
    all_seismograms: bool,
    use_matrix_image: bool,
    causal: bool,
    return_fig: bool,
    engine: Engine | None = None,
) -> tuple["Figure", "Axes", Any] | None:
    """Update the phase pick (t1) for an event.

    Args:
        event_id: UUID of the event being updated.
        iccs: ICCS instance.
        context: If True, plot waveforms with extra context around the taper window.
        all_seismograms: If True, include deselected seismograms in the plot.
        use_matrix_image: If True, pick from the matrix image; otherwise pick from the stack plot.
        causal: If True, use causal (single-pass) instead of zero-phase filtering.
        return_fig: If True, return the figure, axes and widget objects instead of showing the plot.
        engine: Engine to persist the change with. Defaults to the project engine.

    Returns:
        A tuple of (Figure, Axes, widgets) if return_fig is True, otherwise None.
    """

    logger.info("Updating phase pick.")

    result = _update_pick(  # type: ignore[call-overload]
        iccs, context, all_seismograms, use_matrix_image, causal, return_fig=return_fig
    )

    if not return_fig:
        with _evict_cached_iccs_on_failure(event_id), _write_session(engine) as session:
            write_back_seismograms(session, iccs)
            session.commit()
        return None

    logger.debug(_RETURN_FIG_WARNING)
    return result

def update_timewindow(
    event_id: UUID,
    iccs: ICCS,
    context: bool,
    all_seismograms: bool,
    use_matrix_image: bool,
    causal: bool,
    return_fig: bool,
    engine: Engine | None = None,
) -> tuple["Figure", "Axes", Any] | None:
    """Update the cross-correlation time window for the given event.

    Args:
        event_id: UUID of the event being updated.
        iccs: ICCS instance.
        context: If True, plot waveforms with extra context around the taper window.
        all_seismograms: If True, include deselected seismograms in the plot.
        use_matrix_image: If True, pick from the matrix image; otherwise pick from the stack plot.
        causal: If True, use causal (single-pass) instead of zero-phase filtering.
        return_fig: If True, return the figure, axes and widget objects instead of showing the plot.
        engine: Engine to persist the change with. Defaults to the project engine.

    Returns:
        A tuple of (Figure, Axes, widgets) if return_fig is True, otherwise None.
    """

    logger.info(f"Updating time window for event {event_id}.")

    result = _update_timewindow(  # type: ignore[call-overload]
        iccs, context, all_seismograms, use_matrix_image, causal, return_fig=return_fig
    )

    if not return_fig:
        logger.debug(
            f"Saving new time window for event {event_id}: pre={iccs.window_pre}, "
            + f"post={iccs.window_post}"
        )
        with _evict_cached_iccs_on_failure(event_id), _write_session(engine) as session:
            set_event_parameters(
                session,
                event_id,
                {
                    EventParameter.WINDOW_PRE: iccs.window_pre,
                    EventParameter.WINDOW_POST: iccs.window_post,
                },
            )
        return None

    logger.debug(_RETURN_FIG_WARNING)
    return result