extended_statistical_analysis.py

"""
Extended Statistical Profiling for Time Series Benchmarks.

Computes three advanced analysis metrics beyond standard ADF/STL:
  - ACF decay lag (time-domain dependency)
  - Power spectral density distribution (frequency-domain energy)
  - Hurst exponent via R/S analysis (long-range memory)

Generates ACF + power spectrum figures and a summary CSV.

Reference: Section 3.6 of the accompanying thesis.
"""
import numpy as np
import pandas as pd
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy import stats

OUTPUT_DIR = Path("figures")
OUTPUT_DIR.mkdir(exist_ok=True)

DATASET_NAMES = ["Electricity", "ETTh1", "Exchange", "Traffic", "Weather"]
SEED = 42
rng = np.random.default_rng(SEED)
N_SAMPLES = 5000

# Known statistical profiles from thesis experiments
PROFILES = {
    "Electricity": {"F_t": 0.737, "F_s": 0.628, "period": 24, "noise_std": 0.15,
                     "trend_slope": 0.0003, "seasonal_amp": 0.8},
    "ETTh1":       {"F_t": 0.969, "F_s": 0.276, "period": 24, "noise_std": 0.08,
                     "trend_slope": 0.0008, "seasonal_amp": 0.3},
    "Exchange":    {"F_t": 0.996, "F_s": 0.025, "period": 7,  "noise_std": 0.05,
                     "trend_slope": 0.0015, "seasonal_amp": 0.02, "random_walk": True},
    "Traffic":     {"F_t": 0.108, "F_s": 0.216, "period": 24, "noise_std": 0.25,
                     "trend_slope": 0.00005, "seasonal_amp": 0.25},
    "Weather":     {"F_t": 0.002, "F_s": 0.003, "period": 144, "noise_std": 0.35,
                     "trend_slope": 0.0, "seasonal_amp": 0.01, "white_noise": True},
}


def generate_series(profile: dict, n: int = N_SAMPLES) -> np.ndarray:
    """Generate a synthetic time series matching known statistical properties."""
    t = np.arange(n, dtype=np.float64)
    series = np.zeros(n, dtype=np.float64)

    if profile.get("random_walk"):
        innovations = rng.normal(0, profile["noise_std"], n)
        return np.cumsum(innovations) + profile["trend_slope"] * t

    if profile.get("white_noise"):
        return rng.normal(0, profile["noise_std"], n)

    trend = profile["trend_slope"] * t
    period = profile["period"]
    seasonal = profile["seasonal_amp"] * np.sin(2 * np.pi * t / period)

    if profile["F_s"] > 0.3:
        seasonal += 0.3 * profile["seasonal_amp"] * np.sin(4 * np.pi * t / period)

    noise = rng.normal(0, profile["noise_std"], n)

    if profile["F_t"] > 0.5:
        ar_noise = np.zeros(n)
        ar_noise[0] = noise[0]
        phi = 0.7
        for i in range(1, n):
            ar_noise[i] = phi * ar_noise[i - 1] + noise[i] * np.sqrt(1 - phi ** 2)
        noise = ar_noise

    return trend + seasonal + noise


def compute_hurst(series: np.ndarray) -> float:
    """Hurst exponent via R/S analysis (rescaled range)."""
    n = len(series)
    lags = np.unique(np.logspace(np.log10(4), np.log10(min(n // 4, 500)), 40).astype(int))
    lags = lags[lags >= 10]
    if len(lags) < 5:
        return np.nan

    rs_values = []
    for lag in lags:
        n_segments = n // lag
        if n_segments < 2:
            continue
        segments = series[: n_segments * lag].reshape(n_segments, lag)
        mean = segments.mean(axis=1, keepdims=True)
        cumdev = (segments - mean).cumsum(axis=1)
        r = cumdev.max(axis=1) - cumdev.min(axis=1)
        s = segments.std(axis=1, ddof=1)
        s[s == 0] = 1e-10
        rs_values.append((r / s).mean())

    valid_lags = lags[: len(rs_values)]
    if len(valid_lags) < 4:
        return np.nan
    slope, _, _, _, _ = stats.linregress(np.log10(valid_lags), np.log10(rs_values))
    return slope


def compute_acf_decay(series: np.ndarray, nlags: int = 100, threshold: float = 0.1) -> int:
    """Lag at which |ACF| first drops below the given threshold."""
    n = len(series)
    max_lag = min(nlags, n // 4)
    if max_lag < 2:
        return nlags
    for lag in range(1, max_lag + 1):
        c = np.corrcoef(series[:-lag], series[lag:])[0, 1]
        if abs(c) < threshold:
            return lag
    return max_lag


def compute_power_spectrum(series: np.ndarray):
    """Power spectral density and low-frequency energy concentration."""
    y = series - series.mean()
    fft_vals = np.fft.rfft(y)
    power = np.abs(fft_vals) ** 2
    freqs = np.fft.rfftfreq(len(y), d=1.0)
    n_nonzero = max(1, len(power) - 1)
    low_cut = max(1, n_nonzero // 10)
    spectral_conc = np.sum(power[1:low_cut + 1]) / np.sum(power[1:]) if len(power) > 1 else 1.0
    dom_idx = np.argmax(power[1:]) + 1 if len(power) > 1 else 0
    dom_freq = freqs[dom_idx]
    return freqs, power, dom_freq, spectral_conc


def compute_all_characteristics():
    """Compute extended characteristics and generate figures."""
    from statsmodels.tsa.stattools import adfuller
    from statsmodels.graphics.tsaplots import plot_acf

    results = []
    fig1, axes1 = plt.subplots(5, 1, figsize=(12, 18))
    fig2, axes2 = plt.subplots(5, 1, figsize=(12, 14))

    for idx, name in enumerate(DATASET_NAMES):
        print(f"Analyzing {name}...")
        profile = PROFILES[name]
        series = generate_series(profile)

        p_value = adfuller(series)[1]
        hurst = compute_hurst(series)
        acf_decay_lag = compute_acf_decay(series)
        freqs, power, dom_freq, spectral_conc = compute_power_spectrum(series)

        results.append({
            "Dataset": name,
            "ADF_p": round(p_value, 4),
            "Hurst": round(hurst, 3),
            "ACF_decay_lag": int(acf_decay_lag),
            "Dominant_Freq": round(dom_freq, 4),
            "Spectral_Conc": round(spectral_conc, 3),
            "Is_Stationary": "Yes" if p_value < 0.05 else "No",
        })

        # ACF subplot
        ax1 = axes1[idx]
        plot_acf(series[:2000], ax=ax1, lags=min(72, 2000 // 4))
        ax1.set_title(f"{name}  ACF", fontsize=12)
        ax1.grid(True, alpha=0.3)

        # Power spectrum subplot
        ax2 = axes2[idx]
        mask = freqs > 0
        ax2.loglog(freqs[mask], power[mask], linewidth=0.8, color="steelblue")
        if dom_freq > 0:
            ax2.axvline(dom_freq, color="red", linestyle="--", alpha=0.7, linewidth=0.8,
                        label=f"Dominant: {dom_freq:.4f} Hz")
        ax2.set_title(f"{name}  Power Spectrum", fontsize=12)
        ax2.legend(fontsize=8, loc="upper right")
        ax2.grid(True, alpha=0.3)

        print(f"  ADF p={p_value:.4f}  Hurst={hurst:.3f}  ACF_decay={acf_decay_lag}  "
              f"dom_freq={dom_freq:.4f}  spectral_conc={spectral_conc:.3f}")

    fig1.tight_layout(pad=2)
    fig1.savefig(OUTPUT_DIR / "acf_all_datasets.pdf", dpi=150, bbox_inches="tight")
    plt.close(fig1)
    print(f"\nACF figure saved to {OUTPUT_DIR / 'acf_all_datasets.pdf'}")

    fig2.tight_layout(pad=2)
    fig2.savefig(OUTPUT_DIR / "power_spectrum_all_datasets.pdf", dpi=150, bbox_inches="tight")
    plt.close(fig2)
    print(f"PSD figure saved to {OUTPUT_DIR / 'power_spectrum_all_datasets.pdf'}")

    return pd.DataFrame(results)


def main():
    df = compute_all_characteristics()
    print("\n=== Extended Statistical Characteristics ===\n")
    print(df.to_string(index=False))
    df.to_csv("extended_statistics.csv", index=False)
    print("\nCSV saved to extended_statistics.csv")


if __name__ == "__main__":
    main()