"""Public helper module.

Conceptually, this module groups the lightweight styling helpers reused across
the structured figure notebooks.

It exists as a separate unit so recurring axis, legend, guide-line, and
scale-bar conventions remain consistent without introducing a heavy plotting
framework.

It connects notebook plotting code to a small shared visual vocabulary for
figures and annotations.
"""

from __future__ import annotations

from typing import Any, Mapping, Optional, Sequence, Tuple


def add_time_event_lines(
    ax: Any,
    xs: Sequence[float] = (0.0, 1.0),
    *,
    color: str = "0.78",
    linewidth: float = 0.9,
    alpha: float = 1.0,
    zorder: float = 0.0,
) -> None:
    """Draw the shared task-event vertical guide lines on one axis."""
    for x in xs:
        ax.axvline(float(x), color=color, linewidth=linewidth, alpha=alpha, zorder=zorder)


def add_colored_text_legend(
    ax: Any,
    labels: Sequence[str],
    color_map: Mapping[str, Any],
    *,
    x: float = 0.12,
    y_top: float = 0.80,
    dy: float = 0.07,
    fontsize: float = 11.0,
    weight: str = "medium",
    ha: str = "left",
    va: str = "center",
    transform: Optional[Any] = None,
    bbox: bool = False,
    bbox_alpha: float = 0.86,
    bbox_facecolor: str = "white",
    bbox_edgecolor: str = "black",
    bbox_linewidth: float = 0.9,
    text_color_override: Optional[str] = None,
) -> None:
    """Draw a lightweight text legend with one colored label per line.

    Parameters
    ----------
    ax:
        Target matplotlib axis.
    labels:
        Labels shown from top to bottom.
    color_map:
        Mapping from label to matplotlib color.
    x, y_top, dy:
        Axes-coordinate placement controls.
    fontsize, weight, ha, va:
        Text styling controls.
    transform:
        Optional custom transform. Defaults to `ax.transAxes`.
    bbox, bbox_alpha:
        Optional background box behind the text labels.
    bbox_facecolor, bbox_edgecolor, bbox_linewidth:
        Styling controls for the optional legend box.
    text_color_override:
        Optional single text color applied to every label.
    """
    used_transform = ax.transAxes if transform is None else transform
    text_bbox = None
    if bbox:
        text_bbox = dict(
            facecolor=bbox_facecolor,
            edgecolor=bbox_edgecolor,
            linewidth=bbox_linewidth,
            alpha=bbox_alpha,
            pad=0.20,
        )

    for i, label in enumerate(labels):
        color = text_color_override if text_color_override is not None else color_map[label]
        ax.text(
            x,
            y_top - i * dy,
            label,
            transform=used_transform,
            color=color,
            fontsize=fontsize,
            weight=weight,
            ha=ha,
            va=va,
            bbox=text_bbox,
        )


def add_boxed_marker_legend(
    ax: Any,
    labels: Sequence[str],
    color_map: Mapping[str, Any],
    *,
    x: float = 0.10,
    y_top: float = 0.80,
    dy: float = 0.07,
    marker_x: float = 0.115,
    text_x: float = 0.145,
    marker: str = "o",
    marker_size: float = 6.5,
    fontsize: float = 10.8,
    weight: str = "medium",
    text_color: str = "0.12",
    box_pad_x: float = 0.028,
    box_pad_y: float = 0.030,
    box_width: float = 0.22,
    facecolor: str = "white",
    edgecolor: str = "0.78",
    linewidth: float = 0.8,
    alpha: float = 0.92,
    transform: Optional[Any] = None,
) -> None:
    """Draw a framed marker legend with one row per label."""
    from matplotlib.patches import FancyBboxPatch

    used_transform = ax.transAxes if transform is None else transform
    n = len(labels)
    if n == 0:
        return

    box_height = (n - 1) * dy + 2 * box_pad_y
    box_y = y_top - (n - 1) * dy - box_pad_y
    rect = FancyBboxPatch(
        (x - box_pad_x, box_y),
        box_width,
        box_height,
        boxstyle="round,pad=0.015,rounding_size=0.02",
        transform=used_transform,
        facecolor=facecolor,
        edgecolor=edgecolor,
        linewidth=linewidth,
        alpha=alpha,
        zorder=2.0,
    )
    ax.add_patch(rect)

    for i, label in enumerate(labels):
        y = y_top - i * dy
        ax.plot(
            [marker_x],
            [y],
            marker=marker,
            markersize=marker_size,
            markerfacecolor=color_map[label],
            markeredgecolor=color_map[label],
            linestyle="None",
            transform=used_transform,
            clip_on=False,
            zorder=3.0,
        )
        ax.text(
            text_x,
            y,
            label,
            transform=used_transform,
            color=text_color,
            fontsize=fontsize,
            weight=weight,
            ha="left",
            va="center",
            zorder=3.0,
        )


def style_bottom_time_axis(
    ax: Any,
    *,
    ticks: Sequence[float],
    labels: Optional[Sequence[str]] = None,
    xlabel: Optional[str] = None,
    fontsize: float = 10.0,
    tick_width: float = 0.9,
    tick_length: float = 3.0,
    spine_width: float = 1.0,
    color: str = "0.15",
) -> None:
    """Show a clean bottom time axis on otherwise minimalist plots."""
    ax.set_xticks(list(ticks))
    if labels is not None:
        ax.set_xticklabels(list(labels), fontsize=fontsize)
    ax.tick_params(axis="x", width=tick_width, length=tick_length, colors=color, labelsize=fontsize)
    ax.spines["bottom"].set_visible(True)
    ax.spines["bottom"].set_linewidth(spine_width)
    ax.spines["bottom"].set_color(color)
    if xlabel is not None:
        ax.set_xlabel(xlabel, fontsize=fontsize, color=color)


def style_left_value_axis(
    ax: Any,
    *,
    ticks: Sequence[float],
    labels: Optional[Sequence[str]] = None,
    ylabel: Optional[str] = None,
    fontsize: float = 10.0,
    tick_width: float = 0.9,
    tick_length: float = 3.0,
    spine_width: float = 1.0,
    color: str = "0.15",
) -> None:
    """Show a clean left y-axis on otherwise minimalist plots."""
    ax.set_yticks(list(ticks))
    if labels is not None:
        ax.set_yticklabels(list(labels), fontsize=fontsize)
    ax.tick_params(axis="y", width=tick_width, length=tick_length, colors=color, labelsize=fontsize)
    ax.spines["left"].set_visible(True)
    ax.spines["left"].set_linewidth(spine_width)
    ax.spines["left"].set_color(color)
    if ylabel is not None:
        ax.set_ylabel(ylabel, fontsize=fontsize, color=color)


def add_corner_scale_bar(
    ax: Any,
    *,
    x_size: float,
    y_size: float,
    x_label: str,
    y_label: str,
    x_pad_frac: float = 0.02,
    y_pad_frac: float = 0.03,
    x_label_offset_frac: float = 0.025,
    y_label_offset_frac: float = 0.03,
    color: str = "black",
    linewidth: float = 1.8,
    fontsize: float = 10.0,
) -> None:
    """Draw a simple bottom-left scale bar in data coordinates."""
    x0, x1 = ax.get_xlim()
    y0, y1 = ax.get_ylim()
    x_bar_start = x0 + x_pad_frac * (x1 - x0)
    y_bar_start = y0 + y_pad_frac * (y1 - y0)
    x_bar_end = x_bar_start + float(x_size)
    y_bar_end = y_bar_start + float(y_size)

    ax.plot([x_bar_start, x_bar_end], [y_bar_start, y_bar_start], color=color, linewidth=linewidth)
    ax.plot([x_bar_start, x_bar_start], [y_bar_start, y_bar_end], color=color, linewidth=linewidth)

    ax.text(
        (x_bar_start + x_bar_end) / 2.0,
        y_bar_start - x_label_offset_frac * (y1 - y0),
        x_label,
        ha="center",
        va="top",
        fontsize=fontsize,
        color=color,
    )
    ax.text(
        x_bar_start - y_label_offset_frac * (x1 - x0),
        (y_bar_start + y_bar_end) / 2.0,
        y_label,
        ha="center",
        va="center",
        rotation=90,
        fontsize=fontsize,
        color=color,
    )


def add_line_legend(
    container: Any,
    labels: Sequence[str],
    colors: Sequence[Any],
    *,
    linewidth: float = 2.0,
    **legend_kwargs: Any,
) -> Any:
    """Draw a standard line legend on an axis or figure."""
    from matplotlib.lines import Line2D

    handles = [Line2D([0], [0], color=color, linewidth=linewidth) for color in colors]
    return container.legend(handles, list(labels), **legend_kwargs)


def add_scale_bar(
    ax: Any,
    x_len: float,
    y_len: float,
    x_label: str,
    y_label: str,
    *,
    x0: Optional[float] = None,
    y0: Optional[float] = None,
    color: Any = "0.08",
    linewidth: float = 1.6,
    fontsize: float = 9.5,
) -> None:
    """Draw a compact horizontal and vertical scale bar in data coordinates.

    Parameters
    ----------
    ax:
        Target matplotlib axis.
    x_len, y_len:
        Scale-bar lengths expressed in axis data coordinates.
    x_label, y_label:
        Labels shown below and left of the scale bars.
    x0, y0:
        Optional starting coordinate of the scale bar. When omitted, the helper
        anchors the bar near the lower-left corner of the visible axis range.
    color, linewidth, fontsize:
        Shared styling controls for the scale-bar lines and labels.
    """
    xlim = ax.get_xlim()
    ylim = ax.get_ylim()
    if x0 is None:
        x0 = xlim[0] + 0.05 * (xlim[1] - xlim[0])
    if y0 is None:
        y0 = ylim[0] + 0.07 * (ylim[1] - ylim[0])

    ax.plot([x0, x0 + x_len], [y0, y0], color=color, linewidth=linewidth, clip_on=False, solid_capstyle="butt")
    ax.plot([x0, x0], [y0, y0 + y_len], color=color, linewidth=linewidth, clip_on=False, solid_capstyle="butt")
    ax.text(
        x0 + 0.5 * x_len,
        y0 - 0.06 * (ylim[1] - ylim[0]),
        x_label,
        ha="center",
        va="top",
        fontsize=fontsize,
        color=color,
    )
    ax.text(
        x0 - 0.02 * (xlim[1] - xlim[0]),
        y0 + 0.5 * y_len,
        y_label,
        ha="right",
        va="center",
        rotation=90,
        fontsize=fontsize,
        color=color,
    )


def raster_ticks(n_units: int) -> list[int]:
    """Return compact y ticks for dense single-trial raster panels."""
    top = int(max(n_units, 1))
    top_tick = max((top // 50) * 50, 50)
    if top_tick == 50:
        return [0, 25, 50]
    return [0, top_tick // 2, top_tick]


def style_plot_axis(
    ax: Any,
    *,
    facecolor: Optional[str] = None,
    grid_axis: Optional[str] = None,
    grid_color: str = "0.83",
    grid_linewidth: float = 0.7,
    grid_alpha: float = 0.7,
    spine_width: float = 1.0,
    tick_width: float = 1.2,
    tick_length: float = 4.0,
    tick_labelsize: Optional[float] = None,
    hide_top: bool = False,
    hide_right: bool = False,
) -> None:
    """Apply a lightweight reusable axis style for line and summary plots.

    Parameters
    ----------
    ax:
        Target matplotlib axis.
    facecolor:
        Optional axis face color.
    grid_axis:
        Optional axis argument forwarded to `ax.grid`.
    grid_color, grid_linewidth, grid_alpha:
        Grid styling controls.
    spine_width:
        Width applied to visible spines.
    tick_width, tick_length, tick_labelsize:
        Tick styling controls.
    hide_top, hide_right:
        Whether to hide the corresponding spines.
    """
    if facecolor is not None:
        ax.set_facecolor(facecolor)
    if grid_axis is not None:
        ax.grid(axis=grid_axis, color=grid_color, linewidth=grid_linewidth, alpha=grid_alpha)
    for side, spine in ax.spines.items():
        if (side == "top" and hide_top) or (side == "right" and hide_right):
            spine.set_visible(False)
        else:
            spine.set_linewidth(spine_width)
    tick_kwargs = dict(width=tick_width, length=tick_length)
    if tick_labelsize is not None:
        tick_kwargs["labelsize"] = tick_labelsize
    ax.tick_params(**tick_kwargs)