# piker: trading gear for hackers
# Copyright (C) 2018-present  Tyler Goodlet (in stewardship of pikers)

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

'''
Financial time series processing utilities usually
pertaining to OHLCV style sampled data.

Routines are generally implemented in either ``numpy`` or
``polars`` B)

'''
from __future__ import annotations
from functools import partial
from math import (
    ceil,
    floor,
)
import time
from typing import (
    Literal,
    # AsyncGenerator,
    Generator,
)

import numpy as np
import polars as pl
from pendulum import (
    DateTime,
    from_timestamp,
)

from ..toolz.profile import (
    Profiler,
    pg_profile_enabled,
    ms_slower_then,
)
from ..log import (
    get_logger,
    get_console_log,
)
# for "time series processing"
subsys: str = 'piker.tsp'

log = get_logger(name=__name__)
get_console_log = partial(
    get_console_log,
    name=subsys,  # activate for subsys-pkg "downward"
)

# NOTE: union type-defs to handle generic `numpy` and `polars` types
# side-by-side Bo
# |_ TODO: schema spec typing?
#   -[ ] nptyping!
#   -[ ] wtv we can with polars?
Frame = pl.DataFrame | np.ndarray
Seq = pl.Series | np.ndarray


def slice_from_time(
    arr: np.ndarray,
    start_t: float,
    stop_t: float,
    step: float,  # sampler period step-diff

) -> slice:
    '''
    Calculate array indices mapped from a time range and return them in
    a slice.

    Given an input array with an epoch `'time'` series entry, calculate
    the indices which span the time range and return in a slice. Presume
    each `'time'` step increment is uniform and when the time stamp
    series contains gaps (the uniform presumption is untrue) use
    ``np.searchsorted()`` binary search to look up the appropriate
    index.

    '''
    profiler = Profiler(
        msg='slice_from_time()',
        disabled=not pg_profile_enabled(),
        ms_threshold=ms_slower_then,
    )

    times = arr['time']
    t_first = floor(times[0])
    t_last = ceil(times[-1])

    # the greatest index we can return which slices to the
    # end of the input array.
    read_i_max = arr.shape[0]

    # compute (presumed) uniform-time-step index offsets
    i_start_t = floor(start_t)
    read_i_start = floor(((i_start_t - t_first) // step)) - 1

    i_stop_t = ceil(stop_t)

    # XXX: edge case -> always set stop index to last in array whenever
    # the input stop time is detected to be greater then the equiv time
    # stamp at that last entry.
    if i_stop_t >= t_last:
        read_i_stop = read_i_max
    else:
        read_i_stop = ceil((i_stop_t - t_first) // step) + 1

    # always clip outputs to array support
    # for read start:
    # - never allow a start < the 0 index
    # - never allow an end index > the read array len
    read_i_start = min(
        max(0, read_i_start),
        read_i_max - 1,
    )
    read_i_stop = max(
        0,
        min(read_i_stop, read_i_max),
    )

    # check for larger-then-latest calculated index for given start
    # time, in which case we do a binary search for the correct index.
    # NOTE: this is usually the result of a time series with time gaps
    # where it is expected that each index step maps to a uniform step
    # in the time stamp series.
    t_iv_start = times[read_i_start]
    if (
        t_iv_start > i_start_t
    ):
        # do a binary search for the best index mapping to ``start_t``
        # given we measured an overshoot using the uniform-time-step
        # calculation from above.

        # TODO: once we start caching these per source-array,
        # we can just overwrite ``read_i_start`` directly.
        new_read_i_start = np.searchsorted(
            times,
            i_start_t,
            side='left',
        )

        # TODO: minimize binary search work as much as possible:
        # - cache these remap values which compensate for gaps in the
        #   uniform time step basis where we calc a later start
        #   index for the given input ``start_t``.
        # - can we shorten the input search sequence by heuristic?
        #   up_to_arith_start = index[:read_i_start]

        if (
            new_read_i_start <= read_i_start
        ):
            # t_diff = t_iv_start - start_t
            # print(
            #     f"WE'RE CUTTING OUT TIME - STEP:{step}\n"
            #     f'start_t:{start_t} -> 0index start_t:{t_iv_start}\n'
            #     f'diff: {t_diff}\n'
            #     f'REMAPPED START i: {read_i_start} -> {new_read_i_start}\n'
            # )
            read_i_start = new_read_i_start

    t_iv_stop = times[read_i_stop - 1]
    if (
        t_iv_stop > i_stop_t
    ):
        # t_diff = stop_t - t_iv_stop
        # print(
        #     f"WE'RE CUTTING OUT TIME - STEP:{step}\n"
        #     f'calced iv stop:{t_iv_stop} -> stop_t:{stop_t}\n'
        #     f'diff: {t_diff}\n'
        #     # f'SHOULD REMAP STOP: {read_i_start} -> {new_read_i_start}\n'
        # )
        new_read_i_stop = np.searchsorted(
            times[read_i_start:],
            # times,
            i_stop_t,
            side='right',
        )

        if (
            new_read_i_stop <= read_i_stop
        ):
            read_i_stop = read_i_start + new_read_i_stop + 1

    # sanity checks for range size
    # samples = (i_stop_t - i_start_t) // step
    # index_diff = read_i_stop - read_i_start + 1
    # if index_diff > (samples + 3):
    #     breakpoint()

    # read-relative indexes: gives a slice where `shm.array[read_slc]`
    # will be the data spanning the input time range `start_t` ->
    # `stop_t`
    read_slc = slice(
        int(read_i_start),
        int(read_i_stop),
    )

    profiler(
        'slicing complete'
        # f'{start_t} -> {abs_slc.start} | {read_slc.start}\n'
        # f'{stop_t} -> {abs_slc.stop} | {read_slc.stop}\n'
    )

    # NOTE: if caller needs absolute buffer indices they can
    # slice the buffer abs index like so:
    # index = arr['index']
    # abs_indx = index[read_slc]
    # abs_slc = slice(
    #     int(abs_indx[0]),
    #     int(abs_indx[-1]),
    # )

    return read_slc


def get_null_segs(
    frame: Frame,
    period: float,  # sampling step in seconds
    imargin: int = 1,
    col: str = 'time',

) -> tuple[
    # Seq,  # TODO: can we make it an array-type instead?
    list[
        list[int],
    ],
    Seq,
    Frame
] | None:
    '''
    Detect if there are any zero(-epoch stamped) valued
    rows in for the provided `col: str` column; by default
    presume the 'time' field/column.

    Filter to all such zero (time) segments and return
    the corresponding frame zeroed segment's,

      - gap absolute (in buffer terms), inclusive boundary-row
        index endpoints as `absi_zsegs`; each null run is expanded
        by `imargin` and clamped to the input frame
      - abs indices of all rows with zeroed `col` values as
        `absi_zeros`
      - the corresponding frame's row-entries (view) which are
        zeroed for the `col` as `zero_t`

    For an interior run, the default one-row margin makes each
    endpoint a valid datum immediately outside the nulls. At a frame
    edge the missing side is clamped and may remain null; callers that
    require query boundaries must reject or defer that segment.

    Consumers converting a pair to a Python slice must add one to its
    inclusive end index.

    '''
    values: Seq = frame[col]
    zero_pred: Seq = (values == 0)
    tis_zeros: bool = bool(zero_pred.any())

    if not tis_zeros:
        return None

    if imargin < 0:
        raise ValueError('`imargin` must be >= 0')

    if isinstance(frame, np.ndarray):
        zero_t: np.ndarray = frame[zero_pred]
        absi_zeros: np.ndarray = zero_t['index']
        frame_indexes: np.ndarray = frame['index']
        zero_indexes: np.ndarray = absi_zeros
    else:
        zero_t: pl.DataFrame = frame.filter(zero_pred)
        absi_zeros: pl.Series = zero_t['index']
        frame_indexes = frame['index'].to_numpy()
        zero_indexes = absi_zeros.to_numpy()

    # Null-run detection depends on absolute indexes stepping by
    # exactly one per frame row. Without this invariant an
    # index jump between adjacent null rows is indistinguishable from
    # valid data rows separating two null segments.
    if np.any(np.diff(frame_indexes) != 1):
        raise ValueError(
            'OHLCV frame indexes must be contiguous'
        )

    # Scan zero-row absolute indexes for steps larger than one. Each
    # such step begins another contiguous null segment:
    #
    # data  1st zero seg  data  zeros   data    zeros
    # ----  ------------  ----  -----   ------  -----
    # ||||..000000000000..||||..00000...||||||..0000
    # ----  ------------  ----  -----   ------  -----
    #       ^zero_indexes[0]                  ^zero_indexes[-1]
    #                           ^split_at[0]          ^split_at[1]
    #       ^zero_groups[0]     ^zero_groups[1]       ^zero_groups[2]
    #
    # `split_at` entries are frame-relative positions in
    # `zero_indexes`, not absolute buffer indexes. `np.split()` keeps
    # each run together so its first and last absolute indexes can be
    # expanded into the surrounding query boundary rows.
    #
    # Plain arrays also keep the NumPy and Polars paths on
    # exactly the same segment-grouping implementation.
    split_at: np.ndarray = (
        np.flatnonzero(np.diff(zero_indexes) != 1)
        +
        1
    )
    zero_groups: list[np.ndarray] = np.split(
        zero_indexes,
        split_at,
    )

    frame_start: int = int(frame_indexes[0])
    frame_end: int = int(frame_indexes[-1])
    absi_zsegs: list[list[int]] = []
    for group in zero_groups:
        zero_start: int = int(group[0])
        zero_end: int = int(group[-1])
        absi_zsegs.append([
            max(frame_start, zero_start - imargin),
            min(frame_end, zero_end + imargin),
        ])

    log.warning(
        f'Frame has {len(absi_zsegs)} NULL GAPS!?\n'
        f'period: {period}\n'
        f'total null samples: {len(zero_t)}\n'
    )

    return (
        absi_zsegs,  # [start, end] abs slice indices of seg
        absi_zeros,  # all abs indices within all null-segs
        zero_t,  # sliced-view of all null-segment rows-datums
    )


def iter_null_segs(
    timeframe: float,
    frame: Frame|None = None,
    null_segs: tuple|None = None,

) -> Generator[
    tuple[
        int, int,
        int, int,
        float, float,
        float, float,

        # Seq,  # TODO: can we make it an array-type instead?
        # list[
        #     list[int, int],
        # ],
        # Seq,
        # Frame
    ],
    None,
]:
    if null_segs is None:
        null_segs = get_null_segs(
            frame,
            period=timeframe,
        )
    if not null_segs:
        return

    absi_pairs_zsegs: list[list[float, float]]
    izeros: Seq
    zero_t: Frame
    (
        absi_pairs_zsegs,
        izeros,
        zero_t,
    ) = null_segs

    absi_first: int = frame[0]['index']
    for (
        absi_start,
        absi_end,
    ) in absi_pairs_zsegs:

        fi_end: int = absi_end - absi_first
        end_row: Seq = frame[fi_end]
        end_t: float = end_row['time']
        end_dt: DateTime = from_timestamp(end_t)

        fi_start = None
        start_row = None
        start_t = None
        start_dt = None
        if (
            absi_start is not None
            and
            start_t != 0
        ):
            fi_start: int = absi_start - absi_first
            start_row: Seq = frame[fi_start]
            start_t: float = start_row['time']
            start_dt: DateTime = from_timestamp(start_t)

        if absi_start < 0:
            import pdbp
            pdbp.set_trace()

        yield (
            absi_start, absi_end,  # abs indices
            fi_start, fi_end,  # relative "frame" indices
            start_t, end_t,  # epoch times
            start_dt, end_dt,  # dts
        )


def with_dts(
    df: pl.DataFrame,
    time_col: str = 'time',

) -> pl.DataFrame:
    '''
    Insert datetime (casted) columns to a (presumably) OHLC sampled
    time series with an epoch-time column keyed by `time_col: str`.

    '''
    return df.with_columns([
        pl.col(time_col).shift(1).name.suffix('_prev'),
        pl.col(time_col).diff().alias('s_diff'),
        pl.from_epoch(pl.col(time_col)).alias('dt'),
    ]).with_columns([
        pl.from_epoch(
            column=pl.col(f'{time_col}_prev'),
        ).alias('dt_prev'),
        pl.col('dt').diff().alias('dt_diff'),
    ])


t_unit: Literal = Literal[
    'days',
    'hours',
    'minutes',
    'seconds',
    'miliseconds',
    'microseconds',
    'nanoseconds',
]


def detect_time_gaps(
    w_dts: pl.DataFrame,

    time_col: str = 'time',
    # epoch sampling step diff
    expect_period: float = 60,

    # NOTE: legacy stock mkts have venue operating hours
    # and thus gaps normally no more then 1-2 days at
    # a time.
    gap_thresh: float = 1.,

    # TODO: allow passing in a frame of operating hours?
    # -[ ] durations/ranges for faster legit gap checks?
    # XXX -> must be valid ``polars.Expr.dt.<name>``
    # like 'days' which a sane default for venue closures
    # though will detect weekend gaps which are normal :o
    gap_dt_unit: t_unit | None = None,

) -> pl.DataFrame:
    '''
    Filter to OHLC datums which contain sample step gaps.

    For eg. legacy markets which have venue close gaps and/or
    actual missing data segments.

    '''
    # first select by any sample-period (in seconds unit) step size
    # greater then expected.
    step_gaps: pl.DataFrame = w_dts.filter(
        pl.col('s_diff') > expect_period
    )

    if gap_dt_unit is None:
        return step_gaps

    # NOTE: this flag is to indicate that on this (sampling) time
    # scale we expect to only be filtering against larger venue
    # closures-scale time gaps.
    #
    # Map to total_ method since `dt_diff` is a duration type,
    # not datetime - modern polars requires `total_*` methods
    # for duration types (e.g. `total_days()` not `day()`)
    # Ensure plural form for polars API (e.g. 'day' -> 'days')
    unit_plural: str = (
        gap_dt_unit
        if gap_dt_unit.endswith('s')
        else f'{gap_dt_unit}s'
    )
    duration_method: str = f'total_{unit_plural}'
    return step_gaps.filter(
        # Second by an arbitrary dt-unit step size
        getattr(
            pl.col('dt_diff').dt,
            duration_method,
        )().abs() > gap_thresh
    )


def detect_time_ordering_errors(
    w_dts: pl.DataFrame,

) -> pl.DataFrame:
    '''
    Return rows whose timestamp does not follow its predecessor.

    Zero deltas identify duplicate timestamps and negative deltas
    identify out-of-order rows. Neither is a missing-history gap.

    '''
    return w_dts.filter(
        pl.col('s_diff') <= 0
    )


def detect_price_gaps(
    df: pl.DataFrame,
    gt_multiplier: float = 2.,
    price_fields: list[str] = ['high', 'low'],

) -> pl.DataFrame:
    '''
    Detect gaps in clearing price over an OHLC series.

    2 types of gaps generally exist; up gaps and down gaps:

    - UP gap: when any next sample's lo price is strictly greater
      then the current sample's hi price.

    - DOWN gap: when any next sample's hi price is strictly
      less then the current samples lo price.

    '''
    # return df.filter(
    #     pl.col('high') - ) > expect_period,
    # ).select([
    #     pl.dt.datetime(pl.col(time_col).shift(1)).suffix('_previous'),
    #     pl.all(),
    # ]).select([
    #     pl.all(),
    #     (pl.col(time_col) - pl.col(f'{time_col}_previous')).alias('diff'),
    # ])
    ...

# TODO: probably just use the null_segs impl above?
def detect_vlm_gaps(
    df: pl.DataFrame,
    col: str = 'volume',

) -> pl.DataFrame:

    vnull: pl.DataFrame = df.filter(
        pl.col(col) == 0
    )
    return vnull


def dedupe(
    src_df: pl.DataFrame,

    time_gaps: pl.DataFrame | None = None,
    sort: bool = True,
    period: float = 60,

) -> tuple[
    pl.DataFrame,  # with dts
    pl.DataFrame,  # with deduplicated dts (aka gap/repeat removal)
    int,  # len diff between input and deduped
]:
    '''
    Check for time series gaps and if found
    de-duplicate any datetime entries, check for
    a frame height diff and return the newly
    dt-deduplicated frame.

    '''
    wdts: pl.DataFrame = with_dts(src_df)
    ordering_errors = detect_time_ordering_errors(wdts)
    if not ordering_errors.is_empty():
        log.warning(
            f'Found {ordering_errors.height} non-positive '
            f'timestamp step(s) before normalization:\n'
            f'{ordering_errors}'
        )

    # remove duplicated datetime samples/sections
    deduped: pl.DataFrame = src_df.unique(
        # subset=['dt'],
        subset=['time'],
        maintain_order=True,
    )

    # maybe sort on any time field
    if sort:
        deduped = deduped.sort(by='time')

    deduped = with_dts(deduped)

    diff: int = (
        wdts.height
        -
        deduped.height
    )
    return (
        wdts,
        deduped,
        diff,
    )


def sort_diff(
    src_df: pl.DataFrame,
    col: str = 'time',

) -> tuple[
    pl.DataFrame,  # with dts
    pl.DataFrame,  # sorted
    list[int],  # indices of segments that are out-of-order
]:
    ser: pl.Series = src_df[col]
    sortd: pl.DataFrame = ser.sort()
    diff: pl.Series = ser.diff()

    sortd_diff: pl.Series = sortd.diff()
    i_step_diff = (diff != sortd_diff).arg_true()
    frame_reorders: int = i_step_diff.len()
    if frame_reorders:
        log.warn(
            f'Resorted frame on col: {col}\n'
            f'{frame_reorders}'

        )
        # import pdbp; pdbp.set_trace()

# NOTE: thanks to this SO answer for the below conversion routines
# to go from numpy struct-arrays to polars dataframes and back:
# https://stackoverflow.com/a/72054819
def np2pl(array: np.ndarray) -> pl.DataFrame:
    start: float = time.time()

    # XXX: thanks to this SO answer for this conversion tip:
    # https://stackoverflow.com/a/72054819
    df = pl.DataFrame({
        field_name: array[field_name]
        for field_name in array.dtype.fields
    })
    delay: float = round(
        time.time() - start,
        ndigits=6,
    )
    log.info(
        f'numpy -> polars conversion took {delay} secs\n'
        f'polars df: {df}'
    )
    return df


def pl2np(
    df: pl.DataFrame,
    dtype: np.dtype,

) -> np.ndarray:

    # Create numpy struct array of the correct size and dtype
    # and loop through df columns to fill in array fields.
    array = np.empty(
        df.height,
        dtype,
    )
    for field in dtype.fields:
        array[field] = df.get_column(field).to_numpy()

    return array
