"""
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
"""

from datetime import datetime
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": 1,
        "buy_m3": 1,
        "buy_p1": 17,
        "buy_p2": 16,
        "buy_p3": 14,
    }

    # Sell hyperspace params:
    sell_params = {
        "sell_m1": 5,
        "sell_m2": 1,
        "sell_m3": 2,
        "sell_p1": 11,
        "sell_p2": 15,
        "sell_p3": 9,
    }

    # ROI table:
    minimal_roi = {
        "0": 0.088,
        "336": 0.061,
        "485": 0.029,
        "1463": 0
    }

    # Stoploss:
    stoploss = -0.294

    # Trailing stop:
    trailing_stop = True
    trailing_stop_positive = 0.012
    trailing_stop_positive_offset = 0.033
    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 leverage(self, pair: str, current_time: datetime, current_rate: float, proposed_leverage: float, max_leverage: float, entry_tag: str | None, side: str, **kwargs) -> float:
        return min(3.0, max_leverage)

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

        # 1) True‑Range
        dataframe["TR"] = ta.TRANGE(dataframe)

        # 2) ATR for every period that any IntParameter can output
        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:
            atr = ta.ATR(dataframe, timeperiod=p).bfill()   # back‑fill the leading NaNs
            dataframe[f"ATR_{p}"] = atr

        # 3) build Supertrend direction columns
        make_cols: dict[str, pd.Series] = {}
        is_hyperopt = self.config.get("runmode", "normal") == "hyperopt"

        if is_hyperopt:                       # every grid point
            grids = [
                (1, self.buy_m1.range,  self.buy_p1.range,  "buy"),
                (2, self.buy_m2.range,  self.buy_p2.range,  "buy"),
                (3, self.buy_m3.range,  self.buy_p3.range,  "buy"),
                (1, self.sell_m1.range, self.sell_p1.range, "sell"),
                (2, self.sell_m2.range, self.sell_p2.range, "sell"),
                (3, self.sell_m3.range, self.sell_p3.range, "sell"),
            ]
            for idx, m_range, p_range, side in grids:
                for m in m_range:
                    for p in p_range:
                        stx = self._supertrend(           # ← all positional after df
                            dataframe,
                            m,
                            p,
                            f"ATR_{p}"
                        )["STX"]
                        make_cols[f"supertrend_{idx}_{side}_{m}_{p}"] = stx
        else:                                 # just the six the strategy will use
            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:
                sti = self._supertrend(dataframe, m, p, f"ATR_{p}");
                stx = sti["STX"]
                st = sti["ST"]
                make_cols[f"supertrend_{idx}_{side}_{m}_{p}"] = stx
                make_cols[f"supertrend_{idx}_{side}_{m}_{p}_line"] = st

        # 4) join everything in one go (prevents fragmentation warnings)
        return dataframe.join(pd.DataFrame(make_cols, index=dataframe.index))

    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_1_sell_{self.sell_m1.value}_{self.sell_p1.value}"
                ]
                == "down"
            ),
            "exit_long",
        ] = 1

        dataframe.loc[
            (
                dataframe[f"supertrend_1_buy_{self.buy_m1.value}_{self.buy_p1.value}"]
                == "up"
            ),
            "exit_short",
        ] = 1

        return dataframe

    @staticmethod
    def _supertrend(
        df: DataFrame,
        multiplier: int,
        period: int,
        atr_col: str = "ATR"
    ) -> DataFrame:
        """
        Vectorised Supertrend implementation that **never touches rows where
        ATR is NaN**, so no NaNs are propagated into the final result.

        Returns a DataFrame with:
            - ST   : Supertrend line
            - STX  : 'up' / 'down' direction flag
        """
        high, low, close = df["high"].values, df["low"].values, df["close"].values
        atr = df[atr_col].values

        length     = len(df)
        st         = np.full(length, np.nan)
        final_ub   = np.full(length, np.nan)
        final_lb   = np.full(length, np.nan)
        stx        = np.full(length, None, dtype=object)

        hl2        = (high + low) / 2
        basic_ub   = hl2 + multiplier * atr
        basic_lb   = hl2 - multiplier * atr

        # ── find first index where ATR is valid ──────────────────────────────
        start = np.argmax(~np.isnan(atr))
        if np.isnan(atr[start]):        # whole column is NaN → return empty
            return DataFrame({"ST": st, "STX": stx}, index=df.index)

        # seed the first valid row: assume initial trend is down
        final_ub[start] = basic_ub[start]
        final_lb[start] = basic_lb[start]
        st[start]       = final_ub[start]
        stx[start]      = "down"

        # main loop
        for i in range(start + 1, length):
            # final upper / lower bands
            final_ub[i] = (
                basic_ub[i] if (basic_ub[i] < final_ub[i - 1] or close[i - 1] > final_ub[i - 1])
                else final_ub[i - 1]
            )
            final_lb[i] = (
                basic_lb[i] if (basic_lb[i] > final_lb[i - 1] or close[i - 1] < final_lb[i - 1])
                else final_lb[i - 1]
            )

            # Supertrend switch
            if   st[i - 1] == final_ub[i - 1] and close[i]  <= final_ub[i]:
                st[i] = final_ub[i]
            elif st[i - 1] == final_ub[i - 1] and close[i]  >  final_ub[i]:
                st[i] = final_lb[i]
            elif st[i - 1] == final_lb[i - 1] and close[i]  >= final_lb[i]:
                st[i] = final_lb[i]
            else:
                st[i] = final_ub[i]

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

        return DataFrame({"ST": st, "STX": stx}, index=df.index)
        
    @property
    def plot_config(self):
        return {
            'main_plot': {
                f"supertrend_1_buy_{self.buy_m1.value}_{self.buy_p1.value}_line": {'color': 'green'},
                f"supertrend_2_buy_{self.buy_m2.value}_{self.buy_p2.value}_line": {'color': 'green'},
                f"supertrend_3_buy_{self.buy_m3.value}_{self.buy_p3.value}_line": {'color': 'green'},
                f"supertrend_1_sell_{self.sell_m1.value}_{self.sell_p1.value}_line": {'color': 'red'},
                f"supertrend_2_sell_{self.sell_m2.value}_{self.sell_p2.value}_line": {'color': 'red'},
                f"supertrend_3_sell_{self.sell_m3.value}_{self.sell_p3.value}_line": {'color': 'red'},
            }
        }