bar.py

"""Bar chart variants: simple, grouped, and stacked.

Three functions are provided:

* ``bar_plot``       — single-series bar chart with optional value labels.
* ``grouped_bar_plot`` — side-by-side grouped bars for multi-series comparison.
* ``stacked_bar_plot`` — stacked bars for part-to-whole composition.

Example
-------
>>> import numpy as np
>>> from academic_plot import bar_plot, savefig
>>> fig = bar_plot(
...     np.array([85, 92, 78]),
...     ["Method A", "Method B", "Method C"],
...     ylabel="Accuracy (%)",
... )
>>> savefig(fig, "bar_demo")
"""

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 bar_plot(
    values: np.ndarray,
    labels: Sequence[str],
    *,
    xlabel: str | None = None,
    ylabel: str = "Value",
    title: str | None = None,
    color: str | None = None,
    bar_width: float = 0.55,
    bar_edgecolor: str = "white",
    bar_edgewidth: float = 0.5,
    show_values: bool = True,
    value_fontsize: float = 6,
    value_fmt: str = ".2f",
    value_color: str = "#333",
    show_grid: bool = True,
    grid_alpha: float = GRID_ALPHA,
    grid_linewidth: float = GRID_LINEWIDTH,
    grid_color: str = GRID_COLOR,
    xtick_rotation: float = 0,
    figsize: tuple[float, float] = FIGSIZES["single"],
) -> plt.Figure:
    """Simple bar chart with optional value annotations above each bar.

    Parameters
    ----------
    values : np.ndarray
        Bar heights.
    labels : sequence of str
        Category labels on the x-axis.
    xlabel : str or None
        x-axis label.  Often omitted for categorical bar charts.
    ylabel : str
        y-axis label.
    title : str or None
        Subplot title.
    color : str or None
        Bar fill colour.  Falls back to ``PALETTE[0]``.
    bar_width : float
        Bar width in axes coordinates (0–1 range).
    bar_edgecolor : str
        Colour of bar borders.
    bar_edgewidth : float
        Width of bar borders.
    show_values : bool
        Annotate each bar with its numeric value.
    value_fontsize : float
        Font size for value annotations.
    value_fmt : str
        Format string for value annotations (e.g. ``".1f"``, ``".0f"``).
    value_color : str
        Text colour for value annotations.
    show_grid : bool
        Show horizontal grid lines.
    grid_alpha : float
        Grid opacity.
    grid_linewidth : float
        Grid line thickness.
    grid_color : str
        Grid line colour.
    xtick_rotation : float
        Rotation angle for x-tick labels (degrees).
    figsize : tuple[float, float]
        Figure size in inches.

    Returns
    -------
    plt.Figure
    """
    validate_arrays(np.asarray(values), names=["values"])
    fig, ax = plt.subplots(figsize=figsize)
    c = color or PALETTE[0]

    bars = ax.bar(
        labels, values, color=c, width=bar_width,
        edgecolor=bar_edgecolor, linewidth=bar_edgewidth, zorder=Z_ORDER["data"],
    )

    # Annotate each bar with its value
    if show_values:
        for bar in bars:
            h = bar.get_height()
            ax.text(
                bar.get_x() + bar.get_width() / 2,
                h + 0.02 * max(values),
                format(h, value_fmt),
                ha="center", va="bottom", fontsize=value_fontsize,
                color=value_color,
            )

    if xlabel:
        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, axis="y")

    ax.tick_params(axis="x", rotation=xtick_rotation)
    fig.tight_layout(pad=0.3)
    return fig


def grouped_bar_plot(
    data: Sequence[dict],
    labels: Sequence[str],
    *,
    xlabel: str | None = None,
    ylabel: str = "Value",
    title: str | None = None,
    bar_width_ratio: float = 0.7,
    bar_edgecolor: str = "white",
    bar_edgewidth: float = 0.5,
    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["wide_large"],
) -> plt.Figure:
    """Grouped (side-by-side) bar chart for comparing multiple series.

    Each dict in *data* may contain:

    ========== ============ =============================================
    Key        Required     Description
    ========== ============ =============================================
    ``values`` Yes          1-D array of bar heights, length matches *labels*.
    ``label``  No           Legend label (default ``"Series N"``).
    ``color``  No           Hex colour (auto-cycled if omitted).
    ========== ============ =============================================

    Parameters
    ----------
    data : sequence of dict
        Series definitions (see table above).
    labels : sequence of str
        Category labels on the x-axis (shared by all series).
    xlabel : str or None
        x-axis label.
    ylabel : str
        y-axis label.
    title : str or None
        Subplot title.
    bar_width_ratio : float
        Total width of one group relative to 1.0 axis unit.
    bar_edgecolor : str
        Colour of bar borders.
    bar_edgewidth : float
        Width of bar borders.
    show_grid : bool
        Show horizontal grid lines.
    grid_alpha, grid_linewidth, grid_color
        Grid styling.
    legend_loc : str
        Legend location.
    figsize : tuple[float, float]
        Figure size in inches.

    Returns
    -------
    plt.Figure
    """
    n_groups = len(labels)
    n_series = len(data)
    x = np.arange(n_groups)
    width = bar_width_ratio / n_series

    fig, ax = plt.subplots(figsize=figsize)
    for i, ds in enumerate(data):
        c = ds.get("color", PALETTE[i % len(PALETTE)])
        # Centre the group: offset each series symmetrically around x[k]
        offset = (i - (n_series - 1) / 2) * width
        ax.bar(
            x + offset, ds["values"], width, color=c,
            edgecolor=bar_edgecolor, linewidth=bar_edgewidth,
            label=ds.get("label", f"Series {i+1}"), zorder=Z_ORDER["data"],
        )

    ax.set_xticks(x)
    ax.set_xticklabels(labels)
    if xlabel:
        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, axis="y")
    ax.legend(loc=legend_loc)
    fig.tight_layout(pad=0.3)
    return fig


def stacked_bar_plot(
    data: Sequence[dict],
    labels: Sequence[str],
    *,
    xlabel: str | None = None,
    ylabel: str = "Value",
    title: str | None = None,
    bar_width: float = 0.55,
    bar_edgecolor: str = "white",
    bar_edgewidth: float = 0.5,
    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["wide"],
) -> plt.Figure:
    """Stacked bar chart for part-to-whole composition.

    Each dict in *data* may contain:

    ========== ============ =============================================
    Key        Required     Description
    ========== ============ =============================================
    ``values`` Yes          1-D array of segment heights.
    ``label``  No           Legend label.
    ``color``  No           Hex colour (auto-cycled if omitted).
    ========== ============ =============================================

    Parameters
    ----------
    data : sequence of dict
        Segment definitions (see table above).
    labels : sequence of str
        Category labels on the x-axis.
    xlabel : str or None
        x-axis label.
    ylabel : str
        y-axis label.
    title : str or None
        Subplot title.
    bar_width : float
        Width of each bar group.
    bar_edgecolor : str
        Colour of bar borders.
    bar_edgewidth : float
        Width of bar borders.
    show_grid : bool
        Show horizontal grid lines.
    grid_alpha, grid_linewidth, grid_color
        Grid styling.
    legend_loc : str
        Legend location.
    figsize : tuple[float, float]
        Figure size in inches.

    Returns
    -------
    plt.Figure
    """
    x = np.arange(len(labels))
    fig, ax = plt.subplots(figsize=figsize)

    # Accumulate bottom offset so each segment stacks on the previous one
    bottom = np.zeros(len(labels))
    for i, ds in enumerate(data):
        c = ds.get("color", PALETTE[i % len(PALETTE)])
        vals = np.array(ds["values"])
        ax.bar(
            x, vals, bar_width, bottom=bottom, color=c,
            edgecolor=bar_edgecolor, linewidth=bar_edgewidth,
            label=ds.get("label", f"Series {i+1}"), zorder=Z_ORDER["data"],
        )
        bottom += vals  # next segment starts here

    ax.set_xticks(x)
    ax.set_xticklabels(labels)
    if xlabel:
        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, axis="y")
    ax.legend(loc=legend_loc)
    fig.tight_layout(pad=0.3)
    return fig