"""
Supertrend strategy:
* Description: Generate a 3 supertrend indicators for 'buy' strategies & 3 supertrend indicators for 'sell' strategies
               Buys if the 3 'buy' indicators are 'up'
               Sells if the 3 'sell' indicators are 'down'
* Author: @juankysoriano (Juan Carlos Soriano)
* github: https://github.com/juankysoriano/
*** NOTE: This Supertrend strategy is just one of many possible strategies using `Supertrend` as indicator. It should on any case used at your own risk.
          It comes with at least a couple of caveats:
            1. The implementation for the `supertrend` indicator is based on the following discussion: https://github.com/freqtrade/freqtrade-strategies/issues/30 . Concretelly https://github.com/freqtrade/freqtrade-strategies/issues/30#issuecomment-853042401
            2. The implementation for `supertrend` on this strategy is not validated; meaning this that is not proven to match the results by the paper where it was originally introduced or any other trusted academic resources
"""

import logging
from numpy.lib import math
from freqtrade.strategy import IStrategy, IntParameter
from pandas import DataFrame
import talib.abstract as ta
import numpy as np
import pandas as pd


class FSupertrendStrategy(IStrategy):
    # Buy params, Sell params, ROI, Stoploss and Trailing Stop are values generated by 'freqtrade hyperopt --strategy Supertrend --hyperopt-loss ShortTradeDurHyperOptLoss --timerange=20210101- --timeframe=1h --spaces all'
    # It's encourage you find the values that better suites your needs and risk management strategies

    INTERFACE_VERSION: int = 3
    can_short = True
    # Buy hyperspace params:
    buy_params = {
        "buy_m1": 2,
        "buy_m2": 3,
        "buy_m3": 2,
        "buy_p1": 9,
        "buy_p2": 15,
        "buy_p3": 9,
    }

    # Sell hyperspace params:
    sell_params = {
        "sell_m1": 2,
        "sell_m2": 4,
        "sell_m3": 2,
        "sell_p1": 18,
        "sell_p2": 18,
        "sell_p3": 8,
    }

    # ROI table:
    minimal_roi = {
        "0": 0.139,
        "203": 0.046,
        "496": 0.023,
        "950": 0
    }

    # Stoploss:
    stoploss = -0.162

    # Trailing stop:
    trailing_stop = True
    trailing_stop_positive = 0.013
    trailing_stop_positive_offset = 0.03
    trailing_only_offset_is_reached = True

    timeframe = "1h"

    startup_candle_count = 18

    buy_m1 = IntParameter(1, 7, default=1)
    buy_m2 = IntParameter(1, 7, default=3)
    buy_m3 = IntParameter(1, 7, default=4)
    buy_p1 = IntParameter(7, 21, default=14)
    buy_p2 = IntParameter(7, 21, default=10)
    buy_p3 = IntParameter(7, 21, default=10)

    sell_m1 = IntParameter(1, 7, default=1)
    sell_m2 = IntParameter(1, 7, default=3)
    sell_m3 = IntParameter(1, 7, default=4)
    sell_p1 = IntParameter(7, 21, default=14)
    sell_p2 = IntParameter(7, 21, default=10)
    sell_p3 = IntParameter(7, 21, default=10)

    def populate_indicators(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
        # 1) True Range once
        tr = ta.TRANGE(dataframe)
        dataframe["TR"] = tr

        # 2) ATR for every period that might be used
        all_periods = (
            set(self.buy_p1.range)   | set(self.buy_p2.range)   | set(self.buy_p3.range) |
            set(self.sell_p1.range)  | set(self.sell_p2.range)  | set(self.sell_p3.range)
        )
        for p in all_periods:
            dataframe[f"ATR_{p}"] = ta.SMA(tr, timeperiod=p)

        # 3) Collect new columns in a dict
        new_cols: dict[str, Series] = {}

        # Use runmode to decide full-grid vs single-trial
        is_hyperopt = self.config.get("runmode", "normal") == "hyperopt"

        if is_hyperopt:
            # Full grid in hyperopt
            for idx, (m_range, p_range) in enumerate([
                (self.buy_m1.range, self.buy_p1.range),
                (self.buy_m2.range, self.buy_p2.range),
                (self.buy_m3.range, self.buy_p3.range),
            ], start=1):
                for m in m_range:
                    for p in p_range:
                        stx = self._supertrend(
                            dataframe, multiplier=m, period=p,
                            tr_col="TR", atr_col=f"ATR_{p}"
                        )["STX"]
                        new_cols[f"supertrend_{idx}_buy_{m}_{p}"] = stx

            for idx, (m_range, p_range) in enumerate([
                (self.sell_m1.range, self.sell_p1.range),
                (self.sell_m2.range, self.sell_p2.range),
                (self.sell_m3.range, self.sell_p3.range),
            ], start=1):
                for m in m_range:
                    for p in p_range:
                        stx = self._supertrend(
                            dataframe, multiplier=m, period=p,
                            tr_col="TR", atr_col=f"ATR_{p}"
                        )["STX"]
                        new_cols[f"supertrend_{idx}_sell_{m}_{p}"] = stx

        else:
            # Only the six needed for live/backtest
            combos = [
                (1, self.buy_m1.value,  self.buy_p1.value,  "buy"),
                (2, self.buy_m2.value,  self.buy_p2.value,  "buy"),
                (3, self.buy_m3.value,  self.buy_p3.value,  "buy"),
                (1, self.sell_m1.value, self.sell_p1.value, "sell"),
                (2, self.sell_m2.value, self.sell_p2.value, "sell"),
                (3, self.sell_m3.value, self.sell_p3.value, "sell"),
            ]
            for idx, m, p, side in combos:
                stx = self._supertrend(
                    dataframe, multiplier=m, period=p,
                    tr_col="TR", atr_col=f"ATR_{p}"
                )["STX"]
                new_cols[f"supertrend_{idx}_{side}_{m}_{p}"] = stx

        # 4) One big concat to avoid fragmentation
        indicators = pd.DataFrame(new_cols, index=dataframe.index)
        return dataframe.join(indicators)

    def populate_entry_trend(self, dataframe: DataFrame, metadata: dict) -> DataFrame:

        dataframe.loc[
            (
                dataframe[f"supertrend_1_buy_{self.buy_m1.value}_{self.buy_p1.value}"]
                == "up"
            )
            & (
                dataframe[f"supertrend_2_buy_{self.buy_m2.value}_{self.buy_p2.value}"]
                == "up"
            )
            & (
                dataframe[f"supertrend_3_buy_{self.buy_m3.value}_{self.buy_p3.value}"]
                == "up"
            )
            & (  # The three indicators are 'up' for the current candle
                dataframe["volume"] > 0
            ),
            "enter_long",
        ] = 1

        dataframe.loc[
            (
                dataframe[
                    f"supertrend_1_sell_{self.sell_m1.value}_{self.sell_p1.value}"
                ]
                == "down"
            )
            & (
                dataframe[
                    f"supertrend_2_sell_{self.sell_m2.value}_{self.sell_p2.value}"
                ]
                == "down"
            )
            & (
                dataframe[
                    f"supertrend_3_sell_{self.sell_m3.value}_{self.sell_p3.value}"
                ]
                == "down"
            )
            & (  # The three indicators are 'down' for the current candle
                dataframe["volume"] > 0
            ),
            "enter_short",
        ] = 1

        return dataframe

    def populate_exit_trend(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
        dataframe.loc[
            (
                dataframe[
                    f"supertrend_2_sell_{self.sell_m2.value}_{self.sell_p2.value}"
                ]
                == "down"
            ),
            "exit_long",
        ] = 1

        dataframe.loc[
            (
                dataframe[f"supertrend_2_buy_{self.buy_m2.value}_{self.buy_p2.value}"]
                == "up"
            ),
            "exit_short",
        ] = 1

        return dataframe

    """
        Supertrend Indicator; adapted for freqtrade
        from: https://github.com/freqtrade/freqtrade-strategies/issues/30
    """

    def _supertrend(
        self,
        df: DataFrame,
        multiplier: int,
        period: int,
        tr_col: str = "TR",
        atr_col: str = None
    ) -> DataFrame:
        """
        Fast Supertrend calculation using precomputed TR and ATR columns.
        Returns a DataFrame with columns:
        - ST:  the Supertrend value
        - STX: 'up' or 'down' signal
        """
        if atr_col is None:
            atr_col = f"ATR_{period}"

        high = df["high"].values
        low  = df["low"].values
        close = df["close"].values
        tr   = df[tr_col].values
        atr  = df[atr_col].values

        length = len(df)
        final_ub = np.zeros(length)
        final_lb = np.zeros(length)
        st       = np.zeros(length)
        stx      = np.zeros(length, dtype=object)

        # basic bands
        basic_ub = (high + low) / 2 + multiplier * atr
        basic_lb = (high + low) / 2 - multiplier * atr

        # seed first values
        final_ub[:period] = basic_ub[:period]
        final_lb[:period] = basic_lb[:period]
        st[:period] = basic_ub[:period]
        stx[:period] = "down"  # or match your original default

        # rolling calculation
        for i in range(period, length):
            prev_ub = final_ub[i - 1]
            prev_lb = final_lb[i - 1]
            prev_st = st[i - 1]
            prev_close = close[i - 1]

            # final upper band
            if basic_ub[i] < prev_ub or prev_close > prev_ub:
                final_ub[i] = basic_ub[i]
            else:
                final_ub[i] = prev_ub

            # final lower band
            if basic_lb[i] > prev_lb or prev_close < prev_lb:
                final_lb[i] = basic_lb[i]
            else:
                final_lb[i] = prev_lb

            # Supertrend line
            if prev_st == prev_ub and close[i] <= final_ub[i]:
                st[i] = final_ub[i]
            elif prev_st == prev_ub and close[i] > final_ub[i]:
                st[i] = final_lb[i]
            elif prev_st == prev_lb and close[i] >= final_lb[i]:
                st[i] = final_lb[i]
            else:
                st[i] = final_ub[i]

            # direction flag
            stx[i] = "down" if close[i] < st[i] else "up"

        return DataFrame({
            "ST":  st,
            "STX": stx
        }, index=df.index)