fill_between.py

"""Fill-between / confidence-band plots.

Two functions:

* ``confidence_band_plot``       — single series with shaded uncertainty band.
* ``multi_confidence_band_plot`` — multiple series, each with its own band.

Example
-------
>>> import numpy as np
>>> x = np.linspace(0, 10, 80)
>>> y = np.sin(x)
>>> fig = confidence_band_plot(x, y, y - 1, y + 1, xlabel="t", ylabel="Signal")
"""

from __future__ import annotations

from typing import Sequence

import matplotlib.pyplot as plt
import numpy as np

from .style import PALETTE, FIGSIZES, GRID_ALPHA, GRID_LINEWIDTH, GRID_COLOR, Z_ORDER, savefig, apply_grid
from .utils import validate_arrays


def confidence_band_plot(
    x: np.ndarray,
    y: np.ndarray,
    y_low: np.ndarray,
    y_high: np.ndarray,
    *,
    xlabel: str = "x",
    ylabel: str = "y",
    title: str | None = None,
    label: str | None = None,
    band_label: str | None = None,
    color: str | None = None,
    line_linewidth: float = 1.3,
    fill_alpha: float = 0.18,
    show_grid: bool = True,
    grid_alpha: float = GRID_ALPHA,
    grid_linewidth: float = GRID_LINEWIDTH,
    grid_color: str = GRID_COLOR,
    legend_loc: str = "best",
    figsize: tuple[float, float] = FIGSIZES["single"],
) -> plt.Figure:
    """Line plot with a shaded confidence / uncertainty band.

    Parameters
    ----------
    x : np.ndarray
        1-D x coordinates.
    y : np.ndarray
        1-D central (mean) y values.
    y_low : np.ndarray
        Lower bound of the confidence band.
    y_high : np.ndarray
        Upper bound of the confidence band.
    xlabel, ylabel : str
        Axis labels.
    title : str or None
        Subplot title.
    label : str or None
        Legend label for the central line.  Defaults to ``"Mean"``.
    band_label : str or None
        Legend label for the shaded band.  Defaults to ``"95% CI"``.
    color : str or None
        Colour for both the line and the band.  Falls back to ``PALETTE[0]``.
    line_linewidth : float
        Thickness of the central line.
    fill_alpha : float
        Opacity of the shaded band.
    show_grid : bool
        Show background grid.
    grid_alpha, grid_linewidth, grid_color
        Grid styling.
    legend_loc : str
        Legend location.
    figsize : tuple[float, float]
        Figure size in inches.

    Returns
    -------
    plt.Figure
    """
    validate_arrays(x, y, y_low, y_high, names=["x", "y", "y_low", "y_high"])
    fig, ax = plt.subplots(figsize=figsize)
    c = color or PALETTE[0]

    # Shaded region (drawn first so the line sits on top)
    ax.fill_between(
        x, y_low, y_high, color=c, alpha=fill_alpha,
        label=band_label or "95% CI", zorder=Z_ORDER["fill"],
    )
    # Central line
    ax.plot(x, y, color=c, linewidth=line_linewidth,
            label=label or "Mean", zorder=Z_ORDER["data"])

    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    if title:
        ax.set_title(title)
    apply_grid(ax, show=show_grid, alpha=grid_alpha, linewidth=grid_linewidth, color=grid_color)
    ax.minorticks_on()
    ax.legend(loc=legend_loc, frameon=False)

    fig.tight_layout(pad=0.3)
    return fig


def multi_confidence_band_plot(
    datasets: Sequence[dict],
    *,
    xlabel: str = "x",
    ylabel: str = "y",
    title: str | None = None,
    fill_alpha: float = 0.15,
    line_linewidth: float = 1.2,
    show_grid: bool = True,
    grid_alpha: float = GRID_ALPHA,
    grid_linewidth: float = GRID_LINEWIDTH,
    grid_color: str = GRID_COLOR,
    legend_loc: str = "best",
    figsize: tuple[float, float] = FIGSIZES["multi_band"],
) -> plt.Figure:
    """Multiple series, each with its own confidence band.

    Each dict in *datasets* may contain:

    ========== ============ =============================================
    Key        Required     Description
    ========== ============ =============================================
    ``x``      Yes          1-D x coordinates.
    ``y``      Yes          1-D central values.
    ``y_low``  Yes          Lower bound of the confidence band.
    ``y_high`` Yes          Upper bound of the confidence band.
    ``label``  No           Legend label.
    ``color``  No           Hex colour (auto-cycled if omitted).
    ========== ============ =============================================

    Parameters
    ----------
    datasets : sequence of dict
        Series definitions (see table above).
    xlabel, ylabel : str
        Axis labels.
    title : str or None
        Subplot title.
    fill_alpha : float
        Opacity of each shaded band.
    line_linewidth : float
        Thickness of each central line.
    show_grid : bool
        Show background grid.
    grid_alpha, grid_linewidth, grid_color
        Grid styling.
    legend_loc : str
        Legend location.
    figsize : tuple[float, float]
        Figure size in inches.

    Returns
    -------
    plt.Figure
    """
    fig, ax = plt.subplots(figsize=figsize)

    for i, ds in enumerate(datasets):
        validate_arrays(ds["x"], ds["y"], ds["y_low"], ds["y_high"], names=[f"datasets[{i}]['x']", f"datasets[{i}]['y']", f"datasets[{i}]['y_low']", f"datasets[{i}]['y_high']"])
        c = ds.get("color", PALETTE[i % len(PALETTE)])
        lbl = ds.get("label", f"Series {i+1}")

        # Shaded band (behind the line)
        ax.fill_between(
            ds["x"], ds["y_low"], ds["y_high"],
            color=c, alpha=fill_alpha, zorder=Z_ORDER["fill"],
        )
        # Central line
        ax.plot(
            ds["x"], ds["y"], color=c, linewidth=line_linewidth,
            label=lbl, zorder=Z_ORDER["data"],
        )

    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    if title:
        ax.set_title(title)
    apply_grid(ax, show=show_grid, alpha=grid_alpha, linewidth=grid_linewidth, color=grid_color)
    ax.minorticks_on()
    ax.legend(loc=legend_loc, frameon=False)

    fig.tight_layout(pad=0.3)
    return fig