piker/tests/test_tsp_analysis.py

'''
Deterministic time-series analysis regressions.

'''
import numpy as np
import polars as pl
import pytest

from piker.data._source import def_iohlcv_fields
from piker.tsp import _anal
from piker.tsp._anal import (
    dedupe,
    detect_time_gaps,
    detect_time_ordering_errors,
    get_null_segs,
    iter_null_segs,
    np2pl,
    with_dts,
)
from piker.tsp._dedupe_smart import dedupe_ohlcv_smart


def mk_ohlcv(
    times: list[int],
    indexes: list[int]|None = None,
) -> np.ndarray:
    '''
    Build a synthetic OHLCV frame from timestamps.

    '''
    size: int = len(times)
    frame = np.zeros(
        size,
        dtype=np.dtype(def_iohlcv_fields),
    )
    frame['index'] = (
        np.arange(size)
        if indexes is None
        else indexes
    )
    frame['time'] = times
    frame['open'] = np.arange(size) + 10
    frame['high'] = frame['open'] + 2
    frame['low'] = frame['open'] - 2
    frame['close'] = frame['open'] + 1
    frame['volume'] = np.arange(size) + 100
    return frame


def expected_null_segments(
    null_positions: list[int],
    size: int,
    offset: int,
    margin: int,

) -> list[list[int]]:
    '''
    Group relative null positions with a simple state machine.

    '''
    runs: list[list[int]] = []
    for pos in null_positions:
        if (
            not runs
            or
            pos != runs[-1][1] + 1
        ):
            runs.append([pos, pos])
        else:
            runs[-1][1] = pos

    frame_end: int = offset + size - 1
    return [
        [
            max(offset, offset + start - margin),
            min(frame_end, offset + end + margin),
        ]
        for start, end in runs
    ]


def test_dedupe_recomputes_dts_after_sort() -> None:
    '''
    Derived columns describe final sorted, unique rows.

    '''
    src: pl.DataFrame = np2pl(mk_ohlcv([
        60,
        180,
        120,
        180,
    ]))

    _, deduped, diff = dedupe(src)

    assert diff == 1
    assert deduped['time'].to_list() == [60, 120, 180]
    assert deduped['time_prev'].to_list() == [None, 60, 120]
    assert deduped['s_diff'].to_list() == [None, 60, 60]
    assert detect_time_gaps(
        deduped,
        expect_period=60,
    ).is_empty()


def test_negative_delta_is_an_ordering_error(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    '''
    Reversed timestamps are reported as ordering faults, not gaps.

    '''
    src = np2pl(mk_ohlcv([120, 60]))
    wdts: pl.DataFrame = with_dts(src)
    errors = detect_time_ordering_errors(wdts)

    assert detect_time_gaps(
        wdts,
        expect_period=30,
    ).is_empty()
    assert errors['time_prev'].to_list() == [120]
    assert errors['time'].to_list() == [60]
    assert errors['s_diff'].to_list() == [-60]

    warnings: list[str] = []
    monkeypatch.setattr(_anal.log, 'warning', warnings.append)
    dedupe(src)
    assert len(warnings) == 1
    assert 'non-positive timestamp step' in warnings[0]


def test_single_null_row_is_a_segment() -> None:
    '''
    A one-row null segment remains repairable.

    '''
    frame: np.ndarray = mk_ohlcv([
        60,
        120,
        0,
        240,
        300,
    ])

    nulls = get_null_segs(frame, period=60)

    assert nulls is not None
    segments, zero_indexes, zero_rows = nulls
    assert segments == [[1, 3]]
    assert zero_indexes.tolist() == [2]
    assert zero_rows['index'].tolist() == [2]


@pytest.mark.parametrize(
    ('times', 'expected'),
    [
        ([0, 0, 180, 240, 300, 360], [[10, 12]]),
        ([60, 0, 0, 240, 0, 360], [[10, 13], [13, 15]]),
        ([60, 120, 180, 240, 0, 0], [[13, 15]]),
    ],
)
def test_null_segment_boundaries_match_numpy_and_polars(
    times: list[int],
    expected: list[list[int]],
) -> None:
    '''
    Segments have exact, inclusive, margin-expanded endpoints.

    '''
    frame: np.ndarray = mk_ohlcv(
        times,
        indexes=list(range(10, 16)),
    )

    np_nulls = get_null_segs(frame, period=60)
    pl_nulls = get_null_segs(np2pl(frame), period=60)

    assert np_nulls is not None
    assert pl_nulls is not None
    assert np_nulls[0] == expected
    assert pl_nulls[0] == expected
    assert np_nulls[1].tolist() == pl_nulls[1].to_list()


def test_null_segment_grouping_exhaustive() -> None:
    '''
    Every null layout through six rows matches an independent oracle.

    '''
    offset: int = 10
    for size in range(1, 7):
        for mask in range(1, 1 << size):
            null_positions: list[int] = [
                pos
                for pos in range(size)
                if mask & (1 << pos)
            ]
            frame = mk_ohlcv(
                [
                    0
                    if pos in null_positions
                    else (pos + 1) * 60
                    for pos in range(size)
                ],
                indexes=list(range(offset, offset + size)),
            )

            for margin in range(3):
                expected = expected_null_segments(
                    null_positions,
                    size,
                    offset,
                    margin,
                )
                for candidate in (frame, np2pl(frame)):
                    nulls = get_null_segs(
                        candidate,
                        period=60,
                        imargin=margin,
                    )
                    assert nulls is not None
                    assert nulls[0] == expected


def test_null_segment_margin_and_column() -> None:
    '''
    The requested column and index margin define each segment.

    '''
    frame: np.ndarray = mk_ohlcv(
        [60, 120, 180, 240, 300],
        indexes=[10, 11, 12, 13, 14],
    )
    frame['volume'][1:3] = 0

    for candidate in (frame, np2pl(frame)):
        no_margin = get_null_segs(
            candidate,
            period=60,
            imargin=0,
            col='volume',
        )
        wide_margin = get_null_segs(
            candidate,
            period=60,
            imargin=2,
            col='volume',
        )

        assert no_margin is not None
        assert wide_margin is not None
        assert no_margin[0] == [[11, 12]]
        assert wide_margin[0] == [[10, 14]]
        assert list(no_margin[1]) == [11, 12]


def test_null_segment_input_invariants() -> None:
    '''
    Empty results and invalid index/margin inputs are explicit.

    '''
    frame = mk_ohlcv([60, 120, 180])
    assert get_null_segs(frame, period=60) is None

    with pytest.raises(ValueError, match='imargin'):
        get_null_segs(
            mk_ohlcv([60, 0, 180]),
            period=60,
            imargin=-1,
        )

    with pytest.raises(ValueError, match='contiguous'):
        get_null_segs(
            mk_ohlcv(
                [60, 0, 180],
                indexes=[10, 12, 13],
            ),
            period=60,
        )


def test_iter_null_segs_uses_precomputed_segments(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    '''
    Iteration preserves inclusive boundary rows without rescanning.

    '''
    frame = mk_ohlcv([60, 120, 0, 240, 300])
    nulls = get_null_segs(frame, period=60)
    assert nulls is not None

    def fail_rescan(*args, **kwargs):
        raise AssertionError('precomputed segments were ignored')

    monkeypatch.setattr(_anal, 'get_null_segs', fail_rescan)
    segments = list(iter_null_segs(
        timeframe=60,
        frame=frame,
        null_segs=nulls,
    ))

    assert len(segments) == 1
    assert segments[0][:6] == (
        1,
        3,
        1,
        3,
        120,
        240,
    )


def test_smart_dedupe_honors_sort_without_dupes() -> None:
    '''
    The sort option applies even when no duplicates exist.

    '''
    src: pl.DataFrame = np2pl(mk_ohlcv([
        180,
        60,
        120,
    ]))

    _, deduped, diff, _, issues = dedupe_ohlcv_smart(src)

    assert diff == 0
    assert issues is None
    assert deduped['time'].to_list() == [60, 120, 180]
    assert deduped['s_diff'].to_list() == [None, 60, 60]