Ornstein–Uhlenbeck estimation, half-life and Bertram trading bands

"""M13A · Ornstein–Uhlenbeck estimation, half-life and Bertram trading bands. Read notes §13A.7–13A.8 first.

Model: dX = kappa (theta - X) dt + sigma dW, kappa > 0.
Exact discretisation at step dt (M10 eq. 10.18, notes eq. 13A.23):
    X_{t+dt} = a + b X_t + e,  b = exp(-kappa dt),  a = theta (1 - b),
    Var(e) = sigma^2 (1 - b^2) / (2 kappa).

Provided from toolkit.ou (the M10 functions, same conventions as M10):
    simulate_ou(kappa, theta, sigma, dt, n, x0=None, seed=0)  n points, X_0 = x0 (theta if None)
    fit_ar1(x, ddof=2) -> (a, b, s)    OLS of x_{t+1} on [1, x_t], s^2 = SSR / (m - ddof)
    ou_from_ar1(a, b, s, dt) -> (kappa, theta, sigma)
    ou_half_life(kappa) = ln 2 / kappa  (M10 eq. 10.19)

Run:  cd exercises && .venv/bin/python -m pytest modules/M13A-cointegration-stat-arb/tests/test_m13a_ou.py
"""

from __future__ import annotations

import numpy as np
from scipy.optimize import brentq  # noqa: F401  (useful for the optimal band)
from scipy.special import erfcx, erfi  # noqa: F401
from scipy.integrate import quad  # noqa: F401

from toolkit.ou import fit_ar1, ou_from_ar1, ou_half_life, simulate_ou  # noqa: F401  (provided, M10)


def fit_ou(x: np.ndarray, dt: float):
    """Conditional (exact-discretisation) Gaussian MLE of (kappa, theta, sigma) (Prop. 13A.7, eq. 13A.24).

    Regress x_{t+1} on [1, x_t] by OLS to get a, b and the residual variance
    s2 = SSR / n (MLE divisor, n = number of transitions). Then
        kappa = -ln(b)/dt,  theta = a/(1-b),  sigma = sqrt(2 kappa s2 / (1 - b^2)).
    Requires 0 < b < 1; raise ValueError otherwise. Returns (kappa, theta, sigma).
    This is M10's fit_ar1 + ou_from_ar1 with the MLE divisor: which ddof is that?
    """
    raise NotImplementedError


def mean_passage_time(x0: float, level: float, kappa: float, theta: float, sigma: float) -> float:
    """E[first time X hits `level` | X_0 = x0] for the OU process (Prop. 13A.9, eq. 13A.27).

    With z = (x - theta) sqrt(kappa)/sigma:
      x0 < level:  (sqrt(pi)/kappa) * int_{z0}^{zl} erfcx(-z) dz
      x0 > level:  (sqrt(pi)/kappa) * int_{zl}^{z0} erfcx(+z) dz
    where erfcx(u) = exp(u^2) erfc(u) (scipy.special.erfcx). Use scipy.integrate.quad.
    Return 0.0 if x0 == level.
    """
    raise NotImplementedError


def bertram_cycle_time(a: float, m: float, kappa: float, sigma: float, theta: float = 0.0) -> float:
    """Expected length of one trade cycle a -> m -> a (eq. 13A.28):
        (pi/kappa) [erfi((m-theta) sqrt(kappa)/sigma) - erfi((a-theta) sqrt(kappa)/sigma)].
    """
    raise NotImplementedError


def bertram_rate(a: float, m: float, kappa: float, sigma: float, cost: float, theta: float = 0.0) -> float:
    """Expected profit per unit time of the cycle strategy: (m - a - cost) / cycle time (eq. 13A.26)."""
    raise NotImplementedError


def bertram_optimal_symmetric(kappa: float, sigma: float, cost: float, theta: float = 0.0):
    """Optimal symmetric band a = theta - d, m = theta + d maximising bertram_rate (Prop. 13A.11, eq. 13A.29).

    In u = d sqrt(kappa)/sigma and c_hat = cost sqrt(kappa)/sigma the first-order condition is
        (2u - c_hat) exp(u^2) = sqrt(pi) erfi(u),
    with a unique root u* > c_hat/2. Solve it with brentq (bracket: lower end max(c_hat/2, 1e-9);
    double the upper end until the left side exceeds the right). Returns (a*, m*, rate*).
    """
    raise NotImplementedError


def band_backtest(x: np.ndarray, a: float, m: float, cost: float):
    """Discretely monitored one-sided cycle strategy on a path x (notes §13A.8.4).

    Flat -> go long one unit of the spread at the first t with x_t <= a (fill at x_t);
    long -> exit at the first t (at or after entry) with x_t >= m (fill at x_t). Repeat from the
    exit index. A trade still open at the end is dropped.
    Returns dict with
        'pnl':     array of completed-trade P&Ls, x_exit - x_entry - cost
        'entries': int array of entry indices of the completed trades
        'exits':   int array of exit indices
    """
    raise NotImplementedError