from typing import Dict, Optional

import numpy as np
import pandas as pd

from pychanlun.stick import Stick


class Fractal(Stick):

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

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

        stick_df = self.sticks[interval]
        if stick_df is None:
            return

        sticks = list(stick_df.itertuples())
        rows = []

        for index in range(1, len(sticks)):
            prev, cur = sticks[index - 1], sticks[index]
            nxt = sticks[index + 1] if index < len(sticks) - 1 else None

            if self._is_top_fractal(prev, cur, nxt):
                cur = cur._replace(low=np.nan)
            elif self._is_bottom_fractal(prev, cur, nxt):
                cur = cur._replace(high=np.nan)
            else:
                cur = cur._replace(high=np.nan, low=np.nan)

            rows.append(cur)

        self.fractals[interval] = self.to_dataframe(rows, ['high', 'low'])

    @staticmethod
    def _is_top_fractal(prev, cur, nxt):
        return cur.high > prev.high and (nxt is None or cur.high > nxt.high)

    @staticmethod
    def _is_bottom_fractal(prev, cur, nxt):
        return cur.low < prev.low and (nxt is None or cur.low < nxt.low)
