piker/piker/tsp/_annotate.py

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

"""
Time-series (remote) annotation APIs.

"""
from __future__ import annotations
from math import copysign
from typing import (
    Any,
    TYPE_CHECKING,
)

import polars as pl
import tractor

from piker.data._formatters import BGM
from piker.storage import log

if TYPE_CHECKING:
    from piker.ui._remote_ctl import AnnotCtl


async def markup_gaps(
    fqme: str,
    timeframe: float,
    actl: AnnotCtl,
    wdts: pl.DataFrame,
    gaps: pl.DataFrame,

) -> dict[int, dict]:
    '''
    Remote annotate time-gaps in a dt-fielded ts (normally OHLC)
    with rectangles.

    '''
    aids: dict[int] = {}
    for i in range(gaps.height):
        row: pl.DataFrame = gaps[i]

        # the gap's RIGHT-most bar's OPEN value
        # at that time (sample) step.
        iend: int = row['index'][0]
        # dt: datetime = row['dt'][0]
        # dt_prev: datetime = row['dt_prev'][0]
        # dt_end_t: float = dt.timestamp()


        # TODO: can we eventually remove this
        # once we figure out why the epoch cols
        # don't match?
        # TODO: FIX HOW/WHY these aren't matching
        # and are instead off by 4hours (EST
        # vs. UTC?!?!)
        # end_t: float = row['time']
        # assert (
        #     dt.timestamp()
        #     ==
        #     end_t
        # )

        # the gap's LEFT-most bar's CLOSE value
        # at that time (sample) step.
        prev_r: pl.DataFrame = wdts.filter(
            pl.col('index') == iend - 1
        )
        # XXX: probably a gap in the (newly sorted or de-duplicated)
        # dt-df, so we might need to re-index first..
        dt: pl.Series = row['dt']
        dt_prev: pl.Series = row['dt_prev']
        if prev_r.is_empty():

            # XXX, filter out any special ignore cases,
            # - UNIX-epoch stamped datums
            # - first row
            if (
                dt_prev.dt.epoch()[0] == 0
                or
                dt.dt.epoch()[0] == 0
            ):
                log.warning('Skipping row with UNIX epoch timestamp ??')
                continue

            if wdts[0]['index'][0] == iend:  # first row
                log.warning('Skipping first-row (has no previous obvi) !!')
                continue

            # XXX, if the previous-row by shm-index is missing,
            # meaning there is a missing sample (set), get the prior
            # row by df index and attempt to use it?
            i_wdts: pl.DataFrame = wdts.with_row_index(name='i')
            i_row: int = i_wdts.filter(pl.col('index') == iend)['i'][0]
            prev_row_by_i = wdts[i_row]
            prev_r: pl.DataFrame = prev_row_by_i

            # debug any missing pre-row
            if tractor._state.is_debug_mode():
                await tractor.pause()

        istart: int = prev_r['index'][0]
        # TODO: implement px-col width measure
        # and ensure at least as many px-cols
        # shown per rect as configured by user.
        # gap_w: float = abs((iend - istart))
        # if gap_w < 6:
        #     margin: float = 6
        #     iend += margin
        #     istart -= margin

        opn: float = row['open'][0]
        cls: float = prev_r['close'][0]

        # BGM=0.16 is the normal diff from overlap between bars, SO
        # just go slightly "in" from that "between them".
        from_idx: int = BGM - .06  # = .10

        lc: tuple[float, float] = (
            istart + 1 - from_idx,
            cls,
        )
        ro: tuple[float, float] = (
            iend + from_idx,
            opn,
        )

        diff: float = cls - opn
        sgn: float = copysign(1, diff)
        color: str = 'dad_blue'
        # TODO? mks more sense to have up/down coloring?
        # color: str = {
        #     -1: 'lilypad_green',  # up-gap
        #     1: 'wine',  # down-gap
        # }[sgn]

        rect_kwargs: dict[str, Any] = dict(
            fqme=fqme,
            timeframe=timeframe,
            start_pos=lc,
            end_pos=ro,
            color=color,
        )

        # add up/down rects
        aid: int = await actl.add_rect(**rect_kwargs)
        assert aid
        aids[aid] = rect_kwargs
        direction: str = (
            'down' if sgn == 1
            else 'up'
        )
        arrow_kwargs: dict[str, Any] = dict(
            fqme=fqme,
            timeframe=timeframe,
            x=iend,
            y=cls,
            color=color,
            alpha=160,
            pointing=direction,
        )

        aid: int = await actl.add_arrow(
            **arrow_kwargs
        )

    # tell chart to redraw all its
    # graphics view layers Bo
    await actl.redraw(
        fqme=fqme,
        timeframe=timeframe,
    )
    return aids