from __future__ import annotations

import io
import struct

import numpy as np
import pytest

import xy.pyplot as plt


@pytest.fixture(autouse=True)
def _reset_pyplot_state():
    plt.close("all")
    plt.style.use("default")
    yield
    plt.close("all")
    plt.style.use("default")


def test_nested_rc_context_restores_each_level_and_rcdefaults() -> None:
    original = plt.rcParams["lines.linewidth"]
    with plt.rc_context({"lines.linewidth": 2.0}):
        assert plt.rcParams["lines.linewidth"] == 2.0
        with plt.rc_context({"lines.linewidth": 7.0}):
            assert plt.rcParams["lines.linewidth"] == 7.0
        assert plt.rcParams["lines.linewidth"] == 2.0
    assert plt.rcParams["lines.linewidth"] == original
    plt.rcParams["lines.linewidth"] = 9.0
    plt.rcdefaults()
    assert plt.rcParams["lines.linewidth"] == 1.5


def test_style_use_supports_bounded_dicts_and_ordered_lists() -> None:
    plt.style.use({"lines.linewidth": 3.0})
    assert plt.rcParams["lines.linewidth"] == 3.0
    plt.style.use(["default", {"lines.linewidth": 4.0}])
    assert plt.rcParams["lines.linewidth"] == 4.0
    with pytest.raises(NotImplementedError, match=r"unknown\.style\.key"):
        plt.style.use({"unknown.style.key": 1})
    plt.style.use("ggplot")  # stock sheets apply their bounded rcParams subset
    assert plt.rcParams["axes.facecolor"] == "#E5E5E5"
    with pytest.raises(NotImplementedError, match="rcParams dict"):
        plt.style.use("Solarize_Light2")
    plt.rcdefaults()


def test_dark_background_gallery_uses_prop_cycle_length() -> None:
    # Exact source idiom from Matplotlib 3.11's
    # galleries/examples/style_sheets/dark_background.py.
    with plt.style.context("dark_background"):
        length = 6
        x = np.linspace(0, length)
        ncolors = len(plt.rcParams["axes.prop_cycle"])
        shift = np.linspace(0, length, ncolors, endpoint=False)

        _fig, ax = plt.subplots()
        for offset in shift:
            ax.plot(x, np.sin(x + offset), "o-")

        assert ncolors == 10
        assert ncolors == len(plt.rcParams["axes.prop_cycle"].by_key()["color"])
        assert len(ax.get_lines()) == ncolors
        assert ax._theme_tokens["plot_background"] == "black"
        assert ax._theme_tokens["axis_color"] == "white"


def test_figure_facecolor_rcparam_affects_new_figures() -> None:
    with plt.rc_context({"figure.facecolor": "#123456"}):
        assert plt.figure().get_facecolor() == "#123456"


def test_colormap_extremes_alpha_and_reversal_are_preserved() -> None:
    cmap = plt.colormaps["viridis"].with_extremes(
        bad=("red", 0.25), under="#00ff00", over=("blue", 0.75)
    )
    rgba = cmap(np.array([np.nan, -1.0, 0.5, 2.0]))
    np.testing.assert_allclose(rgba[0], [1.0, 0.0, 0.0, 0.25])
    np.testing.assert_allclose(rgba[1], [0.0, 1.0, 0.0, 1.0])
    np.testing.assert_allclose(rgba[3], [0.0, 0.0, 1.0, 0.75])
    forward = plt.colormaps["viridis"](np.array([0.0, 1.0]))
    reverse = plt.colormaps["viridis_r"](np.array([1.0, 0.0]))
    np.testing.assert_allclose(forward, reverse)


def test_savefig_file_objects_support_common_export_options() -> None:
    fig, ax = plt.subplots(figsize=(4, 3), dpi=80)
    ax.plot([0, 1], [0, 1])
    with pytest.raises(ValueError, match="requires format"):
        fig.savefig(io.BytesIO())
    output = io.BytesIO()
    fig.savefig(output, format="png")
    assert output.getvalue().startswith(b"\x89PNG\r\n\x1a\n")
    transparent = io.BytesIO()
    fig.savefig(transparent, format="png", transparent=True)
    transparent.seek(0)
    assert np.asarray(plt.imread(transparent))[0, 0, 3] == 0

    metadata = io.BytesIO()
    fig.savefig(metadata, format="png", metadata={"Author": "xy", "Title": "chart"})
    assert b"Author\x00xy" in metadata.getvalue()

    regular = io.BytesIO()
    tight = io.BytesIO()
    fig.savefig(regular, format="png")
    fig.savefig(tight, format="png", bbox_inches="tight", pad_inches=0)
    regular_size = struct.unpack(">II", regular.getvalue()[16:24])
    tight_size = struct.unpack(">II", tight.getvalue()[16:24])
    assert regular_size == (320, 240)
    # a strict shrink on both axes: equality would mean tight-crop is a no-op
    assert tight_size[0] < regular_size[0] and tight_size[1] < regular_size[1]

    colored = io.BytesIO()
    fig.savefig(colored, format="png", facecolor="#123456")
    colored.seek(0)
    np.testing.assert_array_equal(np.asarray(plt.imread(colored))[0, 0, :3], [18, 52, 86])
    for unsupported_format in ("jpeg", "webp", "pdf"):
        with pytest.raises(NotImplementedError, match=unsupported_format):
            fig.savefig(io.BytesIO(), format=unsupported_format)


def test_savefig_dpi_changes_output_dimensions_without_mutating_figure() -> None:
    fig, ax = plt.subplots(figsize=(4, 3), dpi=80)
    ax.plot([0, 1], [0, 1])

    def dimensions(dpi: int) -> tuple[int, int]:
        output = io.BytesIO()
        fig.savefig(output, format="png", dpi=dpi)
        return struct.unpack(">II", output.getvalue()[16:24])

    low = dimensions(60)
    high = dimensions(120)
    assert low == (240, 180)
    assert high == (480, 360)
    assert high[0] == 2 * low[0]
    assert high[1] == 2 * low[1]
    assert fig.get_dpi() == 80
