Purged k-fold cross-validation with embargo, triple-barrier and meta-labels

"""M13 · Purged k-fold cross-validation with embargo, triple-barrier and meta-labels. Read notes §13.6.

Sample i has a label that depends on information in the closed interval
[t0[i], t1[i]]: t0 = when the feature is observed, t1 = when the label is
resolved (e.g. the end of a 10-day forward return). Samples are sorted by t0.

Run:  cd exercises && .venv/bin/python -m pytest modules/M13-forecasting/tests/test_m13_purged_cv.py
"""

from __future__ import annotations

import numpy as np


def purged_kfold_splits(t0, t1, n_splits: int = 5, embargo_pct: float = 0.0):
    """Return a list of (train_idx, test_idx) integer arrays (positions 0..N-1, sorted).

    1. Test folds: n_splits contiguous blocks of positions, np.array_split(arange(N), n_splits).
    2. Purge: for a test block B let T_start = t0[min B] and T_end = max_{j in B} t1[j].
       A sample i outside B may train only if t1[i] < T_start or t0[i] > T_end.
    3. Embargo: also drop the h = ceil(embargo_pct * N) positions max B + 1, ..., max B + h.
    Raise ValueError if any t1 < t0.
    """
    raise NotImplementedError


def naive_kfold_splits(N: int, n_splits: int = 5):
    """Ordinary unshuffled k-fold: train = complement of each contiguous test block. (Given.)"""
    pos = np.arange(N)
    return [(np.setdiff1d(pos, b), b) for b in np.array_split(pos, n_splits)]


def count_leaks(t0, t1, train_idx, test_idx) -> int:
    """Number of (train, test) pairs whose label intervals [t0, t1] intersect. (Given.)"""
    t0 = np.asarray(t0)
    t1 = np.asarray(t1)
    a0, a1 = t0[train_idx][:, None], t1[train_idx][:, None]
    b0, b1 = t0[test_idx][None, :], t1[test_idx][None, :]
    return int(np.sum((a0 <= b1) & (b0 <= a1)))


def triple_barrier_labels(close, events, pt: float, sl: float, vertical: int):
    """Triple-barrier labels (López de Prado 2018, ch. 3) with fixed log-return barriers.

    For each event position e: upper barrier ln(P_t / P_e) >= pt; lower barrier
    ln(P_t / P_e) <= -sl; vertical barrier at position min(e + vertical, len(close) - 1).
    Scan t = e+1, e+2, ...; the first barrier touched decides the label:
        +1 (upper), -1 (lower), 0 (vertical reached with neither touched).
    Return (labels, t1): int arrays, t1 = position where the label was resolved.
    """
    raise NotImplementedError


def meta_labels(primary_side, tb_labels) -> np.ndarray:
    """Meta-label = 1 iff primary_side (+1/-1) times the triple-barrier label is > 0, else 0."""
    raise NotImplementedError