plot_utils

def render_histogram_png(
    counts: np.ndarray,
    edges: np.ndarray,
    bar_color: str,
    *,
    figsize: tuple[float, float] = (3.0, 2.0),
    dpi: int = 100,
    xlabel: str = "Score",
    title: str | None = None,
    log_y: bool = False,
) -> bytes:
    """Render a histogram as PNG bytes.

    Used for the linear-x score histogram; the source-size distribution has its
    own overlay renderer (:func:`render_size_stratification_png`).

    Args:
        counts: Bin counts from a histogram.
        edges: Bin edges (length ``len(counts) + 1``).
        bar_color: CSS colour string for the bars.
        figsize: Figure size in inches.
        dpi: Output resolution.
        xlabel: X-axis label (defaults to ``"Score"`` for the score histogram).
        title: Optional axes title.
        log_y: Use a logarithmic y-axis — the score distribution is sharply
            peaked, so a linear count axis buries the rare high-score tail (where
            the anomalies live) in a flat line at the bottom.

    Returns:
        PNG image as bytes.
    """
    matplotlib.use("Agg")
    fig = Figure(figsize=figsize, facecolor="black")
    ax = fig.add_subplot(111)
    ax.set_facecolor("black")
    ax.bar(
        edges[:-1],
        counts,
        width=np.diff(edges),
        color=bar_color,
        edgecolor="none",
        align="edge",
    )
    _apply_range_xticks(ax, edges, log_x=False)
    if log_y:
        ax.set_yscale("log")
        # A count of 1 must sit above the axis floor to be visible; 0.5 keeps the
        # single-count bars legible without leaving a large empty gap below them.
        ax.set_ylim(bottom=0.5)
    _style_dark_axes(ax, xlabel=xlabel, ylabel="Count", title=title)
    return _figure_to_png(fig, dpi)

def render_size_stratification_png(
    edges: np.ndarray,
    population_counts: np.ndarray,
    sampled_counts: np.ndarray,
    *,
    population_color: str,
    sampled_color: str,
    unit: str = "pixel",
    stratified: bool = False,
    figsize: tuple[float, float] = (3.0, 2.0),
    dpi: int = 100,
) -> bytes:
    """Overlay the sampled unlabeled pool's sizes on the tile-population sizes.

    The filled bars are the candidate population (every source in the drawn
    tiles); the step line is the sampled unlabeled pool.  Both are normalised to
    a per-bin fraction so their shapes are comparable despite the population
    dwarfing the sample in absolute count — with stratification on, the sampled
    line should be visibly flatter (~uniform in log size) than the small-skewed
    population, which is the visual confirmation that stratification worked.

    Args:
        edges: Shared log-spaced bin edges (length ``len(counts) + 1``).
        population_counts: Per-bin counts of the tile population.
        sampled_counts: Per-bin counts of the sampled unlabeled pool.
        population_color: CSS colour for the population bars.
        sampled_color: CSS colour for the sampled step line.
        unit: Diameter unit for the x-axis label (``"pixel"`` / ``"arcsec"``).
        stratified: Whether size stratification was applied (annotates the title).
        figsize: Figure size in inches.
        dpi: Output resolution.

    Returns:
        PNG image as bytes.
    """
    matplotlib.use("Agg")
    population_counts = np.asarray(population_counts, dtype=float)
    sampled_counts = np.asarray(sampled_counts, dtype=float)

    def _fraction(counts: np.ndarray) -> np.ndarray:
        total = counts.sum()
        return counts / total if total > 0 else counts

    population_fraction = _fraction(population_counts)
    sampled_fraction = _fraction(sampled_counts)

    fig = Figure(figsize=figsize, facecolor="black")
    ax = fig.add_subplot(111)
    ax.set_facecolor("black")
    ax.bar(
        edges[:-1],
        population_fraction,
        width=np.diff(edges),
        color=population_color,
        edgecolor="none",
        align="edge",
        label="In tiles",
    )
    # A step outline (aligned to the same bin edges) reads clearly on top of the
    # filled population bars without a second, occluding set of bars.
    ax.stairs(sampled_fraction, edges, color=sampled_color, linewidth=1.8, label="Sampled")
    _apply_range_xticks(ax, edges, log_x=True)
    title = "Sampled vs tiles (stratified)" if stratified else "Sampled vs tiles"
    _style_dark_axes(ax, xlabel=f"Source diameter ({unit})", ylabel="Fraction", title=title)
    ax.legend(
        facecolor="black",
        edgecolor="#333",
        labelcolor="white",
        fontsize=7,
        loc="upper right",
    )
    return _figure_to_png(fig, dpi)