import pytest
import sys, os
import xarray as xr
import numpy as np

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import process
from process._common import ProcessArgumentInvalid, ProcessArgumentRequired


@pytest.fixture
def generate_data():
    def _construct(
            data = [[[[0.1, 0.15], [0.15, 0.2]], [[0.05, 0.1], [-0.9, 0.05]]]],
            dims = ('t','y','x','band'),
            reduce_by = "band",
            as_list = False
        ):
        if as_list:
            return data

        xrdata = xr.DataArray(
            data,
            dims=dims,
            attrs={'reduce_by': [reduce_by]},
        )
        return xrdata
    return _construct


@pytest.fixture
def execute_array_element_process(generate_data):
    def wrapped(data_arguments={}, index=None, return_nodata=None):
        arguments = {}
        if data_arguments is not None: arguments["data"] = generate_data(**data_arguments)
        if index is not None: arguments["index"] = index
        if return_nodata is not None: arguments["return_nodata"] = return_nodata

        return process.array_element.array_elementEOTask(None, "" , None, {}, "arrayel1").process(arguments)
    return wrapped


###################################
# tests:
###################################

@pytest.mark.parametrize('data,return_nodata,index,expected_result', [
    ([9,8,7,6,5], None, 2, 7),
    (["A","B","C"], None, 0, "A"),
    ([], True, 0, None)
])
def test_examples(execute_array_element_process, data, index, return_nodata, expected_result):
    """
        Test array_element process with examples from https://open-eo.github.io/openeo-api/processreference/#array_element
    """
    data_arguments = {"data": data, "as_list": True}
    result = execute_array_element_process(data_arguments=data_arguments, index=index, return_nodata=return_nodata)
    assert result == expected_result


@pytest.mark.parametrize('data,index,reduce_by,expected_data,expected_dims', [
    ([[[[0.1, 0.15], [0.15, 0.2]], [[0.05, 0.1], [-0.9, 0.05]]]], 0, "band", [[[0.1, 0.15], [0.05, -0.9]]], ('t','y','x')),
    ([[[[0.1, 0.15], [0.15, 0.2]], [[0.05, 0.1], [-0.9, 0.05]]]], 1, "y", [[[0.05, 0.1], [-0.9, 0.05]]], ('t','x','band')),
])
def test_with_xarray(execute_array_element_process, generate_data, data, index, reduce_by, expected_data, expected_dims):
    """
        Test array_element process with xarray.DataArrays
    """
    expected_result = generate_data(data=expected_data, dims=expected_dims, reduce_by=reduce_by)
    result = execute_array_element_process(data_arguments={"data": data, "reduce_by": reduce_by}, index=index)
    xr.testing.assert_allclose(result, expected_result)


def test_with_xarray_out_bounds(execute_array_element_process, generate_data):
    """
        Test array_element process with xarray.DataArrays with out of bounds index
    """
    with pytest.raises(ProcessArgumentInvalid) as ex:
        result = execute_array_element_process(index=5)
    assert ex.value.args[0] == "The argument 'index' in process 'array_element' is invalid: Index out of bounds."


@pytest.mark.parametrize('data_arguments,index,expected_data,expected_dims', [
    ({}, 5, [[[np.nan, np.nan], [np.nan, np.nan]]], ('t','y','x')),
])
def test_with_xarray_out_bounds_return_nodata(execute_array_element_process, generate_data, data_arguments, index, expected_data, expected_dims):
    """
        Test array_element process with xarray.DataArrays with out of bounds index and return_no_data
    """
    expected_result = generate_data(expected_data, dims=expected_dims)
    result = execute_array_element_process(data_arguments=data_arguments, index=index, return_nodata=True)
    xr.testing.assert_equal(result, expected_result)

[((47, 1, 51, 2), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(47, 25, 47, 67): '"""data,return_nodata,index,expected_result"""', (47, 69, 51, 1): "[([9, 8, 7, 6, 5], None, 2, 7), (['A', 'B', 'C'], None, 0, 'A'), ([], True,\n    0, None)]"}, {}), "('data,return_nodata,index,expected_result', [([9, 8,\n    7, 6, 5], None, 2, 7), (['A', 'B', 'C'], None, 0, 'A'), ([], True, 0,\n    None)])", False, 'import pytest\n'), ((61, 1, 64, 2), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(61, 25, 61, 75): '"""data,index,reduce_by,expected_data,expected_dims"""', (61, 77, 64, 1): "[([[[[0.1, 0.15], [0.15, 0.2]], [[0.05, 0.1], [-0.9, 0.05]]]], 0, 'band', [\n    [[0.1, 0.15], [0.05, -0.9]]], ('t', 'y', 'x')), ([[[[0.1, 0.15], [0.15,\n    0.2]], [[0.05, 0.1], [-0.9, 0.05]]]], 1, 'y', [[[0.05, 0.1], [-0.9, \n    0.05]]], ('t', 'x', 'band'))]"}, {}), "('data,index,reduce_by,expected_data,expected_dims',\n    [([[[[0.1, 0.15], [0.15, 0.2]], [[0.05, 0.1], [-0.9, 0.05]]]], 0,\n    'band', [[[0.1, 0.15], [0.05, -0.9]]], ('t', 'y', 'x')), ([[[[0.1, 0.15\n    ], [0.15, 0.2]], [[0.05, 0.1], [-0.9, 0.05]]]], 1, 'y', [[[0.05, 0.1],\n    [-0.9, 0.05]]], ('t', 'x', 'band'))])", False, 'import pytest\n'), ((83, 1, 85, 2), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(83, 25, 83, 75): '"""data_arguments,index,expected_data,expected_dims"""', (83, 77, 85, 1): "[({}, 5, [[[np.nan, np.nan], [np.nan, np.nan]]], ('t', 'y', 'x'))]"}, {}), "('data_arguments,index,expected_data,expected_dims',\n    [({}, 5, [[[np.nan, np.nan], [np.nan, np.nan]]], ('t', 'y', 'x'))])", False, 'import pytest\n'), ((71, 4, 71, 55), 'xarray.testing.assert_allclose', 'xr.testing.assert_allclose', ({(71, 31, 71, 37): 'result', (71, 39, 71, 54): 'expected_result'}, {}), '(result, expected_result)', True, 'import xarray as xr\n'), ((92, 4, 92, 52), 'xarray.testing.assert_equal', 'xr.testing.assert_equal', ({(92, 28, 92, 34): 'result', (92, 36, 92, 51): 'expected_result'}, {}), '(result, expected_result)', True, 'import xarray as xr\n'), ((22, 17, 26, 9), 'xarray.DataArray', 'xr.DataArray', (), '', True, 'import xarray as xr\n'), ((78, 9, 78, 46), 'pytest.raises', 'pytest.raises', ({(78, 23, 78, 45): 'ProcessArgumentInvalid'}, {}), '(ProcessArgumentInvalid)', False, 'import pytest\n'), ((6, 48, 6, 73), 'os.path.abspath', 'os.path.abspath', ({(6, 64, 6, 72): '__file__'}, {}), '(__file__)', False, 'import sys, os\n'), ((39, 15, 39, 89), 'process.array_element.array_elementEOTask', 'process.array_element.array_elementEOTask', ({(39, 57, 39, 61): 'None', (39, 63, 39, 65): '""""""', (39, 68, 39, 72): 'None', (39, 74, 39, 76): '{}', (39, 78, 39, 88): '"""arrayel1"""'}, {}), "(None, '', None, {}, 'arrayel1')", False, 'import process\n')]