"""
Summary statistics, and the weighting that makes the mean mean something.
"""

import numpy as np
import xarray as xr

from dispnc.conventions import normalize
from dispnc.coords import resolve_variable_axes
from dispnc.stats import area_weights, summarize


class Report:
    cell_area = None
    unstructured = None


def test_a_lat_lon_field_is_weighted_by_the_cosine_of_latitude(xios_ds):
    ds, report = normalize(xios_ds)
    axes = [a for a in resolve_variable_axes(ds, 'ps_avg', None)
            if a.dim in ('lat', 'lon')]
    weights, note = area_weights(ds, Report(), axes, (3, 3))
    assert weights is not None and note == 'cos(latitude)'


def test_a_zonal_mean_is_still_weighted(xios_ds):
    """
    The bug: after --reduce mean:lon the field is one dimensional, and the
    weights were dropped with a note claiming there was no latitude axis. An
    unweighted average down a latitude profile is dominated by the poles, which
    is the exact error this module exists to prevent.
    """
    ds, report = normalize(xios_ds)
    lat = next(a for a in resolve_variable_axes(ds, 'ps_avg', None)
               if a.dim == 'lat')
    weights, note = area_weights(ds, Report(), [lat], (lat.size,),
                                 reduced=['lon'])
    assert weights is not None, 'a latitude profile is still weightable'
    assert note == 'cos(latitude)'
    assert weights.shape == (lat.size,)


def test_a_reduction_the_weights_cannot_survive_says_which_one(xios_ds):
    ds, report = normalize(xios_ds)
    weights, note = area_weights(ds, Report(), [], (4,), reduced=['lat'])
    assert weights is None
    assert 'do not survive' in note and 'lat' in note


def test_the_weighted_mean_differs_from_the_plain_one():
    """
    If it did not, the weighting would not be worth doing.
    """
    lats = np.linspace(-89, 89, 30)
    weights = np.cos(np.deg2rad(lats))
    # Cold at both poles and warm at the equator, which is the case weighting
    # exists for: the poles are most of the cells and least of the area.
    values = 200 + 60 * np.cos(np.deg2rad(lats))
    stats = summarize(values, weights)
    assert 'weighted_mean' in stats
    assert not np.isclose(stats['weighted_mean'], stats['mean'])


def test_missing_cells_leave_the_denominator_too():
    """
    The classic way to get a weighted mean wrong.
    """
    values = np.array([1.0, np.nan, 3.0])
    weights = np.array([1.0, 100.0, 1.0])
    stats = summarize(values, weights)
    assert np.isclose(stats['weighted_mean'], 2.0)