from typing import Dict, Optional

import pandas as pd

from pychanlun.stock import Stock


class Stick(Stock):

    def __init__(self, symbol, source: Dict[str, pd.DataFrame]):
        self.sticks: Dict[str, Optional[pd.DataFrame]] = {}
        super().__init__(symbol, source)

    def _generate_interval(self, interval):
        super()._generate_interval(interval)

        source_df = self.sources[interval]
        if source_df is None:
            return

        sources = list(source_df.itertuples())
        rows = []

        prev = cur = None
        index = 0
        while index < len(sources) - 1:
            prev, cur = sources[index], sources[index + 1]
            if self._is_go_up(prev, cur) or self._is_go_down(prev, cur):
                rows.append(prev)
                break
            index += 1

        for i in range(index + 2, len(sources)):
            nxt = sources[i]
            is_go_up = self._is_go_up(prev, cur)

            if self._can_merge_inside(cur, nxt):
                cur = self._update_stick(cur, nxt, is_go_up)
            elif self._can_merge_outside(cur, nxt):
                cur = self._update_stick(nxt, cur, is_go_up)
            else:
                rows.append(cur)
                prev, cur = cur, nxt

        rows.append(cur)
        self.sticks[interval] = self.to_dataframe(rows, ['high', 'low'])

    @staticmethod
    def _is_go_up(cur, nxt):
        return cur.high <= nxt.high and cur.low <= nxt.low

    @staticmethod
    def _is_go_down(cur, nxt):
        return cur.high >= nxt.high and cur.low >= nxt.low

    @staticmethod
    def _can_merge_inside(cur, nxt):
        return cur.high >= nxt.high and cur.low <= nxt.low

    @staticmethod
    def _can_merge_outside(cur, nxt):
        return cur.high <= nxt.high and cur.low >= nxt.low

    @staticmethod
    def _update_stick(cur, nxt, is_go_up):
        if is_go_up:
            return cur._replace(low=nxt.low)
        else:
            return cur._replace(high=nxt.high)