import dask as da

from raptor.utils import BaseIndicator, _ema, _sma, _get_min_or_max


class AroonIndicator(BaseIndicator):

    def __init__(self, close: da.Series, window: int = 25, fillna: bool = False):
        self._close = close
        self._window = window
        self._fillna = fillna
        self._run()

    def _run(self):
        min_periods = 0 if self._fillna else self._window
        rolling_close = self._close.rolling(self._window, min_periods=min_periods)
        self._aroon_up = rolling_close.apply(
            lambda x: float(da.argmax(x) + 1) / self._window * 100, raw=True
        )
        self._aroon_down = rolling_close.apply(
            lambda x: float(da.argmin(x) + 1) / self._window * 100, raw=True
        )

    def aroon_up(self) -> da.Series:
        
        aroon_up_series = self._check_fillna(self._aroon_up, value=0)
        return da.Series(aroon_up_series, name=f"aroon_up_{self._window}")

    def aroon_down(self) -> da.Series:
        
        aroon_down_series = self._check_fillna(self._aroon_down, value=0)
        return da.Series(aroon_down_series, name=f"aroon_down_{self._window}")

    def aroon_indicator(self) -> da.Series:
        
        aroon_diff = self._aroon_up - self._aroon_down
        aroon_diff = self._check_fillna(aroon_diff, value=0)
        return da.Series(aroon_diff, name=f"aroon_ind_{self._window}")


class MACD(BaseIndicator):
    

    def __init__(
        self,
        close: da.Series,
        window_slow: int = 26,
        window_fast: int = 12,
        window_sign: int = 9,
        fillna: bool = False,
    ):
        self._close = close
        self._window_slow = window_slow
        self._window_fast = window_fast
        self._window_sign = window_sign
        self._fillna = fillna
        self._run()

    def _run(self):
        self._emafast = _ema(self._close, self._window_fast, self._fillna)
        self._emaslow = _ema(self._close, self._window_slow, self._fillna)
        self._macd = self._emafast - self._emaslow
        self._macd_signal = _ema(self._macd, self._window_sign, self._fillna)
        self._macd_diff = self._macd - self._macd_signal

    def macd(self) -> da.Series:
        
        macd_series = self._check_fillna(self._macd, value=0)
        return da.Series(
            macd_series, name=f"MACD_{self._window_fast}_{self._window_slow}"
        )

    def macd_signal(self) -> da.Series:
        

        macd_signal_series = self._check_fillna(self._macd_signal, value=0)
        return da.Series(
            macd_signal_series,
            name=f"MACD_sign_{self._window_fast}_{self._window_slow}",
        )

    def macd_diff(self) -> da.Series:
        
        macd_diff_series = self._check_fillna(self._macd_diff, value=0)
        return da.Series(
            macd_diff_series, name=f"MACD_diff_{self._window_fast}_{self._window_slow}"
        )


class EMAIndicator(BaseIndicator):
    
    def __init__(self, close: da.Series, window: int = 14, fillna: bool = False):
        self._close = close
        self._window = window
        self._fillna = fillna

    def ema_indicator(self) -> da.Series:
        
        ema_ = _ema(self._close, self._window, self._fillna)
        return da.Series(ema_, name=f"ema_{self._window}")


class SMAIndicator(BaseIndicator):
    def __init__(self, close: da.Series, window: int, fillna: bool = False):
        self._close = close
        self._window = window
        self._fillna = fillna

    def sma_indicator(self) -> da.Series:
        
        sma_ = _sma(self._close, self._window, self._fillna)
        return da.Series(sma_, name=f"sma_{self._window}")


class WMAIndicator(BaseIndicator):
    

    def __init__(self, close: da.Series, window: int = 9, fillna: bool = False):
        self._close = close
        self._window = window
        self._fillna = fillna
        self._run()

    def _run(self):
        _weight = da.Series(
            [
                i * 2 / (self._window * (self._window + 1))
                for i in range(1, self._window + 1)
            ]
        )

        def weighted_average(weight):
            def _weighted_average(x):
                return (weight * x).sum()

            return _weighted_average

        self._wma = self._close.rolling(self._window).apply(
            weighted_average(_weight), raw=True
        )

    def wma(self) -> da.Series:
        
        wma = self._check_fillna(self._wma, value=0)
        return da.Series(wma, name=f"wma_{self._window}")


class TRIXIndicator(BaseIndicator):
    

    def __init__(self, close: da.Series, window: int = 15, fillna: bool = False):
        self._close = close
        self._window = window
        self._fillna = fillna
        self._run()

    def _run(self):
        ema1 = _ema(self._close, self._window, self._fillna)
        ema2 = _ema(ema1, self._window, self._fillna)
        ema3 = _ema(ema2, self._window, self._fillna)
        self._trix = (ema3 - ema3.shift(1, fill_value=ema3.mean())) / ema3.shift(
            1, fill_value=ema3.mean()
        )
        self._trix *= 100

    def trix(self) -> da.Series:
        trix_series = self._check_fillna(self._trix, value=0)
        return da.Series(trix_series, name=f"trix_{self._window}")





class ADXIndicator(BaseIndicator):

    def __init__(
        self,
        high: da.Series,
        low: da.Series,
        close: da.Series,
        window: int = 14,
        fillna: bool = False,
    ):
        self._high = high
        self._low = low
        self._close = close
        self._window = window
        self._fillna = fillna
        self._run()

    def _run(self):
        if self._window == 0:
            raise ValueError("window may not be 0")

        close_shift = self._close.shift(1)
        pdm = _get_min_or_max(self._high, close_shift, "max")
        pdn = _get_min_or_max(self._low, close_shift, "min")
        diff_directional_movement = pdm - pdn

        self._trs_initial = da.zeros(self._window - 1)
        self._trs = da.zeros(len(self._close) - (self._window - 1))
        self._trs[0] = diff_directional_movement.dropna()[0 : self._window].sum()
        diff_directional_movement = diff_directional_movement.reset_index(drop=True)

        for i in range(1, len(self._trs) - 1):
            self._trs[i] = (
                self._trs[i - 1]
                - (self._trs[i - 1] / float(self._window))
                + diff_directional_movement[self._window + i]
            )

        diff_up = self._high - self._high.shift(1)
        diff_down = self._low.shift(1) - self._low
        pos = abs(((diff_up > diff_down) & (diff_up > 0)) * diff_up)
        neg = abs(((diff_down > diff_up) & (diff_down > 0)) * diff_down)

        self._dip = da.zeros(len(self._close) - (self._window - 1))
        self._dip[0] = pos.dropna()[0 : self._window].sum()

        pos = pos.reset_index(drop=True)

        for i in range(1, len(self._dip) - 1):
            self._dip[i] = (
                self._dip[i - 1]
                - (self._dip[i - 1] / float(self._window))
                + pos[self._window + i]
            )

        self._din = da.zeros(len(self._close) - (self._window - 1))
        self._din[0] = neg.dropna()[0 : self._window].sum()

        neg = neg.reset_index(drop=True)

        for i in range(1, len(self._din) - 1):
            self._din[i] = (
                self._din[i - 1]
                - (self._din[i - 1] / float(self._window))
                + neg[self._window + i]
            )

    def adx(self) -> da.Series:
        
        dip = da.zeros(len(self._trs))
        for i in range(len(self._trs)):
            dip[i] = 100 * (self._dip[i] / self._trs[i])

        din = da.zeros(len(self._trs))
        for i in range(len(self._trs)):
            din[i] = 100 * (self._din[i] / self._trs[i])

        directional_index = 100 * da.abs((dip - din) / (dip + din))

        adx_series = da.zeros(len(self._trs))
        adx_series[self._window] = directional_index[0 : self._window].mean()

        for i in range(self._window + 1, len(adx_series)):
            adx_series[i] = (
                (adx_series[i - 1] * (self._window - 1)) + directional_index[i - 1]
            ) / float(self._window)

        adx_series = da.concatenate((self._trs_initial, adx_series), axis=0)
        adx_series = da.Series(data=adx_series, index=self._close.index)

        adx_series = self._check_fillna(adx_series, value=20)
        return da.Series(adx_series, name="adx")

    def adx_pos(self) -> da.Series:
        dip = da.zeros(len(self._close))
        for i in range(1, len(self._trs) - 1):
            dip[i + self._window] = 100 * (self._dip[i] / self._trs[i])

        adx_pos_series = self._check_fillna(
            da.Series(dip, index=self._close.index), value=20
        )
        return da.Series(adx_pos_series, name="adx_pos")

    def adx_neg(self) -> da.Series:
        din = da.zeros(len(self._close))
        for i in range(1, len(self._trs) - 1):
            din[i + self._window] = 100 * (self._din[i] / self._trs[i])

        adx_neg_series = self._check_fillna(
            da.Series(din, index=self._close.index), value=20
        )
        return da.Series(adx_neg_series, name="adx_neg")



def ema_indicator(close, window=12, fillna=False):
    return EMAIndicator(close=close, window=window, fillna=fillna).ema_indicator()


def sma_indicator(close, window=12, fillna=False):
    return SMAIndicator(close=close, window=window, fillna=fillna).sma_indicator()


def wma_indicator(close, window=9, fillna=False):
    return WMAIndicator(close=close, window=window, fillna=fillna).wma()


def macd(close, window_slow=26, window_fast=12, fillna=False):
    return MACD(
        close=close,
        window_slow=window_slow,
        window_fast=window_fast,
        window_sign=9,
        fillna=fillna,
    ).macd()


def macd_signal(close, window_slow=26, window_fast=12, window_sign=9, fillna=False):
    return MACD(
        close=close,
        window_slow=window_slow,
        window_fast=window_fast,
        window_sign=window_sign,
        fillna=fillna,
    ).macd_signal()


def macd_diff(close, window_slow=26, window_fast=12, window_sign=9, fillna=False):
    return MACD(
        close=close,
        window_slow=window_slow,
        window_fast=window_fast,
        window_sign=window_sign,
        fillna=fillna,
    ).macd_diff()


def adx(high, low, close, window=14, fillna=False):
    return ADXIndicator(
        high=high, low=low, close=close, window=window, fillna=fillna
    ).adx()


def adx_pos(high, low, close, window=14, fillna=False):
    return ADXIndicator(
        high=high, low=low, close=close, window=window, fillna=fillna
    ).adx_pos()


def adx_neg(high, low, close, window=14, fillna=False):
    
    return ADXIndicator(
        high=high, low=low, close=close, window=window, fillna=fillna
    ).adx_neg()







def trix(close, window=15, fillna=False):
    
    return TRIXIndicator(close=close, window=window, fillna=fillna).trix()


def aroon_up(close, window=25, fillna=False):
    
    return AroonIndicator(close=close, window=window, fillna=fillna).aroon_up()


def aroon_down(close, window=25, fillna=False):
    return AroonIndicator(close=close, window=window, fillna=fillna).aroon_down()










