!pip install pandas numpy plotly

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

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)

def volatility_squeeze_strategy(
    df: pd.DataFrame,
    window: int = 20,
    bb_std: float = 2.0,
    kc_mult: float = 1.5,
) -> pd.DataFrame:
    """
    Applies the Bollinger-Keltner Squeeze strategy to financial time series data.

    This function calculates Bollinger Bands and Keltner Channels to identify periods
    of low volatility (a 'squeeze'). It then generates buy or sell signals when
    the squeeze releases, based on the prevailing momentum.

    Args:
        df (pd.DataFrame): Input DataFrame with 'datetime', 'open', 'high', 'low', 'close' columns.
        window (int): The lookback period for calculating EMA, Bollinger Bands, and ATR.
        bb_std (float): The number of standard deviations for Bollinger Bands.
        kc_mult (float): The multiplier for Average True Range (ATR) in Keltner Channels.

    Returns:
        pd.DataFrame: The original DataFrame with added columns for Bollinger Bands,
                      Keltner Channels, squeeze condition, momentum, and trade signals.
    """
    # Ensure data is sorted by datetime for correct rolling calculations
    df = df.copy().sort_values("datetime", ignore_index=True)

    # Calculate Exponential Moving Average (EMA) as the centerline
    df["ema"]      = df["close"].ewm(span=window, adjust=False).mean()

    # Calculate Bollinger Bands
    bb_std_val     = df["close"].rolling(window).std()
    df["bb_upper"] = df["ema"] + bb_std  * bb_std_val # Upper Bollinger Band
    df["bb_lower"] = df["ema"] - bb_std  * bb_std_val # Lower Bollinger Band

    # Calculate True Range (TR) and Average True Range (ATR) for Keltner Channels
    tr = pd.concat([
        df["high"] - df["low"], # High minus Low
        (df["high"] - df["close"].shift(1)).abs(), # High minus previous close (absolute value)
        (df["low"]  - df["close"].shift(1)).abs(),  # Low minus previous close (absolute value)
    ], axis=1).max(axis=1)
    atr            = tr.rolling(window).mean() # Average True Range

    # Calculate Keltner Channels
    df["kc_upper"] = df["ema"] + kc_mult * atr # Upper Keltner Channel
    df["kc_lower"] = df["ema"] - kc_mult * atr # Lower Keltner Channel

    # Determine the squeeze condition: BB fully inside KC
    df["squeeze"]  = (df["bb_upper"] < df["kc_upper"]) & (df["bb_lower"] > df["kc_lower"])

    # Calculate momentum based on close price relative to its moving average
    df["momentum"] = df["close"] - df["close"].rolling(window).mean()

    # Generate trade signals based on squeeze release and momentum direction
    # Signal is 1 for buy, -1 for sell, and 0 if squeeze is active or no clear breakout
    df["signal"]   = np.where(~df["squeeze"] & (df["momentum"] > 0),  1, # Buy if squeeze ends and momentum is positive
                     np.where(~df["squeeze"] & (df["momentum"] < 0), -1, 0)) # Sell if squeeze ends and momentum is negative
    return df

df_signals = volatility_squeeze_strategy(df)

print("--- Signal Distribution ---"); print(df_signals["signal"].value_counts())
print(f"Squeeze bars: {df_signals['squeeze'].sum()}")

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

fig = make_subplots(rows=2, cols=1, shared_xaxes=True,
    subplot_titles=["Price + Bands + Signals", "Momentum + Squeeze Indicator"],
    row_heights=[0.65, 0.35])
fig.add_trace(go.Candlestick(x=df_signals["datetime"],
    open=df_signals["open"], high=df_signals["high"],
    low=df_signals["low"],   close=df_signals["close"], name="Price"), row=1, col=1)
for col, color, name in [("bb_upper","blue","BB Upper"),("bb_lower","blue","BB Lower"),
                          ("kc_upper","orange","KC Upper"),("kc_lower","orange","KC Lower")]:
    fig.add_trace(go.Scatter(x=df_signals["datetime"], y=df_signals[col],
        mode="lines", name=name, line=dict(color=color, width=1, dash="dash")), row=1, col=1)
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)"), row=1, col=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)"), row=1, col=1)
fig.add_trace(go.Bar(x=df_signals["datetime"], y=df_signals["momentum"],
    name="Momentum",
    marker_color=["green" if v >= 0 else "red" for v in df_signals["momentum"]]), row=2, col=1)
fig.add_trace(go.Scatter(x=df_signals["datetime"],
    y=df_signals["squeeze"].astype(int) * df_signals["momentum"].abs().max() * 0.5,
    mode="markers", marker=dict(color=["gray" if s else "rgba(0,0,0,0)" for s in df_signals["squeeze"]], size=4),
    name="Squeeze Active"), row=2, col=1)
fig.update_layout(title_text="Bollinger-Keltner Squeeze Strategy",
    xaxis_rangeslider_visible=False, height=700, yaxis=dict(autorange=True))
fig.show()