from dataclasses import dataclass
from typing import Dict, Optional, Any

import pandas as pd

from pychanlun.fractal import Fractal


@dataclass
class IndexedRow:
    index: int
    row: Any


class Stroke(Fractal):

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

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

        fractal_df = self.fractals[interval]
        if fractal_df is None:
            return

        fractals = list(fractal_df.itertuples())
        rows = []
        temps: list[IndexedRow] = []

        for index in range(len(fractals)):
            cur = fractals[index]

            if self.is_top(cur) or self.is_bottom(cur):
                temps.append(IndexedRow(index, cur))

                is_last = index == len(fractals) - 1
                count = len(temps)
                if count == 2:
                    temps = self._two_fractals(rows, temps, is_last)
                elif count == 3:
                    temps = self._three_fractals(temps, is_last)
                elif count == 4:
                    temps = self._four_fractals(rows, temps, is_last)
                elif count == 5:
                    temps = self._five_fractals(rows, temps, is_last)
                elif count == 6:
                    temps = self._six_fractals(rows, temps)

                if is_last:
                    rows.append(cur)

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

    def _two_fractals(self, rows, temps, is_last):
        cur, nxt_1 = temps[0], temps[1]
        return self._is_stroke(rows, temps, cur, nxt_1, is_last)

    def _three_fractals(self, temps, is_last):
        cur, nxt_2 = temps[0], temps[2]
        return self._is_part_of_stroke(temps, cur, nxt_2, is_last)

    def _four_fractals(self, rows, temps, is_last):
        cur, nxt_3 = temps[0], temps[3]
        return self._is_stroke(rows, temps, cur, nxt_3, is_last)

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

    @staticmethod
    def _six_fractals(rows, temps):
        cur, nxt_5 = temps[0], temps[5]
        rows.append(cur.row)
        return [nxt_5]

    @staticmethod
    def _is_stroke(rows, temps, cur, nxt, is_last):
        if is_last or nxt.index - cur.index >= 4:
            rows.append(cur.row)
            return [nxt]
        return temps

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