"""From-scratch financial-ML cross-validation tools (Lopez de Prado, "Advances
in Financial Machine Learning", ch. 7).

The problem: in finance a label at time i is usually built from a FORWARD window
-- e.g. the return (or triple-barrier outcome) over [i, t1[i]].  Labels of
nearby observations therefore overlap and are highly correlated, so an ordinary
(shuffled or contiguous) k-fold puts information about a test observation into
the training set through its overlapping neighbours.  That leakage inflates the
cross-validated score far above anything achievable live.

Two fixes:
  * PURGING - drop from the training set any observation whose label window
    [j, t1[j]] overlaps the test set's combined label window.
  * EMBARGO - additionally drop a small block of training observations
    immediately AFTER each test fold, to kill leakage through serial
    correlation that purging (which only removes overlaps) misses.

`purged_kfold` yields (train_idx, test_idx) pairs implementing both.
"""
import numpy as np


def purged_kfold(t1, n_splits=5, embargo_pct=0.01):
    """Purged, embargoed k-fold splits.

    Parameters
    ----------
    t1 : array, t1[i] = index at which observation i's label window ends
         (label spans [i, t1[i]]).  For an h-step-ahead label, t1[i] = i + h.
    n_splits : number of contiguous test folds (NOT shuffled -- time order kept).
    embargo_pct : fraction of the sample embargoed after each test fold.

    Yields (train_idx, test_idx) integer arrays.
    """
    t1 = np.asarray(t1)
    n = len(t1)
    idx = np.arange(n)
    embargo = int(n * embargo_pct)
    for test_idx in np.array_split(idx, n_splits):          # contiguous folds, in time
        t_start = test_idx[0]
        t_end = test_idx[-1]
        t_max = int(t1[test_idx].max())                     # furthest reach of any test label
        # purge: a train obs j leaks if its label window [j, t1[j]] overlaps
        # the test window [t_start, t_max]  <=>  j <= t_max AND t1[j] >= t_start
        leaks = (idx <= t_max) & (t1 >= t_start)
        keep = ~leaks
        # embargo: also drop the block just after the test fold
        if embargo > 0:
            keep[t_end + 1: t_end + 1 + embargo] = False
        train_idx = idx[keep]
        yield train_idx, test_idx


def cv_score(clf, X, y, t1, n_splits=5, embargo_pct=0.01, scorer=None):
    """Average a scorer over purged, embargoed folds. scorer(y_true, proba)->float."""
    from sklearn.metrics import roc_auc_score
    scorer = scorer or roc_auc_score
    out = []
    for tr, te in purged_kfold(t1, n_splits, embargo_pct):
        clf.fit(X[tr], y[tr])
        p = clf.predict_proba(X[te])[:, 1]
        out.append(scorer(y[te], p))
    return np.array(out)