"""
Derived fields: differences between variables, and what is refused.
"""

import numpy as np
import pytest
import xarray as xr

from dispnc.derive import DeriveError, apply_difference


def field(units, values=None, dims=('Time', 'lat')):
    data = np.arange(12.0).reshape(3, 4) if values is None else values
    return xr.DataArray(data, dims=dims, attrs={'units': units})


def test_a_difference_is_the_two_fields_subtracted():
    left = field('K', np.full((3, 4), 10.0))
    right = field('K', np.full((3, 4), 4.0))
    out, said = apply_difference(left, right, 'tsurf', 'tsoil')
    assert np.allclose(out.values, 6.0)
    assert said == "'tsurf' minus 'tsoil'"


def test_a_difference_is_marked_as_signed_so_it_lands_centred_on_zero():
    out, _ = apply_difference(field('K'), field('K'), 'a', 'b')
    assert out.attrs['dispnc_signed'] == 1


def test_a_difference_says_on_the_figure_what_it_is():
    """
    A saved map titled 'h2o_ice' that is really h2o_ice minus co2_ice outlives
    the terminal it was described in.
    """
    left = field('K')
    left.attrs['long_name'] = 'Surface temperature'
    out, _ = apply_difference(left, field('K'), 'tsurf', 'tsoil')
    assert out.attrs['long_name'] == 'Surface temperature - tsoil'


def test_subtracting_across_units_is_refused():
    """
    Pa from K produces a number, which is exactly the problem.
    """
    with pytest.raises(DeriveError) as err:
        apply_difference(field('K'), field('Pa'), 'tsurf', 'ps')
    assert 'K' in str(err.value) and 'Pa' in str(err.value)


def test_subtracting_across_grids_is_refused():
    with pytest.raises(DeriveError) as err:
        apply_difference(field('K'), field('K', dims=('Time', 'rlonu')),
                         'tsurf', 'cu')
    assert 'rlonu' in str(err.value)


def test_two_variables_with_no_units_at_all_may_still_be_subtracted():
    out, _ = apply_difference(field(None), field(None), 'a', 'b')
    assert out.shape == (3, 4)


# --- reductions ---------------------------------------------------------------

@pytest.mark.parametrize('op, reference', [
    ('mean', np.nanmean),
    ('std', np.nanstd),
    ('min', np.nanmin),
    ('max', np.nanmax),
    ('sum', np.nansum),
])
def test_a_reduction_gives_what_the_numpy_one_gave(op, reference):
    """
    The decision table, and the guarantee that moving this off numpy and onto
    xarray moved no number. xarray's `std` is ddof=0 like `np.nanstd`, and its
    `sum(skipna=True)` of an all-NaN slice is 0.0 like `np.nansum`.
    """
    from dispnc.derive import apply_reduce, parse_reduce

    values = np.where(np.random.default_rng(1).random((4, 5)) < 0.2, np.nan,
                      np.random.default_rng(2).random((4, 5)) * 100)
    da = xr.DataArray(values, dims=('a', 'b'))

    out, note, reduced = apply_reduce(da, parse_reduce([f'{op}:a'], ['a', 'b']))
    assert np.allclose(out.values, reference(values, axis=0))
    assert note == f'{op} over a'
    assert reduced == ['a']


def test_a_large_reduction_stays_lazy_through_dask():
    """
    The defect: `--reduce mean:time` used to run on an array that had already
    been read whole, where the `-e '{"time": "avg"}'` spelling of the same thing
    held one chunk. Two ways to say one thing, and only one survived a big file.
    """
    from dispnc.derive import apply_reduce, parse_reduce
    from dispnc.io import CHUNK_THRESHOLD_BYTES

    shape = (2400, 64, 64)
    assert np.prod(shape) * 8 > CHUNK_THRESHOLD_BYTES, 'has to be worth chunking'
    big = xr.DataArray(np.zeros(shape), dims=('time', 'lat', 'lon'))

    out, _, _ = apply_reduce(big, parse_reduce(['mean:time'], list(big.dims)))
    assert hasattr(out.data, 'compute'), 'the reduction was not streamed'


def test_a_small_reduction_keeps_the_eager_path_it_had():
    """
    `maybe_chunk` no-ops below its threshold, so a small file pays no dask
    overhead and nothing about its figures moves.
    """
    from dispnc.derive import apply_reduce, parse_reduce

    small = xr.DataArray(np.zeros((4, 5)), dims=('a', 'b'))
    out, _, _ = apply_reduce(small, parse_reduce(['mean:a'], ['a', 'b']))
    assert not hasattr(out.data, 'compute')


def test_a_reduction_over_a_dimension_the_variable_lacks_is_skipped():
    from dispnc.derive import apply_reduce

    da = xr.DataArray(np.zeros((4, 5)), dims=('a', 'b'))
    out, note, reduced = apply_reduce(da, [('mean', ['nope'])])
    assert out.shape == (4, 5) and note == '' and reduced == []


def test_reductions_run_after_the_selection_they_are_combined_with():
    """
    `-e '{"Time": "avg"}' --reduce max:lon` is the maximum of the mean, not the
    mean of the maxima. The two are different numbers.
    """
    from dispnc.pipeline import _slice_and_reduce
    from dispnc.derive import parse_reduce

    # Deliberately irregular: on a linear ramp the two orders coincide, and a
    # test that cannot tell them apart is not testing the order.
    values = np.random.default_rng(3).random((2, 3, 4)) * 100
    da = xr.DataArray(values, dims=('Time', 'lat', 'lon'))
    reductions = parse_reduce(['max:lon'], list(da.dims))

    got, dims, _, reduced = _slice_and_reduce(da, {'Time': 'avg'}, reductions)
    assert np.allclose(got, values.mean(axis=0).max(axis=1))
    assert not np.allclose(got, values.max(axis=2).mean(axis=0))
    assert dims == ['lat'] and reduced == ['lon']