"""Test IterData, which makes iterables behave like Numpy structured arrays."""

import numpy as np
import pytest

from pydap.handlers.lib import IterData
from pydap.model import BaseType, DatasetType, SequenceType


@pytest.fixture
def simple_array():
    arr = np.array(
        [
            (0, 1, 10.0),
            (1, 2, 20.0),
            (2, 3, 30.0),
            (3, 4, 40.0),
            (4, 5, 50.0),
            (5, 6, 60.0),
            (6, 7, 70.0),
            (7, 8, 80.0),
        ],
        dtype=[("byte", "b"), ("int", "i4"), ("float", "f4")],
    )
    return arr


@pytest.fixture
def simple_object(simple_array):
    # add sequence and children for each column
    name = "nameless"
    dataset = DatasetType(name)
    seq = dataset["sequence"] = SequenceType("sequence")
    for var in simple_array.dtype.names:
        seq[var] = BaseType(var)

    obj = IterData(
        [
            (0, 1, 10.0),
            (1, 2, 20.0),
            (2, 3, 30.0),
            (3, 4, 40.0),
            (4, 5, 50.0),
            (5, 6, 60.0),
            (6, 7, 70.0),
            (7, 8, 80.0),
        ],
        seq,
    )
    return obj


def test_iter(simple_array, simple_object):
    assert [tuple(row) for row in simple_array] == [
        (0, 1, 10.0),
        (1, 2, 20.0),
        (2, 3, 30.0),
        (3, 4, 40.0),
        (4, 5, 50.0),
        (5, 6, 60.0),
        (6, 7, 70.0),
        (7, 8, 80.0),
    ]
    assert list(simple_array["byte"]) == [0, 1, 2, 3, 4, 5, 6, 7]

    assert [tuple(row) for row in simple_object] == [
        (0, 1, 10.0),
        (1, 2, 20.0),
        (2, 3, 30.0),
        (3, 4, 40.0),
        (4, 5, 50.0),
        (5, 6, 60.0),
        (6, 7, 70.0),
        (7, 8, 80.0),
    ]
    assert list(simple_object["byte"]) == [0, 1, 2, 3, 4, 5, 6, 7]


def test_dtype(simple_array, simple_object):
    assert simple_array.dtype == np.dtype(
        [("byte", "i1"), ("int", "<i4"), ("float", "<f4")]
    )
    assert simple_array["int"].dtype == "<i4"

    assert simple_object.dtype == np.dtype(
        [("byte", "<i8"), ("int", "<i8"), ("float", "<f8")]
    )
    assert simple_object["int"].dtype == "<i8"


def test_selection(simple_array, simple_object):
    selection = simple_array[simple_array["byte"] > 3]
    assert [tuple(row) for row in selection] == [
        (4, 5, 50.0),
        (5, 6, 60.0),
        (6, 7, 70.0),
        (7, 8, 80.0),
    ]

    selection = simple_object[simple_object["byte"] > 3]
    assert [tuple(row) for row in selection] == [
        (4, 5, 50.0),
        (5, 6, 60.0),
        (6, 7, 70.0),
        (7, 8, 80.0),
    ]


def test_projection(simple_array, simple_object):
    projection = simple_array[1::2]
    assert [tuple(row) for row in projection] == [
        (1, 2, 20.0),
        (3, 4, 40.0),
        (5, 6, 60.0),
        (7, 8, 80.0),
    ]
    assert list(simple_array["byte"][1::2]) == [1, 3, 5, 7]
    assert list(simple_array[1::2]["byte"]) == [1, 3, 5, 7]

    projection = simple_object[1::2]
    assert [tuple(row) for row in projection] == [
        (1, 2, 20.0),
        (3, 4, 40.0),
        (5, 6, 60.0),
        (7, 8, 80.0),
    ]
    assert list(simple_object["byte"][1::2]) == [1, 3, 5, 7]
    assert list(simple_object[1::2]["byte"]) == [1, 3, 5, 7]


def test_combined(simple_array, simple_object):
    filtered = simple_array[simple_array["byte"] > 1]
    filtered = filtered[filtered["byte"] < 6]
    filtered = filtered[::2]
    assert [tuple(row) for row in filtered] == [(2, 3, 30.0), (4, 5, 50.0)]

    filtered = simple_object[simple_object["byte"] > 1]
    filtered = filtered[filtered["byte"] < 6]
    filtered = filtered[::2]
    assert [tuple(row) for row in filtered] == [(2, 3, 30.0), (4, 5, 50.0)]


def test_combined_other(simple_array, simple_object):
    filtered = simple_array[["int", "float"]][simple_array["byte"] < 2]
    assert [tuple(row) for row in filtered] == [(1, 10.0), (2, 20.0)]

    filtered = simple_object[["int", "float"]][simple_object["byte"] < 2]
    assert [tuple(row) for row in filtered] == [(1, 10.0), (2, 20.0)]


@pytest.fixture
def nested_data():
    shallow_data = [(1, 1, 1), (2, 4, 4), (3, 6, 9), (4, 8, 16)]
    deep_data = [
        [(10, 11, 12), (21, 22, 23)],
        [(15, 16, 17)],
        [],
        [(31, 32, 33), (41, 42, 43), (51, 52, 53), (61, 62, 63)],
    ]
    nested = [x + (deep_data[x_id],) for x_id, x in enumerate(shallow_data)]
    return nested


@pytest.fixture
def nested_dtype():
    dtype = np.dtype(
        [
            ("a", "<i8"),
            ("b", "<i8"),
            ("c", "<i8"),
            ("d", np.dtype([("e", "<i8"), ("f", "<i8"), ("g", "<i8")])),
        ]
    )
    return dtype


@pytest.fixture
def nested_object(nested_data):
    name = "nameless"
    dataset = DatasetType(name)
    seq = dataset["nested"] = SequenceType("nested")
    for var in ["a", "b", "c"]:
        seq[var] = BaseType(var)
    seq["d"] = SequenceType("d")
    for var in ["e", "f", "g"]:
        seq["d"][var] = BaseType(var)

    nested = IterData(nested_data, seq)
    return nested


def test_nested_iter(nested_data, nested_object):
    assert [tuple(row) for row in nested_object] == nested_data


def test_nested_iter_child(nested_data, nested_object):
    expected = [row[0] for row in nested_data]
    assert list(nested_object["a"]) == expected


def test_nested_iter_nested_sequence(nested_data, nested_object):
    expected = [row[3] for row in nested_data]
    assert list(nested_object["d"]) == expected


def test_nested_iter_nested_deep_child(nested_data, nested_object):
    expected = [[col[0] for col in row[3]] for row in nested_data]
    assert list(nested_object["d"]["e"]) == expected


def test_nested_dtype(nested_dtype, nested_object):
    assert nested_object.dtype == nested_dtype


def test_nested_selection(nested_data, nested_object):
    selection = nested_object[nested_object["a"] > 2]
    expected = [tuple(row) for row in nested_data if row[0] > 2]
    assert [tuple(row) for row in selection] == expected


def test_nested_projection(nested_data, nested_object):
    projection = nested_data[1::2]
    expected = [
        tuple(row)
        for row_id, row in enumerate(nested_data)
        if row_id in range(1, len(nested_data), 2)
    ]
    assert [tuple(row) for row in projection] == expected

    projection = nested_object[1::2]
    expected = [
        tuple(row)
        for row_id, row in enumerate(nested_data)
        if row_id in range(1, len(nested_data), 2)
    ]
    assert [tuple(row) for row in projection] == expected


def test_nested_combined(nested_data, nested_object):
    filtered = [tuple(row) for row in nested_data if row[0] > 2]
    filtered = [tuple(row) for row in filtered if row[0] < 4]
    filtered = filtered[::2]
    assert [tuple(row) for row in filtered] == [(3, 6, 9, [])]

    filtered = nested_object[nested_object["a"] > 2]
    filtered = filtered[filtered["a"] < 4]
    filtered = filtered[::2]
    assert [tuple(row) for row in filtered] == [(3, 6, 9, [])]
