# 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/>.

'''
Measure deterministic Trend Calculus processing without a feed.

Timings distinguish kernel warmup, batch scanning, scalar commits,
and uncommitted live previews. They exclude network/model latency.
Run from the Piker environment; no credentials or services are used.

'''
import argparse
import json
import platform
from time import perf_counter

import numba
import numpy as np

from piker.fsp._trend import (
    TrendCalculator,
    TrendHistory,
    TrendUpdate,
)


def benchmark(
    rows: int,
    live_rows: int,
) -> dict[str, str|int|float|dict[str, float]]:
    '''
    Return reproducible timings and the retained NumPy state size.

    ``rows`` controls the ``TrendCalculator.extend()`` batch size.
    ``live_rows`` controls timed calls to ``TrendCalculator.push()``
    and previews made through ``TrendCalculator.copy()``. Assertions
    compare directions across all three paths outside timed regions.

    '''
    rng: np.random.Generator = np.random.default_rng(20261004)
    values: np.ndarray = rng.normal(size=rows).cumsum() + 100.
    times: np.ndarray = np.arange(rows, dtype=float) * .25
    calc: TrendCalculator = TrendCalculator()
    start: float = perf_counter()
    calc.extend(values[:8], times[:8])
    calc.push(times[8], values[8])
    warmup: float = perf_counter() - start

    start = perf_counter()
    full: TrendHistory = TrendCalculator().extend(values, times)
    batch: float = perf_counter() - start
    calc = TrendCalculator()
    elapsed: np.ndarray = np.empty(live_rows)
    preview_elapsed: np.ndarray = np.empty(live_rows)
    number: int
    for number in range(live_rows):
        start = perf_counter()
        preview: TrendUpdate = calc.copy().push(
            times[number], values[number],
        )
        preview_elapsed[number] = perf_counter() - start
        start = perf_counter()
        update: TrendUpdate = calc.push(
            times[number], values[number],
        )
        elapsed[number] = perf_counter() - start
        assert preview.trends == update.trends
        assert update.trends == tuple(full.trends[number])

    def percentiles_us(samples: np.ndarray) -> dict[str, float]:
        '''
        Report median and tail operation durations in microseconds.

        '''
        return dict(zip(
            ('p50', 'p95', 'p99'),
            np.percentile(samples * 1_000_000, (50, 95, 99)),
        ))

    array: np.ndarray  # noqa: F842
    return {
        'python': platform.python_version(),
        'machine': platform.machine(),
        'numpy': np.__version__,
        'numba': numba.__version__,
        'rows': rows,
        'live_rows': live_rows,
        'orders': len(calc.windows),
        'kernel_warmup_s': warmup,
        'batch_s': batch,
        'batch_rows_per_s': rows / batch,
        'scalar_commit_us': percentiles_us(elapsed),
        'preview_copy_and_push_us': percentiles_us(preview_elapsed),
        'retained_numpy_state_bytes': sum(
            array.nbytes for array in (
                calc._sizes, calc._bars, calc._counters,
                calc._events, calc._counts,
            )
        ),
    }


def main() -> None:
    '''
    Validate CLI sizes and print benchmark measurements as JSON.

    '''
    parser: argparse.ArgumentParser = argparse.ArgumentParser(
        description=__doc__,
    )
    parser.add_argument('--rows', type=int, default=1_000_000)
    parser.add_argument('--live-rows', type=int, default=20_000)
    args: argparse.Namespace = parser.parse_args()
    if args.rows < 9 or not 1 <= args.live_rows <= args.rows:
        parser.error('require rows >= 9 and 1 <= live-rows <= rows')
    print(json.dumps(benchmark(args.rows, args.live_rows), indent=2))


if __name__ == '__main__':
    main()