!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:
    """
    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 orderblock_detection_basic(
    df: pd.DataFrame,
    atr_window: int = 14,
    impulse_factor: float = 1.5,
    lookback: int = 5,
) -> pd.DataFrame:
    """
    Detect basic bullish and bearish order blocks and generate re-entry signals.

    Core logic
    ----------
    1. Compute ATR to measure average candle size.
    2. For each bar: if the body exceeds impulse_factor × ATR (impulse candle),
       scan the prior lookback bars for the last candle of the opposite direction.
    3. Record the OB zone (high, low of the qualifying candle).
    4. For each bar: if price enters a live OB zone, emit the corresponding signal.

    Parameters
    ----------
    df : pd.DataFrame
        OHLCV DataFrame with columns: open, high, low, close, volume, datetime.
    atr_window : int
        Rolling window for ATR calculation.
    impulse_factor : float
        Multiplier; candle body must exceed this multiple of ATR to be an impulse.
    lookback : int
        Maximum number of bars to look back for the preceding OB candle.

    Returns
    -------
    pd.DataFrame
        Original DataFrame extended with: atr, ob_type, ob_high, ob_low, signal.
    """
    df = df.copy().sort_values("datetime", ignore_index=True)
    df["signal"]  = 0
    df["ob_type"] = "none"
    df["ob_high"] = np.nan
    df["ob_low"]  = np.nan

    # ── 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_window).mean()

    df["body"] = (df["close"] - df["open"]).abs()

    # ── Track active OB zones ─────────────────────────────────────────────────
    bullish_obs = []  # list of (ob_high, ob_low)
    bearish_obs = []

    for i in range(atr_window + lookback, len(df)):
        atr_val = df["atr"].iloc[i]
        if np.isnan(atr_val):
            continue

        body_i = df["body"].iloc[i]
        is_bull = df["close"].iloc[i] > df["open"].iloc[i]  # current candle direction

        # ── Impulse detection ─────────────────────────────────────────────────
        if body_i > impulse_factor * atr_val:
            # Search lookback bars for the last opposite-direction candle
            for j in range(i - 1, max(i - lookback - 1, 0), -1):
                prev_bull = df["close"].iloc[j] > df["open"].iloc[j]
                if is_bull and not prev_bull:
                    # Last bearish candle before bullish impulse → Bullish OB
                    bullish_obs.append((df["high"].iloc[j], df["low"].iloc[j]))
                    break
                elif not is_bull and prev_bull:
                    # Last bullish candle before bearish impulse → Bearish OB
                    bearish_obs.append((df["high"].iloc[j], df["low"].iloc[j]))
                    break

        # ── Signal: price re-enters an OB zone ───────────────────────────────
        c = df["close"].iloc[i]

        for ob_h, ob_l in bullish_obs:
            if ob_l <= c <= ob_h:
                df.at[df.index[i], "signal"]  = 1
                df.at[df.index[i], "ob_type"] = "bullish_ob"
                df.at[df.index[i], "ob_high"] = ob_h
                df.at[df.index[i], "ob_low"]  = ob_l
                break

        if df["signal"].iloc[i] == 0:
            for ob_h, ob_l in bearish_obs:
                if ob_l <= c <= ob_h:
                    df.at[df.index[i], "signal"]  = -1
                    df.at[df.index[i], "ob_type"] = "bearish_ob"
                    df.at[df.index[i], "ob_high"] = ob_h
                    df.at[df.index[i], "ob_low"]  = ob_l
                    break

    return df

df_signals = orderblock_detection_basic(df, atr_window=14, impulse_factor=1.5)

print("--- OB Type Distribution ---")
print(df_signals["ob_type"].value_counts())
print("\n--- 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 = make_subplots(rows=2, cols=1, shared_xaxes=True,
    subplot_titles=["Price + Order Block Signals", "ATR"],
    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="Bullish OB (+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="Bearish OB (-1)"), row=1, col=1)

fig.add_trace(go.Scatter(
    x=df_signals["datetime"], y=df_signals["atr"],
    mode="lines", name="ATR", line=dict(color="orange", width=1)),
    row=2, col=1)

fig.update_layout(
    title_text="Order Block Detection (Basic)",
    xaxis_rangeslider_visible=False,
    height=700, xaxis2_title="Datetime",
)
fig.show()