import warnings; warnings.filterwarnings("ignore")
import pandas as pd
import numpy as np
import plotly.graph_objects as go

!pip install pandas numpy plotly

def generate_data(periods: int) -> pd.DataFrame:
    start_date     = pd.to_datetime("2024-01-01 00:00:00+00:00")
    datetime_index = pd.date_range(start_date, periods=periods, freq="1min", tz="UTC")
    price_data = []; last_close = 42000
    for i in range(periods):
        open_price  = last_close + np.random.normal(0, last_close * 0.0005)
        close_price = open_price + np.random.normal(0, last_close * 0.005)
        body_high   = max(open_price, close_price)
        body_low    = min(open_price, close_price)
        high_price  = max(body_high + abs(np.random.normal(0, last_close * 0.002)), open_price, close_price)
        low_price   = min(body_low  - abs(np.random.normal(0, last_close * 0.002)), open_price, close_price)
        if high_price < low_price: high_price, low_price = low_price, high_price
        price_data.append({"open": max(1,int(open_price)), "high": max(1,int(high_price)),
                            "low":  max(1,int(low_price)),  "close": max(1,int(close_price))})
        last_close = close_price
    df = pd.DataFrame(price_data, index=datetime_index)
    df.index.name = "datetime"
    df["volume"] = np.random.uniform(100.0, 500.0, periods)
    df["datetime"] = df.index.to_series()
    return df.reset_index(drop=True)

df = generate_data(500)
display(df.head())

def supertrend_strategy(
    df:         pd.DataFrame,
    window:     int   = 10,
    multiplier: float = 2.0,
) -> pd.DataFrame:
    df = df.copy().sort_values("datetime", ignore_index=True)

    hl2 = (df["high"] + df["low"]) / 2

    # True Range
    tr = pd.concat([
        df["high"] - df["low"],
        (df["high"] - df["close"].shift(1)).abs(),
        (df["low"]  - df["close"].shift(1)).abs(),
    ], axis=1).max(axis=1)
    atr = tr.rolling(window).mean()

    # Raw bands
    upper_raw = hl2 + multiplier * atr
    lower_raw = hl2 - multiplier * atr

    # Final bands with flip rule
    upper = upper_raw.copy()
    lower = lower_raw.copy()

    for i in range(1, len(df)):
        upper.iloc[i] = upper_raw.iloc[i] if (upper_raw.iloc[i] < upper.iloc[i-1] or \
                         df["close"].iloc[i-1] > upper.iloc[i-1]) else upper.iloc[i-1]
        lower.iloc[i] = lower_raw.iloc[i] if (lower_raw.iloc[i] > lower.iloc[i-1] or \
                         df["close"].iloc[i-1] < lower.iloc[i-1]) else lower.iloc[i-1]

    supertrend = pd.Series(np.nan, index=df.index)
    direction  = pd.Series(0,      index=df.index) # 1 for uptrend, -1 for downtrend

    # Find the first index where ATR (and thus upper/lower bands) is not NaN
    first_valid_idx = atr.first_valid_index()
    if first_valid_idx is None:
        # If no valid ATR, cannot compute Supertrend, return original df with NaN/0 for new columns
        df["supertrend"] = supertrend
        df["signal"] = direction
        return df

    # Initialize the first valid supertrend and direction
    # At the first valid point, if close is above the lower band, assume uptrend, else downtrend.
    if df["close"].iloc[first_valid_idx] > lower.iloc[first_valid_idx]:
        direction.iloc[first_valid_idx] = 1 # Initial direction: Uptrend
        supertrend.iloc[first_valid_idx] = lower.iloc[first_valid_idx] # Supertrend starts at lower band
    else:
        direction.iloc[first_valid_idx] = -1 # Initial direction: Downtrend
        supertrend.iloc[first_valid_idx] = upper.iloc[first_valid_idx] # Supertrend starts at upper band

    # Iterate from the next valid index to calculate Supertrend and generate signals
    for i in range(first_valid_idx + 1, len(df)):
        current_close = df["close"].iloc[i]
        prev_direction = direction.iloc[i-1]
        prev_supertrend = supertrend.iloc[i-1]

        current_upper = upper.iloc[i] # Current calculated upper band
        current_lower = lower.iloc[i] # Current calculated lower band

        # Determine the current direction (signal) and supertrend line
        if prev_direction == 1: # Previous trend was uptrend
            if current_close < prev_supertrend: # Price crossed below previous Supertrend line -> flip to downtrend
                direction.iloc[i] = -1
                supertrend.iloc[i] = current_upper # New Supertrend starts at current upper band
            else: # Price remained above previous Supertrend line -> continue uptrend
                direction.iloc[i] = 1
                supertrend.iloc[i] = max(current_lower, prev_supertrend) # Trail with lower band, cannot move down
        else: # Previous trend was downtrend (-1)
            if current_close > prev_supertrend: # Price crossed above previous Supertrend line -> flip to uptrend
                direction.iloc[i] = 1
                supertrend.iloc[i] = current_lower # New Supertrend starts at current lower band
            else: # Price remained below previous Supertrend line -> continue downtrend
                direction.iloc[i] = -1
                supertrend.iloc[i] = min(current_upper, prev_supertrend) # Trail with upper band, cannot move up

    df["supertrend"] = supertrend
    df["signal"]     = direction

    return df

df_signals = supertrend_strategy(df, window=10, multiplier=2.0)

print("--- Signal Distribution ---")
print(df_signals["signal"].value_counts())

buy_signals  = df_signals[df_signals["signal"] ==  1]
sell_signals = df_signals[df_signals["signal"] == -1]

fig = go.FigureWidget(data=[go.Candlestick(
    x=df_signals["datetime"],
    open=df_signals["open"], high=df_signals["high"],
    low=df_signals["low"],   close=df_signals["close"],
    name="Price"
)])

fig.add_trace(go.Scatter(
    x=df_signals["datetime"], y=df_signals["supertrend"],
    mode="lines", name="Supertrend",
    line=dict(color="purple", width=1.5)
))
fig.add_trace(go.Scatter(
    x=buy_signals["datetime"],  y=buy_signals["low"]  * 0.999,
    mode="markers", marker=dict(symbol="triangle-up",   size=10, color="green"), name="Buy (+1)"))
fig.add_trace(go.Scatter(
    x=sell_signals["datetime"], y=sell_signals["high"] * 1.001,
    mode="markers", marker=dict(symbol="triangle-down", size=10, color="red"),   name="Sell (−1)"))

fig.update_layout(
    title_text="Supertrend Strategy — Signals",
    xaxis_rangeslider_visible=False,
    xaxis_title="Datetime", yaxis_title="Price",
    height=600, yaxis=dict(autorange=True),
)
fig.show()