import numpy as np
import pytest

from matplotlib.testing.decorators import image_comparison
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib import ticker, rcParams


def example_plot(ax, fontsize=12, nodec=False):
    ax.plot([1, 2])
    ax.locator_params(nbins=3)
    if not nodec:
        ax.set_xlabel('x-label', fontsize=fontsize)
        ax.set_ylabel('y-label', fontsize=fontsize)
        ax.set_title('Title', fontsize=fontsize)
    else:
        ax.set_xticklabels('')
        ax.set_yticklabels('')


def example_pcolor(ax, fontsize=12):
    dx, dy = 0.6, 0.6
    y, x = np.mgrid[slice(-3, 3 + dy, dy),
                    slice(-3, 3 + dx, dx)]
    z = (1 - x / 2. + x ** 5 + y ** 3) * np.exp(-x ** 2 - y ** 2)
    pcm = ax.pcolormesh(x, y, z[:-1, :-1], cmap='RdBu_r', vmin=-1., vmax=1.,
                        rasterized=True)
    ax.set_xlabel('x-label', fontsize=fontsize)
    ax.set_ylabel('y-label', fontsize=fontsize)
    ax.set_title('Title', fontsize=fontsize)
    return pcm


@image_comparison(['constrained_layout1.png'])
def test_constrained_layout1():
    """Test constrained_layout for a single subplot"""
    fig = plt.figure(constrained_layout=True)
    ax = fig.add_subplot(111)
    example_plot(ax, fontsize=24)


@image_comparison(['constrained_layout2.png'])
def test_constrained_layout2():
    """Test constrained_layout for 2x2 subplots"""
    fig, axs = plt.subplots(2, 2, constrained_layout=True)
    for ax in axs.flat:
        example_plot(ax, fontsize=24)


@image_comparison(['constrained_layout3.png'])
def test_constrained_layout3():
    """Test constrained_layout for colorbars with subplots"""
    fig, axs = plt.subplots(2, 2, constrained_layout=True)
    for nn, ax in enumerate(axs.flat):
        pcm = example_pcolor(ax, fontsize=24)
        if nn == 3:
            pad = 0.08
        else:
            pad = 0.02  # default
        fig.colorbar(pcm, ax=ax, pad=pad)


@image_comparison(['constrained_layout4'])
def test_constrained_layout4():
    """Test constrained_layout for a single colorbar with subplots"""
    fig, axs = plt.subplots(2, 2, constrained_layout=True)
    for ax in axs.flat:
        pcm = example_pcolor(ax, fontsize=24)
    fig.colorbar(pcm, ax=axs, pad=0.01, shrink=0.6)


@image_comparison(['constrained_layout5.png'], tol=5.e-2)
def test_constrained_layout5():
    """
    Test constrained_layout for a single colorbar with subplots,
    colorbar bottom
    """
    fig, axs = plt.subplots(2, 2, constrained_layout=True)
    for ax in axs.flat:
        pcm = example_pcolor(ax, fontsize=24)
    fig.colorbar(pcm, ax=axs,
                 use_gridspec=False, pad=0.01, shrink=0.6,
                 location='bottom')


@image_comparison(['constrained_layout6.png'])
def test_constrained_layout6():
    """Test constrained_layout for nested gridspecs"""
    fig = plt.figure(constrained_layout=True)
    gs = fig.add_gridspec(1, 2, figure=fig)
    gsl = gs[0].subgridspec(2, 2)
    gsr = gs[1].subgridspec(1, 2)
    axsl = []
    for gs in gsl:
        ax = fig.add_subplot(gs)
        axsl += [ax]
        example_plot(ax, fontsize=12)
    ax.set_xlabel('x-label\nMultiLine')
    axsr = []
    for gs in gsr:
        ax = fig.add_subplot(gs)
        axsr += [ax]
        pcm = example_pcolor(ax, fontsize=12)

    fig.colorbar(pcm, ax=axsr,
                 pad=0.01, shrink=0.99, location='bottom',
                 ticks=ticker.MaxNLocator(nbins=5))


def test_constrained_layout7():
    """Test for proper warning if fig not set in GridSpec"""
    with pytest.warns(
        UserWarning, match=('Calling figure.constrained_layout, but figure '
                            'not setup to do constrained layout')):
        fig = plt.figure(constrained_layout=True)
        gs = gridspec.GridSpec(1, 2)
        gsl = gridspec.GridSpecFromSubplotSpec(2, 2, gs[0])
        gsr = gridspec.GridSpecFromSubplotSpec(1, 2, gs[1])
        for gs in gsl:
            fig.add_subplot(gs)
        # need to trigger a draw to get warning
        fig.draw(fig.canvas.get_renderer())


@image_comparison(['constrained_layout8.png'])
def test_constrained_layout8():
    """Test for gridspecs that are not completely full"""
    fig = plt.figure(figsize=(10, 5), constrained_layout=True)
    gs = gridspec.GridSpec(3, 5, figure=fig)
    axs = []
    for j in [0, 1]:
        if j == 0:
            ilist = [1]
        else:
            ilist = [0, 4]
        for i in ilist:
            ax = fig.add_subplot(gs[j, i])
            axs += [ax]
            pcm = example_pcolor(ax, fontsize=9)
            if i > 0:
                ax.set_ylabel('')
            if j < 1:
                ax.set_xlabel('')
            ax.set_title('')
    ax = fig.add_subplot(gs[2, :])
    axs += [ax]
    pcm = example_pcolor(ax, fontsize=9)

    fig.colorbar(pcm, ax=axs, pad=0.01, shrink=0.6)


@image_comparison(['constrained_layout9.png'])
def test_constrained_layout9():
    """Test for handling suptitle and for sharex and sharey"""
    fig, axs = plt.subplots(2, 2, constrained_layout=True,
                            sharex=False, sharey=False)
    for ax in axs.flat:
        pcm = example_pcolor(ax, fontsize=24)
        ax.set_xlabel('')
        ax.set_ylabel('')
    ax.set_aspect(2.)
    fig.colorbar(pcm, ax=axs, pad=0.01, shrink=0.6)
    fig.suptitle('Test Suptitle', fontsize=28)


@image_comparison(['constrained_layout10.png'])
def test_constrained_layout10():
    """Test for handling legend outside axis"""
    fig, axs = plt.subplots(2, 2, constrained_layout=True)
    for ax in axs.flat:
        ax.plot(np.arange(12), label='This is a label')
    ax.legend(loc='center left', bbox_to_anchor=(0.8, 0.5))


@image_comparison(['constrained_layout11.png'])
def test_constrained_layout11():
    """Test for multiple nested gridspecs"""
    fig = plt.figure(constrained_layout=True, figsize=(13, 3))
    gs0 = gridspec.GridSpec(1, 2, figure=fig)
    gsl = gridspec.GridSpecFromSubplotSpec(1, 2, gs0[0])
    gsl0 = gridspec.GridSpecFromSubplotSpec(2, 2, gsl[1])
    ax = fig.add_subplot(gs0[1])
    example_plot(ax, fontsize=9)
    axs = []
    for gs in gsl0:
        ax = fig.add_subplot(gs)
        axs += [ax]
        pcm = example_pcolor(ax, fontsize=9)
    fig.colorbar(pcm, ax=axs, shrink=0.6, aspect=70.)
    ax = fig.add_subplot(gsl[0])
    example_plot(ax, fontsize=9)


@image_comparison(['constrained_layout11rat.png'])
def test_constrained_layout11rat():
    """Test for multiple nested gridspecs with width_ratios"""
    fig = plt.figure(constrained_layout=True, figsize=(10, 3))
    gs0 = gridspec.GridSpec(1, 2, figure=fig, width_ratios=[6, 1])
    gsl = gridspec.GridSpecFromSubplotSpec(1, 2, gs0[0])
    gsl0 = gridspec.GridSpecFromSubplotSpec(2, 2, gsl[1], height_ratios=[2, 1])
    ax = fig.add_subplot(gs0[1])
    example_plot(ax, fontsize=9)
    axs = []
    for gs in gsl0:
        ax = fig.add_subplot(gs)
        axs += [ax]
        pcm = example_pcolor(ax, fontsize=9)
    fig.colorbar(pcm, ax=axs, shrink=0.6, aspect=70.)
    ax = fig.add_subplot(gsl[0])
    example_plot(ax, fontsize=9)


@image_comparison(['constrained_layout12.png'])
def test_constrained_layout12():
    """Test that very unbalanced labeling still works."""
    fig = plt.figure(constrained_layout=True)

    gs0 = gridspec.GridSpec(6, 2, figure=fig)

    ax1 = fig.add_subplot(gs0[:3, 1])
    ax2 = fig.add_subplot(gs0[3:, 1])

    example_plot(ax1, fontsize=24)
    example_plot(ax2, fontsize=24)

    ax = fig.add_subplot(gs0[0:2, 0])
    example_plot(ax, nodec=True)
    ax = fig.add_subplot(gs0[2:4, 0])
    example_plot(ax, nodec=True)
    ax = fig.add_subplot(gs0[4:, 0])
    example_plot(ax, nodec=True)
    ax.set_xlabel('x-label')


@image_comparison(['constrained_layout13.png'], tol=2.e-2)
def test_constrained_layout13():
    """Test that padding works."""
    fig, axs = plt.subplots(2, 2, constrained_layout=True)
    for ax in axs.flat:
        pcm = example_pcolor(ax, fontsize=12)
        fig.colorbar(pcm, ax=ax, shrink=0.6, aspect=20., pad=0.02)
    fig.set_constrained_layout_pads(w_pad=24./72., h_pad=24./72.)


@image_comparison(['constrained_layout14.png'])
def test_constrained_layout14():
    """Test that padding works."""
    fig, axs = plt.subplots(2, 2, constrained_layout=True)
    for ax in axs.flat:
        pcm = example_pcolor(ax, fontsize=12)
        fig.colorbar(pcm, ax=ax, shrink=0.6, aspect=20., pad=0.02)
    fig.set_constrained_layout_pads(
            w_pad=3./72., h_pad=3./72.,
            hspace=0.2, wspace=0.2)


@image_comparison(['constrained_layout15.png'])
def test_constrained_layout15():
    """Test that rcparams work."""
    rcParams['figure.constrained_layout.use'] = True
    fig, axs = plt.subplots(2, 2)
    for ax in axs.flat:
        example_plot(ax, fontsize=12)


@image_comparison(['constrained_layout16.png'])
def test_constrained_layout16():
    """Test ax.set_position."""
    fig, ax = plt.subplots(constrained_layout=True)
    example_plot(ax, fontsize=12)
    ax2 = fig.add_axes([0.2, 0.2, 0.4, 0.4])


@image_comparison(['constrained_layout17.png'])
def test_constrained_layout17():
    """Test uneven gridspecs"""
    fig = plt.figure(constrained_layout=True)
    gs = gridspec.GridSpec(3, 3, figure=fig)

    ax1 = fig.add_subplot(gs[0, 0])
    ax2 = fig.add_subplot(gs[0, 1:])
    ax3 = fig.add_subplot(gs[1:, 0:2])
    ax4 = fig.add_subplot(gs[1:, -1])

    example_plot(ax1)
    example_plot(ax2)
    example_plot(ax3)
    example_plot(ax4)


def test_constrained_layout18():
    """Test twinx"""
    fig, ax = plt.subplots(constrained_layout=True)
    ax2 = ax.twinx()
    example_plot(ax)
    example_plot(ax2, fontsize=24)
    fig.canvas.draw()
    assert all(ax.get_position().extents == ax2.get_position().extents)


def test_constrained_layout19():
    """Test twiny"""
    fig, ax = plt.subplots(constrained_layout=True)
    ax2 = ax.twiny()
    example_plot(ax)
    example_plot(ax2, fontsize=24)
    ax2.set_title('')
    ax.set_title('')
    fig.canvas.draw()
    assert all(ax.get_position().extents == ax2.get_position().extents)


def test_constrained_layout20():
    """Smoke test cl does not mess up added axes"""
    gx = np.linspace(-5, 5, 4)
    img = np.hypot(gx, gx[:, None])

    fig = plt.figure()
    ax = fig.add_axes([0, 0, 1, 1])
    mesh = ax.pcolormesh(gx, gx, img[:-1, :-1])
    fig.colorbar(mesh)


def test_constrained_layout21():
    """#11035: repeated calls to suptitle should not alter the layout"""
    fig, ax = plt.subplots(constrained_layout=True)

    fig.suptitle("Suptitle0")
    fig.canvas.draw()
    extents0 = np.copy(ax.get_position().extents)

    fig.suptitle("Suptitle1")
    fig.canvas.draw()
    extents1 = np.copy(ax.get_position().extents)

    np.testing.assert_allclose(extents0, extents1)


def test_constrained_layout22():
    """#11035: suptitle should not be include in CL if manually positioned"""
    fig, ax = plt.subplots(constrained_layout=True)

    fig.canvas.draw()
    extents0 = np.copy(ax.get_position().extents)

    fig.suptitle("Suptitle", y=0.5)
    fig.canvas.draw()
    extents1 = np.copy(ax.get_position().extents)

    np.testing.assert_allclose(extents0, extents1)


def test_constrained_layout23():
    """
    Comment in #11035: suptitle used to cause an exception when
    reusing a figure w/ CL with ``clear=True``.
    """

    for i in range(2):
        fig, ax = plt.subplots(num="123", constrained_layout=True, clear=True)
        fig.suptitle("Suptitle{}".format(i))


# This test occasionally fails the image comparison tests, so we mark as
# flaky.  Apparently the constraint solver occasionally doesn't fully
# optimize.  Would be nice if this were more deterministic...
@pytest.mark.timeout(30)
@pytest.mark.flaky(reruns=3)
@image_comparison(['test_colorbar_location.png'],
                  remove_text=True, style='mpl20')
def test_colorbar_location():
    """
    Test that colorbar handling is as expected for various complicated
    cases...
    """

    fig, axs = plt.subplots(4, 5, constrained_layout=True)
    for ax in axs.flat:
        pcm = example_pcolor(ax)
        ax.set_xlabel('')
        ax.set_ylabel('')
    fig.colorbar(pcm, ax=axs[:, 1], shrink=0.4)
    fig.colorbar(pcm, ax=axs[-1, :2], shrink=0.5, location='bottom')
    fig.colorbar(pcm, ax=axs[0, 2:], shrink=0.5, location='bottom')
    fig.colorbar(pcm, ax=axs[-2, 3:], shrink=0.5, location='top')
    fig.colorbar(pcm, ax=axs[0, 0], shrink=0.5, location='left')
    fig.colorbar(pcm, ax=axs[1:3, 2], shrink=0.5, location='right')


def test_hidden_axes():
    # test that if we make an axes not visible that constrained_layout
    # still works.  Note the axes still takes space in the layout
    # (as does a gridspec slot that is empty)
    fig, axs = plt.subplots(2, 2, constrained_layout=True)
    axs[0, 1].set_visible(False)
    fig.canvas.draw()
    extents1 = np.copy(axs[0, 0].get_position().extents)

    np.testing.assert_allclose(
        extents1, [0.045552, 0.548288, 0.47319, 0.982638], rtol=1e-5)