# piker: trading gear for hackers
# Copyright (C) 2018-present Tyler Goodlet (in stewardship of pikers)
# SPDX-License-Identifier: AGPL-3.0-or-later

# This program is free software: you can redistribute it and/or
# modify it under the terms of the GNU Affero General Public License
# as published by the Free Software Foundation, either version 3
# of the License, or (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public
# License along with this program. If not, see
# <https://www.gnu.org/licenses/>.

'''
Causality, streaming, and publication contracts for Trend Calculus.

'''
from collections.abc import AsyncIterator, Callable
from dataclasses import Field, fields
from types import SimpleNamespace
from typing import Any, cast

import numpy as np
import pytest
import trio
from trio.testing import wait_all_tasks_blocked
from trio_typing import TaskStatus

from piker.data.ticktools import FeedQuote
from piker.fsp._api import (
    FspSample,
    FspStream,
    FspYield,
    _load_builtins,
)
from piker.fsp._engine import (
    Cascade,
    QuoteFrame,
    _next_step_output,
    _sample_wakeups,
    _write_sample,
    filter_quotes_by_sym,
)
from piker.fsp._trend import (
    TrendCalculator,
    TrendHistory,
    TrendUpdate,
)
from piker.fsp._trend_fsp import trend_calculus
from tractor.ipc._shm import ShmArray
from tractor.trionics import Lagged


def test_split_preserves_both_reversals() -> None:
    '''
    A net unchanged direction can contain two opposing reversals.

    The final outside window [2, 5] follows the rising window [4, 3].
    Bridging its earlier extremum reveals a top at 4 and a bottom at
    2, both detected on observation 5. A one-pulse or net-sign-only
    implementation loses the top. Assert both pivot coordinates and
    the later detection coordinate, including the historical fields.

    '''
    values: np.ndarray = np.array([1., 2., 4., 3., 2., 5.])
    calc: TrendCalculator = TrendCalculator()
    calc.extend(values[:-1], np.arange(5.))
    update: TrendUpdate = calc.push(5., values[-1])
    assert update.trends == (1, 0, 0, 0)
    assert update.reversals['index'].tolist() == [2, 4]
    assert update.reversals['direction'].tolist() == [-1, 1]
    assert update.reversals['price'].tolist() == [4., 2.]
    assert update.reversals['detected_index'].tolist() == [5, 5]
    assert update.reversals['detected_time'].tolist() == [5., 5.]
    history: TrendHistory = TrendCalculator().extend(
        values, np.arange(6.),
    )
    np.testing.assert_array_equal(
        history.pivot_prices[-1, 0], [4., 2.],
    )
    np.testing.assert_array_equal(
        history.pivot_times[-1, 0], [2., 4.],
    )
    assert np.isnan(history.pivot_prices[:-1]).all()


def test_batch_chunking_and_every_prefix() -> None:
    '''
    Future observations and transport chunking must not alter output.

    Materialize a random walk with ties and a missing observation,
    then compare every prefix and uneven continuation batches with a
    single history scan. Compare pivot detection fields too, because
    equal direction arrays alone can hide backdated reversals.

    '''
    rng: np.random.Generator = np.random.default_rng(18)
    values: np.ndarray = (
        rng.integers(-2, 3, 80).cumsum().astype(float)
    )
    values[31] = np.nan
    times: np.ndarray = np.arange(80.) * 0.125 + 1000.
    indices: np.ndarray = np.arange(80, dtype=np.int64) + 12345
    windows: tuple[int, ...] = (3, 2, 4)
    full: TrendHistory = TrendCalculator(windows).extend(
        values, times, indices,
    )
    cut: int
    field: Field[Any]
    for cut in range(81):
        prefix: TrendHistory = TrendCalculator(windows).extend(
            values[:cut], times[:cut], indices[:cut],
        )
        for field in fields(TrendHistory):
            np.testing.assert_array_equal(
                getattr(prefix, field.name),
                getattr(full, field.name)[:cut],
            )

    calc: TrendCalculator = TrendCalculator(windows)
    start: int = 0
    for cut in (0, 1, 2, 7, 31, 32, 34, 60, 80):
        chunk: TrendHistory = calc.extend(
            values[start:cut], times[start:cut], indices[start:cut],
        )
        for field in fields(TrendHistory):
            np.testing.assert_array_equal(
                getattr(chunk, field.name),
                getattr(full, field.name)[start:cut],
            )
        start = cut


@pytest.mark.parametrize('window', [2, 3, 7])
@pytest.mark.parametrize('direction', [-1, 0, 1])
def test_warmup_flat_and_negative_values(
    window: int,
    direction: int,
) -> None:
    '''
    Startup must not invent a trend from a synthetic zero-price bar.

    Use negative, monotone or flat values through two full windows.
    The first comparable pair establishes readiness and direction;
    no sign change or initial pivot exists before that comparison.

    '''
    values: np.ndarray = (
        np.arange(2 * window, dtype=float) * direction - 100
    )
    result: TrendHistory = TrendCalculator((window,)).extend(
        values, np.arange(float(len(values))),
    )
    assert not result.ready[:-1].any()
    assert result.ready[-1, 0]
    assert not result.trends[:-1].any()
    assert result.trends[-1, 0] == direction
    assert np.isnan(result.pivot_prices).all()


def test_nan_resets_hierarchy_without_future_contamination() -> None:
    '''
    A missing value invalidates every recursive reversal order.

    Build nontrivial history, append NaN, and compare the following
    observations against a newly constructed calculator. Retaining
    even a higher-order partial window would make this replay differ.

    '''
    rng: np.random.Generator = np.random.default_rng(77)
    calc: TrendCalculator = TrendCalculator((2, 2, 2))
    calc.extend(rng.normal(size=200), np.arange(200.))
    update: TrendUpdate = calc.push(200., np.nan)
    assert update.trends == (0, 0, 0)
    assert update.ready == (False, False, False)
    assert not len(update.reversals)
    values: np.ndarray = rng.normal(size=100)
    times: np.ndarray = np.arange(201., 301.)
    actual: TrendHistory = calc.extend(values, times)
    expected: TrendHistory = TrendCalculator((2, 2, 2)).extend(
        values, times,
    )
    field: Field[Any]
    for field in fields(TrendHistory):
        np.testing.assert_array_equal(
            getattr(actual, field.name),
            getattr(expected, field.name),
        )


@pytest.mark.parametrize(
    'time,value,index',
    [
        (3., 1., 14),
        (2., 1., 14),
        (4., np.inf, 14),
        (np.nan, 1., 14),
        (4., 1., 13),
        (4., 1., -1),
        (4., 1., 2 ** 63),
        (4., 1., np.uint64(2 ** 63)),
        (4., 1., True),
        (4., 1., np.bool_(True)),
    ],
)
def test_invalid_sample_is_atomic(
    time: float,
    value: float,
    index: int,
) -> None:
    '''
    Invalid coordinates must not corrupt a partially filled window.

    Seed a calculator with four observations, attempt one rejected
    sample, and compare continuation to an untouched state copy.
    This checks that validation precedes recurrence mutation.

    '''
    calc: TrendCalculator = TrendCalculator()
    calc.extend(np.arange(4.), np.arange(4.), np.arange(10, 14))
    untouched: TrendCalculator = calc.copy()
    with pytest.raises(ValueError):
        calc.push(time, value, index=index)
    actual: TrendUpdate = calc.push(4., -1., index=14)
    expected: TrendUpdate = untouched.push(4., -1., index=14)
    assert actual.trends == expected.trends
    np.testing.assert_array_equal(
        actual.reversals, expected.reversals,
    )


def test_invalid_batch_is_atomic() -> None:
    '''
    A late malformed sample cannot partly advance a batch operation.

    The duplicate final timestamp is detected before the first valid
    row is consumed. A subsequent valid replay matches fresh state.

    '''
    calc: TrendCalculator = TrendCalculator()
    with pytest.raises(ValueError):
        calc.extend(np.arange(4.), np.array([1., 2., 3., 3.]))
    actual: TrendHistory = calc.extend(np.arange(4.), np.arange(4.))
    expected: TrendHistory = TrendCalculator().extend(
        np.arange(4.), np.arange(4.),
    )
    np.testing.assert_array_equal(actual.trends, expected.trends)


def test_last_int64_index_is_exact_and_overflow_is_atomic() -> None:
    '''
    Explicit unsigned coordinates must not wrap in compiled state.

    Observe the last representable coordinate using a NumPy uint64.
    A subsequent automatic push or batch must reject overflow before
    touching the partial window, and retain the exact last index.

    '''
    calc: TrendCalculator = TrendCalculator()
    limit: int = np.iinfo(np.int64).max
    result: TrendUpdate = calc.push(
        1., 2., index=np.uint64(limit),
    )
    assert result.index == limit
    assert isinstance(result.index, int)
    before: np.ndarray = calc._bars.copy()
    with pytest.raises(ValueError, match='int64'):
        calc.push(2., 3.)
    with pytest.raises(ValueError, match='overflow'):
        calc.extend(np.array([3.]), np.array([2.]))
    np.testing.assert_array_equal(calc._bars, before)


class MemorySource:
    '''
    Expose mutable source rows with absolute, nonzero SHM bounds.

    '''
    def __init__(self, values: list[float]) -> None:
        '''
        Seed mutable rows with a nonzero absolute SHM origin.

        '''
        self._first: SimpleNamespace = SimpleNamespace(value=37)
        self._last: SimpleNamespace = SimpleNamespace(
            value=37 + len(values),
        )
        self._array: np.ndarray = np.zeros(128, dtype=[
            ('index', '<i8'),
            ('time', '<f8'),
            ('close', '<f8'),
        ])
        self._array['index'] = np.arange(128)
        self._array['time'] = np.arange(128.) + 1000.
        self.array['close'] = values

    @property
    def array(self) -> np.ndarray:
        '''
        Expose ``MemorySource._array`` between published bounds.

        '''
        return self._array[self._first.value:self._last.value]


def assert_sample_matches(
    sample: FspSample,
    rows: np.ndarray,
) -> None:
    '''
    Compare every FSP field with a fresh source-prefix replay.

    '''
    history: TrendHistory = TrendCalculator().extend(
        rows['close'], rows['time'], rows['index'],
    )
    assert sample.index == rows['index'][-1]
    assert sample.time == rows['time'][-1]
    assert sample.fields['ready'] == history.ready[-1].sum()
    level: int
    for level in range(4):
        order: int = level + 1
        name: str = 'trend' if level == 0 else f'trend_{order}'
        assert sample.fields[name] == history.trends[-1, level]
        side: int
        label: str
        for side, label in enumerate(('top', 'bottom')):
            np.testing.assert_equal(
                sample.fields[f'{label}_{order}'],
                history.pivot_prices[-1, level, side],
            )
            np.testing.assert_equal(
                sample.fields[f'{label}_time_{order}'],
                history.pivot_times[-1, level, side],
            )


def test_fsp_preview_mutations_and_skipped_steps() -> None:
    '''
    Updating an open bar must replace its preview, not consume it.

    Repeatedly mutate the current SHM close across both directions,
    then advance by two rows before the next quote. Each FspSample
    must equal a fresh batch replay, and its absolute row must match
    the source. Quote payloads are deliberately empty wakeups so the
    source state rather than trade arrivals must drive sampling.

    '''
    async def main() -> None:
        '''
        Run the controlled scenario in an isolated Trio lifetime.

        '''
        shm: MemorySource = MemorySource([1., 2., 4., 3., 2., 5.])

        async def quotes() -> AsyncIterator[FeedQuote]:
            '''
            Provide payload-free wakeups for source
            snapshot sampling.

            '''
            while True:
                yield {}

        stream: FspStream = trend_calculus(
            quotes(), cast(ShmArray, shm),
        )
        try:
            historical: FspYield = await anext(stream)
            assert isinstance(historical, dict)
            assert set(historical) == set(trend_calculus.outputs)
            value: np.ndarray|None  # noqa: F842
            assert all(
                value is not None and len(value) == 6
                for value in historical.values()
            )
            close: float
            for close in (100., -20., 5., 5.):
                shm.array['close'][-1] = close
                sample: FspYield = await anext(stream)
                assert isinstance(sample, FspSample)
                assert_sample_matches(sample, shm.array)

            last: int = shm._last.value
            shm._array['close'][last:last + 2] = [-4., -5.]
            shm._last.value += 2
            end: int
            for end in (6, 7, 8):
                sample: FspYield = await anext(stream)
                assert isinstance(sample, FspSample)
                assert_sample_matches(sample, shm.array[:end])
        finally:
            await stream.aclose()

    trio.run(main)


def test_fsp_empty_bootstrap_and_source_revision() -> None:
    '''
    Empty startup must accept later source rows without old state.

    After the first sample, prepend changes the source origin. The
    old generator must stop; restarting it recomputes the revised
    history rather than reusing recurrence state from another origin.

    '''
    async def main() -> None:
        '''
        Run the controlled scenario in an isolated Trio lifetime.

        '''
        shm: MemorySource = MemorySource([])

        async def quotes() -> AsyncIterator[FeedQuote]:
            '''
            Provide payload-free wakeups for source
            snapshot sampling.

            '''
            while True:
                yield {}

        stream: FspStream = trend_calculus(
            quotes(), cast(ShmArray, shm),
        )
        historical: FspYield = await anext(stream)
        assert isinstance(historical, dict)
        value: np.ndarray|None  # noqa: F842
        assert all(
            value is not None and not len(value)
            for value in historical.values()
        )
        shm._array['close'][37] = 1.
        shm._last.value += 1
        sample: FspYield = await anext(stream)
        assert isinstance(sample, FspSample)
        assert_sample_matches(sample, shm.array)
        shm._first.value -= 1
        with pytest.raises(StopAsyncIteration):
            await anext(stream)
        restarted: FspStream = trend_calculus(
            quotes(), cast(ShmArray, shm),
        )
        try:
            await anext(restarted)
            sample: FspYield = await anext(restarted)
            assert isinstance(sample, FspSample)
            assert_sample_matches(sample, shm.array)
        finally:
            await restarted.aclose()

    trio.run(main)


def test_indexed_sample_waits_for_destination_step() -> None:
    '''
    A quote can reach the FSP before the sampler appends destination.

    Block _write_sample with destination one row behind source and
    wait until its task is suspended. Publishing the next destination
    bound and setting its event permits one coherent write. The prior
    row remains unchanged, proving no accidental previous-bar write.

    '''
    async def main() -> None:
        '''
        Run the controlled scenario in an isolated Trio lifetime.

        '''
        array: np.ndarray = np.zeros(16, dtype=[
            ('time', '<f8'),
            ('trend', '<f8'),
            ('ready', '<f8'),
        ])
        dst: SimpleNamespace = SimpleNamespace(
            _first=SimpleNamespace(value=7),
            _last=SimpleNamespace(value=10),
            _array=array,
        )
        casc: Cascade = cast(Cascade, SimpleNamespace(
            dst=SimpleNamespace(rt_shm=dst),
            sample_ready=None,
        ))
        complete: trio.Event = trio.Event()

        async def publish() -> None:
            '''
            Wait for the destination bound, then signal publication.

            '''
            assert await _write_sample(casc, FspSample(
                10, 1010., {'trend': -1., 'ready': 2.},
            ))
            complete.set()

        nursery: trio.Nursery
        async with trio.open_nursery() as nursery:
            nursery.start_soon(publish)
            await wait_all_tasks_blocked()
            assert not complete.is_set()
            assert not array['trend'].any()
            dst._last.value += 1
            assert casc.sample_ready is not None
            casc.sample_ready.set()
            await complete.wait()
        assert tuple(array[9]) == (0., 0., 0.)
        assert tuple(array[10]) == (1010., -1., 2.)

    trio.run(main)


def test_indexed_sample_wait_is_cancellable() -> None:
    '''
    History resync must cancel a compute task waiting on sampling.

    Suspend indexed publication against an empty destination, cancel
    its nursery, then advance that destination. No cancelled sample
    may publish into the new source revision after cancellation.

    '''
    async def main() -> None:
        '''
        Run the controlled scenario in an isolated Trio lifetime.

        '''
        array: np.ndarray = np.zeros(4, dtype=[
            ('time', float), ('trend', float),
        ])
        dst: SimpleNamespace = SimpleNamespace(
            _first=SimpleNamespace(value=0),
            _last=SimpleNamespace(value=0),
            _array=array,
        )
        casc: Cascade = cast(Cascade, SimpleNamespace(
            dst=SimpleNamespace(rt_shm=dst),
            sample_ready=None,
        ))
        nursery: trio.Nursery
        async with trio.open_nursery() as nursery:
            nursery.start_soon(
                _write_sample, casc, FspSample(0, 1., {'trend': 1.}),
            )
            await wait_all_tasks_blocked()
            nursery.cancel_scope.cancel()
        dst._last.value = 1
        assert casc.sample_ready is not None
        casc.sample_ready.set()
        assert not array['trend'].any()

    trio.run(main)


def test_builtin_registration() -> None:
    '''
    Actor-side builtin discovery must resolve the new FSP target.

    '''
    assert _load_builtins()[trend_calculus.ns_path] is trend_calculus


def test_quote_free_sampler_steps_finalize_without_repeating_pivots(
) -> None:
    '''
    A quiet feed still finalizes bars and clears old reversal pulses.

    Start with an open bar containing two provisional reversals and
    advance the source on sampler notifications alone. Mirror the
    destination append protocol and run the actual FSP generator and
    indexed writer. Every stored field must equal causal replay;
    event fields must also be clear before compute catches up.

    '''
    async def main() -> None:
        '''
        Run the controlled scenario in an isolated Trio lifetime.

        '''
        source: MemorySource = MemorySource([1., 2., 4., 3., 2., 5.])
        first: int = source._first.value
        last: int = source._last.value
        name: str
        dtype: list[tuple[str, str]] = [
            ('index', '<i8'),
            ('time', '<f8'),
        ] + [
            (name, '<f8') for name in trend_calculus.outputs
        ]
        array: np.ndarray = np.zeros(128, dtype=dtype)
        array['index'] = np.arange(128)
        destination: SimpleNamespace = SimpleNamespace(
            _first=SimpleNamespace(value=first),
            _last=SimpleNamespace(value=last),
            _array=array,
        )
        casc: Cascade = cast(Cascade, SimpleNamespace(
            dst=SimpleNamespace(rt_shm=destination),
            sample_ready=None,
        ))
        send: trio.MemorySendChannel[QuoteFrame]
        quotes: trio.MemoryReceiveChannel[QuoteFrame]
        send, quotes = trio.open_memory_channel[QuoteFrame](0)
        wakeups: AsyncIterator[QuoteFrame]
        notify: Callable[[], None]|None
        async with send, quotes, _sample_wakeups(
            quotes, 'test.market', True,
        ) as (wakeups, notify):
            stream: FspStream = trend_calculus(
                filter_quotes_by_sym('test.market', wakeups),
                cast(ShmArray, source),
            )
            history: FspYield = await anext(stream)
            assert isinstance(history, dict)
            name: str
            values: np.ndarray
            for name, values in history.items():
                assert values is not None
                array[name][first:last] = values
            array['time'][first:last] = source.array['time']

            async def compute() -> None:
                '''
                Run until cancellation and record completion on exit.

                '''
                try:
                    sample: FspYield
                    async for sample in stream:
                        assert isinstance(sample, FspSample)
                        await _write_sample(casc, sample)
                finally:
                    await stream.aclose()

            nursery: trio.Nursery
            async with trio.open_nursery() as nursery:
                nursery.start_soon(compute)
                await wait_all_tasks_blocked()
                assert array['top_1'][last - 1] == 4.
                assert array['bottom_1'][last - 1] == 2.
                close: float
                for close in (5., 5., -4., -5.):
                    last = source._last.value
                    source._array['close'][last] = close
                    source._last.value += 1
                    array[last:last + 1] = _next_step_output(
                        array[last - 1:last], trend_calculus, False,
                    )
                    array['index'][last] = last
                    array['time'][last] = source.array['time'][-1]
                    destination._last.value += 1
                    assert np.isnan(array['top_1'][last])
                    assert np.isnan(array['bottom_1'][last])
                    if casc.sample_ready is not None:
                        casc.sample_ready.set()
                    casc.sample_ready = trio.Event()
                    assert notify is not None
                    # Coalesce a burst without losing finalized rows.
                    notify()
                    notify()
                    await wait_all_tasks_blocked()
                    offset: int
                    row: np.void
                    for offset, row in enumerate(source.array):
                        name: str
                        sample: FspSample = FspSample(
                            int(row['index']), float(row['time']),
                            {
                                name: array[name][int(row['index'])]
                                for name in trend_calculus.outputs
                            },
                        )
                        assert_sample_matches(
                            sample, source.array[:offset + 1],
                        )
                nursery.cancel_scope.cancel()

    trio.run(main)


def test_snapshot_quote_overrun_resumes_source_wakeups() -> None:
    '''
    Losing irrelevant quote payloads must not kill a SHM processor.

    A broadcast receiver raises Lagged and then remains usable.
    Simulate that protocol, require a catchup notification, and
    deliver another quote afterward to prove the relay resumed.

    '''
    async def main() -> None:
        '''
        Run the controlled scenario in an isolated Trio lifetime.

        '''
        send: trio.MemorySendChannel[QuoteFrame|Exception]
        receive: trio.MemoryReceiveChannel[QuoteFrame|Exception]
        send, receive = trio.open_memory_channel[
            QuoteFrame|Exception
        ](1)

        class Quotes:
            '''
            Inject one overrun without closing the quote receiver.

            '''
            def __aiter__(self) -> 'Quotes':
                '''
                Retain this controlled quote receiver
                across overruns.

                '''
                return self

            async def __anext__(self) -> QuoteFrame:
                '''
                Deliver the scheduled frame or
                injected receiver error.

                '''
                value: QuoteFrame|Exception = await receive.receive()
                if isinstance(value, Exception):
                    raise value
                return value

        wakeups: AsyncIterator[QuoteFrame]
        notify: Callable[[], None]|None
        async with send, receive, _sample_wakeups(
            Quotes(), 'test.market', True,
        ) as (wakeups, notify):
            await send.send(Lagged('synthetic quote overrun'))
            assert 'test.market' in await anext(wakeups)
            await send.send({'test.market': {'ticks': []}})
            assert 'test.market' in await anext(wakeups)
            assert notify is not None

    trio.run(main)


def test_edge_exit_stops_compute_before_source_receiver_closes(
) -> None:
    '''
    Normal sampler EOF must not strand or abruptly close compute.

    Keep the parent nursery alive beyond open_edge and inspect the
    compute task's finalizer immediately after context exit. This
    proves cleanup belongs to the edge scope rather than relying on
    an outer actor cancellation or a closed receiver exception.

    '''
    async def main() -> None:
        '''
        Run the controlled scenario in an isolated Trio lifetime.

        '''
        cleaned: trio.Event = trio.Event()
        nursery: trio.Nursery
        async with trio.open_nursery() as nursery:
            casc: Cascade = cast(Cascade, SimpleNamespace(
                tn=nursery,
                cs=None,
                complete=None,
                bind_func=None,
            ))

            async def compute(
                *,
                task_status: TaskStatus[int] = (
                    trio.TASK_STATUS_IGNORED
                ),
            ) -> None:
                '''
                Run until cancellation and record completion on exit.

                '''
                scope: trio.CancelScope
                with trio.CancelScope() as scope:
                    casc.cs = scope
                    casc.complete = trio.Event()
                    task_status.started(42)
                    try:
                        await trio.sleep_forever()
                    finally:
                        cleaned.set()
                        casc.complete.set()

            index: int
            async with Cascade.open_edge(casc, compute) as index:
                assert index == 42
                assert not cleaned.is_set()
            assert cleaned.is_set()

    trio.run(main)