!pip install pandas numpy plotly scikit-learn ta

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 sklearn.preprocessing import StandardScaler

# Creating a sample DataFrame for demonstration purposes
data = {
    'datetime': pd.to_datetime(pd.date_range(start='2023-01-01', periods=100)),
    'open': np.random.rand(100) * 100 + 100,
    'high': np.random.rand(100) * 100 + 101,
    'low': np.random.rand(100) * 100 + 99,
    'close': np.random.rand(100) * 100 + 100,
    'volume': np.random.rand(100) * 1000 + 1000
}
df = pd.DataFrame(data)

print("Sample DataFrame 'df' created.")
display(df.head())

def prepare_ml_dataset(
    df: pd.DataFrame,
    horizon: int = 5,
    test_size: float = 0.2,
) -> tuple:
    """
    Engineer features from OHLCV data and prepare a supervised ML dataset.

    Core logic
    ----------
    1. Compute a set of lagged technical features from price and volume.
    2. Generate a forward-return binary label: +1 if close[t+horizon] > close[t], else 0.
    3. Drop NaN rows, split chronologically, and standardise via StandardScaler.

    Parameters
    ----------
    df : pd.DataFrame
        OHLCV DataFrame with columns: open, high, low, close, volume, datetime.
    horizon : int
        Number of bars ahead for the prediction target.
    test_size : float
        Fraction of data reserved for the test set (chronological tail).

    Returns
    -------
    tuple
        (X_train, X_test, y_train, y_test, feature_names, scaler)
    """
    df = df.copy().sort_values("datetime", ignore_index=True)

    # ── Feature engineering ──────────────────────────────────────────────────
    df["return_1"]  = df["close"].pct_change(1)           # 1-bar return
    df["return_5"]  = df["close"].pct_change(5)           # 5-bar return
    df["return_10"] = df["close"].pct_change(10)          # 10-bar return

    df["sma_10"] = df["close"].rolling(10).mean()
    df["sma_20"] = df["close"].rolling(20).mean()
    df["sma_ratio"] = df["sma_10"] / df["sma_20"]        # fast/slow MA ratio

    # RSI (14-period)
    delta = df["close"].diff()
    gain  = delta.clip(lower=0).rolling(14).mean()
    loss  = (-delta.clip(upper=0)).rolling(14).mean()
    rs    = gain / loss.replace(0, np.nan)
    df["rsi_14"] = 100 - (100 / (1 + rs))

    # ATR (14-period)
    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_14"] = tr.rolling(14).mean()
    df["atr_pct"] = df["atr_14"] / df["close"]           # normalised ATR

    # Bollinger Band width
    rolling_std   = df["close"].rolling(20).std()
    bb_upper      = df["sma_20"] + 2 * rolling_std
    bb_lower      = df["sma_20"] - 2 * rolling_std
    df["bb_width"] = (bb_upper - bb_lower) / df["sma_20"]

    # Volume ratio
    df["vol_ratio"] = df["volume"] / df["volume"].rolling(20).mean()

    # ── Label: forward return sign ───────────────────────────────────────────
    df["future_return"] = df["close"].shift(-horizon) / df["close"] - 1
    df["label"] = (df["future_return"] > 0).astype(int)  # 1 = up, 0 = down

    # ── Drop NaN rows ────────────────────────────────────────────────────────
    feature_cols = ["return_1", "return_5", "return_10", "sma_ratio",
                    "rsi_14", "atr_pct", "bb_width", "vol_ratio"]
    df.dropna(subset=feature_cols + ["label"], inplace=True)

    X = df[feature_cols].values
    y = df["label"].values

    # ── Chronological train/test split ───────────────────────────────────────
    split = int(len(X) * (1 - test_size))
    X_train, X_test = X[:split], X[split:]
    y_train, y_test = y[:split], y[split:]

    # ── Feature standardisation ──────────────────────────────────────────────
    scaler  = StandardScaler()
    X_train = scaler.fit_transform(X_train)   # fit on training data only
    X_test  = scaler.transform(X_test)        # apply same transform to test data

    return X_train, X_test, y_train, y_test, feature_cols, scaler, df

X_train, X_test, y_train, y_test, feature_names, scaler, df_feat = prepare_ml_dataset(df, horizon=5)

print(f"Training samples : {len(X_train)}")
print(f"Test samples     : {len(X_test)}")
print(f"Features         : {feature_names}")
print(f"Label balance (train): {pd.Series(y_train).value_counts().to_dict()}")

import pandas as pd

feat_df = pd.DataFrame(X_train, columns=feature_names)

fig = make_subplots(rows=2, cols=4,
    subplot_titles=feature_names)

for i, col in enumerate(feature_names):
    row = i // 4 + 1
    c    = i %  4 + 1
    fig.add_trace(go.Histogram(x=feat_df[col], name=col, showlegend=False,
                               marker_color="steelblue"), row=row, col=c)

fig.update_layout(title_text="Scaled Feature Distributions (Training Set)",
                  height=500)
fig.show()