import pandas as pd

from pychanlun.signal import Signal


class Chan(Signal):

    def get_sources(self, interval):
        return self.sources[interval]

    def get_sticks(self, interval):
        source_df = self.sources[interval]
        stick_df = self.sticks[interval]
        if source_df is None or stick_df is None:
            return None

        source_df = source_df[[]].join(stick_df, how='left')
        return source_df

    def get_fractals(self, interval):
        fractal_df = self.fractals[interval]
        if fractal_df is None:
            return None

        fractal_df = fractal_df[fractal_df['high'].notna() | fractal_df['low'].notna()]
        return fractal_df[['high', 'low']]

    def get_strokes(self, interval):
        stroke_df = self.strokes[interval]
        if stroke_df is None:
            return None

        stroke_df['stroke'] = stroke_df['high'].fillna(stroke_df['low'])
        return stroke_df[['stroke']]

    def get_stroke_pivots(self, interval):
        trend_df = self.stroke_pivots[interval]
        if trend_df is None:
            return None

        return self._get_pivots(trend_df)

    def get_stroke_pivot_trends(self, interval):
        trend_df = self.stroke_pivots[interval]
        if trend_df is None:
            return None

        return self._get_trends(trend_df[1:-1])

    def get_stroke_pivot_signals(self, interval):
        signal_df = self.stroke_signals[interval]
        if signal_df is None:
            return None

        return self._get_signals(signal_df)

    def get_segments(self, interval):
        segment_df = self.segments[interval]
        if segment_df is None:
            return None

        segment_df['segment'] = segment_df['high'].fillna(segment_df['low'])
        return segment_df[['segment']]

    def get_segment_pivots(self, interval):
        trend_df = self.segment_pivots[interval]
        if trend_df is None:
            return None

        return self._get_pivots(trend_df)

    def get_segment_pivot_trends(self, interval):
        trend_df = self.segment_pivots[interval]
        if trend_df is None:
            return None

        return self._get_trends(trend_df[1:-1])

    def get_segment_pivot_signals(self, interval):
        signal_df = self.segment_signals[interval]
        if signal_df is None:
            return None

        return self._get_signals(signal_df)

    @staticmethod
    def _get_pivots(df):
        df = pd.DataFrame({
            'datetime': df.index.values[::2],
            'exit': df.index.values[1::2],
            'entry_high': df['high'].values[::2],
            'exit_high': df['high'].values[1::2],
            'entry_low': df['low'].values[::2],
            'exit_low': df['low'].values[1::2],
            'entry_macd': df['macd'].values[::2],
            'exit_macd': df['macd'].values[1::2],
            'trend': df['trend'].values[::2],
            'divergence': df['divergence'].values[::2]
        })
        df['high'] = df['entry_high'].fillna(df['exit_high'])
        df['low'] = df['entry_low'].fillna(df['exit_low'])
        df.set_index('datetime', inplace=True)
        return df[['exit', 'high', 'low', 'entry_macd', 'exit_macd', 'trend', 'divergence']]

    @staticmethod
    def _get_trends(df):
        df = pd.DataFrame({
            'datetime': df.index.values[::2],
            'exit': df.index.values[1::2],
            'entry_price': df['price'].values[::2],
            'exit_price': df['price'].values[1::2]
        })
        df.set_index('datetime', inplace=True)
        return df[['exit', 'entry_price', 'exit_price']]

    @staticmethod
    def _get_signals(df):
        df['price'] = df['high'].fillna(df['low'])
        return df[['price', 'signal']]
