Quantitative Investment: Simulations and Analysis

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from scipy import stats, optimize
from statsmodels.tsa.stattools import adfuller, acf
from statsmodels.tsa.vector_ar.vecm import coint_johansen

rng = np.random.default_rng(42)

plt.rcParams.update({
    'figure.dpi': 150,
    'axes.spines.top': False,
    'axes.spines.right': False,
    'axes.grid': True,
    'grid.alpha': 0.3,
})

def simulate_ou(theta, mu, sigma, T, dt, x0=None, seed=None):
    """Euler-Maruyama discretisation of the OU SDE."""
    rng_ = np.random.default_rng(seed)
    n = int(T / dt)
    x = np.zeros(n)
    x[0] = x0 if x0 is not None else mu
    eps = rng_.standard_normal(n - 1)
    for k in range(n - 1):
        x[k + 1] = x[k] + theta * (mu - x[k]) * dt + sigma * np.sqrt(dt) * eps[k]
    return x


def simulate_trending(alpha, sigma, n, seed=None):
    """AR(1) with alpha>0: persistent/trending increments."""
    rng_ = np.random.default_rng(seed)
    dx = np.zeros(n)
    eps = rng_.standard_normal(n)
    for k in range(1, n):
        dx[k] = alpha * dx[k - 1] + sigma * eps[k]
    return np.cumsum(dx)


T, dt = 5.0, 1 / 252  # 5 years, daily
n = int(T / dt)
t = np.linspace(0, T, n)
sigma = 0.01

x_rw  = np.cumsum(rng.standard_normal(n) * sigma)          # random walk
x_ou  = simulate_ou(theta=5.0, mu=0.0, sigma=sigma, T=T, dt=dt, x0=0.2, seed=0)
x_tr  = simulate_trending(alpha=0.6, sigma=sigma * 0.5, n=n, seed=1)

fig, axes = plt.subplots(1, 3, figsize=(14, 4))

for ax, x, label, color in zip(
    axes,
    [x_rw, x_ou, x_tr],
    ['Random walk ($H=0.5$)', 'Mean-reverting OU ($H<0.5$)', 'Trending AR(1) ($H>0.5$)'],
    ['steelblue', 'tomato', 'seagreen'],
):
    ax.plot(t, x, lw=0.7, color=color, alpha=0.9)
    if label.startswith('Mean'):
        ax.axhline(0, color='k', lw=0.8, ls='--', label='$\\mu$')
        ax.legend(fontsize=9)
    ax.set_title(label, fontsize=10)
    ax.set_xlabel('Time (years)')
    ax.set_ylabel('$x_t$')

fig.suptitle('Three regimes: path comparison', fontsize=12, y=1.01)
fig.tight_layout()
fig.savefig('../markdown/figures/qi_ou_paths.png', bbox_inches='tight')
plt.show()

# ACF and variance-scaling comparison
max_lag = 40
lags_arr = np.arange(1, max_lag + 1)

fig, axes = plt.subplots(1, 2, figsize=(12, 4))

# ACF of increments
ax = axes[0]
for x, label, color in [
    (x_rw, 'Random walk', 'steelblue'),
    (x_ou, 'Mean-reverting OU', 'tomato'),
    (x_tr, 'Trending AR(1)', 'seagreen'),
]:
    dx = np.diff(x)
    acf_vals = acf(dx, nlags=max_lag, fft=True)[1:]
    ax.plot(lags_arr, acf_vals, 'o-', ms=3, lw=1, label=label, color=color)

ax.axhline(0, color='k', lw=0.8)
ax.axhline(1.96 / np.sqrt(n), color='grey', lw=0.8, ls='--', alpha=0.6)
ax.axhline(-1.96 / np.sqrt(n), color='grey', lw=0.8, ls='--', alpha=0.6)
ax.set_xlabel('Lag (days)')
ax.set_ylabel('ACF of increments')
ax.set_title('Autocorrelation of daily returns')
ax.legend(fontsize=8)

# Variance scaling (log-log)
ax = axes[1]
taus = np.unique(np.round(np.logspace(0, 2.5, 30)).astype(int))

for x, label, color in [
    (x_rw, 'Random walk', 'steelblue'),
    (x_ou, 'Mean-reverting OU', 'tomato'),
    (x_tr, 'Trending AR(1)', 'seagreen'),
]:
    var_tau = []
    for tau in taus:
        if tau < len(x):
            diffs = x[tau:] - x[:-tau]
            var_tau.append(np.mean(diffs**2))
        else:
            var_tau.append(np.nan)
    var_tau = np.array(var_tau)
    valid = ~np.isnan(var_tau) & (var_tau > 0)
    ax.loglog(taus[valid], var_tau[valid], 'o-', ms=3, lw=1, label=label, color=color)

# Reference lines for H = 0.3, 0.5, 0.7
tau_ref = np.array([1, taus[-1]])
for H, ls in [(0.3, ':'), (0.5, '--'), (0.7, '-.')]: 
    ax.loglog(tau_ref, tau_ref**(2 * H) * 1e-4, ls=ls, color='grey', lw=0.8, alpha=0.7, label=f'$H={H}$')

ax.set_xlabel('Lag $\\tau$ (days)')
ax.set_ylabel('$\\langle |x_{t+\\tau}-x_t|^2 \\rangle$')
ax.set_title('Variance scaling (log-log): slope = $2H$')
ax.legend(fontsize=7, ncol=2)

fig.tight_layout()
fig.savefig('../markdown/figures/qi_ou_diagnostics.png', bbox_inches='tight')
plt.show()

def hurst_rs(x):
    """Estimate Hurst exponent via rescaled range (R/S) method."""
    n = len(x)
    lags = np.unique(np.round(np.logspace(1, np.log10(n // 4), 20)).astype(int))
    rs_vals = []
    for lag in lags:
        chunks = [x[i:i + lag] for i in range(0, n - lag, lag)]
        rs_chunk = []
        for chunk in chunks:
            mean_c = np.mean(chunk)
            dev = np.cumsum(chunk - mean_c)
            R = np.max(dev) - np.min(dev)
            S = np.std(chunk, ddof=1)
            if S > 0:
                rs_chunk.append(R / S)
        if rs_chunk:
            rs_vals.append(np.mean(rs_chunk))
        else:
            rs_vals.append(np.nan)
    rs_vals = np.array(rs_vals)
    valid = ~np.isnan(rs_vals)
    slope, *_ = np.polyfit(np.log(lags[valid]), np.log(rs_vals[valid]), 1)
    return slope


# ADF p-values and Hurst exponents for varying theta
thetas = np.logspace(-1, 1.5, 20)  # 0.1 to ~30
n_mc = 200
T, dt = 2.0, 1 / 252
n_samp = int(T / dt)

adf_reject_rate = []
hurst_means = []
hurst_stds = []

for theta in thetas:
    halflives_days = np.log(2) / theta * 252
    adf_rejects = 0
    hurst_vals = []
    for seed in range(n_mc):
        x = simulate_ou(theta=theta, mu=0.0, sigma=0.01, T=T, dt=dt, x0=0.0, seed=seed)
        # ADF test
        adf_result = adfuller(x, maxlag=5, regression='c', autolag='AIC')
        if adf_result[1] < 0.05:
            adf_rejects += 1
        # Hurst
        hurst_vals.append(hurst_rs(x))
    adf_reject_rate.append(adf_rejects / n_mc)
    hurst_means.append(np.mean(hurst_vals))
    hurst_stds.append(np.std(hurst_vals))

adf_reject_rate = np.array(adf_reject_rate)
hurst_means = np.array(hurst_means)
hurst_stds = np.array(hurst_stds)
halflives = np.log(2) / thetas * 252

fig, axes = plt.subplots(1, 2, figsize=(12, 4))

ax = axes[0]
ax.semilogx(halflives, adf_reject_rate, 'o-', color='tomato', ms=4, lw=1.5)
ax.axhline(0.05, color='grey', ls='--', lw=1, label='5% size')
ax.axhline(0.8, color='steelblue', ls=':', lw=1, label='80% power')
ax.set_xlabel('Half-life (trading days, log scale)')
ax.set_ylabel('ADF rejection rate at 5%')
ax.set_title('ADF test power vs mean-reversion speed')
ax.legend(fontsize=9)
ax.set_ylim([0, 1.05])

ax = axes[1]
ax.fill_between(halflives, hurst_means - hurst_stds, hurst_means + hurst_stds,
                alpha=0.2, color='seagreen')
ax.semilogx(halflives, hurst_means, 'o-', color='seagreen', ms=4, lw=1.5)
ax.axhline(0.5, color='grey', ls='--', lw=1, label='$H=0.5$ (RW)')
ax.set_xlabel('Half-life (trading days, log scale)')
ax.set_ylabel('Estimated Hurst exponent $\\hat{H}$')
ax.set_title('Hurst exponent vs mean-reversion speed')
ax.legend(fontsize=9)

fig.suptitle(f'OU process, $n={n_samp}$ obs, {n_mc} Monte Carlo paths', fontsize=11)
fig.tight_layout()
fig.savefig('../markdown/figures/qi_adf_hurst.png', bbox_inches='tight')
plt.show()

# Synthetic cointegrated pair
# Common factor: random walk.  Individual noise: OU residuals.
n_days = 504  # 2 years daily
sigma_common = 0.015
sigma_idio = 0.005
theta_spread = 10.0  # mean reversion of spread
beta_true = 2.0      # true cointegrating coefficient

common = np.cumsum(rng.standard_normal(n_days) * sigma_common)
spread = simulate_ou(theta=theta_spread, mu=0.0, sigma=sigma_idio,
                     T=n_days / 252, dt=1 / 252, x0=0.0, seed=7)

# Y = beta * X + spread  (both are log prices)
X = common + rng.standard_normal(n_days) * sigma_idio
Y = beta_true * common + spread + rng.standard_normal(n_days) * sigma_idio

# Recover hedge ratio via OLS
beta_ols, alpha_ols, *_ = stats.linregress(X, Y)
resid_ols = Y - (alpha_ols + beta_ols * X)

# Johansen test
data_johansen = np.column_stack([X, Y])
joh = coint_johansen(data_johansen, det_order=0, k_ar_diff=1)
# First eigenvector: cointegrating vector
coint_vec = joh.evec[:, 0]  # (w1, w2)
spread_johansen = data_johansen @ coint_vec
# Normalise so coefficient on Y is 1
spread_johansen = spread_johansen / coint_vec[1]

print(f'True beta: {beta_true:.3f}')
print(f'OLS beta:  {beta_ols:.3f}')
print(f'Johansen cointegrating vector: {coint_vec}')

adf_resid = adfuller(resid_ols, maxlag=5, regression='c', autolag='AIC')
print(f'\nADF test on OLS residual: stat={adf_resid[0]:.3f}, p={adf_resid[1]:.4f}')
print(f'Johansen trace stats vs 5% critical values:')
for i in range(2):
    print(f'  r<={i}: stat={joh.lr1[i]:.2f}, cv(5%)={joh.cvt[i, 1]:.2f}', 
          '  REJECT' if joh.lr1[i] > joh.cvt[i, 1] else '  fail to reject')

half_life_days = np.log(2) / theta_spread * 252
print(f'\nTrue half-life: {half_life_days:.1f} days')

def backtest_pairs(spread, lookback=60, z_entry=1.5, z_exit=0.3, z_stop=3.0):
    """Z-score pairs trading backtest on a spread series."""
    n = len(spread)
    position = np.zeros(n)  # +1 long spread, -1 short spread
    pos = 0
    pnl = np.zeros(n)
    z = np.zeros(n)

    for t in range(lookback, n):
        window = spread[t - lookback:t]
        mu_hat = np.mean(window)
        sig_hat = np.std(window, ddof=1)
        z[t] = (spread[t] - mu_hat) / sig_hat if sig_hat > 0 else 0.0

        # Entry logic
        if pos == 0:
            if z[t] > z_entry:
                pos = -1  # spread high -> short
            elif z[t] < -z_entry:
                pos = 1   # spread low -> long
        # Exit logic
        elif pos == 1:
            if z[t] > -z_exit or z[t] < -z_stop:
                pos = 0
        elif pos == -1:
            if z[t] < z_exit or z[t] > z_stop:
                pos = 0

        position[t] = pos
        if t > 0:
            pnl[t] = pos * (spread[t] - spread[t - 1])

    return position, np.cumsum(pnl), z


pos_ols, cum_pnl_ols, z_ols = backtest_pairs(resid_ols)

# Performance metrics
daily_pnl_ols = np.diff(np.concatenate([[0], cum_pnl_ols]))
sharpe = np.sqrt(252) * daily_pnl_ols[daily_pnl_ols != 0].mean() / daily_pnl_ols[daily_pnl_ols != 0].std()
mdd = np.min(cum_pnl_ols - np.maximum.accumulate(cum_pnl_ols))
print(f'Annualised Sharpe: {sharpe:.2f}')
print(f'Total P&L: {cum_pnl_ols[-1]:.4f}')
print(f'Max drawdown: {mdd:.4f}')

fig, axes = plt.subplots(3, 1, figsize=(12, 9), sharex=True)

t_arr = np.arange(n_days)

ax = axes[0]
ax.plot(t_arr, resid_ols, lw=0.8, color='steelblue', label='Spread (OLS residual)')
ax.axhline(0, color='k', lw=0.5)
ax2 = ax.twinx()
ax2.plot(t_arr, z_ols, lw=0.6, color='orange', alpha=0.7)
ax2.axhline(1.5, color='red', ls='--', lw=0.7, alpha=0.5)
ax2.axhline(-1.5, color='green', ls='--', lw=0.7, alpha=0.5)
ax2.axhline(0, color='grey', ls=':', lw=0.6)
ax2.set_ylabel('z-score', fontsize=9)
ax.set_ylabel('Spread')
ax.set_title('Spread and z-score')
ax.legend(fontsize=8, loc='upper left')

ax = axes[1]
ax.step(t_arr, pos_ols, where='post', lw=1, color='purple')
ax.set_ylabel('Position')
ax.set_title('Trading position (+1 long, $-1$ short, 0 flat)')
ax.set_yticks([-1, 0, 1])

ax = axes[2]
ax.plot(t_arr, cum_pnl_ols, lw=1.5, color='tomato')
ax.fill_between(t_arr, 0, cum_pnl_ols, alpha=0.15, color='tomato')
ax.set_ylabel('Cumulative P&L')
ax.set_title(f'Cumulative P&L  |  Sharpe={sharpe:.2f}, MDD={mdd:.4f}')
ax.set_xlabel('Trading days')

fig.suptitle('Pairs trading backtest on synthetic cointegrated data', fontsize=12)
fig.tight_layout()
fig.savefig('../markdown/figures/qi_pairs_backtest.png', bbox_inches='tight')
plt.show()

# Synthetic pair with slowly drifting hedge ratio
n_days = 756  # 3 years
sigma_common = 0.015
sigma_eps = 0.005   # observation noise
sigma_omega = 0.002 # state transition noise (beta drift)

# True beta drifts from 1.5 to 2.5 over the sample
beta_true = np.linspace(1.5, 2.5, n_days) + 0.1 * rng.standard_normal(n_days)
X2 = np.cumsum(rng.standard_normal(n_days) * sigma_common)
Y2 = beta_true * X2 + rng.standard_normal(n_days) * sigma_eps


def kalman_pairs(X, Y, sigma_omega, sigma_eps):
    """Kalman filter for time-varying beta in Y_t = beta_t * X_t + eps_t."""
    n = len(X)
    beta = np.zeros(n)
    P = np.zeros(n)  # posterior variance
    beta[0] = Y[0] / X[0] if X[0] != 0 else 1.0
    P[0] = 1.0
    spread = np.zeros(n)

    for t in range(1, n):
        # Predict
        beta_pred = beta[t - 1]
        P_pred = P[t - 1] + sigma_omega**2
        # Update
        innovation = Y[t] - beta_pred * X[t]
        S = X[t]**2 * P_pred + sigma_eps**2
        K = P_pred * X[t] / S
        beta[t] = beta_pred + K * innovation
        P[t] = (1 - K * X[t]) * P_pred
        spread[t] = innovation  # pre-update residual

    return beta, P, spread


beta_kf, P_kf, spread_kf = kalman_pairs(X2, Y2, sigma_omega=sigma_omega, sigma_eps=sigma_eps)

# Fixed OLS hedge (fit on first half, apply to full sample)
half = n_days // 2
beta_ols2, alpha_ols2, *_ = stats.linregress(X2[:half], Y2[:half])
spread_ols2 = Y2 - (alpha_ols2 + beta_ols2 * X2)

fig, axes = plt.subplots(3, 1, figsize=(12, 9), sharex=True)

ax = axes[0]
ax.plot(beta_true, lw=1, color='k', alpha=0.5, label='True $\\beta_t$')
ax.plot(beta_kf, lw=1.2, color='tomato', label='Kalman estimate $\\hat{\\beta}_t$')
ax.axhline(beta_ols2, color='steelblue', ls='--', lw=1, label=f'OLS (in-sample, $\\hat{{\\beta}}={beta_ols2:.2f}$)')
ax.set_ylabel('Hedge ratio $\\beta$')
ax.set_title('Hedge ratio: truth vs Kalman vs OLS')
ax.legend(fontsize=8)

ax = axes[1]
ax.plot(spread_ols2, lw=0.7, color='steelblue', alpha=0.7, label='Spread (OLS hedge)')
ax.plot(spread_kf, lw=0.7, color='tomato', alpha=0.8, label='Spread (Kalman hedge)')
ax.axhline(0, color='k', lw=0.5)
ax.set_ylabel('Spread')
ax.set_title('Spread comparison')
ax.legend(fontsize=8)

# Backtest both
_, cpnl_kf,  _ = backtest_pairs(spread_kf,  lookback=30)
_, cpnl_ols2, _ = backtest_pairs(spread_ols2, lookback=30)
sharpe_kf  = np.sqrt(252) * np.diff(np.concatenate([[0], cpnl_kf ])).mean() / (np.diff(np.concatenate([[0], cpnl_kf ])).std() + 1e-12)
sharpe_ols2 = np.sqrt(252) * np.diff(np.concatenate([[0], cpnl_ols2])).mean() / (np.diff(np.concatenate([[0], cpnl_ols2])).std() + 1e-12)

ax = axes[2]
ax.plot(cpnl_ols2, lw=1.5, color='steelblue', label=f'OLS hedge (SR={sharpe_ols2:.2f})')
ax.plot(cpnl_kf,  lw=1.5, color='tomato',    label=f'Kalman hedge (SR={sharpe_kf:.2f})')
ax.set_ylabel('Cumulative P&L')
ax.set_title('Cumulative P&L comparison (out-of-sample uses OLS fit from first half)')
ax.set_xlabel('Trading days')
ax.legend(fontsize=9)

fig.suptitle('Kalman filter vs fixed OLS hedge ratio (drifting pair)', fontsize=12)
fig.tight_layout()
fig.savefig('../markdown/figures/qi_kalman_pairs.png', bbox_inches='tight')
plt.show()

def ema(x, alpha):
    """Exponential moving average with decay rate alpha."""
    out = np.zeros(len(x))
    out[0] = x[0]
    for t in range(1, len(x)):
        out[t] = alpha * x[t] + (1 - alpha) * out[t - 1]
    return out


def macd_signal(r, alpha_fast, alpha_slow):
    return ema(r, alpha_fast) - ema(r, alpha_slow)


def backtest_tsmom(prices, lookback_days=252, vol_target=0.2):
    """Time-series momentum: sign of past L-day return, vol-scaled position."""
    returns = np.diff(np.log(prices))
    n = len(returns)
    pos = np.zeros(n)
    vol_est = np.zeros(n)
    for t in range(lookback_days, n):
        past_ret = returns[t - lookback_days:t].sum()
        vol = returns[t - 60:t].std() * np.sqrt(252) + 1e-8
        pos[t] = np.sign(past_ret) * vol_target / vol
        vol_est[t] = vol
    pnl = pos * returns
    return np.cumsum(pnl), pos, vol_est


# Generate a trending asset (AR(1) with alpha=0.5)
n_days = 1260  # 5 years
returns_trend = np.zeros(n_days)
eps = rng.standard_normal(n_days) * 0.01
for t in range(1, n_days):
    returns_trend[t] = 0.5 * returns_trend[t - 1] + eps[t]
returns_trend += 0.0002  # slight positive drift
prices_trend = 100 * np.exp(np.cumsum(returns_trend))

# MACD signals
alpha_fast, alpha_slow = 2 / 13, 2 / 63  # 12-day and 62-day EMA
macd = macd_signal(returns_trend, alpha_fast, alpha_slow)

# Backtest TSMOM
cum_pnl_mom, pos_mom, vol_mom = backtest_tsmom(prices_trend, lookback_days=252)

# Sharpe
daily_pnl_mom = np.diff(np.concatenate([[0], cum_pnl_mom]))
sr_mom = np.sqrt(252) * daily_pnl_mom[252:].mean() / (daily_pnl_mom[252:].std() + 1e-8)

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

t_arr = np.arange(n_days)

ax = axes[0, 0]
ax.plot(t_arr, prices_trend, lw=0.8, color='steelblue')
ax.set_ylabel('Price')
ax.set_title('Synthetic trending asset (AR(1) returns, $\\alpha=0.5$)')

ax = axes[0, 1]
ax.plot(t_arr, macd, lw=0.8, color='purple', label='MACD')
ax.axhline(0, color='k', lw=0.5)
ax.fill_between(t_arr, 0, macd, where=macd > 0, alpha=0.2, color='seagreen', label='Bullish')
ax.fill_between(t_arr, 0, macd, where=macd < 0, alpha=0.2, color='tomato', label='Bearish')
ax.set_ylabel('MACD')
ax.set_title('MACD signal (fast EMA 12d − slow EMA 62d)')
ax.legend(fontsize=8)

ax = axes[1, 0]
ax.step(t_arr, pos_mom, where='post', lw=0.8, color='darkorange')
ax.axhline(0, color='k', lw=0.4)
ax.set_ylabel('Position (vol-scaled)')
ax.set_title('TSMOM position (vol-targeted to 20%)')
ax.set_xlabel('Trading days')

ax = axes[1, 1]
ax.plot(t_arr, cum_pnl_mom, lw=1.5, color='tomato')
ax.fill_between(t_arr, 0, cum_pnl_mom, alpha=0.15, color='tomato')
ax.set_ylabel('Cumulative log-return')
ax.set_title(f'TSMOM P&L | Sharpe = {sr_mom:.2f}')
ax.set_xlabel('Trading days')

fig.suptitle('Time-series momentum: signals and backtest', fontsize=12)
fig.tight_layout()
fig.savefig('../markdown/figures/qi_momentum.png', bbox_inches='tight')
plt.show()

# Show connection: AR(1) autocorrelation vs expected MA-crossover profit
alphas = np.linspace(-0.5, 0.8, 30)
n_mc, n_samp = 300, 500

lag1_acf_mean = []
macd_profit_mean = []

for a in alphas:
    acf_vals = []
    profit_vals = []
    for seed in range(n_mc):
        rng_loc = np.random.default_rng(seed)
        r = np.zeros(n_samp)
        eps_loc = rng_loc.standard_normal(n_samp) * 0.01
        for t in range(1, n_samp):
            r[t] = a * r[t - 1] + eps_loc[t]
        acf_vals.append(acf(r, nlags=1, fft=True)[1])
        sig = macd_signal(r, 2 / 5, 2 / 20)
        # Simple: profit when MACD agrees with next-period return
        profit = np.mean(np.sign(sig[20:-1]) * r[21:])
        profit_vals.append(profit)
    lag1_acf_mean.append(np.mean(acf_vals))
    macd_profit_mean.append(np.mean(profit_vals))

fig, ax = plt.subplots(1, 1, figsize=(7, 4))
ax.scatter(lag1_acf_mean, macd_profit_mean, s=30, color='steelblue', alpha=0.8)
# OLS fit
m, b = np.polyfit(lag1_acf_mean, macd_profit_mean, 1)
x_fit = np.linspace(min(lag1_acf_mean), max(lag1_acf_mean), 50)
ax.plot(x_fit, m * x_fit + b, 'r--', lw=1.5, label=f'OLS fit (slope={m:.4f})')
ax.axhline(0, color='k', lw=0.5)
ax.axvline(0, color='k', lw=0.5)
ax.set_xlabel('Lag-1 ACF of returns')
ax.set_ylabel('Expected MACD signal profit')
ax.set_title('MACD profitability scales with return autocorrelation')
ax.legend(fontsize=9)
fig.tight_layout()
fig.savefig('../markdown/figures/qi_acf_profit.png', bbox_inches='tight')
plt.show()

# Synthetic asset universe: 5 assets
N = 5
mu = np.array([0.06, 0.10, 0.08, 0.12, 0.04])   # expected returns (annualised)
sigmas = np.array([0.12, 0.20, 0.15, 0.25, 0.08])  # volatilities

# Correlation matrix (positive definite)
rho = np.array([
    [1.00, 0.40, 0.50, 0.30, 0.10],
    [0.40, 1.00, 0.60, 0.70, 0.05],
    [0.50, 0.60, 1.00, 0.55, 0.15],
    [0.30, 0.70, 0.55, 1.00, 0.00],
    [0.10, 0.05, 0.15, 0.00, 1.00],
])
Sigma = np.outer(sigmas, sigmas) * rho

rf = 0.02  # risk-free rate


def portfolio_stats(w, mu, Sigma, rf):
    ret = w @ mu
    vol = np.sqrt(w @ Sigma @ w)
    sr  = (ret - rf) / vol
    return ret, vol, sr


# Efficient frontier via parametric sweep of lambda
lambdas = np.logspace(-2, 3, 500)
eff_rets, eff_vols = [], []

for lam in lambdas:
    # w* = (1/lambda) Sigma^{-1} (mu - c * 1), c from budget constraint
    Sigma_inv = np.linalg.inv(Sigma)
    ones = np.ones(N)
    # Budget constrained: w = Sigma^{-1}(mu - c*1) / lambda, sum(w)=1
    A = ones @ Sigma_inv @ ones
    B = ones @ Sigma_inv @ mu
    c = (A * (mu @ Sigma_inv @ ones) / lam - A) / (A * A / lam)
    # Direct constrained solve
    c_val = (B - lam) / A  # simplified for unconstrained
    w = Sigma_inv @ (mu - c_val * ones) / lam
    if np.any(np.isnan(w)):
        continue
    ret, vol, _ = portfolio_stats(w, mu, Sigma, rf)
    eff_rets.append(ret)
    eff_vols.append(vol)

eff_rets = np.array(eff_rets)
eff_vols = np.array(eff_vols)

# Minimum variance portfolio
Sigma_inv = np.linalg.inv(Sigma)
ones = np.ones(N)
w_mv = Sigma_inv @ ones / (ones @ Sigma_inv @ ones)
ret_mv, vol_mv, sr_mv = portfolio_stats(w_mv, mu, Sigma, rf)

# Maximum Sharpe (tangency) portfolio — long-short, no budget constraint on sign
w_tan_raw = Sigma_inv @ (mu - rf * ones)
w_tan = w_tan_raw / ones @ w_tan_raw
ret_tan, vol_tan, sr_tan = portfolio_stats(w_tan, mu, Sigma, rf)

# Equal weight
w_ew = ones / N
ret_ew, vol_ew, sr_ew = portfolio_stats(w_ew, mu, Sigma, rf)

# Monte Carlo random portfolios
n_mc_port = 3000
w_rand = rng.dirichlet(np.ones(N), size=n_mc_port)  # long-only random weights
mc_rets = w_rand @ mu
mc_vols = np.sqrt(np.einsum('ij,jk,ik->i', w_rand, Sigma, w_rand))
mc_sr   = (mc_rets - rf) / mc_vols

print(f'Min-variance portfolio:  ret={ret_mv:.3f}, vol={vol_mv:.3f}, SR={sr_mv:.2f}')
print(f'Tangency portfolio:      ret={ret_tan:.3f}, vol={vol_tan:.3f}, SR={sr_tan:.2f}')
print(f'Equal-weight portfolio:  ret={ret_ew:.3f}, vol={vol_ew:.3f}, SR={sr_ew:.2f}')
print(f'\nTangency weights: {np.round(w_tan, 3)}')
print(f'Min-var weights:  {np.round(w_mv, 3)}')

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

ax = axes[0]
sc = ax.scatter(mc_vols, mc_rets, c=mc_sr, cmap='RdYlGn', s=6, alpha=0.5, vmin=-0.2, vmax=1.5)
plt.colorbar(sc, ax=ax, label='Sharpe ratio')

# Efficient frontier (upper part: higher return for given vol)
idx_mv = np.argmin(eff_vols)
ax.plot(eff_vols[idx_mv:], eff_rets[idx_mv:], 'k-', lw=2, label='Efficient frontier')

ax.scatter(vol_mv,  ret_mv,  s=150, marker='*', color='steelblue', zorder=5, label=f'Min-var (SR={sr_mv:.2f})')
ax.scatter(vol_tan, ret_tan, s=150, marker='*', color='gold',      zorder=5, label=f'Tangency (SR={sr_tan:.2f})')
ax.scatter(vol_ew,  ret_ew,  s=100, marker='D', color='tomato',    zorder=5, label=f'Equal weight (SR={sr_ew:.2f})')

# Capital market line
vol_cml = np.linspace(0, 0.35, 100)
ret_cml = rf + sr_tan * vol_cml
ax.plot(vol_cml, ret_cml, 'k--', lw=1, alpha=0.6, label='Capital market line')

# Individual assets
asset_labels = ['Asset A', 'Asset B', 'Asset C', 'Asset D', 'Asset E']
for i, (s, m, lbl) in enumerate(zip(sigmas, mu, asset_labels)):
    ax.scatter(s, m, s=60, marker='o', color='grey', alpha=0.7, zorder=4)
    ax.annotate(lbl, (s, m), textcoords='offset points', xytext=(5, 3), fontsize=7, color='grey')

ax.set_xlabel('Volatility $\\sigma_p$')
ax.set_ylabel('Expected return $\\mu_p$')
ax.set_title('Efficient frontier and random portfolios')
ax.legend(fontsize=8, loc='upper left')
ax.set_xlim([0.05, 0.35])
ax.set_ylim([0.02, 0.16])

# Portfolio weights comparison
ax = axes[1]
x_pos = np.arange(N)
width = 0.28
asset_labels_short = ['A', 'B', 'C', 'D', 'E']
bars1 = ax.bar(x_pos - width, w_mv,  width, label='Min-variance',    color='steelblue', alpha=0.8)
bars2 = ax.bar(x_pos,         w_tan, width, label='Tangency (max SR)', color='gold',     alpha=0.8)
bars3 = ax.bar(x_pos + width, w_ew,  width, label='Equal weight',     color='tomato',   alpha=0.8)
ax.set_xticks(x_pos)
ax.set_xticklabels(asset_labels_short)
ax.set_xlabel('Asset')
ax.set_ylabel('Portfolio weight')
ax.set_title('Portfolio weights: three strategies')
ax.axhline(0, color='k', lw=0.5)
ax.legend(fontsize=8)

fig.suptitle('Markowitz mean-variance optimisation', fontsize=12)
fig.tight_layout()
fig.savefig('../markdown/figures/qi_frontier.png', bbox_inches='tight')
plt.show()

def risk_parity_weights(Sigma, tol=1e-8, max_iter=500):
    """Solve risk-parity weights via iterative proportional approach."""
    N = Sigma.shape[0]
    w = np.ones(N) / N
    for _ in range(max_iter):
        sigma_p = np.sqrt(w @ Sigma @ w)
        rc = w * (Sigma @ w) / sigma_p  # risk contributions
        target = sigma_p / N
        # Gradient-free update: scale weights by sqrt(target/rc)
        w_new = w * np.sqrt(target / (rc + 1e-12))
        w_new = w_new / w_new.sum()
        if np.max(np.abs(w_new - w)) < tol:
            break
        w = w_new
    return w_new


# Simulate T years of daily returns for N=5 assets
T_years = 10
n_total = int(T_years * 252)

# True parameters
mu_true = np.array([0.06, 0.10, 0.08, 0.12, 0.04]) / 252
Sigma_true = Sigma / 252

L = np.linalg.cholesky(Sigma_true)
ret_matrix = (rng.standard_normal((n_total, N)) @ L.T) + mu_true  # (T, N)

# Walk-forward backtest
lookback_days = 252
rebal_freq = 21  # monthly rebalancing

ports = {'EW': [], 'Risk Parity': [], 'Max SR': []}
days_bt = []

for t in range(lookback_days, n_total, rebal_freq):
    end = min(t + rebal_freq, n_total)
    # In-sample covariance
    R_is = ret_matrix[t - lookback_days:t]
    Sig_is = np.cov(R_is.T) * 252
    mu_is  = R_is.mean(axis=0) * 252

    # Portfolio weights
    w_ew_bt  = np.ones(N) / N
    w_rp_bt  = risk_parity_weights(Sig_is / 252)
    Sig_inv_ = np.linalg.pinv(Sig_is)
    w_tan_raw_ = Sig_inv_ @ (mu_is - rf * np.ones(N))
    w_tan_bt = w_tan_raw_ / np.ones(N) @ w_tan_raw_ if np.ones(N) @ w_tan_raw_ > 0 else w_ew_bt

    # Out-of-sample returns
    R_oos = ret_matrix[t:end]
    for name, w in [('EW', w_ew_bt), ('Risk Parity', w_rp_bt), ('Max SR', w_tan_bt)]:
        ports[name].extend((R_oos @ w).tolist())
    days_bt.extend(list(range(t, end)))

days_bt = np.array(days_bt)
port_rets = {name: np.array(v) for name, v in ports.items()}

# Annualised stats
print(f'{"Strategy":<14} {"Ann. Return":>12} {"Ann. Vol":>10} {"Sharpe":>8} {"Max DD":>10}')
print('-' * 58)
for name, r in port_rets.items():
    ret_ann = r.mean() * 252
    vol_ann = r.std() * np.sqrt(252)
    sr = (ret_ann - rf) / vol_ann
    cum = np.cumprod(1 + r)
    mdd = np.min(cum / np.maximum.accumulate(cum)) - 1
    print(f'{name:<14} {ret_ann:>12.3f} {vol_ann:>10.3f} {sr:>8.2f} {mdd:>10.3f}')

colors = {'EW': 'steelblue', 'Risk Parity': 'seagreen', 'Max SR': 'tomato'}

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

# Cumulative wealth
ax = axes[0, 0]
for name, r in port_rets.items():
    cum = np.cumprod(1 + r)
    ax.plot(cum, lw=1.5, color=colors[name], label=name)
ax.set_title('Cumulative wealth (out-of-sample)')
ax.set_xlabel('Trading days')
ax.set_ylabel('Wealth ($1 initial)')
ax.legend(fontsize=9)

# Rolling Sharpe (63-day window)
ax = axes[0, 1]
win = 63
for name, r in port_rets.items():
    roll_sr = np.array([
        (r[max(0, t-win):t].mean() * 252 - rf) / (r[max(0, t-win):t].std() * np.sqrt(252) + 1e-8)
        for t in range(win, len(r))
    ])
    ax.plot(np.arange(win, len(r)), roll_sr, lw=0.9, color=colors[name], label=name, alpha=0.8)
ax.axhline(0, color='k', lw=0.5)
ax.set_title('Rolling 63-day Sharpe ratio')
ax.set_xlabel('Trading days')
ax.set_ylabel('Sharpe ratio')
ax.legend(fontsize=9)

# Drawdown
ax = axes[1, 0]
for name, r in port_rets.items():
    cum = np.cumprod(1 + r)
    dd = cum / np.maximum.accumulate(cum) - 1
    ax.fill_between(np.arange(len(dd)), dd, 0, alpha=0.35, color=colors[name], label=name)
ax.set_title('Drawdown from peak')
ax.set_xlabel('Trading days')
ax.set_ylabel('Drawdown')
ax.legend(fontsize=9)

# Annual return distribution
ax = axes[1, 1]
annual_step = 252
for name, r in port_rets.items():
    ann_rets = [np.prod(1 + r[i:i+annual_step]) - 1 for i in range(0, len(r) - annual_step, annual_step)]
    ax.hist(ann_rets, bins=15, alpha=0.5, color=colors[name], label=name, density=True)
ax.set_title('Annual return distribution')
ax.set_xlabel('Annual return')
ax.set_ylabel('Density')
ax.legend(fontsize=9)

fig.suptitle('Portfolio strategy comparison: EW vs Risk Parity vs Max-SR (walk-forward)', fontsize=12)
fig.tight_layout()
fig.savefig('../markdown/figures/qi_portfolio_comparison.png', bbox_inches='tight')
plt.show()

# Risk contribution breakdown for each strategy
w_rp_final = risk_parity_weights(Sigma / 252)
w_tan_final_raw = np.linalg.inv(Sigma) @ (mu - rf)
w_tan_final = w_tan_final_raw / w_tan_final_raw.sum()
w_ew_final = np.ones(N) / N

def risk_contributions(w, Sigma):
    sigma_p = np.sqrt(w @ Sigma @ w)
    return w * (Sigma @ w) / sigma_p

rc_ew  = risk_contributions(w_ew_final, Sigma)
rc_rp  = risk_contributions(w_rp_final, Sigma)
rc_tan = risk_contributions(w_tan_final, Sigma)

fig, axes = plt.subplots(1, 3, figsize=(13, 4))
x_pos = np.arange(N)
labels = ['A', 'B', 'C', 'D', 'E']

for ax, rc, w, name, color in [
    (axes[0], rc_ew,  w_ew_final,  'Equal weight',    'steelblue'),
    (axes[1], rc_rp,  w_rp_final,  'Risk parity',     'seagreen'),
    (axes[2], rc_tan, w_tan_final, 'Max Sharpe',      'tomato'),
]:
    ax.bar(x_pos - 0.18, w,  0.35, color=color, alpha=0.5, label='Weight')
    ax.bar(x_pos + 0.18, rc, 0.35, color=color, alpha=0.9, label='Risk contrib.')
    ax.axhline(1 / N, color='k', ls='--', lw=0.8, alpha=0.5, label='1/N target')
    ax.set_xticks(x_pos)
    ax.set_xticklabels(labels)
    ax.set_title(name)
    ax.set_ylabel('Weight / Risk contribution')
    ax.legend(fontsize=7)

fig.suptitle('Portfolio weights vs risk contributions (annualised vol units)', fontsize=11)
fig.tight_layout()
fig.savefig('../markdown/figures/qi_risk_contributions.png', bbox_inches='tight')
plt.show()