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

import pandas as pd

from pychanlun.segment import Segment


@dataclass
class Range:
    start: Any
    end: Any
    high: float
    low: float


class Pivot(Segment):

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

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

        source_df = self.sources[interval]
        stroke_df = self.strokes[interval]
        if source_df is not None or stroke_df is not None:
            self.stroke_pivots[interval] = self._generate_pivot_interval(source_df, stroke_df)
        segment_df = self.segments[interval]
        if source_df is not None or segment_df is not None:
            self.segment_pivots[interval] = self._generate_pivot_interval(source_df, segment_df)

    def _generate_pivot_interval(self, source_df, segment_df):
        segment_df = segment_df.copy()
        segment_df['price'] = segment_df['macd'] = segment_df['trend'] = segment_df['divergence'] = 0
        segments = list(segment_df.itertuples())
        rows = []

        index = 0
        while index < len(segments) - 4:
            cur = segments[index]
            range_1 = self.get_range(segments, index + 1)
            range_2 = self.get_range(segments, index + 3)
            pivot_1 = self._get_pivot(range_1, range_2)

            if self._can_start_pivot(pivot_1, cur) and self._has_range_overlap(range_1, range_2):
                step = 0
                while index + step < len(segments) - 7:
                    range_3 = self.get_range(segments, index + step + 5)
                    if self._is_out_of_pivot(pivot_1, range_3):
                        break
                    step += 2

                lst = segments[index + step + 4]
                if self.is_top(range_1.start):
                    range_1.start = range_1.start._replace(high=pivot_1.high)
                    lst = lst._replace(low=pivot_1.low)
                else:
                    range_1.start = range_1.start._replace(low=pivot_1.low)
                    lst = lst._replace(high=pivot_1.high)

                rows.append(range_1.start)
                rows.append(lst)
                index += step + 4
            else:
                index += 1

        self._merge_pivots(rows)
        self._pivot_trend(rows, source_df)

        return self.to_dataframe(rows, ['high', 'low', 'price', 'macd', 'trend', 'divergence'])

    def _merge_pivots(self, rows):
        index = 0
        while index < len(rows) - 3:
            pivot_1 = self.get_range(rows, index)
            pivot_2 = self.get_range(rows, index + 2)

            if self._has_range_overlap(pivot_1, pivot_2):
                if self.is_top(pivot_1.start):
                    pivot_1.start = pivot_1.start._replace(high=min(pivot_1.start.high, pivot_2.start.high))
                    pivot_2.end = pivot_2.end._replace(low=max(pivot_1.end.low, pivot_2.end.low))
                else:
                    pivot_1.start = pivot_1.start._replace(low=max(pivot_1.start.low, pivot_2.start.low))
                    pivot_2.end = pivot_2.end._replace(high=min(pivot_1.end.high, pivot_2.end.high))

                rows[index] = pivot_1.start
                rows[index + 3] = pivot_2.end
                rows.remove(pivot_1.end)
                rows.remove(pivot_2.start)
            else:
                index += 2

    def _pivot_trend(self, rows, source_df):
        level = 0
        for index in range(0, len(rows) - 3, 2):
            pivot_1 = self.get_range(rows, index)
            pivot_2 = self.get_range(rows, index + 2)

            self._set_macd(pivot_1, pivot_2, source_df)
            level = self._set_trend(pivot_1, pivot_2, level)

            rows[index] = pivot_1.start
            rows[index + 1] = pivot_1.end
            rows[index + 2] = pivot_2.start
            rows[index + 3] = pivot_2.end

        if len(rows) > 0:
            pivot_1 = self.get_range(rows, -2)
            self._set_macd(pivot_1, None, source_df)
            rows[-1] = pivot_1.end

    @staticmethod
    def _set_macd(pivot_1, pivot_2, source_df):
        macd_sum = source_df[(pivot_1 is None or source_df.index >= pivot_1.end.Index) &
                             (pivot_2 is None or source_df.index <= pivot_2.start.Index)].macd_dif.sum()
        if pivot_1 is not None:
            pivot_1.end = pivot_1.end._replace(macd=macd_sum)
        if pivot_2 is not None:
            pivot_2.start = pivot_2.start._replace(macd=macd_sum)

    @staticmethod
    def _set_trend(pivot_1, pivot_2, trend):
        if trend == 0:
            pivot_1.start = pivot_1.start._replace(trend=trend)
            pivot_1.end = pivot_1.end._replace(trend=trend + 1)

        if pivot_2.high > pivot_1.high:
            pivot_1.end = pivot_1.end._replace(price=pivot_1.low)
            if trend < 0:
                trend = 0
            trend += 1
            pivot_2.start = pivot_2.start._replace(trend=trend, price=pivot_2.high)
            pivot_2.end = pivot_2.end._replace(trend=trend + 1)
        else:
            pivot_1.end = pivot_1.end._replace(price=pivot_1.high)
            if trend > 0:
                trend = 0
            trend -= 1
            pivot_2.start = pivot_2.start._replace(trend=trend, price=pivot_2.low)
            pivot_2.end = pivot_2.end._replace(trend=trend - 1)

        if pivot_1.start.trend < pivot_2.start.trend > 0:
            pivot_1.start = pivot_1.start._replace(divergence=(1 if pivot_1.start.macd > pivot_1.end.macd else 0))
        if pivot_1.start.trend > pivot_2.start.trend < 0:
            pivot_1.start = pivot_1.start._replace(divergence=(1 if pivot_1.start.macd < pivot_1.end.macd else 0))
        pivot_1.end = pivot_1.end._replace(divergence=pivot_1.start.divergence)
        return trend

    @staticmethod
    def _get_pivot(range_1, range_2):
        return Range(range_1.start, range_2.end, min(range_1.high, range_2.high), max(range_1.low, range_2.low))

    def _can_start_pivot(self, pvt, cur):
        return (self.is_top(cur) and cur.high > pvt.high) or (self.is_bottom(cur) and cur.low < pvt.low)

    @staticmethod
    def _has_range_overlap(range_1, range_2):
        return range_2.high > range_1.low and range_2.low < range_1.high

    @staticmethod
    def _is_out_of_pivot(range_1, range_2):
        return range_2.high < range_1.low or range_2.low > range_1.high

    def get_range(self, rows, index):
        start, end = rows[index], rows[index + 1]
        high = start.high if self.is_top(start) else end.high
        low = start.low if self.is_bottom(start) else end.low
        return Range(start, end, high, low)
