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:
    """
    Generate synthetic OHLCV price data using a geometric random walk.

    Parameters
    ----------
    periods : int
        Number of 1-minute bars to generate.

    Returns
    -------
    pd.DataFrame
        DataFrame with columns: open, high, low, close, volume, datetime.
    """
    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 volatility_adjusted_signal_filter(
    df: pd.DataFrame,
    atr_period: int = 14,
    low_pct: float = 25.0,
    high_pct: float = 75.0,
    base_signal_window: int = 20,
) -> pd.DataFrame:
    """
    Filter a Donchian breakout signal using an ATR-based volatility regime.

    Core logic
    ----------
    1. Compute ATR(atr_period) as the volatility measure.
    2. Classify volatility regime by comparing the current ATR to its
       rolling percentile (low_pct and high_pct thresholds).
    3. Generate a base Donchian breakout signal.
    4. Suppress the base signal whenever volatility is outside the normal regime.

    Parameters
    ----------
    df : pd.DataFrame   OHLCV DataFrame.
    atr_period : int    ATR rolling window.
    low_pct : float     ATR percentile below which volatility is 'low'.
    high_pct : float    ATR percentile above which volatility is 'high'.
    base_signal_window: int  Donchian channel window for the base signal.

    Returns
    -------
    pd.DataFrame with: atr, vol_regime, base_signal, signal.
    """
    df = df.copy().sort_values("datetime", ignore_index=True)

    # ── ATR ──────────────────────────────────────────────────────────────────
    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)
    df["atr"] = tr.rolling(atr_period).mean()

    # ── Volatility regime ─────────────────────────────────────────────────────
    rolling_low  = df["atr"].rolling(100).quantile(low_pct  / 100)
    rolling_high = df["atr"].rolling(100).quantile(high_pct / 100)
    df["vol_regime"] = "normal"
    df.loc[df["atr"] < rolling_low,  "vol_regime"] = "low"
    df.loc[df["atr"] > rolling_high, "vol_regime"] = "high"

    # ── Base signal: Donchian breakout ────────────────────────────────────────
    ch_high = df["high"].rolling(base_signal_window).max().shift(1)
    ch_low  = df["low"].rolling(base_signal_window).min().shift(1)
    df["base_signal"] = np.where(df["close"] > ch_high, 1,
                        np.where(df["close"] < ch_low, -1, 0))

    # ── Filtered signal: suppress outside normal regime ───────────────────────
    df["signal"] = np.where(df["vol_regime"] == "normal", df["base_signal"], 0)

    return df

df_signals = volatility_adjusted_signal_filter(df)
print(df_signals["vol_regime"].value_counts())
print(df_signals["signal"].value_counts())

buy_signals  = df_signals[df_signals["signal"] ==  1]
sell_signals = df_signals[df_signals["signal"] == -1]
regime_colors = {"normal": "rgba(0,200,0,0.08)", "low": "rgba(0,0,200,0.08)", "high": "rgba(200,0,0,0.08)"}
fig = make_subplots(rows=2, cols=1, shared_xaxes=True,
    subplot_titles=["Price + Volatility-Filtered Signals", "ATR + Regime"],
    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=buy_signals["datetime"],  y=buy_signals["low"]  * 0.999,
    mode="markers", marker=dict(symbol="triangle-up",   size=10, color="green"), name="Buy"),  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"), row=1, col=1)
atr_colors = ["green" if r == "normal" else "blue" if r == "low" else "red" for r in df_signals["vol_regime"]]
fig.add_trace(go.Bar(x=df_signals["datetime"], y=df_signals["atr"],
    marker_color=atr_colors, name="ATR"), row=2, col=1)
fig.update_layout(title_text="Volatility-Adjusted Signal Filter",
    xaxis_rangeslider_visible=False, height=700, xaxis2_title="Datetime")
fig.show()