from typing import Dict, Optional

import pandas as pd

from pychanlun.stroke import Stroke, IndexedRow


class Segment(Stroke):

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

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

        stroke_df = self.strokes[interval]
        if stroke_df is None:
            return

        strokes = list(stroke_df.itertuples())
        rows = []
        temps: list[IndexedRow] = []

        for index in range(len(strokes)):
            cur = strokes[index]
            temps.append(IndexedRow(index, cur))

            is_last = index == len(strokes) - 1
            count = len(temps)
            if count == 3:
                temps = self._three_strokes(temps, is_last)
            elif count == 4:
                temps = self._four_strokes(rows, temps, is_last)
            elif count == 5:
                temps = self._five_strokes(rows, temps, is_last)
            elif count >= 6 and count % 2 == 0:
                temps = self._six_strokes(rows, temps, is_last)
            elif count >= 7 and count % 2 != 0:
                temps = self._seven_strokes(rows, temps, is_last)

            if is_last:
                rows.append(cur)

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

    def _three_strokes(self, temps, is_last):
        cur, nex_2 = temps[0], temps[2]
        return self._is_segment(temps, cur, nex_2, is_last)

    def _four_strokes(self, rows, temps, is_last):
        cur, nxt_1, nxt_3 = temps[0], temps[1], temps[3]
        return self._is_part_of_segment(rows, temps, cur, nxt_1, nxt_3, is_last)

    def _five_strokes(self, rows, temps, is_last):
        cur, nxt_4 = temps[0], temps[4]
        return self._is_segment(temps, cur, nxt_4, is_last)

    def _six_strokes(self, rows, temps, is_last):
        cur, nxt_3, nxt_5 = temps[0], temps[-3], temps[-1]
        return self._is_part_of_segment(rows, temps, cur, nxt_3, nxt_5, is_last)

    def _seven_strokes(self, rows, temps, is_last):
        cur, nxt_4, nxt_6 = temps[0], temps[-3], temps[-1]
        if (is_last or self.is_top(cur.row) and nxt_6.row.high >= nxt_4.row.high
                or self.is_bottom(cur.row) and nxt_6.row.low <= nxt_4.row.low):
            rows.append(cur.row)
            if self.is_top(cur.row):
                rows.append(self._lowest_of_middle(temps).row)
            else:
                rows.append(self._highest_of_middle(temps).row)
            return [nxt_6]
        return temps

    def _is_segment(self, temps, cur, nex, is_last):
        if (is_last or self.is_top(cur.row) and nex.row.high >= cur.row.high
                or self.is_bottom(cur.row) and nex.row.low <= cur.row.low):
            return [nex]
        return temps

    def _is_part_of_segment(self, rows, temps, cur, nxt, lst, is_last):
        if (is_last or self.is_top(cur.row) and lst.row.low <= nxt.row.low
                or self.is_bottom(cur.row) and lst.row.high >= nxt.row.high):
            rows.append(cur.row)
            return [lst]
        return temps

    @staticmethod
    def _lowest_of_middle(temps):
        lowest = None
        for index in range(3, len(temps) - 3):
            low = temps[index]
            if lowest is None or low.row.low < lowest.row.low:
                lowest = low
        return lowest

    @staticmethod
    def _highest_of_middle(temps):
        highest = None
        for index in range(3, len(temps) - 3):
            high = temps[index]
            if highest is None or high.row.high > highest.row.high:
                highest = high
        return highest
