import warnings
warnings.filterwarnings("ignore")

!pip install pandas numpy plotly
import pandas as pd
import numpy as np
import plotly.graph_objects as go

def generate_data(periods: int) -> pd.DataFrame:
    """
    Generates synthetic OHLCV data for backtesting purposes.

    Args:
        periods (int): The number of data points (minutes) to generate.

    Returns:
        pd.DataFrame: A DataFrame containing synthetic OHLCV data with a datetime index.
    """
    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)

def atr_volatility_breakout(
    df: pd.DataFrame,
    atr_window: int = 14,
    multiplier: float = 1.5,
) -> pd.DataFrame:
    """
    Calculates ATR volatility breakout signals for a given DataFrame.

    Args:
        df (pd.DataFrame): Input DataFrame with 'datetime', 'high', 'low', 'close' columns.
        atr_window (int): The lookback period for calculating ATR.
        multiplier (float): The multiplier for ATR to define trigger levels.

    Returns:
        pd.DataFrame: DataFrame with 'datetime', 'open', 'high', 'low', 'close', 'volume',
                      'atr', 'upper_trigger', 'lower_trigger', and 'signal' columns.
    """
    df = df.copy().sort_values("datetime", ignore_index=True)

    # Calculate True Range (TR)
    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)

    # Calculate Average True Range (ATR)
    df["atr"] = tr.rolling(atr_window).mean()

    # Define upper and lower trigger levels
    df["upper_trigger"] = df["close"].shift(1) + multiplier * df["atr"]
    df["lower_trigger"] = df["close"].shift(1) - multiplier * df["atr"]

    # Generate trading signals
    df["signal"] = np.where(df["close"] >= df["upper_trigger"],  1,
                          np.where(df["close"] <= df["lower_trigger"], -1, 0))

    return df

df_signals = atr_volatility_breakout(df, atr_window=14, multiplier=1.5)

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["upper_trigger"],
    mode="lines", name="Upper Trigger", line=dict(color="green", width=1, dash="dash")))
fig.add_trace(go.Scatter(x=df_signals["datetime"], y=df_signals["lower_trigger"],
    mode="lines", name="Lower Trigger", line=dict(color="red",   width=1, dash="dash")))
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="ATR Volatility Breakout Strategy",
    xaxis_rangeslider_visible=False, height=600, yaxis=dict(autorange=True))

fig.show()