import matplotlib as mpl
import numpy as np
import pytest
from matplotlib.figure import Figure

import cmyt
from cmyt._utils import cmyt_cmaps, prefix_name

mpl_compare = pytest.mark.skip


@pytest.fixture()
def example_data():
    # generate example data
    prng = np.random.RandomState(0x4D3D3D3)
    noise = prng.random_sample((100, 100))
    x, y = np.mgrid[-50:50, -50:50]
    z = 5 * np.exp(-(x**2 + y**2) / 1000)
    return z + noise


@mpl_compare
@pytest.mark.parametrize("name", cmyt_cmaps)
def test_from_str(name, example_data):
    fig = Figure(figsize=(12, 4))
    axes = fig.subplots(ncols=2)
    pname = prefix_name(name)
    for cmap, ax in zip([pname, f"{pname}_r"], axes, strict=True):
        ax.set_title(cmap)
        im = ax.imshow(example_data, cmap=cmap)
        fig.colorbar(im, ax=ax)
    return fig


@mpl_compare
@pytest.mark.parametrize("name", cmyt_cmaps)
def test_from_obj(name, example_data):
    fig = Figure(figsize=(12, 4))
    axes = fig.subplots(ncols=2)
    for cmap, ax in zip([name, f"{name}_r"], axes, strict=True):
        ax.set_title(cmap)
        im = ax.imshow(example_data, cmap=getattr(cmyt, cmap))
        fig.colorbar(im, ax=ax)
    return fig


@pytest.mark.parametrize("name", cmyt_cmaps)
def test_cmap_name_attr(name):
    assert getattr(cmyt, name).name == name
    assert mpl.colormaps[f"cmyt.{name}"].name == prefix_name(name)