import numpy
import pyarrow
import pytest

from arrakis.block import (
    Freq,
    SeriesBlock,
    Time,
    combine_blocks,
    concatenate_blocks,
    time_as_ns,
)

from ...channel import Channel


@pytest.fixture(scope="module")
def channels():
    return {
        "H1:FKE-TEST_CHANNEL1": Channel(
            "H1:FKE-TEST_CHANNEL1",
            data_type=numpy.dtype(numpy.float64),
            sample_rate=128,
            partition_id="A",
            partition_index=3,
        ),
        "H1:FKE-TEST_CHANNEL2": Channel(
            "H1:FKE-TEST_CHANNEL2",
            data_type=numpy.dtype(numpy.int32),
            sample_rate=32,
            partition_id="B",
            partition_index=5,
        ),
    }


def create_arrays(
    channels,
    duration: int = Time.s // 16,
    gaps: str = "nogap",
):
    arrays = {}
    for channel in channels.values():
        size = duration // (channel.sample_rate * Freq.Hz)
        if gaps == "nogaps":
            array = numpy.random.randint(256, size=size).astype(channel.data_type)
        elif gaps == "gaps":
            mask = numpy.random.randint(2, size=size)
            # assure that there's always at least one gap in the
            # resulting array
            mask[0] = 1
            array = numpy.ma.array(
                numpy.random.randint(256, size=size).astype(channel.data_type),
                mask=mask,
            )
        elif gaps == "fullgap":
            array = numpy.ma.masked_all(size, dtype=channel.data_type)
        else:
            msg = f"unknown gap spec {gaps}"
            raise ValueError(msg)
        arrays[channel.name] = array
    return arrays


@pytest.fixture(scope="module", params=["nogaps", "gaps", "fullgap"])
def arrays(request, channels):
    duration = Time.s // 16
    return create_arrays(channels, duration, request.param)


@pytest.fixture(scope="module")
def block(channels, arrays):
    time = 1187000000 * Time.s
    return SeriesBlock(time, arrays, channels)


##############################


def test_time_as_ns():
    t = time_as_ns(3.14)
    assert t == 3140000000
    assert isinstance(t, int)


def test_block_creation(channels, arrays):
    time_ns = 1187000000 * Time.s
    duration_ns = Time.s // 16
    block = SeriesBlock(time_ns, arrays, channels)

    # cross-check against expected API
    assert block.time_ns == time_ns
    assert block.duration_ns == duration_ns
    for channel, series in block.items():
        assert series.time_ns == time_ns
        assert series.duration_ns == duration_ns
        assert series.data_type == channels[channel].data_type
        assert series.sample_rate == channels[channel].sample_rate
        assert numpy.array_equal(series.data, arrays[channel])
        if isinstance(arrays[channel], numpy.ma.MaskedArray):
            assert all(series.data.mask == arrays[channel].mask)


def test_column_batch_round_trip(block):
    batch = block.to_column_batch()
    rt_block = SeriesBlock.from_column_batch(batch, block.channels)
    assert block.time_ns == rt_block.time_ns
    assert block.duration_ns == rt_block.duration_ns
    assert block.channels == rt_block.channels
    for orig_series, rt_series in zip(block.values(), rt_block.values()):
        assert orig_series.time_ns == rt_series.time_ns
        assert orig_series.duration_ns == rt_series.duration_ns
        assert orig_series.channel == rt_series.channel
        assert numpy.ma.allequal(orig_series.data, rt_series.data)
        assert type(orig_series.data) is type(rt_series.data)
        if isinstance(orig_series.data, numpy.ma.MaskedArray):
            assert type(orig_series.data.mask) is type(rt_series.data.mask)
            assert all(orig_series.data.mask == rt_series.data.mask)


def test_column_batch_schema_round_trip(block):
    # reverse channels in schema
    original_schema = block._generate_column_schema()  # noqa: SLF001
    time_field = original_schema.field(0)  # time field is always first
    channel_fields = [original_schema.field(i) for i in range(1, len(original_schema))]
    reversed_channel_fields = list(reversed(channel_fields))
    schema = pyarrow.schema([time_field, *reversed_channel_fields])

    batch = block.to_column_batch(schema)
    rt_block = SeriesBlock.from_column_batch(batch, block.channels)
    assert block.time_ns == rt_block.time_ns
    assert block.duration_ns == rt_block.duration_ns
    assert block.channels == rt_block.channels
    for channel_name in block.channels:
        orig_series = block[channel_name]
        rt_series = rt_block[channel_name]
        assert orig_series.time_ns == rt_series.time_ns
        assert orig_series.duration_ns == rt_series.duration_ns
        assert orig_series.channel == rt_series.channel
        assert numpy.ma.allequal(orig_series.data, rt_series.data)
        assert type(orig_series.data) is type(rt_series.data)
        if isinstance(orig_series.data, numpy.ma.MaskedArray):
            assert isinstance(orig_series.data.mask, numpy.ndarray)
            assert isinstance(rt_series.data.mask, numpy.ndarray)
            assert type(orig_series.data.mask) is type(rt_series.data.mask)
            assert all(orig_series.data.mask == rt_series.data.mask)


def test_row_batch_round_trip(block):
    index_channel_map = {
        channel.partition_index: channel for channel in block.channels.values()
    }
    (_, batch1), (_, batch2) = block.to_row_batches(block.channels)
    block1 = SeriesBlock.from_row_batch(batch1, index_channel_map)
    block2 = SeriesBlock.from_row_batch(batch2, index_channel_map)
    rt_block = combine_blocks(block1, block2)

    assert block.time_ns == rt_block.time_ns
    assert block.duration_ns == rt_block.duration_ns
    assert block.channels == rt_block.channels
    for orig_series, rt_series in zip(block.values(), rt_block.values()):
        assert orig_series.time_ns == rt_series.time_ns
        assert orig_series.duration_ns == rt_series.duration_ns
        assert orig_series.channel == rt_series.channel
        assert numpy.ma.allequal(orig_series.data, rt_series.data)
        assert type(orig_series.data) is type(rt_series.data)
        if isinstance(orig_series.data, numpy.ma.MaskedArray):
            assert isinstance(orig_series.data.mask, numpy.ndarray)
            assert isinstance(rt_series.data.mask, numpy.ndarray)
            assert type(orig_series.data.mask) is type(rt_series.data.mask)
            assert all(orig_series.data.mask == rt_series.data.mask)


def test_concatenate_blocks(channels, arrays):
    blocks = []
    start_ns = 1187000000 * Time.s
    end_ns = start_ns + 2 * Time.s
    duration_ns = Time.s // 16

    for _i, time_ns in enumerate(range(start_ns, end_ns, duration_ns)):
        block = SeriesBlock(time_ns, arrays, channels)
        blocks.append(block)

    block = concatenate_blocks(*blocks)
    assert block.time_ns == start_ns
    assert block.duration_ns == (end_ns - start_ns)
    assert block.channels == channels
    for channel, series in block.items():
        assert series.time_ns == start_ns
        assert series.duration_ns == (_i + 1) * duration_ns
        assert series.data_type == channels[channel].data_type
        assert series.sample_rate == channels[channel].sample_rate

        ref_series = [block[channel].data for block in blocks]
        if isinstance(arrays[channel], numpy.ma.MaskedArray):
            assert isinstance(series.data, numpy.ma.MaskedArray)
            ref_data = numpy.ma.concatenate(ref_series)
            assert all(series.data.mask == ref_data.mask)
            assert series.data.mask is not False
        else:
            ref_data = numpy.concatenate(ref_series)
            assert all(series.data == ref_data.data)


def test_concatenate_blocks_gaps_nogaps(channels):
    time = 1187000000 * Time.s
    duration = Time.s // 16
    blocks = [
        SeriesBlock(
            time,
            create_arrays(channels, duration, "nogaps"),
            channels,
        ),
        SeriesBlock(
            time + duration,
            create_arrays(channels, duration, "gaps"),
            channels,
        ),
        SeriesBlock(
            time + 2 * duration,
            create_arrays(channels, duration, "fullgap"),
            channels,
        ),
    ]
    block = concatenate_blocks(*blocks)
    for channel, series in block.items():
        assert series.data_type == channels[channel].data_type
        assert series.sample_rate == channels[channel].sample_rate
        assert isinstance(series.data, numpy.ma.MaskedArray)
        assert isinstance(series.data.mask, numpy.ndarray)
