"""
The ways several variables can share one set of axes.

Blending is only one of them, and the one whose cost is highest: where two
layers meet the colour belongs to neither bar. These tests pin down what each of
the others actually puts on the axes.
"""

import matplotlib
matplotlib.use('Agg')

import matplotlib.pyplot as plt
import numpy as np
import pytest
from matplotlib.collections import QuadMesh
from matplotlib.colors import Normalize

from dispnc.colors import CURVE_COLORS, curve_color, curve_style
from dispnc.overlay import Composite, Layer, resolve_thresholds
from dispnc.render import layers as layer_draw


@pytest.fixture
def axes():
    fig, ax = plt.subplots()
    yield fig, ax
    plt.close(fig)


@pytest.fixture
def field():
    y, x = np.mgrid[0:12, 0:20]
    return np.asarray(x + y, dtype=float)


def prepared(values, **kwargs):
    kwargs.setdefault('norm', Normalize(float(np.min(values)), float(np.max(values))))
    return [(Layer('co2_ice', values, 'Reds', **kwargs), values, None)]


def test_blending_puts_a_mesh_and_a_bar_on_the_axes(axes, field):
    fig, ax = axes
    drawn = layer_draw.draw(fig, ax, prepared(field), np.arange(20), np.arange(12),
                            style='blend')
    assert len(drawn) == 1
    assert isinstance(drawn[0][1], QuadMesh)
    # its own bar, and no legend: a bar names its own layer
    assert ax.get_legend() is None


def test_contours_are_lines_and_leave_the_colour_channel_alone(axes, field):
    """
    The point of the style: the overlay is geometry, not paint, so the base
    keeps its own colormap and both fields stay readable where they cross.
    """
    fig, ax = axes
    before = len(ax.collections)
    drawn = layer_draw.draw(fig, ax, prepared(field), np.arange(20), np.arange(12),
                            style='contour')

    assert len(ax.collections) > before, 'nothing was drawn'
    assert not any(isinstance(c, QuadMesh) for c in ax.collections), \
        'a contour layer must not lay down a mesh'
    # No bar, so a legend has to name it instead
    legend = ax.get_legend()
    assert legend is not None
    assert [t.get_text() for t in legend.get_texts()] == ['co2_ice']


def test_a_hatch_marks_where_a_layer_is_above_its_threshold(axes, field, capsys):
    fig, ax = axes
    layer_draw.draw(fig, ax, prepared(field), np.arange(20), np.arange(12),
                    style='hatch')
    legend = ax.get_legend()
    assert legend is not None
    # The label says which way to read it: presence above a level, not amount
    assert '>' in legend.get_texts()[0].get_text()


def test_a_flat_layer_is_reported_rather_than_contoured(axes, capsys):
    fig, ax = axes
    flat = np.ones((12, 20))
    layer_draw.draw(fig, ax, prepared(flat), np.arange(20), np.arange(12),
                    style='contour')
    assert 'nothing to contour' in capsys.readouterr().out


def test_glyphs_are_drawn_without_touching_the_base_colours(axes, field):
    fig, ax = axes
    layer_draw.draw(fig, ax, prepared(field), np.arange(20), np.arange(12),
                    style='glyph')
    assert not any(isinstance(c, QuadMesh) for c in ax.collections)
    assert ax.get_legend() is not None


def test_an_unknown_style_falls_back_to_blending_out_loud(axes, field, capsys):
    fig, ax = axes
    layer_draw.draw(fig, ax, prepared(field), np.arange(20), np.arange(12),
                    style='kaleidoscope')
    assert 'unknown --overlay-style' in capsys.readouterr().out


@pytest.mark.parametrize('style, bars', [
    ('blend', 3),      # the base and one per layer
    ('contour', 1),    # the base alone; the lines are labelled and in the legend
    ('hatch', 1),
    ('glyph', 1),
])
def test_only_blending_asks_for_a_bar_per_layer(style, bars):
    """
    What the other styles buy: the width the extra bars would have taken goes
    back to the plot.
    """
    composite = Composite(layers=[Layer('a', None, 'viridis'),
                                  Layer('b', None, 'Reds'),
                                  Layer('c', None, 'Blues')])
    assert layer_draw.bar_count(composite, style) == bars


def test_no_composite_asks_for_one_bar():
    assert layer_draw.bar_count(None, 'blend') == 1


# --- how a layer's colour is decided -----------------------------------------

def test_the_rota_gives_each_layer_its_own_colour():
    chosen = [curve_color(None, i) for i in range(len(CURVE_COLORS))]
    assert len(set(chosen)) == len(CURVE_COLORS)


def test_a_named_colour_is_taken_as_itself():
    assert curve_color('red', 3) == '#ff0000'


def test_a_named_colormap_becomes_one_colour_of_that_hue():
    """
    `--overlay co2_ice:Reds` means the same thing whichever style draws it: a
    ramp when the layer is a fill, its saturated end when the layer is a line.
    """
    colour = curve_color('Reds', 0)
    red, green, blue = (int(colour[i:i + 2], 16) for i in (1, 3, 5))
    assert red > green and red > blue


def test_something_that_is_neither_says_so(capsys):
    assert curve_color('not-a-thing', 0) == CURVE_COLORS[0]
    assert 'neither a colour nor a colormap' in capsys.readouterr().out


def test_dashes_turn_over_only_once_the_colours_have():
    assert curve_style(0) == curve_style(len(CURVE_COLORS) - 1) == '-'
    assert curve_style(len(CURVE_COLORS)) == '--'


# --- thresholds ---------------------------------------------------------------

def test_a_percentile_is_resolved_against_the_layers_own_values():
    """
    One setting has to mean "the top tenth of each" across fields whose
    magnitudes differ by decades, or it means nothing.
    """
    small = resolve_thresholds(np.linspace(0, 1, 101), 'p90')
    large = resolve_thresholds(np.linspace(0, 1e6, 101), 'p90')
    assert small[0] == pytest.approx(0.9)
    assert large[0] == pytest.approx(9e5)


def test_bare_numbers_are_used_as_given():
    assert resolve_thresholds(np.arange(100.0), '10,20,30') == [10.0, 20.0, 30.0]


def test_thresholds_come_back_sorted_whatever_order_they_were_given():
    assert resolve_thresholds(np.arange(100.0), '30,10,20') == [10.0, 20.0, 30.0]


def test_something_that_is_not_a_threshold_is_reported(capsys):
    assert resolve_thresholds(np.arange(100.0), '10,banana') == [10.0]
    assert 'not a value or a percentile' in capsys.readouterr().out


def test_an_empty_field_has_no_thresholds():
    assert resolve_thresholds(np.full(10, np.nan), 'p90') == []


# --- one panel per variable ---------------------------------------------------

def test_a_grid_hands_out_one_cell_per_variable():
    from dispnc.render.panels import layout

    grid = layout(3)
    try:
        seen = [grid.next_axes()[1] for _ in range(3)]
        assert len(set(id(ax) for ax in seen)) == 3
    finally:
        plt.close(grid.fig)


def test_a_grid_is_near_square_rather_than_one_long_row():
    """
    Eight variables in a row is a figure nobody can read.
    """
    from dispnc.render.panels import layout

    grid = layout(8)
    try:
        rows = {ax.get_subplotspec().rowspan.start for ax in grid.cells}
        columns = {ax.get_subplotspec().colspan.start for ax in grid.cells}
        assert len(rows) == 3 and len(columns) == 3
    finally:
        plt.close(grid.fig)


def test_spare_cells_are_removed_rather_than_left_empty():
    """
    An empty frame reads as a panel whose data failed to draw.
    """
    from dispnc.render.panels import layout

    grid = layout(3)          # a 2x2 grid with one cell over
    try:
        assert len(grid.cells) == 3
        assert len(grid.fig.axes) == 3
    finally:
        plt.close(grid.fig)


def test_the_figure_is_unfinished_until_the_last_cell_is_handed_out():
    """
    Renderers save as their last act and cannot know they are one of several, so
    the grid is what stops the first panel writing the figure over the others.
    """
    from dispnc.figure import finish_figure
    from dispnc.render.panels import layout

    grid = layout(2)
    try:
        grid.next_axes()
        assert grid.fig._dispnc_incomplete is True
        # finish_figure must decline to do anything with it
        assert finish_figure(grid.fig, None) == 0
        assert grid.fig.number in plt.get_fignums()

        grid.next_axes()
        assert grid.fig._dispnc_incomplete is False
    finally:
        plt.close(grid.fig)


def test_a_cell_can_be_reissued_with_a_projection():
    """
    A map needs a GeoAxes and a cartopy axes cannot be converted after the fact,
    so the placeholder is replaced in place.
    """
    import cartopy.crs as ccrs
    from dispnc.render.panels import layout

    grid = layout(1)
    try:
        _, ax = grid.next_axes(projection=ccrs.PlateCarree())
        assert isinstance(getattr(ax, 'projection', None), ccrs.Projection)
        assert len(grid.fig.axes) == 1, 'the placeholder must not be left behind'
    finally:
        plt.close(grid.fig)


def test_the_dateline_cells_share_the_base_mesh_norm():
    """
    A field that is the same value everywhere used to come out with a dark
    column at -180 and another at +180, reading as signal on a flat map.

    Cartopy cannot draw a cell that straddles 180 degrees as part of a QuadMesh,
    so it masks those out and redraws them with `pcolor`. Given no norm - the
    ordinary case - the second collection got a Normalize of its own, and
    `fig.colorbar` widening the degenerate range reached only the first one,
    leaving the dateline cells mapping everything to the bottom of the colormap.
    """
    import cartopy.crs as ccrs

    lons = np.append(np.arange(-180.0, 180.0, 11.25), 180.0)
    lats = np.linspace(90.0, -90.0, 33)
    flat = np.zeros((lats.size, lons.size))

    fig, ax = plt.subplots(subplot_kw={'projection': ccrs.PlateCarree()})
    try:
        mesh, _ = layer_draw.draw_base(fig, ax, flat, [], lons, lats, 'viridis',
                                       None, 'runoff', transform=ccrs.PlateCarree())
        wrapped = getattr(mesh, '_wrapped_collection_fix', None)
        assert wrapped is not None, 'this grid must produce cells to wrap'
        assert wrapped.norm is mesh.norm

        # The bar widens vmin == vmax, and both collections must follow it
        fig.colorbar(mesh)
        assert mesh.get_clim() == wrapped.get_clim()
        assert mesh.get_clim()[0] < mesh.get_clim()[1]
    finally:
        plt.close(fig)


def test_the_cells_a_log_scale_cannot_place_are_painted_under_the_mesh(axes):
    """
    A log scale masks every zero, and a masked cell is a hole - which reads as
    "the field is not here" rather than "the field is zero here". The grey goes
    under the mesh rather than on the colormap's 'bad' entry, because cartopy
    needs that one transparent in order to wrap a quadmesh across the dateline;
    an opaque one made it redraw the wrapped copy over the entire map.
    """
    from matplotlib.colors import LogNorm
    from dispnc.colors import log_colormap

    fig, ax = axes
    values = np.array([[1.0, 10.0], [0.0, 100.0]])
    cmap = log_colormap('viridis', 1, 4)

    mesh, _ = layer_draw.draw_base(fig, ax, values, [], [0, 1], [0, 1], cmap,
                                   LogNorm(vmin=1, vmax=100), 'q')
    underlay = mesh._dispnc_blank_underlay
    assert underlay is not None
    assert cmap.get_bad()[-1] == 0, 'cartopy needs a transparent bad colour'

    painted = ~np.ma.getmaskarray(np.ma.masked_invalid(underlay.get_array()))
    assert painted.ravel().tolist() == [False, False, True, False]


def test_a_plain_colormap_gains_no_underlay(axes, field):
    fig, ax = axes
    mesh, _ = layer_draw.draw_base(fig, ax, field, [], np.arange(20),
                                   np.arange(12), 'viridis', None, 'x')
    assert mesh._dispnc_blank_underlay is None


def test_missing_data_stays_a_hole(axes):
    """
    The zeros get a colour so that they stop being mistaken for missing data,
    which would be no gain at all if the missing data were given it too.
    """
    from dispnc.colors import blank_cells

    assert blank_cells([[0.0, np.nan, 1.0, -2.0]]).ravel().tolist() == \
        [True, False, False, True]


def test_the_underlay_steps_with_the_frames(axes):
    """
    A cap grows and retreats, so an underlay left on frame zero would grey out
    cells the frame on screen has values for.
    """
    from matplotlib.colors import LogNorm
    from dispnc.colors import log_colormap

    fig, ax = axes
    first = np.array([[0.0, 10.0], [1.0, 100.0]])
    second = np.array([[5.0, 10.0], [0.0, 100.0]])

    mesh, _ = layer_draw.draw_base(fig, ax, first, [], [0, 1], [0, 1],
                                   log_colormap('viridis', 1, 4),
                                   LogNorm(vmin=1, vmax=100), 'q')
    layer_draw.refill_blank_cells(mesh, second)

    painted = ~np.ma.getmaskarray(
        np.ma.masked_invalid(mesh._dispnc_blank_underlay.get_array()))
    assert painted.ravel().tolist() == [False, False, True, False]