# Core numerical computing
import numpy as np
import pandas as pd
import warnings
warnings.filterwarnings('ignore')

# Cointegration and time series analysis
import statsmodels.api as sm
from statsmodels.tsa.stattools import coint
from statsmodels.tsa.vector_ar.vecm import coint_johansen

# Market data download (only used for the SPY-IVV-VOO triplet)
import yfinance as yf

# Plotting
import matplotlib.pyplot as plt
import seaborn as sns
sns.set_theme(style="darkgrid")

# color palette (the same used in the R companion)
BLUE, RED, BLACK = "#2980b9", "#c0392b", "black"

def engle_granger_test(y1, y2, autolag="aic"):
    """
    Engle--Granger cointegration test.

    Estimates the static regression y1 = mu + gamma*y2 + eps by OLS and tests
    the residuals eps for a unit root.

    Parameters
    ----------
    y1 : array-like
        Length-T series (e.g., log-prices) of the first asset.
    y2 : array-like
        Length-T series of the second asset.
    autolag : str or None, default "aic"
        Lag selection of the ADF regression on the residuals.

    Returns
    -------
    dict
        With the entries ``mu`` and ``gamma`` (the OLS estimates), ``residuals``,
        ``statistic`` and ``pvalue`` of the test, and ``crit`` (critical values).
    """
    a = np.asarray(y1, dtype=float)
    b = np.asarray(y2, dtype=float)

    ols = sm.OLS(a, sm.add_constant(b)).fit()
    mu, gamma = ols.params

    stat, pvalue, crit = coint(a, b, trend="c", autolag=autolag)

    return {
        "mu": mu,
        "gamma": gamma,
        "residuals": a - mu - gamma * b,
        "statistic": stat,
        "pvalue": pvalue,
        "crit": dict(zip(["1%", "5%", "10%"], crit)),
    }


def report_engle_granger(res):
    """Print a compact summary of an Engle--Granger test."""
    print(f"Estimated relation: y1 = {res['gamma']:.4f} * y2 {res['mu']:+.4f}")
    print(f"Statistic on the residuals: {res['statistic']:.3f}   "
          f"p-value: {res['pvalue']:.4f}")
    print("Critical values: " + ",  ".join(f"{k}: {v:.3f}" for k, v in res["crit"].items()))
    if res["pvalue"] < 0.05:
        print("Conclusion: reject the null of no cointegration at the 5% level "
              "(the series appear to be cointegrated).")
    else:
        print("Conclusion: cannot reject the null of no cointegration at the 5% level "
              "(the series do not appear to be cointegrated).")

def bollinger_zscore(spread, lookback=20):
    """
    Rolling z-score of the spread (Bollinger Bands normalization).

    Parameters
    ----------
    spread : array-like
        Length-T spread series.
    lookback : int, default 20
        Length of the rolling window.

    Returns
    -------
    np.ndarray
        Length-T z-score (NaN for the first ``lookback - 1`` points).
    """
    s = pd.Series(np.asarray(spread, dtype=float))
    mavg = s.rolling(lookback).mean()
    sdev = s.rolling(lookback).std(ddof=0)   # population SD, as TTR::BBands
    return ((s - mavg) / sdev).to_numpy()


def generate_thresholded_signal(z_score, s0=1.0, start_signal_at=0):
    """
    Thresholded (all-in/all-out) sizing signal.

    The signal is +1 when the z-score drops below -s0 (buy the spread), -1 when
    it rises above +s0 (short-sell the spread), and the position is kept while
    the z-score has not reverted to zero (after which it is closed).

    Parameters
    ----------
    z_score : array-like
        Length-T z-score of the spread.
    s0 : float, default 1.0
        Entry threshold.
    start_signal_at : int, default 0
        Index from which the signal starts being generated.

    Returns
    -------
    np.ndarray
        Length-T signal with values in {-1, 0, +1}.
    """
    z = np.asarray(z_score, dtype=float)
    T = len(z)
    signal = np.zeros(T)

    for t in range(start_signal_at, T):
        if z[t] < -s0:                     # buy
            signal[t] = 1
        elif z[t] > s0:                    # short-sell
            signal[t] = -1
        elif t > start_signal_at:          # maintain or close the position
            if signal[t - 1] == 1 and z[t] < 0:
                signal[t] = signal[t - 1]
            elif signal[t - 1] == -1 and z[t] > 0:
                signal[t] = signal[t - 1]

    return signal


def generate_linear_signal(z_score, s0=1.0):
    """
    Linear (contrarian) sizing signal s_t = -[z_t^{score} / s0]_{-1}^{+1}.

    Parameters
    ----------
    z_score : array-like
        Length-T z-score of the spread.
    s0 : float, default 1.0
        Threshold at which the position is fully leveraged.

    Returns
    -------
    np.ndarray
        Length-T signal with values in [-1, +1].
    """
    return np.clip(-np.asarray(z_score, dtype=float) / s0, -1.0, 1.0)


def generate_bb_linear_signal(spread, s0=1.0, lookback=20, start_signal_at=None):
    """
    Linear sizing signal computed from the rolling (Bollinger) z-score.

    Parameters
    ----------
    spread : array-like
        Length-T spread series.
    s0 : float, default 1.0
        Threshold at which the position is fully leveraged.
    lookback : int, default 20
        Length of the rolling window of the z-score.
    start_signal_at : int or None, default None
        Index from which the signal starts being generated (defaults to
        ``lookback``).

    Returns
    -------
    dict
        With the entries ``signal`` and ``z_score``.
    """
    z = bollinger_zscore(spread, lookback)
    if start_signal_at is None:
        start_signal_at = lookback
    signal = generate_linear_signal(z, s0)
    signal[:start_signal_at] = 0
    signal[np.isnan(signal)] = 0
    return {"signal": signal, "z_score": z}


def generate_bb_thresholded_signal(spread, entry_zscore=1.0, exit_zscore=0.0,
                                   lookback=20, start_signal_at=None):
    """
    Thresholded sizing signal computed from the rolling (Bollinger) z-score.

    Parameters
    ----------
    spread : array-like
        Length-T spread series.
    entry_zscore : float, default 1.0
        Entry threshold.
    exit_zscore : float, default 0.0
        Exit threshold (position closed when the z-score reverts to it).
    lookback : int, default 20
        Length of the rolling window of the z-score.
    start_signal_at : int or None, default None
        Index from which the signal starts being generated (defaults to
        ``lookback``).

    Returns
    -------
    dict
        With the entries ``signal`` and ``z_score``.
    """
    z = bollinger_zscore(spread, lookback)
    T = len(z)
    signal = np.zeros(T)

    if start_signal_at is None:
        start_signal_at = lookback
    start_signal_at = max(start_signal_at, lookback)

    for t in range(start_signal_at, T):
        if z[t] < -entry_zscore:           # buy
            signal[t] = 1
        elif z[t] > entry_zscore:          # short-sell
            signal[t] = -1
        elif t > start_signal_at:          # maintain or close the position
            if signal[t - 1] == 1 and z[t] < -exit_zscore:
                signal[t] = signal[t - 1]
            elif signal[t - 1] == -1 and z[t] > exit_zscore:
                signal[t] = signal[t - 1]

    return {"signal": signal, "z_score": z}

def compute_cum_pnl_spread_trading(spread, signal, compounded=False):
    """
    Cumulative P&L of a signal traded on a spread with constant weights.

    Parameters
    ----------
    spread : array-like or pd.Series
        Length-T spread series.
    signal : array-like
        Length-T sizing signal.
    compounded : bool, default False
        If True, compound the returns; otherwise accumulate them.

    Returns
    -------
    pd.Series
        Length-T cumulative P&L (indexed as ``spread`` when it is a Series).
    """
    idx = spread.index if isinstance(spread, pd.Series) else None
    z = np.asarray(spread, dtype=float)
    s = np.asarray(signal, dtype=float)
    T = len(z)

    signal_delayed = np.concatenate(([0.0], s[:-1]))
    spread_ret = np.concatenate(([0.0], np.diff(z)))
    portf_ret = signal_delayed * spread_ret

    cum = np.cumprod(1 + portf_ret) - 1 if compounded else np.cumsum(portf_ret)
    return pd.Series(cum, index=idx)


def compute_cum_pnl_weights(y, w, signal, lag_w=False, compounded=False):
    """
    Cumulative P&L of a signal traded with time-varying weights.

    Parameters
    ----------
    y : array-like or pd.DataFrame
        T x N prices (log-prices) of the assets.
    w : array-like
        T x N weights (with leverage one) of the spread at each period.
    signal : array-like
        Length-T sizing signal.
    lag_w : bool, default False
        If True, the weights held during period t are those of period t-1.
    compounded : bool, default False
        If True, compound the returns; otherwise accumulate them.

    Returns
    -------
    pd.Series
        Length-T cumulative P&L (indexed as ``y`` when it has an index).
    """
    idx = y.index if isinstance(y, pd.DataFrame) else None
    Y = np.asarray(y, dtype=float)
    W = np.asarray(w, dtype=float)
    s = np.asarray(signal, dtype=float)
    T = Y.shape[0]

    if lag_w:
        W = np.vstack([W[:1], W[:-1]])
    signal_delayed = np.concatenate(([0.0], s[:-1]))
    assets_ret = np.vstack([np.zeros((1, Y.shape[1])), np.diff(Y, axis=0)])
    portf_ret = signal_delayed * np.sum(W * assets_ret, axis=1)

    cum = np.cumprod(1 + portf_ret) - 1 if compounded else np.cumsum(portf_ret)
    return pd.Series(cum, index=idx)

def fit_rolling_ls(y1, y2, lookback=2 * 252, every=50):
    """
    Rolling least-squares estimates of the hedge ratio and the intercept.

    Parameters
    ----------
    y1, y2 : pd.Series
        Log-prices of the two assets.
    lookback : int, default 2*252
        Length of the rolling window.
    every : int, default 50
        Number of periods between re-estimations.

    Returns
    -------
    pd.DataFrame
        With the columns ``gamma`` and ``mu``, indexed as ``y1``.
    """
    idx = y1.index if isinstance(y1, pd.Series) else None
    a = np.asarray(y1, dtype=float)
    b = np.asarray(y2, dtype=float)
    T = len(a)

    t_update = list(range(lookback - 1, T, every))
    if t_update[-1] != T - 1:
        t_update.append(T - 1)

    gamma = np.full(T, np.nan)
    mu = np.full(T, np.nan)
    for t in t_update:
        window = slice(t - lookback + 1, t + 1)
        ols = sm.OLS(a[window], sm.add_constant(b[window])).fit()
        mu[t], gamma[t] = ols.params

    # carry the last estimate forward and the first one backward
    est = pd.DataFrame({"mu": mu, "gamma": gamma}).ffill().bfill()
    est.index = idx
    return est


def kalman_hedge_ratio(y1, y2, T_trn, alpha=5e-5):
    """
    Kalman filter for a time-varying hedge ratio (random-walk states).

    State space model

        y1[t] = mu[t] + gamma[t]*y2[t] + eps[t],   eps[t] ~ N(0, var_eps)
        (mu, gamma)[t+1] = (mu, gamma)[t] + eta[t],

    with Q = diag(alpha*var_eps, alpha*var_eps/var_y2). The parameters are
    initialized with the least-squares estimates on the first ``T_trn`` points.

    Parameters
    ----------
    y1, y2 : pd.Series
        Log-prices of the two assets.
    T_trn : int
        Length of the training window used for the initialization.
    alpha : float, default 5e-5
        Forgetting factor of the random-walk states.

    Returns
    -------
    pd.DataFrame
        With the columns ``mu`` and ``gamma`` (the one-step-ahead predictions
        mu[t|t-1] and gamma[t|t-1]), indexed as ``y1``.
    """
    idx = y1.index if isinstance(y1, pd.Series) else None
    a = np.asarray(y1, dtype=float)
    b = np.asarray(y2, dtype=float)
    T = len(a)

    # initial least squares on the training window
    ac = a[:T_trn] - a[:T_trn].mean()
    bc = b[:T_trn] - b[:T_trn].mean()
    gamma_hat = float(bc @ ac) / float(bc @ bc)
    mu_hat = float(np.mean(a[:T_trn] - gamma_hat * b[:T_trn]))

    # Kalman parameters
    eps_hat = a[:T_trn] - mu_hat - gamma_hat * b[:T_trn]
    var_eps = float(np.var(eps_hat, ddof=1))
    var_y2 = float(np.var(b[:T_trn], ddof=1))

    x = np.array([mu_hat, gamma_hat])                                  # state a priori
    V = np.diag([var_eps / T_trn, var_eps / var_y2 / T_trn])           # covariance a priori
    Q = np.diag([alpha * var_eps, alpha * var_eps / var_y2])           # state noise
    R = var_eps                                                        # observation noise

    mu = np.empty(T)
    gamma = np.empty(T)
    for t in range(T):
        mu[t], gamma[t] = x                     # one-step-ahead prediction
        Z = np.array([1.0, b[t]])               # observation matrix
        F = Z @ V @ Z + R                       # prediction variance of y1[t]
        K = V @ Z / F                           # Kalman gain
        x = x + K * (a[t] - Z @ x)              # update with y1[t]
        V = V - np.outer(K, Z @ V) + Q          # transition (identity) and state noise

    est = pd.DataFrame({"mu": mu, "gamma": gamma})
    est.index = idx
    return est


def spread_and_weights(y1, y2, gamma, mu):
    """
    Normalized spread with leverage one and the corresponding weights.

    Computes z[t] = (y1[t] - gamma[t]*y2[t] - mu[t])/(1 + gamma[t]) and the
    weights w[t] = (1, -gamma[t])/(1 + gamma[t]), so that z[t] = w[t]^T y[t].

    Parameters
    ----------
    y1, y2 : array-like or pd.Series
        Log-prices of the two assets.
    gamma, mu : array-like
        Hedge ratio and intercept (possibly time-varying).

    Returns
    -------
    spread : pd.Series
        Length-T normalized spread.
    w : np.ndarray
        T x 2 weights of the spread.
    """
    idx = y1.index if isinstance(y1, pd.Series) else None
    a = np.asarray(y1, dtype=float)
    b = np.asarray(y2, dtype=float)
    T = len(a)
    g = np.full(T, float(gamma)) if np.ndim(gamma) == 0 else np.asarray(gamma, dtype=float)
    m = np.full(T, float(mu)) if np.ndim(mu) == 0 else np.asarray(mu, dtype=float)

    spread = pd.Series((a - g * b - m) / (1 + g), index=idx)
    w = np.column_stack([np.ones_like(g), -g]) / (1 + g)[:, None]
    return spread, w


def plot_pair_trading(spread, z_score, signal, cumret, entry_zscore=1.0,
                      T_trn=None, figsize=(12, 11)):
    """
    Four stacked panels: spread, z-score, signal, and cumulative P&L.

    Parameters
    ----------
    spread : array-like or pd.Series
        Spread series (its index is used as the x-axis when available).
    z_score : array-like
        z-score of the spread.
    signal : array-like
        Sizing signal.
    cumret : pd.Series
        Cumulative P&L.
    entry_zscore : float, default 1.0
        Entry threshold drawn in the z-score panel.
    T_trn : int or None, default None
        If given, marks the end of the training window in the spread panel.
    figsize : tuple, default (12, 11)
        Size of the figure.
    """
    x = spread.index if isinstance(spread, pd.Series) else np.arange(len(spread))
    spread = np.asarray(spread, dtype=float)

    _, axes = plt.subplots(4, 1, figsize=figsize, sharex=True)

    axes[0].plot(x, spread, color=BLUE, lw=0.8)
    axes[0].set_title("Spread")
    if T_trn is not None:
        axes[0].axvline(x[T_trn], color=BLACK, ls="--", lw=1)
        ytop = np.nanmax(spread)
        axes[0].text(x[T_trn], ytop, "training  ", ha="right", va="top", fontsize=9)
        axes[0].text(x[T_trn], ytop, "  out-of-sample", ha="left", va="top", fontsize=9)

    axes[1].plot(x, np.asarray(z_score), color=BLUE, lw=0.8)
    axes[1].axhline(0, color=BLACK, lw=1)
    axes[1].axhline(-entry_zscore, color=BLACK, ls="--", lw=1)
    axes[1].axhline(entry_zscore, color=BLACK, ls="--", lw=1)
    axes[1].set_title("Z-score")

    axes[2].plot(x, np.asarray(signal), color=BLACK, lw=0.8)
    axes[2].set_title("Signal")

    axes[3].plot(cumret.index, cumret.values, color=BLUE, lw=0.8)
    axes[3].set_title("Cumulative return")

    for ax in axes:
        ax.set_xlabel(None)
    plt.tight_layout()
    plt.show()


def plot_comparison(series_a, series_b, label_a, label_b, title, ylim=None,
                    linestyles=("solid", "solid"), figsize=(12, 5)):
    """
    Two time series in the same panel (rolling least squares versus Kalman).

    Parameters
    ----------
    series_a, series_b : pd.Series
        The two series to compare.
    label_a, label_b : str
        Legend labels.
    title : str
        Title of the panel.
    ylim : tuple or None, default None
        Limits of the vertical axis.
    linestyles : tuple of str, default ("solid", "solid")
        Line styles of the two series.
    figsize : tuple, default (12, 5)
        Size of the figure.
    """
    _, ax = plt.subplots(figsize=figsize)
    ax.plot(series_a.index, series_a.values, color=BLUE, lw=0.8,
            ls=linestyles[0], label=label_a)
    ax.plot(series_b.index, series_b.values, color=RED, lw=0.8,
            ls=linestyles[1], label=label_b)
    ax.set_title(title)
    ax.set_xlabel(None)
    if ylim is not None:
        ax.set_ylim(ylim)
    ax.legend(loc="best", frameon=True, fontsize=9)
    plt.tight_layout()
    plt.show()

# generate synthetic data
T = 200
rng = np.random.default_rng(421)
y = np.zeros(T)
for t in range(1, T):
    y[t] = y[t - 1] + rng.normal()

_, ax = plt.subplots(figsize=(12, 4))
ax.plot(np.arange(1, T + 1), y, color=BLUE, lw=1)
ax.set_title("Random walk")
ax.set_xlabel(None)
plt.tight_layout()
plt.show()

# generate synthetic data
T = 200
rng = np.random.default_rng(421)
y = np.zeros(T)
for t in range(1, T):
    y[t] = 0.2 * y[t - 1] + rng.normal()

_, ax = plt.subplots(figsize=(12, 4))
ax.plot(np.arange(1, T + 1), y, color=BLUE, lw=1)
ax.set_title("AR(1) sequence")
ax.set_xlabel(None)
plt.tight_layout()
plt.show()

# generate synthetic data
rng = np.random.default_rng(421)
T = 200
w = rng.normal(0, 0.1, size=T)
x = np.cumsum(w)
y1 = x + rng.normal(0, 0.2, size=T)
y2 = x + rng.normal(0, 0.2, size=T)
spread = y1 - y2
print(f"Correlation of the differences: {np.corrcoef(np.diff(y1), np.diff(y2))[0, 1]:.3f}")

_, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4), gridspec_kw={"width_ratios": [2, 1]})
ax1.plot(np.arange(1, T + 1), y1, label="y1")
ax1.plot(np.arange(1, T + 1), y2, label="y2")
ax1.plot(np.arange(1, T + 1), spread, label="spread")
ax1.set_title("Time series")
ax1.set_xlabel(None)
ax1.legend(frameon=True, fontsize=9)

ax2.plot(np.diff(y1), np.diff(y2), "o", ms=4, color=BLUE, alpha=0.7)
ax2.set_xlim(-0.8, 0.8)
ax2.set_ylim(-0.8, 0.8)
ax2.set_title("Scatter plot of differences")
plt.tight_layout()
plt.show()

# generate synthetic data
rng = np.random.default_rng(421)
T = 200
w = rng.normal(0, 0.3, size=T)
x = np.cumsum(w)
y1 = x + rng.normal(0, 0.05, size=T) + 0.01 * np.arange(T)
y2 = x + rng.normal(0, 0.05, size=T)
spread = y1 - y2
print(f"Correlation of the differences: {np.corrcoef(np.diff(y1), np.diff(y2))[0, 1]:.3f}")

_, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4), gridspec_kw={"width_ratios": [2, 1]})
ax1.plot(np.arange(1, T + 1), y1, label="y1")
ax1.plot(np.arange(1, T + 1), y2, label="y2")
ax1.plot(np.arange(1, T + 1), spread, label="spread")
ax1.set_title("Time series")
ax1.set_xlabel(None)
ax1.legend(frameon=True, fontsize=9)

ax2.plot(np.diff(y1), np.diff(y2), "o", ms=4, color=BLUE, alpha=0.7)
ax2.set_xlim(-0.8, 0.8)
ax2.set_ylim(-0.8, 0.8)
ax2.set_title("Scatter plot of differences")
plt.tight_layout()
plt.show()

# generate synthetic data
rng = np.random.default_rng(421)
T = 200
w = rng.normal(0, 0.1, size=T)
x = np.cumsum(w)
y1 = x + rng.normal(0, 0.2, size=T)
y2 = x + rng.normal(0, 0.2, size=T)

res = engle_granger_test(y1, y2)
report_engle_granger(res)

# plot residuals
_, ax = plt.subplots(figsize=(12, 4))
ax.plot(np.arange(1, T + 1), res["residuals"], color=BLUE, lw=1)
ax.set_title("Residual time series")
ax.set_xlabel(None)
plt.tight_layout()
plt.show()

# generate synthetic data
rng = np.random.default_rng(421)
T = 200
w = rng.normal(0, 0.3, size=T)
x = np.cumsum(w)
y1 = x + rng.normal(0, 0.05, size=T)
y2 = x + rng.normal(0, 0.05, size=T)
y1 = y1 + 0.01 * np.arange(T)

res = engle_granger_test(y1, y2)
report_engle_granger(res)

# plot residuals
_, ax = plt.subplots(figsize=(12, 4))
ax.plot(np.arange(1, T + 1), res["residuals"], color=BLUE, lw=1)
ax.set_title("Residual time series")
ax.set_xlabel(None)
plt.tight_layout()
plt.show()

# load the stored adjusted prices (see the note in the Helper functions section)
prices = pd.read_csv("prices-pairs-trading.csv", index_col="date", parse_dates=True)
# to download fresh data instead:
# prices = yf.download(["EWA", "EWC"], start="2013-01-01", end="2022-12-31",
#                      auto_adjust=True)["Close"]

pair = prices.loc["2016":"2019", ["EWA", "EWC"]]

res = engle_granger_test(pair["EWA"], pair["EWC"])
report_engle_granger(res)

# plot residuals
_, ax = plt.subplots(figsize=(12, 4))
ax.plot(pair.index, res["residuals"], color=BLUE, lw=0.8)
ax.set_title("Residual time series")
ax.set_xlabel(None)
plt.tight_layout()
plt.show()

pair = prices.loc["2017":"2019", ["KO", "PEP"]]

res = engle_granger_test(pair["KO"], pair["PEP"])
report_engle_granger(res)

# plot residuals
_, ax = plt.subplots(figsize=(12, 4))
ax.plot(pair.index, res["residuals"], color=BLUE, lw=0.8)
ax.set_title("Residual time series")
ax.set_xlabel(None)
plt.tight_layout()
plt.show()

triplet = yf.download(["SPY", "IVV", "VOO"], start="2017-01-01", end="2019-12-31",
                      auto_adjust=True, progress=False)["Close"]
triplet = triplet[["SPY", "IVV", "VOO"]].dropna()

# Johansen's trace test
joh = coint_johansen(triplet.values, det_order=0, k_ar_diff=1)
print(pd.DataFrame(
    {"trace statistic": joh.lr1, "95% critical value": joh.cvt[:, 1]},
    index=["r = 0", "r <= 1", "r <= 2"],
).round(3))

# the two cointegration vectors give two residuals
residual1 = pd.Series(triplet.values @ joh.evec[:, 0], index=triplet.index)
residual2 = pd.Series(triplet.values @ joh.evec[:, 1], index=triplet.index)

_, axes = plt.subplots(2, 1, figsize=(12, 7), sharex=True)
axes[0].plot(residual1.index, residual1.values, color=BLUE, lw=0.8)
axes[0].set_title("First residual time series")
axes[1].plot(residual2.index, residual2.values, color=BLUE, lw=0.8)
axes[1].set_title("Second residual time series")
for ax in axes:
    ax.set_xlabel(None)
plt.tight_layout()
plt.show()

# generate synthetic spread (log-prices) as an AR(1)
T = 200
rho = 0.7
z_mu = 0.02   # from real data KO
z_sd = 0.10   # from real data KO
phi = (1 - rho) * z_mu
sd_eps = np.sqrt(1 - rho**2) * z_sd

rng = np.random.default_rng(42)
z = np.empty(T)
z[0] = phi + rng.normal(0, sd_eps)
for t in range(1, T):
    z[t] = phi + rho * z[t - 1] + rng.normal(0, sd_eps)

# compute z-score
z_score = (z - z.mean()) / z.std()

# generate signal and compute P&L
s0 = 1
signal = generate_linear_signal(z_score, s0)
portf_cumret = compute_cum_pnl_spread_trading(z, signal)

plot_pair_trading(z, z_score, signal, portf_cumret, entry_zscore=s0, figsize=(12, 9))

# generate signal and compute P&L
s0 = 1
signal = generate_thresholded_signal(z_score, s0)
portf_cumret = compute_cum_pnl_spread_trading(z, signal)

plot_pair_trading(z, z_score, signal, portf_cumret, entry_zscore=s0, figsize=(12, 9))

y1 = np.log(prices["EWA"])
y2 = np.log(prices["EWC"])
T_trn = round(2 * 252)

# hedge ratio via fixed least squares on the training window
gamma = np.polyfit(y2.values[:T_trn], y1.values[:T_trn], 1)[0]

# spread with leverage one, centered with its training mean
mu = (y1 - gamma * y2).iloc[:T_trn].mean()
spread, w = spread_and_weights(y1, y2, gamma, mu)

# generate sizing signal via thresholded strategy
entry_zscore = 1
res = generate_bb_thresholded_signal(spread, entry_zscore=entry_zscore,
                                     lookback=6 * 21, start_signal_at=T_trn)
z_score, signal = res["z_score"], res["signal"]

# compute P&L
portf_cumret = compute_cum_pnl_spread_trading(spread, signal)

plot_pair_trading(spread, z_score, signal, portf_cumret,
                  entry_zscore=entry_zscore, T_trn=T_trn)

print(f"Final cumulative return: {portf_cumret.iloc[-1]:.2f}")

# hedge ratio via rolling least squares
rolling_ls = fit_rolling_ls(y1, y2, lookback=T_trn, every=1)

# spread with leverage one
spread, w = spread_and_weights(y1, y2, rolling_ls["gamma"], rolling_ls["mu"])

# generate sizing signal via thresholded strategy
res = generate_bb_thresholded_signal(spread, entry_zscore=entry_zscore,
                                     lookback=6 * 21, start_signal_at=T_trn)
z_score, signal = res["z_score"], res["signal"]

# compute P&L (with the weights held during each period)
portf_cumret = compute_cum_pnl_weights(pd.DataFrame({"y1": y1, "y2": y2}), w, signal, lag_w=True)

plot_pair_trading(spread, z_score, signal, portf_cumret,
                  entry_zscore=entry_zscore, T_trn=T_trn)

print(f"Final cumulative return: {portf_cumret.iloc[-1]:.2f}")

y1 = np.log(prices["KO"])
y2 = np.log(prices["PEP"])
T_trn = 2 * 252

# hedge ratio via rolling least squares
rolling_ls = fit_rolling_ls(y1, y2, lookback=T_trn, every=1)
gamma = rolling_ls["gamma"].to_numpy()
mu = rolling_ls["mu"].to_numpy()

# spread with leverage one
spread, w = spread_and_weights(y1, y2, gamma, mu)

# generate sizing signal via thresholded strategy
entry_zscore = 1
res = generate_bb_thresholded_signal(spread, entry_zscore=entry_zscore,
                                     lookback=6 * 21, start_signal_at=T_trn)
z_score, signal = res["z_score"], res["signal"]

# compute P&L (with the weights held during each period)
portf_cumret = compute_cum_pnl_weights(pd.DataFrame({"y1": y1, "y2": y2}), w, signal, lag_w=True)

plot_pair_trading(spread, z_score, signal, portf_cumret,
                  entry_zscore=entry_zscore, T_trn=T_trn)

print(f"Final cumulative return: {portf_cumret.iloc[-1]:.2f}")

# generate sizing signal via thresholded strategy
res = generate_bb_thresholded_signal(spread, entry_zscore=entry_zscore,
                                     lookback=1 * 21, start_signal_at=T_trn)
z_score, signal = res["z_score"], res["signal"]

# compute P&L (with the weights held during each period)
portf_cumret = compute_cum_pnl_weights(pd.DataFrame({"y1": y1, "y2": y2}), w, signal, lag_w=True)

plot_pair_trading(spread, z_score, signal, portf_cumret,
                  entry_zscore=entry_zscore, T_trn=T_trn)

print(f"Final cumulative return: {portf_cumret.iloc[-1]:.2f}")

y1 = np.log(prices["EWA"])
y2 = np.log(prices["EWC"])
T_trn = round(2 * 252)

# hedge ratio via rolling least squares
ewc_rls = fit_rolling_ls(y1, y2, lookback=T_trn, every=1)

# hedge ratio via Kalman filter
ewc_kal = kalman_hedge_ratio(y1, y2, T_trn=T_trn, alpha=5e-5)

plot_comparison(ewc_rls["gamma"], ewc_kal["gamma"], "Rolling LS", "Kalman",
                "Hedge ratio", ylim=(0.5, 1.2))

# compute spreads (with leverage one)
ewc_rls_spread, ewc_rls_w = spread_and_weights(y1, y2, ewc_rls["gamma"], ewc_rls["mu"])
ewc_kal_spread, ewc_kal_w = spread_and_weights(y1, y2, ewc_kal["gamma"], ewc_kal["mu"])

plot_comparison(ewc_rls_spread, ewc_kal_spread, "Rolling LS", "Kalman",
                "Spread", ylim=(-0.05, 0.05))

# generate sizing signal via thresholded strategy
entry_zscore = 1
res = generate_bb_thresholded_signal(ewc_rls_spread, entry_zscore=entry_zscore,
                                     lookback=6 * 21, start_signal_at=T_trn)
ewc_rls_signal = res["signal"]
res = generate_bb_thresholded_signal(ewc_kal_spread, entry_zscore=entry_zscore,
                                     lookback=6 * 21, start_signal_at=T_trn)
ewc_kal_signal = res["signal"]

# compute P&L (with the weights held during each period)
y_ewc = pd.DataFrame({"y1": y1, "y2": y2})
ewc_rls_cum = compute_cum_pnl_weights(y_ewc, ewc_rls_w, ewc_rls_signal, lag_w=True)
ewc_kal_cum = compute_cum_pnl_weights(y_ewc, ewc_kal_w, ewc_kal_signal, lag_w=False)

plot_comparison(ewc_rls_cum, ewc_kal_cum, "Rolling LS", "Kalman", "Cumulative return",
                linestyles=("solid", "dashed"))

print(f"Final cumulative return, rolling LS: {ewc_rls_cum.iloc[-1]:.2f}")
print(f"Final cumulative return, Kalman:     {ewc_kal_cum.iloc[-1]:.2f}")

y1 = np.log(prices["KO"])
y2 = np.log(prices["PEP"])
T_trn = round(2 * 252)

# hedge ratio via rolling least squares
pep_rls = fit_rolling_ls(y1, y2, lookback=T_trn, every=1)

# hedge ratio via Kalman filter
pep_kal = kalman_hedge_ratio(y1, y2, T_trn=T_trn, alpha=5e-5)

plot_comparison(pep_rls["gamma"], pep_kal["gamma"], "Rolling LS", "Kalman",
                "Hedge ratio", ylim=(0.2, 1.0))

# compute spreads (with leverage one)
pep_rls_spread, pep_rls_w = spread_and_weights(y1, y2, pep_rls["gamma"], pep_rls["mu"])
pep_kal_spread, pep_kal_w = spread_and_weights(y1, y2, pep_kal["gamma"], pep_kal["mu"])

plot_comparison(pep_rls_spread, pep_kal_spread, "Rolling LS", "Kalman",
                "Spread", ylim=(-0.06, 0.06))

# generate sizing signal via thresholded strategy
res = generate_bb_thresholded_signal(pep_rls_spread, entry_zscore=entry_zscore,
                                     lookback=6 * 21, start_signal_at=T_trn)
pep_rls_signal = res["signal"]
res = generate_bb_thresholded_signal(pep_kal_spread, entry_zscore=entry_zscore,
                                     lookback=6 * 21, start_signal_at=T_trn)
pep_kal_signal = res["signal"]

# compute P&L (with the weights held during each period)
y_pep = pd.DataFrame({"y1": y1, "y2": y2})
pep_rls_cum = compute_cum_pnl_weights(y_pep, pep_rls_w, pep_rls_signal, lag_w=True)
pep_kal_cum = compute_cum_pnl_weights(y_pep, pep_kal_w, pep_kal_signal, lag_w=False)

plot_comparison(pep_rls_cum, pep_kal_cum, "Rolling LS", "Kalman", "Cumulative return",
                linestyles=("solid", "dashed"))

print(f"Final cumulative return, rolling LS: {pep_rls_cum.iloc[-1]:.2f}")
print(f"Final cumulative return, Kalman:     {pep_kal_cum.iloc[-1]:.2f}")

def turnover_series(signal, w, lag_w=False):
    """
    Daily turnover of a pairs-trading strategy.

    Parameters
    ----------
    signal : array-like
        Length-T sizing signal.
    w : array-like
        T x N weights of the spread.
    lag_w : bool, default False
        If True, the weights held during period t are those of period t-1.

    Returns
    -------
    np.ndarray
        Length-T turnover, sum over the assets of the absolute changes of the
        positions held.
    """
    s = np.asarray(signal, dtype=float)
    W = np.asarray(w, dtype=float)
    if lag_w:
        W = np.vstack([W[:1], W[:-1]])
    positions = np.concatenate(([0.0], s[:-1]))[:, None] * W
    dpos = np.abs(np.diff(positions, axis=0)).sum(axis=1)
    return np.concatenate(([0.0], dpos))


def cumulative_pnl_net_returns(y, w, signal, cost_bps=0.0, lag_w=False):
    """
    Cumulative P&L net of transaction costs (in basis points of the traded amount).

    Parameters
    ----------
    y : array-like
        T x N prices (log-prices) of the assets.
    w : array-like
        T x N weights of the spread.
    signal : array-like
        Length-T sizing signal.
    cost_bps : float, default 0.0
        Transaction cost per unit of traded amount, in basis points.
    lag_w : bool, default False
        If True, the weights held during period t are those of period t-1.

    Returns
    -------
    pd.Series
        Length-T cumulative P&L net of costs.
    """
    cum = compute_cum_pnl_weights(y, w, signal, lag_w=lag_w)
    turnover = turnover_series(signal, w, lag_w=lag_w)
    return cum - cost_bps / 1e4 * np.cumsum(turnover)

# average over the out-of-sample period only (the first two years are the training window)
turnover = pd.Series({
    "EWA-EWC, rolling LS": turnover_series(ewc_rls_signal, ewc_rls_w, lag_w=True)[T_trn:].mean(),
    "EWA-EWC, Kalman": turnover_series(ewc_kal_signal, ewc_kal_w, lag_w=False)[T_trn:].mean(),
    "KO-PEP, rolling LS": turnover_series(pep_rls_signal, pep_rls_w, lag_w=True)[T_trn:].mean(),
    "KO-PEP, Kalman": turnover_series(pep_kal_signal, pep_kal_w, lag_w=False)[T_trn:].mean(),
}, name="average daily turnover (out-of-sample)")
turnover.round(3)

cost_bps = [0, 2, 5, 10]
final = pd.DataFrame({
    "EWA-EWC, rolling LS": [cumulative_pnl_net_returns(y_ewc, ewc_rls_w, ewc_rls_signal,
                                                       cost_bps=c, lag_w=True).iloc[-1]
                            for c in cost_bps],
    "EWA-EWC, Kalman": [cumulative_pnl_net_returns(y_ewc, ewc_kal_w, ewc_kal_signal,
                                                   cost_bps=c, lag_w=False).iloc[-1]
                        for c in cost_bps],
    "KO-PEP, rolling LS": [cumulative_pnl_net_returns(y_pep, pep_rls_w, pep_rls_signal,
                                                      cost_bps=c, lag_w=True).iloc[-1]
                           for c in cost_bps],
    "KO-PEP, Kalman": [cumulative_pnl_net_returns(y_pep, pep_kal_w, pep_kal_signal,
                                                  cost_bps=c, lag_w=False).iloc[-1]
                       for c in cost_bps],
}, index=[f"{c} bps" for c in cost_bps])
final.round(3)