from typing import Dict, Optional

import pandas as pd

from pychanlun.pivot import Pivot


class Signal(Pivot):

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

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

        stroke_df = self.strokes[interval]
        stroke_pivot_df = self.stroke_pivots[interval]
        if stroke_df is not None or stroke_pivot_df is not None:
            self.stroke_signals[interval] = self._generate_signal_interval(stroke_df, stroke_pivot_df)
        segment_df = self.segments[interval]
        segment_pivot_df = self.segment_pivots[interval]
        if segment_df is not None or segment_pivot_df is not None:
            self.segment_signals[interval] = self._generate_signal_interval(segment_df, segment_pivot_df)

    def _generate_signal_interval(self, segment_df, pivot_df):
        segment_df = segment_df.copy()
        segment_df['signal'] = 0
        pivots = list(pivot_df.itertuples())
        rows = []

        for index in range(0, len(pivots) - 3, 2):
            cur = pivots[index]
            pivot = self.get_range(pivots, index + 2)

            if cur.divergence > 0:
                one, two = self._check_signal_one_two(segment_df, pivot)
                if one is not None:
                    rows.append(one)
                if two is not None:
                    rows.append(two)
            three = self._check_signal_three(segment_df, pivot)
            if three is not None:
                rows.append(three)

        return self.to_dataframe(rows, ['high', 'low', 'signal'])

    def _check_signal_one_two(self, segment_df, pivot):
        segments = list(segment_df.loc[pivot.start.Index:].itertuples())
        if len(segments) < 3:
            return None

        nxt_1 = segments[0]
        nxt_3 = segments[2]
        if self.is_top(nxt_1) and nxt_3.high < nxt_1.high:
            nxt_1 = nxt_1._replace(signal=-1)
            nxt_3 = nxt_3._replace(signal=-2)
            return nxt_1, nxt_3
        if self.is_bottom(nxt_1) and nxt_3.low > nxt_1.low:
            nxt_1 = nxt_1._replace(signal=1)
            nxt_3 = nxt_3._replace(signal=2)
            return nxt_1, nxt_3
        return None, None

    def _check_signal_three(self, segment_df, pivot):
        segments = list(segment_df.loc[pivot.end.Index:].itertuples())
        if len(segments) < 3:
            return None

        nxt_2 = segments[1]
        nxt_3 = segments[2]
        if self.is_top(nxt_2) and nxt_3.low > pivot.high:
            nxt_3 = nxt_3._replace(signal=3)
            return nxt_3
        if self.is_bottom(nxt_2) and nxt_3.low < pivot.low:
            nxt_3 = nxt_3._replace(signal=-3)
            return nxt_3
        return None
