import os
import pytest
import spharpy as sp
from pyfar.testing.plot_utils import create_figure, save_and_compare
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import re

"""
For general information on testing plot functions see
https://pyfar-gallery.readthedocs.io/en/latest/contribute/contribution_guidelines.html#testing-plot-functions

Important:
- `create_baseline` and `compare_output` must be ``False`` when pushing
  changes to pyfar.
- `create_baseline` must only be ``True`` if the behavior of a plot function
  changed. In this case it is best practice to recreate only the baseline plots
  of the plot function (plot behavior) that changed.
"""
# global parameters -----------------------------------------------------------
create_baseline = False

# file type used for saving the plots
file_type = "png"

# if true, the plots will be compared to the baseline and an error is raised
# if there are any differences. In any case, differences are written to
# output_path as images
compare_output = False

# path handling
base_path = os.path.join('tests', 'test_plot_data')
baseline_path = os.path.join(base_path, 'baseline')
output_path = os.path.join(base_path, 'output')

if not os.path.isdir(base_path):
    os.mkdir(base_path)
if not os.path.isdir(baseline_path):
    os.mkdir(baseline_path)
if not os.path.isdir(output_path):
    os.mkdir(output_path)

# remove old output files
for file in os.listdir(output_path):
    os.remove(os.path.join(output_path, file))

# the naming scheme of the baseline is as follows:
# <function_name>_<parameter_name>_<parameters>.png

# testing ---------------------------------------------------------------------
@pytest.fixture(autouse=True)
def _close_all_figures():
    """
    Close all matplotlib figures after each test to prevent test
    pollution.
    """
    yield
    plt.close('all')


@pytest.mark.parametrize('function', [
    (sp.plot.scatter),
    (sp.plot.voronoi_cells_sphere)])
def test_sampling_scatter(function):
    """Test the scatter plot with default arguments."""
    coords = sp.samplings.equal_area(n_max=0, n_points=12)

    # do plotting
    filename = f'{function.__name__}_default'
    create_figure()
    function(coords)
    save_and_compare(
        create_baseline, baseline_path, output_path, filename,
        file_type, compare_output)

    match = 'must be a coordinates object.'
    with pytest.raises(ValueError, match=match):
        function('coords')


@pytest.mark.parametrize('function', [
    (sp.plot.balloon),
    (sp.plot.balloon_wireframe),
    (sp.plot.contour),
    (sp.plot.contour_map),
    (sp.plot.pcolor_map),
    (sp.plot.pcolor_sphere)])
def test_spherical_default(function):
    """Test all spherical plots with default arguments."""
    coords = sp.samplings.equal_area(n_max=0, n_points=500)
    data = np.sin(coords.colatitude) * np.cos(coords.azimuth)

    # do plotting
    filename = f'{function.__name__}_default'
    create_figure()
    function(coords, data)
    save_and_compare(
        create_baseline, baseline_path, output_path, filename,
        file_type, compare_output)

    # test error for invalid inputs
    match = 'data must be a 1D array with the same cshape as the coordinates'
    with pytest.raises(ValueError, match=match):
        function(coords, 'data')

    match = 'coordinates must be a coordinates object.'
    with pytest.raises(ValueError, match=match):
        function('coords', data)


@pytest.mark.parametrize('function', [
    (sp.plot.contour),
    (sp.plot.contour_map),
    (sp.plot.pcolor_map)])
@pytest.mark.parametrize('cmap', [plt.get_cmap('plasma'), 'plasma', None])
def test_spherical_cmap(function, cmap):
    """Test all spherical plots with custom cmap argument."""
    coords = sp.samplings.equal_area(n_max=0, n_points=500)
    data = np.sin(coords.colatitude) * np.cos(coords.azimuth)

    # do plotting
    if isinstance(cmap, mpl.colors.Colormap):
        # if a colormap object is passed, use its name for the filename
        cmap_str = 'ColormapObject'
    else:
        # otherwise use the string representation of the colormap
        cmap_str = str(cmap)
    filename = f'{function.__name__}_cmap_{cmap_str}'
    create_figure()
    function(coords, data, cmap=cmap)
    save_and_compare(create_baseline, baseline_path, output_path, filename,
                     file_type, compare_output)

    # test error for invalid inputs
    match = 'cmap'
    with pytest.raises(ValueError, match=match):
        function(coords, data, cmap=5)


@pytest.mark.parametrize('function', [
    (sp.plot.balloon),
    (sp.plot.balloon_wireframe),
    (sp.plot.contour),
    (sp.plot.contour_map),
    (sp.plot.pcolor_map),
    (sp.plot.pcolor_sphere)])
@pytest.mark.parametrize('colorbar', [True, False])
def test_spherical_colorbar(function, colorbar):
    """Test all spherical plots with custom colorbar argument."""
    coords = sp.samplings.equal_area(n_max=0, n_points=500)
    data = np.sin(coords.colatitude) * np.cos(coords.azimuth)

    out = function(coords, data, colorbar=colorbar)
    if colorbar:
        assert isinstance(out[2], mpl.colorbar.Colorbar)
        assert isinstance(out[0], list)
    if not colorbar:
        assert out[2] is None
        assert isinstance(out[0], plt.Axes)

    # do plotting
    filename = f'{function.__name__}_colorbar_{colorbar}'
    create_figure()
    function(coords, data, colorbar=colorbar)
    save_and_compare(create_baseline, baseline_path, output_path, filename,
                     file_type, compare_output)

    # test error for invalid inputs
    match = 'colorbar must be a boolean.'
    with pytest.raises(ValueError, match=match):
        function(coords, data, colorbar='colorbar')


@pytest.mark.parametrize('function', [
    (sp.plot.balloon),
    (sp.plot.balloon_wireframe),
    (sp.plot.contour),
    (sp.plot.contour_map),
    (sp.plot.pcolor_map),
    (sp.plot.pcolor_sphere)])
@pytest.mark.parametrize('limits', [None, (-1.5, 1.5), [-1.5, 1.5]])
def test_spherical_limits(function, limits):
    """Test all spherical plots with custom limits argument."""
    coords = sp.samplings.equal_area(n_max=0, n_points=500)
    data = np.sin(coords.colatitude) * np.cos(coords.azimuth)

    # do plotting
    if limits is None:
        limits_str = 'None'
    else:
        limits_str = f'{type(limits).__name__}'
    filename = f'{function.__name__}_limits_{limits_str}'
    create_figure()
    function(coords, data, limits=limits)
    save_and_compare(create_baseline, baseline_path, output_path, filename,
                     file_type, compare_output)

    # test error for invalid inputs
    match = (
        "limits must be a tuple or list containing the minimum and "
        "maximum values for the colormap or None.")
    with pytest.raises(ValueError, match=match):
        function(coords, data, limits='limits')
    with pytest.raises(ValueError, match=match):
        function(coords, data, limits=[0, 1, 2])


@pytest.mark.parametrize('function', [
    (sp.plot.balloon),
    (sp.plot.balloon_wireframe),
    (sp.plot.pcolor_sphere)])
@pytest.mark.parametrize('cmap_encoding', ['phase', 'magnitude'])
@pytest.mark.parametrize('cmap', [None, 'magma'])
def test_balloon_cmap_encoding(function, cmap_encoding, cmap):
    """Test all balloon like plots with custom arguments."""
    coords = sp.samplings.equal_area(n_max=0, n_points=500)
    data = np.sin(coords.colatitude) * np.cos(coords.azimuth)

    # use the azimuth angle for a continuous phase change
    data = np.abs(data) * np.exp(1j * coords.azimuth)

    # do plotting
    filename = f'{function.__name__}_cmap_encoding_{cmap_encoding}_cmap_{cmap}'
    create_figure()
    function(coords, data, cmap_encoding=cmap_encoding, cmap=cmap)
    save_and_compare(create_baseline, baseline_path, output_path, filename,
                     file_type, compare_output)

@pytest.mark.parametrize('function', [
    (sp.plot.balloon),
    (sp.plot.balloon_wireframe),
    (sp.plot.pcolor_sphere)])
def test_balloon_cmap_encoding_error(function):
    """Test raising errors for invalid cmap_encoding argument."""
    # test error for invalid inputs
    coords = sp.samplings.equal_area(n_max=0, n_points=500)
    data = np.sin(coords.colatitude) * np.cos(coords.azimuth)

    match = "cmap_encoding must be either 'phase' or 'magnitude'."
    with pytest.raises(ValueError, match=match):
        function(coords, data, cmap_encoding='invalid')


@pytest.mark.parametrize('function', [
    (sp.plot.pcolor_map),
    (sp.plot.contour_map),
    ])
@pytest.mark.parametrize('projection', ['mollweide', 'hammer'])
def test_spherical_projection(function, projection):
    """Test spherical plots with custom projection argument."""
    coords = sp.samplings.equal_area(n_max=0, n_points=500)
    data = np.sin(coords.colatitude) * np.cos(coords.azimuth)

    # do plotting
    filename = f'{function.__name__}_projection_{projection}'
    create_figure()
    function(coords, data, projection=projection)
    save_and_compare(create_baseline, baseline_path, output_path, filename,
                     file_type, compare_output)

    # test error for invalid inputs
    match = "The projection of the axis needs to be 'projection'"
    with pytest.raises(ValueError, match=match):
        function(coords, data, projection='projection')


@pytest.mark.parametrize('function', [
    sp.plot.contour_map,
    sp.plot.contour,
    ])
@pytest.mark.parametrize('levels', [
    8,
    (-1, -0.5, 0, 0.5, 1),
    [-1, -0.5, 0, 0.5, 1],
    np.array((-1, -0.5, 0, 0.5, 1)),
    ])
def test_spherical_levels(function, levels):
    """Test contour plots with custom level argument."""
    coords = sp.samplings.equal_area(n_max=0, n_points=500)
    data = np.sin(coords.colatitude) * np.cos(coords.azimuth)

    # do plotting
    filename = f'{function.__name__}_levels_{type(levels).__name__}'
    create_figure()
    function(coords, data, levels=levels)
    save_and_compare(create_baseline, baseline_path, output_path, filename,
                     file_type, compare_output)

    # test error for invalid inputs
    match = 'levels'
    with pytest.raises(ValueError, match=match):
        function(coords, data, levels='levels')


@pytest.mark.parametrize('function', [
    sp.plot.pcolor_map,
    ])
@pytest.mark.parametrize('refine', [
    True, False,
    ])
def test_spherical_refine(function, refine):
    """Test contour plots with custom refine argument."""
    coords = sp.samplings.equal_area(n_max=0, n_points=500)
    data = np.sin(coords.colatitude) * np.cos(coords.azimuth)

    # do plotting
    filename = f'{function.__name__}_refine_{refine}'
    create_figure()
    function(coords, data, refine=refine)
    save_and_compare(create_baseline, baseline_path, output_path, filename,
                     file_type, compare_output)


@pytest.mark.parametrize('function', [
    sp.plot.pcolor_map,
    ])
def test_spherical_refine_error(function):
    """Test contour plots with invalid refine argument."""
    coords = sp.samplings.equal_area(n_max=0, n_points=500)
    data = np.sin(coords.colatitude) * np.cos(coords.azimuth)

    match = 'refine'
    with pytest.raises(ValueError, match=match):
        function(coords, data, refine='refine')


@pytest.mark.parametrize(("function", "projection"), [
    (sp.plot.balloon, '3d'),
    (sp.plot.balloon_wireframe, '3d'),
    (sp.plot.pcolor_sphere, '3d'),
    (sp.plot.contour, 'rectilinear'),
    (sp.plot.contour_map, 'mollweide'),
    (sp.plot.pcolor_map, 'mollweide'),
    ])
def test_data_plots_projection_input_and_return(function, projection):
    """
    Test all spherical plots with ax argument and check if this is
    also returned.
    """
    coords = sp.samplings.equal_area(n_max=0, n_points=500)
    data = np.sin(coords.colatitude) * np.cos(coords.azimuth)
    # do plotting
    create_figure()
    ax = plt.axes(projection=projection)

    ax_out = function(coords, data, ax=ax)[0]

    # check if the returned axis has the correct projection
    assert ax_out[0].name == projection

    # test error for invalid inputs
    match = f"The projection of the axis needs to be '{projection}'"
    with pytest.raises(ValueError, match=match):
        function(coords, data, ax=plt.axes(projection='polar'))

    match = re.escape("If [ax1, ax2] is passed ax2 needs to be of"
                      " 'rectilinear' projection")
    with pytest.raises(ValueError, match=match):
        function(coords, data, ax=[
            plt.axes(projection=projection),
            plt.axes(projection='3d')])

@pytest.mark.parametrize('function', [
    (sp.plot.scatter),
    (sp.plot.voronoi_cells_sphere),
    ])
def test_coordinates_plots_projection_input_and_return(function):
    """Test scatter plots with ax argument and check if this is also returned.
    """
    coords = sp.samplings.equal_area(n_max=0, n_points=12)
    # do plotting
    create_figure()
    ax = plt.axes(projection='3d')

    ax_out = function(coords, ax=ax)

    # check if the returned axis is a 3D axis
    assert ax_out.name == '3d'

    # test error for invalid inputs
    match = "The projection of the axis needs to be '3d'"
    with pytest.raises(ValueError, match=match):
        function(coords, ax=plt.axes())


def test_cmap_phase_twilight():
    """Test the phase twilight colormap."""
    lut = 512
    filename = 'cmap_phase_twilight'
    create_figure(6, .5)

    # do plotting
    colormap = sp.plot.phase_twilight(lut)
    ax = plt.gca()
    gradient = np.linspace(0, 1, lut)
    gradient = np.vstack((gradient, gradient))
    ax.imshow(gradient, aspect='auto', cmap=colormap)
    ax.set_axis_off()

    save_and_compare(
        create_baseline, baseline_path, output_path, filename,
        file_type, compare_output)

    match = 'lut must be a positive integer.'
    with pytest.raises(ValueError, match=match):
        sp.plot.phase_twilight('lut')

    with pytest.raises(ValueError, match=match):
        sp.plot.phase_twilight(0)


@pytest.mark.parametrize(
        ("function"),
        [(sp.plot.balloon),
         (sp.plot.balloon_wireframe),
         (sp.plot.contour),
         (sp.plot.contour_map),
         (sp.plot.pcolor_sphere),
         (sp.plot.pcolor_map)],
)
@pytest.mark.parametrize(
    ("ax"), [(1), ([1, 2, 3]), np.array([1, 2, 3])],
)
def test_ax_parameter_errors(function, ax, equal_area_sampling):
    """Test all warnings related to `ax` parameter."""
    coords, data = equal_area_sampling
    # Test case when `ax`` is wrong type
    match = "ax can be ``None``, a single matplotlib.axes.Axes object or a " \
            "list, tuple or array of two axes"
    with pytest.raises(ValueError, match=match):
        function(coords, data, ax=ax)


@pytest.mark.parametrize(
        ("function"),
        [(sp.plot.balloon),
         (sp.plot.balloon_wireframe),
         (sp.plot.contour),
         (sp.plot.contour_map),
         (sp.plot.pcolor_sphere),
         (sp.plot.pcolor_map)],
)
def test_colorbar_ax_error(function, equal_area_sampling):
    """Test error raised by false colorbar / ax parameter combination."""
    coords, data = equal_area_sampling
    match = "A list of axes can not be used if colorbar is False"
    with pytest.raises(
        ValueError, match=match,
    ):
        function(coords, data, colorbar=False, ax=[plt.gca(), plt.gca()])


@pytest.mark.parametrize(
        ('function', 'projection'),
        [(sp.plot.balloon, '3d'),
         (sp.plot.balloon_wireframe, '3d'),
         (sp.plot.contour, 'rectilinear'),
         (sp.plot.contour_map, 'mollweide'),
         (sp.plot.pcolor_sphere, '3d'),
         (sp.plot.pcolor_map, 'mollweide')],
)
@pytest.mark.parametrize(
    ('ax_option'),
    [('none'), ('single'), ('two')],
)
def test_ax_input(function, projection, ax_option, equal_area_sampling):
    """Test ax input for None, single axis and list of two axes."""
    coords, data = equal_area_sampling

    filename = f'{function.__name__}_ax_{ax_option}'

    if ax_option == 'none':
        function(coords, data)
    if ax_option == 'single':
        ax1 = plt.subplot(111, projection=projection)
        function(coords, data, ax=ax1)
    if ax_option == 'two':
        ax1 = plt.subplot(121, projection=projection)
        ax2 = plt.subplot(122, projection='rectilinear')

        function(coords, data, ax=[ax1, ax2])

        plt.tight_layout()
        plt.subplots_adjust(wspace=0.2)

    save_and_compare(create_baseline, baseline_path, output_path, filename,
                     file_type, compare_output)


@pytest.mark.parametrize(
        ('function', 'kind'),
        [(sp.plot.scatter, 'coordinate_plot'),
         (sp.plot.voronoi_cells_sphere, 'coordinate_plot'),
         (sp.plot.pcolor_sphere, 'data_plot'),
         (sp.plot.pcolor_map, 'data_plot'),
         (sp.plot.balloon, 'data_plot'),
         (sp.plot.balloon_wireframe, 'data_plot'),
         (sp.plot.contour, 'data_plot'),
         (sp.plot.contour_map, 'data_plot')],
)
@pytest.mark.parametrize(
        'style',
        [('light'), ('dark')],
)
def test_pyfar_styles(function, style, kind, equal_area_sampling):
    """Test style-parameter with pyfar plot styles 'light' and 'dark'."""
    coords, data = equal_area_sampling
    filename = f'{function.__name__}_style_{style}'

    if kind == 'coordinate_plot':
        function(coords, style=style)
    elif kind == 'data_plot':
        function(coords, data, style=style)

    save_and_compare(create_baseline, baseline_path, output_path, filename,
                     file_type, compare_output)


@pytest.mark.parametrize(
        ('function', 'kind'),
        [(sp.plot.scatter, 'coordinate_plot'),
         (sp.plot.voronoi_cells_sphere, 'coordinate_plot'),
         (sp.plot.pcolor_sphere, 'data_plot'),
         (sp.plot.pcolor_map, 'data_plot'),
         (sp.plot.balloon, 'data_plot'),
         (sp.plot.balloon_wireframe, 'data_plot'),
         (sp.plot.contour, 'data_plot'),
         (sp.plot.contour_map, 'data_plot')],
)
def test_custom_style(function, kind, equal_area_sampling):
    """Test style parameter with custom value."""
    coords, data = equal_area_sampling
    style = {'axes.facecolor': '#000000'}

    if kind == 'coordinate_plot':
        function(coords, style=style)
    elif kind == 'data_plot':
        function(coords, data, style=style)

    facecolor = mpl.colors.to_hex(plt.gca().patch.get_facecolor())
    assert facecolor == style['axes.facecolor']
