#!/usr/bin/env python3
"""
Averaged time-time (autocorrelation) of historical rates using owlplanner data.

For each rate (S&P 500, Bonds Baa, T-Notes, Inflation), computes autocorrelation
in a sliding ~10-year window, then averages over the full time horizon. Averaging
uses only valid window estimates; sample counts (number of windows contributing
per lag) are used for transparency and can down-weight lags with few samples.
Plot: lag (0 to 10) vs autocorrelation.
"""

import numpy as np
import matplotlib  # noqa: E402
matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
from owlplanner.rates import SP500, BondsBaa, TNotes, Inflation  # noqa: E402

# Window length in years (~10 years). Use 12 so lag 10 has at least 2 points.
WINDOW_YEARS = 12
MAX_LAG = 10  # lags 0, 1, ..., 10


def sliding_autocorr_all_lags(series, window_len, max_lag):
    """
    Compute autocorrelation at lags 0..max_lag in each sliding window.
    Returns: (lags, corr_matrix) where corr_matrix[i, j] is autocorrelation
    at window i, lag j. Lag 0 is always 1.0.
    """
    x = np.asarray(series, dtype=float)
    n = len(x)
    n_windows = n - window_len + 1
    lags = np.arange(0, max_lag + 1, dtype=int)
    corr = np.full((n_windows, len(lags)), np.nan)
    corr[:, 0] = 1.0  # lag 0
    for i in range(n_windows):
        w = x[i : i + window_len]
        if np.std(w) < 1e-12:
            continue
        for j, lag in enumerate(lags):
            if lag == 0:
                continue
            if lag >= window_len - 1:
                continue
            y1 = w[:-lag]
            y2 = w[lag:]
            if np.std(y1) < 1e-12 or np.std(y2) < 1e-12:
                continue
            c = np.corrcoef(y1, y2)[0, 1]
            if np.isfinite(c):
                corr[i, j] = c
    return lags, corr


def averaged_acf(corr, lags, min_count=1):
    """
    Average autocorrelation over the time horizon (over windows).
    Returns (mean_acf, count) per lag. Uses only valid (finite) samples;
    count is the number of windows that contributed to each lag.
    """
    mean_acf = np.nanmean(corr, axis=0)
    count = np.sum(np.isfinite(corr), axis=0)
    # Optionally mask lags with too few samples
    mean_acf = np.where(count >= min_count, mean_acf, np.nan)
    return mean_acf, count


def main():
    series_dict = {
        "S&P 500": np.asarray(SP500),
        "Bonds Baa": np.asarray(BondsBaa),
        "T-Notes": np.asarray(TNotes),
        "Inflation": np.asarray(Inflation),
    }

    fig, axes = plt.subplots(2, 2, figsize=(10, 8), sharex=True)
    axes = axes.ravel()
    colors = plt.cm.tab10(np.linspace(0, 0.8, 4))

    for idx, (ax, (name, series)) in enumerate(zip(axes, series_dict.items())):
        lags, corr = sliding_autocorr_all_lags(
            series, WINDOW_YEARS, MAX_LAG
        )
        mean_acf, count = averaged_acf(corr, lags)
        ax.plot(lags, mean_acf, "o-", color=colors[idx], markersize=6, linewidth=1.5)
        ax.axhline(0, color="gray", linestyle="--", linewidth=0.8)
        ax.set_xlabel("Lag (years)")
        ax.set_ylabel("Autocorrelation")
        n_windows = corr.shape[0]
        ax.set_title(f"{name}\n(n={n_windows} windows)")
        ax.set_xticks(lags)
        ax.set_ylim(-1.05, 1.05)
        ax.grid(True, alpha=0.3, axis="y")

    fig.suptitle(
        "Averaged time–time correlation of historical rates\n"
        f"(sliding {WINDOW_YEARS}-year window, averaged over time; normalized by window std)",
        fontsize=11,
    )
    plt.tight_layout()
    outpath = "time_time_correlation_averaged.png"
    plt.savefig(outpath, dpi=150, bbox_inches="tight")
    print(f"Saved {outpath}")


if __name__ == "__main__":
    main()
