test_plot.py

from unittest import mock
from unittest.mock import patch

import numpy as np
import pandas as pd
import pytest
from cycler import V
from pandas.core.arrays.arrow.accessors import pa

import ultraplot as uplt
from ultraplot.internals.warnings import UltraPlotWarning


@pytest.mark.mpl_image_compare
def test_seaborn_lineplot_legend_hue_only():
    """
    Regression test: seaborn lineplot on UltraPlot axes should not add spurious
    legend entries like 'y'/'ymin'. Only hue categories should appear unless the user
    explicitly labels helper bands.
    """
    import seaborn as sns

    fig, ax = uplt.subplots()
    df = pd.DataFrame(
        {
            "xcol": np.concatenate([np.arange(10)] * 2),
            "ycol": np.concatenate([np.arange(10), 1.5 * np.arange(10)]),
            "hcol": ["h1"] * 10 + ["h2"] * 10,
        }
    )

    with ax.external():
        sns.lineplot(data=df, x="xcol", y="ycol", hue="hcol", ax=ax)

    # Create (or refresh) legend and collect labels
    leg = ax.legend()
    labels = {t.get_text() for t in leg.get_texts()}

    # Should contain only hue levels; must not contain inferred 'y' or CI helpers
    assert "y" not in labels
    assert "ymin" not in labels
    assert {"h1", "h2"}.issubset(labels)
    return fig


"""
This file is used to test base properties of ultraplot.axes.plot. For higher order plotting related functions, please use 1d and 2plots
"""


def test_external_preserves_explicit_label():
    """
    In external mode, explicit labels must still be respected even when autolabels are disabled.
    """
    fig, ax = uplt.subplots()
    ax.set_external(True)
    (h,) = ax.plot([0, 1, 2], [0, 1, 0], label="X")
    leg = ax.legend(h, loc="best")
    labels = [t.get_text() for t in leg.get_texts()]
    assert "X" in labels


def test_external_disables_autolabels_no_label():
    """
    In external mode, if no labels are provided, autolabels are disabled and a placeholder is used.
    """
    fig, ax = uplt.subplots()
    ax.set_external(True)
    (h,) = ax.plot([0, 1, 2], [0, 1, 0])
    # Explicitly pass the handle so we test the label on that artist
    leg = ax.legend(h, loc="best")
    labels = [t.get_text() for t in leg.get_texts()]
    # With no explicit labels and autolabels disabled, a placeholder is used
    assert (not labels) or (labels[0] in ("_no_label", ""))


def test_parse_level_lim_accepts_list_input():
    """
    Ensure list inputs are converted before checking ndim in _parse_level_lim.
    """
    fig, ax = uplt.subplots()
    vmin, vmax, _ = ax[0]._parse_level_lim([[1, 2], [3, 4]])
    assert vmin == 1
    assert vmax == 4
    uplt.close(fig)


def test_error_shading_explicit_label_external():
    """
    Explicit label on fill_between should be preserved in legend entries.
    """
    fig, ax = uplt.subplots()
    ax.set_external(True)
    x = np.linspace(0, 2 * np.pi, 50)
    y = np.sin(x)
    patch = ax.fill_between(x, y - 0.5, y + 0.5, alpha=0.3, label="Band")
    leg = ax.legend([patch], loc="best")
    labels = [t.get_text() for t in leg.get_texts()]
    assert "Band" in labels


def test_graph_nodes_kw():
    """Test the graph method by setting keywords for nodes"""
    import networkx as nx

    g = nx.path_graph(5)
    labels_in = {node: node for node in range(2)}

    fig, ax = uplt.subplots()
    nodes, edges, labels = ax.graph(g, nodes=[0, 1], labels=labels_in)

    # Expecting 2 nodes 1 edge
    assert len(edges.get_offsets()) == 1
    assert len(nodes.get_offsets()) == 2
    assert len(labels) == len(labels_in)


def test_graph_edges_kw():
    """Test the graph method by setting keywords for nodes"""
    import networkx as nx

    g = nx.path_graph(5)
    edges_in = [(0, 1)]

    fig, ax = uplt.subplots()
    nodes, edges, labels = ax.graph(g, edges=edges_in)

    # Expecting 2 nodes 1 edge
    assert len(edges.get_offsets()) == 1
    assert len(nodes.get_offsets()) == 5
    assert labels == False


def test_graph_input():
    """
    Test graph input methods. We allow for graphs, adjacency matrices, and edgelists.
    """
    import networkx as nx

    g = nx.path_graph(5)
    A = nx.to_numpy_array(g)
    el = np.array(g.edges())
    fig, ax = uplt.subplots()
    # Test input methods
    ax.graph(g)  # Graphs
    ax.graph(A)  # Adjcency matrices
    ax.graph(el)  # edgelists
    with pytest.raises(TypeError):
        ax.graph("invalid_input")


def test_graph_on_3d_projection():
    """
    Ensure graph plotting is available on 3D axes.
    """
    import networkx as nx

    g = nx.path_graph(5)
    _, axs = uplt.subplots(proj="3d")
    ax = axs[0]
    nodes, edges, labels = ax.graph(g)
    assert callable(getattr(ax, "graph", None))
    assert nodes is not False
    assert edges is not False
    assert labels is False


def test_graph_layout_input():
    """
    Test if layout is in a [0, 1] x [0, 1] box
    """
    import networkx as nx

    g = nx.path_graph(5)
    circular = nx.circular_layout(g)
    layouts = [None, nx.spring_layout, circular, "forceatlas2", "spring_layout"]
    fig, ax = uplt.subplots(ncols=len(layouts))
    for axi, layout in zip(ax[1:], layouts):
        axi.graph(g, layout=layout)


def test_graph_rescale():
    """
    Graphs can be normalized such that the node size is the same independnt of the fig size
    """
    import networkx as nx

    g = nx.path_graph(5)
    layout = nx.spring_layout(g)
    # shift layout outside the box
    layout = {node: np.array(pos) + np.array([10, 10]) for node, pos in layout.items()}

    fig, ax = uplt.subplots()
    nodes1 = ax.graph(g, layout=layout, rescale=True)[0]

    xlim_scaled = np.array(ax.get_xlim())
    ylim_scaled = np.array(ax.get_ylim())

    fig, ax = uplt.subplots()
    nodes2 = ax.graph(g, layout=layout, rescale=False)[0]

    for x, y in nodes1.get_offsets():
        assert x >= 0 and x <= 1
        assert y >= 0 and y <= 1
    for x, y in nodes2.get_offsets():
        assert x > 1
        assert y > 1


def test_violin_labels():
    """
    Test the labels functionality of violinplot and violinploth.
    """
    labels = "hello world !".split()
    fig, ax = uplt.subplots()
    bodies = ax.violinplot(y=[1, 2, 3], labels=labels)
    for label, body in zip(labels, bodies):
        assert body.get_label() == label

    # # Also test the horizontal ticks
    bodies = ax.violinploth(x=[1, 2, 3], labels=labels)
    ytick_labels = ax.get_yticklabels()
    for label, body in zip(labels, bodies):
        assert body.get_label() == label

    # Labels are padded if they are shorter than the data
    shorter_labels = [labels[0]]
    with pytest.warns(UltraPlotWarning):
        bodies = ax.violinplot(y=[[1, 2, 3], [2, 3, 4]], labels=shorter_labels)
        assert len(bodies) == 3
        assert bodies[0].get_label() == shorter_labels[0]


@pytest.mark.parametrize(
    "mpl_version, expected_key, expected_value",
    [
        ("3.10.0", "orientation", "vertical"),
        ("3.9.0", "vert", True),
    ],
)
def test_violinplot_mpl_versions(
    mpl_version: str,
    expected_key: str,
    expected_value: bool | str,
):
    """
    Test specific logic for violinplot to ensure that past and current versions work as expected.
    """
    fig, ax = uplt.subplots()
    with mock.patch("ultraplot.axes.plot._version_mpl", new=mpl_version):
        with mock.patch.object(ax.axes, "_call_native") as mock_call:
            # Note: implicit testing of labels passing. It should work
            ax.violinplot(y=[1, 2, 3], vert=True)

            mock_call.assert_called_once()
            _, kwargs = mock_call.call_args
            assert kwargs[expected_key] == expected_value
            if expected_key == "orientation":
                assert "vert" not in kwargs
            else:
                assert "orientation" not in kwargs


def test_violinplot_hatches():
    """
    Test the input on the hatches parameter. Either a singular or a list of strings. When a list is provided, it must be of the same length as the number of violins.
    """
    # should be ok
    fig, ax = uplt.subplots()
    ax.violinplot(y=[1, 2, 3], vert=True, hatch="x")

    with pytest.raises(ValueError):
        ax.violinplot(y=[1, 2, 3], vert=True, hatches=["x", "o"])


@pytest.mark.parametrize(
    "mpl_version, expected_key, expected_value",
    [
        ("3.10.0", "orientation", "vertical"),
        ("3.9.0", "vert", True),
    ],
)
def test_boxplot_mpl_versions(
    mpl_version: str,
    expected_key: str,
    expected_value: bool | str,
):
    """
    Test specific logic for violinplot to ensure that past and current versions work as expected.
    """
    fig, ax = uplt.subplots()
    with mock.patch("ultraplot.axes.plot._version_mpl", new=mpl_version):
        with mock.patch.object(ax.axes, "_call_native") as mock_call:
            # Note: implicit testing of labels passing. It should work
            ax.boxplot(y=[1, 2, 3], vert=True)

            mock_call.assert_called_once()
            _, kwargs = mock_call.call_args
            assert kwargs[expected_key] == expected_value
            if expected_key == "orientation":
                assert "vert" not in kwargs
            else:
                assert "orientation" not in kwargs


def test_quiver_discrete_colors(rng):
    """
    Edge case where colors are discrete for quiver plots
    """
    X = np.array([0, 1, 2])
    Y = np.array([0, 1, 2])

    U = np.array([1, 1, 0])
    V = np.array([0, 1, 1])

    colors = ["r", "g", "b"]

    fig, ax = uplt.subplots()
    q = ax.quiver(X, Y, U, V, color=colors, infer_rgb=True)
    expectations = [uplt.colors.mcolors.to_rgba(color) for color in colors]
    facecolors = q.get_facecolors()
    for expectation, facecolor in zip(expectations, facecolors):
        assert np.allclose(
            facecolor, expectation, 0.1
        ), f"Expected {expectation} but got {facecolor}"
    C = ["#ff0000", "#00ff00", "#0000ff"]
    ax.quiver(X - 1, Y, U, V, color=C, infer_rgb=True)

    # pass rgba values
    C = rng.random((3, 4))
    ax.quiver(X - 2, Y, U, V, C)
    ax.quiver(X - 3, Y, U, V, color="red", infer_rgb=True)
    uplt.close(fig)


def test_setting_log_with_rc():
    """
    Test setting log scale with rc context manager
    """
    import re

    x, y = np.linspace(0, 1e6, 10), np.linspace(0, 1e6, 10)

    def check_ticks(axis, target=True):
        pattern = r"\$\\mathdefault\{10\^\{(\d+)\}\}\$"
        for tick in axis.get_ticklabels():
            match = re.match(pattern, tick.get_text())
            expectation = False
            if match:
                expectation = True
            assert expectation == target

    def reset(ax):
        ax.set_xscale("linear")
        ax.set_yscale("linear")

    funcs = [
        "semilogx",
        "semilogy",
        "loglog",
    ]
    conditions = [
        ["x"],
        ["y"],
        ["x", "y"],
    ]

    with uplt.rc.context({"formatter.log": True}):
        fig, ax = uplt.subplots()
        for func, targets in zip(funcs, conditions):
            reset(ax)
            # Call the function
            getattr(ax, func)(x, y)
            # Check if the formatter is set
            for target in targets:
                axi = getattr(ax, f"{target}axis")
                check_ticks(axi, target=True)

    with uplt.rc.context({"formatter.log": False}):
        fig, ax = uplt.subplots()
        for func, targets in zip(funcs, conditions):
            reset(ax)
            getattr(ax, func)(x, y)
            for target in targets:
                axi = getattr(ax, f"{target}axis")
                check_ticks(axi, target=False)

    uplt.close(fig)


def test_format_log_scale_preserves_log_formatter():
    """
    Test that setting a log scale preserves the log formatter when enabled.
    """
    x = np.linspace(1, 1e6, 10)
    log_formatter = uplt.constructor.Formatter("log")
    log_formatter_type = type(log_formatter)

    with uplt.rc.context({"formatter.log": True}):
        fig, ax = uplt.subplots()
        ax.plot(x, x)
        ax.format(yscale="log")
        assert isinstance(ax.yaxis.get_major_formatter(), log_formatter_type)
        ax.set_yscale("log")
        assert isinstance(ax.yaxis.get_major_formatter(), log_formatter_type)

    with uplt.rc.context({"formatter.log": False}):
        fig, ax = uplt.subplots()
        ax.plot(x, x)
        ax.format(yscale="log")
        assert not isinstance(ax.yaxis.get_major_formatter(), log_formatter_type)
        ax.set_yscale("log")
        assert not isinstance(ax.yaxis.get_major_formatter(), log_formatter_type)

    uplt.close(fig)


def test_shading_pcolor(rng):
    """
    Pcolormesh by default adjusts the plot by
    getting the edges of the data for x and y.
    This creates a conflict when shading is used
    such as nearest and Gouraud.
    """
    nx, ny = 5, 7
    x = np.linspace(0, 5, nx)
    y = np.linspace(0, 4, ny)
    X, Y = np.meshgrid(x, y)
    Z = rng.random((nx, ny)).T
    fig, ax = uplt.subplots()

    results = []

    # Custom wrapper to capture return values
    def wrapped_parse_2d_args(x, y, z, *args, **kwargs):
        out = original_parse_2d_args(x, y, z, *args, **kwargs)
        results.append(out[:3])  # Capture x, y, z only
        return out

    # Save original method
    original_parse_2d_args = ax[0]._parse_2d_args
    shadings = ["flat", "nearest", "gouraud"]

    with patch.object(ax[0], "_parse_2d_args", side_effect=wrapped_parse_2d_args):
        for shading in shadings:
            ax.pcolormesh(X, Y, Z, shading=shading)

    # Now check results
    for i, (shading, (x, y, z)) in enumerate(zip(shadings, results)):
        assert x.shape[0] == y.shape[0]
        assert x.shape[1] == y.shape[1]
        if shading == "flat":
            assert x.shape[0] == z.shape[0] + 1
            assert x.shape[1] == z.shape[1] + 1
        else:
            assert x.shape[0] == z.shape[0]
            assert x.shape[1] == z.shape[1]
    uplt.close(fig)


def test_cycle_with_singular_column(rng):
    """
    While parsing singular columns the ncycle attribute should
    be ignored.
    """

    cycle = "qual1"
    # Create mock data that triggers the cycle
    # when plot directly but is is ignored when plot in
    # a loop
    data = rng.random((3, 6))

    fig, ax = uplt.subplots()
    active_cycle = ax[0]._active_cycle
    original_init = uplt.constructor.Cycle.__init__
    with mock.patch.object(
        uplt.constructor.Cycle,
        "__init__",
        wraps=uplt.constructor.Cycle.__init__,
        autospec=True,
        side_effect=original_init,
    ) as mocked:
        ax[0]._active_cycle = active_cycle  # reset the cycler
        ax.plot(data, cycle=cycle)
        assert mocked.call_args.kwargs["N"] == 6

        ax[0]._active_cycle = active_cycle  # reset the cycler
        for col in data.T:
            ax.plot(col, cycle=cycle)
            assert "N" not in mocked.call_args.kwargs
    uplt.close(fig)


def test_colorbar_center_levels(rng):
    """
    Allow centering of the colorbar ticks to the center
    """
    data = rng.random((10, 10)) * 2 - 1
    expectation = np.linspace(-1, 1, uplt.rc["cmap.levels"])
    fig, ax = uplt.subplots(ncols=2)
    for axi, center_levels in zip(ax, [False, True]):
        h = axi.pcolormesh(data, colorbar="r", center_levels=center_levels)
        cbar = axi._colorbar_dict[("right", "center")]
        if center_levels:
            deltas = cbar.get_ticks() - expectation
            assert np.all(np.allclose(deltas, 0))
            # For centered levels we are off by 1;
            # We have 1 more boundary bin than the expectation
            assert len(cbar.norm.boundaries) == expectation.size + 1
            w = np.diff(expectation)[0]
            # We check if the expectation is a center for the
            # the boundary
            assert expectation[0] - w * 0.5 == cbar.norm.boundaries[0]
        axi.set_title(f"{center_levels=}")
    uplt.close(fig)


def test_center_labels_colormesh_data_type(rng):
    """
    Test if how center_levels respond for discrete of continuous data
    """
    data = rng.random((10, 10)) * 2 - 1
    fig, ax = uplt.subplots(ncols=2)
    for axi, discrete in zip(ax, [True, False]):
        axi.pcolormesh(
            data,
            discrete=discrete,
            center_levels=True,
            colorbar="r",
        )

    uplt.close(fig)


def test_pie_labeled_series_in_dataframes():
    """
    Test an edge case where labeled indices cause
    labels to be grouped and passed on to pie chart
    which does not support this way of parsing.

    This only occurs when dataframes are passed.
    See https://github.com/Ultraplot/UltraPlot/issues/259
    """
    data = pd.DataFrame(
        index=pd.Index(list("abcd"), name="x"),
        data=dict(y=range(1, 5)),
    )
    fig, ax = uplt.subplots()
    wedges, texts = ax.pie("y", data=data)
    for text, index in zip(texts, data.index):
        assert text.get_text() == index
    uplt.close(fig)


def test_color_parsing_for_none():
    """
    Ensure that none is not parsed to white
    """
    fig, ax = uplt.subplots()
    ax.scatter(0.4, 0.5, 100, fc="none", ec="k", alpha=0.2)
    ax.scatter(0.5, 0.5, 100, fc="none", ec="k")
    ax.scatter(0.6, 0.5, 100, fc="none", ec="k", alpha=1)
    for artist in ax[0].collections:
        assert artist.get_facecolor().shape[0] == 0
    uplt.close(fig)


@pytest.mark.mpl_image_compare
def test_inhomogeneous_violin(rng):
    """
    Test that inhomogeneous violin plots work correctly.
    """
    fig, ax = uplt.subplots()
    data = [rng.normal(size=100), rng.normal(size=200)]
    violins = ax.violinplot(data, vert=True, labels=["A", "B"])
    assert len(violins) == 2
    for violin in violins:
        assert violin.get_paths()  # Ensure paths are created
    return fig


@pytest.mark.mpl_image_compare
def test_curved_quiver(rng):
    # Create a grid
    x = np.linspace(-4, 4, 20)
    y = np.linspace(-3, 3, 20)
    X, Y = np.meshgrid(x, y)

    # Define a rotational vector field (circular flow)
    U = -Y
    V = X
    speed = np.sqrt(U**2 + V**2)

    # Create a figure and axes
    fig, axs = uplt.subplots(ncols=3, sharey=True, figsize=(12, 4))

    # Left plot: matplotlib's streamplot
    axs[0].streamplot(X, Y, U, V, color=speed)
    axs[0].set_title("streamplot (native)")

    # Middle plot: quiver
    axs[1].quiver(X, Y, U, V, speed)
    axs[1].set_title("quiver")

    # Right plot: curved_quiver
    m = axs[2].curved_quiver(
        X, Y, U, V, color=speed, arrow_at_end=True, scale=2.0, grains=10
    )
    axs[2].set_title("curved_quiver")
    fig.colorbar(m.lines, ax=axs[:], label="speed")
    return fig


def test_validate_vector_shapes_pass():
    """
    Test that vector shapes match the grid shape using CurvedQuiverSolver.
    """
    from ultraplot.axes.plot_types.curved_quiver import _CurvedQuiverGrid

    x = np.linspace(0, 1, 3)
    y = np.linspace(0, 1, 3)
    grid = _CurvedQuiverGrid(x, y)
    u = np.ones(grid.shape)
    v = np.ones(grid.shape)
    assert u.shape == grid.shape
    assert v.shape == grid.shape


def test_validate_vector_shapes_fail():
    """
    Test that assertion fails when u and v do not match the grid shape using CurvedQuiverSolver.
    """
    from ultraplot.axes.plot_types.curved_quiver import (
        CurvedQuiverSolver,
        _CurvedQuiverGrid,
    )

    x = np.linspace(0, 1, 3)
    y = np.linspace(0, 1, 3)
    grid = _CurvedQuiverGrid(x, y)
    u = np.ones((2, 2))
    v = np.ones(grid.shape)
    with pytest.raises(AssertionError):
        assert u.shape == grid.shape


def test_normalize_magnitude():
    """
    Test that magnitude normalization returns a normalized array with max value 1.0 and correct shape.
    """
    u = np.array([[1, 2], [3, 4]])
    v = np.array([[4, 3], [2, 1]])
    mag = np.sqrt(u**2 + v**2)
    mag_norm = mag / np.max(mag)
    assert np.allclose(np.max(mag_norm), 1.0)
    assert mag_norm.shape == u.shape


def test_generate_start_points():
    """
    Test that CurvedQuiverSolver.gen_starting_points returns valid grid coordinates for seed points,
    and that grid.within_grid detects points outside the grid boundaries.
    """
    from ultraplot.axes.plot_types.curved_quiver import CurvedQuiverSolver

    x = np.linspace(0, 1, 5)
    y = np.linspace(0, 1, 5)
    grains = 5
    solver = CurvedQuiverSolver(x, y, density=5)
    sp2 = solver.gen_starting_points(x, y, grains)
    assert sp2.shape[1] == 2
    # Should detect if outside boundaries
    bad_points = np.array([[10, 10]])
    grid = solver.grid
    for pt in bad_points:
        assert not grid.within_grid(pt[0], pt[1])


def test_calculate_trajectories():
    """
    Test that CurvedQuiverSolver.get_integrator returns callable for each seed point
    and returns lists of trajectories and edges of correct length.
    """
    from ultraplot.axes.plot_types.curved_quiver import CurvedQuiverSolver

    x = np.linspace(0, 1, 5)
    y = np.linspace(0, 1, 5)
    u = np.ones((5, 5))
    v = np.ones((5, 5))
    mag = np.sqrt(u**2 + v**2)
    solver = CurvedQuiverSolver(x, y, density=5)
    integrator = solver.get_integrator(
        u, v, minlength=0.1, resolution=1.0, magnitude=mag
    )
    seeds = solver.gen_starting_points(x, y, grains=2)
    results = [integrator(pt[0], pt[1]) for pt in seeds]
    assert len(results) == seeds.shape[0]


@pytest.mark.mpl_image_compare
def test_curved_quiver_multicolor_lines():
    """
    Test that curved_quiver handles color arrays and returns a lines object.
    """
    x = np.linspace(0, 1, 5)
    y = np.linspace(0, 1, 5)
    X, Y = np.meshgrid(x, y)
    U = np.ones_like(X)
    V = np.ones_like(Y)
    speed = np.sqrt(U**2 + V**2)

    fig, ax = uplt.subplots()
    m = ax.curved_quiver(X, Y, U, V, color=speed)
    from matplotlib.collections import LineCollection

    assert isinstance(m.lines, LineCollection)
    assert m.lines.get_array().size > 0  # we have colors set
    assert m.lines.get_cmap() is not None
    return fig


@pytest.mark.mpl_image_compare
@pytest.mark.parametrize(
    "cmap",
    (
        "k",  # color
        "viridis",  # built-in
        "viko",  # bundled with ultraplot
    ),
)
def test_curved_quiver_color_and_cmap(rng, cmap):
    """
    Check that we can pass colors or colormaps
    """
    x = np.linspace(0, 1, 5)
    y = np.linspace(0, 1, 5)
    X, Y = np.meshgrid(x, y)
    U = np.ones_like(X)
    V = np.ones_like(Y)

    # Deal with color or cmap
    color = rng.random(X.shape)
    if cmap == "k":
        cmap = None
        color = "k"

    fig, ax = uplt.subplots()
    ax.curved_quiver(X, Y, U, V, color=color, cmap=cmap)
    return fig


@pytest.mark.mpl_image_compare
def test_sankey_basic():
    """
    Basic sanity check for Sankey diagrams.
    """
    fig, ax = uplt.subplots()
    diagram = ax.sankey(
        flows=[1.0, -0.6, -0.4],
        labels=["in", "out_a", "out_b"],
        orientations=[0, 1, -1],
        trunklength=1.1,
    )
    assert getattr(diagram, "patch", None) is not None
    assert getattr(diagram, "flows", None) is not None
    return fig


@pytest.mark.mpl_image_compare
def test_sankey_layered_nodes_flows():
    """
    Check that layered sankey accepts nodes and flows.
    """
    fig, ax = uplt.subplots()
    nodes = ["Budget", "Ops", "R&D", "Marketing"]
    flows = [
        ("Budget", "Ops", 5),
        ("Budget", "R&D", 3),
        ("Budget", "Marketing", 2),
    ]
    diagram = ax.sankey(nodes=nodes, flows=flows)
    assert len(diagram.nodes) == len(nodes)
    assert len(diagram.flows) == len(flows)
    return fig


@pytest.mark.mpl_image_compare
def test_sankey_layered_labels_and_style():
    """
    Check that style presets and label boxes are accepted.
    """
    fig, ax = uplt.subplots()
    nodes = ["Budget", "Ops", "R&D", "Marketing"]
    flows = [
        ("Budget", "Ops", 5),
        ("Budget", "R&D", 3),
        ("Budget", "Marketing", 2),
    ]
    diagram = ax.sankey(
        nodes=nodes,
        flows=flows,
        style="budget",
        flow_labels=True,
        value_format="{:.1f}",
        node_label_box=True,
    )
    flow_label_keys = [key for key in diagram.labels if isinstance(key, tuple)]
    assert flow_label_keys
    return fig


def test_sankey_invalid_flows():
    """Validate error handling for malformed flow inputs."""
    from ultraplot.axes.plot_types import sankey as sankey_mod

    with pytest.raises(ValueError):
        sankey_mod._normalize_flows(None)
    with pytest.raises(ValueError):
        sankey_mod._normalize_flows([("A", "B", -1)])
    with pytest.raises(ValueError):
        sankey_mod._normalize_flows([("A", "B", 0)])


def test_sankey_cycle_layers_error():
    """Cycles in the graph should raise a clear error."""
    from ultraplot.axes.plot_types import sankey as sankey_mod

    flows = [
        {"source": "A", "target": "B", "value": 1.0},
        {"source": "B", "target": "A", "value": 1.0},
    ]
    with pytest.raises(ValueError):
        sankey_mod._assign_layers(flows, ["A", "B"], None)


def test_sankey_flow_label_frac_alternates():
    """Label fractions should alternate around the midpoint."""
    from ultraplot.axes.plot_types import sankey as sankey_mod

    base = 0.5
    assert sankey_mod._flow_label_frac(0, 2, base) == 0.25
    assert sankey_mod._flow_label_frac(1, 2, base) == 0.75
    frac0 = sankey_mod._flow_label_frac(0, 3, base)
    frac1 = sankey_mod._flow_label_frac(1, 3, base)
    frac2 = sankey_mod._flow_label_frac(2, 3, base)
    assert 0.05 <= frac0 <= 0.95
    assert 0.05 <= frac1 <= 0.95
    assert 0.05 <= frac2 <= 0.95
    assert frac0 < base < frac1


def test_sankey_node_labels_outside_auto():
    """Auto outside labels should flip to the left/right on edge layers."""
    fig, ax = uplt.subplots()
    diagram = ax.sankey(
        nodes=["A", "B", "C"],
        flows=[("A", "B", 2.0), ("B", "C", 2.0)],
        node_labels=True,
        flow_labels=False,
    )
    label_a = diagram.labels["A"]
    label_c = diagram.labels["C"]
    node_a = diagram.nodes["A"]
    node_c = diagram.nodes["C"]
    ax_a, _ = label_a.get_position()
    ax_c, _ = label_c.get_position()
    assert ax_a < node_a.get_x()
    assert ax_c > node_c.get_x() + node_c.get_width()
    uplt.close(fig)


def test_sankey_flow_other_creates_other_node():
    """Small flows should be aggregated into an 'Other' node when requested."""
    fig, ax = uplt.subplots()
    diagram = ax.sankey(
        flows=[("A", "X", 0.2), ("A", "Y", 2.0)],
        flow_other=0.5,
        other_label="Other",
        node_labels=True,
    )
    assert "Other" in diagram.nodes
    assert "Other" in diagram.labels
    uplt.close(fig)


def test_sankey_unknown_style_error():
    """Unknown style presets should raise."""
    from ultraplot.axes.plot_types import sankey as sankey_mod

    with pytest.raises(ValueError):
        sankey_mod._apply_style(
            "nope",
            flow_cycle=["C0"],
            node_facecolor="0.7",
            flow_alpha=0.8,
            flow_curvature=0.5,
            node_label_box=False,
            node_label_kw={},
        )


def test_sankey_links_parameter_uses_layered():
    """Links should force layered sankey even with numeric flows input."""
    fig, ax = uplt.subplots()
    diagram = ax.sankey(
        flows=[1.0, -1.0],
        links=[("A", "B", 1.0)],
        node_labels=False,
        flow_labels=False,
    )
    assert "A" in diagram.nodes
    assert "B" in diagram.nodes
    assert diagram.layout["scale"] > 0
    uplt.close(fig)


def test_sankey_tuple_flows_use_layered():
    """Tuple flows without nodes should trigger layered sankey."""
    fig, ax = uplt.subplots()
    diagram = ax.sankey(flows=[("A", "B", 1.0)])
    assert "A" in diagram.nodes
    assert "B" in diagram.nodes
    uplt.close(fig)


def test_sankey_dict_flows_use_layered():
    """Dict flows should trigger layered sankey."""
    fig, ax = uplt.subplots()
    diagram = ax.sankey(flows=[{"source": "A", "target": "B", "value": 1.0}])
    assert "A" in diagram.nodes
    assert "B" in diagram.nodes
    assert "nodes" in diagram.layout
    uplt.close(fig)


def test_sankey_mixed_flow_formats_layered():
    """Mixed dict/tuple flows should still render in layered mode."""
    fig, ax = uplt.subplots()
    flows = [
        {"source": "A", "target": "B", "value": 1.0},
        ("B", "C", 2.0),
    ]
    diagram = ax.sankey(flows=flows)
    assert set(diagram.nodes.keys()) == {"A", "B", "C"}
    assert len(diagram.flows) == 2
    uplt.close(fig)


def test_sankey_numpy_flows_use_matplotlib():
    """1D numeric flows should use Matplotlib Sankey."""
    import numpy as np

    fig, ax = uplt.subplots()
    diagram = ax.sankey(flows=np.array([1.0, -1.0]))
    assert hasattr(diagram, "patch")
    assert not hasattr(diagram, "layout")
    uplt.close(fig)


def test_sankey_matplotlib_kwargs_passthrough():
    """Matplotlib sankey should pass patch kwargs through."""
    from matplotlib.colors import to_rgba

    fig, ax = uplt.subplots()
    diagram = ax.sankey(
        flows=[1.0, -1.0],
        orientations=[0, 0],
        facecolor="red",
        edgecolor="blue",
        linewidth=1.5,
    )
    assert np.allclose(diagram.patch.get_facecolor(), to_rgba("red"))
    assert np.allclose(diagram.patch.get_edgecolor(), to_rgba("blue"))
    assert diagram.patch.get_linewidth() == 1.5
    uplt.close(fig)


def test_sankey_matplotlib_connect_none():
    """Matplotlib sankey should allow connect=None."""
    fig, ax = uplt.subplots()
    diagram = ax.sankey(
        flows=[1.0, -1.0],
        orientations=[0, 0],
        connect=None,
    )
    assert hasattr(diagram, "patch")
    uplt.close(fig)


def test_sankey_normalize_nodes_dict_order_and_labels():
    """Node dict inputs should preserve order and resolve labels."""
    from ultraplot.axes.plot_types import sankey as sankey_mod

    nodes = {"A": {"label": "Alpha"}, "B": {"label": "Beta"}}
    flows = [{"source": "A", "target": "B", "value": 1.0}]
    node_map, order = sankey_mod._normalize_nodes(nodes, flows)
    assert order == ["A", "B"]
    assert node_map["A"]["label"] == "Alpha"
    assert node_map["B"]["label"] == "Beta"


def test_sankey_layer_order_missing_raises():
    """layer_order must include every layer."""
    from ultraplot.axes.plot_types import sankey as sankey_mod

    flows = [
        {"source": "A", "target": "B", "value": 1.0},
        {"source": "B", "target": "C", "value": 1.0},
    ]
    with pytest.raises(ValueError):
        sankey_mod._validate_layer_order([0], flows, ["A", "B", "C"], None)


def test_sankey_label_box_dict_copy():
    """Label box dicts should be copied so callers can reuse input."""
    from ultraplot.axes.plot_types import sankey as sankey_mod

    box = {"boxstyle": "round", "facecolor": "white"}
    resolved = sankey_mod._label_box(box)
    assert resolved == box
    assert resolved is not box


def test_sankey_label_box_default():
    """node_label_box=True should create a default box style."""
    from ultraplot.axes.plot_types import sankey as sankey_mod

    resolved = sankey_mod._label_box(True)
    assert resolved["boxstyle"].startswith("round")
    assert resolved["facecolor"] == "white"


def test_ribbon_smoke():
    """Smoke test for top-aligned ribbon flow diagrams."""
    import pandas as pd

    records = [
        ("E1", "P1", "T1"),
        ("E1", "P2", "T2"),
        ("E1", "P3", "T2"),
        ("E2", "P1", "T1"),
        ("E2", "P2", "T1"),
        ("E2", "P3", "T3"),
        ("E3", "P1", "T2"),
        ("E3", "P2", "T2"),
        ("E3", "P3", "T3"),
    ]
    data = pd.DataFrame(records, columns=["id", "period", "topic"])

    fig, ax = uplt.subplots()
    artists = ax.ribbon(
        data,
        id_col="id",
        period_col="period",
        topic_col="topic",
        period_order=["P1", "P2", "P3"],
        topic_order=["T1", "T2", "T3"],
        group_map={"T1": "G1", "T2": "G1", "T3": "G2"},
        group_order=["G1", "G2"],
    )
    assert artists["node_patches"]
    assert artists["flow_patches"]
    uplt.close(fig)


def test_sankey_assign_flow_colors_group_cycle():
    """Group cycle should be used for flow colors."""
    from ultraplot.axes.plot_types import sankey as sankey_mod

    flows = [
        {"source": "A", "target": "B", "value": 1.0, "group": "g1", "color": None},
        {"source": "A", "target": "C", "value": 1.0, "group": "g2", "color": None},
    ]
    color_map = sankey_mod._assign_flow_colors(
        flows, flow_cycle=None, group_cycle=["C0", "C1"]
    )
    assert color_map["g1"] == "C0"
    assert color_map["g2"] == "C1"
    assert flows[0]["color"] == "C0"
    assert flows[1]["color"] == "C1"


def test_sankey_assign_flow_colors_preserves_explicit():
    """Explicit flow colors should be preserved."""
    from ultraplot.axes.plot_types import sankey as sankey_mod

    flows = [
        {"source": "A", "target": "B", "value": 1.0, "group": "g1", "color": "red"}
    ]
    color_map = sankey_mod._assign_flow_colors(flows, flow_cycle=None, group_cycle=None)
    assert flows[0]["color"] == "red"
    assert color_map == {}


def test_sankey_node_dict_missing_id_raises():
    """Node dicts must include id or name."""
    from ultraplot.axes.plot_types import sankey as sankey_mod

    flows = [{"source": "A", "target": "B", "value": 1.0}]
    with pytest.raises(ValueError):
        sankey_mod._normalize_nodes([{"label": "missing"}], flows)


def test_sankey_node_order_missing_nodes_raises():
    """node_order must include all flow endpoints."""
    from ultraplot.axes.plot_types import sankey as sankey_mod

    flows = [{"source": "A", "target": "B", "value": 1.0}]
    with pytest.raises(ValueError):
        sankey_mod._ensure_nodes(["A"], flows, node_order=["A"])


def test_sankey_flow_other_multiple_sources():
    """flow_other should aggregate per source."""
    from ultraplot.axes.plot_types import sankey as sankey_mod

    flows = [
        {"source": "A", "target": "X", "value": 0.2, "label": None, "color": None},
        {"source": "A", "target": "Y", "value": 0.1, "label": None, "color": None},
        {"source": "B", "target": "Z", "value": 0.3, "label": None, "color": None},
        {"source": "B", "target": "W", "value": 2.0, "label": None, "color": None},
    ]
    result = sankey_mod._apply_flow_other(flows, 0.5, "Other")
    others = [flow for flow in result if flow["target"] == "Other"]
    assert len(others) == 2
    sums = {flow["source"]: flow["value"] for flow in others}
    assert np.isclose(sums["A"], 0.3)
    assert np.isclose(sums["B"], 0.3)


def test_sankey_flow_label_text_callable():
    """Callable value_format should be used for flow labels."""
    from ultraplot.axes.plot_types import sankey as sankey_mod

    flow = {"value": 1.234, "label": None}
    text = sankey_mod._flow_label_text(flow, lambda v: f"{v:.1f}")
    assert text == "1.2"


def test_radar_chart_smoke():
    """Smoke test for pyCirclize radar chart wrapper."""
    try:
        from ultraplot.axes.plot_types.circlize import _import_pycirclize

        _import_pycirclize()
    except ImportError:
        pytest.skip("pycirclize is not available")

    import pandas as pd

    df = pd.DataFrame({"A": [1, 2], "B": [3, 4], "C": [2, 1]}, index=["set1", "set2"])
    fig, ax = uplt.subplots(proj="polar")
    circos = ax.radar_chart(df, vmin=0, vmax=4, fill=False, marker_size=3)
    assert hasattr(circos, "plotfig")
    uplt.close(fig)


def test_chord_diagram_smoke():
    """Smoke test for pyCirclize chord diagram wrapper."""
    try:
        from ultraplot.axes.plot_types.circlize import _import_pycirclize

        _import_pycirclize()
    except ImportError:
        pytest.skip("pycirclize is not available")

    import pandas as pd

    df = pd.DataFrame(
        [[5, 2, 1], [2, 6, 3], [1, 3, 4]],
        index=["A", "B", "C"],
        columns=["A", "B", "C"],
    )
    fig, axs = uplt.subplots(proj="polar")
    ax = axs[0]
    circos = ax.chord_diagram(df, ticks_interval=None)
    assert hasattr(circos, "plotfig")
    uplt.close(fig)


def test_phylogeny_smoke():
    """Smoke test for pyCirclize phylogeny wrapper."""
    try:
        from ultraplot.axes.plot_types.circlize import _import_pycirclize

        _import_pycirclize()
    except ImportError:
        pytest.skip("pycirclize is not available")

    fig, axs = uplt.subplots(proj="polar")
    ax = axs[0]
    circos, treeviz = ax.phylogeny("((A,B),C);", leaf_label_size=8)
    assert hasattr(circos, "plotfig")
    assert treeviz is not None
    uplt.close(fig)


def test_circos_bed_smoke(tmp_path):
    """Smoke test for BED-based circlize wrapper."""
    try:
        from ultraplot.axes.plot_types.circlize import _import_pycirclize

        _import_pycirclize()
    except ImportError:
        pytest.skip("pycirclize is not available")

    bed_path = tmp_path / "mini.bed"
    bed_path.write_text("chr1\t0\t100\nchr2\t0\t120\n", encoding="utf-8")

    fig, axs = uplt.subplots(proj="polar")
    ax = axs[0]
    circos = ax.circos_bed(bed_path, plot=True)
    assert len(circos.sectors) == 2
    uplt.close(fig)


def test_circos_builder_smoke():
    """Smoke test for general Circos wrapper."""
    try:
        from ultraplot.axes.plot_types.circlize import _import_pycirclize

        _import_pycirclize()
    except ImportError:
        pytest.skip("pycirclize is not available")

    fig, axs = uplt.subplots(proj="polar")
    ax = axs[0]
    circos = ax.circos({"A": 10, "B": 12}, plot=True)
    assert len(circos.sectors) == 2
    uplt.close(fig)


def test_circos_unshares_axes():
    """Circos wrappers should unshare axes if they were shared."""
    try:
        from ultraplot.axes.plot_types.circlize import _import_pycirclize

        _import_pycirclize()
    except ImportError:
        pytest.skip("pycirclize is not available")

    fig, axs = uplt.subplots(ncols=2, proj="polar", share="all")
    ax = axs[0]
    x_siblings = list(ax._shared_axes["x"].get_siblings(ax))
    y_siblings = list(ax._shared_axes["y"].get_siblings(ax))
    if len(x_siblings) == 1 and len(y_siblings) == 1:
        pytest.skip("polar axes are not shared in this configuration")
    ax.circos({"A": 10, "B": 12}, plot=False)
    x_siblings = list(ax._shared_axes["x"].get_siblings(ax))
    y_siblings = list(ax._shared_axes["y"].get_siblings(ax))
    assert len(x_siblings) == 1
    assert len(y_siblings) == 1
    uplt.close(fig)


def test_circos_delegation_subplots():
    """SubplotGrid should delegate circos calls for singleton grids."""
    try:
        from ultraplot.axes.plot_types.circlize import _import_pycirclize

        _import_pycirclize()
    except ImportError:
        pytest.skip("pycirclize is not available")

    fig, axs = uplt.subplots(proj="polar")
    circos = axs.circos({"A": 10, "B": 12}, plot=False)
    assert len(circos.sectors) == 2
    uplt.close(fig)


def test_histogram_norms():
    """
    Check that all histograms-like plotting functions
    use the sum of the weights.
    """
    rng = np.random.default_rng(seed=100)
    x, y = rng.normal(size=(2, 100))
    w = rng.uniform(size=100)

    fig, axs = uplt.subplots()
    _, _, bars = axs.hist(x, weights=w, bins=5)
    tot_weights = np.sum([bar.get_height() for bar in bars])
    np.testing.assert_allclose(tot_weights, np.sum(w))

    fig, axs = uplt.subplots()
    _, _, _, qm = axs.hist2d(x, y, weights=w, bins=5)
    tot_weights = np.sum(qm.get_array())
    np.testing.assert_allclose(tot_weights, np.sum(w))

    fig, axs = uplt.subplots()
    pc = axs.hexbin(x, y, weights=w, gridsize=5)
    tot_weights = np.sum(pc.get_array())
    np.testing.assert_allclose(tot_weights, np.sum(w))

    # check that a different reduce_C_function produces
    # a different result
    fig, axs = uplt.subplots()
    pc = axs.hexbin(x, y, weights=w, gridsize=5, reduce_C_function=np.max)
    tot_weights = np.sum(pc.get_array())
    # check they are not equal and that the different is not
    # due to floating point errors
    assert tot_weights != np.sum(w)
    assert not np.allclose(tot_weights, np.sum(w))