!pip install pandas numpy plotly scipy

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
from scipy.stats import linregress

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 flag_pennant_detection(
    df: pd.DataFrame,
    pole_bars: int = 10,
    flag_bars: int = 10,
    pole_threshold: float = 1.5,
) -> pd.DataFrame:
    """
    Detect Flag and Pennant continuation patterns following a strong directional move.

    Core logic
    ----------
    1. Compute ATR over pole_bars to normalise the flagpole magnitude.
    2. Scan each bar: if the price move over the prior pole_bars exceeds
       pole_threshold × ATR, a flagpole is identified.
    3. Over the subsequent flag_bars, compute linear regression slopes of highs
       and lows to determine consolidation shape (parallel = flag, converging = pennant).
    4. Emit a directional signal matching the flagpole direction.

    Parameters
    ----------
    df : pd.DataFrame
        OHLCV DataFrame with columns: open, high, low, close, volume, datetime.
    pole_bars : int
        Number of bars used to measure the flagpole move.
    flag_bars : int
        Number of bars following the pole used to evaluate the consolidation.
    pole_threshold : float
        ATR multiplier; flagpole moves must exceed this multiple of ATR.

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

    # ── ATR (Average True Range) ─────────────────────────────────────────────
    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(pole_bars).mean()

    close = df["close"].values
    highs = df["high"].values
    lows  = df["low"].values

    for i in range(pole_bars, len(df) - flag_bars):
        atr_val = df["atr"].iloc[i]
        if np.isnan(atr_val) or atr_val == 0:
            continue

        # ── Flagpole identification ──────────────────────────────────────────
        pole_move = close[i] - close[i - pole_bars]
        if abs(pole_move) < pole_threshold * atr_val:
            continue  # insufficient impulse

        pole_direction = 1 if pole_move > 0 else -1  # +1 bullish, -1 bearish

        # ── Consolidation regression ─────────────────────────────────────────
        flag_slice = np.arange(flag_bars)
        slope_h, *_ = linregress(flag_slice, highs[i: i + flag_bars])
        slope_l, *_ = linregress(flag_slice, lows[i:  i + flag_bars])

        norm_h = slope_h / np.mean(highs[i: i + flag_bars])
        norm_l = slope_l / np.mean(lows[i:  i + flag_bars])

        signal_bar = min(i + flag_bars, len(df) - 1)

        # Pennant: slopes converge (opposite signs after impulse)
        if norm_h < 0 and norm_l > 0:
            df.at[signal_bar, "pattern"] = f"{'bull' if pole_direction == 1 else 'bear'}_pennant"
            df.at[signal_bar, "signal"]  = pole_direction

        # Flag: both slopes move against the pole (retracement channel)
        elif pole_direction == 1 and norm_h < 0 and norm_l < 0:
            df.at[signal_bar, "pattern"] = "bull_flag"
            df.at[signal_bar, "signal"]  = 1
        elif pole_direction == -1 and norm_h > 0 and norm_l > 0:
            df.at[signal_bar, "pattern"] = "bear_flag"
            df.at[signal_bar, "signal"]  = -1

    return df

df_signals = flag_pennant_detection(df, pole_bars=10, flag_bars=10, pole_threshold=1.5)

print("--- Pattern Distribution ---")
print(df_signals["pattern"].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 + Flag/Pennant 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="Bull (+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="Bear (-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="Flag and Pennant Detection",
    xaxis_rangeslider_visible=False,
    height=700, xaxis2_title="Datetime",
)
fig.show()