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

For each rate (S&P 500, Bonds Baa, T-Notes, Inflation) separately, computes
autocorrelation in a sliding ~10-year window, normalized by the standard
deviation within each window. Results are plotted with matplotlib.
"""

import numpy as np
import matplotlib  # noqa: E402
matplotlib.use("Agg")  # non-interactive backend for script
import matplotlib.pyplot as plt  # noqa: E402
from owlplanner.rates import SP500, BondsBaa, TNotes, Inflation, FROM  # noqa: E402

# Window length in years (~10 years)
WINDOW_YEARS = 10
# Need at least 2 points per segment for corrcoef, so max lag = window - 2
MAX_LAG = WINDOW_YEARS - 2


def sliding_autocorr(series, window_len, max_lag):
    """
    Compute autocorrelation at lags 1..max_lag in each sliding window.
    Correlation is normalized by the sample std within each window (Pearson).
    Returns: (center_indices, lags, corr_matrix) where corr_matrix[i, j] is
    autocorrelation at window i, lag j (lag index 0 = lag 1).
    """
    x = np.asarray(series, dtype=float)
    n = len(x)
    n_windows = n - window_len + 1
    center_idx = np.arange(n_windows) + (window_len - 1) / 2.0
    lags = np.arange(1, max_lag + 1, dtype=int)
    corr = np.full((n_windows, len(lags)), np.nan)
    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 >= 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 center_idx, lags, corr


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

    fig, axes = plt.subplots(2, 2, figsize=(10, 8), sharex=True)
    axes = axes.ravel()
    colors = plt.cm.viridis(np.linspace(0.2, 0.8, min(MAX_LAG, 5)))

    for ax, (name, series) in zip(axes, series_dict.items()):
        center_idx, lags, corr = sliding_autocorr(
            series, WINDOW_YEARS, MAX_LAG
        )
        center_years = years[0] + center_idx
        # Plot a few lags (1, 2, 3, 5, 9) to avoid clutter
        plot_lags = [1, 2, 3, 5]
        if MAX_LAG >= 9:
            plot_lags.append(9)
        plot_lags = [lag for lag in plot_lags if lag <= MAX_LAG]
        for k, lag in enumerate(plot_lags):
            j = np.where(lags == lag)[0][0]
            ax.plot(
                center_years,
                corr[:, j],
                label=f"lag {lag}",
                color=colors[k % len(colors)],
                alpha=0.9,
            )
        ax.axhline(0, color="gray", linestyle="--", linewidth=0.8)
        ax.set_title(name)
        ax.set_ylabel("Autocorrelation")
        ax.legend(loc="best", fontsize=8)
        ax.set_ylim(-1.05, 1.05)
        ax.grid(True, alpha=0.3)

    axes[0].set_xlabel("")
    axes[1].set_xlabel("")
    axes[2].set_xlabel("Center year of 10-year window")
    axes[3].set_xlabel("Center year of 10-year window")
    fig.suptitle(
        "Time–time correlation of historical rates\n"
        "(sliding 10-year window, normalized by window std)",
        fontsize=11,
    )
    plt.tight_layout()
    outpath = "time_time_correlation.png"
    plt.savefig(outpath, dpi=150, bbox_inches="tight")
    print(f"Saved {outpath}")


if __name__ == "__main__":
    main()
