"""
Colormap and normalization choices.
"""

import numpy as np
import pytest
from matplotlib.colors import LogNorm, Normalize, SymLogNorm

from dispnc.colors import (DIVERGING, SEQUENTIAL, choose_style, spans_decades,
                        straddles_zero, symmetric_limit)


def test_explicit_colormap_is_passed_through():
    cmap, norm, _ = choose_style(np.arange(10.0), 'x', cmap='jet')
    assert cmap == 'jet' and norm is None


def test_default_is_perceptual_sequential():
    cmap, _, _ = choose_style(np.linspace(200, 300, 50), 'tsurf')
    assert cmap == SEQUENTIAL


def test_anomaly_is_diverging_and_symmetric():
    data = np.linspace(-3, 7, 100)
    cmap, norm, _ = choose_style(data, 'ps_avg', signed=True)
    assert cmap == DIVERGING
    assert norm.vmin == pytest.approx(-norm.vmax)


def test_tendency_name_that_straddles_zero_is_diverging():
    data = np.linspace(-5, 5, 100)
    cmap, norm, _ = choose_style(data, 'd_co2ice')
    assert cmap == DIVERGING
    assert norm.vmin == pytest.approx(-norm.vmax)


def test_tendency_name_that_never_crosses_zero_is_not_centred():
    """
    The data has the final say: a d_ field that is entirely positive must not be
    centred on zero.
    """
    cmap, norm, _ = choose_style(np.linspace(1, 10, 50), 'd_h2oice')
    assert cmap == SEQUENTIAL


def test_mostly_positive_field_with_a_sliver_of_negative_is_not_diverging():
    data = np.concatenate([np.linspace(0.5, 10, 199), [-0.001]])
    assert not straddles_zero(data)
    cmap, _, _ = choose_style(data, 'd_runoff')
    assert cmap == SEQUENTIAL


def test_symmetric_limit_ignores_outliers():
    data = np.concatenate([np.linspace(-2, 2, 999), [1e6]])
    assert symmetric_limit(data) < 10


def test_log_scale_for_a_field_spanning_decades():
    data = np.logspace(-8, 0, 500)
    assert spans_decades(data)
    cmap, norm, _ = choose_style(data, 'q')
    assert isinstance(norm, LogNorm)


def test_field_that_merely_touches_zero_is_not_log_scaled():
    """
    A wind magnitude vanishes at the poles. Keying the log test on the smallest
    positive value made the whole map one flat colour.
    """
    lat = np.linspace(-90, 90, 200)
    magnitude = 20 * np.abs(np.cos(np.deg2rad(lat)))
    assert not spans_decades(magnitude)
    cmap, norm, _ = choose_style(magnitude, 'wind')
    assert not isinstance(norm, LogNorm)


def test_explicit_limits_are_honoured():
    _, norm, _ = choose_style(np.arange(100.0), 'x', vmin=10, vmax=20)
    assert (norm.vmin, norm.vmax) == (10, 20)


def test_norm_log_falls_back_when_no_positive_values(capsys):
    _, norm, _ = choose_style(np.linspace(-5, -1, 10), 'x', cmap='viridis', norm='log')
    assert norm is None
    assert 'falling back to linear' in capsys.readouterr().out


def test_an_explicit_colormap_keeps_the_zero_centring_of_a_difference():
    """
    `--anomaly time -c RdBu_r` is the natural thing to type, and it used to hand
    back norm=None: matplotlib then autoscaled to the raw asymmetric range and
    put zero at 17% of the bar. Naming a palette says which colours, not what
    the middle colour means.
    """
    data = np.array([-2.0, 0.5, 10.0])
    cmap, norm, _ = choose_style(data, 'ps_avg', cmap='RdBu_r', signed=True)
    assert cmap == 'RdBu_r'
    assert norm.vmin == pytest.approx(-norm.vmax)


def test_an_explicit_colormap_still_overrides_the_name_heuristic():
    """
    The counterpart: guessing from the *name* that a variable is signed must not
    override what the user typed, so 'd_co2ice' with -c jet gets no norm.
    """
    cmap, norm, _ = choose_style(np.linspace(-5, 5, 100), 'd_co2ice', cmap='jet')
    assert cmap == 'jet' and norm is None


def test_explicit_limits_beat_the_centring_of_a_difference():
    _, norm, _ = choose_style(np.linspace(-3, 7, 100), 'ps_avg', cmap='jet',
                              signed=True, vmin=0, vmax=5)
    assert (norm.vmin, norm.vmax) == (0, 5)


def test_a_log_scale_reports_the_cells_it_cannot_draw(capsys):
    """
    A log scale has nowhere on its bar for a zero or a negative, and says how
    many it had to paint instead. d_h2oice on one PEM step is 31% negative.
    """
    data = np.concatenate([np.linspace(-232, -1, 335), np.logspace(0, 3, 754)])
    choose_style(data, 'd_h2oice', cmap='viridis', norm='log')
    out = capsys.readouterr().out
    assert 'cannot place 335 of 1089 cells (31%)' in out
    assert '--norm symlog' in out


def test_a_log_scale_paints_the_cells_it_cannot_place(capsys):
    """
    The counterpart to the printed note: the colormap comes back carrying a
    colour for those cells, so the figure has no holes in it, and a sentence
    saying what the colour means.
    """
    data = np.concatenate([np.linspace(-232, -1, 335), np.logspace(0, 3, 754)])
    cmap, _, _ = choose_style(data, 'd_h2oice', cmap='viridis', norm='log')
    assert cmap._dispnc_blank_note
    assert '31%' in cmap._dispnc_blank_note


def test_a_log_scale_over_positive_values_says_nothing(capsys):
    choose_style(np.logspace(0, 4, 100), 'q', cmap='viridis', norm='log')
    assert 'cannot draw' not in capsys.readouterr().out


def test_a_field_that_is_mostly_zero_is_not_log_scaled():
    """
    A CO2 ice cap is zero over 97% of a PEM step and spans decades across the
    3% that is left. Percentiles of the positive values could not see that, so
    the automatic choice was a log scale that blanked almost the whole map -
    a picture of where the variable is rather than of what it is.
    """
    sparse = np.concatenate([np.zeros(1056), np.logspace(-4, 2, 33)])
    assert not spans_decades(sparse)
    cmap, norm, why = choose_style(sparse, 'co2ice')
    assert not isinstance(norm, LogNorm)
    assert cmap == SEQUENTIAL


def test_a_field_with_a_minority_of_zeros_gets_a_symlog_scale():
    """
    The counterpart: h2oice is zero over 15% of the same step, and the decades
    across the rest are exactly what a log scale is for. A plain LogNorm would
    have to leave that 15% off its bar, so the automatic choice is the log-like
    scale that has a place for a zero.
    """
    mostly = np.concatenate([np.zeros(158), np.logspace(-4, 2, 931)])
    assert spans_decades(mostly)
    cmap, norm, why = choose_style(mostly, 'h2oice')
    assert isinstance(norm, SymLogNorm)
    assert norm.vmin <= 0 <= norm.vmax
    assert cmap == SEQUENTIAL          # nothing blanked, so nothing to paint
    assert 'symlog' in why


def test_a_handful_of_zeros_keeps_the_log_scale_and_is_painted():
    """
    Below SYMLOG_SWITCH the zeros are too few to be worth the linear window a
    symlog puts in the middle of the bar, so the log scale stays and those cells
    are given a colour instead.
    """
    barely = np.concatenate([np.zeros(10), np.logspace(-4, 2, 1079)])
    cmap, norm, _ = choose_style(barely, 'h2oice')
    assert isinstance(norm, LogNorm)
    assert cmap._dispnc_blank_note


def test_a_log_scale_with_nothing_to_hide_keeps_the_plain_colormap():
    cmap, norm, _ = choose_style(np.logspace(-4, 2, 1089), 'h2oice')
    assert isinstance(norm, LogNorm)
    assert cmap == SEQUENTIAL


def test_an_explicit_log_scale_is_still_honoured_on_a_sparse_field(capsys):
    """
    Declining to *choose* a log scale is not the same as refusing one that was
    asked for, and a --norm log is never quietly upgraded to a symlog: it keeps
    its LogNorm, says what it could not place, and paints those cells.
    """
    sparse = np.concatenate([np.zeros(1056), np.logspace(-4, 2, 33)])
    cmap, norm, _ = choose_style(sparse, 'co2ice', cmap='viridis', norm='log')
    assert isinstance(norm, LogNorm)
    assert cmap._dispnc_blank_note
    assert 'cannot place 1056 of 1089 cells (97%)' in capsys.readouterr().out