"""
Tests for the select_slice Field method. They focus on consecutive slicing on different axes
to verify that the multiindex is correctly sorted after each slice.
It is also checked that the function throws the exceptions it is supposed to.
"""
import pytest
import re
from canopy.tests.test_data.registry import load_test_data

# TODO: have slices calculated dynamically instead of hardcoded so the tests work for different datasets

def test_x():
    """Slice x-axis consecutively"""

    field = load_test_data('anpp_spain_1990_2010.out.gz')
    x_axis = field.grid.axis_names[0]
    f1 = field.select_slice({x_axis:(-5,0)})
    f1 = field.select_slice({x_axis:(-3,-1)})
    f2 = field.select_slice({x_axis:(-3,-1)})

    assert f1.data.equals(f2.data)


def test_y():
    """Slice y-axis consecutively"""

    field = load_test_data('anpp_spain_1990_2010.out.gz')
    y_axis = field.grid.axis_names[1]
    f1 = field.select_slice({y_axis:(35,40)})
    f1 = field.select_slice({y_axis:(37,39)})
    f2 = field.select_slice({y_axis:(37,39)})

    assert f1.data.equals(f2.data)


def test_time():
    """Slice time-axis consecutively"""

    field = load_test_data('anpp_spain_1990_2010.out.gz')
    f1 = field.select_slice({'time':(1993,2005)})
    f1 = field.select_slice({'time':(1993,2001)})
    f2 = field.select_slice({'time':(1993,2001)})

    assert f1.data.equals(f2.data)


def test_x_y_time():
    """Slice first x, then y, then time"""

    field = load_test_data('anpp_spain_1990_2010.out.gz')
    x_axis, y_axis = field.grid.axis_names
    f_x = field.select_slice({x_axis:(-5,0)})
    f_x_y = f_x.select_slice({y_axis:(35,40)})
    f_x_y_time = f_x_y.select_slice({'time':(1995,2005)})
    f_all = field.select_slice({x_axis:(-5,0), y_axis:(35,40), 'time':(1995,2005)})

    assert f_x_y_time.data.equals(f_all.data)


def test_y_x_time():
    """Slice first y, then x, then time"""

    field = load_test_data('anpp_spain_1990_2010.out.gz')
    x_axis, y_axis = field.grid.axis_names
    f_y = field.select_slice({y_axis:(35,40)})
    f_y_x = f_y.select_slice({x_axis:(-5,0)})
    f_y_x_time = f_y_x.select_slice({'time':(1995,2005)})
    f_all = field.select_slice({x_axis:(-5,0), y_axis:(35,40), 'time':(1995,2005)})

    assert f_y_x_time.data.equals(f_all.data)


def test_x_time_y():
    """Slice first x, then time, then y"""

    field = load_test_data('anpp_spain_1990_2010.out.gz')
    x_axis, y_axis = field.grid.axis_names
    f_x = field.select_slice({x_axis:(-5,0)})
    f_x_time = f_x.select_slice({'time':(1995,2005)})
    f_x_time_y = f_x_time.select_slice({y_axis:(35,40)})
    f_all = field.select_slice({x_axis:(-5,0), y_axis:(35,40), 'time':(1995,2005)})

    assert f_x_time_y.data.equals(f_all.data)


def test_y_time_x():
    """Slice first y, then time, then x"""

    field = load_test_data('anpp_spain_1990_2010.out.gz')
    x_axis, y_axis = field.grid.axis_names
    f_y = field.select_slice({y_axis:(35,40)})
    f_y_time = f_y.select_slice({'time':(1995,2005)})
    f_y_time_x = f_y_time.select_slice({x_axis:(-5,0)})
    f_all = field.select_slice({x_axis:(-5,0), y_axis:(35,40), 'time':(1995,2005)})

    assert f_y_time_x.data.equals(f_all.data)


def test_reduced_x():
    """Slice y and time in an x-reduced grid"""

    field = load_test_data('anpp_spain_1990_2010.out.gz')
    x_axis, y_axis = field.grid.axis_names
    f_x = field.reduce_grid('av', x_axis)
    f_x_sliced = f_x.select_slice({y_axis:(35,40)})
    f_x_sliced = f_x_sliced.select_slice({'time':(1990,2000)})
    f_x_sliced = f_x_sliced.select_slice({y_axis:(37,39)})
    f_x_sliced = f_x_sliced.select_slice({'time':(1993,1995)})
    f_x_all = f_x.select_slice({y_axis:(37,39), 'time':(1993,1995)})

    assert f_x_sliced.data.equals(f_x_all.data)


def test_reduced_y():
    """Slice x and time in an y-reduced grid"""

    field = load_test_data('anpp_spain_1990_2010.out.gz')
    x_axis, y_axis = field.grid.axis_names
    f_y = field.reduce_grid('av', y_axis)
    f_y_sliced = f_y.select_slice({x_axis:(-5,1)})
    f_y_sliced = f_y_sliced.select_slice({'time':(1990,2000)})
    f_y_sliced = f_y_sliced.select_slice({x_axis:(-3,0)})
    f_y_sliced = f_y_sliced.select_slice({'time':(1993,1995)})
    f_y_all = f_y.select_slice({x_axis:(-3,0), 'time':(1993,1995)})

    assert f_y_sliced.data.equals(f_y_all.data)


def test_reduced_both():
    """Slice time consecutively when both axes are reduced"""

    field = load_test_data('anpp_spain_1990_2010.out.gz')
    x_axis, y_axis = field.grid.axis_names
    f_y_x = field.reduce_grid('av')
    f_y_x_time = f_y_x.select_slice({'time':(1995,2005)})
    f_y_x_time = f_y_x_time.select_slice({'time':(1997,2003)})
    f_all = f_y_x.select_slice({'time':(1997,2003)})

    assert f_y_x_time.data.equals(f_all.data)


def test_unrecognized_axis_keys():
    """Pass keys different than grid axes"""

    field = load_test_data('anpp_spain_1990_2010.out.gz')
    x_axis, y_axis = field.grid.axis_names
    bad_axis = 'probably_bad_axis'
    slices = {x_axis:(1,2), bad_axis:(2,3)}

    with pytest.raises(ValueError, match=re.escape(f"Axes {[bad_axis]} not recognized")):
        _ = field.select_slice(slices)


def test_slice_reduced_x():
    """Try to slice reduced x-axis"""

    field = load_test_data('anpp_spain_1990_2010.out.gz')
    x_axis, y_axis = field.grid.axis_names
    f_x = field.reduce_grid('av', x_axis)

    with pytest.raises(ValueError, match=f"Cannot slice reduced axis '{x_axis}'"):
        _ = f_x.select_slice({x_axis:(-3,0)})


def test_slice_reduced_y():
    """Try to slice reduced y-axis"""

    field = load_test_data('anpp_spain_1990_2010.out.gz')
    x_axis, y_axis = field.grid.axis_names
    f_y = field.reduce_grid('av', y_axis)

    with pytest.raises(ValueError, match=f"Cannot slice reduced axis '{y_axis}'"):
        _ = f_y.select_slice({y_axis:(-3,0)})


def test_bad_tuple_1():
    """Pass a tuple with 1 element"""

    field = load_test_data('anpp_spain_1990_2010.out.gz')

    with pytest.raises(ValueError, match=re.escape("Slice tuple for dimension 'time' must have two elements: (start, end).")):
        _ = field.select_slice({'time': (1999,)})


def test_bad_tuple_2():
    """Pass a tuple with >2 elements"""

    field = load_test_data('anpp_spain_1990_2010.out.gz')

    with pytest.raises(ValueError, match=re.escape("Slice tuple for dimension 'time' must have two elements: (start, end).")):
        _ = field.select_slice({'time': (1999, 2000, 2001)})

