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

!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 vwap_reversion_strategy(
    df:         pd.DataFrame,
    band_window:int   = 20,
    std_bands:  float = 1.5,
) -> pd.DataFrame:
    """Calculates VWAP, VWAP bands, and trading signals based on a reversion strategy.

    Args:
        df (pd.DataFrame): Input DataFrame containing 'high', 'low', 'close', 'volume', and 'datetime' columns.
        band_window (int): The rolling window for calculating the standard deviation of typical price from VWAP.
        std_bands (float): Multiplier for the standard deviation to set the width of the VWAP bands.

    Returns:
        pd.DataFrame: The original DataFrame with added 'vwap', 'vwap_upper', 'vwap_lower',
                      'vwap_distance_pct', and 'signal' columns.
    """
    # Create a copy to avoid modifying the original DataFrame and sort by datetime
    df = df.copy().sort_values("datetime", ignore_index=True)

    # Calculate Typical Price (TP)
    tp = (df["high"] + df["low"] + df["close"]) / 3

    # Calculate Cumulative VWAP (session-style)
    # VWAP = Sum(Typical Price * Volume) / Sum(Volume)
    df["vwap"] = (tp * df["volume"]).cumsum() / df["volume"].cumsum()

    # Calculate the rolling standard deviation of the typical price from VWAP
    # This is used to dynamically adjust the width of the VWAP bands based on recent volatility
    deviation = (tp - df["vwap"]).rolling(band_window).std()

    # Calculate the upper and lower VWAP bands
    # Upper band = VWAP + (std_bands * deviation)
    # Lower band = VWAP - (std_bands * deviation)
    df["vwap_upper"] = df["vwap"] + std_bands * deviation
    df["vwap_lower"] = df["vwap"] - std_bands * deviation

    # Calculate the percentage distance of the close price from VWAP
    # Useful for understanding how far price has diverged from the average
    df["vwap_distance_pct"] = (df["close"] - df["vwap"]) / df["vwap"] * 100

    # Generate trading signals:
    #   1: Buy signal (close price falls below the lower VWAP band)
    #  -1: Sell signal (close price rises above the upper VWAP band)
    #   0: No signal (close price is between the bands)
    df["signal"] = np.where(df["close"] < df["vwap_lower"],  1,
                   np.where(df["close"] > df["vwap_upper"], -1, 0))

    return df

df_signals = vwap_reversion_strategy(df, band_window=20, std_bands=1.5)

print("--- Signal Distribution ---")
print(df_signals["signal"].value_counts())
print("\n--- VWAP Distance Statistics ---")
print(df_signals["vwap_distance_pct"].describe().round(4))

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 + VWAP Bands + Signals", "VWAP Distance (%)"],
    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)

fig.add_trace(go.Scatter(x=df_signals["datetime"], y=df_signals["vwap"],
    mode="lines", name="VWAP", line=dict(color="blue", width=1.5)), row=1, col=1)
fig.add_trace(go.Scatter(x=df_signals["datetime"], y=df_signals["vwap_upper"],
    mode="lines", name="VWAP Upper", line=dict(color="orange", width=1, dash="dash")), row=1, col=1)
fig.add_trace(go.Scatter(x=df_signals["datetime"], y=df_signals["vwap_lower"],
    mode="lines", name="VWAP Lower", line=dict(color="orange", 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.Scatter(
    x=df_signals["datetime"], y=df_signals["vwap_distance_pct"],
    mode="lines", name="VWAP Distance %", line=dict(color="purple", width=1)), row=2, col=1)
fig.add_hline(y=0, line_dash="dot", line_color="gray", row=2, col=1)

fig.update_layout(
    title_text="VWAP Reversion Strategy",
    xaxis_rangeslider_visible=False,
    height=700, yaxis=dict(autorange=True),
    xaxis2_title="Datetime",
)
fig.show()