#  Standard library
import warnings
import random
from collections import defaultdict

#  Numerical / data manipulation
import numpy as np
import pandas as pd

#  Machine learning
from sklearn.linear_model import SGDClassifier       # supports partial_fit (online learning)
from sklearn.ensemble import RandomForestClassifier  # batch baseline for comparison
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import (
    accuracy_score,
    precision_score,
    recall_score,
    f1_score,
    roc_auc_score,
)
from sklearn.pipeline import Pipeline

#  Visualisation
import matplotlib.pyplot as plt
import matplotlib.dates as mdates

#  Reproducibility & display settings
RANDOM_SEED = 42
np.random.seed(RANDOM_SEED)
random.seed(RANDOM_SEED)
warnings.filterwarnings("ignore")
pd.set_option("display.float_format", "{:.4f}".format)

#  Pipeline configuration constants
N_TOTAL_BARS    = 2_000   # total number of synthetic price bars to generate
WINDOW_SIZE     = 200     # number of bars in each training window
STEP_SIZE       = 50      # bars to advance between successive training cycles
LOOKBACK        = 20      # feature look-back period (rolling statistics)
FORWARD_RETURN  = 5       # bars ahead used to define the binary label
THRESHOLD       = 0.0     # label: 1 if forward_return > threshold, else 0
TEST_RATIO      = 0.20    # fraction of each window reserved for evaluation

print("Environment initialised. Configuration loaded.")

def generate_market_data(
    n_bars: int = N_TOTAL_BARS,
    regime_length: int = 300,
    seed: int = RANDOM_SEED,
) -> pd.DataFrame:
    """
    Generate synthetic OHLCV market data with alternating volatility regimes.

    Parameters
    ----------
    n_bars : int
        Total number of daily bars to generate.
    regime_length : int
        Approximate number of bars per volatility regime before switching.
    seed : int
        Random seed for reproducibility.

    Returns
    -------
    pd.DataFrame
        DataFrame with columns [open, high, low, close, volume] and a
        DatetimeIndex at daily frequency.
    """
    rng = np.random.default_rng(seed)

    #  Regime parameters: alternate between low and high volatility
    low_vol  = 0.01   # daily return std in calm regime
    high_vol = 0.03   # daily return std in turbulent regime
    drift    = 0.0003 # small upward drift (annualised ~7%)

    closes = [100.0]  # starting price
    regimes = []      # track regime label per bar

    current_vol = low_vol
    bars_in_regime = 0

    for _ in range(n_bars - 1):
        bars_in_regime += 1

        # Switch regime probabilistically after regime_length bars
        if bars_in_regime > regime_length and rng.random() < 0.05:
            current_vol = high_vol if current_vol == low_vol else low_vol
            bars_in_regime = 0  # reset counter after regime switch

        regimes.append(current_vol)

        # Geometric Brownian Motion step: dS = S*(mu*dt + sigma*dW)
        daily_return = drift + current_vol * rng.standard_normal()
        new_close = closes[-1] * (1 + daily_return)
        closes.append(max(new_close, 0.01))  # floor at near-zero to avoid negative prices

    regimes.append(current_vol)  # append last regime label

    closes = np.array(closes)

    #  Derive OHLV from close prices using realistic intraday spread
    intraday_range = closes * np.abs(rng.normal(0.005, 0.002, n_bars))
    opens  = closes * (1 + rng.normal(0, 0.003, n_bars))       # open near prior close
    highs  = np.maximum(closes, opens) + intraday_range * 0.6
    lows   = np.minimum(closes, opens) - intraday_range * 0.4
    volume = rng.integers(500_000, 5_000_000, n_bars).astype(float)

    #  Assemble DataFrame with a DatetimeIndex
    dates = pd.bdate_range(start="2018-01-01", periods=n_bars)  # business days only

    df = pd.DataFrame(
        {
            "open":   opens,
            "high":   highs,
            "low":    lows,
            "close":  closes,
            "volume": volume,
            "regime_vol": regimes,  # store for later inspection
        },
        index=dates,
    )

    return df


#  Execute and inspect
market_data = generate_market_data()
print(f"Dataset shape: {market_data.shape}")
print(f"Date range  : {market_data.index[0].date()} → {market_data.index[-1].date()}")
market_data.head()

def engineer_features(df: pd.DataFrame, lookback: int = LOOKBACK) -> pd.DataFrame:
    """
    Derive momentum, volatility, mean-reversion, and volume features
    from raw OHLCV market data.

    Parameters
    ----------
    df : pd.DataFrame
        Raw market data with columns [open, high, low, close, volume].
    lookback : int
        Rolling window length (in bars) for all feature computations.

    Returns
    -------
    pd.DataFrame
        Original DataFrame augmented with feature columns; NaN rows removed.
    """
    feat = df.copy()

    #  Log returns: preferred over arithmetic returns for stationarity
    feat["log_return"] = np.log(feat["close"] / feat["close"].shift(1))

    #  Momentum features
    # Rate of change over lookback period
    feat["roc"] = feat["close"].pct_change(periods=lookback)

    # Z-score of close relative to rolling mean and std (normalised momentum)
    rolling_mean = feat["close"].rolling(lookback).mean()
    rolling_std  = feat["close"].rolling(lookback).std()
    feat["zscore_close"] = (feat["close"] - rolling_mean) / rolling_std

    # Exponential moving average crossover signal
    ema_fast = feat["close"].ewm(span=lookback // 2, adjust=False).mean()
    ema_slow = feat["close"].ewm(span=lookback,      adjust=False).mean()
    feat["ema_cross"] = (ema_fast - ema_slow) / ema_slow  # normalised spread

    #  Volatility features
    # Realised volatility: rolling std of log returns
    feat["realised_vol"] = feat["log_return"].rolling(lookback).std()

    # Garman-Klass volatility estimator (uses OHLC, more efficient than close-to-close)
    hl_log = np.log(feat["high"] / feat["low"])
    co_log = np.log(feat["close"] / feat["open"])
    feat["gk_vol"] = np.sqrt(
        0.5 * hl_log**2 - (2 * np.log(2) - 1) * co_log**2
    ).rolling(lookback).mean()

    #  Mean-reversion features
    # Distance of close from rolling maximum (measures drawdown from peak)
    rolling_max = feat["close"].rolling(lookback).max()
    feat["dist_from_max"] = (feat["close"] - rolling_max) / rolling_max

    # Distance of close from rolling minimum
    rolling_min = feat["close"].rolling(lookback).min()
    feat["dist_from_min"] = (feat["close"] - rolling_min) / rolling_min

    #  Volume features
    # Relative volume: ratio of current volume to rolling average
    feat["rel_volume"] = feat["volume"] / feat["volume"].rolling(lookback).mean()

    # Volume-weighted momentum: log_return scaled by relative volume
    feat["vol_momentum"] = feat["log_return"] * feat["rel_volume"]

    #  Drop NaN rows produced by rolling/shift operations
    feat.dropna(inplace=True)

    return feat


#  Execute and inspect
featured_data = engineer_features(market_data)

FEATURE_COLS = [
    "roc", "zscore_close", "ema_cross",
    "realised_vol", "gk_vol",
    "dist_from_max", "dist_from_min",
    "rel_volume", "vol_momentum",
]

print(f"Feature matrix shape: {featured_data.shape}")
print(f"Feature columns     : {FEATURE_COLS}")
featured_data[FEATURE_COLS].describe().T

def construct_labels(
    df: pd.DataFrame,
    forward_bars: int = FORWARD_RETURN,
    threshold: float = THRESHOLD,
) -> pd.DataFrame:
    """
    Construct a binary directional label from forward close-to-close returns.

    Parameters
    ----------
    df : pd.DataFrame
        Feature-augmented market data with a 'close' column.
    forward_bars : int
        Number of bars ahead used to compute the forward return.
    threshold : float
        Minimum forward return required to assign label 1. Default: 0.0
        (any positive return qualifies).

    Returns
    -------
    pd.DataFrame
        DataFrame with an additional 'label' column (int: 0 or 1) and
        all NaN rows removed.
    """
    labelled = df.copy()

    # Compute forward return: percentage change from current close to future close
    labelled["forward_return"] = (
        labelled["close"].shift(-forward_bars) - labelled["close"]
    ) / labelled["close"]

    # Binary label: 1 if forward return exceeds threshold, otherwise 0
    labelled["label"] = (labelled["forward_return"] > threshold).astype(int)

    # Remove the last `forward_bars` rows where forward return cannot be computed
    labelled.dropna(subset=["forward_return"], inplace=True)

    return labelled


#  Execute and inspect label distribution
labelled_data = construct_labels(featured_data)

label_counts = labelled_data["label"].value_counts(normalize=True)
print(f"Total labelled samples : {len(labelled_data)}")
print(f"Label distribution     :")
print(f"  Class 0 (flat/down)  : {label_counts[0]:.2%}")
print(f"  Class 1 (up)         : {label_counts[1]:.2%}")

def sliding_window_iterator(
    df: pd.DataFrame,
    window_size: int = WINDOW_SIZE,
    step_size: int   = STEP_SIZE,
    test_ratio: float = TEST_RATIO,
    feature_cols: list = None,
):
    """
    Generate sequential (train, test) splits using a sliding window.

    Parameters
    ----------
    df : pd.DataFrame
        Full labelled dataset ordered by time.
    window_size : int
        Total number of bars in each window.
    step_size : int
        Number of bars to advance the window on each iteration.
    test_ratio : float
        Fraction of each window allocated to the test split.
    feature_cols : list
        Column names to use as model input features.

    Yields
    ------
    dict with keys:
        cycle       : int   — iteration index (0-based)
        X_train     : ndarray — training features
        y_train     : ndarray — training labels
        X_test      : ndarray — test features
        y_test      : ndarray — test labels
        train_start : Timestamp — start date of training slice
        test_end    : Timestamp — end date of test slice
    """
    if feature_cols is None:
        feature_cols = FEATURE_COLS

    n_rows     = len(df)
    train_size = int(window_size * (1 - test_ratio))  # bars in training split
    start_idx  = 0
    cycle      = 0

    while start_idx + window_size <= n_rows:
        end_idx = start_idx + window_size  # exclusive upper bound

        window = df.iloc[start_idx:end_idx]  # extract current window

        # Partition window into temporal train / test splits
        train_window = window.iloc[:train_size]
        test_window  = window.iloc[train_size:]

        yield {
            "cycle":       cycle,
            "X_train":     train_window[feature_cols].values,
            "y_train":     train_window["label"].values,
            "X_test":      test_window[feature_cols].values,
            "y_test":      test_window["label"].values,
            "train_start": window.index[0],
            "test_end":    window.index[-1],
        }

        start_idx += step_size  # advance window by step_size
        cycle += 1


#  Preview window counts
total_cycles = sum(1 for _ in sliding_window_iterator(labelled_data))
print(f"Total training cycles : {total_cycles}")
print(f"Window size           : {WINDOW_SIZE} bars  |  Step size: {STEP_SIZE} bars")
print(f"Train / test split    : {int(WINDOW_SIZE*(1-TEST_RATIO))} / {int(WINDOW_SIZE*TEST_RATIO)} bars per cycle")

from sklearn.utils.class_weight import compute_class_weight

def build_online_model() -> dict:
    """
    Initialise an online learning model and its associated scaler.

    The SGDClassifier with log-loss approximates logistic regression
    trained via stochastic gradient descent, making it compatible with
    incremental updates via partial_fit.

    Returns
    -------
    dict
        {'model': SGDClassifier, 'scaler': StandardScaler}
        Both objects are uninitialised and ready for first partial_fit call.
    """
    model = SGDClassifier(
        loss="log_loss",      # produces calibrated probability estimates
        penalty="l2",         # L2 regularisation to prevent weight explosion
        alpha=0.001,          # regularisation strength
        max_iter=1,           # single pass per batch (online regime)
        warm_start=False,     # weights are preserved between partial_fit calls by design
        # class_weight="balanced",  # Removed as per error, will handle via sample_weight in partial_fit
        random_state=RANDOM_SEED,
    )

    scaler = StandardScaler()  # will be updated incrementally with partial_fit

    return {"model": model, "scaler": scaler}


def update_online_model(
    online_components: dict,
    X_train: np.ndarray,
    y_train: np.ndarray,
    X_test: np.ndarray,
    y_test: np.ndarray,
) -> dict:
    """
    Perform one incremental update on the online model and evaluate on the test split.

    Parameters
    ----------
    online_components : dict
        Dictionary containing 'model' and 'scaler' from build_online_model().
    X_train, y_train : ndarray
        Training features and labels for the current window.
    X_test, y_test : ndarray
        Test features and labels for evaluation.

    Returns
    -------
    dict
        Evaluation metrics: accuracy, precision, recall, f1, auc.
    """
    model  = online_components["model"]
    scaler = online_components["scaler"]

    # Incrementally update scaler statistics with the new training batch
    scaler.partial_fit(X_train)

    # Normalise training and test data using the updated scaler
    X_train_scaled = scaler.transform(X_train)
    X_test_scaled  = scaler.transform(X_test)

    # Compute class weights for the current training batch and convert to sample weights
    # This addresses the ValueError: 'balanced' class_weight is not supported for partial_fit.
    classes_in_batch = np.unique(y_train)
    sample_weight = None
    if len(classes_in_batch) > 1: # Ensure both classes are present to compute balanced weights
        class_weights_dict = dict(zip(
            classes_in_batch,
            compute_class_weight(
                class_weight='balanced',
                classes=classes_in_batch,
                y=y_train
            )
        ))
        sample_weight = np.array([class_weights_dict[label] for label in y_train])

    # Provide all class labels on the first call; SGD requires this for initialisation
    # Ensure all possible classes are always passed to partial_fit at least once
    # even if not all classes are present in the current batch.
    all_possible_classes = np.array([0, 1])
    model.partial_fit(X_train_scaled, y_train, classes=all_possible_classes, sample_weight=sample_weight)

    # Generate predictions and probability estimates on the held-out test split
    y_pred      = model.predict(X_test_scaled)
    y_prob      = model.predict_proba(X_test_scaled)[:, 1]  # probability of class 1

    return _compute_metrics(y_test, y_pred, y_prob)


print("Online model components defined.")

def update_batch_model(
    X_train: np.ndarray,
    y_train: np.ndarray,
    X_test: np.ndarray,
    y_test: np.ndarray,
) -> dict:
    """
    Fit a RandomForestClassifier from scratch on the current window
    and evaluate it on the test split.

    Parameters
    ----------
    X_train, y_train : ndarray
        Training data for the current window.
    X_test, y_test : ndarray
        Test data for out-of-sample evaluation.

    Returns
    -------
    dict
        Evaluation metrics: accuracy, precision, recall, f1, auc.
    """
    # Fit a fresh scaler on this window's training data
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled  = scaler.transform(X_test)

    # Instantiate and train a RandomForest from scratch on this window
    model = RandomForestClassifier(
        n_estimators=50,           # moderate tree count for speed
        max_depth=4,               # shallow trees reduce overfitting on small windows
        class_weight="balanced",
        random_state=RANDOM_SEED,
        n_jobs=-1,                 # parallelise across available CPU cores
    )
    model.fit(X_train_scaled, y_train)

    # Evaluate on test split
    y_pred = model.predict(X_test_scaled)
    y_prob = model.predict_proba(X_test_scaled)[:, 1]

    return _compute_metrics(y_test, y_pred, y_prob)


print("Batch baseline model defined.")

def _compute_metrics(
    y_true: np.ndarray,
    y_pred: np.ndarray,
    y_prob: np.ndarray,
) -> dict:
    """
    Compute binary classification evaluation metrics.

    Parameters
    ----------
    y_true : ndarray
        Ground-truth binary labels.
    y_pred : ndarray
        Predicted binary labels.
    y_prob : ndarray
        Predicted probabilities for class 1.

    Returns
    -------
    dict
        Keys: accuracy, precision, recall, f1, auc.
    """
    # Guard against edge cases where only one class is present in y_true
    try:
        auc = roc_auc_score(y_true, y_prob)
    except ValueError:
        auc = float("nan")  # AUC undefined when only one class present in batch

    return {
        "accuracy":  accuracy_score(y_true, y_pred),
        "precision": precision_score(y_true, y_pred, zero_division=0),
        "recall":    recall_score(y_true, y_pred, zero_division=0),
        "f1":        f1_score(y_true, y_pred, zero_division=0),
        "auc":       auc,
    }


print("Metric utility function defined.")

def detect_drift(
    auc_series: list,
    n_baseline: int = 5,
    sensitivity: float = 1.5,
) -> list:
    """
    Detect performance drift by comparing current AUC to a rolling baseline.

    A drift event is flagged when the current AUC is more than `sensitivity`
    standard deviations below the mean of the previous `n_baseline` cycles.

    Parameters
    ----------
    auc_series : list of float
        AUC values recorded across all completed cycles.
    n_baseline : int
        Number of prior cycles used to compute the rolling reference.
    sensitivity : float
        Z-score threshold below which drift is flagged.

    Returns
    -------
    list of bool
        True at each index where drift is detected; False otherwise.
        First `n_baseline` entries are always False (insufficient history).
    """
    drift_flags = []

    for i, current_auc in enumerate(auc_series):
        if i < n_baseline:
            # Insufficient history to compute a reliable baseline
            drift_flags.append(False)
            continue

        # Rolling window of preceding AUC values (exclude current)
        baseline_window = [v for v in auc_series[i - n_baseline:i] if not np.isnan(v)]

        if len(baseline_window) < 2:
            drift_flags.append(False)
            continue

        baseline_mean = np.mean(baseline_window)
        baseline_std  = np.std(baseline_window, ddof=1)  # sample std

        if baseline_std == 0:
            drift_flags.append(False)
            continue

        # Z-score of current AUC relative to baseline distribution
        z_score = (current_auc - baseline_mean) / baseline_std

        # Drift: current performance is significantly below the rolling mean
        drift_flags.append(z_score < -sensitivity)

    return drift_flags


print("Drift detection function defined.")

def run_continual_learning_pipeline(
    labelled_df: pd.DataFrame,
    feature_cols: list = None,
) -> pd.DataFrame:
    """
    Execute the full continual learning pipeline across all sliding windows.

    Parameters
    ----------
    labelled_df : pd.DataFrame
        Fully featured and labelled market data.
    feature_cols : list
        Feature column names. Defaults to FEATURE_COLS if None.

    Returns
    -------
    pd.DataFrame
        Per-cycle results including metrics for online and batch models,
        date boundaries, and drift flags.
    """
    if feature_cols is None:
        feature_cols = FEATURE_COLS

    # Initialise the online model once — it persists and updates across all cycles
    online_components = build_online_model()

    results = []  # accumulate per-cycle records

    print(f"{'Cycle':>6}  {'Train Start':>12}  {'Test End':>12}  "
          f"{'Online AUC':>10}  {'Batch AUC':>9}")
    print("-" * 62)

    for window in sliding_window_iterator(labelled_df, feature_cols=feature_cols):
        cycle   = window["cycle"]
        X_train = window["X_train"]
        y_train = window["y_train"]
        X_test  = window["X_test"]
        y_test  = window["y_test"]

        #  Online model: incremental update
        online_metrics = update_online_model(
            online_components, X_train, y_train, X_test, y_test
        )

        #  Batch model: full retrain from scratch
        batch_metrics = update_batch_model(X_train, y_train, X_test, y_test)

        #  Record results
        record = {
            "cycle":       cycle,
            "train_start": window["train_start"],
            "test_end":    window["test_end"],
        }

        # Prefix metric keys to distinguish model sources
        for k, v in online_metrics.items():
            record[f"online_{k}"] = v
        for k, v in batch_metrics.items():
            record[f"batch_{k}"] = v

        results.append(record)

        # Print progress every 5 cycles
        if cycle % 5 == 0:
            print(
                f"{cycle:>6}  "
                f"{str(window['train_start'].date()):>12}  "
                f"{str(window['test_end'].date()):>12}  "
                f"{online_metrics['auc']:>10.4f}  "
                f"{batch_metrics['auc']:>9.4f}"
            )

    results_df = pd.DataFrame(results)

    #  Apply drift detection to the online AUC series
    results_df["drift_detected"] = detect_drift(
        results_df["online_auc"].tolist()
    )

    n_drift = results_df["drift_detected"].sum()
    print("-" * 62)
    print(f"\nPipeline complete. Total cycles: {len(results_df)}")
    print(f"Drift events detected          : {n_drift}")

    return results_df


#  Execute the pipeline
results_df = run_continual_learning_pipeline(labelled_data)

def summarise_results(results_df: pd.DataFrame) -> pd.DataFrame:
    """
    Compute summary statistics comparing online and batch model performance
    across all training cycles.

    Parameters
    ----------
    results_df : pd.DataFrame
        Per-cycle results from run_continual_learning_pipeline().

    Returns
    -------
    pd.DataFrame
        Summary table with mean and std for key metrics.
    """
    metrics = ["accuracy", "precision", "recall", "f1", "auc"]

    rows = []
    for model_prefix, model_label in [("online", "Online (SGD)"), ("batch", "Batch (RandomForest)")]:
        row = {"Model": model_label}
        for m in metrics:
            col = f"{model_prefix}_{m}"
            row[f"{m}_mean"] = results_df[col].mean()
            row[f"{m}_std"]  = results_df[col].std()
        rows.append(row)

    summary = pd.DataFrame(rows).set_index("Model")
    return summary


summary_df = summarise_results(results_df)
print("\nCross-Cycle Performance Summary")
print("=" * 60)
summary_df

def plot_pipeline_results(results_df: pd.DataFrame) -> None:
    """
    Generate a four-panel visualisation of continual learning pipeline results.

    Parameters
    ----------
    results_df : pd.DataFrame
        Per-cycle results from run_continual_learning_pipeline().
    """
    cycles      = results_df["cycle"].values
    drift_mask  = results_df["drift_detected"].values
    drift_cycles= cycles[drift_mask]

    #  Colour palette
    C_ONLINE = "#2196F3"  # blue  — online model
    C_BATCH  = "#FF5722"  # orange — batch model
    C_DRIFT  = "#E91E63"  # red   — drift events
    C_GRID   = "#EEEEEE"  # light grey grid

    fig, axes = plt.subplots(2, 2, figsize=(16, 10))
    fig.suptitle(
        "Continual Learning Pipeline — Performance Across Training Cycles",
        fontsize=14, fontweight="bold", y=1.01
    )

    #  Panel 1: AUC over cycles
    ax = axes[0, 0]
    ax.plot(cycles, results_df["online_auc"], color=C_ONLINE, lw=1.8,
            label="Online (SGD)", zorder=3)
    ax.plot(cycles, results_df["batch_auc"],  color=C_BATCH,  lw=1.8,
            label="Batch (RandomForest)", zorder=3)

    # Mark drift events with vertical dashed lines
    for dc in drift_cycles:
        ax.axvline(dc, color=C_DRIFT, lw=0.8, alpha=0.6, linestyle="--")

    ax.axhline(0.5, color="grey", lw=0.8, linestyle=":", label="Random baseline (AUC=0.5)")
    ax.set_title("ROC-AUC per Cycle", fontweight="bold")
    ax.set_xlabel("Training Cycle")
    ax.set_ylabel("AUC")
    ax.legend(fontsize=8)
    ax.set_ylim(0.3, 1.0)
    ax.yaxis.grid(True, color=C_GRID)
    ax.set_axisbelow(True)

    #  Panel 2: F1 score over cycles
    ax = axes[0, 1]
    ax.plot(cycles, results_df["online_f1"], color=C_ONLINE, lw=1.8,
            label="Online (SGD)")
    ax.plot(cycles, results_df["batch_f1"],  color=C_BATCH,  lw=1.8,
            label="Batch (RandomForest)")

    # Rolling 5-cycle moving average to smooth short-term variance
    online_f1_ma = results_df["online_f1"].rolling(5, min_periods=1).mean()
    batch_f1_ma  = results_df["batch_f1"].rolling(5, min_periods=1).mean()
    ax.plot(cycles, online_f1_ma, color=C_ONLINE, lw=2.5, linestyle="--", alpha=0.5,
            label="Online 5-cycle MA")
    ax.plot(cycles, batch_f1_ma,  color=C_BATCH,  lw=2.5, linestyle="--", alpha=0.5,
            label="Batch 5-cycle MA")

    ax.set_title("F1 Score per Cycle", fontweight="bold")
    ax.set_xlabel("Training Cycle")
    ax.set_ylabel("F1 Score")
    ax.legend(fontsize=8)
    ax.set_ylim(0.0, 1.0)
    ax.yaxis.grid(True, color=C_GRID)
    ax.set_axisbelow(True)

    #  Panel 3: Online vs Batch AUC scatter
    ax = axes[1, 0]
    non_drift = ~drift_mask

    ax.scatter(
        results_df.loc[non_drift, "batch_auc"],
        results_df.loc[non_drift, "online_auc"],
        color=C_ONLINE, alpha=0.6, s=30, label="Normal cycles", zorder=3
    )
    if drift_mask.any():
        ax.scatter(
            results_df.loc[drift_mask, "batch_auc"],
            results_df.loc[drift_mask, "online_auc"],
            color=C_DRIFT, alpha=0.9, s=60, marker="X", label="Drift cycles", zorder=4
        )

    # Diagonal: y = x (parity line between models)
    lim_min = min(results_df[["online_auc", "batch_auc"]].min()) - 0.02
    lim_max = max(results_df[["online_auc", "batch_auc"]].max()) + 0.02
    ax.plot([lim_min, lim_max], [lim_min, lim_max], "k--", lw=1, alpha=0.5,
            label="Parity (Online = Batch)")

    ax.set_title("Online vs Batch AUC per Cycle", fontweight="bold")
    ax.set_xlabel("Batch AUC")
    ax.set_ylabel("Online AUC")
    ax.legend(fontsize=8)
    ax.set_xlim(lim_min, lim_max)
    ax.set_ylim(lim_min, lim_max)
    ax.set_aspect("equal")
    ax.yaxis.grid(True, color=C_GRID)
    ax.xaxis.grid(True, color=C_GRID)
    ax.set_axisbelow(True)

    #  Panel 4: Drift event timeline
    ax = axes[1, 1]
    ax.fill_between(cycles, results_df["online_auc"], 0.5,
                    where=(results_df["online_auc"] >= 0.5),
                    alpha=0.3, color=C_ONLINE, label="AUC > 0.5")
    ax.fill_between(cycles, results_df["online_auc"], 0.5,
                    where=(results_df["online_auc"] < 0.5),
                    alpha=0.3, color="grey", label="AUC < 0.5")
    ax.plot(cycles, results_df["online_auc"], color=C_ONLINE, lw=1.5, zorder=3)

    # Highlight drift cycles as vertical red bands
    for dc in drift_cycles:
        ax.axvspan(dc - 0.5, dc + 0.5, color=C_DRIFT, alpha=0.4)

    ax.axhline(0.5, color="grey", lw=0.8, linestyle=":")
    ax.set_title("Online AUC Timeline with Drift Events", fontweight="bold")
    ax.set_xlabel("Training Cycle")
    ax.set_ylabel("Online AUC")
    ax.set_ylim(0.3, 1.0)

    # Custom legend entry for drift bands
    from matplotlib.patches import Patch
    legend_elements = [
        Patch(facecolor=C_ONLINE, alpha=0.5, label="AUC > 0.5"),
        Patch(facecolor="grey",   alpha=0.3, label="AUC < 0.5"),
        Patch(facecolor=C_DRIFT,  alpha=0.4, label="Drift event"),
    ]
    ax.legend(handles=legend_elements, fontsize=8)
    ax.yaxis.grid(True, color=C_GRID)
    ax.set_axisbelow(True)

    plt.tight_layout()
    plt.savefig("continual_learning_results.png", dpi=150, bbox_inches="tight")
    plt.show()
    print("Figure saved to continual_learning_results.png")


#  Execute visualisation
plot_pipeline_results(results_df)