"""Structural Break Detection — classical tests and ML-based classification."""

import logging
import warnings
from typing import Any

warnings.filterwarnings("ignore")
# ml4t-engineer decorates some feature functions with per-call INFO timing logs
# (coefficient_of_variation sets log_data_quality=True), which otherwise leak
# into committed cell outputs as library chatter. Disable INFO so the outputs
# show only the feature values; WARNING and above still surface.
logging.disable(logging.INFO)

from datetime import datetime

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import polars as pl
import ruptures as rpt
from arch.unitroot import ZivotAndrews
from IPython.display import display
from lightgbm import LGBMClassifier
from ml4t.engineer.features.statistics import (
    coefficient_of_variation,
    rolling_drift,
    rolling_kl_divergence,
    rolling_wasserstein,
)
from scipy import stats
from scipy.spatial.distance import jensenshannon
from sklearn.metrics import roc_auc_score, roc_curve
from sklearn.model_selection import cross_val_predict, cross_val_score

from data import load_etfs
from utils.reproducibility import set_global_seeds

# Production defaults — Papermill injects overrides for CI
START_DATE = "2000-01-01"
END_DATE = "2024-12-31"
SEED = 42

set_global_seeds(SEED)

etfs = load_etfs(symbols=["SPY"])

START = datetime.strptime(START_DATE, "%Y-%m-%d")
END = datetime.strptime(END_DATE, "%Y-%m-%d")

spy = (
    etfs.filter((pl.col("timestamp") >= START.date()) & (pl.col("timestamp") <= END.date()))
    .sort("timestamp")
    .select(["timestamp", "close"])
)

spy_pd = spy.to_pandas().set_index("timestamp")
spy_pd.index = pd.DatetimeIndex(spy_pd.index)
spy_pd["returns"] = np.log(spy_pd["close"]).diff() * 100

print(
    f"SPY: {len(spy_pd):,} observations ({spy_pd.index.min().date()} to {spy_pd.index.max().date()})"
)

def run_zivot_andrews(series: pd.Series, name: str) -> dict:
    """Run Zivot-Andrews test and return results with break location."""
    series = series.dropna()
    za = ZivotAndrews(series, lags=12, trend="ct", trim=0.15)

    return {
        "series": name,
        "za_stat": za.stat,
        "za_pval": za.pvalue,
        "critical_1pct": za.critical_values["1%"],
        "critical_5pct": za.critical_values["5%"],
        "reject_h0": za.pvalue < 0.05,
    }


za_result = run_zivot_andrews(spy_pd["close"], "SPY Levels")
print("=== Zivot-Andrews Test: SPY Levels ===")
print(f"Statistic: {za_result['za_stat']:.4f}")
print(f"P-value: {za_result['za_pval']:.4f}")
print(f"5% critical: {za_result['critical_5pct']:.4f}")
print(f"Reject H0: {za_result['reject_h0']}")

def detect_multiple_breaks(series: np.ndarray, n_breaks: int = 3, model: str = "l2") -> dict:
    """Detect multiple break points using Pelt or dynamic programming.

    Args:
        series: 1D array of values
        n_breaks: Expected number of breaks
        model: Cost model ('l2' for mean shift, 'rbf' for distribution shift)
    """
    # Pelt with penalty (automatic n_breaks)
    algo_pelt = rpt.Pelt(model=model, min_size=60).fit(series)
    breaks_pelt = algo_pelt.predict(pen=np.log(len(series)) * series.var())

    # Binary segmentation (fixed n_breaks)
    algo_binseg = rpt.Binseg(model=model, min_size=60).fit(series)
    breaks_binseg = algo_binseg.predict(n_bkps=n_breaks)

    return {
        "pelt_breaks": breaks_pelt[:-1],  # Remove last (=len)
        "binseg_breaks": breaks_binseg[:-1],
    }


prices = spy_pd["close"].dropna().values
breaks = detect_multiple_breaks(prices, n_breaks=4)

print(f"PELT breaks: {len(breaks['pelt_breaks'])} detected")
print(f"BinSeg breaks: {len(breaks['binseg_breaks'])} detected (requested 4)")

fig, axes = plt.subplots(2, 1, figsize=(14, 8), sharex=True)

dates = spy_pd["close"].dropna().index

# PELT results
ax = axes[0]
ax.plot(dates, prices, linewidth=0.8)
for bp in breaks["pelt_breaks"]:
    if bp < len(dates):
        ax.axvline(dates[bp], color="red", linestyle="--", alpha=0.7)
        ax.text(
            dates[bp],
            prices.max() * 0.95,
            dates[bp].strftime("%Y-%m"),
            rotation=45,
            fontsize=8,
            color="red",
        )
ax.set_title(f"PELT: {len(breaks['pelt_breaks'])} Break Points Detected")
ax.set_ylabel("Price")

# BinSeg results
ax = axes[1]
ax.plot(dates, prices, linewidth=0.8)
for bp in breaks["binseg_breaks"]:
    if bp < len(dates):
        ax.axvline(dates[bp], color="red", linestyle="--", alpha=0.7)
        ax.text(
            dates[bp],
            prices.max() * 0.95,
            dates[bp].strftime("%Y-%m"),
            rotation=45,
            fontsize=8,
            color="red",
        )
ax.set_title("Binary Segmentation: 4 Break Points")
ax.set_ylabel("Price")

plt.tight_layout()
plt.show()

def compute_break_features(
    series: np.ndarray, breaks: list, dates: pd.DatetimeIndex
) -> pd.DataFrame:
    """Convert break points into time-varying features."""
    n = len(series)
    features = pd.DataFrame(index=dates[:n])

    # Binary: has a break been detected before this point?
    features["za_break_detected"] = 0
    for bp in breaks:
        if bp < n:
            features.iloc[bp:, features.columns.get_loc("za_break_detected")] = 1

    # Days since last break
    features["days_since_break"] = n  # Default: no break
    for i in range(n):
        past_breaks = [bp for bp in breaks if bp <= i]
        if past_breaks:
            features.iloc[i, features.columns.get_loc("days_since_break")] = i - max(past_breaks)

    # Pre/post break means (using most recent break)
    features["pre_break_mean"] = np.nan
    features["post_break_mean"] = np.nan
    for bp in sorted(breaks):
        if bp < n and bp > 50:
            pre_mean = series[max(0, bp - 252) : bp].mean()
            post_end = min(bp + 252, n)
            post_mean = series[bp:post_end].mean()
            features.iloc[bp:, features.columns.get_loc("pre_break_mean")] = pre_mean
            features.iloc[bp:, features.columns.get_loc("post_break_mean")] = post_mean

    return features


break_features = compute_break_features(prices, breaks["pelt_breaks"], dates)
display(break_features.dropna().head(10))

def cusum_statistic(series: np.ndarray, target: float | None = None) -> np.ndarray:
    """Compute one-sided CUSUM statistic.

    CUSUM accumulates deviations from the target mean.
    Sustained positive CUSUM indicates upward shift; negative indicates downward.
    """
    if target is None:
        target = series.mean()
    return np.cumsum(series - target)


returns = spy_pd["returns"].dropna()
cusum = cusum_statistic(returns.values)

fig, axes = plt.subplots(2, 1, figsize=(14, 6), sharex=True)

ax = axes[0]
ax.plot(returns.index, returns.values, linewidth=0.3, alpha=0.5)
ax.set_title("SPY Returns (%)")
ax.set_ylabel("Return (%)")

ax = axes[1]
ax.plot(returns.index, cusum, linewidth=0.8)
ax.axhline(0, color="red", linestyle="--", linewidth=0.5)
ax.set_title("CUSUM of Returns (Cumulative Deviation from Mean)")
ax.set_ylabel("Cumulative Sum")

plt.tight_layout()
plt.show()

def mosum_statistic(
    series: np.ndarray, bandwidth: int = 50, target: float | None = None
) -> np.ndarray:
    """Compute MOSUM (moving sum) monitoring statistic.

    Unlike CUSUM which accumulates from the start, MOSUM uses a fixed-width
    sliding window — better at detecting abrupt, localized shifts.

    Parameters
    ----------
    series : array
        Input time series.
    bandwidth : int
        Window width (h). Larger h → smoother, less sensitive to short breaks.
    target : float, optional
        Reference mean. Defaults to full-sample mean.

    Returns
    -------
    mosum : array
        Standardized MOSUM statistic (same length as series, NaN-padded).
    """
    if target is None:
        target = series.mean()
    sigma = series.std()
    n = len(series)
    mosum = np.full(n, np.nan)
    for t in range(bandwidth, n):
        window_sum = np.sum(series[t - bandwidth + 1 : t + 1] - target)
        mosum[t] = window_sum / (sigma * np.sqrt(bandwidth))
    return mosum

# Compare CUSUM and MOSUM on the same returns series
mosum_50 = mosum_statistic(returns.values, bandwidth=50)
mosum_20 = mosum_statistic(returns.values, bandwidth=20)

fig, axes = plt.subplots(3, 1, figsize=(14, 7), sharex=True)

ax = axes[0]
ax.plot(returns.index, returns.values, linewidth=0.3, alpha=0.5)
ax.set_title("SPY Returns (%)")
ax.set_ylabel("Return (%)")

ax = axes[1]
ax.plot(returns.index, cusum, linewidth=0.8, label="CUSUM")
ax.axhline(0, color="red", linestyle="--", linewidth=0.5)
ax.set_title("CUSUM — Accumulates Drift from Start")
ax.set_ylabel("Cumulative Sum")
ax.legend(loc="upper left")

ax = axes[2]
ax.plot(returns.index, mosum_50, linewidth=0.8, label="MOSUM (h=50)")
ax.plot(returns.index, mosum_20, linewidth=0.8, alpha=0.6, label="MOSUM (h=20)")
ax.axhline(0, color="red", linestyle="--", linewidth=0.5)
ax.set_title("MOSUM — Detects Localized Shifts")
ax.set_ylabel("Standardized MOSUM")
ax.legend(loc="upper left")

plt.tight_layout()
plt.show()

def generate_break_series(
    n: int = 500, break_type: str = "mean_shift", magnitude: float = 1.0, seed: int | None = None
) -> tuple[np.ndarray, bool, int]:
    """Generate a time series with or without a structural break at midpoint.

    Returns (series, has_break, break_idx).
    """
    rng = np.random.RandomState(seed)
    mid = n // 2

    if break_type == "none":
        series = rng.randn(n)
        return series, False, mid

    elif break_type == "mean_shift":
        pre = rng.randn(mid)
        post = rng.randn(n - mid) + magnitude
        return np.concatenate([pre, post]), True, mid

    elif break_type == "var_shift":
        pre = rng.randn(mid)
        post = rng.randn(n - mid) * (1 + magnitude)
        return np.concatenate([pre, post]), True, mid

    elif break_type == "trend_shift":
        t = np.arange(n, dtype=float)
        pre_trend = 0.001 * t[:mid]
        post_trend = 0.001 * t[:mid][-1] + magnitude * 0.01 * (t[mid:] - t[mid])
        noise = rng.randn(n) * 0.5
        return np.concatenate([pre_trend, post_trend]) + noise, True, mid

    elif break_type == "autocorr_shift":
        pre = np.zeros(mid)
        for i in range(1, mid):
            pre[i] = 0.2 * pre[i - 1] + rng.randn()
        post = np.zeros(n - mid)
        for i in range(1, n - mid):
            post[i] = (0.2 + magnitude * 0.5) * post[i - 1] + rng.randn()
        return np.concatenate([pre, post]), True, mid

    msg = f"Unknown break_type: {break_type}"
    raise ValueError(msg)


# Generate examples
rng = np.random.RandomState(42)
n_samples = 200
series_list = []
labels = []

for i in range(n_samples):
    if i < n_samples // 2:
        # No break
        s, has_break, mid = generate_break_series(500, "none", seed=i)
    else:
        # Random break type
        break_types = ["mean_shift", "var_shift", "trend_shift", "autocorr_shift"]
        bt = break_types[(i - n_samples // 2) % len(break_types)]
        mag = rng.uniform(0.3, 2.0)
        s, has_break, mid = generate_break_series(500, bt, magnitude=mag, seed=i)
    series_list.append(s)
    labels.append(int(has_break))

labels = np.array(labels)
print(
    f"Generated {n_samples} series: {(labels == 0).sum()} no-break, {(labels == 1).sum()} with break"
)

# Visualize examples
fig, axes = plt.subplots(2, 4, figsize=(16, 6))

types = ["none", "mean_shift", "var_shift", "trend_shift"]
for i, bt in enumerate(types):
    s, _, mid = generate_break_series(500, bt, magnitude=1.0, seed=99)
    ax = axes[0, i]
    ax.plot(s, linewidth=0.5)
    ax.axvline(mid, color="red", linestyle="--", alpha=0.7)
    ax.set_title(bt.replace("_", " ").title())

    # Show pre/post distributions
    ax = axes[1, i]
    ax.hist(s[:mid], bins=30, alpha=0.5, density=True, label="Pre")
    ax.hist(s[mid:], bins=30, alpha=0.5, density=True, label="Post")
    ax.legend(fontsize=8)
    ax.set_title("Pre vs Post Distribution")

plt.suptitle("Structural Break Types", fontweight="bold")
plt.tight_layout()
plt.show()

def location_shift_features(series: np.ndarray, boundary: int) -> dict:
    """Compute location shift features comparing pre/post boundary.

    Based on Alphabot features f01-f02: Welch t-test on absolute values
    at multiple window sizes, with Fisher aggregation.
    """
    features = {}
    pre = series[:boundary]
    post = series[boundary:]

    # Welch t-test on raw values
    t_stat, p_val = stats.ttest_ind(pre, post, equal_var=False)
    features["welch_t_pval"] = p_val
    features["welch_t_stat"] = abs(t_stat)

    # Welch t-test on absolute values (captures scale + location shifts)
    for window in [50, 100, 250]:
        w = min(window, len(pre), len(post))
        if w < 10:
            features[f"local_t_abs_{window}"] = 0.5
            continue
        pre_w = np.abs(pre[-w:])
        post_w = np.abs(post[:w])
        _, p = stats.ttest_ind(pre_w, post_w, equal_var=False)
        features[f"local_t_abs_{window}"] = p

    # Fisher aggregation across window sizes (Alphabot f02)
    p_values = [features.get(f"local_t_abs_{w}", 0.5) for w in [50, 100]]
    p_values = [max(p, 1e-300) for p in p_values]  # Avoid log(0)
    fisher_stat = -2 * sum(np.log(p) for p in p_values)
    features["fisher_location"] = 1 - stats.chi2.cdf(fisher_stat, df=2 * len(p_values))

    # Mean difference (standardized)
    pooled_std = np.sqrt((pre.var() + post.var()) / 2)
    if pooled_std > 0:
        features["mean_diff_std"] = abs(pre.mean() - post.mean()) / pooled_std
    else:
        features["mean_diff_std"] = 0.0

    return features

def scale_shift_features(series: np.ndarray, boundary: int) -> dict:
    """Compute scale shift features: variance ratio, Levene, Fligner-Killeen."""
    features = {}
    pre = series[:boundary]
    post = series[boundary:]

    # F-test variance ratio
    var_ratio = post.var(ddof=1) / max(pre.var(ddof=1), 1e-10)
    features["var_ratio"] = var_ratio
    features["log_var_ratio"] = abs(np.log(max(var_ratio, 1e-10)))

    # F-test p-value (Alphabot f09)
    f_stat = max(var_ratio, 1.0 / max(var_ratio, 1e-10))
    df1 = len(post) - 1
    df2 = len(pre) - 1
    f_pval = 2 * (1 - stats.f.cdf(f_stat, df1, df2))
    features["f_test_pval"] = f_pval
    features["neg_log10_f_pval"] = -np.log10(max(f_pval, 1e-300))

    # Fligner-Killeen test (robust to non-normality, Alphabot f05)
    _, fk_pval = stats.fligner(pre, post)
    features["fligner_pval"] = fk_pval

    # Levene test
    _, levene_pval = stats.levene(pre, post)
    features["levene_pval"] = levene_pval

    return features

def distribution_shift_features(series: np.ndarray, boundary: int, n_bins: int = 50) -> dict:
    """Compute distribution shift features: KS, JSD, Hellinger, Wasserstein."""
    features = {}
    pre = series[:boundary]
    post = series[boundary:]

    # Kolmogorov-Smirnov test
    ks_stat, ks_pval = stats.ks_2samp(pre, post)
    features["ks_stat"] = ks_stat
    features["ks_pval"] = ks_pval

    # Jensen-Shannon divergence (Alphabot f08)
    # Histogram-based
    all_data = np.concatenate([pre, post])
    bins = np.linspace(all_data.min() - 0.1, all_data.max() + 0.1, n_bins + 1)
    pre_hist, _ = np.histogram(pre, bins=bins, density=True)
    post_hist, _ = np.histogram(post, bins=bins, density=True)
    # Add small epsilon to avoid zeros
    pre_hist = pre_hist + 1e-10
    post_hist = post_hist + 1e-10
    pre_hist = pre_hist / pre_hist.sum()
    post_hist = post_hist / post_hist.sum()
    features["jsd"] = jensenshannon(pre_hist, post_hist) ** 2  # Squared JSD

    # Hellinger distance (Alphabot f11)
    hellinger = np.sqrt(1 - np.sum(np.sqrt(pre_hist * post_hist)))
    features["hellinger"] = hellinger

    # Wasserstein distance (Earth mover's distance)
    features["wasserstein"] = stats.wasserstein_distance(pre, post)

    return features

def dependence_shift_features(series: np.ndarray, boundary: int) -> dict:
    """Compute dependence shift features: autocorrelation change, BDS proxy."""
    features = {}
    pre = series[:boundary]
    post = series[boundary:]

    # Lag-1 autocorrelation difference
    if len(pre) > 2 and len(post) > 2:
        ac_pre = np.corrcoef(pre[:-1], pre[1:])[0, 1]
        ac_post = np.corrcoef(post[:-1], post[1:])[0, 1]
        features["autocorr_diff"] = abs(ac_post - ac_pre)
        features["autocorr_pre"] = ac_pre
        features["autocorr_post"] = ac_post
    else:
        features["autocorr_diff"] = 0.0
        features["autocorr_pre"] = 0.0
        features["autocorr_post"] = 0.0

    # Variance of squared values (proxy for volatility clustering change)
    sq_pre = pre**2
    sq_post = post**2
    if len(sq_pre) > 2 and len(sq_post) > 2:
        ac_sq_pre = np.corrcoef(sq_pre[:-1], sq_pre[1:])[0, 1]
        ac_sq_post = np.corrcoef(sq_post[:-1], sq_post[1:])[0, 1]
        features["sq_autocorr_diff"] = abs(ac_sq_post - ac_sq_pre)
    else:
        features["sq_autocorr_diff"] = 0.0

    return features

def alignment_features(series: np.ndarray, boundary: int) -> dict:
    """Compute change point alignment features (Alphabot f03)."""
    features = {}
    n = len(series)

    # CUSUM-based change point
    cusum = np.cumsum(series - series.mean())
    cusum_cp = np.argmax(np.abs(cusum))
    features["cusum_dist_to_boundary"] = abs(cusum_cp - boundary) / n
    features["cusum_max"] = abs(cusum[cusum_cp]) / (series.std() * np.sqrt(n))

    # Ruptures Pelt change point — Pelt/optimization can occasionally fail on
    # short, near-constant series; fall back to "no detected break" (distance=1).
    try:
        algo = rpt.Pelt(model="l2", min_size=20).fit(series)
        bkps = algo.predict(pen=np.log(n) * series.var())
        bkps = [b for b in bkps if b < n]
        if bkps:
            closest = min(bkps, key=lambda b: abs(b - boundary))
            features["ruptures_dist_to_boundary"] = abs(closest - boundary) / n
        else:
            features["ruptures_dist_to_boundary"] = 1.0
    except (ValueError, RuntimeError):
        features["ruptures_dist_to_boundary"] = 1.0

    return features

def compute_all_break_features(series: np.ndarray, boundary: int) -> dict:
    """Compute all structural break features for a single series."""
    features = {}
    features.update(location_shift_features(series, boundary))
    features.update(scale_shift_features(series, boundary))
    features.update(distribution_shift_features(series, boundary))
    features.update(dependence_shift_features(series, boundary))
    features.update(alignment_features(series, boundary))
    return features


# Build feature matrix
feature_rows = []
for i, series in enumerate(series_list):
    mid = len(series) // 2
    feat = compute_all_break_features(series, mid)
    feature_rows.append(feat)

X = pd.DataFrame(feature_rows)
y = labels

print(f"Feature matrix: {X.shape}")
print(f"Features: {list(X.columns)}")
X.head()

# Handle any NaN/inf values
X = X.replace([np.inf, -np.inf], np.nan).fillna(0)

# Cross-validated AUC
clf = LGBMClassifier(n_estimators=100, max_depth=3, random_state=42, verbose=-1)
cv_scores = cross_val_score(clf, X, y, cv=5, scoring="roc_auc")

print("=== Structural Break Classification ===")
print(f"5-Fold CV AUC: {cv_scores.mean():.4f} (+/- {cv_scores.std():.4f})")

# Fit on full data for feature importance
clf.fit(X, y)
importances = pd.Series(clf.feature_importances_, index=X.columns).sort_values(ascending=False)

print("\nTop 10 Features by Importance:")
for feat, imp in importances.head(10).items():
    print(f"  {feat:<30} {imp:.4f}")

fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# Feature importance
ax = axes[0]
importances.head(15).plot.barh(ax=ax)
ax.set_title("Feature Importance (GBM)")
ax.set_xlabel("Importance")

# ROC curve via cross-validation predictions
y_prob = cross_val_predict(clf, X, y, cv=5, method="predict_proba")[:, 1]
fpr, tpr, _ = roc_curve(y, y_prob)
ax = axes[1]
ax.plot(fpr, tpr, linewidth=2, label=f"AUC = {roc_auc_score(y, y_prob):.3f}")
ax.plot([0, 1], [0, 1], "k--", linewidth=0.5)
ax.set_xlabel("False Positive Rate")
ax.set_ylabel("True Positive Rate")
ax.set_title("ROC Curve — Break Detection")
ax.legend()

plt.tight_layout()
plt.show()

family_map = {
    "location": [
        "welch_t_pval",
        "welch_t_stat",
        "local_t_abs_50",
        "local_t_abs_100",
        "local_t_abs_250",
        "fisher_location",
        "mean_diff_std",
    ],
    "scale": [
        "var_ratio",
        "log_var_ratio",
        "f_test_pval",
        "neg_log10_f_pval",
        "fligner_pval",
        "levene_pval",
    ],
    "distribution": ["ks_stat", "ks_pval", "jsd", "hellinger", "wasserstein"],
    "dependence": ["autocorr_diff", "autocorr_pre", "autocorr_post", "sq_autocorr_diff"],
    "alignment": ["cusum_dist_to_boundary", "cusum_max", "ruptures_dist_to_boundary"],
}

family_importance = {}
for family, feats in family_map.items():
    existing = [f for f in feats if f in importances.index]
    family_importance[family] = importances[existing].sum()

family_series = pd.Series(family_importance).sort_values(ascending=True)

fig, ax = plt.subplots(figsize=(8, 4))
family_series.plot.barh(ax=ax, color=["#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd"])
ax.set_title("Feature Family Contribution to Break Detection")
ax.set_xlabel("Total Importance")
plt.tight_layout()
plt.show()

def analyze_period(
    returns: pd.Series, center_date: str, window: int = 250, label: str = ""
) -> dict[str, Any] | None:
    """Compute break features around a specific date."""
    center = pd.Timestamp(center_date)
    mask = (returns.index >= center - pd.Timedelta(days=window * 2)) & (
        returns.index <= center + pd.Timedelta(days=window * 2)
    )
    period_returns = returns[mask].dropna().values

    if len(period_returns) < 100:
        return None

    mid = len(period_returns) // 2
    features = compute_all_break_features(period_returns, mid)
    features["period"] = label
    features["center_date"] = center_date
    features["n_obs"] = len(period_returns)
    return features


spy_returns = spy_pd["returns"].dropna()

periods = [
    ("2008-09-15", "2008 Financial Crisis"),
    ("2020-03-11", "COVID-19 Crash"),
    ("2011-08-05", "US Debt Downgrade"),
    ("2015-06-15", "Mid-2015 (Quiet)"),
    ("2017-06-15", "Mid-2017 (Quiet)"),
    ("2019-06-15", "Mid-2019 (Quiet)"),
]

period_results = []
for date, label in periods:
    result = analyze_period(spy_returns, date, label=label)
    if result is not None:
        period_results.append(result)

period_df = pd.DataFrame(period_results)

# Key metrics comparison
key_cols = [
    "period",
    "ks_stat",
    "jsd",
    "wasserstein",
    "mean_diff_std",
    "var_ratio",
    "fisher_location",
]
display(period_df[key_cols])

# Predict break probability for each period
period_X = period_df[[c for c in X.columns if c in period_df.columns]]
period_X = period_X.replace([np.inf, -np.inf], np.nan).fillna(0)
period_df["break_probability"] = clf.predict_proba(period_X)[:, 1]

display(
    period_df[["period", "break_probability"]].sort_values("break_probability", ascending=False)
)

fig, ax = plt.subplots(figsize=(10, 5))
colors = [
    "red" if "Crisis" in p or "COVID" in p or "Downgrade" in p else "blue"
    for p in period_df["period"]
]
ax.barh(period_df["period"], period_df["break_probability"], color=colors, alpha=0.7)
ax.axvline(0.5, color="black", linestyle="--", linewidth=0.5)
ax.set_xlabel("Break Probability")
ax.set_title("Structural Break Detection — SPY Around Key Dates")
plt.tight_layout()
plt.show()

# Load SPY for ml4t-engineer demonstration
spy_break = (
    etfs.filter(pl.col("symbol") == "SPY")
    .select(["timestamp", "close"])
    .sort("timestamp")
    .with_columns(returns=pl.col("close").pct_change())
    .drop_nulls()
)

# Coefficient of variation (std / |mean|) is only well defined on a
# strictly-positive series: raw daily returns have a near-zero mean, so the
# denominator collapses and CV explodes into the thousands - meaningless as a
# regime gate. Apply it to the *magnitude* of returns instead, where the mean
# is safely positive; CV(|returns|) then reads as the relative variability of
# move size, spiking when the volatility regime itself becomes unstable.
spy_break = spy_break.with_columns(
    cv=coefficient_of_variation(pl.col("returns").abs(), window=50),
    kl_div=rolling_kl_divergence("returns", window=100),
    w_dist=rolling_wasserstein("returns", window=100),
    drift=rolling_drift("returns", window=100),
)

print("=== ml4t-engineer: Break Detection Features ===")
for col in ["cv", "kl_div", "w_dist", "drift"]:
    vals = spy_break[col].drop_nulls()
    print(f"  {col:<10}: mean={vals.mean():.4f}, std={vals.std():.4f}")