# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/10_plot.ipynb.

# %% auto 0
__all__ = ['plot_waveform', 'plot_spectrogram']

# %% ../../nbs/10_plot.ipynb 2
from pathlib import Path
from typing import assert_never

import librosa
import numpy as np
import matplotlib.pyplot as plt

from bioacoustics_augmentation.utils.audio import Spectrogram

# %% ../../nbs/10_plot.ipynb 4
def plot_waveform(
    waveform: np.ndarray,
    sample_rate: int,
) -> None:
    """
    Plots a waveform.

    Args:
        waveform: The audio waveform to plot.
        sample_rate: The sample rate of the audio waveform.
        title: The title of the plot.
    """
    librosa.display.waveshow(
        waveform,
        sr=sample_rate,
    )

# %% ../../nbs/10_plot.ipynb 5
def plot_spectrogram(
    spectrogram: np.ndarray | Spectrogram,
    *,
    sample_rate: int = 22050,
    fmax: int | None = None,
) -> None:
    """
    Plots a spectrogram.

    Args:
        spectrogram: The spectrogram to plot.
        sample_rate: Sample rate for proper frequency axis scaling.
        fmax: Maximum frequency to display (optional).
    """
    data: np.ndarray = spectrogram.data if isinstance(spectrogram, Spectrogram) else spectrogram

    if isinstance(spectrogram, np.ndarray):
        librosa.display.specshow(
            data,
            x_axis='time',
            y_axis='mel',
            fmax=fmax,
            sr=sample_rate,
        )
    elif isinstance(spectrogram, Spectrogram):
        librosa.display.specshow(
            spectrogram.data, 
            sr=spectrogram.metadata.sample_rate, 
            hop_length=spectrogram.metadata.hop_length, 
            x_axis='time', 
            y_axis='mel',
            fmax=fmax,
        )
    else:
        assert_never()
