# Copyright (c) 2023 CNES
#
# All rights reserved. Use of this source code is governed by a
# BSD-style license that can be found in the LICENSE file.
"""
Test partitioning by sequence.
==============================
"""
from __future__ import annotations

from typing import Iterator
import pickle

import dask.array.core
import numpy
import pytest
import xarray

from zcollection.type_hints import ArrayLike

from . import data
from .. import Sequence, get_codecs
from ... import dataset
# pylint: disable=unused-import # Need to import for fixtures
from ...tests.cluster import dask_client, dask_cluster

# pylint: enable=unused-import # Need to import for fixtures


def test_construction() -> None:
    """Test the sequence constructor."""
    assert isinstance(Sequence(('a', 'b')), Sequence)
    assert len(Sequence(('a', 'b'))) == 2
    with pytest.raises(ValueError):
        Sequence(('a', 'b'), (0, ))  # type: ignore
    with pytest.raises(ValueError):
        Sequence((), ())
    with pytest.raises(ValueError):
        Sequence(('a', 'b'), dtype=('c', 'd'))
    with pytest.raises(ValueError):
        Sequence(('a', 'b'), dtype=('float32', 'int32'))
    with pytest.raises(TypeError):
        Sequence(('a', 'b'), dtype='int32')
    partitioning = Sequence(('a', 'b'))
    partition_keys = partitioning.parse('a=1/b=2')
    assert partitioning.encode(partition_keys) == (1, 2)
    with pytest.raises(ValueError):
        partitioning.encode((('A', 1), ('b', 2)))
    assert partitioning.decode((1, 2)) == (('a', 1), ('b', 2))
    assert partition_keys == (('a', 1), ('b', 2))
    with pytest.raises(ValueError):
        partitioning.parse('a=1/b=2/c=3')
    with pytest.raises(ValueError):
        partitioning.parse('field=1')


@pytest.mark.parametrize('delayed', [False, True])
def test_split_dataset(
    dask_client,  # pylint: disable=redefined-outer-name,unused-argument
    delayed: bool,
) -> None:
    """Test the split_dataset method."""
    repeatability = 5
    xds = data.create_test_sequence(repeatability, 20, 10)
    partitioning = Sequence(('cycle_number', 'pass_number'))

    cycle_number = 1
    pass_number = 1

    assert partitioning.dtype() == (
        ('cycle_number', 'int64'),
        ('pass_number', 'int64'),
    )

    # Build the test dataset
    zds = dataset.Dataset.from_xarray(xds)
    if not delayed:
        zds = zds.compute()

    iterator = partitioning.split_dataset(zds, 'num_lines')
    assert isinstance(iterator, Iterator)

    for partition, indexer in iterator:
        subset = zds.isel(indexer)
        expected = (f'cycle_number={cycle_number}',
                    f'pass_number={pass_number}')
        assert expected == partition
        assert numpy.all(
            xds.where((xds.cycle_number == cycle_number)
                      & (xds.pass_number == pass_number),
                      drop=True).observation ==
            subset.variables['observation'].array)

        partition_keys = partitioning.parse('/'.join(partition))
        assert partition_keys == (('cycle_number', cycle_number),
                                  ('pass_number', pass_number))
        assert partitioning.decode(
            partitioning.encode(partition_keys)) == partition_keys
        assert partitioning.join(partition_keys, '/') == '/'.join(partition)

        pass_number += 1
        if pass_number > repeatability:
            pass_number = 1
            cycle_number += 1

    xds['cycle_number'] = xarray.DataArray(numpy.array(
        [xds['cycle_number'].values] * 2).T,
                                           dims=('num_lines', 'nump_pixels'))
    zds = dataset.Dataset.from_xarray(xds)
    if not delayed:
        zds = zds.compute()
    with pytest.raises(ValueError):
        list(partitioning.split_dataset(zds, 'num_lines'))


def test_config() -> None:
    """Test the configuration of the Sequence class."""
    partitioning = Sequence(('cycle_number', 'pass_number'))
    config = partitioning.get_config()
    partitioning = get_codecs(config)  # type: ignore[assignment]
    assert isinstance(partitioning, Sequence)


def test_pickle() -> None:
    """Test the pickling of the Date class."""
    partitioning = Sequence(('cycle_number', 'pass_number'))
    other = pickle.loads(pickle.dumps(partitioning))
    assert isinstance(other, Sequence)
    assert other.variables == ('cycle_number', 'pass_number')


# pylint: disable=protected-access
@pytest.mark.parametrize('delayed', [False, True])
def test_multiple_sequence(
    dask_client,  # pylint: disable=redefined-outer-name,unused-argument
    delayed: bool,
) -> None:
    """Test the creation of a sequence with multiple variables."""
    arrays = {
        '_a': numpy.array([], dtype='i8'),
        '_b': numpy.array([], dtype='i8'),
        '_c': numpy.array([], dtype='i8')
    }
    for _a in range(5):
        for _b in range(5):
            arrays['_a'] = numpy.concatenate(
                (arrays['_a'], numpy.full((5, ), _a, dtype='i8')))
            arrays['_b'] = numpy.concatenate(
                (arrays['_b'], numpy.full((5, ), _b, dtype='i8')))
            arrays['_c'] = numpy.concatenate(
                (arrays['_c'], numpy.arange(5, dtype='i8')))
    partitioning = Sequence(('_a', '_b', '_c'))
    chunks: str = (10, )  # type: ignore[assignment]
    if delayed:
        variables: dict[str, ArrayLike] = {  # type: ignore[assignment]
            '_a': dask.array.core.from_array(arrays['_a'], chunks=chunks),
            '_b': dask.array.core.from_array(arrays['_b'], chunks=chunks),
            '_c': dask.array.core.from_array(arrays['_c'], chunks=chunks)
        }
    else:
        variables = arrays  # type: ignore[assignment]
    _a = 0
    _b = 0
    _c = 0
    for idx, item in enumerate(
            partitioning._split(variables)):  # type: ignore[arg-type]
        assert item[0] == (('_a', _a), ('_b', _b), ('_c', _c))
        _c += 1
        if _c > 4:
            _c = 0
            _b += 1
        if _b > 4:
            _b = 0
            _a += 1
        assert item[1] == slice(idx, idx + 1)

    numpy.random.shuffle(arrays['_c'])
    variables['_c'] = dask.array.core.from_array(  # type: ignore[assignment]
        arrays['_c'], chunks=chunks) if delayed else arrays['_c']

    with pytest.raises(ValueError):
        list(partitioning._split(variables))  # type: ignore[arg-type]

    del variables['_c']
    del variables['_b']
    partitioning = Sequence(('_a', '_b', '_c'))

    _a = 0
    for idx, item in enumerate(
            partitioning._split(variables)):  # type: ignore[arg-type]
        assert item[0] == (('_a', _a), )
        _a += 1
        assert item[1] == slice(idx * 25, idx * 25 + 25)
    # pylint: enable=protected-access


def test_values_must_be_integer() -> None:
    """Test that the values must be integer."""
    values = numpy.arange(0, 100, dtype='f8')
    partitioning = Sequence(('values', ))
    # pylint: disable=protected-access
    with pytest.raises(TypeError):
        list(partitioning._split({'values': values}))
    # pylint: enable=protected-access
