ISO Plot¶

def create_subplots(
    self,
    nrows: int = 1,
    ncols: int = 1,
    figsize: tuple[int, int] = (5, 5),
    subplot_by: str | None = None,
    subplot_datas: list[pd.DataFrame] | None = None,
    subplot_titles: list[str] | None = None,
    *,
    adjust_figsize: bool = True,
    auto_allocate_axes: bool = False,
    **kwargs,
) -> ISOPlot:
    """
    Create subplots for the circumplex plot.

    Parameters
    ----------
    nrows
        Number of rows in the subplot grid, by default 1
    ncols
        Number of columns in the subplot grid, by default 1
    figsize
        Size of the figure (width, height), by default (5, 5)
    subplot_by
        Column name to create subplots by unique values, by default None
    subplot_datas
        List of dataframes for each subplot, by default None
    subplot_titles
        List of titles for each subplot, by default None
    adjust_figsize
        Whether to adjust the figure size based on nrows/ncols, by default True
    auto_allocate_axes
        Whether to automatically determine nrows/ncols based on data,
        by default False
    **kwargs :
        Additional parameters for plt.subplots

    Returns
    -------
    :
        The current plot instance for chaining

    Examples
    --------
    Create a basic subplot grid:

    >>> import pandas as pd
    >>> import numpy as np
    >>> rng = np.random.default_rng(42)
    >>> data = pd.DataFrame(
    ...    np.c_[rng.multivariate_normal([0.2, 0.15], [[0.1, 0], [0, 0.2]], 100),
    ...          rng.integers(1, 3, 100)],
    ...    columns=['ISOPleasant', 'ISOEventful', 'Group'])
    >>> plot = ISOPlot(data=data).create_subplots(nrows=2, ncols=2)
    >>> len(plot.subplot_contexts) == 4
    True
    >>> plot.close()  # Clean up

    Create subplots by a column in the data:

    >>> plot = (ISOPlot(data=data)
    ...         .create_subplots(nrows=1, ncols=2, subplot_by='Group'))
    >>> len(plot.subplot_contexts) == 2
    True
    >>> plot.close()  # Clean up

    Create subplots with auto-allocation of axes:

    >>> plot = (ISOPlot(data=data)
    ...        .create_subplots(subplot_by='Group', auto_allocate_axes=True))
    >>> len(plot.subplot_contexts) == 2
    True
    >>> plot.close()  # Clean up

    """
    # Set up subplot params
    self.subplots_params.update(
        nrows=nrows,
        ncols=ncols,
        figsize=figsize,
        subplot_by=subplot_by,
        adjust_figsize=adjust_figsize,
        auto_allocate_axes=auto_allocate_axes,
        **kwargs,
    )
    # Create a list of dataframes and titles for each subplot
    # based on the unique values in the specified column
    if self.subplots_params.subplot_by:
        logger.debug(
            "Creating subplots by unique values "
            f"in {self.subplots_params.subplot_by}."
        )
        subplot_datas, subplot_titles, n_subplots_by = self._setup_subplot_by(
            self.subplots_params.subplot_by, subplot_datas, subplot_titles
        )
    else:
        n_subplots_by = -1

    if subplot_titles and self.subplots_params.auto_allocate_axes:
        # Attempt to allocate axes based on the number of subplots
        self.subplots_params.nrows, self.subplots_params.ncols = (
            self._allocate_subplot_axes(subplot_titles)
        )

    if adjust_figsize:
        self.subplots_params.figsize = (
            self.subplots_params.ncols * self.subplots_params.figsize[0],
            self.subplots_params.nrows * self.subplots_params.figsize[1],
        )

    logger.debug(f"Subplot parameters: {self.subplots_params}")

    # Create the figure and axes
    self.figure, self.axes = plt.subplots(
        **self.subplots_params.as_plt_subplots_args()
    )

    # If subplot_datas or subplot_titles are provided, validate them
    if subplot_datas is not None or subplot_titles is not None:
        self._validate_subplots_datas(subplot_datas, subplot_titles)

    # Create PlotContext objects for each subplot
    self.subplot_contexts = []

    for i, ax in enumerate(self.yield_axes_objects()):
        if i >= self._naxes:
            break
        if subplot_by and i >= n_subplots_by:
            logger.debug(f"Created {i + 1} subplots for {subplot_by}.")
            break
        # Get data and title for this subplot if available
        data = (
            subplot_datas[i] if subplot_datas and i < len(subplot_datas) else None
        )
        title = (
            subplot_titles[i]
            if subplot_titles and i < len(subplot_titles)
            else None
        )

        context = self.main_context.create_child(data=data, title=title, ax=ax)
        self.subplot_contexts.append(context)

    return self

def show(self) -> None:
    """
    Show the figure.

    This method is a wrapper around plt.show() to display the figure.

    """
    if self.figure is None:
        msg = (
            "No figure object provided. "
            "Please create a figure using create_subplots() first."
        )
        raise ValueError(msg)
    if self._has_subplots:
        plt.tight_layout()
    plt.show()

def close(self, fig: int | str | Figure | None = None) -> None:
    """
    Close the figure.

    This method is a wrapper around plt.close() to close the figure.

    """
    if fig is None:
        fig = self.figure
        if fig is None:
            msg = (
                "No figure object provided. "
                "Please create a figure using create_subplots() first."
            )
            raise ValueError(msg)
    plt.close(fig)

def savefig(self, *args: Any, **kwargs: Any) -> None:
    """
    Save the figure.

    This method is a wrapper around plt.savefig() to save the figure.

    """
    if self.figure is None:
        msg = (
            "No figure object provided. "
            "Please create a figure using create_subplots() first."
        )
        raise ValueError(msg)
    self.figure.savefig(*args, **kwargs)

def get_figure(self) -> Figure | SubFigure:
    """
    Get the figure object.

    Returns
    -------
    :
        The matplotlib Figure or SubFigure object associated with this plot.


    Raises
    ------
    ValueError
        If the figure object does not exist.
    TypeError
        If the figure object is not a valid Figure or SubFigure.

    """
    if self.figure is None:
        msg = (
            "No figure object provided. "
            "Please create a figure using create_subplots() first."
        )
        raise ValueError(msg)
    if isinstance(self.figure, Figure | SubFigure):
        return self.figure
    msg = "Invalid figure object. Please provide a valid Figure or SubFigure."
    raise TypeError(msg)

def get_axes(self) -> Axes | np.ndarray:
    """
    Get the axes object.

    Returns
    -------
    :
        The matplotlib Axes object or array of Axes associated with this plot.

    Raises
    ------
    ValueError
        If the axes object does not exist.
    TypeError
        If the axes object is not a valid Axes or ndarray of Axes.

    """
    self._check_for_axes()
    if isinstance(self.axes, Axes | np.ndarray):
        return self.axes
    msg = "Invalid axes object. Please provide a valid Axes or ndarray of Axes."
    raise TypeError(msg)

def get_single_axes(self, ax_idx: int | tuple[int, int] | None = None) -> Axes:
    """
    Get a specific axes object.

    Parameters
    ----------
    ax_idx
        The index of the axes to get. If None, returns the first axes.
        Can be an integer for flattened access or a tuple of (row, col).

    Returns
    -------
    :
        The requested matplotlib Axes object

    Raises
    ------
    ValueError
        If the axes object does not exist or the index is invalid.
    TypeError
        If the axes object is not a valid Axes or ndarray of Axes.

    """
    self._check_for_axes()

    def validate_tuple_axes_index(
        nrows: int, ncols: int, naxes: int, ax_idx: tuple[int, int]
    ) -> None:
        """
        Validate the tuple axes index.

        This checks the `ax_idx` types and compares the implied number of axes
        with the actual number of axes in the figure.
        """
        if (
            len(ax_idx) != 2  # noqa: PLR2004
            or not isinstance(ax_idx[0], int)
            or not isinstance(ax_idx[1], int)
            or ax_idx[0] < 0
            or ax_idx[1] < 0
        ):
            msg = (
                "Invalid axes index provided. "
                "Expected a tuple of 2 positive integers."
            )
            raise ValueError(msg)

        if ax_idx[0] >= (nrows - 1) or ax_idx[1] >= (ncols - 1):
            msg = (
                "Invalid axes index provided."
                f" The figure contains {nrows} rows and {ncols} columns. "
                f"ax_idx implied {ax_idx[0] + 1} rows and {ax_idx[1] + 1} columns."
            )
            raise ValueError(msg)

        idx_implied_n_axes = (ax_idx[0] + 1) * (ax_idx[1] + 1)
        if naxes < idx_implied_n_axes:
            msg = (
                "Invalid axes index provided."
                f" The figure contains {naxes} axes. "
                f"ax_idx implied {idx_implied_n_axes} axes."
            )
            raise ValueError(msg)

    def validate_int_axes_index(naxes: int, ax_idx: int) -> None:
        """
        Validate the integer axes index.

        This checks the `ax_idx` type and compares the implied number of axes
        with the actual number of axes in the figure.
        """
        if not isinstance(ax_idx, int) or ax_idx < 0:
            msg = "Invalid axes index provided. Expected a positive integer."
            raise ValueError(msg)

        if (ax_idx + 1) > naxes:
            msg = (
                "Invalid axes index provided."
                f" The figure contains {naxes} axes. "
                f"ax_idx implied {ax_idx + 1} axes."
            )
            raise ValueError(msg)

    if isinstance(self.axes, np.ndarray) and ax_idx is not None:
        if isinstance(ax_idx, tuple):
            validate_tuple_axes_index(self._nrows, self._ncols, self._naxes, ax_idx)
            return self.axes[ax_idx[0], ax_idx[1]]

        validate_int_axes_index(self._naxes, ax_idx)
        return self.axes.flatten()[ax_idx]

    if isinstance(self.axes, Axes) and (ax_idx == 0 or ax_idx is None):
        return self.axes

    msg = "Invalid axes index provided."
    raise ValueError(msg)

def yield_axes_objects(self) -> Generator[Axes, None, None]:
    """
    Generate a sequence of axes objects to iterate over.

    This method is a helper to iterate over all axes in the figure,
    whether the figure contains a single Axes object or an array of Axes objects.

    Yields
    ------
    :
        Individual matplotlib Axes objects from the current figure.

    """
    if isinstance(self.axes, np.ndarray):
        yield from self.axes.flatten()
    elif isinstance(self.axes, Axes):
        yield self.axes

def add_layer(
    self,
    layer_class: type[Layer],
    data: pd.DataFrame | None = None,
    *,
    on_axis: int | tuple[int, int] | list[int] | None = None,
    **params: Any,
) -> ISOPlot:
    """
    Add a visualization layer, optionally targeting specific subplot(s).

    Parameters
    ----------
    layer_class
        The type of layer to add
    on_axis
        Target specific axis/axes:

        - int: Index of subplot (flattened)
        - tuple: (row, col) coordinates
        - list: Multiple indices to apply the layer to
        - None: Apply to all subplots (default)
    data
        Custom data for this specific layer, overriding context data
    **params
        Parameters for the layer

    Returns
    -------
    :
        The current plot instance for chaining

    Examples
    --------
    Add a scatter layer to all subplots:

    >>> import pandas as pd
    >>> import numpy as np
    >>> from soundscapy.plotting.layers import ScatterLayer
    >>> rng = np.random.default_rng(42)
    >>> data = pd.DataFrame(
    ...    np.c_[rng.multivariate_normal([0.2, 0.15], [[0.1, 0], [0, 0.2]], 100),
    ...          rng.integers(1, 3, 100)],
    ...    columns=['ISOPleasant', 'ISOEventful', 'Group'])
    >>> # Will create 2x2 subplots all with the same data
    >>> plot = (ISOPlot(data=data)
    ...         .create_subplots(nrows=2, ncols=2)
    ...         .add_layer(ScatterLayer)
    ...         .style())
    >>> plot.show() # doctest: +SKIP
    >>> all(len(ctx.layers) == 1 for ctx in plot.subplot_contexts)
        True
    >>> plot.close()  # Clean up

    Add a layer to a specific subplot:

    >>> plot = (ISOPlot(data=data)
    ...         .create_subplots(nrows=2, ncols=2)
    ...         .add_layer(ScatterLayer, on_axis=0)
    ...         .style())
    >>> plot.show() # doctest: +SKIP
    >>> len(plot.subplot_contexts[0].layers) == 1
    True
    >>> all(len(ctx.layers) == 0 for ctx in plot.subplot_contexts[1:])
    True
    >>> plot.close()

    Add a layer to multiple subplots:

    >>> plot = (ISOPlot(data=data)
    ...            .create_subplots(nrows=2, ncols=2)
    ...            .add_layer(ScatterLayer, on_axis=[0, 2])
    ...            .style())
    >>> plot.show() # doctest: +SKIP
    >>> len(plot.subplot_contexts[0].layers) == 1
    True
    >>> len(plot.subplot_contexts[2].layers) == 1
    True
    >>> len(plot.subplot_contexts[1].layers) == 0
    True
    >>> plot.close()

    Add a layer with custom data to a specific subplot:
    >>> custom_data = pd.DataFrame({
    ...     'ISOPleasant': rng.normal(0.2, 0.1, 50),
    ...     'ISOEventful': rng.normal(0.15, 0.2, 50),
    ... })
    >>> plot = (ISOPlot(data=data)
    ...        .create_subplots(nrows=2, ncols=2)
    ...        .add_layer(ScatterLayer) # Add to all subplots
    ...        # Add a layer with custom data to the first subplot
    ...        .add_layer(ScatterLayer, data=data.iloc[:50], on_axis=0, color='red')
    ...        # Add a layer with custom data to the second subplot
    ...        .add_layer(ScatterLayer, data=custom_data, on_axis=1)
    ...        .style())
    >>> plot.show() # doctest: +SKIP
    >>> plot.close()

    """
    # TODO(MitchellAcoustics): Need to handle legend/label creation
    #                          for new data added to a specific subplot
    # Create the layer instance
    layer = layer_class(custom_data=data, **params)

    # Check if we have axes to render on
    self._check_for_axes()

    # If no subplots created yet, add to main context
    if not self.subplot_contexts:
        if self.main_context.ax is None:
            # Get the single axis and assign it to main context
            if isinstance(self.axes, Axes):
                self.main_context.ax = self.axes
            elif isinstance(self.axes, np.ndarray) and self.axes.size > 0:
                self.main_context.ax = self.axes.flatten()[0]

        # Add layer to main context
        self.main_context.layers.append(layer)
        # Render the layer immediately
        layer.render(self.main_context)
        return self

    # Handle various axis targeting options
    target_contexts = self._resolve_target_contexts(on_axis)
    logger.debug(f"N target contexts: {len(target_contexts)}")

    # Add the layer to each target context and render it
    for i, context in enumerate(target_contexts):
        if data is not None and i >= self.subplots_params.n_subplots_by > 0:
            # If custom data is provided, use it for the specific subplot
            break
        context.layers.append(layer)
        layer.render(context)

    return self

def add_scatter(
    self,
    data: pd.DataFrame | None = None,
    *,
    on_axis: int | tuple[int, int] | list[int] | None = None,
    **params: Any,
) -> ISOPlot:
    """
    Add a scatter layer to specific subplot(s).

    Parameters
    ----------
    on_axis
        Target specific axis/axes
    data
        Custom data for this specific scatter plot
    **params
        Parameters for the scatter plot

    Returns
    -------
    :
        The current plot instance for chaining

    Examples
    --------
    Add a scatter layer to all subplots:

    >>> import pandas as pd
    >>> import numpy as np
    >>> rng = np.random.default_rng(42)
    >>> data = pd.DataFrame(
    ...    np.c_[rng.multivariate_normal([0.2, 0.15], [[0.1, 0], [0, 0.2]], 100),
    ...          rng.integers(1, 3, 100)],
    ...    columns=['ISOPleasant', 'ISOEventful', 'Group'])
    >>> plot = (ISOPlot(data=data)
    ...           .create_subplots(nrows=2, ncols=1)
    ...           .add_scatter(s=50, alpha=0.7, hue='Group')
    ...           .style())
    >>> plot.show() # doctest: +SKIP
    >>> all(len(ctx.layers) == 1 for ctx in plot.subplot_contexts)
    True
    >>> plot.close()  # Clean up

    Add a scatter layer with custom data to a specific subplot:

    >>> custom_data = pd.DataFrame({
    ...     'ISOPleasant': rng.normal(0.2, 0.1, 50),
    ...     'ISOEventful': rng.normal(0.15, 0.2, 50),
    ... })
    >>> plot = (ISOPlot(data=data)
    ...            .create_subplots(nrows=2, ncols=1)
    ...            .add_scatter(hue='Group')
    ...            .add_scatter(on_axis=0, data=custom_data, color='red')
    ...            .style())
    >>> plot.show() # doctest: +SKIP
    >>> plot.subplot_contexts[0].layers[1].custom_data is custom_data
    True
    >>> plot.close()  # Clean up

    """
    # Merge default scatter parameters with provided ones
    # Remove data from scatter_params to avoid conflict
    scatter_params = self._scatter_params.copy()
    scatter_params.drop("data")
    scatter_params.update(**params)

    return self.add_layer(
        ScatterLayer,
        data=data,
        on_axis=on_axis,
        **scatter_params.as_dict(drop=["data"]),
    )